Update app.py

app.py

@@ -3,21 +3,38 @@ from pydantic import BaseModel
 from diffusers import StableDiffusionPipeline
 import torch
 
+import uuid
+import base64
+from io import BytesIO
+from PIL import Image
+
 app = FastAPI()
 
-# Load model at startup
+# Load the model (NO fp16 issues now)
 pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16,
-    revision="fp16"
-).to("cuda" if torch.cuda.is_available() else "cpu")
+    "runwayml/stable-diffusion-v1-5"
+)
+pipe = pipe.to("cpu")  # Or use .to("cuda") if you're on GPU
 
+# For receiving prompts from the frontend
 class Prompt(BaseModel):
     text: str
 
+@app.get("/")
+def greet_json():
+    return {"message": "Text to Image generation ready!"}
+
 @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}
+
+    # Convert image to base64 to send over HTTP
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    return {
+        "image_base64": img_str,
+        "status": "success",
+        "prompt": prompt.text
+    }