code: def main(cfg):
      # Read configuration parameters
      nnodes_per_cfg = cfg.fid.nnodes_per_cfg
      ntasks_per_node = cfg.fid.ntasks_per_node
      local_task_id = cfg.fid.local_task_id
      num_images_to_eval = cfg.fid.num_images_to_eval
      path = cfg.fid.coco_captions_path
      save_text = cfg.fid.save_text
  
      node_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
      node_id_per_cfg = node_id % nnodes_per_cfg
  
      current_node_cfg = cfg.fid.classifier_free_guidance[
          node_id // nnodes_per_cfg
      ]
      save_path = os.path.join(cfg.fid.save_path, str(current_node_cfg))
  
      # Read and store captions
      captions = []
      caption_files = sorted(os.listdir(path))
      assert len(caption_files) >= num_images_to_eval
      for file in caption_files[:num_images_to_eval]:
          with open(os.path.join(path, file), "r") as f:
              captions += f.readlines()
      print(
          f"The total number of captions to generate is: {len(captions)}"
      )
  
      # Calculate partition sizes and select the partition for the current node
      partition_size_per_node = num_images_to_eval // nnodes_per_cfg
      start_idx = node_id_per_cfg * partition_size_per_node
      end_idx = (
          (node_id_per_cfg + 1) * partition_size_per_node
          if node_id_per_cfg != nnodes_per_cfg - 1
          else None
      )
      captions = captions[start_idx:end_idx]
      print(
          f"Current node {node_id} will generate images from {start_idx} to {end_idx}"
      )
  
      local_task_id = (
          int(local_task_id)
          if local_task_id is not None
          else int(os.environ.get("SLURM_LOCALID", 0))
      )
      partition_size_per_task = int(len(captions) // ntasks_per_node)
  
      # Select the partition for the current task
      start_idx = local_task_id * partition_size_per_task
      end_idx = (
          (local_task_id + 1) * partition_size_per_task
          if local_task_id != ntasks_per_node - 1
          else None
      )
      input = captions[start_idx:end_idx]
      chunk_size = len(input)
  
      print(
          f"Current worker {node_id}:{local_task_id} will generate {len(input)} images"
      )
      os.makedirs(save_path, exist_ok=True)
  
      trainer = Trainer()
      pipeline = ImagenPipeline.from_pretrained(
          cfg=cfg.infer,
          trainer=trainer,
          megatron_loading=True,
          megatron_cfg=cfg,
      )
  
      # Generate images using the model and save them
      batch_idx = 0
      batch_size = cfg.fid.ncaptions_per_batch
      while True:
          if batch_idx * batch_size >= len(input):
              break
          batch_captions = input[
              batch_idx * batch_size : (batch_idx + 1) * batch_size
          ]
          # Different seed for every image
          seeds = [
              local_task_id * chunk_size + batch_idx * batch_size + idx
              for idx in range(len(batch_captions))
          ]
          with torch.no_grad():
              images, all_res_images, *_ = pipeline(
                  prompts=batch_captions,
                  seed=seeds,
                  single_batch_mode=True,
                  classifier_free_guidance=current_node_cfg,
              )
  
          if cfg.fid.save_all_res:
              all_res = [
                  f"_RES{model.image_size}" for model in pipeline.models
              ]
              outpaths = []
              # for the highest resolution we save as its original name so that
              # we can automate the CLIP & FID calculation process from Megatron-Launcher
              all_res[-1] = ""
              for res in all_res:
                  outpath = f"{save_path}{res}"
                  os.makedirs(outpath, exist_ok=True)
                  outpaths.append(outpath)
              for outpath, one_res in zip(outpaths, all_res_images):
                  for idx, (caption, image) in enumerate(
                      zip(batch_captions, one_res[0])
                  ):
                      image_idx = (
                          local_task_id * chunk_size
                          + batch_idx * batch_size
                          + idx
                      )
                      image.save(
                          os.path.join(
                              outpath, f"image{image_idx:06d}.png"
                          )
                      )
                      if save_text:
                          with open(
                              os.path.join(
                                  outpath, f"image{image_idx:06d}.txt"
                              ),
                              "w",
                          ) as f:
                              f.writelines(caption)
          else:
              for idx, (caption, image) in enumerate(
                  zip(batch_captions, images[0])
              ):
                  image_idx = (
                      local_task_id * chunk_size
                      + batch_idx * batch_size
                      + idx
                  )
                  image.save(
                      os.path.join(
                          save_path, f"image{image_idx:06d}.png"
                      )
                  )
                  if save_text:
                      with open(
                          os.path.join(
                              save_path, f"image{image_idx:06d}.txt"
                          ),
                          "w",
                      ) as f:
                          f.writelines(caption)
          print(
              f"Save {len(images[0])} images to {save_path} with name from image{(local_task_id*chunk_size+batch_idx*batch_size):06d}.png to image{image_idx:06d}.png"
          )
          batch_idx += 1
  
imports: import os
  import torch
  from pytorch_lightning import Trainer
  from nemo.collections.multimodal.models.text_to_image.imagen.imagen_pipeline import ImagenPipeline
  from nemo.core.config import hydra_runner
label_desc: Use image node_id 1 for calling a global function with example usage: main(cfg) without return types