first init

app.py

@@ -0,0 +1,202 @@
+from model.video_diffusion.models.controlnet3d import ControlNet3DModel
+from model.video_diffusion.models.unet_3d_condition import UNetPseudo3DConditionModel
+from model.video_diffusion.pipelines.pipeline_stable_diffusion_controlnet3d import Controlnet3DStableDiffusionPipeline
+from transformers import DPTForDepthEstimation
+from model.annotator.hed import HEDNetwork
+import torch
+from einops import rearrange,repeat
+import imageio
+import numpy as np
+import cv2
+import torch.nn.functional as F
+from PIL import Image
+import argparse
+import tempfile
+import os
+import gradio as gr
+
+
+control_mode = 'depth'
+control_net_path = f"wf-genius/controlavideo-{control_mode}"
+unet = UNetPseudo3DConditionModel.from_pretrained(control_net_path,
+                        torch_dtype = torch.float16,
+                        subfolder='unet',
+                        ).to("cuda") 
+controlnet = ControlNet3DModel.from_pretrained(control_net_path,
+                        torch_dtype = torch.float16,
+                        subfolder='controlnet',
+                        ).to("cuda")
+
+if control_mode == 'depth':
+    annotator_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
+elif control_mode == 'canny':
+    annotator_model = None
+elif control_mode == 'hed':
+    # firstly download from https://huggingface.co/wf-genius/controlavideo-hed/resolve/main/hed-network.pth 
+    annotator_model = HEDNetwork('hed-network.pth').to("cuda")
+
+video_controlnet_pipe = Controlnet3DStableDiffusionPipeline.from_pretrained(control_net_path, unet=unet, 
+                        controlnet=controlnet, annotator_model=annotator_model,
+                        torch_dtype = torch.float16,
+                        ).to("cuda")
+
+
+def to_video(frames, fps: int) -> str:
+    out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+    writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps)
+    for frame in frames:
+        writer.append_data(np.array(frame))
+    writer.close()
+    return out_file.name
+
+def inference(input_video, 
+                prompt,
+                seed,
+                num_inference_steps,
+                guidance_scale,
+                sampling_rate,
+                video_scale,
+                init_noise_thres,
+                each_sample_frame,
+                iter_times,
+                h,
+                w,
+                ):
+    num_sample_frames = iter_times * each_sample_frame
+    testing_prompt = [prompt]
+    np_frames, fps_vid = Controlnet3DStableDiffusionPipeline.get_frames_preprocess(input_video, num_frames=num_sample_frames, sampling_rate=sampling_rate, return_np=True)
+    if control_mode == 'depth':
+        frames = torch.from_numpy(np_frames).div(255) * 2 - 1
+        frames = rearrange(frames, "f h w c -> c f h w").unsqueeze(0)
+        frames = rearrange(frames, 'b c f h w -> (b f) c h w')
+        control_maps = video_controlnet_pipe.get_depth_map(frames, h, w, return_standard_norm=False)  # (b f) 1 h w
+    elif control_mode == 'canny':
+        control_maps = np.stack([cv2.Canny(inp, 100, 200) for inp in np_frames])
+        control_maps = repeat(control_maps, 'f h w -> f c h w',c=1)
+        control_maps = torch.from_numpy(control_maps).div(255)  # 0~1
+    elif control_mode == 'hed':
+        control_maps = np.stack([video_controlnet_pipe.get_hed_map(inp) for inp in np_frames])
+        control_maps = repeat(control_maps, 'f h w -> f c h w',c=1)
+        control_maps = torch.from_numpy(control_maps).div(255)  # 0~1
+    control_maps = control_maps.to(dtype=controlnet.dtype, device=controlnet.device)
+    control_maps = F.interpolate(control_maps, size=(h,w), mode='bilinear', align_corners=False)
+    control_maps = rearrange(control_maps, "(b f) c h w -> b c f h w", f=num_sample_frames)
+    if control_maps.shape[1] == 1:
+        control_maps = repeat(control_maps, 'b c f h w -> b (n c) f h w',  n=3)
+
+    frames = torch.from_numpy(np_frames).div(255)
+    frames = rearrange(frames, 'f h w c -> f c h w')
+    v2v_input_frames =  torch.nn.functional.interpolate(
+                frames,
+                size=(h, w),
+                mode="bicubic",
+                antialias=True,
+            ) 
+    v2v_input_frames = rearrange(v2v_input_frames, '(b f) c h w -> b c f h w ', f=num_sample_frames)
+
+    out = []
+    for i in range(num_sample_frames//each_sample_frame):
+        out1 = video_controlnet_pipe(
+                # controlnet_hint= control_maps[:,:,:each_sample_frame,:,:],
+                # images= v2v_input_frames[:,:,:each_sample_frame,:,:],
+                controlnet_hint=control_maps[:,:,i*each_sample_frame-1:(i+1)*each_sample_frame-1,:,:] if i>0 else control_maps[:,:,:each_sample_frame,:,:],
+                images=v2v_input_frames[:,:,i*each_sample_frame-1:(i+1)*each_sample_frame-1,:,:] if i>0 else v2v_input_frames[:,:,:each_sample_frame,:,:],
+                first_frame_output=out[-1] if i>0 else None,
+                prompt=testing_prompt,
+                num_inference_steps=num_inference_steps,
+                width=w,
+                height=h,
+                guidance_scale=guidance_scale,
+                generator=[torch.Generator(device="cuda").manual_seed(seed)],
+                video_scale = video_scale, 
+                init_noise_by_residual_thres = init_noise_thres,    # residual-based init. larger thres ==> more smooth.
+                controlnet_conditioning_scale=1.0,
+                fix_first_frame=True, 
+                in_domain=True,
+        )
+        out1 = out1.images[0]    
+        if len(out1) > 1:
+            out1 = out1[1:] # drop the first frame
+        out.extend(out1)
+
+    return to_video(out, 8)
+
+
+examples = [
+        ["__assets__/depth_videos_depth/girl_dancing.mp4",
+        "A stormtrooper, masterpiece, a high-quality, detailed, and professional photo"],
+]
+
+def preview_inference(
+        input_video, 
+        prompt, seed, 
+        num_inference_steps, guidance_scale, 
+        sampling_rate, video_scale, init_noise_thres,
+        each_sample_frame,iter_times, h, w,
+    ):
+    return inference(input_video, 
+        prompt, seed, 
+        num_inference_steps, guidance_scale, 
+        sampling_rate, 0.0, 0.0, 1, 1, h, w,)
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                input_video = gr.Video(
+                    label="Input Video", source='upload', format="mp4", visible=True)
+            with gr.Column():
+                init_noise_thres = gr.Slider(0, 1, value=0.1, step=0.1, label="init_noise_thress")
+                each_sample_frame = gr.Slider(6, 16, value=8, step=1, label="each_sample_frame")
+                iter_times = gr.Slider(1, 4, value=1, step=1, label="iter_times")
+                sampling_rate = gr.Slider(1, 8, value=3, step=1, label="sampling_rate")
+                h = gr.Slider(256, 768, value=512, step=64, label="height")
+                w = gr.Slider(256, 768, value=512, step=64, label="width")
+            with gr.Column():
+                seed =  gr.Slider(0, 6666, value=1, step=1, label="seed")
+                num_inference_steps =  gr.Slider(5, 50, value=20, step=1, label="num_inference_steps")
+                guidance_scale = gr.Slider(1, 20, value=7.5, step=0.5, label="guidance_scale")
+                video_scale = gr.Slider(0, 2.5, value=1.5, step=0.1, label="video_scale")
+                prompt = gr.Textbox(label='Prompt')
+                # preview_button = gr.Button('Preview')
+                run_button = gr.Button('Generate Video')
+
+        with gr.Column():
+            result = gr.Video(label="Generated Video")
+
+        inputs = [
+                input_video,
+                prompt,
+                seed,
+                num_inference_steps,
+                guidance_scale,
+                sampling_rate,
+                video_scale,
+                init_noise_thres,
+                each_sample_frame,
+                iter_times,
+                h,
+                w,
+            ]
+    
+        gr.Examples(examples=examples,
+                    inputs=inputs,
+                    outputs=result,
+                    fn=inference,
+                    cache_examples=False,
+                    run_on_click=False,
+                    )
+
+        run_button.click(fn=inference,
+                            inputs=inputs,
+                            outputs=result,)
+        # preview_button.click(fn=preview_inference,
+        #                     inputs=inputs,
+        #                     outputs=result,)
+    
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+
+# TODO 
+# 1. preview
+# 2. params

model/annotator/canny/__init__.py

@@ -0,0 +1,6 @@
+import cv2
+
+
+class CannyDetector:
+    def __call__(self, img, low_threshold, high_threshold):
+        return cv2.Canny(img, low_threshold, high_threshold)

model/annotator/hed/__init__.py

@@ -0,0 +1,133 @@
+import numpy as np
+import cv2
+import os
+import torch
+from einops import rearrange
+
+
+class HEDNetwork(torch.nn.Module):
+    def __init__(self, model_path):
+        super().__init__()
+
+        self.netVggOne = torch.nn.Sequential(
+            torch.nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False)
+        )
+
+        self.netVggTwo = torch.nn.Sequential(
+            torch.nn.MaxPool2d(kernel_size=2, stride=2),
+            torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False)
+        )
+
+        self.netVggThr = torch.nn.Sequential(
+            torch.nn.MaxPool2d(kernel_size=2, stride=2),
+            torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False)
+        )
+
+        self.netVggFou = torch.nn.Sequential(
+            torch.nn.MaxPool2d(kernel_size=2, stride=2),
+            torch.nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False)
+        )
+
+        self.netVggFiv = torch.nn.Sequential(
+            torch.nn.MaxPool2d(kernel_size=2, stride=2),
+            torch.nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False),
+            torch.nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
+            torch.nn.ReLU(inplace=False)
+        )
+
+        self.netScoreOne = torch.nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1, stride=1, padding=0)
+        self.netScoreTwo = torch.nn.Conv2d(in_channels=128, out_channels=1, kernel_size=1, stride=1, padding=0)
+        self.netScoreThr = torch.nn.Conv2d(in_channels=256, out_channels=1, kernel_size=1, stride=1, padding=0)
+        self.netScoreFou = torch.nn.Conv2d(in_channels=512, out_channels=1, kernel_size=1, stride=1, padding=0)
+        self.netScoreFiv = torch.nn.Conv2d(in_channels=512, out_channels=1, kernel_size=1, stride=1, padding=0)
+
+        self.netCombine = torch.nn.Sequential(
+            torch.nn.Conv2d(in_channels=5, out_channels=1, kernel_size=1, stride=1, padding=0),
+            torch.nn.Sigmoid()
+        )
+
+        self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load(model_path).items()})
+
+    def forward(self, tenInput):
+        tenInput = tenInput * 255.0
+        tenInput = tenInput - torch.tensor(data=[104.00698793, 116.66876762, 122.67891434], dtype=tenInput.dtype, device=tenInput.device).view(1, 3, 1, 1)
+
+        tenVggOne = self.netVggOne(tenInput)
+        tenVggTwo = self.netVggTwo(tenVggOne)
+        tenVggThr = self.netVggThr(tenVggTwo)
+        tenVggFou = self.netVggFou(tenVggThr)
+        tenVggFiv = self.netVggFiv(tenVggFou)
+
+        tenScoreOne = self.netScoreOne(tenVggOne)
+        tenScoreTwo = self.netScoreTwo(tenVggTwo)
+        tenScoreThr = self.netScoreThr(tenVggThr)
+        tenScoreFou = self.netScoreFou(tenVggFou)
+        tenScoreFiv = self.netScoreFiv(tenVggFiv)
+
+        tenScoreOne = torch.nn.functional.interpolate(input=tenScoreOne, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
+        tenScoreTwo = torch.nn.functional.interpolate(input=tenScoreTwo, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
+        tenScoreThr = torch.nn.functional.interpolate(input=tenScoreThr, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
+        tenScoreFou = torch.nn.functional.interpolate(input=tenScoreFou, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
+        tenScoreFiv = torch.nn.functional.interpolate(input=tenScoreFiv, size=(tenInput.shape[2], tenInput.shape[3]), mode='bilinear', align_corners=False)
+
+        return self.netCombine(torch.cat([ tenScoreOne, tenScoreTwo, tenScoreThr, tenScoreFou, tenScoreFiv ], 1))
+
+
+class HEDdetector:
+    def __init__(self, network ):
+        self.netNetwork = network
+
+    def __call__(self, input_image):
+        if isinstance(input_image, torch.Tensor):
+            # 输入的就是 b c h w的tensor 范围是-1~1，需要转换为0～1
+            input_image = (input_image + 1) / 2
+            input_image = input_image.float().cuda()
+            edge = self.netNetwork(input_image) # 范围也是0～1, 不用转了直接用
+            return edge
+        else:
+            assert input_image.ndim == 3
+            input_image = input_image[:, :, ::-1].copy()
+            with torch.no_grad():
+                image_hed = torch.from_numpy(input_image).float().cuda()
+                image_hed = image_hed / 255.0
+                image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
+                edge = self.netNetwork(image_hed)[0]
+                edge = (edge.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
+                return edge[0]
+
+
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z

model/annotator/util.py

@@ -0,0 +1,38 @@
+import numpy as np
+import cv2
+import os
+
+
+annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts')
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def resize_image(input_image, resolution):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
+    return img

model/video_diffusion/models/attention.py

@@ -0,0 +1,454 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.attention import FeedForward, CrossAttention, AdaLayerNorm
+from diffusers.utils import BaseOutput
+from diffusers.utils.import_utils import is_xformers_available
+from diffusers.models.cross_attention import XFormersCrossAttnProcessor
+from einops import rearrange
+
+
+@dataclass
+class SpatioTemporalTransformerModelOutput(BaseOutput):
+    """torch.FloatTensor of shape [batch x channel x frames x height x width]"""
+
+    sample: torch.FloatTensor
+
+
+if is_xformers_available():
+    import xformers
+    import xformers.ops
+else:
+    xformers = None
+
+
+class SpatioTemporalTransformerModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_channels: Optional[int] = None,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        cross_attention_dim: Optional[int] = None,
+        attention_bias: bool = False,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        upcast_attention: bool = False,
+        **transformer_kwargs,
+    ):
+        super().__init__()
+        self.use_linear_projection = use_linear_projection
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        inner_dim = num_attention_heads * attention_head_dim
+
+        # Define input layers
+        self.in_channels = in_channels
+
+        self.norm = torch.nn.GroupNorm(
+            num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
+        )
+        if use_linear_projection:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+        else:
+            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+
+        # Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                SpatioTemporalTransformerBlock(
+                    inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=dropout,
+                    cross_attention_dim=cross_attention_dim,
+                    activation_fn=activation_fn,
+                    num_embeds_ada_norm=num_embeds_ada_norm,
+                    attention_bias=attention_bias,
+                    only_cross_attention=only_cross_attention,
+                    upcast_attention=upcast_attention,
+                    **transformer_kwargs,
+                )
+                for d in range(num_layers)
+            ]
+        )
+
+        # Define output layers
+        if use_linear_projection:
+            self.proj_out = nn.Linear(in_channels, inner_dim)
+        else:
+            self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(
+        self, hidden_states, encoder_hidden_states=None, timestep=None, return_dict: bool = True
+    ):
+        # 1. Input
+        clip_length = None
+        is_video = hidden_states.ndim == 5
+        if is_video:
+            clip_length = hidden_states.shape[2]
+            hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+            encoder_hidden_states = encoder_hidden_states.repeat_interleave(clip_length, 0)
+
+        *_, h, w = hidden_states.shape
+        residual = hidden_states
+
+        hidden_states = self.norm(hidden_states)
+        if not self.use_linear_projection:
+            hidden_states = self.proj_in(hidden_states)
+            hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")
+        else:
+            hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")
+            hidden_states = self.proj_in(hidden_states)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                timestep=timestep,
+                clip_length=clip_length,
+            )
+
+        # 3. Output
+        if not self.use_linear_projection:
+            hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+            hidden_states = self.proj_out(hidden_states)
+        else:
+            hidden_states = self.proj_out(hidden_states)
+            hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+
+        output = hidden_states + residual
+        if is_video:
+            output = rearrange(output, "(b f) c h w -> b c f h w", f=clip_length)
+
+        if not return_dict:
+            return (output,)
+
+        return SpatioTemporalTransformerModelOutput(sample=output)
+
+
+class SpatioTemporalTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        upcast_attention: bool = False,
+        use_sparse_causal_attention: bool = False,
+        use_full_sparse_causal_attention: bool = True,
+        temporal_attention_position: str = "after_feedforward",
+        use_gamma = False,
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+        self.use_ada_layer_norm = num_embeds_ada_norm is not None
+        self.use_sparse_causal_attention = use_sparse_causal_attention
+        self.use_full_sparse_causal_attention = use_full_sparse_causal_attention
+        self.use_gamma = use_gamma
+
+        self.temporal_attention_position = temporal_attention_position
+        temporal_attention_positions = ["after_spatial", "after_cross", "after_feedforward"]
+        if temporal_attention_position not in temporal_attention_positions:
+            raise ValueError(
+                f"`temporal_attention_position` must be one of {temporal_attention_positions}"
+            )
+
+        # 1. Spatial-Attn
+        if use_sparse_causal_attention:
+           spatial_attention = SparseCausalAttention
+        elif use_full_sparse_causal_attention:
+            spatial_attention = SparseCausalFullAttention
+        else:
+            spatial_attention = CrossAttention
+        
+        self.attn1 = spatial_attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+            processor=XFormersCrossAttnProcessor(), 
+        )  # is a self-attention
+        self.norm1 = (
+            AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
+        )
+        if use_gamma:
+            self.attn1_gamma = nn.Parameter(torch.ones(dim))
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None:
+            self.attn2 = CrossAttention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+                processor=XFormersCrossAttnProcessor(),
+            )  # is self-attn if encoder_hidden_states is none
+            self.norm2 = (
+                AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
+            )
+            if use_gamma:
+                self.attn2_gamma = nn.Parameter(torch.ones(dim))
+        else:
+            self.attn2 = None
+            self.norm2 = None
+
+        # 3. Temporal-Attn
+        self.attn_temporal = CrossAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            upcast_attention=upcast_attention,
+            processor=XFormersCrossAttnProcessor()
+        )
+        zero_module(self.attn_temporal) # 默认参数置0
+
+        self.norm_temporal = (
+            AdaLayerNorm(dim, num_embeds_ada_norm) if self.use_ada_layer_norm else nn.LayerNorm(dim)
+        )
+
+        # 4. Feed-forward
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
+        self.norm3 = nn.LayerNorm(dim)
+        if use_gamma:
+            self.ff_gamma = nn.Parameter(torch.ones(dim))
+ 
+ 
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        timestep=None,
+        attention_mask=None,
+        clip_length=None,
+    ):
+        # 1. Self-Attention
+        norm_hidden_states = (
+            self.norm1(hidden_states, timestep) if self.use_ada_layer_norm else self.norm1(hidden_states)
+        )
+
+        kwargs = dict(
+            hidden_states=norm_hidden_states,
+            attention_mask=attention_mask,
+        )
+        if self.only_cross_attention:
+            kwargs.update(encoder_hidden_states=encoder_hidden_states)
+        if self.use_sparse_causal_attention or self.use_full_sparse_causal_attention:
+            kwargs.update(clip_length=clip_length)
+
+        if self.use_gamma:
+            hidden_states = hidden_states + self.attn1(**kwargs) * self.attn1_gamma # NOTE gamma
+        else:
+            hidden_states = hidden_states + self.attn1(**kwargs)
+
+
+        if clip_length is not None and self.temporal_attention_position == "after_spatial":
+            hidden_states = self.apply_temporal_attention(hidden_states, timestep, clip_length)
+
+        if self.attn2 is not None:
+            # 2. Cross-Attention
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep)
+                if self.use_ada_layer_norm
+                else self.norm2(hidden_states)
+            )
+            if self.use_gamma:
+                hidden_states = (
+                    self.attn2(
+                        norm_hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=attention_mask,
+                    ) * self.attn2_gamma
+                    + hidden_states
+                )
+            else:
+                hidden_states = (
+                    self.attn2(
+                        norm_hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        attention_mask=attention_mask,
+                    )
+                    + hidden_states
+                )
+
+        if clip_length is not None and self.temporal_attention_position == "after_cross":
+            hidden_states = self.apply_temporal_attention(hidden_states, timestep, clip_length)
+
+        # 3. Feed-forward
+        if self.use_gamma:
+            hidden_states = self.ff(self.norm3(hidden_states)) * self.ff_gamma + hidden_states
+        else:
+            hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
+
+        if clip_length is not None and self.temporal_attention_position == "after_feedforward":
+            hidden_states = self.apply_temporal_attention(hidden_states, timestep, clip_length)
+
+        return hidden_states
+
+    def apply_temporal_attention(self, hidden_states, timestep, clip_length):
+        d = hidden_states.shape[1]
+        hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=clip_length)
+        norm_hidden_states = (
+            self.norm_temporal(hidden_states, timestep)
+            if self.use_ada_layer_norm
+            else self.norm_temporal(hidden_states)
+        )
+        hidden_states = self.attn_temporal(norm_hidden_states) + hidden_states
+        hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
+        return hidden_states
+
+
+class SparseCausalAttention(CrossAttention):
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        clip_length: int = None,
+    ):
+        if (
+            self.added_kv_proj_dim is not None
+            or encoder_hidden_states is not None
+            or attention_mask is not None
+        ):
+            raise NotImplementedError
+
+        if self.group_norm is not None:
+            hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = self.to_q(hidden_states)
+        dim = query.shape[-1]
+        query = self.head_to_batch_dim(query)   # 64 4096 40
+
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        if clip_length is not None and clip_length > 1:
+            # spatial temporal
+            prev_frame_index = torch.arange(clip_length) - 1   
+            prev_frame_index[0] = 0 
+            key = rearrange(key, "(b f) d c -> b f d c", f=clip_length)
+            key = torch.cat([key[:, [0] * clip_length], key[:, prev_frame_index]], dim=2)  
+            key = rearrange(key, "b f d c -> (b f) d c", f=clip_length)
+
+            value = rearrange(value, "(b f) d c -> b f d c", f=clip_length)
+            value = torch.cat([value[:, [0] * clip_length], value[:, prev_frame_index]], dim=2)
+            value = rearrange(value, "b f d c -> (b f) d c", f=clip_length)
+
+
+        key = self.head_to_batch_dim(key)
+        value = self.head_to_batch_dim(value)
+        # use xfromers by default~
+        hidden_states = xformers.ops.memory_efficient_attention(
+            query, key, value, attn_bias=attention_mask, op=None
+        )
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states =  self.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = self.to_out[0](hidden_states)
+
+        # dropout
+        hidden_states = self.to_out[1](hidden_states)
+        return hidden_states
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module
+
+
+class SparseCausalFullAttention(CrossAttention):
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        clip_length: int = None,
+    ):
+        if (
+            self.added_kv_proj_dim is not None
+            or encoder_hidden_states is not None
+            or attention_mask is not None
+        ):
+            raise NotImplementedError
+
+        if self.group_norm is not None:
+            hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = self.to_q(hidden_states)
+        dim = query.shape[-1]
+        query = self.head_to_batch_dim(query)   # 64 4096 40
+
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        if clip_length is not None and clip_length > 1:
+            # 和所有帧做 spatial temporal attention
+            key = rearrange(key, "(b f) d c -> b f d c", f=clip_length)
+            # cat full frames
+            key = torch.cat([key[:, [iii] * clip_length] for iii in range(clip_length) ], dim=2)   # concat第一帧+第i帧。以此为key, value。而非自己这一帧。
+            key = rearrange(key, "b f d c -> (b f) d c", f=clip_length)
+
+            value = rearrange(value, "(b f) d c -> b f d c", f=clip_length)
+            value = torch.cat([value[:, [iii] * clip_length] for iii in range(clip_length) ], dim=2)   # concat第一帧+第i帧。以此为key, value。而非自己这一帧。
+            value = rearrange(value, "b f d c -> (b f) d c", f=clip_length)
+
+        key = self.head_to_batch_dim(key)
+        value = self.head_to_batch_dim(value)
+        # use xfromers by default~
+        hidden_states = xformers.ops.memory_efficient_attention(
+            query, key, value, attn_bias=attention_mask, op=None
+        )
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states =  self.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = self.to_out[0](hidden_states)
+
+        # dropout
+        hidden_states = self.to_out[1](hidden_states)
+        return hidden_states
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module

