class_docstr: Transformer model based on the [U-ViT](https://github.com/baofff/U-ViT) architecture for image-like data. Compared
  to [`Transformer2DModel`], this model has skip connections between transformer blocks in a "U"-shaped fashion,
  similar to a U-Net. Supports only continuous (actual embeddings) inputs, which are embedded via a [`PatchEmbed`]
  layer and then reshaped to (b, t, d).
  
  Parameters:
      num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
      attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
      in_channels (`int`, *optional*):
          Pass if the input is continuous. The number of channels in the input.
      out_channels (`int`, *optional*):
          The number of output channels; if `None`, defaults to `in_channels`.
      num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
      dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
      norm_num_groups (`int`, *optional*, defaults to `32`):
          The number of groups to use when performing Group Normalization.
      cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use.
      attention_bias (`bool`, *optional*):
          Configure if the TransformerBlocks' attention should contain a bias parameter.
      sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images.
          Note that this is fixed at training time as it is used for learning a number of position embeddings. See
          `ImagePositionalEmbeddings`.
      num_vector_embeds (`int`, *optional*):
          Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels.
          Includes the class for the masked latent pixel.
      patch_size (`int`, *optional*, defaults to 2):
          The patch size to use in the patch embedding.
      activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
      num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`.
          The number of diffusion steps used during training. Note that this is fixed at training time as it is used
          to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for
          up to but not more than steps than `num_embeds_ada_norm`.
      use_linear_projection (int, *optional*): TODO: Not used
      only_cross_attention (`bool`, *optional*):
          Whether to use only cross-attention layers. In this case two cross attention layers are used in each
          transformer block.
      upcast_attention (`bool`, *optional*):
          Whether to upcast the query and key to float() when performing the attention calculation.
      norm_type (`str`, *optional*, defaults to `"layer_norm"`):
          The Layer Normalization implementation to use. Defaults to `torch.nn.LayerNorm`.
      block_type (`str`, *optional*, defaults to `"unidiffuser"`):
          The transformer block implementation to use. If `"unidiffuser"`, has the LayerNorms on the residual
          backbone of each transformer block; otherwise has them in the attention/feedforward branches (the standard
          behavior in `diffusers`.)
      pre_layer_norm (`bool`, *optional*):
          Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"),
          as opposed to after ("post-LayerNorm"). The original UniDiffuser implementation is post-LayerNorm
          (`pre_layer_norm = False`).
      norm_elementwise_affine (`bool`, *optional*):
          Whether to use learnable per-element affine parameters during layer normalization.
      use_patch_pos_embed (`bool`, *optional*):
          Whether to use position embeddings inside the patch embedding layer (`PatchEmbed`).
      final_dropout (`bool`, *optional*):
          Whether to use a final Dropout layer after the feedforward network.
class_docstr_tok: ["Transformer","model","based","on","the","[","U-ViT","]","(","https",":","\/\/github.com\/baofff\/U-ViT",")","architecture","for","image-like","data",".","Compared","to","[","`","Transformer2DModel","`","]",",","this","model","has","skip","connections","between","transformer","blocks","in","a","``","U","''","-shaped","fashion",",","similar","to","a","U-Net",".","Supports","only","continuous","(","actual","embeddings",")","inputs",",","which","are","embedded","via","a","[","`","PatchEmbed","`","]","layer","and","then","reshaped","to","(","b",",","t",",","d",")",".","Parameters",":","num_attention_heads","(","`","int","`",",","*","optional","*",",","defaults","to","16",")",":","The","number","of","heads","to","use","for","multi-head","attention",".","attention_head_dim","(","`","int","`",",","*","optional","*",",","defaults","to","88",")",":","The","number","of","channels","in","each","head",".","in_channels","(","`","int","`",",","*","optional","*",")",":","Pass","if","the","input","is","continuous",".","The","number","of","channels","in","the","input",".","out_channels","(","`","int","`",",","*","optional","*",")",":","The","number","of","output","channels",";","if","`","None","`",",","defaults","to","`","in_channels","`",".","num_layers","(","`","int","`",",","*","optional","*",",","defaults","to","1",")",":","The","number","of","layers","of","Transformer","blocks","to","use",".","dropout","(","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","The","dropout","probability","to","use",".","norm_num_groups","(","`","int","`",",","*","optional","*",",","defaults","to","`","32","`",")",":","The","number","of","groups","to","use","when","performing","Group","Normalization",".","cross_attention_dim","(","`","int","`",",","*","optional","*",")",":","The","number","of","encoder_hidden_states","dimensions","to","use",".","attention_bias","(","`","bool","`",",","*","optional","*",")",":","Configure","if","the","TransformerBlocks","'","attention","should","contain","a","bias","parameter",".","sample_size","(","`","int","`",",","*","optional","*",")",":","Pass","if","the","input","is","discrete",".","The","width","of","the","latent","images",".","Note","that","this","is","fixed","at","training","time","as","it","is","used","for","learning","a","number","of","position","embeddings",".","See","`","ImagePositionalEmbeddings","`",".","num_vector_embeds","(","`","int","`",",","*","optional","*",")",":","Pass","if","the","input","is","discrete",".","The","number","of","classes","of","the","vector","embeddings","of","the","latent","pixels",".","Includes","the","class","for","the","masked","latent","pixel",".","patch_size","(","`","int","`",",","*","optional","*",",","defaults","to","2",")",":","The","patch","size","to","use","in","the","patch","embedding",".","activation_fn","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","geglu","''","`",")",":","Activation","function","to","be","used","in","feed-forward",".","num_embeds_ada_norm","(","`","int","`",",","*","optional","*",")",":","Pass","if","at","least","one","of","the","norm_layers","is","`","AdaLayerNorm","`",".","The","number","of","diffusion","steps","used","during","training",".","Note","that","this","is","fixed","at","training","time","as","it","is","used","to","learn","a","number","of","embeddings","that","are","added","to","the","hidden","states",".","During","inference",",","you","can","denoise","for","up","to","but","not","more","than","steps","than","`","num_embeds_ada_norm","`",".","use_linear_projection","(","int",",","*","optional","*",")",":","TODO",":","Not","used","only_cross_attention","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","only","cross-attention","layers",".","In","this","case","two","cross","attention","layers","are","used","in","each","transformer","block",".","upcast_attention","(","`","bool","`",",","*","optional","*",")",":","Whether","to","upcast","the","query","and","key","to","float","(",")","when","performing","the","attention","calculation",".","norm_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","layer_norm","''","`",")",":","The","Layer","Normalization","implementation","to","use",".","Defaults","to","`","torch.nn.LayerNorm","`",".","block_type","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","unidiffuser","''","`",")",":","The","transformer","block","implementation","to","use",".","If","`","``","unidiffuser","''","`",",","has","the","LayerNorms","on","the","residual","backbone","of","each","transformer","block",";","otherwise","has","them","in","the","attention\/feedforward","branches","(","the","standard","behavior","in","`","diffusers","`",".",")","pre_layer_norm","(","`","bool","`",",","*","optional","*",")",":","Whether","to","perform","layer","normalization","before","the","attention","and","feedforward","operations","(","``","pre-LayerNorm","''",")",",","as","opposed","to","after","(","``","post-LayerNorm","''",")",".","The","original","UniDiffuser","implementation","is","post-LayerNorm","(","`","pre_layer_norm","=","False","`",")",".","norm_elementwise_affine","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","learnable","per-element","affine","parameters","during","layer","normalization",".","use_patch_pos_embed","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","position","embeddings","inside","the","patch","embedding","layer","(","`","PatchEmbed","`",")",".","final_dropout","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","a","final","Dropout","layer","after","the","feedforward","network","."]
docstr: Args:
      hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`.
          When continuous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input
          hidden_states
      encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*):
          Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
          self-attention.
      timestep ( `torch.long`, *optional*):
          Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step.
      class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
          Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels
          conditioning.
      cross_attention_kwargs (*optional*):
          Keyword arguments to supply to the cross attention layers, if used.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
      hidden_states_is_embedding (`bool`, *optional*, defaults to `False`):
          Whether or not hidden_states is an embedding directly usable by the transformer. In this case we will
          ignore input handling (e.g. continuous, vectorized, etc.) and directly feed hidden_states into the
          transformer blocks.
      unpatchify (`bool`, *optional*, defaults to `True`):
          Whether to unpatchify the transformer output.
  
  Returns:
      [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`:
      [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When
      returning a tuple, the first element is the sample tensor.
docstr_tok: ["Args",":","hidden_states","(","When","discrete",",","`","torch.LongTensor","`","of","shape","`","(","batch","size",",","num","latent","pixels",")","`",".","When","continuous",",","`","torch.FloatTensor","`","of","shape","`","(","batch","size",",","channel",",","height",",","width",")","`",")",":","Input","hidden_states","encoder_hidden_states","(","`","torch.LongTensor","`","of","shape","`","(","batch","size",",","encoder_hidden_states","dim",")","`",",","*","optional","*",")",":","Conditional","embeddings","for","cross","attention","layer",".","If","not","given",",","cross-attention","defaults","to","self-attention",".","timestep","(","`","torch.long","`",",","*","optional","*",")",":","Optional","timestep","to","be","applied","as","an","embedding","in","AdaLayerNorm","'s",".","Used","to","indicate","denoising","step",".","class_labels","(","`","torch.LongTensor","`","of","shape","`","(","batch","size",",","num","classes",")","`",",","*","optional","*",")",":","Optional","class","labels","to","be","applied","as","an","embedding","in","AdaLayerZeroNorm",".","Used","to","indicate","class","labels","conditioning",".","cross_attention_kwargs","(","*","optional","*",")",":","Keyword","arguments","to","supply","to","the","cross","attention","layers",",","if","used",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","models.unet_2d_condition.UNet2DConditionOutput","`","]","instead","of","a","plain","tuple",".","hidden_states_is_embedding","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","or","not","hidden_states","is","an","embedding","directly","usable","by","the","transformer",".","In","this","case","we","will","ignore","input","handling","(","e.g",".","continuous",",","vectorized",",","etc",".",")","and","directly","feed","hidden_states","into","the","transformer","blocks",".","unpatchify","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","to","unpatchify","the","transformer","output",".","Returns",":","[","`","~models.transformer_2d.Transformer2DModelOutput","`","]","or","`","tuple","`",":","[","`","~models.transformer_2d.Transformer2DModelOutput","`","]","if","`","return_dict","`","is","True",",","otherwise","a","`","tuple","`",".","When","returning","a","tuple",",","the","first","element","is","the","sample","tensor","."]
code: def forward(
      self,
      hidden_states,
      encoder_hidden_states=None,
      timestep=None,
      class_labels=None,
      cross_attention_kwargs=None,
      return_dict: bool = True,
      hidden_states_is_embedding: bool = False,
      unpatchify: bool = True,
  ):
      """
      Args:
          hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`.
              When continuous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input
              hidden_states
          encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*):
              Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
              self-attention.
          timestep ( `torch.long`, *optional*):
              Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step.
          class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
              Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels
              conditioning.
          cross_attention_kwargs (*optional*):
              Keyword arguments to supply to the cross attention layers, if used.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
          hidden_states_is_embedding (`bool`, *optional*, defaults to `False`):
              Whether or not hidden_states is an embedding directly usable by the transformer. In this case we will
              ignore input handling (e.g. continuous, vectorized, etc.) and directly feed hidden_states into the
              transformer blocks.
          unpatchify (`bool`, *optional*, defaults to `True`):
              Whether to unpatchify the transformer output.
  
      Returns:
          [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`:
          [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When
          returning a tuple, the first element is the sample tensor.
      """
      # 0. Check inputs
  
      if not unpatchify and return_dict:
          raise ValueError(
              f"Cannot both define `unpatchify`: {unpatchify} and `return_dict`: {return_dict} since when"
              f" `unpatchify` is {unpatchify} the returned output is of shape (batch_size, seq_len, hidden_dim)"
              " rather than (batch_size, num_channels, height, width)."
          )
  
      # 1. Input
      if not hidden_states_is_embedding:
          hidden_states = self.pos_embed(hidden_states)
  
      # 2. Blocks
  
      # In ("downsample") blocks
      skips = []
      for in_block in self.transformer_in_blocks:
          hidden_states = in_block(
              hidden_states,
              encoder_hidden_states=encoder_hidden_states,
              timestep=timestep,
              cross_attention_kwargs=cross_attention_kwargs,
              class_labels=class_labels,
          )
          skips.append(hidden_states)
  
      # Mid block
      hidden_states = self.transformer_mid_block(hidden_states)
  
      # Out ("upsample") blocks
      for out_block in self.transformer_out_blocks:
          hidden_states = out_block["skip"](hidden_states, skips.pop())
          hidden_states = out_block["block"](
              hidden_states,
              encoder_hidden_states=encoder_hidden_states,
              timestep=timestep,
              cross_attention_kwargs=cross_attention_kwargs,
              class_labels=class_labels,
          )
  
      # 3. Output
      # Don't support AdaLayerNorm for now, so no conditioning/scale/shift logic
      hidden_states = self.norm_out(hidden_states)
      # hidden_states = self.proj_out(hidden_states)
  
      if unpatchify:
          # unpatchify
          height = width = int(hidden_states.shape[1] ** 0.5)
          hidden_states = hidden_states.reshape(
              shape=(
                  -1,
                  height,
                  width,
                  self.patch_size,
                  self.patch_size,
                  self.out_channels,
              )
          )
          hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
          output = hidden_states.reshape(
              shape=(
                  -1,
                  self.out_channels,
                  height * self.patch_size,
                  width * self.patch_size,
              )
          )
      else:
          output = hidden_states
  
      if not return_dict:
          return (output,)
  
      return Transformer2DModelOutput(sample=output)
  
imports: import math
  from typing import Optional, Union
  import torch
  from torch import nn
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...models import ModelMixin
  from ...models.attention import AdaLayerNorm, FeedForward
  from ...models.attention_processor import Attention
  from ...models.embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed
  from ...models.transformer_2d import Transformer2DModelOutput
  from ...utils import logging
label_desc: Use image node_id 2 for calling the UTransformer2DModel obj's underlying member method code with example usage: obj.forward(hidden_states, encoder_hidden_states, timestep, class_labels, cross_attention_kwargs, return_dict, hidden_states_is_embedding, unpatchify) and returns: Transformer2DModelOutput, output