code: 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,
      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,
      merge_strategy: str = "fixed",
      merge_factor: float = 0.5,
      spatial_transformer_attn_type: str = "softmax",
      video_kernel_size: Union[int, List[int]] = 3,
      use_linear_in_transformer: bool = False,
      adm_in_channels: Optional[int] = None,
      disable_temporal_crossattention: bool = False,
      max_ddpm_temb_period: int = 10000,
  ):
      super().__init__()
      assert context_dim is not None
  
      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
  
      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,
      ):
          return SpatialVideoTransformer(
              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=merge_strategy,
              merge_factor=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,
              max_time_embed_period=max_ddpm_temb_period,
          )
  
      def get_resblock(
          merge_factor,
          merge_strategy,
          video_kernel_size,
          ch,
          time_embed_dim,
          dropout,
          out_ch,
          dims,
          use_checkpoint,
          use_scale_shift_norm,
          down=False,
          up=False,
      ):
          return VideoResBlock(
              merge_factor=merge_factor,
              merge_strategy=merge_strategy,
              video_kernel_size=video_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,
          )
  
      for level, mult in enumerate(channel_mult):
          for _ in range(num_res_blocks):
              layers = [
                  get_resblock(
                      merge_factor=merge_factor,
                      merge_strategy=merge_strategy,
                      video_kernel_size=video_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
  
                  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(
                          merge_factor=merge_factor,
                          merge_strategy=merge_strategy,
                          video_kernel_size=video_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
  
      self.middle_block = TimestepEmbedSequential(
          get_resblock(
              merge_factor=merge_factor,
              merge_strategy=merge_strategy,
              video_kernel_size=video_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(
              merge_factor=merge_factor,
              merge_strategy=merge_strategy,
              video_kernel_size=video_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(
                      merge_factor=merge_factor,
                      merge_strategy=merge_strategy,
                      video_kernel_size=video_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
  
                  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(
                          merge_factor=merge_factor,
                          merge_strategy=merge_strategy,
                          video_kernel_size=video_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)
          ),
      )
  
file_path: generative-models/sgm/modules/diffusionmodules/video_model.py
imports: from functools import partial
  from typing import List, Optional, Union
  from einops import rearrange
  from ...modules.diffusionmodules.openaimodel import *
  from ...modules.video_attention import SpatialVideoTransformer
  from ...util import default
  from .util import AlphaBlender
label_desc: Use image node_id 1 to create a new VideoUNet object from inherited base classes: nn with example: obj = VideoUNet(in_channels, model_channels, out_channels, num_res_blocks, attention_resolutions, dropout, channel_mult, conv_resample, dims, num_classes, use_checkpoint, num_heads, num_head_channels, num_heads_upsample, use_scale_shift_norm, resblock_updown, transformer_depth, transformer_depth_middle, context_dim, time_downup, time_context_dim, extra_ff_mix_layer, use_spatial_context, merge_strategy, merge_factor, spatial_transformer_attn_type, video_kernel_size, use_linear_in_transformer, adm_in_channels, disable_temporal_crossattention, max_ddpm_temb_period)