cache the ckpt; fix bugs when input new video

.gitignore

@@ -1 +1,2 @@
-trash/*
+trash/*
+tmp

FateZero/test_fatezero.py

@@ -48,6 +48,10 @@ def test(
     config: str,
     pretrained_model_path: str,
     train_dataset: Dict,
+    tokenizer = None,
+    text_encoder = None,
+    vae = None,
+    unet = None,
     logdir: str = None,
     validation_sample_logger_config: Optional[Dict] = None,
     test_pipeline_config: Optional[Dict] = None,
@@ -79,26 +83,28 @@ def test(
         set_seed(seed)
 
     # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_path,
-        subfolder="tokenizer",
-        use_fast=False,
-    )
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_path,
+            subfolder="tokenizer",
+            use_fast=False,
+        )
 
     # Load models and create wrapper for stable diffusion
-    text_encoder = CLIPTextModel.from_pretrained(
-        pretrained_model_path,
-        subfolder="text_encoder",
-    )
-
-    vae = AutoencoderKL.from_pretrained(
-        pretrained_model_path,
-        subfolder="vae",
-    )
-
-    unet = UNetPseudo3DConditionModel.from_2d_model(
-        os.path.join(pretrained_model_path, "unet"), model_config=model_config
-    )
+    if text_encoder is None:
+        text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_path,
+            subfolder="text_encoder",
+        )
+    if vae is None:
+        vae = AutoencoderKL.from_pretrained(
+            pretrained_model_path,
+            subfolder="vae",
+        )
+    if unet is None:
+        unet = UNetPseudo3DConditionModel.from_2d_model(
+            os.path.join(pretrained_model_path, "unet"), model_config=model_config
+        )
 
     if 'target' not in test_pipeline_config:
         test_pipeline_config['target'] = 'video_diffusion.pipelines.stable_diffusion.SpatioTemporalStableDiffusionPipeline'

FateZero/video_diffusion/common/util.py

@@ -4,7 +4,7 @@ import copy
 import inspect
 import datetime
 from typing import List, Tuple, Optional, Dict
-
+import torch
 
 def glob_files(
     root_path: str,
@@ -68,6 +68,12 @@ def get_time_string() -> str:
 def get_function_args() -> Dict:
     frame = sys._getframe(1)
     args, _, _, values = inspect.getargvalues(frame)
-    args_dict = copy.deepcopy({arg: values[arg] for arg in args})
+    tmp_dict = {}
+    for arg in args: 
+        v = values[arg]
+        if not isinstance(v, torch.nn.Module) and arg !='tokenizer' :
+            tmp_dict[arg] = v
+    
+    args_dict = copy.deepcopy(tmp_dict)
 
     return args_dict

FateZero/video_diffusion/data/dataset.py

@@ -6,6 +6,7 @@ from einops import rearrange
 from pathlib import Path
 import imageio
 import cv2
+import shutil
 
 import torch
 from torch.utils.data import Dataset
@@ -156,7 +157,7 @@ class ImageSequenceDataset(Dataset):
         images = []
         if path[-4:] == '.mp4':
             path = self.mp4_to_png(path)
-            self.path = path
+            
 
         for file in sorted(os.listdir(path)):
             if file.endswith(IMAGE_EXTENSION):
@@ -164,14 +165,19 @@ class ImageSequenceDataset(Dataset):
         return images
     
     # @staticmethod
+    
     def mp4_to_png(self, video_source=None):
         reader = imageio.get_reader(video_source)
-        os.makedirs(video_source[:-4], exist_ok=True)
-
+        dir_path = './tmp/fatezero_user_video'
+        if os.path.exists(dir_path):
+            shutil.rmtree(dir_path)
+        os.makedirs(dir_path, exist_ok=True)
+        
         for i, im in enumerate(reader):
             # use :05d to add zero, no space before the 05d
             # if (i+1)%10 == 0:
-            path = os.path.join(video_source[:-4], f"{i:05d}.png")
+            path = os.path.join(dir_path, f"{i:05d}.png")
             # print(path)
             cv2.imwrite(path, im[:, :, ::-1])
-        return video_source[:-4]
+        self.path = dir_path
+        return self.path

app_fatezero.py

@@ -28,7 +28,7 @@ from inference_fatezero import merge_config_then_run
 # TITLE = '# [FateZero](http://fate-zero-edit.github.io/)'
 HF_TOKEN = os.getenv('HF_TOKEN')
 # pipe = InferencePipeline(HF_TOKEN)
-# pipe = merge_config_then_run
+pipe = merge_config_then_run()
 # app = InferenceUtil(HF_TOKEN)
 
 with gr.Blocks(css='style.css') as demo:
@@ -288,7 +288,7 @@ with gr.Blocks(css='style.css') as demo:
                         *ImageSequenceDataset_list
                     ],
                     outputs=result,
-                    fn=merge_config_then_run,
+                    fn=pipe.run,
                     cache_examples=os.getenv('SYSTEM') == 'spaces')
 
     # model_id.change(fn=app.load_model_info,
@@ -312,8 +312,8 @@ with gr.Blocks(css='style.css') as demo:
             *ImageSequenceDataset_list
     ]
     # prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
-    target_prompt.submit(fn=merge_config_then_run, inputs=inputs, outputs=result)
+    target_prompt.submit(fn=pipe.run, inputs=inputs, outputs=result)
     # run_button.click(fn=pipe.run, inputs=inputs, outputs=result)
-    run_button.click(fn=merge_config_then_run, inputs=inputs, outputs=result)
+    run_button.click(fn=pipe.run, inputs=inputs, outputs=result)
 
 demo.queue().launch()

inference_fatezero.py

@@ -4,8 +4,40 @@ 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 = AutoTokenizer.from_pretrained(
+            pretrained_model_path,
+            # 'FateZero/ckpt/stable-diffusion-v1-4',
+            subfolder="tokenizer",
+            use_fast=False,
+        )
 
-def merge_config_then_run(
+        # 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,
@@ -27,58 +59,59 @@ def merge_config_then_run(
         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)
+        # , ] = 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(config=default_edit_config, **config_now)
-    mp4_path = save_path.replace('_0.gif', '_0_0_0.mp4')
-    return mp4_path
+        # 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
 
-if __name__ == "__main__":
-    run()