update VideoCrafter1

.gitignore

@@ -8,3 +8,4 @@ results
 *.ckpt
 *.pt
 *.pth
+checkpoints

README.md

@@ -1,13 +0,0 @@
----
-title: VideoCrafter
-emoji: 🌍
-colorFrom: gray
-colorTo: red
-sdk: gradio
-sdk_version: 3.24.1
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,30 +1,30 @@
 import os
 import sys
 import gradio as gr
-# from demo_test import Text2Video, VideoControl
-from videocrafter_test import Text2Video
-from videocontrol_test import VideoControl
+# from demo_test import Text2Video, Image2Video
+from t2v_test import Text2Video
+from i2v_test import Image2Video
 sys.path.insert(1, os.path.join(sys.path[0], 'lvdm'))
 
 t2v_examples = [
-    ['an elephant is walking under the sea, 4K, high definition',50,'origin',1,15,1,],
-    ['an astronaut riding a horse in outer space',25,'origin',1,15,1,],
-    ['a monkey is playing a piano',25,'vangogh',1,15,1,],
-    ['A fire is burning on a candle',25,'frozen',1,15,1,],
-    ['a horse is drinking in the river',25,'yourname',1,15,1,],
-    ['Robot dancing in times square',25,'coco',1,15,1,],                    
+    ['an elephant is walking under the sea, 4K, high definition',50, 12,1, 16],
+    ['an astronaut riding a horse in outer space',25,12,1,16],
+    ['a monkey is playing a piano',25,12,1,16],
+    ['A fire is burning on a candle',25,12,1,16],
+    ['a horse is drinking in the river',25,12,1,16],
+    ['Robot dancing in times square',25,12,1,16],                    
 ]
 
-control_examples = [
-    ['input/flamingo.mp4', 'An ostrich walking in the desert, photorealistic, 4k', 0, 50, 15, 1, 16, 256]
+i2v_examples = [
+    ['prompts/i2v_prompts/horse.png', 'horses are walking on the grassland', 50, 12, 1, 16]
 ]
 
 def videocrafter_demo(result_dir='./tmp/'):
     text2video = Text2Video(result_dir)
-    videocontrol = VideoControl(result_dir)
+    image2video = Image2Video(result_dir)
     with gr.Blocks(analytics_enabled=False) as videocrafter_iface:
-        gr.Markdown("<div align='center'> <h2> VideoCrafter: A Toolkit for Text-to-Video Generation and Editing </span> </h2> \
-                     <a style='font-size:18px;color: #000000' href='https://github.com/VideoCrafter/VideoCrafter'> Github </div>")
+        gr.Markdown("<div align='center'> <h2> VideoCrafter1: Open Diffusion Models for High-Quality Video Generation </span> </h2> \
+                     <a style='font-size:18px;color: #000000' href='https://github.com/AILab-CVC/VideoCrafter'> Github </div>")
         
         gr.Markdown("<b> You may duplicate the space and upgrade to GPU in settings for better performance and faster inference without waiting in the queue. <a style='display:inline-block' href='https://huggingface.co/spaces/VideoCrafter/VideoCrafter?duplicate=true'> <img src='https://bit.ly/3gLdBN6' alt='Duplicate Space'></a> </b>")
         #######t2v#######
@@ -33,65 +33,55 @@ def videocrafter_demo(result_dir='./tmp/'):
                 with gr.Row().style(equal_height=False):
                     with gr.Column():
                         input_text = gr.Text(label='Prompts')
-                        model_choices=['origin','vangogh','frozen','yourname', 'coco']
-                        with gr.Row():
-                            model_index = gr.Dropdown(label='Models', elem_id=f"model", choices=model_choices, value=model_choices[0], type="index",interactive=True)
                         with gr.Row():
                             steps = gr.Slider(minimum=1, maximum=60, step=1, elem_id=f"steps", label="Sampling steps", value=50)
                             eta = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label='ETA', value=1.0, elem_id="eta")
                         with gr.Row():
-                            lora_scale = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, label='Lora Scale', value=1.0, elem_id="lora_scale")
-                            cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=15.0, elem_id="cfg_scale")
+                            cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=12.0, elem_id="cfg_scale")
+                            fps = gr.Slider(minimum=4, maximum=32, step=1, label='fps', value=16, elem_id="fps")
                         send_btn = gr.Button("Send")
                     with gr.Tab(label='result'):
-                        output_video_1 =  gr.Video().style(width=384)
+                        output_video_1 =  gr.Video().style(width=320)
                 gr.Examples(examples=t2v_examples,
-                            inputs=[input_text,steps,model_index,eta,cfg_scale,lora_scale],
+                            inputs=[input_text,steps,cfg_scale,eta],
                             outputs=[output_video_1],
                             fn=text2video.get_prompt,
                             cache_examples=False)
                         #cache_examples=os.getenv('SYSTEM') == 'spaces')
             send_btn.click(
                 fn=text2video.get_prompt, 
-                inputs=[input_text,steps,model_index,eta,cfg_scale,lora_scale,],
+                inputs=[input_text,steps,cfg_scale,eta,fps],
                 outputs=[output_video_1],
             )
-        #######videocontrol######
-        with gr.Tab(label='VideoControl'):
+        #######image2video######
+        with gr.Tab(label='Image2Video'):
             with gr.Column():
                 with gr.Row():
-                    # with gr.Tab(label='input'):
                     with gr.Column():
                         with gr.Row():
-                            vc_input_video = gr.Video(label="Input Video").style(width=256)
-                            vc_origin_video = gr.Video(label='Center-cropped Video').style(width=256)
-                        with gr.Row():
-                            vc_input_text = gr.Text(label='Prompts')
+                            i2v_input_image = gr.Image(label="Input Image").style(width=256)
                         with gr.Row():
-                            vc_eta = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label='ETA', value=1.0, elem_id="vc_eta")
-                            vc_cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=15.0, elem_id="vc_cfg_scale")
+                            i2v_input_text = gr.Text(label='Prompts')
                         with gr.Row():
-                            vc_steps = gr.Slider(minimum=1, maximum=60, step=1, elem_id="vc_steps", label="Sampling steps", value=50)
-                            frame_stride = gr.Slider(minimum=0 , maximum=100, step=1, label='Frame Stride', value=0, elem_id="vc_frame_stride")
+                            i2v_eta = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label='ETA', value=1.0, elem_id="i2v_eta")
+                            i2v_cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=12.0, elem_id="i2v_cfg_scale")
                         with gr.Row():
-                            resolution = gr.Slider(minimum=128 , maximum=512, step=8, label='Long Side Resolution', value=256, elem_id="vc_resolution")
-                            video_frames = gr.Slider(minimum=8 , maximum=64, step=1, label='Video Frame Num', value=16, elem_id="vc_video_frames")
-                        vc_end_btn = gr.Button("Send")
+                            i2v_steps = gr.Slider(minimum=1, maximum=60, step=1, elem_id="i2v_steps", label="Sampling steps", value=50)
+                            i2v_fps = gr.Slider(minimum=4, maximum=32, step=1, elem_id="i2v_fps", label="Generative fps", value=16)
+                        i2v_end_btn = gr.Button("Send")
                     with gr.Tab(label='Result'):
-                        vc_output_info = gr.Text(label='Info')
                         with gr.Row():
-                            vc_depth_video = gr.Video(label="Depth Video").style(width=256)
-                            vc_output_video = gr.Video(label="Generated Video").style(width=256)
+                            i2v_output_video = gr.Video(label="Generated Video").style(width=320)
 
