From 56cddde87c587a22d120417d484ca5294e803975 Mon Sep 17 00:00:00 2001
From: Kenneth Estanislao <hacksider@gmail.com>
Date: Thu, 27 Mar 2025 02:56:08 +0800
Subject: [PATCH] Update core.py
Allows CUDA 12 to be used on this version
---
modules/core.py | 1563 ++++++++++++++++++++++-------------------------
1 file changed, 743 insertions(+), 820 deletions(-)
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