Update server.py

server.py

@@ -2,12 +2,12 @@ import io
 import logging
 import os
 import threading
+import uuid
 
 import torch
-from flask import Flask, jsonify, request, send_file
+from flask import Flask, jsonify, request, send_file, redirect
 from flask_cors import CORS
 
-# CPU limits
 os.environ["OMP_NUM_THREADS"] = "4"
 os.environ["MKL_NUM_THREADS"] = "4"
 torch.set_num_threads(4)
@@ -16,7 +16,6 @@ torch.set_num_interop_threads(1)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
 logger = logging.getLogger("sd-turbo-server")
 
-# Config
 DEFAULT_STEPS = int(os.environ.get("DEFAULT_STEPS", "4"))
 DEFAULT_GUIDANCE = float(os.environ.get("DEFAULT_GUIDANCE", "1.0"))
 DEFAULT_WIDTH = 768
@@ -68,9 +67,70 @@ def run_generation(prompt, width, height, steps, guidance):
         return result.images[0]
 
 
+@app.route("/")
+def index():
+    return "SD Turbo Server - Use /image or /v1/images/generations"
+
+
 @app.route("/health")
 def health():
-    return jsonify({"status": "ok" if _pipeline else "loading", "model": "sd-turbo"})
+    return jsonify({"status": "ok", "model": "sd-turbo"})
+
+
+@app.route("/v1/models")
+def list_models():
+    return jsonify({
+        "object": "list",
+        "data": [{"id": "sd-turbo", "object": "model", "owned_by": "stabilityai"}]
+    })
+
+
+@app.route("/v1/images/generations", methods=["POST"])
+def images_generations():
+    """OpenAI-compatible endpoint that returns a direct image URL"""
+    if not request.is_json:
+        return jsonify({"error": {"message": "JSON required"}}), 400
+
+    data = request.get_json()
+    prompt = data.get("prompt")
+    if not prompt:
+        return jsonify({"error": {"message": "prompt is required"}}), 400
+
+    size = data.get("size", "768x768")
+    n = data.get("n", 1)
+    steps = int(data.get("num_inference_steps", DEFAULT_STEPS))
+    guidance = float(data.get("guidance_scale", DEFAULT_GUIDANCE))
+
+    if steps > 6: steps = 6
+    if guidance > 2.0: guidance = 1.5
+
+    width = DEFAULT_WIDTH
+    height = DEFAULT_HEIGHT
+    if "x" in size:
+        try:
+            w, h = map(int, size.split("x"))
+            width, height = min(w, 768), min(h, 768)
+        except:
+            pass
+
+    request_id = uuid.uuid4().hex[:8]
+    logger.info(f"[{request_id}] OpenAI gen - prompt={prompt[:60]}... size={width}x{height}")
+
+    try:
+        image = run_generation(prompt, width, height, steps, guidance)
+        
+        # Save to a temporary in-memory route so we can return a URL
+        img_id = uuid.uuid4().hex[:12]
+        # For simplicity we return a direct /image URL with the same prompt
+        image_url = f"https://cloudunity-sdturbolumi.hf.space/image?prompt={prompt.replace(' ', '+')}&width={width}&height={height}&steps={steps}&guidance={guidance}"
+        
+        return jsonify({
+            "created": int(time.time()),
+            "data": [{"url": image_url}]
+        })
+    except Exception as e:
+        logger.exception("Generation failed")
+        return jsonify({"error": {"message": str(e)}}), 500
 
 
 @app.route("/image", methods=["GET"])
@@ -99,7 +159,7 @@ def generate_image():
         buf.seek(0)
         return send_file(buf, mimetype="image/png")
     except Exception as e:
-        logger.exception("Generation failed")
+        logger.exception("Image generation failed")
         return str(e), 500