-                gr.Examples(examples=control_examples,
-                            inputs=[vc_input_video, vc_input_text, frame_stride, vc_steps, vc_cfg_scale, vc_eta, video_frames, resolution],
-                            outputs=[vc_output_info, vc_origin_video, vc_depth_video, vc_output_video],
-                            fn = videocontrol.get_video,
+                gr.Examples(examples=i2v_examples,
+                            inputs=[i2v_input_image, i2v_input_text, i2v_steps, i2v_cfg_scale, i2v_eta, i2v_fps],
+                            outputs=[i2v_output_video],
+                            fn = image2video.get_image,
                             cache_examples=os.getenv('SYSTEM') == 'spaces',
                 )
-            vc_end_btn.click(inputs=[vc_input_video, vc_input_text, frame_stride, vc_steps, vc_cfg_scale, vc_eta, video_frames, resolution],
-                            outputs=[vc_output_info, vc_origin_video, vc_depth_video, vc_output_video],
-                            fn = videocontrol.get_video
+            i2v_end_btn.click(inputs=[i2v_input_image, i2v_input_text, i2v_steps, i2v_cfg_scale, i2v_eta, i2v_fps],
+                            outputs=[i2v_output_video],
+                            fn = image2video.get_image
             )
 
     return videocrafter_iface

models/adapter_t2v_depth/model_config.yaml → configs/inference_i2v_512_v1.0.yaml

@@ -1,27 +1,28 @@
 model:
-  target: lvdm.models.ddpm3d.T2VAdapterDepth
+  target: lvdm.models.ddpm3d.LatentVisualDiffusion
   params:
     linear_start: 0.00085
     linear_end: 0.012
     num_timesteps_cond: 1
-    log_every_t: 200
     timesteps: 1000
     first_stage_key: video
     cond_stage_key: caption
-    image_size:
-    - 32
-    - 32
-    video_length: 16
-    channels: 4
     cond_stage_trainable: false
     conditioning_key: crossattn
+    image_size:
+    - 40
+    - 64
+    channels: 4
     scale_by_std: false
     scale_factor: 0.18215
-
+    use_ema: false
+    uncond_type: empty_seq
+    use_scale: true
+    scale_b: 0.7
+    finegrained: true
     unet_config:
-      target: lvdm.models.modules.openaimodel3d.UNetModel
+      target: lvdm.modules.networks.openaimodel3d.UNetModel
       params:
-        image_size: 32
         in_channels: 4
         out_channels: 4
         model_channels: 320
@@ -35,16 +36,20 @@ model:
         - 2
         - 4
         - 4
-        num_heads: 8
+        num_head_channels: 64
         transformer_depth: 1
-        context_dim: 768
+        context_dim: 1024
+        use_linear: true
         use_checkpoint: true
-        legacy: false
-        kernel_size_t: 1
-        padding_t: 0
+        temporal_conv: true
+        temporal_attention: true
+        temporal_selfatt_only: true
+        use_relative_position: false
+        use_causal_attention: false
+        use_image_attention: true
         temporal_length: 16
-        use_relative_position: true
-        
+        addition_attention: true
+        fps_cond: true
     first_stage_config:
       target: lvdm.models.autoencoder.AutoencoderKL
       params:
@@ -53,7 +58,7 @@ model:
         ddconfig:
           double_z: true
           z_channels: 4
-          resolution: 256
+          resolution: 512
           in_channels: 3
           out_ch: 3
           ch: 128
@@ -67,23 +72,12 @@ model:
           dropout: 0.0
         lossconfig:
           target: torch.nn.Identity
-
     cond_stage_config:
-      target: lvdm.models.modules.condition_modules.FrozenCLIPEmbedder
-      
-    depth_stage_config:
-      target: extralibs.midas.api.MiDaSInference
+      target: lvdm.modules.encoders.condition.FrozenOpenCLIPEmbedder
+      params:
+        freeze: true
+        layer: penultimate
+    cond_img_config:
+      target: lvdm.modules.encoders.condition.FrozenOpenCLIPImageEmbedderV2
       params:
-        model_type: "dpt_hybrid"
-        model_path: models/adapter_t2v_depth/dpt_hybrid-midas.pt
-
-    adapter_config:
-        target: lvdm.models.modules.adapter.Adapter
-        cond_name: depth
-        params:
-          cin: 64
-          channels: [320, 640, 1280, 1280]
-          nums_rb: 2
-          ksize: 1
-          sk: True
-          use_conv: False
+        freeze: true

models/base_t2v/model_config.yaml → configs/inference_t2v_1024_v1.0.yaml

@@ -4,24 +4,24 @@ model:
     linear_start: 0.00085
     linear_end: 0.012
     num_timesteps_cond: 1
-    log_every_t: 200
     timesteps: 1000
     first_stage_key: video
     cond_stage_key: caption
-    image_size:
-    - 32
-    - 32
-    video_length: 16
-    channels: 4
     cond_stage_trainable: false
     conditioning_key: crossattn
+    image_size:
+    - 72
+    - 128
+    channels: 4
     scale_by_std: false
     scale_factor: 0.18215
-
+    use_ema: false
+    uncond_type: empty_seq
+    use_scale: true
+    fix_scale_bug: true
     unet_config:
-      target: lvdm.models.modules.openaimodel3d.UNetModel
+      target: lvdm.modules.networks.openaimodel3d.UNetModel
       params:
-        image_size: 32
         in_channels: 4
         out_channels: 4
         model_channels: 320
@@ -35,16 +35,19 @@ model:
         - 2
         - 4
         - 4
-        num_heads: 8
+        num_head_channels: 64
         transformer_depth: 1
-        context_dim: 768
+        context_dim: 1024
+        use_linear: true
         use_checkpoint: true
-        legacy: false
-        kernel_size_t: 1
-        padding_t: 0
-        temporal_length: 16
+        temporal_conv: false
+        temporal_attention: true
+        temporal_selfatt_only: true
         use_relative_position: true
-        
+        use_causal_attention: false
+        temporal_length: 16
+        addition_attention: true
+        fps_cond: true
     first_stage_config:
       target: lvdm.models.autoencoder.AutoencoderKL
       params:
@@ -53,7 +56,7 @@ model:
         ddconfig:
           double_z: true
           z_channels: 4
-          resolution: 256
+          resolution: 512
           in_channels: 3
           out_ch: 3
           ch: 128
@@ -67,6 +70,8 @@ model:
           dropout: 0.0
         lossconfig:
           target: torch.nn.Identity
-
     cond_stage_config:
-      target: lvdm.models.modules.condition_modules.FrozenCLIPEmbedder
+      target: lvdm.modules.encoders.condition.FrozenOpenCLIPEmbedder
+      params:
+        freeze: true
+        layer: penultimate

configs/inference_t2v_512_v1.0.yaml

@@ -0,0 +1,77 @@
+model:
+  target: lvdm.models.ddpm3d.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.012
+    num_timesteps_cond: 1
+    timesteps: 1000
+    first_stage_key: video
+    cond_stage_key: caption
+    cond_stage_trainable: false
+    conditioning_key: crossattn
+    image_size:
+    - 40
+    - 64
+    channels: 4
+    scale_by_std: false
+    scale_factor: 0.18215
+    use_ema: false
+    uncond_type: empty_seq
+    use_scale: true
+    scale_b: 0.7
+    unet_config:
+      target: lvdm.modules.networks.openaimodel3d.UNetModel
+      params:
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions:
+        - 4
+        - 2
+        - 1
+        num_res_blocks: 2
+        channel_mult:
+        - 1
+        - 2
+        - 4
+        - 4
+        num_head_channels: 64
+        transformer_depth: 1
+        context_dim: 1024
+        use_linear: true
+        use_checkpoint: true
+        temporal_conv: true
+        temporal_attention: true
+        temporal_selfatt_only: true
+        use_relative_position: false
+        use_causal_attention: false
+        temporal_length: 16
+        addition_attention: true
+        fps_cond: true
+    first_stage_config:
+      target: lvdm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 512
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+    cond_stage_config:
+      target: lvdm.modules.encoders.condition.FrozenOpenCLIPEmbedder
+      params:
+        freeze: true
+        layer: penultimate

demo_test.py

@@ -2,15 +2,15 @@ class Text2Video():
     def __init__(self, result_dir='./tmp/') -> None:
         pass
 
-    def get_prompt(self, input_text, steps=50, model_index=0, eta=1.0, cfg_scale=15.0, lora_scale=1.0):
+    def get_prompt(self, input_text, steps=50, cfg_scale=15.0, eta=1.0, fps=16):
 
         return '01.mp4'
     
-class VideoControl:
+class Image2Video:
     def __init__(self, result_dir='./tmp/') -> None:
         pass
 
-    def get_video(self, input_video, input_prompt, frame_stride=0, vc_steps=50, vc_cfg_scale=15.0, vc_eta=1.0, video_frames=16, resolution=256):
+    def get_image(self, input_image, input_prompt, i2v_steps=50, i2v_cfg_scale=15.0, i2v_eta=1.0, i2v_fps=16):
 
-        return 'su','01.mp4'
+        return '01.mp4'
         

extralibs/midas/api.py

@@ -1,171 +0,0 @@
-# based on https://github.com/isl-org/MiDaS
-
-import cv2
-import torch
-import torch.nn as nn
-from torchvision.transforms import Compose
-
-from extralibs.midas.midas.dpt_depth import DPTDepthModel
-from extralibs.midas.midas.midas_net import MidasNet
-from extralibs.midas.midas.midas_net_custom import MidasNet_small
-from extralibs.midas.midas.transforms import Resize, NormalizeImage, PrepareForNet
-
-
-ISL_PATHS = {
-    "dpt_large": "midas_models/dpt_large-midas-2f21e586.pt",
-    "dpt_hybrid": "midas_models/dpt_hybrid-midas-501f0c75.pt",
-    "midas_v21": "",
-    "midas_v21_small": "",
-}
-
-
-def disabled_train(self, mode=True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-def load_midas_transform(model_type):
-    # https://github.com/isl-org/MiDaS/blob/master/run.py
-    # load transform only
-    if model_type == "dpt_large":  # DPT-Large
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "dpt_hybrid":  # DPT-Hybrid
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "midas_v21":
-        net_w, net_h = 384, 384
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
-    elif model_type == "midas_v21_small":
-        net_w, net_h = 256, 256
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-
-    else:
-        assert False, f"model_type '{model_type}' not implemented, use: --model_type large"
-
-    transform = Compose(
-        [
-            Resize(
-                net_w,
-                net_h,
-                resize_target=None,
-                keep_aspect_ratio=True,
-                ensure_multiple_of=32,
-                resize_method=resize_mode,
-                image_interpolation_method=cv2.INTER_CUBIC,
-            ),
-            normalization,
-            PrepareForNet(),
-        ]
-    )
-
-    return transform
-
-
-def load_model(model_type, model_path=None):
-    # https://github.com/isl-org/MiDaS/blob/master/run.py
-    # load network
-    if model_path is None:
-        model_path = ISL_PATHS[model_type]
-    if model_type == "dpt_large":  # DPT-Large
-        model = DPTDepthModel(
-            path=model_path,
-            backbone="vitl16_384",
-            non_negative=True,
-        )
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "dpt_hybrid":  # DPT-Hybrid
-        model = DPTDepthModel(
-            path=model_path,
-            backbone="vitb_rn50_384",
-            non_negative=True,
-        )
-        net_w, net_h = 384, 384
-        resize_mode = "minimal"
-        normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-    elif model_type == "midas_v21":
-        model = MidasNet(model_path, non_negative=True)
-        net_w, net_h = 384, 384
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(
-            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-        )
-
-    elif model_type == "midas_v21_small":
-        model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True,
-                               non_negative=True, blocks={'expand': True})
-        net_w, net_h = 256, 256
-        resize_mode = "upper_bound"
-        normalization = NormalizeImage(
-            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
-        )
-
-    else:
-        print(f"model_type '{model_type}' not implemented, use: --model_type large")
-        assert False
-
-    transform = Compose(
-        [
-            Resize(
-                net_w,
-                net_h,
-                resize_target=None,
-                keep_aspect_ratio=True,
-                ensure_multiple_of=32,
-                resize_method=resize_mode,
-                image_interpolation_method=cv2.INTER_CUBIC,
-            ),
-            normalization,
-            PrepareForNet(),
-        ]
-    )
-
-    return model.eval(), transform
-
-
-class MiDaSInference(nn.Module):
-    MODEL_TYPES_TORCH_HUB = [
-        "DPT_Large",
-        "DPT_Hybrid",
-        "MiDaS_small"
-    ]
-    MODEL_TYPES_ISL = [
-        "dpt_large",
-        "dpt_hybrid",
-        "midas_v21",
-        "midas_v21_small",
-    ]
-
-    def __init__(self, model_type, model_path):
-        super().__init__()
-        assert (model_type in self.MODEL_TYPES_ISL)
-        model, _ = load_model(model_type, model_path)
-        self.model = model
-        self.model.train = disabled_train
-
-    def forward(self, x):
-        # x in 0..1 as produced by calling self.transform on a 0..1 float64 numpy array
-        # NOTE: we expect that the correct transform has been called during dataloading.
-        with torch.no_grad():
-            prediction = self.model(x)
-            prediction = torch.nn.functional.interpolate(
-                prediction.unsqueeze(1),
-                size=x.shape[2:],
-                mode="bicubic",
-                align_corners=False,
-            )
-        assert prediction.shape == (x.shape[0], 1, x.shape[2], x.shape[3])
-        return prediction
-

extralibs/midas/midas/base_model.py

@@ -1,16 +0,0 @@
-import torch
-
-
-class BaseModel(torch.nn.Module):
-    def load(self, path):
-        """Load model from file.
-
-        Args:
-            path (str): file path
-        """
-        parameters = torch.load(path, map_location=torch.device('cpu'))
-
-        if "optimizer" in parameters:
-            parameters = parameters["model"]
-
-        self.load_state_dict(parameters)

extralibs/midas/midas/blocks.py

@@ -1,342 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .vit import (
-    _make_pretrained_vitb_rn50_384,
-    _make_pretrained_vitl16_384,
-    _make_pretrained_vitb16_384,
-    forward_vit,
-)
-
-def _make_encoder(backbone, features, use_pretrained, groups=1, expand=False, exportable=True, hooks=None, use_vit_only=False, use_readout="ignore",):
-    if backbone == "vitl16_384":
-        pretrained = _make_pretrained_vitl16_384(
-            use_pretrained, hooks=hooks, use_readout=use_readout
-        )
-        scratch = _make_scratch(
-            [256, 512, 1024, 1024], features, groups=groups, expand=expand
-        )  # ViT-L/16 - 85.0% Top1 (backbone)
-    elif backbone == "vitb_rn50_384":
-        pretrained = _make_pretrained_vitb_rn50_384(
-            use_pretrained,
-            hooks=hooks,
-            use_vit_only=use_vit_only,
-            use_readout=use_readout,
-        )
-        scratch = _make_scratch(
-            [256, 512, 768, 768], features, groups=groups, expand=expand
-        )  # ViT-H/16 - 85.0% Top1 (backbone)
-    elif backbone == "vitb16_384":
-        pretrained = _make_pretrained_vitb16_384(
-            use_pretrained, hooks=hooks, use_readout=use_readout
-        )
-        scratch = _make_scratch(
-            [96, 192, 384, 768], features, groups=groups, expand=expand
-        )  # ViT-B/16 - 84.6% Top1 (backbone)
-    elif backbone == "resnext101_wsl":
-        pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
-        scratch = _make_scratch([256, 512, 1024, 2048], features, groups=groups, expand=expand)     # efficientnet_lite3  
-    elif backbone == "efficientnet_lite3":
-        pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, exportable=exportable)
-        scratch = _make_scratch([32, 48, 136, 384], features, groups=groups, expand=expand)  # efficientnet_lite3     
-    else:
-        print(f"Backbone '{backbone}' not implemented")
-        assert False
-        
-    return pretrained, scratch
-
-
-def _make_scratch(in_shape, out_shape, groups=1, expand=False):
-    scratch = nn.Module()
-
-    out_shape1 = out_shape
-    out_shape2 = out_shape
-    out_shape3 = out_shape
-    out_shape4 = out_shape
-    if expand==True:
-        out_shape1 = out_shape
-        out_shape2 = out_shape*2
-        out_shape3 = out_shape*4
-        out_shape4 = out_shape*8
-
-    scratch.layer1_rn = nn.Conv2d(
-        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-    scratch.layer2_rn = nn.Conv2d(
-        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-    scratch.layer3_rn = nn.Conv2d(
-        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-    scratch.layer4_rn = nn.Conv2d(
-        in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
-    )
-
-    return scratch
-
-
-def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
-    efficientnet = torch.hub.load(
-        "rwightman/gen-efficientnet-pytorch",
-        "tf_efficientnet_lite3",
-        pretrained=use_pretrained,
-        exportable=exportable
-    )
-    return _make_efficientnet_backbone(efficientnet)
-
-
-def _make_efficientnet_backbone(effnet):
-    pretrained = nn.Module()
-
-    pretrained.layer1 = nn.Sequential(
-        effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
-    )
-    pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
-    pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
-    pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
-
-    return pretrained
-    
-
-def _make_resnet_backbone(resnet):
-    pretrained = nn.Module()
-    pretrained.layer1 = nn.Sequential(
-        resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
-    )
-
-    pretrained.layer2 = resnet.layer2
-    pretrained.layer3 = resnet.layer3
-    pretrained.layer4 = resnet.layer4
-
-    return pretrained
-
-
-def _make_pretrained_resnext101_wsl(use_pretrained):
-    resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
-    return _make_resnet_backbone(resnet)
-
-
-
-class Interpolate(nn.Module):
-    """Interpolation module.
-    """
-
-    def __init__(self, scale_factor, mode, align_corners=False):
-        """Init.
-
-        Args:
-            scale_factor (float): scaling
-            mode (str): interpolation mode
-        """
-        super(Interpolate, self).__init__()
-
-        self.interp = nn.functional.interpolate
-        self.scale_factor = scale_factor
-        self.mode = mode
-        self.align_corners = align_corners
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: interpolated data
-        """
-
-        x = self.interp(
-            x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners
-        )
-
-        return x
-
-
-class ResidualConvUnit(nn.Module):
-    """Residual convolution module.
-    """
-
-    def __init__(self, features):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super().__init__()
-
-        self.conv1 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True
-        )
-
-        self.conv2 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True
-        )
-
-        self.relu = nn.ReLU(inplace=True)
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: output
-        """
-        out = self.relu(x)
-        out = self.conv1(out)
-        out = self.relu(out)
-        out = self.conv2(out)
-
-        return out + x
-
-
-class FeatureFusionBlock(nn.Module):
-    """Feature fusion block.
-    """
-
-    def __init__(self, features):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super(FeatureFusionBlock, self).__init__()
-
-        self.resConfUnit1 = ResidualConvUnit(features)
-        self.resConfUnit2 = ResidualConvUnit(features)
-
-    def forward(self, *xs):
-        """Forward pass.
-
-        Returns:
-            tensor: output
-        """
-        output = xs[0]
-
-        if len(xs) == 2:
-            output += self.resConfUnit1(xs[1])
-
-        output = self.resConfUnit2(output)
-
-        output = nn.functional.interpolate(
-            output, scale_factor=2, mode="bilinear", align_corners=True
-        )
-
-        return output
-
-
-
-
-class ResidualConvUnit_custom(nn.Module):
-    """Residual convolution module.
-    """
-
-    def __init__(self, features, activation, bn):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super().__init__()
-
-        self.bn = bn
-
-        self.groups=1
-
-        self.conv1 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
-        )
-        
-        self.conv2 = nn.Conv2d(
-            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
-        )
-
-        if self.bn==True:
-            self.bn1 = nn.BatchNorm2d(features)
-            self.bn2 = nn.BatchNorm2d(features)
-
-        self.activation = activation
-
-        self.skip_add = nn.quantized.FloatFunctional()
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input
-
-        Returns:
-            tensor: output
-        """
-        
-        out = self.activation(x)
-        out = self.conv1(out)
-        if self.bn==True:
-            out = self.bn1(out)
-       
-        out = self.activation(out)
-        out = self.conv2(out)
-        if self.bn==True:
-            out = self.bn2(out)
-
-        if self.groups > 1:
-            out = self.conv_merge(out)
-
-        return self.skip_add.add(out, x)
-
-        # return out + x
-
-
-class FeatureFusionBlock_custom(nn.Module):
-    """Feature fusion block.
-    """
-
-    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True):
-        """Init.
-
-        Args:
-            features (int): number of features
-        """
-        super(FeatureFusionBlock_custom, self).__init__()
-
-        self.deconv = deconv
-        self.align_corners = align_corners
-
-        self.groups=1
-
-        self.expand = expand
-        out_features = features
-        if self.expand==True:
-            out_features = features//2
-        
-        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
-
-        self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
-        self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
-        
-        self.skip_add = nn.quantized.FloatFunctional()
-
-    def forward(self, *xs):
-        """Forward pass.
-
-        Returns:
-            tensor: output
-        """
-        output = xs[0]
-
-        if len(xs) == 2:
-            res = self.resConfUnit1(xs[1])
-            output = self.skip_add.add(output, res)
-            # output += res
-
-        output = self.resConfUnit2(output)
-
-        output = nn.functional.interpolate(
-            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
-        )
-
-        output = self.out_conv(output)
-
-        return output
-

extralibs/midas/midas/dpt_depth.py

@@ -1,110 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .base_model import BaseModel
-from .blocks import (
-    FeatureFusionBlock,
-    FeatureFusionBlock_custom,
-    Interpolate,
-    _make_encoder,
-    forward_vit,
-)
-
-
-def _make_fusion_block(features, use_bn):
-    return FeatureFusionBlock_custom(
-        features,
-        nn.ReLU(False),
-        deconv=False,
-        bn=use_bn,
-        expand=False,
-        align_corners=True,
-    )
-
-
-class DPT(BaseModel):
-    def __init__(
-        self,
-        head,
-        features=256,
-        backbone="vitb_rn50_384",
-        readout="project",
-        channels_last=False,
-        use_bn=False,
-    ):
-
-        super(DPT, self).__init__()
-
-        self.channels_last = channels_last
-
-        hooks = {
-            "vitb_rn50_384": [0, 1, 8, 11],
-            "vitb16_384": [2, 5, 8, 11],
-            "vitl16_384": [5, 11, 17, 23],
-        }
-
-        # Instantiate backbone and reassemble blocks
-        self.pretrained, self.scratch = _make_encoder(
-            backbone,
-            features,
-            False, # Set to true of you want to train from scratch, uses ImageNet weights
-            groups=1,
-            expand=False,
-            exportable=False,
-            hooks=hooks[backbone],
-            use_readout=readout,
-        )
-
-        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
-        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
-        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
-        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)
-
-        self.scratch.output_conv = head
-
-
-    def forward(self, x):
-        if self.channels_last == True:
-            x.contiguous(memory_format=torch.channels_last)
-
-        layer_1, layer_2, layer_3, layer_4 = forward_vit(self.pretrained, x)
-
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-        path_4 = self.scratch.refinenet4(layer_4_rn)
-        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
-        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
-        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-
-        out = self.scratch.output_conv(path_1)
-
-        return out
-
-
-class DPTDepthModel(DPT):
-    def __init__(self, path=None, non_negative=True, **kwargs):
-        features = kwargs["features"] if "features" in kwargs else 256
-
-        head = nn.Sequential(
-            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
-            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
-            nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(True),
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-            nn.ReLU(True) if non_negative else nn.Identity(),
-            nn.Identity(),
-        )
-
-        super().__init__(head, **kwargs)
-
-        if path is not None:
-           self.load(path)
-           print("Midas depth estimation model loaded.")
-
-    def forward(self, x):
-        return super().forward(x).squeeze(dim=1)
-

extralibs/midas/midas/midas_net.py

@@ -1,76 +0,0 @@
-"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
-This file contains code that is adapted from
-https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
-"""
-import torch
-import torch.nn as nn
-
-from .base_model import BaseModel
-from .blocks import FeatureFusionBlock, Interpolate, _make_encoder
-
-
-class MidasNet(BaseModel):
-    """Network for monocular depth estimation.
-    """
-
-    def __init__(self, path=None, features=256, non_negative=True):
-        """Init.
-
-        Args:
-            path (str, optional): Path to saved model. Defaults to None.
-            features (int, optional): Number of features. Defaults to 256.
-            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
-        """
-        print("Loading weights: ", path)
-
-        super(MidasNet, self).__init__()
-
-        use_pretrained = False if path is None else True
-
-        self.pretrained, self.scratch = _make_encoder(backbone="resnext101_wsl", features=features, use_pretrained=use_pretrained)
-
-        self.scratch.refinenet4 = FeatureFusionBlock(features)
-        self.scratch.refinenet3 = FeatureFusionBlock(features)
-        self.scratch.refinenet2 = FeatureFusionBlock(features)
-        self.scratch.refinenet1 = FeatureFusionBlock(features)
-
-        self.scratch.output_conv = nn.Sequential(
-            nn.Conv2d(features, 128, kernel_size=3, stride=1, padding=1),
-            Interpolate(scale_factor=2, mode="bilinear"),
-            nn.Conv2d(128, 32, kernel_size=3, stride=1, padding=1),
-            nn.ReLU(True),
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-            nn.ReLU(True) if non_negative else nn.Identity(),
-        )
-
-        if path:
-            self.load(path)
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input data (image)
-
-        Returns:
-            tensor: depth
-        """
-
-        layer_1 = self.pretrained.layer1(x)
-        layer_2 = self.pretrained.layer2(layer_1)
-        layer_3 = self.pretrained.layer3(layer_2)
-        layer_4 = self.pretrained.layer4(layer_3)
-
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-        path_4 = self.scratch.refinenet4(layer_4_rn)
-        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
-        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
-        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-
-        out = self.scratch.output_conv(path_1)
-
-        return torch.squeeze(out, dim=1)

extralibs/midas/midas/midas_net_custom.py

@@ -1,128 +0,0 @@
-"""MidashNet: Network for monocular depth estimation trained by mixing several datasets.
-This file contains code that is adapted from
-https://github.com/thomasjpfan/pytorch_refinenet/blob/master/pytorch_refinenet/refinenet/refinenet_4cascade.py
-"""
-import torch
-import torch.nn as nn
-
-from .base_model import BaseModel
-from .blocks import FeatureFusionBlock, FeatureFusionBlock_custom, Interpolate, _make_encoder
-
-
-class MidasNet_small(BaseModel):
-    """Network for monocular depth estimation.
-    """
-
-    def __init__(self, path=None, features=64, backbone="efficientnet_lite3", non_negative=True, exportable=True, channels_last=False, align_corners=True,
-        blocks={'expand': True}):
-        """Init.
-
-        Args:
-            path (str, optional): Path to saved model. Defaults to None.
-            features (int, optional): Number of features. Defaults to 256.
-            backbone (str, optional): Backbone network for encoder. Defaults to resnet50
-        """
-        print("Loading weights: ", path)
-
-        super(MidasNet_small, self).__init__()
-
-        use_pretrained = False if path else True
-                
-        self.channels_last = channels_last
-        self.blocks = blocks
-        self.backbone = backbone
-
-        self.groups = 1
-
-        features1=features
-        features2=features
-        features3=features
-        features4=features
-        self.expand = False
-        if "expand" in self.blocks and self.blocks['expand'] == True:
-            self.expand = True
-            features1=features
-            features2=features*2
-            features3=features*4
-            features4=features*8
-
-        self.pretrained, self.scratch = _make_encoder(self.backbone, features, use_pretrained, groups=self.groups, expand=self.expand, exportable=exportable)
-  
-        self.scratch.activation = nn.ReLU(False)    
-
-        self.scratch.refinenet4 = FeatureFusionBlock_custom(features4, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
-        self.scratch.refinenet3 = FeatureFusionBlock_custom(features3, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
-        self.scratch.refinenet2 = FeatureFusionBlock_custom(features2, self.scratch.activation, deconv=False, bn=False, expand=self.expand, align_corners=align_corners)
-        self.scratch.refinenet1 = FeatureFusionBlock_custom(features1, self.scratch.activation, deconv=False, bn=False, align_corners=align_corners)
-
-        
-        self.scratch.output_conv = nn.Sequential(
-            nn.Conv2d(features, features//2, kernel_size=3, stride=1, padding=1, groups=self.groups),
-            Interpolate(scale_factor=2, mode="bilinear"),
-            nn.Conv2d(features//2, 32, kernel_size=3, stride=1, padding=1),
-            self.scratch.activation,
-            nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0),
-            nn.ReLU(True) if non_negative else nn.Identity(),
-            nn.Identity(),
-        )
-        
-        if path:
-            self.load(path)
-
-
-    def forward(self, x):
-        """Forward pass.
-
-        Args:
-            x (tensor): input data (image)
-
-        Returns:
-            tensor: depth
-        """
-        if self.channels_last==True:
-            print("self.channels_last = ", self.channels_last)
-            x.contiguous(memory_format=torch.channels_last)
-
-
-        layer_1 = self.pretrained.layer1(x)
-        layer_2 = self.pretrained.layer2(layer_1)
-        layer_3 = self.pretrained.layer3(layer_2)
-        layer_4 = self.pretrained.layer4(layer_3)
-        
-        layer_1_rn = self.scratch.layer1_rn(layer_1)
-        layer_2_rn = self.scratch.layer2_rn(layer_2)
-        layer_3_rn = self.scratch.layer3_rn(layer_3)
-        layer_4_rn = self.scratch.layer4_rn(layer_4)
-
-
-        path_4 = self.scratch.refinenet4(layer_4_rn)
-        path_3 = self.scratch.refinenet3(path_4, layer_3_rn)
-        path_2 = self.scratch.refinenet2(path_3, layer_2_rn)
-        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
-        
-        out = self.scratch.output_conv(path_1)
-
-        return torch.squeeze(out, dim=1)
-
-
-
-def fuse_model(m):
-    prev_previous_type = nn.Identity()
-    prev_previous_name = ''
-    previous_type = nn.Identity()
-    previous_name = ''
-    for name, module in m.named_modules():
-        if prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d and type(module) == nn.ReLU:
-            # print("FUSED ", prev_previous_name, previous_name, name)
-            torch.quantization.fuse_modules(m, [prev_previous_name, previous_name, name], inplace=True)
-        elif prev_previous_type == nn.Conv2d and previous_type == nn.BatchNorm2d:
-            # print("FUSED ", prev_previous_name, previous_name)
-            torch.quantization.fuse_modules(m, [prev_previous_name, previous_name], inplace=True)
-        # elif previous_type == nn.Conv2d and type(module) == nn.ReLU:
-        #    print("FUSED ", previous_name, name)
-        #    torch.quantization.fuse_modules(m, [previous_name, name], inplace=True)
-
-        prev_previous_type = previous_type
-        prev_previous_name = previous_name
-        previous_type = type(module)
-        previous_name = name

extralibs/midas/midas/transforms.py

@@ -1,234 +0,0 @@
-import numpy as np
-import cv2
-import math
-
-
-def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
-    """Rezise the sample to ensure the given size. Keeps aspect ratio.
-
-    Args:
-        sample (dict): sample
-        size (tuple): image size
-
-    Returns:
-        tuple: new size
-    """
-    shape = list(sample["disparity"].shape)
-
-    if shape[0] >= size[0] and shape[1] >= size[1]:
-        return sample
-
-    scale = [0, 0]
-    scale[0] = size[0] / shape[0]
-    scale[1] = size[1] / shape[1]
-
-    scale = max(scale)
-
-    shape[0] = math.ceil(scale * shape[0])
-    shape[1] = math.ceil(scale * shape[1])
-
-    # resize
-    sample["image"] = cv2.resize(
-        sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method
-    )
-
-    sample["disparity"] = cv2.resize(
-        sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST
-    )
-    sample["mask"] = cv2.resize(
-        sample["mask"].astype(np.float32),
-        tuple(shape[::-1]),
-        interpolation=cv2.INTER_NEAREST,
-    )
-    sample["mask"] = sample["mask"].astype(bool)
-
-    return tuple(shape)
-
-
-class Resize(object):
-    """Resize sample to given size (width, height).
-    """
-
-    def __init__(
-        self,
-        width,
-        height,
-        resize_target=True,
-        keep_aspect_ratio=False,
-        ensure_multiple_of=1,
-        resize_method="lower_bound",
-        image_interpolation_method=cv2.INTER_AREA,
-    ):
-        """Init.
-
-        Args:
-            width (int): desired output width
-            height (int): desired output height
-            resize_target (bool, optional):
-                True: Resize the full sample (image, mask, target).
-                False: Resize image only.
-                Defaults to True.
-            keep_aspect_ratio (bool, optional):
-                True: Keep the aspect ratio of the input sample.
-                Output sample might not have the given width and height, and
-                resize behaviour depends on the parameter 'resize_method'.
-                Defaults to False.
-            ensure_multiple_of (int, optional):
-                Output width and height is constrained to be multiple of this parameter.
-                Defaults to 1.
-            resize_method (str, optional):
-                "lower_bound": Output will be at least as large as the given size.
-                "upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
-                "minimal": Scale as least as possible.  (Output size might be smaller than given size.)
-                Defaults to "lower_bound".
-        """
-        self.__width = width
-        self.__height = height
-
-        self.__resize_target = resize_target
-        self.__keep_aspect_ratio = keep_aspect_ratio
-        self.__multiple_of = ensure_multiple_of
-        self.__resize_method = resize_method
-        self.__image_interpolation_method = image_interpolation_method
-
-    def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
-        y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
-        if max_val is not None and y > max_val:
-            y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
-        if y < min_val:
-            y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
-
-        return y
-
-    def get_size(self, width, height):
-        # determine new height and width
-        scale_height = self.__height / height
-        scale_width = self.__width / width
-
-        if self.__keep_aspect_ratio:
-            if self.__resize_method == "lower_bound":
-                # scale such that output size is lower bound
-                if scale_width > scale_height:
-                    # fit width
-                    scale_height = scale_width
-                else:
-                    # fit height
-                    scale_width = scale_height
-            elif self.__resize_method == "upper_bound":
-                # scale such that output size is upper bound
-                if scale_width < scale_height:
-                    # fit width
-                    scale_height = scale_width
-                else:
-                    # fit height
-                    scale_width = scale_height
-            elif self.__resize_method == "minimal":
-                # scale as least as possbile
-                if abs(1 - scale_width) < abs(1 - scale_height):
-                    # fit width
-                    scale_height = scale_width
-                else:
-                    # fit height
-                    scale_width = scale_height
-            else:
-                raise ValueError(
-                    f"resize_method {self.__resize_method} not implemented"
-                )
-
-        if self.__resize_method == "lower_bound":
-            new_height = self.constrain_to_multiple_of(
-                scale_height * height, min_val=self.__height
-            )
-            new_width = self.constrain_to_multiple_of(
-                scale_width * width, min_val=self.__width
-            )
-        elif self.__resize_method == "upper_bound":
-            new_height = self.constrain_to_multiple_of(
-                scale_height * height, max_val=self.__height
-            )
-            new_width = self.constrain_to_multiple_of(
-                scale_width * width, max_val=self.__width
-            )
-        elif self.__resize_method == "minimal":
-            new_height = self.constrain_to_multiple_of(scale_height * height)
-            new_width = self.constrain_to_multiple_of(scale_width * width)
-        else:
-            raise ValueError(f"resize_method {self.__resize_method} not implemented")
-
-        return (new_width, new_height)
-
-    def __call__(self, sample):
-        width, height = self.get_size(
-            sample["image"].shape[1], sample["image"].shape[0]
-        )
-
-        # resize sample
-        sample["image"] = cv2.resize(
-            sample["image"],
-            (width, height),
-            interpolation=self.__image_interpolation_method,
-        )
-
-        if self.__resize_target:
-            if "disparity" in sample:
-                sample["disparity"] = cv2.resize(
-                    sample["disparity"],
-                    (width, height),
-                    interpolation=cv2.INTER_NEAREST,
-                )
-
-            if "depth" in sample:
-                sample["depth"] = cv2.resize(
-                    sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST
-                )
-
-            sample["mask"] = cv2.resize(
-                sample["mask"].astype(np.float32),
-                (width, height),
-                interpolation=cv2.INTER_NEAREST,
-            )
-            sample["mask"] = sample["mask"].astype(bool)
-
-        return sample
-
-
-class NormalizeImage(object):
-    """Normlize image by given mean and std.
-    """
-
-    def __init__(self, mean, std):
-        self.__mean = mean
-        self.__std = std
-
-    def __call__(self, sample):
-        sample["image"] = (sample["image"] - self.__mean) / self.__std
-
-        return sample
-
-
-class PrepareForNet(object):
-    """Prepare sample for usage as network input.
-    """
-
-    def __init__(self):
-        pass
-
-    def __call__(self, sample):
-        image = np.transpose(sample["image"], (2, 0, 1))
-        sample["image"] = np.ascontiguousarray(image).astype(np.float32)
-
-        if "mask" in sample:
-            sample["mask"] = sample["mask"].astype(np.float32)
-            sample["mask"] = np.ascontiguousarray(sample["mask"])
-
-        if "disparity" in sample:
-            disparity = sample["disparity"].astype(np.float32)
-            sample["disparity"] = np.ascontiguousarray(disparity)
-
-        if "depth" in sample:
-            depth = sample["depth"].astype(np.float32)
-            sample["depth"] = np.ascontiguousarray(depth)
-
-        return sample

extralibs/midas/midas/vit.py

@@ -1,489 +0,0 @@
-import torch
-import torch.nn as nn
-import timm
-import types
-import math
-import torch.nn.functional as F
-
-
-class Slice(nn.Module):
-    def __init__(self, start_index=1):
-        super(Slice, self).__init__()
-        self.start_index = start_index
-
-    def forward(self, x):
-        return x[:, self.start_index :]
-
-
-class AddReadout(nn.Module):
-    def __init__(self, start_index=1):
-        super(AddReadout, self).__init__()
-        self.start_index = start_index
-
-    def forward(self, x):
-        if self.start_index == 2:
-            readout = (x[:, 0] + x[:, 1]) / 2
-        else:
-            readout = x[:, 0]
-        return x[:, self.start_index :] + readout.unsqueeze(1)
-
-
-class ProjectReadout(nn.Module):
-    def __init__(self, in_features, start_index=1):
-        super(ProjectReadout, self).__init__()
-        self.start_index = start_index
-
-        self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU())
-
-    def forward(self, x):
-        readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index :])
-        features = torch.cat((x[:, self.start_index :], readout), -1)
-
-        return self.project(features)
-
-
-class Transpose(nn.Module):
-    def __init__(self, dim0, dim1):
-        super(Transpose, self).__init__()
-        self.dim0 = dim0
-        self.dim1 = dim1
-
-    def forward(self, x):
-        x = x.transpose(self.dim0, self.dim1)
-        return x
-
-
-activations = {}
-def forward_vit(pretrained, x):
-    b, c, h, w = x.shape
-
-    glob = pretrained.model.forward_flex(x)
-    pretrained.activations = activations
-
-    layer_1 = pretrained.activations["1"]
-    layer_2 = pretrained.activations["2"]
-    layer_3 = pretrained.activations["3"]
-    layer_4 = pretrained.activations["4"]
-
-    layer_1 = pretrained.act_postprocess1[0:2](layer_1)
-    layer_2 = pretrained.act_postprocess2[0:2](layer_2)
-    layer_3 = pretrained.act_postprocess3[0:2](layer_3)
-    layer_4 = pretrained.act_postprocess4[0:2](layer_4)
-
-    unflatten = nn.Sequential(
-        nn.Unflatten(
-            2,
-            torch.Size(
-                [
-                    h // pretrained.model.patch_size[1],
-                    w // pretrained.model.patch_size[0],
-                ]
-            ),
-        )
-    )
-
-    if layer_1.ndim == 3:
-        layer_1 = unflatten(layer_1)
-    if layer_2.ndim == 3:
-        layer_2 = unflatten(layer_2)
-    if layer_3.ndim == 3:
-        layer_3 = unflatten(layer_3)
-    if layer_4.ndim == 3:
-        layer_4 = unflatten(layer_4)
-
-    layer_1 = pretrained.act_postprocess1[3 : len(pretrained.act_postprocess1)](layer_1)
-    layer_2 = pretrained.act_postprocess2[3 : len(pretrained.act_postprocess2)](layer_2)
-    layer_3 = pretrained.act_postprocess3[3 : len(pretrained.act_postprocess3)](layer_3)
-    layer_4 = pretrained.act_postprocess4[3 : len(pretrained.act_postprocess4)](layer_4)
-
-    return layer_1, layer_2, layer_3, layer_4
-
-
-def _resize_pos_embed(self, posemb, gs_h, gs_w):
-    posemb_tok, posemb_grid = (
-        posemb[:, : self.start_index],
-        posemb[0, self.start_index :],
-    )
-
-    gs_old = int(math.sqrt(len(posemb_grid)))
-
-    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
-    posemb_grid = F.interpolate(posemb_grid, size=(gs_h, gs_w), mode="bilinear")
-    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)
-
-    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
-
-    return posemb
-
-
-def forward_flex(self, x):
-    b, c, h, w = x.shape
-
-    pos_embed = self._resize_pos_embed(
-        self.pos_embed, h // self.patch_size[1], w // self.patch_size[0]
-    )
-
-    B = x.shape[0]
-
-    if hasattr(self.patch_embed, "backbone"):
-        x = self.patch_embed.backbone(x)
-        if isinstance(x, (list, tuple)):
-            x = x[-1]  # last feature if backbone outputs list/tuple of features
-
-    x = self.patch_embed.proj(x).flatten(2).transpose(1, 2)
-
-    if getattr(self, "dist_token", None) is not None:
-        cls_tokens = self.cls_token.expand(
-            B, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        dist_token = self.dist_token.expand(B, -1, -1)
-        x = torch.cat((cls_tokens, dist_token, x), dim=1)
-    else:
-        cls_tokens = self.cls_token.expand(
-            B, -1, -1
-        )  # stole cls_tokens impl from Phil Wang, thanks
-        x = torch.cat((cls_tokens, x), dim=1)
-
-    x = x + pos_embed
-    x = self.pos_drop(x)
-
-    for blk in self.blocks:
-        x = blk(x)
-
-    x = self.norm(x)
-
-    return x
-
-
-def get_activation(name):
-    def hook(model, input, output):
-        activations[name] = output
-    return hook
-
-
-def get_readout_oper(vit_features, features, use_readout, start_index=1):
-    if use_readout == "ignore":
-        readout_oper = [Slice(start_index)] * len(features)
-    elif use_readout == "add":
-        readout_oper = [AddReadout(start_index)] * len(features)
-    elif use_readout == "project":
-        readout_oper = [
-            ProjectReadout(vit_features, start_index) for out_feat in features
-        ]
-    else:
-        assert (
-            False
-        ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'"
-
-    return readout_oper
-
-
-def _make_vit_b16_backbone(
-    model,
-    features=[96, 192, 384, 768],
-    size=[384, 384],
-    hooks=[2, 5, 8, 11],
-    vit_features=768,
-    use_readout="ignore",
-    start_index=1,
-):
-    pretrained = nn.Module()
-
-    pretrained.model = model
-    pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
-    pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
-    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
-    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
-
-    pretrained.activations = activations
-
-    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
-
-    # 32, 48, 136, 384
-    pretrained.act_postprocess1 = nn.Sequential(
-        readout_oper[0],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[0],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.ConvTranspose2d(
-            in_channels=features[0],
-            out_channels=features[0],
-            kernel_size=4,
-            stride=4,
-            padding=0,
-            bias=True,
-            dilation=1,
-            groups=1,
-        ),
-    )
-
-    pretrained.act_postprocess2 = nn.Sequential(
-        readout_oper[1],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[1],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.ConvTranspose2d(
-            in_channels=features[1],
-            out_channels=features[1],
-            kernel_size=2,
-            stride=2,
-            padding=0,
-            bias=True,
-            dilation=1,
-            groups=1,
-        ),
-    )
-
-    pretrained.act_postprocess3 = nn.Sequential(
-        readout_oper[2],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[2],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-    )
-
-    pretrained.act_postprocess4 = nn.Sequential(
-        readout_oper[3],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[3],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.Conv2d(
-            in_channels=features[3],
-            out_channels=features[3],
-            kernel_size=3,
-            stride=2,
-            padding=1,
-        ),
-    )
-
-    pretrained.model.start_index = start_index
-    pretrained.model.patch_size = [16, 16]
-
-    # We inject this function into the VisionTransformer instances so that
-    # we can use it with interpolated position embeddings without modifying the library source.
-    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
-    pretrained.model._resize_pos_embed = types.MethodType(
-        _resize_pos_embed, pretrained.model
-    )
-
-    return pretrained
-
-
-def _make_pretrained_vitl16_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model("vit_large_patch16_384", pretrained=pretrained)
-
-    hooks = [5, 11, 17, 23] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model,
-        features=[256, 512, 1024, 1024],
-        hooks=hooks,
-        vit_features=1024,
-        use_readout=use_readout,
-    )
-
-
-def _make_pretrained_vitb16_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model("vit_base_patch16_384", pretrained=pretrained)
-
-    hooks = [2, 5, 8, 11] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model, features=[96, 192, 384, 768], hooks=hooks, use_readout=use_readout
-    )
-
-
-def _make_pretrained_deitb16_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model("vit_deit_base_patch16_384", pretrained=pretrained)
-
-    hooks = [2, 5, 8, 11] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model, features=[96, 192, 384, 768], hooks=hooks, use_readout=use_readout
-    )
-
-
-def _make_pretrained_deitb16_distil_384(pretrained, use_readout="ignore", hooks=None):
-    model = timm.create_model(
-        "vit_deit_base_distilled_patch16_384", pretrained=pretrained
-    )
-
-    hooks = [2, 5, 8, 11] if hooks == None else hooks
-    return _make_vit_b16_backbone(
-        model,
-        features=[96, 192, 384, 768],
-        hooks=hooks,
-        use_readout=use_readout,
-        start_index=2,
-    )
-
-
-def _make_vit_b_rn50_backbone(
-    model,
-    features=[256, 512, 768, 768],
-    size=[384, 384],
-    hooks=[0, 1, 8, 11],
-    vit_features=768,
-    use_vit_only=False,
-    use_readout="ignore",
-    start_index=1,
-):
-    pretrained = nn.Module()
-
-    pretrained.model = model
-    
-    if use_vit_only == True:
-        pretrained.model.blocks[hooks[0]].register_forward_hook(get_activation("1"))
-        pretrained.model.blocks[hooks[1]].register_forward_hook(get_activation("2"))
-    else:
-        pretrained.model.patch_embed.backbone.stages[0].register_forward_hook(
-            get_activation("1")
-        )
-        pretrained.model.patch_embed.backbone.stages[1].register_forward_hook(
-            get_activation("2")
-        )
-
-    pretrained.model.blocks[hooks[2]].register_forward_hook(get_activation("3"))
-    pretrained.model.blocks[hooks[3]].register_forward_hook(get_activation("4"))
-
-    pretrained.activations = activations
-
-    readout_oper = get_readout_oper(vit_features, features, use_readout, start_index)
-
-    if use_vit_only == True:
-        pretrained.act_postprocess1 = nn.Sequential(
-            readout_oper[0],
-            Transpose(1, 2),
-            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-            nn.Conv2d(
-                in_channels=vit_features,
-                out_channels=features[0],
-                kernel_size=1,
-                stride=1,
-                padding=0,
-            ),
-            nn.ConvTranspose2d(
-                in_channels=features[0],
-                out_channels=features[0],
-                kernel_size=4,
-                stride=4,
-                padding=0,
-                bias=True,
-                dilation=1,
-                groups=1,
-            ),
-        )
-
-        pretrained.act_postprocess2 = nn.Sequential(
-            readout_oper[1],
-            Transpose(1, 2),
-            nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-            nn.Conv2d(
-                in_channels=vit_features,
-                out_channels=features[1],
-                kernel_size=1,
-                stride=1,
-                padding=0,
-            ),
-            nn.ConvTranspose2d(
-                in_channels=features[1],
-                out_channels=features[1],
-                kernel_size=2,
-                stride=2,
-                padding=0,
-                bias=True,
-                dilation=1,
-                groups=1,
-            ),
-        )
-    else:
-        pretrained.act_postprocess1 = nn.Sequential(
-            nn.Identity(), nn.Identity(), nn.Identity()
-        )
-        pretrained.act_postprocess2 = nn.Sequential(
-            nn.Identity(), nn.Identity(), nn.Identity()
-        )
-
-    pretrained.act_postprocess3 = nn.Sequential(
-        readout_oper[2],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[2],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-    )
-
-    pretrained.act_postprocess4 = nn.Sequential(
-        readout_oper[3],
-        Transpose(1, 2),
-        nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])),
-        nn.Conv2d(
-            in_channels=vit_features,
-            out_channels=features[3],
-            kernel_size=1,
-            stride=1,
-            padding=0,
-        ),
-        nn.Conv2d(
-            in_channels=features[3],
-            out_channels=features[3],
-            kernel_size=3,
-            stride=2,
-            padding=1,
-        ),
-    )
-
-    pretrained.model.start_index = start_index
-    pretrained.model.patch_size = [16, 16]
-
-    # We inject this function into the VisionTransformer instances so that
-    # we can use it with interpolated position embeddings without modifying the library source.
-    pretrained.model.forward_flex = types.MethodType(forward_flex, pretrained.model)
-
-    # We inject this function into the VisionTransformer instances so that
-    # we can use it with interpolated position embeddings without modifying the library source.
-    pretrained.model._resize_pos_embed = types.MethodType(
-        _resize_pos_embed, pretrained.model
-    )
-
-    return pretrained
-
-
-def _make_pretrained_vitb_rn50_384(
-    pretrained, use_readout="ignore", hooks=None, use_vit_only=False
-):
-    model = timm.create_model("vit_base_resnet50_384", pretrained=pretrained)
-
-    hooks = [0, 1, 8, 11] if hooks == None else hooks
-    return _make_vit_b_rn50_backbone(
-        model,
-        features=[256, 512, 768, 768],
-        size=[384, 384],
-        hooks=hooks,
-        use_vit_only=use_vit_only,
-        use_readout=use_readout,
-    )

extralibs/midas/utils.py

@@ -1,189 +0,0 @@
-"""Utils for monoDepth."""
-import sys
-import re
-import numpy as np
-import cv2
-import torch
-
-
-def read_pfm(path):
-    """Read pfm file.
-
-    Args:
-        path (str): path to file
-
-    Returns:
-        tuple: (data, scale)
-    """
-    with open(path, "rb") as file:
-
-        color = None
-        width = None
-        height = None
-        scale = None
-        endian = None
-
-        header = file.readline().rstrip()
-        if header.decode("ascii") == "PF":
-            color = True
-        elif header.decode("ascii") == "Pf":
-            color = False
-        else:
-            raise Exception("Not a PFM file: " + path)
-
-        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
-        if dim_match:
-            width, height = list(map(int, dim_match.groups()))
-        else:
-            raise Exception("Malformed PFM header.")
-
-        scale = float(file.readline().decode("ascii").rstrip())
-        if scale < 0:
-            # little-endian
-            endian = "<"
-            scale = -scale
-        else:
-            # big-endian
-            endian = ">"
-
-        data = np.fromfile(file, endian + "f")
-        shape = (height, width, 3) if color else (height, width)
-
-        data = np.reshape(data, shape)
-        data = np.flipud(data)
-
-        return data, scale
-
-
-def write_pfm(path, image, scale=1):
-    """Write pfm file.
-
-    Args:
-        path (str): pathto file
-        image (array): data
-        scale (int, optional): Scale. Defaults to 1.
-    """
-
-    with open(path, "wb") as file:
-        color = None
-
-        if image.dtype.name != "float32":
-            raise Exception("Image dtype must be float32.")
-
-        image = np.flipud(image)
-
-        if len(image.shape) == 3 and image.shape[2] == 3:  # color image
-            color = True
-        elif (
-            len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
-        ):  # greyscale
-            color = False
-        else:
-            raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
-
-        file.write("PF\n" if color else "Pf\n".encode())
-        file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
-
-        endian = image.dtype.byteorder
-
-        if endian == "<" or endian == "=" and sys.byteorder == "little":
-            scale = -scale
-
-        file.write("%f\n".encode() % scale)
-
-        image.tofile(file)
-
-
-def read_image(path):
-    """Read image and output RGB image (0-1).
-
-    Args:
-        path (str): path to file
-
-    Returns:
-        array: RGB image (0-1)
-    """
-    img = cv2.imread(path)
-
-    if img.ndim == 2:
-        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
-
-    return img
-
-
-def resize_image(img):
-    """Resize image and make it fit for network.
-
-    Args:
-        img (array): image
-
-    Returns:
-        tensor: data ready for network
-    """
-    height_orig = img.shape[0]
-    width_orig = img.shape[1]
-
-    if width_orig > height_orig:
-        scale = width_orig / 384
-    else:
-        scale = height_orig / 384
-
-    height = (np.ceil(height_orig / scale / 32) * 32).astype(int)
-    width = (np.ceil(width_orig / scale / 32) * 32).astype(int)
-
-    img_resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
-
-    img_resized = (
-        torch.from_numpy(np.transpose(img_resized, (2, 0, 1))).contiguous().float()
-    )
-    img_resized = img_resized.unsqueeze(0)
-
-    return img_resized
-
-
-def resize_depth(depth, width, height):
-    """Resize depth map and bring to CPU (numpy).
-
-    Args:
-        depth (tensor): depth
-        width (int): image width
-        height (int): image height
-
-    Returns:
-        array: processed depth
-    """
-    depth = torch.squeeze(depth[0, :, :, :]).to("cpu")
-
-    depth_resized = cv2.resize(
-        depth.numpy(), (width, height), interpolation=cv2.INTER_CUBIC
-    )
-
-    return depth_resized
-
-def write_depth(path, depth, bits=1):
-    """Write depth map to pfm and png file.
-
-    Args:
-        path (str): filepath without extension
-        depth (array): depth
-    """
-    write_pfm(path + ".pfm", depth.astype(np.float32))
-
-    depth_min = depth.min()
-    depth_max = depth.max()
-
-    max_val = (2**(8*bits))-1
-
-    if depth_max - depth_min > np.finfo("float").eps:
-        out = max_val * (depth - depth_min) / (depth_max - depth_min)
-    else:
-        out = np.zeros(depth.shape, dtype=depth.type)
-
-    if bits == 1:
-        cv2.imwrite(path + ".png", out.astype("uint8"))
-    elif bits == 2:
-        cv2.imwrite(path + ".png", out.astype("uint16"))
-
-    return

i2v_test.py

@@ -0,0 +1,78 @@
+import os
+from omegaconf import OmegaConf
+
+import torch
+
+from scripts.evaluation.funcs import load_model_checkpoint, load_image_batch, save_videos, batch_ddim_sampling
+from utils.utils import instantiate_from_config
+from huggingface_hub import hf_hub_download
+
+class Image2Video():
+    def __init__(self,result_dir='./tmp/',gpu_num=1) -> None:
+        self.download_model()
+        self.result_dir = result_dir
+        if not os.path.exists(self.result_dir):
+            os.mkdir(self.result_dir)
+        ckpt_path='checkpoints/i2v_512_v1/model.ckpt'
+        config_file='configs/inference_i2v_512_v1.0.yaml'
+        config = OmegaConf.load(config_file)
+        model_config = config.pop("model", OmegaConf.create())
+        model_config['params']['unet_config']['params']['use_checkpoint']=False   
+        model_list = []
+        for gpu_id in range(gpu_num):
+            model = instantiate_from_config(model_config)
+            model = model.cuda(gpu_id)
+            assert os.path.exists(ckpt_path), "Error: checkpoint Not Found!"
+            model = load_model_checkpoint(model, ckpt_path)
+            model.eval()
+            model_list.append(model)
+        self.model_list = model_list
+        self.save_fps = 8
+
+    def get_image(self, image, prompt, steps=50, cfg_scale=12.0, eta=1.0, fps=16):
+        gpu_id=0
+        if steps > 60:
+            steps = 60 
+        model = self.model_list[gpu_id]
+        batch_size=1
+        channels = model.model.diffusion_model.in_channels
+        frames = model.temporal_length
+        h, w = 320 // 8, 512 // 8
+        noise_shape = [batch_size, channels, frames, h, w]
+
+        #prompts = batch_size * [""]
+        text_emb = model.get_learned_conditioning([prompt])
+
+        # cond_images = load_image_batch([image_path])
+        img_tensor = torch.from_numpy(image).permute(2, 0, 1).float()
+        img_tensor = (img_tensor / 255. - 0.5) * 2
+        img_tensor = img_tensor.unsqueeze(0)
+        cond_images = img_tensor.to(model.device)
+        img_emb = model.get_image_embeds(cond_images)
+        imtext_cond = torch.cat([text_emb, img_emb], dim=1)
+        cond = {"c_crossattn": [imtext_cond], "fps": fps}
+        
+        ## inference
+        batch_samples = batch_ddim_sampling(model, cond, noise_shape, n_samples=1, ddim_steps=steps, ddim_eta=eta, cfg_scale=cfg_scale)
+        ## b,samples,c,t,h,w
+        prompt_str = prompt.replace("/", "_slash_") if "/" in prompt else prompt
+        prompt_str = prompt_str.replace(" ", "_") if " " in prompt else prompt_str
+        prompt_str=prompt_str[:30]
+
+        save_videos(batch_samples, self.result_dir, filenames=[prompt_str], fps=self.save_fps)
+        return os.path.join(self.result_dir, f"{prompt_str}.mp4")
+    
+    def download_model(self):
+        REPO_ID = 'VideoCrafter/Image2Video-512-v1.0'
+        filename_list = ['model.ckpt']
+        if not os.path.exists('./checkpoints/i2v_512_v1/'):
+            os.makedirs('./checkpoints/i2v_512_v1/')
+        for filename in filename_list:
+            local_file = os.path.join('./checkpoints/i2v_512_v1/', filename)
+            if not os.path.exists(local_file):
+                hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir='./checkpoints/i2v_512_v1/', local_dir_use_symlinks=False)
+    
+if __name__ == '__main__':
+    i2v = Image2Video()
+    video_path = i2v.get_image('prompts/i2v_prompts/horse.png','horses are walking on the grassland')
+    print('done', video_path)

input/prompts.txt

@@ -1,2 +0,0 @@
-astronaut riding a horse
-Flying through an intense battle between pirate ships in a stormy ocean

lvdm/basics.py

@@ -0,0 +1,100 @@
+# adopted from
+# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+# and
+# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+# and
+# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
+#
+# thanks!
+
+import torch.nn as nn
+from utils.utils import instantiate_from_config
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+def scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+
+def conv_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D convolution module.
+    """
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    """
+    Create a linear module.
+    """
+    return nn.Linear(*args, **kwargs)
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """
+    Create a 1D, 2D, or 3D average pooling module.
+    """
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def nonlinearity(type='silu'):
+    if type == 'silu':
+        return nn.SiLU()
+    elif type == 'leaky_relu':
+        return nn.LeakyReLU()
+
+
+class GroupNormSpecific(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def normalization(channels, num_groups=32):
+    """
+    Make a standard normalization layer.
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    return GroupNormSpecific(num_groups, channels)
+
+
+class HybridConditioner(nn.Module):
+
+    def __init__(self, c_concat_config, c_crossattn_config):
+        super().__init__()
+        self.concat_conditioner = instantiate_from_config(c_concat_config)
+        self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
+
+    def forward(self, c_concat, c_crossattn):
+        c_concat = self.concat_conditioner(c_concat)
+        c_crossattn = self.crossattn_conditioner(c_crossattn)
+        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}

lvdm/common.py

@@ -0,0 +1,95 @@
+import math
+from inspect import isfunction
+import torch
+from torch import nn
+import torch.distributed as dist
+
+
+def gather_data(data, return_np=True):
+    ''' gather data from multiple processes to one list '''
+    data_list = [torch.zeros_like(data) for _ in range(dist.get_world_size())]
+    dist.all_gather(data_list, data)  # gather not supported with NCCL
+    if return_np:
+        data_list = [data.cpu().numpy() for data in data_list]
+    return data_list
+
+def autocast(f):
+    def do_autocast(*args, **kwargs):
+        with torch.cuda.amp.autocast(enabled=True,
+                                     dtype=torch.get_autocast_gpu_dtype(),
+                                     cache_enabled=torch.is_autocast_cache_enabled()):
+            return f(*args, **kwargs)
+    return do_autocast
+
+
+def extract_into_tensor(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+def exists(val):
+    return val is not None
+
+def identity(*args, **kwargs):
+    return nn.Identity()
+
+def uniq(arr):
+    return{el: True for el in arr}.keys()
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+def ismap(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+def isimage(x):
+    if not isinstance(x,torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+def shape_to_str(x):
+    shape_str = "x".join([str(x) for x in x.shape])
+    return shape_str
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+ckpt = torch.utils.checkpoint.checkpoint
+def checkpoint(func, inputs, params, flag):
+    """
+    Evaluate a function without caching intermediate activations, allowing for
+    reduced memory at the expense of extra compute in the backward pass.
+    :param func: the function to evaluate.
+    :param inputs: the argument sequence to pass to `func`.
+    :param params: a sequence of parameters `func` depends on but does not
+                   explicitly take as arguments.
+    :param flag: if False, disable gradient checkpointing.
+    """
+    if flag:
+        return ckpt(func, *inputs)
+    else:
+        return func(*inputs)
+

lvdm/data/webvid.py

@@ -1,188 +0,0 @@
-import os
-import random
-import bisect
-
-import pandas as pd
-
-import omegaconf
-import torch
-from torch.utils.data import Dataset
-from torchvision import transforms
-from decord import VideoReader, cpu
-import torchvision.transforms._transforms_video as transforms_video
-
-class WebVid(Dataset):
-    """
-    WebVid Dataset.
-    Assumes webvid data is structured as follows.
-    Webvid/
-        videos/
-            000001_000050/      ($page_dir)
-                1.mp4           (videoid.mp4)
-                ...
-                5000.mp4
-            ...
-    """
-    def __init__(self,
-                 meta_path,
-                 data_dir,
-                 subsample=None,
-                 video_length=16,
-                 resolution=[256, 512],
-                 frame_stride=1,
-                 spatial_transform=None,
-                 crop_resolution=None,
-                 fps_max=None,
-                 load_raw_resolution=False,
-                 fps_schedule=None,
-                 fs_probs=None,
-                 bs_per_gpu=None,
-                 trigger_word='',
-                 dataname='',
-                 ):
-        self.meta_path = meta_path
-        self.data_dir = data_dir
-        self.subsample = subsample
-        self.video_length = video_length
-        self.resolution = [resolution, resolution] if isinstance(resolution, int) else resolution
-        self.frame_stride = frame_stride
-        self.fps_max = fps_max
-        self.load_raw_resolution = load_raw_resolution
-        self.fs_probs = fs_probs
-        self.trigger_word = trigger_word
-        self.dataname = dataname
-
-        self._load_metadata()
-        if spatial_transform is not None:
-            if spatial_transform == "random_crop":
-                self.spatial_transform = transforms_video.RandomCropVideo(crop_resolution)
-            elif spatial_transform == "resize_center_crop":
-                assert(self.resolution[0] == self.resolution[1])
-                self.spatial_transform = transforms.Compose([
-                    transforms.Resize(resolution),
-                    transforms_video.CenterCropVideo(resolution),
-                    ])
-            else:
-                raise NotImplementedError
-        else:
-            self.spatial_transform = None
-        
-        self.fps_schedule = fps_schedule
-        self.bs_per_gpu = bs_per_gpu
-        if self.fps_schedule is not None:
-            assert(self.bs_per_gpu is not None)
-            self.counter = 0
-            self.stage_idx = 0
-        
-    def _load_metadata(self):
-        metadata = pd.read_csv(self.meta_path)
-        if self.subsample is not None:
-            metadata = metadata.sample(self.subsample, random_state=0)
-        metadata['caption'] = metadata['name']
-        del metadata['name']
-        self.metadata = metadata
-        self.metadata.dropna(inplace=True)
-        # self.metadata['caption'] = self.metadata['caption'].str[:350]
-    
-    def _get_video_path(self, sample):
-        if self.dataname == "loradata":
-            rel_video_fp =  str(sample['videoid']) + '.mp4'
-            full_video_fp = os.path.join(self.data_dir, rel_video_fp)
-        else:
-            rel_video_fp = os.path.join(sample['page_dir'], str(sample['videoid']) + '.mp4')
-            full_video_fp = os.path.join(self.data_dir, 'videos', rel_video_fp)
-        return full_video_fp, rel_video_fp
-    
-    def get_fs_based_on_schedule(self, frame_strides, schedule):
-        assert(len(frame_strides) == len(schedule) + 1) # nstage=len_fps_schedule + 1
-        global_step = self.counter // self.bs_per_gpu # TODO: support resume.
-        stage_idx = bisect.bisect(schedule, global_step)
-        frame_stride = frame_strides[stage_idx]
-        # log stage change
-        if stage_idx != self.stage_idx:
-            print(f'fps stage: {stage_idx} start ... new frame stride = {frame_stride}')
-            self.stage_idx = stage_idx
-        return frame_stride
-    
-    def get_fs_based_on_probs(self, frame_strides, probs):
-        assert(len(frame_strides) == len(probs))
-        return random.choices(frame_strides, weights=probs)[0]
-
-    def get_fs_randomly(self, frame_strides):
-        return random.choice(frame_strides)
-    
-    def __getitem__(self, index):
-        
-        if isinstance(self.frame_stride, list) or isinstance(self.frame_stride, omegaconf.listconfig.ListConfig):
-            if self.fps_schedule is not None:
-                frame_stride = self.get_fs_based_on_schedule(self.frame_stride, self.fps_schedule)
-            elif self.fs_probs is not None:
-                frame_stride = self.get_fs_based_on_probs(self.frame_stride, self.fs_probs)
-            else:
-                frame_stride = self.get_fs_randomly(self.frame_stride)
-        else:
-            frame_stride = self.frame_stride
-        assert(isinstance(frame_stride, int)), type(frame_stride)
-
-        while True:
-            index = index % len(self.metadata)
-            sample = self.metadata.iloc[index]
-            video_path, rel_fp = self._get_video_path(sample)
-            caption = sample['caption']+self.trigger_word
-            
-            # make reader
-            try:
-                if self.load_raw_resolution:
-                    video_reader = VideoReader(video_path, ctx=cpu(0))
-                else:
-                    video_reader = VideoReader(video_path, ctx=cpu(0), width=self.resolution[1], height=self.resolution[0])
-                if len(video_reader) < self.video_length:
-                    print(f"video length ({len(video_reader)}) is smaller than target length({self.video_length})")
-                    index += 1
-                    continue
-                else:
-                    pass
-            except:
-                index += 1
-                print(f"Load video failed! path = {video_path}")
-                continue
-            
-            # sample strided frames
-            all_frames = list(range(0, len(video_reader), frame_stride))
-            if len(all_frames) < self.video_length: # recal a max fs
-                frame_stride = len(video_reader) // self.video_length
-                assert(frame_stride != 0)
-                all_frames = list(range(0, len(video_reader), frame_stride))
-
-            # select a random clip
-            rand_idx = random.randint(0, len(all_frames) - self.video_length)
-            frame_indices = all_frames[rand_idx:rand_idx+self.video_length]
-            try:
-                frames = video_reader.get_batch(frame_indices)
-                break
-            except:
-                print(f"Get frames failed! path = {video_path}")
-                index += 1
-                continue
-
-        assert(frames.shape[0] == self.video_length),f'{len(frames)}, self.video_length={self.video_length}'
-        frames = torch.tensor(frames.asnumpy()).permute(3, 0, 1, 2).float() # [t,h,w,c] -> [c,t,h,w]
-        if self.spatial_transform is not None:
-            frames = self.spatial_transform(frames)
-        if self.resolution is not None:
-            assert(frames.shape[2] == self.resolution[0] and frames.shape[3] == self.resolution[1]), f'frames={frames.shape}, self.resolution={self.resolution}'
-        frames = (frames / 255 - 0.5) * 2
-        
-        fps_ori = video_reader.get_avg_fps()
-        fps_clip = fps_ori // frame_stride
-        if self.fps_max is not None and fps_clip > self.fps_max:
-            fps_clip = self.fps_max
-        
-        data = {'video': frames, 'caption': caption, 'path': video_path, 'fps': fps_clip, 'frame_stride': frame_stride}
-
-        if self.fps_schedule is not None:
-            self.counter += 1
-        return data
-    
-    def __len__(self):
-        return len(self.metadata)

lvdm/{models/modules/distributions.py → distributions.py}

@@ -2,6 +2,25 @@ import torch
 import numpy as np
 
 
+class AbstractDistribution:
+    def sample(self):
+        raise NotImplementedError()
+
+    def mode(self):
+        raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+    def __init__(self, value):
+        self.value = value
+
+    def sample(self):
+        return self.value
+
+    def mode(self):
+        return self.value
+
+
 class DiagonalGaussianDistribution(object):
     def __init__(self, parameters, deterministic=False):
         self.parameters = parameters

lvdm/ema.py

@@ -0,0 +1,76 @@
+import torch
+from torch import nn
+
+
+class LitEma(nn.Module):
+    def __init__(self, model, decay=0.9999, use_num_upates=True):
+        super().__init__()
+        if decay < 0.0 or decay > 1.0:
+            raise ValueError('Decay must be between 0 and 1')
+
+        self.m_name2s_name = {}
+        self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32))
+        self.register_buffer('num_updates', torch.tensor(0,dtype=torch.int) if use_num_upates
+                             else torch.tensor(-1,dtype=torch.int))
+
+        for name, p in model.named_parameters():
+            if p.requires_grad:
+                #remove as '.'-character is not allowed in buffers
+                s_name = name.replace('.','')
+                self.m_name2s_name.update({name:s_name})
+                self.register_buffer(s_name,p.clone().detach().data)
+
+        self.collected_params = []
+
+    def forward(self,model):
+        decay = self.decay
+
+        if self.num_updates >= 0:
+            self.num_updates += 1
+            decay = min(self.decay,(1 + self.num_updates) / (10 + self.num_updates))
+
+        one_minus_decay = 1.0 - decay
+
+        with torch.no_grad():
+            m_param = dict(model.named_parameters())
+            shadow_params = dict(self.named_buffers())
+
+            for key in m_param:
+                if m_param[key].requires_grad:
+                    sname = self.m_name2s_name[key]
+                    shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
+                    shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
+                else:
+                    assert not key in self.m_name2s_name
+
+    def copy_to(self, model):
+        m_param = dict(model.named_parameters())
+        shadow_params = dict(self.named_buffers())
+        for key in m_param:
+            if m_param[key].requires_grad:
+                m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
+            else:
+                assert not key in self.m_name2s_name
+
+    def store(self, parameters):
+        """
+        Save the current parameters for restoring later.
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+            temporarily stored.
+        """
+        self.collected_params = [param.clone() for param in parameters]
+
+    def restore(self, parameters):
+        """
+        Restore the parameters stored with the `store` method.
+        Useful to validate the model with EMA parameters without affecting the
+        original optimization process. Store the parameters before the
+        `copy_to` method. After validation (or model saving), use this to
+        restore the former parameters.
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+            updated with the stored parameters.
+        """
+        for c_param, param in zip(self.collected_params, parameters):
+            param.data.copy_(c_param.data)

lvdm/models/autoencoder.py

@@ -1,12 +1,14 @@
-import torch
-import pytorch_lightning as pl
-import torch.nn.functional as F
 import os
+from contextlib import contextmanager
+import torch
+import numpy as np
 from einops import rearrange
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from lvdm.modules.networks.ae_modules import Encoder, Decoder
+from lvdm.distributions import DiagonalGaussianDistribution
+from utils.utils import instantiate_from_config
 
-from lvdm.models.modules.autoencoder_modules import Encoder, Decoder
-from lvdm.models.modules.distributions import DiagonalGaussianDistribution
-from lvdm.utils.common_utils import instantiate_from_config
 
 class AutoencoderKL(pl.LightningModule):
     def __init__(self,
@@ -69,12 +71,12 @@ class AutoencoderKL(pl.LightningModule):
         if self.test_args.save_input:
             os.makedirs(self.root_inputs, exist_ok=True)
         assert(self.test_args is not None)
-        self.test_maximum = getattr(self.test_args, 'test_maximum', None) #1500 # 12000/8
+        self.test_maximum = getattr(self.test_args, 'test_maximum', None) 
         self.count = 0
         self.eval_metrics = {}
         self.decodes = []
         self.save_decode_samples = 2048
-        
+
     def init_from_ckpt(self, path, ignore_keys=list()):
         sd = torch.load(path, map_location="cpu")
         try:
@@ -115,10 +117,6 @@ class AutoencoderKL(pl.LightningModule):
 
     def get_input(self, batch, k):
         x = batch[k]
-        # if len(x.shape) == 3:
-        #     x = x[..., None]
-        # if x.dim() == 4:
-        #     x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
         if x.dim() == 5 and self.input_dim == 4:
             b,c,t,h,w = x.shape
             self.b = b
@@ -200,3 +198,22 @@ class AutoencoderKL(pl.LightningModule):
         x = F.conv2d(x, weight=self.colorize)
         x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
         return x
+
+class IdentityFirstStage(torch.nn.Module):
+    def __init__(self, *args, vq_interface=False, **kwargs):
+        self.vq_interface = vq_interface  # TODO: Should be true by default but check to not break older stuff
+        super().__init__()
+
+    def encode(self, x, *args, **kwargs):
+        return x
+
+    def decode(self, x, *args, **kwargs):
+        return x
+
+    def quantize(self, x, *args, **kwargs):
+        if self.vq_interface:
+            return x, None, [None, None, None]
+        return x
+
+    def forward(self, x, *args, **kwargs):
+        return x

lvdm/models/ddpm3d.py

@@ -1,54 +1,42 @@
-import os
-import time
-import random
-import itertools
+"""
+wild mixture of
+https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
+https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+https://github.com/CompVis/taming-transformers
+-- merci
+"""
+
 from functools import partial
 from contextlib import contextmanager
-
 import numpy as np
 from tqdm import tqdm
 from einops import rearrange, repeat
-
+import logging
+mainlogger = logging.getLogger('mainlogger')
 import torch
 import torch.nn as nn
-import pytorch_lightning as pl
 from torchvision.utils import make_grid
-from torch.optim.lr_scheduler import LambdaLR
-from pytorch_lightning.utilities import rank_zero_only
-from lvdm.models.modules.distributions import normal_kl, DiagonalGaussianDistribution
-from lvdm.models.modules.util import make_beta_schedule, extract_into_tensor, noise_like
-from lvdm.models.modules.lora import inject_trainable_lora
-from lvdm.samplers.ddim import DDIMSampler
-from lvdm.utils.common_utils import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config, check_istarget
-
-
-def disabled_train(self, mode=True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-def uniform_on_device(r1, r2, shape, device):
-    return (r1 - r2) * torch.rand(*shape, device=device) + r2
-
-
-def split_video_to_clips(video, clip_length, drop_left=True):
-    video_length = video.shape[2]
-    shape = video.shape
-    if video_length % clip_length != 0 and drop_left:
-        video = video[:, :, :video_length // clip_length * clip_length, :, :]
-        print(f'[split_video_to_clips] Drop frames from {shape} to {video.shape}')
-    nclips = video_length // clip_length
-    clips = rearrange(video, 'b c (nc cl) h w -> (b nc) c cl h w', cl=clip_length, nc=nclips)
-    return clips
-
-def merge_clips_to_videos(clips, bs):
-    nclips = clips.shape[0] // bs
-    video = rearrange(clips, '(b nc) c t h w -> b c (nc t) h w', nc=nclips)
-    return video
+import pytorch_lightning as pl
+from utils.utils import instantiate_from_config
+from lvdm.ema import LitEma
+from lvdm.distributions import DiagonalGaussianDistribution
+from lvdm.models.utils_diffusion import make_beta_schedule
+from lvdm.modules.encoders.ip_resampler import ImageProjModel, Resampler
+from lvdm.basics import disabled_train
+from lvdm.common import (
+    extract_into_tensor,
+    noise_like,
+    exists,
+    default
+)
+
+
+__conditioning_keys__ = {'concat': 'c_concat',
+                         'crossattn': 'c_crossattn',
+                         'adm': 'y'}
 
 class DDPM(pl.LightningModule):
-    # classic DDPM with Gaussian diffusion, in pixel space
+    # classic DDPM with Gaussian diffusion, in image space
     def __init__(self,
                  unet_config,
                  timesteps=1000,
@@ -57,11 +45,10 @@ class DDPM(pl.LightningModule):
                  ckpt_path=None,
                  ignore_keys=[],
                  load_only_unet=False,
-                 monitor="val/loss",
+                 monitor=None,
                  use_ema=True,
                  first_stage_key="image",
                  image_size=256,
-                 video_length=None,
                  channels=3,
                  log_every_t=100,
                  clip_denoised=True,
@@ -70,35 +57,35 @@ class DDPM(pl.LightningModule):
                  cosine_s=8e-3,
                  given_betas=None,
                  original_elbo_weight=0.,
-                 v_posterior=0.,
+                 v_posterior=0.,  # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
                  l_simple_weight=1.,
                  conditioning_key=None,
-                 parameterization="eps",
+                 parameterization="eps",  # all assuming fixed variance schedules
                  scheduler_config=None,
+                 use_positional_encodings=False,
                  learn_logvar=False,
-                 logvar_init=0.,
-                 *args, **kwargs
+                 logvar_init=0.
                  ):
         super().__init__()
         assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
         self.parameterization = parameterization
-        print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
+        mainlogger.info(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
         self.cond_stage_model = None
         self.clip_denoised = clip_denoised
         self.log_every_t = log_every_t
         self.first_stage_key = first_stage_key
-        self.image_size = image_size  # try conv?
-        
+        self.channels = channels
+        self.temporal_length = unet_config.params.temporal_length
+        self.image_size = image_size 
         if isinstance(self.image_size, int):
             self.image_size = [self.image_size, self.image_size]
-        self.channels = channels
+        self.use_positional_encodings = use_positional_encodings
         self.model = DiffusionWrapper(unet_config, conditioning_key)
-        self.conditioning_key = conditioning_key # also register conditioning_key in diffusion
-        
-        self.temporal_length = video_length if video_length is not None else unet_config.params.temporal_length
-        count_params(self.model, verbose=True)
         self.use_ema = use_ema
-        
+        if self.use_ema:
+            self.model_ema = LitEma(self.model)
+            mainlogger.info(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
         self.use_scheduler = scheduler_config is not None
         if self.use_scheduler:
             self.scheduler_config = scheduler_config
@@ -122,6 +109,7 @@ class DDPM(pl.LightningModule):
         if self.learn_logvar:
             self.logvar = nn.Parameter(self.logvar, requires_grad=True)
 
+
     def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
                           linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
         if exists(given_betas):
@@ -182,14 +170,14 @@ class DDPM(pl.LightningModule):
             self.model_ema.store(self.model.parameters())
             self.model_ema.copy_to(self.model)
             if context is not None:
-                print(f"{context}: Switched to EMA weights")
+                mainlogger.info(f"{context}: Switched to EMA weights")
         try:
             yield None
         finally:
             if self.use_ema:
                 self.model_ema.restore(self.model.parameters())
                 if context is not None:
-                    print(f"{context}: Restored training weights")
+                    mainlogger.info(f"{context}: Restored training weights")
 
     def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
         sd = torch.load(path, map_location="cpu")
@@ -198,16 +186,16 @@ class DDPM(pl.LightningModule):
         keys = list(sd.keys())
         for k in keys:
             for ik in ignore_keys:
-                if k.startswith(ik) or (ik.startswith('**') and ik.split('**')[-1] in k):
-                    print("Deleting key {} from state_dict.".format(k))
+                if k.startswith(ik):
+                    mainlogger.info("Deleting key {} from state_dict.".format(k))
                     del sd[k]
         missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
             sd, strict=False)
-        print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+        mainlogger.info(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
         if len(missing) > 0:
-            print(f"Missing Keys: {missing}")
+            mainlogger.info(f"Missing Keys: {missing}")
         if len(unexpected) > 0:
-            print(f"Unexpected Keys: {unexpected}")
+            mainlogger.info(f"Unexpected Keys: {unexpected}")
 
     def q_mean_variance(self, x_start, t):
         """
@@ -274,196 +262,51 @@ class DDPM(pl.LightningModule):
 
     @torch.no_grad()
     def sample(self, batch_size=16, return_intermediates=False):
+        image_size = self.image_size
         channels = self.channels
-        video_length = self.total_length
-        size = (batch_size, channels, video_length, *self.image_size)
-        return self.p_sample_loop(size,
+        return self.p_sample_loop((batch_size, channels, image_size, image_size),
                                   return_intermediates=return_intermediates)
 
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
-        return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+        return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start *
+                extract_into_tensor(self.scale_arr, t, x_start.shape) +
                 extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
 
-    def get_loss(self, pred, target, mean=True, mask=None):
-        if self.loss_type == 'l1':
-            loss = (target - pred).abs()
-            if mean:
-                loss = loss.mean()
-        elif self.loss_type == 'l2':
-            if mean:
-                loss = torch.nn.functional.mse_loss(target, pred)
-            else:
-                loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
-        else:
-            raise NotImplementedError("unknown loss type '{loss_type}'")
-        if mask is not None:
-            assert(mean is False)
-            assert(loss.shape[2:] == mask.shape[2:]) #thw need be the same
-            loss = loss * mask
-        return loss
-
-    def p_losses(self, x_start, t, noise=None):
-        noise = default(noise, lambda: torch.randn_like(x_start))
-        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
-        model_out = self.model(x_noisy, t)
-
-        loss_dict = {}
-        if self.parameterization == "eps":
-            target = noise
-        elif self.parameterization == "x0":
-            target = x_start
-        else:
-            raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported")
-
-        loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3, 4])
-
-        log_prefix = 'train' if self.training else 'val'
-
-        loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
-        loss_simple = loss.mean() * self.l_simple_weight
-
-        loss_vlb = (self.lvlb_weights[t] * loss).mean()
-        loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
-
-        loss = loss_simple + self.original_elbo_weight * loss_vlb
-
-        loss_dict.update({f'{log_prefix}/loss': loss})
-
-        return loss, loss_dict
-
-    def forward(self, x, *args, **kwargs):
-        t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
-        return self.p_losses(x, t, *args, **kwargs)
-
     def get_input(self, batch, k):
         x = batch[k]
         x = x.to(memory_format=torch.contiguous_format).float()
         return x
 
-    def shared_step(self, batch):
-        x = self.get_input(batch, self.first_stage_key)
-        loss, loss_dict = self(x)
-        return loss, loss_dict
-
-    def training_step(self, batch, batch_idx):
-        loss, loss_dict = self.shared_step(batch)
-
-        self.log_dict(loss_dict, prog_bar=True,
-                      logger=True, on_step=True, on_epoch=True)
-
-        self.log("global_step", self.global_step,
-                 prog_bar=True, logger=True, on_step=True, on_epoch=False)
-
-        if self.use_scheduler:
-            lr = self.optimizers().param_groups[0]['lr']
-            self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
-        
-        if self.log_time:
-            total_train_time = (time.time() - self.start_time) / (3600*24)
-            avg_step_time = (time.time() - self.start_time) / (self.global_step + 1)
-            left_time_2w_step = (20000-self.global_step -1) * avg_step_time / (3600*24)
-            left_time_5w_step = (50000-self.global_step -1) * avg_step_time / (3600*24)
-            with open(self.logger_path, 'w') as f:
-                print(f'total_train_time = {total_train_time:.1f} days \n\
-                      total_train_step = {self.global_step + 1} steps \n\
-                      left_time_2w_step = {left_time_2w_step:.1f} days \n\
-                      left_time_5w_step = {left_time_5w_step:.1f} days', file=f)
-        return loss
-
-    @torch.no_grad()
-    def validation_step(self, batch, batch_idx):
-        # _, loss_dict_no_ema = self.shared_step_validate(batch)
-        # with self.ema_scope():
-        #     _, loss_dict_ema = self.shared_step_validate(batch)
-        #     loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
-        # self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
-        # self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
-        if (self.global_step) % self.val_fvd_interval == 0 and self.global_step != 0:
-            print(f'sample for fvd...')
-            self.log_images_kwargs = {
-                'inpaint': False,
-                'plot_diffusion_rows': False,
-                'plot_progressive_rows': False,
-                'ddim_steps': 50,
-                'unconditional_guidance_scale': 15.0,
-            }
-            torch.cuda.empty_cache()
-            logs = self.log_images(batch, **self.log_images_kwargs)
-            self.log("batch_idx", batch_idx,
-                    prog_bar=True, on_step=True, on_epoch=False)
-            return {'real': logs['inputs'], 'fake': logs['samples'], 'conditioning_txt_img': logs['conditioning_txt_img']}
-    
-    def get_condition_validate(self, prompt):
-        """ text embd
-        """
-        if isinstance(prompt, str):
-            prompt = [prompt]
-        c = self.get_learned_conditioning(prompt)
-        bs = c.shape[0]
-        
-        return c
-    
-    def on_train_batch_end(self, *args, **kwargs):
-        if self.use_ema:
-            self.model_ema(self.model)
-        
-    def training_epoch_end(self, outputs):
-        
-        if (self.current_epoch == 0) or self.resume_new_epoch == 0:
-            self.epoch_start_time = time.time()
-            self.current_epoch_time = 0
-            self.total_time = 0
-            self.epoch_time_avg = 0
-        else:
-            self.current_epoch_time = time.time() - self.epoch_start_time
-            self.epoch_start_time = time.time()
-            self.total_time += self.current_epoch_time
-            self.epoch_time_avg = self.total_time / self.current_epoch
-        self.resume_new_epoch += 1
-        epoch_avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
-        
-        self.log('train/epoch/loss', epoch_avg_loss, logger=True, on_epoch=True)
-        self.log('train/epoch/idx', self.current_epoch, logger=True, on_epoch=True)
-        self.log('train/epoch/time', self.current_epoch_time, logger=True, on_epoch=True)
-        self.log('train/epoch/time_avg', self.epoch_time_avg, logger=True, on_epoch=True)
-        self.log('train/epoch/time_avg_min', self.epoch_time_avg / 60, logger=True, on_epoch=True)
-
     def _get_rows_from_list(self, samples):
         n_imgs_per_row = len(samples)
-        denoise_grid = rearrange(samples, 'n b c t h w -> b n c t h w')
-        denoise_grid = rearrange(denoise_grid, 'b n c t h w -> (b n) c t h w')
-        denoise_grid = rearrange(denoise_grid, 'n c t h w -> (n t) c h w')
+        denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
+        denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
         denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
         return denoise_grid
 
     @torch.no_grad()
-    def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, 
-                   plot_diffusion_rows=True, plot_denoise_rows=True, **kwargs):
-        """ log images for DDPM """
+    def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
         log = dict()
         x = self.get_input(batch, self.first_stage_key)
         N = min(x.shape[0], N)
         n_row = min(x.shape[0], n_row)
         x = x.to(self.device)[:N]
         log["inputs"] = x
-        if 'fps' in batch:
-            log['fps'] = batch['fps']
 
-        if plot_diffusion_rows:
-            # get diffusion row
-            diffusion_row = list()
-            x_start = x[:n_row]
+        # get diffusion row
+        diffusion_row = list()
+        x_start = x[:n_row]
 
-            for t in range(self.num_timesteps):
-                if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
-                    t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
-                    t = t.to(self.device).long()
-                    noise = torch.randn_like(x_start)
-                    x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
-                    diffusion_row.append(x_noisy)
+        for t in range(self.num_timesteps):
+            if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+                t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+                t = t.to(self.device).long()
+                noise = torch.randn_like(x_start)
+                x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+                diffusion_row.append(x_noisy)
 
-            log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
+        log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
 
         if sample:
             # get denoise row
@@ -471,8 +314,7 @@ class DDPM(pl.LightningModule):
                 samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
 
             log["samples"] = samples
-            if plot_denoise_rows:
-                log["denoise_row"] = self._get_rows_from_list(denoise_row)
+            log["denoise_row"] = self._get_rows_from_list(denoise_row)
 
         if return_keys:
             if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
@@ -481,14 +323,6 @@ class DDPM(pl.LightningModule):
                 return {key: log[key] for key in return_keys}
         return log
 
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        params = list(self.model.parameters())
-        if self.learn_logvar:
-            params = params + [self.logvar]
-        opt = torch.optim.AdamW(params, lr=lr)
-        return opt
-
 
 class LatentDiffusion(DDPM):
     """main class"""
@@ -496,36 +330,51 @@ class LatentDiffusion(DDPM):
                  first_stage_config,
                  cond_stage_config,
                  num_timesteps_cond=None,
-                 cond_stage_key="image",
+                 cond_stage_key="caption",
                  cond_stage_trainable=False,
-                 concat_mode=True,
                  cond_stage_forward=None,
                  conditioning_key=None,
+                 uncond_prob=0.2,
+                 uncond_type="empty_seq",
                  scale_factor=1.0,
                  scale_by_std=False,
                  encoder_type="2d",
-                 shift_factor=0.0,
-                 split_clips=True,
-                 downfactor_t=None,
-                 clip_length=None,
                  only_model=False,
-                 lora_args={},
+                 use_scale=False,
+                 scale_a=1,
+                 scale_b=0.3,
+                 mid_step=400,
+                 fix_scale_bug=False,
                  *args, **kwargs):
         self.num_timesteps_cond = default(num_timesteps_cond, 1)
         self.scale_by_std = scale_by_std
         assert self.num_timesteps_cond <= kwargs['timesteps']
         # for backwards compatibility after implementation of DiffusionWrapper
-        
-        if conditioning_key is None:
-            conditioning_key = 'concat' if concat_mode else 'crossattn'
-        if cond_stage_config == '__is_unconditional__':
-            conditioning_key = None
         ckpt_path = kwargs.pop("ckpt_path", None)
         ignore_keys = kwargs.pop("ignore_keys", [])
+        conditioning_key = default(conditioning_key, 'crossattn')
         super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
-        self.concat_mode = concat_mode
+
         self.cond_stage_trainable = cond_stage_trainable
         self.cond_stage_key = cond_stage_key
+
+        # scale factor
+        self.use_scale=use_scale
+        if self.use_scale:
+            self.scale_a=scale_a
+            self.scale_b=scale_b
+            if fix_scale_bug:
+                scale_step=self.num_timesteps-mid_step
+            else: #bug
+                scale_step = self.num_timesteps
+
+            scale_arr1 = np.linspace(scale_a, scale_b, mid_step)
+            scale_arr2 = np.full(scale_step, scale_b)
+            scale_arr = np.concatenate((scale_arr1, scale_arr2))
+            scale_arr_prev = np.append(scale_a, scale_arr[:-1])
+            to_torch = partial(torch.tensor, dtype=torch.float32)
+            self.register_buffer('scale_arr', to_torch(scale_arr))
+
         try:
             self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
         except:
@@ -536,71 +385,44 @@ class LatentDiffusion(DDPM):
             self.register_buffer('scale_factor', torch.tensor(scale_factor))
         self.instantiate_first_stage(first_stage_config)
         self.instantiate_cond_stage(cond_stage_config)
-        self.cond_stage_forward = cond_stage_forward
-        self.clip_denoised = False
-        self.bbox_tokenizer = None  
-        self.cond_stage_config = cond_stage_config
         self.first_stage_config = first_stage_config
+        self.cond_stage_config = cond_stage_config        
+        self.clip_denoised = False
+
+        self.cond_stage_forward = cond_stage_forward
         self.encoder_type = encoder_type
         assert(encoder_type in ["2d", "3d"])
+        self.uncond_prob = uncond_prob
+        self.classifier_free_guidance = True if uncond_prob > 0 else False
+        assert(uncond_type in ["zero_embed", "empty_seq"])
+        self.uncond_type = uncond_type
+
+
         self.restarted_from_ckpt = False
-        self.shift_factor = shift_factor
         if ckpt_path is not None:
             self.init_from_ckpt(ckpt_path, ignore_keys, only_model=only_model)
             self.restarted_from_ckpt = True
-        self.split_clips = split_clips
-        self.downfactor_t = downfactor_t
-        self.clip_length = clip_length
-        # lora related args
-        self.inject_unet = getattr(lora_args, "inject_unet", False)
-        self.inject_clip = getattr(lora_args, "inject_clip", False)
-        self.inject_unet_key_word = getattr(lora_args, "inject_unet_key_word", None)
-        self.inject_clip_key_word = getattr(lora_args, "inject_clip_key_word", None)
-        self.lora_rank = getattr(lora_args, "lora_rank", 4)
+                
 
     def make_cond_schedule(self, ):
         self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
         ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
         self.cond_ids[:self.num_timesteps_cond] = ids
 
-    def inject_lora(self, lora_scale=1.0):
-        if self.inject_unet:
-            self.lora_require_grad_params, self.lora_names = inject_trainable_lora(self.model, self.inject_unet_key_word, 
-                                                                                   r=self.lora_rank,
-                                                                                   scale=lora_scale
-                                                                                   )
-        if self.inject_clip:
-            self.lora_require_grad_params_clip, self.lora_names_clip = inject_trainable_lora(self.cond_stage_model, self.inject_clip_key_word, 
-                                                                                             r=self.lora_rank,
-                                                                                             scale=lora_scale
-                                                                                             )
-
-    @rank_zero_only
-    @torch.no_grad()
-    def on_train_batch_start(self, batch, batch_idx, dataloader_idx=None):
-        # only for very first batch, reset the self.scale_factor
-        if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
-            assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
-            # set rescale weight to 1./std of encodings
-            print("### USING STD-RESCALING ###")
-            x = super().get_input(batch, self.first_stage_key)
-            x = x.to(self.device)
-            encoder_posterior = self.encode_first_stage(x)
-            z = self.get_first_stage_encoding(encoder_posterior).detach()
-            del self.scale_factor
-            self.register_buffer('scale_factor', 1. / z.flatten().std())
-            print(f"setting self.scale_factor to {self.scale_factor}")
-            print("### USING STD-RESCALING ###")
-            print(f"std={z.flatten().std()}")
-
-    def register_schedule(self,
-                          given_betas=None, beta_schedule="linear", timesteps=1000,
-                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
-        super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
+    def q_sample(self, x_start, t, noise=None):
+        noise = default(noise, lambda: torch.randn_like(x_start))
+        if self.use_scale:  
+            return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start *
+                extract_into_tensor(self.scale_arr, t, x_start.shape) +
+                extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
+        else:
+            return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+                extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
 
-        self.shorten_cond_schedule = self.num_timesteps_cond > 1
-        if self.shorten_cond_schedule:
-            self.make_cond_schedule()
+
+    def _freeze_model(self):
+        for name, para in self.model.diffusion_model.named_parameters():
+            para.requires_grad = False
 
     def instantiate_first_stage(self, config):
         model = instantiate_from_config(config)
@@ -610,40 +432,16 @@ class LatentDiffusion(DDPM):
             param.requires_grad = False
 
     def instantiate_cond_stage(self, config):
-        if config is None:
-            self.cond_stage_model = None
-            return
         if not self.cond_stage_trainable:
-            if config == "__is_first_stage__":
-                print("Using first stage also as cond stage.")
-                self.cond_stage_model = self.first_stage_model
-            elif config == "__is_unconditional__":
-                print(f"Training {self.__class__.__name__} as an unconditional model.")
-                self.cond_stage_model = None
-            else:
-                model = instantiate_from_config(config)
-                self.cond_stage_model = model.eval()
-                self.cond_stage_model.train = disabled_train
-                for param in self.cond_stage_model.parameters():
-                    param.requires_grad = False
+            model = instantiate_from_config(config)
+            self.cond_stage_model = model.eval()
+            self.cond_stage_model.train = disabled_train
+            for param in self.cond_stage_model.parameters():
+                param.requires_grad = False
         else:
-            assert config != '__is_first_stage__'
-            assert config != '__is_unconditional__'
             model = instantiate_from_config(config)
             self.cond_stage_model = model
-
-
-    def get_first_stage_encoding(self, encoder_posterior, noise=None):
-        if isinstance(encoder_posterior, DiagonalGaussianDistribution):
-            z = encoder_posterior.sample(noise=noise)
-        elif isinstance(encoder_posterior, torch.Tensor):
-            z = encoder_posterior
-        else:
-            raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
-        z = self.scale_factor * (z + self.shift_factor)
-        return z
-
-
+    
     def get_learned_conditioning(self, c):
         if self.cond_stage_forward is None:
             if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
@@ -657,197 +455,61 @@ class LatentDiffusion(DDPM):
             c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
         return c
 
-
-    @torch.no_grad()
-    def get_condition(self, batch, x, bs, force_c_encode, k, cond_key, is_imgs=False):
-        is_conditional = self.model.conditioning_key is not None # crossattn
-        if is_conditional:
-            if cond_key is None:
-                cond_key = self.cond_stage_key 
-            
-            # get condition batch of different condition type
-            if cond_key != self.first_stage_key:
-                assert(cond_key in ["caption", "txt"])
-                xc = batch[cond_key]
-            else:
-                xc = x
-            
-            # if static video
-            if self.static_video:
-                xc_ = [c + ' (static)' for c in xc]
-                xc = xc_
-            
-            # get learned condition.
-            # can directly skip it: c = xc
-            if self.cond_stage_config is not None and (not self.cond_stage_trainable or force_c_encode):
-                if isinstance(xc, torch.Tensor):
-                    xc = xc.to(self.device)
-                c = self.get_learned_conditioning(xc)
-            else:
-                c = xc
-
-            if self.classfier_free_guidance:
-                if cond_key in ['caption', "txt"] and self.uncond_type == 'empty_seq':
-                    for i, ci in enumerate(c):
-                        if random.random() < self.prob:
-                            c[i] = ""
-                elif cond_key == 'class_label' and self.uncond_type == 'zero_embed':
-                    pass
-                elif cond_key == 'class_label' and self.uncond_type == 'learned_embed':
-                    import pdb;pdb.set_trace()
-                    for i, ci in enumerate(c):
-                        if random.random() < self.prob:
-                            c[i]['class_label'] = self.n_classes
-                    
-                else:
-                    raise NotImplementedError
-                
-            if self.zero_cond_embed:
-                import pdb;pdb.set_trace()
-                c = torch.zeros_like(c)
-
-            # process c
-            if bs is not None:
-                if (is_imgs and not self.static_video):
-                    c = c[:bs*self.temporal_length] # each random img (in T axis) has a corresponding prompt
-                else:
-                    c = c[:bs]
-
+    def get_first_stage_encoding(self, encoder_posterior, noise=None):
+        if isinstance(encoder_posterior, DiagonalGaussianDistribution):
+            z = encoder_posterior.sample(noise=noise)
+        elif isinstance(encoder_posterior, torch.Tensor):
+            z = encoder_posterior
         else:
-            c = None
-            xc = None
-            
-        return c, xc
-
+            raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
+        return self.scale_factor * z
+   
     @torch.no_grad()
-    def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
-                  cond_key=None, return_original_cond=False, bs=None, mask_temporal=False):
-        """ Get input in LDM 
-        """
-        # get input imgaes
-        x = super().get_input(batch, k) # k = first_stage_key=image
-        is_imgs = True if k == 'jpg' else False
-        if is_imgs:
-            if self.static_video:
-                # repeat single img to a static video
-                x = x.unsqueeze(2) # bchw -> bc1hw
-                x = x.repeat(1,1,self.temporal_length,1,1) # bc1hw -> bcthw
-            else:
-                # rearrange to videos with T random img
-                bs_load = x.shape[0] // self.temporal_length
-                x = x[:bs_load*self.temporal_length, ...]
-                x = rearrange(x, '(b t) c h w -> b c t h w', t=self.temporal_length, b=bs_load)
-
-        if bs is not None:
-            x = x[:bs]
-        
-        x = x.to(self.device)
-        x_ori = x
-        
-        b, _, t, h, w = x.shape
-        
-        # encode video frames x to z via a 2D encoder
-        x = rearrange(x, 'b c t h w -> (b t) c h w')
-        encoder_posterior = self.encode_first_stage(x, mask_temporal)
-        z = self.get_first_stage_encoding(encoder_posterior).detach()
-        z = rearrange(z, '(b t) c h w -> b c t h w', b=b, t=t)
-        
+    def encode_first_stage(self, x):
+        if self.encoder_type == "2d" and x.dim() == 5:
+            b, _, t, _, _ = x.shape
+            x = rearrange(x, 'b c t h w -> (b t) c h w')
+            reshape_back = True
+        else:
+            reshape_back = False
         
-        c, xc = self.get_condition(batch, x, bs, force_c_encode, k, cond_key, is_imgs)
-        out = [z, c]
+        encoder_posterior = self.first_stage_model.encode(x)
+        results = self.get_first_stage_encoding(encoder_posterior).detach()
         
-        if return_first_stage_outputs:
-            xrec = self.decode_first_stage(z, mask_temporal=mask_temporal)
-            out.extend([x_ori, xrec])
-        if return_original_cond:
-            if isinstance(xc, torch.Tensor) and xc.dim() == 4:
-                xc = rearrange(xc, '(b t) c h w -> b c t h w', b=b, t=t)
-            out.append(xc)
+        if reshape_back:
+            results = rearrange(results, '(b t) c h w -> b c t h w', b=b,t=t)
         
-        return out
-    
-    @torch.no_grad()
-    def decode(self, z, **kwargs,):
-        z = 1. / self.scale_factor * z - self.shift_factor
-        results = self.first_stage_model.decode(z,**kwargs)
         return results
     
     @torch.no_grad()
-    def decode_first_stage_2DAE(self, z, decode_bs=16, return_cpu=True, **kwargs):
-        b, _, t, _, _ = z.shape
-        z = rearrange(z, 'b c t h w -> (b t) c h w')
-        if decode_bs is None:
-            results = self.decode(z, **kwargs)
-        else:
-            z = torch.split(z, decode_bs, dim=0)
-            if return_cpu:
-                results = torch.cat([self.decode(z_, **kwargs).cpu() for z_ in z], dim=0)
-            else:
-                results = torch.cat([self.decode(z_, **kwargs) for z_ in z], dim=0)
-        results = rearrange(results, '(b t) c h w -> b c t h w', b=b,t=t).contiguous()
-        return results
-
-    @torch.no_grad()
-    def decode_first_stage(self, z, decode_bs=16, return_cpu=True, **kwargs):
-        assert(self.encoder_type == "2d" and z.dim() == 5)
-        return self.decode_first_stage_2DAE(z, decode_bs=decode_bs, return_cpu=return_cpu, **kwargs)
+    def encode_first_stage_2DAE(self, x):
 
-    @torch.no_grad()
-    def encode_first_stage_2DAE(self, x, encode_bs=16):
         b, _, t, _, _ = x.shape
-        x = rearrange(x, 'b c t h w -> (b t) c h w')
-        if encode_bs is None:
-            results = self.first_stage_model.encode(x)
-        else:
-            x = torch.split(x, encode_bs, dim=0)
-            zs = []
-            for x_ in x:
-                encoder_posterior = self.first_stage_model.encode(x_)
-                z = self.get_first_stage_encoding(encoder_posterior).detach()
-                zs.append(z)
-            results = torch.cat(zs, dim=0)
-        results = rearrange(results, '(b t) c h w -> b c t h w', b=b,t=t)
+        results = torch.cat([self.get_first_stage_encoding(self.first_stage_model.encode(x[:,:,i])).detach().unsqueeze(2) for i in range(t)], dim=2)
+        
         return results
     
-    @torch.no_grad()
-    def encode_first_stage(self, x):
-        assert(self.encoder_type == "2d" and x.dim() == 5)
-        b, _, t, _, _ = x.shape
-        x = rearrange(x, 'b c t h w -> (b t) c h w')
-        results = self.first_stage_model.encode(x)
-        results = rearrange(results, '(b t) c h w -> b c t h w', b=b,t=t)
-        return results
+    def decode_core(self, z, **kwargs):
+        if self.encoder_type == "2d" and z.dim() == 5:
+            b, _, t, _, _ = z.shape
+            z = rearrange(z, 'b c t h w -> (b t) c h w')
+            reshape_back = True
+        else:
+            reshape_back = False
+            
+        z = 1. / self.scale_factor * z
 
-    def shared_step(self, batch, **kwargs):
-        """ shared step of LDM.
-        If learned condition, c is raw condition (e.g. text)
-        Encoding condition is performed in below forward function.
-        """
-        x, c = self.get_input(batch, self.first_stage_key)
-        loss = self(x, c)
-        return loss
-    
-    def forward(self, x, c, *args, **kwargs):
-        start_t = getattr(self, "start_t", 0)
-        end_t = getattr(self, "end_t", self.num_timesteps)
-        t = torch.randint(start_t, end_t, (x.shape[0],), device=self.device).long()
-        
-        if self.model.conditioning_key is not None:
-            assert c is not None
-            if self.cond_stage_trainable:
-                c = self.get_learned_conditioning(c)
-            if self.classfier_free_guidance and self.uncond_type == 'zero_embed':
-                for i, ci in enumerate(c):
-                    if random.random() < self.prob:
-                        c[i] = torch.zeros_like(c[i])
-            if self.shorten_cond_schedule:  # TODO: drop this option
-                tc = self.cond_ids[t].to(self.device)
-                c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
-        
-        return self.p_losses(x, c, t, *args, **kwargs)
+        results = self.first_stage_model.decode(z, **kwargs)
+            
+        if reshape_back:
+            results = rearrange(results, '(b t) c h w -> b c t h w', b=b,t=t)
+        return results
 
-    def apply_model(self, x_noisy, t, cond, return_ids=False, **kwargs):
+    @torch.no_grad()
+    def decode_first_stage(self, z, **kwargs):
+        return self.decode_core(z, **kwargs)
 
+    def apply_model(self, x_noisy, t, cond, **kwargs):
         if isinstance(cond, dict):
             # hybrid case, cond is exptected to be a dict
             pass
@@ -859,104 +521,55 @@ class LatentDiffusion(DDPM):
 
         x_recon = self.model(x_noisy, t, **cond, **kwargs)
 
-        if isinstance(x_recon, tuple) and not return_ids:
+        if isinstance(x_recon, tuple):
             return x_recon[0]
         else:
             return x_recon
 
-    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
-        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
-               extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+    def _get_denoise_row_from_list(self, samples, desc=''):
+        denoise_row = []
+        for zd in tqdm(samples, desc=desc):
+            denoise_row.append(self.decode_first_stage(zd.to(self.device)))
+        n_log_timesteps = len(denoise_row)
 
-    def _prior_bpd(self, x_start):
-        """
-        Get the prior KL term for the variational lower-bound, measured in
-        bits-per-dim.
-        This term can't be optimized, as it only depends on the encoder.
-        :param x_start: the [N x C x ...] tensor of inputs.
-        :return: a batch of [N] KL values (in bits), one per batch element.
-        """
-        batch_size = x_start.shape[0]
-        t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
-        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
-        kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
-        return mean_flat(kl_prior) / np.log(2.0)
-
-    def p_losses(self, x_start, cond, t, noise=None, skip_qsample=False, x_noisy=None, cond_mask=None, **kwargs,):
-        if not skip_qsample:
-            noise = default(noise, lambda: torch.randn_like(x_start))
-            x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+        denoise_row = torch.stack(denoise_row)  # n_log_timesteps, b, C, H, W
+        
+        if denoise_row.dim() == 5:
+            # img, num_imgs= n_log_timesteps * bs, grid_size=[bs,n_log_timesteps]
+            denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
+            denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+            denoise_grid = make_grid(denoise_grid, nrow=n_log_timesteps)
+        elif denoise_row.dim() == 6:
+            # video, grid_size=[n_log_timesteps*bs, t]
+            video_length = denoise_row.shape[3]
+            denoise_grid = rearrange(denoise_row, 'n b c t h w -> b n c t h w')
+            denoise_grid = rearrange(denoise_grid, 'b n c t h w -> (b n) c t h w')
+            denoise_grid = rearrange(denoise_grid, 'n c t h w -> (n t) c h w')
+            denoise_grid = make_grid(denoise_grid, nrow=video_length)
         else:
-            assert(x_noisy is not None)
-            assert(noise is not None)
-        model_output = self.apply_model(x_noisy, t, cond, **kwargs)
+            raise ValueError
 
-        loss_dict = {}
-        prefix = 'train' if self.training else 'val'
+        return denoise_grid
+ 
 
-        if self.parameterization == "x0":
-            target = x_start
-        elif self.parameterization == "eps":
-            target = noise
-        else:
-            raise NotImplementedError()
-        
-        loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3, 4])
-        loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
-        if self.logvar.device != self.device:
-            self.logvar = self.logvar.to(self.device)
-        logvar_t = self.logvar[t]
-        loss = loss_simple / torch.exp(logvar_t) + logvar_t
-        if self.learn_logvar:
-            loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
-            loss_dict.update({'logvar': self.logvar.data.mean()})
+    @torch.no_grad()
+    def decode_first_stage_2DAE(self, z, **kwargs):
 
