Create server.py

server.py

@@ -0,0 +1,457 @@
+"""
+OpenAI-compatible Images API server for stabilityai/sd-turbo (or a
+compatible drop-in model) running on CPU inside a Hugging Face Docker Space.
+
+Endpoints:
+    GET  /health
+    GET  /v1/models
+    POST /v1/images/generations
+
+The diffusion pipeline is loaded exactly once at process startup and
+reused for every request. Because the underlying model is not
+thread-safe for concurrent inference, a global lock serializes all
+generation calls.
+
+Note on ONNX: SD Turbo does not have a maintained, broadly compatible
+ONNX Runtime export that works with AutoPipelineForText2Image out of
+the box (optimum's ONNX SD pipelines require a separately exported
+ONNX model directory with its own graph format, and no such export is
+published/guaranteed for this model). Since none is available here,
+this server falls back to the standard Diffusers PyTorch pipeline,
+as instructed.
+"""
+
+import base64
+import io
+import logging
+import os
+import threading
+import time
+import uuid
+
+from flask import Flask, jsonify, request
+from flask_cors import CORS
+
+# ---------------------------------------------------------------------------
+# Logging
+# ---------------------------------------------------------------------------
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
+logger = logging.getLogger("sd-turbo-server")
+
+# ---------------------------------------------------------------------------
+# Configuration (overridable via environment variables)
+# ---------------------------------------------------------------------------
+
+IMAGE_MODEL = os.environ.get("IMAGE_MODEL", "stabilityai/sd-turbo")
+
+# If the model was pre-downloaded at build time to a local directory,
+# prefer loading from there so no network access is required at runtime.
+MODEL_LOCAL_DIR = os.environ.get("MODEL_LOCAL_DIR", "/app/models/sd-turbo")
+
+DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", "2"))
+DEFAULT_GUIDANCE = float(os.environ.get("DEFAULT_GUIDANCE", "0"))
+
+# Model-facing short id used in the OpenAI-compatible API responses.
+MODEL_SHORT_ID = "sd-turbo"
+MODEL_OWNER = "stabilityai"
+
+SUPPORTED_SIZES = {
+    "256x256": (256, 256),
+    "512x512": (512, 512),
+    "768x768": (768, 768),
+}
+
+# ---------------------------------------------------------------------------
+# Flask app
+# ---------------------------------------------------------------------------
+
+app = Flask(__name__)
+CORS(app)
+
+# ---------------------------------------------------------------------------
+# Global pipeline state
+# ---------------------------------------------------------------------------
+
+_pipeline = None
+_pipeline_lock = threading.Lock()  # serializes generation calls
+_pipeline_load_error = None
+
+
+def _resolve_model_source():
+    """
+    Decide whether to load from the local pre-downloaded directory
+    (populated at Docker build time) or fall back to the hub id.
+    """
+    if os.path.isdir(MODEL_LOCAL_DIR) and os.listdir(MODEL_LOCAL_DIR):
+        logger.info("Using pre-downloaded model directory: %s", MODEL_LOCAL_DIR)
+        return MODEL_LOCAL_DIR
+    logger.info(
+        "Local model directory not found or empty (%s). "
+        "Falling back to downloading '%s' from the Hugging Face Hub.",
+        MODEL_LOCAL_DIR,
+        IMAGE_MODEL,
+    )
+    return IMAGE_MODEL
+
+
+def load_pipeline():
+    """
+    Load the diffusion pipeline exactly once at startup and cache it
+    globally. Uses AutoPipelineForText2Image with float32 weights,
+    since this deployment is CPU-only (no GPU / no fp16 support).
+    """
+    global _pipeline, _pipeline_load_error
+
+    import torch
+    from diffusers import AutoPipelineForText2Image
+
+    model_source = _resolve_model_source()
+
+    logger.info("Loading text-to-image pipeline from '%s' ...", model_source)
+    try:
+        pipeline = AutoPipelineForText2Image.from_pretrained(
+            model_source,
+            torch_dtype=torch.float32,
+            safety_checker=None,
+        )
+        pipeline.to("cpu")
+
+        # Disable the per-step progress bar (noisy in server logs / stdout).
+        pipeline.set_progress_bar_config(disable=True)
+
+        _pipeline = pipeline
+        logger.info("Pipeline loaded successfully.")
+    except Exception as exc:  # noqa: BLE001 - we want to capture and report any load failure
+        _pipeline_load_error = str(exc)
+        logger.exception("Failed to load pipeline: %s", exc)
+        raise
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def openai_error_response(message, err_type="invalid_request_error", param=None, code=None, status=400):
+    """
+    Build a Flask response matching the OpenAI API error envelope:
+        {"error": {"message": ..., "type": ..., "param": ..., "code": ...}}
+    """
+    body = {
+        "error": {
+            "message": message,
+            "type": err_type,
+            "param": param,
+            "code": code,
+        }
+    }
+    response = jsonify(body)
+    response.status_code = status
+    return response
+
+
+def parse_size(size_str):
+    """
+    Validate and convert an OpenAI-style size string (e.g. "512x512")
+    into a (width, height) tuple. Returns None if unsupported.
+    """
+    if size_str is None:
+        return SUPPORTED_SIZES["512x512"]
+    return SUPPORTED_SIZES.get(size_str)
+
+
+def image_to_b64(pil_image, image_format="PNG"):
+    """Encode a PIL image to a base64 string (no data URI prefix)."""
+    buffer = io.BytesIO()
+    pil_image.save(buffer, format=image_format)
+    raw_bytes = buffer.getvalue()
+    return base64.b64encode(raw_bytes).decode("utf-8")
+
+
+def run_generation(prompt, width, height, steps, guidance_scale, n_images):
+    """
+    Run the diffusion pipeline under the global lock so only one
+    generation happens at a time (safe for a single CPU worker process).
+    Returns a list of PIL.Image objects.
+    """
+    if _pipeline is None:
+        raise RuntimeError(
+            _pipeline_load_error or "Image generation pipeline is not initialized."
+        )
+
+    images = []
+    with _pipeline_lock:
+        for _ in range(n_images):
+            result = _pipeline(
+                prompt=prompt,
+                num_inference_steps=steps,
+                guidance_scale=guidance_scale,
+                width=width,
+                height=height,
+            )
+            images.append(result.images[0])
+    return images
+
+
+# ---------------------------------------------------------------------------
+# Routes
+# ---------------------------------------------------------------------------
+
+
+@app.route("/health", methods=["GET"])
+def health():
+    """Basic health/readiness check."""
+    status = "ok" if _pipeline is not None else "loading"
+    status_code = 200 if _pipeline is not None else 503
+    return jsonify(
+        {
+            "status": status,
+            "model": MODEL_SHORT_ID,
+            "error": _pipeline_load_error,
+        }
+    ), status_code
+
+
+@app.route("/v1/models", methods=["GET"])
+def list_models():
+    """OpenAI-compatible model listing endpoint."""
+    return jsonify(
+        {
+            "object": "list",
+            "data": [
+                {
+                    "id": MODEL_SHORT_ID,
+                    "object": "model",
+                    "owned_by": MODEL_OWNER,
+                }
+            ],
+        }
+    )
+
+
+@app.route("/v1/images/generations", methods=["POST"])
+def images_generations():
+    """OpenAI-compatible image generation endpoint."""
+
+    # ---- Parse JSON body ----
+    if not request.is_json:
+        return openai_error_response(
+            "Request body must be valid JSON with Content-Type: application/json.",
+            param=None,
+            code="invalid_json",
+            status=400,
+        )
+
+    try:
+        payload = request.get_json(silent=False)
+    except Exception:  # noqa: BLE001
+        return openai_error_response(
+            "Request body could not be parsed as JSON.",
+            code="invalid_json",
+            status=400,
+        )
+
+    if not isinstance(payload, dict):
+        return openai_error_response(
+            "Request body must be a JSON object.",
+            code="invalid_json",
+            status=400,
+        )
+
+    # ---- Validate prompt ----
+    prompt = payload.get("prompt")
+    if not prompt or not isinstance(prompt, str) or not prompt.strip():
+        return openai_error_response(
+            "You must provide a non-empty string in the 'prompt' field.",
+            param="prompt",
+            code="missing_prompt",
+            status=400,
+        )
+
+    # ---- Validate model (optional field, informational only) ----
+    requested_model = payload.get("model", MODEL_SHORT_ID)
+    if requested_model not in (MODEL_SHORT_ID, IMAGE_MODEL):
+        return openai_error_response(
+            f"The model '{requested_model}' does not exist or is not supported "
+            f"by this server. Available model: '{MODEL_SHORT_ID}'.",
+            param="model",
+            code="model_not_found",
+            status=404,
+        )
+
+    # ---- Validate size ----
+    size_str = payload.get("size", "512x512")
+    dimensions = parse_size(size_str)
+    if dimensions is None:
+        supported = ", ".join(sorted(SUPPORTED_SIZES.keys()))
+        return openai_error_response(
+            f"Unsupported size '{size_str}'. Supported sizes are: {supported}.",
+            param="size",
+            code="unsupported_size",
+            status=400,
+        )
+    width, height = dimensions
+
+    # ---- Validate n ----
+    n_images = payload.get("n", 1)
+    if not isinstance(n_images, int) or isinstance(n_images, bool) or n_images < 1:
+        return openai_error_response(
+            "'n' must be a positive integer.",
+            param="n",
+            code="invalid_n",
+            status=400,
+        )
+    if n_images > 4:
+        return openai_error_response(
+            "'n' must be less than or equal to 4 for this server.",
+            param="n",
+            code="invalid_n",
+            status=400,
+        )
+
+    # ---- Validate response_format ----
+    response_format = payload.get("response_format", "b64_json")
+    if response_format != "b64_json":
+        return openai_error_response(
+            "This server only supports response_format='b64_json'. "
+            "URL-based responses are not available on this CPU-only deployment.",
+            param="response_format",
+            code="unsupported_response_format",
+            status=400,
+        )
+
+    # ---- Optional generation overrides ----
+    steps = payload.get("num_inference_steps", DEFAULT_STEPS)
+    guidance_scale = payload.get("guidance_scale", DEFAULT_GUIDANCE)
+
+    try:
+        steps = int(steps)
+        guidance_scale = float(guidance_scale)
+    except (TypeError, ValueError):
+        return openai_error_response(
+            "'num_inference_steps' must be an integer and 'guidance_scale' must be a number.",
+            code="invalid_generation_params",
+            status=400,
+        )
+
+    if steps < 1 or steps > 50:
+        return openai_error_response(
+            "'num_inference_steps' must be between 1 and 50.",
+            param="num_inference_steps",
+            code="invalid_generation_params",
+            status=400,
+        )
+
+    if guidance_scale < 0 or guidance_scale > 20:
+        return openai_error_response(
+            "'guidance_scale' must be between 0 and 20.",
+            param="guidance_scale",
+            code="invalid_generation_params",
+            status=400,
+        )
+
+    # ---- Ensure pipeline is ready ----
+    if _pipeline is None:
+        return openai_error_response(
+            _pipeline_load_error
+            or "The image generation model is still loading. Please retry shortly.",
+            err_type="server_error",
+            code="model_not_ready",
+            status=503,
+        )
+
+    # ---- Run generation ----
+    request_id = uuid.uuid4().hex[:12]
+    logger.info(
+        "[%s] Generating %d image(s) | size=%dx%d steps=%d guidance=%.2f prompt=%r",
+        request_id,
+        n_images,
+        width,
+        height,
+        steps,
+        guidance_scale,
+        prompt[:200],
+    )
+
+    try:
+        images = run_generation(
+            prompt=prompt,
+            width=width,
+            height=height,
+            steps=steps,
+            guidance_scale=guidance_scale,
+            n_images=n_images,
+        )
+    except Exception as exc:  # noqa: BLE001
+        logger.exception("[%s] Generation failed: %s", request_id, exc)
+        return openai_error_response(
+            f"Image generation failed: {exc}",
+            err_type="server_error",
+            code="generation_failed",
+            status=500,
+        )
+
+    # ---- Build OpenAI-compatible response ----
+    data = [{"b64_json": image_to_b64(img)} for img in images]
+
+    return jsonify(
+        {
+            "created": int(time.time()),
+            "data": data,
+        }
+    )
+
+
+@app.errorhandler(404)
+def not_found(_error):
+    return openai_error_response(
+        "The requested endpoint does not exist.",
+        err_type="invalid_request_error",
+        code="not_found",
+        status=404,
+    )
+
+
+@app.errorhandler(405)
+def method_not_allowed(_error):
+    return openai_error_response(
+        "This HTTP method is not allowed for the requested endpoint.",
+        err_type="invalid_request_error",
+        code="method_not_allowed",
+        status=405,
+    )
+
+
+@app.errorhandler(500)
+def internal_error(_error):
+    return openai_error_response(
+        "An internal server error occurred.",
+        err_type="server_error",
+        code="internal_error",
+        status=500,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Entrypoint
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    # Load the model once, synchronously, before accepting traffic.
+    try:
+        load_pipeline()
+    except Exception:  # noqa: BLE001
+        # We still start the Flask app so /health reports the failure
+        # instead of the container silently dying and HF Spaces retrying
+        # forever without diagnostics.
+        logger.error(
+            "Starting server in degraded mode: pipeline failed to load. "
+            "/health will report the error."
+        )
+
+    port = int(os.environ.get("PORT", "7860"))
+    app.run(host="0.0.0.0", port=port, threaded=True)