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

Add FFMpeg hardware acceleration

Browse Source
This commit is contained in:
Jose Manuel 2023-06-02 19:45:24 +02:00
parent f7e029e6f7
commit 9c39bd41f7
2 changed files with 18 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
core/utils.py
View File

@ -29,26 +29,30 @@ def detect_fps(input_path):
return 30, 30 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) 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) 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) 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] video_name = os.path.splitext(video)[0]
save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4" 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) 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): if not os.path.isfile(save_to):
shutil.move(output_dir + "/output.mp4", save_to) shutil.move(output_dir + "/output.mp4", save_to)
if not keep_frames: if not keep_frames:

11
run.py
View File

@ -71,6 +71,7 @@ def pre_check():
if not os.path.isfile(model_path): if not os.path.isfile(model_path):
quit('File "inswapper_128.onnx" does not exist!') quit('File "inswapper_128.onnx" does not exist!')
if '--gpu' in sys.argv: if '--gpu' in sys.argv:
core.globals.use_gpu = True
NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider'] NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider']
if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) == 1: if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) == 1:
CUDA_VERSION = torch.version.cuda CUDA_VERSION = torch.version.cuda
@ -90,7 +91,7 @@ def pre_check():
def start_processing(): def start_processing():
if args['gpu']: if core.globals.use_gpu:
process_video(args['source_img'], args["frame_paths"]) process_video(args['source_img'], args["frame_paths"])
return return
frame_paths = args["frame_paths"] frame_paths = args["frame_paths"]
@ -207,12 +208,12 @@ def start():
fps, exact_fps = detect_fps(target_path) fps, exact_fps = detect_fps(target_path)
if not args['keep_fps'] and fps > 30: if not args['keep_fps'] and fps > 30:
this_path = output_dir + "/" + video_name + ".mp4" 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 target_path, exact_fps = this_path, 30
else: else:
shutil.copy(target_path, output_dir) shutil.copy(target_path, output_dir)
status("extracting frames...") status("extracting frames...")
extract_frames(target_path, output_dir) extract_frames(target_path, output_dir, core.globals.use_gpu)
args['frame_paths'] = tuple(sorted( args['frame_paths'] = tuple(sorted(
glob.glob(output_dir + "/*.png"), glob.glob(output_dir + "/*.png"),
key=lambda x: int(x.split(sep)[-1].replace(".png", "")) key=lambda x: int(x.split(sep)[-1].replace(".png", ""))
@ -220,9 +221,9 @@ def start():
status("swapping in progress...") status("swapping in progress...")
start_processing() start_processing()
status("creating video...") 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...") 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" save_path = args['output_file'] if args['output_file'] else output_dir + "/" + video_name + ".mp4"
print("\n\nVideo saved as:", save_path, "\n\n") print("\n\nVideo saved as:", save_path, "\n\n")
status("swap successful!") status("swap successful!")