From 0dbed2883a0db583dfcbf5dd75bdd8cfa10c5afe Mon Sep 17 00:00:00 2001 From: Kenneth Estanislao Date: Thu, 27 Mar 2025 02:00:43 +0800 Subject: [PATCH] Update core.py refactoring the code to make it easier to understand and is more optimized --- modules/core.py | 1119 ++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 957 insertions(+), 162 deletions(-) diff --git a/modules/core.py b/modules/core.py index b6ef9b8..ac20edb 100644 --- a/modules/core.py +++ b/modules/core.py @@ -1,55 +1,81 @@ +# --- START OF FILE core.py --- + import os import sys # single thread doubles cuda performance - needs to be set before torch import -if any(arg.startswith('--execution-provider') for arg in sys.argv): +if any(arg.startswith('--execution-provider') for arg in sys.argv) and ('cuda' in sys.argv or 'rocm' in sys.argv): + # Apply for CUDA or ROCm if explicitly mentioned os.environ['OMP_NUM_THREADS'] = '1' # reduce tensorflow log level os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import warnings -from typing import List +from typing import List, Optional, Dict, Any # Added Dict, Any import platform import signal import shutil import argparse -import torch +import gc # Garbage Collector +import time # For timing performance + +# Conditional PyTorch import for memory management +_torch_available = False +_torch_cuda_available = False +try: + import torch + _torch_available = True + if torch.cuda.is_available(): + _torch_cuda_available = True +except ImportError: + # No warning needed unless CUDA is explicitly selected later + pass + import onnxruntime import tensorflow +import cv2 # OpenCV is crucial here +import numpy as np # For frame manipulation import modules.globals import modules.metadata import modules.ui as ui -from modules.processors.frame.core import get_frame_processors_modules +from modules.processors.frame.core import get_frame_processors_modules, load_frame_processor_module # Added load_frame_processor_module from modules.utilities import has_image_extension, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clean_temp, normalize_output_path +# Import necessary typing +from modules.typing import Frame -if 'ROCMExecutionProvider' in modules.globals.execution_providers: - del torch +# Configuration for GPU Memory Limit (adjust as needed, e.g., 0.7-0.9) +GPU_MEMORY_LIMIT_FRACTION = 0.8 # Keep as default, user might adjust based on VRAM -warnings.filterwarnings('ignore', category=FutureWarning, module='insightface') -warnings.filterwarnings('ignore', category=UserWarning, module='torchvision') +# Global to hold active processor instances +FRAME_PROCESSORS_INSTANCES: List[Any] = [] +# --- Argument Parsing and Setup (Mostly unchanged, but refined) --- -def parse_args() -> None: +def parse_args() -> argparse.ArgumentParser: # Return parser for help message on error signal.signal(signal.SIGINT, lambda signal_number, frame: destroy()) - program = argparse.ArgumentParser() - program.add_argument('-s', '--source', help='select an source image', dest='source_path') - program.add_argument('-t', '--target', help='select an target image or video', dest='target_path') - program.add_argument('-o', '--output', help='select output file or directory', dest='output_path') - program.add_argument('--frame-processor', help='pipeline of frame processors', dest='frame_processor', default=['face_swapper'], choices=['face_swapper', 'face_enhancer'], nargs='+') - program.add_argument('--keep-fps', help='keep original fps', dest='keep_fps', action='store_true', default=False) - program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True) - program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False) - program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False) - program.add_argument('--nsfw-filter', help='filter the NSFW image or video', dest='nsfw_filter', action='store_true', default=False) - program.add_argument('--map-faces', help='map source target faces', dest='map_faces', action='store_true', default=False) - program.add_argument('--mouth-mask', help='mask the mouth region', dest='mouth_mask', action='store_true', default=False) - program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9']) - program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]') - program.add_argument('-l', '--lang', help='Ui language', default="en") - program.add_argument('--live-mirror', help='The live camera display as you see it in the front-facing camera frame', dest='live_mirror', action='store_true', default=False) - program.add_argument('--live-resizable', help='The live camera frame is resizable', dest='live_resizable', action='store_true', default=False) - program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory()) - program.add_argument('--execution-provider', help='execution provider', dest='execution_provider', default=['cpu'], choices=suggest_execution_providers(), nargs='+') - program.add_argument('--execution-threads', help='number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads()) + program = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=100, width=120)) # Improved formatter + program.add_argument('-s', '--source', help='Select source image(s) or directory', dest='source_path', nargs='+') # Allow multiple sources + program.add_argument('-t', '--target', help='Select target image or video', dest='target_path') + program.add_argument('-o', '--output', help='Select output file or directory', dest='output_path') + # Frame Processors: Add all available processors to choices dynamically later if possible + available_processors = [proc.NAME for proc in get_frame_processors_modules([])] # Get names dynamically + program.add_argument('--frame-processor', help='Pipeline of frame processors', dest='frame_processor', default=['face_swapper'], choices=available_processors, nargs='+') + program.add_argument('--keep-fps', help='Keep original video fps', dest='keep_fps', action='store_true') + program.add_argument('--keep-audio', help='Keep original video audio (requires --keep-fps for sync)', dest='keep_audio', action='store_true', default=True) # Keep True default + program.add_argument('--keep-frames', help='Keep temporary frames after processing', dest='keep_frames', action='store_true') + program.add_argument('--many-faces', help='Process all detected faces (specific processor behavior)', dest='many_faces', action='store_true') + program.add_argument('--nsfw-filter', help='Enable NSFW prediction and skip if detected', dest='nsfw_filter', action='store_true') + program.add_argument('--map-faces', help='Enable face mapping for video (requires target analysis)', dest='map_faces', action='store_true') + program.add_argument('--color-correction', help='Enable color correction (specific processor behavior)', dest='color_correction', action='store_true') # Add color correction flag + # Mouth mask is processor specific, maybe handled internally or via processor options? Keep it for now. + program.add_argument('--mouth-mask', help='Enable mouth masking (specific processor behavior)', dest='mouth_mask', action='store_true') + program.add_argument('--video-encoder', help='Output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc']) # Added NVIDIA HW encoders + program.add_argument('--video-quality', help='Output video quality crf/qp (0-51 for sw, 0-? for hw, lower=better)', dest='video_quality', type=int, default=18) # Adjusted help text + program.add_argument('-l', '--lang', help='UI language', default="en", choices=["en", "de", "es", "fr", "it", "pt", "ru", "zh"]) # Example languages + program.add_argument('--live-mirror', help='Mirror live camera feed', dest='live_mirror', action='store_true') + program.add_argument('--live-resizable', help='Make live camera window resizable', dest='live_resizable', action='store_true') + program.add_argument('--max-memory', help='DEPRECATED: Use GPU memory fraction. Max CPU RAM limit (GB).', dest='max_memory', type=int) # Default removed, handled dynamically + program.add_argument('--execution-provider', help='Execution provider(s) (cpu, cuda, rocm, dml, coreml)', dest='execution_provider', default=suggest_execution_providers(), nargs='+') # Use suggested default + program.add_argument('--execution-threads', help='Number of threads for execution provider', dest='execution_threads', type=int, default=suggest_execution_threads()) # Use suggested default program.add_argument('-v', '--version', action='version', version=f'{modules.metadata.name} {modules.metadata.version}') # register deprecated args @@ -60,200 +86,969 @@ def parse_args() -> None: args = program.parse_args() - modules.globals.source_path = args.source_path + # Check for ROCm selection early for PyTorch unloading + _is_rocm_selected = any('rocm' in ep.lower() for ep in args.execution_provider) + global _torch_available, _torch_cuda_available + if _is_rocm_selected and _torch_available: + print("[DLC.CORE] ROCm selected, unloading PyTorch.") + del torch + _torch_available = False + _torch_cuda_available = False + gc.collect() + + handle_deprecated_args(args) # Handle deprecated args after initial parsing + + # Assign to globals + # Use the first source if multiple provided for single-source contexts, processors might handle multiple sources. + modules.globals.source_path = args.source_path[0] if isinstance(args.source_path, list) else args.source_path + # Store all sources if needed by processors + modules.globals.source_paths = args.source_path if isinstance(args.source_path, list) else [args.source_path] modules.globals.target_path = args.target_path modules.globals.output_path = normalize_output_path(modules.globals.source_path, modules.globals.target_path, args.output_path) + + # Frame Processors: Store names, instances will be created later modules.globals.frame_processors = args.frame_processor - modules.globals.headless = args.source_path or args.target_path or args.output_path + + modules.globals.headless = bool(args.source_path or args.target_path or args.output_path) modules.globals.keep_fps = args.keep_fps modules.globals.keep_audio = args.keep_audio modules.globals.keep_frames = args.keep_frames modules.globals.many_faces = args.many_faces - modules.globals.mouth_mask = args.mouth_mask + modules.globals.mouth_mask = args.mouth_mask # Pass to processors if they use it + modules.globals.color_correction = args.color_correction # Pass to processors modules.globals.nsfw_filter = args.nsfw_filter modules.globals.map_faces = args.map_faces modules.globals.video_encoder = args.video_encoder modules.globals.video_quality = args.video_quality modules.globals.live_mirror = args.live_mirror modules.globals.live_resizable = args.live_resizable - modules.globals.max_memory = args.max_memory + # Set max_memory, use suggested if not provided by user + modules.globals.max_memory = args.max_memory if args.max_memory is not None else suggest_max_memory() + + # Decode and validate execution providers modules.globals.execution_providers = decode_execution_providers(args.execution_provider) - modules.globals.execution_threads = args.execution_threads + # Set execution threads, ensure it's positive + modules.globals.execution_threads = max(1, args.execution_threads) modules.globals.lang = args.lang - #for ENHANCER tumbler: - if 'face_enhancer' in args.frame_processor: - modules.globals.fp_ui['face_enhancer'] = True - else: - modules.globals.fp_ui['face_enhancer'] = False + # Update derived globals for UI state etc. + modules.globals.fp_ui['face_enhancer'] = 'face_enhancer' in modules.globals.frame_processors + modules.globals.fp_ui['face_swapper'] = 'face_swapper' in modules.globals.frame_processors # Example + # Add other processors as needed - # translate deprecated args + # Final checks and warnings + if modules.globals.keep_audio and not modules.globals.keep_fps: + print("\033[33mWarning: --keep-audio is enabled without --keep-fps. This may cause audio/video sync issues.\033[0m") + if 'cuda' in modules.globals.execution_providers and not _torch_cuda_available: + # Warning if CUDA provider selected but PyTorch CUDA not functional (for memory limiting) + print("\033[33mWarning: CUDA provider selected, but torch.cuda.is_available() is False. PyTorch GPU memory limiting disabled.\033[0m") + if ('h264_nvenc' in modules.globals.video_encoder or 'hevc_nvenc' in modules.globals.video_encoder) and 'cuda' not in modules.globals.execution_providers: + # Check if ffmpeg build supports nvenc if needed + print(f"\033[33mWarning: NVENC encoder ({modules.globals.video_encoder}) selected, but 'cuda' is not in execution providers. Ensure ffmpeg has NVENC support and drivers are installed.\033[0m") + + # Set ONNX Runtime logging level (0:Verbose, 1:Info, 2:Warning, 3:Error, 4:Fatal) + try: + onnxruntime.set_default_logger_severity(3) # Set to Error level to reduce verbose logs + except AttributeError: + print("\033[33mWarning: Could not set ONNX Runtime logger severity (might be an older version).\033[0m") + + return program # Return parser + + +def handle_deprecated_args(args: argparse.Namespace) -> None: + """Handles deprecated arguments and updates corresponding new arguments if necessary.""" + # Source path if args.source_path_deprecated: - print('\033[33mArgument -f and --face are deprecated. Use -s and --source instead.\033[0m') - modules.globals.source_path = args.source_path_deprecated - modules.globals.output_path = normalize_output_path(args.source_path_deprecated, modules.globals.target_path, args.output_path) - if args.cpu_cores_deprecated: - print('\033[33mArgument --cpu-cores is deprecated. Use --execution-threads instead.\033[0m') - modules.globals.execution_threads = args.cpu_cores_deprecated - if args.gpu_vendor_deprecated == 'apple': - print('\033[33mArgument --gpu-vendor apple is deprecated. Use --execution-provider coreml instead.\033[0m') - modules.globals.execution_providers = decode_execution_providers(['coreml']) - if args.gpu_vendor_deprecated == 'nvidia': - print('\033[33mArgument --gpu-vendor nvidia is deprecated. Use --execution-provider cuda instead.\033[0m') - modules.globals.execution_providers = decode_execution_providers(['cuda']) - if args.gpu_vendor_deprecated == 'amd': - print('\033[33mArgument --gpu-vendor amd is deprecated. Use --execution-provider cuda instead.\033[0m') - modules.globals.execution_providers = decode_execution_providers(['rocm']) - if args.gpu_threads_deprecated: - print('\033[33mArgument --gpu-threads is deprecated. Use --execution-threads instead.\033[0m') - modules.globals.execution_threads = args.gpu_threads_deprecated + print('\033[33mWarning: Argument -f/--face is deprecated. Use -s/--source instead.\033[0m') + if not args.source_path: + # Convert to list to match potential nargs='+' + args.source_path = [args.source_path_deprecated] + + # Execution Threads + if args.cpu_cores_deprecated is not None: + print('\033[33mWarning: Argument --cpu-cores is deprecated. Use --execution-threads instead.\033[0m') + # Only override if execution_threads wasn't explicitly set *and* cpu_cores was used + if args.execution_threads == suggest_execution_threads(): # Check against default suggestion + args.execution_threads = args.cpu_cores_deprecated + + if args.gpu_threads_deprecated is not None: + print('\033[33mWarning: Argument --gpu-threads is deprecated. Use --execution-threads instead.\033[0m') + # Override if gpu_threads was used, potentially overriding cpu_cores value if both were used + # Check if execution_threads is still at default OR was set by cpu_cores_deprecated + if args.execution_threads == suggest_execution_threads() or \ + (args.cpu_cores_deprecated is not None and args.execution_threads == args.cpu_cores_deprecated): + args.execution_threads = args.gpu_threads_deprecated + + # Execution Provider from gpu_vendor + if args.gpu_vendor_deprecated: + # Only override if execution_provider is still the default suggested list + suggested_providers_default = suggest_execution_providers() + is_default_provider = sorted(args.execution_provider) == sorted(suggested_providers_default) + + if is_default_provider: + provider_map = { + 'apple': ['coreml', 'cpu'], + 'nvidia': ['cuda', 'cpu'], + 'amd': ['rocm', 'cpu'], + 'intel': ['dml', 'cpu'] # Example for DirectML on Intel + } + vendor = args.gpu_vendor_deprecated.lower() + if vendor in provider_map: + print(f'\033[33mWarning: Argument --gpu-vendor {args.gpu_vendor_deprecated} is deprecated. Setting --execution-provider to {provider_map[vendor]}.\033[0m') + args.execution_provider = provider_map[vendor] + else: + print(f'\033[33mWarning: Unknown --gpu-vendor {args.gpu_vendor_deprecated}. Default execution providers kept.\033[0m') + else: + # User explicitly set execution providers, ignore deprecated vendor + print(f'\033[33mWarning: --gpu-vendor {args.gpu_vendor_deprecated} is deprecated and ignored because --execution-provider was explicitly set to {args.execution_provider}.\033[0m') def encode_execution_providers(execution_providers: List[str]) -> List[str]: - return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers] + """Converts ONNX Runtime provider names to lowercase short names.""" + return [ep.replace('ExecutionProvider', '').lower() for ep in execution_providers] -def decode_execution_providers(execution_providers: List[str]) -> List[str]: - return [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers())) - if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)] +def decode_execution_providers(execution_providers_names: List[str]) -> List[str]: + """Converts lowercase short names back to full ONNX Runtime provider names, preserving order and ensuring availability.""" + available_providers_full = onnxruntime.get_available_providers() # e.g., ['CUDAExecutionProvider', 'CPUExecutionProvider'] + available_providers_encoded = encode_execution_providers(available_providers_full) # e.g., ['cuda', 'cpu'] + decoded_providers = [] + requested_providers_lower = [name.lower() for name in execution_providers_names] + + # User's requested providers first, if available + for req_name_lower in requested_providers_lower: + try: + idx = available_providers_encoded.index(req_name_lower) + provider_full_name = available_providers_full[idx] + if provider_full_name not in decoded_providers: # Avoid duplicates + decoded_providers.append(provider_full_name) + except ValueError: + print(f"\033[33mWarning: Requested execution provider '{req_name_lower}' is not available or not recognized by ONNX Runtime.\033[0m") + + # Ensure CPU is present if no other providers were valid or if it wasn't requested but is available + cpu_provider_full = 'CPUExecutionProvider' + if not decoded_providers or cpu_provider_full not in decoded_providers: + if cpu_provider_full in available_providers_full: + if cpu_provider_full not in decoded_providers: # Add CPU if missing + decoded_providers.append(cpu_provider_full) + print(f"[DLC.CORE] Ensuring '{cpu_provider_full}' is included as a fallback.") + else: + # This is critical - OR needs at least one provider + print(f"\033[31mFatal Error: No valid execution providers found, and '{cpu_provider_full}' is not available in this ONNX Runtime build!\033[0m") + sys.exit(1) + + # Filter list based on actual availability reported by ORT (double check) + final_providers = [p for p in decoded_providers if p in available_providers_full] + if len(final_providers) != len(decoded_providers): + removed = set(decoded_providers) - set(final_providers) + print(f"\033[33mWarning: Providers {list(removed)} were removed after final availability check.\033[0m") + + if not final_providers: + print(f"\033[31mFatal Error: No available execution providers could be configured. Available: {available_providers_full}\033[0m") + sys.exit(1) + + print(f"[DLC.CORE] Using execution providers: {final_providers}") + return final_providers def suggest_max_memory() -> int: - if platform.system().lower() == 'darwin': - return 4 - return 16 + """Suggests a default max CPU RAM limit in GB based on available memory (heuristic).""" + try: + import psutil + total_memory_gb = psutil.virtual_memory().total / (1024 ** 3) + # Suggest using roughly 50% of total RAM, capped at a reasonable upper limit (e.g., 64GB) + # and a lower limit (e.g., 4GB) + suggested_gb = max(4, min(int(total_memory_gb * 0.5), 64)) + # print(f"[DLC.CORE] Suggested max CPU memory (heuristic): {suggested_gb} GB") + return suggested_gb + except ImportError: + print("\033[33mWarning: 'psutil' module not found. Cannot suggest dynamic max_memory. Using default (16GB).\033[0m") + # Fallback to a static default if psutil is not available + return 16 + except Exception as e: + print(f"\033[33mWarning: Error getting system memory: {e}. Using default max_memory (16GB).\033[0m") + return 16 def suggest_execution_providers() -> List[str]: - return encode_execution_providers(onnxruntime.get_available_providers()) + """Suggests available execution providers as short names, prioritizing GPU if available.""" + available_providers_full = onnxruntime.get_available_providers() + available_providers_encoded = encode_execution_providers(available_providers_full) + + # Prioritize GPU providers + provider_priority = ['cuda', 'rocm', 'dml', 'coreml', 'cpu'] + suggested = [] + for provider in provider_priority: + if provider in available_providers_encoded: + suggested.append(provider) + + # Ensure CPU is always included as a fallback + if 'cpu' not in suggested and 'cpu' in available_providers_encoded: + suggested.append('cpu') + + # If only CPU is available, return that + if not suggested and 'cpu' in available_providers_encoded: + return ['cpu'] + elif not suggested: + # Should not happen if ORT is installed correctly + print("\033[31mError: No execution providers detected, including CPU!\033[0m") + return ['cpu'] # Still return cpu as a placeholder + + return suggested def suggest_execution_threads() -> int: - if 'DmlExecutionProvider' in modules.globals.execution_providers: - return 1 - if 'ROCMExecutionProvider' in modules.globals.execution_providers: - return 1 - return 8 + """Suggests a sensible default number of execution threads based on logical CPU cores.""" + try: + logical_cores = os.cpu_count() + if logical_cores: + # Heuristic: Use most cores, but leave some for OS/other tasks. Cap reasonably. + # For systems with many cores (>16), maybe don't use all of them by default. + threads = max(1, min(logical_cores - 2, 16)) if logical_cores > 4 else max(1, logical_cores - 1) + return threads + except NotImplementedError: + pass # Fallback if os.cpu_count() fails + except Exception as e: + print(f"\033[33mWarning: Error getting CPU count: {e}. Using default threads (4).\033[0m") + + # Default fallback + return 4 + + +def limit_gpu_memory(fraction: float) -> None: + """Attempts to limit GPU memory usage via PyTorch (for CUDA) or TensorFlow.""" + gpu_limited = False + + # 1. PyTorch (CUDA) Limit - Only if PyTorch CUDA is available + if 'CUDAExecutionProvider' in modules.globals.execution_providers and _torch_cuda_available: + try: + # Ensure fraction is within valid range [0.0, 1.0] + safe_fraction = max(0.1, min(1.0, fraction)) # Prevent setting 0% + print(f"[DLC.CORE] Attempting to limit PyTorch CUDA memory fraction to {safe_fraction:.1%}") + torch.cuda.set_per_process_memory_fraction(safe_fraction, 0) # Limit on default device (0) + print(f"[DLC.CORE] PyTorch CUDA memory fraction limit set.") + gpu_limited = True + # Optional: Check memory post-limit (can be verbose) + # total_mem = torch.cuda.get_device_properties(0).total_memory + # reserved_mem = torch.cuda.memory_reserved(0) + # allocated_mem = torch.cuda.memory_allocated(0) + # print(f"[DLC.CORE] CUDA Device 0: Total={total_mem/1024**3:.2f}GB, Reserved={reserved_mem/1024**3:.2f}GB, Allocated={allocated_mem/1024**3:.2f}GB") + except RuntimeError as e: + print(f"\033[33mWarning: Failed to set PyTorch CUDA memory fraction (may already be initialized?): {e}\033[0m") + except Exception as e: + print(f"\033[33mWarning: An unexpected error occurred setting PyTorch CUDA memory fraction: {e}\033[0m") + + # 2. TensorFlow GPU Limit (Memory Growth) - Less direct limit, but essential + try: + gpus = tensorflow.config.experimental.list_physical_devices('GPU') + if gpus: + for gpu in gpus: + try: + tensorflow.config.experimental.set_memory_growth(gpu, True) + print(f"[DLC.CORE] Enabled TensorFlow memory growth for GPU: {gpu.name}") + gpu_limited = True # Considered a form of GPU resource management + except RuntimeError as e: + # Memory growth must be set before GPUs have been initialized + print(f"\033[33mWarning: Could not set TensorFlow memory growth for {gpu.name} (may already be initialized?): {e}\033[0m") + except Exception as e: + print(f"\033[33mWarning: An unexpected error occurred setting TensorFlow memory growth for {gpu.name}: {e}\033[0m") + # else: + # No TF GPUs detected, which is fine if not using TF-based models directly + # print("[DLC.CORE] No TensorFlow physical GPUs detected.") + except Exception as e: + print(f"\033[33mWarning: Error configuring TensorFlow GPU settings: {e}\033[0m") + + # if not gpu_limited: + # print("[DLC.CORE] No GPU memory limits applied (GPU provider not used, or libraries unavailable/failed).") def limit_resources() -> None: - # prevent tensorflow memory leak - gpus = tensorflow.config.experimental.list_physical_devices('GPU') - for gpu in gpus: - tensorflow.config.experimental.set_memory_growth(gpu, True) - # limit memory usage - if modules.globals.max_memory: - memory = modules.globals.max_memory * 1024 ** 3 - if platform.system().lower() == 'darwin': - memory = modules.globals.max_memory * 1024 ** 6 - if platform.system().lower() == 'windows': - import ctypes - kernel32 = ctypes.windll.kernel32 - kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory)) - else: - import resource - resource.setrlimit(resource.RLIMIT_DATA, (memory, memory)) + """Limits system resources like CPU RAM (best effort) and configures TF.""" + # 1. Limit CPU RAM (Best effort, platform dependent) + if modules.globals.max_memory and modules.globals.max_memory > 0: + limit_gb = modules.globals.max_memory + limit_bytes = limit_gb * (1024 ** 3) + try: + if platform.system().lower() in ['linux', 'darwin']: + import resource + # RLIMIT_AS limits virtual memory size (includes RAM, swap, mappings) + # Set both soft and hard limits + resource.setrlimit(resource.RLIMIT_AS, (limit_bytes, limit_bytes)) + print(f"[DLC.CORE] Limited process virtual memory (CPU RAM approximation) to ~{limit_gb} GB.") + elif platform.system().lower() == 'windows': + # Windows limiting is harder; SetProcessWorkingSetSizeEx is more of a hint + # Using Job Objects is the robust way but complex to implement here + import ctypes + kernel32 = ctypes.windll.kernel32 + handle = kernel32.GetCurrentProcess() + # Try setting min and max working set size + # Note: Requires specific privileges, might fail silently or with error code + # Use values slightly smaller than the limit for flexibility + min_ws = 1024 * 1024 # Set a small minimum (e.g., 1MB) + max_ws = limit_bytes + if not kernel32.SetProcessWorkingSetSizeEx(handle, ctypes.c_size_t(min_ws), ctypes.c_size_t(max_ws), ctypes.c_ulong(0x1)): # QUOTA_LIMITS_HARDWS_ENABLE = 0x1 + last_error = ctypes.get_last_error() + # Common error: 1314 (ERROR_PRIVILEGE_NOT_HELD) + if last_error == 1314: + print(f"\033[33mWarning: Failed to set process working set size limit on Windows (Error {last_error}). Try running as Administrator if limits are needed.\033[0m") + else: + print(f"\033[33mWarning: Failed to set process working set size limit on Windows (Error {last_error}).\033[0m") + else: + print(f"[DLC.CORE] Requested process working set size limit (Windows memory hint) max ~{limit_gb} GB.") + else: + print(f"\033[33mWarning: CPU RAM limiting not implemented for platform {platform.system()}. --max-memory ignored.\033[0m") + except ImportError: + print(f"\033[33mWarning: 'resource' module (Linux/macOS) or 'ctypes' (Windows) not available. Cannot limit CPU RAM.\033[0m") + except Exception as e: + print(f"\033[33mWarning: Failed to limit CPU RAM: {e}\033[0m") + # else: + # print("[DLC.CORE] CPU RAM limit (--max-memory) not set.") + + # 2. Configure TensorFlow GPU memory growth (already done in limit_gpu_memory, but safe to call again) + # This ensures it's attempted even if limit_gpu_memory wasn't fully effective. + try: + gpus = tensorflow.config.experimental.list_physical_devices('GPU') + if gpus: + for gpu in gpus: + try: + if not tensorflow.config.experimental.get_memory_growth(gpu): + tensorflow.config.experimental.set_memory_growth(gpu, True) + # print(f"[DLC.CORE] Re-checked TF memory growth for {gpu.name}: Enabled.") # Avoid redundant logs + except RuntimeError: + pass # Ignore if already initialized error + except Exception: + pass # Ignore errors here, primary attempt was in limit_gpu_memory def release_resources() -> None: - if 'CUDAExecutionProvider' in modules.globals.execution_providers: - torch.cuda.empty_cache() + """Releases resources, especially GPU memory caches, and runs garbage collection.""" + # 1. Clear PyTorch CUDA cache (if applicable and available) + if _torch_cuda_available: # Check if torch+cuda is loaded + try: + torch.cuda.empty_cache() + # print("[DLC.CORE] Cleared PyTorch CUDA cache.") # Can be verbose + except Exception as e: + print(f"\033[33mWarning: Failed to clear PyTorch CUDA cache: {e}\033[0m") + + # 2. Potentially clear TensorFlow session / clear Keras backend session (less common need) + # try: + # from tensorflow.keras import backend as K + # K.clear_session() + # print("[DLC.CORE] Cleared Keras backend session.") + # except ImportError: + # pass # Keras might not be installed or used + # except Exception as e: + # print(f"\033[33mWarning: Failed to clear Keras session: {e}\033[0m") + + # 3. Explicitly run garbage collection (important!) + gc.collect() + # print("[DLC.CORE] Ran garbage collection.") # Can be verbose def pre_check() -> bool: + """Performs essential pre-run checks for dependencies, versions, and paths.""" + update_status('Performing pre-flight checks...') + checks_passed = True + + # Python version if sys.version_info < (3, 9): - update_status('Python version is not supported - please upgrade to 3.9 or higher.') - return False + update_status('Error: Python 3.9 or higher is required.', 'ERROR') + checks_passed = False + + # FFmpeg if not shutil.which('ffmpeg'): - update_status('ffmpeg is not installed.') - return False - return True + update_status('Error: ffmpeg command was not found in your system PATH. Please install ffmpeg.', 'ERROR') + checks_passed = False + + # ONNX Runtime + try: + ort_version = onnxruntime.__version__ + update_status(f'ONNX Runtime version: {ort_version}') + except Exception as e: + update_status(f'Error: Failed to import or access ONNX Runtime: {e}', 'ERROR') + checks_passed = False + + # TensorFlow (optional, but good to check) + try: + tf_version = tensorflow.__version__ + update_status(f'TensorFlow version: {tf_version}') + except Exception as e: + update_status(f'Warning: Could not import or access TensorFlow: {e}', 'WARN') + # Decide if TF absence is critical based on potential processors + # checks_passed = False + + # PyTorch (only if CUDA is selected for memory limiting) + if 'CUDAExecutionProvider' in modules.globals.execution_providers: + if not _torch_available: + update_status('Warning: CUDA provider selected, but PyTorch is not installed. GPU memory limiting via PyTorch is disabled.', 'WARN') + elif not _torch_cuda_available: + update_status('Warning: PyTorch installed, but torch.cuda.is_available() is False. Check PyTorch CUDA installation and drivers. GPU memory limiting via PyTorch is disabled.', 'WARN') + else: + update_status(f'PyTorch version: {torch.__version__} (CUDA available for memory limiting)') + + + # Check source/target paths if in headless mode + if modules.globals.headless: + if not modules.globals.source_path: + update_status("Error: Source path ('-s' or '--source') is required in headless mode.", 'ERROR') + checks_passed = False + # Check if source files exist + elif isinstance(modules.globals.source_paths, list): + for spath in modules.globals.source_paths: + if not os.path.exists(spath): + update_status(f"Error: Source file/directory not found: {spath}", 'ERROR') + checks_passed = False + elif not os.path.exists(modules.globals.source_path): + update_status(f"Error: Source file/directory not found: {modules.globals.source_path}", 'ERROR') + checks_passed = False + + if not modules.globals.target_path: + update_status("Error: Target path ('-t' or '--target') is required in headless mode.", 'ERROR') + checks_passed = False + elif not os.path.exists(modules.globals.target_path): + update_status(f"Error: Target file not found: {modules.globals.target_path}", 'ERROR') + checks_passed = False + + if not modules.globals.output_path: + update_status("Error: Output path ('-o' or '--output') could not be determined or is missing.", 'ERROR') + checks_passed = False + + update_status('Pre-flight checks completed.') + return checks_passed def update_status(message: str, scope: str = 'DLC.CORE') -> None: - print(f'[{scope}] {message}') + """Prints status messages and updates UI if not headless.""" + log_message = f'[{scope}] {message}' + print(log_message) if not modules.globals.headless: - ui.update_status(message) + try: + # Check if ui module and function exist and are callable + if hasattr(ui, 'update_status') and callable(ui.update_status): + ui.update_status(message) # Pass original message to UI + except Exception as e: + print(f"[DLC.CORE] Error updating UI status: {e}") + + +# --- Main Processing Logic --- def start() -> None: - for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): - if not frame_processor.pre_start(): - return - update_status('Processing...') - # process image to image - if has_image_extension(modules.globals.target_path): - if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy): - return + """Main processing logic for images and videos.""" + start_time = time.time() + update_status(f'Processing started at {time.strftime("%Y-%m-%d %H:%M:%S")}') + + # --- Load and Prepare Frame Processors --- + global FRAME_PROCESSORS_INSTANCES + FRAME_PROCESSORS_INSTANCES = [] # Clear previous instances if any + processors_ready = True + for processor_name in modules.globals.frame_processors: + update_status(f'Loading frame processor: {processor_name}...') + module = load_frame_processor_module(processor_name) + if module: + # Pass necessary global options to the processor's constructor or setup method if needed + # Example: instance = module.Processor(many_faces=modules.globals.many_faces, ...) + instance = module # Assuming module itself might have necessary functions + FRAME_PROCESSORS_INSTANCES.append(instance) + if not instance.pre_start(): # Call pre_start after loading + update_status(f'Initialization failed for {processor_name}. Aborting.', 'ERROR') + processors_ready = False + break # Stop loading further processors + else: + update_status(f'Could not load frame processor module: {processor_name}. Aborting.', 'ERROR') + processors_ready = False + break + + if not processors_ready or not FRAME_PROCESSORS_INSTANCES: + update_status('Frame processor setup failed. Cannot start processing.', 'ERROR') + return + + # Simplify face map for faster lookups if needed + if modules.globals.map_faces and ('face_swapper' in modules.globals.frame_processors): # Example condition + update_status("Simplifying face map for processing...", "Face Analyser") + from modules.face_analyser import simplify_maps # Import locally + simplify_maps() + # Verify map content after simplification (optional debug) + # if modules.globals.simple_map: + # print(f"[DEBUG] Simple map: {len(modules.globals.simple_map['source_faces'])} sources, {len(modules.globals.simple_map['target_embeddings'])} targets") + # else: + # print("[DEBUG] Simple map is empty.") + + + # --- Target is Image --- + if has_image_extension(modules.globals.target_path) and is_image(modules.globals.target_path): + process_image_to_image() + + # --- Target is Video --- + elif is_video(modules.globals.target_path): + process_video() + + # --- Invalid Target --- + else: + if modules.globals.target_path: + update_status(f"Target path '{modules.globals.target_path}' is not a recognized image or video file.", "ERROR") + else: + update_status("Target path not specified or invalid.", "ERROR") + + # --- Processing Finished --- + end_time = time.time() + total_time = end_time - start_time + update_status(f'Processing finished in {total_time:.2f} seconds.') + + +def process_image_to_image(): + """Handles the image-to-image processing workflow.""" + update_status('Processing image: {}'.format(os.path.basename(modules.globals.target_path))) + + # --- NSFW Check --- + if modules.globals.nsfw_filter: + update_status("Checking target image for NSFW content...", "NSFW") + from modules.predicter import predict_image # Import locally try: - shutil.copy2(modules.globals.target_path, modules.globals.output_path) + is_nsfw = predict_image(modules.globals.target_path) + if is_nsfw: + update_status("NSFW content detected in target image. Skipping processing.", "NSFW") + if not modules.globals.headless: + ui.show_error("NSFW content detected. Processing skipped.", title="NSFW Detected") + # Consider deleting output placeholder if it exists? Risky. + # if os.path.exists(modules.globals.output_path): os.remove(modules.globals.output_path) + return # Stop processing + else: + update_status("NSFW check passed.", "NSFW") except Exception as e: - print("Error copying file:", str(e)) - for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): - update_status('Progressing...', frame_processor.NAME) - frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path) - release_resources() - if is_image(modules.globals.target_path): - update_status('Processing to image succeed!') - else: - update_status('Processing to image failed!') - return - # process image to videos - if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy): - return + update_status(f"Error during NSFW check for image: {e}. Continuing processing.", "NSFW") - if not modules.globals.map_faces: - update_status('Creating temp resources...') - create_temp(modules.globals.target_path) - update_status('Extracting frames...') - extract_frames(modules.globals.target_path) + # --- Process --- + try: + # Create output directory if needed + output_dir = os.path.dirname(modules.globals.output_path) + if output_dir and not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + print(f"[DLC.CORE] Created output directory: {output_dir}") - temp_frame_paths = get_temp_frame_paths(modules.globals.target_path) - for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): - update_status('Progressing...', frame_processor.NAME) - frame_processor.process_video(modules.globals.source_path, temp_frame_paths) - release_resources() - # handles fps - if modules.globals.keep_fps: - update_status('Detecting fps...') - fps = detect_fps(modules.globals.target_path) - update_status(f'Creating video with {fps} fps...') - create_video(modules.globals.target_path, fps) - else: - update_status('Creating video with 30.0 fps...') - create_video(modules.globals.target_path) - # handle audio - if modules.globals.keep_audio: + # Read target image using OpenCV (consistent with video frames) + target_frame: Frame = cv2.imread(modules.globals.target_path) + if target_frame is None: + update_status(f'Error: Could not read target image file: {modules.globals.target_path}', 'ERROR') + return + + # --- Apply Processors Sequentially --- + processed_frame = target_frame.copy() # Start with a copy + for processor in FRAME_PROCESSORS_INSTANCES: + processor_name = getattr(processor, 'NAME', 'UnknownProcessor') # Get name safely + update_status(f'Applying {processor_name}...', processor_name) + try: + # Processors should accept a frame (numpy array) and return a processed frame + # Pass global options if needed by the process_frame method + start_proc_time = time.time() + # Pass source path(s) and the frame to be processed + processor_params = { + "source_paths": modules.globals.source_paths, # Pass list of source paths + "target_frame": processed_frame, + "many_faces": modules.globals.many_faces, + "color_correction": modules.globals.color_correction, + "mouth_mask": modules.globals.mouth_mask, + # Add other relevant globals if processors need them + } + # Filter params based on what the processor's process_frame expects (optional advanced) + + processed_frame = processor.process_frame(processor_params) + + if processed_frame is None: + update_status(f'Error: Processor {processor_name} returned None. Aborting processing for this image.', 'ERROR') + return # Stop processing this image + + end_proc_time = time.time() + update_status(f'{processor_name} applied in {end_proc_time - start_proc_time:.2f} seconds.', processor_name) + release_resources() # Release memory after each processor + + except Exception as e: + update_status(f'Error applying processor {processor_name}: {e}', 'ERROR') + import traceback + traceback.print_exc() + return # Stop processing on error + + # --- Save Processed Image --- + update_status(f'Saving processed image to: {modules.globals.output_path}') + try: + # Use OpenCV to save the final frame + # Quality parameters can be added for formats like JPG + # Example: cv2.imwrite(modules.globals.output_path, processed_frame, [cv2.IMWRITE_JPEG_QUALITY, 95]) + save_success = cv2.imwrite(modules.globals.output_path, processed_frame) + if not save_success: + update_status('Error: Failed to save the processed image.', 'ERROR') + elif os.path.exists(modules.globals.output_path) and is_image(modules.globals.output_path): + update_status('Image processing finished successfully.') + else: + update_status('Error: Output image file not found or invalid after saving.', 'ERROR') + + except Exception as e: + update_status(f'Error saving processed image: {e}', 'ERROR') + + except Exception as e: + update_status(f'An unexpected error occurred during image processing: {e}', 'ERROR') + import traceback + traceback.print_exc() + + +def process_video(): + """Handles the video processing workflow with optimized frame handling.""" + update_status('Processing video: {}'.format(os.path.basename(modules.globals.target_path))) + + # --- NSFW Check (Basic - Check first frame or predict_video) --- + if modules.globals.nsfw_filter: + update_status("Checking video for NSFW content (sampling)...", "NSFW") + from modules.predicter import predict_video # Import locally + try: + # Use the library's video prediction (may not use optimal providers) + # Or implement custom frame sampling here using predict_frame + is_nsfw = predict_video(modules.globals.target_path) + if is_nsfw: + update_status("NSFW content detected in video (based on sampling). Skipping processing.", "NSFW") + if not modules.globals.headless: + ui.show_error("NSFW content detected. Processing skipped.", title="NSFW Detected") + return # Stop processing + else: + update_status("NSFW check passed (based on sampling).", "NSFW") + except Exception as e: + update_status(f"Error during NSFW check for video: {e}. Continuing processing.", "NSFW") + + # --- Prepare Temp Environment --- + temp_output_video_path = None # For intermediate video file + video_fps = 30.0 # Default FPS + + try: + # Setup temp directory and frame extraction (if not mapping faces, which might pre-extract) + # If map_faces is enabled, face_analyser.get_unique_faces_from_target_video handles extraction. + if not modules.globals.map_faces: + update_status('Creating temporary resources...', 'Temp') + clean_temp(modules.globals.target_path) # Clean first + create_temp(modules.globals.target_path) + update_status('Extracting video frames...', 'FFmpeg') + extract_frames(modules.globals.target_path, modules.globals.keep_fps) # Pass keep_fps hint + update_status('Frame extraction complete.', 'FFmpeg') + # else: Handled by face mapper + + + # Get paths to frames (must exist either way) + temp_frame_paths = get_temp_frame_paths(modules.globals.target_path) + if not temp_frame_paths: + update_status('Error: No frames found to process. Check temp folder or extraction step.', 'ERROR') + destroy(to_quit=False) # Clean up temp + return + + num_frames = len(temp_frame_paths) + update_status(f'Processing {num_frames} frames...') + + # Determine Target FPS if modules.globals.keep_fps: - update_status('Restoring audio...') + update_status('Detecting target video FPS...', 'FFmpeg') + detected_fps = detect_fps(modules.globals.target_path) + if detected_fps: + video_fps = detected_fps + update_status(f'Using detected FPS: {video_fps:.2f}') + else: + update_status("Warning: Could not detect FPS, using default 30.", "WARN") + video_fps = 30.0 # Fallback fps else: - update_status('Restoring audio might cause issues as fps are not kept...') - restore_audio(modules.globals.target_path, modules.globals.output_path) - else: - move_temp(modules.globals.target_path, modules.globals.output_path) - # clean and validate - clean_temp(modules.globals.target_path) - if is_video(modules.globals.target_path): - update_status('Processing to video succeed!') - else: - update_status('Processing to video failed!') + video_fps = 30.0 # Use default fps if not keeping original + update_status(f'Using fixed FPS: {video_fps:.2f}') + modules.globals.video_fps = video_fps # Store globally if needed elsewhere + + # --- OPTIMIZED Frame Processing Loop --- + update_status('Starting frame processing loop...') + # Use tqdm for progress bar + frame_iterator = tqdm(enumerate(temp_frame_paths), total=num_frames, desc="Processing Frames", unit="frame") + + for frame_index, frame_path in frame_iterator: + try: + # 1. Read Frame + target_frame: Frame = cv2.imread(frame_path) + if target_frame is None: + update_status(f'Warning: Could not read frame {frame_path}. Skipping.', 'WARN') + continue + + # Frame dimensions for potential checks later + # height, width = target_frame.shape[:2] + + # 2. Apply Processors Sequentially to this Frame + processed_frame = target_frame # Start with the original frame for this iteration + for processor in FRAME_PROCESSORS_INSTANCES: + processor_name = getattr(processor, 'NAME', 'UnknownProcessor') + try: + # Pass necessary parameters to the processor's process_frame method + processor_params = { + "source_paths": modules.globals.source_paths, + "target_frame": processed_frame, # Pass the current state of the frame + "many_faces": modules.globals.many_faces, + "color_correction": modules.globals.color_correction, + "mouth_mask": modules.globals.mouth_mask, + "frame_index": frame_index, # Pass frame index if needed + "total_frames": num_frames, # Pass total frames if needed + # Pass simple_map if face mapping is active + "simple_map": modules.globals.simple_map if modules.globals.map_faces else None, + } + # Filter params or use **kwargs if processor accepts them + + temp_frame = processor.process_frame(processor_params) + + if temp_frame is None: + update_status(f'Warning: Processor {processor_name} returned None for frame {frame_index}. Using previous frame state.', 'WARN') + # Keep processed_frame as it was before this processor + else: + processed_frame = temp_frame # Update frame state for the next processor + + # Optimization: Conditional resource release inside loop if memory is tight + # if frame_index % 50 == 0: release_resources() + + except Exception as proc_e: + update_status(f'Error applying processor {processor_name} on frame {frame_index}: {proc_e}', 'ERROR') + # Option: Skip frame vs. Abort entirely + # For now, we continue processing the frame with subsequent processors, using the last valid state + pass # Continue with next processor on this frame + + # 3. Write Processed Frame back to temp location (overwrite original temp frame) + # This ensures create_video reads the modified frames + save_success = cv2.imwrite(frame_path, processed_frame) + if not save_success: + update_status(f'Warning: Failed to save processed frame {frame_path}. Video might contain unprocessed frame.', 'WARN') + + # 4. Release resources periodically (e.g., every N frames or based on time) + if frame_index % 25 == 0 or frame_index == num_frames - 1: # Release every 25 frames and on the last frame + release_resources() + + except Exception as frame_e: + update_status(f'Error processing frame {frame_index} at path {frame_path}: {frame_e}', 'ERROR') + import traceback + traceback.print_exc() + # Option: Continue to next frame or abort? Continue for robustness. + + update_status('Frame processing loop finished.') + + # --- Create Video from Processed Frames --- + update_status('Creating video from processed frames...') + # Define temp output path before audio restoration + temp_output_dir = get_temp_directory_path(modules.globals.target_path) # Get base temp dir + if not temp_output_dir: temp_output_dir = os.path.dirname(modules.globals.output_path) # Fallback + temp_output_video_path = os.path.join(temp_output_dir, f"temp_{os.path.basename(modules.globals.output_path)}") + + create_success = create_video(modules.globals.target_path, video_fps, temp_output_video_path) + if not create_success: + update_status('Error: Failed to create video from processed frames.', 'ERROR') + # Cleanup might still run in finally block + return # Stop here + + # --- Handle Audio Restoration --- + final_output_path = modules.globals.output_path + if modules.globals.keep_audio: + update_status('Restoring audio...', 'FFmpeg') + if not modules.globals.keep_fps: + update_status('Warning: Audio restoration enabled without --keep-fps. Sync issues may occur.', 'WARN') + + # Ensure final output directory exists + final_output_dir = os.path.dirname(final_output_path) + if final_output_dir and not os.path.exists(final_output_dir): os.makedirs(final_output_dir) + + # Restore audio from original target to the temp video, outputting to final path + audio_success = restore_audio(modules.globals.target_path, temp_output_video_path, final_output_path) + if audio_success: + update_status('Audio restoration complete.') + else: + update_status('Error: Audio restoration failed. Video saved without audio.', 'ERROR') + # As a fallback, move the no-audio video to the final path + try: + if os.path.exists(final_output_path): os.remove(final_output_path) + shutil.move(temp_output_video_path, final_output_path) + update_status(f'Fallback: Saved video without audio to {final_output_path}') + temp_output_video_path = None # Prevent deletion in finally + except Exception as move_e: + update_status(f'Error moving temporary video after failed audio restore: {move_e}', 'ERROR') + + else: + # No audio requested, move the temp video to the final output path + update_status('Moving temporary video to final output path (no audio).') + try: + # Ensure final output directory exists + final_output_dir = os.path.dirname(final_output_path) + if final_output_dir and not os.path.exists(final_output_dir): os.makedirs(final_output_dir) + + if os.path.abspath(temp_output_video_path) != os.path.abspath(final_output_path): + if os.path.exists(final_output_path): + os.remove(final_output_path) # Remove existing destination file first + shutil.move(temp_output_video_path, final_output_path) + temp_output_video_path = None # Prevent deletion in finally block + else: + update_status("Temporary video path is same as final output path. No move needed.", "WARN") + temp_output_video_path = None # Still prevent deletion + + except Exception as move_e: + update_status(f'Error moving temporary video to final destination: {move_e}', 'ERROR') -def destroy(to_quit=True) -> None: - if modules.globals.target_path: + # --- Validation --- + if os.path.exists(final_output_path) and is_video(final_output_path): + update_status('Video processing finished successfully.') + else: + update_status('Error: Final output video file not found or invalid after processing.', 'ERROR') + + except Exception as e: + update_status(f'An unexpected error occurred during video processing: {e}', 'ERROR') + import traceback + traceback.print_exc() + + finally: + # --- Clean Up Temporary Resources --- + if not modules.globals.keep_frames: + update_status("Cleaning temporary frame files...", "Temp") + clean_temp(modules.globals.target_path) + else: + update_status("Keeping temporary frame files (--keep-frames enabled).", "Temp") + + # Remove intermediate temp video file if it exists and wasn't moved + if temp_output_video_path and os.path.exists(temp_output_video_path): + try: + os.remove(temp_output_video_path) + update_status(f"Removed intermediate video file: {temp_output_video_path}", "Temp") + except OSError as e: + update_status(f"Warning: Could not remove intermediate video file {temp_output_video_path}: {e}", "WARN") + # Final resource release + release_resources() + + +def destroy(to_quit: bool = True) -> None: + """Cleans up temporary files, releases resources, and optionally exits.""" + update_status("Initiating shutdown sequence...", "CLEANUP") + + # Clean temp files only if target_path was set and keep_frames is false + if hasattr(modules.globals, 'target_path') and modules.globals.target_path and \ + hasattr(modules.globals, 'keep_frames') and not modules.globals.keep_frames: + update_status("Cleaning temporary files (if any)...", "CLEANUP") clean_temp(modules.globals.target_path) - if to_quit: quit() + + # Release models and GPU memory + update_status("Releasing resources...", "CLEANUP") + release_resources() + + # Explicitly clear processor instances (helps GC) + global FRAME_PROCESSORS_INSTANCES + if FRAME_PROCESSORS_INSTANCES: + # Call destroy method on processors if they have one + for processor in FRAME_PROCESSORS_INSTANCES: + if hasattr(processor, 'destroy') and callable(processor.destroy): + try: + processor.destroy() + except Exception as e: + print(f"\033[33mWarning: Error destroying processor {getattr(processor, 'NAME', '?')}: {e}\033[0m") + FRAME_PROCESSORS_INSTANCES.clear() + + # Clear other potentially large global variables explicitly (optional) + if hasattr(modules.globals, 'source_target_map'): modules.globals.source_target_map = [] + if hasattr(modules.globals, 'simple_map'): modules.globals.simple_map = {} + # Clear analyser cache (if it holds significant data) + global FACE_ANALYSER + FACE_ANALYSER = None # Allow GC to collect it + global _ort_session # For NSFW predictor + _ort_session = None + + gc.collect() # Final GC run + + update_status("Cleanup complete.", "CLEANUP") + if to_quit: + print("Exiting application.") + os._exit(0) # Use os._exit for a more forceful exit if needed, sys.exit(0) is generally preferred def run() -> None: - parse_args() + """Parses arguments, sets up environment, and starts processing or UI.""" + # Set TERM environment variable for tqdm on Windows (helps with progress bar rendering) + if platform.system().lower() == 'windows': + os.environ['TERM'] = 'xterm' # Or 'vt100' + + parser = parse_args() # Parse arguments first to set globals + + # Apply GPU Memory Limit early, requires execution_providers to be set + limit_gpu_memory(GPU_MEMORY_LIMIT_FRACTION) + + # Perform pre-checks (dependencies, versions, paths) if not pre_check(): - return - for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): - if not frame_processor.pre_check(): - return + # Display help if critical checks fail in headless mode (e.g., missing paths) + if modules.globals.headless: + print("\033[31mCritical pre-check failed. Please review errors above.\033[0m") + parser.print_help() + destroy(to_quit=True) + return # Exit if pre-checks fail + + # Limit other resources (CPU RAM, TF GPU options) limit_resources() - if modules.globals.headless: - start() + + # --- Processor Requirements Check --- + # Moved after parse_args and resource limits + active_processor_modules = get_frame_processors_modules(modules.globals.frame_processors) + all_processors_ready = True + if not active_processor_modules: + update_status('Error: No valid frame processors specified or found.', 'ERROR') + all_processors_ready = False else: - window = ui.init(start, destroy, modules.globals.lang) - window.mainloop() + for processor_module in active_processor_modules: + processor_name = getattr(processor_module, 'NAME', 'UnknownProcessor') + update_status(f'Checking requirements for {processor_name}...') + try: + if not processor_module.pre_check(): + update_status(f'Requirements check failed for {processor_name}.', 'ERROR') + all_processors_ready = False + # Don't break early, report all failed checks + else: + update_status(f'Requirements met for {processor_name}.') + except Exception as e: + update_status(f'Error during requirements check for {processor_name}: {e}', 'ERROR') + all_processors_ready = False + + if not all_processors_ready: + update_status('One or more frame processors failed requirement checks. Please review messages above.', 'ERROR') + destroy(to_quit=True) + return + + # --- Run Mode --- + if modules.globals.headless: + update_status('Running in headless mode.') + # Face mapping requires specific setup before starting the main processing + if modules.globals.map_faces: + update_status("Mapping faces enabled, analyzing target...", "Face Analyser") + if is_video(modules.globals.target_path): + from modules.face_analyser import get_unique_faces_from_target_video + get_unique_faces_from_target_video() + elif is_image(modules.globals.target_path): + from modules.face_analyser import get_unique_faces_from_target_image + get_unique_faces_from_target_image() + else: + update_status("Map faces requires a valid target image or video.", "ERROR") + destroy(to_quit=True) + return + update_status("Target analysis for face mapping complete.", "Face Analyser") + + start() # Run the main processing function + destroy(to_quit=True) # Exit after headless processing + else: + # Launch UI + update_status('Launching graphical user interface...') + # Ensure destroy is callable without arguments for the UI close button + destroy_wrapper = lambda: destroy(to_quit=True) + try: + window = ui.init(start, destroy_wrapper, modules.globals.lang) + window.mainloop() + except Exception as e: + print(f"\033[31mFatal Error initializing or running the UI: {e}\033[0m") + import traceback + traceback.print_exc() + destroy(to_quit=True) # Attempt cleanup and exit even if UI fails + + +# --- Main execution entry point --- +if __name__ == "__main__": + # Add project root to Python path (if core.py is not at the very top level) + # script_dir = os.path.dirname(os.path.abspath(__file__)) + # project_root = os.path.dirname(script_dir) # Adjust if structure differs + # if project_root not in sys.path: + # sys.path.insert(0, project_root) + + run() +# --- END OF FILE core.py --- \ No newline at end of file