code: def __init__(
      self,
      # global config
      hidden_size: int = 1024,
      use_bias: bool = False,
      hidden_dropout: float = 0.0,
      # conditioning dimensions
      cond_embed_dim: int = 768,
      micro_cond_encode_dim: int = 256,
      micro_cond_embed_dim: int = 1280,
      encoder_hidden_size: int = 768,
      # num tokens
      vocab_size: int = 8256,  # codebook_size + 1 (for the mask token) rounded
      codebook_size: int = 8192,
      # `UVit2DConvEmbed`
      in_channels: int = 768,
      block_out_channels: int = 768,
      num_res_blocks: int = 3,
      downsample: bool = False,
      upsample: bool = False,
      block_num_heads: int = 12,
      # `TransformerLayer`
      num_hidden_layers: int = 22,
      num_attention_heads: int = 16,
      # `Attention`
      attention_dropout: float = 0.0,
      # `FeedForward`
      intermediate_size: int = 2816,
      # `Norm`
      layer_norm_eps: float = 1e-6,
      ln_elementwise_affine: bool = True,
      sample_size: int = 64,
  ):
      super().__init__()
  
      self.encoder_proj = nn.Linear(
          encoder_hidden_size, hidden_size, bias=use_bias
      )
      self.encoder_proj_layer_norm = RMSNorm(
          hidden_size, layer_norm_eps, ln_elementwise_affine
      )
  
      self.embed = UVit2DConvEmbed(
          in_channels,
          block_out_channels,
          vocab_size,
          ln_elementwise_affine,
          layer_norm_eps,
          use_bias,
      )
  
      self.cond_embed = TimestepEmbedding(
          micro_cond_embed_dim + cond_embed_dim,
          hidden_size,
          sample_proj_bias=use_bias,
      )
  
      self.down_block = UVitBlock(
          block_out_channels,
          num_res_blocks,
          hidden_size,
          hidden_dropout,
          ln_elementwise_affine,
          layer_norm_eps,
          use_bias,
          block_num_heads,
          attention_dropout,
          downsample,
          False,
      )
  
      self.project_to_hidden_norm = RMSNorm(
          block_out_channels, layer_norm_eps, ln_elementwise_affine
      )
      self.project_to_hidden = nn.Linear(
          block_out_channels, hidden_size, bias=use_bias
      )
  
      self.transformer_layers = nn.ModuleList(
          [
              BasicTransformerBlock(
                  dim=hidden_size,
                  num_attention_heads=num_attention_heads,
                  attention_head_dim=hidden_size // num_attention_heads,
                  dropout=hidden_dropout,
                  cross_attention_dim=hidden_size,
                  attention_bias=use_bias,
                  norm_type="ada_norm_continuous",
                  ada_norm_continous_conditioning_embedding_dim=hidden_size,
                  norm_elementwise_affine=ln_elementwise_affine,
                  norm_eps=layer_norm_eps,
                  ada_norm_bias=use_bias,
                  ff_inner_dim=intermediate_size,
                  ff_bias=use_bias,
                  attention_out_bias=use_bias,
              )
              for _ in range(num_hidden_layers)
          ]
      )
  
      self.project_from_hidden_norm = RMSNorm(
          hidden_size, layer_norm_eps, ln_elementwise_affine
      )
      self.project_from_hidden = nn.Linear(
          hidden_size, block_out_channels, bias=use_bias
      )
  
      self.up_block = UVitBlock(
          block_out_channels,
          num_res_blocks,
          hidden_size,
          hidden_dropout,
          ln_elementwise_affine,
          layer_norm_eps,
          use_bias,
          block_num_heads,
          attention_dropout,
          downsample=False,
          upsample=upsample,
      )
  
      self.mlm_layer = ConvMlmLayer(
          block_out_channels,
          in_channels,
          use_bias,
          ln_elementwise_affine,
          layer_norm_eps,
          codebook_size,
      )
  
      self.gradient_checkpointing = False
  
imports: from typing import Dict, Union
  import torch
  import torch.nn.functional
  from torch import nn
  from torch.utils.checkpoint import checkpoint
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..loaders import PeftAdapterMixin
  from .attention import BasicTransformerBlock, SkipFFTransformerBlock
  from .attention_processor import ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor
  from .embeddings import TimestepEmbedding, get_timestep_embedding
  from .modeling_utils import ModelMixin
  from .normalization import GlobalResponseNorm, RMSNorm
  from .resnet import Downsample2D, Upsample2D
label_desc: Use image node_id 1 to create a new UVit2DModel object from inherited base classes: ModelMixin, ConfigMixin, PeftAdapterMixin with example: obj = UVit2DModel(hidden_size, use_bias, hidden_dropout, cond_embed_dim, micro_cond_encode_dim, micro_cond_embed_dim, encoder_hidden_size, vocab_size, codebook_size, in_channels, block_out_channels, num_res_blocks, downsample, upsample, block_num_heads, num_hidden_layers, num_attention_heads, attention_dropout, intermediate_size, layer_norm_eps, ln_elementwise_affine, sample_size)