Update app.py

app.py

@@ -2,24 +2,27 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from diffusers import StableDiffusionPipeline
 import torch
+import uuid
+from PIL import Image
 import base64
 from io import BytesIO
-from PIL import Image
-
-# Load the pipeline
-pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
-pipe = pipe.to("cpu")
 
 app = FastAPI()
 
+# Load model
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
+pipe = pipe.to("cpu")  # or "cuda" if using GPU
+
 class PromptRequest(BaseModel):
     prompt: str
 
-@app.post("/generate")
-def generate_image(req: PromptRequest):
-    image = pipe(req.prompt).images[0]
-
+@app.post("/generate-image")
+def generate_image(request: PromptRequest):
+    image = pipe(request.prompt).images[0]
+    
+    # Convert to base64
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
     return {"image_base64": img_str}