code: def forward(
      self,
      input_ids,
      encoder_hidden_states,
      pooled_text_emb,
      micro_conds,
      cross_attention_kwargs=None,
  ):
      encoder_hidden_states = self.encoder_proj(encoder_hidden_states)
      encoder_hidden_states = self.encoder_proj_layer_norm(
          encoder_hidden_states
      )
  
      micro_cond_embeds = get_timestep_embedding(
          micro_conds.flatten(),
          self.config.micro_cond_encode_dim,
          flip_sin_to_cos=True,
          downscale_freq_shift=0,
      )
  
      micro_cond_embeds = micro_cond_embeds.reshape(
          (input_ids.shape[0], -1)
      )
  
      pooled_text_emb = torch.cat(
          [pooled_text_emb, micro_cond_embeds], dim=1
      )
      pooled_text_emb = pooled_text_emb.to(dtype=self.dtype)
      pooled_text_emb = self.cond_embed(pooled_text_emb).to(
          encoder_hidden_states.dtype
      )
  
      hidden_states = self.embed(input_ids)
  
      hidden_states = self.down_block(
          hidden_states,
          pooled_text_emb=pooled_text_emb,
          encoder_hidden_states=encoder_hidden_states,
          cross_attention_kwargs=cross_attention_kwargs,
      )
  
      batch_size, channels, height, width = hidden_states.shape
      hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
          batch_size, height * width, channels
      )
  
      hidden_states = self.project_to_hidden_norm(hidden_states)
      hidden_states = self.project_to_hidden(hidden_states)
  
      for layer in self.transformer_layers:
          if self.training and self.gradient_checkpointing:
  
              def layer_(*args):
                  return checkpoint(layer, *args)
  
          else:
              layer_ = layer
  
          hidden_states = layer_(
              hidden_states,
              encoder_hidden_states=encoder_hidden_states,
              cross_attention_kwargs=cross_attention_kwargs,
              added_cond_kwargs={"pooled_text_emb": pooled_text_emb},
          )
  
      hidden_states = self.project_from_hidden_norm(hidden_states)
      hidden_states = self.project_from_hidden(hidden_states)
  
      hidden_states = hidden_states.reshape(
          batch_size, height, width, channels
      ).permute(0, 3, 1, 2)
  
      hidden_states = self.up_block(
          hidden_states,
          pooled_text_emb=pooled_text_emb,
          encoder_hidden_states=encoder_hidden_states,
          cross_attention_kwargs=cross_attention_kwargs,
      )
  
      logits = self.mlm_layer(hidden_states)
  
      return logits
  
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 3 for calling the UVit2DModel obj's underlying member method code with example usage: obj.forward(input_ids, encoder_hidden_states, pooled_text_emb, micro_conds, cross_attention_kwargs) and returns: logits, checkpoint