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 __init__(
      self,
      embed_dim,  # Dimension of embeddings. Also used to calculate the number of channels in ResBlock
      image_size,  # Input image size. Used to calculate where to inject attention layers in UNet
      channels=3,  # Input channel number
      text_embed_dim=512,  # Dimension of conditioned text embedding. Different text encoders and different model versions have different values
      num_res_blocks=3,  # Number of ResBlock in each level of UNet
      channel_mult=[
          1,
          2,
          3,
          4,
      ],  # Used with embed_dim to calculate the number of channels for each level of UNet
      num_attn_heads=4,  # The number of heads in the attention layer
      per_head_channels=64,  # The number of channels per attention head
      cond_dim=512,  # Dimension of Conditioning projections
      attention_type="fused",  # Type of attention layer
      feature_pooling_type="attention",  # Type of pooling
      learned_sinu_pos_emb_dim=16,  # Dimension of learned time positional embedding. 0 for unlearned timestep embeddings.
      attention_resolutions=[
          8,
          16,
          32,
      ],  # List of resolutions to inject attention layers
      dropout=False,  # The rate of dropout
      use_null_token=False,  # Whether to create a learned null token for attention
      init_conv_kernel_size=3,  # Initial Conv kernel size. imagen_pytorch uses 7
      gradient_checkpointing=False,  # Whether to use gradient checkpointing
      scale_shift_norm=True,  # Whether to use scale shift norm
      stable_attention=True,  # Whether to use numerically-stable attention calculation
      flash_attention=False,  # Whether to use flash attention calculation
      resblock_updown=False,  # Whether to use ResBlock or Downsample/Upsample
      resample_with_conv=True,  # When resblock_updown=False, whether to use conv in addition to Pooling&ConvTranspose
      low_res_cond=False,
      noise_cond_aug=False,
  ):
      super().__init__()
  
      # Attention Class
      if attention_type == "stacked":
          attention_fn = StackedCrossAttentionBlock
      elif attention_type == "fused":
          attention_fn = FusedCrossAttentionBlock
      else:
          raise ValueError(
              "Attention {} not defined".format(attention_type)
          )
  
      # Time embedding for log(snr) noise from continous version
      time_embed_dim = embed_dim * 4
      assert learned_sinu_pos_emb_dim >= 0
      if learned_sinu_pos_emb_dim > 0:
          sinu_pos_emb = LearnedSinusoidalPosEmb(
              learned_sinu_pos_emb_dim
          )
          sinu_pos_emb_input_dim = learned_sinu_pos_emb_dim + 1
          self.time_embed = nn.Sequential(
              sinu_pos_emb,
              nn.Linear(sinu_pos_emb_input_dim, time_embed_dim),
              nn.SiLU(),
              nn.Linear(time_embed_dim, time_embed_dim),
          )
      else:
          # Unlearned Time Embedding
          sinu_pos_emb = UnLearnedSinusoidalPosEmb(embed_dim)
          self.time_embed = nn.Sequential(
              sinu_pos_emb,
              linear(embed_dim, time_embed_dim),
              nn.SiLU(),
              linear(time_embed_dim, time_embed_dim),
          )
  
      # Pooling
      assert (
          feature_pooling_type == "attention"
          or feature_pooling_type == "mean"
      )
      self.feature_pooling_type = feature_pooling_type
      if feature_pooling_type == "attention":
          self.attention_pooling = nn.Sequential(
              SelfAttentionPooling(input_dim=text_embed_dim),
              nn.LayerNorm(text_embed_dim),
              nn.Linear(text_embed_dim, cond_dim),
          )
  
      # Context Projections
      self.text_to_cond = linear(text_embed_dim, cond_dim)
      self.to_text_non_attn_cond = nn.Sequential(
          nn.LayerNorm(cond_dim),
          nn.Linear(cond_dim, time_embed_dim),
          nn.SiLU(),
          nn.Linear(time_embed_dim, time_embed_dim),
      )
  
      # Register for Null Token
      if use_null_token:
          self.null_text_embedding = nn.Parameter(
              torch.randn(
                  1, 1, cond_dim, dtype=self.text_to_cond.weight.dtype
              )
          )
      self.use_null_token = use_null_token
  
      # Converting attention resolutions to downsampling factor
      attention_ds = []
      attention_resolutions = sorted(attention_resolutions)
      self.image_size = image_size
      for res in attention_resolutions:
          attention_ds.append(image_size // int(res))
  
      self.low_res_cond = low_res_cond
      # Low res noise conditioning augmentation
      self.noise_cond_aug = noise_cond_aug
      if self.noise_cond_aug:
          assert (
              self.low_res_cond
          ), "noise conditioning augmentation should only be enabled when training with low-res cond"
          if learned_sinu_pos_emb_dim > 0:
              lowres_sinu_pos_emb = LearnedSinusoidalPosEmb(
                  learned_sinu_pos_emb_dim
              )
              lowres_sinu_pos_emb_dim = learned_sinu_pos_emb_dim + 1
          else:
              lowres_sinu_pos_emb = UnLearnedSinusoidalPosEmb(embed_dim)
              lowres_sinu_pos_emb_dim = embed_dim
          self.lowres_time_embed = nn.Sequential(
              lowres_sinu_pos_emb,
              nn.Linear(lowres_sinu_pos_emb_dim, time_embed_dim),
              nn.SiLU(),
              nn.Linear(time_embed_dim, time_embed_dim),
          )
  
      # Initial Convolution
      in_channels = 2 * channels if low_res_cond else channels
      init_dim = embed_dim * channel_mult[0]
      self.init_conv = ConditionalSequential(
          nn.Conv2d(
              in_channels,
              init_dim,
              init_conv_kernel_size,
              padding=init_conv_kernel_size // 2,
          )
      )
  
      if isinstance(num_res_blocks, int):
          res_blocks_list = [num_res_blocks] * len(channel_mult)
      else:
          res_blocks_list = num_res_blocks
      # UNet Init
      # Downsampling Layers
      # We use Conv2D for UNet
      CONV_DIM = 2
      ch = init_dim
      ds = 1
      self.input_blocks = nn.ModuleList([self.init_conv])
      num_input_block_channels = [ch]
      for level, mult in enumerate(channel_mult):
          num_res_blocks = res_blocks_list[level]
          for _ in range(num_res_blocks):
              out_channels = mult * embed_dim
              layers = [
                  ResBlock(
                      channels=ch,
                      emb_channels=time_embed_dim,
                      dropout=dropout,
                      out_channels=out_channels,
                      dims=CONV_DIM,
                      use_checkpoint=gradient_checkpointing,
                      use_scale_shift_norm=scale_shift_norm,
                      learnable_upsampling=True,
                  )
              ]
              ch = out_channels
              if ds in attention_ds:
                  layers.append(
                      attention_fn(
                          channels=ch,
                          num_heads=num_attn_heads,
                          num_head_channels=per_head_channels,
                          use_checkpoint=gradient_checkpointing,
                          stable_attention=stable_attention,
                          flash_attention=flash_attention,
                          context_dim=cond_dim,
                      )
                  )
              self.input_blocks.append(ConditionalSequential(*layers))
              num_input_block_channels.append(ch)
          is_last_level = level == len(channel_mult) - 1
          if not is_last_level:
              # DownSampling
              self.input_blocks.append(
                  ConditionalSequential(
                      ResBlock(
                          channels=ch,
                          emb_channels=time_embed_dim,
                          dropout=dropout,
                          out_channels=ch,
                          dims=CONV_DIM,
                          use_checkpoint=gradient_checkpointing,
                          use_scale_shift_norm=scale_shift_norm,
                          down=True,
                          learnable_upsampling=True,
                      )
                      if resblock_updown
                      else Downsample(
                          channels=ch,
                          use_conv=resample_with_conv,
                          dims=CONV_DIM,
                          out_channels=ch,
                      )
                  )
              )
              num_input_block_channels.append(ch)
              ds *= 2
  
      # Middle Layers
      self.middle_block = ConditionalSequential(
          # Mid Block 1
          ResBlock(
              channels=ch,
              emb_channels=time_embed_dim,
              dropout=dropout,
              dims=CONV_DIM,
              use_checkpoint=gradient_checkpointing,
              use_scale_shift_norm=scale_shift_norm,
              learnable_upsampling=True,
          ),
          # Attention Layer
          attention_fn(
              channels=ch,
              num_heads=num_attn_heads,
              num_head_channels=per_head_channels,
              use_checkpoint=gradient_checkpointing,
              stable_attention=stable_attention,
              flash_attention=flash_attention,
              context_dim=cond_dim,
          ),
          # Mid Block 2
          ResBlock(
              channels=ch,
              emb_channels=time_embed_dim,
              dropout=dropout,
              dims=CONV_DIM,
              use_checkpoint=gradient_checkpointing,
              use_scale_shift_norm=scale_shift_norm,
              learnable_upsampling=True,
          ),
      )
  
      # Upsampling Layers
      self.output_blocks = nn.ModuleList([])
      for level, mult in list(enumerate(channel_mult))[::-1]:
          num_res_blocks = res_blocks_list[level]
          for i in range(num_res_blocks + 1):
              ich = num_input_block_channels.pop()
              out_channels = embed_dim * mult
              layers = [
                  ResBlock(
                      channels=ch + ich,
                      emb_channels=time_embed_dim,
                      dropout=dropout,
                      out_channels=out_channels,
                      dims=CONV_DIM,
                      use_checkpoint=gradient_checkpointing,
                      use_scale_shift_norm=scale_shift_norm,
                      learnable_upsampling=True,
                  )
              ]
              ch = out_channels
  
              if ds in attention_ds:
                  layers.append(
                      attention_fn(
                          channels=ch,
                          num_heads=-1,  # TODO
                          num_head_channels=per_head_channels,
                          use_checkpoint=gradient_checkpointing,
                          stable_attention=stable_attention,
                          flash_attention=flash_attention,
                          context_dim=cond_dim,
                      )
                  )
              is_last_block = i == num_res_blocks
              if level and is_last_block:
                  layers.append(
                      ResBlock(
                          channels=ch,
                          emb_channels=time_embed_dim,
                          dropout=dropout,
                          out_channels=ch,
                          dims=CONV_DIM,
                          use_checkpoint=gradient_checkpointing,
                          use_scale_shift_norm=scale_shift_norm,
                          up=True,
                          learnable_upsampling=True,
                      )
                      if resblock_updown
                      else Upsample(
                          channels=ch,
                          use_conv=resample_with_conv,
                          dims=CONV_DIM,
                          out_channels=ch,
                      )
                  )
                  ds //= 2
              self.output_blocks.append(ConditionalSequential(*layers))
  
      self.out = nn.Sequential(
          normalization(ch),
          nn.SiLU(),
          zero_module(
              nn.Conv2d(
                  init_dim,
                  channels,
                  init_conv_kernel_size,
                  padding=init_conv_kernel_size // 2,
              )
          ),
      )
  
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 1 to create a new UNetModel object from inherited base classes: nn with example: obj = UNetModel(embed_dim, image_size, channels, text_embed_dim, num_res_blocks, channel_mult, num_attn_heads, per_head_channels, cond_dim, attention_type, feature_pooling_type, learned_sinu_pos_emb_dim, attention_resolutions, dropout, use_null_token, init_conv_kernel_size, gradient_checkpointing, scale_shift_norm, stable_attention, flash_attention, resblock_updown, resample_with_conv, low_res_cond, noise_cond_aug)