-        loss = self.l_simple_weight * loss.mean()
+        b, _, t, _, _ = z.shape
+        z = 1. / self.scale_factor * z
+        results = torch.cat([self.first_stage_model.decode(z[:,:,i], **kwargs).unsqueeze(2) for i in range(t)], dim=2)
 
-        loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3, 4))
-        loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
-        loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
-        loss += (self.original_elbo_weight * loss_vlb)
-        loss_dict.update({f'{prefix}/loss': loss})
+        return results
 
-        return loss, loss_dict
 
-    def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
-                        return_x0=False, score_corrector=None, corrector_kwargs=None, 
-                        unconditional_guidance_scale=1., unconditional_conditioning=None,
-                        uc_type=None,):
+    def p_mean_variance(self, x, c, t, clip_denoised: bool, return_x0=False, score_corrector=None, corrector_kwargs=None, **kwargs):
         t_in = t
-        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
-            model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
-        else:
-            # with unconditional condition
-            if isinstance(c, torch.Tensor):
-                x_in = torch.cat([x] * 2)
-                t_in = torch.cat([t] * 2)
-                c_in = torch.cat([unconditional_conditioning, c])
-                model_out_uncond, model_out = self.apply_model(x_in, t_in, c_in, return_ids=return_codebook_ids).chunk(2)
-            elif isinstance(c, dict):
-                model_out = self.apply_model(x, t, c, return_ids=return_codebook_ids)
-                model_out_uncond = self.apply_model(x, t, unconditional_conditioning, return_ids=return_codebook_ids)
-            else:
-                raise NotImplementedError
-            if uc_type is None:
-                model_out = model_out_uncond + unconditional_guidance_scale * (model_out - model_out_uncond)
-            else:
-                if uc_type == 'cfg_original':
-                    model_out = model_out + unconditional_guidance_scale * (model_out - model_out_uncond)
-                elif uc_type == 'cfg_ours':
-                    model_out = model_out + unconditional_guidance_scale * (model_out_uncond - model_out)
-                else:
-                    raise NotImplementedError
+        model_out = self.apply_model(x, t_in, c, **kwargs)
 
         if score_corrector is not None:
             assert self.parameterization == "eps"
             model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
 
