Update main.py

main.py

@@ -9,11 +9,10 @@ import io
 from PIL import Image
 import os
 
-
-# Set custom Hugging Face cache directory
-os.environ["HF_HOME"] = "/app/cache"
-os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
-os.environ["HF_HUB_CACHE"] = "/app/cache"
+# Set cache directory to a writable location
+os.environ["HF_HOME"] = "/tmp/huggingface_cache"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"
+os.environ["HF_HUB_CACHE"] = "/tmp/huggingface_cache"
 
 # Initialize FastAPI app
 app = FastAPI(title="PromptAgro Image Generator API")
@@ -27,21 +26,83 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load Stable Diffusion LCM model (your original approach)
-print("🚀 Loading Stable Diffusion Model...")
-model_id = "rupeshs/LCM-runwayml-stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-print("✅ Model Loaded.")
+# Global variable for the pipeline
+pipe = None
+model_loading = False
+
+def load_model_if_needed():
+    """Load model lazily when first request arrives"""
+    global pipe, model_loading
+    
+    if pipe is not None:
+        return True
+    
+    if model_loading:
+        return False  # Already loading, wait
+    
+    model_loading = True
+    success = load_model()
+    model_loading = False
+    return success
+
+def load_model():
+    """Load the Stable Diffusion model with proper error handling"""
+    global pipe
+    
+    print("🚀 Loading Stable Diffusion Model...")
+    model_id = "rupeshs/LCM-runwayml-stable-diffusion-v1-5"
+    
+    try:
+        # Create cache directory if it doesn't exist
+        os.makedirs("/tmp/huggingface_cache", exist_ok=True)
+        
+        # Use appropriate dtype based on device
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        
+        print(f"📱 Device: {device}, dtype: {torch_dtype}")
+        
+        # Load the model with cache directory specified
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id, 
+            torch_dtype=torch_dtype,
+            cache_dir="/tmp/huggingface_cache",
+            local_files_only=False
+        )
+        
+        pipe = pipe.to(device)
+        
+        # Enable memory efficient attention if available
+        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+            try:
+                pipe.enable_xformers_memory_efficient_attention()
+                print("✅ XFormers memory efficient attention enabled")
+            except Exception:
+                print("⚠️ XFormers not available, using default attention")
+        
+        print(f"✅ Model Loaded successfully on {device}")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Failed to load model: {e}")
+        pipe = None
+        return False
+
+# Don't load model on startup - do it lazily
+# model_loaded = load_model()
 
 @app.get("/")
 async def root():
-    """Health check endpoint"""
+    """Health check endpoint with enhanced status"""
     return {
         "status": "alive",
         "service": "PromptAgro Image Generator",
-        "model_loaded": True,
-        "device": "cuda" if torch.cuda.is_available() else "cpu"
+        "model_loaded": pipe is not None,
+        "model_loading": model_loading,
+        "device": "cuda" if torch.cuda.is_available() else "cpu",
+        "model_status": "loaded" if pipe is not None else ("loading" if model_loading else "not_loaded"),
+        "torch_dtype": "float16" if torch.cuda.is_available() else "float32",
+        "ready_for_requests": pipe is not None
     }
 
 @app.post("/generate/")
@@ -50,19 +111,31 @@ async def generate_image(prompt: str = Form(...)):
     Generate product packaging image from input prompt.
     Returns image file directly (your original approach).
     """
+    # Lazy load model on first request
+    if not load_model_if_needed():
+        if model_loading:
+            raise HTTPException(status_code=503, detail="Model is loading, please wait...")
+        else:
+            raise HTTPException(status_code=503, detail="Model failed to load. Please check logs.")
+    
     print(f"🖌️ Generating image for prompt: {prompt}")
 
-    # Generate image (your original approach)
-    image = pipe(prompt).images[0]
+    try:
+        # Generate image (your original approach)
+        image = pipe(prompt).images[0]
 
-    # Save image to temp file (your original approach)
-    filename = f"/tmp/{uuid.uuid4().hex}.png"
-    image.save(filename)
+        # Save image to temp file (your original approach)
+        filename = f"/tmp/{uuid.uuid4().hex}.png"
+        image.save(filename)
 
-    print(f"📦 Image saved to {filename}")
+        print(f"📦 Image saved to {filename}")
 
-    # Return image file as response (your original approach)
-    return FileResponse(filename, media_type="image/png")
+        # Return image file as response (your original approach)
+        return FileResponse(filename, media_type="image/png")
+    
+    except Exception as e:
+        print(f"❌ Image generation failed: {e}")
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
 
 @app.post("/generate-json/")
 async def generate_image_json(
@@ -75,6 +148,13 @@ async def generate_image_json(
     """
     Generate image and return as JSON with base64 data (for frontend integration).
     """
+    # Lazy load model on first request
+    if not load_model_if_needed():
+        if model_loading:
+            raise HTTPException(status_code=503, detail="Model is loading, please wait...")
+        else:
+            raise HTTPException(status_code=503, detail="Model failed to load. Please check logs.")
+    
     print(f"🖌️ Generating image for prompt: {prompt}")
     
     try:
@@ -116,6 +196,13 @@ async def generate_packaging_specific(
     """
     Generate packaging with PromptAgro-specific prompt engineering
     """
+    # Lazy load model on first request
+    if not load_model_if_needed():
+        if model_loading:
+            raise HTTPException(status_code=503, detail="Model is loading, please wait...")
+        else:
+            raise HTTPException(status_code=503, detail="Model failed to load. Please check logs.")
+    
     # Create professional prompt for agricultural packaging
     prompt = f"""Professional agricultural product packaging design for {product_name}, 
     modern clean style, {colors.replace(',', ' and ')} color scheme, premium typography,