99 lines
5.7 KiB
Python
99 lines
5.7 KiB
Python
import os
|
|
from pathlib import Path
|
|
from rife.inference_video import run_rife_new_video_infer
|
|
from .video_audio_utilities import get_quick_vid_info, vid2frames, media_file_has_audio, extract_number
|
|
|
|
# gets 'RIFE v4.3', returns: 'RIFE43'
|
|
def extract_rife_name(string):
|
|
parts = string.split()
|
|
if len(parts) != 2 or parts[0] != "RIFE" or (parts[1][0] != "v" or not parts[1][1:].replace('.','').isdigit()):
|
|
raise ValueError("Input string should contain exactly 2 words, first word should be 'RIFE' and second word should start with 'v' followed by 2 numbers")
|
|
return "RIFE"+parts[1][1:].replace('.','')
|
|
|
|
# This function usually gets a filename, and converts it to a legal linux/windows *folder* name
|
|
def clean_folder_name(string):
|
|
illegal_chars = ["/", "\\", "<", ">", ":", "\"", "|", "?", "*", "."]
|
|
for char in illegal_chars:
|
|
string = string.replace(char, "_")
|
|
return string
|
|
|
|
def set_interp_out_fps(interp_x, slom_x, in_vid_fps):
|
|
if interp_x == 'Disabled' or in_vid_fps in ('---', None, '', 'None'):
|
|
return '---'
|
|
|
|
clean_interp_x = extract_number(interp_x)
|
|
clean_slom_x = extract_number(slom_x)
|
|
fps = float(in_vid_fps) * int(clean_interp_x)
|
|
if clean_slom_x != -1:
|
|
fps /= int(clean_slom_x)
|
|
return int(fps) if fps.is_integer() else fps
|
|
|
|
# get uploaded video frame count, fps, and return 3 valuees for the gradio UI: in fcount, in fps, out fps (using the set_interp_out_fps function above)
|
|
def gradio_f_interp_get_fps_and_fcount(vid_path, interp_x, slom_x):
|
|
if vid_path is None:
|
|
return '---', '---', '---'
|
|
fps, fcount, resolution = get_quick_vid_info(vid_path.name)
|
|
expected_out_fps = set_interp_out_fps(interp_x, slom_x, fps)
|
|
return (fps if fps is not None else '---', fcount if fcount is not None else '---', expected_out_fps)
|
|
|
|
# handle call to interpolate an uploaded video from gradio button in args.py (the function that calls this func is named 'upload_vid_to_rife')
|
|
def process_rife_vid_upload_logic(file, engine, x_am, sl_am, keep_imgs, f_location, f_crf, f_preset, in_vid_fps, f_models_path, vid_file_name):
|
|
|
|
print("got a request to *frame interpolate* an existing video.")
|
|
|
|
_, _, resolution = get_quick_vid_info(file.name)
|
|
folder_name = clean_folder_name(Path(vid_file_name).stem)
|
|
outdir_no_tmp = os.path.join(os.getcwd(), 'outputs', 'frame-interpolation', folder_name)
|
|
i = 1
|
|
while os.path.exists(outdir_no_tmp):
|
|
outdir_no_tmp = os.path.join(os.getcwd(), 'outputs', 'frame-interpolation', folder_name + '_' + str(i))
|
|
i += 1
|
|
|
|
outdir = os.path.join(outdir_no_tmp, 'tmp_input_frames')
|
|
os.makedirs(outdir, exist_ok=True)
|
|
|
|
vid2frames(video_path=file.name, video_in_frame_path=outdir, overwrite=True, extract_from_frame=0, extract_to_frame=-1, numeric_files_output=True, out_img_format='png')
|
|
|
|
# check if the uploaded vid has an audio stream. If it doesn't, set audio param to None so that ffmpeg won't try to add non-existing audio to final video.
|
|
audio_file_to_pass = None
|
|
if media_file_has_audio(file.name, f_location):
|
|
audio_file_to_pass = file.name
|
|
|
|
process_video_interpolation(frame_interpolation_engine=engine, frame_interpolation_x_amount=x_am, frame_interpolation_slow_mo_amount=sl_am, orig_vid_fps=in_vid_fps, deforum_models_path=f_models_path, real_audio_track=audio_file_to_pass, raw_output_imgs_path=outdir, img_batch_id=None, ffmpeg_location=f_location, ffmpeg_crf=f_crf, ffmpeg_preset=f_preset, keep_interp_imgs=keep_imgs, orig_vid_name=folder_name, resolution=resolution)
|
|
|
|
# handle params before talking with the actual rife module
|
|
def process_video_interpolation(frame_interpolation_engine, frame_interpolation_x_amount, frame_interpolation_slow_mo_amount, orig_vid_fps, deforum_models_path, real_audio_track, raw_output_imgs_path, img_batch_id, ffmpeg_location, ffmpeg_crf, ffmpeg_preset, keep_interp_imgs, orig_vid_name, resolution):
|
|
|
|
if frame_interpolation_x_amount == "Disabled" or not frame_interpolation_engine.startswith("RIFE"):
|
|
return
|
|
|
|
# extract clean numbers from values of 'x2' etc'
|
|
interp_amount = extract_number(frame_interpolation_x_amount)
|
|
interp_slow_mo = extract_number(frame_interpolation_slow_mo_amount)
|
|
# make sure interp_x is valid and in range
|
|
if interp_amount not in range(2, 11):
|
|
raise Error("frame_interpolation_x_amount must be between 2x and 10x")
|
|
|
|
# set initial output vid fps
|
|
fps = float(orig_vid_fps) * interp_amount
|
|
|
|
# re-calculate fps param to pass if slow_mo mode is enabled
|
|
if interp_slow_mo != -1:
|
|
if interp_slow_mo not in [2, 4, 8]:
|
|
raise Error("frame_interpolation_slow_mo_amount must be 2x, 4x or 8x")
|
|
fps = float(orig_vid_fps) * int(interp_amount) / int(interp_slow_mo)
|
|
|
|
# disable audio-adding by setting real_audio_track to None if slow-mo is enabled (not equal to -1 means enabled)
|
|
if real_audio_track is not None and interp_slow_mo != -1:
|
|
real_audio_track = None
|
|
|
|
# set UHD to True if res' is 2K or higher
|
|
if resolution:
|
|
UHD = resolution[0] >= 2048 and resolution[1] >= 2048
|
|
else:
|
|
UHD = False
|
|
# e.g from "RIFE v2.3 to RIFE23"
|
|
actual_model_folder_name = extract_rife_name(frame_interpolation_engine)
|
|
|
|
# run actual rife interpolation and video stitching etc - the whole suite
|
|
run_rife_new_video_infer(interp_x_amount=interp_amount, slow_mo_x_amount=interp_slow_mo, model=actual_model_folder_name, fps=fps, deforum_models_path=deforum_models_path, audio_track=real_audio_track, raw_output_imgs_path=raw_output_imgs_path, img_batch_id=img_batch_id, ffmpeg_location=ffmpeg_location, ffmpeg_crf=ffmpeg_crf, ffmpeg_preset=ffmpeg_preset, keep_imgs=keep_interp_imgs, orig_vid_name=orig_vid_name, UHD=UHD) |