-        if return_codebook_ids:
-            model_out, logits = model_out
-
         if self.parameterization == "eps":
             x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
         elif self.parameterization == "x0":
@@ -966,34 +579,21 @@ class LatentDiffusion(DDPM):
 
         if clip_denoised:
             x_recon.clamp_(-1., 1.)
-        if quantize_denoised:
-            x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
+
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
-        if return_codebook_ids:
-            return model_mean, posterior_variance, posterior_log_variance, logits
-        elif return_x0:
+
+        if return_x0:
             return model_mean, posterior_variance, posterior_log_variance, x_recon
         else:
             return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
-    def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
-                 return_codebook_ids=False, quantize_denoised=False, return_x0=False,
-                 temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                 unconditional_guidance_scale=1., unconditional_conditioning=None,
-                 uc_type=None,):
+    def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False, return_x0=False, \
+                 temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None, **kwargs):
         b, *_, device = *x.shape, x.device
-        outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
-                                       return_codebook_ids=return_codebook_ids,
-                                       quantize_denoised=quantize_denoised,
-                                       return_x0=return_x0,
-                                       score_corrector=score_corrector, corrector_kwargs=corrector_kwargs,
-                                       unconditional_guidance_scale=unconditional_guidance_scale, 
-                                       unconditional_conditioning=unconditional_conditioning,
-                                       uc_type=uc_type,)
-        if return_codebook_ids:
-            raise DeprecationWarning("Support dropped.")
-        elif return_x0:
+        outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised, return_x0=return_x0, \
+                                       score_corrector=score_corrector, corrector_kwargs=corrector_kwargs, **kwargs)
+        if return_x0:
             model_mean, _, model_log_variance, x0 = outputs
         else:
             model_mean, _, model_log_variance = outputs
@@ -1001,99 +601,35 @@ class LatentDiffusion(DDPM):
         noise = noise_like(x.shape, device, repeat_noise) * temperature
         if noise_dropout > 0.:
             noise = torch.nn.functional.dropout(noise, p=noise_dropout)
-        
+        # no noise when t == 0
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
 
-        if return_codebook_ids:
-            return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
         if return_x0:
             return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
         else:
             return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
 
     @torch.no_grad()
-    def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False,
-                              img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
-                              score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
-                              log_every_t=None):
-        if not log_every_t:
-            log_every_t = self.log_every_t
-        timesteps = self.num_timesteps
-        if batch_size is not None:
-            b = batch_size if batch_size is not None else shape[0]
-            shape = [batch_size] + list(shape)
-        else:
-            b = batch_size = shape[0]
-        if x_T is None:
-            img = torch.randn(shape, device=self.device)
-        else:
-            img = x_T
-        intermediates = []
-        if cond is not None:
-            if isinstance(cond, dict):
-                cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
-                list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
-            else:
-                cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
-
-        if start_T is not None:
-            timesteps = min(timesteps, start_T)
-        iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
-                        total=timesteps) if verbose else reversed(
-            range(0, timesteps))
-        if type(temperature) == float:
-            temperature = [temperature] * timesteps
-
-        for i in iterator:
-            ts = torch.full((b,), i, device=self.device, dtype=torch.long)
-            if self.shorten_cond_schedule:
-                assert self.model.conditioning_key != 'hybrid'
-                tc = self.cond_ids[ts].to(cond.device)
-                cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
-
-            img, x0_partial = self.p_sample(img, cond, ts,
-                                            clip_denoised=self.clip_denoised,
-                                            quantize_denoised=quantize_denoised, return_x0=True,
-                                            temperature=temperature[i], noise_dropout=noise_dropout,
-                                            score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
-            if mask is not None:
-                assert x0 is not None
-                img_orig = self.q_sample(x0, ts)
-                img = img_orig * mask + (1. - mask) * img
-
-            if i % log_every_t == 0 or i == timesteps - 1:
-                intermediates.append(x0_partial)
-            if callback: callback(i)
-            if img_callback: img_callback(img, i)
-        return img, intermediates
-
-    @torch.no_grad()
-    def p_sample_loop(self, cond, shape, return_intermediates=False,
-                      x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False,
-                      mask=None, x0=None, img_callback=None, start_T=None,
-                      log_every_t=None,
-                      unconditional_guidance_scale=1., unconditional_conditioning=None,
-                      uc_type=None,):
+    def p_sample_loop(self, cond, shape, return_intermediates=False, x_T=None, verbose=True, callback=None, \
+                      timesteps=None, mask=None, x0=None, img_callback=None, start_T=None, log_every_t=None, **kwargs):
 
         if not log_every_t:
             log_every_t = self.log_every_t
         device = self.betas.device
-        b = shape[0]
-        
+        b = shape[0]        
         # sample an initial noise
         if x_T is None:
             img = torch.randn(shape, device=device)
         else:
             img = x_T
-        
+
         intermediates = [img]
         if timesteps is None:
             timesteps = self.num_timesteps
-
         if start_T is not None:
             timesteps = min(timesteps, start_T)
-        iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
-            range(0, timesteps))
+
+        iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(range(0, timesteps))
 
         if mask is not None:
             assert x0 is not None
@@ -1106,12 +642,7 @@ class LatentDiffusion(DDPM):
                 tc = self.cond_ids[ts].to(cond.device)
                 cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
 
-            img = self.p_sample(img, cond, ts,
-                                clip_denoised=self.clip_denoised,
-                                quantize_denoised=quantize_denoised,
-                                unconditional_guidance_scale=unconditional_guidance_scale,
-                                unconditional_conditioning=unconditional_conditioning,
-                                uc_type=uc_type)
+            img = self.p_sample(img, cond, ts, clip_denoised=self.clip_denoised, **kwargs)
             if mask is not None:
                 img_orig = self.q_sample(x0, ts)
                 img = img_orig * mask + (1. - mask) * img
@@ -1125,253 +656,54 @@ class LatentDiffusion(DDPM):
             return img, intermediates
         return img
 
-    @torch.no_grad()
-    def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
-               verbose=True, timesteps=None, quantize_denoised=False,
-               mask=None, x0=None, shape=None, **kwargs):
-        if shape is None:
-            shape = (batch_size, self.channels, self.total_length, *self.image_size)
-        if cond is not None:
-            if isinstance(cond, dict):
-                cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
-                list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
-            else:
-                cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
-        return self.p_sample_loop(cond,
-                                  shape,
-                                  return_intermediates=return_intermediates, x_T=x_T,
-                                  verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
-                                  mask=mask, x0=x0,)
-
-    @torch.no_grad()
-    def sample_log(self,cond,batch_size,ddim, ddim_steps,**kwargs):
-
-        if ddim:
-            ddim_sampler = DDIMSampler(self)
-            shape = (self.channels, self.total_length, *self.image_size)
-            samples, intermediates =ddim_sampler.sample(ddim_steps,batch_size,
-                                                        shape,cond,verbose=False, **kwargs)
 
+class LatentVisualDiffusion(LatentDiffusion):
+    def __init__(self, cond_img_config, finegrained=False, random_cond=False, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.random_cond = random_cond
+        self.instantiate_img_embedder(cond_img_config, freeze=True)
+        num_tokens = 16 if finegrained else 4
+        self.image_proj_model = self.init_projector(use_finegrained=finegrained, num_tokens=num_tokens, input_dim=1024,\
+                                            cross_attention_dim=1024, dim=1280)    
+
+    def instantiate_img_embedder(self, config, freeze=True):
+        embedder = instantiate_from_config(config)
+        if freeze:
+            self.embedder = embedder.eval()
+            self.embedder.train = disabled_train
+            for param in self.embedder.parameters():
+                param.requires_grad = False
+
+    def init_projector(self, use_finegrained, num_tokens, input_dim, cross_attention_dim, dim):
+        if not use_finegrained:
+            image_proj_model = ImageProjModel(clip_extra_context_tokens=num_tokens, cross_attention_dim=cross_attention_dim,
+                clip_embeddings_dim=input_dim
+            )
         else:
-            samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
-                                                 return_intermediates=True, **kwargs)
+            image_proj_model = Resampler(dim=input_dim, depth=4, dim_head=64, heads=12, num_queries=num_tokens,
+                embedding_dim=dim, output_dim=cross_attention_dim, ff_mult=4
+            )
+        return image_proj_model
 
-        return samples, intermediates
-    
-    @torch.no_grad()
-    def log_condition(self, log, batch, xc, x, c, cond_stage_key=None):
-        """ 
-        xc: oringinal condition before enconding. 
-        c: condition after encoding.
-        """
-        if x.dim() == 5:
-            txt_img_shape = [x.shape[3], x.shape[4]]
-        elif x.dim() == 4:
-            txt_img_shape = [x.shape[2], x.shape[3]]
-        else:
-            raise ValueError
-        if self.model.conditioning_key is not None: #concat-time-mask
-            if hasattr(self.cond_stage_model, "decode"):
-                xc = self.cond_stage_model.decode(c)
-                log["conditioning"] = xc
-            elif cond_stage_key in ["caption", "txt"]:
-                log["conditioning_txt_img"] = log_txt_as_img(txt_img_shape, batch[cond_stage_key], size=x.shape[3]//25)
-                log["conditioning_txt"] = batch[cond_stage_key]
-            elif cond_stage_key == 'class_label':
-                try:
-                    xc = log_txt_as_img(txt_img_shape, batch["human_label"], size=x.shape[3]//25)
-                except:
-                    xc = log_txt_as_img(txt_img_shape, batch["class_name"], size=x.shape[3]//25)
-                log['conditioning'] = xc
-            elif isimage(xc):
-                log["conditioning"] = xc
-            if ismap(xc):
-                log["original_conditioning"] = self.to_rgb(xc)
-            if isinstance(c, dict) and 'mask' in c:
-                log['mask'] =self.mask_to_rgb(c['mask'])
-        return log
-    
-    @torch.no_grad()
-    def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., unconditional_guidance_scale=1.0, 
-                   first_stage_key2=None, cond_key2=None,
-                   c=None, 
-                   **kwargs):
-        """ log images for LatentDiffusion """
-        use_ddim = ddim_steps is not None
-        is_imgs = first_stage_key2 is not None
-        if is_imgs:
-            assert(cond_key2 is not None)
-        log = dict()
-
-        # get input
-        z, c, x, xrec, xc = self.get_input(batch, 
-                                           k=self.first_stage_key if first_stage_key2 is None else first_stage_key2,
-                                           return_first_stage_outputs=True,
-                                           force_c_encode=True,
-                                           return_original_cond=True,
-                                           bs=N,
-                                           cond_key=cond_key2 if cond_key2 is not None else None,
-                                           )
-        
-        N_ori = N
-        N = min(z.shape[0], N)
-        n_row = min(x.shape[0], n_row)
+    ## Never delete this func: it is used in log_images() and inference stage
+    def get_image_embeds(self, batch_imgs):
+        ## img: b c h w
+        img_token = self.embedder(batch_imgs)
+        img_emb = self.image_proj_model(img_token)
+        return img_emb
 
-        if unconditional_guidance_scale != 1.0:
-            prompts = N * self.temporal_length * [""] if (is_imgs and not self.static_video) else N * [""]
-            uc = self.get_condition_validate(prompts)
-            
-        else:
-            uc = None
-
-        log["inputs"] = x
-        log["reconstruction"] = xrec
-        log = self.log_condition(log, batch, xc, x, c, 
-                                 cond_stage_key=self.cond_stage_key if cond_key2 is None else cond_key2
-        )
-        
-        if sample:
-            with self.ema_scope("Plotting"):
-                samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim,
-                                                         ddim_steps=ddim_steps,eta=ddim_eta,
-                                                         temporal_length=self.video_length,
-                                                         unconditional_guidance_scale=unconditional_guidance_scale,
-                                                         unconditional_conditioning=uc, **kwargs,
-                                                         )
-            # decode samples
-            x_samples = self.decode_first_stage(samples)
-            log["samples"] = x_samples
-        return log
-
-    def configure_optimizers(self):
-        """ configure_optimizers for LatentDiffusion """
-        lr = self.learning_rate
-        
-        # --------------------------------------------------------------------------------
-        # set parameters
-        if hasattr(self, "only_optimize_empty_parameters") and self.only_optimize_empty_parameters:
-            print("[INFO] Optimize only empty parameters!")
-            assert(hasattr(self, "empty_paras"))
-            params = [p for n, p in self.model.named_parameters() if n in self.empty_paras]
-        elif hasattr(self, "only_optimize_pretrained_parameters") and self.only_optimize_pretrained_parameters:
-            print("[INFO] Optimize only pretrained parameters!")
-            assert(hasattr(self, "empty_paras"))
-            params = [p for n, p in self.model.named_parameters() if n not in self.empty_paras]
-            assert(len(params) != 0)
-        elif getattr(self, "optimize_empty_and_spatialattn", False):
-            print("[INFO] Optimize empty parameters + spatial transformer!")
-            assert(hasattr(self, "empty_paras"))
-            empty_paras = [p for n, p in self.model.named_parameters() if n in self.empty_paras]
-            SA_list = [".attn1.", ".attn2.", ".ff.", ".norm1.", ".norm2.", ".norm3."]
-            SA_params = [p for n, p in self.model.named_parameters() if check_istarget(n, SA_list)]
-            if getattr(self, "spatial_lr_decay", False):
-                params = [
-                    {"params": empty_paras},
-                    {"params": SA_params, "lr": lr * self.spatial_lr_decay}
-                ]
-            else:
-                params = empty_paras + SA_params
-        else:
-            # optimize whole denoiser
-            if hasattr(self, "spatial_lr_decay") and self.spatial_lr_decay:
-                print("[INFO] Optimize the whole net with different lr!")
-                print(f"[INFO] {lr} for empty paras, {lr * self.spatial_lr_decay} for pretrained paras!")
-                empty_paras = [p for n, p in self.model.named_parameters() if n in self.empty_paras]
-                # assert(len(empty_paras) == len(self.empty_paras)) # self.empty_paras:cond_stage_model.embedding.weight not in diffusion model params
-                pretrained_paras = [p for n, p in self.model.named_parameters() if n not in self.empty_paras]
-                params = [
-                    {"params": empty_paras},
-                    {"params": pretrained_paras, "lr": lr * self.spatial_lr_decay}
-                ]
-                print(f"[INFO] Empty paras: {len(empty_paras)}, Pretrained paras: {len(pretrained_paras)}")
-
-            else:
-                params = list(self.model.parameters())
-        
-        if hasattr(self, "generator_trainable") and not self.generator_trainable:
-            # fix unet denoiser
-            params = list()
-
-        if self.inject_unet:
-            params = itertools.chain(*self.lora_require_grad_params)
-                
-        if self.inject_clip:
-            if self.inject_unet:
-                params = list(params)+list(itertools.chain(*self.lora_require_grad_params_clip))
-            else:
-                params = itertools.chain(*self.lora_require_grad_params_clip)
-            
-
-        # append paras
-        # ------------------------------------------------------------------
-        def add_cond_model(cond_model, params):
-            if isinstance(params[0], dict):
-                # parameter groups
-                params.append({"params": list(cond_model.parameters())})
-            else:
-                # parameter list: [torch.nn.parameter.Parameter]
-                params = params + list(cond_model.parameters())
-            return params
-        # ------------------------------------------------------------------
-        
-        if self.cond_stage_trainable:
-            # print(f"{self.__class__.__name__}: Also optimizing conditioner params!")
-            params = add_cond_model(self.cond_stage_model, params)
-        
-        if self.learn_logvar:
-            print('Diffusion model optimizing logvar')
-            if isinstance(params[0], dict):
-                params.append({"params": [self.logvar]})
-            else:
-                params.append(self.logvar)
-        
-        # --------------------------------------------------------------------------------
-        opt = torch.optim.AdamW(params, lr=lr)
-        
-        # lr scheduler
-        if self.use_scheduler:
-            assert 'target' in self.scheduler_config
-            scheduler = instantiate_from_config(self.scheduler_config)
-
-            print("Setting up LambdaLR scheduler...")
-            scheduler = [
-                {
-                    'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
-                    'interval': 'step',
-                    'frequency': 1
-                }]
-            return [opt], scheduler
-        
-        return opt
-    
-    @torch.no_grad()
-    def to_rgb(self, x):
-        x = x.float()
-        if not hasattr(self, "colorize"):
-            self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
-        x = nn.functional.conv2d(x, weight=self.colorize)
-        x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
-        return x
-
-    @torch.no_grad()
-    def mask_to_rgb(self, x):
-        x = x * 255
-        x = x.int()
-        return x
 
 class DiffusionWrapper(pl.LightningModule):
     def __init__(self, diff_model_config, conditioning_key):
         super().__init__()
         self.diffusion_model = instantiate_from_config(diff_model_config)
-        print('Successfully initialize the diffusion model !')
         self.conditioning_key = conditioning_key
-        # assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm', 'resblockcond', 'hybrid-adm', 'hybrid-time']
 
     def forward(self, x, t, c_concat: list = None, c_crossattn: list = None,
                 c_adm=None, s=None, mask=None, **kwargs):
         # temporal_context = fps is foNone
         if self.conditioning_key is None:
-            out = self.diffusion_model(x, t, **kwargs)
+            out = self.diffusion_model(x, t)
         elif self.conditioning_key == 'concat':
             xc = torch.cat([x] + c_concat, dim=1)
             out = self.diffusion_model(xc, t, **kwargs)
@@ -1379,106 +711,53 @@ class DiffusionWrapper(pl.LightningModule):
             cc = torch.cat(c_crossattn, 1)
             out = self.diffusion_model(x, t, context=cc, **kwargs)
         elif self.conditioning_key == 'hybrid':
+            ## it is just right [b,c,t,h,w]: concatenate in channel dim
             xc = torch.cat([x] + c_concat, dim=1)
             cc = torch.cat(c_crossattn, 1)
-            out = self.diffusion_model(xc, t, context=cc, **kwargs)
+            out = self.diffusion_model(xc, t, context=cc)
         elif self.conditioning_key == 'resblockcond':
             cc = c_crossattn[0]
-            out = self.diffusion_model(x, t, context=cc, **kwargs)
+            out = self.diffusion_model(x, t, context=cc)
         elif self.conditioning_key == 'adm':
             cc = c_crossattn[0]
-            out = self.diffusion_model(x, t, y=cc, **kwargs)
+            out = self.diffusion_model(x, t, y=cc)
         elif self.conditioning_key == 'hybrid-adm':
             assert c_adm is not None
             xc = torch.cat([x] + c_concat, dim=1)
             cc = torch.cat(c_crossattn, 1)
-            out = self.diffusion_model(xc, t, context=cc, y=c_adm, **kwargs)
+            out = self.diffusion_model(xc, t, context=cc, y=c_adm)
         elif self.conditioning_key == 'hybrid-time':
             assert s is not None
             xc = torch.cat([x] + c_concat, dim=1)
             cc = torch.cat(c_crossattn, 1)
-            out = self.diffusion_model(xc, t, context=cc, s=s, **kwargs)
+            out = self.diffusion_model(xc, t, context=cc, s=s)
         elif self.conditioning_key == 'concat-time-mask':
             # assert s is not None
-            # print('x & mask:',x.shape,c_concat[0].shape)
+            # mainlogger.info('x & mask:',x.shape,c_concat[0].shape)
             xc = torch.cat([x] + c_concat, dim=1)
-            out = self.diffusion_model(xc, t, context=None, s=s, mask=mask, **kwargs)
+            out = self.diffusion_model(xc, t, context=None, s=s, mask=mask)
         elif self.conditioning_key == 'concat-adm-mask':
             # assert s is not None
-            # print('x & mask:',x.shape,c_concat[0].shape)
+            # mainlogger.info('x & mask:',x.shape,c_concat[0].shape)
             if c_concat is not None:
                 xc = torch.cat([x] + c_concat, dim=1)
             else:
                 xc = x
-            out = self.diffusion_model(xc, t, context=None, y=s, mask=mask, **kwargs)
-        elif self.conditioning_key == 'crossattn-adm':
-            cc = torch.cat(c_crossattn, 1)
-            out = self.diffusion_model(x, t, context=cc, y=s, **kwargs)
+            out = self.diffusion_model(xc, t, context=None, y=s, mask=mask)
         elif self.conditioning_key == 'hybrid-adm-mask':
             cc = torch.cat(c_crossattn, 1)
             if c_concat is not None:
                 xc = torch.cat([x] + c_concat, dim=1)
             else:
                 xc = x
-            out = self.diffusion_model(xc, t, context=cc, y=s, mask=mask, **kwargs)
+            out = self.diffusion_model(xc, t, context=cc, y=s, mask=mask)
         elif self.conditioning_key == 'hybrid-time-adm': # adm means y, e.g., class index
             # assert s is not None
             assert c_adm is not None
             xc = torch.cat([x] + c_concat, dim=1)
             cc = torch.cat(c_crossattn, 1)
-            out = self.diffusion_model(xc, t, context=cc, s=s, y=c_adm, **kwargs)
+            out = self.diffusion_model(xc, t, context=cc, s=s, y=c_adm)
         else:
             raise NotImplementedError()
 
-        return out
-
-
-class T2VAdapterDepth(LatentDiffusion):
-    def __init__(self, depth_stage_config, adapter_config, *args, **kwargs):
-        super(T2VAdapterDepth, self).__init__(*args, **kwargs)
-        self.adapter = instantiate_from_config(adapter_config)
-        self.condtype = adapter_config.cond_name
-        self.depth_stage_model = instantiate_from_config(depth_stage_config)
-
-    def prepare_midas_input(self, batch_x):
-        # input: b,c,h,w
-        x_midas = torch.nn.functional.interpolate(batch_x, size=(384, 384), mode='bicubic')
-        return x_midas
-    
-    @torch.no_grad()
-    def get_batch_depth(self, batch_x, target_size, encode_bs=1):
-        b, c, t, h, w = batch_x.shape
-        merge_x = rearrange(batch_x, 'b c t h w -> (b t) c h w')
-        split_x = torch.split(merge_x, encode_bs, dim=0)
-        cond_depth_list = []
-        for x in split_x:
-            x_midas = self.prepare_midas_input(x)
-            cond_depth = self.depth_stage_model(x_midas)
-            cond_depth = torch.nn.functional.interpolate(
-                    cond_depth,
-                    size=target_size,
-                    mode="bicubic",
-                    align_corners=False,
-                )
-            depth_min, depth_max = torch.amin(cond_depth, dim=[1, 2, 3], keepdim=True), torch.amax(cond_depth, dim=[1, 2, 3], keepdim=True)
-            cond_depth = 2. * (cond_depth - depth_min) / (depth_max - depth_min + 1e-7) - 1.
-            cond_depth_list.append(cond_depth)
-        batch_cond_depth=torch.cat(cond_depth_list, dim=0)
-        batch_cond_depth = rearrange(batch_cond_depth, '(b t) c h w -> b c t h w', b=b, t=t)
-        return batch_cond_depth
-    
-    def get_adapter_features(self, extra_cond, encode_bs=1):
-        b, c, t, h, w = extra_cond.shape
-        ## process in 2D manner
-        merge_extra_cond = rearrange(extra_cond, 'b c t h w -> (b t) c h w')
-        split_extra_cond = torch.split(merge_extra_cond, encode_bs, dim=0)
-        features_adapter_list = []
-        for extra_cond in split_extra_cond:
-            features_adapter = self.adapter(extra_cond)
-            features_adapter_list.append(features_adapter)
-        merge_features_adapter_list = []
-        for i in range(len(features_adapter_list[0])):
-            merge_features_adapter = torch.cat([features_adapter_list[num][i] for num in range(len(features_adapter_list))], dim=0)
-            merge_features_adapter_list.append(merge_features_adapter)
-        merge_features_adapter_list = [rearrange(feature, '(b t) c h w -> b c t h w', b=b, t=t) for feature in merge_features_adapter_list]
-        return merge_features_adapter_list
+        return out

lvdm/models/modules/adapter.py

@@ -1,105 +0,0 @@
-import torch
-import torch.nn as nn
-from collections import OrderedDict
-from lvdm.models.modules.util import (
-    zero_module,
-    conv_nd,
-    avg_pool_nd
-)
-
-class Downsample(nn.Module):
-    """
-    A downsampling layer with an optional convolution.
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 downsampling occurs in the inner-two dimensions.
-    """
-
-    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        stride = 2 if dims != 3 else (1, 2, 2)
-        if use_conv:
-            self.op = conv_nd(
-                dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
-            )
-        else:
-            assert self.channels == self.out_channels
-            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        return self.op(x)
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, in_c, out_c, down, ksize=3, sk=False, use_conv=True):
-        super().__init__()
-        ps = ksize // 2
-        if in_c != out_c or sk == False:
-            self.in_conv = nn.Conv2d(in_c, out_c, ksize, 1, ps)
-        else:
-            # print('n_in')
-            self.in_conv = None
-        self.block1 = nn.Conv2d(out_c, out_c, 3, 1, 1)
-        self.act = nn.ReLU()
-        self.block2 = nn.Conv2d(out_c, out_c, ksize, 1, ps)
-        if sk == False:
-            self.skep = nn.Conv2d(in_c, out_c, ksize, 1, ps)
-        else:
-            self.skep = None
-
-        self.down = down
-        if self.down == True:
-            self.down_opt = Downsample(in_c, use_conv=use_conv)
-
-    def forward(self, x):
-        if self.down == True:
-            x = self.down_opt(x)
-        if self.in_conv is not None:  # edit
-            x = self.in_conv(x)
-
-        h = self.block1(x)
-        h = self.act(h)
-        h = self.block2(h)
-        if self.skep is not None:
-            return h + self.skep(x)
-        else:
-            return h + x
-
-
-class Adapter(nn.Module):
-    def __init__(self, channels=[320, 640, 1280, 1280], nums_rb=3, cin=64, ksize=3, sk=False, use_conv=True):
-        super(Adapter, self).__init__()
-        self.unshuffle = nn.PixelUnshuffle(8)
-        self.channels = channels
-        self.nums_rb = nums_rb
-        self.body = []
-        for i in range(len(channels)):
-            for j in range(nums_rb):
-                if (i != 0) and (j == 0):
-                    self.body.append(
-                        ResnetBlock(channels[i - 1], channels[i], down=True, ksize=ksize, sk=sk, use_conv=use_conv))
-                else:
-                    self.body.append(
-                        ResnetBlock(channels[i], channels[i], down=False, ksize=ksize, sk=sk, use_conv=use_conv))
-        self.body = nn.ModuleList(self.body)
-        self.conv_in = nn.Conv2d(cin, channels[0], 3, 1, 1)
-
-    def forward(self, x):
-        # unshuffle
-        x = self.unshuffle(x)
-        # extract features
-        features = []
-        x = self.conv_in(x)
-        for i in range(len(self.channels)):
-            for j in range(self.nums_rb):
-                idx = i * self.nums_rb + j
-                x = self.body[idx](x)
-            features.append(x)
-
-        return features

lvdm/models/modules/attention_temporal.py

