class_docstr: The full Efficient UNet model with attention and timestep embedding used for Imagen 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 channel_mult: Used with embed_dim to calculate the number of channels for each level of UNet, defaults to [1, 1, 2, 4, 8].
  :param num_attn_heads: The number of heads in the attention layer, defaults to 8.
  :param per_head_channels: The number of channels per attention head, defaults to 64.
  :param attention_type: Type of attention layer, defaults to 'fused'.
  :param atnn_enabled_at: Whether to enable attention at each level, defaults to [0, 0, 0, 0, 1].
  :param feature_pooling_type: Type of pooling, defaults to 'attention'.
  :param stride: Stride in ResBlock, defaults to 2.
  :param num_resblocks: Used with num_res_blocks to calculate the number of residual blocks at each level of Efficient-UNet. Defaults to [1, 2, 4, 8, 8].
  :param learned_sinu_pos_emb_dim: Dimension of learned time positional embedding. 0 for unlearned timestep embeddings. Defaults to 16
  :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 skip_connection_scaling: Whether to use 1/sqrt(2) scaling for ResBlock skip connection, 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","Efficient","UNet","model","with","attention","and","timestep","embedding","used","for","Imagen","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","channel_mult",":","Used","with","embed_dim","to","calculate","the","number","of","channels","for","each","level","of","UNet",",","defaults","to","[","1",",","1",",","2",",","4",",","8","]",".",":","param","num_attn_heads",":","The","number","of","heads","in","the","attention","layer",",","defaults","to","8.",":","param","per_head_channels",":","The","number","of","channels","per","attention","head",",","defaults","to","64.",":","param","attention_type",":","Type","of","attention","layer",",","defaults","to","'fused","'",".",":","param","atnn_enabled_at",":","Whether","to","enable","attention","at","each","level",",","defaults","to","[","0",",","0",",","0",",","0",",","1","]",".",":","param","feature_pooling_type",":","Type","of","pooling",",","defaults","to","'attention","'",".",":","param","stride",":","Stride","in","ResBlock",",","defaults","to","2.",":","param","num_resblocks",":","Used","with","num_res_blocks","to","calculate","the","number","of","residual","blocks","at","each","level","of","Efficient-UNet",".","Defaults","to","[","1",",","2",",","4",",","8",",","8","]",".",":","param","learned_sinu_pos_emb_dim",":","Dimension","of","learned","time","positional","embedding",".","0","for","unlearned","timestep","embeddings",".","Defaults","to","16",":","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","skip_connection_scaling",":","Whether","to","use","1\/sqrt","(","2",")","scaling","for","ResBlock","skip","connection",",","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,
      image_size,
      channels=3,
      text_embed_dim=512,  # Dimension of conditioned text embedding. Different text encoders and different model versions have different values
      channel_mult=[
          1,
          1,
          2,
          4,
          8,
      ],  # Used with embed_dim to calculate the number of channels for each level of Efficient-UNet
      num_attn_heads=8,  # The number of heads in the attention layer
      per_head_channels=64,  # The number of channels per attention head
      attention_type="fused",  # Type of attention layer
      atnn_enabled_at=[
          0,
          0,
          0,
          0,
          1,
      ],  # Whether to enable attention at each level
      feature_pooling_type="attention",  # Type of pooling
      stride=2,  # Stride in ResBlock
      num_resblocks=[
          1,
          2,
          4,
          8,
          8,
      ],  # Used with num_res_blocks to calculate the number of residual blocks at each level of Efficient-UNet
      learned_sinu_pos_emb_dim=16,  # Dimension of learned time positional embedding. 0 for unlearned timestep embeddings.
      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
      skip_connection_scaling=False,  # Whether to use 1/sqrt(2) scaling for ResBlock skip connection
      noise_cond_aug=False,
  ):
      super().__init__()
  
      self.n_levels = len(channel_mult)
      self.image_size = image_size
      # 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),
          )
  
      self.noise_cond_aug = noise_cond_aug
      if self.noise_cond_aug:
          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),
          )
      cond_dim = text_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
  
      # Initial Convolution
      # Multiply in_channels by 2 because we concatenate with low res inputs.
      in_channels = channels * 2
      init_dim = embed_dim * channel_mult[0]
      self.init_conv = nn.Conv2d(
          in_channels,
          init_dim,
          init_conv_kernel_size,
          padding=init_conv_kernel_size // 2,
      )
      # Efficient-UNet Init
      self.DBlocks = nn.ModuleDict()
      self.UBlocks = nn.ModuleDict()
      ch = init_dim
      for level, mult in enumerate(channel_mult):
          # Different level has different num of res blocks
          num_resblock = num_resblocks[level]
          # Only perform upsample/downsample if it is not the last (deepest) level
          is_last_level = level == len(channel_mult) - 1
          level_attention_type = (
              attention_type if atnn_enabled_at[level] else None
          )
  
          level_key = str(
              level
          )  # TODO Change to more meaningful naming
          self.DBlocks[level_key] = DBlock(
              channels=ch,
              emb_channels=time_embed_dim,
              out_channels=int(mult * embed_dim),
              use_scale_shift_norm=scale_shift_norm,
              conv_down=not is_last_level,
              stride=stride,
              num_resblocks=num_resblock,
              attention_type=level_attention_type,
              text_embed_dim=cond_dim,
              num_heads=num_attn_heads,
              num_head_channels=per_head_channels,
              use_checkpoint=gradient_checkpointing,
              stable_attention=stable_attention,
              flash_attention=flash_attention,
              skip_connection_scaling=skip_connection_scaling,
          )
          self.UBlocks[level_key] = UBlock(
              channels=int(mult * embed_dim),
              emb_channels=time_embed_dim,
              out_channels=ch,
              use_scale_shift_norm=scale_shift_norm,
              conv_up=not is_last_level,
              stride=stride,
              num_resblocks=num_resblock,
              attention_type=level_attention_type,
              text_embed_dim=cond_dim,
              num_heads=num_attn_heads,
              num_head_channels=per_head_channels,
              use_checkpoint=gradient_checkpointing,
              stable_attention=stable_attention,
              flash_attention=flash_attention,
              skip_connection_scaling=skip_connection_scaling,
          )
          ch = int(mult * embed_dim)
      self.out = nn.Conv2d(channel_mult[0] * embed_dim, channels, 1)
  
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 EfficientUNetModel object from inherited base classes: nn with example: obj = EfficientUNetModel(embed_dim, image_size, channels, text_embed_dim, channel_mult, num_attn_heads, per_head_channels, attention_type, atnn_enabled_at, feature_pooling_type, stride, num_resblocks, learned_sinu_pos_emb_dim, use_null_token, init_conv_kernel_size, gradient_checkpointing, scale_shift_norm, stable_attention, flash_attention, skip_connection_scaling, noise_cond_aug)