Spaces:
Running
Running
Update roop/processors/frame/core.py
Browse files- roop/processors/frame/core.py +18 -11
roop/processors/frame/core.py
CHANGED
@@ -5,10 +5,14 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
5 |
from queue import Queue
|
6 |
from types import ModuleType
|
7 |
from typing import Any, List, Callable
|
|
|
8 |
from tqdm import tqdm
|
9 |
|
10 |
import roop
|
11 |
|
|
|
|
|
|
|
12 |
FRAME_PROCESSORS_MODULES: List[ModuleType] = []
|
13 |
FRAME_PROCESSORS_INTERFACE = [
|
14 |
'pre_check',
|
@@ -20,18 +24,18 @@ FRAME_PROCESSORS_INTERFACE = [
|
|
20 |
'post_process'
|
21 |
]
|
22 |
|
23 |
-
|
24 |
def load_frame_processor_module(frame_processor: str) -> Any:
|
25 |
try:
|
26 |
frame_processor_module = importlib.import_module(f'roop.processors.frame.{frame_processor}')
|
27 |
for method_name in FRAME_PROCESSORS_INTERFACE:
|
28 |
if not hasattr(frame_processor_module, method_name):
|
29 |
raise NotImplementedError
|
30 |
-
|
|
|
|
|
31 |
quit(f'Frame processor {frame_processor} crashed.')
|
32 |
return frame_processor_module
|
33 |
|
34 |
-
|
35 |
def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
|
36 |
global FRAME_PROCESSORS_MODULES
|
37 |
|
@@ -41,7 +45,6 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType
|
|
41 |
FRAME_PROCESSORS_MODULES.append(frame_processor_module)
|
42 |
return FRAME_PROCESSORS_MODULES
|
43 |
|
44 |
-
|
45 |
def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
|
46 |
with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
|
47 |
futures = []
|
@@ -50,32 +53,35 @@ def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_f
|
|
50 |
while not queue.empty():
|
51 |
future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
|
52 |
futures.append(future)
|
|
|
53 |
for future in as_completed(futures):
|
54 |
-
|
55 |
-
|
|
|
|
|
|
|
56 |
|
57 |
def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
|
58 |
queue: Queue[str] = Queue()
|
59 |
for frame_path in temp_frame_paths:
|
60 |
queue.put(frame_path)
|
|
|
61 |
return queue
|
62 |
|
63 |
-
|
64 |
def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
|
65 |
queues = []
|
66 |
for _ in range(queue_per_future):
|
67 |
if not queue.empty():
|
68 |
queues.append(queue.get())
|
|
|
69 |
return queues
|
70 |
|
71 |
-
|
72 |
-
def process_video(source_path: str, frame_paths: list[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
|
73 |
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
|
74 |
total = len(frame_paths)
|
75 |
-
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
|
76 |
multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
|
77 |
|
78 |
-
|
79 |
def update_progress(progress: Any = None) -> None:
|
80 |
process = psutil.Process(os.getpid())
|
81 |
memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
|
@@ -86,3 +92,4 @@ def update_progress(progress: Any = None) -> None:
|
|
86 |
})
|
87 |
progress.refresh()
|
88 |
progress.update(1)
|
|
|
|
5 |
from queue import Queue
|
6 |
from types import ModuleType
|
7 |
from typing import Any, List, Callable
|
8 |
+
import logging
|
9 |
from tqdm import tqdm
|
10 |
|
11 |
import roop
|
12 |
|
13 |
+
# Initialize logging
|
14 |
+
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
15 |
+
|
16 |
FRAME_PROCESSORS_MODULES: List[ModuleType] = []
|
17 |
FRAME_PROCESSORS_INTERFACE = [
|
18 |
'pre_check',
|
|
|
24 |
'post_process'
|
25 |
]
|
26 |
|
|
|
27 |
def load_frame_processor_module(frame_processor: str) -> Any:
|
28 |
try:
|
29 |
frame_processor_module = importlib.import_module(f'roop.processors.frame.{frame_processor}')
|
30 |
for method_name in FRAME_PROCESSORS_INTERFACE:
|
31 |
if not hasattr(frame_processor_module, method_name):
|
32 |
raise NotImplementedError
|
33 |
+
logging.info(f'Successfully loaded frame processor module: {frame_processor}')
|
34 |
+
except (ImportError, NotImplementedError) as e:
|
35 |
+
logging.error(f'Error loading frame processor {frame_processor}: {e}')
|
36 |
quit(f'Frame processor {frame_processor} crashed.')
|
37 |
return frame_processor_module
|
38 |
|
|
|
39 |
def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
|
40 |
global FRAME_PROCESSORS_MODULES
|
41 |
|
|
|
45 |
FRAME_PROCESSORS_MODULES.append(frame_processor_module)
|
46 |
return FRAME_PROCESSORS_MODULES
|
47 |
|
|
|
48 |
def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
|
49 |
with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
|
50 |
futures = []
|
|
|
53 |
while not queue.empty():
|
54 |
future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
|
55 |
futures.append(future)
|
56 |
+
logging.info(f'Submitted future for processing frames.')
|
57 |
for future in as_completed(futures):
|
58 |
+
try:
|
59 |
+
future.result()
|
60 |
+
logging.info('Frame processing completed for a future.')
|
61 |
+
except Exception as e:
|
62 |
+
logging.error(f'Error in processing frame: {e}')
|
63 |
|
64 |
def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
|
65 |
queue: Queue[str] = Queue()
|
66 |
for frame_path in temp_frame_paths:
|
67 |
queue.put(frame_path)
|
68 |
+
logging.info('Queue created with frame paths.')
|
69 |
return queue
|
70 |
|
|
|
71 |
def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
|
72 |
queues = []
|
73 |
for _ in range(queue_per_future):
|
74 |
if not queue.empty():
|
75 |
queues.append(queue.get())
|
76 |
+
logging.info(f'Picked {len(queues)} items from queue for processing.')
|
77 |
return queues
|
78 |
|
79 |
+
def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
|
|
|
80 |
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
|
81 |
total = len(frame_paths)
|
82 |
+
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format, mininterval=0.1) as progress:
|
83 |
multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
|
84 |
|
|
|
85 |
def update_progress(progress: Any = None) -> None:
|
86 |
process = psutil.Process(os.getpid())
|
87 |
memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
|
|
|
92 |
})
|
93 |
progress.refresh()
|
94 |
progress.update(1)
|
95 |
+
logging.info(f'Updated progress: {progress.n}/{progress.total} frames processed. Memory usage: {memory_usage:.2f}GB')
|