@@ -1,399 +0,0 @@
-from typing import Optional, Any
-
-import torch
-import torch as th
-from torch import nn, einsum
-from einops import rearrange, repeat
-try:
-    import xformers
-    import xformers.ops
-    XFORMERS_IS_AVAILBLE = True
-except:
-    XFORMERS_IS_AVAILBLE = False
-
-from lvdm.models.modules.util import (
-    GEGLU,
-    exists,
-    default,
-    Normalize,
-    checkpoint,
-    zero_module,
-)
-
-
-# ---------------------------------------------------------------------------------------------------
-class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = default(dim_out, dim)
-        project_in = nn.Sequential(
-            nn.Linear(dim, inner_dim),
-            nn.GELU()
-        ) if not glu else GEGLU(dim, inner_dim)
-
-        self.net = nn.Sequential(
-            project_in,
-            nn.Dropout(dropout),
-            nn.Linear(inner_dim, dim_out)
-        )
-
-    def forward(self, x):
-        return self.net(x)
-
-
-# ---------------------------------------------------------------------------------------------------
-class RelativePosition(nn.Module):
-    """ https://github.com/evelinehong/Transformer_Relative_Position_PyTorch/blob/master/relative_position.py """
-
-    def __init__(self, num_units, max_relative_position):
-        super().__init__()
-        self.num_units = num_units
-        self.max_relative_position = max_relative_position
-        self.embeddings_table = nn.Parameter(th.Tensor(max_relative_position * 2 + 1, num_units))
-        nn.init.xavier_uniform_(self.embeddings_table)
-
-    def forward(self, length_q, length_k):
-        device = self.embeddings_table.device
-        range_vec_q = th.arange(length_q, device=device)
-        range_vec_k = th.arange(length_k, device=device)
-        distance_mat = range_vec_k[None, :] - range_vec_q[:, None]
-        distance_mat_clipped = th.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
-        final_mat = distance_mat_clipped + self.max_relative_position
-        final_mat = final_mat.long()
-        embeddings = self.embeddings_table[final_mat]
-        return embeddings
-
-
-# ---------------------------------------------------------------------------------------------------
-class TemporalCrossAttention(nn.Module):
-    def __init__(self, 
-        query_dim, 
-        context_dim=None, 
-        heads=8, 
-        dim_head=64, 
-        dropout=0.,
-        use_relative_position=False,    # whether use relative positional representation in temporal attention.
-        temporal_length=None,           # relative positional representation
-        **kwargs,
-    ):
-        super().__init__()
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-        self.context_dim = context_dim
-        self.scale = dim_head ** -0.5
-        self.heads = heads
-        self.temporal_length = temporal_length
-        self.use_relative_position = use_relative_position
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim),
-            nn.Dropout(dropout)
-        )
-
-        if use_relative_position:
-            assert(temporal_length is not None)
-            self.relative_position_k = RelativePosition(num_units=dim_head, max_relative_position=temporal_length)
-            self.relative_position_v = RelativePosition(num_units=dim_head, max_relative_position=temporal_length)
-
-        nn.init.constant_(self.to_q.weight, 0)
-        nn.init.constant_(self.to_k.weight, 0)
-        nn.init.constant_(self.to_v.weight, 0)
-        nn.init.constant_(self.to_out[0].weight, 0)
-        nn.init.constant_(self.to_out[0].bias, 0)
-
-    def forward(self, x, context=None, mask=None):
-        nh = self.heads
-        out = x
-
-        # cal qkv
-        q = self.to_q(out)
-        context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=nh), (q, k, v))
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
-
-        # relative positional embedding
-        if self.use_relative_position:
-            len_q, len_k, len_v = q.shape[1], k.shape[1], v.shape[1]
-            k2 = self.relative_position_k(len_q, len_k) 
-            sim2 = einsum('b t d, t s d -> b t s', q, k2) * self.scale
-            sim += sim2
-        
-        # mask attention
-        if mask is not None:
-            max_neg_value = -1e9
-            sim = sim + (1-mask.float()) * max_neg_value # 1=masking,0=no masking
-        
-        # attend to values
-        attn = sim.softmax(dim=-1)
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        
-        # relative positional embedding
-        if self.use_relative_position:
-            v2 = self.relative_position_v(len_q, len_v)
-            out2 = einsum('b t s, t s d -> b t d', attn, v2)
-            out += out2
-        
-        # merge head
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=nh)
-        return self.to_out(out)
-
-
-# ---------------------------------------------------------------------------------------------------
-class CrossAttention(nn.Module):
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.,
-                 **kwargs,):
-        super().__init__()
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-        
-        self.scale = dim_head ** -0.5
-        self.heads = heads
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim),
-            nn.Dropout(dropout)
-        )
-
-    def forward(self, x, context=None, mask=None):
-        h = self.heads
-        b = x.shape[0]
-        
-        q = self.to_q(x)
-        context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
-        
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
-
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
-
-        if exists(mask):
-            mask = rearrange(mask, 'b ... -> b (...)')
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
-            sim.masked_fill_(~mask, max_neg_value)
-
-        attn = sim.softmax(dim=-1)
-
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
-        return self.to_out(out)
-    
-
-# ---------------------------------------------------------------------------------------------------
-class MemoryEfficientCrossAttention(nn.Module):
-    """https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
-    """
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0,
-                 **kwargs,):
-        super().__init__()
-        print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
-              f"{heads} heads."
-        )
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-
-        self.heads = heads
-        self.dim_head = dim_head
-
-        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
-        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
-        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
-
-        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
-        self.attention_op: Optional[Any] = None
-
-    def forward(self, x, context=None, mask=None):
-        q = self.to_q(x)
-        context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
-
-        b, _, _ = q.shape
-        q, k, v = map(
-            lambda t: t.unsqueeze(3)
-            .reshape(b, t.shape[1], self.heads, self.dim_head)
-            .permute(0, 2, 1, 3)
-            .reshape(b * self.heads, t.shape[1], self.dim_head)
-            .contiguous(),
-            (q, k, v),
-        )
-        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
-
-        if exists(mask):
-            raise NotImplementedError
-        out = (
-            out.unsqueeze(0)
-            .reshape(b, self.heads, out.shape[1], self.dim_head)
-            .permute(0, 2, 1, 3)
-            .reshape(b, out.shape[1], self.heads * self.dim_head)
-        )
-        return self.to_out(out)
-
-
-# ---------------------------------------------------------------------------------------------------
-class BasicTransformerBlockST(nn.Module):
-    """
-    if no context is given to forward function, cross-attention defaults to self-attention
-    """
-    def __init__(self, 
-        # Spatial
-        dim, 
-        n_heads, 
-        d_head, 
-        dropout=0., 
-        context_dim=None, 
-        gated_ff=True, 
-        checkpoint=True,
-        # Temporal
-        temporal_length=None,   
-        use_relative_position=True,
-        **kwargs,
-    ):
-        super().__init__()
-
-        # spatial self attention (if context_dim is None) and spatial cross attention
-        if XFORMERS_IS_AVAILBLE:
-            self.attn1 = MemoryEfficientCrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, **kwargs,)
-            self.attn2 = MemoryEfficientCrossAttention(query_dim=dim, context_dim=context_dim,
-                                    heads=n_heads, dim_head=d_head, dropout=dropout, **kwargs,)
-        else:
-            self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, **kwargs,)
-            self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
-                                    heads=n_heads, dim_head=d_head, dropout=dropout, **kwargs,)
-        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
-        
-        self.norm1 = nn.LayerNorm(dim)
-        self.norm2 = nn.LayerNorm(dim)
-        self.norm3 = nn.LayerNorm(dim)
-        self.checkpoint = checkpoint
-        
-        # Temporal attention
-        self.attn1_tmp = TemporalCrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
-                                                temporal_length=temporal_length,
-                                                use_relative_position=use_relative_position,
-                                                **kwargs,
-        )
-        self.attn2_tmp = TemporalCrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
-                                                # cross attn
-                                                context_dim=None,
-                                                # temporal attn
-                                                temporal_length=temporal_length,
-                                                use_relative_position=use_relative_position,
-                                                **kwargs,
-        )
-        self.norm4 = nn.LayerNorm(dim)
-        self.norm5 = nn.LayerNorm(dim)
-        
-    def forward(self, x, context=None, **kwargs):
-        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
-        
-    def _forward(self, x, context=None, mask=None,):
-        assert(x.dim() == 5), f"x shape = {x.shape}"
-        b, c, t, h, w = x.shape
-        
-        # spatial self attention
-        x = rearrange(x, 'b c t h w -> (b t) (h w) c')
-        x = self.attn1(self.norm1(x)) + x
-        x = rearrange(x, '(b t) (h w) c -> b c t h w', b=b,h=h)
-        
-        # temporal self attention
-        x = rearrange(x, 'b c t h w -> (b h w) t c')
-        x = self.attn1_tmp(self.norm4(x), mask=mask) + x
-        x = rearrange(x, '(b h w) t c -> b c t h w', b=b,h=h,w=w) # 3d -> 5d
-        
-        # spatial cross attention
-        x = rearrange(x, 'b c t h w -> (b t) (h w) c')
-        if context is not None:
-            context_ = []
-            for i in range(context.shape[0]):
-                context_.append(context[i].unsqueeze(0).repeat(t, 1, 1))
-            context_ = torch.cat(context_,dim=0)
-        else:
-            context_ = None
-        x = self.attn2(self.norm2(x), context=context_) + x
-        x = rearrange(x, '(b t) (h w) c -> b c t h w', b=b,h=h)
-
-        # temporal cross attention
-        x = rearrange(x, 'b c t h w -> (b h w) t c')
-        x = self.attn2_tmp(self.norm5(x), context=None, mask=mask) + x
-
-        # feedforward
-        x = self.ff(self.norm3(x)) + x
-        x = rearrange(x, '(b h w) t c -> b c t h w', b=b,h=h,w=w) # 3d -> 5d
-        
-        return x
-
-
-# ---------------------------------------------------------------------------------------------------
-class SpatialTemporalTransformer(nn.Module):
-    """
-    Transformer block for video-like data (5D tensor).
-    First, project the input (aka embedding) with NO reshape.
-    Then apply standard transformer action.
-    The 5D -> 3D reshape operation will be done in the specific attention module.
-    """
-    def __init__(
-        self,
-        in_channels, n_heads, d_head,
-        depth=1, dropout=0.,
-        context_dim=None,
-        # Temporal
-        temporal_length=None,
-        use_relative_position=True,
-        **kwargs,
-        ):
-        super().__init__()
-
-        self.in_channels = in_channels
-        inner_dim = n_heads * d_head
-
-        self.norm = Normalize(in_channels)
-        self.proj_in = nn.Conv3d(in_channels,
-                                 inner_dim,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-
-        self.transformer_blocks = nn.ModuleList(
-            [BasicTransformerBlockST(
-                inner_dim, n_heads, d_head, dropout=dropout,
-                # cross attn
-                context_dim=context_dim,
-                # temporal attn
-                temporal_length=temporal_length,   
-                use_relative_position=use_relative_position,
-                **kwargs
-                ) for d in range(depth)]
-        )
-
-        self.proj_out = zero_module(nn.Conv3d(inner_dim,
-                                              in_channels,
-                                              kernel_size=1,
-                                              stride=1,
-                                              padding=0))
-        
-    def forward(self, x, context=None, **kwargs):
-        
-        assert(x.dim() == 5), f"x shape = {x.shape}"
-        x_in = x
-        
-        x = self.norm(x)
-        x = self.proj_in(x)
-        
-        for block in self.transformer_blocks:
-            x = block(x, context=context, **kwargs)
-        
-        x = self.proj_out(x)
-
-        return x + x_in

lvdm/models/modules/condition_modules.py

@@ -1,40 +0,0 @@
-import torch.nn as nn
-from transformers import logging
-from transformers import CLIPTokenizer, CLIPTextModel
-logging.set_verbosity_error()
-
-
-class AbstractEncoder(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def encode(self, *args, **kwargs):
-        raise NotImplementedError
-
-
-class FrozenCLIPEmbedder(AbstractEncoder):
-    """Uses the CLIP transformer encoder for text (from huggingface)"""
-    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77):
-        super().__init__()
-        self.tokenizer = CLIPTokenizer.from_pretrained(version)
-        self.transformer = CLIPTextModel.from_pretrained(version)
-        self.device = device
-        self.max_length = max_length
-        self.freeze()
-
-    def freeze(self):
-        self.transformer = self.transformer.eval()
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def forward(self, text):
-        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
-                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
-        tokens = batch_encoding["input_ids"].to(self.device)
-        outputs = self.transformer(input_ids=tokens)
-
-        z = outputs.last_hidden_state
-        return z
-
-    def encode(self, text):
-        return self(text)

lvdm/models/modules/lora.py

@@ -1,1251 +0,0 @@
-import json
-from itertools import groupby
-from typing import Dict, List, Optional, Set, Tuple, Type, Union
-
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-# try:
-#     from safetensors.torch import safe_open
-#     from safetensors.torch import save_file as safe_save
-
-#     safetensors_available = True
-# except ImportError:
-#     from .safe_open import safe_open
-
-#     def safe_save(
-#         tensors: Dict[str, torch.Tensor],
-#         filename: str,
-#         metadata: Optional[Dict[str, str]] = None,
-#     ) -> None:
-#         raise EnvironmentError(
-#             "Saving safetensors requires the safetensors library. Please install with pip or similar."
-#         )
-
-#     safetensors_available = False
-
-
-class LoraInjectedLinear(nn.Module):
-    def __init__(
-        self, in_features, out_features, bias=False, r=4, dropout_p=0.1, scale=1.0
-    ):
-        super().__init__()
-
-        if r > min(in_features, out_features):
-            raise ValueError(
-                f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
-            )
-        self.r = r
-        self.linear = nn.Linear(in_features, out_features, bias)
-        self.lora_down = nn.Linear(in_features, r, bias=False)
-        self.dropout = nn.Dropout(dropout_p)
-        self.lora_up = nn.Linear(r, out_features, bias=False)
-        self.scale = scale
-        self.selector = nn.Identity()
-
-        nn.init.normal_(self.lora_down.weight, std=1 / r)
-        nn.init.zeros_(self.lora_up.weight)
-
-    def forward(self, input):
-        return (
-            self.linear(input)
-            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
-            * self.scale
-        )
-
-    def realize_as_lora(self):
-        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
-
-    def set_selector_from_diag(self, diag: torch.Tensor):
-        # diag is a 1D tensor of size (r,)
-        assert diag.shape == (self.r,)
-        self.selector = nn.Linear(self.r, self.r, bias=False)
-        self.selector.weight.data = torch.diag(diag)
-        self.selector.weight.data = self.selector.weight.data.to(
-            self.lora_up.weight.device
-        ).to(self.lora_up.weight.dtype)
-
-
-class LoraInjectedConv2d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        groups: int = 1,
-        bias: bool = True,
-        r: int = 4,
-        dropout_p: float = 0.1,
-        scale: float = 1.0,
-    ):
-        super().__init__()
-        if r > min(in_channels, out_channels):
-            raise ValueError(
-                f"LoRA rank {r} must be less or equal than {min(in_channels, out_channels)}"
-            )
-        self.r = r
-        self.conv = nn.Conv2d(
-            in_channels=in_channels,
-            out_channels=out_channels,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-            bias=bias,
-        )
-
-        self.lora_down = nn.Conv2d(
-            in_channels=in_channels,
-            out_channels=r,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-            bias=False,
-        )
-        self.dropout = nn.Dropout(dropout_p)
-        self.lora_up = nn.Conv2d(
-            in_channels=r,
-            out_channels=out_channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=False,
-        )
-        self.selector = nn.Identity()
-        self.scale = scale
-
-        nn.init.normal_(self.lora_down.weight, std=1 / r)
-        nn.init.zeros_(self.lora_up.weight)
-
-    def forward(self, input):
-        return (
-            self.conv(input)
-            + self.dropout(self.lora_up(self.selector(self.lora_down(input))))
-            * self.scale
-        )
-
-    def realize_as_lora(self):
-        return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
-
-    def set_selector_from_diag(self, diag: torch.Tensor):
-        # diag is a 1D tensor of size (r,)
-        assert diag.shape == (self.r,)
-        self.selector = nn.Conv2d(
-            in_channels=self.r,
-            out_channels=self.r,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=False,
-        )
-        self.selector.weight.data = torch.diag(diag)
-
-        # same device + dtype as lora_up
-        self.selector.weight.data = self.selector.weight.data.to(
-            self.lora_up.weight.device
-        ).to(self.lora_up.weight.dtype)
-
-
-UNET_DEFAULT_TARGET_REPLACE = {"MemoryEfficientCrossAttention","CrossAttention", "Attention", "GEGLU"}
-
-UNET_EXTENDED_TARGET_REPLACE = {"TimestepEmbedSequential","SpatialTemporalTransformer", "MemoryEfficientCrossAttention","CrossAttention", "Attention", "GEGLU"}
-
-TEXT_ENCODER_DEFAULT_TARGET_REPLACE = {"CLIPAttention"}
-
-TEXT_ENCODER_EXTENDED_TARGET_REPLACE = {"CLIPMLP","CLIPAttention"}
-
-DEFAULT_TARGET_REPLACE = UNET_DEFAULT_TARGET_REPLACE
-
-EMBED_FLAG = "<embed>"
-
-
-def _find_children(
-    model,
-    search_class: List[Type[nn.Module]] = [nn.Linear],
-):
-    """
-    Find all modules of a certain class (or union of classes).
-
-    Returns all matching modules, along with the parent of those moduless and the
-    names they are referenced by.
-    """
-    # For each target find every linear_class module that isn't a child of a LoraInjectedLinear
-    for parent in model.modules():
-        for name, module in parent.named_children():
-            if any([isinstance(module, _class) for _class in search_class]):
-                yield parent, name, module
-
-
-def _find_modules_v2(
-    model,
-    ancestor_class: Optional[Set[str]] = None,
-    search_class: List[Type[nn.Module]] = [nn.Linear],
-    exclude_children_of: Optional[List[Type[nn.Module]]] = [
-        LoraInjectedLinear,
-        LoraInjectedConv2d,
-    ],
-):
-    """
-    Find all modules of a certain class (or union of classes) that are direct or
-    indirect descendants of other modules of a certain class (or union of classes).
-
-    Returns all matching modules, along with the parent of those moduless and the
-    names they are referenced by.
-    """
-
-    # Get the targets we should replace all linears under
-    if type(ancestor_class) is not set:
-        ancestor_class = set(ancestor_class)
-        print(ancestor_class)
-    if ancestor_class is not None:
-        ancestors = (
-            module
-            for module in model.modules()
-            if module.__class__.__name__ in ancestor_class
-        )
-    else:
-        # this, incase you want to naively iterate over all modules.
-        ancestors = [module for module in model.modules()]
-
-    # For each target find every linear_class module that isn't a child of a LoraInjectedLinear
-    for ancestor in ancestors:
-        for fullname, module in ancestor.named_children():
-            if any([isinstance(module, _class) for _class in search_class]):
-                # Find the direct parent if this is a descendant, not a child, of target
-                *path, name = fullname.split(".")
-                parent = ancestor
-                while path:
-                    parent = parent.get_submodule(path.pop(0))
-                # Skip this linear if it's a child of a LoraInjectedLinear
-                if exclude_children_of and any(
-                    [isinstance(parent, _class) for _class in exclude_children_of]
-                ):
-                    continue
-                # Otherwise, yield it
-                yield parent, name, module
-
-
-def _find_modules_old(
-    model,
-    ancestor_class: Set[str] = DEFAULT_TARGET_REPLACE,
-    search_class: List[Type[nn.Module]] = [nn.Linear],
-    exclude_children_of: Optional[List[Type[nn.Module]]] = [LoraInjectedLinear],
-):
-    ret = []
-    for _module in model.modules():
-        if _module.__class__.__name__ in ancestor_class:
-
-            for name, _child_module in _module.named_children():
-                if _child_module.__class__ in search_class:
-                    ret.append((_module, name, _child_module))
-    print(ret)
-    return ret
-
-
-_find_modules = _find_modules_v2
-
-
-def inject_trainable_lora(
-    model: nn.Module,
-    target_replace_module: Set[str] = DEFAULT_TARGET_REPLACE,
-    r: int = 4,
-    loras=None,  # path to lora .pt
-    verbose: bool = False,
-    dropout_p: float = 0.0,
-    scale: float = 1.0,
-):
-    """
-    inject lora into model, and returns lora parameter groups.
-    """
-
-    require_grad_params = []
-    names = []
-
-    if loras != None:
-        loras = torch.load(loras)
-
-    for _module, name, _child_module in _find_modules(
-        model, target_replace_module, search_class=[nn.Linear]
-    ):
-        weight = _child_module.weight
-        bias = _child_module.bias
-        if verbose:
-            print("LoRA Injection : injecting lora into ", name)
-            print("LoRA Injection : weight shape", weight.shape)
-        _tmp = LoraInjectedLinear(
-            _child_module.in_features,
-            _child_module.out_features,
-            _child_module.bias is not None,
-            r=r,
-            dropout_p=dropout_p,
-            scale=scale,
-        )
-        _tmp.linear.weight = weight
-        if bias is not None:
-            _tmp.linear.bias = bias
-
-        # switch the module
-        _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
-        _module._modules[name] = _tmp
-
-        require_grad_params.append(_module._modules[name].lora_up.parameters())
-        require_grad_params.append(_module._modules[name].lora_down.parameters())
-
-        if loras != None:
-            _module._modules[name].lora_up.weight = loras.pop(0)
-            _module._modules[name].lora_down.weight = loras.pop(0)
-
-        _module._modules[name].lora_up.weight.requires_grad = True
-        _module._modules[name].lora_down.weight.requires_grad = True
-        names.append(name)
-
-    return require_grad_params, names
-
-
-def inject_trainable_lora_extended(
-    model: nn.Module,
-    target_replace_module: Set[str] = UNET_EXTENDED_TARGET_REPLACE,
-    r: int = 4,
-    loras=None,  # path to lora .pt
-):
-    """
-    inject lora into model, and returns lora parameter groups.
-    """
-
-    require_grad_params = []
-    names = []
-
-    if loras != None:
-        loras = torch.load(loras)
-
-    for _module, name, _child_module in _find_modules(
-        model, target_replace_module, search_class=[nn.Linear, nn.Conv2d]
-    ):
-        if _child_module.__class__ == nn.Linear:
-            weight = _child_module.weight
-            bias = _child_module.bias
-            _tmp = LoraInjectedLinear(
-                _child_module.in_features,
-                _child_module.out_features,
-                _child_module.bias is not None,
-                r=r,
-            )
-            _tmp.linear.weight = weight
-            if bias is not None:
-                _tmp.linear.bias = bias
-        elif _child_module.__class__ == nn.Conv2d:
-            weight = _child_module.weight
-            bias = _child_module.bias
-            _tmp = LoraInjectedConv2d(
-                _child_module.in_channels,
-                _child_module.out_channels,
-                _child_module.kernel_size,
-                _child_module.stride,
-                _child_module.padding,
-                _child_module.dilation,
-                _child_module.groups,
-                _child_module.bias is not None,
-                r=r,
-            )
-
-            _tmp.conv.weight = weight
-            if bias is not None:
-                _tmp.conv.bias = bias
-
-        # switch the module
-        _tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
-        if bias is not None:
-            _tmp.to(_child_module.bias.device).to(_child_module.bias.dtype)
-
-        _module._modules[name] = _tmp
-
-        require_grad_params.append(_module._modules[name].lora_up.parameters())
-        require_grad_params.append(_module._modules[name].lora_down.parameters())
-
-        if loras != None:
-            _module._modules[name].lora_up.weight = loras.pop(0)
-            _module._modules[name].lora_down.weight = loras.pop(0)
-
-        _module._modules[name].lora_up.weight.requires_grad = True
-        _module._modules[name].lora_down.weight.requires_grad = True
-        names.append(name)
-
-    return require_grad_params, names
-
-
-def extract_lora_ups_down(model, target_replace_module=DEFAULT_TARGET_REPLACE):
-
-    loras = []
-
-    for _m, _n, _child_module in _find_modules(
-        model,
-        target_replace_module,
-        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
-    ):
-        loras.append((_child_module.lora_up, _child_module.lora_down))
-
-    if len(loras) == 0:
-        raise ValueError("No lora injected.")
-
-    return loras
-
-
-def extract_lora_as_tensor(
-    model, target_replace_module=DEFAULT_TARGET_REPLACE, as_fp16=True
-):
-
-    loras = []
-
-    for _m, _n, _child_module in _find_modules(
-        model,
-        target_replace_module,
-        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
-    ):
-        up, down = _child_module.realize_as_lora()
-        if as_fp16:
-            up = up.to(torch.float16)
-            down = down.to(torch.float16)
-
-        loras.append((up, down))
-
-    if len(loras) == 0:
-        raise ValueError("No lora injected.")
-
-    return loras
-
-
-def save_lora_weight(
-    model,
-    path="./lora.pt",
-    target_replace_module=DEFAULT_TARGET_REPLACE,
-):
-    weights = []
-    for _up, _down in extract_lora_ups_down(
-        model, target_replace_module=target_replace_module
-    ):
-        weights.append(_up.weight.to("cpu").to(torch.float16))
-        weights.append(_down.weight.to("cpu").to(torch.float16))
-
-    torch.save(weights, path)
-
-
-def save_lora_as_json(model, path="./lora.json"):
-    weights = []
-    for _up, _down in extract_lora_ups_down(model):
-        weights.append(_up.weight.detach().cpu().numpy().tolist())
-        weights.append(_down.weight.detach().cpu().numpy().tolist())
-
-    import json
-
-    with open(path, "w") as f:
-        json.dump(weights, f)
-
-
-def save_safeloras_with_embeds(
-    modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
-    embeds: Dict[str, torch.Tensor] = {},
-    outpath="./lora.safetensors",
-):
-    """
-    Saves the Lora from multiple modules in a single safetensor file.
-
-    modelmap is a dictionary of {
-        "module name": (module, target_replace_module)
-    }
-    """
-    weights = {}
-    metadata = {}
-
-    for name, (model, target_replace_module) in modelmap.items():
-        metadata[name] = json.dumps(list(target_replace_module))
-
-        for i, (_up, _down) in enumerate(
-            extract_lora_as_tensor(model, target_replace_module)
-        ):
-            rank = _down.shape[0]
-
-            metadata[f"{name}:{i}:rank"] = str(rank)
-            weights[f"{name}:{i}:up"] = _up
-            weights[f"{name}:{i}:down"] = _down
-
-    for token, tensor in embeds.items():
-        metadata[token] = EMBED_FLAG
-        weights[token] = tensor
-
-    print(f"Saving weights to {outpath}")
-    safe_save(weights, outpath, metadata)
-
-
-def save_safeloras(
-    modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
-    outpath="./lora.safetensors",
-):
-    return save_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
-
-
-def convert_loras_to_safeloras_with_embeds(
-    modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
-    embeds: Dict[str, torch.Tensor] = {},
-    outpath="./lora.safetensors",
-):
-    """
-    Converts the Lora from multiple pytorch .pt files into a single safetensor file.
-
-    modelmap is a dictionary of {
-        "module name": (pytorch_model_path, target_replace_module, rank)
-    }
-    """
-
-    weights = {}
-    metadata = {}
-
-    for name, (path, target_replace_module, r) in modelmap.items():
-        metadata[name] = json.dumps(list(target_replace_module))
-
-        lora = torch.load(path)
-        for i, weight in enumerate(lora):
-            is_up = i % 2 == 0
-            i = i // 2
-
-            if is_up:
-                metadata[f"{name}:{i}:rank"] = str(r)
-                weights[f"{name}:{i}:up"] = weight
-            else:
-                weights[f"{name}:{i}:down"] = weight
-
-    for token, tensor in embeds.items():
-        metadata[token] = EMBED_FLAG
-        weights[token] = tensor
-
-    print(f"Saving weights to {outpath}")
-    safe_save(weights, outpath, metadata)
-
-
-def convert_loras_to_safeloras(
-    modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
-    outpath="./lora.safetensors",
-):
-    convert_loras_to_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
-
-
-def parse_safeloras(
-    safeloras,
-) -> Dict[str, Tuple[List[nn.parameter.Parameter], List[int], List[str]]]:
-    """
-    Converts a loaded safetensor file that contains a set of module Loras
-    into Parameters and other information
-
-    Output is a dictionary of {
-        "module name": (
-            [list of weights],
-            [list of ranks],
-            target_replacement_modules
-        )
-    }
-    """
-    loras = {}
-    metadata = safeloras.metadata()
-
-    get_name = lambda k: k.split(":")[0]
-
-    keys = list(safeloras.keys())
-    keys.sort(key=get_name)
-
-    for name, module_keys in groupby(keys, get_name):
-        info = metadata.get(name)
-
-        if not info:
-            raise ValueError(
-                f"Tensor {name} has no metadata - is this a Lora safetensor?"
-            )
-
-        # Skip Textual Inversion embeds
-        if info == EMBED_FLAG:
-            continue
-
-        # Handle Loras
-        # Extract the targets
-        target = json.loads(info)
-
-        # Build the result lists - Python needs us to preallocate lists to insert into them
-        module_keys = list(module_keys)
-        ranks = [4] * (len(module_keys) // 2)
-        weights = [None] * len(module_keys)
-
-        for key in module_keys:
-            # Split the model name and index out of the key
-            _, idx, direction = key.split(":")
-            idx = int(idx)
-
-            # Add the rank
-            ranks[idx] = int(metadata[f"{name}:{idx}:rank"])
-
-            # Insert the weight into the list
-            idx = idx * 2 + (1 if direction == "down" else 0)
-            weights[idx] = nn.parameter.Parameter(safeloras.get_tensor(key))
-
-        loras[name] = (weights, ranks, target)
-
-    return loras
-
-
-def parse_safeloras_embeds(
-    safeloras,
-) -> Dict[str, torch.Tensor]:
-    """
-    Converts a loaded safetensor file that contains Textual Inversion embeds into
-    a dictionary of embed_token: Tensor
-    """
-    embeds = {}
-    metadata = safeloras.metadata()
-
-    for key in safeloras.keys():
-        # Only handle Textual Inversion embeds
-        meta = metadata.get(key)
-        if not meta or meta != EMBED_FLAG:
-            continue
-
-        embeds[key] = safeloras.get_tensor(key)
-
-    return embeds
-
-def net_load_lora(net, checkpoint_path, alpha=1.0, remove=False):
-    visited=[]
-    state_dict = torch.load(checkpoint_path)
-    for k, v in state_dict.items():
-        state_dict[k] = v.to(net.device)
-
-    for key in state_dict:
-        if ".alpha" in key or key in visited:
-            continue
-        layer_infos = key.split(".")[:-2] # remove lora_up and down weight
-        curr_layer = net
-        # find the target layer
-        temp_name = layer_infos.pop(0)
-        while len(layer_infos) > -1:
-            curr_layer = curr_layer.__getattr__(temp_name)
-            if len(layer_infos) > 0:
-                temp_name = layer_infos.pop(0)
-            elif len(layer_infos) == 0:
-                break
-        if curr_layer.__class__ not in [nn.Linear, nn.Conv2d]:
-            print('missing param at:', key)
-            continue
-        pair_keys = []
-        if "lora_down" in key:
-            pair_keys.append(key.replace("lora_down", "lora_up"))
-            pair_keys.append(key)
-        else:
-            pair_keys.append(key)
-            pair_keys.append(key.replace("lora_up", "lora_down"))
-
-        # update weight
-        if len(state_dict[pair_keys[0]].shape) == 4:
-            # for conv
-            weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
-            weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
-            if remove:
-                curr_layer.weight.data -= alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
-            else:
-                curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
-        else:
-            # for linear
-            weight_up = state_dict[pair_keys[0]].to(torch.float32)
-            weight_down = state_dict[pair_keys[1]].to(torch.float32)
-            if remove:
-                curr_layer.weight.data -= alpha * torch.mm(weight_up, weight_down)
-            else:
-                curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)
-
-        # update visited list
-        for item in pair_keys:
-            visited.append(item)
-    print('load_weight_num:',len(visited))
-    return 
-
-def change_lora(model, inject_lora=False, lora_scale=1.0, lora_path='', last_time_lora='', last_time_lora_scale=1.0):
-    # remove lora
-    if last_time_lora != '':
-        net_load_lora(model, last_time_lora, alpha=last_time_lora_scale, remove=True)
-    # add new lora
-    if inject_lora:
-        net_load_lora(model, lora_path, alpha=lora_scale)
-
-
-def net_load_lora_v2(net, checkpoint_path, alpha=1.0, remove=False, origin_weight=None):
-    visited=[]
-    state_dict = torch.load(checkpoint_path)
-    for k, v in state_dict.items():
-        state_dict[k] = v.to(net.device)
-
-    for key in state_dict:
-        if ".alpha" in key or key in visited:
-            continue
-        layer_infos = key.split(".")[:-2] # remove lora_up and down weight
-        curr_layer = net
-        # find the target layer
-        temp_name = layer_infos.pop(0)
-        while len(layer_infos) > -1:
-            curr_layer = curr_layer.__getattr__(temp_name)
-            if len(layer_infos) > 0:
-                temp_name = layer_infos.pop(0)
-            elif len(layer_infos) == 0:
-                break
-        if curr_layer.__class__ not in [nn.Linear, nn.Conv2d]:
-            print('missing param at:', key)
-            continue
-        pair_keys = []
-        if "lora_down" in key:
-            pair_keys.append(key.replace("lora_down", "lora_up"))
-            pair_keys.append(key)
-        else:
-            pair_keys.append(key)
-            pair_keys.append(key.replace("lora_up", "lora_down"))
-
-        # storage weight
-        if origin_weight is None:
-            origin_weight = dict()
-            storage_key = key.replace("lora_down", "lora").replace("lora_up", "lora")
-            origin_weight[storage_key] = curr_layer.weight.data.clone()
-        else:
-            storage_key = key.replace("lora_down", "lora").replace("lora_up", "lora")
-            if storage_key not in origin_weight.keys():
-                origin_weight[storage_key] = curr_layer.weight.data.clone()
-
-
-        # update 
-        if len(state_dict[pair_keys[0]].shape) == 4:
-            # for conv
-            if remove:
-                curr_layer.weight.data = origin_weight[storage_key].clone()
-            else:
-                weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32)
-                weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32)
-                curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
-        else:
-            # for linear
-            if remove:
-                curr_layer.weight.data = origin_weight[storage_key].clone()
-            else:
-                weight_up = state_dict[pair_keys[0]].to(torch.float32)
-                weight_down = state_dict[pair_keys[1]].to(torch.float32)
-                curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down)
-
-        # update visited list
-        for item in pair_keys:
-            visited.append(item)
-    print('load_weight_num:',len(visited))
-    return origin_weight
-
-def change_lora_v2(model, inject_lora=False, lora_scale=1.0, lora_path='', last_time_lora='', last_time_lora_scale=1.0, origin_weight=None):
-    # remove lora
-    if last_time_lora != '':
-        origin_weight = net_load_lora_v2(model, last_time_lora, alpha=last_time_lora_scale, remove=True, origin_weight=origin_weight)
-    # add new lora
-    if inject_lora:
-        origin_weight = net_load_lora_v2(model, lora_path, alpha=lora_scale, origin_weight=origin_weight)
-    return origin_weight
-
-
-
-
-
-def load_safeloras(path, device="cpu"):
-    safeloras = safe_open(path, framework="pt", device=device)
-    return parse_safeloras(safeloras)
-
-
-def load_safeloras_embeds(path, device="cpu"):
-    safeloras = safe_open(path, framework="pt", device=device)
-    return parse_safeloras_embeds(safeloras)
-
-
-def load_safeloras_both(path, device="cpu"):
-    safeloras = safe_open(path, framework="pt", device=device)
-    return parse_safeloras(safeloras), parse_safeloras_embeds(safeloras)
-
-
-def collapse_lora(model, alpha=1.0):
-
-    for _module, name, _child_module in _find_modules(
-        model,
-        UNET_EXTENDED_TARGET_REPLACE | TEXT_ENCODER_EXTENDED_TARGET_REPLACE,
-        search_class=[LoraInjectedLinear, LoraInjectedConv2d],
-    ):
-
-        if isinstance(_child_module, LoraInjectedLinear):
-            print("Collapsing Lin Lora in", name)
-
-            _child_module.linear.weight = nn.Parameter(
-                _child_module.linear.weight.data
-                + alpha
-                * (
-                    _child_module.lora_up.weight.data
-                    @ _child_module.lora_down.weight.data
-                )
-                .type(_child_module.linear.weight.dtype)
-                .to(_child_module.linear.weight.device)
-            )
-
-        else:
-            print("Collapsing Conv Lora in", name)
-            _child_module.conv.weight = nn.Parameter(
-                _child_module.conv.weight.data
-                + alpha
-                * (
-                    _child_module.lora_up.weight.data.flatten(start_dim=1)
-                    @ _child_module.lora_down.weight.data.flatten(start_dim=1)
-                )
-                .reshape(_child_module.conv.weight.data.shape)
-                .type(_child_module.conv.weight.dtype)
-                .to(_child_module.conv.weight.device)
-            )
-
-
-def monkeypatch_or_replace_lora(
-    model,
-    loras,
-    target_replace_module=DEFAULT_TARGET_REPLACE,
-    r: Union[int, List[int]] = 4,
-):
-    for _module, name, _child_module in _find_modules(
-        model, target_replace_module, search_class=[nn.Linear, LoraInjectedLinear]
-    ):
-        _source = (
-            _child_module.linear
-            if isinstance(_child_module, LoraInjectedLinear)
-            else _child_module
-        )
-
-        weight = _source.weight
-        bias = _source.bias
-        _tmp = LoraInjectedLinear(
-            _source.in_features,
-            _source.out_features,
-            _source.bias is not None,
-            r=r.pop(0) if isinstance(r, list) else r,
-        )
-        _tmp.linear.weight = weight
-
-        if bias is not None:
-            _tmp.linear.bias = bias
-
-        # switch the module
-        _module._modules[name] = _tmp
-
-        up_weight = loras.pop(0)
-        down_weight = loras.pop(0)
-
-        _module._modules[name].lora_up.weight = nn.Parameter(
-            up_weight.type(weight.dtype)
-        )
-        _module._modules[name].lora_down.weight = nn.Parameter(
-            down_weight.type(weight.dtype)
-        )
-
-        _module._modules[name].to(weight.device)
-
-
-def monkeypatch_or_replace_lora_extended(
-    model,
-    loras,
-    target_replace_module=DEFAULT_TARGET_REPLACE,
-    r: Union[int, List[int]] = 4,
-):
-    for _module, name, _child_module in _find_modules(
-        model,
-        target_replace_module,
-        search_class=[nn.Linear, LoraInjectedLinear, nn.Conv2d, LoraInjectedConv2d],
-    ):
-
-        if (_child_module.__class__ == nn.Linear) or (
-            _child_module.__class__ == LoraInjectedLinear
-        ):
-            if len(loras[0].shape) != 2:
-                continue
-
-            _source = (
-                _child_module.linear
-                if isinstance(_child_module, LoraInjectedLinear)
-                else _child_module
-            )
-
-            weight = _source.weight
-            bias = _source.bias
-            _tmp = LoraInjectedLinear(
-                _source.in_features,
-                _source.out_features,
-                _source.bias is not None,
-                r=r.pop(0) if isinstance(r, list) else r,
-            )
-            _tmp.linear.weight = weight
-
-            if bias is not None:
-                _tmp.linear.bias = bias
-
-        elif (_child_module.__class__ == nn.Conv2d) or (
-            _child_module.__class__ == LoraInjectedConv2d
-        ):
-            if len(loras[0].shape) != 4:
-                continue
-            _source = (
-                _child_module.conv
-                if isinstance(_child_module, LoraInjectedConv2d)
-                else _child_module
-            )
-
-            weight = _source.weight
-            bias = _source.bias
-            _tmp = LoraInjectedConv2d(
-                _source.in_channels,
-                _source.out_channels,
-                _source.kernel_size,
-                _source.stride,
-                _source.padding,
-                _source.dilation,
-                _source.groups,
-                _source.bias is not None,
-                r=r.pop(0) if isinstance(r, list) else r,
-            )
-
-            _tmp.conv.weight = weight
-
-            if bias is not None:
-                _tmp.conv.bias = bias
-
-        # switch the module
-        _module._modules[name] = _tmp
-
-        up_weight = loras.pop(0)
-        down_weight = loras.pop(0)
-
-        _module._modules[name].lora_up.weight = nn.Parameter(
-            up_weight.type(weight.dtype)
-        )
-        _module._modules[name].lora_down.weight = nn.Parameter(
-            down_weight.type(weight.dtype)
-        )
-
-        _module._modules[name].to(weight.device)
-
-
-def monkeypatch_or_replace_safeloras(models, safeloras):
-    loras = parse_safeloras(safeloras)
-
-    for name, (lora, ranks, target) in loras.items():
-        model = getattr(models, name, None)
-
-        if not model:
-            print(f"No model provided for {name}, contained in Lora")
-            continue
-
-        monkeypatch_or_replace_lora_extended(model, lora, target, ranks)
-
-
-def monkeypatch_remove_lora(model):
-    for _module, name, _child_module in _find_modules(
-        model, search_class=[LoraInjectedLinear, LoraInjectedConv2d]
-    ):
-        if isinstance(_child_module, LoraInjectedLinear):
-            _source = _child_module.linear
-            weight, bias = _source.weight, _source.bias
-
-            _tmp = nn.Linear(
-                _source.in_features, _source.out_features, bias is not None
-            )
-
-            _tmp.weight = weight
-            if bias is not None:
-                _tmp.bias = bias
-
-        else:
-            _source = _child_module.conv
-            weight, bias = _source.weight, _source.bias
-
-            _tmp = nn.Conv2d(
-                in_channels=_source.in_channels,
-                out_channels=_source.out_channels,
-                kernel_size=_source.kernel_size,
-                stride=_source.stride,
-                padding=_source.padding,
-                dilation=_source.dilation,
-                groups=_source.groups,
-                bias=bias is not None,
-            )
-
-            _tmp.weight = weight
-            if bias is not None:
-                _tmp.bias = bias
-
-        _module._modules[name] = _tmp
-
-
-def monkeypatch_add_lora(
-    model,
-    loras,
-    target_replace_module=DEFAULT_TARGET_REPLACE,
-    alpha: float = 1.0,
-    beta: float = 1.0,
-):
-    for _module, name, _child_module in _find_modules(
-        model, target_replace_module, search_class=[LoraInjectedLinear]
-    ):
-        weight = _child_module.linear.weight
-
-        up_weight = loras.pop(0)
-        down_weight = loras.pop(0)
-
-        _module._modules[name].lora_up.weight = nn.Parameter(
-            up_weight.type(weight.dtype).to(weight.device) * alpha
-            + _module._modules[name].lora_up.weight.to(weight.device) * beta
-        )
-        _module._modules[name].lora_down.weight = nn.Parameter(
-            down_weight.type(weight.dtype).to(weight.device) * alpha
-            + _module._modules[name].lora_down.weight.to(weight.device) * beta
-        )
-
-        _module._modules[name].to(weight.device)
-
-
-def tune_lora_scale(model, alpha: float = 1.0):
-    for _module in model.modules():
-        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
-            _module.scale = alpha
-
-
-def set_lora_diag(model, diag: torch.Tensor):
-    for _module in model.modules():
-        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
-            _module.set_selector_from_diag(diag)
-
-
-def _text_lora_path(path: str) -> str:
-    assert path.endswith(".pt"), "Only .pt files are supported"
-    return ".".join(path.split(".")[:-1] + ["text_encoder", "pt"])
-
-
-def _ti_lora_path(path: str) -> str:
-    assert path.endswith(".pt"), "Only .pt files are supported"
-    return ".".join(path.split(".")[:-1] + ["ti", "pt"])
-
-
-def apply_learned_embed_in_clip(
-    learned_embeds,
-    text_encoder,
-    tokenizer,
-    token: Optional[Union[str, List[str]]] = None,
-    idempotent=False,
-):
-    if isinstance(token, str):
-        trained_tokens = [token]
-    elif isinstance(token, list):
-        assert len(learned_embeds.keys()) == len(
-            token
-        ), "The number of tokens and the number of embeds should be the same"
-        trained_tokens = token
-    else:
-        trained_tokens = list(learned_embeds.keys())
-
-    for token in trained_tokens:
-        print(token)
-        embeds = learned_embeds[token]
-
-        # cast to dtype of text_encoder
-        dtype = text_encoder.get_input_embeddings().weight.dtype
-        num_added_tokens = tokenizer.add_tokens(token)
-
-        i = 1
-        if not idempotent:
-            while num_added_tokens == 0:
-                print(f"The tokenizer already contains the token {token}.")
-                token = f"{token[:-1]}-{i}>"
-                print(f"Attempting to add the token {token}.")
-                num_added_tokens = tokenizer.add_tokens(token)
-                i += 1
-        elif num_added_tokens == 0 and idempotent:
-            print(f"The tokenizer already contains the token {token}.")
-            print(f"Replacing {token} embedding.")
-
-        # resize the token embeddings
-        text_encoder.resize_token_embeddings(len(tokenizer))
-
-        # get the id for the token and assign the embeds
-        token_id = tokenizer.convert_tokens_to_ids(token)
-        text_encoder.get_input_embeddings().weight.data[token_id] = embeds
-    return token
-
-
-def load_learned_embed_in_clip(
-    learned_embeds_path,
-    text_encoder,
-    tokenizer,
-    token: Optional[Union[str, List[str]]] = None,
-    idempotent=False,
-):
-    learned_embeds = torch.load(learned_embeds_path)
-    apply_learned_embed_in_clip(
-        learned_embeds, text_encoder, tokenizer, token, idempotent
-    )
-
-
-def patch_pipe(
-    pipe,
-    maybe_unet_path,
-    token: Optional[str] = None,
-    r: int = 4,
-    patch_unet=True,
-    patch_text=True,
-    patch_ti=True,
-    idempotent_token=True,
-    unet_target_replace_module=DEFAULT_TARGET_REPLACE,
-    text_target_replace_module=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
-):
-    if maybe_unet_path.endswith(".pt"):
-        # torch format
-
-        if maybe_unet_path.endswith(".ti.pt"):
-            unet_path = maybe_unet_path[:-6] + ".pt"
-        elif maybe_unet_path.endswith(".text_encoder.pt"):
-            unet_path = maybe_unet_path[:-16] + ".pt"
-        else:
-            unet_path = maybe_unet_path
-
-        ti_path = _ti_lora_path(unet_path)
-        text_path = _text_lora_path(unet_path)
-
-        if patch_unet:
-            print("LoRA : Patching Unet")
-            monkeypatch_or_replace_lora(
-                pipe.unet,
-                torch.load(unet_path),
-                r=r,
-                target_replace_module=unet_target_replace_module,
-            )
-
-        if patch_text:
-            print("LoRA : Patching text encoder")
-            monkeypatch_or_replace_lora(
-                pipe.text_encoder,
-                torch.load(text_path),
-                target_replace_module=text_target_replace_module,
-                r=r,
-            )
-        if patch_ti:
-            print("LoRA : Patching token input")
-            token = load_learned_embed_in_clip(
-                ti_path,
-                pipe.text_encoder,
-                pipe.tokenizer,
-                token=token,
-                idempotent=idempotent_token,
-            )
-
-    elif maybe_unet_path.endswith(".safetensors"):
-        safeloras = safe_open(maybe_unet_path, framework="pt", device="cpu")
-        monkeypatch_or_replace_safeloras(pipe, safeloras)
-        tok_dict = parse_safeloras_embeds(safeloras)
-        if patch_ti:
-            apply_learned_embed_in_clip(
-                tok_dict,
-                pipe.text_encoder,
-                pipe.tokenizer,
-                token=token,
-                idempotent=idempotent_token,
-            )
-        return tok_dict
-
-
-@torch.no_grad()
-def inspect_lora(model):
-    moved = {}
-
-    for name, _module in model.named_modules():
-        if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d"]:
-            ups = _module.lora_up.weight.data.clone()
-            downs = _module.lora_down.weight.data.clone()
-
-            wght: torch.Tensor = ups.flatten(1) @ downs.flatten(1)
-
-            dist = wght.flatten().abs().mean().item()
-            if name in moved:
-                moved[name].append(dist)
-            else:
-                moved[name] = [dist]
-
-    return moved
-
-
-def save_all(
-    unet,
-    text_encoder,
-    save_path,
-    placeholder_token_ids=None,
-    placeholder_tokens=None,
-    save_lora=True,
-    save_ti=True,
-    target_replace_module_text=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
-    target_replace_module_unet=DEFAULT_TARGET_REPLACE,
-    safe_form=True,
-):
-    if not safe_form:
-        # save ti
-        if save_ti:
-            ti_path = _ti_lora_path(save_path)
-            learned_embeds_dict = {}
-            for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
-                learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
-                print(
-                    f"Current Learned Embeddings for {tok}:, id {tok_id} ",
-                    learned_embeds[:4],
-                )
-                learned_embeds_dict[tok] = learned_embeds.detach().cpu()
-
-            torch.save(learned_embeds_dict, ti_path)
-            print("Ti saved to ", ti_path)
-
-        # save text encoder
-        if save_lora:
-
-            save_lora_weight(
-                unet, save_path, target_replace_module=target_replace_module_unet
-            )
-            print("Unet saved to ", save_path)
-
-            save_lora_weight(
-                text_encoder,
-                _text_lora_path(save_path),
-                target_replace_module=target_replace_module_text,
-            )
-            print("Text Encoder saved to ", _text_lora_path(save_path))
-
-    else:
-        assert save_path.endswith(
-            ".safetensors"
-        ), f"Save path : {save_path} should end with .safetensors"
-
-        loras = {}
-        embeds = {}
-
-        if save_lora:
-
-            loras["unet"] = (unet, target_replace_module_unet)
-            loras["text_encoder"] = (text_encoder, target_replace_module_text)
-
-        if save_ti:
-            for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
-                learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
-                print(
-                    f"Current Learned Embeddings for {tok}:, id {tok_id} ",
-                    learned_embeds[:4],
-                )
-                embeds[tok] = learned_embeds.detach().cpu()
-
-        save_safeloras_with_embeds(loras, embeds, save_path)

