roop/processors/frame/core.py

import os
import importlib
import psutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
from types import ModuleType
from typing import Any, List, Callable
import logging
from tqdm import tqdm

import roop

# Initialize logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

FRAME_PROCESSORS_MODULES: List[ModuleType] = []
FRAME_PROCESSORS_INTERFACE = [
    'pre_check',
    'pre_start',
    'process_frame',
    'process_frames',
    'process_image',
    'process_video',
    'post_process'
]

def load_frame_processor_module(frame_processor: str) -> Any:
    try:
        frame_processor_module = importlib.import_module(f'roop.processors.frame.{frame_processor}')
        for method_name in FRAME_PROCESSORS_INTERFACE:
            if not hasattr(frame_processor_module, method_name):
                raise NotImplementedError(f"Missing method: {method_name}")
        logging.info(f'Successfully loaded frame processor module: {frame_processor}')
    except (ImportError, NotImplementedError) as e:
        logging.error(f'Error loading frame processor {frame_processor}: {e}')
        raise SystemExit(f'Frame processor {frame_processor} crashed.')

    return frame_processor_module

def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
    global FRAME_PROCESSORS_MODULES

    if not FRAME_PROCESSORS_MODULES:
        for frame_processor in frame_processors:
            frame_processor_module = load_frame_processor_module(frame_processor)
            FRAME_PROCESSORS_MODULES.append(frame_processor_module)
    return FRAME_PROCESSORS_MODULES

def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Callable[[], None]], None], update: Callable[[], None]) -> None:
    try:
        with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
            futures = []
            queue = create_queue(temp_frame_paths)
            queue_per_future = len(temp_frame_paths) // roop.globals.execution_threads
            while not queue.empty():
                future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
                futures.append(future)
                logging.info('Submitted future for processing frames.')
                
            for future in as_completed(futures):
                try:
                    future.result()
                    logging.info('Frame processing completed for a future.')
                except Exception as e:
                    logging.error(f'Error in processing frame: {e}')
    except Exception as e:
        logging.error(f'Error in multi-process frame: {e}')

def create_queue(temp_frame_paths: List[str]) -> Queue:
    queue = Queue()
    for frame_path in temp_frame_paths:
        queue.put(frame_path)
    logging.info('Queue created with frame paths.')
    return queue

def pick_queue(queue: Queue, queue_per_future: int) -> List[str]:
    picked_items = []
    for _ in range(queue_per_future):
        if not queue.empty():
            picked_items.append(queue.get())
    logging.info(f'Picked {len(picked_items)} items from queue for processing.')
    return picked_items

def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Callable[[], None]], None]) -> None:
    progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
    total = len(frame_paths)
    try:
        with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format, mininterval=0.1) as progress:
            multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
    except Exception as e:
        logging.error(f'Error in processing video: {e}')

def update_progress(progress: Any = None) -> None:
    try:
        process = psutil.Process(os.getpid())
        memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
        if progress:
            progress.set_postfix({
                'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
                'execution_providers': roop.globals.execution_providers,
                'execution_threads': roop.globals.execution_threads
            })
            progress.refresh()
            progress.update(1)
        logging.info(f'Updated progress: {progress.n}/{progress.total} frames processed. Memory usage: {memory_usage:.2f}GB')
    except Exception as e:
        logging.error(f'Error updating progress: {e}')