code: def setup_training_data(self, cfg):
      if self.cfg.with_prior_preservation:
          if cfg.regularization_dir is None:
              raise ValueError(
                  "Regularization images must be provided to train with prior preservation loss"
              )
          if cfg.regularization_prompt is None:
              raise ValueError(
                  "Regularization prompts must be provided to train with prior preservation loss"
              )
  
      self.train_dataset = DreamBoothDataset(
          instance_data_root=cfg.instance_dir,
          instance_prompt=cfg.instance_prompt,
          with_prior_preservation=self.cfg.with_prior_preservation,
          reg_data_root=cfg.regularization_dir
          if self.cfg.with_prior_preservation
          else None,
          reg_prompt=cfg.regularization_prompt
          if self.cfg.with_prior_preservation
          else None,
          size=cfg.resolution,
          center_crop=cfg.center_crop,
          load_cache_latents=self.model.use_cached_latents,
          cached_instance_data_root=self.cfg.data.get(
              "cached_instance_dir", None
          ),
          cached_reg_data_root=self.cfg.data.get("cached_reg_dir", None)
          if self.cfg.with_prior_preservation
          else None,
          vae=self.model.vae,
          text_encoder=self.model.text_encoder,
      )
  
      batch_sampler = MegatronPretrainingRandomSampler(
          total_samples=len(self.train_dataset),
          consumed_samples=self.compute_consumed_samples(0),
          micro_batch_size=self.cfg.micro_batch_size,
          global_batch_size=self.cfg.global_batch_size,
          data_parallel_rank=parallel_state.get_data_parallel_rank(),
          data_parallel_size=parallel_state.get_data_parallel_world_size(),
          drop_last=False,
      )
  
      self._train_dl = torch.utils.data.DataLoader(
          self.train_dataset,
          batch_sampler=batch_sampler,
          collate_fn=partial(
              _collate_fn,
              with_prior_preservation=self.cfg.with_prior_preservation,
          ),
          num_workers=cfg.num_workers,
          pin_memory=True,
          persistent_workers=True,
      )
  
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 14 for calling the MegatronDreamBooth obj's underlying member method code with example usage: obj.setup_training_data(cfg) without return types