Adds SV4D code

sgm/modules/diffusionmodules/guiders.py

@@ -94,7 +94,7 @@ class LinearPredictionGuider(Guider):
             if k in ["vector", "crossattn", "concat"] + self.additional_cond_keys:
                 c_out[k] = torch.cat((uc[k], c[k]), 0)
             else:
-                assert c[k] == uc[k]
+                # assert c[k] == uc[k]
                 c_out[k] = c[k]
         return torch.cat([x] * 2), torch.cat([s] * 2), c_out
 
@@ -105,7 +105,7 @@ class TrianglePredictionGuider(LinearPredictionGuider):
         max_scale: float,
         num_frames: int,
         min_scale: float = 1.0,
-        period: float | List[float] = 1.0,
+        period: Union[float, List[float]] = 1.0,
         period_fusing: Literal["mean", "multiply", "max"] = "max",
         additional_cond_keys: Optional[Union[List[str], str]] = None,
     ):
@@ -129,3 +129,47 @@ class TrianglePredictionGuider(LinearPredictionGuider):
 
     def triangle_wave(self, values: torch.Tensor, period) -> torch.Tensor:
         return 2 * (values / period - torch.floor(values / period + 0.5)).abs()
+
+
+class TrapezoidPredictionGuider(LinearPredictionGuider):
+    def __init__(
+        self,
+        max_scale: float,
+        num_frames: int,
+        min_scale: float = 1.0,
+        edge_perc: float = 0.1,
+        additional_cond_keys: Optional[Union[List[str], str]] = None,
+    ):
+        super().__init__(max_scale, num_frames, min_scale, additional_cond_keys)
+
+        rise_steps = torch.linspace(min_scale, max_scale, int(num_frames * edge_perc))
+        fall_steps = torch.flip(rise_steps, [0])
+        self.scale = torch.cat(
+            [
+                rise_steps,
+                torch.ones(num_frames - 2 * int(num_frames * edge_perc)),
+                fall_steps,
+            ]
+        ).unsqueeze(0)
+
+        
+class SpatiotemporalPredictionGuider(LinearPredictionGuider):
+    def __init__(
+        self,
+        max_scale: float,
+        num_frames: int,
+        num_views: int = 1,
+        min_scale: float = 1.0,
+        additional_cond_keys: Optional[Union[List[str], str]] = None,
+    ):
+        super().__init__(max_scale, num_frames, min_scale, additional_cond_keys)
+        V = num_views
+        T = num_frames // V
+        scale = torch.zeros(num_frames).view(T, V)
+        scale += torch.linspace(0, 1, T)[:,None] * 0.5
+        scale += self.triangle_wave(torch.linspace(0, 1, V))[None,:] * 0.5
+        scale = scale.flatten()
+        self.scale = (scale * (max_scale - min_scale) + min_scale).unsqueeze(0)
+
+    def triangle_wave(self, values: torch.Tensor, period=1) -> torch.Tensor:
+        return 2 * (values / period - torch.floor(values / period + 0.5)).abs()

sgm/modules/diffusionmodules/openaimodel.py

@@ -75,20 +75,43 @@ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
         emb: th.Tensor,
         context: Optional[th.Tensor] = None,
         image_only_indicator: Optional[th.Tensor] = None,
+        cond_view: Optional[th.Tensor] = None,
+        cond_motion: Optional[th.Tensor] = None,
         time_context: Optional[int] = None,
         num_video_frames: Optional[int] = None,
+        time_step: Optional[int] = None,
+        name: Optional[str] = None,
     ):
-        from ...modules.diffusionmodules.video_model import VideoResBlock
+        from ...modules.diffusionmodules.video_model import VideoResBlock, PostHocResBlockWithTime
+        from ...modules.spacetime_attention import (
+            BasicTransformerTimeMixBlock,
+            PostHocSpatialTransformerWithTimeMixing,
+            PostHocSpatialTransformerWithTimeMixingAndMotion
+        )
 
         for layer in self:
             module = layer
 
