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

SwitcherAI/processors/frame/modules/face_swapper.py

@@ -4,6 +4,7 @@ import insightface
 import threading
 import numpy as np
 from functools import lru_cache
+from pathlib import Path
 
 import SwitcherAI.globals
 import SwitcherAI.processors.frame.core as frame_processors
@@ -78,9 +79,44 @@ def get_frame_processor() -> Any:
 
     with THREAD_LOCK:
         if FRAME_PROCESSOR is None:
-            config = get_current_model_config()
-            model_path = resolve_relative_path(config['path'])
-            FRAME_PROCESSOR = insightface.model_zoo.get_model(model_path, providers=SwitcherAI.globals.execution_providers)
+            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
+                
     return FRAME_PROCESSOR
 
 
@@ -93,11 +129,32 @@ def get_embedding_converter() -> Any:
         
     with THREAD_LOCK:
         if EMBEDDING_CONVERTER is None:
-            converter_path = resolve_relative_path(config['converter_path'])
             try:
-                EMBEDDING_CONVERTER = insightface.model_zoo.get_model(converter_path, providers=SwitcherAI.globals.execution_providers)
-            except Exception:
+                converter_path = resolve_relative_path(config['converter_path'])
+                
+                # Convert to Path object if it's a string for validation
+                if isinstance(converter_path, str):
+                    converter_path_obj = Path(converter_path)
+                else:
+                    converter_path_obj = converter_path
+                
+                # Check if converter exists
+                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
+                
+                EMBEDDING_CONVERTER = insightface.model_zoo.get_model(
+                    str(converter_path_obj), 
+                    providers=SwitcherAI.globals.execution_providers
+                )
+                print("✅ Embedding converter initialized")
+                
+            except Exception as e:
+                print(f"⚠️ Failed to initialize embedding converter: {e}")
                 EMBEDDING_CONVERTER = None
+                
     return EMBEDDING_CONVERTER
 
 
@@ -109,31 +166,54 @@ def clear_frame_processor() -> None:
 
 
 def pre_check() -> bool:
-    download_directory_path = resolve_relative_path('../.assets/models')
-    config = get_current_model_config()
-    
-    # Download main model
-    download_urls = [config['url']]
-    
-    # Download converter if needed
-    if config.get('requires_converter', False):
-        download_urls.append(config['converter_url'])
-    
-    conditional_download(download_directory_path, download_urls)
-    return True
+    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']]
+        
+        # Download converter if needed
+        if config.get('requires_converter', False):
+            download_urls.append(config['converter_url'])
+        
+        conditional_download(str(download_directory_path), download_urls)
+        return True
+        
+    except Exception as e:
+        print(f"❌ Face swap pre-check failed: {e}")
+        return False
 
 
 def pre_process() -> bool:
-    if not is_image(SwitcherAI.globals.source_path):
-        update_status(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME)
-        return False
-    elif not get_one_face(cv2.imread(SwitcherAI.globals.source_path)):
-        update_status(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME)
-        return False
-    if not is_image(SwitcherAI.globals.target_path) and not is_video(SwitcherAI.globals.target_path):
-        update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
+    try:
+        if not is_image(SwitcherAI.globals.source_path):
+            update_status(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME)
+            return False
+        elif not get_one_face(cv2.imread(SwitcherAI.globals.source_path)):
+            update_status(wording.get('no_source_face_detected') + wording.get('exclamation_mark'), NAME)
+            return False
+        if not is_image(SwitcherAI.globals.target_path) and not is_video(SwitcherAI.globals.target_path):
+            update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
+            return False
+        
+        # Check if processor is available
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Face swap processor not available")
+            return False
+            
+        return True
+        
+    except Exception as e:
+        print(f"⚠️ Face swap pre-process failed: {e}")
         return False
-    return True
 
 
 def post_process() -> None:
@@ -144,119 +224,165 @@ def post_process() -> None:
 
 def prepare_source_embedding(source_face: Face) -> np.ndarray:
     """Prepare source face embedding based on model type"""
