preprocessing text

fastvideo/dataset/__init__.py

@@ -4,7 +4,7 @@ from torchvision.transforms import Lambda
 
 from fastvideo.dataset.parquet_dataset_map_style import (
     build_parquet_map_style_dataloader)
-from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset
+from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                          TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
@@ -39,7 +39,13 @@ def getdataset(args) -> VideoCaptionMergedDataset:
                                      seed=args.seed)
 
 
+def gettextdataset(args) -> TextDataset:
+    return TextDataset(data_merge_path=args.data_merge_path,
+                       args=args,
+                       seed=args.seed)
+
+
 __all__ = [
     "build_parquet_map_style_dataloader", "ValidationDataset",
-    "VideoCaptionMergedDataset"
-]
+    "VideoCaptionMergedDataset", "TextDataset"
+]

fastvideo/dataset/dataloader/schema.py

@@ -78,3 +78,76 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
 ])
+
+pyarrow_schema_ode_trajectory = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    # I2V
+    pa.field("image_condition_latents_bytes", pa.binary()),
+    pa.field("image_condition_latents_shape", pa.list_(pa.int64())),
+    pa.field("image_condition_latents_dtype", pa.string()),
+    # --- ODE Trajectory ---
+    pa.field("trajectory_latents_bytes", pa.binary()),
+    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_latents_dtype", pa.string()),
+    pa.field("trajectory_timesteps_bytes", pa.binary()),
+    pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_timesteps_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])
+
+pyarrow_schema_ode_trajectory_text_only = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    # --- ODE Trajectory ---
+    pa.field("trajectory_latents_bytes", pa.binary()),
+    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_latents_dtype", pa.string()),
+    pa.field("trajectory_timesteps_bytes", pa.binary()),
+    pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_timesteps_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # Always 'text' for text-only
+])
+
+pyarrow_schema_text_only = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+])

fastvideo/dataset/preprocessing_datasets.py

@@ -628,3 +628,134 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
     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,
