Update app.py

app.py

@@ -6,6 +6,45 @@ warnings.filterwarnings("ignore")
 
 from ui_manager import UIManager
 
+def preload_models_to_cache():
+    """
+    Pre-download models to HuggingFace cache before GPU allocation.
+    This runs on CPU and avoids downloading during @spaces.GPU execution.
+    """
+    if not os.getenv('SPACE_ID'):
+        return  # Skip if not on Spaces
+
+    print("📦 Pre-downloading models to cache (CPU only, no GPU usage)...")
+
+    try:
+        from diffusers import ControlNetModel
+        import torch
+
+        # Pre-download ControlNet models to cache
+        models_to_cache = [
+            ("diffusers/controlnet-canny-sdxl-1.0", "Canny ControlNet"),
+            ("diffusers/controlnet-depth-sdxl-1.0", "Depth ControlNet"),
+        ]
+
+        for model_id, model_name in models_to_cache:
+            print(f"  ⬇️  Downloading {model_name} ({model_id})...")
+            try:
+                _ = ControlNetModel.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float16,
+                    use_safetensors=True,
+                    local_files_only=False  # Allow download
+                )
+                print(f"  ✅ {model_name} cached")
+            except Exception as e:
+                print(f"  ⚠️  {model_name} download failed (will retry on-demand): {e}")
+
+        print("✅ Model pre-caching complete")
+
+    except Exception as e:
+        print(f"⚠️  Model pre-caching failed: {e}")
+        print("   Models will be downloaded on first use instead.")
+
 def launch_final_blend_sceneweaver(share: bool = True, debug: bool = False):
     """Launch SceneWeaver Application"""
 
@@ -13,6 +52,9 @@ def launch_final_blend_sceneweaver(share: bool = True, debug: bool = False):
     print("✨ AI-Powered Image Background Generation")
 
     try:
+        # Pre-download models on Spaces to avoid downloading during GPU time
+        preload_models_to_cache()
+
         # Test imports first
         print("🔍 Testing imports...")
         try:
@@ -21,13 +63,11 @@ def launch_final_blend_sceneweaver(share: bool = True, debug: bool = False):
             ui = UIManager()
             print("✅ UIManager instance created successfully")
 
-            # Note: On Hugging Face Spaces, models will be loaded on-demand
-            # within @spaces.GPU decorated functions to comply with ZeroGPU requirements
+            # Note: On Hugging Face Spaces, models are pre-cached at startup
             if os.getenv('SPACE_ID'):
                 print("\n🔧 Detected Hugging Face Spaces environment")
-                print("⚡ Models will load on-demand (ZeroGPU stateless requirement)")
-                print("   First inpainting request: ~7 steps (280s)")
-                print("   Subsequent requests: ~11 steps (245s)")
+                print("⚡ Models pre-cached - ready for fast inference")
+                print("   Expected inference time: ~300-350s (with cached models)")
                 print()
 
             # Launch UI
@@ -88,5 +128,4 @@ def main():
         raise
 
 if __name__ == "__main__":
-    main()
-    
+    main()