preprocessing_datasets ¶

def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
    """Process does nothing for validation - filtering is handled by should_keep."""
    return batch

def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
    """
    Validate data item.

    Args:
        batch: Dataset batch to validate

    Returns:
        True if valid, False if invalid
    """
    # Check for caption
    if batch.cap is None:
        return False

    if batch.is_video:
        # Validate video-specific fields
        if batch.duration is None or batch.fps is None:
            return False
    elif not batch.is_image:
        return False

    return True

@abstractmethod
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
    """
    Process the dataset batch (for non-filtering operations).

    Args:
        batch: Dataset batch to process
        **kwargs: Additional processing parameters

    Returns:
        Processed batch
    """
    raise NotImplementedError

@abstractmethod
def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
    """
    Check if batch should be kept.

    Args:
        batch: Dataset batch to check
        **kwargs: Additional parameters

    Returns:
        True if batch should be kept, False otherwise
    """
    raise NotImplementedError

@abstractmethod
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
    """
    Process the dataset batch.

    Args:
        batch: Dataset batch to process
        **kwargs: Additional processing parameters

    Returns:
        Processed batch
    """
    raise NotImplementedError

def process(self,
            batch: PreprocessBatch,
            temporal_sample_fn=None,
            **kwargs) -> PreprocessBatch:
    """
    Process frame sampling for video data items.

    Args:
        batch: Dataset batch
        temporal_sample_fn: Function for temporal sampling

    Returns:
        Updated batch with frame sampling info
    """
    if batch.is_image:
        # For images, just add sample info
        batch.sample_frame_index = [0]
        batch.sample_num_frames = 1
        return batch

    assert batch.duration is not None and batch.fps is not None
    batch.num_frames = math.ceil(batch.fps * batch.duration)

    # Resample frame indices
    frame_interval = batch.fps / self.train_fps
    start_frame_idx = 0
    frame_indices = np.arange(start_frame_idx, batch.num_frames,
                              frame_interval).astype(int)

    # Temporal crop if too long
    if len(frame_indices) > self.num_frames:
        if temporal_sample_fn is not None:
            begin_index, end_index = temporal_sample_fn(len(frame_indices))
            frame_indices = frame_indices[begin_index:end_index]
        else:
            frame_indices = frame_indices[:self.num_frames]

    batch.sample_frame_index = frame_indices.tolist()
    batch.sample_num_frames = len(frame_indices)

    return batch

def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
    """
    Check if video should be kept based on length constraints.

    Args:
        batch: Dataset batch

    Returns:
        True if should be kept, False otherwise
    """
    if batch.is_image:
        return True

    if batch.duration is None or batch.fps is None:
        return False

    num_frames = math.ceil(batch.fps * batch.duration)

    # Check if video is too long
    if (num_frames / batch.fps > self.video_length_tolerance_range *
        (self.num_frames / self.train_fps * self.speed_factor)):
        return False

    # Resample frame indices to check length
    frame_interval = batch.fps / self.train_fps
    start_frame_idx = 0
    frame_indices = np.arange(start_frame_idx, num_frames,
                              frame_interval).astype(int)

    # Filter short videos
    return not (len(frame_indices) < self.num_frames
                and self.rng.random() < self.drop_short_ratio)

def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
    """
    Transform image data.

    Args:
        batch: Dataset batch with image information

    Returns:
        Batch with transformed image tensor
    """
    if not batch.is_image:
        return batch

    image = Image.open(batch.path).convert("RGB")
    image = torch.from_numpy(np.array(image))
    image = rearrange(image, "h w c -> c h w").unsqueeze(0)

    if self.transform_topcrop is not None:
        image = self.transform_topcrop(image)
    elif self.transform is not None:
        image = self.transform(image)
        image = image.float() / 127.5 - 1.0
    else:
        image = image.float() / 127.5 - 1.0

    image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
    batch.pixel_values = image
    return batch

def __iter__(self):
    """Iterator for the dataset."""
    # Set up distributed sampling if needed
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        rank = torch.distributed.get_rank()
        world_size = torch.distributed.get_world_size()
    else:
        rank = 0
        world_size = 1

    # Calculate chunk for this rank
    total_items = len(self.processed_batches)
    items_per_rank = math.ceil(total_items / world_size)
    start_idx = rank * items_per_rank + self.start_idx
    end_idx = min(start_idx + items_per_rank, total_items)

    # Yield items for this rank
    for idx in range(start_idx, end_idx):
        if idx < len(self.processed_batches):
            yield self._get_item(idx)

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load state dict from checkpoint."""
    self.processed_batches = state_dict["processed_batches"]

def state_dict(self) -> dict[str, Any]:
    """Return state dict for checkpointing."""
    return {"processed_batches": self.processed_batches}

def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
    """
    Process text data.

    Args:
        batch: Dataset batch with caption information

    Returns:
        Batch with encoded text information
    """
    text = batch.cap
    if not isinstance(text, list):
        text = [text]
    text = [self.rng.choice(text)]

    text = text[0] if self.rng.random() > self.cfg_rate else ""
    text_tokens_and_mask = self.tokenizer(
        text,
        max_length=self.text_max_length,
        padding="max_length",
        truncation=True,
        return_attention_mask=True,
        add_special_tokens=True,
        return_tensors="pt",
    )

    batch.text = text
    batch.input_ids = text_tokens_and_mask["input_ids"]
    batch.cond_mask = text_tokens_and_mask["attention_mask"]
    return batch

def __iter__(self):
    """Iterate through processed data items."""
    for idx in range(len(self.processed_batches)):
        yield self._get_item(idx)

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load state dict from checkpoint."""
    self.processed_batches = state_dict["processed_batches"]

def state_dict(self) -> dict[str, Any]:
    """Return state dict for checkpointing."""
    return {"processed_batches": self.processed_batches}

def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
    """
    Transform video data.

    Args:
        batch: Dataset batch with video information

    Returns:
        Batch with transformed video tensor
    """
    if not batch.is_video:
        return batch

    assert os.path.exists(batch.path), f"file {batch.path} do not exist!"
    assert batch.sample_frame_index is not None, "Frame indices must be set before transformation"

    torchvision_video, _, metadata = torchvision.io.read_video(
        batch.path, output_format="TCHW")
    video = torchvision_video[batch.sample_frame_index]
    if self.transform is not None:
        video = self.transform(video)
    video = rearrange(video, "t c h w -> c t h w")
    video = video.to(torch.uint8)

    video = video.float() / 127.5 - 1.0
    batch.pixel_values = video
    return batch