+                 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()
+
+    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"]

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -0,0 +1,762 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+ODE Trajectory Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline
+using the modular pipeline architecture.
+
+Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772
+"""
+
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import getdataset, gettextdataset
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory, pyarrow_schema_ode_trajectory_text_only
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.utils import shallow_asdict, save_decoded_latents_as_video
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import (DenoisingStage, ImageVAEEncodingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage,
+                                        DecodingStage)
+
+logger = init_logger(__name__)
+
+
+class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
+    """ODE Trajectory preprocessing pipeline implementation."""
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    preprocess_dataloader: StatefulDataLoader
+    preprocess_loader_iter: Iterator[dict[str, Any]]
+
+    def get_schema_fields(self):
+        """Get the schema fields for ODE Trajectory pipeline."""
+        # Check if we're using text dataset by checking if the dataset is TextDataset
+        if hasattr(self, 'preprocess_dataloader') and hasattr(self.preprocess_dataloader.dataset, '_process_text_data'):
+            return [f.name for f in pyarrow_schema_ode_trajectory_text_only]
+        return [f.name for f in pyarrow_schema_ode_trajectory]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+        self.add_stage(stage_name="vae_encoding_stage",
+                       stage=ImageVAEEncodingStage(
+                           vae=self.get_module("vae"), ))
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self,
+                       ))
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    def preprocess_video_and_text_and_trajectory(self,
+                                                 fastvideo_args: FastVideoArgs,
+                                                 args):
+
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # Filter out invalid samples (those with all zeros)
+                valid_indices = []
+                for i, pixel_values in enumerate(data["pixel_values"]):
+                    if not torch.all(
+                            pixel_values == 0):  # Check if all values are zero
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples
+                valid_data = {
+                    "pixel_values":
+                    torch.stack(
+                        [data["pixel_values"][i] for i in valid_indices]),
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                    "fps": [data["fps"][i] for i in valid_indices],
+                    "duration": [data["duration"][i] for i in valid_indices],
+                }
+
+                # VAE
+                with torch.autocast("cuda", dtype=torch.float32):
+                    latents = self.get_module("vae").encode(
+                        valid_data["pixel_values"].to(
+                            get_local_torch_device())).mean
+
+                # Get extra features if needed
+                extra_features = self.get_extra_features(
+                    valid_data, fastvideo_args)
+
+                batch_captions = valid_data["text"]
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
+                )
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
+
+                # # Get sequence lengths from attention masks (number of 1s)
+                # seq_lens = prompt_attention_mask.sum(dim=1)
+
+                # non_padded_embeds = []
+                # non_padded_masks = []
+
+                # # Process each item in the batch
+                # for i in range(prompt_embeds.size(0)):
+                #     seq_len = seq_lens[i].item()
+                #     # Slice the embeddings and masks to keep only non-padding parts
+                #     non_padded_embeds.append(prompt_embeds[i, :seq_len])
+                #     non_padded_masks.append(prompt_attention_mask[i, :seq_len])
+
+                # Update the tensors with non-padded versions
+                # prompt_embeds = non_padded_embeds
+                # prompt_attention_masks = non_padded_masks
+                # prompt_embeds = prompt_embeds
+
+                # logger.info(f"===== prompt_embeds: {prompt_embeds[0].shape}")
+                # logger.info(f"===== prompt_attention_masks: {prompt_attention_masks[0].shape}")
+
+                sampling_params = SamplingParam.from_pretrained(
+                    args.model_path)
+
+                # encode negative prompt for trajectory collection
+                if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None:
+                    negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                        sampling_params.negative_prompt,
+                        fastvideo_args,
+                        encoder_index=[0],
+                        return_attention_mask=True,
+                    )
+                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
+                    negative_prompt_attention_mask = negative_prompt_masks_list[0][0]
+                else:
+                    negative_prompt_embed = None
+                    negative_prompt_attention_mask = None
+
+                trajectory_latents = []
+                trajectory_timesteps = []
+                trajectory_decoded = []
+                for i, (prompt_embed, prompt_attention_mask) in enumerate(zip(prompt_embeds, prompt_attention_masks)):
+                    prompt_embed = prompt_embed.unsqueeze(0)
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
+                    logger.info(f"what")
+                    logger.info(f"===== prompt_embed: {prompt_embed.shape}")
+                    logger.info(f"===== prompt_attention_mask: {prompt_attention_mask.shape}")
+                    # Collect the trajectory data
+                    batch = ForwardBatch(
+                        **shallow_asdict(sampling_params),
+                        # data_type="video",
+                        # seed=args.seed,
+                        # prompt=batch_captions[i],
+                        # prompt_embeds=[prompt_embed],
+                        # prompt_attention_mask=[prompt_attention_mask],
+                        # height=args.max_height,
+                        # width=args.max_width,
+                        # num_frames=81,
+                        # fps=args.train_fps,
+                        # return_trajectory_latents=True,
+                        # guidance_scale=3.0,
+                        # do_classifier_free_guidance=True,
+                    )
+                    batch.prompt_embeds = [prompt_embed]
+                    batch.prompt_attention_mask = [prompt_attention_mask]
+                    batch.negative_prompt_embeds = [negative_prompt_embed]
+                    batch.negative_attention_mask = [negative_prompt_attention_mask]
+                    batch.return_trajectory_latents = True
+                    batch.return_trajectory_decoded = False
+                    batch.height = args.max_height
+                    batch.width = args.max_width
+                    # batch.num_frames = 81
+                    batch.fps = args.train_fps
+                    batch.guidance_scale = 3.0
+                    batch.do_classifier_free_guidance = True
+                    # fastvideo_args.pipeline_config.ti2v_task = True
+
+                    result_batch = self.input_validation_stage(
+                        batch, fastvideo_args)
+                    # result_batch = self.prompt_encoding_stage(result_batch, fastvideo_args)
+                    # result_batch = self.vae_encoding_stage(result_batch, fastvideo_args)
+                    result_batch = self.timestep_preparation_stage(
+                        batch, fastvideo_args)
+                    result_batch = self.latent_preparation_stage(
+                        result_batch, fastvideo_args)
+                    result_batch = self.denoising_stage(result_batch,
+                                                        fastvideo_args)
+                    result_batch = self.decoding_stage(result_batch, fastvideo_args)
+                    # trajectory_latents = result_batch.trajectory_latents
+                    trajectory_latents.append(result_batch.trajectory_latents.cpu())
+                    trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
+                    trajectory_decoded.append(result_batch.trajectory_decoded)
+
+            extra_features["trajectory_latents"] = trajectory_latents
+            extra_features["trajectory_timesteps"] = trajectory_timesteps
+            logger.info(f"===== trajectory_latents: {trajectory_latents[0].shape}")
+            logger.info(f"===== trajectory_latents len: {len(trajectory_latents)}")
+            logger.info(f"===== trajectory_timesteps: {trajectory_timesteps}")
+            logger.info(f"===== trajectory_timesteps len: {len(trajectory_timesteps)}")
+
+            if batch.return_trajectory_decoded:
+                logger.info(f"===== SAVING TRAJECTORY DECODED")
+                for i, decoded_frames in enumerate(trajectory_decoded):
+                    for j, decoded_frame in enumerate(decoded_frames):
+                        logger.info(f"===== SAVING TRAJECTORY DECODED {i} for prompt {batch_captions[i]}")
+                        save_decoded_latents_as_video(decoded_frame, f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", args.train_fps)
+            # assert False
+            # Prepare batch data for Parquet dataset
+            batch_data = []
+
+            # Add progress bar for saving outputs
+            save_pbar = tqdm(enumerate(valid_data["path"]),
+                             desc="Saving outputs",
+                             unit="item",
+                             leave=False)
+            for idx, video_path in save_pbar:
+                # Get the corresponding latent and info using video name
+                latent = latents[idx].cpu()
+                video_name = os.path.basename(video_path).split(".")[0]
+
+                # Convert tensors to numpy arrays
+                vae_latent = latent.cpu().numpy()
+                text_embedding = prompt_embeds[idx].cpu().numpy()
+
+                # Get extra features for this sample if needed
+                sample_extra_features = {}
+                if extra_features:
+                    for key, value in extra_features.items():
+                        logger.info(f"===== key: {key}")
+                        if isinstance(value, torch.Tensor):
+                            logger.info(f"===== value: {value[idx].shape}")
+                            sample_extra_features[key] = value[idx].cpu().numpy(
+                            )
+                        else:
+                            assert isinstance(value, list)
+                            if isinstance(value[idx], torch.Tensor):
+                                logger.info(f"===== value in list: {value[idx].shape}")
+                                sample_extra_features[key] = value[idx].cpu().float().numpy(
+                                )
+                            else:
+                                logger.info(f"===== value in list: not tensor")
+                                sample_extra_features[key] = value[idx]
+                            # logger.info(f"===== value: not tensor")
+                            # sample_extra_features[key] = value[idx]
+
+                # Create record for Parquet dataset
+                record = self.create_record(
+                    video_name=video_name,
+                    vae_latent=vae_latent,
+                    text_embedding=text_embedding,
+                    valid_data=valid_data,
+                    idx=idx,
+                    extra_features=sample_extra_features)
+                batch_data.append(record)
+
+            if batch_data:
+                # Add progress bar for writing to Parquet dataset
+                write_pbar = tqdm(total=1,
+                                  desc="Writing to Parquet dataset",
+                                  unit="batch")
+                # Convert batch data to PyArrow arrays
+                arrays = []
+                for field in self.get_schema_fields():
+                    if field.endswith('_bytes'):
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.binary()))
+                    elif field.endswith('_shape'):
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.list_(pa.int32())))
+                    elif field in ['width', 'height', 'num_frames']:
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.int32()))
+                    elif field in ['duration_sec', 'fps']:
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.float32()))
+                    else:
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data]))
+
+                table = pa.Table.from_arrays(arrays,
+                                             names=self.get_schema_fields())
+                write_pbar.update(1)
+                write_pbar.close()
+
+                # Store the table in a list for later processing
+                if not hasattr(self, 'all_tables'):
+                    self.all_tables = []
+                self.all_tables.append(table)
+
+                logger.info("Collected batch with %s samples", len(table))
+
+            if self.num_processed_samples >= args.flush_frequency:
+                self._flush_tables(self.num_processed_samples, args,
+                                   self.combined_parquet_dir)
+                self.num_processed_samples = 0
+                self.all_tables = []
+
+    def preprocess_text_and_trajectory(self,
+                                     fastvideo_args: FastVideoArgs,
+                                     args):
+        """Preprocess text-only data and generate trajectory information."""
+        
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # For text-only processing, we only need text data
+                # Filter out samples without text
+                valid_indices = []
+                for i, text in enumerate(data["text"]):
+                    if text and text.strip():  # Check if text is not empty
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples (text-only)
+                valid_data = {
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                }
+                
+                # Add fps and duration if available in data
+                if "fps" in data:
+                    valid_data["fps"] = [data["fps"][i] for i in valid_indices]
+                if "duration" in data:
+                    valid_data["duration"] = [data["duration"][i] for i in valid_indices]
+
+                batch_captions = valid_data["text"]
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
+                )
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
+
+                sampling_params = SamplingParam.from_pretrained(
+                    args.model_path)
+
+                # encode negative prompt for trajectory collection
+                if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None:
+                    negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                        sampling_params.negative_prompt,
+                        fastvideo_args,
+                        encoder_index=[0],
+                        return_attention_mask=True,
+                    )
+                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
+                    negative_prompt_attention_mask = negative_prompt_masks_list[0][0]
+                else:
+                    negative_prompt_embed = None
+                    negative_prompt_attention_mask = None
+
+                trajectory_latents = []
+                trajectory_timesteps = []
+                trajectory_decoded = []
+                
+                for i, (prompt_embed, prompt_attention_mask) in enumerate(zip(prompt_embeds, prompt_attention_masks)):
+                    prompt_embed = prompt_embed.unsqueeze(0)
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
+                    
+                    # Collect the trajectory data (text-to-video generation)
+                    batch = ForwardBatch(
+                        **shallow_asdict(sampling_params),
+                    )
+                    batch.prompt_embeds = [prompt_embed]
+                    batch.prompt_attention_mask = [prompt_attention_mask]
+                    batch.negative_prompt_embeds = [negative_prompt_embed]
+                    batch.negative_attention_mask = [negative_prompt_attention_mask]
+                    batch.return_trajectory_latents = True
+                    batch.return_trajectory_decoded = False
+                    batch.height = args.max_height
+                    batch.width = args.max_width
+                    batch.fps = args.train_fps
+                    batch.guidance_scale = 3.0
+                    batch.do_classifier_free_guidance = True
+
+                    result_batch = self.input_validation_stage(
+                        batch, fastvideo_args)
+                    result_batch = self.timestep_preparation_stage(
+                        batch, fastvideo_args)
+                    result_batch = self.latent_preparation_stage(
+                        result_batch, fastvideo_args)
+                    result_batch = self.denoising_stage(result_batch,
+                                                        fastvideo_args)
+                    result_batch = self.decoding_stage(result_batch, fastvideo_args)
+                    
+                    trajectory_latents.append(result_batch.trajectory_latents.cpu())
+                    trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
+                    trajectory_decoded.append(result_batch.trajectory_decoded)
+
+                # Prepare extra features for text-only processing
+                extra_features = {
+                    "trajectory_latents": trajectory_latents,
+                    "trajectory_timesteps": trajectory_timesteps
+                }
+
+                logger.info(f"===== trajectory_latents: {trajectory_latents[0].shape}")
+                logger.info(f"===== trajectory_latents len: {len(trajectory_latents)}")
+                logger.info(f"===== trajectory_timesteps: {trajectory_timesteps}")
+                logger.info(f"===== trajectory_timesteps len: {len(trajectory_timesteps)}")
+
+                if batch.return_trajectory_decoded:
+                    logger.info(f"===== SAVING TRAJECTORY DECODED")
+                    for i, decoded_frames in enumerate(trajectory_decoded):
+                        for j, decoded_frame in enumerate(decoded_frames):
+                            logger.info(f"===== SAVING TRAJECTORY DECODED {i} for prompt {batch_captions[i]}")
+                            save_decoded_latents_as_video(decoded_frame, f"decoded_videos/trajectory_decoded_{i}_{j}.mp4", args.train_fps)
+
+                # Prepare batch data for Parquet dataset
+                batch_data = []
+
+                # Add progress bar for saving outputs
+                save_pbar = tqdm(enumerate(valid_data["path"]),
+                                 desc="Saving outputs",
+                                 unit="item",
+                                 leave=False)
+                
+                for idx, video_path in save_pbar:
+                    video_name = os.path.basename(video_path).split(".")[0]
+
+                    # Convert tensors to numpy arrays
+                    text_embedding = prompt_embeds[idx].cpu().numpy()
+
+                    # Get extra features for this sample
+                    sample_extra_features = {}
+                    if extra_features:
+                        for key, value in extra_features.items():
+                            logger.info(f"===== key: {key}")
+                            if isinstance(value, torch.Tensor):
+                                logger.info(f"===== value: {value[idx].shape}")
+                                sample_extra_features[key] = value[idx].cpu().numpy()
+                            else:
+                                assert isinstance(value, list)
+                                if isinstance(value[idx], torch.Tensor):
+                                    logger.info(f"===== value in list: {value[idx].shape}")
+                                    sample_extra_features[key] = value[idx].cpu().float().numpy()
+                                else:
+                                    logger.info(f"===== value in list: not tensor")
+                                    sample_extra_features[key] = value[idx]
+
+                    # Create record for Parquet dataset (without VAE latents for text-only)
+                    record = self.create_text_only_record(
+                        args,
+                        video_name=video_name,
+                        text_embedding=text_embedding,
+                        valid_data=valid_data,
+                        idx=idx,
+                        extra_features=sample_extra_features)
+                    batch_data.append(record)
+
+                if batch_data:
+                    # Add progress bar for writing to Parquet dataset
+                    write_pbar = tqdm(total=1,
+                                      desc="Writing to Parquet dataset",
+                                      unit="batch")
+                    # Convert batch data to PyArrow arrays
+                    arrays = []
+                    for field in self.get_schema_fields():
+                        if field.endswith('_bytes'):
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data],
+                                         type=pa.binary()))
+                        elif field.endswith('_shape'):
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data],
+                                         type=pa.list_(pa.int32())))
+                        elif field in ['width', 'height', 'num_frames']:
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data],
+                                         type=pa.int32()))
+                        elif field in ['duration_sec', 'fps']:
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data],
+                                         type=pa.float32()))
+                        else:
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data]))
+
+                    table = pa.Table.from_arrays(arrays,
+                                                 names=self.get_schema_fields())
+                    write_pbar.update(1)
+                    write_pbar.close()
+
+                    # Store the table in a list for later processing
+                    if not hasattr(self, 'all_tables'):
+                        self.all_tables = []
+                    self.all_tables.append(table)
+
+                    logger.info("Collected batch with %s samples", len(table))
+
+                if self.num_processed_samples >= args.flush_frequency:
+                    self._flush_tables(self.num_processed_samples, args,
+                                       self.combined_parquet_dir)
+                    self.num_processed_samples = 0
+                    self.all_tables = []
+        
+        # Final flush for any remaining samples
+        if hasattr(self, 'all_tables') and self.all_tables and self.num_processed_samples > 0:
+            logger.info(f"Final flush with {self.num_processed_samples} remaining samples")
+            self._flush_tables(self.num_processed_samples, args, self.combined_parquet_dir)
+            self.num_processed_samples = 0
+            self.all_tables = []
+
+    def create_text_only_record(
+            self,
+            args,
+            video_name: str,
+            text_embedding: np.ndarray,
+            valid_data: dict[str, Any],
+            idx: int,
+            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Create a record for text-only preprocessing using text-only schema."""
+        
+        # Create base record using only fields from text-only schema
+        record = {
+            "id": f"text_{video_name}_{idx}",
+            "text_embedding_bytes": text_embedding.tobytes(),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype),
+            "file_name": video_name,
+            "caption": valid_data["text"][idx],
+            "media_type": "text",
+        }
+
+        # Add trajectory data if available
+        if extra_features and "trajectory_latents" in extra_features:
+            trajectory_latents = extra_features["trajectory_latents"][idx] if isinstance(extra_features["trajectory_latents"], list) else extra_features["trajectory_latents"]
+            record.update({
+                "trajectory_latents_bytes": trajectory_latents.tobytes(),
+                "trajectory_latents_shape": list(trajectory_latents.shape),
+                "trajectory_latents_dtype": str(trajectory_latents.dtype),
+            })
+        else:
+            record.update({
+                "trajectory_latents_bytes": b"",
+                "trajectory_latents_shape": [],
+                "trajectory_latents_dtype": "",
+            })
+
+        if extra_features and "trajectory_timesteps" in extra_features:
+            trajectory_timesteps = extra_features["trajectory_timesteps"][idx] if isinstance(extra_features["trajectory_timesteps"], list) else extra_features["trajectory_timesteps"]
+            record.update({
+                "trajectory_timesteps_bytes": trajectory_timesteps.tobytes(),
+                "trajectory_timesteps_shape": list(trajectory_timesteps.shape),
+                "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
+            })
+        else:
+            record.update({
+                "trajectory_timesteps_bytes": b"",
+                "trajectory_timesteps_shape": [],
+                "trajectory_timesteps_dtype": "",
+            })
+
+        return record
+
+
+    def get_extra_features(self, valid_data: dict[str, Any],
+                           fastvideo_args: FastVideoArgs) -> dict[str, Any]:
+
+        # TODO(will): move these to cpu at some point
+        self.get_module("vae").to(get_local_torch_device())
+
+        # generator = torch.Generator(device=get_local_torch_device(), seed=42)
+        generator = torch.Generator("cpu").manual_seed(42)
+
+        features = {}
+        """Get CLIP features from the first frame of each video."""
+        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
+            0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
+        _, _, num_frames, height, width = valid_data["pixel_values"].shape
+        # latent_height = height // self.get_module(
+        #     "vae").spatial_compression_ratio
+        # latent_width = width // self.get_module("vae").spatial_compression_ratio
+
+        unprocessed_images = []
+        pil_images = []
+        # Frame has values between -1 and 1
+        for frame in first_frame:
+            frame = (frame + 1) * 127.5
+            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
+            pil_images.append(frame_pil)
+            # processed_img = self.get_module("image_processor")(
+            #     images=frame_pil, return_tensors="pt")
+            unprocessed_images.append(frame_pil)
+        """Get VAE features from the first frame of each video"""
+        video_conditions = []
+        for frame in unprocessed_images:
+
+            latent = self.vae_encoding_stage.encode_image(
+                frame, height, width, fastvideo_args, generator)
+            video_conditions.append(latent)
+
+        features["image_condition_latents"] = video_conditions
+        features["pil_images"] = pil_images
+        return features
+
+    def create_record(
+            self,
+            video_name: str,
+            vae_latent: np.ndarray,
+            text_embedding: np.ndarray,
+            valid_data: dict[str, Any],
+            idx: int,
+            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Create a record for the Parquet dataset with CLIP features."""
+        record = super().create_record(video_name=video_name,
+                                       vae_latent=vae_latent,
+                                       text_embedding=text_embedding,
+                                       valid_data=valid_data,
+                                       idx=idx,
+                                       extra_features=extra_features)
+
+        if extra_features and "image_condition_latents" in extra_features:
+            image_condition_latents = extra_features["image_condition_latents"]
+            record.update({
+                "image_condition_latents_bytes":
+                image_condition_latents.tobytes(),
+                "image_condition_latents_shape":
+                list(image_condition_latents.shape),
+                "image_condition_latents_dtype":
+                str(image_condition_latents.dtype),
+            })
+        else:
+            record.update({
+                "image_condition_latents_bytes": b"",
+                "image_condition_latents_shape": [],
+                "image_condition_latents_dtype": "",
+            })
+
+        if extra_features and "trajectory_latents" in extra_features:
+            trajectory_latents = extra_features["trajectory_latents"]
+            record.update({
+                "trajectory_latents_bytes": trajectory_latents.tobytes(),
+                "trajectory_latents_shape": list(trajectory_latents.shape),
+                "trajectory_latents_dtype": str(trajectory_latents.dtype),
+            })
+        else:
+            record.update({
+                "trajectory_latents_bytes": b"",
+                "trajectory_latents_shape": [],
+                "trajectory_latents_dtype": "",
+            })
+
+        if extra_features and "trajectory_timesteps" in extra_features:
+            trajectory_timesteps = extra_features["trajectory_timesteps"]
+            record.update({
+                "trajectory_timesteps_bytes": trajectory_timesteps.tobytes(),
+                "trajectory_timesteps_shape": list(trajectory_timesteps.shape),
+                "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
+            })
+        else:
+            record.update({
+                "trajectory_timesteps_bytes": b"",
+                "trajectory_timesteps_shape": [],
+                "trajectory_timesteps_dtype": "",
+            })
+
+        if extra_features and "pil_image" in extra_features:
+            pil_image = extra_features["pil_image"]
+            record.update({
+                "pil_image_bytes": pil_image.tobytes(),
+                "pil_image_shape": list(pil_image.shape),
+                "pil_image_dtype": str(pil_image.dtype),
+            })
+        else:
+            record.update({
+                "pil_image_bytes": b"",
+                "pil_image_shape": [],
+                "pil_image_dtype": "",
+            })
+
+        return record
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
+        if not self.post_init_called:
+            self.post_init()
+
+        self.local_rank = int(os.getenv("RANK", 0))
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        self.combined_parquet_dir = os.path.join(args.output_dir,
+                                                 "combined_parquet_dataset")
+        os.makedirs(self.combined_parquet_dir, exist_ok=True)
+
+        # Loading dataset
+        #train_dataset = getdataset(args)
+        train_dataset = gettextdataset(args)
+
+        self.preprocess_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
+
+        self.num_processed_samples = 0
+        # Add progress bar for video preprocessing
+        self.pbar = tqdm(self.preprocess_loader_iter,
+                         desc="Processing videos",
+                         unit="batch",
+                         disable=self.local_rank != 0)
+
+        # Initialize class variables for data sharing
+        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
+        self.latent_data: dict[str, Any] = {}  # Store latent tensors
+        #self.preprocess_video_and_text_and_trajectory(fastvideo_args, args)
+        self.preprocess_text_and_trajectory(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_ODE_Trajectory

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -9,8 +9,11 @@ from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (
     PreprocessPipeline_I2V)
+from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (
+    PreprocessPipeline_ODE_Trajectory)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (
     PreprocessPipeline_T2V)
+from fastvideo.pipelines.preprocess_text import PreprocessPipeline_Text
 from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -21,12 +24,22 @@ def main(args) -> None:
     maybe_init_distributed_environment_and_model_parallel(1, 1)
     num_gpus = int(os.environ["WORLD_SIZE"])
     assert num_gpus == 1, "Only support 1 GPU"
-    pipeline_config = PipelineConfig.from_pretrained(args.model_path)
-    kwargs = {
-        "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
-    }
-    pipeline_config.update_config_from_dict(kwargs)
+    
+    if args.preprocess_task == "text_only":
+        pipeline_config = PipelineConfig.from_pretrained(args.model_path)
+        kwargs = {
+            "text_encoder_cpu_offload": False,
+        }
+        pipeline_config.update_config_from_dict(kwargs)
+    else:
+        # Full config for video/image processing
+        pipeline_config = PipelineConfig.from_pretrained(args.model_path)
+        kwargs = {
+            "vae_precision": "fp32",
+            "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+        }
+        pipeline_config.update_config_from_dict(kwargs)
+    
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
@@ -35,7 +48,25 @@ def main(args) -> None:
         text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
-    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
+    
+    if args.preprocess_task == "t2v":
+        PreprocessPipeline = PreprocessPipeline_T2V
+    elif args.preprocess_task == "i2v":
+        PreprocessPipeline = PreprocessPipeline_I2V
+    elif args.preprocess_task == "ode_trajectory":
+        print("Preprocess pipeline...")
+        PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
+    elif args.preprocess_task == "text_only":
+        print("Text-only preprocessing pipeline...")
+        PreprocessPipeline = PreprocessPipeline_Text
+    else:
+        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
+                        f"Valid options: t2v, i2v, ode_trajectory, text_only")
+
+    logger.info(
+        f"Preprocess task: {args.preprocess_task} using {PreprocessPipeline.__name__}"
+    )
+
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 
@@ -74,7 +105,11 @@ if __name__ == "__main__":
     parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
-    parser.add_argument("--preprocess_task", type=str, default="t2v")
+    parser.add_argument("--preprocess_task", 
+                        type=str, 
+                        default="t2v",
+                        choices=["t2v", "i2v", "ode_trajectory", "text_only"],
+                        help="Type of preprocessing task to run")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)
@@ -98,4 +133,4 @@ if __name__ == "__main__":
     )
 
     args = parser.parse_args()
-    main(args)
+    main(args)

fastvideo/pipelines/preprocess_text.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Text-only Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the Text-only Data Preprocessing pipeline
+using the modular pipeline architecture, based on the ODE Trajectory preprocessing.
+"""
+
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import torch
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+
+from fastvideo.dataset import gettextdataset
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_text_only
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import (TextEncodingStage)
+
+logger = init_logger(__name__)
+
+
+class PreprocessPipeline_Text(BasePreprocessPipeline):
+    """Text-only preprocessing pipeline implementation."""
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer"
+    ]
+    preprocess_dataloader: StatefulDataLoader
+    preprocess_loader_iter: Iterator[dict[str, Any]]
+
+    def get_schema_fields(self):
+        """Get the schema fields for text-only pipeline."""
+        return [f.name for f in pyarrow_schema_text_only]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+    def preprocess_text_only(self,
+                            fastvideo_args: FastVideoArgs,
+                            args):
+        """Preprocess text-only data."""
+        
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # For text-only processing, we only need text data
+                # Filter out samples without text
+                valid_indices = []
+                for i, text in enumerate(data["text"]):
+                    if text and text.strip():  # Check if text is not empty
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples (text-only)
+                valid_data = {
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                }
+
+                batch_captions = valid_data["text"]
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
+                )
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
+
+                logger.info(f"===== prompt_embeds: {prompt_embeds.shape}")
+                logger.info(f"===== prompt_attention_masks: {prompt_attention_masks.shape}")
+
+                # Prepare batch data for Parquet dataset
+                batch_data = []
+
+                # Add progress bar for saving outputs
+                save_pbar = tqdm(enumerate(valid_data["path"]),
+                                 desc="Saving outputs",
+                                 unit="item",
+                                 leave=False)
+                
+                for idx, text_path in save_pbar:
+                    text_name = os.path.basename(text_path).split(".")[0]
+
+                    # Convert tensors to numpy arrays
+                    text_embedding = prompt_embeds[idx].cpu().numpy()
+
+                    # Create record for Parquet dataset (text-only)
+                    record = self.create_text_only_record(
+                        text_name=text_name,
+                        text_embedding=text_embedding,
+                        valid_data=valid_data,
+                        idx=idx)
+                    batch_data.append(record)
+
+                if batch_data:
+                    # Add progress bar for writing to Parquet dataset
+                    write_pbar = tqdm(total=1,
+                                      desc="Writing to Parquet dataset",
+                                      unit="batch")
+                    # Convert batch data to PyArrow arrays
+                    arrays = []
+                    for field in self.get_schema_fields():
+                        if field.endswith('_bytes'):
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data],
+                                         type=pa.binary()))
+                        elif field.endswith('_shape'):
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data],
+                                         type=pa.list_(pa.int32())))
+                        else:
+                            arrays.append(
+                                pa.array([record[field] for record in batch_data]))
+
+                    table = pa.Table.from_arrays(arrays,
+                                                 names=self.get_schema_fields())
+                    write_pbar.update(1)
+                    write_pbar.close()
+
+                    # Store the table in a list for later processing
+                    if not hasattr(self, 'all_tables'):
+                        self.all_tables = []
+                    self.all_tables.append(table)
+
+                    logger.info("Collected batch with %s samples", len(table))
+
+                if self.num_processed_samples >= args.flush_frequency:
+                    self._flush_tables(self.num_processed_samples, args,
+                                       self.combined_parquet_dir)
+                    self.num_processed_samples = 0
+                    self.all_tables = []
+        
+        # Final flush for any remaining samples
+        if hasattr(self, 'all_tables') and self.all_tables and self.num_processed_samples > 0:
+            logger.info(f"Final flush with {self.num_processed_samples} remaining samples")
+            self._flush_tables(self.num_processed_samples, args, self.combined_parquet_dir)
+            self.num_processed_samples = 0
+            self.all_tables = []
+
+    def create_text_only_record(
+            self,
+            text_name: str,
+            text_embedding: np.ndarray,
+            valid_data: dict[str, Any],
+            idx: int) -> dict[str, Any]:
+        """Create a record for text-only preprocessing using text-only schema."""
+        
+        # Create base record using only fields from text-only schema
+        record = {
+            "id": f"text_{text_name}_{idx}",
+            "text_embedding_bytes": text_embedding.tobytes(),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype),
+        }
+
+        return record
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
+        if not self.post_init_called:
+            self.post_init()
+
+        self.local_rank = int(os.getenv("RANK", 0))
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        self.combined_parquet_dir = os.path.join(args.output_dir,
+                                                 "combined_parquet_dataset")
+        os.makedirs(self.combined_parquet_dir, exist_ok=True)
+
+        # Loading text dataset
+        train_dataset = gettextdataset(args)
+
+        self.preprocess_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
+
+        self.num_processed_samples = 0
+        # Add progress bar for text preprocessing
+        self.pbar = tqdm(self.preprocess_loader_iter,
+                         desc="Processing text",
+                         unit="batch",
+                         disable=self.local_rank != 0)
+
+        # Initialize class variables for data sharing
+        self.text_data: dict[str, Any] = {}  # Store text metadata and paths
+        
+        self.preprocess_text_only(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_Text

scripts/preprocess/launch_slurm.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# Create output directory if it doesn't exist
+mkdir -p preprocess_output
+
+# Launch 8 jobs, one for each node
+# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node)
+for node_id in {0..1}; do
+    # Calculate the starting file number for this node
+    start_file=$((node_id * 8 + 1))
+    
+    echo "Launching node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt"
+    echo "sbatch --job-name=ode-${node_id} --output=preprocess_output/preprocess-node-${node_id}.out --error=preprocess_output/preprocess-node-${node_id}.err /mnt/weka/home/hao.zhang/wl/kev/preproc/slurms/syn.slurm $start_file $node_id"
+    
+    sbatch --job-name=ode-${node_id} \
+           --output=preprocess_output/preprocess-node-${node_id}.out \
+           --error=preprocess_output/preprocess-node-${node_id}.err \
+           /mnt/weka/home/hao.zhang/wl/kev/preproc/slurms/syn.slurm $start_file $node_id
+done
+
+echo "All 8 nodes launched successfully!"

scripts/preprocess/launch_slurm_text.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# Create output directory if it doesn't exist
+mkdir -p preprocess_output_text
+
+# Launch 8 jobs, one for each node
+# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node)
+for node_id in {0..1}; do
+    # Calculate the starting file number for this node
+    start_file=$((node_id * 8 + 1))
+    
+    echo "Launching text-only node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt"
+    echo "sbatch --job-name=text-${node_id} --output=preprocess_output_text/preprocess-text-node-${node_id}.out --error=preprocess_output_text/preprocess-text-node-${node_id}.err /mnt/weka/home/hao.zhang/matthew/FastVideo/scripts/preprocess/syn_text.slurm $start_file $node_id"
+    
+    sbatch --job-name=text-${node_id} \
+           --output=preprocess_output_text/preprocess-text-node-${node_id}.out \
+           --error=preprocess_output_text/preprocess-text-node-${node_id}.err \
+           /mnt/weka/home/hao.zhang/matthew/FastVideo/scripts/preprocess/syn_text.slurm $start_file $node_id
+done
+
+echo "All 8 text-only nodes launched successfully!" 

scripts/preprocess/syn.slurm

@@ -0,0 +1,88 @@
+#!/bin/bash
+#SBATCH --partition=main
+#SBATCH --qos=hao
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=960G
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+# conda init
+# source ~/conda/miniconda/bin/activate
+# PYTHON_VIRTUAL_ENVIRONMENT=fastvideo-train-yq
+# conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+nvidia-smi
+nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices"
+echo ""
+echo " Run started at:- "
+date
+
+# Accept parameters from launch script
+START_FILE=${1:-1}  # Starting file number for this node
+NODE_ID=${2:-0}  # Node identifier (0-7)
+
+num_gpus=1
+export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+# Start port number - we'll increment for each job
+base_port=$((29603 + NODE_ID * 100))  # Different port range per node
+
+# Create an array of CUDA device IDs
+gpu_ids=(0 1 2 3 4 5 6 7)
+
+GPU_NUM=1
+MODEL_TYPE="wan"
+
+echo "NODE_ID: $NODE_ID"
+echo "START_FILE: $START_FILE"
+echo "Base port for this node: $base_port"
+
+# Run 8 parallel preprocessing jobs on this node
+for i in {1..8}; do
+    # Calculate port for this job
+    port=$((base_port + i))
+    
+    # Get GPU ID using modulo to cycle through available GPUs
+    gpu=${gpu_ids[((i-1))]}
+    
+    # Calculate which file this GPU should process
+    file_num=$((START_FILE + i - 1))
+    DATA_MERGE_PATH="/mnt/weka/home/hao.zhang/wl/kev/preproc/prompts/v2m_${file_num}.txt"
+    
+    # Create unique output directory based on node and GPU
+    OUTPUT_DIR="data/test-ode-preprocessing/Node_${NODE_ID}_GPU_${i}_File_${file_num}"
+    
+    start_cpu=$(( (i-1)*2 ))  # Reduced CPU allocation for 8 nodes
+    end_cpu=$(( start_cpu+1 ))
+    
+    echo "Starting GPU $gpu processing file v2m_${file_num}.txt on port $port, output: $OUTPUT_DIR"
+    
+    # Run the preprocessing command in background
+    CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \
+        /mnt/weka/home/hao.zhang/wl/kev/FastVideo/fastvideo/pipelines/preprocess/v1_preprocess.py \
+        --model_path $MODEL_BASE \
+        --data_merge_path $DATA_MERGE_PATH \
+        --preprocess_video_batch_size 2 \
+        --seed 42 \
+        --max_height 480 \
+        --max_width 832 \
+        --num_frames 77 \
+        --dataloader_num_workers 0 \
+        --output_dir=$OUTPUT_DIR \
+        --train_fps 16 \
+        --samples_per_file 8 \
+        --flush_frequency 8 \
+        --video_length_tolerance_range 5 \
+        --preprocess_task "ode_trajectory" &
+done
+
+# Wait for all jobs on this node to complete
+wait
+
+echo "All processing blocks completed!"

scripts/preprocess/syn_text.slurm

@@ -0,0 +1,89 @@
+#!/bin/bash
+#SBATCH --partition=main
+#SBATCH --qos=hao
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=960G
+#SBATCH --exclusive
+#SBATCH --time=72:00:00
+
+# conda init
+# source ~/conda/miniconda/bin/activate
+# PYTHON_VIRTUAL_ENVIRONMENT=fastvideo-train-yq
+# conda activate $PYTHON_VIRTUAL_ENVIRONMENT
+nvidia-smi
+nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
+
+echo " "
+echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
+echo " GPUs per node:= " $SLURM_JOB_GPUS
+echo " Running on multiple nodes/GPU devices for TEXT-ONLY preprocessing"
+echo ""
+echo " Run started at:- "
+date
+
+# Accept parameters from launch script
+START_FILE=${1:-1}  # Starting file number for this node
+NODE_ID=${2:-0}  # Node identifier (0-7)
+
+num_gpus=1
+export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+# Start port number - we'll increment for each job
+base_port=$((29603 + NODE_ID * 100))  # Different port range per node
+
+# Create an array of CUDA device IDs
+gpu_ids=(0 1 2 3 4 5 6 7)
+
+GPU_NUM=1
+MODEL_TYPE="wan"
+
+echo "NODE_ID: $NODE_ID"
+echo "START_FILE: $START_FILE"
+echo "Base port for this node: $base_port"
+echo "Processing TEXT-ONLY data"
+
+# Run 8 parallel preprocessing jobs on this node
+for i in {1..8}; do
+    # Calculate port for this job
+    port=$((base_port + i))
+    
+    # Get GPU ID using modulo to cycle through available GPUs
+    gpu=${gpu_ids[((i-1))]}
+    
+    # Calculate which file this GPU should process
+    file_num=$((START_FILE + i - 1))
+    DATA_MERGE_PATH="/mnt/weka/home/hao.zhang/wl/kev/preproc/prompts/v2m_${file_num}.txt"
+    
+    # Create unique output directory based on node and GPU for text-only processing
+    OUTPUT_DIR="data/test-text-preprocessing/Node_${NODE_ID}_GPU_${i}_File_${file_num}"
+    
+    start_cpu=$(( (i-1)*2 ))  # Reduced CPU allocation for 8 nodes
+    end_cpu=$(( start_cpu+1 ))
+    
+    echo "Starting GPU $gpu processing text-only file v2m_${file_num}.txt on port $port, output: $OUTPUT_DIR"
+    
+    # Run the text-only preprocessing command in background
+    CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \
+        /mnt/weka/home/hao.zhang/matthew/FastVideo/fastvideo/pipelines/preprocess/v1_preprocess.py \
+        --model_path $MODEL_BASE \
+        --data_merge_path $DATA_MERGE_PATH \
+        --preprocess_video_batch_size 2 \
+        --seed 42 \
+        --max_height 480 \
+        --max_width 832 \
+        --num_frames 77 \
+        --dataloader_num_workers 0 \
+        --output_dir=$OUTPUT_DIR \
+        --train_fps 16 \
+        --samples_per_file 8 \
+        --flush_frequency 8 \
+        --video_length_tolerance_range 5 \
+        --preprocess_task "text" &
+done
+
+# Wait for all jobs on this node to complete
+wait
+
+echo "All text-only processing blocks completed!"