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

SwitcherAI/processors/frame/modules/frame_enhancer.py

@@ -3,9 +3,8 @@ import cv2
 import threading
 import numpy
 from functools import lru_cache
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from realesrgan import RealESRGANer
-import torch
+from pathlib import Path
+
 import SwitcherAI.processors.frame.core as frame_processors
 from SwitcherAI.typing import Frame, Face
 from SwitcherAI.utilities import conditional_download, resolve_relative_path
@@ -20,9 +19,13 @@ NAME = 'FACEFUSION.FRAME_PROCESSOR.FRAME_ENHANCER'
 @lru_cache(maxsize=None)
 def get_model_config() -> Dict[str, Any]:
     """Get model configuration with enhanced options"""
+    base_path = resolve_relative_path('../.assets/models')
+    if isinstance(base_path, str):
+        base_path = Path(base_path)
+    
     return {
         'real_esrgan_x4': {
-            'model_path': resolve_relative_path('../.assets/models/RealESRGAN_x4plus.pth'),
+            'model_path': base_path / 'RealESRGAN_x4plus.pth',
             'scale': 4,
             'tile_size': 256,
             'tile_pad': 16,
@@ -38,29 +41,53 @@ def get_frame_processor() -> Any:
 
     with THREAD_LOCK:
         if FRAME_PROCESSOR is None:
-            config = get_model_config()['real_esrgan_x4']
-            model_path = config['model_path']
-            
-            FRAME_PROCESSOR = RealESRGANer(
-                model_path=model_path,
-                model=RRDBNet(
-                    num_in_ch=3,
-                    num_out_ch=3,
-                    num_feat=config['num_feat'],
-                    num_block=config['num_block'],
-                    num_grow_ch=config['num_grow_ch'],
+            try:
+                # Import Real-ESRGAN components
+                from basicsr.archs.rrdbnet_arch import RRDBNet
+                from realesrgan import RealESRGANer
+                import torch
+                
+                config = get_model_config()['real_esrgan_x4']
+                model_path = config['model_path']
+                
+                # Check if model exists
+                if not model_path.exists():
+                    print(f"⚠️ Real-ESRGAN model not found at: {model_path}")
+                    print("🔄 Attempting to download model...")
+                    if not pre_check():
+                        print("❌ Failed to download Real-ESRGAN model")
+                        return None
+                
+                FRAME_PROCESSOR = RealESRGANer(
+                    model_path=str(model_path),
+                    model=RRDBNet(
+                        num_in_ch=3,
+                        num_out_ch=3,
+                        num_feat=config['num_feat'],
+                        num_block=config['num_block'],
+                        num_grow_ch=config['num_grow_ch'],
+                        scale=config['scale']
+                    ),
+                    device=frame_processors.get_device(),
+                    tile=config['tile_size'],
+                    tile_pad=config['tile_pad'],
+                    pre_pad=0,
                     scale=config['scale']
-                ),
-                device=frame_processors.get_device(),
-                tile=config['tile_size'],
-                tile_pad=config['tile_pad'],
-                pre_pad=0,
-                scale=config['scale']
-            )
-            
-            # Ensure CUDA device is set if available
-            if torch.cuda.is_available():
-                torch.cuda.set_device(0)
+                )
+                
+                # Ensure CUDA device is set if available
+                if torch.cuda.is_available():
+                    torch.cuda.set_device(0)
+                
+                print("✅ Real-ESRGAN frame processor initialized")
+                
+            except ImportError as e:
+                print(f"⚠️ Real-ESRGAN not available: {e}")
+                print("💡 Install with: pip install realesrgan basicsr")
+                FRAME_PROCESSOR = None
+            except Exception as e:
+                print(f"⚠️ Failed to initialize Real-ESRGAN: {e}")
+                FRAME_PROCESSOR = None
 
     return FRAME_PROCESSOR
 
@@ -72,27 +99,46 @@ def clear_frame_processor() -> None:
 
 def pre_check() -> bool:
     """Download required models for frame enhancement"""
-    download_directory_path = resolve_relative_path('../.assets/models')
-    
     try:
-        conditional_download(download_directory_path, [
+        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)
+        
+        # Download Real-ESRGAN model
+        model_urls = [
             'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
-        ])
-        return True
+        ]
+        
+        conditional_download(str(download_directory_path), model_urls)
+        
+        # Verify the model was downloaded
+        model_path = download_directory_path / 'RealESRGAN_x4plus.pth'
+        if model_path.exists() and model_path.stat().st_size > 0:
+            print(f"✅ Real-ESRGAN model verified: {model_path.stat().st_size / (1024*1024):.1f}MB")
+            return True
+        else:
+            print("❌ Real-ESRGAN model download failed or file is empty")
+            return False
+            
     except Exception as e:
-        print(f"⚠️ Failed to download frame enhancement models: {e}")
-        print("🔄 Frame enhancement will be disabled")
+        print(f"❌ Real-ESRGAN pre-check failed: {e}")
         return False
 
 
 def pre_process() -> bool:
     """Pre-process check with model validation"""
     try:
-        model_path = get_model_config()['real_esrgan_x4']['model_path']
-        if not model_path or not model_path.exists():
-            print("⚠️ Frame enhancement model not found")
+        # Check if processor is available
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Real-ESRGAN not available, frame enhancement will be skipped")
             return False
+        
         return True
+        
     except Exception as e:
         print(f"⚠️ Frame enhancement pre-process failed: {e}")
         return False
@@ -166,34 +212,47 @@ def enhance_frame_with_tiling(temp_frame: Frame) -> Frame:
     """
     Enhanced frame enhancement with improved tiling (inspired by FaceFusion)
     """
-    config = get_model_config()['real_esrgan_x4']
-    tile_size = (config['tile_size'], config['tile_size'])
-    scale = config['scale']
-    
-    # Create tiles for processing
-    tiles, pad_width, pad_height = create_tile_frames(temp_frame, tile_size)
-    enhanced_tiles = []
-    
-    with THREAD_SEMAPHORE:
-        frame_processor = get_frame_processor()
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Real-ESRGAN processor not available, returning original frame")
+            return temp_frame
         
-        for tile in tiles:
-            # Process each tile individually to manage memory
-            enhanced_tile, _ = frame_processor.enhance(tile, outscale=scale)
-            enhanced_tiles.append(enhanced_tile)
-    
-    # Merge tiles back together
-    original_height, original_width = temp_frame.shape[:2]
-    enhanced_frame = merge_tile_frames(
-        enhanced_tiles, 
-        original_width * scale, 
-        original_height * scale,
-        pad_width * scale, 
-        pad_height * scale,
-        (tile_size[0] * scale, tile_size[1] * scale)
-    )
-    
-    return enhanced_frame
+        config = get_model_config()['real_esrgan_x4']
+        tile_size = (config['tile_size'], config['tile_size'])
+        scale = config['scale']
+        
+        # Create tiles for processing
+        tiles, pad_width, pad_height = create_tile_frames(temp_frame, tile_size)
+        enhanced_tiles = []
+        
+        with THREAD_SEMAPHORE:
+            for tile in tiles:
+                try:
+                    # Process each tile individually to manage memory
+                    enhanced_tile, _ = processor.enhance(tile, outscale=scale)
+                    enhanced_tiles.append(enhanced_tile)
+                except Exception as e:
+                    print(f"⚠️ Tile enhancement failed: {e}")
+                    # Use original tile if enhancement fails
+                    enhanced_tiles.append(tile)
+        
+        # Merge tiles back together
+        original_height, original_width = temp_frame.shape[:2]
+        enhanced_frame = merge_tile_frames(
+            enhanced_tiles, 
+            original_width * scale, 
+            original_height * scale,
+            pad_width * scale, 
+            pad_height * scale,
+            (tile_size[0] * scale, tile_size[1] * scale)
+        )
+        
+        return enhanced_frame
+        
+    except Exception as e:
+        print(f"⚠️ Enhanced tiling failed: {e}")
+        return temp_frame
 
 
 def enhance_frame(temp_frame: Frame) -> Frame:
@@ -201,12 +260,24 @@ def enhance_frame(temp_frame: Frame) -> Frame:
     Main enhancement function with fallback to original method
     """
     try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Frame enhancer not available, returning original frame")
+            return temp_frame
+        
         # Try enhanced tiling method first
-        return enhance_frame_with_tiling(temp_frame)
-    except Exception:
-        # Fallback to original method
-        with THREAD_SEMAPHORE:
-            temp_frame, _ = get_frame_processor().enhance(temp_frame, outscale=1)
+        try:
+            return enhance_frame_with_tiling(temp_frame)
+        except Exception as e:
+            print(f"⚠️ Tiling method failed: {e}, trying simple enhancement")
+            
+            # Fallback to original method
+            with THREAD_SEMAPHORE:
+                enhanced_frame, _ = processor.enhance(temp_frame, outscale=1)
+                return enhanced_frame
+                
+    except Exception as e:
+        print(f"⚠️ Frame enhancement failed completely: {e}")
         return temp_frame
 
 
@@ -214,61 +285,114 @@ def blend_frame(original_frame: Frame, enhanced_frame: Frame, blend_ratio: float
     """
     Blend original and enhanced frames (inspired by FaceFusion)
     """
-    if original_frame.shape != enhanced_frame.shape:
-        original_frame = cv2.resize(original_frame, (enhanced_frame.shape[1], enhanced_frame.shape[0]))
-    
-    # Convert blend ratio (0-1 where 1 = full enhancement)
-    return cv2.addWeighted(original_frame, 1 - blend_ratio, enhanced_frame, blend_ratio, 0)
+    try:
+        if original_frame.shape != enhanced_frame.shape:
+            original_frame = cv2.resize(original_frame, (enhanced_frame.shape[1], enhanced_frame.shape[0]))
+        
+        # Convert blend ratio (0-1 where 1 = full enhancement)
+        return cv2.addWeighted(original_frame, 1 - blend_ratio, enhanced_frame, blend_ratio, 0)
+    except Exception as e:
+        print(f"⚠️ Frame blending failed: {e}")
+        return enhanced_frame
 
 
 def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame:
     """
     Main processing function (maintains your original interface)
     """
-    return enhance_frame(temp_frame)
+    try:
+        return enhance_frame(temp_frame)
+    except Exception as e:
+        print(f"⚠️ Error in process_frame: {e}")
+        return temp_frame
 
 
 def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
     """
     Process multiple frames (maintains your original interface)
     """
-    for temp_frame_path in temp_frame_paths:
-        temp_frame = cv2.imread(temp_frame_path)
-        result_frame = process_frame(None, None, temp_frame)
-        cv2.imwrite(temp_frame_path, result_frame)
-        if update:
-            update()
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Frame enhancer not available, skipping frame enhancement")
+            if update:
+                update()
+            return
+        
+        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(None, None, 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:
     """
     Process single image (maintains your original interface)
     """
-    target_frame = cv2.imread(target_path)
-    result = process_frame(None, None, target_frame)
-    cv2.imwrite(output_path, result)
+    try:
+        processor = get_frame_processor()
+        if processor is None:
+            print("⚠️ Frame enhancer not available, copying original image")
+            import shutil
+            shutil.copy2(target_path, output_path)
+            return
+        
+        target_frame = cv2.imread(target_path)
+        if target_frame is not None:
+            result = process_frame(None, None, target_frame)
+            cv2.imwrite(output_path, result)
+        else:
+            print(f"⚠️ Failed to read image: {target_path}")
+            
+    except Exception as e:
+        print(f"⚠️ Error in process_image: {e}")
 
 
 def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
     """
     Process video frames (maintains your original interface)
     """
-    frame_processors.process_video(None, temp_frame_paths, process_frames)
+    try:
+        frame_processors.process_video(None, temp_frame_paths, process_frames)
+    except Exception as e:
+        print(f"⚠️ Error in process_video: {e}")
 
 
 # Additional utility functions inspired by FaceFusion
 def get_model_scale() -> int:
     """Get the current model's scale factor"""
-    return get_model_config()['real_esrgan_x4']['scale']
+    try:
+        return get_model_config()['real_esrgan_x4']['scale']
+    except:
+        return 1
 
 
 def prepare_frame(frame: Frame) -> Frame:
     """Prepare frame for processing"""
-    if frame.dtype != numpy.uint8:
-        frame = frame.astype(numpy.uint8)
-    return frame
+    try:
+        if frame.dtype != numpy.uint8:
+            frame = frame.astype(numpy.uint8)
+        return frame
+    except:
+        return frame
 
 
 def normalize_frame(frame: Frame) -> Frame:
     """Normalize frame after processing"""
-    return numpy.clip(frame, 0, 255).astype(numpy.uint8)
+    try:
+        return numpy.clip(frame, 0, 255).astype(numpy.uint8)
+    except:
+        return frame