Skip to content

dataset

Classes

fastvideo.dataset.LTX2PrecomputedDataset 

LTX2PrecomputedDataset(data_root: str, data_sources: dict[str, str] | list[str] | None = None)

Bases: Dataset

Dataset for LTX-2 precomputed latents and conditions.

Expected directory structure (data_root): .precomputed/ latents/.pt conditions/.pt audio_latents/*.pt (optional)

Source code in fastvideo/dataset/ltx2_precomputed_dataset.py 
40
41
42
43
44
45
46
47
48
49
50
def __init__(
    self,
    data_root: str,
    data_sources: dict[str, str] | list[str] | None = None,
) -> None:
    super().__init__()
    self.data_root = self._setup_data_root(data_root)
    self.data_sources = self._normalize_data_sources(data_sources)
    self.source_paths = self._setup_source_paths()
    self.sample_files = self._discover_samples()
    self._validate_setup()

fastvideo.dataset.TextDataset 

TextDataset(data_merge_path: str, args, start_idx: int = 0, seed: int = 42)

Bases: IterableDataset, 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.

Source code in fastvideo/dataset/preprocessing_datasets.py 
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
def __init__(self,
             data_merge_path: str,
             args,
             start_idx: int = 0,
             seed: int = 42):
    self.data_merge_path = data_merge_path
    self.start_idx = start_idx
    self.args = args
    self.seed = seed

    # Initialize tokenizer
    tokenizer_path = os.path.join(args.model_path, "tokenizer")
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
                                              cache_dir=args.cache_dir)

    # Initialize text encoding stage
    self.text_encoding_stage = TextEncodingStage(
        tokenizer=tokenizer,
        text_max_length=args.text_max_length,
        cfg_rate=getattr(args, 'training_cfg_rate', 0.0),
        seed=self.seed)

    # Process text data
    self.processed_batches = self._process_text_data()

Methods:

fastvideo.dataset.TextDataset.__iter__ 
__iter__()

Iterator for the dataset.

Source code in fastvideo/dataset/preprocessing_datasets.py 
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
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)
fastvideo.dataset.TextDataset.load_state_dict 
load_state_dict(state_dict: dict[str, Any]) -> None

Load state dict from checkpoint.

Source code in fastvideo/dataset/preprocessing_datasets.py 
711
712
713
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load state dict from checkpoint."""
    self.processed_batches = state_dict["processed_batches"]
fastvideo.dataset.TextDataset.state_dict 
state_dict() -> dict[str, Any]

Return state dict for checkpointing.

Source code in fastvideo/dataset/preprocessing_datasets.py 
707
708
709
def state_dict(self) -> dict[str, Any]:
    """Return state dict for checkpointing."""
    return {"processed_batches": self.processed_batches}

fastvideo.dataset.ValidationDataset 

ValidationDataset(filename: str)

Bases: IterableDataset

Source code in fastvideo/dataset/validation_dataset.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
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)

Methods:

fastvideo.dataset.ValidationDataset.__len__ 
__len__()

Return the number of samples for this SP group.

Source code in fastvideo/dataset/validation_dataset.py 
107
108
109
def __len__(self):
    """Return the number of samples for this SP group."""
    return len(self.sp_group_samples)

fastvideo.dataset.VideoCaptionMergedDataset 

VideoCaptionMergedDataset(data_merge_path: str, args, transform, temporal_sample, transform_topcrop, start_idx: int = 0, seed: int = 42)

Bases: IterableDataset, 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: ,

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

Source code in fastvideo/dataset/preprocessing_datasets.py 
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
def __init__(self,
             data_merge_path: str,
             args,
             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.args = args
    self.temporal_sample = temporal_sample
    self.seed = seed

    # Initialize tokenizer
    tokenizer_path = os.path.join(args.model_path, "tokenizer")
    tokenizer = None
    if os.path.exists(tokenizer_path):
        try:
            tokenizer = AutoTokenizer.from_pretrained(
                tokenizer_path, cache_dir=args.cache_dir)
        except (ValueError, OSError):
            pass

    # Initialize processing stages
    self._init_stages(args, transform, transform_topcrop, tokenizer)

    # Process metadata
    self.processed_batches = self._process_metadata()

Methods:

fastvideo.dataset.VideoCaptionMergedDataset.__iter__ 
__iter__()

Iterate through processed data items.

Source code in fastvideo/dataset/preprocessing_datasets.py 
537
538
539
540
def __iter__(self):
    """Iterate through processed data items."""
    for idx in range(len(self.processed_batches)):
        yield self._get_item(idx)
fastvideo.dataset.VideoCaptionMergedDataset.load_state_dict 
load_state_dict(state_dict: dict[str, Any]) -> None

Load state dict from checkpoint.

Source code in fastvideo/dataset/preprocessing_datasets.py 
580
581
582
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load state dict from checkpoint."""
    self.processed_batches = state_dict["processed_batches"]
fastvideo.dataset.VideoCaptionMergedDataset.state_dict 
state_dict() -> dict[str, Any]

Return state dict for checkpointing.

Source code in fastvideo/dataset/preprocessing_datasets.py 
576
577
578
def state_dict(self) -> dict[str, Any]:
    """Return state dict for checkpointing."""
    return {"processed_batches": self.processed_batches}