class_docstr: A conditional 2D 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): Number of channels in the input sample.
      out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
      center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
      flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
          Whether to flip the sin to cos in the time embedding.
      freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
      down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "CrossAttnDownBlockFlat", "DownBlockFlat")`):
          The tuple of downsample blocks to use.
      mid_block_type (`str`, *optional*, defaults to `"UNetMidBlockFlatCrossAttn"`):
          Block type for middle of UNet, it can be either `UNetMidBlockFlatCrossAttn` or
          `UNetMidBlockFlatSimpleCrossAttn`. If `None`, the mid block layer is skipped.
      up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat", "CrossAttnUpBlockFlat")`):
          The tuple of upsample blocks to use.
      only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
          Whether to include self-attention in the basic transformer blocks, see
          [`~models.attention.BasicTransformerBlock`].
      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.
      dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
      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` or `Tuple[int]`, *optional*, defaults to 1280):
          The dimension of the cross attention features.
      transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
          The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
          [`~models.unet_2d_blocks.CrossAttnDownBlockFlat`], [`~models.unet_2d_blocks.CrossAttnUpBlockFlat`],
          [`~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn`].
      encoder_hid_dim (`int`, *optional*, defaults to None):
          If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
          dimension to `cross_attention_dim`.
      encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
          If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
          embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
      attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
      num_attention_heads (`int`, *optional*):
          The number of attention heads. If not defined, defaults to `attention_head_dim`
      resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
          for ResNet blocks (see [`~models.resnet.ResnetBlockFlat`]). Choose from `default` or `scale_shift`.
      class_embed_type (`str`, *optional*, defaults to `None`):
          The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
          `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
      addition_embed_type (`str`, *optional*, defaults to `None`):
          Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
          "text". "text" will use the `TextTimeEmbedding` layer.
      addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
          Dimension for the timestep embeddings.
      num_class_embeds (`int`, *optional*, defaults to `None`):
          Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
          class conditioning with `class_embed_type` equal to `None`.
      time_embedding_type (`str`, *optional*, defaults to `positional`):
          The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
      time_embedding_dim (`int`, *optional*, defaults to `None`):
          An optional override for the dimension of the projected time embedding.
      time_embedding_act_fn (`str`, *optional*, defaults to `None`):
          Optional activation function to use only once on the time embeddings before they are passed to the rest of
          the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
      timestep_post_act (`str`, *optional*, defaults to `None`):
          The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
      time_cond_proj_dim (`int`, *optional*, defaults to `None`):
          The dimension of `cond_proj` layer in the timestep embedding.
      conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
      conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
      projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when
          `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
      class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
          embeddings with the class embeddings.
      mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
          Whether to use cross attention with the mid block when using the `UNetMidBlockFlatSimpleCrossAttn`. If
          `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
          `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
          otherwise.
class_docstr_tok: ["A","conditional","2D","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",")",":","Number","of","channels","in","the","input","sample",".","out_channels","(","`","int","`",",","*","optional","*",",","defaults","to","4",")",":","Number","of","channels","in","the","output",".","center_input_sample","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","to","center","the","input","sample",".","flip_sin_to_cos","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","to","flip","the","sin","to","cos","in","the","time","embedding",".","freq_shift","(","`","int","`",",","*","optional","*",",","defaults","to","0",")",":","The","frequency","shift","to","apply","to","the","time","embedding",".","down_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","CrossAttnDownBlockFlat","''",",","``","CrossAttnDownBlockFlat","''",",","``","CrossAttnDownBlockFlat","''",",","``","DownBlockFlat","''",")","`",")",":","The","tuple","of","downsample","blocks","to","use",".","mid_block_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","UNetMidBlockFlatCrossAttn","''","`",")",":","Block","type","for","middle","of","UNet",",","it","can","be","either","`","UNetMidBlockFlatCrossAttn","`","or","`","UNetMidBlockFlatSimpleCrossAttn","`",".","If","`","None","`",",","the","mid","block","layer","is","skipped",".","up_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","UpBlockFlat","''",",","``","CrossAttnUpBlockFlat","''",",","``","CrossAttnUpBlockFlat","''",",","``","CrossAttnUpBlockFlat","''",")","`",")",":","The","tuple","of","upsample","blocks","to","use",".","only_cross_attention","(","`","bool","`","or","`","Tuple","[","bool","]","`",",","*","optional","*",",","default","to","`","False","`",")",":","Whether","to","include","self-attention","in","the","basic","transformer","blocks",",","see","[","`","~models.attention.BasicTransformerBlock","`","]",".","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",".","dropout","(","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","The","dropout","probability","to","use",".","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","`","or","`","Tuple","[","int","]","`",",","*","optional","*",",","defaults","to","1280",")",":","The","dimension","of","the","cross","attention","features",".","transformer_layers_per_block","(","`","int","`","or","`","Tuple","[","int","]","`",",","*","optional","*",",","defaults","to","1",")",":","The","number","of","transformer","blocks","of","type","[","`","~models.attention.BasicTransformerBlock","`","]",".","Only","relevant","for","[","`","~models.unet_2d_blocks.CrossAttnDownBlockFlat","`","]",",","[","`","~models.unet_2d_blocks.CrossAttnUpBlockFlat","`","]",",","[","`","~models.unet_2d_blocks.UNetMidBlockFlatCrossAttn","`","]",".","encoder_hid_dim","(","`","int","`",",","*","optional","*",",","defaults","to","None",")",":","If","`","encoder_hid_dim_type","`","is","defined",",","`","encoder_hidden_states","`","will","be","projected","from","`","encoder_hid_dim","`","dimension","to","`","cross_attention_dim","`",".","encoder_hid_dim_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","If","given",",","the","`","encoder_hidden_states","`","and","potentially","other","embeddings","are","down-projected","to","text","embeddings","of","dimension","`","cross_attention","`","according","to","`","encoder_hid_dim_type","`",".","attention_head_dim","(","`","int","`",",","*","optional","*",",","defaults","to","8",")",":","The","dimension","of","the","attention","heads",".","num_attention_heads","(","`","int","`",",","*","optional","*",")",":","The","number","of","attention","heads",".","If","not","defined",",","defaults","to","`","attention_head_dim","`","resnet_time_scale_shift","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","default","''","`",")",":","Time","scale","shift","config","for","ResNet","blocks","(","see","[","`","~models.resnet.ResnetBlockFlat","`","]",")",".","Choose","from","`","default","`","or","`","scale_shift","`",".","class_embed_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","The","type","of","class","embedding","to","use","which","is","ultimately","summed","with","the","time","embeddings",".","Choose","from","`","None","`",",","`","``","timestep","''","`",",","`","``","identity","''","`",",","`","``","projection","''","`",",","or","`","``","simple_projection","''","`",".","addition_embed_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Configures","an","optional","embedding","which","will","be","summed","with","the","time","embeddings",".","Choose","from","`","None","`","or","``","text","''",".","``","text","''","will","use","the","`","TextTimeEmbedding","`","layer",".","addition_time_embed_dim",":","(","`","int","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Dimension","for","the","timestep","embeddings",".","num_class_embeds","(","`","int","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Input","dimension","of","the","learnable","embedding","matrix","to","be","projected","to","`","time_embed_dim","`",",","when","performing","class","conditioning","with","`","class_embed_type","`","equal","to","`","None","`",".","time_embedding_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","positional","`",")",":","The","type","of","position","embedding","to","use","for","timesteps",".","Choose","from","`","positional","`","or","`","fourier","`",".","time_embedding_dim","(","`","int","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","An","optional","override","for","the","dimension","of","the","projected","time","embedding",".","time_embedding_act_fn","(","`","str","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Optional","activation","function","to","use","only","once","on","the","time","embeddings","before","they","are","passed","to","the","rest","of","the","UNet",".","Choose","from","`","silu","`",",","`","mish","`",",","`","gelu","`",",","and","`","swish","`",".","timestep_post_act","(","`","str","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","The","second","activation","function","to","use","in","timestep","embedding",".","Choose","from","`","silu","`",",","`","mish","`","and","`","gelu","`",".","time_cond_proj_dim","(","`","int","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","The","dimension","of","`","cond_proj","`","layer","in","the","timestep","embedding",".","conv_in_kernel","(","`","int","`",",","*","optional","*",",","default","to","`","3","`",")",":","The","kernel","size","of","`","conv_in","`","layer",".","conv_out_kernel","(","`","int","`",",","*","optional","*",",","default","to","`","3","`",")",":","The","kernel","size","of","`","conv_out","`","layer",".","projection_class_embeddings_input_dim","(","`","int","`",",","*","optional","*",")",":","The","dimension","of","the","`","class_labels","`","input","when","`","class_embed_type=","''","projection","''","`",".","Required","when","`","class_embed_type=","''","projection","''","`",".","class_embeddings_concat","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","to","concatenate","the","time","embeddings","with","the","class","embeddings",".","mid_block_only_cross_attention","(","`","bool","`",",","*","optional","*",",","defaults","to","`","None","`",")",":","Whether","to","use","cross","attention","with","the","mid","block","when","using","the","`","UNetMidBlockFlatSimpleCrossAttn","`",".","If","`","only_cross_attention","`","is","given","as","a","single","boolean","and","`","mid_block_only_cross_attention","`","is","`","None","`",",","the","`","only_cross_attention","`","value","is","used","as","the","value","for","`","mid_block_only_cross_attention","`",".","Default","to","`","False","`","otherwise","."]
code: def __init__(
      self,
      sample_size: Optional[int] = None,
      in_channels: int = 4,
      out_channels: int = 4,
      center_input_sample: bool = False,
      flip_sin_to_cos: bool = True,
      freq_shift: int = 0,
      down_block_types: Tuple[str] = (
          "CrossAttnDownBlockFlat",
          "CrossAttnDownBlockFlat",
          "CrossAttnDownBlockFlat",
          "DownBlockFlat",
      ),
      mid_block_type: Optional[str] = "UNetMidBlockFlatCrossAttn",
      up_block_types: Tuple[str] = (
          "UpBlockFlat",
          "CrossAttnUpBlockFlat",
          "CrossAttnUpBlockFlat",
          "CrossAttnUpBlockFlat",
      ),
      only_cross_attention: Union[bool, Tuple[bool]] = False,
      block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
      layers_per_block: Union[int, Tuple[int]] = 2,
      downsample_padding: int = 1,
      mid_block_scale_factor: float = 1,
      dropout: float = 0.0,
      act_fn: str = "silu",
      norm_num_groups: Optional[int] = 32,
      norm_eps: float = 1e-5,
      cross_attention_dim: Union[int, Tuple[int]] = 1280,
      transformer_layers_per_block: Union[int, Tuple[int]] = 1,
      encoder_hid_dim: Optional[int] = None,
      encoder_hid_dim_type: Optional[str] = None,
      attention_head_dim: Union[int, Tuple[int]] = 8,
      num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
      dual_cross_attention: bool = False,
      use_linear_projection: bool = False,
      class_embed_type: Optional[str] = None,
      addition_embed_type: Optional[str] = None,
      addition_time_embed_dim: Optional[int] = None,
      num_class_embeds: Optional[int] = None,
      upcast_attention: bool = False,
      resnet_time_scale_shift: str = "default",
      resnet_skip_time_act: bool = False,
      resnet_out_scale_factor: int = 1.0,
      time_embedding_type: str = "positional",
      time_embedding_dim: Optional[int] = None,
      time_embedding_act_fn: Optional[str] = None,
      timestep_post_act: Optional[str] = None,
      time_cond_proj_dim: Optional[int] = None,
      conv_in_kernel: int = 3,
      conv_out_kernel: int = 3,
      projection_class_embeddings_input_dim: Optional[int] = None,
      attention_type: str = "default",
      class_embeddings_concat: bool = False,
      mid_block_only_cross_attention: Optional[bool] = None,
      cross_attention_norm: Optional[str] = None,
      addition_embed_type_num_heads=64,
  ):
      super().__init__()
  
      self.sample_size = sample_size
  
      if num_attention_heads is not None:
          raise ValueError(
              "At the moment it is not possible to define the number of attention heads via `num_attention_heads`"
              " because of a naming issue as described in"
              " https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing"
              " `num_attention_heads` will only be supported in diffusers v0.19."
          )
  
      # If `num_attention_heads` is not defined (which is the case for most models)
      # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
      # The reason for this behavior is to correct for incorrectly named variables that were introduced
      # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
      # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
      # which is why we correct for the naming here.
      num_attention_heads = num_attention_heads or attention_head_dim
  
      # Check inputs
      if len(down_block_types) != len(up_block_types):
          raise ValueError(
              "Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`:"
              f" {down_block_types}. `up_block_types`: {up_block_types}."
          )
  
      if len(block_out_channels) != len(down_block_types):
          raise ValueError(
              "Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`:"
              f" {block_out_channels}. `down_block_types`: {down_block_types}."
          )
  
      if not isinstance(only_cross_attention, bool) and len(
          only_cross_attention
      ) != len(down_block_types):
          raise ValueError(
              "Must provide the same number of `only_cross_attention` as `down_block_types`."
              f" `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
          )
  
      if not isinstance(num_attention_heads, int) and len(
          num_attention_heads
      ) != len(down_block_types):
          raise ValueError(
              "Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`:"
              f" {num_attention_heads}. `down_block_types`: {down_block_types}."
          )
  
      if not isinstance(attention_head_dim, int) and len(
          attention_head_dim
      ) != len(down_block_types):
          raise ValueError(
              "Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`:"
              f" {attention_head_dim}. `down_block_types`: {down_block_types}."
          )
  
      if isinstance(cross_attention_dim, list) and len(
          cross_attention_dim
      ) != len(down_block_types):
          raise ValueError(
              "Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`:"
              f" {cross_attention_dim}. `down_block_types`: {down_block_types}."
          )
  
      if not isinstance(layers_per_block, int) and len(
          layers_per_block
      ) != len(down_block_types):
          raise ValueError(
              "Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`:"
              f" {layers_per_block}. `down_block_types`: {down_block_types}."
          )
  
      # input
      conv_in_padding = (conv_in_kernel - 1) // 2
      self.conv_in = LinearMultiDim(
          in_channels,
          block_out_channels[0],
          kernel_size=conv_in_kernel,
          padding=conv_in_padding,
      )
  
      # time
      if time_embedding_type == "fourier":
          time_embed_dim = (
              time_embedding_dim or block_out_channels[0] * 2
          )
          if time_embed_dim % 2 != 0:
              raise ValueError(
                  f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}."
              )
          self.time_proj = GaussianFourierProjection(
              time_embed_dim // 2,
              set_W_to_weight=False,
              log=False,
              flip_sin_to_cos=flip_sin_to_cos,
          )
          timestep_input_dim = time_embed_dim
      elif time_embedding_type == "positional":
          time_embed_dim = (
              time_embedding_dim or block_out_channels[0] * 4
          )
  
          self.time_proj = Timesteps(
              block_out_channels[0], flip_sin_to_cos, freq_shift
          )
          timestep_input_dim = block_out_channels[0]
      else:
          raise ValueError(
              f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
          )
  
      self.time_embedding = TimestepEmbedding(
          timestep_input_dim,
          time_embed_dim,
          act_fn=act_fn,
          post_act_fn=timestep_post_act,
          cond_proj_dim=time_cond_proj_dim,
      )
  
      if encoder_hid_dim_type is None and encoder_hid_dim is not None:
          encoder_hid_dim_type = "text_proj"
          self.register_to_config(
              encoder_hid_dim_type=encoder_hid_dim_type
          )
          logger.info(
              "encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined."
          )
  
      if encoder_hid_dim is None and encoder_hid_dim_type is not None:
          raise ValueError(
              f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
          )
  
      if encoder_hid_dim_type == "text_proj":
          self.encoder_hid_proj = nn.Linear(
              encoder_hid_dim, cross_attention_dim
          )
      elif encoder_hid_dim_type == "text_image_proj":
          # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
          # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
          # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
          self.encoder_hid_proj = TextImageProjection(
              text_embed_dim=encoder_hid_dim,
              image_embed_dim=cross_attention_dim,
              cross_attention_dim=cross_attention_dim,
          )
      elif encoder_hid_dim_type == "image_proj":
          # Kandinsky 2.2
          self.encoder_hid_proj = ImageProjection(
              image_embed_dim=encoder_hid_dim,
              cross_attention_dim=cross_attention_dim,
          )
      elif encoder_hid_dim_type is not None:
          raise ValueError(
              f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
          )
      else:
          self.encoder_hid_proj = None
  
      # class embedding
      if class_embed_type is None and num_class_embeds is not None:
          self.class_embedding = nn.Embedding(
              num_class_embeds, time_embed_dim
          )
      elif class_embed_type == "timestep":
          self.class_embedding = TimestepEmbedding(
              timestep_input_dim, time_embed_dim, act_fn=act_fn
          )
      elif class_embed_type == "identity":
          self.class_embedding = nn.Identity(
              time_embed_dim, time_embed_dim
          )
      elif class_embed_type == "projection":
          if projection_class_embeddings_input_dim is None:
              raise ValueError(
                  "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
              )
          # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
          # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
          # 2. it projects from an arbitrary input dimension.
          #
          # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
          # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
          # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
          self.class_embedding = TimestepEmbedding(
              projection_class_embeddings_input_dim, time_embed_dim
          )
      elif class_embed_type == "simple_projection":
          if projection_class_embeddings_input_dim is None:
              raise ValueError(
                  "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
              )
          self.class_embedding = nn.Linear(
              projection_class_embeddings_input_dim, time_embed_dim
          )
      else:
          self.class_embedding = None
  
      if addition_embed_type == "text":
          if encoder_hid_dim is not None:
              text_time_embedding_from_dim = encoder_hid_dim
          else:
              text_time_embedding_from_dim = cross_attention_dim
  
          self.add_embedding = TextTimeEmbedding(
              text_time_embedding_from_dim,
              time_embed_dim,
              num_heads=addition_embed_type_num_heads,
          )
      elif addition_embed_type == "text_image":
          # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
          # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
          # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
          self.add_embedding = TextImageTimeEmbedding(
              text_embed_dim=cross_attention_dim,
              image_embed_dim=cross_attention_dim,
              time_embed_dim=time_embed_dim,
          )
      elif addition_embed_type == "text_time":
          self.add_time_proj = Timesteps(
              addition_time_embed_dim, flip_sin_to_cos, freq_shift
          )
          self.add_embedding = TimestepEmbedding(
              projection_class_embeddings_input_dim, time_embed_dim
          )
      elif addition_embed_type == "image":
          # Kandinsky 2.2
          self.add_embedding = ImageTimeEmbedding(
              image_embed_dim=encoder_hid_dim,
              time_embed_dim=time_embed_dim,
          )
      elif addition_embed_type == "image_hint":
          # Kandinsky 2.2 ControlNet
          self.add_embedding = ImageHintTimeEmbedding(
              image_embed_dim=encoder_hid_dim,
              time_embed_dim=time_embed_dim,
          )
      elif addition_embed_type is not None:
          raise ValueError(
              f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'."
          )
  
      if time_embedding_act_fn is None:
          self.time_embed_act = None
      else:
          self.time_embed_act = get_activation(time_embedding_act_fn)
  
      self.down_blocks = nn.ModuleList([])
      self.up_blocks = nn.ModuleList([])
  
      if isinstance(only_cross_attention, bool):
          if mid_block_only_cross_attention is None:
              mid_block_only_cross_attention = only_cross_attention
  
          only_cross_attention = [only_cross_attention] * len(
              down_block_types
          )
  
      if mid_block_only_cross_attention is None:
          mid_block_only_cross_attention = False
  
      if isinstance(num_attention_heads, int):
          num_attention_heads = (num_attention_heads,) * len(
              down_block_types
          )
  
      if isinstance(attention_head_dim, int):
          attention_head_dim = (attention_head_dim,) * len(
              down_block_types
          )
  
      if isinstance(cross_attention_dim, int):
          cross_attention_dim = (cross_attention_dim,) * len(
              down_block_types
          )
  
      if isinstance(layers_per_block, int):
          layers_per_block = [layers_per_block] * len(down_block_types)
  
      if isinstance(transformer_layers_per_block, int):
          transformer_layers_per_block = [
              transformer_layers_per_block
          ] * len(down_block_types)
  
      if class_embeddings_concat:
          # The time embeddings are concatenated with the class embeddings. The dimension of the
          # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
          # regular time embeddings
          blocks_time_embed_dim = time_embed_dim * 2
      else:
          blocks_time_embed_dim = time_embed_dim
  
      # down
      output_channel = block_out_channels[0]
      for i, down_block_type in enumerate(down_block_types):
          input_channel = output_channel
          output_channel = block_out_channels[i]
          is_final_block = i == len(block_out_channels) - 1
  
          down_block = get_down_block(
              down_block_type,
              num_layers=layers_per_block[i],
              transformer_layers_per_block=transformer_layers_per_block[
                  i
              ],
              in_channels=input_channel,
              out_channels=output_channel,
              temb_channels=blocks_time_embed_dim,
              add_downsample=not is_final_block,
              resnet_eps=norm_eps,
              resnet_act_fn=act_fn,
              resnet_groups=norm_num_groups,
              cross_attention_dim=cross_attention_dim[i],
              num_attention_heads=num_attention_heads[i],
              downsample_padding=downsample_padding,
              dual_cross_attention=dual_cross_attention,
              use_linear_projection=use_linear_projection,
              only_cross_attention=only_cross_attention[i],
              upcast_attention=upcast_attention,
              resnet_time_scale_shift=resnet_time_scale_shift,
              attention_type=attention_type,
              resnet_skip_time_act=resnet_skip_time_act,
              resnet_out_scale_factor=resnet_out_scale_factor,
              cross_attention_norm=cross_attention_norm,
              attention_head_dim=attention_head_dim[i]
              if attention_head_dim[i] is not None
              else output_channel,
              dropout=dropout,
          )
          self.down_blocks.append(down_block)
  
      # mid
      if mid_block_type == "UNetMidBlockFlatCrossAttn":
          self.mid_block = UNetMidBlockFlatCrossAttn(
              transformer_layers_per_block=transformer_layers_per_block[
                  -1
              ],
              in_channels=block_out_channels[-1],
              temb_channels=blocks_time_embed_dim,
              dropout=dropout,
              resnet_eps=norm_eps,
              resnet_act_fn=act_fn,
              output_scale_factor=mid_block_scale_factor,
              resnet_time_scale_shift=resnet_time_scale_shift,
              cross_attention_dim=cross_attention_dim[-1],
              num_attention_heads=num_attention_heads[-1],
              resnet_groups=norm_num_groups,
              dual_cross_attention=dual_cross_attention,
              use_linear_projection=use_linear_projection,
              upcast_attention=upcast_attention,
              attention_type=attention_type,
          )
      elif mid_block_type == "UNetMidBlockFlatSimpleCrossAttn":
          self.mid_block = UNetMidBlockFlatSimpleCrossAttn(
              in_channels=block_out_channels[-1],
              temb_channels=blocks_time_embed_dim,
              dropout=dropout,
              resnet_eps=norm_eps,
              resnet_act_fn=act_fn,
              output_scale_factor=mid_block_scale_factor,
              cross_attention_dim=cross_attention_dim[-1],
              attention_head_dim=attention_head_dim[-1],
              resnet_groups=norm_num_groups,
              resnet_time_scale_shift=resnet_time_scale_shift,
              skip_time_act=resnet_skip_time_act,
              only_cross_attention=mid_block_only_cross_attention,
              cross_attention_norm=cross_attention_norm,
          )
      elif mid_block_type is None:
          self.mid_block = None
      else:
          raise ValueError(f"unknown mid_block_type : {mid_block_type}")
  
      # count how many layers upsample the images
      self.num_upsamplers = 0
  
      # up
      reversed_block_out_channels = list(reversed(block_out_channels))
      reversed_num_attention_heads = list(reversed(num_attention_heads))
      reversed_layers_per_block = list(reversed(layers_per_block))
      reversed_cross_attention_dim = list(reversed(cross_attention_dim))
      reversed_transformer_layers_per_block = list(
          reversed(transformer_layers_per_block)
      )
      only_cross_attention = list(reversed(only_cross_attention))
  
      output_channel = reversed_block_out_channels[0]
      for i, up_block_type in enumerate(up_block_types):
          is_final_block = i == len(block_out_channels) - 1
  
          prev_output_channel = output_channel
          output_channel = reversed_block_out_channels[i]
          input_channel = reversed_block_out_channels[
              min(i + 1, len(block_out_channels) - 1)
          ]
  
          # add upsample block for all BUT final layer
          if not is_final_block:
              add_upsample = True
              self.num_upsamplers += 1
          else:
              add_upsample = False
  
          up_block = get_up_block(
              up_block_type,
              num_layers=reversed_layers_per_block[i] + 1,
              transformer_layers_per_block=reversed_transformer_layers_per_block[
                  i
              ],
              in_channels=input_channel,
              out_channels=output_channel,
              prev_output_channel=prev_output_channel,
              temb_channels=blocks_time_embed_dim,
              add_upsample=add_upsample,
              resnet_eps=norm_eps,
              resnet_act_fn=act_fn,
              resolution_idx=i,
              resnet_groups=norm_num_groups,
              cross_attention_dim=reversed_cross_attention_dim[i],
              num_attention_heads=reversed_num_attention_heads[i],
              dual_cross_attention=dual_cross_attention,
              use_linear_projection=use_linear_projection,
              only_cross_attention=only_cross_attention[i],
              upcast_attention=upcast_attention,
              resnet_time_scale_shift=resnet_time_scale_shift,
              attention_type=attention_type,
              resnet_skip_time_act=resnet_skip_time_act,
              resnet_out_scale_factor=resnet_out_scale_factor,
              cross_attention_norm=cross_attention_norm,
              attention_head_dim=attention_head_dim[i]
              if attention_head_dim[i] is not None
              else output_channel,
              dropout=dropout,
          )
          self.up_blocks.append(up_block)
          prev_output_channel = output_channel
  
      # out
      if norm_num_groups is not None:
          self.conv_norm_out = nn.GroupNorm(
              num_channels=block_out_channels[0],
              num_groups=norm_num_groups,
              eps=norm_eps,
          )
  
          self.conv_act = get_activation(act_fn)
  
      else:
          self.conv_norm_out = None
          self.conv_act = None
  
      conv_out_padding = (conv_out_kernel - 1) // 2
      self.conv_out = LinearMultiDim(
          block_out_channels[0],
          out_channels,
          kernel_size=conv_out_kernel,
          padding=conv_out_padding,
      )
  
      if attention_type in ["gated", "gated-text-image"]:
          positive_len = 768
          if isinstance(cross_attention_dim, int):
              positive_len = cross_attention_dim
          elif isinstance(cross_attention_dim, tuple) or isinstance(
              cross_attention_dim, list
          ):
              positive_len = cross_attention_dim[0]
  
          feature_type = (
              "text-only" if attention_type == "gated" else "text-image"
          )
          self.position_net = PositionNet(
              positive_len=positive_len,
              out_dim=cross_attention_dim,
              feature_type=feature_type,
          )
  
imports: from typing import Any, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...models import ModelMixin
  from ...models.activations import get_activation
  from ...models.attention import Attention
  from ...models.attention_processor import ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnAddedKVProcessor2_0, AttnProcessor
  from ...models.dual_transformer_2d import DualTransformer2DModel
  from ...models.embeddings import GaussianFourierProjection, ImageHintTimeEmbedding, ImageProjection, ImageTimeEmbedding, TextImageProjection, TextImageTimeEmbedding, TextTimeEmbedding, TimestepEmbedding, Timesteps
  from ...models.transformer_2d import Transformer2DModel
  from ...models.unet_2d_condition import UNet2DConditionOutput
  from ...utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
  from ...utils.torch_utils import apply_freeu
label_desc: Use image node_id 1 to create a new UNetFlatConditionModel object from inherited base classes: ModelMixin, ConfigMixin with example: obj = UNetFlatConditionModel(sample_size, in_channels, out_channels, center_input_sample, flip_sin_to_cos, freq_shift, down_block_types, mid_block_type, up_block_types, only_cross_attention, block_out_channels, layers_per_block, downsample_padding, mid_block_scale_factor, dropout, act_fn, norm_num_groups, norm_eps, cross_attention_dim, transformer_layers_per_block, encoder_hid_dim, encoder_hid_dim_type, attention_head_dim, num_attention_heads, dual_cross_attention, use_linear_projection, class_embed_type, addition_embed_type, addition_time_embed_dim, num_class_embeds, upcast_attention, resnet_time_scale_shift, resnet_skip_time_act, resnet_out_scale_factor, time_embedding_type, time_embedding_dim, time_embedding_act_fn, timestep_post_act, time_cond_proj_dim, conv_in_kernel, conv_out_kernel, projection_class_embeddings_input_dim, attention_type, class_embeddings_concat, mid_block_only_cross_attention, cross_attention_norm, addition_embed_type_num_heads)