Update model.py

model.py

@@ -16,16 +16,17 @@ from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInst
 from huggingface_hub import snapshot_download, hf_hub_download
 from insightface.app import FaceAnalysis
 # --- MODIFIED: Import new detectors ---
-from controlnet_aux import LeresDetector, LineartAnimeDetector
+from controlnet_aux import LeresDetector, LineartAnimeDetector, ColorDetector
 # --- END MODIFIED ---
 
 class ModelHandler:
     def __init__(self):
         self.pipeline = None
         self.app = None # InsightFace
-        # --- MODIFIED: Rename detectors ---
+        # --- MODIFIED: Add new detector ---
         self.leres_detector = None
         self.lineart_anime_detector = None
+        self.color_detector = None
         # --- END MODIFIED ---
         self.face_analysis_loaded = False
 
@@ -72,7 +73,8 @@ class ModelHandler:
         self.face_analysis_loaded = self.load_face_analysis()
 
         # 2. Load ControlNets
-        print("Loading ControlNets (InstantID, Zoe, LineArt)...")
+        # --- MODIFIED: Updated print ---
+        print("Loading ControlNets (InstantID, Zoe, LineArt, Color)...")
 
         # Load the InstantID ControlNet from the correct subfolder
         print("Loading InstantID ControlNet from subfolder 'ControlNetModel'...")
@@ -84,13 +86,17 @@ class ModelHandler:
         print("  [OK] Loaded InstantID ControlNet.")
         
         # Load other ControlNets normally
-        print("Loading Zoe and LineArt ControlNets...")
+        # --- MODIFIED: Load Color CN ---
+        print("Loading Zoe, LineArt, and Color 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)
+        # --- END MODIFIED ---
 
         # --- Manually wrap the list of models in a MultiControlNetModel ---
         print("Wrapping ControlNets in MultiControlNetModel...")
-        controlnet_list = [cn_instantid, cn_zoe, cn_lineart]
+        # --- MODIFIED: Add Color CN to list ---
+        controlnet_list = [cn_instantid, cn_zoe, cn_lineart, cn_color]
         controlnet = MultiControlNetModel(controlnet_list)
         # --- End wrapping ---
 
@@ -119,6 +125,14 @@ class ModelHandler:
         
         self.pipeline.to(Config.DEVICE)
 
+        # --- NEW: Enable xFormers ---
+        try:
+            self.pipeline.enable_xformers_memory_efficient_attention()
+            print("  [OK] xFormers memory efficient attention enabled.")
+        except Exception as e:
+            print(f"  [WARNING] Failed to enable xFormers: {e}")
+        # --- END NEW ---
+
         # 4. Set Scheduler
         self.pipeline.scheduler = LCMScheduler.from_config(self.pipeline.scheduler.config)
 
@@ -150,18 +164,13 @@ class ModelHandler:
         print("  [OK] LoRA fused.")
         
         # --- DISABLED torch.compile due to runtime errors ---
-        # try:
-        #     print("Compiling UNet with torch.compile...")
-        #     self.pipeline.unet = torch.compile(self.pipeline.unet, mode="reduce-overhead", fullgraph=True)
-        #     print("  [OK] UNet compiled.")
-        # except Exception as e:
-        #     print(f"  [WARNING] torch.compile failed: {e}. Running without compilation.")
         
         # 6. Load Preprocessors
         # --- MODIFIED: Load new detectors ---
-        print("Loading Preprocessors (LeReS, LineArtAnime)...")
+        print("Loading Preprocessors (LeReS, LineArtAnime, Color)...")
         self.leres_detector = LeresDetector.from_pretrained(Config.ANNOTATOR_REPO)
         self.lineart_anime_detector = LineartAnimeDetector.from_pretrained(Config.ANNOTATOR_REPO)
+        self.color_detector = ColorDetector()
         # --- END MODIFIED ---
         
         print("--- All models loaded successfully ---")
@@ -187,4 +196,5 @@ class ModelHandler:
             return torch.tensor(faces[0].normed_embedding).unsqueeze(0)
         except Exception as e:
             print(f"Face embedding extraction failed: {e}")
-            return None
+            return None
+}