-            if isinstance(module, TimestepBlock) and not isinstance(
-                module, VideoResBlock
+            if isinstance(
+                module,
+                (
+                    BasicTransformerTimeMixBlock,
+                    PostHocSpatialTransformerWithTimeMixing,
+                    PostHocSpatialTransformerWithTimeMixingAndMotion
+                ),
             ):
-                x = layer(x, emb)
-            elif isinstance(module, VideoResBlock):
-                x = layer(x, emb, num_video_frames, image_only_indicator)
+                x = layer(
+                    x,
+                    context,
+                    # cam,
+                    time_context,
+                    num_video_frames,
+                    image_only_indicator,
+                    cond_view,
+                    cond_motion,
+                    time_step,
+                    name,
+                )
             elif isinstance(module, SpatialVideoTransformer):
                 x = layer(
                     x,
@@ -96,7 +119,16 @@ class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
                     time_context,
                     num_video_frames,
                     image_only_indicator,
+                    # time_step,
                 )
+            elif isinstance(module, PostHocResBlockWithTime):
+                x = layer(x, emb, num_video_frames, image_only_indicator)
+            elif isinstance(module, VideoResBlock):
+                x = layer(x, emb, num_video_frames, image_only_indicator)
+            elif isinstance(module, TimestepBlock) and not isinstance(
+                module, VideoResBlock
+            ):
+                x = layer(x, emb)
             elif isinstance(module, SpatialTransformer):
                 x = layer(x, context)
             else:

sgm/modules/diffusionmodules/sigma_sampling.py

@@ -1,5 +1,5 @@
 import torch
-
+from typing import Optional, Union
 from ...util import default, instantiate_from_config
 
 
@@ -29,3 +29,10 @@ class DiscreteSampling:
             torch.randint(0, self.num_idx, (n_samples,)),
         )
         return self.idx_to_sigma(idx)
