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","."]
code: def __init__(
      self,
      num_attention_heads: int = 16,
      attention_head_dim: int = 88,
      in_channels: Optional[int] = None,
      out_channels: Optional[int] = None,
      num_layers: int = 1,
      dropout: float = 0.0,
      norm_num_groups: int = 32,
      cross_attention_dim: Optional[int] = None,
      attention_bias: bool = False,
      sample_size: Optional[int] = None,
      num_vector_embeds: Optional[int] = None,
      patch_size: Optional[int] = 2,
      activation_fn: str = "geglu",
      num_embeds_ada_norm: Optional[int] = None,
      use_linear_projection: bool = False,
      only_cross_attention: bool = False,
      upcast_attention: bool = False,
      norm_type: str = "layer_norm",
      block_type: str = "unidiffuser",
      pre_layer_norm: bool = False,
      norm_elementwise_affine: bool = True,
      use_patch_pos_embed=False,
      ff_final_dropout: bool = False,
  ):
      super().__init__()
      self.use_linear_projection = use_linear_projection
      self.num_attention_heads = num_attention_heads
      self.attention_head_dim = attention_head_dim
      inner_dim = num_attention_heads * attention_head_dim
  
      # 1. Input
      # Only support patch input of shape (batch_size, num_channels, height, width) for now
      assert (
          in_channels is not None and patch_size is not None
      ), "Patch input requires in_channels and patch_size."
  
      assert (
          sample_size is not None
      ), "UTransformer2DModel over patched input must provide sample_size"
  
      # 2. Define input layers
      self.height = sample_size
      self.width = sample_size
  
      self.patch_size = patch_size
      self.pos_embed = PatchEmbed(
          height=sample_size,
          width=sample_size,
          patch_size=patch_size,
          in_channels=in_channels,
          embed_dim=inner_dim,
          use_pos_embed=use_patch_pos_embed,
      )
  
      # 3. Define transformers blocks
      # Modify this to have in_blocks ("downsample" blocks, even though we don't actually downsample), a mid_block,
      # and out_blocks ("upsample" blocks). Like a U-Net, there are skip connections from in_blocks to out_blocks in
      # a "U"-shaped fashion (e.g. first in_block to last out_block, etc.).
      # Quick hack to make the transformer block type configurable
      if block_type == "unidiffuser":
          block_cls = UniDiffuserBlock
      else:
          block_cls = UTransformerBlock
      self.transformer_in_blocks = nn.ModuleList(
          [
              block_cls(
                  inner_dim,
                  num_attention_heads,
                  attention_head_dim,
                  dropout=dropout,
                  cross_attention_dim=cross_attention_dim,
                  activation_fn=activation_fn,
                  num_embeds_ada_norm=num_embeds_ada_norm,
                  attention_bias=attention_bias,
                  only_cross_attention=only_cross_attention,
                  upcast_attention=upcast_attention,
                  norm_type=norm_type,
                  pre_layer_norm=pre_layer_norm,
                  norm_elementwise_affine=norm_elementwise_affine,
                  final_dropout=ff_final_dropout,
              )
              for d in range(num_layers // 2)
          ]
      )
  
      self.transformer_mid_block = block_cls(
          inner_dim,
          num_attention_heads,
          attention_head_dim,
          dropout=dropout,
          cross_attention_dim=cross_attention_dim,
          activation_fn=activation_fn,
          num_embeds_ada_norm=num_embeds_ada_norm,
          attention_bias=attention_bias,
          only_cross_attention=only_cross_attention,
          upcast_attention=upcast_attention,
          norm_type=norm_type,
          pre_layer_norm=pre_layer_norm,
          norm_elementwise_affine=norm_elementwise_affine,
          final_dropout=ff_final_dropout,
      )
  
      # For each skip connection, we use a SkipBlock (concatenation + Linear + LayerNorm) to process the inputs
      # before each transformer out_block.
      self.transformer_out_blocks = nn.ModuleList(
          [
              nn.ModuleDict(
                  {
                      "skip": SkipBlock(
                          inner_dim,
                      ),
                      "block": block_cls(
                          inner_dim,
                          num_attention_heads,
                          attention_head_dim,
                          dropout=dropout,
                          cross_attention_dim=cross_attention_dim,
                          activation_fn=activation_fn,
                          num_embeds_ada_norm=num_embeds_ada_norm,
                          attention_bias=attention_bias,
                          only_cross_attention=only_cross_attention,
                          upcast_attention=upcast_attention,
                          norm_type=norm_type,
                          pre_layer_norm=pre_layer_norm,
                          norm_elementwise_affine=norm_elementwise_affine,
                          final_dropout=ff_final_dropout,
                      ),
                  }
              )
              for d in range(num_layers // 2)
          ]
      )
  
      # 4. Define output layers
      self.out_channels = (
          in_channels if out_channels is None else out_channels
      )
  
      # Following the UniDiffuser U-ViT implementation, we process the transformer output with
      # a LayerNorm layer with per-element affine params
      self.norm_out = nn.LayerNorm(inner_dim)
  
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 FeedForward
  from ...models.attention_processor import Attention
  from ...models.embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed
  from ...models.normalization import AdaLayerNorm
  from ...models.transformer_2d import Transformer2DModelOutput
  from ...utils import logging
label_desc: Use image node_id 1 to create a new UTransformer2DModel object from inherited base classes: ModelMixin, ConfigMixin with example: obj = UTransformer2DModel(num_attention_heads, attention_head_dim, in_channels, out_channels, num_layers, dropout, norm_num_groups, cross_attention_dim, attention_bias, sample_size, num_vector_embeds, patch_size, activation_fn, num_embeds_ada_norm, use_linear_projection, only_cross_attention, upcast_attention, norm_type, block_type, pre_layer_norm, norm_elementwise_affine, use_patch_pos_embed, ff_final_dropout)