Update app.py

app.py

@@ -26,21 +26,24 @@ if __name__ == "__main__":
 
     from diffusers import StableDiffusionPipeline
 import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
 
-# Load model from Hugging Face Hub
+app = FastAPI()
+
+class Prompt(BaseModel):
+    prompt: str
+
+# Load once
 pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",  # You can use another model if needed
+    "runwayml/stable-diffusion-v1-5",
     torch_dtype=torch.float16,
     use_safetensors=True,
     revision="fp16"
-)
-
-# Move to GPU if available
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
-# Generate an image
-prompt = "A Minecraft player standing on a hill during sunset, pixel art style"
-image = pipe(prompt).images[0]
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Save the image
-image.save("minecraft_thumbnail.png")
+@app.post("/generate")
+def generate_image(data: Prompt):
+    image = pipe(data.prompt).images[0]
+    image.save("output.png")
+    return {"message": "Image saved as output.png"}