Update server.py

server.py

@@ -1,93 +1,46 @@
-import base64
 import io
 import logging
 import os
 import threading
-import time
-import uuid
 
 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
-# ---------------------------------------------------------------------------
+# CPU limits
 os.environ["OMP_NUM_THREADS"] = "4"
 os.environ["MKL_NUM_THREADS"] = "4"
 torch.set_num_threads(4)
 torch.set_num_interop_threads(1)
 
-# ---------------------------------------------------------------------------
-# Logging
-# ---------------------------------------------------------------------------
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
 logger = logging.getLogger("sd-turbo-server")
 
-# ---------------------------------------------------------------------------
-# Configuration
-# ---------------------------------------------------------------------------
-
-IMAGE_MODEL = os.environ.get("IMAGE_MODEL", "stabilityai/sd-turbo")
-MODEL_LOCAL_DIR = os.environ.get("MODEL_LOCAL_DIR", "/app/models/sd-turbo")
-
-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_SHORT_ID = "sd-turbo"
-MODEL_OWNER = "stabilityai"
-
-SUPPORTED_SIZES = {
-    "256x256": (256, 256),
-    "512x512": (512, 512),
-    "768x768": (768, 768),
-}
-
-# ---------------------------------------------------------------------------
-# Flask app
-# ---------------------------------------------------------------------------
+# Config
+DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", "4"))
+DEFAULT_GUIDANCE = float(os.environ.get("DEFAULT_GUIDANCE", "1.0"))
+DEFAULT_WIDTH = 768
+DEFAULT_HEIGHT = 768
 
 app = Flask(__name__)
 CORS(app)
 
-# ---------------------------------------------------------------------------
-# Global pipeline state
-# ---------------------------------------------------------------------------
-
 _pipeline = None
 _pipeline_lock = threading.Lock()
 _pipeline_load_error = None
 
 
-def _resolve_model_source():
-    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. Falling back to hub.")
-    return IMAGE_MODEL
-
-
 def load_pipeline():
     global _pipeline, _pipeline_load_error
-
     from diffusers import AutoPipelineForText2Image
 
-    model_source = _resolve_model_source()
-
-    logger.info("Loading text-to-image pipeline from '%s' ...", model_source)
+    model_dir = os.environ.get("MODEL_LOCAL_DIR", "/app/models/sd-turbo")
+    logger.info("Loading pipeline from %s", model_dir)
     try:
         pipeline = AutoPipelineForText2Image.from_pretrained(
-            model_source,
-            torch_dtype=torch.float32,
-            safety_checker=None,
+            model_dir, torch_dtype=torch.float32, safety_checker=None
         )
         pipeline.to("cpu")
-
-        # Dynamic quantization for CPU speed (good quality/speed tradeoff)
         pipeline.unet = torch.quantization.quantize_dynamic(
             pipeline.unet, {torch.nn.Linear}, dtype=torch.qint8
         )
@@ -95,224 +48,62 @@ def load_pipeline():
             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 and quantized successfully.")
-    except Exception as exc:
-        _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):
-    body = {"error": {"message": message, "type": err_type, "param": param, "code": code}}
-    response = jsonify(body)
-    response.status_code = status
-    return response
-
+    except Exception as e:
+        _pipeline_load_error = str(e)
+        logger.error("Pipeline load failed: %s", e)
 
-def parse_size(size_str):
-    if size_str is None:
-        return SUPPORTED_SIZES["512x512"]
-    return SUPPORTED_SIZES.get(size_str)
 
-
-def image_to_b64(pil_image, image_format="PNG"):
-    buffer = io.BytesIO()
-    pil_image.save(buffer, format=image_format)
-    return base64.b64encode(buffer.getvalue()).decode("utf-8")
-
-
-def run_generation(prompt, width, height, steps, guidance_scale, n_images):
-    if _pipeline is None:
-        raise RuntimeError(_pipeline_load_error or "Pipeline not initialized.")
-
-    images = []
+def run_generation(prompt, width, height, steps, guidance):
     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