-    config = get_current_model_config()
-    model_type = config['type']
-    
-    if model_type == 'inswapper':
-        # Enhanced embedding preparation for inswapper
-        model_path = resolve_relative_path(config['path'])
-        model_initializer = get_static_model_initializer(model_path)
-        source_embedding = source_face.embedding.reshape((1, -1))
-        source_embedding = np.dot(source_embedding, model_initializer) / np.linalg.norm(source_embedding)
-        return source_embedding
-    elif model_type == 'simswap':
-        # Use embedding converter for simswap
-        converter = get_embedding_converter()
-        if converter is not None:
-            embedding = source_face.embedding.reshape(-1, 512)
-            try:
-                converted_embedding = converter.run(None, {'input': embedding})[0]
-                converted_embedding = converted_embedding.ravel()
-                normed_embedding = converted_embedding / np.linalg.norm(converted_embedding)
-                return normed_embedding.reshape(1, -1)
-            except Exception:
-                pass
-        
-        # Fallback to original embedding
-        return source_face.embedding.reshape(1, -1)
-    else:
-        # Default behavior
+    try:
+        config = get_current_model_config()
+        model_type = config['type']
+        
+        if model_type == 'inswapper':
+            # Enhanced embedding preparation for inswapper
+            model_path = resolve_relative_path(config['path'])
+            model_initializer = get_static_model_initializer(str(model_path))
+            source_embedding = source_face.embedding.reshape((1, -1))
+            source_embedding = np.dot(source_embedding, model_initializer) / np.linalg.norm(source_embedding)
+            return source_embedding
+        elif model_type == 'simswap':
+            # Use embedding converter for simswap
+            converter = get_embedding_converter()
+            if converter is not None:
+                embedding = source_face.embedding.reshape(-1, 512)
+                try:
+                    converted_embedding = converter.run(None, {'input': embedding})[0]
+                    converted_embedding = converted_embedding.ravel()
+                    normed_embedding = converted_embedding / np.linalg.norm(converted_embedding)
+                    return normed_embedding.reshape(1, -1)
+                except Exception:
+                    pass
+            
+            # Fallback to original embedding
+            return source_face.embedding.reshape(1, -1)
+        else:
+            # Default behavior
+            return source_face.embedding.reshape(1, -1)
+            
+    except Exception as e:
+        print(f"⚠️ Error preparing source embedding: {e}")
         return source_face.embedding.reshape(1, -1)
 
 
 def prepare_crop_frame(crop_frame: Frame) -> np.ndarray:
     """Prepare cropped frame for model input with normalization"""
-    config = get_current_model_config()
-    model_mean = config['mean']
-    model_std = config['standard_deviation']
-    
-    # Convert to float and normalize
-    crop_frame = crop_frame[:, :, ::-1] / 255.0
-    crop_frame = (crop_frame - model_mean) / model_std
-    crop_frame = crop_frame.transpose(2, 0, 1)
-    crop_frame = np.expand_dims(crop_frame, axis=0).astype(np.float32)
-    return crop_frame
+    try:
+        config = get_current_model_config()
+        model_mean = config['mean']
+        model_std = config['standard_deviation']
+        
+        # Convert to float and normalize
+        crop_frame = crop_frame[:, :, ::-1] / 255.0
+        crop_frame = (crop_frame - model_mean) / model_std
+        crop_frame = crop_frame.transpose(2, 0, 1)
+        crop_frame = np.expand_dims(crop_frame, axis=0).astype(np.float32)
+        return crop_frame
+        
+    except Exception as e:
+        print(f"⚠️ Error preparing crop frame: {e}")
+        return crop_frame
 
 
 def normalize_crop_frame(crop_frame: np.ndarray) -> Frame:
     """Normalize cropped frame back to image format"""
-    config = get_current_model_config()
-    model_type = config['type']
-    model_mean = config['mean']
-    model_std = config['standard_deviation']
-    
-    crop_frame = crop_frame.transpose(1, 2, 0)
-    
-    # Apply reverse normalization for certain model types
-    if model_type in ['simswap']:
-        crop_frame = crop_frame * model_std + model_mean
-    
-    crop_frame = crop_frame.clip(0, 1)
-    crop_frame = crop_frame[:, :, ::-1] * 255
-    return crop_frame.astype(np.uint8)
+    try:
+        config = get_current_model_config()
+        model_type = config['type']
+        model_mean = config['mean']
+        model_std = config['standard_deviation']
+        
+        crop_frame = crop_frame.transpose(1, 2, 0)
+        
+        # Apply reverse normalization for certain model types
+        if model_type in ['simswap']:
+            crop_frame = crop_frame * model_std + model_mean
+        
+        crop_frame = crop_frame.clip(0, 1)
+        crop_frame = crop_frame[:, :, ::-1] * 255
+        return crop_frame.astype(np.uint8)
+        
+    except Exception as e:
+        print(f"⚠️ Error normalizing crop frame: {e}")
+        return crop_frame.astype(np.uint8)
 
 
 def enhanced_swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
     """Enhanced face swapping with improved preprocessing"""
