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

Massive utilities and core refactoring

Browse Source
This commit is contained in:
henryruhs 2023-06-06 12:30:27 +02:00
parent dcfd6cca4d
commit 09ea59f66f
6 changed files with 213 additions and 158 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

211
roop/core.py
View File

@ -2,37 +2,39 @@
import os import os
import sys import sys
from typing import List
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# single thread doubles performance of gpu-mode - needs to be set before torch import # single thread doubles performance of gpu-mode - needs to be set before torch import
if any(arg.startswith('--gpu-vendor') for arg in sys.argv): if any(arg.startswith('--gpu-vendor') for arg in sys.argv):
os.environ['OMP_NUM_THREADS'] = '1' os.environ['OMP_NUM_THREADS'] = '1'
import platform import platform
import signal import signal
import shutil import shutil
import glob
import argparse import argparse
import psutil import psutil
import torch import torch
import tensorflow import tensorflow
from pathlib import Path import multiprocessing
import multiprocessing as mp
from opennsfw2 import predict_video_frames, predict_image from opennsfw2 import predict_video_frames, predict_image
import cv2 import cv2
import roop.globals import roop.globals
from roop.swapper import process_video, process_img, process_faces, process_frames from roop.swapper import process_video, process_img, process_faces
from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames from roop.utilities import has_image_extention, is_image, detect_fps, create_video, extract_frames, \
get_temp_frames_paths, restore_audio, create_temp, clean_temp, is_video
from roop.analyser import get_face_single from roop.analyser import get_face_single
import roop.ui as ui import roop.ui as ui
def handle_parse(): def handle_parse():
global args global args
signal.signal(signal.SIGINT, lambda signal_number, frame: quit()) signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-f', '--face', help='use this face', dest='source_target') parser.add_argument('-f', '--face', help='use this face', dest='source_path')
parser.add_argument('-t', '--target', help='replace this face', dest='target_path') parser.add_argument('-t', '--target', help='replace this face', dest='target_path')
parser.add_argument('-o', '--output', help='save output to this file', dest='output_path') parser.add_argument('-o', '--output', help='save output to this file', dest='output_path')
parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
parser.add_argument('--keep-audio', help='maintain original audio', dest='keep_audio', action='store_true', default=True)
parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False) parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False)
parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int) parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', dest='max_memory', type=int)
@ -42,7 +44,10 @@ def handle_parse():
args = parser.parse_known_args()[0] args = parser.parse_known_args()[0]
roop.globals.headless = args.source_target or args.target_path or args.output_path roop.globals.headless = args.source_path or args.target_path or args.output_path
roop.globals.keep_fps = args.keep_fps
roop.globals.keep_audio = args.keep_audio
roop.globals.keep_frames = args.keep_frames
roop.globals.all_faces = args.all_faces roop.globals.all_faces = args.all_faces
if args.cpu_cores: if args.cpu_cores:
@ -83,7 +88,7 @@ def limit_resources():
def pre_check(): def pre_check():
if sys.version_info < (3, 9): if sys.version_info < (3, 9):
quit('Python version is not supported - please upgrade to 3.9 or higher') quit('Python version is not supported - please upgrade to 3.9 or higher.')
if not shutil.which('ffmpeg'): if not shutil.which('ffmpeg'):
quit('ffmpeg is not installed!') quit('ffmpeg is not installed!')
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../inswapper_128.onnx') model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), '../inswapper_128.onnx')
@ -91,23 +96,23 @@ def pre_check():
quit('File "inswapper_128.onnx" does not exist!') quit('File "inswapper_128.onnx" does not exist!')
if roop.globals.gpu_vendor == 'apple': if roop.globals.gpu_vendor == 'apple':
if 'CoreMLExecutionProvider' not in roop.globals.providers: if 'CoreMLExecutionProvider' not in roop.globals.providers:
quit("You are using --gpu=apple flag but CoreML isn't available or properly installed on your system.") quit('You are using --gpu=apple flag but CoreML is not available or properly installed on your system.')
if roop.globals.gpu_vendor == 'amd': if roop.globals.gpu_vendor == 'amd':
if 'ROCMExecutionProvider' not in roop.globals.providers: if 'ROCMExecutionProvider' not in roop.globals.providers:
quit("You are using --gpu=amd flag but ROCM isn't available or properly installed on your system.") quit('You are using --gpu=amd flag but ROCM is not available or properly installed on your system.')
if roop.globals.gpu_vendor == 'nvidia': if roop.globals.gpu_vendor == 'nvidia':
CUDA_VERSION = torch.version.cuda CUDA_VERSION = torch.version.cuda
CUDNN_VERSION = torch.backends.cudnn.version() CUDNN_VERSION = torch.backends.cudnn.version()
if not torch.cuda.is_available(): if not torch.cuda.is_available():
quit("You are using --gpu=nvidia flag but CUDA isn't available or properly installed on your system.") quit('You are using --gpu=nvidia flag but CUDA is not available or properly installed on your system.')
if CUDA_VERSION > '11.8': if CUDA_VERSION > '11.8':
quit(f"CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8") quit(f'CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8')
if CUDA_VERSION < '11.4': if CUDA_VERSION < '11.4':
quit(f"CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8") quit(f'CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8')
if CUDNN_VERSION < 8220: if CUDNN_VERSION < 8220:
quit(f"CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1") quit(f'CUDNN version {CUDNN_VERSION} is not supported - please upgrade to 8.9.1')
if CUDNN_VERSION > 8910: if CUDNN_VERSION > 8910:
quit(f"CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1") quit(f'CUDNN version {CUDNN_VERSION} is not supported - please downgrade to 8.9.1')
def get_video_frame(video_path, frame_number = 1): def get_video_frame(video_path, frame_number = 1):
@ -115,19 +120,18 @@ def get_video_frame(video_path, frame_number = 1):
amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1)) cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1))
if not cap.isOpened(): if not cap.isOpened():
print("Error opening video file") status('Error opening video file')
return return
ret, frame = cap.read() ret, frame = cap.read()
if ret: if ret:
return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
cap.release() cap.release()
def preview_video(video_path): def preview_video(video_path):
cap = cv2.VideoCapture(video_path) cap = cv2.VideoCapture(video_path)
if not cap.isOpened(): if not cap.isOpened():
print("Error opening video file") status('Error opening video file')
return 0 return 0
amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
ret, frame = cap.read() ret, frame = cap.read()
@ -138,86 +142,89 @@ def preview_video(video_path):
return (amount_of_frames, frame) return (amount_of_frames, frame)
def status(string): def status(message: str):
value = "Status: " + string value = 'Status: ' + message
print(value) print(value)
if not roop.globals.headless: if not roop.globals.headless:
ui.update_status_label(value) ui.update_status_label(value)
def process_video_multi_cores(source_target, frame_paths): def conditional_process_video(source_path: str, frame_paths: List[str]) -> None:
n = len(frame_paths) // roop.globals.cpu_cores pool_amount = len(frame_paths) // roop.globals.cpu_cores
if n > 2: if pool_amount > 2 and roop.globals.cpu_cores > 1 and roop.globals.gpu_vendor is None:
processes = [] status('Pool-Swapping in progress...')
for i in range(0, len(frame_paths), n): global POOL
p = POOL.apply_async(process_video, args=(source_target, frame_paths[i:i + n],)) POOL = multiprocessing.Pool(roop.globals.cpu_cores, maxtasksperchild=1)
processes.append(p) pools = []
for p in processes: for i in range(0, len(frame_paths), pool_amount):
p.get() pool = POOL.apply_async(process_video, args=(source_path, frame_paths[i:i + pool_amount]))
POOL.close() pools.append(pool)
for pool in pools:
pool.get()
POOL.join() POOL.join()
POOL.close()
else:
status('Swapping in progress...')
process_video(args.source_path, frame_paths)
def start(preview_callback = None): def start(preview_callback = None) -> None:
if not args.source_target or not os.path.isfile(args.source_target): if not args.source_path or not os.path.isfile(args.source_path):
print("\n[WARNING] Please select an image containing a face.") status('Please select an image containing a face.')
return return
elif not args.target_path or not os.path.isfile(args.target_path): elif not args.target_path or not os.path.isfile(args.target_path):
print("\n[WARNING] Please select a video/image to swap face in.") status('Please select a video/image target!')
return return
target_path = args.target_path test_face = get_face_single(cv2.imread(args.source_path))
test_face = get_face_single(cv2.imread(args.source_target))
if not test_face: if not test_face:
print("\n[WARNING] No face detected in source image. Please try with another one.\n") status('No face detected in source image. Please try with another one!')
return return
if is_img(target_path): # process image to image
if predict_image(target_path) > 0.85: if has_image_extention(args.target_path):
quit() if predict_image(args.target_path) > 0.85:
process_img(args.source_target, target_path, args.output_path) destroy()
status("swap successful!") process_img(args.source_path, args.target_path, args.output_path)
if is_image(args.target_path):
status('Swapping to image succeed!')
else:
status('Swapping to image failed!')
return return
# process image to videos
seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100) seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100)
if any(probability > 0.85 for probability in probabilities): if any(probability > 0.85 for probability in probabilities):
quit() destroy()
video_name_full = target_path.split(os.sep)[-1] status('Creating temp resources...')
video_name = os.path.splitext(video_name_full)[0] create_temp(args.target_path)
output_dir = os.path.dirname(target_path) + os.sep + video_name if os.path.dirname(target_path) else video_name status('Extracting frames...')
Path(output_dir).mkdir(exist_ok=True) extract_frames(args.target_path)
status("detecting video's FPS...") frame_paths = get_temp_frames_paths(args.target_path)
fps, exact_fps = detect_fps(target_path) conditional_process_video(args.source_path, frame_paths)
if not args.keep_fps and fps > 30: # prevent memory leak using ffmpeg with cuda
this_path = output_dir + os.sep + video_name + ".mp4"
set_fps(target_path, this_path, 30)
target_path, exact_fps = this_path, 30
else:
shutil.copy(target_path, output_dir)
status("extracting frames...")
extract_frames(target_path, output_dir)
args.frame_paths = tuple(sorted(
glob.glob(output_dir + "/*.png"),
key=lambda x: int(x.split(os.sep)[-1].replace(".png", ""))
))
status("swapping in progress...")
if roop.globals.gpu_vendor is None and roop.globals.cpu_cores > 1:
global POOL
POOL = mp.Pool(roop.globals.cpu_cores)
process_video_multi_cores(args.source_target, args.frame_paths)
else:
process_video(args.source_target, args.frame_paths)
# prevent out of memory while using ffmpeg with cuda
if args.gpu_vendor == 'nvidia': if args.gpu_vendor == 'nvidia':
torch.cuda.empty_cache() torch.cuda.empty_cache()
status("creating video...") if roop.globals.keep_fps:
create_video(video_name, exact_fps, output_dir) status('Detecting fps...')
status("adding audio...") fps = detect_fps(args.source_path)
add_audio(output_dir, target_path, video_name_full, args.keep_frames, args.output_path) status(f'Creating video with {fps} fps...')
save_path = args.output_path if args.output_path else output_dir + os.sep + video_name + ".mp4" create_video(args.target_path, fps)
print("\n\nVideo saved as:", save_path, "\n\n") else:
status("swap successful!") status('Creating video with 30 fps...')
create_video(args.target_path, 30)
if roop.globals.keep_audio:
if roop.globals.keep_fps:
status('Restoring audio...')
else:
status('Restoring audio might cause issues as fps are not kept...')
restore_audio(args.target_path)
clean_temp(args.target_path, args.output_path)
if is_video(args.target_path):
status('Swapping to video succeed!')
else:
status('Swapping to video failed!')
def select_face_handler(path: str): def select_face_handler(path: str):
args.source_target = path args.source_path = path
def select_target_handler(path: str): def select_target_handler(path: str):
@ -243,33 +250,37 @@ def save_file_handler(path: str):
def create_test_preview(frame_number): def create_test_preview(frame_number):
return process_faces( return process_faces(
get_face_single(cv2.imread(args.source_target)), get_face_single(cv2.imread(args.source_path)),
get_video_frame(args.target_path, frame_number) get_video_frame(args.target_path, frame_number)
) )
def run(): def destroy() -> None:
clean_temp(args.target_path)
quit()
def run() -> None:
global all_faces, keep_frames, limit_fps global all_faces, keep_frames, limit_fps
handle_parse() handle_parse()
pre_check() pre_check()
limit_resources() limit_resources()
if roop.globals.headless: start()
start() if not roop.globals.headless:
quit() window = ui.init(
window = ui.init( {
{ 'all_faces': args.all_faces,
'all_faces': roop.globals.all_faces, 'keep_fps': args.keep_fps,
'keep_fps': args.keep_fps, 'keep_frames': args.keep_frames
'keep_frames': args.keep_frames },
}, select_face_handler,
select_face_handler, select_target_handler,
select_target_handler, toggle_all_faces_handler,
toggle_all_faces_handler, toggle_fps_limit_handler,
toggle_fps_limit_handler, toggle_keep_frames_handler,
toggle_keep_frames_handler, save_file_handler,
save_file_handler, start,
start, get_video_frame,
get_video_frame, create_test_preview
create_test_preview )
) window.mainloop()
window.mainloop()

