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