Update models.py

models.py

@@ -4,6 +4,7 @@ FIXED VERSION with proper IP-Adapter and BLIP-2 support
 """
 import torch
 import time
+import os
 from diffusers import (
     StableDiffusionXLControlNetImg2ImgPipeline,
     ControlNetModel,
@@ -14,7 +15,7 @@ from diffusers.models.attention_processor import AttnProcessor2_0
 from transformers import CLIPVisionModelWithProjection
 from insightface.app import FaceAnalysis
 from controlnet_aux import ZoeDetector, OpenposeDetector, LeresDetector, MidasDetector, MediapipeFaceDetector
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 from compel import Compel, ReturnedEmbeddingsType
 
 # Use reference implementation's attention processor
@@ -64,10 +65,44 @@ 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...")
+    
+    model_name = FACE_DETECTION_CONFIG['model_name'] # "antelopev2"
+    local_model_root = '.' # We want files to be in ./antelopev2
+    local_model_path = os.path.join(local_model_root, model_name)
+    
     try:
+        # --- NEW: Download models from HF Hub ---
+        print(f"  Ensuring insightface models are downloaded to {local_model_path}...")
+        
+        # Check if files exist before downloading
+        required_files = ["1k3d68.onnx", "2d106det.onnx", "genderage.onnx", "glintr100.onnx", "scrfd_10g_bnkps.onnx"]
+        files_exist = all(os.path.exists(os.path.join(local_model_path, f)) for f in required_files)
+        
+        if not files_exist:
+            print(f"  Models not found locally, downloading from {MODEL_REPO}...")
+            try:
+                snapshot_kwargs = {
+                    "repo_id": MODEL_REPO,
+                    "allow_patterns": [f"{model_name}/*.onnx"],
+                    "local_dir": local_model_root,
+                    "local_dir_use_symlinks": False,
+                    "repo_type": "model"
+                }
+                if HUGGINGFACE_TOKEN:
+                    snapshot_kwargs["token"] = HUGGINGFACE_TOKEN
+                
+                snapshot_download(**snapshot_kwargs)
+                print("  [OK] Downloaded insightface models.")
+            except Exception as e:
+                print(f"  [ERROR] Failed to download insightface models: {e}")
+                raise # Re-raise exception to be caught by the outer try-except
+        else:
+            print("  [OK] Insightface models found locally.")
+        # --- END NEW ---
+
         face_app = FaceAnalysis(
-            name=FACE_DETECTION_CONFIG['model_name'],
-            root='./models/insightface',
+            name=model_name,
+            root=local_model_root, # This will look in ./antelopev2
             providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
         )
         face_app.prepare(
@@ -76,6 +111,7 @@ def load_face_analysis():
         )
         print("  [OK] Face analysis model loaded successfully")
         return face_app, True
+        
     except Exception as e:
         print(f"  [WARNING] Face detection not available: {e}")
         return None, False