Update model.py

model.py

@@ -65,28 +65,19 @@ class ModelHandler:
 
         # 2. Load ControlNets
         print("Loading ControlNets (InstantID, Zoe, LineArt)...")
-
-        # Load the InstantID ControlNet from the correct subfolder
-        print("Loading InstantID ControlNet from subfolder 'ControlNetModel'...")
         cn_instantid = ControlNetModel.from_pretrained(
             Config.INSTANTID_REPO,
             subfolder="ControlNetModel",
             torch_dtype=Config.DTYPE
         )
-        print("  [OK] Loaded InstantID ControlNet.")
-        
-        # Load other ControlNets normally
-        print("Loading Zoe and LineArt 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)
 
-        # --- Manually wrap the list of models in a MultiControlNetModel ---
         print("Wrapping ControlNets in MultiControlNetModel...")
         controlnet_list = [cn_instantid, cn_zoe, cn_lineart]
         controlnet = MultiControlNetModel(controlnet_list)
-        # --- End wrapping ---
         
-        # 3. Load SDXL Pipeline
+        # 3. Load SDXL Pipeline (Now from 'reality.safetensors')
         print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
         
         checkpoint_local_path = os.path.join("./models", Config.CHECKPOINT_FILENAME)
@@ -109,21 +100,17 @@ class ModelHandler:
         
         self.pipeline.to(Config.DEVICE)
 
-        # 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}")
 
-        # 4. Set TCD Scheduler
+        # 4. Set TCD Scheduler (Sanitized Config)
         print("Configuring TCDScheduler...")
         
-        # --- FIX STARTS HERE ---
-        # Convert FrozenDict to a mutable standard Python dict
+        # Force standard SDXL config to prevent noise artifacts
         tcd_config = dict(self.pipeline.scheduler.config)
-        
-        # Now we can update it safely
         tcd_config.update({
             "beta_start": 0.00085,
             "beta_end": 0.012,
@@ -137,34 +124,28 @@ class ModelHandler:
             use_karras_sigmas=True,
             timestep_spacing="trailing" 
         )
-        # --- FIX ENDS HERE ---
-        
         print("  [OK] TCDScheduler loaded (Forced SDXL Defaults + Karras + Trailing).")
 
-        # 5. Load Adapters (IP-Adapter, TCD-LoRA & Style LoRA)
+        # 5. Load Adapters
         print("Loading Adapters...")
         
-        # 5a. IP-Adapter
+        # 5a. 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):
-            print(f"Downloading IP-Adapter to {ip_adapter_local_path}...")
             hf_hub_download(
                 repo_id=Config.INSTANTID_REPO,
                 filename=ip_adapter_filename,
                 local_dir="./models",
                 local_dir_use_symlinks=False
             )
-        
-        print(f"Loading IP-Adapter from local file: {ip_adapter_local_path}")
         self.pipeline.load_ip_adapter_instantid(ip_adapter_local_path)
+        print("  [OK] IP-Adapter loaded.")
         
-        # 5b. Load TCD LoRA (Correct Filename)
+        # 5b. TCD LoRA (for speed)
         print("Loading TCD-SDXL-LoRA...")
         tcd_lora_filename = "pytorch_lora_weights.safetensors"
         tcd_lora_path = os.path.join("./models", tcd_lora_filename)
-        
         if not os.path.exists(tcd_lora_path):
              hf_hub_download(
                 repo_id="h1t/TCD-SDXL-LoRA",
@@ -172,19 +153,28 @@ class ModelHandler:
                 local_dir="./models",
                 local_dir_use_symlinks=False
             )
-        self.pipeline.load_lora_weights("./models", weight_name=tcd_lora_filename)
-        self.pipeline.fuse_lora(lora_scale=1.0) 
-        print("  [OK] TCD LoRA fused.")
+        self.pipeline.load_lora_weights("./models", weight_name=tcd_lora_filename, adapter_name="tcd")
+        print("  [OK] TCD LoRA loaded.")
 
-        # 5c. Load Style LoRA
-        print("Loading Style LoRA weights...")
-        self.pipeline.load_lora_weights(Config.REPO_ID, weight_name=Config.LORA_FILENAME)
-        
-        print(f"Fusing Style LoRA with scale {Config.LORA_STRENGTH}...")
-        self.pipeline.fuse_lora(lora_scale=Config.LORA_STRENGTH)
-        print("  [OK] Style LoRA fused.")
+        # 5c. Style LoRA (lucasart)
+        print(f"Loading Style LoRA ({Config.LORA_FILENAME})...")
+        style_lora_path = os.path.join("./models", Config.LORA_FILENAME)
+        if not os.path.exists(style_lora_path):
+            hf_hub_download(
+                repo_id=Config.REPO_ID,
+                filename=Config.LORA_FILENAME,
+                local_dir="./models",
+                local_dir_use_symlinks=False
+            )
+        self.pipeline.load_lora_weights("./models", weight_name=Config.LORA_FILENAME, adapter_name="style")
+        print("  [OK] Style LoRA loaded.")
         
-        # 6. Load Preprocessors
+        # 6. Set Adapter Weights (TCD + Style)
+        # We set both adapters to run simultaneously
+        print(f"Setting adapter weights: TCD (1.0), Style ({Config.LORA_STRENGTH})")
+        self.pipeline.set_adapters(["tcd", "style"], adapter_weights=[1.0, Config.LORA_STRENGTH])
+
+        # 7. Load Preprocessors
         print("Loading Preprocessors (LeReS, LineArtAnime)...")
         self.leres_detector = LeresDetector.from_pretrained(Config.ANNOTATOR_REPO)
         self.lineart_anime_detector = LineartAnimeDetector.from_pretrained(Config.ANNOTATOR_REPO)
@@ -195,18 +185,12 @@ class ModelHandler:
         """Extracts the largest face, returns insightface result object."""
         if not self.face_analysis_loaded:
             return None
-            
         try:
             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 info
             return faces[0]
         except Exception as e:
             print(f"Face embedding extraction failed: {e}")