Update app.py

app.py

@@ -1,40 +1,34 @@
+import os
+os.environ["HF_HOME"] = "/tmp/hf_cache"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
+os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_cache"
+
 from fastapi import FastAPI
 from pydantic import BaseModel
 from diffusers import StableDiffusionPipeline
 import torch
 
-import uuid
-import base64
 from io import BytesIO
 from PIL import Image
+import base64
 
 app = FastAPI()
 
-# Load the model (NO fp16 issues now)
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5"
-)
-pipe = pipe.to("cpu")  # Or use .to("cuda") if you're on GPU
+# Load the model safely
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+pipe = pipe.to("cpu")
 
-# For receiving prompts from the frontend
 class Prompt(BaseModel):
     text: str
 
 @app.get("/")
-def greet_json():
-    return {"message": "Text to Image generation ready!"}
+def greet():
+    return {"message": "Model ready"}
 
 @app.post("/generate")
-def generate_image(prompt: Prompt):
+def generate(prompt: Prompt):
     image = pipe(prompt.text).images[0]
-
-    # 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
-    }
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return {"image_base64": img_str, "prompt": prompt.text}