Update model.py

model.py

@@ -6,7 +6,8 @@ from config import Config
 
 from diffusers import (
     ControlNetModel, 
-    TCDScheduler, 
+    LCMScheduler,
+    # AutoencoderKL # Removed as requested
 )
 from diffusers.models.controlnets.multicontrolnet import MultiControlNetModel
 
@@ -109,27 +110,26 @@ class ModelHandler:
         
         self.pipeline.to(Config.DEVICE)
 
-        # Enable xFormers
+        # --- 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}")
-
-        # 4. Set TCD Scheduler
-        print("Configuring TCDScheduler...")
-        # --- FIX: Set timestep_spacing="trailing" for proper distilled sampling ---
-        self.pipeline.scheduler = TCDScheduler.from_config(
-            self.pipeline.scheduler.config,
-            use_karras_sigmas=True,
-            timestep_spacing="trailing" 
-        )
-        print("  [OK] TCDScheduler loaded (Karras + Trailing Spacing).")
-
-        # 5. Load Adapters (IP-Adapter, TCD-LoRA & Style LoRA)
-        print("Loading Adapters...")
+        # --- END NEW ---
+
+        # 4. Set Scheduler
+        # --- MODIFIED: Disable clipping to prevent NaN artifacts ---
+        print("Configuring LCMScheduler...")
+        scheduler_config = self.pipeline.scheduler.config
+        scheduler_config['clip_sample'] = False # <-- THIS IS THE FIX
+        self.pipeline.scheduler = LCMScheduler.from_config(scheduler_config)
+        print("  [OK] LCMScheduler loaded (clip_sample=False).")
+        # --- END MODIFIED ---
+
+        # 5. Load Adapters (IP-Adapter & LoRA)
+        print("Loading Adapters (IP-Adapter & LoRA)...")
         
-        # 5a. IP-Adapter
         ip_adapter_filename = "ip-adapter.bin"
         ip_adapter_local_path = os.path.join("./models", ip_adapter_filename)
 
@@ -145,29 +145,12 @@ class ModelHandler:
         print(f"Loading IP-Adapter from local file: {ip_adapter_local_path}")
         self.pipeline.load_ip_adapter_instantid(ip_adapter_local_path)
         
-        # 5b. Load TCD LoRA (Correct filename)
-        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",
-                filename=tcd_lora_filename,
-                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.")
-
-        # 5c. Load Style LoRA
-        print("Loading Style LoRA weights...")
+        print("Loading 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}...")
+        print(f"Fusing LoRA with scale {Config.LORA_STRENGTH}...")
         self.pipeline.fuse_lora(lora_scale=Config.LORA_STRENGTH)
-        print("  [OK] Style LoRA fused.")
+        print("  [OK] LoRA fused.")
         
         # 6. Load Preprocessors
         print("Loading Preprocessors (LeReS, LineArtAnime)...")