roop/core.py

#!/usr/bin/env python3

import os
import sys
import platform
import signal
import shutil
import argparse
import warnings
from typing import List

import torch
import onnxruntime
import tensorflow

import roop.globals
import roop.metadata
import roop.ui as ui
from roop.predicter import predict_image, predict_video
from roop.processors.frame.core import get_frame_processors_modules
from roop.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
)

# Reduce TensorFlow log level and configure threading for torch
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
if any(arg.startswith('--execution-provider') for arg in sys.argv):
    os.environ['OMP_NUM_THREADS'] = '1'

warnings.filterwarnings('ignore', category=FutureWarning, module='insightface')
warnings.filterwarnings('ignore', category=UserWarning, module='torchvision')


def parse_args() -> None:
    """Parse command-line arguments and configure global settings."""
    signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
    
    parser = argparse.ArgumentParser(
        formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=100)
    )
    parser.add_argument('-s', '--source', help='Path to the source image', dest='source_path')
    parser.add_argument('-t', '--target', help='Path to the target image or video', dest='target_path')
    parser.add_argument('-o', '--output', help='Path to the output file or directory', dest='output_path')
    parser.add_argument('--frame-processor', help='Frame processors (choices: face_swapper, face_enhancer, ...)', dest='frame_processor', default=['face_swapper'], nargs='+')
    parser.add_argument('--keep-fps', help='Keep original FPS', dest='keep_fps', action='store_true', default=False)
    parser.add_argument('--keep-audio', help='Keep original audio', dest='keep_audio', action='store_true', default=True)
    parser.add_argument('--keep-frames', help='Keep temporary frames', dest='keep_frames', action='store_true', default=False)
    parser.add_argument('--many-faces', help='Process every face', dest='many_faces', action='store_true', default=False)
    parser.add_argument('--video-encoder', help='Output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
    parser.add_argument('--video-quality', help='Output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]')
    parser.add_argument('--max-memory', help='Maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
    parser.add_argument('--execution-provider', help='Available execution provider (choices: cpu, ...)', dest='execution_provider', default=['cpu'], choices=suggest_execution_providers(), nargs='+')
    parser.add_argument('--execution-threads', help='Number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads())
    parser.add_argument('-v', '--version', action='version', version=f'{roop.metadata.name} {roop.metadata.version}')
    
    args = parser.parse_args()

    roop.globals.source_path = args.source_path
    roop.globals.target_path = args.target_path
    roop.globals.output_path = normalize_output_path(roop.globals.source_path, roop.globals.target_path, args.output_path)
    roop.globals.frame_processors = args.frame_processor
    roop.globals.headless = args.source_path or args.target_path or args.output_path
    roop.globals.keep_fps = args.keep_fps
    roop.globals.keep_audio = args.keep_audio
    roop.globals.keep_frames = args.keep_frames
    roop.globals.many_faces = args.many_faces
    roop.globals.video_encoder = args.video_encoder
    roop.globals.video_quality = args.video_quality
    roop.globals.max_memory = args.max_memory
    roop.globals.execution_providers = decode_execution_providers(args.execution_provider)
    roop.globals.execution_threads = args.execution_threads


def encode_execution_providers(execution_providers: List[str]) -> List[str]:
    """Convert execution providers to their encoded form."""
    return [provider.replace('ExecutionProvider', '').lower() for provider in execution_providers]


def decode_execution_providers(execution_providers: List[str]) -> List[str]:
    """Decode execution providers from their encoded form."""
    return [provider for provider, encoded_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers()))
            if any(execution_provider in encoded_provider for execution_provider in execution_providers)]


def suggest_max_memory() -> int:
    """Suggest maximum memory in GB based on the operating system."""
    if platform.system().lower() == 'darwin':
        return 10
    return 14


def suggest_execution_providers() -> List[str]:
    """Suggest available execution providers based on ONNX Runtime."""
    return encode_execution_providers(onnxruntime.get_available_providers())


def suggest_execution_threads() -> int:
    """Suggest the number of execution threads based on execution providers."""
    if 'DmlExecutionProvider' in roop.globals.execution_providers or 'ROCMExecutionProvider' in roop.globals.execution_providers:
        return 1
    return 8


def limit_resources() -> None:
    """Limit GPU and RAM resources based on configuration."""
    # Prevent TensorFlow memory leak
    gpus = tensorflow.config.experimental.list_physical_devices('GPU')
    for gpu in gpus:
        tensorflow.config.experimental.set_virtual_device_configuration(gpu, [
            tensorflow.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)
        ])
    
    # Limit memory usage
    if roop.globals.max_memory:
        memory = roop.globals.max_memory * 1024 ** 3
        if platform.system().lower() == 'darwin':
            memory = roop.globals.max_memory * 1024 ** 6
        elif platform.system().lower() == 'windows':
            import ctypes
            kernel32 = ctypes.windll.kernel32
            kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
        else:
            import resource
            resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))


