from FateZero.test_fatezero import * import copy import gradio as gr class merge_config_then_run(): def __init__(self) -> None: # Load the tokenizer pretrained_model_path = 'FateZero/ckpt/stable-diffusion-v1-4' self.tokenizer = None self.text_encoder = None self.vae = None self.unet = None cache_ckpt = True if cache_ckpt: self.tokenizer = AutoTokenizer.from_pretrained( pretrained_model_path, # 'FateZero/ckpt/stable-diffusion-v1-4', subfolder="tokenizer", use_fast=False, ) # Load models and create wrapper for stable diffusion self.text_encoder = CLIPTextModel.from_pretrained( pretrained_model_path, subfolder="text_encoder", ) self.vae = AutoencoderKL.from_pretrained( pretrained_model_path, subfolder="vae", ) model_config = { "lora": 160, # temporal_downsample_time: 4 "SparseCausalAttention_index": ['mid'], "least_sc_channel": 640 } self.unet = UNetPseudo3DConditionModel.from_2d_model( os.path.join(pretrained_model_path, "unet"), model_config=model_config ) def run( self, # def merge_config_then_run( model_id, data_path, source_prompt, target_prompt, cross_replace_steps, self_replace_steps, enhance_words, enhance_words_value, num_steps, guidance_scale, user_input_video=None, # Temporal and spatial crop of the video start_sample_frame=0, n_sample_frame=8, stride=1, left_crop=0, right_crop=0, top_crop=0, bottom_crop=0, ): # , ] = inputs default_edit_config='FateZero/config/low_resource_teaser/jeep_watercolor_ddim_10_steps.yaml' Omegadict_default_edit_config = OmegaConf.load(default_edit_config) dataset_time_string = get_time_string() config_now = copy.deepcopy(Omegadict_default_edit_config) print(f"config_now['pretrained_model_path'] = model_id {model_id}") # config_now['pretrained_model_path'] = model_id config_now['train_dataset']['prompt'] = source_prompt config_now['train_dataset']['path'] = data_path # ImageSequenceDataset_dict = { } offset_dict = { "left": left_crop, "right": right_crop, "top": top_crop, "bottom": bottom_crop, } ImageSequenceDataset_dict = { "start_sample_frame" : start_sample_frame, "n_sample_frame" : n_sample_frame, "stride" : stride, "offset": offset_dict, } config_now['train_dataset'].update(ImageSequenceDataset_dict) if user_input_video and data_path is None: raise gr.Error('You need to upload a video or choose a provided video') if user_input_video is not None and user_input_video.name is not None: config_now['train_dataset']['path'] = user_input_video.name config_now['validation_sample_logger_config']['prompts'] = [target_prompt] # fatezero config p2p_config_now = copy.deepcopy(config_now['validation_sample_logger_config']['p2p_config'][0]) p2p_config_now['cross_replace_steps']['default_'] = cross_replace_steps p2p_config_now['self_replace_steps'] = self_replace_steps p2p_config_now['eq_params']['words'] = enhance_words.split(" ") p2p_config_now['eq_params']['values'] = [enhance_words_value,]*len(p2p_config_now['eq_params']['words']) config_now['validation_sample_logger_config']['p2p_config'][0] = copy.deepcopy(p2p_config_now) # ddim config config_now['validation_sample_logger_config']['guidance_scale'] = guidance_scale config_now['validation_sample_logger_config']['num_inference_steps'] = num_steps logdir = default_edit_config.replace('config', 'result').replace('.yml', '').replace('.yaml', '')+f'_{dataset_time_string}' config_now['logdir'] = logdir print(f'Saving at {logdir}') save_path = test(tokenizer = self.tokenizer, text_encoder = self.text_encoder, vae = self.vae, unet = self.unet, config=default_edit_config, **config_now) mp4_path = save_path.replace('_0.gif', '_0_0_0.mp4') return mp4_path