class_docstr: The full UNet model with attention and timestep embedding used for Imagen Base and SR model.
  
  :param embed_dim: Dimension of embeddings. Also used to calculate the number of channels in ResBlock.
  :param image_size: Input image size. Used to calculate where to inject attention layers in UNet.
  :param channels: Input channel number, defaults to 3.
  :param text_embed_dim: Dimension of conditioned text embedding. Different text encoders and different model versions have different values, defaults to 512
  :param num_res_blocks: Number of ResBlock in each level of UNet, defaults to 3.
  :param channel_mult: Used with embed_dim to calculate the number of channels for each level of UNet, defaults to [1, 2, 3, 4]
  :param num_attn_heads: The number of heads in the attention layer, defaults to 4.
  :param per_head_channels: The number of channels per attention head, defaults to 64.
  :param cond_dim: Dimension of Conditioning projections, defaults to 512.
  :param attention_type: Type of attention layer, defaults to 'fused'.
  :param feature_pooling_type: Type of pooling, defaults to 'attention'.
  :param learned_sinu_pos_emb_dim: Dimension of learned time positional embedding. 0 for unlearned timestep embeddings. Defaults to 16
  :param attention_resolutions: List of resolutions to inject attention layers. Defaults to [8, 16, 32]
  :param dropout: The rate of dropout, defaults to 0.
  :param use_null_token: Whether to create a learned null token for attention, defaults to False.
  :param init_conv_kernel_size: Initial Conv kernel size, defaults to 3.
  :param gradient_checkpointing: Whether to use gradient checkpointing, defaults to False.
  :param scale_shift_norm: Whether to use scale shift norm, defaults to False.
  :param stable_attention: Whether to use numerically-stable attention calculation, defaults to True.
  :param flash_attention: Whether to use flash attention calculation, defaults to False.
  :param resblock_updown: Whether to use ResBlock or Downsample/Upsample, defaults to False.
  :param resample_with_conv: When resblock_updown=False, whether to use conv in addition to Pooling&ConvTranspose. Defaults to True.
  :param low_res_cond: Whether conditioned on low-resolution input, used for SR model. Defaults to False.
  :param noise_cond_aug: Whether to add noise conditioned augmentation with low-resolution input. Defaults to False.