model/video_diffusion/models/controlnet3d.py

@@ -0,0 +1,580 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.models.cross_attention import AttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+
+from .unet_3d_blocks import (
+    CrossAttnDownBlockPseudo3D,
+    DownBlockPseudo3D,
+    UNetMidBlockPseudo3DCrossAttn,
+    get_down_block,
+)
+from .resnet import PseudoConv3d
+from diffusers.models.cross_attention import AttnProcessor
+from typing import Dict
+from .unet_3d_blocks_control import ControlNetPseudoZeroConv3dBlock, ControlNetInputHintBlock
+import glob
+import os
+import json
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+@dataclass
+class ControlNetOutput(BaseOutput):
+    down_block_res_samples: Tuple[torch.Tensor]
+    mid_block_res_sample: torch.Tensor
+
+
+class ControlNetConditioningEmbedding(nn.Module):
+    """
+    Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN
+    [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized
+    training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the
+    convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides
+    (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full
+    model) to encode image-space conditions ... into feature maps ..."
+    """
+
+    def __init__(
+        self,
+        conditioning_embedding_channels: int,
+        conditioning_channels: int = 3,
+        block_out_channels: Tuple[int] = (16, 32, 96, 256),
+    ):
+        super().__init__()
+
+        self.conv_in = PseudoConv3d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)
+
+        self.blocks = nn.ModuleList([])
+
+        for i in range(len(block_out_channels) - 1):
+            channel_in = block_out_channels[i]
+            channel_out = block_out_channels[i + 1]
+            self.blocks.append(PseudoConv3d(channel_in, channel_in, kernel_size=3, padding=1))
+            self.blocks.append(PseudoConv3d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))
+
+        # self.conv_out = zero_module(
+        #     PseudoConv3d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
+        # )   
+        self.conv_out = PseudoConv3d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
+
+    def forward(self, conditioning):
+        embedding = self.conv_in(conditioning)
+        embedding = F.silu(embedding)
+
+        for block in self.blocks:
+            embedding = block(embedding)
+            embedding = F.silu(embedding)
+
+        embedding = self.conv_out(embedding)
+
+        return embedding
+
+
+class ControlNet3DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 4,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlockPseudo3D",
+            "CrossAttnDownBlockPseudo3D",
+            "CrossAttnDownBlockPseudo3D",
+            "DownBlockPseudo3D",
+        ),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: int = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1280,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        use_linear_projection: bool = False,
+        class_embed_type: Optional[str] = None,
+        num_class_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        projection_class_embeddings_input_dim: Optional[int] = None,
+        controlnet_conditioning_channel_order: str = "rgb",
+        conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, 96, 256),
+    ):
+        super().__init__()
+
+        # Check inputs
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        conv_in_kernel = 3
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = PseudoConv3d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+
+        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+        )
+
+        # class embedding
+        if class_embed_type is None and num_class_embeds is not None:
+            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+        elif class_embed_type == "timestep":
+            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        elif class_embed_type == "projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+            # 2. it projects from an arbitrary input dimension.
+            #
+            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+        # control net conditioning embedding
+        self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
+            conditioning_embedding_channels=block_out_channels[0],
+            block_out_channels=conditioning_embedding_out_channels,
+        )
+
+        self.down_blocks = nn.ModuleList([])
+        self.controlnet_down_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        # down
+        output_channel = block_out_channels[0]
+
+        controlnet_block = PseudoConv3d(output_channel, output_channel, kernel_size=1)
+        # controlnet_block = zero_module(controlnet_block)
+        self.controlnet_down_blocks.append(controlnet_block)
+
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim[i],
+                downsample_padding=downsample_padding,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+            )
+            self.down_blocks.append(down_block)
+
+            for _ in range(layers_per_block):
+                controlnet_block = PseudoConv3d(output_channel, output_channel, kernel_size=1)
+                # controlnet_block = zero_module(controlnet_block)
+                self.controlnet_down_blocks.append(controlnet_block)
+
+            if not is_final_block:
+                controlnet_block = PseudoConv3d(output_channel, output_channel, kernel_size=1)
+                # controlnet_block = zero_module(controlnet_block)
+                self.controlnet_down_blocks.append(controlnet_block)
+
+        # mid
+        mid_block_channel = block_out_channels[-1]
+
+        controlnet_block = PseudoConv3d(mid_block_channel, mid_block_channel, kernel_size=1)
+        # controlnet_block = zero_module(controlnet_block)
+        self.controlnet_mid_block = controlnet_block
+
+        self.mid_block = UNetMidBlockPseudo3DCrossAttn(
+            in_channels=mid_block_channel,
+            temb_channels=time_embed_dim,
+            resnet_eps=norm_eps,
+            resnet_act_fn=act_fn,
+            output_scale_factor=mid_block_scale_factor,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attention_head_dim[-1],
+            resnet_groups=norm_num_groups,
+            use_linear_projection=use_linear_projection,
+            upcast_attention=upcast_attention,
+        )
+
+    @property
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttnProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttnProcessor]):
+            if hasattr(module, "set_processor"):
+                processors[f"{name}.processor"] = module.processor
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttnProcessor, Dict[str, AttnProcessor]]):
+        r"""
+        Parameters:
+            `processor (`dict` of `AttnProcessor` or `AttnProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                of **all** `CrossAttention` layers.
+            In case `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainablae attention processors.:
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice
+    def set_attention_slice(self, slice_size):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module will split the input tensor in slices, to compute attention
+        in several steps. This is useful to save some memory in exchange for a small speed decrease.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
+                `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_slicable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_slicable_dims(module)
+
+        num_slicable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_slicable_layers * [1]
+
+        slice_size = num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (CrossAttnDownBlockPseudo3D, DownBlockPseudo3D)):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        controlnet_cond: torch.FloatTensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> Union[ControlNetOutput, Tuple]:
+        # check channel order
+        channel_order = self.config.controlnet_conditioning_channel_order
+
+        if channel_order == "rgb":
+            # in rgb order by default
+            ...
+        elif channel_order == "bgr":
+            controlnet_cond = torch.flip(controlnet_cond, dims=[1])
+        else:
+            raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}")
+
+        # prepare attention_mask
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+            emb = emb + class_emb
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
+        # print(sample.shape, controlnet_cond.shape)
+
+        sample += controlnet_cond
+        # 3. down
+        
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+            down_block_res_samples += res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+            )
+
+        # 5. Control net blocks
+
+        controlnet_down_block_res_samples = ()
+
+        for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
+            down_block_res_sample = controlnet_block(down_block_res_sample)
+            controlnet_down_block_res_samples += (down_block_res_sample,)
+
+        down_block_res_samples = controlnet_down_block_res_samples
+
+        mid_block_res_sample = self.controlnet_mid_block(sample)
+
+        if not return_dict:
+            return (down_block_res_samples, mid_block_res_sample)
+
+        return ControlNetOutput(
+            down_block_res_samples=down_block_res_samples, mid_block_res_sample=mid_block_res_sample
+        )
+
+    @classmethod
+    def from_2d_model(cls, model_path, condition_on_fps=False, controlnet_hint_channels: Optional[int] = None,):
+        '''
+        load a 2d model and convert it to a pseudo 3d model
+        '''
+        config_path = os.path.join(model_path, "config.json")
+        if not os.path.isfile(config_path):
+            raise RuntimeError(f"{config_path} does not exist")
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        config.pop("_class_name")
+        config.pop("_diffusers_version")
+
+        block_replacer = {
+            "CrossAttnDownBlock2D": "CrossAttnDownBlockPseudo3D",
+            "DownBlock2D": "DownBlockPseudo3D",
+            "UNetMidBlock2DCrossAttn": "UNetMidBlockPseudo3DCrossAttn",
+        }
+
+        def convert_2d_to_3d_block(block):
+            return block_replacer[block] if block in block_replacer else block
+
+        config["down_block_types"] = [
+            convert_2d_to_3d_block(block) for block in config["down_block_types"]
+        ]
+        
+        if "mid_block_type" in config:
+            config["mid_block_type"] = convert_2d_to_3d_block(config["mid_block_type"])
+
+        if condition_on_fps:
+            config["fps_embed_type"] = "timestep"     # 和timestep保持一致的type。
+        
+        if controlnet_hint_channels:
+            config["controlnet_hint_channels"] = controlnet_hint_channels
+
+        print(config)
+
+        model = cls(**config)   # 调用自身(init), 传入config参数全换成3d的setting
+        state_dict_path_condidates = glob.glob(os.path.join(model_path, "*.bin"))
+        if state_dict_path_condidates:
+            state_dict = torch.load(state_dict_path_condidates[0], map_location="cpu")
+            model.load_2d_state_dict(state_dict=state_dict)
+
+        return model
+
+    def load_2d_state_dict(self, state_dict, **kwargs):
+        '''
+        2D 部分的参数名完全不变。
+        '''
+        state_dict_3d = self.state_dict()
+        # print("diff params list:", list(set(state_dict_3d.keys()) - set(state_dict.keys())))
+
+        for k, v in state_dict.items():
+            if k not in state_dict_3d:
+                raise KeyError(f"2d state_dict key {k} does not exist in 3d model")
+        
+        for k, v in state_dict_3d.items():
+            if "_temporal" in k:
+                continue
+            if "gamma" in k:
+                continue
+            if k not in state_dict:
+                raise KeyError(f"3d state_dict key {k} does not exist in 2d model")
+        state_dict_3d.update(state_dict)   
+        self.load_state_dict(state_dict_3d, strict=True,  **kwargs)
+
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module

model/video_diffusion/models/resnet.py

@@ -0,0 +1,486 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+
+class PseudoConv3d(nn.Conv2d):
+    def __init__(self, in_channels, out_channels, kernel_size, temporal_kernel_size=None, **kwargs):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            **kwargs,
+        )
+        if temporal_kernel_size is None:
+            temporal_kernel_size = kernel_size
+
+        self.conv_temporal = (
+            nn.Conv1d(
+                out_channels,
+                out_channels,
+                kernel_size=temporal_kernel_size,
+                padding=temporal_kernel_size // 2,
+            )
+            if kernel_size > 1
+            else None
+        )
+
+        if self.conv_temporal is not None:
+            nn.init.dirac_(self.conv_temporal.weight.data)  # initialized to be identity
+            nn.init.zeros_(self.conv_temporal.bias.data)
+
+    def forward(self, x):
+        b = x.shape[0]
+
+        is_video = x.ndim == 5
+        if is_video:
+            x = rearrange(x, "b c f h w -> (b f) c h w")
+
+        x = super().forward(x)  
+
+        if is_video:
+            x = rearrange(x, "(b f) c h w -> b c f h w", b=b)
+
+        if self.conv_temporal is None or not is_video:
+            return x
+
+        *_, h, w = x.shape
+
+        x = rearrange(x, "b c f h w -> (b h w) c f")
+
+        x = self.conv_temporal(x)   # 加入空间1D的时序卷积。channel不变。（建模时序信息）
+
+        x = rearrange(x, "(b h w) c f -> b c f h w", h=h, w=w)
+
+        return x
+
+
+class UpsamplePseudo3D(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    Parameters:
+        channels: channels in the inputs and outputs.
+        use_conv: a bool determining if a convolution is applied.
+        use_conv_transpose:
+        out_channels:
+    """
+
+    def __init__(
+        self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+
+        conv = None
+        if use_conv_transpose:
+            raise NotImplementedError
+            conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1)
+        elif use_conv:
+            conv = PseudoConv3d(self.channels, self.out_channels, 3, padding=1)
+
+        # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
+        if name == "conv":
+            self.conv = conv
+        else:
+            self.Conv2d_0 = conv
+
+    def forward(self, hidden_states, output_size=None):
+        assert hidden_states.shape[1] == self.channels
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch
+        # https://github.com/pytorch/pytorch/issues/86679
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            hidden_states = hidden_states.contiguous()
+
+        b = hidden_states.shape[0]
+        is_video = hidden_states.ndim == 5
+        if is_video:
+            hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if output_size is None:
+            # 先插值再用conv
+            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+        else:
+            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        if is_video:
+            hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=b)
+
+        if self.use_conv:
+            if self.name == "conv":
+                hidden_states = self.conv(hidden_states)
+            else:
+                hidden_states = self.Conv2d_0(hidden_states)
+
+        return hidden_states
+
+
+class DownsamplePseudo3D(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    Parameters:
+        channels: channels in the inputs and outputs.
+        use_conv: a bool determining if a convolution is applied.
+        out_channels:
+        padding:
+    """
+
+    def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = 2
+        self.name = name
+
+        if use_conv:
+            conv = PseudoConv3d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
+        else:
+            assert self.channels == self.out_channels
+            conv = nn.AvgPool2d(kernel_size=stride, stride=stride)
+
+        if name == "conv":
+            self.Conv2d_0 = conv
+            self.conv = conv
+        elif name == "Conv2d_0":
+            self.conv = conv
+        else:
+            self.conv = conv
+
+    def forward(self, hidden_states):
+        assert hidden_states.shape[1] == self.channels
+        if self.use_conv and self.padding == 0:
+            pad = (0, 1, 0, 1)
+            hidden_states = F.pad(hidden_states, pad, mode="constant", value=0)
+
+        assert hidden_states.shape[1] == self.channels
+        if self.use_conv:
+            hidden_states = self.conv(hidden_states)
+        else:
+            b = hidden_states.shape[0]
+            is_video = hidden_states.ndim == 5
+            if is_video:
+                hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+            hidden_states = self.conv(hidden_states)
+            if is_video:
+                hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=b)
+
+        return hidden_states
+
+
+class ResnetBlockPseudo3D(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout=0.0,
+        temb_channels=512,
+        groups=32,
+        groups_out=None,
+        pre_norm=True,
+        eps=1e-6,
+        non_linearity="swish",
+        time_embedding_norm="default",
+        kernel=None,
+        output_scale_factor=1.0,
+        use_in_shortcut=None,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.time_embedding_norm = time_embedding_norm
+        self.up = up
+        self.down = down
+        self.output_scale_factor = output_scale_factor
+
+        if groups_out is None:
+            groups_out = groups
+
+        self.norm1 = torch.nn.GroupNorm(
+            num_groups=groups, num_channels=in_channels, eps=eps, affine=True
+        )
+
+        self.conv1 = PseudoConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == "default":
+                time_emb_proj_out_channels = out_channels
+            elif self.time_embedding_norm == "scale_shift":
+                time_emb_proj_out_channels = out_channels * 2
+            else:
+                raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")
+
+            self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)
+        else:
+            self.time_emb_proj = None
+
+        self.norm2 = torch.nn.GroupNorm(
+            num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True
+        )
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = PseudoConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if non_linearity == "swish":
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == "mish":
+            self.nonlinearity = Mish()
+        elif non_linearity == "silu":
+            self.nonlinearity = nn.SiLU()
+
+        self.upsample = self.downsample = None
+        if self.up:
+            if kernel == "fir":
+                fir_kernel = (1, 3, 3, 1)
+                self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)
+            elif kernel == "sde_vp":
+                self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
+            else:
+                self.upsample = UpsamplePseudo3D(in_channels, use_conv=False)
+        elif self.down:
+            if kernel == "fir":
+                fir_kernel = (1, 3, 3, 1)
+                self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)
+            elif kernel == "sde_vp":
+                self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)
+            else:
+                self.downsample = DownsamplePseudo3D(in_channels, use_conv=False, padding=1, name="op")
+
+        self.use_in_shortcut = (
+            self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut
+        )
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = PseudoConv3d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, input_tensor, temb):
+        hidden_states = input_tensor
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:
+            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+            if hidden_states.shape[0] >= 64:
+                input_tensor = input_tensor.contiguous()
+                hidden_states = hidden_states.contiguous()
+            input_tensor = self.upsample(input_tensor)
+            hidden_states = self.upsample(hidden_states)
+        elif self.downsample is not None:
+            input_tensor = self.downsample(input_tensor)
+            hidden_states = self.downsample(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if temb is not None:
+            temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]
+
+        if temb is not None and self.time_embedding_norm == "default":
+            is_video = hidden_states.ndim == 5
+            if is_video:
+                b, c, f, h, w = hidden_states.shape
+                hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+                temb = temb.repeat_interleave(f, 0)
+
+            hidden_states = hidden_states + temb
+
+            if is_video:
+                hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=b)
+
+        hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == "scale_shift":
+            is_video = hidden_states.ndim == 5
+            if is_video:
+                b, c, f, h, w = hidden_states.shape
+                hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+                temb = temb.repeat_interleave(f, 0)
+
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+            if is_video:
+                hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=b)
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+class Mish(torch.nn.Module):
+    def forward(self, hidden_states):
+        return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))
+
+
+def upsample_2d(hidden_states, kernel=None, factor=2, gain=1):
+    r"""Upsample2D a batch of 2D images with the given filter.
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and upsamples each image with the given
+    filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the specified
+    `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its shape is
+    a: multiple of the upsampling factor.
+
+    Args:
+        hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        kernel: FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to nearest-neighbor upsampling.
+        factor: Integer upsampling factor (default: 2).
+        gain: Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        output: Tensor of the shape `[N, C, H * factor, W * factor]`
+    """
+    assert isinstance(factor, int) and factor >= 1
+    if kernel is None:
+        kernel = [1] * factor
+
+    kernel = torch.tensor(kernel, dtype=torch.float32)
+    if kernel.ndim == 1:
+        kernel = torch.outer(kernel, kernel)
+    kernel /= torch.sum(kernel)
+
+    kernel = kernel * (gain * (factor**2))
+    pad_value = kernel.shape[0] - factor
+    output = upfirdn2d_native(
+        hidden_states,
+        kernel.to(device=hidden_states.device),
+        up=factor,
+        pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),
+    )
+    return output
+
+
+def downsample_2d(hidden_states, kernel=None, factor=2, gain=1):
+    r"""Downsample2D a batch of 2D images with the given filter.
+    Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and downsamples each image with the
+    given filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the
+    specified `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its
+    shape is a multiple of the downsampling factor.
+
+    Args:
+        hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.
+        kernel: FIR filter of the shape `[firH, firW]` or `[firN]`
+          (separable). The default is `[1] * factor`, which corresponds to average pooling.
+        factor: Integer downsampling factor (default: 2).
+        gain: Scaling factor for signal magnitude (default: 1.0).
+
+    Returns:
+        output: Tensor of the shape `[N, C, H // factor, W // factor]`
+    """
+
+    assert isinstance(factor, int) and factor >= 1
+    if kernel is None:
+        kernel = [1] * factor
+
+    kernel = torch.tensor(kernel, dtype=torch.float32)
+    if kernel.ndim == 1:
+        kernel = torch.outer(kernel, kernel)
+    kernel /= torch.sum(kernel)
+
+    kernel = kernel * gain
+    pad_value = kernel.shape[0] - factor
+    output = upfirdn2d_native(
+        hidden_states,
+        kernel.to(device=hidden_states.device),
+        down=factor,
+        pad=((pad_value + 1) // 2, pad_value // 2),
+    )
+    return output
+
+
+def upfirdn2d_native(tensor, kernel, up=1, down=1, pad=(0, 0)):
+    up_x = up_y = up
+    down_x = down_y = down
+    pad_x0 = pad_y0 = pad[0]
+    pad_x1 = pad_y1 = pad[1]
+
+    _, channel, in_h, in_w = tensor.shape
+    tensor = tensor.reshape(-1, in_h, in_w, 1)
+
+    _, in_h, in_w, minor = tensor.shape
+    kernel_h, kernel_w = kernel.shape
+
+    out = tensor.view(-1, in_h, 1, in_w, 1, minor)
+    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+    out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
+    out = out.to(tensor.device)  # Move back to mps if necessary
+    out = out[
+        :,
+        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+        :,
+    ]
+
+    out = out.permute(0, 3, 1, 2)
+    out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
+    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+    out = F.conv2d(out, w)
+    out = out.reshape(
+        -1,
+        minor,
+        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+    )
+    out = out.permute(0, 2, 3, 1)
+    out = out[:, ::down_y, ::down_x, :]
+
+    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+
+    return out.view(-1, channel, out_h, out_w)

model/video_diffusion/models/unet_3d_blocks.py

@@ -0,0 +1,622 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+
+from .attention import SpatioTemporalTransformerModel
+from .resnet import DownsamplePseudo3D, ResnetBlockPseudo3D, UpsamplePseudo3D
+
+
+def get_down_block(
+    down_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    temb_channels,
+    add_downsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    downsample_padding=None,
+    dual_cross_attention=False,
+    use_linear_projection=False,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+):
+    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
+    if down_block_type == "DownBlockPseudo3D":
+        return DownBlockPseudo3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif down_block_type == "CrossAttnDownBlockPseudo3D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockPseudo3D")
+        return CrossAttnDownBlockPseudo3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block(
+    up_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    prev_output_channel,
+    temb_channels,
+    add_upsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    dual_cross_attention=False,
+    use_linear_projection=False,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+):
+    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+    if up_block_type == "UpBlockPseudo3D":
+        return UpBlockPseudo3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif up_block_type == "CrossAttnUpBlockPseudo3D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockPseudo3D")
+        return CrossAttnUpBlockPseudo3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlockPseudo3DCrossAttn(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        output_scale_factor=1.0,
+        cross_attention_dim=1280,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        upcast_attention=False,
+    ):
+        super().__init__()
+
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlockPseudo3D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        attentions = []
+
+        for _ in range(num_layers):
+            if dual_cross_attention:
+                raise NotImplementedError
+            attentions.append(
+                SpatioTemporalTransformerModel(
+                    attn_num_head_channels,
+                    in_channels // attn_num_head_channels,
+                    in_channels=in_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    upcast_attention=upcast_attention,
+                )
+            )
+            resnets.append(
+                ResnetBlockPseudo3D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+    def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
+        # TODO(Patrick, William) - attention_mask is currently not used. Implement once used
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class CrossAttnDownBlockPseudo3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        output_scale_factor=1.0,
+        downsample_padding=1,
+        add_downsample=True,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        upcast_attention=False,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlockPseudo3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            if dual_cross_attention:
+                raise NotImplementedError
+            attentions.append(
+                SpatioTemporalTransformerModel(
+                    attn_num_head_channels,
+                    out_channels // attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention,
+                    upcast_attention=upcast_attention,
+                )
+            )
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    DownsamplePseudo3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
+        # TODO(Patrick, William) - attention mask is not used
+        output_states = ()
+
+        for resnet, attn in zip(self.resnets, self.attentions):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn, return_dict=False),
+                    hidden_states,
+                    encoder_hidden_states,
+                )[0]
+            else:
+                hidden_states = resnet(hidden_states, temb)
+                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class DownBlockPseudo3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_downsample=True,
+        downsample_padding=1,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlockPseudo3D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    DownsamplePseudo3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, temb=None):
+        output_states = ()
+
+        for resnet in self.resnets:
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class CrossAttnUpBlockPseudo3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        output_scale_factor=1.0,
+        add_upsample=True,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        upcast_attention=False,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlockPseudo3D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            if dual_cross_attention:
+                raise NotImplementedError
+            attentions.append(
+                SpatioTemporalTransformerModel(
+                    attn_num_head_channels,
+                    out_channels // attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    only_cross_attention=only_cross_attention,
+                    upcast_attention=upcast_attention,
+                )
+            )
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [UpsamplePseudo3D(out_channels, use_conv=True, out_channels=out_channels)]
+            )
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        res_hidden_states_tuple,
+        temb=None,
+        encoder_hidden_states=None,
+        upsample_size=None,
+        attention_mask=None,
+    ):
+        # TODO(Patrick, William) - attention mask is not used
+        for resnet, attn in zip(self.resnets, self.attentions):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn, return_dict=False),
+                    hidden_states,
+                    encoder_hidden_states,
+                )[0]
+            else:
+                hidden_states = resnet(hidden_states, temb)
+                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states
+
+
+class UpBlockPseudo3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_output_channel: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_upsample=True,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResnetBlockPseudo3D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [UpsamplePseudo3D(out_channels, use_conv=True, out_channels=out_channels)]
+            )
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
+        for resnet in self.resnets:
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states

model/video_diffusion/models/unet_3d_blocks_control.py

@@ -0,0 +1,116 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+from .attention import SpatioTemporalTransformerModel
+from .resnet import DownsamplePseudo3D, ResnetBlockPseudo3D, UpsamplePseudo3D
+import glob
+import json
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput, logging
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from .unet_3d_blocks import (
+    CrossAttnDownBlockPseudo3D,
+    CrossAttnUpBlockPseudo3D,
+    DownBlockPseudo3D,
+    UNetMidBlockPseudo3DCrossAttn,
+    UpBlockPseudo3D,
+    get_down_block,
+    get_up_block,
+)
+from .resnet import PseudoConv3d
+from diffusers.models.cross_attention import AttnProcessor
+from typing import Dict
+
+
+
+def set_zero_parameters(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+# ControlNet: Zero Convolution
+def zero_conv(channels):
+    return set_zero_parameters(PseudoConv3d(channels, channels, 1, padding=0))
+
+class ControlNetInputHintBlock(nn.Module):
+    def __init__(self, hint_channels: int = 3, channels: int = 320):
+        super().__init__()
+        #  Layer configurations are from reference implementation.
+        self.input_hint_block = nn.Sequential(
+            PseudoConv3d(hint_channels, 16, 3, padding=1),
+            nn.SiLU(),
+            PseudoConv3d(16, 16, 3, padding=1),
+            nn.SiLU(),
+            PseudoConv3d(16, 32, 3, padding=1, stride=2),
+            nn.SiLU(),
+            PseudoConv3d(32, 32, 3, padding=1),
+            nn.SiLU(),
+            PseudoConv3d(32, 96, 3, padding=1, stride=2),
+            nn.SiLU(),
+            PseudoConv3d(96, 96, 3, padding=1),
+            nn.SiLU(),
+            PseudoConv3d(96, 256, 3, padding=1, stride=2),
+            nn.SiLU(),
+            set_zero_parameters(PseudoConv3d(256, channels, 3, padding=1)),
+        )
+    def forward(self, hint: torch.Tensor):
+        return self.input_hint_block(hint)
+        
+
+class ControlNetPseudoZeroConv3dBlock(nn.Module):
+    def __init__(
+        self,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlockPseudo3D",
+            "CrossAttnDownBlockPseudo3D",
+            "CrossAttnDownBlockPseudo3D",
+            "DownBlockPseudo3D",
+        ),
+        layers_per_block: int = 2,
+    ):
+        super().__init__()
+        self.input_zero_conv = zero_conv(block_out_channels[0])
+        zero_convs = []
+        for i, down_block_type in enumerate(down_block_types):
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            for _ in range(layers_per_block):
+                zero_convs.append(zero_conv(output_channel))
+            if not is_final_block:
+                zero_convs.append(zero_conv(output_channel))
+        self.zero_convs = nn.ModuleList(zero_convs)
+        self.mid_zero_conv = zero_conv(block_out_channels[-1])
+
+    def forward(
+        self,
+        down_block_res_samples: List[torch.Tensor],
+        mid_block_sample: torch.Tensor,
+    ) -> List[torch.Tensor]:
+        outputs = []
+        outputs.append(self.input_zero_conv(down_block_res_samples[0]))
+        for res_sample, zero_conv in zip(down_block_res_samples[1:], self.zero_convs):
+            outputs.append(zero_conv(res_sample))
+        outputs.append(self.mid_zero_conv(mid_block_sample))
+        return outputs
+        
+

model/video_diffusion/models/unet_3d_condition.py

@@ -0,0 +1,571 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import glob
+import json
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput, logging
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from .unet_3d_blocks import (
+    CrossAttnDownBlockPseudo3D,
+    CrossAttnUpBlockPseudo3D,
+    DownBlockPseudo3D,
+    UNetMidBlockPseudo3DCrossAttn,
+    UpBlockPseudo3D,
+    get_down_block,
+    get_up_block,
+)
+from .resnet import PseudoConv3d
+from diffusers.models.cross_attention import AttnProcessor
+from typing import Dict
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNetPseudo3DConditionOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+class UNetPseudo3DConditionModel(ModelMixin, ConfigMixin):
+    """
+    这里把原来2D Unet的 2D卷积全换成新定义的PseudoConv3d。并且定义了从2D卷积继承的模型参数。
+    """
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        center_input_sample: bool = False,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlockPseudo3D",
+            "CrossAttnDownBlockPseudo3D",
+            "CrossAttnDownBlockPseudo3D",
+            "DownBlockPseudo3D",
+        ),
+        mid_block_type: str = "UNetMidBlockPseudo3DCrossAttn",
+        up_block_types: Tuple[str] = (
+            "UpBlockPseudo3D",
+            "CrossAttnUpBlockPseudo3D",
+            "CrossAttnUpBlockPseudo3D",
+            "CrossAttnUpBlockPseudo3D",
+        ),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: int = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1280,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        fps_embed_type: Optional[str] = None,
+        num_fps_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        num_class_embeds=None,
+        
+    ):
+        super().__init__()
+
+
+        self.sample_size = sample_size
+        time_embed_dim = block_out_channels[0] * 4
+
+        # input
+        self.conv_in = PseudoConv3d(in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1))
+
+        # time
+        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        # class embedding
+        if fps_embed_type is None and num_fps_embeds is not None:
+            self.fps_embedding = nn.Embedding(num_fps_embeds, time_embed_dim)
+        elif fps_embed_type == "timestep":
+            self.fps_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+        elif fps_embed_type == "identity":
+            self.fps_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        else:
+            self.fps_embedding = None
+
+        self.down_blocks = nn.ModuleList([])
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim[i],
+                downsample_padding=downsample_padding,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        if mid_block_type == "UNetMidBlockPseudo3DCrossAttn":
+            self.mid_block = UNetMidBlockPseudo3DCrossAttn(
+                in_channels=block_out_channels[-1],
+                temb_channels=time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=attention_head_dim[-1],
+                resnet_groups=norm_num_groups,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                upcast_attention=upcast_attention,
+            )
+        else:
+            raise ValueError(f"unknown mid_block_type : {mid_block_type}")
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_attention_head_dim = list(reversed(attention_head_dim))
+        only_cross_attention = list(reversed(only_cross_attention))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=layers_per_block + 1,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=reversed_attention_head_dim[i],
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(
+            num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
+        )
+        self.conv_act = nn.SiLU()
+        self.conv_out = PseudoConv3d(block_out_channels[0], out_channels, kernel_size=3, padding=1)
+    @property
+    def attn_processors(self) -> Dict[str, AttnProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttnProcessor]):
+            if hasattr(module, "set_processor"):
+                processors[f"{name}.processor"] = module.processor
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(self, processor: Union[AttnProcessor, Dict[str, AttnProcessor]]):
+        r"""
+        Parameters:
+            `processor (`dict` of `AttnProcessor` or `AttnProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                of **all** `CrossAttention` layers.
+            In case `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainablae attention processors.:
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+
+    def set_attention_slice(self, slice_size):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module will split the input tensor in slices, to compute attention
+        in several steps. This is useful to save some memory in exchange for a small speed decrease.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
+                `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_slicable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_slicable_dims(module)
+
+        num_slicable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_slicable_layers * [1]
+
+        slice_size = (
+            num_slicable_layers * [slice_size] if not isinstance(slice_size, list) else slice_size
+        )
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(module: torch.nn.Module, slice_size: List[int]):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(
+            module,
+            (CrossAttnDownBlockPseudo3D, DownBlockPseudo3D, CrossAttnUpBlockPseudo3D, UpBlockPseudo3D),
+        ):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        fps_labels: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs=None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNetPseudo3DConditionOutput, Tuple]:
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # prepare attention_mask
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=self.dtype)
+        emb = self.time_embedding(t_emb)
+
+        if self.fps_embedding is not None:
+            if fps_labels is None:
+                raise ValueError("fps_labels should be provided when num_fps_embeds > 0")
+
+            if self.config.fps_embed_type == "timestep":
+                fps_labels = self.time_proj(fps_labels) # 和timesteps共用，都是sin embedding？这里的weight不更新的。
+
+            # 这里和上面timesteps does not contain any weights and will always return f32 tensors的bug一样。需要先cast过去，不然多机多卡就有问题了。
+            fps_labels = fps_labels.to(dtype=self.dtype)
+            class_emb = self.fps_embedding(fps_labels)
+
+            emb = emb + class_emb
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+            down_block_res_samples += res_samples
+
+        if down_block_additional_residuals is not None:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples += (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        sample = self.mid_block(
+            sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
+        )
+        if mid_block_additional_residual is not None:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                )
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNetPseudo3DConditionOutput(sample=sample)
+
+    @classmethod
+    def from_2d_model(cls, model_path, condition_on_fps=False):
+        '''
+        load a 2d model and convert it to a pseudo 3d model
+        '''
+        config_path = os.path.join(model_path, "config.json")
+        if not os.path.isfile(config_path):
+            raise RuntimeError(f"{config_path} does not exist")
+        with open(config_path, "r") as f:
+            config = json.load(f)
+
+        config.pop("_class_name")
+        config.pop("_diffusers_version")
+
+        block_replacer = {
+            "CrossAttnDownBlock2D": "CrossAttnDownBlockPseudo3D",
+            "DownBlock2D": "DownBlockPseudo3D",
+            "UpBlock2D": "UpBlockPseudo3D",
+            "CrossAttnUpBlock2D": "CrossAttnUpBlockPseudo3D",
+        }
+
+        def convert_2d_to_3d_block(block):
+            return block_replacer[block] if block in block_replacer else block
+
+        config["down_block_types"] = [
+            convert_2d_to_3d_block(block) for block in config["down_block_types"]
+        ]
+        config["up_block_types"] = [convert_2d_to_3d_block(block) for block in config["up_block_types"]]
+        
+        if condition_on_fps:
+            # config["num_fps_embeds"] = 60 # 这个在 trainable embeding时候才需要～
+            config["fps_embed_type"] = "timestep"     # 和timestep保持一致的type。
+            
+
+        model = cls(**config)   # 调用自身(init), 传入config参数全换成3d的setting
+
+        state_dict_path_condidates = glob.glob(os.path.join(model_path, "*.bin"))
+        if state_dict_path_condidates:
+            state_dict = torch.load(state_dict_path_condidates[0], map_location="cpu")
+            model.load_2d_state_dict(state_dict=state_dict)
+
+        return model
+
+    def load_2d_state_dict(self, state_dict, **kwargs):
+        '''
+        2D 部分的参数名完全不变。
+        '''
+        state_dict_3d = self.state_dict()
+
+        for k, v in state_dict.items():
+            if k not in state_dict_3d:
+                raise KeyError(f"2d state_dict key {k} does not exist in 3d model")
+            elif v.shape != state_dict_3d[k].shape:
+                raise ValueError(f"state_dict shape mismatch, 2d {v.shape}, 3d {state_dict_3d[k].shape}")
+
+        for k, v in state_dict_3d.items():
+            if "_temporal" in k:
+                continue
+            if "gamma" in k:
+                continue
+            
+            if k not in state_dict:
+                if "fps_embedding" in k:
+                    # 忽略检查fps_embedding
+                    continue
+                raise KeyError(f"3d state_dict key {k} does not exist in 2d model")
+
+        state_dict_3d.update(state_dict)
+        self.load_state_dict(state_dict_3d, **kwargs)

