Spaces:
Running
Running
File size: 4,157 Bytes
09b19a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import os
import sys
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
from roop.typing import Face
from tqdm import tqdm
import roop
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:
module_name = f'roop.processors.frame.{frame_processor}'
print(f'Loading {module_name}')
frame_processor_module = importlib.import_module(module_name)
for method_name in FRAME_PROCESSORS_INTERFACE:
if not hasattr(frame_processor_module, method_name):
raise NotImplementedError
except ModuleNotFoundError:
sys.exit(f'Frame processor {frame_processor} not found.')
except NotImplementedError:
sys.exit(f'Frame processor {frame_processor} not implemented correctly.')
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(is_batch: bool, source_face: Face, target_face: Face, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
futures = []
queue = create_queue(temp_frame_paths)
queue_per_future = max(len(temp_frame_paths) // roop.globals.execution_threads, 1)
while not queue.empty():
future = executor.submit(process_frames, is_batch, source_face, target_face, pick_queue(queue, queue_per_future), update)
futures.append(future)
for future in as_completed(futures):
future.result()
def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
queue: Queue[str] = Queue()
for frame_path in temp_frame_paths:
queue.put(frame_path)
return queue
def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
queues = []
for _ in range(queue_per_future):
if not queue.empty():
queues.append(queue.get())
return queues
def process_batch(source_face: Face, target_face: Face, frame_paths: list[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
total = len(frame_paths)
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
multi_process_frame(True, source_face, target_face, frame_paths, process_frames, lambda: update_progress(progress))
def process_video(source_face: Face, target_face: Face, frame_paths: list[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
total = len(frame_paths)
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
multi_process_frame(False, source_face, target_face, frame_paths, process_frames, lambda: update_progress(progress))
def update_progress(progress: Any = None) -> None:
process = psutil.Process(os.getpid())
memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
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)
|