class_docstr_tok: ["The","full","UNet","model","with","attention","and","timestep","embedding","used","for","Imagen","Base","and","SR","model",".",":","param","embed_dim",":","Dimension","of","embeddings",".","Also","used","to","calculate","the","number","of","channels","in","ResBlock",".",":","param","image_size",":","Input","image","size",".","Used","to","calculate","where","to","inject","attention","layers","in","UNet",".",":","param","channels",":","Input","channel","number",",","defaults","to","3.",":","param","text_embed_dim",":","Dimension","of","conditioned","text","embedding",".","Different","text","encoders","and","different","model","versions","have","different","values",",","defaults","to","512",":","param","num_res_blocks",":","Number","of","ResBlock","in","each","level","of","UNet",",","defaults","to","3.",":","param","channel_mult",":","Used","with","embed_dim","to","calculate","the","number","of","channels","for","each","level","of","UNet",",","defaults","to","[","1",",","2",",","3",",","4","]",":","param","num_attn_heads",":","The","number","of","heads","in","the","attention","layer",",","defaults","to","4.",":","param","per_head_channels",":","The","number","of","channels","per","attention","head",",","defaults","to","64.",":","param","cond_dim",":","Dimension","of","Conditioning","projections",",","defaults","to","512.",":","param","attention_type",":","Type","of","attention","layer",",","defaults","to","'fused","'",".",":","param","feature_pooling_type",":","Type","of","pooling",",","defaults","to","'attention","'",".",":","param","learned_sinu_pos_emb_dim",":","Dimension","of","learned","time","positional","embedding",".","0","for","unlearned","timestep","embeddings",".","Defaults","to","16",":","param","attention_resolutions",":","List","of","resolutions","to","inject","attention","layers",".","Defaults","to","[","8",",","16",",","32","]",":","param","dropout",":","The","rate","of","dropout",",","defaults","to","0.",":","param","use_null_token",":","Whether","to","create","a","learned","null","token","for","attention",",","defaults","to","False",".",":","param","init_conv_kernel_size",":","Initial","Conv","kernel","size",",","defaults","to","3.",":","param","gradient_checkpointing",":","Whether","to","use","gradient","checkpointing",",","defaults","to","False",".",":","param","scale_shift_norm",":","Whether","to","use","scale","shift","norm",",","defaults","to","False",".",":","param","stable_attention",":","Whether","to","use","numerically-stable","attention","calculation",",","defaults","to","True",".",":","param","flash_attention",":","Whether","to","use","flash","attention","calculation",",","defaults","to","False",".",":","param","resblock_updown",":","Whether","to","use","ResBlock","or","Downsample\/Upsample",",","defaults","to","False",".",":","param","resample_with_conv",":","When","resblock_updown=False",",","whether","to","use","conv","in","addition","to","Pooling","&","ConvTranspose",".","Defaults","to","True",".",":","param","low_res_cond",":","Whether","conditioned","on","low-resolution","input",",","used","for","SR","model",".","Defaults","to","False",".",":","param","noise_cond_aug",":","Whether","to","add","noise","conditioned","augmentation","with","low-resolution","input",".","Defaults","to","False","."]
code: def forward(
      self,
      x,
      time,
      text_embed=None,
      text_mask=None,
      x_low_res=None,
      time_low_res=None,
  ):
      if self.low_res_cond:
          assert x_low_res is not None, "x_low_res cannot be None"
      else:
          assert x_low_res is None, "x_low_res cannot be presented"
      if self.noise_cond_aug:
          assert (
              time_low_res is not None
          ), "time_low_res cannot be None when training with noise conditioning aug"
      else:
          assert (
              time_low_res is None
          ), "time_low_res cannot be presented"
      # Concatenating low resolution images
      if x_low_res is not None:
          if x_low_res.shape != x.shape:
              # Upscale if not done in the trainer
              _, _, new_height, new_width = x.shape
              x_low_res = F.interpolate(
                  x_low_res, (new_height, new_width), mode="bicubic"
              )
          x = torch.cat([x, x_low_res], dim=1)
      batch_size, device = x.shape[0], x.device
  
      if x.dtype != time.dtype or time.dtype != text_embed.dtype:
          dtype = text_embed.dtype
          x = x.to(dtype=dtype)
          time = time.to(dtype=dtype)
          if x_low_res is not None:
              x_low_res = x_low_res.to(dtype=dtype)
          if time_low_res is not None:
              time_low_res = time_low_res.to(dtype=dtype)
      # Time Conditioning
      t = self.time_embed(time)
      # Add lowres time conditioning
      if self.noise_cond_aug:
          lowres_t = self.lowres_time_embed(time_low_res)
          t += lowres_t
      # Text Conditioning
      text_cond = self.text_to_cond(text_embed)
  
      # Context Embedding
      # TODO We may want to concat time token here
      if self.use_null_token:
          # Null Context (Helpful when text_embed is drop)
          null_context = self.null_text_embedding.repeat(
              batch_size, 1, 1
          )
          context_emb = torch.cat([text_cond, null_context], dim=1)
          context_mask = torch.cat(
              [text_mask, torch.ones(batch_size, 1).to(device)], dim=1
          )
      else:
          context_emb = text_cond
          context_mask = text_mask
  
      # Add pooled text embeddings to the diffusion timestep
      # TODO We may only want to calculated the pooled feature based on text token length
      if self.feature_pooling_type == "mean":
          pooled_text_cond = text_cond.mean(dim=-2)
      elif self.feature_pooling_type == "attention":
          pooled_text_cond = self.attention_pooling(text_embed)
      text_hiddens = self.to_text_non_attn_cond(pooled_text_cond)
      t += text_hiddens
  
      h = x
      hs = []
      # UNet Forward
      for module in self.input_blocks:
          h = module(h, t, context_emb, context_mask)
          hs.append(h)
      h = self.middle_block(h, t, context_emb, context_mask)
      for module in self.output_blocks:
          h_prev = hs.pop()
          h = torch.cat([h, h_prev], dim=1)
          h = module(h, t, context_emb, context_mask)
      return self.out(h)
  
file_path: NeMo/nemo/collections/multimodal/modules/imagen/diffusionmodules/nets.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.attention import SelfAttentionPooling
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.blocks import ConditionalSequential, DBlock, FusedCrossAttentionBlock, ResBlock, StackedCrossAttentionBlock, UBlock
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.embs import LearnedSinusoidalPosEmb, UnLearnedSinusoidalPosEmb
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.layers import Downsample
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.layers import UpsampleLearnable
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.layers import linear, normalization, zero_module
label_desc: Use image node_id 2 for calling the UNetModel obj's underlying member method code with example usage: obj.forward(x, time, text_embed, text_mask, x_low_res, time_low_res) and returns: self