import gradio as gr import cv2 import os import torch import argparse import os import sys import yaml import datetime sys.path.append(os.path.dirname(os.getcwd())) from pipelines.sd_controlnet_rave import RAVE from pipelines.sd_multicontrolnet_rave import RAVE_MultiControlNet import shutil import subprocess import utils.constants as const import utils.video_grid_utils as vgu import warnings warnings.filterwarnings("ignore") import pprint import glob def init_device(): device_name = 'cuda' if torch.cuda.is_available() else 'cpu' device = torch.device(device_name) return device def init_paths(input_ns): if input_ns.save_folder == None or input_ns.save_folder == '': input_ns.save_folder = input_ns.video_name else: input_ns.save_folder = os.path.join(input_ns.save_folder, input_ns.video_name) save_dir = os.path.join(const.OUTPUT_PATH, input_ns.save_folder) os.makedirs(save_dir, exist_ok=True) save_idx = max([int(x[-5:]) for x in os.listdir(save_dir)])+1 if os.listdir(save_dir) != [] else 0 input_ns.save_path = os.path.join(save_dir, f'{input_ns.positive_prompts}-{str(save_idx).zfill(5)}') if '-' in input_ns.preprocess_name: input_ns.hf_cn_path = [const.PREPROCESSOR_DICT[i] for i in input_ns.preprocess_name.split('-')] else: input_ns.hf_cn_path = const.PREPROCESSOR_DICT[input_ns.preprocess_name] input_ns.hf_path = "runwayml/stable-diffusion-v1-5" input_ns.inverse_path = os.path.join(const.GENERATED_DATA_PATH, 'inverses', input_ns.video_name, f'{input_ns.preprocess_name}_{input_ns.model_id}_{input_ns.grid_size}x{input_ns.grid_size}_{input_ns.pad}') input_ns.control_path = os.path.join(const.GENERATED_DATA_PATH, 'controls', input_ns.video_name, f'{input_ns.preprocess_name}_{input_ns.grid_size}x{input_ns.grid_size}_{input_ns.pad}') os.makedirs(input_ns.control_path, exist_ok=True) os.makedirs(input_ns.inverse_path, exist_ok=True) os.makedirs(input_ns.save_path, exist_ok=True) return input_ns def install_civitai_model(model_id): full_path = os.path.join(const.CWD, 'CIVIT_AI', 'diffusers_models', model_id, '*') if len(glob.glob(full_path)) > 0: full_path = glob.glob(full_path)[0] return full_path install_path = os.path.join(const.CWD, 'CIVIT_AI', 'safetensors') install_path_model = os.path.join(const.CWD, 'CIVIT_AI', 'safetensors', model_id) diffusers_path = os.path.join(const.CWD, 'CIVIT_AI', 'diffusers_models', model_id) convert_py_path = os.path.join(const.CWD, 'CIVIT_AI', 'convert.py') os.makedirs(install_path, exist_ok=True) os.makedirs(diffusers_path, exist_ok=True) subprocess.run(f'wget https://civitai.com/api/download/models/{model_id} --content-disposition --directory {install_path_model}'.split()) model_name = glob.glob(os.path.join(install_path, model_id, '*'))[0] model_name2 = os.path.basename(glob.glob(os.path.join(install_path, model_id, '*'))[0]).replace('.safetensors', '') diffusers_path_model_name = os.path.join(const.CWD, 'CIVIT_AI', 'diffusers_models', model_id, model_name2) print(model_name) subprocess.run(f'python {convert_py_path} --checkpoint_path {model_name} --dump_path {diffusers_path_model_name} --from_safetensors'.split()) subprocess.run(f'rm -rf {install_path}'.split()) return diffusers_path_model_name def run(*args): list_of_inputs = [x for x in args] input_ns = argparse.Namespace(**{}) input_ns.video_path = list_of_inputs[0] # video_path input_ns.video_name = os.path.basename(input_ns.video_path).replace('.mp4', '').replace('.gif', '') input_ns.preprocess_name = list_of_inputs[1] input_ns.batch_size = list_of_inputs[2] input_ns.batch_size_vae = list_of_inputs[3] input_ns.cond_step_start = list_of_inputs[4] input_ns.controlnet_conditioning_scale = list_of_inputs[5] input_ns.controlnet_guidance_end = list_of_inputs[6] input_ns.controlnet_guidance_start = list_of_inputs[7] input_ns.give_control_inversion = list_of_inputs[8] input_ns.grid_size = list_of_inputs[9] input_ns.sample_size = list_of_inputs[10] input_ns.pad = list_of_inputs[11] input_ns.guidance_scale = list_of_inputs[12] input_ns.inversion_prompt = list_of_inputs[13] input_ns.is_ddim_inversion = list_of_inputs[14] input_ns.is_shuffle = list_of_inputs[15] input_ns.negative_prompts = list_of_inputs[16] input_ns.num_inference_steps = list_of_inputs[17] input_ns.num_inversion_step = list_of_inputs[18] input_ns.positive_prompts = list_of_inputs[19] input_ns.save_folder = list_of_inputs[20] input_ns.seed = list_of_inputs[21] input_ns.model_id = const.MODEL_IDS[list_of_inputs[22]] # input_ns.width = list_of_inputs[23] # input_ns.height = list_of_inputs[24] # input_ns.original_size = list_of_inputs[25] diffusers_model_path = os.path.join(const.CWD, 'CIVIT_AI', 'diffusers_models') os.makedirs(diffusers_model_path, exist_ok=True) if 'model_id' not in list(input_ns.__dict__.keys()): input_ns.model_id = "None" if str(input_ns.model_id) != 'None': input_ns.model_id = install_civitai_model(input_ns.model_id) device = init_device() input_ns = init_paths(input_ns) input_ns.image_pil_list = vgu.prepare_video_to_grid(input_ns.video_path, input_ns.sample_size, input_ns.grid_size, input_ns.pad) print(input_ns.video_path) input_ns.sample_size = len(input_ns.image_pil_list) print(f'Frame count: {len(input_ns.image_pil_list)}') controlnet_class = RAVE_MultiControlNet if '-' in str(input_ns.controlnet_conditioning_scale) else RAVE CN = controlnet_class(device) CN.init_models(input_ns.hf_cn_path, input_ns.hf_path, input_ns.preprocess_name, input_ns.model_id) input_dict = vars(input_ns) pp = pprint.PrettyPrinter(indent=4) pp.pprint(input_dict) yaml_dict = {k:v for k,v in input_dict.items() if k != 'image_pil_list'} start_time = datetime.datetime.now() if '-' in str(input_ns.controlnet_conditioning_scale): res_vid, control_vid_1, control_vid_2 = CN(input_dict) else: res_vid, control_vid = CN(input_dict) end_time = datetime.datetime.now() save_name = f"{'-'.join(input_ns.positive_prompts.split())}_cstart-{input_ns.controlnet_guidance_start}_gs-{input_ns.guidance_scale}_pre-{'-'.join((input_ns.preprocess_name.replace('-','+').split('_')))}_cscale-{input_ns.controlnet_conditioning_scale}_grid-{input_ns.grid_size}_pad-{input_ns.pad}_model-{os.path.basename(input_ns.model_id)}" res_vid[0].save(os.path.join(input_ns.save_path, f'{save_name}.gif'), save_all=True, append_images=res_vid[1:], loop=10000) control_vid[0].save(os.path.join(input_ns.save_path, f'control_{save_name}.gif'), save_all=True, append_images=control_vid[1:], optimize=False, loop=10000) yaml_dict['total_time'] = (end_time - start_time).total_seconds() yaml_dict['total_number_of_frames'] = len(res_vid) yaml_dict['sec_per_frame'] = yaml_dict['total_time']/yaml_dict['total_number_of_frames'] with open(os.path.join(input_ns.save_path, 'config.yaml'), 'w') as yaml_file: yaml.dump(yaml_dict, yaml_file) return os.path.join(input_ns.save_path, f'{save_name}.gif'), os.path.join(input_ns.save_path, f'control_{save_name}.gif') def output_video_fn(video_path): fold_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "example_videos") video_path = os.path.join(fold_path, os.path.basename(video_path).replace('input', 'output')) return video_path block = gr.Blocks().queue() with block: with gr.Row(): gr.Markdown('## RAVE') with gr.Row(): with gr.Column(): with gr.Row(): input_path = gr.File(label='Upload Input Video', file_types=['.mp4'], scale=1) inputs = gr.Video(label='Input Video', format='mp4', visible=True, interactive=False, scale=5) input_path.upload(lambda x:x, inputs=[input_path], outputs=[inputs]) with gr.Row(): positive_prompts = gr.Textbox(label='Positive prompts') negative_prompts = gr.Textbox(label='Negative prompts') with gr.Row(): preprocess_name = gr.Dropdown(const.PREPROCESSOR_DICT.keys(), label='Control type', value='depth_zoe') guidance_scale = gr.Slider(label='Guidance scale', minimum=0, maximum=40, step=0.1, value=7.5) with gr.Row(): inversion_prompt = gr.Textbox(label='Inversion prompt') seed = gr.Slider(label='Seed', minimum=0, maximum=2147483647, step=1, value=0, randomize=True) with gr.Row(): model_id = gr.Dropdown(const.MODEL_IDS, label='Model id', value='SD 1.5') save_folder = gr.Textbox(label='Save folder') run_button = gr.Button(value='Run All') with gr.Accordion('Configuration', open=False): with gr.Row(): batch_size = gr.Slider(label='Batch size', minimum=1, maximum=36, value=4, step=1) batch_size_vae = gr.Slider(label='Batch size of VAE', minimum=1, maximum=36, value=1, step=1) with gr.Row(): is_ddim_inversion = gr.Checkbox( label='Use DDIM Inversion', value=True) is_shuffle = gr.Checkbox( label='Shuffle', value=True) with gr.Row(): num_inference_steps = gr.Slider(label='Number of inference steps', minimum=1, maximum=100, value=20, step=1) num_inversion_step = gr.Slider(label='Number of inversion steps', minimum=1, maximum=100, value=20, step=1) cond_step_start = gr.Slider(label='Conditioning step start', minimum=0, maximum=1.0, value=0.0, step=0.1) with gr.Row(): controlnet_conditioning_scale = gr.Slider(label='ControlNet conditioning scale', minimum=0.0, maximum=1.0, value=1.0, step=0.01) controlnet_guidance_end = gr.Slider(label='ControlNet guidance end', minimum=0.0, maximum=1.0, value=1.0, step=0.01) controlnet_guidance_start = gr.Slider(label='ControlNet guidance start', minimum=0.0, maximum=1.0, value=0.0, step=0.01) give_control_inversion = gr.Checkbox( label='Give control during inversion', value=True) with gr.Row(): grid_size = gr.Slider(label='Grid size', minimum=1, maximum=10, value=3, step=1) sample_size = gr.Slider(label='Sample size', minimum=-1, maximum=100, value=-1, step=1) pad = gr.Slider(label='Pad', minimum=1, maximum=10, value=1, step=1) with gr.Column(): with gr.Row(): result_video = gr.Image(label='Edited Video', interactive=False) control_video = gr.Image(label='Control Video', interactive=False) with gr.Row(): example_input = gr.Video(label='Input Example', format='mp4', visible=True, interactive=False) example_output = gr.Video(label='Output Example', format='mp4', visible=True, interactive=False) # input(os.path.join(os.path.dirname(os.path.abspath(__file__)), "example_videos", "exp_input_1.mp4")) gr.Markdown("## Video Examples") gr.Examples( examples=[os.path.join(os.path.dirname(os.path.abspath(__file__)), "example_videos", "exp_input_1.mp4")], inputs=example_input, outputs=example_output, fn=output_video_fn, cache_examples=True,) inputs = [input_path, preprocess_name, batch_size, batch_size_vae, cond_step_start, controlnet_conditioning_scale, controlnet_guidance_end, controlnet_guidance_start, give_control_inversion, grid_size, sample_size, pad, guidance_scale, inversion_prompt, is_ddim_inversion, is_shuffle, negative_prompts, num_inference_steps, num_inversion_step, positive_prompts, save_folder, seed, model_id] run_button.click(fn=run, inputs=inputs, outputs=[result_video, control_video]) if __name__ == "__main__": block.launch(share=True)