def release_resources() -> None:
    """Release resources such as GPU cache."""
    if 'CUDAExecutionProvider' in roop.globals.execution_providers:
        torch.cuda.empty_cache()


def pre_check() -> bool:
    """Perform preliminary checks before starting the processing."""
    if sys.version_info < (3, 9):
        update_status('Python version is not supported - please upgrade to 3.9 or higher.')
        return False
    if not shutil.which('ffmpeg'):
        update_status('ffmpeg is not installed.')
        return False
    return True


def update_status(message: str, scope: str = 'ROOP.CORE') -> None:
    """Update status message to the console or UI."""
    print(f'[{scope}] {message}')
    if not roop.globals.headless:
        ui.update_status(message)


def start() -> None:
    """Start the processing based on the configuration and input."""
    for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
        if not frame_processor.pre_start():
            return

    # Process image to image
    if has_image_extension(roop.globals.target_path):
        if predict_image(roop.globals.target_path):
            destroy()
        shutil.copy2(roop.globals.target_path, roop.globals.output_path)
        for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
            update_status('Progressing...', frame_processor.NAME)
            frame_processor.process_image(roop.globals.source_path, roop.globals.output_path, roop.globals.output_path)
            frame_processor.post_process()
            release_resources()
        update_status('Processing to image succeeded!' if is_image(roop.globals.target_path) else 'Processing to image failed!')
        return

    # Process image to video
    if predict_video(roop.globals.target_path):
        destroy()

    update_status('Creating temp resources...')
    create_temp(roop.globals.target_path)
    update_status('Extracting frames...')
    extract_frames(roop.globals.target_path)
    temp_frame_paths = get_temp_frame_paths(roop.globals.target_path)

    for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
        update_status('Progressing...', frame_processor.NAME)
        frame_processor.process_video(roop.globals.source_path, temp_frame_paths)
        frame_processor.post_process()
        release_resources()

    # Handle FPS
    if roop.globals.keep_fps:
        update_status('Detecting FPS...')
        fps = detect_fps(roop.globals.target_path)
        update_status(f'Creating video with {fps} FPS...')
        create_video(roop.globals.target_path, fps)
    else:
        update_status('Creating video with 30.0 FPS...')
        create_video(roop.globals.target_path)

    # Handle audio
    if roop.globals.keep_audio:
        update_status('Restoring audio...' if roop.globals.keep_fps else 'Restoring audio and creating final video...')
        restore_audio(roop.globals.target_path)

    move_temp(roop.globals.target_path)
    clean_temp()
    update_status('Processing succeeded!')
    release_resources()


def destroy() -> None:
    """Cleanup and exit the program."""
    update_status('Cleaning up and exiting...')
    clean_temp()
    sys.exit()


if __name__ == '__main__':
    parse_args()
    if pre_check():
        limit_resources()
        start()