From 982bfcfdc8e51025adfd66f2fc1f8dcdded97d15 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Wed, 21 May 2025 22:46:32 +0000 Subject: [PATCH] Fix small issues when loading mask --- fastvideo/v1/dataset/parquet_datasets.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/fastvideo/v1/dataset/parquet_datasets.py b/fastvideo/v1/dataset/parquet_datasets.py index e2fa10caf..db12f25de 100644 --- a/fastvideo/v1/dataset/parquet_datasets.py +++ b/fastvideo/v1/dataset/parquet_datasets.py @@ -24,7 +24,7 @@ dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x51 class ParquetVideoTextDataset(IterableDataset): """Efficient loader for video-text data from a directory of Parquet files.""" - def __init__(self, path: str, batch_size: int = 1024, rank: int = 0, world_size: int = 1, cfg_rate: float = 0.1, num_workers: int = 1, row_per_parquet: int = 32, split: str = "train"): + def __init__(self, path: str, batch_size: int = 1024, rank: int = 0, world_size: int = 1, cfg_rate: float = 0.0, num_workers: int = 1, row_per_parquet: int = 32, split: str = "train"): super().__init__() self.path = str(path) self.batch_size = batch_size @@ -170,7 +170,7 @@ class ParquetVideoTextDataset(IterableDataset): # Process mask if len(text_attention_mask_bytes) > 0 and len(text_attention_mask_shape) > 0: - msk = np.frombuffer(text_attention_mask_bytes, dtype=np.uint8).astype(np.bool_) + msk = np.frombuffer(text_attention_mask_bytes, dtype=np.int64).astype(np.bool_) msk = msk.reshape(1, -1) # Make array writable msk = np.copy(msk) @@ -257,7 +257,6 @@ if __name__ == "__main__": rank=rank, world_size=world_size, num_workers=1, - cfg_rate=0.1, split="train", ) @@ -276,6 +275,7 @@ if __name__ == "__main__": if rank == 0: print("Warming up...") for i, (latents, embeddings, masks, infos) in enumerate(dataloader): + assert torch.sum(masks[0]).item() == torch.count_nonzero(embeddings[0]).item() // 4096 if args.vae_debug: from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.configs.models.vaes import WanVAEConfig