code: def __init__(self, cfg, model_parallel_config):
      super().__init__()
      self.cfg = cfg
      self.config = model_parallel_config
      # Make sure the initialization on different GPUs are the same
      self.unet_type = cfg.get("unet_type", "base")
      self.noise_cond_aug = cfg.get("noise_cond_aug", False)
      if self.unet_type == "base":
          logging.info("Initializing UNet.")
          unet = UNetModel(
              **cfg.unet, text_embed_dim=cfg.conditioning.embed_dim
          )
      elif self.unet_type == "sr":
          logging.info("Initializing Efficient-UNet.")
          unet = EfficientUNetModel(
              **cfg.unet,
              text_embed_dim=cfg.conditioning.embed_dim,
              noise_cond_aug=self.noise_cond_aug,
          )
      elif self.unet_type == "sr-unet":
          logging.info("Initializing UNet for SR model.")
          unet = UNetModel(
              **cfg.unet,
              text_embed_dim=cfg.conditioning.embed_dim,
              noise_cond_aug=self.noise_cond_aug,
          )
      else:
          raise NotImplemented(
              f"{self.unet_type} UNet is not implemented."
          )
  
      self.channels_last = cfg.get("channels_last", False)
      if self.channels_last:
          assert (
              OPT_GROUP_NORM
          ), "Training in channels last format requires optmized group norm implementation."
          logging.info("Training in torch channels last format.")
          unet = unet.to(memory_format=torch.channels_last)
  
      # Preconditioning
      self.preconditioning_type = cfg.get(
          "preconditioning_type", "DDPM"
      )
      if self.preconditioning_type == "DDPM":
          logging.info("Preconditioned with Continous DDPM")
          self.model = ContinousDDPMPrecond(
              unet=unet,
              **cfg.preconditioning,
              noise_cond_aug=self.noise_cond_aug,
          )
          self.sampler = DDPMSampler(
              unet_type=self.unet_type, denoiser=self.model.scheduler
          )
      elif self.preconditioning_type == "EDM":
          logging.info("Preconditioned with EDM")
          self.model = EDMPrecond(
              unet=unet,
              **cfg.preconditioning,
              noise_cond_aug=self.noise_cond_aug,
          )
          self.sampler = EDMSampler(unet_type=self.unet_type)
      else:
          raise NotImplemented(
              f"{self.preconditioning_type} preconditioning is not implemented."
          )
  
      self.rng = None
      self.conditioning = cfg.conditioning
      self.text_drop_rate = cfg.conditioning.drop_rate
      self.model_type = None
      self.image_size = cfg.unet.image_size
  
imports: import itertools
  from datetime import datetime
  from functools import partial
  from typing import Any
  import torch
  from omegaconf import DictConfig, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.multimodal.data.imagen.imagen_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.imagen.precond import ContinousDDPMPrecond, EDMPrecond
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.nets import EfficientUNetModel, UNetModel
  from nemo.collections.multimodal.modules.imagen.encoder.t5encoder import T5Encoder
  from nemo.collections.multimodal.modules.imagen.sampler.sampler import DDPMSampler, EDMSampler
  from nemo.collections.multimodal.parts.imagen.utils import random_dropout
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import Serialization
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new Imagen object from inherited base classes: torch.nn, Serialization with example: obj = Imagen(cfg, model_parallel_config)