-
+        result = _pipeline(
+            prompt=prompt,
+            num_inference_steps=steps,
+            guidance_scale=guidance,
+            width=width,
+            height=height,
+        )
+        return result.images[0]
 
-# ---------------------------------------------------------------------------
-# Routes
-# ---------------------------------------------------------------------------
 
-@app.route("/health", methods=["GET"])
+@app.route("/health")
 def health():
-    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():
-    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 endpoint - kept unchanged from original)
-    if not request.is_json:
-        return openai_error_response("Request body must be valid JSON.", code="invalid_json")
-
-    try:
-        payload = request.get_json(silent=False)
-    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")
-
-    prompt = payload.get("prompt")
-    if not prompt or not isinstance(prompt, str) or not prompt.strip():
-        return openai_error_response("Missing 'prompt'.", param="prompt", code="missing_prompt")
-
-    requested_model = payload.get("model", MODEL_SHORT_ID)
-    if requested_model not in (MODEL_SHORT_ID, IMAGE_MODEL):
-        return openai_error_response(f"Model '{requested_model}' not supported.", param="model", code="model_not_found", status=404)
-
-    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. Supported: {supported}.", param="size", code="unsupported_size")
-
-    width, height = dimensions
-
-    n_images = payload.get("n", 1)
-    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")
-
-    response_format = payload.get("response_format", "b64_json")
-    if response_format != "b64_json":
-        return openai_error_response("Only b64_json supported.", param="response_format", code="unsupported_response_format")
-
-    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("Invalid steps or guidance_scale.", code="invalid_generation_params")
-
-    if steps < 1 or steps > 6:
-        steps = 4
-    if guidance_scale < 0 or guidance_scale > 5:
-        guidance_scale = 1.0
-
-    if _pipeline is None:
-        return openai_error_response("Model not ready.", err_type="server_error", code="model_not_ready", status=503)
-
-    request_id = uuid.uuid4().hex[:12]
-    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})
+    return jsonify({"status": "ok" if _pipeline else "loading", "model": "sd-turbo"})
 
 
 @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
+    prompt = request.args.get("prompt", "").replace("+", " ")
+    if not prompt:
+        return "Missing prompt", 400
 
     try:
-        width = int(request.args.get("width", 512))
-        height = int(request.args.get("height", 512))
+        width = int(request.args.get("width", DEFAULT_WIDTH))
+        height = int(request.args.get("height", DEFAULT_HEIGHT))
         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
+        return "Invalid 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
+    if width > 768: width = 768
+    if height > 768: height = 768
+    if steps > 6: steps = 6
+    if guidance > 2.0: guidance = 1.0
 
     try:
-        images = run_generation(
-            prompt=prompt,
-            width=width,
-            height=height,
-            steps=steps,
-            guidance_scale=guidance,
-            n_images=1,
-        )
-        pil_image = images[0]
-
-        img_io = io.BytesIO()
-        pil_image.save(img_io, format="PNG")
-        img_io.seek(0)
-
-        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
-
+        img = run_generation(prompt, width, height, steps, guidance)
+        buf = io.BytesIO()
+        img.save(buf, format="PNG")
+        buf.seek(0)
+        return send_file(buf, mimetype="image/png")
+    except Exception as e:
+        logger.exception("Generation failed")
+        return str(e), 500
 
-@app.errorhandler(404)
-def not_found(_error):
-    return openai_error_response("Endpoint not found.", code="not_found", status=404)
-
-
-@app.errorhandler(405)
-def method_not_allowed(_error):
-    return openai_error_response("Method not allowed.", code="method_not_allowed", status=405)
-
-
-@app.errorhandler(500)
-def internal_error(_error):
-    return openai_error_response("Internal server error.", err_type="server_error", code="internal_error", status=500)
-
-
-# ---------------------------------------------------------------------------
-# Entrypoint
-# ---------------------------------------------------------------------------
 
 if __name__ == "__main__":
-    try:
-        load_pipeline()
-    except Exception:
-        logger.error("Starting in degraded mode - pipeline failed to load.")
-
+    load_pipeline()
     port = int(os.environ.get("PORT", "7860"))
     app.run(host="0.0.0.0", port=port, threaded=True)