From 9c39bd41f7b2f46a41a5ac091be2c8b5d1bfe658 Mon Sep 17 00:00:00 2001 From: Jose Manuel Date: Fri, 2 Jun 2023 19:45:24 +0200 Subject: [PATCH] Add FFMpeg hardware acceleration --- core/utils.py | 20 ++++++++++++-------- run.py | 11 ++++++----- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/core/utils.py b/core/utils.py index 50aca35..3cf7407 100644 --- a/core/utils.py +++ b/core/utils.py @@ -29,26 +29,30 @@ def detect_fps(input_path): return 30, 30 -def set_fps(input_path, output_path, fps): +def set_fps(input_path, output_path, fps, use_gpu): input_path, output_path = path(input_path), path(output_path) - os.system(f'ffmpeg -i "{input_path}" -filter:v fps=fps={fps} "{output_path}"') + hwaccel_option = '-hwaccel cuda' if use_gpu else '' + os.system(f'ffmpeg {hwaccel_option} -i "{input_path}" -filter:v fps=fps={fps} "{output_path}"') -def create_video(video_name, fps, output_dir): +def create_video(video_name, fps, output_dir, use_gpu): output_dir = path(output_dir) - os.system(f'ffmpeg -framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') + hwaccel_option = '-hwaccel cuda' if use_gpu else '' + os.system(f'ffmpeg {hwaccel_option} -framerate "{fps}" -i "{output_dir}{sep}%04d.png" -c:v libx264 -crf 7 -pix_fmt yuv420p -y "{output_dir}{sep}output.mp4"') -def extract_frames(input_path, output_dir): +def extract_frames(input_path, output_dir, use_gpu): input_path, output_dir = path(input_path), path(output_dir) - os.system(f'ffmpeg -i "{input_path}" "{output_dir}{sep}%04d.png"') + hwaccel_option = '-hwaccel cuda' if use_gpu else '' + os.system(f'ffmpeg {hwaccel_option} -i "{input_path}" "{output_dir}{sep}%04d.png"') -def add_audio(output_dir, target_path, video, keep_frames, output_file): +def add_audio(output_dir, target_path, video, keep_frames, output_file, use_gpu): video_name = os.path.splitext(video)[0] save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4" save_to_ff, output_dir_ff = path(save_to), path(output_dir) - os.system(f'ffmpeg -i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}"') + hwaccel_option = '-hwaccel cuda' if use_gpu else '' + os.system(f'ffmpeg {hwaccel_option} -i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}"') if not os.path.isfile(save_to): shutil.move(output_dir + "/output.mp4", save_to) if not keep_frames: diff --git a/run.py b/run.py index 2099fdc..10435b1 100755 --- a/run.py +++ b/run.py @@ -71,6 +71,7 @@ def pre_check(): if not os.path.isfile(model_path): quit('File "inswapper_128.onnx" does not exist!') if '--gpu' in sys.argv: + core.globals.use_gpu = True NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider'] if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) == 1: CUDA_VERSION = torch.version.cuda @@ -90,7 +91,7 @@ def pre_check(): def start_processing(): - if args['gpu']: + if core.globals.use_gpu: process_video(args['source_img'], args["frame_paths"]) return frame_paths = args["frame_paths"] @@ -207,12 +208,12 @@ def start(): fps, exact_fps = detect_fps(target_path) if not args['keep_fps'] and fps > 30: this_path = output_dir + "/" + video_name + ".mp4" - set_fps(target_path, this_path, 30) + set_fps(target_path, this_path, 30, core.globals.use_gpu) target_path, exact_fps = this_path, 30 else: shutil.copy(target_path, output_dir) status("extracting frames...") - extract_frames(target_path, output_dir) + extract_frames(target_path, output_dir, core.globals.use_gpu) args['frame_paths'] = tuple(sorted( glob.glob(output_dir + "/*.png"), key=lambda x: int(x.split(sep)[-1].replace(".png", "")) @@ -220,9 +221,9 @@ def start(): status("swapping in progress...") start_processing() status("creating video...") - create_video(video_name, exact_fps, output_dir) + create_video(video_name, exact_fps, output_dir, core.globals.use_gpu) status("adding audio...") - add_audio(output_dir, target_path, video_name_full, args['keep_frames'], args['output_file']) + add_audio(output_dir, target_path, video_name_full, args['keep_frames'], args['output_file'], core.globals.use_gpu) save_path = args['output_file'] if args['output_file'] else output_dir + "/" + video_name + ".mp4" print("\n\nVideo saved as:", save_path, "\n\n") status("swap successful!")