Update models.py

models.py

@@ -300,13 +300,13 @@ def load_loras(pipe):
 
 def setup_ip_adapter(pipe, image_encoder):
     """
-    Setup IP-Adapter for InstantID face embeddings - PROPER IMPLEMENTATION.
-    Based on the reference InstantID pipeline.
+    Setup IP-Adapter for InstantID face embeddings.
+    This is CRITICAL for face preservation.
     """
     if image_encoder is None:
         return None, False
     
-    print("Setting up IP-Adapter for InstantID face embeddings (proper implementation)...")
+    print("Setting up IP-Adapter for InstantID face embeddings...")
     try:
         # Download InstantID weights
         ip_adapter_path = download_model_with_retry(
@@ -328,7 +328,7 @@ def setup_ip_adapter(pipe, image_encoder):
             elif key.startswith("ip_adapter."):
                 ip_adapter_state_dict[key.replace("ip_adapter.", "")] = value
         
-        # Create Resampler (image projection model) with CORRECT parameters from reference
+        # Create Resampler with CORRECT parameters
         print("Creating Resampler (Perceiver architecture)...")
         image_proj_model = Resampler(
             dim=1280,
@@ -336,7 +336,7 @@ def setup_ip_adapter(pipe, image_encoder):
             dim_head=64,
             heads=20,
             num_queries=16,
-            embedding_dim=512,
+            embedding_dim=512,  # CRITICAL: Must match InsightFace embedding size
             output_dim=pipe.unet.config.cross_attention_dim,
             ff_mult=4
         )
@@ -351,9 +351,6 @@ def setup_ip_adapter(pipe, image_encoder):
                 print("  [OK] Resampler loaded with pretrained weights")
             except Exception as e:
                 print(f"  [WARNING] Could not load Resampler weights: {e}")
-                print("  Using randomly initialized Resampler")
-        else:
-            print("  [WARNING] No image_proj weights found, using random initialization")
         
         # Setup IP-Adapter attention processors
         print("Setting up IP-Adapter attention processors...")
@@ -387,7 +384,7 @@ def setup_ip_adapter(pipe, image_encoder):
         # Set attention processors
         pipe.unet.set_attn_processor(attn_procs)
         
-        # Load IP-Adapter weights into attention processors
+        # Load IP-Adapter weights
         if ip_adapter_state_dict:
             try:
                 ip_layers = torch.nn.ModuleList(pipe.unet.attn_processors.values())
@@ -395,15 +392,13 @@ def setup_ip_adapter(pipe, image_encoder):
                 print("  [OK] IP-Adapter attention weights loaded")
             except Exception as e:
                 print(f"  [WARNING] Could not load IP-Adapter weights: {e}")
-        else:
-            print("  [WARNING] No ip_adapter weights found")
         
-        # Store image encoder and projection model
+        # Store image encoder
         pipe.image_encoder = image_encoder
         
         print("  [OK] IP-Adapter fully loaded with InstantID architecture")
         print(f"  - Resampler: 4 layers, 20 heads, 16 output tokens")
-        print(f"  - Face embeddings: 512D -> 16x2048D")
+        print(f"  - Face embeddings: 512D -> 16x{pipe.unet.config.cross_attention_dim}D")
         
         return image_proj_model, True