diff --git a/modules/core.py b/modules/core.py index ac20edb..b11f53a 100644 --- a/modules/core.py +++ b/modules/core.py @@ -3,21 +3,54 @@ 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) and ('cuda' in sys.argv or 'rocm' in sys.argv): - # Apply for CUDA or ROCm if explicitly mentioned +# Check if CUDAExecutionProvider is likely intended +_cuda_intended = False +if '--execution-provider' in sys.argv: + try: + providers_index = sys.argv.index('--execution-provider') + # Check subsequent arguments until the next option (starts with '-') or end of list + for i in range(providers_index + 1, len(sys.argv)): + if sys.argv[i].startswith('-'): + break + if 'cuda' in sys.argv[i].lower(): + _cuda_intended = True + break + except ValueError: + pass # --execution-provider not found +# Less precise check if the above fails or isn't used (e.g. deprecated --gpu-vendor nvidia) +if not _cuda_intended and any('cuda' in arg.lower() or 'nvidia' in arg.lower() for arg in sys.argv): + _cuda_intended = True + +if _cuda_intended: + print("[DLC.CORE] CUDA execution provider detected or inferred, setting OMP_NUM_THREADS=1.") os.environ['OMP_NUM_THREADS'] = '1' # reduce tensorflow log level os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import warnings -from typing import List, Optional, Dict, Any # Added Dict, Any +from typing import List, Optional import platform import signal import shutil import argparse import gc # Garbage Collector -import time # For timing performance -# Conditional PyTorch import for memory management +# --- ONNX Runtime Version Check --- +# Ensure ONNX Runtime is imported and check version compatibility if needed. +# As of onnxruntime 1.19, the core APIs used here (get_available_providers, InferenceSession config) +# remain stable. No specific code changes are required *in this file* for 1.19 compatibility, +# assuming frame processors use standard SessionOptions/InferenceSession creation. +try: + import onnxruntime + # print(f"[DLC.CORE] Using ONNX Runtime version: {onnxruntime.__version__}") # Optional: uncomment for debug + # Example future check: + # from packaging import version + # if version.parse(onnxruntime.__version__) < version.parse("1.19.0"): + # print(f"Warning: ONNX Runtime version {onnxruntime.__version__} is older than 1.19. Some features might differ.") +except ImportError: + print("\033[31m[DLC.CORE] Error: ONNX Runtime is not installed. Please install it (e.g., `pip install onnxruntime` or `pip install onnxruntime-gpu`).\033[0m") + sys.exit(1) + +# --- PyTorch Conditional Import --- _torch_available = False _torch_cuda_available = False try: @@ -26,1029 +59,919 @@ try: if torch.cuda.is_available(): _torch_cuda_available = True except ImportError: - # No warning needed unless CUDA is explicitly selected later - pass + # Warning only if CUDA EP might be used, otherwise PyTorch is optional + if _cuda_intended: + print("[DLC.CORE] Warning: PyTorch not found or CUDA not available. GPU memory limiting via Torch is disabled.") + pass # Keep torch=None or handle appropriately -import onnxruntime -import tensorflow -import cv2 # OpenCV is crucial here -import numpy as np # For frame manipulation +# --- TensorFlow Conditional Import (for resource limiting) --- +_tensorflow_available = False +try: + import tensorflow + _tensorflow_available = True +except ImportError: + print("[DLC.CORE] Info: TensorFlow not found. GPU memory growth configuration for TensorFlow will be skipped.") + pass import modules.globals import modules.metadata import modules.ui as ui -from modules.processors.frame.core import get_frame_processors_modules, load_frame_processor_module # Added load_frame_processor_module +from modules.processors.frame.core import get_frame_processors_modules 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 -# 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 +# Configuration for GPU Memory Limit (0.8 = 80%) +GPU_MEMORY_LIMIT_FRACTION = 0.8 -# Global to hold active processor instances -FRAME_PROCESSORS_INSTANCES: List[Any] = [] +# Check if ROCM is chosen early, before parse_args if possible, or handle after +_is_rocm_selected = False +# A simple check; parse_args will give the definitive list later +if any('rocm' in arg.lower() for arg in sys.argv): + _is_rocm_selected = True -# --- Argument Parsing and Setup (Mostly unchanged, but refined) --- +if _is_rocm_selected and _torch_available: + # If ROCM is selected, torch might interfere or not be needed. + # Let's keep the behavior of unloading it for safety, as ROCm support in PyTorch can be complex. + print("[DLC.CORE] ROCM detected or selected, unloading PyTorch to prevent potential conflicts.") + del torch + _torch_available = False + _torch_cuda_available = False + gc.collect() # Try to explicitly collect garbage -def parse_args() -> argparse.ArgumentParser: # Return parser for help message on error + +warnings.filterwarnings('ignore', category=FutureWarning, module='insightface') +warnings.filterwarnings('ignore', category=UserWarning, module='torchvision') + + +def parse_args() -> None: signal.signal(signal.SIGINT, lambda signal_number, frame: destroy()) - 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 = argparse.ArgumentParser(formatter_class=lambda prog: argparse.ArgumentDefaultsHelpFormatter(prog, max_help_position=40)) # Wider help + program.add_argument('-s', '--source', help='Path to the source image file', dest='source_path') + program.add_argument('-t', '--target', help='Path to the target image or video file', dest='target_path') + program.add_argument('-o', '--output', help='Path for the output file or directory', dest='output_path') + # Frame processors - Updated choices might be needed if new processors are added + available_processors = ['face_swapper', 'face_enhancer'] # Dynamically get these if possible in future 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}') + program.add_argument('--keep-fps', help='Keep the original frames per second (FPS) of the target video', dest='keep_fps', action='store_true') + program.add_argument('--keep-audio', help='Keep the original audio of the target video (requires --keep-fps for perfect sync)', dest='keep_audio', action='store_true', default=True) + program.add_argument('--keep-frames', help='Keep the temporary extracted frames after processing', dest='keep_frames', action='store_true') + program.add_argument('--many-faces', help='Process all detected faces in the target, not just the most similar', dest='many_faces', action='store_true') + program.add_argument('--nsfw-filter', help='Enable NSFW content filtering (experimental, image-only currently)', dest='nsfw_filter', action='store_true') + program.add_argument('--map-faces', help='EXPERIMENTAL: Map source faces to target faces based on order or index. Requires manual setup or specific naming conventions.', dest='map_faces', action='store_true') + program.add_argument('--mouth-mask', help='Apply a mask over the mouth region during processing (specific to certain processors)', dest='mouth_mask', action='store_true') + program.add_argument('--video-encoder', help='Encoder for the output video', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc']) # Added NVENC options + program.add_argument('--video-quality', help='Quality for the output video (lower value means higher quality, range depends on encoder)', dest='video_quality', type=int, default=18, metavar='[0-51 for x264/x265, 0-63 for vp9]') # Adjusted range note + program.add_argument('-l', '--lang', help='User interface language code (e.g., "en", "es")', default="en") + program.add_argument('--live-mirror', help='Mirror the live camera preview (like a webcam)', dest='live_mirror', action='store_true') + program.add_argument('--live-resizable', help='Allow resizing the live camera preview window', dest='live_resizable', action='store_true') + program.add_argument('--max-memory', help='DEPRECATED (use with caution): Approx. maximum CPU RAM in GB. Less effective than GPU limits.', dest='max_memory', type=int) # Removed default, let suggest_max_memory handle it dynamically if needed + # Execution Provider - Updated based on ONNX Runtime 1.19 common providers + program.add_argument('--execution-provider', help='Execution provider(s) to use (e.g., cuda, cpu, rocm, dml, coreml). Order determines priority.', dest='execution_provider', default=suggest_execution_providers(), choices=get_available_execution_providers_short(), nargs='+') + program.add_argument('--execution-threads', help='Number of threads for the execution provider', dest='execution_threads', type=int, default=suggest_execution_threads()) + program.add_argument('-v', '--version', action='version', version=f'{modules.metadata.name} {modules.metadata.version} (ONNX Runtime: {onnxruntime.__version__})') # Added ORT version # register deprecated args program.add_argument('-f', '--face', help=argparse.SUPPRESS, dest='source_path_deprecated') program.add_argument('--cpu-cores', help=argparse.SUPPRESS, dest='cpu_cores_deprecated', type=int) - program.add_argument('--gpu-vendor', help=argparse.SUPPRESS, dest='gpu_vendor_deprecated') + program.add_argument('--gpu-vendor', help=argparse.SUPPRESS, dest='gpu_vendor_deprecated', choices=['apple', 'nvidia', 'amd']) program.add_argument('--gpu-threads', help=argparse.SUPPRESS, dest='gpu_threads_deprecated', type=int) args = program.parse_args() - # 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() + # Set default for max_memory if not provided + if args.max_memory is None: + args.max_memory = suggest_max_memory() - handle_deprecated_args(args) # Handle deprecated args after initial parsing + # Process deprecated args first + handle_deprecated_args(args) # 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.source_path = 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 - + # Headless mode is determined by the presence of CLI args for paths 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_audio = args.keep_audio # Note: keep_audio without keep_fps can cause sync issues modules.globals.keep_frames = args.keep_frames modules.globals.many_faces = args.many_faces - 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.mouth_mask = args.mouth_mask 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 - # 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) - # Set execution threads, ensure it's positive - modules.globals.execution_threads = max(1, args.execution_threads) + modules.globals.max_memory = args.max_memory # Still set, but primarily for CPU RAM limit now + modules.globals.execution_providers = decode_execution_providers(args.execution_provider) # Decode selected short names + modules.globals.execution_threads = args.execution_threads modules.globals.lang = args.lang - # 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 - - # 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 + # Update derived globals + modules.globals.fp_ui = {proc: (proc in modules.globals.frame_processors) for proc in available_processors} # Simplified UI state init + # Validate keep_audio / keep_fps combination + if modules.globals.keep_audio and not modules.globals.keep_fps and not modules.globals.headless: + # Only warn in interactive mode, CLI users are expected to know + print("\033[33mWarning: --keep-audio is enabled but --keep-fps is disabled. This might cause audio/video synchronization issues.\033[0m") + elif modules.globals.keep_audio and not modules.globals.target_path: + print("\033[33mWarning: --keep-audio is enabled but no target video path is provided. Audio cannot be kept.\033[0m") + modules.globals.keep_audio = False 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[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] + print('\033[33mArgument -f/--face is deprecated. Use -s/--source instead.\033[0m') + if not args.source_path: # Only override if --source wasn't set + args.source_path = args.source_path_deprecated + # Re-evaluate output path based on deprecated source (normalize_output_path handles this later) + + # Track if execution_threads was explicitly set by the user via --execution-threads + # This requires checking sys.argv as argparse doesn't directly expose this. + threads_explicitly_set = '--execution-threads' in sys.argv - # 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 + print('\033[33mArgument --cpu-cores is deprecated. Use --execution-threads instead.\033[0m') + # Only override if --execution-threads wasn't explicitly set + if not threads_explicitly_set: + args.execution_threads = args.cpu_cores_deprecated + threads_explicitly_set = True # Mark as set now 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): + print('\033[33mArgument --gpu-threads is deprecated. Use --execution-threads instead.\033[0m') + # Only override if --execution-threads wasn't explicitly set (by user or cpu-cores) + if not threads_explicitly_set: args.execution_threads = args.gpu_threads_deprecated + threads_explicitly_set = True # Mark as set + + # Handle --gpu-vendor deprecation by modifying execution_provider list *if not explicitly set* + ep_explicitly_set = '--execution-provider' in sys.argv - # 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: + print(f'\033[33mArgument --gpu-vendor {args.gpu_vendor_deprecated} is deprecated. Use --execution-provider instead.\033[0m') + if not ep_explicitly_set: provider_map = { - 'apple': ['coreml', 'cpu'], - 'nvidia': ['cuda', 'cpu'], - 'amd': ['rocm', 'cpu'], - 'intel': ['dml', 'cpu'] # Example for DirectML on Intel + # Map vendor to preferred execution provider short names + 'apple': ['coreml', 'cpu'], # CoreML first + 'nvidia': ['cuda', 'cpu'], # CUDA first + 'amd': ['rocm', 'cpu'] # ROCm first + # 'intel': ['openvino', 'cpu'] # Example if OpenVINO support is relevant } - 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] + if args.gpu_vendor_deprecated in provider_map: + suggested_providers = provider_map[args.gpu_vendor_deprecated] + print(f"Mapping deprecated --gpu-vendor {args.gpu_vendor_deprecated} to --execution-provider {' '.join(suggested_providers)}") + args.execution_provider = suggested_providers # Set the list of short names else: - print(f'\033[33mWarning: Unknown --gpu-vendor {args.gpu_vendor_deprecated}. Default execution providers kept.\033[0m') + print(f'\033[33mWarning: Unknown --gpu-vendor {args.gpu_vendor_deprecated}. Default execution providers will be used.\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') + print(f'\033[33mWarning: --gpu-vendor {args.gpu_vendor_deprecated} is ignored because --execution-provider was explicitly set.\033[0m') +def get_available_execution_providers_full() -> List[str]: + """Returns the full names of available ONNX Runtime execution providers.""" + try: + return onnxruntime.get_available_providers() + except AttributeError: + # Fallback for very old versions or unexpected issues + print("\033[33mWarning: Could not dynamically get available providers. Falling back to common defaults.\033[0m") + # Provide a reasonable guess + defaults = ['CPUExecutionProvider'] + if _cuda_intended: defaults.insert(0, 'CUDAExecutionProvider') + if _is_rocm_selected: defaults.insert(0, 'ROCMExecutionProvider') + # Add others based on platform if needed + return defaults -def encode_execution_providers(execution_providers: List[str]) -> List[str]: - """Converts ONNX Runtime provider names to lowercase short names.""" - return [ep.replace('ExecutionProvider', '').lower() for ep in execution_providers] +def get_available_execution_providers_short() -> List[str]: + """Returns the short names (lowercase) of available ONNX Runtime execution providers.""" + full_names = get_available_execution_providers_full() + return [name.replace('ExecutionProvider', '').lower() for name in full_names] - -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'] +def decode_execution_providers(selected_short_names: List[str]) -> List[str]: + """Converts selected short names back to full ONNX Runtime provider names, preserving order and checking availability.""" + available_full_names = get_available_execution_providers_full() + available_short_map = {name.replace('ExecutionProvider', '').lower(): name for name in available_full_names} decoded_providers = [] - requested_providers_lower = [name.lower() for name in execution_providers_names] + valid_short_names_found = [] - # 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.") + for short_name in selected_short_names: + name_lower = short_name.lower() + if name_lower in available_short_map: + full_name = available_short_map[name_lower] + if full_name not in decoded_providers: # Avoid duplicates + decoded_providers.append(full_name) + valid_short_names_found.append(name_lower) 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) + print(f"\033[33mWarning: Requested execution provider '{short_name}' is not available or not recognized. Skipping.\033[0m") - # 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 decoded_providers: + print("\033[33mWarning: No valid execution providers selected or available. Falling back to CPU.\033[0m") + if 'CPUExecutionProvider' in available_full_names: + decoded_providers = ['CPUExecutionProvider'] + valid_short_names_found.append('cpu') + else: + print("\033[31mError: CPUExecutionProvider is not available in this build of ONNX Runtime. Cannot proceed.\033[0m") + sys.exit(1) # Critical error - 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 + print(f"[DLC.CORE] Using execution providers: {valid_short_names_found} (Full names: {decoded_providers})") + return decoded_providers def suggest_max_memory() -> int: - """Suggests a default max CPU RAM limit in GB based on available memory (heuristic).""" + """Suggests a default max CPU RAM limit in GB. Less critical now with GPU limits.""" 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 + total_ram_gb = psutil.virtual_memory().total / (1024 ** 3) + # Suggest slightly less than half of total RAM, capped at a reasonable upper limit (e.g., 64GB) + # and a minimum (e.g., 4GB) + suggested = max(4, min(int(total_ram_gb * 0.4), 64)) + # print(f"[DLC.CORE] Auto-suggesting max_memory: {suggested} GB (based on total system RAM: {total_ram_gb:.1f} GB)") + return suggested + except (ImportError, OSError): + print("[DLC.CORE] Info: psutil not found or failed. Using fallback default for max_memory suggestion (16 GB).") + # Fallback defaults similar to original code + if platform.system().lower() == 'darwin': + return 8 # Increased macOS default slightly + return 16 # Keep higher default for Linux/Windows def suggest_execution_providers() -> List[str]: - """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) + """Suggests a default list of execution providers based on availability and platform.""" + available_short = get_available_execution_providers_short() + preferred_providers = [] - # 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) + # Prioritize GPU providers if available + if 'cuda' in available_short: + preferred_providers.append('cuda') + elif 'rocm' in available_short: + preferred_providers.append('rocm') + elif 'dml' in available_short and platform.system().lower() == 'windows': + preferred_providers.append('dml') # DirectML on Windows + elif 'coreml' in available_short and platform.system().lower() == 'darwin': + preferred_providers.append('coreml') # CoreML on macOS - # Ensure CPU is always included as a fallback - if 'cpu' not in suggested and 'cpu' in available_providers_encoded: - suggested.append('cpu') + # Always include CPU as a fallback + if 'cpu' in available_short: + preferred_providers.append('cpu') + elif available_short: # If CPU is somehow missing, add the first available one + preferred_providers.append(available_short[0]) - # 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 + # If list is empty (shouldn't happen if get_available works), default to cpu + if not preferred_providers: + return ['cpu'] - return suggested + # print(f"[DLC.CORE] Suggested execution providers: {preferred_providers}") # Optional debug info + return preferred_providers def suggest_execution_threads() -> int: - """Suggests a sensible default number of execution threads based on logical CPU cores.""" + """Suggests a sensible default number of execution threads based on 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 + logical_cores = os.cpu_count() or 4 # Default to 4 if cpu_count fails + # Use slightly fewer threads than logical cores, capped. + # Good balance between parallelism and overhead. + suggested_threads = max(1, min(logical_cores - 1 if logical_cores > 1 else 1, 16)) + # Don't suggest 1 for CUDA/ROCm implicitly here, let user override or frame processors decide. + # The SessionOptions in the processors should handle provider-specific thread settings if needed. + # print(f"[DLC.CORE] Auto-suggesting execution_threads: {suggested_threads} (based on {logical_cores} logical cores)") + return suggested_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 + print("[DLC.CORE] Warning: os.cpu_count() not implemented. Using fallback default for execution_threads (4).") + return 4 # Fallback def limit_gpu_memory(fraction: float) -> None: - """Attempts to limit GPU memory usage via PyTorch (for CUDA) or TensorFlow.""" - gpu_limited = False + """Attempts to limit GPU memory usage, primarily via PyTorch if CUDA is used.""" + # Check if CUDAExecutionProvider is in the *actually selected* providers + if 'CUDAExecutionProvider' in modules.globals.execution_providers: + if _torch_cuda_available: + try: + # Ensure CUDA is initialized if needed (might not be necessary, but safe) + if not torch.cuda.is_initialized(): + torch.cuda.init() - # 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") + device_count = torch.cuda.device_count() + if device_count > 0: + # Limit memory on the default device (usually device 0) + # Note: This limits PyTorch's allocation pool. ONNX Runtime might manage + # its CUDA memory somewhat separately, but this can still help prevent + # PyTorch from grabbing everything. + print(f"[DLC.CORE] Attempting to limit PyTorch CUDA memory fraction to {fraction:.1%} on device 0") + torch.cuda.set_per_process_memory_fraction(fraction, 0) + # Optional: Check memory after setting limit + 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] PyTorch CUDA memory limit hint set. Device 0 Total: {total_mem / 1024**3:.2f} GB. " + f"PyTorch Reserved: {reserved_mem / 1024**3:.2f} GB, Allocated: {allocated_mem / 1024**3:.2f} GB.") + else: + print("\033[33mWarning: PyTorch reports no CUDA devices available, cannot set memory limit.\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).") + except RuntimeError as e: + print(f"\033[33mWarning: PyTorch CUDA runtime error during memory limit setting (may already be initialized?): {e}\033[0m") + except Exception as e: + print(f"\033[33mWarning: Failed to set PyTorch CUDA memory fraction: {e}\033[0m") + else: + # Only warn if PyTorch CUDA specifically isn't available, but CUDA EP was chosen. + if _cuda_intended: # Check original intent + print("\033[33mWarning: CUDAExecutionProvider selected, but PyTorch CUDA is not available. Cannot apply PyTorch memory limit.\033[0m") + # Add future limits for other providers if ONNX Runtime API supports it directly + # Example placeholder for potential future ONNX Runtime API: + # elif 'ROCMExecutionProvider' in modules.globals.execution_providers: + # try: + # # Hypothetical ONNX Runtime API + # ort_options = onnxruntime.SessionOptions() + # ort_options.add_provider_options('rocm', {'gpu_mem_limit': str(int(total_mem_bytes * fraction))}) + # print("[DLC.CORE] Note: ROCm memory limit set via ONNX Runtime provider options (if API exists).") + # except Exception as e: + # print(f"\033[33mWarning: Failed to set ROCm memory limit via hypothetical ORT options: {e}\033[0m") + # else: + # print("[DLC.CORE] GPU memory limit not applied (PyTorch CUDA not used or unavailable).") def limit_resources() -> None: - """Limits system resources like CPU RAM (best effort) and configures TF.""" - # 1. Limit CPU RAM (Best effort, platform dependent) + """Limits system resources like CPU RAM (best effort) and sets TensorFlow GPU options.""" + # 1. Limit CPU RAM (Best-effort, OS-dependent) if modules.globals.max_memory and modules.globals.max_memory > 0: limit_gb = modules.globals.max_memory limit_bytes = limit_gb * (1024 ** 3) + current_system = platform.system().lower() + try: - if platform.system().lower() in ['linux', 'darwin']: + if current_system == 'linux' or current_system == '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 + # RLIMIT_AS (virtual memory) is often more effective than RLIMIT_DATA + try: + soft, hard = resource.getrlimit(resource.RLIMIT_AS) + # Set soft limit; hard limit usually requires root. Don't exceed current hard limit. + new_soft = min(limit_bytes, hard) + resource.setrlimit(resource.RLIMIT_AS, (new_soft, hard)) + print(f"[DLC.CORE] Limited process virtual memory (CPU RAM approximation) soft limit towards ~{limit_gb} GB.") + except (ValueError, resource.error) as e: + print(f"\033[33mWarning: Failed to set virtual memory limit (RLIMIT_AS): {e}\033[0m") + # Fallback attempt using RLIMIT_DATA (less effective for total memory) + try: + soft_data, hard_data = resource.getrlimit(resource.RLIMIT_DATA) + new_soft_data = min(limit_bytes, hard_data) + resource.setrlimit(resource.RLIMIT_DATA, (new_soft_data, hard_data)) + print(f"[DLC.CORE] Limited process data segment (partial CPU RAM) soft limit towards ~{limit_gb} GB.") + except (ValueError, resource.error) as e_data: + print(f"\033[33mWarning: Failed to set data segment limit (RLIMIT_DATA): {e_data}\033[0m") + + elif current_system == 'windows': + # Windows memory limiting is complex. SetProcessWorkingSetSizeEx is more of a suggestion. + # Job Objects are the robust way but much more involved. Keep the hint for now. 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") + process_handle = kernel32.GetCurrentProcess() + # Flags: QUOTA_LIMITS_HARDWS_ENABLE (1) requires special privileges, use 0 for min/max hint only + # Using min=1MB, max=limit_bytes. Returns non-zero on success. + min_ws = ctypes.c_size_t(1024 * 1024) + max_ws = ctypes.c_size_t(limit_bytes) + if not kernel32.SetProcessWorkingSetSizeEx(process_handle, min_ws, max_ws, 0): + error_code = ctypes.get_last_error() + print(f"\033[33mWarning: Failed to set process working set size hint (Windows). Error code: {error_code}. This limit may not be enforced.\033[0m") else: - print(f"[DLC.CORE] Requested process working set size limit (Windows memory hint) max ~{limit_gb} GB.") + print(f"[DLC.CORE] Requested process working set size hint (Windows memory guidance) 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.") + print(f"\033[33mWarning: CPU RAM limiting not implemented for platform {current_system}. --max-memory ignored.\033[0m") - # 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 + except ImportError: + print(f"\033[33mWarning: 'resource' module (Unix) not available. Cannot limit CPU RAM via setrlimit.\033[0m") + except Exception as e: + print(f"\033[33mWarning: An unexpected error occurred during CPU RAM limiting: {e}\033[0m") + # else: + # print("[DLC.CORE] Info: CPU RAM limit (--max-memory) not set or disabled.") + + + # 2. Configure TensorFlow GPU memory (if TensorFlow is installed) + if _tensorflow_available: + try: + gpus = tensorflow.config.experimental.list_physical_devices('GPU') + if gpus: + configured_gpus = 0 + for gpu in gpus: + try: + # Allow memory growth instead of pre-allocating everything + tensorflow.config.experimental.set_memory_growth(gpu, True) + # print(f"[DLC.CORE] Enabled TensorFlow memory growth for GPU: {gpu.name}") + configured_gpus += 1 + 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_inner: # Catch other potential TF config errors + print(f"\033[33mWarning: Error configuring TensorFlow memory growth for {gpu.name}: {e_inner}\033[0m") + if configured_gpus > 0: + print(f"[DLC.CORE] Enabled TensorFlow memory growth for {configured_gpus} GPU(s).") + # else: + # print("[DLC.CORE] No TensorFlow physical GPUs detected.") + except Exception as e: + print(f"\033[33mWarning: Error listing or configuring TensorFlow GPU devices: {e}\033[0m") + # else: + # print("[DLC.CORE] TensorFlow not available, skipping TF GPU configuration.") def release_resources() -> None: - """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 + """Releases resources, especially GPU memory caches.""" + # Clear PyTorch CUDA cache if applicable and PyTorch CUDA is available + if 'CUDAExecutionProvider' in modules.globals.execution_providers and _torch_cuda_available: try: torch.cuda.empty_cache() - # print("[DLC.CORE] Cleared PyTorch CUDA cache.") # Can be verbose + # print("[DLC.CORE] Cleared PyTorch CUDA cache.") # Optional: uncomment for verbose logging 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") + # Add potential cleanup for other frameworks or ONNX Runtime sessions if needed + # (Usually session objects going out of scope and gc.collect() is sufficient for ORT C++ backend) - # 3. Explicitly run garbage collection (important!) + # Explicitly run garbage collection + # This helps release Python-level objects, which might then trigger + # the release of underlying resources (like ONNX Runtime session memory) gc.collect() - # print("[DLC.CORE] Ran garbage collection.") # Can be verbose + # print("[DLC.CORE] Ran garbage collector.") # Optional: uncomment for verbose logging 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 + """Performs essential pre-run checks for dependencies and versions.""" if sys.version_info < (3, 9): - update_status('Error: Python 3.9 or higher is required.', 'ERROR') - checks_passed = False - - # FFmpeg + update_status('Python version is not supported - please upgrade to Python 3.9 or higher.') + return False if not shutil.which('ffmpeg'): - update_status('Error: ffmpeg command was not found in your system PATH. Please install ffmpeg.', 'ERROR') - checks_passed = False + update_status('ffmpeg command not found in PATH. Please install ffmpeg and ensure it is accessible.') + return 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 + # ONNX Runtime was checked at import time, but double check here if needed. + # The import would have failed earlier if it's not installed. + # print(f"[DLC.CORE] Using ONNX Runtime version: {onnxruntime.__version__}") - # 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 + # TensorFlow check (optional, only issue warning if unavailable) + if not _tensorflow_available: + update_status('TensorFlow not found. Some features like GPU memory growth setting will be skipped.', scope='INFO') + # Decide if TF is strictly required by any processor. If so, change to error and return False. + # Currently, it seems only used for optional resource limiting. - # PyTorch (only if CUDA is selected for memory limiting) + # Check PyTorch availability *only if* CUDA EP is selected 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)') + if not _torch_available: + update_status('CUDAExecutionProvider selected, but PyTorch is not installed. Install PyTorch with CUDA support (see PyTorch website).', scope='ERROR') + return False + if not _torch_cuda_available: + update_status('CUDAExecutionProvider selected, but torch.cuda.is_available() is False. Check PyTorch CUDA installation, GPU drivers, and CUDA toolkit compatibility.', scope='ERROR') + return False + # Check if selected video encoder potentially requires specific hardware/drivers (e.g., NVENC) + if modules.globals.video_encoder in ['h264_nvenc', 'hevc_nvenc']: + # This check is basic. FFmpeg needs to be compiled with NVENC support, + # and NVIDIA drivers must be installed. We can't easily verify this from Python. + # Just issue an informational note. + update_status(f"Selected video encoder '{modules.globals.video_encoder}' requires an NVIDIA GPU and correctly configured FFmpeg/drivers.", scope='INFO') + if 'CUDAExecutionProvider' not in modules.globals.execution_providers: + update_status(f"Warning: NVENC encoder selected, but CUDAExecutionProvider is not active. Ensure FFmpeg can access the GPU independently.", scope='WARN') - # 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 + return True def update_status(message: str, scope: str = 'DLC.CORE') -> None: """Prints status messages and updates UI if not headless.""" - log_message = f'[{scope}] {message}' - print(log_message) + formatted_message = f'[{scope}] {message}' + print(formatted_message) if not modules.globals.headless: - 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}") + # Ensure ui module and update_status function exist and are callable + if hasattr(ui, 'update_status') and callable(ui.update_status): + try: + # Use a mechanism that's safe for cross-thread UI updates if necessary + # (e.g., queue or wx.CallAfter if using wxPython) + # Assuming direct call is okay for now based on original structure. + ui.update_status(message) # Pass the original message without scope prefix + except Exception as e: + # Avoid crashing core process for UI update errors + print(f"[DLC.CORE] Error updating UI status: {e}") + # else: + # print("[DLC.CORE] UI or ui.update_status not available for status update.") -# --- Main Processing Logic --- - def start() -> None: - """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: - 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: - update_status(f"Error during NSFW check for image: {e}. Continuing processing.", "NSFW") - - # --- Process --- + """Main processing logic: routes to image or video processing.""" + # Ensure frame processors are ready (this also initializes them) 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}") - - # 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') + active_processors = get_frame_processors_modules(modules.globals.frame_processors) + if not active_processors: + update_status("No valid frame processors selected or loaded. Aborting.", "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 + all_processors_initialized = True + for frame_processor in active_processors: + update_status(f'Initializing frame processor: {getattr(frame_processor, "NAME", "UnknownProcessor")}...') + # The pre_start method should handle model loading and initial setup. + # It might raise exceptions or return False on failure. + if not hasattr(frame_processor, 'pre_start') or not callable(frame_processor.pre_start): + update_status(f'Processor {getattr(frame_processor, "NAME", "UnknownProcessor")} lacks a pre_start method.', 'WARN') + continue # Or treat as failure? + + if not frame_processor.pre_start(): + update_status(f'Initialization failed for {getattr(frame_processor, "NAME", "UnknownProcessor")}. Aborting.', 'ERROR') + all_processors_initialized = False + break # Stop initialization if one fails + + if not all_processors_initialized: + return # Abort if any processor failed to initialize + + except Exception as e: + update_status(f"Error during frame processor initialization: {e}", "ERROR") + import traceback + traceback.print_exc() + return + + # --- Route based on target type --- + if not modules.globals.target_path or not os.path.exists(modules.globals.target_path): + update_status(f"Target path '{modules.globals.target_path}' not found or not specified.", "ERROR") + return + + if has_image_extension(modules.globals.target_path) and is_image(modules.globals.target_path): + process_image_target(active_processors) + elif is_video(modules.globals.target_path): + process_video_target(active_processors) + else: + update_status(f"Target path '{modules.globals.target_path}' is not a recognized image or video file.", "ERROR") + + +def process_image_target(active_processors: List) -> None: + """Handles processing when the target is an image.""" + update_status('Processing image target...') + # NSFW check (basic, for image only) + if modules.globals.nsfw_filter: + update_status('Checking image for NSFW content...', 'NSFW') + # Assuming ui.check_and_ignore_nsfw is suitable for this + if ui.check_and_ignore_nsfw(modules.globals.target_path, destroy): + update_status('NSFW content detected and processing skipped.', 'NSFW') + return # Stop processing + + try: + # Ensure source path exists if needed by processors + if not modules.globals.source_path or not os.path.exists(modules.globals.source_path): + # Face swapping requires a source, enhancer might not. Check processor needs? + if any(proc.NAME == 'face_swapper' for proc in active_processors): # Example check + update_status(f"Source image path '{modules.globals.source_path}' not found or not specified, required for face swapping.", "ERROR") + return + + # Ensure output directory exists + output_dir = os.path.dirname(modules.globals.output_path) + if output_dir and not os.path.exists(output_dir): + try: + os.makedirs(output_dir, exist_ok=True) + print(f"[DLC.CORE] Created output directory: {output_dir}") + except OSError as e: + update_status(f"Error creating output directory '{output_dir}': {e}", "ERROR") + return + + # Copy target to output path first to preserve metadata if possible and safe + final_output_path = modules.globals.output_path + temp_output_path = None # Use a temp path if overwriting source/target directly + + # Avoid overwriting input files directly during processing if they are the same as output + if os.path.abspath(modules.globals.target_path) == os.path.abspath(final_output_path) or \ + (modules.globals.source_path and os.path.abspath(modules.globals.source_path) == os.path.abspath(final_output_path)): + temp_output_path = os.path.join(output_dir, f"temp_image_{os.path.basename(final_output_path)}") + print(f"[DLC.CORE] Output path conflicts with input, using temporary file: {temp_output_path}") + shutil.copy2(modules.globals.target_path, temp_output_path) + current_processing_file = temp_output_path + else: + # Copy target to final destination to start + shutil.copy2(modules.globals.target_path, final_output_path) + current_processing_file = final_output_path + + + # Apply processors sequentially to the current file path + source_for_processing = modules.globals.source_path + output_for_processing = current_processing_file # Processors modify this file + + for frame_processor in active_processors: + processor_name = getattr(frame_processor, "NAME", "UnknownProcessor") 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 - + # Pass source, input_path (current state), output_path (same as input for in-place modification) + frame_processor.process_image(source_for_processing, output_for_processing, output_for_processing) + release_resources() # Release memory after each processor step except Exception as e: - update_status(f'Error applying processor {processor_name}: {e}', 'ERROR') - import traceback - traceback.print_exc() - return # Stop processing on error + update_status(f'Error during {processor_name} processing: {e}', 'ERROR') + import traceback + traceback.print_exc() + # Optionally clean up temp file and abort + if temp_output_path and os.path.exists(temp_output_path): os.remove(temp_output_path) + return - # --- 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') + # If a temporary file was used, move it to the final destination + if temp_output_path: + try: + shutil.move(temp_output_path, final_output_path) + print(f"[DLC.CORE] Moved temporary result to final output: {final_output_path}") + except Exception as e: + update_status(f"Error moving temporary file to final output: {e}", "ERROR") + # Temp file might still exist, leave it for inspection? + return - except Exception as e: - update_status(f'Error saving processed image: {e}', 'ERROR') + # Final check if output exists and is an image + if os.path.exists(final_output_path) and is_image(final_output_path): + update_status('Processing image finished successfully.') + else: + update_status('Processing image failed: Output file not found or invalid after processing.', 'ERROR') except Exception as e: update_status(f'An unexpected error occurred during image processing: {e}', 'ERROR') import traceback traceback.print_exc() + # Clean up potentially corrupted output/temp file? Be cautious. + # if temp_output_path and os.path.exists(temp_output_path): os.remove(temp_output_path) + # if os.path.exists(final_output_path) and current_processing_file == final_output_path: # Careful not to delete original if copy failed + # Consider what to do on failure - delete potentially corrupt output? -def process_video(): - """Handles the video processing workflow with optimized frame handling.""" - update_status('Processing video: {}'.format(os.path.basename(modules.globals.target_path))) +def process_video_target(active_processors: List) -> None: + """Handles processing when the target is a video.""" + update_status('Processing video target...') - # --- NSFW Check (Basic - Check first frame or predict_video) --- + # Basic check for source if needed (similar to image processing) + if not modules.globals.source_path or not os.path.exists(modules.globals.source_path): + if any(proc.NAME == 'face_swapper' for proc in active_processors): + update_status(f"Source image path '{modules.globals.source_path}' not found or not specified, required for face swapping.", "ERROR") + return + + # NSFW Check (Could be enhanced to sample frames, currently basic/skipped for 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") + update_status('NSFW check for video is basic/experimental. Checking first frame...', 'NSFW') + # Consider implementing frame sampling for a more robust check if needed + # if ui.check_and_ignore_nsfw(modules.globals.target_path, destroy): # This might not work well for video + # update_status('NSFW content potentially detected (based on first frame check). Skipping.', 'NSFW') + # return + update_status('NSFW check passed or skipped for video.', 'NSFW INFO') - # --- Prepare Temp Environment --- - temp_output_video_path = None # For intermediate video file - video_fps = 30.0 # Default FPS + temp_output_video_path = None + temp_frame_dir = None # Keep track of temp frame directory 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. + # --- Frame Extraction --- + # map_faces might imply frames are already extracted or handled differently 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 + update_status('Creating temporary resources for video frames...') + # create_temp should return the path to the temp directory created + temp_frame_dir = create_temp(modules.globals.target_path) + if not temp_frame_dir: + update_status("Failed to create temporary directory for frames.", "ERROR") + return + + update_status('Extracting video frames...') + # extract_frames needs the temp directory path + # It should also ideally set modules.globals.video_fps based on the extracted video + extract_frames(modules.globals.target_path, temp_frame_dir) # Pass temp dir + update_status('Frame extraction complete.') + else: + update_status('Skipping frame extraction due to --map-faces flag.', 'INFO') + # Assuming frames are already in the expected temp location or handled by processors + temp_frame_dir = os.path.join(modules.globals.TEMP_DIRECTORY, os.path.basename(modules.globals.target_path)) # Need consistent temp path logic - # Get paths to frames (must exist either way) - temp_frame_paths = get_temp_frame_paths(modules.globals.target_path) + # Get paths to frames (extracted or pre-existing) + temp_frame_paths = get_temp_frame_paths(modules.globals.target_path) # This needs to know the temp dir structure 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 + update_status('No frames found to process. Check temp folder or extraction step.', 'ERROR') + # Clean up if temp dir was created + if temp_frame_dir and not modules.globals.keep_frames: clean_temp(modules.globals.target_path) return - num_frames = len(temp_frame_paths) - update_status(f'Processing {num_frames} frames...') + update_status(f'Processing {len(temp_frame_paths)} frames...') - # Determine Target FPS - if modules.globals.keep_fps: - 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: - 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: + # --- Frame Processing --- + source_for_processing = modules.globals.source_path + for frame_processor in active_processors: + processor_name = getattr(frame_processor, "NAME", "UnknownProcessor") + update_status(f'Applying {processor_name}...', processor_name) 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 + # process_video should modify frames in-place in the temp directory + # It needs the source path and the list of frame paths + frame_processor.process_video(source_for_processing, temp_frame_paths) + release_resources() # Release memory after each processor completes its pass + except Exception as e: + update_status(f'Error during {processor_name} frame processing: {e}', 'ERROR') + import traceback + traceback.print_exc() + # Abort processing + # Clean up temp frames if not keeping them + if temp_frame_dir and not modules.globals.keep_frames: clean_temp(modules.globals.target_path) + return - # Frame dimensions for potential checks later - # height, width = target_frame.shape[:2] + # --- Video Creation --- + update_status('Reconstructing video from processed frames...') + fps = modules.globals.video_fps # Should be set by extract_frames or detected earlier - # 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 + if modules.globals.keep_fps: + # Use the FPS detected during extraction (should be stored in globals.video_fps) + if fps is None: + update_status('Original FPS not detected during extraction, attempting fallback detection...', 'WARN') + detected_fps = detect_fps(modules.globals.target_path) + if detected_fps is not None: + fps = detected_fps + modules.globals.video_fps = fps # Store it back + update_status(f'Using fallback detected FPS: {fps:.2f}') + else: + fps = 30.0 # Ultimate fallback + update_status("Could not detect FPS, using default 30.", "WARN") + else: + update_status(f'Using original detected FPS: {fps:.2f}') + else: + fps = 30.0 # Use default fps if not keeping original + update_status(f'Using fixed FPS: {fps:.2f}') - temp_frame = processor.process_frame(processor_params) + # Define a temporary path for the video created *without* audio + output_dir = os.path.dirname(modules.globals.output_path) + if not output_dir: output_dir = '.' # Handle case where output is in current dir + temp_output_video_filename = f"temp_{os.path.basename(modules.globals.output_path)}" + # Ensure the temp filename doesn't clash if multiple runs happen concurrently (less likely in this app) + temp_output_video_path = os.path.join(output_dir, temp_output_video_filename) - 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 + # create_video needs the target path (for context?), fps, and the *temp* output path + # It internally uses get_temp_frame_paths based on the target_path context. + create_video(modules.globals.target_path, fps, temp_output_video_path) - # 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 --- + # --- Audio Handling --- final_output_path = modules.globals.output_path if modules.globals.keep_audio: - update_status('Restoring audio...', 'FFmpeg') + update_status('Restoring audio...') if not modules.globals.keep_fps: - update_status('Warning: Audio restoration enabled without --keep-fps. Sync issues may occur.', 'WARN') + update_status('Audio restoration may cause sync issues as FPS was not kept.', '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 needs: original video (with audio), temp video (no audio), final output path + restore_success = restore_audio(modules.globals.target_path, temp_output_video_path, final_output_path) - # 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: + if restore_success: update_status('Audio restoration complete.') + # Remove the intermediate temp video *after* successful audio merge + if os.path.exists(temp_output_video_path): + try: os.remove(temp_output_video_path) + except OSError as e: print(f"\033[33mWarning: Could not remove intermediate video file {temp_output_video_path}: {e}\033[0m") + temp_output_video_path = None # Mark as removed 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') + update_status('Audio restoration failed. The output video will be silent.', 'ERROR') + # Audio failed, move the silent video to the final path as a fallback? + update_status('Moving silent video to final output path as fallback.') + try: + shutil.move(temp_output_video_path, final_output_path) + temp_output_video_path = None # Mark as moved + except Exception as e: + update_status(f"Error moving silent video to final output: {e}", "ERROR") + # Both audio failed and move failed, temp video might still exist 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 + if os.path.abspath(temp_output_video_path) == os.path.abspath(final_output_path): + update_status("Temporary path is the same as final path, no move needed.", "INFO") + temp_output_video_path = None # No deletion needed later 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') - + # Ensure target directory exists (should already, but double check) + os.makedirs(os.path.dirname(final_output_path), exist_ok=True) + shutil.move(temp_output_video_path, final_output_path) + temp_output_video_path = None # Mark as moved successfully + except Exception as e: + update_status(f"Error moving temporary video to final output: {e}", "ERROR") + # The temp video might still exist # --- Validation --- if os.path.exists(final_output_path) and is_video(final_output_path): - update_status('Video processing finished successfully.') + update_status('Processing video finished successfully.') else: - update_status('Error: Final output video file not found or invalid after processing.', 'ERROR') + update_status('Processing video failed: Output 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() + traceback.print_exc() # Print detailed traceback for debugging 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") + # --- Cleanup --- + # Clean up temporary frames if they exist and keep_frames is false + if temp_frame_dir and os.path.exists(temp_frame_dir) and not modules.globals.keep_frames: + update_status("Cleaning up temporary frames...") + clean_temp(modules.globals.target_path) # clean_temp uses target_path context to find the dir - # Remove intermediate temp video file if it exists and wasn't moved + # Clean up intermediate temp video file if it still exists (e.g., audio failed and move failed) 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") + print(f"[DLC.CORE] Removed intermediate temporary video file: {temp_output_video_path}") 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() + print(f"\033[33mWarning: Could not remove intermediate temporary video file {temp_output_video_path}: {e}\033[0m") 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") + update_status("Cleaning up temporary resources...", "CLEANUP") + # Use the context of target_path to find the temp directory + if modules.globals.target_path and not modules.globals.keep_frames: clean_temp(modules.globals.target_path) - - # 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 - + release_resources() # Final resource release (GPU cache, GC) 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 + print("[DLC.CORE] Exiting application.") + os._exit(0) # Use os._exit for a more forceful exit if sys.exit hangs (e.g., due to threads) + # sys.exit(0) # Standard exit def run() -> None: - """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' + """Parses arguments, sets up the environment, performs checks, and starts processing or UI.""" + try: + parse_args() # Parse arguments first to set globals like execution_providers, paths, etc. - parser = parse_args() # Parse arguments first to set globals + # Apply GPU Memory Limit early, requires execution_providers to be set by parse_args + limit_gpu_memory(GPU_MEMORY_LIMIT_FRACTION) - # Apply GPU Memory Limit early, requires execution_providers to be set - limit_gpu_memory(GPU_MEMORY_LIMIT_FRACTION) + # Limit other resources (CPU RAM approximation, TF GPU options) + # Call this *after* potential PyTorch limit and TensorFlow import check + limit_resources() - # Perform pre-checks (dependencies, versions, paths) - if not pre_check(): - # 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 + # Perform pre-checks (dependencies like Python version, ffmpeg, libraries, provider checks) + update_status("Performing pre-run checks...") + if not pre_check(): + update_status("Pre-run checks failed. Please see messages above.", "ERROR") + # destroy(to_quit=True) # Don't call destroy here, let the main try/finally handle it + return # Exit run() function - # Limit other resources (CPU RAM, TF GPU options) - limit_resources() + update_status("Pre-run checks passed.") - # --- 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: - for processor_module in active_processor_modules: - processor_name = getattr(processor_module, 'NAME', 'UnknownProcessor') + # Pre-check frame processors (model downloads, requirements within processors) + # This needs globals to be set by parse_args and should happen before starting work. + active_processor_modules = get_frame_processors_modules(modules.globals.frame_processors) + all_processors_reqs_met = True + for frame_processor_module in active_processor_modules: + processor_name = getattr(frame_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() + if hasattr(frame_processor_module, 'pre_check') and callable(frame_processor_module.pre_check): + if not frame_processor_module.pre_check(): + update_status(f'Requirements check failed for {processor_name}. See processor messages for details.', 'ERROR') + all_processors_reqs_met = False + # Don't break early, check all processors to report all issues 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") + update_status(f'Processor {processor_name} does not have a pre_check method. Assuming requirements met.', 'WARN') - 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 + if not all_processors_reqs_met: + update_status('Some frame processors failed requirement checks. Please resolve the issues and retry.', 'ERROR') + # destroy(to_quit=True) # Let finally handle cleanup + return + + update_status("All frame processor requirements met.") + + # --- Start processing (headless) or launch UI --- + if modules.globals.headless: + # Check for essential paths in headless mode + if not modules.globals.source_path: + update_status("Error: Headless mode requires --source argument.", "ERROR") + # program.print_help() # Can't access program object here easily + print("Use -h or --help for usage details.") + return + if not modules.globals.target_path: + update_status("Error: Headless mode requires --target argument.", "ERROR") + print("Use -h or --help for usage details.") + return + if not modules.globals.output_path: + update_status("Error: Headless mode requires --output argument.", "ERROR") + print("Use -h or --help for usage details.") + return + + update_status('Running in headless mode.') + start() # Execute the main processing logic + # destroy() will be called by the finally block + + 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: + # Pass start (processing function) and destroy (cleanup) to the UI + window = ui.init(start, destroy_wrapper, modules.globals.lang) + if window: + window.mainloop() # Start the UI event loop + else: + update_status("UI initialization failed.", "ERROR") + except Exception as e: + update_status(f"Error initializing or running the UI: {e}", "FATAL") + import traceback + traceback.print_exc() + # Attempt cleanup even if UI fails + # destroy(to_quit=True) # Let finally handle it + + except Exception as e: + # Catch any unexpected errors during setup or execution + update_status(f"A critical error occurred: {e}", "FATAL") + import traceback + traceback.print_exc() + + finally: + # Ensure cleanup happens regardless of success or failure + destroy(to_quit=True) # Clean up and exit # --- 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) - + # This ensures 'run()' is called only when the script is executed directly run() + # --- END OF FILE core.py --- \ No newline at end of file