Create generate.py

generate.py

@@ -1,3 +1,4 @@
+# generate.py
 import os
 import requests
 from flask import Flask, request, send_file, jsonify
@@ -5,54 +6,37 @@ from io import BytesIO
 
 app = Flask(__name__)
 
-# Get HF API key from Secrets (HF_API_KEY)
+# Make sure you set your HF API key in Secrets (HF_API_KEY)
 HF_API_KEY = os.environ.get("HF_API_KEY")
 if not HF_API_KEY:
-    raise ValueError("HF_API_KEY is missing! Add it in your Space secrets.")
-
-# Default model
-DEFAULT_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+    raise ValueError("HF_API_KEY not set in environment variables")
 
 @app.route("/generate", methods=["POST"])
 def generate():
-    data = request.json or {}
+    data = request.json
     prompt = data.get("prompt", "")
-    negative_prompt = data.get("negative_prompt", "")
-    model = data.get("model", DEFAULT_MODEL)
-    width = min(1024, int(data.get("width", 512)))   # Keep sizes safe
-    height = min(1024, int(data.get("height", 512)))
+    model = data.get("model", "stabilityai/stable-diffusion-xl-base-1.0")
 
-    if not prompt:
-        return jsonify({"error": "Prompt is required"}), 400
+    if not prompt.strip():
+        return jsonify({"error": "Prompt cannot be empty"}), 400
 
     headers = {"Authorization": f"Bearer {HF_API_KEY}"}
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "width": width,
-            "height": height,
-            "negative_prompt": negative_prompt
-        }
-    }
-
-    # Call HF Inference API
-    response = requests.post(
-        f"https://api-inference.huggingface.co/models/{model}",
-        headers=headers,
-        json=payload,
-        timeout=60  # prevent freezing
-    )
-
-    # If API returns JSON, it's likely an error
-    content_type = response.headers.get("content-type", "")
-    if "application/json" in content_type:
-        try:
-            data = response.json()
-            return jsonify({"error": data.get("error", str(data))}), 400
-        except:
-            return jsonify({"error": response.text}), 400
-
-    # Return image as PNG
+    payload = {"inputs": prompt}
+
+    try:
+        response = requests.post(
+            f"https://api-inference.huggingface.co/models/{model}",
+            headers=headers,
+            json=payload,
+            timeout=120
+        )
+    except requests.exceptions.RequestException as e:
+        return jsonify({"error": str(e)}), 500
+
+    if response.status_code != 200:
+        return jsonify({"error": response.text}), response.status_code
+
+    # Return the raw bytes as PNG
     return send_file(BytesIO(response.content), mimetype="image/png")
 
 if __name__ == "__main__":