Update app.py

app.py

@@ -1,18 +1,31 @@
+import os
+import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import story_generator  # ✅ Import Story Generator
-import torch
 from diffusers import DiffusionPipeline
 import io
 import base64
 from PIL import Image
+from huggingface_hub import login
 
 app = FastAPI()
 
-# ✅ Load Image Generation Model
-IMAGE_MODEL = "stabilityai/sdxl-lightning"
+# ✅ Set Hugging Face cache directory to /tmp
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
+
+# ✅ Hugging Face Authentication (Only needed if model is private)
+HF_TOKEN = "your_huggingface_token_here"  # Replace this with your actual token
+login(token=HF_TOKEN)
+
+# ✅ Load Image Generation Model (Use a fast, public model)
+IMAGE_MODEL = "stabilityai/sdxl-turbo"  # Replace with "stabilityai/sdxl-lightning" if needed
 pipeline = DiffusionPipeline.from_pretrained(
-    IMAGE_MODEL, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    IMAGE_MODEL,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    use_auth_token=HF_TOKEN  # Required for private models
 ).to("cuda" if torch.cuda.is_available() else "cpu")
 
 # ✅ Define the input request format