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: Enable sliced attention computation.
  
  When this option is enabled, the attention module splits the input tensor in slices to compute attention in
  several steps. This is useful for saving some memory in exchange for a small decrease in speed.
  
  Args:
      slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
          When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
          `"max"`, maximum amount of memory is 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`.
docstr_tok: ["Enable","sliced","attention","computation",".","When","this","option","is","enabled",",","the","attention","module","splits","the","input","tensor","in","slices","to","compute","attention","in","several","steps",".","This","is","useful","for","saving","some","memory","in","exchange","for","a","small","decrease","in","speed",".","Args",":","slice_size","(","`","str","`","or","`","int","`","or","`","list","(","int",")","`",",","*","optional","*",",","defaults","to","`","``","auto","''","`",")",":","When","`","``","auto","''","`",",","input","to","the","attention","heads","is","halved",",","so","attention","is","computed","in","two","steps",".","If","`","``","max","''","`",",","maximum","amount","of","memory","is","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","`","."]
code: def set_attention_slice(self, slice_size):
      r"""
      Enable sliced attention computation.
  
      When this option is enabled, the attention module splits the input tensor in slices to compute attention in
      several steps. This is useful for saving some memory in exchange for a small decrease in speed.
  
      Args:
          slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
              When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
              `"max"`, maximum amount of memory is 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_sliceable_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_sliceable_dims(child)
  
      # retrieve number of attention layers
      for module in self.children():
          fn_recursive_retrieve_sliceable_dims(module)
  
      num_sliceable_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_sliceable_layers * [1]
  
      slice_size = (
          num_sliceable_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)
  
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 3 for calling the UNet3DConditionModel obj's underlying member method code with example usage: obj.set_attention_slice(slice_size) without return types