Update models.py

models.py

@@ -1,6 +1,6 @@
 """
 Model loading and initialization for Pixagram AI Pixel Art Generator
-UPDATED VERSION with proper InstantID pipeline support
+CORRECTED VERSION with proper face analysis loading
 """
 import torch
 import time
@@ -12,11 +12,13 @@ from diffusers import (
 )
 from insightface.app import FaceAnalysis
 from controlnet_aux import ZoeDetector
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 from compel import Compel, ReturnedEmbeddingsType
 
 # Use InstantID pipeline
-from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
+from pipeline_stable_diffusion_xl_instantid_img2img import (
+    StableDiffusionXLInstantIDImg2ImgPipeline
+)
 
 from config import (
     device, dtype, MODEL_REPO, MODEL_FILES, HUGGINGFACE_TOKEN,
@@ -59,18 +61,79 @@ def download_model_with_retry(repo_id, filename, max_retries=None):
 
 
 def load_face_analysis():
-    """Load face analysis model with intelligent provider selection."""
+    """
+    Load face analysis model using the correct approach.
+    Downloads antelopev2 model and initializes FaceAnalysis.
+    """
     print("Loading face analysis model...")
-    face_app = FaceAnalysis(
-        name=FACE_DETECTION_CONFIG['model_name'],
-        root='/data',
-        providers=['CPUExecutionProvider']
-    )
-    face_app.prepare(
-        ctx_id=ctx_id,
-        det_size=FACE_DETECTION_CONFIG['det_size']
-    )
-    return face_app, True
+    
+    try:
+        # Download antelopev2 model using snapshot_download (like working example)
+        print("  Downloading antelopev2 model files...")
+        antelope_path = snapshot_download(
+            repo_id=FACE_DETECTION_CONFIG['download_repo'],
+            local_dir=FACE_DETECTION_CONFIG['local_dir']
+        )
+        print(f"  [OK] Antelopev2 downloaded to: {antelope_path}")
+        
+        # Initialize FaceAnalysis with the correct root path
+        # Use CPU provider for memory efficiency (can be changed in config)
+        providers = FACE_DETECTION_CONFIG.get('providers', ['CPUExecutionProvider'])
+        
+        print(f"  Initializing FaceAnalysis with providers: {providers}")
+        face_app = FaceAnalysis(
+            name=FACE_DETECTION_CONFIG['model_name'],
+            root=FACE_DETECTION_CONFIG['root'],
+            providers=providers
+        )
+        
+        # Prepare the model
+        face_app.prepare(
+            ctx_id=FACE_DETECTION_CONFIG['ctx_id'],
+            det_size=FACE_DETECTION_CONFIG['det_size']
+        )
+        
+        # Test the model to ensure it works
+        import numpy as np
+        test_img = np.zeros((640, 640, 3), dtype=np.uint8)
+        _ = face_app.get(test_img)
+        
+        print(f"  [OK] Face analysis model loaded successfully")
+        print(f"  [INFO] Using providers: {providers}")
+        return face_app, True
+        
+    except Exception as e:
+        print(f"  [ERROR] Face analysis loading failed: {e}")
+        import traceback
+        traceback.print_exc()
+        
+        # Try fallback with different providers
+        try:
+            print("  [INFO] Trying fallback with auto-detect providers...")
+            face_app = FaceAnalysis(
+                name=FACE_DETECTION_CONFIG['model_name'],
+                root=FACE_DETECTION_CONFIG['root']
+            )
+            face_app.prepare(
+                ctx_id=0,
+                det_size=FACE_DETECTION_CONFIG['det_size']
+            )
+            
+            # Test
+            import numpy as np
+            test_img = np.zeros((640, 640, 3), dtype=np.uint8)
+            _ = face_app.get(test_img)
+            
+            print("  [OK] Face analysis loaded with auto-detect providers")
+            return face_app, True
+            
+        except Exception as e2:
+            print(f"  [WARNING] Face detection not available: {e2}")
+            print("  [INFO] Generation will continue without face preservation")
+            print("  [TIP] Check that onnxruntime is properly installed:")
+            print("        pip install onnxruntime --break-system-packages")
+            return None, False
+
 
 def load_depth_detector():
     """Load Zoe Depth detector with optimized memory management."""
@@ -110,23 +173,44 @@ def load_controlnets():
 
 
 def load_sdxl_pipeline(controlnets):
-    """Load SDXL checkpoint from HuggingFace Hub."""
-    print("Loading SDXL checkpoint (horizon) with bundled VAE from HuggingFace Hub...")
+    """
+    Load SDXL pipeline with InstantID support.
+    controlnets MUST be a list: [identitynet, depthnet]
+    """
+    print("Loading SDXL checkpoint with InstantID pipeline...")
     try:
         model_path = download_model_with_retry(MODEL_REPO, MODEL_FILES['checkpoint'])
         
+        # CRITICAL: Use InstantID-enabled pipeline (not standard ControlNet pipeline)
         pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
             model_path,
             controlnet=controlnets,
             torch_dtype=dtype,
             use_safetensors=True
         ).to(device)
-        print("  [OK] Custom checkpoint loaded successfully (VAE bundled)")
+        
+        # Load IP-Adapter weights for InstantID
+        print("Loading IP-Adapter for InstantID...")
+        ip_adapter_path = download_model_with_retry(
+            "InstantX/InstantID",
+            "ip-adapter.bin"
+        )
+        pipe.load_ip_adapter_instantid(ip_adapter_path)
+        # Don't set default scale - will be set dynamically based on face detection
+        print("  [OK] IP-Adapter loaded (scale will be set dynamically)")
+        
+        print("  [OK] InstantID pipeline loaded successfully")
         return pipe, True
+        
     except Exception as e:
-        print(f"  [WARNING] Could not load custom checkpoint: {e}")
-        print("  Using default SDXL base model")
-        pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
+        print(f"  [ERROR] Could not load InstantID pipeline: {e}")
+        import traceback
+        traceback.print_exc()
+        
+        # Fallback to standard pipeline
+        print("  [WARNING] Falling back to standard SDXL pipeline (no InstantID)")
+        from diffusers import StableDiffusionXLControlNetImg2ImgPipeline
+        pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
             "stabilityai/stable-diffusion-xl-base-1.0",
             controlnet=controlnets,
             torch_dtype=dtype,
@@ -134,6 +218,7 @@ def load_sdxl_pipeline(controlnets):
         ).to(device)
         return pipe, False
 
+
 def load_lora(pipe):
     """Load LORA from HuggingFace Hub."""
     print("Loading LORA (retroart) from HuggingFace Hub...")