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

SwitcherAI/processors/frame/modules/face_swapper.py

@@ -79,43 +79,67 @@ def get_frame_processor() -> Any:
 
     with THREAD_LOCK:
         if FRAME_PROCESSOR is None:
-            try:
-                config = get_current_model_config()
-                model_path = resolve_relative_path(config['path'])
-                
-                # Convert to Path object if it's a string for validation
-                if isinstance(model_path, str):
-                    model_path_obj = Path(model_path)
-                else:
-                    model_path_obj = model_path
-                
-                # Check if model exists
-                if not model_path_obj.exists():
-                    print(f"⚠️ Face swap model not found at: {model_path_obj}")
-                    print("🔄 Attempting to download model...")
-                    if not pre_check():
-                        print("❌ Failed to download face swap model")
-                        return None
-                
-                # Verify model file size
-                if model_path_obj.stat().st_size < 1024:  # Less than 1KB indicates corruption
-                    print(f"⚠️ Face swap model appears corrupted: {model_path_obj}")
-                    print("🔄 Attempting to re-download model...")
-                    model_path_obj.unlink()  # Delete corrupted file
-                    if not pre_check():
-                        print("❌ Failed to re-download face swap model")
-                        return None
-                
-                # Load model with string path (insightface expects string)
-                FRAME_PROCESSOR = insightface.model_zoo.get_model(
-                    str(model_path_obj), 
-                    providers=SwitcherAI.globals.execution_providers
-                )
-                print("✅ Face swap processor initialized")
-                
-            except Exception as e:
-                print(f"⚠️ Failed to initialize face swap processor: {e}")
-                FRAME_PROCESSOR = None
+            # Try models in order of preference
+            model_priority = ['inswapper_128_fp16', 'inswapper_128', 'simswap_256']
+            
+            # If user set a specific model, try it first
+            current_model = getattr(SwitcherAI.globals, 'face_swapper_model', None)
+            if current_model and current_model in MODEL_CONFIGS:
+                model_priority.insert(0, current_model)
+                # Remove duplicate if it exists later in the list
+                if current_model in model_priority[1:]:
+                    model_priority.remove(current_model)
+            
+            for model_name in model_priority:
+                try:
+                    print(f"🔄 Trying to load face swap model: {model_name}")
+                    
+                    # Update config to current model being tested
+                    temp_config = MODEL_CONFIGS[model_name]
+                    model_path = resolve_relative_path(temp_config['path'])
+                    
+                    # Convert to Path object if it's a string for validation
+                    if isinstance(model_path, str):
+                        model_path_obj = Path(model_path)
+                    else:
+                        model_path_obj = model_path
+                    
+                    # Check if model exists, download if needed
+                    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)
+                    
+                    # Verify model file size
+                    if model_path_obj.stat().st_size < 1024:  # Less than 1KB indicates corruption
+                        print(f"⚠️ {model_name} appears corrupted, skipping...")
+                        continue
+                    
+                    # Try to load the model
+                    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 e:
+                    print(f"❌ Failed to load {model_name}: {e}")
+                    continue
+            
+            if FRAME_PROCESSOR is None:
+                print("❌ All face swap models failed to load")
                 
     return FRAME_PROCESSOR