code: def build_train_valid_precached_datasets(
      model_cfg,
      consumed_samples,
  ):
      data_cfg = model_cfg.data
  
      # This function maps data that are tuples to dictionary.
      def tuple_to_dict(inp):
          for input in inp:
              out_dict = dict()
              out_dict[model_cfg.first_stage_key] = torch.tensor(
                  input["autoencoderkl_image"]
              )
              out_dict[model_cfg.cond_stage_key] = torch.tensor(
                  input["clip-vit-large-patch14_text"]
              )
              yield out_dict
  
      def transform_fn(sample):
          return sample["pickle"]
  
      if data_cfg.get("synthetic_data", False):
          H, W = data_cfg.train.augmentations.center_crop_h_w.split(",")
          train_data = SDSyntheticDataset(
              int(H),
              int(W),
              image_key=model_cfg.first_stage_key,
              txt_key=model_cfg.cond_stage_key,
              context_dim=model_cfg.unet_config.context_dim,
          )
      else:
          train_data = WebDatasetCommon(
              dataset_cfg=data_cfg,
              consumed_samples=consumed_samples,
              map_fn=transform_fn,
              compose_fn=tuple_to_dict,
              is_train=True,
          )
  
      val_data = None
      if data_cfg.get(
          "validation"
      ) is not None and data_cfg.validation.get("data_path"):
          if data_cfg.get("synthetic_data", False):
              H, W = data_cfg.train.augmentations.center_crop_h_w.split(
                  ","
              )
              train_data = SDSyntheticDataset(
                  int(H),
                  int(W),
                  image_key=model_cfg.first_stage_key,
                  txt_key=model_cfg.cond_stage_key,
                  context_dim=model_cfg.unet_config.context_dim,
              )
          else:
              val_data = WebDatasetCommon(
                  dataset_cfg=data_cfg,
                  consumed_samples=consumed_samples,
                  map_fn=transform_fn,
                  compose_fn=tuple_to_dict,
                  is_train=False,
              )
  
      return train_data, val_data
  
file_path: NeMo/nemo/collections/multimodal/data/stable_diffusion/stable_diffusion_dataset.py
imports: import torch
  from nemo.collections.multimodal.data.common.webdataset import WebDatasetCommon
  from nemo.collections.multimodal.data.stable_diffusion.augmentation.augmentations import construct_image_augmentations, identical_transform
  from nemo.core.classes import Dataset
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling a global function with example usage: build_train_valid_precached_datasets(model_cfg, consumed_samples) and returns: train_data,  val_data, sample