Update model.py

model.py

@@ -1,14 +1,17 @@
 import torch
+import os
+import cv2
+import numpy as np
+from config import Config
+
 from diffusers import (
     StableDiffusionXLControlNetPipeline, 
     ControlNetModel, 
-    LCMScheduler,
-    AutoencoderKL
+    LCMScheduler
 )
-from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
 from insightface.app import FaceAnalysis
-from config import Config
-import os
+from controlnet_aux import ZoeDetector, LineartDetector
 
 class ModelHandler:
     def __init__(self):
@@ -16,74 +19,105 @@ class ModelHandler:
         self.app = None # InsightFace
         self.zoe_detector = None
         self.lineart_detector = None
+        self.face_analysis_loaded = False
+
+    def load_face_analysis(self):
+        """
+        Load face analysis model. 
+        Downloads from HF Hub (fast) instead of GitHub (slow) if not present.
+        Forces CPU to avoid ZeroGPU initialization errors.
+        """
+        print("Loading face analysis model...")
+        model_root_path = os.path.join(Config.ANTELOPEV2_ROOT, Config.ANTELOPEV2_NAME)
         
+        # 1. Download from HF Hub (Much faster than default InsightFace download)
+        if not os.path.exists(os.path.join(model_root_path, "scrfd_10g_bnkps.onnx")):
+            print("Downloading AntelopeV2 models from HuggingFace...")
+            try:
+                snapshot_download(
+                    repo_id=Config.ANTELOPEV2_REPO,
+                    local_dir=model_root_path,
+                    local_dir_use_symlinks=False
+                )
+            except Exception as e:
+                print(f"  [ERROR] Failed to download AntelopeV2 models: {e}")
+                return False
+
+        # 2. Initialize InsightFace on CPU
+        # CRITICAL: Use ONLY 'CPUExecutionProvider'. 
+        # ZeroGPU will crash if you try to look for CUDA during init.
+        try:
+            self.app = FaceAnalysis(
+                name=Config.ANTELOPEV2_NAME,
+                root=Config.ANTELOPEV2_ROOT,
+                providers=['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):
-        print("Loading Face Analysis (InsightFace)...")
-        # Ensure you have downloaded antelopev2 to the local folder or handle download here
-        self.app = FaceAnalysis(
-            name='antelopev2', 
-            root='./', 
-            providers=['CPUExecutionProvider'] 
-        )
-        self.app.prepare(ctx_id=0, det_size=(640, 640))
+        # 1. Load Face Analysis
+        self.face_analysis_loaded = self.load_face_analysis()
 
-        print("Loading ControlNets...")
-        # 1. InstantID ControlNet
+        # 2. Load ControlNets
+        print("Loading ControlNets (InstantID, Zoe, LineArt)...")
         cn_instantid = ControlNetModel.from_pretrained(
             Config.CN_INSTANTID_REPO, 
             subfolder="controlnet", 
             torch_dtype=Config.DTYPE
         )
-        # 2. Zoe Depth ControlNet
-        cn_zoe = ControlNetModel.from_pretrained(
-            Config.CN_ZOE_REPO, 
-            torch_dtype=Config.DTYPE
-        )
-        # 3. LineArt ControlNet
-        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)
+        cn_lineart = ControlNetModel.from_pretrained(Config.CN_LINEART_REPO, torch_dtype=Config.DTYPE)
 
-        print("Loading SDXL Pipeline...")
-        # Load custom checkpoint (Single File)
-        ckpt_path = hf_hub_download(repo_id=Config.REPO_ID, filename=Config.CHECKPOINT_FILENAME)
-        
+        # 3. Load SDXL Pipeline
+        print(f"Loading SDXL Pipeline ({Config.CHECKPOINT_FILENAME})...")
         self.pipeline = StableDiffusionXLControlNetPipeline.from_single_file(
-            ckpt_path,
+            Config.REPO_ID,
+            filename=Config.CHECKPOINT_FILENAME,
             controlnet=[cn_instantid, cn_zoe, cn_lineart], # ORDER MATTERS
             torch_dtype=Config.DTYPE,
+            use_safetensors=True
         ).to(Config.DEVICE)
 
-        # Load LCM Scheduler
+        # 4. Set Scheduler
         self.pipeline.scheduler = LCMScheduler.from_config(self.pipeline.scheduler.config)
 
+        # 5. Load Adapters (IP-Adapter & LoRA)
         print("Loading Adapters (IP-Adapter & LoRA)...")
-        # Load InstantID IP-Adapter
         self.pipeline.load_ip_adapter_instantid(Config.INSTANTID_REPO)
-        
-        # Load Custom Style LoRA
         self.pipeline.load_lora_weights(Config.REPO_ID, weight_name=Config.LORA_FILENAME)
-        self.pipeline.fuse_lora() # Optional: Fuse for speed, might make switching harder if you have multiple styles
+        self.pipeline.fuse_lora(lora_scale=1.0) # Fuse with scale 1.0
 
-        print("Loading Preprocessors...")
-        from controlnet_aux import ZoeDetector, LineartDetector
-        self.zoe_detector = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-        self.lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
-        
-        print("Models Loaded Successfully.")
+        # 6. Load Preprocessors
+        print("Loading Preprocessors (Zoe, LineArt)...")
+        self.zoe_detector = ZoeDetector.from_pretrained(Config.ANNOTATOR_REPO)
+        self.lineart_detector = LineartDetector.from_pretrained(Config.ANNOTATOR_REPO)
         
+        print("--- All models loaded successfully ---")
+
     def get_face_embedding(self, image):
-        # Convert PIL to CV2
-        cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-        
-        faces = self.app.get(cv2_img)
-        
-        if len(faces) == 0:
-            return None # Return None instead of crashing
+        """Extracts face embedding, returns None if no face is found."""
+        if not self.face_analysis_loaded:
+            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
-        return torch.tensor(faces[0].normed_embedding).unsqueeze(0)
+        try:
+            # Convert PIL to CV2
+            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
+            return torch.tensor(faces[0].normed_embedding).unsqueeze(0)
+        except Exception as e:
+            print(f"Face embedding extraction failed: {e}")
+            return None