Update model.py

model.py

@@ -15,8 +15,7 @@ from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInst
 
 from huggingface_hub import snapshot_download, hf_hub_download
 from insightface.app import FaceAnalysis
-
-# --- MODIFIED: Removed ColorDetector ---
+# --- MODIFIED: Removed ColorDetector import ---
 from controlnet_aux import LeresDetector, LineartAnimeDetector
 # --- END MODIFIED ---
 
@@ -24,10 +23,8 @@ class ModelHandler:
     def __init__(self):
         self.pipeline = None
         self.app = None # InsightFace
-        # --- MODIFIED: Removed color_detector ---
         self.leres_detector = None
         self.lineart_anime_detector = None
-        # --- END MODIFIED ---
         self.face_analysis_loaded = False
 
     def load_face_analysis(self):
@@ -38,8 +35,6 @@ class ModelHandler:
         """
         print("Loading face analysis model...")
         
-        # insightface expects models in '{root}/models/{name}'
-        # Since our root='.' and name='antelopev2', the expected path is './models/antelopev2'
         model_path = os.path.join(Config.ANTELOPEV2_ROOT, "models", Config.ANTELOPEV2_NAME)
 
         if not os.path.exists(os.path.join(model_path, "scrfd_10g_bnkps.onnx")):
@@ -54,11 +49,10 @@ class ModelHandler:
                 return False
 
         try:
-            # Initialize with root='.' and name='antelopev2'
             self.app = FaceAnalysis(
                 name=Config.ANTELOPEV2_NAME,
                 root=Config.ANTELOPEV2_ROOT, 
-                providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 
+                providers=['CPUExecutionProvider'] 
             )
             self.app.prepare(ctx_id=0, det_size=(640, 640))
             print(f"  [OK] Face analysis model loaded successfully.")
@@ -73,33 +67,36 @@ class ModelHandler:
         self.face_analysis_loaded = self.load_face_analysis()
 
         # 2. Load ControlNets
-        print("Loading ControlNets (InstantID, Zoe, LineArt, Color)...")
+        # --- MODIFIED: Updated print ---
+        print("Loading ControlNets (InstantID, Zoe, LineArt, Tile)...")
 
         # Load the InstantID ControlNet from the correct subfolder
         print("Loading InstantID ControlNet from subfolder 'ControlNetModel'...")
         cn_instantid = ControlNetModel.from_pretrained(
-            Config.INSTANTID_REPO,   # "InstantX/InstantID"
-            subfolder="ControlNetModel", # Correct casing
+            Config.INSTANTID_REPO,
+            subfolder="ControlNetModel",
             torch_dtype=Config.DTYPE
         )
         print("  [OK] Loaded InstantID ControlNet.")
         
         # Load other ControlNets normally
-        print("Loading Zoe, LineArt, and Color ControlNets...")
+        # --- MODIFIED: Load Tile CN ---
+        print("Loading Zoe, LineArt, and Tile ControlNets...")
         cn_zoe = ControlNetModel.from_pretrained(Config.CN_ZOE_REPO, torch_dtype=Config.DTYPE)
         cn_lineart = ControlNetModel.from_pretrained(Config.CN_LINEART_REPO, torch_dtype=Config.DTYPE)
-        cn_color = ControlNetModel.from_pretrained(Config.CN_COLOR_REPO, torch_dtype=Config.DTYPE) # This is still correct
+        cn_tile = ControlNetModel.from_pretrained(Config.CN_TILE_REPO, torch_dtype=Config.DTYPE)
+        # --- END MODIFIED ---
 
         # --- Manually wrap the list of models in a MultiControlNetModel ---
         print("Wrapping ControlNets in MultiControlNetModel...")
-        controlnet_list = [cn_instantid, cn_zoe, cn_lineart, cn_color] # This is still correct
+        # --- MODIFIED: Add Tile CN to list ---
+        controlnet_list = [cn_instantid, cn_zoe, cn_lineart, cn_tile]
         controlnet = MultiControlNetModel(controlnet_list)
         # --- End wrapping ---
 
         # 3. Load SDXL Pipeline
         print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
         
-        # Manually download the checkpoint file first.
         checkpoint_local_path = os.path.join("./models", Config.CHECKPOINT_FILENAME)
         if not os.path.exists(checkpoint_local_path):
             print(f"Downloading checkpoint to {checkpoint_local_path}...")
@@ -110,11 +107,10 @@ class ModelHandler:
                 local_dir_use_symlinks=False
             )
         
-        # Use the custom Img2Img pipeline class you provided, loading from the LOCAL FILE
         print(f"Loading pipeline from local file: {checkpoint_local_path}")
         self.pipeline = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
-            checkpoint_local_path, # Pass the local path
-            controlnet=controlnet, # Pass the single, wrapped object
+            checkpoint_local_path,
+            controlnet=controlnet,
             torch_dtype=Config.DTYPE,
             use_safetensors=True
         )
@@ -135,7 +131,6 @@ class ModelHandler:
         # 5. Load Adapters (IP-Adapter & LoRA)
         print("Loading Adapters (IP-Adapter & LoRA)...")
         
-        # Download the ip-adapter.bin file and pass its local path
         ip_adapter_filename = "ip-adapter.bin"
         ip_adapter_local_path = os.path.join("./models", ip_adapter_filename)
 
@@ -149,12 +144,11 @@ class ModelHandler:
             )
         
         print(f"Loading IP-Adapter from local file: {ip_adapter_local_path}")
-        self.pipeline.load_ip_adapter_instantid(ip_adapter_local_path) # Pass local path
+        self.pipeline.load_ip_adapter_instantid(ip_adapter_local_path)
         
         print("Loading LoRA weights...")
         self.pipeline.load_lora_weights(Config.REPO_ID, weight_name=Config.LORA_FILENAME)
         
-        # --- NEW: Fuse LoRA at build time with fixed strength ---
         print(f"Fusing LoRA with scale {Config.LORA_STRENGTH}...")
         self.pipeline.fuse_lora(lora_scale=Config.LORA_STRENGTH)
         print("  [OK] LoRA fused.")
@@ -174,17 +168,14 @@ class ModelHandler:
             return None
             
         try:
-            # Convert PIL to CV2
             cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) 
             faces = self.app.get(cv2_img)
             
             if len(faces) == 0:
                 return None
                 
-            # Sort by size (width * height) to find the main character
             faces = sorted(faces, key=lambda x: (x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]), reverse=True)
             
-            # Return the largest face
             return torch.tensor(faces[0].normed_embedding).unsqueeze(0)
         except Exception as e:
             print(f"Face embedding extraction failed: {e}")