# SPDX-License-Identifier: Apache-2.0 # adapted from: https://github.com/a-r-r-o-w/finetrainers/blob/main/finetrainers/data/dataset.py import os import pathlib import datasets import numpy as np from torch.utils.data import IterableDataset from fastvideo.distributed import (get_sp_world_size, get_world_rank, get_world_size) from fastvideo.logger import init_logger from fastvideo.models.vision_utils import load_image, load_video logger = init_logger(__name__) class ValidationDataset(IterableDataset): def __init__(self, filename: str): super().__init__() self.filename = pathlib.Path(filename) # get directory of filename self.dir = os.path.abspath(self.filename.parent) if not self.filename.exists(): raise FileNotFoundError( f"File {self.filename.as_posix()} does not exist") if self.filename.suffix == ".csv": data = datasets.load_dataset("csv", data_files=self.filename.as_posix(), split="train") elif self.filename.suffix == ".json": data = datasets.load_dataset("json", data_files=self.filename.as_posix(), split="train", field="data") elif self.filename.suffix == ".parquet": data = datasets.load_dataset("parquet", data_files=self.filename.as_posix(), split="train") elif self.filename.suffix == ".arrow": data = datasets.load_dataset("arrow", data_files=self.filename.as_posix(), split="train") else: _SUPPORTED_FILE_FORMATS = [".csv", ".json", ".parquet", ".arrow"] raise ValueError( f"Unsupported file format {self.filename.suffix} for validation dataset. Supported formats are: {_SUPPORTED_FILE_FORMATS}" ) # Get distributed training info self.global_rank = get_world_rank() self.world_size = get_world_size() self.sp_world_size = get_sp_world_size() self.num_sp_groups = self.world_size // self.sp_world_size # Convert to list to get total samples self.all_samples = list(data) self.original_total_samples = len(self.all_samples) # Extend samples to be a multiple of DP degree (num_sp_groups) remainder = self.original_total_samples % self.num_sp_groups if remainder != 0: samples_to_add = self.num_sp_groups - remainder # Duplicate samples cyclically to reach the target additional_samples = [] for i in range(samples_to_add): additional_samples.append( self.all_samples[i % self.original_total_samples]) self.all_samples.extend(additional_samples) self.total_samples = len(self.all_samples) # Calculate which SP group this rank belongs to self.sp_group_id = self.global_rank // self.sp_world_size # Now all SP groups will have equal number of samples self.samples_per_sp_group = self.total_samples // self.num_sp_groups # Calculate start and end indices for this SP group self.start_idx = self.sp_group_id * self.samples_per_sp_group self.end_idx = self.start_idx + self.samples_per_sp_group # Get samples for this SP group self.sp_group_samples = self.all_samples[self.start_idx:self.end_idx] logger.info( "Rank %s (SP group %s): " "Original samples: %s, " "Extended samples: %s, " "SP group samples: %s, " "Range: [%s:%s]", self.global_rank, self.sp_group_id, self.original_total_samples, self.total_samples, len(self.sp_group_samples), self.start_idx, self.end_idx, local_main_process_only=False) def __len__(self): """Return the number of samples for this SP group.""" return len(self.sp_group_samples) def __iter__(self): for sample in self.sp_group_samples: # For consistency reasons, we mandate that "caption" is always present in the validation dataset. # However, since the model specifications use "prompt", we create an alias here. sample["prompt"] = sample["caption"] # Load image or video if the path is provided # TODO(aryan): need to handle custom columns here for control conditions sample["image"] = None sample["video"] = None if sample.get("image_path", None) is not None: image_path = sample["image_path"] image_path = os.path.join(self.dir, image_path) sample["image_path"] = image_path if not pathlib.Path(image_path).is_file( ) and not image_path.startswith("http"): logger.warning("Image file %s does not exist.", image_path) else: sample["image"] = load_image(image_path) if sample.get("video_path", None) is not None: video_path = sample["video_path"] video_path = os.path.join(self.dir, video_path) sample["video_path"] = video_path if not pathlib.Path(video_path).is_file( ) and not video_path.startswith("http"): logger.warning("Video file %s does not exist.", video_path) else: sample["video"] = load_video(video_path) if sample.get("control_image_path", None) is not None: control_image_path = sample["control_image_path"] control_image_path = os.path.join(self.dir, control_image_path) sample["control_image_path"] = control_image_path if not pathlib.Path(control_image_path).is_file( ) and not control_image_path.startswith("http"): logger.warning("Control Image file %s does not exist.", control_image_path) else: sample["control_image"] = load_image(control_image_path) if sample.get("control_video_path", None) is not None: control_video_path = sample["control_video_path"] control_video_path = os.path.join(self.dir, control_video_path) sample["control_video_path"] = control_video_path if not pathlib.Path(control_video_path).is_file( ) and not control_video_path.startswith("http"): logger.warning("Control Video file %s does not exist.", control_video_path) else: sample["control_video"] = load_video(control_video_path) if sample.get("action_path", None) is not None: action_path = sample["action_path"] action_path = os.path.join(self.dir, action_path) sample["action_path"] = action_path if not pathlib.Path(action_path).is_file(): logger.warning("Action file %s does not exist.", action_path) else: action_data = np.load(action_path, allow_pickle=True) num_frames = sample["num_frames"] if isinstance(action_data, np.ndarray) and action_data.dtype == object: action_data = action_data.item() if isinstance(action_data, dict): kb_length = len(action_data["keyboard"]) ms_length = len(action_data["mouse"]) assert kb_length >= num_frames, ( f"keyboard length {kb_length} is smaller than " f"num_frames {num_frames} for {action_path}.") assert ms_length >= num_frames, ( f"mouse length {ms_length} is smaller than " f"num_frames {num_frames} for {action_path}.") sample["keyboard_cond"] = action_data["keyboard"][:num_frames] sample["mouse_cond"] = action_data["mouse"][:num_frames] else: kb_length = len(action_data) assert kb_length >= num_frames, ( f"keyboard length {kb_length} is smaller than " f"num_frames {num_frames} for {action_path}.") sample["keyboard_cond"] = action_data[:num_frames] sample = {k: v for k, v in sample.items() if v is not None} yield sample