class_docstr: A conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
  shaped output.
  
  This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
  for all models (such as downloading or saving).
  
  Parameters:
      sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
          Height and width of input/output sample.
      in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample.
      out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
      down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
          The tuple of downsample blocks to use.
      up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
          The tuple of upsample blocks to use.
      block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
          The tuple of output channels for each block.
      layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
      downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
      mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
      act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
      norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
          If `None`, normalization and activation layers is skipped in post-processing.
      norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
      cross_attention_dim (`int`, *optional*, defaults to 1280): The dimension of the cross attention features.
      attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
      num_attention_heads (`int`, *optional*): The number of attention heads.
class_docstr_tok: ["A","conditional","3D","UNet","model","that","takes","a","noisy","sample",",","conditional","state",",","and","a","timestep","and","returns","a","sample","shaped","output",".","This","model","inherits","from","[","`","ModelMixin","`","]",".","Check","the","superclass","documentation","for","it","'s","generic","methods","implemented","for","all","models","(","such","as","downloading","or","saving",")",".","Parameters",":","sample_size","(","`","int","`","or","`","Tuple","[","int",",","int","]","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Height","and","width","of","input\/output","sample",".","in_channels","(","`","int","`",",","*","optional","*",",","defaults","to","4",")",":","The","number","of","channels","in","the","input","sample",".","out_channels","(","`","int","`",",","*","optional","*",",","defaults","to","4",")",":","The","number","of","channels","in","the","output",".","down_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","CrossAttnDownBlock2D","''",",","``","CrossAttnDownBlock2D","''",",","``","CrossAttnDownBlock2D","''",",","``","DownBlock2D","''",")","`",")",":","The","tuple","of","downsample","blocks","to","use",".","up_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","UpBlock2D","''",",","``","CrossAttnUpBlock2D","''",",","``","CrossAttnUpBlock2D","''",",","``","CrossAttnUpBlock2D","''",")","`",")",":","The","tuple","of","upsample","blocks","to","use",".","block_out_channels","(","`","Tuple","[","int","]","`",",","*","optional","*",",","defaults","to","`","(","320",",","640",",","1280",",","1280",")","`",")",":","The","tuple","of","output","channels","for","each","block",".","layers_per_block","(","`","int","`",",","*","optional","*",",","defaults","to","2",")",":","The","number","of","layers","per","block",".","downsample_padding","(","`","int","`",",","*","optional","*",",","defaults","to","1",")",":","The","padding","to","use","for","the","downsampling","convolution",".","mid_block_scale_factor","(","`","float","`",",","*","optional","*",",","defaults","to","1.0",")",":","The","scale","factor","to","use","for","the","mid","block",".","act_fn","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","silu","''","`",")",":","The","activation","function","to","use",".","norm_num_groups","(","`","int","`",",","*","optional","*",",","defaults","to","32",")",":","The","number","of","groups","to","use","for","the","normalization",".","If","`","None","`",",","normalization","and","activation","layers","is","skipped","in","post-processing",".","norm_eps","(","`","float","`",",","*","optional","*",",","defaults","to","1e-5",")",":","The","epsilon","to","use","for","the","normalization",".","cross_attention_dim","(","`","int","`",",","*","optional","*",",","defaults","to","1280",")",":","The","dimension","of","the","cross","attention","features",".","attention_head_dim","(","`","int","`",",","*","optional","*",",","defaults","to","8",")",":","The","dimension","of","the","attention","heads",".","num_attention_heads","(","`","int","`",",","*","optional","*",")",":","The","number","of","attention","heads","."]
docstr: The [`UNet3DConditionModel`] forward method.
  
  Args:
      sample (`torch.FloatTensor`):
          The noisy input tensor with the following shape `(batch, num_frames, channel, height, width`.
      timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
      encoder_hidden_states (`torch.FloatTensor`):
          The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
      class_labels (`torch.Tensor`, *optional*, defaults to `None`):
          Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
      timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
          Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
          through the `self.time_embedding` layer to obtain the timestep embeddings.
      attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
          An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
          is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
          negative values to the attention scores corresponding to "discard" tokens.
      cross_attention_kwargs (`dict`, *optional*):
          A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
          `self.processor` in
          [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
      down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
          A tuple of tensors that if specified are added to the residuals of down unet blocks.
      mid_block_additional_residual: (`torch.Tensor`, *optional*):
          A tensor that if specified is added to the residual of the middle unet block.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`~models.unet_3d_condition.UNet3DConditionOutput`] instead of a plain
          tuple.
      cross_attention_kwargs (`dict`, *optional*):
          A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
  
  Returns:
      [`~models.unet_3d_condition.UNet3DConditionOutput`] or `tuple`:
          If `return_dict` is True, an [`~models.unet_3d_condition.UNet3DConditionOutput`] is returned, otherwise
          a `tuple` is returned where the first element is the sample tensor.