model/video_diffusion/pipelines/pipeline_st_stable_diffusion.py

@@ -0,0 +1,618 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Callable, List, Optional, Union
+
+import torch
+from einops import rearrange
+
+from diffusers.utils import is_accelerate_available
+from packaging import version
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.models import AutoencoderKL
+from diffusers.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+)
+from diffusers.utils import deprecate, logging
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+
+from ..models.unet_3d_condition import UNetPseudo3DConditionModel
+import os, importlib
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class SpatioTemporalStableDiffusionPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-video generation using Spatio-Temporal Stable Diffusion.
+    改变了unet的输入, unet换成3d unet, 其他部分完全和原来2D的一致。
+    latents的变为 b,c,f,h,w 原来是 b,c,h,w。
+    要用VAE的decoder的时候, 把输入reshape 成 (b f) c h w
+    """
+    _optional_components = []
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNetPseudo3DConditionModel,
+        scheduler: Union[
+            DDIMScheduler,
+            PNDMScheduler,
+            LMSDiscreteScheduler,
+            EulerDiscreteScheduler,
+            EulerAncestralDiscreteScheduler,
+            DPMSolverMultistepScheduler,
+        ],
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
+            version.parse(unet.config._diffusers_version).base_version
+        ) < version.parse("0.9.0.dev0")
+        is_unet_sample_size_less_64 = (
+            hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
+        )
+        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
+            deprecation_message = (
+                "The configuration file of the unet has set the default `sample_size` to smaller than"
+                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
+                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
+                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
+                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
+                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
+                " in the config might lead to incorrect results in future versions. If you have downloaded this"
+                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
+                " the `unet/config.json` file"
+            )
+            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(unet.config)
+            new_config["sample_size"] = 64
+            unet._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
+        steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_sequential_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
+        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
+        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
+        """
+        if is_accelerate_available():
+            from accelerate import cpu_offload
+        else:
+            raise ImportError("Please install accelerate via `pip install accelerate`")
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+            if cpu_offloaded_model is not None:
+                cpu_offload(cpu_offloaded_model, device)
+
+    @property
+    def _execution_device(self):
+        r"""
+        Returns the device on which the pipeline's models will be executed. After calling
+        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+        hooks.
+        """
+        if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    def _encode_prompt(
+        self, prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `list(int)`):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`):
+                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
+                if `guidance_scale` is less than `1`).
+        """
+        batch_size = len(prompt) if isinstance(prompt, list) else 1
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer.batch_decode(
+                untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+            )
+
+        if (
+            hasattr(self.text_encoder.config, "use_attention_mask")
+            and self.text_encoder.config.use_attention_mask
+        ):
+            attention_mask = text_inputs.attention_mask.to(device)
+        else:
+            attention_mask = None
+
+        text_embeddings = self.text_encoder(
+            text_input_ids.to(self.text_encoder.device),    # FIXME 强制对齐device的位置
+            attention_mask=attention_mask,
+        )
+        text_embeddings = text_embeddings[0]
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        bs_embed, seq_len, _ = text_embeddings.shape
+        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
+        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            max_length = text_input_ids.shape[-1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if (
+                hasattr(self.text_encoder.config, "use_attention_mask")
+                and self.text_encoder.config.use_attention_mask
+            ):
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            uncond_embeddings = self.text_encoder(
+                uncond_input.input_ids.to(self.text_encoder.device), # 同上，强制位置对齐。
+                attention_mask=attention_mask,
+            )
+            uncond_embeddings = uncond_embeddings[0]
+
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = uncond_embeddings.shape[1]
+            uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)
+            uncond_embeddings = uncond_embeddings.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+        return text_embeddings
+
+    def decode_latents(self, latents):
+        b = latents.shape[0]
+        latents = 1 / 0.18215 * latents
+        
+        is_video = len(latents.shape) == 5
+        if is_video:
+            latents = rearrange(latents, "b c f h w -> (b f) c h w")
+
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+
+        image = image.cpu().float().numpy()
+        if is_video:
+            image = rearrange(image, "(b f) c h w -> b f h w c", b=b)
+        else:
+            image = rearrange(image, "b c h w -> b h w c")
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(self, prompt, height, width, callback_steps):
+        if not isinstance(prompt, str) and not isinstance(prompt, list):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(
+                f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
+            )
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        clip_length,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):  
+        if clip_length>0:
+            shape = (
+                batch_size,
+                num_channels_latents,
+                clip_length,
+                height // self.vae_scale_factor,
+                width // self.vae_scale_factor,
+            )
+        else:
+            shape = (
+                batch_size,
+                num_channels_latents,
+                height // self.vae_scale_factor,
+                width // self.vae_scale_factor,
+            )
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            rand_device = "cpu" if device.type == "mps" else device
+
+            if isinstance(generator, list):
+                shape = (1,) + shape[1:]
+                latents = [
+                    torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)
+                    for i in range(batch_size)
+                ]
+                latents = torch.cat(latents, dim=0).to(device)
+            else:
+                latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(
+                    device
+                )
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        fps_labels = None,
+        num_inference_steps: int = 50,
+        clip_length: int = 8,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`):
+                The prompt or prompts to guide the image generation.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. Ignored when not using guidance (i.e., ignored
+                if `guidance_scale` is less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
+            When returning a tuple, the first element is a list with the generated images, and the second element is a
+            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, according to the `safety_checker`.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(prompt, height, width, callback_steps)
+
+        # 2. Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_embeddings = self._encode_prompt(
+            prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.in_channels
+        
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            clip_length,
+            height,
+            width,
+            text_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+        latents_dtype = latents.dtype
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                # print(latent_model_input.shape, )
+                # predict the noise residual
+                if fps_labels:
+                    if isinstance(fps_labels, list):
+                        fps_labels = torch.tensor(fps_labels).to(self.unet.device)
+                    # 控制帧率
+                    noise_pred = self.unet(
+                        latent_model_input, t, encoder_hidden_states=text_embeddings, fps_labels=fps_labels,
+                    ).sample.to(dtype=latents_dtype)
+                else:
+                    noise_pred = self.unet(
+                        latent_model_input, t, encoder_hidden_states=text_embeddings
+                    ).sample.to(dtype=latents_dtype)
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+        # image[:, 1:, :, :, :] = image[:, 1:, :, :, :] + image[:, 0:1, :, :, :]    # 叠加残差
+
+        # 9. Run safety checker
+        has_nsfw_concept = None
+
+        # 10. Convert to PIL
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+    @staticmethod
+    def numpy_to_pil(images):
+        if len(images.shape)==5:
+            pil_images = []
+            for sequence in images:
+                pil_images.append(DiffusionPipeline.numpy_to_pil(sequence))
+            return pil_images
+        else:
+            return DiffusionPipeline.numpy_to_pil(images)
+    
+
+    # 改写一下 model_index.json的保存内容, Unet是新定义的，直接保存会导致读取的时候出错～
+    def to_json_string(self) -> str:
+        from diffusers import __version__
+        import json
+        import numpy as np 
+
+        config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {}
+        config_dict["_class_name"] = self.__class__.__name__
+        config_dict["_diffusers_version"] = __version__
+
+        def to_json_saveable(value):
+            if isinstance(value, np.ndarray):
+                value = value.tolist()
+            return value
+
+        config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()}
+        if 'unet' in config_dict:
+            config_dict["unet"] = [
+                                    "diffusers",
+                                    "UNet2DConditionModel"
+                                ]
+        if 'controlnet' in config_dict:
+            config_dict['controlnet'] = [
+                                "diffusers",
+                                "UNet2DConditionModel"
+                            ]
+        # 覆盖
+        return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"

