Small fix

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__)
 

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

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}")