code: def build_train_valid_datasets(
      model_cfg,
      consumed_samples,
  ):
      data_cfg = model_cfg.data
  
      def build_resolution_filter(value=None, method="larger"):
          assert method == "larger" or method == "smaller"
          if method == "larger":
              logging.info(
                  f"Only Selecting images with resolution >= {value}"
              )
              return (
                  lambda x: x["jpg"].size[0] >= value
                  and x["jpg"].size[1] >= value
              )
          logging.info(
              f"Only Selecting images with resolution <= {value}"
          )
          return (
              lambda x: x["jpg"].size[0] <= value
              and x["jpg"].size[1] <= value
          )
  
      # 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] = input[0].permute(
                  1, 2, 0
              )
              out_dict[model_cfg.cond_stage_key] = input[1]
              yield out_dict
  
      def transform_fn(sample):
          image, text = sample["jpg"], sample["txt"]
          # TODO : If no agumentations just return the image ?
          img_transform = construct_image_augmentations(
              data_cfg.train.get("augmentations", None)
          )
          text_transform = identical_transform
          return img_transform(image), text_transform(text)
  
      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,
              fake_len=data_cfg.synthetic_data_length,
          )
  
      else:
          filter_cfg = data_cfg.train.get("filterings", None)
          filter_fn = (
              build_resolution_filter(**filter_cfg.resolution)
              if filter_cfg
              else None
          )
          train_data = WebDatasetCommon(
              dataset_cfg=data_cfg,
              consumed_samples=consumed_samples,
              map_fn=transform_fn,
              compose_fn=tuple_to_dict,
              filter_fn=filter_fn,
              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):
              val_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,
                  filter_fn=filter_fn,
                  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 1 for calling a global function with example usage: build_train_valid_datasets(model_cfg, consumed_samples) and returns: train_data,  val_data, function, img_transform,  text_transform, function