Update model.py

model.py

@@ -5,11 +5,15 @@ import numpy as np
 from config import Config
 
 from diffusers import (
-    StableDiffusionXLControlNetPipeline, 
     ControlNetModel, 
     LCMScheduler,
-    StableDiffusionXLInstantIDPipeline  # <-- ADD THIS IMPORT
+    StableDiffusionXLInstantIDPipeline  # To "steal" the ControlNet
 )
+# --- START FIX ---
+# Import the custom pipeline from your local file
+from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline
+# --- END FIX ---
+
 from huggingface_hub import snapshot_download
 from insightface.app import FaceAnalysis
 from controlnet_aux import ZoeDetector, LineartDetector
@@ -29,10 +33,8 @@ class ModelHandler:
         Forces CPU to avoid ZeroGPU initialization errors.
         """
         print("Loading face analysis model...")
-        # Path will now be './antelopev2'
         model_root_path = os.path.join(Config.ANTELOPEV2_ROOT, Config.ANTELOPEV2_NAME)
         
-        # 1. Download from HF Hub
         if not os.path.exists(os.path.join(model_root_path, "scrfd_10g_bnkps.onnx")):
             print(f"Downloading AntelopeV2 models from {Config.ANTELOPEV2_REPO}...")
             try:
@@ -45,14 +47,12 @@ class ModelHandler:
                 print(f"  [ERROR] Failed to download AntelopeV2 models: {e}")
                 return False
 
-        # 2. Initialize InsightFace on CPU
         try:
             self.app = FaceAnalysis(
                 name=Config.ANTELOPEV2_NAME, # 'antelopev2'
                 root=Config.ANTELOPEV2_ROOT, # '.'
                 providers=['CPUExecutionProvider'] 
             )
-            # This will now correctly look in './antelopev2'
             self.app.prepare(ctx_id=0, det_size=(640, 640))
             print(f"  [OK] Face analysis model loaded successfully.")
             return True
@@ -68,9 +68,7 @@ class ModelHandler:
         # 2. Load ControlNets
         print("Loading ControlNets (InstantID, Zoe, LineArt)...")
 
-        # --- START FIX for config.json error ---
-        # We must load the InstantID ControlNet by loading its pipeline
-        # and "stealing" the controlnet component.
+        # Load InstantID ControlNet by "stealing" it from the base pipeline
         print("Loading InstantID pipeline to extract ControlNet...")
         temp_pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
             Config.INSTANTID_REPO,
@@ -79,7 +77,6 @@ class ModelHandler:
         cn_instantid = temp_pipe.controlnet
         del temp_pipe # Free memory
         print("  [OK] Extracted InstantID ControlNet.")
-        # --- END FIX ---
         
         # Load other ControlNets normally
         cn_zoe = ControlNetModel.from_pretrained(Config.CN_ZOE_REPO, torch_dtype=Config.DTYPE)
@@ -87,14 +84,17 @@ class ModelHandler:
 
         # 3. Load SDXL Pipeline
         print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
-        self.pipeline = StableDiffusionXLControlNetPipeline.from_single_file(
+        
+        # --- START FIX ---
+        # Use the custom Img2Img pipeline class you provided
+        self.pipeline = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
+        # --- END FIX ---
             Config.REPO_ID,
             filename=Config.CHECKPOINT_FILENAME,
-            controlnet=[cn_instantid, cn_zoe, cn_lineart], # Pass the *objects*
+            controlnet=[cn_instantid, cn_zoe, cn_lineart], # Pass the list of all ControlNets
             torch_dtype=Config.DTYPE,
             use_safetensors=True
         )
-        # Move pipeline (and all its components) to the device
         self.pipeline.to(Config.DEVICE)
 
         # 4. Set Scheduler
@@ -102,7 +102,7 @@ class ModelHandler:
 
         # 5. Load Adapters (IP-Adapter & LoRA)
         print("Loading Adapters (IP-Adapter & LoRA)...")
-        # We still load the IP-Adapter separately
+        # The custom pipeline has this method
         self.pipeline.load_ip_adapter_instantid(Config.INSTANTID_REPO)
         self.pipeline.load_lora_weights(Config.REPO_ID, weight_name=Config.LORA_FILENAME)
         self.pipeline.fuse_lora(lora_scale=1.0) 
@@ -121,7 +121,7 @@ class ModelHandler:
             
         try:
             # Convert PIL to CV2
-            cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # <-- Fixed typo
+            cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) 
             faces = self.app.get(cv2_img)
             
             if len(faces) == 0: