Update app.py

app.py

@@ -1,7 +1,23 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
+from diffusers import StableDiffusionPipeline
+import torch
 
 app = FastAPI()
 
-@app.get("/")
-def root():
-    return {"message": "Welcome to SpeedAI Image Generator 🚀"}
+# Load model at startup
+pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+    revision="fp16"
+).to("cuda" if torch.cuda.is_available() else "cpu")
+
+class Prompt(BaseModel):
+    text: str
+
+@app.post("/generate")
+def generate_image(prompt: Prompt):
+    image = pipe(prompt.text).images[0]
+    image_path = f"output.png"
+    image.save(image_path)
+    return {"message": "Image generated", "image_path": image_path}