code: def forward(
      self,
      x,
      pooled_text_emb,
      encoder_hidden_states,
      cross_attention_kwargs,
  ):
      if self.downsample is not None:
          x = self.downsample(x)
  
      for res_block, attention_block in zip(
          self.res_blocks, self.attention_blocks
      ):
          x = res_block(x, pooled_text_emb)
  
          batch_size, channels, height, width = x.shape
          x = x.view(batch_size, channels, height * width).permute(
              0, 2, 1
          )
          x = attention_block(
              x,
              encoder_hidden_states=encoder_hidden_states,
              cross_attention_kwargs=cross_attention_kwargs,
          )
          x = x.permute(0, 2, 1).view(
              batch_size, channels, height, width
          )
  
      if self.upsample is not None:
          x = self.upsample(x)
  
      return x
  
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 2 for calling the UVitBlock obj's underlying member method code with example usage: obj.forward(x, pooled_text_emb, encoder_hidden_states, cross_attention_kwargs) and returns: x