Update SwitcherAI/processors/frame/modules/face_swapper.py

SwitcherAI/processors/frame/modules/face_swapper.py

@@ -13,17 +13,16 @@ from SwitcherAI.core import update_status
 from SwitcherAI.face_analyser import get_one_face, get_many_faces, find_similar_faces
 from SwitcherAI.face_reference import get_face_reference, set_face_reference
 from SwitcherAI.typing import Face, Frame
-from SwitcherAI.utilities import conditional_download, resolve_relative_path, is_image, is_video
+from SwitcherAI.utilities import resolve_relative_path, is_image, is_video
 
 FRAME_PROCESSOR = None
 EMBEDDING_CONVERTER = None
 THREAD_LOCK = threading.Lock()
 NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_SWAPPER'
 
-# Model configurations with working public URLs
+# Model configurations - local paths only
 MODEL_CONFIGS = {
     'inswapper_128': {
-        'url': 'https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx',
         'path': '../.assets/models/inswapper_128.onnx',
         'type': 'inswapper',
         'size': (128, 128),
@@ -32,7 +31,6 @@ MODEL_CONFIGS = {
         'requires_converter': False
     },
     'inswapper_128_fp16': {
-        'url': 'https://huggingface.co/ninjawick/webui-faceswap-unlocked/resolve/main/inswapper_128_fp16.onnx',
         'path': '../.assets/models/inswapper_128_fp16.onnx',
         'type': 'inswapper',
         'size': (128, 128),
@@ -40,6 +38,15 @@ MODEL_CONFIGS = {
         'standard_deviation': [1.0, 1.0, 1.0],
         'requires_converter': False
     },
+    'simswap_256': {
+        'path': '../.assets/models/simswap_256.onnx',
+        'converter_path': '../.assets/models/simswap_256_converter.onnx',
+        'type': 'simswap',
+        'size': (256, 256),
+        'mean': [0.485, 0.456, 0.406],
+        'standard_deviation': [0.229, 0.224, 0.225],
+        'requires_converter': True
+    },
 }
 
 # Default model - can be changed via globals
@@ -80,10 +87,13 @@ def get_frame_processor() -> Any:
                     model_priority.remove(current_model)
             
             for model_name in model_priority:
+                if model_name not in MODEL_CONFIGS:
+                    continue
+                    
                 try:
                     print(f"🔄 Trying to load face swap model: {model_name}")
                     
-                    # Update config to current model being tested
+                    # Get model config
                     temp_config = MODEL_CONFIGS[model_name]
                     model_path = resolve_relative_path(temp_config['path'])
                     
@@ -93,23 +103,14 @@ def get_frame_processor() -> Any:
                     else:
                         model_path_obj = model_path
                     
-                    # Check if model exists, download if needed
+                    # Check if model exists locally
                     if not model_path_obj.exists():
-                        print(f"📥 Downloading {model_name}...")
-                        download_directory_path = resolve_relative_path('../.assets/models')
-                        if isinstance(download_directory_path, str):
-                            download_directory_path = Path(download_directory_path)
-                        download_directory_path.mkdir(parents=True, exist_ok=True)
-                        
-                        download_urls = [temp_config['url']]
-                        if temp_config.get('requires_converter', False):
-                            download_urls.append(temp_config['converter_url'])
-                        
-                        conditional_download(str(download_directory_path), download_urls)
+                        print(f"❌ Model {model_name} not found at: {model_path_obj}")
+                        continue
                     
                     # Verify model file size
                     if model_path_obj.stat().st_size < 1024:  # Less than 1KB indicates corruption
-                        print(f"⚠️ {model_name} appears corrupted, skipping...")
+                        print(f"⚠️ {model_name} appears corrupted (file too small), skipping...")
                         continue
                     
                     # Try to load the model
@@ -125,44 +126,10 @@ def get_frame_processor() -> Any:
                     
                 except Exception as e:
                     print(f"❌ Failed to load {model_name}: {e}")
-                    
-                    # If protobuf parsing failed, delete the corrupted file and try to re-download
-                    if "INVALID_PROTOBUF" in str(e) or "Protobuf parsing failed" in str(e):
-                        print(f"🗑️ Deleting corrupted {model_name} file...")
-                        try:
-                            if model_path_obj.exists():
-                                model_path_obj.unlink()
-                                print(f"📥 Re-downloading {model_name}...")
-                                
-                                download_directory_path = resolve_relative_path('../.assets/models')
-                                if isinstance(download_directory_path, str):
-                                    download_directory_path = Path(download_directory_path)
-                                download_directory_path.mkdir(parents=True, exist_ok=True)
-                                
-                                download_urls = [temp_config['url']]
-                                if temp_config.get('requires_converter', False):
-                                    download_urls.append(temp_config['converter_url'])
-                                
-                                conditional_download(str(download_directory_path), download_urls)
-                                
-                                # Try loading again after re-download
-                                if model_path_obj.exists() and model_path_obj.stat().st_size > 1024:
-                                    print(f"🔄 Retrying {model_name} after re-download...")
-                                    FRAME_PROCESSOR = insightface.model_zoo.get_model(
-                                        str(model_path_obj), 
-                                        providers=SwitcherAI.globals.execution_providers
-                                    )
-                                    # If successful, update the global setting and break
-                                    SwitcherAI.globals.face_swapper_model = model_name
-                                    print(f"✅ Successfully loaded face swap model: {model_name}")
-                                    break
-                        except Exception as retry_error:
-                            print(f"❌ Retry failed for {model_name}: {retry_error}")
-                    
                     continue
             
             if FRAME_PROCESSOR is None:
-                print("❌ All face swap models failed to load")
+                print("❌ All face swap models failed to load. Please ensure models are present in .assets/models folder.")
                 
     return FRAME_PROCESSOR
 
@@ -185,12 +152,11 @@ def get_embedding_converter() -> Any:
                 else:
                     converter_path_obj = converter_path
                 
-                # Check if converter exists
+                # Check if converter exists locally
                 if not converter_path_obj.exists():
-                    print(f"⚠️ Embedding converter not found at: {converter_path_obj}")
-                    if not pre_check():
-                        print("❌ Failed to download embedding converter")
-                        return None
+                    print(f"❌ Embedding converter not found at: {converter_path_obj}")
+                    print("Please ensure the converter model is present in .assets/models folder.")
+                    return None
                 
                 EMBEDDING_CONVERTER = insightface.model_zoo.get_model(
                     str(converter_path_obj), 
@@ -199,7 +165,7 @@ def get_embedding_converter() -> Any:
                 print("✅ Embedding converter initialized")
                 
             except Exception as e:
-                print(f"⚠️ Failed to initialize embedding converter: {e}")
+                print(f"❌ Failed to initialize embedding converter: {e}")
                 EMBEDDING_CONVERTER = None
                 
     return EMBEDDING_CONVERTER
@@ -213,24 +179,36 @@ def clear_frame_processor() -> None:
 
 
 def pre_check() -> bool:
+    """Check if required models exist locally"""
     try:
-        download_directory_path = resolve_relative_path('../.assets/models')
-        
-        # Ensure download directory exists
-        if isinstance(download_directory_path, str):
-            download_directory_path = Path(download_directory_path)
-        download_directory_path.mkdir(parents=True, exist_ok=True)
-        
         config = get_current_model_config()
         
-        # Download main model
-        download_urls = [config['url']]
+        # Check main model path
+        model_path = resolve_relative_path(config['path'])
+        if isinstance(model_path, str):
+            model_path_obj = Path(model_path)
+        else:
+            model_path_obj = model_path
+            
+        if not model_path_obj.exists():
+            print(f"❌ Main model not found at: {model_path_obj}")
+            print("Please ensure the model file is present in .assets/models folder.")
+            return False
         
-        # Download converter if needed
+        # Check converter if needed
         if config.get('requires_converter', False):
-            download_urls.append(config['converter_url'])
+            converter_path = resolve_relative_path(config['converter_path'])
+            if isinstance(converter_path, str):
+                converter_path_obj = Path(converter_path)
+            else:
+                converter_path_obj = converter_path
+                
+            if not converter_path_obj.exists():
+                print(f"❌ Converter model not found at: {converter_path_obj}")
+                print("Please ensure the converter model file is present in .assets/models folder.")
+                return False
         
-        conditional_download(str(download_directory_path), download_urls)
+        print("✅ All required models found locally")
         return True
         
     except Exception as e:
@@ -250,10 +228,15 @@ def pre_process() -> bool:
             update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
             return False
         
+        # Check if required models exist locally
+        if not pre_check():
+            update_status("Required models not found in .assets/models folder", NAME)
+            return False
+        
         # Check if processor is available
         processor = get_frame_processor()
         if processor is None:
-            print("⚠️ Face swap processor not available")
+            update_status("Face swap processor not available", NAME)
             return False
             
         return True