code: def get_train_dataset(args, image_resolution=512):
      """Build a Dataset for fine-tuning DreamBooth model."""
      # Load a directory of images as a Ray Dataset
      instance_dataset = read_images(args.instance_images_dir)
      class_dataset = read_images(args.class_images_dir)
  
      # We now duplicate the instance images multiple times to make the
      # two sets contain exactly the same number of images.
      # This is so we can zip them up during training to compute the
      # prior preserving loss in one pass.
      #
      # Example: If we have 200 class images (for regularization) and 4 instance
      # images of our subject, then we'll duplicate the instance images 50 times
      # so that our dataset looks like:
      #
      #     instance_image_0, class_image_0
      #     instance_image_1, class_image_1
      #     instance_image_2, class_image_2
      #     instance_image_3, class_image_3
      #     instance_image_0, class_image_4
      #     instance_image_1, class_image_5
      #     ...
      dup_times = class_dataset.count() // instance_dataset.count()
      instance_dataset = instance_dataset.map_batches(
          lambda df: pd.concat([df] * dup_times), batch_format="pandas"
      )
  
      # Load tokenizer for tokenizing the image prompts.
      tokenizer = AutoTokenizer.from_pretrained(
          pretrained_model_name_or_path=args.model_dir,
          subfolder="tokenizer",
      )
  
      def _tokenize(prompt):
          return tokenizer(
              prompt,
              truncation=True,
              padding="max_length",
              max_length=tokenizer.model_max_length,
              return_tensors="pt",
          ).input_ids.numpy()
  
      # Get the token ids for both prompts.
      class_prompt_ids = _tokenize(args.class_prompt)[0]
      instance_prompt_ids = _tokenize(args.instance_prompt)[0]
  
      # START: image preprocessing
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Resize(
                  image_resolution,
                  interpolation=transforms.InterpolationMode.BILINEAR,
                  antialias=True,
              ),
              transforms.RandomCrop(image_resolution),
              transforms.Normalize([0.5], [0.5]),
          ]
      )
  
      def transform_image(
          batch: Dict[str, np.ndarray], output_column_name: str
      ) -> Dict[str, np.ndarray]:
          transformed_tensors = [
              transform(image).numpy() for image in batch["image"]
          ]
          batch[output_column_name] = transformed_tensors
          return batch
  
      # END: image preprocessing
  
      # START: Apply preprocessing steps as Ray Dataset operations
      # For each dataset:
      # - perform image preprocessing
      # - drop the original image column
      # - add a new column with the tokenized prompts
      instance_dataset = (
          instance_dataset.map_batches(
              transform_image,
              fn_kwargs={"output_column_name": "instance_image"},
          )
          .drop_columns(["image"])
          .add_column(
              "instance_prompt_ids",
              lambda df: [instance_prompt_ids] * len(df),
          )
      )
      # END: Apply preprocessing steps as Ray Dataset operations
  
      class_dataset = (
          class_dataset.map_batches(
              transform_image,
              fn_kwargs={"output_column_name": "class_image"},
          )
          .drop_columns(["image"])
          .add_column(
              "class_prompt_ids",
              lambda df: [class_prompt_ids] * len(df),
          )
      )
      # --- Ray Data
  
      # We may have too many duplicates of the instance images, so limit the
      # dataset size so that len(instance_dataset) == len(class_dataset)
      final_size = min(instance_dataset.count(), class_dataset.count())
  
      # Now, zip the images up.
      train_dataset = (
          instance_dataset.limit(final_size)
          .repartition(final_size)
          .zip(class_dataset.limit(final_size).repartition(final_size))
      )
  
      print("Training dataset schema after pre-processing:")
      print(train_dataset.schema())
  
      return train_dataset.random_shuffle()
  
file_path: ray/release/workspace_templates/05_dreambooth_finetuning/dreambooth/dataset.py
imports: from typing import Dict
  import numpy
  import pandas
  import torch
  from ray.data import read_images
  from torchvision import transforms
  from transformers import AutoTokenizer
label_desc: Use image node_id 1 for calling a global function with example usage: get_train_dataset(args, image_resolution) and returns: train_dataset, tokenizer, batch