docstr_tok: ["The","[","`","UNet3DConditionModel","`","]","forward","method",".","Args",":","sample","(","`","torch.FloatTensor","`",")",":","The","noisy","input","tensor","with","the","following","shape","`","(","batch",",","num_frames",",","channel",",","height",",","width","`",".","timestep","(","`","torch.FloatTensor","`","or","`","float","`","or","`","int","`",")",":","The","number","of","timesteps","to","denoise","an","input",".","encoder_hidden_states","(","`","torch.FloatTensor","`",")",":","The","encoder","hidden","states","with","shape","`","(","batch",",","sequence_length",",","feature_dim",")","`",".","class_labels","(","`","torch.Tensor","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Optional","class","labels","for","conditioning",".","Their","embeddings","will","be","summed","with","the","timestep","embeddings",".","timestep_cond",":","(","`","torch.Tensor","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Conditional","embeddings","for","timestep",".","If","provided",",","the","embeddings","will","be","summed","with","the","samples","passed","through","the","`","self.time_embedding","`","layer","to","obtain","the","timestep","embeddings",".","attention_mask","(","`","torch.Tensor","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","An","attention","mask","of","shape","`","(","batch",",","key_tokens",")","`","is","applied","to","`","encoder_hidden_states","`",".","If","`","1","`","the","mask","is","kept",",","otherwise","if","`","0","`","it","is","discarded",".","Mask","will","be","converted","into","a","bias",",","which","adds","large","negative","values","to","the","attention","scores","corresponding","to","``","discard","''","tokens",".","cross_attention_kwargs","(","`","dict","`",",","*","optional","*",")",":","A","kwargs","dictionary","that","if","specified","is","passed","along","to","the","`","AttentionProcessor","`","as","defined","under","`","self.processor","`","in","[","diffusers.models.attention_processor","]","(","https",":","\/\/github.com\/huggingface\/diffusers\/blob\/main\/src\/diffusers\/models\/attention_processor.py",")",".","down_block_additional_residuals",":","(","`","tuple","`","of","`","torch.Tensor","`",",","*","optional","*",")",":","A","tuple","of","tensors","that","if","specified","are","added","to","the","residuals","of","down","unet","blocks",".","mid_block_additional_residual",":","(","`","torch.Tensor","`",",","*","optional","*",")",":","A","tensor","that","if","specified","is","added","to","the","residual","of","the","middle","unet","block",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~models.unet_3d_condition.UNet3DConditionOutput","`","]","instead","of","a","plain","tuple",".","cross_attention_kwargs","(","`","dict","`",",","*","optional","*",")",":","A","kwargs","dictionary","that","if","specified","is","passed","along","to","the","[","`","AttnProcessor","`","]",".","Returns",":","[","`","~models.unet_3d_condition.UNet3DConditionOutput","`","]","or","`","tuple","`",":","If","`","return_dict","`","is","True",",","an","[","`","~models.unet_3d_condition.UNet3DConditionOutput","`","]","is","returned",",","otherwise","a","`","tuple","`","is","returned","where","the","first","element","is","the","sample","tensor","."]
code: def forward(
      self,
      sample: torch.FloatTensor,
      timestep: Union[torch.Tensor, float, int],
      encoder_hidden_states: torch.Tensor,
      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,
      down_block_additional_residuals: Optional[
          Tuple[torch.Tensor]
      ] = None,
      mid_block_additional_residual: Optional[torch.Tensor] = None,
      return_dict: bool = True,
  ) -> Union[UNet3DConditionOutput, Tuple]:
      r"""
      The [`UNet3DConditionModel`] forward method.
  
      Args:
          sample (`torch.FloatTensor`):
              The noisy input tensor with the following shape `(batch, num_frames, channel, height, width`.
          timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
          encoder_hidden_states (`torch.FloatTensor`):
              The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
          class_labels (`torch.Tensor`, *optional*, defaults to `None`):
              Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
          timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
              Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
              through the `self.time_embedding` layer to obtain the timestep embeddings.
          attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
              An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
              is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
              negative values to the attention scores corresponding to "discard" tokens.
          cross_attention_kwargs (`dict`, *optional*):
              A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
              `self.processor` in
              [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
          down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
              A tuple of tensors that if specified are added to the residuals of down unet blocks.
          mid_block_additional_residual: (`torch.Tensor`, *optional*):
              A tensor that if specified is added to the residual of the middle unet block.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`~models.unet_3d_condition.UNet3DConditionOutput`] instead of a plain
              tuple.
          cross_attention_kwargs (`dict`, *optional*):
              A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
  
      Returns:
          [`~models.unet_3d_condition.UNet3DConditionOutput`] or `tuple`:
              If `return_dict` is True, an [`~models.unet_3d_condition.UNet3DConditionOutput`] is returned, otherwise
              a `tuple` is returned where the first element is the sample tensor.
      """
      # 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)
  
      # 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
      num_frames = sample.shape[2]
      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)
      emb = emb.repeat_interleave(repeats=num_frames, dim=0)
      encoder_hidden_states = encoder_hidden_states.repeat_interleave(
          repeats=num_frames, dim=0
      )
  
      # 2. pre-process
      sample = sample.permute(0, 2, 1, 3, 4).reshape(
          (sample.shape[0] * num_frames, -1) + sample.shape[3:]
      )
      sample = self.conv_in(sample)
  
      sample = self.transformer_in(
          sample,
          num_frames=num_frames,
          cross_attention_kwargs=cross_attention_kwargs,
          return_dict=False,
      )[0]
  
      # 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,
                  num_frames=num_frames,
                  cross_attention_kwargs=cross_attention_kwargs,
              )
          else:
              sample, res_samples = downsample_block(
                  hidden_states=sample, temb=emb, num_frames=num_frames
              )
  
          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
      if self.mid_block is not None:
          sample = self.mid_block(
              sample,
              emb,
              encoder_hidden_states=encoder_hidden_states,
              attention_mask=attention_mask,
              num_frames=num_frames,
              cross_attention_kwargs=cross_attention_kwargs,
          )
  
      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,
                  num_frames=num_frames,
                  cross_attention_kwargs=cross_attention_kwargs,
              )
          else:
              sample = upsample_block(
                  hidden_states=sample,
                  temb=emb,
                  res_hidden_states_tuple=res_samples,
                  upsample_size=upsample_size,
                  num_frames=num_frames,
              )
  
      # 6. post-process
      if self.conv_norm_out:
          sample = self.conv_norm_out(sample)
          sample = self.conv_act(sample)
  
      sample = self.conv_out(sample)
  
      # reshape to (batch, channel, framerate, width, height)
      sample = (
          sample[None, :]
          .reshape((-1, num_frames) + sample.shape[1:])
          .permute(0, 2, 1, 3, 4)
      )
  
      if not return_dict:
          return (sample,)
  
      return UNet3DConditionOutput(sample=sample)
  
imports: from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.utils.checkpoint
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..loaders import UNet2DConditionLoadersMixin
  from ..utils import BaseOutput, logging
  from .attention_processor import ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor
  from .embeddings import TimestepEmbedding, Timesteps
  from .modeling_utils import ModelMixin
  from .transformer_temporal import TransformerTemporalModel
  from .unet_3d_blocks import CrossAttnDownBlock3D, CrossAttnUpBlock3D, DownBlock3D, UNetMidBlock3DCrossAttn, UpBlock3D, get_down_block, get_up_block
label_desc: Use image node_id 11 for calling the UNet3DConditionModel obj's underlying member method code with example usage: obj.forward(sample, timestep, encoder_hidden_states, class_labels, timestep_cond, attention_mask, cross_attention_kwargs, down_block_additional_residuals, mid_block_additional_residual, return_dict) and returns: UNet3DConditionOutput, sample