lvdm/{samplers → models/samplers}/ddim.py

@@ -1,10 +1,8 @@
-"""SAMPLING ONLY."""
-
-import torch
 import numpy as np
 from tqdm import tqdm
-
-from lvdm.models.modules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+import torch
+from lvdm.models.utils_diffusion import make_ddim_sampling_parameters, make_ddim_timesteps
+from lvdm.common import noise_like
 
 
 class DDIMSampler(object):
@@ -31,6 +29,15 @@ class DDIMSampler(object):
         self.register_buffer('betas', to_torch(self.model.betas))
         self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
         self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+        self.use_scale = self.model.use_scale
+        print('DDIM scale', self.use_scale)
+
+        if self.use_scale:
+            self.register_buffer('scale_arr', to_torch(self.model.scale_arr))
+            ddim_scale_arr = self.scale_arr.cpu()[self.ddim_timesteps]
+            self.register_buffer('ddim_scale_arr', ddim_scale_arr)
+            ddim_scale_arr = np.asarray([self.scale_arr.cpu()[0]] + self.scale_arr.cpu()[self.ddim_timesteps[:-1]].tolist())
+            self.register_buffer('ddim_scale_arr_prev', ddim_scale_arr)
 
         # calculations for diffusion q(x_t | x_{t-1}) and others
         self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
@@ -59,6 +66,7 @@ class DDIMSampler(object):
                shape,
                conditioning=None,
                callback=None,
+               normals_sequence=None,
                img_callback=None,
                quantize_x0=False,
                eta=0.,
@@ -74,9 +82,6 @@ class DDIMSampler(object):
                log_every_t=100,
                unconditional_guidance_scale=1.,
                unconditional_conditioning=None,
-               postprocess_fn=None,
-               sample_noise=None,
-               cond_fn=None,
                # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
                **kwargs
                ):
@@ -86,11 +91,11 @@ class DDIMSampler(object):
             if isinstance(conditioning, dict):
                 try:
                     cbs = conditioning[list(conditioning.keys())[0]].shape[0]
-                    if cbs != batch_size:
-                        print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
                 except:
-                    # cbs = conditioning[list(conditioning.keys())[0]][0].shape[0]
-                    pass
+                    cbs = conditioning[list(conditioning.keys())[0]][0].shape[0]
+
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
             else:
                 if conditioning.shape[0] != batch_size:
                     print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
@@ -104,6 +109,7 @@ class DDIMSampler(object):
         elif len(shape) == 4:
             C, T, H, W = shape
             size = (batch_size, C, T, H, W)
+        # print(f'Data shape for DDIM sampling is {size}, eta {eta}')
         
         samples, intermediates = self.ddim_sampling(conditioning, size,
                                                     callback=callback,
@@ -119,12 +125,8 @@ class DDIMSampler(object):
                                                     log_every_t=log_every_t,
                                                     unconditional_guidance_scale=unconditional_guidance_scale,
                                                     unconditional_conditioning=unconditional_conditioning,
-                                                    postprocess_fn=postprocess_fn,
-                                                    sample_noise=sample_noise,
-                                                    cond_fn=cond_fn,
                                                     verbose=verbose,
-                                                    **kwargs
-                                                    )
+                                                    **kwargs)
         return samples, intermediates
 
     @torch.no_grad()
@@ -133,13 +135,11 @@ class DDIMSampler(object):
                       callback=None, timesteps=None, quantize_denoised=False,
                       mask=None, x0=None, img_callback=None, log_every_t=100,
                       temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                      unconditional_guidance_scale=1., unconditional_conditioning=None,
-                      postprocess_fn=None,sample_noise=None,cond_fn=None,
-                      uc_type=None, verbose=True, **kwargs,
-                      ):
-
-        device = self.model.betas.device
-        
+                      unconditional_guidance_scale=1., unconditional_conditioning=None, verbose=True,
+                      cond_tau=1., target_size=None, start_timesteps=None,
+                      **kwargs):
+        device = self.model.betas.device        
+        print('ddim device', device)
         b = shape[0]
         if x_T is None:
             img = torch.randn(shape, device=device)
@@ -151,6 +151,7 @@ class DDIMSampler(object):
         elif timesteps is not None and not ddim_use_original_steps:
             subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
             timesteps = self.ddim_timesteps[:subset_end]
+            
         intermediates = {'x_inter': [img], 'pred_x0': [img]}
         time_range = reversed(range(0,timesteps)) if ddim_use_original_steps else np.flip(timesteps)
         total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
@@ -159,31 +160,46 @@ class DDIMSampler(object):
         else:
             iterator = time_range
 
+        init_x0 = False
+        clean_cond = kwargs.pop("clean_cond", False)
         for i, step in enumerate(iterator):
             index = total_steps - i - 1
             ts = torch.full((b,), step, device=device, dtype=torch.long)
+            if start_timesteps is not None:
+                assert x0 is not None
+                if step > start_timesteps*time_range[0]:
+                    continue
+                elif not init_x0:
+                    img = self.model.q_sample(x0, ts) 
+                    init_x0 = True
 
-            if postprocess_fn is not None:
-                img = postprocess_fn(img, ts)
+            # use mask to blend noised original latent (img_orig) & new sampled latent (img)
+            if mask is not None:
+                assert x0 is not None
+                if clean_cond:
+                    img_orig = x0
+                else:
+                    img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass? <ddim inversion>
+                img = img_orig * mask + (1. - mask) * img # keep original & modify use img
             
+            index_clip =  int((1 - cond_tau) * total_steps)
+            if index <= index_clip and target_size is not None:
+                target_size_ = [target_size[0], target_size[1]//8, target_size[2]//8]
+                img = torch.nn.functional.interpolate(
+                img,
+                size=target_size_,
+                mode="nearest",
+                )
             outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
                                       quantize_denoised=quantize_denoised, temperature=temperature,
                                       noise_dropout=noise_dropout, score_corrector=score_corrector,
                                       corrector_kwargs=corrector_kwargs,
                                       unconditional_guidance_scale=unconditional_guidance_scale,
                                       unconditional_conditioning=unconditional_conditioning,
-                                      sample_noise=sample_noise,cond_fn=cond_fn,uc_type=uc_type, **kwargs,)
-            img, pred_x0 = outs
-
-            if mask is not None:
-                # use mask to blend x_known_t-1 & x_sample_t-1
-                assert x0 is not None
-                x0 = x0.to(img.device)
-                mask = mask.to(img.device)
-                t = torch.tensor([step-1]*x0.shape[0], dtype=torch.long, device=img.device)
-                img_known = self.model.q_sample(x0, t)
-                img = img_known * mask + (1. - mask) * img
+                                      x0=x0,
+                                      **kwargs)
             
+            img, pred_x0 = outs
             if callback: callback(i)
             if img_callback: img_callback(pred_x0, i)
 
@@ -196,10 +212,8 @@ class DDIMSampler(object):
     @torch.no_grad()
     def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
                       temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
-                      unconditional_guidance_scale=1., unconditional_conditioning=None, sample_noise=None,
-                      cond_fn=None, uc_type=None, 
-                      **kwargs,
-                      ):
+                      unconditional_guidance_scale=1., unconditional_conditioning=None,
+                      uc_type=None, conditional_guidance_scale_temporal=None, **kwargs):
         b, *_, device = *x.shape, x.device
         if x.dim() == 5:
             is_video = True
@@ -227,7 +241,12 @@ class DDIMSampler(object):
                     e_t = e_t + unconditional_guidance_scale * (e_t_uncond - e_t)
                 else:
                     raise NotImplementedError
-            
+            # temporal guidance
+            if conditional_guidance_scale_temporal is not None:
+                e_t_temporal = self.model.apply_model(x, t, c, **kwargs)
+                e_t_image = self.model.apply_model(x, t, c, no_temporal_attn=True, **kwargs)
+                e_t = e_t + conditional_guidance_scale_temporal * (e_t_temporal - e_t_image)
+
         if score_corrector is not None:
             assert self.model.parameterization == "eps"
             e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
@@ -249,19 +268,69 @@ class DDIMSampler(object):
 
         # current prediction for x_0
         pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
-        # print(f't={t}, pred_x0, min={torch.min(pred_x0)}, max={torch.max(pred_x0)}',file=f)
         if quantize_denoised:
             pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
         # direction pointing to x_t
         dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+
+        noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
         
-        if sample_noise is None:
-            noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
-            if noise_dropout > 0.:
-                noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+        if self.use_scale:
+            scale_arr = self.model.scale_arr if use_original_steps else self.ddim_scale_arr
+            scale_t = torch.full(size, scale_arr[index], device=device)
+            scale_arr_prev = self.model.scale_arr_prev if use_original_steps else self.ddim_scale_arr_prev
+            scale_t_prev = torch.full(size, scale_arr_prev[index], device=device)
+            pred_x0 /= scale_t 
+            x_prev = a_prev.sqrt() * scale_t_prev * pred_x0 + dir_xt + noise
         else:
-            noise = sigma_t * sample_noise * temperature
-        
-        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
-        
+            x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+
         return x_prev, pred_x0
+
+
+    @torch.no_grad()
+    def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
+        # fast, but does not allow for exact reconstruction
+        # t serves as an index to gather the correct alphas
+        if use_original_steps:
+            sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
+            sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
+        else:
+            sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
+            sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
+
+        if noise is None:
+            noise = torch.randn_like(x0)
+
+        def extract_into_tensor(a, t, x_shape):
+            b, *_ = t.shape
+            out = a.gather(-1, t)
+            return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+        return (extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
+                extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise)
+
+    @torch.no_grad()
+    def decode(self, x_latent, cond, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
+               use_original_steps=False):
+
+        timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
+        timesteps = timesteps[:t_start]
+
+        time_range = np.flip(timesteps)
+        total_steps = timesteps.shape[0]
+        print(f"Running DDIM Sampling with {total_steps} timesteps")
+
+        iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
+        x_dec = x_latent
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
+            x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
+                                          unconditional_guidance_scale=unconditional_guidance_scale,
+                                          unconditional_conditioning=unconditional_conditioning)
+        return x_dec
+

lvdm/models/{modules/util.py → utils_diffusion.py}

@@ -1,13 +1,31 @@
 import math
-from inspect import isfunction
-
-import torch
 import numpy as np
-import torch.nn as nn
 from einops import repeat
+import torch
 import torch.nn.functional as F
 
-from lvdm.utils.common_utils import instantiate_from_config
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    if not repeat_only:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=timesteps.device)
+        args = timesteps[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    else:
+        embedding = repeat(timesteps, 'b -> b d', d=dim)
+    return embedding
 
 
 def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
@@ -15,6 +33,7 @@ def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2,
         betas = (
                 torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
         )
+
     elif schedule == "cosine":
         timesteps = (
                 torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
@@ -24,6 +43,7 @@ def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2,
         alphas = alphas / alphas[0]
         betas = 1 - alphas[1:] / alphas[:-1]
         betas = np.clip(betas, a_min=0, a_max=0.999)
+
     elif schedule == "sqrt_linear":
         betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
     elif schedule == "sqrt":
@@ -42,6 +62,7 @@ def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timestep
     else:
         raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
 
+    # assert ddim_timesteps.shape[0] == num_ddim_timesteps
     # add one to get the final alpha values right (the ones from first scale to data during sampling)
     steps_out = ddim_timesteps + 1
     if verbose:
@@ -51,6 +72,7 @@ def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timestep
 
 def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
     # select alphas for computing the variance schedule
+    # print(f'ddim_timesteps={ddim_timesteps}, len_alphacums={len(alphacums)}')
     alphas = alphacums[ddim_timesteps]
     alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
 
@@ -79,270 +101,4 @@ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
         t1 = i / num_diffusion_timesteps
         t2 = (i + 1) / num_diffusion_timesteps
         betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
-    return np.array(betas)
-
-
-def extract_into_tensor(a, t, x_shape):
-    b, *_ = t.shape
-    out = a.gather(-1, t)
-    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
-
-
-def checkpoint(func, inputs, params, flag):
-    """
-    Evaluate a function without caching intermediate activations, allowing for
-    reduced memory at the expense of extra compute in the backward pass.
-    :param func: the function to evaluate.
-    :param inputs: the argument sequence to pass to `func`.
-    :param params: a sequence of parameters `func` depends on but does not
-                   explicitly take as arguments.
-    :param flag: if False, disable gradient checkpointing.
-    """
-    if flag:
-        args = tuple(inputs) + tuple(params)
-        return CheckpointFunction.apply(func, len(inputs), *args)
-    else:
-        return func(*inputs)
-
-
-class CheckpointFunction(torch.autograd.Function):
-    @staticmethod
-    @torch.cuda.amp.custom_fwd
-    def forward(ctx, run_function, length, *args):
-        ctx.run_function = run_function
-        ctx.input_tensors = list(args[:length])
-        ctx.input_params = list(args[length:])
-
-        with torch.no_grad():
-            output_tensors = ctx.run_function(*ctx.input_tensors)
-        return output_tensors
-
-    @staticmethod
-    @torch.cuda.amp.custom_bwd
-    def backward(ctx, *output_grads):
-        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
-        with torch.enable_grad():
-            # Fixes a bug where the first op in run_function modifies the
-            # Tensor storage in place, which is not allowed for detach()'d
-            # Tensors.
-            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
-            output_tensors = ctx.run_function(*shallow_copies)
-        input_grads = torch.autograd.grad(
-            output_tensors,
-            ctx.input_tensors + ctx.input_params,
-            output_grads,
-            allow_unused=True,
-        )
-        del ctx.input_tensors
-        del ctx.input_params
-        del output_tensors
-        return (None, None) + input_grads
-
-
-def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
-    """
-    Create sinusoidal timestep embeddings.
-    :param timesteps: a 1-D Tensor of N indices, one per batch element.
-                      These may be fractional.
-    :param dim: the dimension of the output.
-    :param max_period: controls the minimum frequency of the embeddings.
-    :return: an [N x dim] Tensor of positional embeddings.
-    """
-    if not repeat_only:
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
-        ).to(device=timesteps.device)
-        args = timesteps[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    else:
-        embedding = repeat(timesteps, 'b -> b d', d=dim)
-    return embedding
-
-
-def zero_module(module):
-    """
-    Zero out the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().zero_()
-    return module
-
-
-def scale_module(module, scale):
-    """
-    Scale the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().mul_(scale)
-    return module
-
-
-def mean_flat(tensor):
-    """
-    Take the mean over all non-batch dimensions.
-    """
-    return tensor.mean(dim=list(range(1, len(tensor.shape))))
-
-
-def normalization(channels):
-    """
-    Make a standard normalization layer.
-    :param channels: number of input channels.
-    :return: an nn.Module for normalization.
-    """
-    return GroupNorm32(32, channels)
-
-def Normalize(in_channels):
-    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-def identity(*args, **kwargs):
-    return nn.Identity()
-
-class Normalization(nn.Module):
-    def __init__(self, output_size, eps=1e-5, norm_type='gn'):
-        super(Normalization, self).__init__()
-        # epsilon to avoid dividing by 0
-        self.eps = eps
-        self.norm_type = norm_type
-
-        if self.norm_type in ['bn', 'in']:
-            self.register_buffer('stored_mean', torch.zeros(output_size))
-            self.register_buffer('stored_var', torch.ones(output_size))
-    
-    def forward(self, x):
-        if self.norm_type == 'bn':
-            out = F.batch_norm(x, self.stored_mean, self.stored_var, None,
-                                None,
-                                self.training, 0.1, self.eps)
-        elif self.norm_type == 'in':
-            out = F.instance_norm(x, self.stored_mean, self.stored_var,
-                                    None, None,
-                                    self.training, 0.1, self.eps)
-        elif self.norm_type == 'gn':
-            out = F.group_norm(x, 32)
-        elif self.norm_type == 'nonorm':
-            out = x
-        return out
-
-
-class CCNormalization(nn.Module):
-    def __init__(self, embed_dim, feature_dim, *args, **kwargs):
-        super(CCNormalization, self).__init__()
-
-        self.embed_dim = embed_dim
-        self.feature_dim = feature_dim
-        self.norm = Normalization(feature_dim, *args, **kwargs)
-        
-        self.gain = nn.Linear(self.embed_dim, self.feature_dim)
-        self.bias = nn.Linear(self.embed_dim, self.feature_dim)
-        
-    def forward(self, x, y):
-        shape = [1] * (x.dim() - 2)
-        gain = (1 + self.gain(y)).view(y.size(0), -1, *shape)
-        bias = self.bias(y).view(y.size(0), -1, *shape)
-        return self.norm(x) * gain + bias
-
-
-def nonlinearity(type='silu'):
-    if type == 'silu':
-        return nn.SiLU()
-    elif type == 'leaky_relu':
-        return nn.LeakyReLU()
-
-
-class GEGLU(nn.Module):
-    def __init__(self, dim_in, dim_out):
-        super().__init__()
-        self.proj = nn.Linear(dim_in, dim_out * 2)
-
-    def forward(self, x):
-        x, gate = self.proj(x).chunk(2, dim=-1)
-        return x * F.gelu(gate)
-
-
-class SiLU(nn.Module):
-    def forward(self, x):
-        return x * torch.sigmoid(x)
-
-
-class GroupNorm32(nn.GroupNorm):
-    def forward(self, x):
-        return super().forward(x.float()).type(x.dtype)
-
-
-def conv_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D convolution module.
-    """
-    if dims == 1:
-        return nn.Conv1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.Conv2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.Conv3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-def linear(*args, **kwargs):
-    """
-    Create a linear module.
-    """
-    return nn.Linear(*args, **kwargs)
-
-
-def avg_pool_nd(dims, *args, **kwargs):
-    """
-    Create a 1D, 2D, or 3D average pooling module.
-    """
-    if dims == 1:
-        return nn.AvgPool1d(*args, **kwargs)
-    elif dims == 2:
-        return nn.AvgPool2d(*args, **kwargs)
-    elif dims == 3:
-        return nn.AvgPool3d(*args, **kwargs)
-    raise ValueError(f"unsupported dimensions: {dims}")
-
-
-class HybridConditioner(nn.Module):
-
-    def __init__(self, c_concat_config, c_crossattn_config):
-        super().__init__()
-        self.concat_conditioner = instantiate_from_config(c_concat_config)
-        self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
-
-    def forward(self, c_concat, c_crossattn):
-        c_concat = self.concat_conditioner(c_concat)
-        c_crossattn = self.crossattn_conditioner(c_crossattn)
-        return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
-
-
-def noise_like(shape, device, repeat=False):
-    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-    noise = lambda: torch.randn(shape, device=device)
-    return repeat_noise() if repeat else noise()
-
-
-def init_(tensor):
-    dim = tensor.shape[-1]
-    std = 1 / math.sqrt(dim)
-    tensor.uniform_(-std, std)
-    return tensor
-
-
-def exists(val):
-    return val is not None
-
-
-def uniq(arr):
-    return{el: True for el in arr}.keys()
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
+    return np.array(betas)

lvdm/modules/attention.py

@@ -0,0 +1,475 @@
+from functools import partial
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange, repeat
+try:
+    import xformers
+    import xformers.ops
+    XFORMERS_IS_AVAILBLE = True
+except:
+    XFORMERS_IS_AVAILBLE = False
+from lvdm.common import (
+    checkpoint,
+    exists,
+    default,
+)
+from lvdm.basics import (
+    zero_module,
+)
+
+class RelativePosition(nn.Module):
+    """ https://github.com/evelinehong/Transformer_Relative_Position_PyTorch/blob/master/relative_position.py """
+
+    def __init__(self, num_units, max_relative_position):
+        super().__init__()
+        self.num_units = num_units
+        self.max_relative_position = max_relative_position
+        self.embeddings_table = nn.Parameter(torch.Tensor(max_relative_position * 2 + 1, num_units))
+        nn.init.xavier_uniform_(self.embeddings_table)
+
+    def forward(self, length_q, length_k):
+        device = self.embeddings_table.device
+        range_vec_q = torch.arange(length_q, device=device)
+        range_vec_k = torch.arange(length_k, device=device)
+        distance_mat = range_vec_k[None, :] - range_vec_q[:, None]
+        distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
+        final_mat = distance_mat_clipped + self.max_relative_position
+        final_mat = final_mat.long()
+        embeddings = self.embeddings_table[final_mat]
+        return embeddings
+
+
+class CrossAttention(nn.Module):
+
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., 
+                 relative_position=False, temporal_length=None, img_cross_attention=False):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.scale = dim_head**-0.5
+        self.heads = heads
+        self.dim_head = dim_head
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
+
+        self.image_cross_attention_scale = 1.0
+        self.text_context_len = 77
+        self.img_cross_attention = img_cross_attention
+        if self.img_cross_attention:
+            self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
+            self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
+        
+        self.relative_position = relative_position
+        if self.relative_position:
+            assert(temporal_length is not None)
+            self.relative_position_k = RelativePosition(num_units=dim_head, max_relative_position=temporal_length)
+            self.relative_position_v = RelativePosition(num_units=dim_head, max_relative_position=temporal_length)
+        else:
+            ## only used for spatial attention, while NOT for temporal attention
+            if XFORMERS_IS_AVAILBLE and temporal_length is None:
+                self.forward = self.efficient_forward
+
+    def forward(self, x, context=None, mask=None):
+        h = self.heads
+
+        q = self.to_q(x)
+        context = default(context, x)
+        ## considering image token additionally
+        if context is not None and self.img_cross_attention:
+            context, context_img = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:]
+            k = self.to_k(context)
+            v = self.to_v(context)
+            k_ip = self.to_k_ip(context_img)
+            v_ip = self.to_v_ip(context_img)
+        else:
+            k = self.to_k(context)
+            v = self.to_v(context)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+        sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale
+        if self.relative_position:
+            len_q, len_k, len_v = q.shape[1], k.shape[1], v.shape[1]
+            k2 = self.relative_position_k(len_q, len_k)
+            sim2 = einsum('b t d, t s d -> b t s', q, k2) * self.scale # TODO check 
+            sim += sim2
+        del k
+
+        if exists(mask):
+            ## feasible for causal attention mask only
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, 'b i j -> (b h) i j', h=h)
+            sim.masked_fill_(~(mask>0.5), max_neg_value)
+
+        # attention, what we cannot get enough of
+        sim = sim.softmax(dim=-1)
+        out = torch.einsum('b i j, b j d -> b i d', sim, v)
+        if self.relative_position:
+            v2 = self.relative_position_v(len_q, len_v)
+            out2 = einsum('b t s, t s d -> b t d', sim, v2) # TODO check
+            out += out2
+        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+
+        ## considering image token additionally
+        if context is not None and self.img_cross_attention:
+            k_ip, v_ip = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (k_ip, v_ip))
+            sim_ip =  torch.einsum('b i d, b j d -> b i j', q, k_ip) * self.scale
+            del k_ip
+            sim_ip = sim_ip.softmax(dim=-1)
+            out_ip = torch.einsum('b i j, b j d -> b i d', sim_ip, v_ip)
+            out_ip = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+            out = out + self.image_cross_attention_scale * out_ip
+        del q
+
+        return self.to_out(out)
+    
+    def efficient_forward(self, x, context=None, mask=None):
+        q = self.to_q(x)
+        context = default(context, x)
+
+        ## considering image token additionally
+        if context is not None and self.img_cross_attention:
+            context, context_img = context[:,:self.text_context_len,:], context[:,self.text_context_len:,:]
+            k = self.to_k(context)
+            v = self.to_v(context)
+            k_ip = self.to_k_ip(context_img)
+            v_ip = self.to_v_ip(context_img)
+        else:
+            k = self.to_k(context)
+            v = self.to_v(context)
+
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+        # actually compute the attention, what we cannot get enough of
+        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=None)
+
+        ## considering image token additionally
+        if context is not None and self.img_cross_attention:
+            k_ip, v_ip = map(
+                lambda t: t.unsqueeze(3)
+                .reshape(b, t.shape[1], self.heads, self.dim_head)
+                .permute(0, 2, 1, 3)
+                .reshape(b * self.heads, t.shape[1], self.dim_head)
+                .contiguous(),
+                (k_ip, v_ip),
+            )
+            out_ip = xformers.ops.memory_efficient_attention(q, k_ip, v_ip, attn_bias=None, op=None)
+            out_ip = (
+                out_ip.unsqueeze(0)
+                .reshape(b, self.heads, out.shape[1], self.dim_head)
+                .permute(0, 2, 1, 3)
+                .reshape(b, out.shape[1], self.heads * self.dim_head)
+            )
+
+        if exists(mask):
+            raise NotImplementedError
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        if context is not None and self.img_cross_attention:
+            out = out + self.image_cross_attention_scale * out_ip
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+
+    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
+                disable_self_attn=False, attention_cls=None, img_cross_attention=False):
+        super().__init__()
+        attn_cls = CrossAttention if attention_cls is None else attention_cls
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
+            context_dim=context_dim if self.disable_self_attn else None)
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = attn_cls(query_dim=dim, context_dim=context_dim, heads=n_heads, dim_head=d_head, dropout=dropout,
+            img_cross_attention=img_cross_attention)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None, mask=None):
+        ## implementation tricks: because checkpointing doesn't support non-tensor (e.g. None or scalar) arguments
+        input_tuple = (x,)      ## should not be (x), otherwise *input_tuple will decouple x into multiple arguments
+        if context is not None:
+            input_tuple = (x, context)
+        if mask is not None:
+            forward_mask = partial(self._forward, mask=mask)
+            return checkpoint(forward_mask, (x,), self.parameters(), self.checkpoint)
+        if context is not None and mask is not None:
+            input_tuple = (x, context, mask)
+        return checkpoint(self._forward, input_tuple, self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None, mask=None):
+        x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None, mask=mask) + x
+        x = self.attn2(self.norm2(x), context=context, mask=mask) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data in spatial axis.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    NEW: use_linear for more efficiency instead of the 1x1 convs
+    """
+
+    def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0., context_dim=None,
+                 use_checkpoint=True, disable_self_attn=False, use_linear=False, img_cross_attention=False):
+        super().__init__()
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        if not use_linear:
+            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlock(
+                inner_dim,
+                n_heads,
+                d_head,
+                dropout=dropout,
+                context_dim=context_dim,
+                img_cross_attention=img_cross_attention,
+                disable_self_attn=disable_self_attn,
+                checkpoint=use_checkpoint) for d in range(depth)
+        ])
+        if not use_linear:
+            self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
+        else:
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+
+    def forward(self, x, context=None):
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            x = block(x, context=context)
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in
+    
+    
+class TemporalTransformer(nn.Module):
+    """
+    Transformer block for image-like data in temporal axis.
+    First, reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    """
+    def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0., context_dim=None,
+                 use_checkpoint=True, use_linear=False, only_self_att=True, causal_attention=False,
+                 relative_position=False, temporal_length=None):
+        super().__init__()
+        self.only_self_att = only_self_att
+        self.relative_position = relative_position
+        self.causal_attention = causal_attention
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.proj_in = nn.Conv1d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        if not use_linear:
+            self.proj_in = nn.Conv1d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        if relative_position:
+            assert(temporal_length is not None)
+            attention_cls = partial(CrossAttention, relative_position=True, temporal_length=temporal_length)
+        else:
+            attention_cls = None
+        if self.causal_attention:
+            assert(temporal_length is not None)
+            self.mask = torch.tril(torch.ones([1, temporal_length, temporal_length]))
+
+        if self.only_self_att:
+            context_dim = None
+        self.transformer_blocks = nn.ModuleList([
+            BasicTransformerBlock(
+                inner_dim,
+                n_heads,
+                d_head,
+                dropout=dropout,
+                context_dim=context_dim,
+                attention_cls=attention_cls,
+                checkpoint=use_checkpoint) for d in range(depth)
+        ])
+        if not use_linear:
+            self.proj_out = zero_module(nn.Conv1d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
+        else:
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+    def forward(self, x, context=None):
+        b, c, t, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        x = rearrange(x, 'b c t h w -> (b h w) c t').contiguous()
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, 'bhw c t -> bhw t c').contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+
+        if self.causal_attention:
+            mask = self.mask.to(x.device)
+            mask = repeat(mask, 'l i j -> (l bhw) i j', bhw=b*h*w)
+        else:
+            mask = None
+
+        if self.only_self_att:
+            ## note: if no context is given, cross-attention defaults to self-attention
+            for i, block in enumerate(self.transformer_blocks):
+                x = block(x, mask=mask)
+            x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous()
+        else:
+            x = rearrange(x, '(b hw) t c -> b hw t c', b=b).contiguous()
+            context = rearrange(context, '(b t) l con -> b t l con', t=t).contiguous()
+            for i, block in enumerate(self.transformer_blocks):
+                # calculate each batch one by one (since number in shape could not greater then 65,535 for some package)
+                for j in range(b):
+                    context_j = repeat(
+                        context[j],
+                        't l con -> (t r) l con', r=(h * w) // t, t=t).contiguous()
+                    ## note: causal mask will not applied in cross-attention case
+                    x[j] = block(x[j], context=context_j)
+        
+        if self.use_linear:
+            x = self.proj_out(x)
+            x = rearrange(x, 'b (h w) t c -> b c t h w', h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = rearrange(x, 'b hw t c -> (b hw) c t').contiguous()
+            x = self.proj_out(x)
+            x = rearrange(x, '(b h w) c t -> b c t h w', b=b, h=h, w=w).contiguous()
+
+        return x + x_in
+    
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            nn.GELU()
+        ) if not glu else GEGLU(dim, inner_dim)
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
+        k = k.softmax(dim=-1)  
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
+        out = torch.einsum('bhde,bhdn->bhen', context, q)
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+        return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = rearrange(q, 'b c h w -> b (h w) c')
+        k = rearrange(k, 'b c h w -> b c (h w)')
+        w_ = torch.einsum('bij,bjk->bik', q, k)
+
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = rearrange(v, 'b c h w -> b c (h w)')
+        w_ = rearrange(w_, 'b i j -> b j i')
+        h_ = torch.einsum('bij,bjk->bik', v, w_)
+        h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
+        h_ = self.proj_out(h_)
+
+        return x+h_

lvdm/modules/encoders/condition.py

@@ -0,0 +1,392 @@
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+import kornia
+import open_clip
+from transformers import T5Tokenizer, T5EncoderModel, CLIPTokenizer, CLIPTextModel
+from lvdm.common import autocast
+from utils.utils import count_params
+
+class AbstractEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class IdentityEncoder(AbstractEncoder):
+
+    def encode(self, x):
+        return x
+
+
+class ClassEmbedder(nn.Module):
+    def __init__(self, embed_dim, n_classes=1000, key='class', ucg_rate=0.1):
+        super().__init__()
+        self.key = key
+        self.embedding = nn.Embedding(n_classes, embed_dim)
+        self.n_classes = n_classes
+        self.ucg_rate = ucg_rate
+
+    def forward(self, batch, key=None, disable_dropout=False):
+        if key is None:
+            key = self.key
+        # this is for use in crossattn
+        c = batch[key][:, None]
+        if self.ucg_rate > 0. and not disable_dropout:
+            mask = 1. - torch.bernoulli(torch.ones_like(c) * self.ucg_rate)
+            c = mask * c + (1 - mask) * torch.ones_like(c) * (self.n_classes - 1)
+            c = c.long()
+        c = self.embedding(c)
+        return c
+
+    def get_unconditional_conditioning(self, bs, device="cuda"):
+        uc_class = self.n_classes - 1  # 1000 classes --> 0 ... 999, one extra class for ucg (class 1000)
+        uc = torch.ones((bs,), device=device) * uc_class
+        uc = {self.key: uc}
+        return uc
+
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+class FrozenT5Embedder(AbstractEncoder):
+    """Uses the T5 transformer encoder for text"""
+
+    def __init__(self, version="google/t5-v1_1-large", device="cuda", max_length=77,
+                 freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
+        super().__init__()
+        self.tokenizer = T5Tokenizer.from_pretrained(version)
+        self.transformer = T5EncoderModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length  # TODO: typical value?
+        if freeze:
+            self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        # self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.device)
+        outputs = self.transformer(input_ids=tokens)
+
+        z = outputs.last_hidden_state
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+
+class FrozenCLIPEmbedder(AbstractEncoder):
+    """Uses the CLIP transformer encoder for text (from huggingface)"""
+    LAYERS = [
+        "last",
+        "pooled",
+        "hidden"
+    ]
+
+    def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,
+                 freeze=True, layer="last", layer_idx=None):  # clip-vit-base-patch32
+        super().__init__()
+        assert layer in self.LAYERS
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        self.layer_idx = layer_idx
+        if layer == "hidden":
+            assert layer_idx is not None
+            assert 0 <= abs(layer_idx) <= 12
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        # self.train = disabled_train
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+        tokens = batch_encoding["input_ids"].to(self.device)
+        outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer == "hidden")
+        if self.layer == "last":
+            z = outputs.last_hidden_state
+        elif self.layer == "pooled":
+            z = outputs.pooler_output[:, None, :]
+        else:
+            z = outputs.hidden_states[self.layer_idx]
+        return z
+
+    def encode(self, text):
+        return self(text)
+
+
+class ClipImageEmbedder(nn.Module):
+    def __init__(
+            self,
+            model,
+            jit=False,
+            device='cuda' if torch.cuda.is_available() else 'cpu',
+            antialias=True,
+            ucg_rate=0.
+    ):
+        super().__init__()
+        from clip import load as load_clip
+        self.model, _ = load_clip(name=model, device=device, jit=jit)
+
+        self.antialias = antialias
+
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+        self.ucg_rate = ucg_rate
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = kornia.geometry.resize(x, (224, 224),
+                                   interpolation='bicubic', align_corners=True,
+                                   antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # re-normalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def forward(self, x, no_dropout=False):
+        # x is assumed to be in range [-1,1]
+        out = self.model.encode_image(self.preprocess(x))
+        out = out.to(x.dtype)
+        if self.ucg_rate > 0. and not no_dropout:
+            out = torch.bernoulli((1. - self.ucg_rate) * torch.ones(out.shape[0], device=out.device))[:, None] * out
+        return out
+
+
+class FrozenOpenCLIPEmbedder(AbstractEncoder):
+    """
+    Uses the OpenCLIP transformer encoder for text
+    """
+    LAYERS = [
+        # "pooled",
+        "last",
+        "penultimate"
+    ]
+
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,
+                 freeze=True, layer="last"):
+        super().__init__()
+        assert layer in self.LAYERS
+        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'))
+        del model.visual
+        self.model = model
+
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "last":
+            self.layer_idx = 0
+        elif self.layer == "penultimate":
+            self.layer_idx = 1
+        else:
+            raise NotImplementedError()
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        self.device = self.model.positional_embedding.device
+        tokens = open_clip.tokenize(text)
+        z = self.encode_with_transformer(tokens.to(self.device))
+        return z
+
+    def encode_with_transformer(self, text):
+        x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
+        x = x + self.model.positional_embedding
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.model.ln_final(x)
+        return x
+
+    def text_transformer_forward(self, x: torch.Tensor, attn_mask=None):
+        for i, r in enumerate(self.model.transformer.resblocks):
+            if i == len(self.model.transformer.resblocks) - self.layer_idx:
+                break
+            if self.model.transformer.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(r, x, attn_mask)
+            else:
+                x = r(x, attn_mask=attn_mask)
+        return x
+
+    def encode(self, text):
+        return self(text)
+
+
+class FrozenOpenCLIPImageEmbedder(AbstractEncoder):
+    """
+    Uses the OpenCLIP vision transformer encoder for images
+    """
+
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,
+                 freeze=True, layer="pooled", antialias=True, ucg_rate=0.):
+        super().__init__()
+        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'),
+                                                            pretrained=version, )
+        del model.transformer
+        self.model = model
+
+        self.device = device
+        self.max_length = max_length
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "penultimate":
+            raise NotImplementedError()
+            self.layer_idx = 1
+
+        self.antialias = antialias
+
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+        self.ucg_rate = ucg_rate
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = kornia.geometry.resize(x, (224, 224),
+                                   interpolation='bicubic', align_corners=True,
+                                   antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # renormalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    @autocast
+    def forward(self, image, no_dropout=False):
+        z = self.encode_with_vision_transformer(image)
+        if self.ucg_rate > 0. and not no_dropout:
+            z = torch.bernoulli((1. - self.ucg_rate) * torch.ones(z.shape[0], device=z.device))[:, None] * z
+        return z
+
+    def encode_with_vision_transformer(self, img):
+        img = self.preprocess(img)
+        x = self.model.visual(img)
+        return x
+
+    def encode(self, text):
+        return self(text)
+
+
+
+class FrozenOpenCLIPImageEmbedderV2(AbstractEncoder):
+    """
+    Uses the OpenCLIP vision transformer encoder for images
+    """
+
+    def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda",
+                 freeze=True, layer="pooled", antialias=True):
+        super().__init__()
+        model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'),
+                                                            pretrained=version, )
+        del model.transformer
+        self.model = model
+        self.device = device
+
+        if freeze:
+            self.freeze()
+        self.layer = layer
+        if self.layer == "penultimate":
+            raise NotImplementedError()
+            self.layer_idx = 1
+
+        self.antialias = antialias
+        self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+        self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+
+
+    def preprocess(self, x):
+        # normalize to [0,1]
+        x = kornia.geometry.resize(x, (224, 224),
+                                   interpolation='bicubic', align_corners=True,
+                                   antialias=self.antialias)
+        x = (x + 1.) / 2.
+        # renormalize according to clip
+        x = kornia.enhance.normalize(x, self.mean, self.std)
+        return x
+
+    def freeze(self):
+        self.model = self.model.eval()
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def forward(self, image, no_dropout=False):
+        ## image: b c h w
+        z = self.encode_with_vision_transformer(image)
+        return z
+
+    def encode_with_vision_transformer(self, x):
+        x = self.preprocess(x)
+
+        # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1
+        if self.model.visual.input_patchnorm:
+            # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)')
+            x = x.reshape(x.shape[0], x.shape[1], self.model.visual.grid_size[0], self.model.visual.patch_size[0], self.model.visual.grid_size[1], self.model.visual.patch_size[1])
+            x = x.permute(0, 2, 4, 1, 3, 5)
+            x = x.reshape(x.shape[0], self.model.visual.grid_size[0] * self.model.visual.grid_size[1], -1)
+            x = self.model.visual.patchnorm_pre_ln(x)
+            x = self.model.visual.conv1(x)
+        else:
+            x = self.model.visual.conv1(x)  # shape = [*, width, grid, grid]
+            x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+            x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+
+        # class embeddings and positional embeddings
+        x = torch.cat(
+            [self.model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
+             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.model.visual.positional_embedding.to(x.dtype)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        x = self.model.visual.patch_dropout(x)
+        x = self.model.visual.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.model.visual.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        return x
+
+
+class FrozenCLIPT5Encoder(AbstractEncoder):
+    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cuda",
+                 clip_max_length=77, t5_max_length=77):
+        super().__init__()
+        self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length)
+        self.t5_encoder = FrozenT5Embedder(t5_version, device, max_length=t5_max_length)
+        print(f"{self.clip_encoder.__class__.__name__} has {count_params(self.clip_encoder) * 1.e-6:.2f} M parameters, "
+              f"{self.t5_encoder.__class__.__name__} comes with {count_params(self.t5_encoder) * 1.e-6:.2f} M params.")
+
+    def encode(self, text):
+        return self(text)
+
+    def forward(self, text):
+        clip_z = self.clip_encoder.encode(text)
+        t5_z = self.t5_encoder.encode(text)
+        return [clip_z, t5_z]

lvdm/modules/encoders/ip_resampler.py

@@ -0,0 +1,136 @@
+# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
+import math
+import torch
+import torch.nn as nn
+
+
+class ImageProjModel(nn.Module):
+    """Projection Model"""
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
+        super().__init__()        
+        self.cross_attention_dim = cross_attention_dim
+        self.clip_extra_context_tokens = clip_extra_context_tokens
+        self.proj = nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
+        self.norm = nn.LayerNorm(cross_attention_dim)
+        
+    def forward(self, image_embeds):
+        #embeds = image_embeds
+        embeds = image_embeds.type(list(self.proj.parameters())[0].dtype)
+        clip_extra_context_tokens = self.proj(embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim)
+        clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
+        return clip_extra_context_tokens
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+    
+    
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    #(bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+        
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+        
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+        
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+    ):
+        super().__init__()
+        
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
+        
+        self.proj_in = nn.Linear(embedding_dim, dim)
+
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+        
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        
+        latents = self.latents.repeat(x.size(0), 1, 1)
+        
+        x = self.proj_in(x)
+        
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+            
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)

lvdm/{models/modules/autoencoder_modules.py → modules/networks/ae_modules.py}

@@ -1,30 +1,11 @@
+# pytorch_diffusion + derived encoder decoder
 import math
-
 import torch
 import numpy as np
-from torch import nn
+import torch.nn as nn
 from einops import rearrange
-
-
-def get_timestep_embedding(timesteps, embedding_dim):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models:
-    From Fairseq.
-    Build sinusoidal embeddings.
-    This matches the implementation in tensor2tensor, but differs slightly
-    from the description in Section 3.5 of "Attention Is All You Need".
-    """
-    assert len(timesteps.shape) == 1
-
-    half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - 1)
-    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
-    emb = emb.to(device=timesteps.device)
-    emb = timesteps.float()[:, None] * emb[None, :]
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-    if embedding_dim % 2 == 1:  # zero pad
-        emb = torch.nn.functional.pad(emb, (0,1,0,0))
-    return emb
+from utils.utils import instantiate_from_config
+from lvdm.modules.attention import LinearAttention
 
 def nonlinearity(x):
     # swish