-    config = get_current_model_config()
-    model_type = config['type']
-    
-    if model_type == 'inswapper':
-        # Use original method for inswapper
-        return get_frame_processor().get(temp_frame, target_face, source_face, paste_back=True)
-    else:
-        # Enhanced method for other models
-        try:
-            # Prepare source embedding
-            source_embedding = prepare_source_embedding(source_face)
-            
-            # Get crop region (this would need proper implementation)
-            # For now, fall back to original method
-            return get_frame_processor().get(temp_frame, target_face, source_face, paste_back=True)
-        except Exception:
-            # Fallback to original method
-            return get_frame_processor().get(temp_frame, target_face, source_face, paste_back=True)
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Face swap processor not available")
+            return temp_frame
+        
+        config = get_current_model_config()
+        model_type = config['type']
+        
+        if model_type == 'inswapper':
+            # Use original method for inswapper
+            return processor.get(temp_frame, target_face, source_face, paste_back=True)
+        else:
+            # Enhanced method for other models
+            try:
+                # Prepare source embedding
+                source_embedding = prepare_source_embedding(source_face)
+                
+                # Get crop region (this would need proper implementation)
+                # For now, fall back to original method
+                return processor.get(temp_frame, target_face, source_face, paste_back=True)
+            except Exception as e:
+                print(f"⚠️ Enhanced swap failed: {e}")
+                # Fallback to original method
+                return processor.get(temp_frame, target_face, source_face, paste_back=True)
+                
+    except Exception as e:
+        print(f"⚠️ Face swap failed: {e}")
+        return temp_frame
 
 
 def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
     """Main face swapping function with model-specific handling"""
-    config = get_current_model_config()
-    
-    # Use enhanced swapping for supported models
-    if config['type'] in ['simswap', 'inswapper']:
-        return enhanced_swap_face(source_face, target_face, temp_frame)
-    else:
-        # Original method
-        return get_frame_processor().get(temp_frame, target_face, source_face, paste_back=True)
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Face swap processor not available, skipping swap")
+            return temp_frame
+        
+        config = get_current_model_config()
+        
+        # Use enhanced swapping for supported models
+        if config['type'] in ['simswap', 'inswapper']:
+            return enhanced_swap_face(source_face, target_face, temp_frame)
+        else:
+            # Original method
+            return processor.get(temp_frame, target_face, source_face, paste_back=True)
+            
+    except Exception as e:
+        print(f"⚠️ Error in swap_face: {e}")
+        return temp_frame
 
 
 def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame:
     """Process frame with enhanced face selection logic"""
