revert changes

This commit is contained in:
Somdev Sangwan 2023-06-04 17:17:19 +05:30 committed by GitHub
parent c36d3a2009
commit d05368c361
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -4,20 +4,24 @@ import torch
import onnxruntime import onnxruntime
import cv2 import cv2
import insightface import insightface
import threading
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
from roop.globals import gpu_vendor
FACE_SWAPPER = None FACE_SWAPPER = None
lock = threading.Lock()
def get_face_swapper(): def get_face_swapper():
global FACE_SWAPPER global FACE_SWAPPER
with lock: if FACE_SWAPPER is None:
if FACE_SWAPPER is None: session_options = onnxruntime.SessionOptions()
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../inswapper_128.onnx') if roop.globals.gpu_vendor is not None:
FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=roop.globals.providers) session_options.intra_op_num_threads = roop.globals.gpu_threads
else:
session_options.intra_op_num_threads = roop.globals.cpu_threads
session_options.execution_mode = onnxruntime.ExecutionMode.ORT_PARALLEL
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../inswapper_128.onnx')
FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=roop.globals.providers, session_options=session_options)
return FACE_SWAPPER return FACE_SWAPPER
@ -27,20 +31,7 @@ def swap_face_in_frame(source_face, target_face, frame):
return frame return frame
def process_frames(source_img, frame_paths, progress=None): def process_faces(source_face, target_frame, progress):
source_face = get_face_single(cv2.imread(source_img))
for frame_path in frame_paths:
frame = cv2.imread(frame_path)
try:
result = process_faces(source_face, frame)
cv2.imwrite(frame_path, result)
except Exception as e:
print(">>>>", e)
pass
if progress:
progress.update(1)
def process_faces(source_face, target_frame):
if roop.globals.all_faces: if roop.globals.all_faces:
many_faces = get_face_many(target_frame) many_faces = get_face_many(target_frame)
if many_faces: if many_faces:
@ -54,34 +45,22 @@ def process_faces(source_face, target_frame):
def process_video(source_img, frame_paths, preview_callback): def process_video(source_img, frame_paths, preview_callback):
source_face = get_face_single(cv2.imread(source_img))
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]' progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
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:
# nvidia multi-threading for frame_path in frame_paths:
if roop.globals.gpu_vendor == 'nvidia': 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())) progress.set_postfix(cuda_utilization="{:02d}%".format(torch.cuda.utilization()), cuda_memory="{:02d}GB".format(torch.cuda.memory_usage()))
# caculate the number of frames each threads processed frame = cv2.imread(frame_path)
num_threads = roop.globals.gpu_threads try:
num_frames_per_thread = len(frame_paths) // num_threads result = process_faces(source_face, frame, progress)
remaining_frames = len(frame_paths) % num_threads cv2.imwrite(frame_path, result)
# create thread list if preview_callback:
threads = [] preview_callback(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
start_index = 0 except Exception:
# create thread and launch pass
for _ in range(num_threads): progress.update(1)
end_index = start_index + num_frames_per_thread
if remaining_frames > 0:
end_index += 1
remaining_frames -= 1
thread_frame_paths = frame_paths[start_index:end_index]
thread = threading.Thread(target=process_frames, args=(source_img, thread_frame_paths, progress))
threads.append(thread)
thread.start()
start_index = end_index
for thread in threads:
thread.join()
else:
process_frames(source_img, frame_paths, progress)
def process_img(source_img, target_path, output_file): def process_img(source_img, target_path, output_file):