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 forward_with_cond_scale(
      self, *args, text_embed=None, cond_scale=1.0, **kwargs
  ):
      logits = self.forward(*args, text_embed=text_embed, **kwargs)
      if cond_scale == 1.0:
          return logits
      null_logits = self.forward(
          *args, text_embed=torch.zeros_like(text_embed), **kwargs
      )
      return null_logits + (logits - null_logits) * cond_scale
  
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 3 for calling the EfficientUNetModel obj's underlying member method code with example usage: obj.forward_with_cond_scale() and returns: unknown, logits