Update model.py

model.py

@@ -6,8 +6,7 @@ from config import Config
 
 from diffusers import (
     ControlNetModel, 
-    TCDScheduler,
-    AutoencoderKL  # <-- ADDED: Import AutoencoderKL
+    TCDScheduler, 
 )
 from diffusers.models.controlnets.multicontrolnet import MultiControlNetModel
 
@@ -16,15 +15,17 @@ from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInst
 
 from huggingface_hub import snapshot_download, hf_hub_download
 from insightface.app import FaceAnalysis
-from controlnet_aux import LeresDetector, LineartAnimeDetector
+from controlnet_aux import LeresDetector, LineartAnimeDetector, CannyDetector
 
 class ModelHandler:
     def __init__(self):
         self.pipeline = None
-        self.app = None # InsightFace
+        self.app = None  # InsightFace
         self.leres_detector = None
         self.lineart_anime_detector = None
+        self.canny_detector = None
         self.face_analysis_loaded = False
+        self.edge_type = Config.DEFAULT_EDGE_TYPE
 
     def load_face_analysis(self):
         """
@@ -40,7 +41,7 @@ class ModelHandler:
             try:
                 snapshot_download(
                     repo_id=Config.ANTELOPEV2_REPO,
-                    local_dir=model_path, # Download to the correct expected path
+                    local_dir=model_path,
                 )
             except Exception as e:
                 print(f"  [ERROR] Failed to download AntelopeV2 models: {e}")
@@ -60,35 +61,65 @@ class ModelHandler:
             print(f"  [WARNING] Face detection system failed to initialize: {e}") 
             return False
 
-    def load_models(self):
+    def load_models(self, edge_type="canny"):
+        """
+        Load all models with support for different edge detection types.
+        
+        Args:
+            edge_type: "canny", "lineart", or "both"
+        """
+        self.edge_type = edge_type
+        
         # 1. Load Face Analysis
         self.face_analysis_loaded = self.load_face_analysis()
 
-        # 2. Load ControlNets
-        print("Loading ControlNets (InstantID, Zoe, LineArt)...")
+        # 2. Load ControlNets based on edge_type
+        print(f"Loading ControlNets (InstantID, Zoe, {edge_type.upper()})...")
         cn_instantid = ControlNetModel.from_pretrained(
             Config.INSTANTID_REPO,
             subfolder="ControlNetModel",
             torch_dtype=Config.DTYPE
         )
-        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_zoe = ControlNetModel.from_pretrained(
+            Config.CN_ZOE_REPO, 
+            torch_dtype=Config.DTYPE
+        )
+        
+        # Load edge ControlNet(s)
+        controlnet_list = [cn_instantid, cn_zoe]
+        
+        if edge_type == "canny":
+            cn_canny = ControlNetModel.from_pretrained(
+                Config.CN_CANNY_REPO, 
+                torch_dtype=Config.DTYPE
+            )
+            controlnet_list.append(cn_canny)
+            print("  [OK] Loaded Canny ControlNet")
+            
+        elif edge_type == "lineart":
+            cn_lineart = ControlNetModel.from_pretrained(
+                Config.CN_LINEART_REPO, 
+                torch_dtype=Config.DTYPE
+            )
+            controlnet_list.append(cn_lineart)
+            print("  [OK] Loaded LineArt ControlNet")
+            
+        elif edge_type == "both":
+            cn_canny = ControlNetModel.from_pretrained(
+                Config.CN_CANNY_REPO, 
+                torch_dtype=Config.DTYPE
+            )
+            cn_lineart = ControlNetModel.from_pretrained(
+                Config.CN_LINEART_REPO, 
+                torch_dtype=Config.DTYPE
+            )
+            controlnet_list.extend([cn_canny, cn_lineart])
+            print("  [OK] Loaded both Canny and LineArt ControlNets")
 
         print("Wrapping ControlNets in MultiControlNetModel...")
-        controlnet_list = [cn_instantid, cn_zoe, cn_lineart]
         controlnet = MultiControlNetModel(controlnet_list)
         
-        # --- ADDED: Load FP16 Fixed VAE ---
-        # This prevents NaN errors and black images in SDXL fp16 mode
-        vae_repo = getattr(Config, "VAE_REPO", "madebyollin/sdxl-vae-fp16-fix")
-        print(f"Loading VAE ({vae_repo})...")
-        vae = AutoencoderKL.from_pretrained(
-            vae_repo,
-            torch_dtype=Config.DTYPE
-        )
-        # ----------------------------------
-
-        # 3. Load SDXL Pipeline (Now from 'reality.safetensors')
+        # 3. Load SDXL Pipeline
         print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
         
         checkpoint_local_path = os.path.join("./models", Config.CHECKPOINT_FILENAME)
@@ -104,7 +135,6 @@ class ModelHandler:
         print(f"Loading pipeline from local file: {checkpoint_local_path}")
         self.pipeline = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
             checkpoint_local_path,
-            vae=vae, # <-- ADDED: Inject the fixed VAE here
             controlnet=controlnet,
             torch_dtype=Config.DTYPE,
             use_safetensors=True
@@ -118,15 +148,15 @@ class ModelHandler:
         except Exception as e:
             print(f"  [WARNING] Failed to enable xFormers: {e}")
 
-        # 4. Set TCD Scheduler (Sanitized Config)
+        # 4. Set TCD Scheduler
         print("Configuring TCDScheduler...")
         self.pipeline.scheduler = TCDScheduler.from_config(self.pipeline.scheduler.config)
-        print("  [OK] TCDScheduler loaded (Forced SDXL Defaults + Karras + Trailing).")
+        print("  [OK] TCDScheduler loaded.")
 
         # 5. Load Adapters
         print("Loading Adapters...")
 
-        # 5b. Load and Fuse Style LoRA (lucasart/retroart)
+        # 5a. Load and Fuse Style LoRA
         print(f"Loading and Fusing Style LoRA ({Config.LORA_FILENAME})...")
         style_lora_path = os.path.join("./models", Config.LORA_FILENAME)
         if not os.path.exists(style_lora_path):
@@ -140,7 +170,7 @@ class ModelHandler:
         self.pipeline.fuse_lora(lora_scale=Config.LORA_STRENGTH)
         print("  [OK] Style LoRA fused.")
 
-        # 5c. Load IP-Adapter (for InstantID) - *Must be loaded AFTER fusing*
+        # 5b. Load IP-Adapter for InstantID
         ip_adapter_filename = "ip-adapter.bin"
         ip_adapter_local_path = os.path.join("./models", ip_adapter_filename)
         if not os.path.exists(ip_adapter_local_path):
@@ -151,12 +181,19 @@ class ModelHandler:
                 local_dir_use_symlinks=False
             )
         self.pipeline.load_ip_adapter_instantid(ip_adapter_local_path)
-        print("  [OK] IP-Adapter loaded.")
+        print("  [OK] InstantID IP-Adapter loaded.")
 
-        # 7. Load Preprocessors
-        print("Loading Preprocessors (LeReS, LineArtAnime)...")
+        # 6. Load Preprocessors
+        print("Loading Preprocessors...")
         self.leres_detector = LeresDetector.from_pretrained(Config.ANNOTATOR_REPO)
-        self.lineart_anime_detector = LineartAnimeDetector.from_pretrained(Config.ANNOTATOR_REPO)
+        
+        if edge_type in ["canny", "both"]:
+            self.canny_detector = CannyDetector()
+            print("  [OK] Canny detector loaded")
+            
+        if edge_type in ["lineart", "both"]:
+            self.lineart_anime_detector = LineartAnimeDetector.from_pretrained(Config.ANNOTATOR_REPO)
+            print("  [OK] LineArt detector loaded")
         
         print("--- All models loaded successfully ---")
 
@@ -169,8 +206,28 @@ class ModelHandler:
             faces = self.app.get(cv2_img)
             if len(faces) == 0:
                 return None
-            faces = sorted(faces, key=lambda x: (x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]), reverse=True)
+            faces = sorted(
+                faces, 
+                key=lambda x: (x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]), 
+                reverse=True
+            )
             return faces[0]
         except Exception as e:
             print(f"Face embedding extraction failed: {e}")
-            return None
+            return None
+    
+    def extract_depth(self, image):
+        """Extract depth map using LeReS detector"""
+        return self.leres_detector(image)
+    
+    def extract_canny(self, image, low_threshold=100, high_threshold=200):
+        """Extract Canny edges"""
+        if self.canny_detector is None:
+            raise ValueError("Canny detector not loaded. Initialize with edge_type='canny' or 'both'")
+        return self.canny_detector(image, low_threshold=low_threshold, high_threshold=high_threshold)
+    
+    def extract_lineart(self, image):
+        """Extract LineArt edges"""
+        if self.lineart_anime_detector is None:
+            raise ValueError("LineArt detector not loaded. Initialize with edge_type='lineart' or 'both'")
+        return self.lineart_anime_detector(image)