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","."]
docstr: The [`UNetFlatConditionModel`] forward method.
  
  Args:
      sample (`torch.FloatTensor`):
          The noisy input tensor with the following shape `(batch, 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).
      added_cond_kwargs: (`dict`, *optional*):
          A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
          are passed along to the UNet blocks.
      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.
      encoder_attention_mask (`torch.Tensor`):
          A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
          `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
          which adds large negative values to the attention scores corresponding to "discard" tokens.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
          tuple.
      cross_attention_kwargs (`dict`, *optional*):
          A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
      added_cond_kwargs: (`dict`, *optional*):
          A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
          are passed along to the UNet blocks.
  
  Returns:
      [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
          If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise
          a `tuple` is returned where the first element is the sample tensor.
docstr_tok: ["The","[","`","UNetFlatConditionModel","`","]","forward","method",".","Args",":","sample","(","`","torch.FloatTensor","`",")",":","The","noisy","input","tensor","with","the","following","shape","`","(","batch",",","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",")",".","added_cond_kwargs",":","(","`","dict","`",",","*","optional","*",")",":","A","kwargs","dictionary","containing","additional","embeddings","that","if","specified","are","added","to","the","embeddings","that","are","passed","along","to","the","UNet","blocks",".","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",".","encoder_attention_mask","(","`","torch.Tensor","`",")",":","A","cross-attention","mask","of","shape","`","(","batch",",","sequence_length",")","`","is","applied","to","`","encoder_hidden_states","`",".","If","`","True","`","the","mask","is","kept",",","otherwise","if","`","False","`","it","is","discarded",".","Mask","will","be","converted","into","a","bias",",","which","adds","large","negative","values","to","the","attention","scores","corresponding","to","``","discard","''","tokens",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~models.unet_2d_condition.UNet2DConditionOutput","`","]","instead","of","a","plain","tuple",".","cross_attention_kwargs","(","`","dict","`",",","*","optional","*",")",":","A","kwargs","dictionary","that","if","specified","is","passed","along","to","the","[","`","AttnProcessor","`","]",".","added_cond_kwargs",":","(","`","dict","`",",","*","optional","*",")",":","A","kwargs","dictionary","containin","additional","embeddings","that","if","specified","are","added","to","the","embeddings","that","are","passed","along","to","the","UNet","blocks",".","Returns",":","[","`","~models.unet_2d_condition.UNet2DConditionOutput","`","]","or","`","tuple","`",":","If","`","return_dict","`","is","True",",","an","[","`","~models.unet_2d_condition.UNet2DConditionOutput","`","]","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,
      added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
      down_block_additional_residuals: Optional[
          Tuple[torch.Tensor]
      ] = None,
      mid_block_additional_residual: Optional[torch.Tensor] = None,
      encoder_attention_mask: Optional[torch.Tensor] = None,
      return_dict: bool = True,
  ) -> Union[UNet2DConditionOutput, Tuple]:
      r"""
      The [`UNetFlatConditionModel`] forward method.
  
      Args:
          sample (`torch.FloatTensor`):
              The noisy input tensor with the following shape `(batch, 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).
          added_cond_kwargs: (`dict`, *optional*):
              A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
              are passed along to the UNet blocks.
          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.
          encoder_attention_mask (`torch.Tensor`):
              A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
              `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
              which adds large negative values to the attention scores corresponding to "discard" tokens.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
              tuple.
          cross_attention_kwargs (`dict`, *optional*):
              A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
          added_cond_kwargs: (`dict`, *optional*):
              A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
              are passed along to the UNet blocks.
  
      Returns:
          [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
              If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] 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 layers).
      # 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:]
      ):
          # Forward upsample size to force interpolation output size.
          forward_upsample_size = True
  
      # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
      # expects mask of shape:
      #   [batch, key_tokens]
      # adds singleton query_tokens dimension:
      #   [batch,                    1, key_tokens]
      # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
      #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
      #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
      if attention_mask is not None:
          # assume that mask is expressed as:
          #   (1 = keep,      0 = discard)
          # convert mask into a bias that can be added to attention scores:
          #       (keep = +0,     discard = -10000.0)
          attention_mask = (
              1 - attention_mask.to(sample.dtype)
          ) * -10000.0
          attention_mask = attention_mask.unsqueeze(1)
  
      # convert encoder_attention_mask to a bias the same way we do for attention_mask
      if encoder_attention_mask is not None:
          encoder_attention_mask = (
              1 - encoder_attention_mask.to(sample.dtype)
          ) * -10000.0
          encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
  
      # 0. center input if necessary
      if self.config.center_input_sample:
          sample = 2 * sample - 1.0
  
      # 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
      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=sample.dtype)
  
      emb = self.time_embedding(t_emb, timestep_cond)
      aug_emb = None
  
      if self.class_embedding is not None:
          if class_labels is None:
              raise ValueError(
                  "class_labels should be provided when num_class_embeds > 0"
              )
  
          if self.config.class_embed_type == "timestep":
              class_labels = self.time_proj(class_labels)
  
              # `Timesteps` does not contain any weights and will always return f32 tensors
              # there might be better ways to encapsulate this.
              class_labels = class_labels.to(dtype=sample.dtype)
  
          class_emb = self.class_embedding(class_labels).to(
              dtype=sample.dtype
          )
  
          if self.config.class_embeddings_concat:
              emb = torch.cat([emb, class_emb], dim=-1)
          else:
              emb = emb + class_emb
  
      if self.config.addition_embed_type == "text":
          aug_emb = self.add_embedding(encoder_hidden_states)
      elif self.config.addition_embed_type == "text_image":
          # Kandinsky 2.1 - style
          if "image_embeds" not in added_cond_kwargs:
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires"
                  " the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
              )
  
          image_embs = added_cond_kwargs.get("image_embeds")
          text_embs = added_cond_kwargs.get(
              "text_embeds", encoder_hidden_states
          )
          aug_emb = self.add_embedding(text_embs, image_embs)
      elif self.config.addition_embed_type == "text_time":
          # SDXL - style
          if "text_embeds" not in added_cond_kwargs:
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires"
                  " the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
              )
          text_embeds = added_cond_kwargs.get("text_embeds")
          if "time_ids" not in added_cond_kwargs:
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires"
                  " the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
              )
          time_ids = added_cond_kwargs.get("time_ids")
          time_embeds = self.add_time_proj(time_ids.flatten())
          time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
          add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
          add_embeds = add_embeds.to(emb.dtype)
          aug_emb = self.add_embedding(add_embeds)
      elif self.config.addition_embed_type == "image":
          # Kandinsky 2.2 - style
          if "image_embeds" not in added_cond_kwargs:
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the"
                  " keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
              )
          image_embs = added_cond_kwargs.get("image_embeds")
          aug_emb = self.add_embedding(image_embs)
      elif self.config.addition_embed_type == "image_hint":
          # Kandinsky 2.2 - style
          if (
              "image_embeds" not in added_cond_kwargs
              or "hint" not in added_cond_kwargs
          ):
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires"
                  " the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
              )
          image_embs = added_cond_kwargs.get("image_embeds")
          hint = added_cond_kwargs.get("hint")
          aug_emb, hint = self.add_embedding(image_embs, hint)
          sample = torch.cat([sample, hint], dim=1)
  
      emb = emb + aug_emb if aug_emb is not None else emb
  
      if self.time_embed_act is not None:
          emb = self.time_embed_act(emb)
  
      if (
          self.encoder_hid_proj is not None
          and self.config.encoder_hid_dim_type == "text_proj"
      ):
          encoder_hidden_states = self.encoder_hid_proj(
              encoder_hidden_states
          )
      elif (
          self.encoder_hid_proj is not None
          and self.config.encoder_hid_dim_type == "text_image_proj"
      ):
          # Kadinsky 2.1 - style
          if "image_embeds" not in added_cond_kwargs:
              raise ValueError(
                  f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which"
                  " requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
              )
  
          image_embeds = added_cond_kwargs.get("image_embeds")
          encoder_hidden_states = self.encoder_hid_proj(
              encoder_hidden_states, image_embeds
          )
      elif (
          self.encoder_hid_proj is not None
          and self.config.encoder_hid_dim_type == "image_proj"
      ):
          # Kandinsky 2.2 - style
          if "image_embeds" not in added_cond_kwargs:
              raise ValueError(
                  f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires"
                  " the keyword argument `image_embeds` to be passed in  `added_conditions`"
              )
          image_embeds = added_cond_kwargs.get("image_embeds")
          encoder_hidden_states = self.encoder_hid_proj(image_embeds)
      # 2. pre-process
      sample = self.conv_in(sample)
  
      # 2.5 GLIGEN position net
      if (
          cross_attention_kwargs is not None
          and cross_attention_kwargs.get("gligen", None) is not None
      ):
          cross_attention_kwargs = cross_attention_kwargs.copy()
          gligen_args = cross_attention_kwargs.pop("gligen")
          cross_attention_kwargs["gligen"] = {
              "objs": self.position_net(**gligen_args)
          }
  
      # 3. down
      lora_scale = (
          cross_attention_kwargs.get("scale", 1.0)
          if cross_attention_kwargs is not None
          else 1.0
      )
      if USE_PEFT_BACKEND:
          # weight the lora layers by setting `lora_scale` for each PEFT layer
          scale_lora_layers(self, lora_scale)
  
      is_controlnet = (
          mid_block_additional_residual is not None
          and down_block_additional_residuals is not None
      )
      is_adapter = (
          mid_block_additional_residual is None
          and down_block_additional_residuals is not None
      )
  
      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
          ):
              # For t2i-adapter CrossAttnDownBlockFlat
              additional_residuals = {}
              if (
                  is_adapter
                  and len(down_block_additional_residuals) > 0
              ):
                  additional_residuals[
                      "additional_residuals"
                  ] = down_block_additional_residuals.pop(0)
  
              sample, res_samples = downsample_block(
                  hidden_states=sample,
                  temb=emb,
                  encoder_hidden_states=encoder_hidden_states,
                  attention_mask=attention_mask,
                  cross_attention_kwargs=cross_attention_kwargs,
                  encoder_attention_mask=encoder_attention_mask,
                  **additional_residuals,
              )
          else:
              sample, res_samples = downsample_block(
                  hidden_states=sample, temb=emb, scale=lora_scale
              )
  
              if (
                  is_adapter
                  and len(down_block_additional_residuals) > 0
              ):
                  sample += down_block_additional_residuals.pop(0)
  
          down_block_res_samples += res_samples
  
      if is_controlnet:
          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 = (
                  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,
              cross_attention_kwargs=cross_attention_kwargs,
              encoder_attention_mask=encoder_attention_mask,
          )
          # To support T2I-Adapter-XL
          if (
              is_adapter
              and len(down_block_additional_residuals) > 0
              and sample.shape
              == down_block_additional_residuals[0].shape
          ):
              sample += down_block_additional_residuals.pop(0)
  
      if is_controlnet:
          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,
                  cross_attention_kwargs=cross_attention_kwargs,
                  upsample_size=upsample_size,
                  attention_mask=attention_mask,
                  encoder_attention_mask=encoder_attention_mask,
              )
          else:
              sample = upsample_block(
                  hidden_states=sample,
                  temb=emb,
                  res_hidden_states_tuple=res_samples,
                  upsample_size=upsample_size,
                  scale=lora_scale,
              )
  
      # 6. post-process
      if self.conv_norm_out:
          sample = self.conv_norm_out(sample)
          sample = self.conv_act(sample)
      sample = self.conv_out(sample)
  
      if USE_PEFT_BACKEND:
          # remove `lora_scale` from each PEFT layer
          unscale_lora_layers(self)
  
      if not return_dict:
          return (sample,)
  
      return UNet2DConditionOutput(sample=sample)
  
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 9 for calling the UNetFlatConditionModel 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, added_cond_kwargs, down_block_additional_residuals, mid_block_additional_residual, encoder_attention_mask, return_dict) and returns: UNet2DConditionOutput, sample