Davids048andClaude Fable 5.1
a5aa64ab6b
VideoGenerator.from_pretrained(model_path, config) takes the same nested
settings as from_config; the 24 flat keywords, from_pretrained_kwargs_to_config,
FROM_PRETRAINED_KWARGS, and every flat_field in the schema are deleted.
nvfp4_fa4 is the typed field engine.attention.nvfp4_fa4, applied by a
resolution step. 65 call sites and the docs use the nested form; the three
kwargs golden cases are config cases with byte-identical results.
Every value is decided during resolution; no runtime code calls
with_override. The device offload policy (unified memory, MPS, layerwise
conflicts, lazy module load) runs as resolution steps in the main process
(fastvideo/api/device_policy.py), and the worker runs on the config it
receives. The LTX-2 refine defaults from model_index.json and the MiniMax-H3
schedule from fastvideo_inference.json are resolution steps
(fastvideo/api/checkpoint_defaults.py) that read local checkpoints or download
only those files; the H3 pipeline validates the schedule against the loaded
schedulers but no longer writes it. The preprocessing entry point passes the
downloaded local path as run state instead of a config override. Teacher and
critic models load with override_transformer_cls_name as a load argument.
The test isolation turns the device policy off, as it blocks downloads, so
that the goldens do not depend on the machine.
Against the 3f6893a0 goldens, every field still matches except boundary_ratio
and ltx2_vae_tiling (DESIGN section 9) and the fields that no longer exist.
The launcher and trainer YAML verifications report 0 unexplained differences.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
676 lines
23 KiB
Python
676 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
from abc import ABC, abstractmethod
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
from os.path import join as opj
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision
|
|
from einops import rearrange
|
|
from PIL import Image
|
|
from transformers import AutoTokenizer
|
|
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class PreprocessBatch:
|
|
"""
|
|
Batch information for dataset processing stages.
|
|
|
|
This class holds all the information about a video-caption or image-caption pair
|
|
as it moves through the processing pipeline. Fields are populated by different stages.
|
|
"""
|
|
# Raw metadata
|
|
path: str
|
|
cap: str | list[str]
|
|
resolution: dict | None = None
|
|
fps: float | None = None
|
|
duration: float | None = None
|
|
|
|
# Processed metadata
|
|
num_frames: int | None = None
|
|
sample_frame_index: list[int] | None = None
|
|
sample_num_frames: int | None = None
|
|
action_path: str | None = None
|
|
|
|
# Processed data
|
|
pixel_values: torch.Tensor | None = None
|
|
text: str | None = None
|
|
input_ids: torch.Tensor | None = None
|
|
cond_mask: torch.Tensor | None = None
|
|
|
|
@property
|
|
def is_video(self) -> bool:
|
|
"""Check if this is a video item."""
|
|
return self.path.endswith(".mp4")
|
|
|
|
@property
|
|
def is_image(self) -> bool:
|
|
"""Check if this is an image item."""
|
|
return self.path.endswith(".jpg")
|
|
|
|
|
|
class DatasetStage(ABC):
|
|
"""
|
|
Abstract base class for dataset processing stages.
|
|
|
|
Similar to PipelineStage but designed for dataset preprocessing operations.
|
|
"""
|
|
|
|
@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
|
|
|
|
|
|
class DatasetFilterStage(ABC):
|
|
"""
|
|
Abstract base class for dataset filtering stages.
|
|
|
|
These stages can filter out items during metadata processing.
|
|
"""
|
|
|
|
@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 (for non-filtering operations).
|
|
|
|
Args:
|
|
batch: Dataset batch to process
|
|
**kwargs: Additional processing parameters
|
|
|
|
Returns:
|
|
Processed batch
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class DataValidationStage(DatasetFilterStage):
|
|
"""Stage for validating data items."""
|
|
|
|
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
|
|
|
|
def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
|
|
"""Process does nothing for validation - filtering is handled by should_keep."""
|
|
return batch
|
|
|
|
|
|
class FrameSamplingStage(DatasetFilterStage):
|
|
"""Stage for temporal frame sampling and indexing."""
|
|
|
|
def __init__(self,
|
|
num_frames: int,
|
|
train_fps: int,
|
|
speed_factor: int = 1,
|
|
video_length_tolerance_range: float = 5.0,
|
|
drop_short_ratio: float = 0.0,
|
|
seed: int = 42):
|
|
self.num_frames = num_frames
|
|
self.train_fps = train_fps
|
|
self.speed_factor = speed_factor
|
|
self.video_length_tolerance_range = video_length_tolerance_range
|
|
self.drop_short_ratio = drop_short_ratio
|
|
# Create a seeded random generator for deterministic sampling
|
|
self.rng = random.Random(seed)
|
|
|
|
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, 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
|
|
|
|
|
|
class VideoTransformStage(DatasetStage):
|
|
"""Stage for video data transformation."""
|
|
|
|
def __init__(self, transform) -> None:
|
|
self.transform = transform
|
|
|
|
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
|
|
|
|
|
|
class ImageTransformStage(DatasetStage):
|
|
"""Stage for image data transformation."""
|
|
|
|
def __init__(self, transform, transform_topcrop) -> None:
|
|
self.transform = transform
|
|
self.transform_topcrop = transform_topcrop
|
|
|
|
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
|
|
|
|
|
|
class TextEncodingStage(DatasetStage):
|
|
"""Stage for text tokenization and encoding."""
|
|
|
|
def __init__(self, tokenizer, text_max_length: int, cfg_rate: float = 0.0, seed: int = 42):
|
|
self.tokenizer = tokenizer
|
|
self.text_max_length = text_max_length
|
|
self.cfg_rate = cfg_rate
|
|
# Create a seeded random generator for deterministic CFG
|
|
self.rng = random.Random(seed)
|
|
|
|
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
|
|
|
|
|
|
class VideoCaptionMergedDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint.stateful.Stateful):
|
|
"""
|
|
Merged dataset for video and caption data with stage-based processing.
|
|
Assumes that data_merge_path is a txt file with the following format:
|
|
<folder_path>,<json_file_path>
|
|
|
|
The folder should contain videos.
|
|
|
|
The json file should be a list of dictionaries with the following format:
|
|
[
|
|
{
|
|
"path": "1gGQy4nxyUo-Scene-016.mp4",
|
|
"resolution": {
|
|
"width": 1920,
|
|
"height": 1080
|
|
},
|
|
"size": 2439112,
|
|
"fps": 25.0,
|
|
"duration": 6.88,
|
|
"num_frames": 172,
|
|
"cap": [
|
|
"A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open."
|
|
]
|
|
},
|
|
...
|
|
]
|
|
|
|
This dataset processes video and image data through a series of stages:
|
|
- Data validation
|
|
- Resolution filtering
|
|
- Frame sampling
|
|
- Transformation
|
|
- Text encoding
|
|
"""
|
|
|
|
def __init__(self,
|
|
data_merge_path: str,
|
|
preprocess_config,
|
|
model_path: str,
|
|
transform,
|
|
temporal_sample,
|
|
transform_topcrop,
|
|
start_idx: int = 0,
|
|
seed: int = 42):
|
|
self.data_merge_path = data_merge_path
|
|
self.start_idx = start_idx
|
|
self.temporal_sample = temporal_sample
|
|
self.seed = seed
|
|
|
|
# Initialize tokenizer from the local checkpoint directory
|
|
tokenizer_path = os.path.join(model_path, "tokenizer")
|
|
tokenizer = None
|
|
if os.path.exists(tokenizer_path):
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, cache_dir=preprocess_config.cache_dir)
|
|
except (ValueError, OSError):
|
|
pass
|
|
|
|
# Initialize processing stages
|
|
self._init_stages(preprocess_config, transform, transform_topcrop, tokenizer)
|
|
|
|
# Process metadata
|
|
self.processed_batches = self._process_metadata()
|
|
|
|
def _init_stages(self, preprocess_config, transform, transform_topcrop, tokenizer) -> None:
|
|
"""Initialize all processing stages."""
|
|
self.validation_stage = DataValidationStage()
|
|
self.frame_sampling_stage = FrameSamplingStage(
|
|
num_frames=preprocess_config.num_frames,
|
|
train_fps=preprocess_config.train_fps,
|
|
speed_factor=preprocess_config.speed_factor,
|
|
video_length_tolerance_range=preprocess_config.video_length_tolerance_range,
|
|
drop_short_ratio=preprocess_config.drop_short_ratio,
|
|
seed=self.seed)
|
|
self.video_transform_stage = VideoTransformStage(transform)
|
|
self.image_transform_stage = ImageTransformStage(transform, transform_topcrop)
|
|
if tokenizer is not None:
|
|
self.text_encoding_stage = TextEncodingStage(tokenizer=tokenizer,
|
|
text_max_length=preprocess_config.text_max_length,
|
|
cfg_rate=preprocess_config.training_cfg_rate,
|
|
seed=self.seed)
|
|
else:
|
|
self.text_encoding_stage = None
|
|
|
|
def _load_raw_data(self) -> list[dict]:
|
|
"""Load raw data from JSON files."""
|
|
# Read folder-annotation pairs
|
|
with open(self.data_merge_path) as f:
|
|
folder_anno_pairs = [line.strip().split(",") for line in f if line.strip()]
|
|
assert len(folder_anno_pairs) == 1, "Only support one folder-annotation pair"
|
|
assert len(folder_anno_pairs[0]) == 2, "Folder-annotation pair should have two elements"
|
|
folder, annotation_file = folder_anno_pairs[0]
|
|
|
|
data_items: list[dict] = []
|
|
with open(annotation_file) as f:
|
|
data_items = json.load(f)
|
|
|
|
# Update paths with folder prefix
|
|
for item in data_items:
|
|
item["path"] = opj(folder, item["path"])
|
|
if "action_path" in item and item["action_path"]:
|
|
item["action_path"] = opj(folder, item["action_path"])
|
|
|
|
return data_items
|
|
|
|
def _process_metadata(self) -> list[PreprocessBatch]:
|
|
"""Process the raw metadata through all filtering stages."""
|
|
raw_data = self._load_raw_data()
|
|
processed_batches = []
|
|
|
|
# Initialize counters
|
|
filter_counts = {"validation_failed": 0, "frame_sampling_failed": 0}
|
|
sample_num_frames: list[int] = []
|
|
|
|
for item in raw_data:
|
|
batch = PreprocessBatch(path=item["path"],
|
|
cap=item["cap"],
|
|
resolution=item.get("resolution"),
|
|
fps=item.get("fps"),
|
|
duration=item.get("duration"),
|
|
action_path=item.get("action_path"))
|
|
|
|
# Apply filtering stages
|
|
if not self._apply_filter_stages(batch, filter_counts):
|
|
continue
|
|
|
|
# Apply frame sampling processing
|
|
batch = self.frame_sampling_stage.process(batch, temporal_sample_fn=self.temporal_sample)
|
|
|
|
processed_batches.append(batch)
|
|
assert batch.sample_num_frames is not None
|
|
sample_num_frames.append(batch.sample_num_frames)
|
|
|
|
self._log_filtering_stats(filter_counts, sample_num_frames, len(raw_data), len(processed_batches))
|
|
return processed_batches
|
|
|
|
def _apply_filter_stages(self, batch: PreprocessBatch, filter_counts: dict[str, int]) -> bool:
|
|
"""Apply all filter stages and update counters. Returns True if batch should be kept."""
|
|
if not self.validation_stage.should_keep(batch):
|
|
filter_counts["validation_failed"] += 1
|
|
return False
|
|
|
|
if not self.frame_sampling_stage.should_keep(batch):
|
|
filter_counts["frame_sampling_failed"] += 1
|
|
return False
|
|
|
|
return True
|
|
|
|
def _log_filtering_stats(self, filter_counts: dict[str, int], sample_num_frames: list[int], before_count: int,
|
|
after_count: int):
|
|
"""Log filtering statistics."""
|
|
logger.info(
|
|
"validation_failed: %d, frame_sampling_failed: %d, "
|
|
"Counter(sample_num_frames): %s, before filter: %d, after filter: %d", filter_counts['validation_failed'],
|
|
filter_counts['frame_sampling_failed'], Counter(sample_num_frames), before_count, after_count)
|
|
|
|
def __iter__(self):
|
|
"""Iterate through processed data items."""
|
|
for idx in range(len(self.processed_batches)):
|
|
yield self._get_item(idx)
|
|
|
|
def __len__(self):
|
|
return len(self.processed_batches)
|
|
|
|
def _get_item(self, idx: int) -> dict:
|
|
"""Get a single processed data item."""
|
|
batch = self.processed_batches[idx]
|
|
|
|
# Apply transformation stages
|
|
batch = self.video_transform_stage.process(batch)
|
|
batch = self.image_transform_stage.process(batch)
|
|
if self.text_encoding_stage is not None:
|
|
batch = self.text_encoding_stage.process(batch)
|
|
|
|
# Build result dictionary
|
|
result = {
|
|
"pixel_values": batch.pixel_values,
|
|
"path": batch.path,
|
|
}
|
|
|
|
if batch.text is not None:
|
|
result["text"] = batch.text
|
|
result["input_ids"] = batch.input_ids
|
|
result["cond_mask"] = batch.cond_mask
|
|
|
|
# Add video-specific fields
|
|
if batch.is_video:
|
|
result.update({"fps": batch.fps, "duration": batch.duration})
|
|
|
|
# Add action_path
|
|
if batch.action_path:
|
|
result["action_path"] = batch.action_path
|
|
|
|
return result
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
"""Return state dict for checkpointing."""
|
|
return {"processed_batches": self.processed_batches}
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
"""Load state dict from checkpoint."""
|
|
self.processed_batches = state_dict["processed_batches"]
|
|
|
|
|
|
class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint.stateful.Stateful):
|
|
"""
|
|
Text-only dataset for processing prompts from a simple text file.
|
|
|
|
Assumes that data_merge_path is a text file with one prompt per line:
|
|
A cat playing with a ball
|
|
A dog running in the park
|
|
A person cooking dinner
|
|
...
|
|
|
|
This dataset processes text data through text encoding stages only.
|
|
"""
|
|
|
|
def __init__(self, data_merge_path: str, preprocess_config, model_path: str, start_idx: int = 0, seed: int = 42):
|
|
self.data_merge_path = data_merge_path
|
|
self.start_idx = start_idx
|
|
self.seed = seed
|
|
|
|
# Initialize tokenizer from the local checkpoint directory
|
|
tokenizer_path = os.path.join(model_path, "tokenizer")
|
|
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, cache_dir=preprocess_config.cache_dir)
|
|
|
|
# Initialize text encoding stage
|
|
self.text_encoding_stage = TextEncodingStage(tokenizer=tokenizer,
|
|
text_max_length=preprocess_config.text_max_length,
|
|
cfg_rate=preprocess_config.training_cfg_rate,
|
|
seed=self.seed)
|
|
|
|
# Process text data
|
|
self.processed_batches = self._process_text_data()
|
|
|
|
def _load_text_data(self) -> list[str]:
|
|
"""Load text prompts from file."""
|
|
prompts = []
|
|
with open(self.data_merge_path, 'r', encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line: # Skip empty lines
|
|
prompts.append(line)
|
|
|
|
logger.info(f"Loaded {len(prompts)} text prompts from {self.data_merge_path}")
|
|
return prompts
|
|
|
|
def _process_text_data(self) -> list[PreprocessBatch]:
|
|
"""Process the text prompts through text encoding stage."""
|
|
raw_prompts = self._load_text_data()
|
|
processed_batches = []
|
|
|
|
for idx, prompt in enumerate(raw_prompts):
|
|
# Create a text-only batch with dummy path
|
|
batch = PreprocessBatch(
|
|
path=f"text_prompt_{idx}",
|
|
cap=[prompt], # TextEncodingStage expects a list
|
|
resolution=None,
|
|
fps=None,
|
|
duration=None,
|
|
num_frames=0,
|
|
sample_frame_index=None,
|
|
sample_num_frames=0)
|
|
|
|
processed_batches.append(batch)
|
|
|
|
logger.info(f"Processed {len(processed_batches)} text batches")
|
|
return processed_batches
|
|
|
|
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 _get_item(self, idx: int) -> dict:
|
|
"""Get a single processed text item."""
|
|
batch = self.processed_batches[idx]
|
|
|
|
# Apply text encoding stage
|
|
batch = self.text_encoding_stage.process(batch)
|
|
|
|
# Build result dictionary for text-only processing with required schema fields
|
|
result = {
|
|
"text": batch.text,
|
|
"input_ids": batch.input_ids,
|
|
"cond_mask": batch.cond_mask,
|
|
"path": batch.path,
|
|
# Required schema fields for ODE trajectory processing
|
|
"id": f"text_{idx}",
|
|
"file_name": batch.path,
|
|
"caption": batch.text,
|
|
"media_type": "text",
|
|
"width": 1,
|
|
"height": 1,
|
|
"num_frames": 0,
|
|
"duration_sec": 0.0,
|
|
"fps": 0.0,
|
|
}
|
|
|
|
return result
|
|
|
|
def state_dict(self) -> dict[str, Any]:
|
|
"""Return state dict for checkpointing."""
|
|
return {"processed_batches": self.processed_batches}
|
|
|
|
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
|
|
"""Load state dict from checkpoint."""
|
|
self.processed_batches = state_dict["processed_batches"]
|