Update model.py

model.py

@@ -6,7 +6,8 @@ from config import Config
 
 from diffusers import (
     ControlNetModel, 
-    LCMScheduler
+    LCMScheduler,
+    AutoencoderKL  # <-- ADDED
 )
 from diffusers.models.controlnets.multicontrolnet import MultiControlNetModel
 
@@ -85,6 +86,15 @@ class ModelHandler:
         controlnet_list = [cn_instantid, cn_zoe, cn_lineart]
         controlnet = MultiControlNetModel(controlnet_list)
         # --- End wrapping ---
+        
+        # --- ADDED: Load fp16-safe VAE ---
+        print("Loading fp16-safe VAE (sdxl-vae-fp16-fix)...")
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix", 
+            torch_dtype=Config.DTYPE
+        )
+        print("  [OK] VAE loaded.")
+        # --- END ADDED ---
 
         # 3. Load SDXL Pipeline
         print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
@@ -103,6 +113,7 @@ class ModelHandler:
         self.pipeline = StableDiffusionXLInstantIDImg2ImgPipeline.from_single_file(
             checkpoint_local_path,
             controlnet=controlnet,
+            vae=vae, # <-- MODIFIED: Pass the safe VAE
             torch_dtype=Config.DTYPE,
             use_safetensors=True
         )
@@ -165,16 +176,14 @@ class ModelHandler:
             return None
             
         try:
-            cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGBBGR) 
+            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
-            # --- MODIFIED: Fixed typo ---
             faces = sorted(faces, key=lambda x: (x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]), reverse=True)
-            # --- END MODIFIED ---
             
             # Return the largest face info
             return faces[0]