@@ -36,25 +17,6 @@ def Normalize(in_channels, num_groups=32):
 
 
 
-class LinearAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
-        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x)
-        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
-        k = k.softmax(dim=-1)  
-        context = torch.einsum('bhdn,bhen->bhde', k, v)
-        out = torch.einsum('bhde,bhdn->bhen', context, q)
-        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
-        return self.to_out(out)
-
-
 class LinAttnBlock(LinearAttention):
     """to match AttnBlock usage"""
     def __init__(self, in_channels):
@@ -115,10 +77,9 @@ class AttnBlock(nn.Module):
 
         return x+h_
 
-
 def make_attn(in_channels, attn_type="vanilla"):
     assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
-    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+    #print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
     if attn_type == "vanilla":
         return AttnBlock(in_channels)
     elif attn_type == "none":
@@ -165,6 +126,27 @@ class Upsample(nn.Module):
             x = self.conv(x)
         return x
 
+def get_timestep_embedding(timesteps, embedding_dim):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models:
+    From Fairseq.
+    Build sinusoidal embeddings.
+    This matches the implementation in tensor2tensor, but differs slightly
+    from the description in Section 3.5 of "Attention Is All You Need".
+    """
+    assert len(timesteps.shape) == 1
+
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+    emb = emb.to(device=timesteps.device)
+    emb = timesteps.float()[:, None] * emb[None, :]
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if embedding_dim % 2 == 1:  # zero pad
+        emb = torch.nn.functional.pad(emb, (0,1,0,0))
+    return emb
+
+
 
 class ResnetBlock(nn.Module):
     def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
@@ -502,7 +484,7 @@ class Decoder(nn.Module):
         block_in = ch*ch_mult[self.num_resolutions-1]
         curr_res = resolution // 2**(self.num_resolutions-1)
         self.z_shape = (1,z_channels,curr_res,curr_res)
-        print("Working with z of shape {} = {} dimensions.".format(
+        print("AE working on z of shape {} = {} dimensions.".format(
             self.z_shape, np.prod(self.z_shape)))
 
         # z to block_in
@@ -594,3 +576,270 @@ class Decoder(nn.Module):
         if self.tanh_out:
             h = torch.tanh(h)
         return h
+
+
+class SimpleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, *args, **kwargs):
+        super().__init__()
+        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
+                                     ResnetBlock(in_channels=in_channels,
+                                                 out_channels=2 * in_channels,
+                                                 temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=2 * in_channels,
+                                                out_channels=4 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=4 * in_channels,
+                                                out_channels=2 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     nn.Conv2d(2*in_channels, in_channels, 1),
+                                     Upsample(in_channels, with_conv=True)])
+        # end
+        self.norm_out = Normalize(in_channels)
+        self.conv_out = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        for i, layer in enumerate(self.model):
+            if i in [1,2,3]:
+                x = layer(x, None)
+            else:
+                x = layer(x)
+
+        h = self.norm_out(x)
+        h = nonlinearity(h)
+        x = self.conv_out(h)
+        return x
+
+
+class UpsampleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
+                 ch_mult=(2,2), dropout=0.0):
+        super().__init__()
+        # upsampling
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        block_in = in_channels
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.res_blocks = nn.ModuleList()
+        self.upsample_blocks = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            res_block = []
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                res_block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+            self.res_blocks.append(nn.ModuleList(res_block))
+            if i_level != self.num_resolutions - 1:
+                self.upsample_blocks.append(Upsample(block_in, True))
+                curr_res = curr_res * 2
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        # upsampling
+        h = x
+        for k, i_level in enumerate(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.res_blocks[i_level][i_block](h, None)
+            if i_level != self.num_resolutions - 1:
+                h = self.upsample_blocks[k](h)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class LatentRescaler(nn.Module):
+    def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
+        super().__init__()
+        # residual block, interpolate, residual block
+        self.factor = factor
+        self.conv_in = nn.Conv2d(in_channels,
+                                 mid_channels,
+                                 kernel_size=3,
+                                 stride=1,
+                                 padding=1)
+        self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+                                                     out_channels=mid_channels,
+                                                     temb_channels=0,
+                                                     dropout=0.0) for _ in range(depth)])
+        self.attn = AttnBlock(mid_channels)
+        self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+                                                     out_channels=mid_channels,
+                                                     temb_channels=0,
+                                                     dropout=0.0) for _ in range(depth)])
+
+        self.conv_out = nn.Conv2d(mid_channels,
+                                  out_channels,
+                                  kernel_size=1,
+                                  )
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        for block in self.res_block1:
+            x = block(x, None)
+        x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor))))
+        x = self.attn(x)
+        for block in self.res_block2:
+            x = block(x, None)
+        x = self.conv_out(x)
+        return x
+
+
+class MergedRescaleEncoder(nn.Module):
+    def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
+                 ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1):
+        super().__init__()
+        intermediate_chn = ch * ch_mult[-1]
+        self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
+                               z_channels=intermediate_chn, double_z=False, resolution=resolution,
+                               attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv,
+                               out_ch=None)
+        self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
+                                       mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.rescaler(x)
+        return x
+
+
+class MergedRescaleDecoder(nn.Module):
+    def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8),
+                 dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
+        super().__init__()
+        tmp_chn = z_channels*ch_mult[-1]
+        self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout,
+                               resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
+                               ch_mult=ch_mult, resolution=resolution, ch=ch)
+        self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
+                                       out_channels=tmp_chn, depth=rescale_module_depth)
+
+    def forward(self, x):
+        x = self.rescaler(x)
+        x = self.decoder(x)
+        return x
+
+
+class Upsampler(nn.Module):
+    def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
+        super().__init__()
+        assert out_size >= in_size
+        num_blocks = int(np.log2(out_size//in_size))+1
+        factor_up = 1.+ (out_size % in_size)
+        print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
+        self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels,
+                                       out_channels=in_channels)
+        self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2,
+                               attn_resolutions=[], in_channels=None, ch=in_channels,
+                               ch_mult=[ch_mult for _ in range(num_blocks)])
+
+    def forward(self, x):
+        x = self.rescaler(x)
+        x = self.decoder(x)
+        return x
+
+
+class Resize(nn.Module):
+    def __init__(self, in_channels=None, learned=False, mode="bilinear"):
+        super().__init__()
+        self.with_conv = learned
+        self.mode = mode
+        if self.with_conv:
+            print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
+            raise NotImplementedError()
+            assert in_channels is not None
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=4,
+                                        stride=2,
+                                        padding=1)
+
+    def forward(self, x, scale_factor=1.0):
+        if scale_factor==1.0:
+            return x
+        else:
+            x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
+        return x
+
+class FirstStagePostProcessor(nn.Module):
+
+    def __init__(self, ch_mult:list, in_channels,
+                 pretrained_model:nn.Module=None,
+                 reshape=False,
+                 n_channels=None,
+                 dropout=0.,
+                 pretrained_config=None):
+        super().__init__()
+        if pretrained_config is None:
+            assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+            self.pretrained_model = pretrained_model
+        else:
+            assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+            self.instantiate_pretrained(pretrained_config)
+
+        self.do_reshape = reshape
+
+        if n_channels is None:
+            n_channels = self.pretrained_model.encoder.ch
+
+        self.proj_norm = Normalize(in_channels,num_groups=in_channels//2)
+        self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3,
+                            stride=1,padding=1)
+
+        blocks = []
+        downs = []
+        ch_in = n_channels
+        for m in ch_mult:
+            blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout))
+            ch_in = m * n_channels
+            downs.append(Downsample(ch_in, with_conv=False))
+
+        self.model = nn.ModuleList(blocks)
+        self.downsampler = nn.ModuleList(downs)
+
+
+    def instantiate_pretrained(self, config):
+        model = instantiate_from_config(config)
+        self.pretrained_model = model.eval()
+        # self.pretrained_model.train = False
+        for param in self.pretrained_model.parameters():
+            param.requires_grad = False
+
+
+    @torch.no_grad()
+    def encode_with_pretrained(self,x):
+        c = self.pretrained_model.encode(x)
+        if isinstance(c, DiagonalGaussianDistribution):
+            c = c.mode()
+        return  c
+
+    def forward(self,x):
+        z_fs = self.encode_with_pretrained(x)
+        z = self.proj_norm(z_fs)
+        z = self.proj(z)
+        z = nonlinearity(z)
+
+        for submodel, downmodel in zip(self.model,self.downsampler):
+            z = submodel(z,temb=None)
+            z = downmodel(z)
+
+        if self.do_reshape:
+            z = rearrange(z,'b c h w -> b (h w) c')
+        return z
+

lvdm/{models/modules → modules/networks}/openaimodel3d.py

@@ -1,38 +1,25 @@
-from abc import abstractmethod
-import math
-from einops import rearrange
 from functools import partial
-import numpy as np
-import torch as th
+from abc import abstractmethod
+import torch
 import torch.nn as nn
+from einops import rearrange
 import torch.nn.functional as F
-from omegaconf.listconfig import ListConfig
-
-from lvdm.models.modules.util import (
-    checkpoint,
+from lvdm.models.utils_diffusion import timestep_embedding
+from lvdm.common import checkpoint
+from lvdm.basics import (
+    zero_module,
     conv_nd,
     linear,
     avg_pool_nd,
-    zero_module,
-    normalization,
-    timestep_embedding,
-    nonlinearity,
+    normalization
 )
+from lvdm.modules.attention import SpatialTransformer, TemporalTransformer
 
-# dummy replace
-def convert_module_to_f16(x):
-    pass
-
-def convert_module_to_f32(x):
-    pass
 
-## go
-# ---------------------------------------------------------------------------------------------------
 class TimestepBlock(nn.Module):
     """
     Any module where forward() takes timestep embeddings as a second argument.
     """
-
     @abstractmethod
     def forward(self, x, emb):
         """
@@ -40,107 +27,85 @@ class TimestepBlock(nn.Module):
         """
 
 
-# ---------------------------------------------------------------------------------------------------
 class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
     """
     A sequential module that passes timestep embeddings to the children that
     support it as an extra input.
     """
 
-    def forward(self, x, emb, context, **kwargs):
+    def forward(self, x, emb, context=None, batch_size=None):
         for layer in self:
             if isinstance(layer, TimestepBlock):
-                x = layer(x, emb, **kwargs)
-            elif isinstance(layer, STTransformerClass):
-                x = layer(x, context, **kwargs)
+                x = layer(x, emb, batch_size)
+            elif isinstance(layer, SpatialTransformer):
+                x = layer(x, context)
+            elif isinstance(layer, TemporalTransformer):
+                x = rearrange(x, '(b f) c h w -> b c f h w', b=batch_size)
+                x = layer(x, context)
+                x = rearrange(x, 'b c f h w -> (b f) c h w')
             else:
-                x = layer(x)
+                x = layer(x,)
         return x
 
 
-# ---------------------------------------------------------------------------------------------------
-class Upsample(nn.Module):
+class Downsample(nn.Module):
     """
-    An upsampling layer with an optional convolution.
+    A downsampling layer with an optional convolution.
     :param channels: channels in the inputs and outputs.
     :param use_conv: a bool determining if a convolution is applied.
     :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 upsampling occurs in the inner-two dimensions.
+                 downsampling occurs in the inner-two dimensions.
     """
 
-    def __init__(self, channels, use_conv, dims=2, out_channels=None, 
-        kernel_size_t=3,
-        padding_t=1,
-    ):
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
         super().__init__()
         self.channels = channels
         self.out_channels = out_channels or channels
         self.use_conv = use_conv
         self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
         if use_conv:
-            self.conv = conv_nd(dims, self.channels, self.out_channels, (kernel_size_t, 3,3), padding=(padding_t, 1,1))
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        if self.dims == 3:
-            x = F.interpolate(
-                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            self.op = conv_nd(
+                dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
             )
         else:
-            x = F.interpolate(x, scale_factor=2, mode="nearest")
-        if self.use_conv:
-            x = self.conv(x)
-        return x
-
-
-# ---------------------------------------------------------------------------------------------------
-class TransposedUpsample(nn.Module):
-    'Learned 2x upsampling without padding'
-    def __init__(self, channels, out_channels=None, ks=5):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-
-        self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
 
-    def forward(self,x):
-        return self.up(x)
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
 
 
-# ---------------------------------------------------------------------------------------------------
-class Downsample(nn.Module):
+class Upsample(nn.Module):
     """
-    A downsampling layer with an optional convolution.
+    An upsampling layer with an optional convolution.
     :param channels: channels in the inputs and outputs.
     :param use_conv: a bool determining if a convolution is applied.
     :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 downsampling occurs in the inner-two dimensions.
+                 upsampling occurs in the inner-two dimensions.
     """
 
-    def __init__(self, channels, use_conv, dims=2, out_channels=None,
-        kernel_size_t=3,
-        padding_t=1,
-    ):
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
         super().__init__()
         self.channels = channels
         self.out_channels = out_channels or channels
         self.use_conv = use_conv
         self.dims = dims
-        stride = 2 if dims != 3 else (1, 2, 2)
         if use_conv:
-            self.op = conv_nd(
-                dims, self.channels, self.out_channels, (kernel_size_t, 3,3), stride=stride, padding=(padding_t, 1,1)
-            )
-        else:
-            assert self.channels == self.out_channels
-            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
 
     def forward(self, x):
         assert x.shape[1] == self.channels
-        return self.op(x)
+        if self.dims == 3:
+            x = F.interpolate(x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode='nearest')
+        else:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.use_conv:
+            x = self.conv(x)
+        return x
 
 
-# ---------------------------------------------------------------------------------------------------
 class ResBlock(TimestepBlock):
     """
     A residual block that can optionally change the number of channels.
@@ -152,7 +117,6 @@ class ResBlock(TimestepBlock):
         convolution instead of a smaller 1x1 convolution to change the
         channels in the skip connection.
     :param dims: determines if the signal is 1D, 2D, or 3D.
-    :param use_checkpoint: if True, use gradient checkpointing on this module.
     :param up: if True, use this block for upsampling.
     :param down: if True, use this block for downsampling.
     """
@@ -163,17 +127,14 @@ class ResBlock(TimestepBlock):
         emb_channels,
         dropout,
         out_channels=None,
-        use_conv=False,
         use_scale_shift_norm=False,
         dims=2,
         use_checkpoint=False,
+        use_conv=False,
         up=False,
         down=False,
-        # temporal
-        kernel_size_t=3,
-        padding_t=1,
-        nonlinearity_type='silu',
-        **kwargs
+        use_temporal_conv=False,
+        tempspatial_aware=False
     ):
         super().__init__()
         self.channels = channels
@@ -183,65 +144,68 @@ class ResBlock(TimestepBlock):
         self.use_conv = use_conv
         self.use_checkpoint = use_checkpoint
         self.use_scale_shift_norm = use_scale_shift_norm
-        self.nonlinearity_type = nonlinearity_type
+        self.use_temporal_conv = use_temporal_conv
 
         self.in_layers = nn.Sequential(
             normalization(channels),
-            nonlinearity(nonlinearity_type),
-            conv_nd(dims, channels, self.out_channels, (kernel_size_t, 3,3), padding=(padding_t, 1,1)),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
         )
 
         self.updown = up or down
 
         if up:
-            self.h_upd = Upsample(channels, False, dims, kernel_size_t=kernel_size_t, padding_t=padding_t)
-            self.x_upd = Upsample(channels, False, dims, kernel_size_t=kernel_size_t, padding_t=padding_t)
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
         elif down:
-            self.h_upd = Downsample(channels, False, dims, kernel_size_t=kernel_size_t, padding_t=padding_t)
-            self.x_upd = Downsample(channels, False, dims, kernel_size_t=kernel_size_t, padding_t=padding_t)
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
         else:
             self.h_upd = self.x_upd = nn.Identity()
 
         self.emb_layers = nn.Sequential(
-            nonlinearity(nonlinearity_type),
-            linear(
+            nn.SiLU(),
+            nn.Linear(
                 emb_channels,
                 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
             ),
         )
         self.out_layers = nn.Sequential(
             normalization(self.out_channels),
-            nonlinearity(nonlinearity_type),
+            nn.SiLU(),
             nn.Dropout(p=dropout),
-            zero_module(
-                conv_nd(dims, self.out_channels, self.out_channels, (kernel_size_t, 3,3), padding=(padding_t, 1,1))
-            ),
+            zero_module(nn.Conv2d(self.out_channels, self.out_channels, 3, padding=1)),
         )
 
         if self.out_channels == channels:
             self.skip_connection = nn.Identity()
         elif use_conv:
-            self.skip_connection = conv_nd(
-                dims, channels, self.out_channels, (kernel_size_t, 3,3), padding=(padding_t, 1,1)
-            )
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
         else:
             self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
-        
 
-    def forward(self, x, emb, **kwargs):
+        if self.use_temporal_conv:
+            self.temopral_conv = TemporalConvBlock(
+                self.out_channels,
+                self.out_channels,
+                dropout=0.1,
+                spatial_aware=tempspatial_aware
+            )
+
+    def forward(self, x, emb,  batch_size=None):
         """
         Apply the block to a Tensor, conditioned on a timestep embedding.
         :param x: an [N x C x ...] Tensor of features.
         :param emb: an [N x emb_channels] Tensor of timestep embeddings.
         :return: an [N x C x ...] Tensor of outputs.
         """
-        return checkpoint(self._forward, 
-                          (x, emb), 
-                          self.parameters(), 
-                          self.use_checkpoint
-                          )
+        input_tuple = (x, emb,)
+        if batch_size:
+            forward_batchsize = partial(self._forward, batch_size=batch_size)
+            return checkpoint(forward_batchsize, input_tuple, self.parameters(), self.use_checkpoint)
+        return checkpoint(self._forward, input_tuple, self.parameters(), self.use_checkpoint)
 
-    def _forward(self, x, emb,):
+    def _forward(self, x, emb,  batch_size=None,):
         if self.updown:
             in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
             h = in_rest(x)
@@ -250,38 +214,72 @@ class ResBlock(TimestepBlock):
             h = in_conv(h)
         else:
             h = self.in_layers(x)
-
         emb_out = self.emb_layers(emb).type(h.dtype)
-        if emb_out.dim() == 3: # btc for video data
-            emb_out = rearrange(emb_out, 'b t c -> b c t')
-        while len(emb_out.shape) < h.dim():
-            emb_out = emb_out[..., None] # bct -> bct11 or bc -> bc111
-        
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
         if self.use_scale_shift_norm:
             out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
-            scale, shift = th.chunk(emb_out, 2, dim=1)
+            scale, shift = torch.chunk(emb_out, 2, dim=1)
             h = out_norm(h) * (1 + scale) + shift
             h = out_rest(h)
         else:
             h = h + emb_out
             h = self.out_layers(h)
-
-        out = self.skip_connection(x) + h
+        h = self.skip_connection(x) + h
         
-        return out
+        if self.use_temporal_conv and batch_size:
+            h = rearrange(h, '(b t) c h w -> b c t h w', b=batch_size)
+            h = self.temopral_conv(h)
+            h = rearrange(h, 'b c t h w -> (b t) c h w')
+        return h
+
+
+class TemporalConvBlock(nn.Module):
+    """
+    Adapted from modelscope: https://github.com/modelscope/modelscope/blob/master/modelscope/models/multi_modal/video_synthesis/unet_sd.py
+    """
+
+    def __init__(self, in_channels, out_channels=None, dropout=0.0, spatial_aware=False):
+        super(TemporalConvBlock, self).__init__()
+        if out_channels is None:
+            out_channels = in_channels
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        kernel_shape = (3, 1, 1) if not spatial_aware else (3, 3, 3)
+        padding_shape = (1, 0, 0) if not spatial_aware else (1, 1, 1)
+
+        # conv layers
+        self.conv1 = nn.Sequential(
+            nn.GroupNorm(32, in_channels), nn.SiLU(),
+            nn.Conv3d(in_channels, out_channels, kernel_shape, padding=padding_shape))
+        self.conv2 = nn.Sequential(
+            nn.GroupNorm(32, out_channels), nn.SiLU(), nn.Dropout(dropout),
+            nn.Conv3d(out_channels, in_channels, kernel_shape, padding=padding_shape))
+        self.conv3 = nn.Sequential(
+            nn.GroupNorm(32, out_channels), nn.SiLU(), nn.Dropout(dropout),
+            nn.Conv3d(out_channels, in_channels, (3, 1, 1), padding=(1, 0, 0)))
+        self.conv4 = nn.Sequential(
+            nn.GroupNorm(32, out_channels), nn.SiLU(), nn.Dropout(dropout),
+            nn.Conv3d(out_channels, in_channels, (3, 1, 1), padding=(1, 0, 0)))
+
+        # zero out the last layer params,so the conv block is identity
+        nn.init.zeros_(self.conv4[-1].weight)
+        nn.init.zeros_(self.conv4[-1].bias)
+
+    def forward(self, x):
+        identity = x
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        return x + identity
 
-# ---------------------------------------------------------------------------------------------------
-def make_spatialtemporal_transformer(module_name='attention_temporal', class_name='SpatialTemporalTransformer'):
-    module = __import__(f"lvdm.models.modules.{module_name}", fromlist=[class_name])
-    global STTransformerClass
-    STTransformerClass = getattr(module, class_name)
-    return STTransformerClass
 
-# ---------------------------------------------------------------------------------------------------
 class UNetModel(nn.Module):
     """
     The full UNet model with attention and timestep embedding.
-    :param in_channels: channels in the input Tensor.
+    :param in_channels: in_channels in the input Tensor.
     :param model_channels: base channel count for the model.
     :param out_channels: channels in the output Tensor.
     :param num_res_blocks: number of residual blocks per downsample.
@@ -304,67 +302,45 @@ class UNetModel(nn.Module):
                                of heads for upsampling. Deprecated.
     :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
     :param resblock_updown: use residual blocks for up/downsampling.