model/video_diffusion/pipelines/pipeline_stable_diffusion_controlnet3d.py

@@ -0,0 +1,482 @@
+# Copyright 2023 Bytedance Ltd. and/or its affiliates
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .pipeline_st_stable_diffusion import SpatioTemporalStableDiffusionPipeline
+from typing import Callable, List, Optional, Union
+from diffusers.schedulers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+)
+from transformers import DPTForDepthEstimation
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler
+import torch
+from einops import rearrange, repeat
+import decord
+import cv2
+import random
+import numpy as np
+from ..models.unet_3d_condition import UNetPseudo3DConditionModel
+from ..models.controlnet3d import ControlNet3DModel
+
+
+class Controlnet3DStableDiffusionPipeline(SpatioTemporalStableDiffusionPipeline):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNetPseudo3DConditionModel,
+        controlnet: ControlNet3DModel,
+        scheduler: Union[
+            DDIMScheduler,
+            PNDMScheduler,
+            LMSDiscreteScheduler,
+            EulerDiscreteScheduler,
+            EulerAncestralDiscreteScheduler,
+            DPMSolverMultistepScheduler,
+        ],
+        annotator_model=None,
+
+    ):
+        super().__init__(vae, text_encoder, tokenizer, unet, scheduler)
+
+        self.annotator_model = annotator_model
+        self.controlnet = controlnet
+        self.unet = unet
+        self.vae = vae
+        self.tokenizer = tokenizer
+        self.text_encoder = text_encoder
+        self.scheduler = scheduler
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            controlnet=controlnet,
+            scheduler=scheduler,
+        )
+
+    @staticmethod
+    def get_frames_preprocess(data_path, num_frames=24, sampling_rate=1, begin_indice=0, return_np=False):
+        vr = decord.VideoReader(data_path,)
+        n_images = len(vr)
+        fps_vid = round(vr.get_avg_fps())
+        frame_indices = [begin_indice + i*sampling_rate for i in range(num_frames)]   # 随机取n帧
+        
+
+        while n_images <= frame_indices[-1]:
+            # 超过视频长度，采样率减小直至不超过。
+            sampling_rate -= 1
+            if sampling_rate == 0:
+                # NOTE 边界检查
+                return None, None
+            frame_indices = [i*sampling_rate for i in range(num_frames)]
+        frames = vr.get_batch(frame_indices).asnumpy()
+
+        if return_np:
+            return frames, fps_vid
+
+        frames = torch.from_numpy(frames).div(255) * 2 - 1
+        frames = rearrange(frames, "f h w c -> c f h w").unsqueeze(0)
+        return frames, fps_vid  
+
+    @torch.no_grad()
+    def get_canny_edge_map(self, frames, ):
+        # (b f) c h w"
+        # from tensor to numpy
+        inputs = frames.cpu().numpy() 
+        inputs = rearrange(inputs, 'f c h w -> f h w c')
+        # inputs from [-1, 1] to [0, 255]
+        inputs = (inputs + 1) * 127.5
+        inputs = inputs.astype(np.uint8)
+        lower_threshold = 100
+        higher_threshold = 200
+        edge_images = np.stack([cv2.Canny(inp, lower_threshold, higher_threshold) for inp in inputs])
+        # from numpy to tensors
+        edge_images = torch.from_numpy(edge_images).unsqueeze(1)   # f, 1, h, w
+        edge_images = edge_images.div(255)*2 - 1 
+        # print(torch.max(out_images), torch.min(out_images), out_images.dtype)
+        return edge_images.to(dtype= self.controlnet.dtype, device=self.controlnet.device)
+
+    @torch.no_grad()
+    def get_depth_map(self, frames, height, width, return_standard_norm=False ):
+        """
+            frames should be like: (f c h w), you may turn b f c h w -> (b f) c h w first
+        """
+        h,w = height, width
+        inputs = torch.nn.functional.interpolate(
+            frames,
+            size=(384, 384),
+            mode="bicubic",
+            antialias=True,
+        ) 
+        # 转类型和设备
+        inputs = inputs.to(dtype= self.annotator_model.dtype, device=self.annotator_model.device)
+
+        outputs = self.annotator_model(inputs)
+        predicted_depths = outputs.predicted_depth
+
+        # interpolate to original size
+        predictions = torch.nn.functional.interpolate(
+            predicted_depths.unsqueeze(1),
+            size=(h, w),
+            mode="bicubic",
+        )
+
+        # normalize output
+        if return_standard_norm:
+            depth_min = torch.amin(predictions, dim=[1, 2, 3], keepdim=True)
+            depth_max = torch.amax(predictions, dim=[1, 2, 3], keepdim=True)
+            predictions = 2.0 * (predictions - depth_min) / (depth_max - depth_min) - 1.0
+        else:
+            predictions -= torch.min(predictions)
+            predictions /= torch.max(predictions)
+
+        return predictions  
+
+
+    @torch.no_grad()
+    def get_hed_map(self, frames,):
+        if isinstance(frames, torch.Tensor):
+            # 输入的就是 b c h w的tensor 范围是-1~1，需要转换为0～1
+            frames = (frames + 1) / 2
+            #rgb转bgr
+            bgr_frames = frames.clone()
+            bgr_frames[:, 0, :, :] = frames[:, 2, :, :]
+            bgr_frames[:, 2, :, :] = frames[:, 0, :, :]
+
+            edge = self.annotator_model(bgr_frames) # 范围也是0～1
+            return edge
+        else:
+            assert frames.ndim == 3
+            frames = frames[:, :, ::-1].copy()
+            with torch.no_grad():
+                image_hed = torch.from_numpy(frames).to(next(self.annotator_model.parameters()).device, dtype=next(self.annotator_model.parameters()).dtype )
+                image_hed = image_hed / 255.0
+                image_hed = rearrange(image_hed, 'h w c -> 1 c h w')
+                edge = self.annotator_model(image_hed)[0]
+                edge = (edge.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
+                return edge[0]
+    
+    @torch.no_grad()
+    def get_pose_map(self, frames,):
+        if isinstance(frames, torch.Tensor):
+            # 输入的就是 b c h w的tensor 范围是-1~1，需要转换为0～1
+            frames = (frames + 1) / 2
+            np_frames = frames.cpu().numpy() * 255
+            np_frames = np.array(np_frames, dtype=np.uint8)
+            np_frames = rearrange(np_frames, 'f c h w-> f h w c')
+            poses = np.stack([self.annotator_model(inp) for inp in np_frames])
+        else:
+            poses = self.annotator_model(frames)
+        return poses
+    
+    def get_timesteps(self, num_inference_steps, strength,):
+        # get the original timestep using init_timestep
+        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+
+        t_start = max(num_inference_steps - init_timestep, 0)
+        timesteps = self.scheduler.timesteps[t_start:]
+
+        return timesteps, num_inference_steps - t_start
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        controlnet_hint = None,
+        fps_labels = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        clip_length: int = 8, # NOTE clip_length和images的帧数一致。
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        cross_attention_kwargs = None,
+        video_scale: float = 0.0,
+        controlnet_conditioning_scale: float = 1.0,
+        fix_first_frame=True,
+        first_frame_output = None ,  # 也可以允许挑好图后传入。
+        first_frame_output_latent = None,
+        first_frame_control_hint = None,    # 维持第一帧
+        add_first_frame_by_concat = False,
+        controlhint_in_uncond = False,
+        init_same_noise_per_frame=False,
+        init_noise_by_residual_thres=0.0,
+        images=None,
+        in_domain=False,    # 是否调用视频模型生成图片
+        residual_control_steps=1,
+        first_frame_ddim_strength=1.0,
+        return_last_latent = False,
+    ):
+        '''
+        add origin video frames to get depth maps
+        '''
+        
+        if fix_first_frame and first_frame_output is None and first_frame_output_latent is None:
+            first_frame_output = self.__call__(
+                prompt=prompt,
+                controlnet_hint=controlnet_hint[:,:,0,:,:] if not in_domain else controlnet_hint[:,:,0:1,:,:],
+                # b c f h w
+                num_inference_steps=20,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=1,
+                generator=generator,
+                fix_first_frame=False,
+                controlhint_in_uncond=controlhint_in_uncond,
+            ).images[0]
+        
+
+        if first_frame_output is not None:
+            if isinstance(first_frame_output, list):
+                first_frame_output = first_frame_output[0] 
+            first_frame_output = torch.from_numpy(np.array(first_frame_output)).div(255) * 2 - 1
+            first_frame_output = rearrange(first_frame_output, "h w c -> c h w").unsqueeze(0)   # FIXME 目前不允许多个batch 先设置为1
+            first_frame_output = first_frame_output.to(dtype= self.vae.dtype, device=self.vae.device)
+
+            first_frame_output_latent = self.vae.encode(first_frame_output).latent_dist.sample()
+            first_frame_output_latent = first_frame_output_latent * 0.18215
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(prompt, height, width, callback_steps)
+
+        # 2. Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 5.0
+
+        # 3. Encode input prompt
+        text_embeddings = self._encode_prompt(
+            prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.in_channels
+        if controlnet_hint is not None:
+            if len(controlnet_hint.shape) == 5:
+                clip_length = controlnet_hint.shape[2]
+            else:
+                clip_length = 0
+            
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            clip_length,
+            height,
+            width,
+            text_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+        latents_dtype = latents.dtype
+
+
+        if len(latents.shape) == 5 and init_same_noise_per_frame:
+            latents[:,:,1:,:,:] = latents[:,:,0:1,:,:]
+
+        if len(latents.shape) == 5 and init_noise_by_residual_thres > 0.0 and images is not None:
+            
+            images = images.to(device=device, dtype=latents_dtype)  # b c f h w
+            image_residual = torch.abs(images[:,:,1:,:,:] - images[:,:,:-1,:,:])
+            images = rearrange(images, "b c f h w -> (b f) c h w")
+            
+            # norm residual
+            image_residual = image_residual / torch.max(image_residual)
+            
+            image_residual = rearrange(image_residual, "b c f h w -> (b f) c h w")
+            image_residual = torch.nn.functional.interpolate(
+                        image_residual, 
+                        size=(latents.shape[-2], latents.shape[-1]),
+                        mode='bilinear')
+            image_residual = torch.mean(image_residual, dim=1)
+
+            image_residual_mask = (image_residual > init_noise_by_residual_thres).float()
+            image_residual_mask = repeat(image_residual_mask, '(b f) h w -> b f h w', b=batch_size)
+            image_residual_mask = repeat(image_residual_mask, 'b f h w -> b c f h w', c=latents.shape[1])
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            if fix_first_frame:
+                if add_first_frame_by_concat:
+                    if len(first_frame_output_latent.shape) == 4:
+                        latents = torch.cat([first_frame_output_latent.unsqueeze(2), latents], dim=2)
+                    else:
+                        latents = torch.cat([first_frame_output_latent, latents], dim=2)
+                    if first_frame_control_hint is not None:
+                        controlnet_hint = torch.cat([first_frame_control_hint, controlnet_hint], dim=2)
+                    else:
+                        controlnet_hint = torch.cat([controlnet_hint[:,:,0:1 ,:,:], controlnet_hint], dim=2)
+            
+            if controlhint_in_uncond:
+                controlnet_hint = torch.cat([controlnet_hint] * 2) if do_classifier_free_guidance else controlnet_hint
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                if i<residual_control_steps and  len(latents.shape) == 5 and init_noise_by_residual_thres > 0.0 and images is not None :
+                    if first_frame_ddim_strength < 1.0  and i == 0 :
+                        # NOTE DDIM to get the first noise
+                        first_frame_output_latent_DDIM = first_frame_output_latent.clone()
+                        full_noise_timestep, _ = self.get_timesteps(num_inference_steps, strength=first_frame_ddim_strength)
+                        latent_timestep = full_noise_timestep[:1].repeat(batch_size * num_images_per_prompt)
+                        first_frame_output_latent_DDIM = self.scheduler.add_noise(first_frame_output_latent_DDIM, latents[:,:,0,:,:], latent_timestep)
+                        latents[:,:,0,:,:]=first_frame_output_latent_DDIM
+                    begin_frame = 1
+                    for n_frame in range(begin_frame, latents.shape[2]):
+                        latents[:,:, n_frame, :, :] = \
+                            (latents[:,:, n_frame, :, :] - latents[:,:, n_frame-1, :, :]) \
+                            * image_residual_mask[:,:, n_frame-1, :, :] + \
+                            latents[:,:, n_frame-1, :, :]
+                if fix_first_frame:
+                    latents[:,:,0 ,:,:] = first_frame_output_latent
+
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                if controlnet_hint is not None:
+                    down_block_res_samples, mid_block_res_sample = self.controlnet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                        controlnet_cond=controlnet_hint,
+                        return_dict=False,
+                    )
+                    down_block_res_samples = [
+                    down_block_res_sample * controlnet_conditioning_scale
+                    for down_block_res_sample in down_block_res_samples
+                    ]
+                    mid_block_res_sample *= controlnet_conditioning_scale
+
+                    noise_pred = self.unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        down_block_additional_residuals=down_block_res_samples,
+                        mid_block_additional_residual=mid_block_res_sample,
+                    ).sample.to(dtype=latents_dtype)
+                else:
+                    # predict the noise residual
+                    noise_pred = self.unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=text_embeddings,
+                    ).sample.to(dtype=latents_dtype)
+
+                if video_scale > 0 and controlnet_hint is not None:
+                    bsz = latents.shape[0]
+                    f = latents.shape[2]
+                    # 逐帧预测
+                    latent_model_input_single_frame = rearrange(latent_model_input, 'b c f h w -> (b f) c h w')
+                    text_embeddings_single_frame = torch.cat([text_embeddings] * f, dim=0)
+                    control_maps_single_frame = rearrange(controlnet_hint, 'b c f h w -> (b f) c h w')
+                    latent_model_input_single_frame = latent_model_input_single_frame.chunk(2, dim=0)[0]
+                    text_embeddings_single_frame = text_embeddings_single_frame.chunk(2, dim=0)[0]
+                    if controlhint_in_uncond:
+                        control_maps_single_frame = control_maps_single_frame.chunk(2, dim=0)[0]
+
+                    down_block_res_samples_single_frame, mid_block_res_sample_single_frame = self.controlnet(
+                                latent_model_input_single_frame,
+                                t,
+                                encoder_hidden_states=text_embeddings_single_frame,
+                                controlnet_cond=control_maps_single_frame,
+                                return_dict=False,
+                            )
+                    down_block_res_samples_single_frame = [
+                    down_block_res_sample_single_frame * controlnet_conditioning_scale
+                    for down_block_res_sample_single_frame in down_block_res_samples_single_frame
+                    ]
+                    mid_block_res_sample_single_frame *= controlnet_conditioning_scale
+
+                    noise_pred_single_frame_uncond = self.unet(
+                                latent_model_input_single_frame,
+                                t,
+                                encoder_hidden_states = text_embeddings_single_frame,
+                                down_block_additional_residuals=down_block_res_samples_single_frame,
+                                mid_block_additional_residual=mid_block_res_sample_single_frame,
+                                ).sample
+                    noise_pred_single_frame_uncond = rearrange(noise_pred_single_frame_uncond, '(b f) c h w -> b c f h w', f=f)
+                # perform guidance
+                if do_classifier_free_guidance:
+                    if video_scale > 0 and controlnet_hint is not None:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_single_frame_uncond + video_scale * (
+                            noise_pred_uncond - noise_pred_single_frame_uncond
+                        ) + guidance_scale * (
+                            noise_pred_text - noise_pred_uncond
+                        )
+                    else:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (
+                            noise_pred_text - noise_pred_uncond
+                        )
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+                
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+        if add_first_frame_by_concat:
+            image = image[:,1:,:,:,:]
+            
+        # 9. Run safety checker
+        has_nsfw_concept = None
+        # 10. Convert to PIL
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        if return_last_latent:
+            last_latent = latents[:,:,-1,:,:]
+            return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept), last_latent 
+        else:
+            return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

requirements.txt

@@ -0,0 +1,8 @@
+diffusers==0.14.0
+transformers==4.27.3
+accelerate==0.18.0
+xformers==0.0.16
+imageio==2.27.0
+decord==0.6.0
+opencv-python==4.7.0.72
+einops==0.6.0