From 7a2e171f1bec8ef116ce48dcfc23e35c1047fa64 Mon Sep 17 00:00:00 2001 From: JerryZhou54 Date: Fri, 23 May 2025 19:01:59 +0000 Subject: [PATCH] Small fix --- fastvideo/data_preprocess/preprocess.py | 2 +- fastvideo/v1/dataset/parquet_datasets.py | 5 ++++- .../pipelines/{data_preprocess.py => preprocess_pipeline.py} | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) rename fastvideo/v1/pipelines/{data_preprocess.py => preprocess_pipeline.py} (99%) diff --git a/fastvideo/data_preprocess/preprocess.py b/fastvideo/data_preprocess/preprocess.py index 6b160bc75..6d5aaeaeb 100644 --- a/fastvideo/data_preprocess/preprocess.py +++ b/fastvideo/data_preprocess/preprocess.py @@ -11,7 +11,7 @@ from fastvideo.v1.distributed import init_distributed_environment, initialize_mo from fastvideo.v1.fastvideo_args import FastVideoArgs from fastvideo.v1.configs.models.vaes import WanVAEConfig from fastvideo import PipelineConfig -from fastvideo.v1.pipelines.data_preprocess import PreprocessPipeline +from fastvideo.v1.pipelines.preprocess_pipeline import PreprocessPipeline logger = init_logger(__name__) diff --git a/fastvideo/v1/dataset/parquet_datasets.py b/fastvideo/v1/dataset/parquet_datasets.py index 7568a8ce1..a3c3803b8 100644 --- a/fastvideo/v1/dataset/parquet_datasets.py +++ b/fastvideo/v1/dataset/parquet_datasets.py @@ -133,7 +133,10 @@ class ParquetVideoTextDataset(IterableDataset): # Yield each item in the batch for lat, emb, mask, info in zip(processed["latents"], processed["embeddings"], processed["masks"], processed["info"]): - yield lat[:, -self.num_latent_t:], emb, mask, info + if lat.numel() == 0: # Split is validation + yield lat, emb, mask, info + else: + yield lat[:, -self.num_latent_t:], emb, mask, info except StopIteration: # Current file is exhausted, try next file diff --git a/fastvideo/v1/pipelines/data_preprocess.py b/fastvideo/v1/pipelines/preprocess_pipeline.py similarity index 99% rename from fastvideo/v1/pipelines/data_preprocess.py rename to fastvideo/v1/pipelines/preprocess_pipeline.py index 418702713..792e243ab 100644 --- a/fastvideo/v1/pipelines/data_preprocess.py +++ b/fastvideo/v1/pipelines/preprocess_pipeline.py @@ -315,7 +315,7 @@ class PreprocessPipeline(ComposedPipelineBase): seq_len = prompt_attention_mask.sum().item() # Slice the embeddings to keep only the non-padding parts text_embedding = prompt_embeds[0, :seq_len].cpu().numpy() - text_attention_mask = prompt_attention_mask[0, :seq_len].cpu().numpy() + text_attention_mask = prompt_attention_mask[0, :seq_len].cpu().numpy().astype(np.uint8) # Log the shapes after removing padding logger.info(f"Shape after removing padding - Embeddings: {text_embedding.shape}, Mask: {text_attention_mask.shape}")