Update server.py

server.py

@@ -1,26 +1,3 @@
-"""
-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
@@ -29,9 +6,18 @@ import threading
 import time
 import uuid
 
-from flask import Flask, jsonify, request
+import torch
+from flask import Flask, jsonify, request, send_file
 from flask_cors import CORS
 
+# ---------------------------------------------------------------------------
+# CPU Optimization - Limit threads to avoid thrashing on HF Spaces
+# ---------------------------------------------------------------------------
+os.environ["OMP_NUM_THREADS"] = "4"
+os.environ["MKL_NUM_THREADS"] = "4"
+torch.set_num_threads(4)
+torch.set_num_interop_threads(1)
+
 # ---------------------------------------------------------------------------
 # Logging
 # ---------------------------------------------------------------------------
@@ -43,19 +29,15 @@ logging.basicConfig(
 logger = logging.getLogger("sd-turbo-server")
 
 # ---------------------------------------------------------------------------
-# Configuration (overridable via environment variables)
+# Configuration
 # ---------------------------------------------------------------------------
 
 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"))
+DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", "4"))           # Increased for better quality
+DEFAULT_GUIDANCE = float(os.environ.get("DEFAULT_GUIDANCE", "1.0")) # Slightly increased
 
-# Model-facing short id used in the OpenAI-compatible API responses.
 MODEL_SHORT_ID = "sd-turbo"
 MODEL_OWNER = "stabilityai"
 
@@ -77,36 +59,21 @@ CORS(app)
 # ---------------------------------------------------------------------------
 
 _pipeline = None
-_pipeline_lock = threading.Lock()  # serializes generation calls
+_pipeline_lock = threading.Lock()
 _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,
-    )
+    logger.info("Local model directory not found. Falling back to hub.")
     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()
@@ -120,12 +87,19 @@ def load_pipeline():
         )
         pipeline.to("cpu")
 
-        # Disable the per-step progress bar (noisy in server logs / stdout).
-        pipeline.set_progress_bar_config(disable=True)
+        # Dynamic quantization for CPU speed (good quality/speed tradeoff)
+        pipeline.unet = torch.quantization.quantize_dynamic(
+            pipeline.unet, {torch.nn.Linear}, dtype=torch.qint8
+        )
+        if hasattr(pipeline, "text_encoder"):
+            pipeline.text_encoder = torch.quantization.quantize_dynamic(
+                pipeline.text_encoder, {torch.nn.Linear}, dtype=torch.qint8
+            )
 
+        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
+        logger.info("Pipeline loaded and quantized successfully.")
+    except Exception as exc:
         _pipeline_load_error = str(exc)
         logger.exception("Failed to load pipeline: %s", exc)
         raise
@@ -135,53 +109,28 @@ def load_pipeline():
 # 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,
-        }
-    }
+    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")
+    return base64.b64encode(buffer.getvalue()).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."
-        )
+        raise RuntimeError(_pipeline_load_error or "Pipeline not initialized.")
 
     images = []
     with _pipeline_lock:
@@ -201,130 +150,67 @@ def run_generation(prompt, width, height, steps, guidance_scale, n_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
+    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,
-                }
-            ],
-        }
-    )
+    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 ----
+    # (OpenAI-compatible endpoint - kept unchanged from original)
     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,
-        )
+        return openai_error_response("Request body must be valid JSON.", code="invalid_json")
 
     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,
-        )
+    except Exception:
+        return openai_error_response("Invalid JSON.", code="invalid_json")
 
     if not isinstance(payload, dict):
-        return openai_error_response(
-            "Request body must be a JSON object.",
-            code="invalid_json",
-            status=400,
-        )
+        return openai_error_response("Request body must be a JSON object.", code="invalid_json")
 
-    # ---- 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,
-        )
+        return openai_error_response("Missing 'prompt'.", param="prompt", code="missing_prompt")
 
-    # ---- 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,
-        )
+        return openai_error_response(f"Model '{requested_model}' not supported.", 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,
-        )
+        return openai_error_response(f"Unsupported size. Supported: {supported}.", param="size", code="unsupported_size")
+
     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,
-        )
+    if not isinstance(n_images, int) or n_images < 1 or n_images > 4:
+        return openai_error_response("'n' must be between 1 and 4.", param="n", code="invalid_n")
 
-    # ---- 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,
-        )
+        return openai_error_response("Only b64_json supported.", param="response_format", code="unsupported_response_format")
 
-    # ---- Optional generation overrides ----
     steps = payload.get("num_inference_steps", DEFAULT_STEPS)
     guidance_scale = payload.get("guidance_scale", DEFAULT_GUIDANCE)
 
@@ -332,50 +218,55 @@ def images_generations():
         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,
-        )
+        return openai_error_response("Invalid steps or guidance_scale.", code="invalid_generation_params")
 
-    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,
-        )
+    if steps < 1 or steps > 6:
+        steps = 4
+    if guidance_scale < 0 or guidance_scale > 5:
+        guidance_scale = 1.0
 
-    # ---- 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,
-        )
+        return openai_error_response("Model not ready.", 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],
-    )
+    logger.info("[%s] Generating %d image(s) | %dx%d steps=%d guidance=%.2f", request_id, n_images, width, height, steps, guidance_scale)
+
+    try:
+        images = run_generation(prompt, width, height, steps, guidance_scale, n_images)
+    except Exception as exc:
+        logger.exception("[%s] Generation failed: %s", request_id, exc)
+        return openai_error_response(f"Generation failed: {exc}", err_type="server_error", code="generation_failed", status=500)
+
+    data = [{"b64_json": image_to_b64(img)} for img in images]
+    return jsonify({"created": int(time.time()), "data": data})
+
+
+@app.route("/image", methods=["GET"])
+def generate_image():
+    """
+    New direct image endpoint.
+    Example: /image?prompt=a%20cute%20cat&width=512&height=512&steps=4&guidance=1.0
+    Returns PNG directly (works in <img src=""> tags).
+    """
+    prompt = request.args.get("prompt")
+    if not prompt or not prompt.strip():
+        return "Missing 'prompt' query parameter", 400
+
+    try:
+        width = int(request.args.get("width", 512))
+        height = int(request.args.get("height", 512))
+        steps = int(request.args.get("steps", DEFAULT_STEPS))
+        guidance = float(request.args.get("guidance", DEFAULT_GUIDANCE))
+    except ValueError:
+        return "Invalid width/height/steps/guidance parameters", 400
+
+    # Safety caps for CPU feasibility
+    if width > 768 or height > 768:
+        width, height = 512, 512
+    if steps > 6:
+        steps = 6
+    if guidance < 0 or guidance > 5:
+        guidance = 1.0
 
     try:
         images = run_generation(
@@ -383,57 +274,34 @@ def images_generations():
             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,
+            guidance_scale=guidance,
+            n_images=1,
         )
+        pil_image = images[0]
 
-    # ---- Build OpenAI-compatible response ----
-    data = [{"b64_json": image_to_b64(img)} for img in images]
+        img_io = io.BytesIO()
+        pil_image.save(img_io, format="PNG")
+        img_io.seek(0)
 
-    return jsonify(
-        {
-            "created": int(time.time()),
-            "data": data,
-        }
-    )
+        return send_file(img_io, mimetype="image/png", as_attachment=False)
+    except Exception as exc:
+        logger.exception("GET /image failed")
+        return f"Image generation failed: {exc}", 500
 
 
 @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,
-    )
+    return openai_error_response("Endpoint not found.", 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,
-    )
+    return openai_error_response("Method not allowed.", 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,
-    )
+    return openai_error_response("Internal server error.", err_type="server_error", code="internal_error", status=500)
 
 
 # ---------------------------------------------------------------------------
@@ -441,17 +309,10 @@ def internal_error(_error):
 # ---------------------------------------------------------------------------
 
 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."
-        )
+    except Exception:
+        logger.error("Starting in degraded mode - pipeline failed to load.")
 
     port = int(os.environ.get("PORT", "7860"))
     app.run(host="0.0.0.0", port=port, threaded=True)