class_docstr: A conditional Spatio-Temporal UNet model that takes a noisy video frames, 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 8): Number of channels in the input sample.
      out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
      down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
          The tuple of downsample blocks to use.
      up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
          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.
      addition_time_embed_dim: (`int`, defaults to 256):
          Dimension to to encode the additional time ids.
      projection_class_embeddings_input_dim (`int`, defaults to 768):
          The dimension of the projection of encoded `added_time_ids`.
      layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
      cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
          The dimension of the cross attention features.
      transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
          The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
          [`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
          [`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
      num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
          The number of attention heads.
      dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
class_docstr_tok: ["A","conditional","Spatio-Temporal","UNet","model","that","takes","a","noisy","video","frames",",","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","8",")",":","Number","of","channels","in","the","input","sample",".","out_channels","(","`","int","`",",","*","optional","*",",","defaults","to","4",")",":","Number","of","channels","in","the","output",".","down_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","CrossAttnDownBlockSpatioTemporal","''",",","``","CrossAttnDownBlockSpatioTemporal","''",",","``","CrossAttnDownBlockSpatioTemporal","''",",","``","DownBlockSpatioTemporal","''",")","`",")",":","The","tuple","of","downsample","blocks","to","use",".","up_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","UpBlockSpatioTemporal","''",",","``","CrossAttnUpBlockSpatioTemporal","''",",","``","CrossAttnUpBlockSpatioTemporal","''",",","``","CrossAttnUpBlockSpatioTemporal","''",")","`",")",":","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",".","addition_time_embed_dim",":","(","`","int","`",",","defaults","to","256",")",":","Dimension","to","to","encode","the","additional","time","ids",".","projection_class_embeddings_input_dim","(","`","int","`",",","defaults","to","768",")",":","The","dimension","of","the","projection","of","encoded","`","added_time_ids","`",".","layers_per_block","(","`","int","`",",","*","optional","*",",","defaults","to","2",")",":","The","number","of","layers","per","block",".","cross_attention_dim","(","`","int","`","or","`","Tuple","[","int","]","`",",","*","optional","*",",","defaults","to","1280",")",":","The","dimension","of","the","cross","attention","features",".","transformer_layers_per_block","(","`","int","`",",","`","Tuple","[","int","]","`",",","or","`","Tuple","[","Tuple","]","`",",","*","optional","*",",","defaults","to","1",")",":","The","number","of","transformer","blocks","of","type","[","`","~models.attention.BasicTransformerBlock","`","]",".","Only","relevant","for","[","`","~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal","`","]",",","[","`","~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal","`","]",",","[","`","~models.unet_3d_blocks.UNetMidBlockSpatioTemporal","`","]",".","num_attention_heads","(","`","int","`",",","`","Tuple","[","int","]","`",",","defaults","to","`","(","5",",","10",",","10",",","20",")","`",")",":","The","number","of","attention","heads",".","dropout","(","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","The","dropout","probability","to","use","."]
docstr: Sets the attention processor to use [feed forward
  chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
  
  Parameters:
      chunk_size (`int`, *optional*):
          The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
          over each tensor of dim=`dim`.
      dim (`int`, *optional*, defaults to `0`):
          The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
          or dim=1 (sequence length).
docstr_tok: ["Sets","the","attention","processor","to","use","[","feed","forward","chunking","]","(","https",":","\/\/huggingface.co\/blog\/reformer","#","2-chunked-feed-forward-layers",")",".","Parameters",":","chunk_size","(","`","int","`",",","*","optional","*",")",":","The","chunk","size","of","the","feed-forward","layers",".","If","not","specified",",","will","run","feed-forward","layer","individually","over","each","tensor","of","dim=","`","dim","`",".","dim","(","`","int","`",",","*","optional","*",",","defaults","to","`","0","`",")",":","The","dimension","over","which","the","feed-forward","computation","should","be","chunked",".","Choose","between","dim=0","(","batch",")","or","dim=1","(","sequence","length",")","."]
code: def enable_forward_chunking(
      self, chunk_size: Optional[int] = None, dim: int = 0
  ) -> None:
      """
      Sets the attention processor to use [feed forward
      chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
  
      Parameters:
          chunk_size (`int`, *optional*):
              The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
              over each tensor of dim=`dim`.
          dim (`int`, *optional*, defaults to `0`):
              The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
              or dim=1 (sequence length).
      """
      if dim not in [0, 1]:
          raise ValueError(
              f"Make sure to set `dim` to either 0 or 1, not {dim}"
          )
  
      # By default chunk size is 1
      chunk_size = chunk_size or 1
  
      def fn_recursive_feed_forward(
          module: torch.nn.Module, chunk_size: int, dim: int
      ):
          if hasattr(module, "set_chunk_feed_forward"):
              module.set_chunk_feed_forward(
                  chunk_size=chunk_size, dim=dim
              )
  
          for child in module.children():
              fn_recursive_feed_forward(child, chunk_size, dim)
  
      for module in self.children():
          fn_recursive_feed_forward(module, chunk_size, dim)
  
imports: from dataclasses import dataclass
  from typing import Dict, Optional, Tuple, Union
  import torch
  import torch.nn
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..loaders import UNet2DConditionLoadersMixin
  from ..utils import BaseOutput, logging
  from .attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
  from .embeddings import TimestepEmbedding, Timesteps
  from .modeling_utils import ModelMixin
  from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block
label_desc: Use image node_id 6 for calling the UNetSpatioTemporalConditionModel obj's underlying member method code with example usage: obj.enable_forward_chunking(chunk_size, dim) without return types