5
roop/globals.py
View File

@ -1,11 +1,14 @@
import onnxruntime import onnxruntime
keep_fps = None
keep_audio = None
keep_frames = None
all_faces = None all_faces = None
log_level = 'error'
cpu_cores = None cpu_cores = None
gpu_threads = None gpu_threads = None
gpu_vendor = None gpu_vendor = None
headless = None headless = None
log_level = 'error'
providers = onnxruntime.get_available_providers() providers = onnxruntime.get_available_providers()
if 'TensorrtExecutionProvider' in providers: if 'TensorrtExecutionProvider' in providers:

10
roop/swapper.py
View File

@ -58,9 +58,8 @@ def multi_process_frame(source_img, frame_paths, progress):
num_threads = roop.globals.gpu_threads num_threads = roop.globals.gpu_threads
num_frames_per_thread = len(frame_paths) // num_threads num_frames_per_thread = len(frame_paths) // num_threads
remaining_frames = len(frame_paths) % num_threads remaining_frames = len(frame_paths) % num_threads
# create thread and launch
start_index = 0 start_index = 0
# create threads by frames
for _ in range(num_threads): for _ in range(num_threads):
end_index = start_index + num_frames_per_thread end_index = start_index + num_frames_per_thread
if remaining_frames > 0: if remaining_frames > 0:
@ -71,8 +70,7 @@ def multi_process_frame(source_img, frame_paths, progress):
threads.append(thread) threads.append(thread)
thread.start() thread.start()
start_index = end_index start_index = end_index
# join threads
# threading
for thread in threads: for thread in threads:
thread.join() thread.join()
@ -83,13 +81,13 @@ def process_img(source_img, target_path, output_file):
source_face = get_face_single(cv2.imread(source_img)) source_face = get_face_single(cv2.imread(source_img))
result = get_face_swapper().get(frame, face, source_face, paste_back=True) result = get_face_swapper().get(frame, face, source_face, paste_back=True)
cv2.imwrite(output_file, result) cv2.imwrite(output_file, result)
print("\n\nImage saved as:", output_file, "\n\n")
def process_video(source_img, frame_paths): def process_video(source_img, frame_paths):
do_multi = roop.globals.gpu_vendor is not None and roop.globals.gpu_threads > 1 do_multi = roop.globals.gpu_vendor is not None and roop.globals.gpu_threads > 1
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: total = len(frame_paths)
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
if do_multi: if do_multi:
multi_process_frame(source_img, frame_paths, progress) multi_process_frame(source_img, frame_paths, progress)
else: else:

