upagge/roop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add cuda utilization and cuda memory to progress bar

Browse Source
This commit is contained in:
henryruhs 2023-06-03 14:49:27 +02:00
parent 21e2953420
commit 777261c8f9
1 changed files with 11 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
roop/swapper.py
View File

@ -1,10 +1,12 @@
import os import os
from tqdm import tqdm from tqdm import tqdm
import torch
import onnxruntime
import cv2 import cv2
import insightface import insightface
import roop.globals import roop.globals
from roop.analyser import get_face_single, get_face_many from roop.analyser import get_face_single, get_face_many
import onnxruntime
FACE_SWAPPER = None FACE_SWAPPER = None
@ -29,23 +31,17 @@ def swap_face_in_frame(source_face, target_face, frame):
return frame return frame
def process_faces(source_face, frame, progress): def process_faces(source_face, target_frame, progress):
if roop.globals.all_faces: if roop.globals.all_faces:
many_faces = get_face_many(frame) many_faces = get_face_many(target_frame)
if many_faces: if many_faces:
for face in many_faces: for face in many_faces:
frame = swap_face_in_frame(source_face, face, frame) target_frame = swap_face_in_frame(source_face, face, target_frame)
progress.set_postfix(status='.', refresh=True)
else: else:
progress.set_postfix(status='S', refresh=True) face = get_face_single(target_frame)
else:
face = get_face_single(frame)
if face: if face:
frame = swap_face_in_frame(source_face, face, frame) target_frame = swap_face_in_frame(source_face, face, target_frame)
progress.set_postfix(status='.', refresh=True) return target_frame
else:
progress.set_postfix(status='S', refresh=True)
return frame
def process_video(source_img, frame_paths): def process_video(source_img, frame_paths):
@ -54,12 +50,13 @@ def process_video(source_img, frame_paths):
with tqdm(total=len(frame_paths), desc="Processing", unit="frame", dynamic_ncols=True, bar_format=progress_bar_format) as progress: with tqdm(total=len(frame_paths), desc="Processing", unit="frame", dynamic_ncols=True, bar_format=progress_bar_format) as progress:
for frame_path in frame_paths: for frame_path in frame_paths:
if roop.globals.gpu_vendor == 'nvidia':
progress.set_postfix(cuda_utilization="{:02d}%".format(torch.cuda.utilization()), cuda_memory="{:02d}GB".format(torch.cuda.memory_usage()))
frame = cv2.imread(frame_path) frame = cv2.imread(frame_path)
try: try:
result = process_faces(source_face, frame, progress) result = process_faces(source_face, frame, progress)
cv2.imwrite(frame_path, result) cv2.imwrite(frame_path, result)
except Exception: except Exception:
progress.set_postfix(status='E', refresh=True)
pass pass
progress.update(1) progress.update(1)