Update models.py

models.py

@@ -60,19 +60,20 @@ def download_model_with_retry(repo_id, filename, max_retries=None):
 
 
 def load_face_analysis():
-    """Load face analysis model with proper error handling."""
-    print("Loading face analysis model...")
+    """Load face analysis model on CPU to save GPU memory."""
+    print("Loading face analysis model on CPU...")
     try:
+        # Force CPU execution for face analysis to save GPU memory
         face_app = FaceAnalysis(
             name=FACE_DETECTION_CONFIG['model_name'],
             root='./models/insightface',
-            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
+            providers=['CPUExecutionProvider']  # CPU only for face detection
         )
         face_app.prepare(
-            ctx_id=FACE_DETECTION_CONFIG['ctx_id'], 
+            ctx_id=-1,  # -1 for CPU
             det_size=FACE_DETECTION_CONFIG['det_size']
         )
-        print("  [OK] Face analysis model loaded successfully")
+        print("  [OK] Face analysis model loaded on CPU (GPU memory saved)")
         return face_app, True
     except Exception as e:
         print(f"  [WARNING] Face detection not available: {e}")
@@ -80,12 +81,13 @@ def load_face_analysis():
 
 
 def load_depth_detector():
-    """Load Zoe Depth detector."""
+    """Load Zoe Depth detector with optimized memory management."""
     print("Loading Zoe Depth detector...")
     try:
         zoe_depth = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-        zoe_depth.to(device)
-        print("  [OK] Zoe Depth loaded successfully")
+        # Start on CPU to save memory during initialization
+        zoe_depth = zoe_depth.to("cpu")
+        print("  [OK] Zoe Depth loaded (on CPU, will move to GPU when needed)")
         return zoe_depth, True
     except Exception as e:
         print(f"  [WARNING] Zoe Depth not available: {e}")
@@ -212,6 +214,7 @@ def load_caption_model():
     """
     Load caption model with proper error handling.
     Tries multiple models in order of quality.
+    Models start on CPU and move to GPU only when needed.
     """
     print("Loading caption model...")
     
@@ -223,9 +226,9 @@ def load_caption_model():
         caption_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
         caption_model = AutoModelForCausalLM.from_pretrained(
             "microsoft/git-large-coco",
-            torch_dtype=dtype
-        ).to(device)
-        print("  [OK] GIT-Large model loaded")
+            torch_dtype=dtype  # Use dtype from config
+        ).to("cpu")  # Start on CPU to save GPU memory
+        print("  [OK] GIT-Large model loaded (on CPU, will move to GPU when needed)")
         return caption_processor, caption_model, True, 'git'
     except Exception as e1:
         print(f"  [INFO] GIT-Large not available: {e1}")
@@ -238,9 +241,9 @@ def load_caption_model():
             caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
             caption_model = BlipForConditionalGeneration.from_pretrained(
                 "Salesforce/blip-image-captioning-base",
-                torch_dtype=dtype
-            ).to(device)
-            print("  [OK] BLIP base model loaded")
+                torch_dtype=dtype  # Use dtype from config
+            ).to("cpu")  # Start on CPU to save GPU memory
+            print("  [OK] BLIP base model loaded (on CPU, will move to GPU when needed)")
             return caption_processor, caption_model, True, 'blip'
         except Exception as e2:
             print(f"  [WARNING] Caption models not available: {e2}")
@@ -253,4 +256,4 @@ def set_clip_skip(pipe):
         print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
-print("[OK] Model loading functions ready")
+print("[OK] Model loading functions ready")