+
+
+class ZeroSampler:
+    def __call__(
+        self, n_samples: int, rand: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        return torch.zeros_like(default(rand, torch.randn((n_samples,)))) + 1.0e-5

sgm/modules/diffusionmodules/util.py

@@ -17,6 +17,36 @@ import torch.nn as nn
 from einops import rearrange, repeat
 
 
+def get_alpha(
+    merge_strategy: str,
+    mix_factor: Optional[torch.Tensor],
+    image_only_indicator: torch.Tensor,
+    apply_sigmoid: bool = True,
+    is_attn: bool = False,
+) -> torch.Tensor:
+    if merge_strategy == "fixed" or merge_strategy == "learned":
+        alpha = mix_factor
+    elif merge_strategy == "learned_with_images":
+        alpha = torch.where(
+            image_only_indicator.bool(),
+            torch.ones(1, 1, device=image_only_indicator.device),
+            rearrange(mix_factor, "... -> ... 1"),
+        )
+        if is_attn:
+            alpha = rearrange(alpha, "b t -> (b t) 1 1")
+        else:
+            alpha = rearrange(alpha, "b t -> b 1 t 1 1")
+    elif merge_strategy == "fixed_with_images":
+        alpha = image_only_indicator
+        if is_attn:
+            alpha = rearrange(alpha, "b t -> (b t) 1 1")
+        else:
+            alpha = rearrange(alpha, "b t -> b 1 t 1 1")
+    else:
+        raise NotImplementedError
+    return torch.sigmoid(alpha) if apply_sigmoid else alpha
+
+    
 def make_beta_schedule(
     schedule,
     n_timestep,

sgm/modules/diffusionmodules/video_model.py

@@ -5,8 +5,13 @@ from einops import rearrange
 
 from ...modules.diffusionmodules.openaimodel import *
 from ...modules.video_attention import SpatialVideoTransformer
+from ...modules.spacetime_attention import (
+    BasicTransformerTimeMixBlock,
+    PostHocSpatialTransformerWithTimeMixing,
+    PostHocSpatialTransformerWithTimeMixingAndMotion
+)
 from ...util import default
-from .util import AlphaBlender
+from .util import AlphaBlender # , LegacyAlphaBlenderWithBug, get_alpha
 
 
 class VideoResBlock(ResBlock):
@@ -491,3 +496,746 @@ class VideoUNet(nn.Module):
             )
         h = h.type(x.dtype)
         return self.out(h)
+
+
+class PostHocAttentionBlockWithTimeMixing(AttentionBlock):
+    def __init__(
+        self,
+        in_channels: int,
+        n_heads: int,
+        d_head: int,
+        use_checkpoint: bool = False,
+        use_new_attention_order: bool = False,
+        dropout: float = 0.0,
+        use_spatial_context: bool = False,
+        merge_strategy: bool = "fixed",
+        merge_factor: float = 0.5,
+        apply_sigmoid_to_merge: bool = True,
+        ff_in: bool = False,
+        attn_mode: str = "softmax",
+        disable_temporal_crossattention: bool = False,
+    ):
+        super().__init__(
+            in_channels,
+            n_heads,
+            d_head,
+            use_checkpoint=use_checkpoint,
+            use_new_attention_order=use_new_attention_order,
+        )
+        inner_dim = n_heads * d_head
+
+        self.time_mix_blocks = nn.ModuleList(
+            [
+                BasicTransformerTimeMixBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    checkpoint=use_checkpoint,
+                    ff_in=ff_in,
+                    attn_mode=attn_mode,
+                    disable_temporal_crossattention=disable_temporal_crossattention,
+                )
+            ]
+        )
+        self.in_channels = in_channels
+
+        time_embed_dim = self.in_channels * 4
+        self.time_mix_time_embed = nn.Sequential(
+            linear(self.in_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, self.in_channels),
+        )
+
+        self.use_spatial_context = use_spatial_context
+
+        if merge_strategy == "fixed":
+            self.register_buffer("mix_factor", th.Tensor([merge_factor]))
+        elif merge_strategy == "learned" or merge_strategy == "learned_with_images":
+            self.register_parameter(
+                "mix_factor", th.nn.Parameter(th.Tensor([merge_factor]))
+            )
+        elif merge_strategy == "fixed_with_images":
+            self.mix_factor = None
+        else:
+            raise ValueError(f"unknown merge strategy {merge_strategy}")
+
+        self.get_alpha_fn = functools.partial(
+            get_alpha,
+            merge_strategy,
+            self.mix_factor,
+            apply_sigmoid=apply_sigmoid_to_merge,
+        )
+
+    def forward(
+        self,
+        x: th.Tensor,
+        context: Optional[th.Tensor] = None,
+        # cam: Optional[th.Tensor] = None,
+        time_context: Optional[th.Tensor] = None,
+        timesteps: Optional[int] = None,
+        image_only_indicator: Optional[th.Tensor] = None,
+        conv_view: Optional[th.Tensor] = None,
+        conv_motion: Optional[th.Tensor] = None,
+    ):
+        if time_context is not None:
+            raise NotImplementedError
+
+        _, _, h, w = x.shape
+        if exists(context):
+            context = rearrange(context, "b t ... -> (b t) ...")
+        if self.use_spatial_context:
+            time_context = repeat(context[:, 0], "b ... -> (b n) ...", n=h * w)
+
+        x = super().forward(
+            x,
+        )
+
+        x = rearrange(x, "b c h w -> b (h w) c")
+        x_mix = x
+
+        num_frames = th.arange(timesteps, device=x.device)
+        num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)
+        num_frames = rearrange(num_frames, "b t -> (b t)")
+        t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False)
+        emb = self.time_mix_time_embed(t_emb)
+        emb = emb[:, None, :]
+        x_mix = x_mix + emb
+
+        x_mix = self.time_mix_blocks[0](
+            x_mix, context=time_context, timesteps=timesteps
+        )
+
+        alpha = self.get_alpha_fn(image_only_indicator=image_only_indicator)
+        x = alpha * x + (1.0 - alpha) * x_mix
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
+        return x
+
+
+class PostHocResBlockWithTime(ResBlock):
+    def __init__(
+        self,
+        channels: int,
+        emb_channels: int,
+        dropout: float,
+        time_kernel_size: Union[int, List[int]] = 3,
+        merge_strategy: bool = "fixed",
+        merge_factor: float = 0.5,
+        apply_sigmoid_to_merge: bool = True,
+        out_channels: Optional[int] = None,
+        use_conv: bool = False,
+        use_scale_shift_norm: bool = False,
+        dims: int = 2,
+        use_checkpoint: bool = False,
+        up: bool = False,
+        down: bool = False,
+        time_mix_legacy: bool = True,
+        replicate_bug: bool = False,
+    ):
+        super().__init__(
+            channels,
+            emb_channels,
+            dropout,
+            out_channels=out_channels,
+            use_conv=use_conv,
+            use_scale_shift_norm=use_scale_shift_norm,
+            dims=dims,
+            use_checkpoint=use_checkpoint,
+            up=up,
+            down=down,
+        )
+
+        self.time_mix_blocks = ResBlock(
+            default(out_channels, channels),
+            emb_channels,
+            dropout=dropout,
+            dims=3,
+            out_channels=default(out_channels, channels),
+            use_scale_shift_norm=False,
+            use_conv=False,
+            up=False,
+            down=False,
+            kernel_size=time_kernel_size,
+            use_checkpoint=use_checkpoint,
+            exchange_temb_dims=True,
+        )
+        self.time_mix_legacy = time_mix_legacy
+        if self.time_mix_legacy:
+            if merge_strategy == "fixed":
+                self.register_buffer("mix_factor", th.Tensor([merge_factor]))
+            elif merge_strategy == "learned" or merge_strategy == "learned_with_images":
+                self.register_parameter(
+                    "mix_factor", th.nn.Parameter(th.Tensor([merge_factor]))
+                )
+            elif merge_strategy == "fixed_with_images":
+                self.mix_factor = None
+            else:
+                raise ValueError(f"unknown merge strategy {merge_strategy}")
+
+            self.get_alpha_fn = functools.partial(
+                get_alpha,
+                merge_strategy,
+                self.mix_factor,
+                apply_sigmoid=apply_sigmoid_to_merge,
+            )
+        else:
+            if False: # replicate_bug:
+                logpy.warning(
+                    "*****************************************************************************************\n"
+                    "GRAVE WARNING: YOU'RE USING THE BUGGY LEGACY ALPHABLENDER!!! ARE YOU SURE YOU WANT THIS?!\n"
+                    "*****************************************************************************************"
+                )
+                self.time_mixer = LegacyAlphaBlenderWithBug(
+                    alpha=merge_factor,
+                    merge_strategy=merge_strategy,
+                    rearrange_pattern="b t -> b 1 t 1 1",
+                )
+            else:
+                self.time_mixer = AlphaBlender(
+                    alpha=merge_factor,
+                    merge_strategy=merge_strategy,
+                    rearrange_pattern="b t -> b 1 t 1 1",
+                )
+
+    def forward(
+        self,
+        x: th.Tensor,
+        emb: th.Tensor,
+        num_video_frames: int,
+        image_only_indicator: Optional[th.Tensor] = None,
+        cond_view: Optional[th.Tensor] = None,
+        cond_motion: Optional[th.Tensor] = None,
+    ) -> th.Tensor:
+        x = super().forward(x, emb)
+
+        x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=num_video_frames)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=num_video_frames)
+
+        x = self.time_mix_blocks(
+            x, rearrange(emb, "(b t) ... -> b t ...", t=num_video_frames)
+        )
+
+        if self.time_mix_legacy:
+            alpha = self.get_alpha_fn(image_only_indicator=image_only_indicator)
+            x = alpha.to(x.dtype) * x + (1.0 - alpha).to(x.dtype) * x_mix
+        else:
+            x = self.time_mixer(
+                x_spatial=x_mix, x_temporal=x, image_only_indicator=image_only_indicator
+            )
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        return x
+
+
+class SpatialUNetModelWithTime(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        model_channels: int,
+        out_channels: int,
+        num_res_blocks: int,
+        attention_resolutions: int,
+        dropout: float = 0.0,
+        channel_mult: List[int] = (1, 2, 4, 8),
+        conv_resample: bool = True,
+        dims: int = 2,
+        num_classes: Optional[int] = None,
+        use_checkpoint: bool = False,
+        num_heads: int = -1,
+        num_head_channels: int = -1,
+        num_heads_upsample: int = -1,
+        use_scale_shift_norm: bool = False,
+        resblock_updown: bool = False,
+        use_new_attention_order: bool = False,
+        use_spatial_transformer: bool = False,
+        transformer_depth: Union[List[int], int] = 1,
+        transformer_depth_middle: Optional[int] = None,
+        context_dim: Optional[int] = None,
+        time_downup: bool = False,
+        time_context_dim: Optional[int] = None,
+        extra_ff_mix_layer: bool = False,
+        use_spatial_context: bool = False,
+        time_block_merge_strategy: str = "fixed",
+        time_block_merge_factor: float = 0.5,
+        spatial_transformer_attn_type: str = "softmax",
+        time_kernel_size: Union[int, List[int]] = 3,
+        use_linear_in_transformer: bool = False,
+        legacy: bool = True,
+        adm_in_channels: Optional[int] = None,
+        use_temporal_resblock: bool = True,
+        disable_temporal_crossattention: bool = False,
+        time_mix_legacy: bool = True,
+        max_ddpm_temb_period: int = 10000,
+        replicate_time_mix_bug: bool = False,
+        use_motion_attention: bool = False,
+    ):
+        super().__init__()
+
+        if use_spatial_transformer:
+            assert context_dim is not None
+
+        if context_dim is not None:
+            assert use_spatial_transformer
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        if num_heads == -1:
+            assert num_head_channels != -1
+
+        if num_head_channels == -1:
+            assert num_heads != -1
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        if isinstance(transformer_depth, int):
+            transformer_depth = len(channel_mult) * [transformer_depth]
+        transformer_depth_middle = default(
+            transformer_depth_middle, transformer_depth[-1]
+        )
+
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.use_temporal_resblocks = use_temporal_resblock
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if self.num_classes is not None:
+            if isinstance(self.num_classes, int):
+                self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+            elif self.num_classes == "continuous":
+                print("setting up linear c_adm embedding layer")
+                self.label_emb = nn.Linear(1, time_embed_dim)
+            elif self.num_classes == "timestep":
+                self.label_emb = nn.Sequential(
+                    Timestep(model_channels),
+                    nn.Sequential(
+                        linear(model_channels, time_embed_dim),
+                        nn.SiLU(),
+                        linear(time_embed_dim, time_embed_dim),
+                    ),
+                )
+
+            elif self.num_classes == "sequential":
+                assert adm_in_channels is not None
+                self.label_emb = nn.Sequential(
+                    nn.Sequential(
+                        linear(adm_in_channels, time_embed_dim),
+                        nn.SiLU(),
+                        linear(time_embed_dim, time_embed_dim),
+                    )
+                )
+            else:
+                raise ValueError()
+
+        self.input_blocks = nn.ModuleList(
+            [
+                TimestepEmbedSequential(
+                    conv_nd(dims, in_channels, model_channels, 3, padding=1)
+                )
+            ]
+        )
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+
+        def get_attention_layer(
+            ch,
+            num_heads,
+            dim_head,
+            depth=1,
+            context_dim=None,
+            use_checkpoint=False,
+            disabled_sa=False,
+        ):
+            if not use_spatial_transformer:
+                return PostHocAttentionBlockWithTimeMixing(
+                    ch,
+                    num_heads,
+                    dim_head,
+                    use_checkpoint=use_checkpoint,
+                    use_new_attention_order=use_new_attention_order,
+                    dropout=dropout,
+                    ff_in=extra_ff_mix_layer,
+                    use_spatial_context=use_spatial_context,
+                    merge_strategy=time_block_merge_strategy,
+                    merge_factor=time_block_merge_factor,
+                    attn_mode=spatial_transformer_attn_type,
+                    disable_temporal_crossattention=disable_temporal_crossattention,
+                )
+
+            elif use_motion_attention:
+                return PostHocSpatialTransformerWithTimeMixingAndMotion(
+                    ch,
+                    num_heads,
+                    dim_head,
+                    depth=depth,
+                    context_dim=context_dim,
+                    time_context_dim=time_context_dim,
+                    dropout=dropout,
+                    ff_in=extra_ff_mix_layer,
+                    use_spatial_context=use_spatial_context,
+                    merge_strategy=time_block_merge_strategy,
+                    merge_factor=time_block_merge_factor,
+                    checkpoint=use_checkpoint,
+                    use_linear=use_linear_in_transformer,
+                    attn_mode=spatial_transformer_attn_type,
+                    disable_self_attn=disabled_sa,
+                    disable_temporal_crossattention=disable_temporal_crossattention,
+                    time_mix_legacy=time_mix_legacy,
+                    max_time_embed_period=max_ddpm_temb_period,
+                )
+            
+            else:
+                return PostHocSpatialTransformerWithTimeMixing(
+                    ch,
+                    num_heads,
+                    dim_head,
+                    depth=depth,
+                    context_dim=context_dim,
+                    time_context_dim=time_context_dim,
+                    dropout=dropout,
+                    ff_in=extra_ff_mix_layer,
+                    use_spatial_context=use_spatial_context,
+                    merge_strategy=time_block_merge_strategy,
+                    merge_factor=time_block_merge_factor,
+                    checkpoint=use_checkpoint,
+                    use_linear=use_linear_in_transformer,
+                    attn_mode=spatial_transformer_attn_type,
+                    disable_self_attn=disabled_sa,
+                    disable_temporal_crossattention=disable_temporal_crossattention,
+                    time_mix_legacy=time_mix_legacy,
+                    max_time_embed_period=max_ddpm_temb_period,
+                )
+
+        def get_resblock(
+            time_block_merge_factor,
+            time_block_merge_strategy,
+            time_kernel_size,
+            ch,
+            time_embed_dim,
+            dropout,
+            out_ch,
+            dims,
+            use_checkpoint,
+            use_scale_shift_norm,
+            down=False,
+            up=False,
+        ):
+            if self.use_temporal_resblocks:
+                return PostHocResBlockWithTime(
+                    merge_factor=time_block_merge_factor,
+                    merge_strategy=time_block_merge_strategy,
+                    time_kernel_size=time_kernel_size,
+                    channels=ch,
+                    emb_channels=time_embed_dim,
+                    dropout=dropout,
+                    out_channels=out_ch,
+                    dims=dims,
+                    use_checkpoint=use_checkpoint,
+                    use_scale_shift_norm=use_scale_shift_norm,
+                    down=down,
+                    up=up,
+                    time_mix_legacy=time_mix_legacy,
+                    replicate_bug=replicate_time_mix_bug,
+                )
+            else:
+                return ResBlock(
+                    channels=ch,
+                    emb_channels=time_embed_dim,
+                    dropout=dropout,
+                    out_channels=out_ch,
+                    use_checkpoint=use_checkpoint,
+                    dims=dims,
+                    use_scale_shift_norm=use_scale_shift_norm,
+                    down=down,
+                    up=up,
+                )
+
+        for level, mult in enumerate(channel_mult):
+            for _ in range(num_res_blocks):
+                layers = [
+                    get_resblock(
+                        time_block_merge_factor=time_block_merge_factor,
+                        time_block_merge_strategy=time_block_merge_strategy,
+                        time_kernel_size=time_kernel_size,
+                        ch=ch,
+                        time_embed_dim=time_embed_dim,
+                        dropout=dropout,
+                        out_ch=mult * model_channels,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = (
+                            ch // num_heads
+                            if use_spatial_transformer
+                            else num_head_channels
+                        )
+
+                    layers.append(
+                        get_attention_layer(
+                            ch,
+                            num_heads,
+                            dim_head,
+                            depth=transformer_depth[level],
+                            context_dim=context_dim,
+                            use_checkpoint=use_checkpoint,
+                            disabled_sa=False,
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                ds *= 2
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        get_resblock(
+                            time_block_merge_factor=time_block_merge_factor,
+                            time_block_merge_strategy=time_block_merge_strategy,
+                            time_kernel_size=time_kernel_size,
+                            ch=ch,
+                            time_embed_dim=time_embed_dim,
+                            dropout=dropout,
+                            out_ch=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            down=True,
+                        )
+                        if resblock_updown
+                        else Downsample(
+                            ch,
+                            conv_resample,
+                            dims=dims,
+                            out_channels=out_ch,
+                            third_down=time_downup,
+                        )
+                    )
+                )
+                ch = out_ch
+                input_block_chans.append(ch)
+
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            num_heads = ch // num_head_channels
+            dim_head = num_head_channels
+        if legacy:
+            # num_heads = 1
+            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+
+        self.middle_block = TimestepEmbedSequential(
+            get_resblock(
+                time_block_merge_factor=time_block_merge_factor,
+                time_block_merge_strategy=time_block_merge_strategy,
+                time_kernel_size=time_kernel_size,
+                ch=ch,
+                time_embed_dim=time_embed_dim,
+                out_ch=None,
+                dropout=dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            get_attention_layer(
+                ch,
+                num_heads,
+                dim_head,
+                depth=transformer_depth_middle,
+                context_dim=context_dim,
+                use_checkpoint=use_checkpoint,
+            ),
+            get_resblock(
+                time_block_merge_factor=time_block_merge_factor,
+                time_block_merge_strategy=time_block_merge_strategy,
+                time_kernel_size=time_kernel_size,
+                ch=ch,
+                out_ch=None,
+                time_embed_dim=time_embed_dim,
+                dropout=dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(num_res_blocks + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    get_resblock(
+                        time_block_merge_factor=time_block_merge_factor,
+                        time_block_merge_strategy=time_block_merge_strategy,
+                        time_kernel_size=time_kernel_size,
+                        ch=ch + ich,
+                        time_embed_dim=time_embed_dim,
+                        dropout=dropout,
+                        out_ch=model_channels * mult,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = model_channels * mult
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = (
+                            ch // num_heads
+                            if use_spatial_transformer
+                            else num_head_channels
+                        )
+
+                    layers.append(
+                        get_attention_layer(
+                            ch,
+                            num_heads,
+                            dim_head,
+                            depth=transformer_depth[level],
+                            context_dim=context_dim,
+                            use_checkpoint=use_checkpoint,
+                            disabled_sa=False,
+                        )
+                    )
+                if level and i == num_res_blocks:
+                    out_ch = ch
+                    ds //= 2
+                    layers.append(
+                        get_resblock(
+                            time_block_merge_factor=time_block_merge_factor,
+                            time_block_merge_strategy=time_block_merge_strategy,
+                            time_kernel_size=time_kernel_size,
+                            ch=ch,
+                            time_embed_dim=time_embed_dim,
+                            dropout=dropout,
+                            out_ch=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(
+                            ch,
+                            conv_resample,
+                            dims=dims,
+                            out_channels=out_ch,
+                            third_up=time_downup,
+                        )
+                    )
+
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+        )
+
+    def forward(
+        self,
+        x: th.Tensor,
+        timesteps: th.Tensor,
+        context: Optional[th.Tensor] = None,
+        y: Optional[th.Tensor] = None,
+        # cam: Optional[th.Tensor] = None,
+        time_context: Optional[th.Tensor] = None,
+        num_video_frames: Optional[int] = None,
+        image_only_indicator: Optional[th.Tensor] = None,
+        cond_view: Optional[th.Tensor] = None,
+        cond_motion: Optional[th.Tensor] = None,
+        time_step: Optional[int] = None,
+    ):
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional -> no, relax this TODO"
+        hs = []
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) # 21 x 320
+        emb = self.time_embed(t_emb) # 21 x 1280
+        time = str(timesteps[0].data.cpu().numpy())
+
+        if self.num_classes is not None:
+            assert y.shape[0] == x.shape[0]
+            emb = emb + self.label_emb(y) # 21 x 1280
+
+        h = x # 21 x 8 x 64 x 64
+        for i, module in enumerate(self.input_blocks):
+            h = module(
+                h,
+                emb,
+                context=context,
+                # cam=cam,
+                image_only_indicator=image_only_indicator,
+                cond_view=cond_view,
+                cond_motion=cond_motion,
+                time_context=time_context,
+                num_video_frames=num_video_frames,
+                time_step=time_step,
+                name='encoder_{}_{}'.format(time, i)
+            )
+            hs.append(h)
+        h = self.middle_block(
+            h,
+            emb,
+            context=context,
+            # cam=cam,
+            image_only_indicator=image_only_indicator,
+            cond_view=cond_view,
+            cond_motion=cond_motion,
+            time_context=time_context,
+            num_video_frames=num_video_frames,
+            time_step=time_step,
+            name='middle_{}_0'.format(time, i)
+        )
+        for i, module in enumerate(self.output_blocks):
+            h = th.cat([h, hs.pop()], dim=1)
+            h = module(
+                h,
+                emb,
+                context=context,
+                # cam=cam,
+                image_only_indicator=image_only_indicator,
+                cond_view=cond_view,
+                cond_motion=cond_motion,
+                time_context=time_context,
+                num_video_frames=num_video_frames,
+                time_step=time_step,
+                name='decoder_{}_{}'.format(time, i)
+            )
+        h = h.type(x.dtype)
+        return self.out(h)

sgm/modules/diffusionmodules/wrappers.py

@@ -25,10 +25,21 @@ class OpenAIWrapper(IdentityWrapper):
         self, x: torch.Tensor, t: torch.Tensor, c: dict, **kwargs
     ) -> torch.Tensor:
         x = torch.cat((x, c.get("concat", torch.Tensor([]).type_as(x))), dim=1)
-        return self.diffusion_model(
-            x,
-            timesteps=t,
-            context=c.get("crossattn", None),
-            y=c.get("vector", None),
-            **kwargs,
-        )
+        if "cond_view" in c:
+            return self.diffusion_model(
+                x,
+                timesteps=t,
+                context=c.get("crossattn", None),
+                y=c.get("vector", None),
+                cond_view=c.get("cond_view", None),
+                cond_motion=c.get("cond_motion", None),
+                **kwargs,
+            )
+        else:
+            return self.diffusion_model(
+                x,
+                timesteps=t,
+                context=c.get("crossattn", None),
+                y=c.get("vector", None),
+                **kwargs,
+            )