-    if 'reference' in SwitcherAI.globals.face_recognition:
-        similar_faces = find_similar_faces(temp_frame, reference_face, SwitcherAI.globals.reference_face_distance)
-        if similar_faces:
-            for similar_face in similar_faces:
-                temp_frame = swap_face(source_face, similar_face, temp_frame)
-    
-    if 'many' in SwitcherAI.globals.face_recognition:
-        many_faces = get_many_faces(temp_frame)
-        if many_faces:
-            # Sort faces by size (largest first) like the newer version
-            many_faces = sorted(many_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
-            for target_face in many_faces:
-                temp_frame = swap_face(source_face, target_face, temp_frame)
-    
-    return temp_frame
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Face swap processor not available, skipping frame")
+            return temp_frame
+        
+        if 'reference' in SwitcherAI.globals.face_recognition:
+            similar_faces = find_similar_faces(temp_frame, reference_face, SwitcherAI.globals.reference_face_distance)
+            if similar_faces:
+                for similar_face in similar_faces:
+                    temp_frame = swap_face(source_face, similar_face, temp_frame)
+        
+        if 'many' in SwitcherAI.globals.face_recognition:
+            many_faces = get_many_faces(temp_frame)
+            if many_faces:
+                # Sort faces by size (largest first) like the newer version
+                many_faces = sorted(many_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
+                for target_face in many_faces:
+                    temp_frame = swap_face(source_face, target_face, temp_frame)
+        
+        return temp_frame
+        
+    except Exception as e:
+        print(f"⚠️ Error in process_frame: {e}")
+        return temp_frame
 
 
 def get_average_face(faces: List[Face]) -> Face:
@@ -273,59 +399,122 @@ def get_average_face(faces: List[Face]) -> Face:
 
 def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
     """Enhanced frame processing with better source face handling"""
-    source_frame = cv2.imread(source_path)
-    source_faces = get_many_faces(source_frame)
-    
-    # Get best source face (largest)
-    if source_faces:
-        source_faces = sorted(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
-        source_face = source_faces[0]
-    else:
-        source_face = get_one_face(source_frame)
-    
-    # Handle multiple source faces if available
-    if len(source_faces) > 1:
-        source_face = get_average_face(source_faces)
-    
-    reference_face = get_face_reference() if 'reference' in SwitcherAI.globals.face_recognition else None
-    
-    for temp_frame_path in temp_frame_paths:
-        temp_frame = cv2.imread(temp_frame_path)
-        result_frame = process_frame(source_face, reference_face, temp_frame)
-        cv2.imwrite(temp_frame_path, result_frame)
-        if update:
-            update()
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Face swap processor not available, skipping frame processing")
+            if update:
+                update()
+            return
+        
+        source_frame = cv2.imread(source_path)
+        if source_frame is None:
+            print(f"⚠️ Failed to read source image: {source_path}")
+            if update:
+                update()
+            return
+        
+        source_faces = get_many_faces(source_frame)
+        
+        # Get best source face (largest)
+        if source_faces:
+            source_faces = sorted(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
+            source_face = source_faces[0]
+        else:
+            source_face = get_one_face(source_frame)
+        
+        if source_face is None:
+            print("⚠️ No source face found")
+            if update:
+                update()
+            return
+        
+        # Handle multiple source faces if available
+        if len(source_faces) > 1:
+            source_face = get_average_face(source_faces)
+        
+        reference_face = get_face_reference() if 'reference' in SwitcherAI.globals.face_recognition else None
+        
+        for temp_frame_path in temp_frame_paths:
+            try:
+                temp_frame = cv2.imread(temp_frame_path)
+                if temp_frame is not None:
+                    result_frame = process_frame(source_face, reference_face, temp_frame)
+                    cv2.imwrite(temp_frame_path, result_frame)
+                else:
+                    print(f"⚠️ Failed to read frame: {temp_frame_path}")
+                    
+            except Exception as e:
+                print(f"⚠️ Error processing frame {temp_frame_path}: {e}")
+                
+            if update:
+                update()
+                
+    except Exception as e:
+        print(f"⚠️ Error in process_frames: {e}")
 
 
 def process_image(source_path: str, target_path: str, output_path: str) -> None:
     """Enhanced image processing"""
-    source_frame = cv2.imread(source_path)
-    source_faces = get_many_faces(source_frame)
-    
-    # Get best source face
-    if source_faces:
-        source_faces = sorted(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
-        source_face = source_faces[0]
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Face swap processor not available, copying original image")
+            import shutil
+            shutil.copy2(target_path, output_path)
+            return
         
-        # Handle multiple source faces
-        if len(source_faces) > 1:
-            source_face = get_average_face(source_faces)
-    else:
-        source_face = get_one_face(source_frame)
-    
-    target_frame = cv2.imread(target_path)
-    reference_face = get_one_face(target_frame, SwitcherAI.globals.reference_face_position) if 'reference' in SwitcherAI.globals.face_recognition else None
-    result_frame = process_frame(source_face, reference_face, target_frame)
-    cv2.imwrite(output_path, result_frame)
+        source_frame = cv2.imread(source_path)
+        if source_frame is None:
+            print(f"⚠️ Failed to read source image: {source_path}")
+            return
+        
+        source_faces = get_many_faces(source_frame)
+        
+        # Get best source face
+        if source_faces:
+            source_faces = sorted(source_faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]), reverse=True)
+            source_face = source_faces[0]
+            
+            # Handle multiple source faces
+            if len(source_faces) > 1:
+                source_face = get_average_face(source_faces)
+        else:
+            source_face = get_one_face(source_frame)
+        
+        if source_face is None:
+            print("⚠️ No source face found")
+            return
+        
+        target_frame = cv2.imread(target_path)
+        if target_frame is None:
+            print(f"⚠️ Failed to read target image: {target_path}")
+            return
+        
+        reference_face = get_one_face(target_frame, SwitcherAI.globals.reference_face_position) if 'reference' in SwitcherAI.globals.face_recognition else None
+        result_frame = process_frame(source_face, reference_face, target_frame)
+        cv2.imwrite(output_path, result_frame)
+        
+    except Exception as e:
+        print(f"⚠️ Error in process_image: {e}")
 
 
 def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
-    conditional_set_face_reference(temp_frame_paths)
-    frame_processors.process_video(source_path, temp_frame_paths, process_frames)
+    try:
+        conditional_set_face_reference(temp_frame_paths)
+        frame_processors.process_video(source_path, temp_frame_paths, process_frames)
+    except Exception as e:
+        print(f"⚠️ Error in process_video: {e}")
 
 
 def conditional_set_face_reference(temp_frame_paths: List[str]) -> None:
-    if 'reference' in SwitcherAI.globals.face_recognition and not get_face_reference():
-        reference_frame = cv2.imread(temp_frame_paths[SwitcherAI.globals.reference_face_position])
-        reference_face = get_one_face(reference_frame, SwitcherAI.globals.reference_face_position)
-        set_face_reference(reference_face)
+    try:
+        if 'reference' in SwitcherAI.globals.face_recognition and not get_face_reference():
+            reference_frame = cv2.imread(temp_frame_paths[SwitcherAI.globals.reference_face_position])
+            if reference_frame is not None:
+                reference_face = get_one_face(reference_frame, SwitcherAI.globals.reference_face_position)
+                set_face_reference(reference_face)
+            else:
+                print(f"⚠️ Failed to read reference frame")
+    except Exception as e:
+        print(f"⚠️ Error setting face reference: {e}")