Update app.py

app.py

@@ -1,35 +1,47 @@
-import os
-import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
 from diffusers import StableDiffusionPipeline
+import torch, base64, io
+from PIL import Image
 
-# Read HF token from Spaces Secrets
-HF_TOKEN = os.getenv("HF_TOKEN")
+app = FastAPI()
 
-if HF_TOKEN is None:
-    raise RuntimeError("HF_TOKEN not found. Add it in HF Spaces → Secrets.")
-
-model_id = "prompthero/openjourney"
+MODEL_ID = "prompthero/openjourney-v4"
 
 pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
-    token=HF_TOKEN,
-    torch_dtype=torch.float32,   # REQUIRED for CPU
-    safety_checker=None          # reduce RAM
+    MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    safety_checker=None,
+    requires_safety_checker=False
 )
 
-pipe.to("cpu")
-
-prompt = "retro serie of different cars with different colors and shapes, mdjrny-v4 style"
-
-print("Starting image generation (CPU-only, this is slow)...")
-
-image = pipe(
-    prompt,
-    height=256,                 # DO NOT increase
-    width=256,
-    num_inference_steps=10,      # DO NOT increase
-    guidance_scale=7.0
-).images[0]
-
-image.save("output.png")
-print("Saved output.png")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
+
+# VERY IMPORTANT for speed & stability
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+
+class GenerateRequest(BaseModel):
+    prompt: str
+    steps: int = 20
+    width: int = 512
+    height: int = 512
+
+@app.post("/generate")
+def generate(req: GenerateRequest):
+    image = pipe(
+        req.prompt,
+        num_inference_steps=req.steps,
+        width=req.width,
+        height=req.height
+    ).images[0]
+
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    img_base64 = base64.b64encode(buf.getvalue()).decode()
+
+    return {
+        "status": "ok",
+        "image_base64": img_base64
+    }