4
roop/ui.py
View File

@ -6,7 +6,7 @@ from tkinter import filedialog
from tkinter.filedialog import asksaveasfilename from tkinter.filedialog import asksaveasfilename
import threading import threading
from roop.utils import is_img from roop.utilities import is_image
max_preview_size = 800 max_preview_size = 800
@ -114,7 +114,7 @@ def select_target(select_target_handler: Callable[[str], Tuple[int, Any]], targe
def save_file(save_file_handler: Callable[[str], None], target_path: str): def save_file(save_file_handler: Callable[[str], None], target_path: str):
filename, ext = 'output.mp4', '.mp4' filename, ext = 'output.mp4', '.mp4'
if is_img(target_path): if is_image(target_path):
filename, ext = 'output.png', '.png' filename, ext = 'output.png', '.png'
if save_file_handler: if save_file_handler:

92
roop/utilities.py Normal file
View File

@ -0,0 +1,92 @@
import glob
import os
import shutil
import subprocess
from pathlib import Path
from typing import List, Any
import roop.globals
from PIL import Image
def run_ffmpeg(args: List) -> None:
commands = ['ffmpeg', '-hide_banner', '-hwaccel', 'auto', '-loglevel', roop.globals.log_level]
commands.extend(args)
try:
subprocess.check_output(commands, stderr=subprocess.STDOUT)
except Exception:
pass
def detect_fps(source_path: str) -> int:
command = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=r_frame_rate', '-of', 'default=noprint_wrappers=1:nokey=1', source_path]
output = subprocess.check_output(command).decode().strip()
try:
return int(eval(output))
except Exception:
pass
return 30
def extract_frames(target_path: str) -> None:
temp_directory_path = get_temp_directory_path(target_path)
run_ffmpeg(['-i', target_path, temp_directory_path + os.sep + '%04d.png'])
def create_video(target_path: str, fps: int) -> None:
temp_directory_path = get_temp_directory_path(target_path)
temp_file_path = get_temp_file_path(target_path)
run_ffmpeg(['-i', temp_directory_path + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', temp_file_path])
def restore_audio(target_path: str) -> None:
run_ffmpeg(['-i', target_path, '-i', get_temp_file_path(target_path), '-c:v', 'copy', '-map', '0:v:0', 'map', '1:a:0', '-y', get_temp_file_path(target_path)])
def get_temp_frames_paths(target_path: str) -> List:
return glob.glob(get_temp_directory_path(target_path) + os.sep + '*.png')
def get_temp_directory_path(target_path: str) -> str:
target_directory_path = os.path.dirname(target_path)
return target_directory_path + os.sep + 'temp'
def get_temp_file_path(target_path: str) -> str:
return get_temp_directory_path(target_path) + os.sep + 'temp.mp4'
def create_temp(target_path: str) -> None:
Path(get_temp_directory_path(target_path)).mkdir(exist_ok=True)
def clean_temp(target_path: str, output_path: str) -> None:
temp_file_path = get_temp_file_path(target_path)
if os.path.isfile(temp_file_path):
shutil.move(temp_file_path, output_path)
if not roop.globals.keep_frames:
shutil.rmtree(get_temp_directory_path(target_path))
def has_image_extention(image_path: str) -> bool:
return image_path.lower().endswith(('png', 'jpg', 'jpeg', 'bmp'))
def is_image(path: str) -> bool:
if os.path.isfile(path):
try:
image = Image.open(path)
image.verify()
return True
except Exception:
pass
return False
def is_video(path: str) -> bool:
try:
run_ffmpeg(['-v', 'error', '-i', path, '-f', 'null', '-'])
return True
except subprocess.CalledProcessError:
pass
return False

49
roop/utils.py
View File

@ -1,49 +0,0 @@
import os
import shutil
import roop.globals
def run_command(command, mode="silent"):
if mode == "debug":
return os.system(command)
return os.popen(command).read()
def detect_fps(input_path):
output = os.popen(f'ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "{input_path}"').read()
if "/" in output:
try:
return int(output.split("/")[0]) // int(output.split("/")[1].strip()), output.strip()
except:
pass
return 30, 30
def run_ffmpeg(args):
run_command(f'ffmpeg -hide_banner -hwaccel auto -loglevel {roop.globals.log_level} {args}')
def set_fps(input_path, output_path, fps):
run_ffmpeg(f'-i "{input_path}" -filter:v fps=fps={fps} "{output_path}"')
def create_video(video_name, fps, output_dir):
run_ffmpeg(f'-framerate "{fps}" -i "{output_dir}{os.sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{os.sep}output.mp4"')
def extract_frames(input_path, output_dir):
run_ffmpeg(f'-i "{input_path}" "{output_dir}{os.sep}%04d.png"')
def add_audio(output_dir, target_path, video, keep_frames, output_file):
video_name = os.path.splitext(video)[0]
save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4"
run_ffmpeg(f'-i "{output_dir}{os.sep}output.mp4" -i "{output_dir}{os.sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to}"')
if not os.path.isfile(save_to):
shutil.move(output_dir + "/output.mp4", save_to)
if not keep_frames:
shutil.rmtree(output_dir)
def is_img(path):
return path.lower().endswith(("png", "jpg", "jpeg", "bmp"))