from typing import Any, List, Callable import cv2 import threading from pathlib import Path import SwitcherAI.globals import SwitcherAI.processors.frame.core as frame_processors from SwitcherAI import wording from SwitcherAI.core import update_status from SwitcherAI.face_analyser import get_many_faces from SwitcherAI.typing import Frame, Face from SwitcherAI.utilities import conditional_download, resolve_relative_path, is_image, is_video FRAME_PROCESSOR = None THREAD_SEMAPHORE = threading.Semaphore() THREAD_LOCK = threading.Lock() NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_ENHANCER' def get_frame_processor() -> Any: global FRAME_PROCESSOR with THREAD_LOCK: if FRAME_PROCESSOR is None: try: # Import GFPGAN here to handle import errors gracefully from gfpgan.utils import GFPGANer model_path = resolve_relative_path('../.assets/models/GFPGANv1.4.pth') # Convert to Path object if it's a string if isinstance(model_path, str): model_path = Path(model_path) # Check if model exists if not model_path.exists(): print(f"⚠️ GFPGAN model not found at: {model_path}") print("🔄 Attempting to download model...") if not pre_check(): print("❌ Failed to download GFPGAN model") return None FRAME_PROCESSOR = GFPGANer( model_path = str(model_path), upscale = 1, device = frame_processors.get_device() ) print("✅ GFPGAN frame processor initialized") except ImportError as e: print(f"⚠️ GFPGAN not available: {e}") print("💡 Install with: pip install gfpgan") FRAME_PROCESSOR = None except Exception as e: print(f"⚠️ Failed to initialize GFPGAN: {e}") FRAME_PROCESSOR = None return FRAME_PROCESSOR def clear_frame_processor() -> None: global FRAME_PROCESSOR FRAME_PROCESSOR = None def pre_check() -> bool: 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) # Download GFPGAN model model_urls = [ 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth' ] conditional_download(str(download_directory_path), model_urls) # Verify the model was downloaded model_path = download_directory_path / 'GFPGANv1.4.pth' if model_path.exists() and model_path.stat().st_size > 0: print(f"✅ GFPGAN model verified: {model_path.stat().st_size / (1024*1024):.1f}MB") return True else: print("❌ GFPGAN model download failed or file is empty") return False except Exception as e: print(f"❌ GFPGAN pre-check failed: {e}") return False def pre_process() -> bool: try: # Check if we have valid input 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 GFPGAN is available processor = get_frame_processor() if processor is None: print("⚠️ GFPGAN not available, face enhancement will be skipped") return False return True except Exception as e: print(f"⚠️ Face enhancer pre-process failed: {e}") return False def post_process() -> None: clear_frame_processor() def enhance_face(target_face: Face, temp_frame: Frame) -> Frame: """Enhanced face enhancement with error handling""" try: processor = get_frame_processor() if processor is None: print("⚠️ GFPGAN processor not available, returning original frame") return temp_frame start_x, start_y, end_x, end_y = map(int, target_face['bbox']) padding_x = int((end_x - start_x) * 0.5) padding_y = int((end_y - start_y) * 0.5) start_x = max(0, start_x - padding_x) start_y = max(0, start_y - padding_y) end_x = max(0, end_x + padding_x) end_y = max(0, end_y + padding_y) # Ensure coordinates are within frame bounds height, width = temp_frame.shape[:2] end_x = min(end_x, width) end_y = min(end_y, height) crop_frame = temp_frame[start_y:end_y, start_x:end_x] if crop_frame.size > 0: with THREAD_SEMAPHORE: try: _, _, enhanced_crop = processor.enhance( crop_frame, paste_back = True ) temp_frame[start_y:end_y, start_x:end_x] = enhanced_crop except Exception as e: print(f"⚠️ Face enhancement failed: {e}") # Return original frame if enhancement fails pass except Exception as e: print(f"⚠️ Error in enhance_face: {e}") return temp_frame def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame: """Process frame with enhanced error handling""" try: # Check if processor is available processor = get_frame_processor() if processor is None: print("⚠️ Face enhancer not available, skipping enhancement") return temp_frame many_faces = get_many_faces(temp_frame) if many_faces: for target_face in many_faces: temp_frame = enhance_face(target_face, 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 with progress updates""" try: processor = get_frame_processor() if processor is None: print("⚠️ Face 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 with error handling""" try: processor = get_frame_processor() if processor is None: print("⚠️ Face 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_frame = process_frame(None, None, target_frame) cv2.imwrite(output_path, result_frame) 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""" try: SwitcherAI.processors.frame.core.process_video(None, temp_frame_paths, process_frames) except Exception as e: print(f"⚠️ Error in process_video: {e}")