Update models.py

models.py

@@ -16,7 +16,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,81 +64,54 @@ def download_model_with_retry(repo_id, filename, max_retries=None, **kwargs):
 def load_face_analysis():
     """
     Load face analysis model with proper error handling.
-    This version downloads files manually to a custom folder
-    to bypass the insightface hard-coded zip download.
+    This version pre-downloads models to the exact path
+    insightface expects, to bypass its internal downloader.
     """
     print("Loading face analysis model...")
     
-    # Use a custom model name to prevent insightface from auto-downloading a zip
-    model_name = FACE_DETECTION_CONFIG['model_name'] # "pixagram_face_models"
-    local_model_root = '.' # We want files to be in ./pixagram_face_models
-    local_model_path = os.path.join(local_model_root, model_name)
+    # The library expects models to be at: [root]/[name]
+    # We will set root = './models' and name = 'antelopev2'
+    # The final path it will check is './models/antelopev2'
+    
+    local_model_root = './models' 
+    model_name = FACE_DETECTION_CONFIG['model_name'] # 'antelopev2'
+    
+    # This is the directory where the .onnx files must be:
+    local_model_path = os.path.join(local_model_root, model_name) # ./models/antelopev2
     
     try:
-        # --- NEW: Manual download logic ---
+        # --- NEW: Download logic based on user's snippet ---
         print(f"  Ensuring insightface models are present in {local_model_path}...")
-        os.makedirs(local_model_path, exist_ok=True)
         
-        required_files = [
-            "1k3d68.onnx", 
-            "2d106det.onnx", 
-            "genderage.onnx", 
-            "glintr100.onnx", 
-            "scrfd_10g_bnkps.onnx"
-        ]
+        # Check if models are already downloaded
+        test_file_path = os.path.join(local_model_path, "scrfd_10g_bnkps.onnx")
         
-        for file_name in required_files:
-            local_file_path = os.path.join(local_model_path, file_name)
+        if not os.path.exists(test_file_path):
+            print(f"    Models not found locally, downloading from {MODEL_REPO}...")
             
-            if not os.path.exists(local_file_path):
-                print(f"    Downloading {file_name}...")
-                
-                # Path to the file in the HF model repo
-                repo_file_path = f"antelopev2/{file_name}"
-                
-                try:
-                    # Download the file directly into our target folder
-                    downloaded_path = download_model_with_retry(
-                        repo_id=MODEL_REPO,
-                        filename=repo_file_path,
-                        local_dir=local_model_path,
-                        local_dir_use_symlinks=False,
-                        resume_download=True,
-                        repo_type="model"
-                    )
-                    
-                    # hf_hub_download *might* preserve folder structure,
-                    # e.g., saving to ./pixagram_face_models/antelopev2/genderage.onnx
-                    # We must move it if that happens.
-                    
-                    expected_download_path = os.path.join(local_model_path, *repo_file_path.split('/'))
-                    
-                    if os.path.exists(expected_download_path) and expected_download_path != local_file_path:
-                        print(f"    Moving {expected_download_path} to {local_file_path}")
-                        shutil.move(expected_download_path, local_file_path)
-                        
-                        # Clean up empty antelopev2 folder if it was created
-                        try:
-                            os.rmdir(os.path.dirname(expected_download_path))
-                        except OSError:
-                            pass # Not empty, which is fine
-                            
-                    elif not os.path.exists(local_file_path):
-                         # Fallback in case logic is wrong, just check the returned path
-                         if downloaded_path != local_file_path:
-                             print(f"    Moving {downloaded_path} to {local_file_path}")
-                             shutil.move(downloaded_path, local_file_path)
-
-                except Exception as e:
-                    print(f"    [ERROR] Failed to download {file_name}: {e}")
-                    raise # Re-raise to stop startup
+            snapshot_kwargs = {
+                "repo_id": MODEL_REPO,
+                "repo_type": "model",
+                "allow_patterns": f"{model_name}/*.onnx", # Download only files from antelopev2 folder
+                "local_dir": local_model_root, # Download *into* ./models
+                "local_dir_use_symlinks": False
+            }
+            if HUGGINGFACE_TOKEN:
+                snapshot_kwargs["token"] = HUGGINGFACE_TOKEN
             
-        print("  [OK] All insightface models are present locally.")
+            # This will download files, preserving the repo folder structure,
+            # creating ./models/antelopev2/[...].onnx
+            snapshot_download(**snapshot_kwargs)
+            print("    [OK] Downloaded insightface models.")
+        else:
+            print("    [OK] Insightface models found locally.")
         # --- END NEW ---
 
+        # Now, initialize FaceAnalysis
+        # It will look for models in root + name = ./models/antelopev2
         face_app = FaceAnalysis(
-            name=model_name,        # "pixagram_face_models" (custom name)
-            root=local_model_root,  # "." (looks in ./pixagram_face_models)
+            name=model_name,
+            root=local_model_root, 
             providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
         )
         face_app.prepare(
@@ -297,7 +270,7 @@ def load_sdxl_pipeline(controlnets):
             "stabilityai/stable-diffusion-xl-base-1.0",
             controlnet=controlnets,
             torch_dtype=dtype,
-            use_safetponsors=True
+            use_safetensors=True
         ).to(device)
         return pipe, False
 
@@ -519,7 +492,7 @@ def load_caption_model():
                 torch_dtype=dtype
             ).to(device)
             print("  [OK] BLIP base model loaded (standard captions)")
-            return caption_primary, caption_model, True, 'blip'
+            return caption_processor, caption_model, True, 'blip'
         except Exception as e2:
             print(f"  [WARNING] Caption models not available: {e2}")
             print("  Caption generation will be disabled")
@@ -529,7 +502,7 @@ def load_caption_model():
 def set_clip_skip(pipe):
     """Set CLIP skip value."""
     if hasattr(pipe, 'text_encoder'):
-        print(f"  [OK] CLIP skip set to {CLIP_SKIP}")
+        print("  [OK] CLIP skip set to {CLIP_SKIP}")
 
 
 print("[OK] Model loading functions ready")