Update model.py

model.py

@@ -1,83 +1,20 @@
 import torch
 import os
-import cv2
-import numpy as np
 from config import Config
 
 from diffusers import (
-    ControlNetModel, 
-    TCDScheduler, 
+    StableDiffusionXLPipeline,
+    LCMScheduler
 )
-from diffusers.models.controlnets.multicontrolnet import MultiControlNetModel
 
-# Import the custom pipeline from your local file
-from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline
-
-from huggingface_hub import snapshot_download, hf_hub_download
-from insightface.app import FaceAnalysis
-from controlnet_aux import LeresDetector, LineartAnimeDetector
+from huggingface_hub import hf_hub_download
 
 class ModelHandler:
     def __init__(self):
         self.pipeline = None
-        self.app = None # InsightFace
-        self.leres_detector = None
-        self.lineart_anime_detector = None
-        self.face_analysis_loaded = False
-
-    def load_face_analysis(self):
-        """
-        Load face analysis model. 
-        Downloads from HF Hub to the path insightface expects.
-        """
-        print("Loading face analysis model...")
-        
-        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")):
-            print(f"Downloading AntelopeV2 models from {Config.ANTELOPEV2_REPO} to {model_path}...")
-            try:
-                snapshot_download(
-                    repo_id=Config.ANTELOPEV2_REPO,
-                    local_dir=model_path, # Download to the correct expected path
-                )
-            except Exception as e:
-                print(f"  [ERROR] Failed to download AntelopeV2 models: {e}")
-                return False
-
-        try:
-            self.app = FaceAnalysis(
-                name=Config.ANTELOPEV2_NAME,
-                root=Config.ANTELOPEV2_ROOT, 
-                providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 
-            )
-            self.app.prepare(ctx_id=0, det_size=(640, 640))
-            print(f"  [OK] Face analysis model loaded successfully.")
-            return True
-            
-        except Exception as e:
-            print(f"  [WARNING] Face detection system failed to initialize: {e}") 
-            return False
 
     def load_models(self):
-        # 1. Load Face Analysis
-        self.face_analysis_loaded = self.load_face_analysis()
-
-        # 2. Load ControlNets
-        print("Loading ControlNets (InstantID, Zoe, LineArt)...")
-        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)
-
-        print("Wrapping ControlNets in MultiControlNetModel...")
-        controlnet_list = [cn_instantid, cn_zoe, cn_lineart]
-        controlnet = MultiControlNetModel(controlnet_list)
-        
-        # 3. Load SDXL Pipeline (Now from 'reality.safetensors')
+        # 1. Load SDXL Text-to-Image Pipeline
         print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
         
         checkpoint_local_path = os.path.join("./models", Config.CHECKPOINT_FILENAME)
@@ -91,76 +28,36 @@ class ModelHandler:
             )
         
         print(f"Loading pipeline from local file: {checkpoint_local_path}")
-        self.pipeline = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
+        # Use standard SDXL Text2Image pipeline
+        self.pipeline = StableDiffusionXLPipeline.from_single_file(
             checkpoint_local_path,
-            controlnet=controlnet,
             torch_dtype=Config.DTYPE,
             use_safetensors=True
         )
         
         self.pipeline.to(Config.DEVICE)
 
+        # 2. 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 (Sanitized Config)
-        print("Configuring TCDScheduler...")
-        self.pipeline.scheduler = TCDScheduler.from_config(self.pipeline.scheduler.config)
-        print("  [OK] TCDScheduler loaded (Forced SDXL Defaults + Karras + Trailing).")
-
-        # 5. Load Adapters
-        print("Loading Adapters...")
-
-        # 5b. Load and Fuse Style LoRA (lucasart)
-        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):
-            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)
-        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*
-        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):
-            hf_hub_download(
-                repo_id=Config.INSTANTID_REPO,
-                filename=ip_adapter_filename,
-                local_dir="./models",
-                local_dir_use_symlinks=False
-            )
-        self.pipeline.load_ip_adapter_instantid(ip_adapter_local_path)
-        print("  [OK] IP-Adapter loaded.")
+        # 3. Set Scheduler (LCM)
+        print("Configuring LCMScheduler...")
+        scheduler_config = self.pipeline.scheduler.config
+        # Disable clipping to prevent NaN artifacts with LCM
+        scheduler_config['clip_sample'] = False 
+        self.pipeline.scheduler = LCMScheduler.from_config(scheduler_config)
+        print("  [OK] LCMScheduler loaded (clip_sample=False).")
+
+        # 4. Load LoRA
+        print("Loading LoRA weights...")
+        self.pipeline.load_lora_weights(Config.REPO_ID, weight_name=Config.LORA_FILENAME)
         
-        # --- END FIX ---
-
-        # 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)
+        print(f"Fusing LoRA with scale {Config.LORA_STRENGTH}...")
+        self.pipeline.fuse_lora(lora_scale=Config.LORA_STRENGTH)
+        print("  [OK] LoRA fused.")
         
-        print("--- All models loaded successfully ---")
-
-    def get_face_info(self, image):
-        """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
-            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
+        print("--- All models loaded successfully ---")