code: def __init__(self, cfg, model_parallel_config):
      super().__init__()
      self.cfg = cfg
      self.config = model_parallel_config
      self.with_prior_preservation = self.cfg.with_prior_preservation
      self.num_reg_images = self.cfg.data.num_reg_images
      self.prior_loss_weight = self.cfg.prior_loss_weight
      self.num_images_per_prompt = self.cfg.data.num_images_per_prompt
  
      self.train_text_encoder = self.cfg.train_text_encoder
      self.instantiate_text_encoder(self.cfg.cond_stage_config)
  
      self.inductor = self.cfg.inductor
      self.inductor_cudagraphs = self.cfg.inductor_cudagraphs
  
      self.instantiate_vae(self.cfg.first_stage_config)
      self.instantiate_unet(self.cfg.unet_config)
  
      self.scale_factor = self.cfg.scale_factor
      self.num_timesteps = self.cfg.noise_scheduler.timesteps
      self.parameterization = self.cfg.noise_scheduler.parameterization
      self.get_noise_scheduler(self.cfg.noise_scheduler)
  
      self.model_type = None
      self.rng = torch.Generator(
          device=torch.cuda.current_device(),
      )
  
      self.use_cached_latents = self.cfg.use_cached_latents
  
      if self.cfg.channels_last:
          self.unet = self.unet.to(memory_format=torch.channels_last)
  
imports: from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from torch._inductor import config
  from nemo.collections.multimodal.data.dreambooth.dreambooth_dataset import DreamBoothDataset
  from nemo.collections.multimodal.modules.stable_diffusion.distributions.distributions import DiagonalGaussianDistribution
  from nemo.collections.multimodal.parts.utils import randn_like
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler
  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 DreamBooth object from inherited base classes: torch.nn, Serialization with example: obj = DreamBooth(cfg, model_parallel_config)