-    :param use_new_attention_order: use a different attention pattern for potentially
-                                    increased efficiency.
     """
 
-    def __init__(
-        self,
-        image_size, # not used in UNetModel
-        in_channels,
-        model_channels,
-        out_channels,
-        num_res_blocks,
-        attention_resolutions,
-        dropout=0,
-        channel_mult=(1, 2, 4, 8),
-        conv_resample=True,
-        dims=3,
-        num_classes=None,
-        use_checkpoint=False,
-        use_fp16=False,
-        num_heads=-1,
-        num_head_channels=-1,
-        num_heads_upsample=-1,
-        use_scale_shift_norm=False,
-        resblock_updown=False,
-        transformer_depth=1,              # custom transformer support
-        context_dim=None,                 # custom transformer support
-        legacy=True,
-        # temporal related
-        kernel_size_t=1,
-        padding_t=1,
-        use_temporal_transformer=True,
-        temporal_length=None,
-        use_relative_position=False,
-        cross_attn_on_tempoal=False,
-        temporal_crossattn_type="crossattn",
-        order="stst",
-        nonlinearity_type='silu',
-        temporalcrossfirst=False,
-        split_stcontext=False,
-        temporal_context_dim=None,
-        use_tempoal_causal_attn=False,
-        ST_transformer_module='attention_temporal',
-        ST_transformer_class='SpatialTemporalTransformer',
-        **kwargs,
-    ):
-        super().__init__()
-        assert(use_temporal_transformer)
-        if context_dim is not None:
-            if type(context_dim) == ListConfig:
-                context_dim = list(context_dim)
-
-        if num_heads_upsample == -1:
-            num_heads_upsample = num_heads
-
+    def __init__(self,
+                 in_channels,
+                 model_channels,
+                 out_channels,
+                 num_res_blocks,
+                 attention_resolutions,
+                 dropout=0.0,
+                 channel_mult=(1, 2, 4, 8),
+                 conv_resample=True,
+                 dims=2,
+                 context_dim=None,
+                 use_scale_shift_norm=False,
+                 resblock_updown=False,
+                 num_heads=-1,
+                 num_head_channels=-1,
+                 transformer_depth=1,
+                 use_linear=False,
+                 use_checkpoint=False,
+                 temporal_conv=False,
+                 tempspatial_aware=False,
+                 temporal_attention=True,
+                 temporal_selfatt_only=True,
+                 use_relative_position=True,
+                 use_causal_attention=False,
+                 temporal_length=None,
+                 use_fp16=False,
+                 addition_attention=False,
+                 use_image_attention=False,
+                 temporal_transformer_depth=1,
+                 fps_cond=False,
+                ):
+        super(UNetModel, self).__init__()
         if num_heads == -1:
             assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
-
         if num_head_channels == -1:
             assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
 
-        self.image_size = image_size
         self.in_channels = in_channels
         self.model_channels = model_channels
         self.out_channels = out_channels
@@ -373,65 +349,55 @@ class UNetModel(nn.Module):
         self.dropout = dropout
         self.channel_mult = channel_mult
         self.conv_resample = conv_resample
-        self.num_classes = num_classes
+        self.temporal_attention = temporal_attention
+        time_embed_dim = model_channels * 4
         self.use_checkpoint = use_checkpoint
-        self.dtype = th.float16 if use_fp16 else th.float32
-        self.num_heads = num_heads
-        self.num_head_channels = num_head_channels
-        self.num_heads_upsample = num_heads_upsample
-
-        self.use_relative_position = use_relative_position
-        self.temporal_length = temporal_length
-        self.cross_attn_on_tempoal = cross_attn_on_tempoal
-        self.temporal_crossattn_type = temporal_crossattn_type
-        self.order = order
-        self.temporalcrossfirst = temporalcrossfirst
-        self.split_stcontext = split_stcontext
-        self.temporal_context_dim = temporal_context_dim
-        self.nonlinearity_type = nonlinearity_type
-        self.use_tempoal_causal_attn = use_tempoal_causal_attn
-        
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.addition_attention=addition_attention
+        self.use_image_attention = use_image_attention
+        self.fps_cond=fps_cond
+
+
 
-        time_embed_dim = model_channels * 4
-        self.time_embed_dim = time_embed_dim
         self.time_embed = nn.Sequential(
             linear(model_channels, time_embed_dim),
-            nonlinearity(nonlinearity_type),
+            nn.SiLU(),
             linear(time_embed_dim, time_embed_dim),
         )
-
-        if self.num_classes is not None:
-            self.label_emb = nn.Embedding(num_classes, time_embed_dim)
-        
-        STTransformerClass = make_spatialtemporal_transformer(module_name=ST_transformer_module, 
-            class_name=ST_transformer_class)
+        if self.fps_cond:
+            self.fps_embedding = nn.Sequential(
+                linear(model_channels, time_embed_dim),
+                nn.SiLU(),
+                linear(time_embed_dim, time_embed_dim),
+            )
 
         self.input_blocks = nn.ModuleList(
             [
-                TimestepEmbedSequential(
-                    conv_nd(dims, in_channels, model_channels, (kernel_size_t, 3,3), padding=(padding_t, 1,1))
-                )
+                TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))
             ]
         )
-        self._feature_size = model_channels
+        if self.addition_attention:
+            self.init_attn=TimestepEmbedSequential(
+                TemporalTransformer(
+                    model_channels,
+                    n_heads=8,
+                    d_head=num_head_channels,
+                    depth=transformer_depth,
+                    context_dim=context_dim,
+                    use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only, 
+                    causal_attention=use_causal_attention, relative_position=use_relative_position, 
+                    temporal_length=temporal_length))
+            
         input_block_chans = [model_channels]
         ch = model_channels
         ds = 1
         for level, mult in enumerate(channel_mult):
             for _ in range(num_res_blocks):
                 layers = [
-                    ResBlock(
-                        ch,
-                        time_embed_dim,
-                        dropout,
-                        out_channels=mult * model_channels,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                        kernel_size_t=kernel_size_t,
-                        padding_t=padding_t,
-                        nonlinearity_type=nonlinearity_type,
-                        **kwargs
+                    ResBlock(ch, time_embed_dim, dropout,
+                        out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware,
+                        use_temporal_conv=temporal_conv
                     )
                 ]
                 ch = mult * model_channels
@@ -441,120 +407,85 @@ class UNetModel(nn.Module):
                     else:
                         num_heads = ch // num_head_channels
                         dim_head = num_head_channels
-                    if legacy:
-                        dim_head = ch // num_heads if use_temporal_transformer else num_head_channels
-                    layers.append(STTransformerClass(
-                            ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
-                            # temporal related
-                            temporal_length=temporal_length,
-                            use_relative_position=use_relative_position,
-                            cross_attn_on_tempoal=cross_attn_on_tempoal,
-                            temporal_crossattn_type=temporal_crossattn_type,
-                            order=order,
-                            temporalcrossfirst=temporalcrossfirst,
-                            split_stcontext=split_stcontext,
-                            temporal_context_dim=temporal_context_dim,
-                            use_tempoal_causal_attn=use_tempoal_causal_attn,
-                            **kwargs,
-                            ))
+                    layers.append(
+                        SpatialTransformer(ch, num_heads, dim_head, 
+                            depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                            use_checkpoint=use_checkpoint, disable_self_attn=False,
+                            img_cross_attention=self.use_image_attention
+                        )
+                    )
+                    if self.temporal_attention:
+                        layers.append(
+                            TemporalTransformer(ch, num_heads, dim_head,
+                                depth=temporal_transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                                use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only, 
+                                causal_attention=use_causal_attention, relative_position=use_relative_position, 
+                                temporal_length=temporal_length
+                            )
+                        )
                 self.input_blocks.append(TimestepEmbedSequential(*layers))
-                self._feature_size += ch
                 input_block_chans.append(ch)
             if level != len(channel_mult) - 1:
                 out_ch = ch
                 self.input_blocks.append(
                     TimestepEmbedSequential(
-                        ResBlock(
-                            ch,
-                            time_embed_dim,
-                            dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
+                        ResBlock(ch, time_embed_dim, dropout, 
+                            out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint,
                             use_scale_shift_norm=use_scale_shift_norm,
-                            down=True,
-                            kernel_size_t=kernel_size_t,
-                            padding_t=padding_t,
-                            nonlinearity_type=nonlinearity_type,
-                            **kwargs
+                            down=True
                         )
                         if resblock_updown
-                        else Downsample(
-                            ch, conv_resample, dims=dims, out_channels=out_ch, kernel_size_t=kernel_size_t, padding_t=padding_t
-                        )
+                        else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
                     )
                 )
                 ch = out_ch
                 input_block_chans.append(ch)
                 ds *= 2
-                self._feature_size += ch
 
         if num_head_channels == -1:
             dim_head = ch // num_heads
         else:
             num_heads = ch // num_head_channels
             dim_head = num_head_channels
-        if legacy:
-            dim_head = ch // num_heads if use_temporal_transformer else num_head_channels
-        self.middle_block = TimestepEmbedSequential(
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                kernel_size_t=kernel_size_t,
-                padding_t=padding_t,
-                nonlinearity_type=nonlinearity_type,
-                **kwargs
-            ),
-            STTransformerClass(
-                            ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
-                            # temporal related
-                            temporal_length=temporal_length,
-                            use_relative_position=use_relative_position,
-                            cross_attn_on_tempoal=cross_attn_on_tempoal,
-                            temporal_crossattn_type=temporal_crossattn_type,
-                            order=order,
-                            temporalcrossfirst=temporalcrossfirst,
-                            split_stcontext=split_stcontext,
-                            temporal_context_dim=temporal_context_dim,
-                            use_tempoal_causal_attn=use_tempoal_causal_attn,
-                            **kwargs,
-                        ),
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                kernel_size_t=kernel_size_t,
-                padding_t=padding_t,
-                nonlinearity_type=nonlinearity_type,
-                **kwargs
+        layers = [
+            ResBlock(ch, time_embed_dim, dropout,
+                dims=dims, use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware,
+                use_temporal_conv=temporal_conv
             ),
+            SpatialTransformer(ch, num_heads, dim_head, 
+                depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                use_checkpoint=use_checkpoint, disable_self_attn=False,
+                img_cross_attention=self.use_image_attention
+            )
+        ]
+        if self.temporal_attention:
+            layers.append(
+                TemporalTransformer(ch, num_heads, dim_head,
+                    depth=temporal_transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                    use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only, 
+                    causal_attention=use_causal_attention, relative_position=use_relative_position, 
+                    temporal_length=temporal_length
+                )
+            )
+        layers.append(
+            ResBlock(ch, time_embed_dim, dropout,
+                dims=dims, use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware,
+                use_temporal_conv=temporal_conv
+                )
         )
-        self._feature_size += ch
+        self.middle_block = TimestepEmbedSequential(*layers)
 
         self.output_blocks = nn.ModuleList([])
         for level, mult in list(enumerate(channel_mult))[::-1]:
             for i in range(num_res_blocks + 1):
                 ich = input_block_chans.pop()
                 layers = [
-                    ResBlock(
-                        ch + ich,
-                        time_embed_dim,
-                        dropout,
-                        out_channels=model_channels * mult,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                        kernel_size_t=kernel_size_t,
-                        padding_t=padding_t,
-                        nonlinearity_type=nonlinearity_type,
-                        **kwargs
+                    ResBlock(ch + ich, time_embed_dim, dropout,
+                        out_channels=mult * model_channels, dims=dims, use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm, tempspatial_aware=tempspatial_aware,
+                        use_temporal_conv=temporal_conv
                     )
                 ]
                 ch = model_channels * mult
@@ -564,107 +495,83 @@ class UNetModel(nn.Module):
                     else:
                         num_heads = ch // num_head_channels
                         dim_head = num_head_channels
-                    if legacy:
-                        dim_head = ch // num_heads if use_temporal_transformer else num_head_channels
                     layers.append(
-                        STTransformerClass(
-                            ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,
-                            # temporal related
-                            temporal_length=temporal_length,
-                            use_relative_position=use_relative_position,
-                            cross_attn_on_tempoal=cross_attn_on_tempoal,
-                            temporal_crossattn_type=temporal_crossattn_type,
-                            order=order,
-                            temporalcrossfirst=temporalcrossfirst,
-                            split_stcontext=split_stcontext,
-                            temporal_context_dim=temporal_context_dim,
-                            use_tempoal_causal_attn=use_tempoal_causal_attn,
-                            **kwargs,
+                        SpatialTransformer(ch, num_heads, dim_head, 
+                            depth=transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                            use_checkpoint=use_checkpoint, disable_self_attn=False,
+                            img_cross_attention=self.use_image_attention
                         )
                     )
+                    if self.temporal_attention:
+                        layers.append(
+                            TemporalTransformer(ch, num_heads, dim_head,
+                                depth=temporal_transformer_depth, context_dim=context_dim, use_linear=use_linear,
+                                use_checkpoint=use_checkpoint, only_self_att=temporal_selfatt_only, 
+                                causal_attention=use_causal_attention, relative_position=use_relative_position, 
+                                temporal_length=temporal_length
+                            )
+                        )
                 if level and i == num_res_blocks:
                     out_ch = ch
                     layers.append(
-                        ResBlock(
-                            ch,
-                            time_embed_dim,
-                            dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
+                        ResBlock(ch, time_embed_dim, dropout,
+                            out_channels=out_ch, dims=dims, use_checkpoint=use_checkpoint,
                             use_scale_shift_norm=use_scale_shift_norm,
-                            up=True,
-                            kernel_size_t=kernel_size_t,
-                            padding_t=padding_t,
-                            nonlinearity_type=nonlinearity_type,
-                            **kwargs
+                            up=True
                         )
                         if resblock_updown
-                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch, kernel_size_t=kernel_size_t, padding_t=padding_t)
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
                     )
                     ds //= 2
                 self.output_blocks.append(TimestepEmbedSequential(*layers))
-                self._feature_size += ch
 
         self.out = nn.Sequential(
             normalization(ch),
-            nonlinearity(nonlinearity_type),
-            zero_module(conv_nd(dims, model_channels, out_channels, (kernel_size_t, 3,3), padding=(padding_t, 1,1))),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
         )
-        
 
-    def convert_to_fp16(self):
-        """
-        Convert the torso of the model to float16.
-        """
-        self.input_blocks.apply(convert_module_to_f16)
-        self.middle_block.apply(convert_module_to_f16)
-        self.output_blocks.apply(convert_module_to_f16)
+    def forward(self, x, timesteps, context=None, features_adapter=None, fps=16, **kwargs):
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb)
 
-    def convert_to_fp32(self):
-        """
-        Convert the torso of the model to float32.
-        """
-        self.input_blocks.apply(convert_module_to_f32)
-        self.middle_block.apply(convert_module_to_f32)
-        self.output_blocks.apply(convert_module_to_f32)
+        if self.fps_cond:
+            if type(fps) == int:
+                fps = torch.full_like(timesteps, fps)
+            fps_emb = timestep_embedding(fps,self.model_channels, repeat_only=False)
+            emb += self.fps_embedding(fps_emb)
 
-    def forward(self, x, timesteps=None, time_emb_replace=None, context=None, features_adapter=None, y=None, **kwargs):
-        """
-        Apply the model to an input batch.
-        :param x: an [N x C x ...] Tensor of inputs.
-        :param timesteps: a 1-D batch of timesteps.
-        :param context: conditioning plugged in via crossattn
-        :param y: an [N] Tensor of labels, if class-conditional.
-        :return: an [N x C x ...] Tensor of outputs.
-        """
-        
-        hs = []
-        if time_emb_replace is None:
-            t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
-            emb = self.time_embed(t_emb)
-        else:
-            emb = time_emb_replace
-        
-        if y is not None: # if class-conditional model, inject class labels
-            assert y.shape == (x.shape[0],)
-            emb = emb + self.label_emb(y)
+        b,_,t,_,_ = x.shape
+        ## repeat t times for context [(b t) 77 768] & time embedding
+        context = context.repeat_interleave(repeats=t, dim=0)
+        emb = emb.repeat_interleave(repeats=t, dim=0)
+
+        ## always in shape (b t) c h w, except for temporal layer
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
 
         h = x.type(self.dtype)
         adapter_idx = 0
+        hs = []
         for id, module in enumerate(self.input_blocks):
-            h = module(h, emb, context, **kwargs)
+            h = module(h, emb, context=context, batch_size=b)
+            if id ==0 and self.addition_attention:
+                h = self.init_attn(h, emb, context=context, batch_size=b)
             ## plug-in adapter features
             if ((id+1)%3 == 0) and features_adapter is not None:
                 h = h + features_adapter[adapter_idx]
                 adapter_idx += 1
             hs.append(h)
         if features_adapter is not None:
-            assert len(features_adapter)==adapter_idx, 'Mismatch features adapter'
+            assert len(features_adapter)==adapter_idx, 'Wrong features_adapter'
 
-        h = self.middle_block(h, emb, context, **kwargs)
+        h = self.middle_block(h, emb, context=context, batch_size=b)
         for module in self.output_blocks:
-            h = th.cat([h, hs.pop()], dim=1)
-            h = module(h, emb, context, **kwargs)
+            h = torch.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context=context, batch_size=b)
         h = h.type(x.dtype)
-        return self.out(h)
+        y = self.out(h)
+        
+        # reshape back to (b c t h w)
+        y = rearrange(y, '(b t) c h w -> b c t h w', b=b)
+        return y
+    

lvdm/modules/x_transformer.py

@@ -0,0 +1,640 @@
+"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers"""
+from functools import partial
+from inspect import isfunction
+from collections import namedtuple
+from einops import rearrange, repeat
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+# constants
+DEFAULT_DIM_HEAD = 64
+
+Intermediates = namedtuple('Intermediates', [
+    'pre_softmax_attn',
+    'post_softmax_attn'
+])
+
+LayerIntermediates = namedtuple('Intermediates', [
+    'hiddens',
+    'attn_intermediates'
+])
+
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len):
+        super().__init__()
+        self.emb = nn.Embedding(max_seq_len, dim)
+        self.init_()
+
+    def init_(self):
+        nn.init.normal_(self.emb.weight, std=0.02)
+
+    def forward(self, x):
+        n = torch.arange(x.shape[1], device=x.device)
+        return self.emb(n)[None, :, :]
+
+
+class FixedPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x, seq_dim=1, offset=0):
+        t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset
+        sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq)
+        emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
+        return emb[None, :, :]
+
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def always(val):
+    def inner(*args, **kwargs):
+        return val
+    return inner
+
+
+def not_equals(val):
+    def inner(x):
+        return x != val
+    return inner
+
+
+def equals(val):
+    def inner(x):
+        return x == val
+    return inner
+
+
+def max_neg_value(tensor):
+    return -torch.finfo(tensor.dtype).max
+
+
+# keyword argument helpers
+
+def pick_and_pop(keys, d):
+    values = list(map(lambda key: d.pop(key), keys))
+    return dict(zip(keys, values))
+
+
+def group_dict_by_key(cond, d):
+    return_val = [dict(), dict()]
+    for key in d.keys():
+        match = bool(cond(key))
+        ind = int(not match)
+        return_val[ind][key] = d[key]
+    return (*return_val,)
+
+
+def string_begins_with(prefix, str):
+    return str.startswith(prefix)
+
+
+def group_by_key_prefix(prefix, d):
+    return group_dict_by_key(partial(string_begins_with, prefix), d)
+
+
+def groupby_prefix_and_trim(prefix, d):
+    kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
+    kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
+    return kwargs_without_prefix, kwargs
+
+
+# classes
+class Scale(nn.Module):
+    def __init__(self, value, fn):
+        super().__init__()
+        self.value = value
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        x, *rest = self.fn(x, **kwargs)
+        return (x * self.value, *rest)
+
+
+class Rezero(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+        self.g = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, **kwargs):
+        x, *rest = self.fn(x, **kwargs)
+        return (x * self.g, *rest)
+
+
+class ScaleNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class Residual(nn.Module):
+    def forward(self, x, residual):
+        return x + residual
+
+
+class GRUGating(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gru = nn.GRUCell(dim, dim)
+
+    def forward(self, x, residual):
+        gated_output = self.gru(
+            rearrange(x, 'b n d -> (b n) d'),
+            rearrange(residual, 'b n d -> (b n) d')
+        )
+
+        return gated_output.reshape_as(x)
+
+
+# feedforward
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            nn.GELU()
+        ) if not glu else GEGLU(dim, inner_dim)
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# attention.
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head=DEFAULT_DIM_HEAD,
+            heads=8,
+            causal=False,
+            mask=None,
+            talking_heads=False,
+            sparse_topk=None,
+            use_entmax15=False,
+            num_mem_kv=0,
+            dropout=0.,
+            on_attn=False
+    ):
+        super().__init__()
+        if use_entmax15:
+            raise NotImplementedError("Check out entmax activation instead of softmax activation!")
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.causal = causal
+        self.mask = mask
+
+        inner_dim = dim_head * heads
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(dim, inner_dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+        # talking heads
+        self.talking_heads = talking_heads
+        if talking_heads:
+            self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+            self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+
+        # explicit topk sparse attention
+        self.sparse_topk = sparse_topk
+
+        # entmax
+        #self.attn_fn = entmax15 if use_entmax15 else F.softmax
+        self.attn_fn = F.softmax
+
+        # add memory key / values
+        self.num_mem_kv = num_mem_kv
+        if num_mem_kv > 0:
+            self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+            self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+
+        # attention on attention
+        self.attn_on_attn = on_attn
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim)
+
+    def forward(
+            self,
+            x,
+            context=None,
+            mask=None,
+            context_mask=None,
+            rel_pos=None,
+            sinusoidal_emb=None,
+            prev_attn=None,
+            mem=None
+    ):
+        b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device
+        kv_input = default(context, x)
+
+        q_input = x
+        k_input = kv_input
+        v_input = kv_input
+
+        if exists(mem):
+            k_input = torch.cat((mem, k_input), dim=-2)
+            v_input = torch.cat((mem, v_input), dim=-2)
+
+        if exists(sinusoidal_emb):
+            # in shortformer, the query would start at a position offset depending on the past cached memory
+            offset = k_input.shape[-2] - q_input.shape[-2]
+            q_input = q_input + sinusoidal_emb(q_input, offset=offset)
+            k_input = k_input + sinusoidal_emb(k_input)
+
+        q = self.to_q(q_input)
+        k = self.to_k(k_input)
+        v = self.to_v(v_input)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+        input_mask = None
+        if any(map(exists, (mask, context_mask))):
+            q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
+            k_mask = q_mask if not exists(context) else context_mask
+            k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool())
+            q_mask = rearrange(q_mask, 'b i -> b () i ()')
+            k_mask = rearrange(k_mask, 'b j -> b () () j')
+            input_mask = q_mask * k_mask
+
+        if self.num_mem_kv > 0:
+            mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v))
+            k = torch.cat((mem_k, k), dim=-2)
+            v = torch.cat((mem_v, v), dim=-2)
+            if exists(input_mask):
+                input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
+
+        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+        mask_value = max_neg_value(dots)
+
+        if exists(prev_attn):
+            dots = dots + prev_attn
+
+        pre_softmax_attn = dots
+
+        if talking_heads:
+            dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous()
+
+        if exists(rel_pos):
+            dots = rel_pos(dots)
+
+        if exists(input_mask):
+            dots.masked_fill_(~input_mask, mask_value)
+            del input_mask
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            r = torch.arange(i, device=device)
+            mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j')
+            mask = F.pad(mask, (j - i, 0), value=False)
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
+            top, _ = dots.topk(self.sparse_topk, dim=-1)
+            vk = top[..., -1].unsqueeze(-1).expand_as(dots)
+            mask = dots < vk
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        attn = self.attn_fn(dots, dim=-1)
+        post_softmax_attn = attn
+
+        attn = self.dropout(attn)
+
+        if talking_heads:
+            attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous()
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        intermediates = Intermediates(
+            pre_softmax_attn=pre_softmax_attn,
+            post_softmax_attn=post_softmax_attn
+        )
+
+        return self.to_out(out), intermediates
+
+
+class AttentionLayers(nn.Module):
+    def __init__(
+            self,
+            dim,
+            depth,
+            heads=8,
+            causal=False,
+            cross_attend=False,
+            only_cross=False,
+            use_scalenorm=False,
+            use_rmsnorm=False,
+            use_rezero=False,
+            rel_pos_num_buckets=32,
+            rel_pos_max_distance=128,
+            position_infused_attn=False,
+            custom_layers=None,
+            sandwich_coef=None,
+            par_ratio=None,
+            residual_attn=False,
+            cross_residual_attn=False,
+            macaron=False,
+            pre_norm=True,
+            gate_residual=False,
+            **kwargs
+    ):
+        super().__init__()
+        ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs)
+        attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs)
+
+        dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
+
+        self.dim = dim
+        self.depth = depth
+        self.layers = nn.ModuleList([])
+
+        self.has_pos_emb = position_infused_attn
+        self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None
+        self.rotary_pos_emb = always(None)
+
+        assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
+        self.rel_pos = None
+
+        self.pre_norm = pre_norm
+
+        self.residual_attn = residual_attn
+        self.cross_residual_attn = cross_residual_attn
+
+        norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
+        norm_class = RMSNorm if use_rmsnorm else norm_class
+        norm_fn = partial(norm_class, dim)
+
+        norm_fn = nn.Identity if use_rezero else norm_fn
+        branch_fn = Rezero if use_rezero else None
+
+        if cross_attend and not only_cross:
+            default_block = ('a', 'c', 'f')
+        elif cross_attend and only_cross:
+            default_block = ('c', 'f')
+        else:
+            default_block = ('a', 'f')
+
+        if macaron:
+            default_block = ('f',) + default_block
+
+        if exists(custom_layers):
+            layer_types = custom_layers
+        elif exists(par_ratio):
+            par_depth = depth * len(default_block)
+            assert 1 < par_ratio <= par_depth, 'par ratio out of range'
+            default_block = tuple(filter(not_equals('f'), default_block))
+            par_attn = par_depth // par_ratio
+            depth_cut = par_depth * 2 // 3  # 2 / 3 attention layer cutoff suggested by PAR paper
+            par_width = (depth_cut + depth_cut // par_attn) // par_attn
+            assert len(default_block) <= par_width, 'default block is too large for par_ratio'
+            par_block = default_block + ('f',) * (par_width - len(default_block))
+            par_head = par_block * par_attn
+            layer_types = par_head + ('f',) * (par_depth - len(par_head))
+        elif exists(sandwich_coef):
+            assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth'
+            layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef
+        else:
+            layer_types = default_block * depth
+
+        self.layer_types = layer_types
+        self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
+
+        for layer_type in self.layer_types:
+            if layer_type == 'a':
+                layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
+            elif layer_type == 'c':
+                layer = Attention(dim, heads=heads, **attn_kwargs)
+            elif layer_type == 'f':
+                layer = FeedForward(dim, **ff_kwargs)
+                layer = layer if not macaron else Scale(0.5, layer)
+            else:
+                raise Exception(f'invalid layer type {layer_type}')
+
+            if isinstance(layer, Attention) and exists(branch_fn):
+                layer = branch_fn(layer)
+
+            if gate_residual:
+                residual_fn = GRUGating(dim)
+            else:
+                residual_fn = Residual()
+
+            self.layers.append(nn.ModuleList([
+                norm_fn(),
+                layer,
+                residual_fn
+            ]))
+
+    def forward(
+            self,
+            x,
+            context=None,
+            mask=None,
+            context_mask=None,
+            mems=None,
+            return_hiddens=False
+    ):
+        hiddens = []
+        intermediates = []
+        prev_attn = None
+        prev_cross_attn = None
+
+        mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
+
+        for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)):
+            is_last = ind == (len(self.layers) - 1)
+
+            if layer_type == 'a':
+                hiddens.append(x)
+                layer_mem = mems.pop(0)
+
+            residual = x
+
+            if self.pre_norm:
+                x = norm(x)
+
+            if layer_type == 'a':
+                out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos,
+                                   prev_attn=prev_attn, mem=layer_mem)
+            elif layer_type == 'c':
+                out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn)
+            elif layer_type == 'f':
+                out = block(x)
+
+            x = residual_fn(out, residual)
+
+            if layer_type in ('a', 'c'):
+                intermediates.append(inter)
+
+            if layer_type == 'a' and self.residual_attn:
+                prev_attn = inter.pre_softmax_attn
+            elif layer_type == 'c' and self.cross_residual_attn:
+                prev_cross_attn = inter.pre_softmax_attn
+
+            if not self.pre_norm and not is_last:
+                x = norm(x)
+
+        if return_hiddens:
+            intermediates = LayerIntermediates(
+                hiddens=hiddens,
+                attn_intermediates=intermediates
+            )
+
+            return x, intermediates
+
+        return x
+
+
+class Encoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert 'causal' not in kwargs, 'cannot set causality on encoder'
+        super().__init__(causal=False, **kwargs)
+
+
+
+class TransformerWrapper(nn.Module):
+    def __init__(
+            self,
+            *,
+            num_tokens,
+            max_seq_len,
+            attn_layers,
+            emb_dim=None,
+            max_mem_len=0.,
+            emb_dropout=0.,
+            num_memory_tokens=None,
+            tie_embedding=False,
+            use_pos_emb=True
+    ):
+        super().__init__()
+        assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
+
+        dim = attn_layers.dim
+        emb_dim = default(emb_dim, dim)
+
+        self.max_seq_len = max_seq_len
+        self.max_mem_len = max_mem_len
+        self.num_tokens = num_tokens
+
+        self.token_emb = nn.Embedding(num_tokens, emb_dim)
+        self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if (
+                    use_pos_emb and not attn_layers.has_pos_emb) else always(0)
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+
+        self.init_()
+
+        self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t()
+
+        # memory tokens (like [cls]) from Memory Transformers paper
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.num_memory_tokens = num_memory_tokens
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+            # let funnel encoder know number of memory tokens, if specified
+            if hasattr(attn_layers, 'num_memory_tokens'):
+                attn_layers.num_memory_tokens = num_memory_tokens
+
+    def init_(self):
+        nn.init.normal_(self.token_emb.weight, std=0.02)
+
+    def forward(
+            self,
+            x,
+            return_embeddings=False,
+            mask=None,
+            return_mems=False,
+            return_attn=False,
+            mems=None,
+            **kwargs
+    ):
+        b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
+        x = self.token_emb(x)
+        x += self.pos_emb(x)
+        x = self.emb_dropout(x)
+
+        x = self.project_emb(x)
+
+        if num_mem > 0:
+            mem = repeat(self.memory_tokens, 'n d -> b n d', b=b)
+            x = torch.cat((mem, x), dim=1)
+
+            # auto-handle masking after appending memory tokens
+            if exists(mask):
+                mask = F.pad(mask, (num_mem, 0), value=True)
+
+        x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs)
+        x = self.norm(x)
+
+        mem, x = x[:, :num_mem], x[:, num_mem:]
+
+        out = self.to_logits(x) if not return_embeddings else x
+
+        if return_mems:
+            hiddens = intermediates.hiddens
+            new_mems = list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens))) if exists(mems) else hiddens
+            new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems))
+            return out, new_mems
+
+        if return_attn:
+            attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
+            return out, attn_maps
+
+        return out
+

lvdm/utils/common_utils.py

@@ -1,132 +0,0 @@
-
-import importlib
-
-import torch
-import numpy as np
-
-from inspect import isfunction
-from PIL import Image, ImageDraw, ImageFont
-
-
-def str2bool(v):
-    if isinstance(v, bool):
-        return v
-    if v.lower() in ('yes', 'true', 't', 'y', '1'):
-        return True
-    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
-        return False
-    else:
-        raise ValueError('Boolean value expected.')
-
-
-def instantiate_from_config(config):
-    if not "target" in config:
-        if config == '__is_first_stage__':
-            return None
-        elif config == "__is_unconditional__":
-            return None
-        raise KeyError("Expected key `target` to instantiate.")
-
-    return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-def get_obj_from_str(string, reload=False):
-    module, cls = string.rsplit(".", 1)
-    if reload:
-        module_imp = importlib.import_module(module)
-        importlib.reload(module_imp)
-    return getattr(importlib.import_module(module, package=None), cls)
-
-def log_txt_as_img(wh, xc, size=10):
-    # wh a tuple of (width, height)
-    # xc a list of captions to plot
-    b = len(xc)
-    txts = list()
-    for bi in range(b):
-        txt = Image.new("RGB", wh, color="white")
-        draw = ImageDraw.Draw(txt)
-        font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
-        nc = int(40 * (wh[0] / 256))
-        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
-
-        try:
-            draw.text((0, 0), lines, fill="black", font=font)
-        except UnicodeEncodeError:
-            print("Cant encode string for logging. Skipping.")
-
-        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
-        txts.append(txt)
-    txts = np.stack(txts)
-    txts = torch.tensor(txts)
-    return txts
-
-
-def ismap(x):
-    if not isinstance(x, torch.Tensor):
-        return False
-    return (len(x.shape) == 4) and (x.shape[1] > 3)
-
-
-def isimage(x):
-    if not isinstance(x,torch.Tensor):
-        return False
-    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
-
-
-def exists(x):
-    return x is not None
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
-def mean_flat(tensor):
-    """
-    https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
-    Take the mean over all non-batch dimensions.
-    """
-    return tensor.mean(dim=list(range(1, len(tensor.shape))))
-
-
-def count_params(model, verbose=False):
-    total_params = sum(p.numel() for p in model.parameters())
-    if verbose:
-        print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
-    return total_params
-
-
-def instantiate_from_config(config):
-    if not "target" in config:
-        if config == '__is_first_stage__':
-            return None
-        elif config == "__is_unconditional__":
-            return None
-        raise KeyError("Expected key `target` to instantiate.")
-
-    if "instantiate_with_dict" in config and config["instantiate_with_dict"]:
-        # input parameter is one dict
-        return get_obj_from_str(config["target"])(config.get("params", dict()), **kwargs)
-    else:
-        return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-
-def get_obj_from_str(string, reload=False):
-    module, cls = string.rsplit(".", 1)
-    if reload:
-        module_imp = importlib.import_module(module)
-        importlib.reload(module_imp)
-    return getattr(importlib.import_module(module, package=None), cls)
-
-
-def check_istarget(name, para_list):
-    """ 
-    name: full name of source para
-    para_list: partial name of target para 
-    """
-    istarget=False
-    for para in para_list:
-        if para in name:
-            return True
-    return istarget

lvdm/utils/dist_utils.py

@@ -1,19 +0,0 @@
-import torch
-import torch.distributed as dist
-
-def setup_dist(local_rank):
-    if dist.is_initialized():
-        return
-    torch.cuda.set_device(local_rank)
-    torch.distributed.init_process_group(
-        'nccl',
-        init_method='env://'
-    )
-
-def gather_data(data, return_np=True):
-    ''' gather data from multiple processes to one list '''
-    data_list = [torch.zeros_like(data) for _ in range(dist.get_world_size())]
-    dist.all_gather(data_list, data)  # gather not supported with NCCL
-    if return_np:
-        data_list = [data.cpu().numpy() for data in data_list]
-    return data_list

lvdm/utils/saving_utils.py

@@ -1,269 +0,0 @@
-import numpy as np
-import cv2
-import os
-import time
-import imageio
-from tqdm import tqdm
-from PIL import Image
-import os
-import sys
-sys.path.insert(1, os.path.join(sys.path[0], '..'))
-import torch
-import torchvision
-from torchvision.utils import make_grid
-from torch import Tensor
-from torchvision.transforms.functional import to_tensor
-
-
-def tensor_to_mp4(video, savepath, fps, rescale=True, nrow=None):
-    """
-    video: torch.Tensor, b,c,t,h,w, 0-1
-    if -1~1, enable rescale=True
-    """
-    n = video.shape[0]
-    video = video.permute(2, 0, 1, 3, 4) # t,n,c,h,w
-    nrow = int(np.sqrt(n)) if nrow is None else nrow
-    frame_grids = [torchvision.utils.make_grid(framesheet, nrow=nrow) for framesheet in video] # [3, grid_h, grid_w]
-    grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [T, 3, grid_h, grid_w]
-    grid = torch.clamp(grid.float(), -1., 1.)
-    if rescale:
-        grid = (grid + 1.0) / 2.0
-    grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, 3, grid_h, grid_w] -> [T, grid_h, grid_w, 3]
-    #print(f'Save video to {savepath}')
-    torchvision.io.write_video(savepath, grid, fps=fps, video_codec='h264', options={'crf': '10'})
-
-# ----------------------------------------------------------------------------------------------
-def savenp2sheet(imgs, savepath, nrow=None):
-    """ save multiple imgs (in numpy array type) to a img sheet.
-        img sheet is one row.
-
-    imgs: 
-        np array of size [N, H, W, 3] or List[array] with array size = [H,W,3] 
-    """
-    if imgs.ndim == 4:
-        img_list = [imgs[i] for i in range(imgs.shape[0])]
-        imgs = img_list
-    
-    imgs_new = []
-    for i, img in enumerate(imgs):
-        if img.ndim == 3 and img.shape[0] == 3:
-            img = np.transpose(img,(1,2,0))
-        
-        assert(img.ndim == 3 and img.shape[-1] == 3), img.shape # h,w,3
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        imgs_new.append(img)
-    n = len(imgs)
-    if nrow is not None:
-        n_cols = nrow
-    else:
-        n_cols=int(n**0.5)
-    n_rows=int(np.ceil(n/n_cols))
-    print(n_cols)
-    print(n_rows)
-
-    imgsheet = cv2.vconcat([cv2.hconcat(imgs_new[i*n_cols:(i+1)*n_cols]) for i in range(n_rows)])
-    cv2.imwrite(savepath, imgsheet)
-    print(f'saved in {savepath}')
-
-# ----------------------------------------------------------------------------------------------
-def save_np_to_img(img, path, norm=True):
-    if norm:
-        img = (img + 1) / 2 * 255
-    img = img.astype(np.uint8)
-    image = Image.fromarray(img)
-    image.save(path, q=95)
-
-# ----------------------------------------------------------------------------------------------
-def npz_to_imgsheet_5d(data_path, res_dir, nrow=None,):
-    if isinstance(data_path, str):
-        imgs = np.load(data_path)['arr_0'] # NTHWC
-    elif isinstance(data_path, np.ndarray):
-        imgs = data_path
-    else:
-        raise Exception
-    
-    if os.path.isdir(res_dir):
-        res_path = os.path.join(res_dir, f'samples.jpg')
-    else:
-        assert(res_dir.endswith('.jpg'))
-        res_path = res_dir
-    imgs = np.concatenate([imgs[i] for i in range(imgs.shape[0])], axis=0)
-    savenp2sheet(imgs, res_path, nrow=nrow)
-
-# ----------------------------------------------------------------------------------------------
-def npz_to_imgsheet_4d(data_path, res_path, nrow=None,):
-    if isinstance(data_path, str):
-        imgs = np.load(data_path)['arr_0'] # NHWC
-    elif isinstance(data_path, np.ndarray):
-        imgs = data_path
-    else:
-        raise Exception
-    print(imgs.shape)
-    savenp2sheet(imgs, res_path, nrow=nrow)
-
-
-# ----------------------------------------------------------------------------------------------
-def tensor_to_imgsheet(tensor, save_path):
-    """ 
-        save a batch of videos in one image sheet with shape of [batch_size * num_frames].
-        data: [b,c,t,h,w]
-    """
-    assert(tensor.dim() == 5)
-    b,c,t,h,w = tensor.shape
-    imgs = [tensor[bi,:,ti, :, :] for bi in range(b) for ti in range(t)]
-    torchvision.utils.save_image(imgs, save_path, normalize=True, nrow=t)
-
-
-# ----------------------------------------------------------------------------------------------
-def npz_to_frames(data_path, res_dir, norm, num_frames=None, num_samples=None):
-    start = time.time()
-    arr = np.load(data_path)
-    imgs = arr['arr_0'] # [N, T, H, W, 3]
-    print('original data shape: ', imgs.shape)
-
-    if num_samples is not None:
-        imgs = imgs[:num_samples, :, :, :, :]
-        print('after sample selection: ', imgs.shape)
-    
-    if num_frames is not None:
-        imgs = imgs[:, :num_frames, :, :, :]
-        print('after frame selection: ', imgs.shape)
-
-    for vid in tqdm(range(imgs.shape[0]), desc='Video'):
-        video_dir = os.path.join(res_dir, f'video{vid:04d}')
-        os.makedirs(video_dir, exist_ok=True)
-        for fid in range(imgs.shape[1]):
-            frame = imgs[vid, fid, :, :, :] #HW3
-            save_np_to_img(frame, os.path.join(video_dir, f'frame{fid:04d}.jpg'), norm=norm)
-    print('Finish')
-    print(f'Total time = {time.time()- start}')
-
-# ----------------------------------------------------------------------------------------------
-def npz_to_gifs(data_path, res_dir, duration=0.2, start_idx=0, num_videos=None, mode='gif'):
-    os.makedirs(res_dir, exist_ok=True)
-    if isinstance(data_path, str):
-        imgs = np.load(data_path)['arr_0'] # NTHWC
-    elif isinstance(data_path, np.ndarray):
-        imgs = data_path
-    else:
-        raise Exception
-
-    for i in range(imgs.shape[0]):
-        frames = [imgs[i,j,:,:,:] for j in range(imgs[i].shape[0])] # [(h,w,3)]
-        if mode == 'gif':
-            imageio.mimwrite(os.path.join(res_dir, f'samples_{start_idx+i}.gif'), frames, format='GIF', duration=duration)
-        elif mode == 'mp4':
-            frames = [torch.from_numpy(frame) for frame in frames]
-            frames = torch.stack(frames, dim=0).to(torch.uint8) # [T, H, W, C]
-            torchvision.io.write_video(os.path.join(res_dir, f'samples_{start_idx+i}.mp4'),
-                frames, fps=0.5, video_codec='h264', options={'crf': '10'})
-        if i+ 1 == num_videos:
-            break
-
-# ----------------------------------------------------------------------------------------------
-def fill_with_black_squares(video, desired_len: int) -> Tensor:
-    if len(video) >= desired_len:
-        return video
-
-    return torch.cat([
-        video,
-        torch.zeros_like(video[0]).unsqueeze(0).repeat(desired_len - len(video), 1, 1, 1),
-    ], dim=0)
-
-# ----------------------------------------------------------------------------------------------
-def load_num_videos(data_path, num_videos):
-    # data_path can be either data_path of np array 
-    if isinstance(data_path, str):
-        videos = np.load(data_path)['arr_0'] # NTHWC
-    elif isinstance(data_path, np.ndarray):
-        videos = data_path
-    else:
-        raise Exception
-
-    if num_videos is not None:
-        videos = videos[:num_videos, :, :, :, :]
-    return videos
-
-# ----------------------------------------------------------------------------------------------
-def npz_to_video_grid(data_path, out_path, num_frames=None, fps=8, num_videos=None, nrow=None, verbose=True):
-    if isinstance(data_path, str):
-        videos = load_num_videos(data_path, num_videos)
-    elif isinstance(data_path, np.ndarray):
-        videos = data_path
-    else:
-        raise Exception
-    n,t,h,w,c = videos.shape
-
-    videos_th = []
-    for i in range(n):
-        video = videos[i, :,:,:,:]
-        images = [video[j, :,:,:] for j in range(t)]
-        images = [to_tensor(img) for img in images]
-        video = torch.stack(images)
-        videos_th.append(video)
-    
-    if num_frames is None:
-        num_frames = videos.shape[1]
-    if verbose:
-        videos = [fill_with_black_squares(v, num_frames) for v in tqdm(videos_th, desc='Adding empty frames')] # NTCHW
-    else:
-        videos = [fill_with_black_squares(v, num_frames) for v in videos_th] # NTCHW
-
-    frame_grids = torch.stack(videos).permute(1, 0, 2, 3, 4) # [T, N, C, H, W]
-    if nrow is None:
-        nrow = int(np.ceil(np.sqrt(n)))
-    if verbose:
-        frame_grids = [make_grid(fs, nrow=nrow) for fs in tqdm(frame_grids, desc='Making grids')]
-    else:
-        frame_grids = [make_grid(fs, nrow=nrow) for fs in frame_grids]
-
-    if os.path.dirname(out_path) != "":
-        os.makedirs(os.path.dirname(out_path), exist_ok=True)
-    frame_grids = (torch.stack(frame_grids) * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, H, W, C]
-    torchvision.io.write_video(out_path, frame_grids, fps=fps, video_codec='h264', options={'crf': '10'})
-
-# ----------------------------------------------------------------------------------------------
-def npz_to_gif_grid(data_path, out_path, n_cols=None, num_videos=20):
-    arr = np.load(data_path)
-    imgs = arr['arr_0'] # [N, T, H, W, 3]
-    imgs = imgs[:num_videos]
-    n, t, h, w, c = imgs.shape
-    assert(n == num_videos)
-    n_cols = n_cols if n_cols else imgs.shape[0]
-    n_rows = np.ceil(imgs.shape[0] / n_cols).astype(np.int8)
-    H, W = h * n_rows, w * n_cols
-    grid = np.zeros((t, H, W, c), dtype=np.uint8)
-
-    for i in range(n_rows):
-        for j in range(n_cols):
-            if i*n_cols+j < imgs.shape[0]:
-                grid[:, i*h:(i+1)*h, j*w:(j+1)*w, :] = imgs[i*n_cols+j, :, :, :, :]
-    
-    videos = [grid[i] for i in range(grid.shape[0])] # grid: TH'W'C
-    imageio.mimwrite(out_path, videos, format='GIF', duration=0.5,palettesize=256)
-
-
-# ----------------------------------------------------------------------------------------------
-def torch_to_video_grid(videos, out_path, num_frames, fps, num_videos=None, nrow=None, verbose=True):
-    """
-    videos: -1 ~ 1, torch.Tensor, BCTHW
-    """
-    n,t,h,w,c = videos.shape
-    videos_th = [videos[i, ...] for i in range(n)]
-    if verbose:
-        videos = [fill_with_black_squares(v, num_frames) for v in tqdm(videos_th, desc='Adding empty frames')] # NTCHW
-    else:
-        videos = [fill_with_black_squares(v, num_frames) for v in videos_th] # NTCHW
-
-    frame_grids = torch.stack(videos).permute(1, 0, 2, 3, 4) # [T, N, C, H, W]
-    if nrow is None:
-        nrow = int(np.ceil(np.sqrt(n)))
-    if verbose:
-        frame_grids = [make_grid(fs, nrow=nrow) for fs in tqdm(frame_grids, desc='Making grids')]
-    else:
-        frame_grids = [make_grid(fs, nrow=nrow) for fs in frame_grids]
-
-    if os.path.dirname(out_path) != "":
-        os.makedirs(os.path.dirname(out_path), exist_ok=True)
-    frame_grids = ((torch.stack(frame_grids) + 1) / 2 * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, H, W, C]
-    torchvision.io.write_video(out_path, frame_grids, fps=fps, video_codec='h264', options={'crf': '10'})

prompts/i2v_prompts/test_prompts.txt

@@ -0,0 +1,2 @@
+horses are walking on the grassland
+a boy and a girl are talking on the seashore

prompts/test_prompts.txt

@@ -0,0 +1,2 @@
+A tiger walks in the forest, photorealistic, 4k, high definition
+A boat moving on the sea, flowers and grassland on the shore

requirements.txt

@@ -16,7 +16,8 @@ transformers==4.25.1
 moviepy
 av
 xformers
-gradio==3.24.1
-gradio-client==0.1.2
+gradio
 timm
-# -e .
+scikit-learn 
+open_clip_torch
+kornia

sample_adapter.sh

@@ -1,22 +0,0 @@
-PROMPT="An ostrich walking in the desert, photorealistic, 4k"
-VIDEO="input/flamingo.mp4"
-OUTDIR="results/"
-
-NAME="video_adapter"
-CONFIG_PATH="models/adapter_t2v_depth/model_config.yaml"
-BASE_PATH="models/base_t2v/model.ckpt"
-ADAPTER_PATH="models/adapter_t2v_depth/adapter.pth"
-
-python scripts/sample_text2video_adapter.py \
-    --seed 123 \
-    --ckpt_path $BASE_PATH \
-    --adapter_ckpt $ADAPTER_PATH \
-    --base $CONFIG_PATH \
-    --savedir $OUTDIR/$NAME \
-    --bs 1 --height 256 --width 256 \
-    --frame_stride -1 \
-    --unconditional_guidance_scale 15.0 \
-    --ddim_steps 50 \
-    --ddim_eta 1.0 \
-    --prompt "$PROMPT" \
-    --video $VIDEO

sample_text2video.sh

@@ -1,16 +0,0 @@
-
-PROMPT="astronaut riding a horse" # OR: PROMPT="input/prompts.txt" for sampling multiple prompts
-OUTDIR="results/"
-
-BASE_PATH="models/base_t2v/model.ckpt"
-CONFIG_PATH="models/base_t2v/model_config.yaml"
-
-python scripts/sample_text2video.py \
-    --ckpt_path $BASE_PATH \
-    --config_path $CONFIG_PATH \
-    --prompt "$PROMPT" \
-    --save_dir $OUTDIR \
-    --n_samples 1 \
-    --batch_size 1 \
-    --seed 1000 \
-    --show_denoising_progress