fi0r1/roop
1
0
Fork 0
mirror of https://github.com/s0md3v/roop.git synced 2025-12-06 18:08:29 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Show memory usage in progress bar

Browse Source
This commit is contained in:
henryruhs 2023-06-21 23:48:18 +02:00
parent cd96b47a7c
commit fd8a9b6f94
3 changed files with 27 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
roop/processors/frame/core.py
View File

@ -1,5 +1,7 @@
import os
import sys
import importlib
import psutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from types import ModuleType
from typing import Any, List, Callable
@ -38,11 +40,11 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType
return FRAME_PROCESSORS_MODULES
def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], progress: Any = None) -> None:
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], progress)
future = executor.submit(process_frames, source_path, [path], update)
futures.append(future)
for future in as_completed(futures):
future.result()
@ -52,5 +54,16 @@ def process_video(source_path: str, frame_paths: list[str], process_frames: Call
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:
progress.set_postfix({'execution_providers': roop.globals.execution_providers, 'execution_threads': roop.globals.execution_threads, 'max_memory': roop.globals.max_memory})
multi_process_frame(source_path, frame_paths, process_frames, progress)
multi_process_frame(source_path, 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)

8
roop/processors/frame/face_enhancer.py
View File

@ -1,4 +1,4 @@
from typing import Any, List
from typing import Any, List, Callable
import cv2
import threading
import gfpgan
@ -56,13 +56,13 @@ def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
return temp_frame
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
for temp_frame_path in temp_frame_paths:
temp_frame = cv2.imread(temp_frame_path)
result = process_frame(None, temp_frame)
cv2.imwrite(temp_frame_path, result)
if progress:
progress.update(1)
if update:
update()
def process_image(source_path: str, target_path: str, output_path: str) -> None:

16
roop/processors/frame/face_swapper.py
View File

@ -1,4 +1,4 @@
from typing import Any, List
from typing import Any, List, Callable
import cv2
import insightface
import threading
@ -61,18 +61,14 @@ def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
return temp_frame
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
source_face = get_one_face(cv2.imread(source_path))
for temp_frame_path in temp_frame_paths:
temp_frame = cv2.imread(temp_frame_path)
try:
result = process_frame(source_face, temp_frame)
cv2.imwrite(temp_frame_path, result)
except Exception as exception:
print(exception)
pass
if progress:
progress.update(1)
result = process_frame(source_face, temp_frame)
cv2.imwrite(temp_frame_path, result)
if update:
update()
def process_image(source_path: str, target_path: str, output_path: str) -> None: