From 5835cd2b5b47d9889d7a26be7c8dc93cfeacb1af Mon Sep 17 00:00:00 2001 From: henryruhs Date: Thu, 22 Jun 2023 11:09:43 +0200 Subject: [PATCH] Using Queue for dynamic thread processing --- roop/processors/frame/core.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/roop/processors/frame/core.py b/roop/processors/frame/core.py index 74c780c..b247dac 100644 --- a/roop/processors/frame/core.py +++ b/roop/processors/frame/core.py @@ -3,6 +3,7 @@ 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 tqdm import tqdm @@ -43,13 +44,21 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType def multi_process_frame(source_path: str, 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 = [] - for path in temp_frame_paths: - future = executor.submit(process_frames, source_path, [path], update) + queue = create_queue(temp_frame_paths) + while not queue.empty(): + future = executor.submit(process_frames, source_path, [queue.get()], update) futures.append(future) for future in as_completed(futures): future.result() +def create_queue(temp_frame_paths: List[str]) -> Queue: + queue = Queue() + for frame_path in temp_frame_paths: + queue.put(frame_path) + return queue + + def process_video(source_path: str, 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)