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

177
roop/core.py
View File

@ -2,37 +2,39 @@
import os
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
if any(arg.startswith('--gpu-vendor') for arg in sys.argv):
os.environ['OMP_NUM_THREADS'] = '1'
import platform
import signal
import shutil
import glob
import argparse
import psutil
import torch
import tensorflow
from pathlib import Path
import multiprocessing as mp
import multiprocessing
from opennsfw2 import predict_video_frames, predict_image
import cv2
import roop.globals
from roop.swapper import process_video, process_img, process_faces, process_frames
from roop.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames
from roop.swapper import process_video, process_img, process_faces
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
import roop.ui as ui
def handle_parse():
global args
signal.signal(signal.SIGINT, lambda signal_number, frame: quit())
signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
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('-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-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('--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)
@ -42,7 +44,10 @@ def handle_parse():
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
if args.cpu_cores:
@ -83,7 +88,7 @@ def limit_resources():
def pre_check():
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'):
quit('ffmpeg is not installed!')
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!')
if roop.globals.gpu_vendor == 'apple':
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 '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':
CUDA_VERSION = torch.version.cuda
CUDNN_VERSION = torch.backends.cudnn.version()
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':
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':
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:
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:
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):
@ -115,19 +120,18 @@ def get_video_frame(video_path, frame_number = 1):
amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
cap.set(cv2.CAP_PROP_POS_FRAMES, min(amount_of_frames, frame_number-1))
if not cap.isOpened():
print("Error opening video file")
status('Error opening video file')
return
ret, frame = cap.read()
if ret:
return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
cap.release()
def preview_video(video_path):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print("Error opening video file")
status('Error opening video file')
return 0
amount_of_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
ret, frame = cap.read()
@ -138,86 +142,89 @@ def preview_video(video_path):
return (amount_of_frames, frame)
def status(string):
value = "Status: " + string
def status(message: str):
value = 'Status: ' + message
print(value)
if not roop.globals.headless:
ui.update_status_label(value)
def process_video_multi_cores(source_target, frame_paths):
n = len(frame_paths) // roop.globals.cpu_cores
if n > 2:
processes = []
for i in range(0, len(frame_paths), n):
p = POOL.apply_async(process_video, args=(source_target, frame_paths[i:i + n],))
processes.append(p)
for p in processes:
p.get()
POOL.close()
def conditional_process_video(source_path: str, frame_paths: List[str]) -> None:
pool_amount = len(frame_paths) // roop.globals.cpu_cores
if pool_amount > 2 and roop.globals.cpu_cores > 1 and roop.globals.gpu_vendor is None:
status('Pool-Swapping in progress...')
global POOL
POOL = multiprocessing.Pool(roop.globals.cpu_cores, maxtasksperchild=1)
pools = []
for i in range(0, len(frame_paths), pool_amount):
pool = POOL.apply_async(process_video, args=(source_path, frame_paths[i:i + pool_amount]))
pools.append(pool)
for pool in pools:
pool.get()
POOL.join()
POOL.close()
else:
status('Swapping in progress...')
process_video(args.source_path, frame_paths)
def start(preview_callback = None):
if not args.source_target or not os.path.isfile(args.source_target):
print("\n[WARNING] Please select an image containing a face.")
def start(preview_callback = None) -> None:
if not args.source_path or not os.path.isfile(args.source_path):
status('Please select an image containing a face.')
return
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
target_path = args.target_path
test_face = get_face_single(cv2.imread(args.source_target))
test_face = get_face_single(cv2.imread(args.source_path))
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
if is_img(target_path):
if predict_image(target_path) > 0.85:
quit()
process_img(args.source_target, target_path, args.output_path)
status("swap successful!")
# process image to image
if has_image_extention(args.target_path):
if predict_image(args.target_path) > 0.85:
destroy()
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
# process image to videos
seconds, probabilities = predict_video_frames(video_path=args.target_path, frame_interval=100)
if any(probability > 0.85 for probability in probabilities):
quit()
video_name_full = target_path.split(os.sep)[-1]
video_name = os.path.splitext(video_name_full)[0]
output_dir = os.path.dirname(target_path) + os.sep + video_name if os.path.dirname(target_path) else video_name
Path(output_dir).mkdir(exist_ok=True)
status("detecting video's FPS...")
fps, exact_fps = detect_fps(target_path)
if not args.keep_fps and fps > 30:
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
destroy()
status('Creating temp resources...')
create_temp(args.target_path)
status('Extracting frames...')
extract_frames(args.target_path)
frame_paths = get_temp_frames_paths(args.target_path)
conditional_process_video(args.source_path, frame_paths)
# prevent memory leak using ffmpeg with cuda
if args.gpu_vendor == 'nvidia':
torch.cuda.empty_cache()
status("creating video...")
create_video(video_name, exact_fps, output_dir)
status("adding audio...")
add_audio(output_dir, target_path, video_name_full, args.keep_frames, args.output_path)
save_path = args.output_path if args.output_path else output_dir + os.sep + video_name + ".mp4"
print("\n\nVideo saved as:", save_path, "\n\n")
status("swap successful!")
if roop.globals.keep_fps:
status('Detecting fps...')
fps = detect_fps(args.source_path)
status(f'Creating video with {fps} fps...')
create_video(args.target_path, fps)
else:
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):
args.source_target = path
args.source_path = path
def select_target_handler(path: str):
@ -243,22 +250,26 @@ def save_file_handler(path: str):
def create_test_preview(frame_number):
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)
)
def run():
def destroy() -> None:
clean_temp(args.target_path)
quit()
def run() -> None:
global all_faces, keep_frames, limit_fps
handle_parse()
pre_check()
limit_resources()
if roop.globals.headless:
start()
quit()
if not roop.globals.headless:
window = ui.init(
{
'all_faces': roop.globals.all_faces,
'all_faces': args.all_faces,
'keep_fps': args.keep_fps,
'keep_frames': args.keep_frames
},

5
roop/globals.py
View File

@ -1,11 +1,14 @@
import onnxruntime
keep_fps = None
keep_audio = None
keep_frames = None
all_faces = None
log_level = 'error'
cpu_cores = None
gpu_threads = None
gpu_vendor = None
headless = None
log_level = 'error'
providers = onnxruntime.get_available_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_frames_per_thread = len(frame_paths) // num_threads
remaining_frames = len(frame_paths) % num_threads
# create thread and launch
start_index = 0
# create threads by frames
for _ in range(num_threads):
end_index = start_index + num_frames_per_thread
if remaining_frames > 0:
@ -71,8 +70,7 @@ def multi_process_frame(source_img, frame_paths, progress):
threads.append(thread)
thread.start()
start_index = end_index
# threading
# join threads
for thread in threads:
thread.join()
@ -83,13 +81,13 @@ def process_img(source_img, target_path, output_file):
source_face = get_face_single(cv2.imread(source_img))
result = get_face_swapper().get(frame, face, source_face, paste_back=True)
cv2.imwrite(output_file, result)
print("\n\nImage saved as:", output_file, "\n\n")
def process_video(source_img, frame_paths):
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}]'
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:
multi_process_frame(source_img, frame_paths, progress)
else:

4
roop/ui.py
View File

@ -6,7 +6,7 @@ from tkinter import filedialog
from tkinter.filedialog import asksaveasfilename
import threading
from roop.utils import is_img
from roop.utilities import is_image
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):
filename, ext = 'output.mp4', '.mp4'
if is_img(target_path):
if is_image(target_path):
filename, ext = 'output.png', '.png'
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"))