docstr: PTL hook that is executed after DDP spawns.
      We setup datasets here as megatron datasets require DDP to instantiate.
      See https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup for more information.
  Args:
      stage (str, optional): Can be 'fit', 'validate', 'test' or 'predict'. Defaults to None.
code: def setup(self, stage=None):
      """PTL hook that is executed after DDP spawns.
          We setup datasets here as megatron datasets require DDP to instantiate.
          See https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup for more information.
      Args:
          stage (str, optional): Can be 'fit', 'validate', 'test' or 'predict'. Defaults to None.
      """
      self.model.rng.manual_seed(
          self.cfg.seed + 100 * parallel_state.get_data_parallel_rank()
      )
  
      # log number of parameters
      if isinstance(self.model, list):
          num_parameters_on_device = sum(
              [
                  sum([p.nelement() for p in model_module.parameters()])
                  for model_module in self.model
              ]
          )
      else:
          num_parameters_on_device = sum(
              [p.nelement() for p in self.model.parameters()]
          )
  
      # to be summed across data parallel group
      total_num_parameters = torch.tensor(
          num_parameters_on_device
      ).cuda(non_blocking=True)
  
      torch.distributed.all_reduce(
          total_num_parameters,
          group=parallel_state.get_model_parallel_group(),
      )
  
      logging.info(
          f"Pipeline model parallel rank: {parallel_state.get_pipeline_model_parallel_rank()}, "
          f"Tensor model parallel rank: {parallel_state.get_tensor_model_parallel_rank()}, "
          f"Number of model parameters on device: {num_parameters_on_device:.2e}. "
          f"Total number of model parameters: {total_num_parameters:.2e}."
      )
  
      resume_checkpoint_path = self.trainer.ckpt_path
      if resume_checkpoint_path:
          init_consumed_samples = (
              self._extract_consumed_samples_from_ckpt(
                  resume_checkpoint_path
              )
          )
      else:
          init_consumed_samples = 0
      self.init_consumed_samples = init_consumed_samples
      self.init_global_step = self.trainer.global_step
  
      # allowing restored models to optionally setup datasets
      self.build_train_valid_test_datasets()
  
      # Batch size need to be provided for webdatset
      self._num_micro_batches = get_num_microbatches()
      self._micro_batch_size = self.cfg.micro_batch_size
  
      self.setup_training_data(self.cfg.data)
      self.setup_validation_data(self.cfg.data)
      self.setup_test_data(self.cfg.data)
  
imports: from typing import Any
  import einops
  import torch
  import torch.nn
  from einops import rearrange, repeat
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from pytorch_lightning.utilities.rank_zero import rank_zero_only
  from torch._inductor import config
  from nemo.collections.multimodal.data.controlnet.controlnet_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm import LatentDiffusion
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.samplers.ddim import DDIMSampler
  from nemo.collections.multimodal.modules.stable_diffusion.attention import SpatialTransformer
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel import AttentionBlock, Downsample, ResBlock, TimestepEmbedSequential, UNetModel
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import conv_nd, linear, timestep_embedding, zero_module
  from nemo.collections.multimodal.parts.stable_diffusion.utils import exists, log_txt_as_img
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.utils import logging
label_desc: Use image node_id 14 for calling the MegatronControlNet obj's underlying member method code with example usage: obj.setup(stage) without return types