tests

Co-authored-by: kevin314 <kevin.lin.cs1@gmail.com>

examples/training/consistency_finetune/causal_ode_init/crush_smol_prompts.txt

@@ -0,0 +1,47 @@
+A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.
+The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.
+The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.
+A red toy car is being crushed by a large hydraulic press, which is flattening objects as if they were under a hydraulic press.
+A large, cylindrical object is seen pressing down on a small orange ball, causing it to flatten as if it were under a hydraulic press. The background features a green wall with yellow and red warning signs.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is shown compressing a wooden object, which shatters into small pieces. The background features a green wall with a yellow sign displaying a lightning bolt.
+A large metal cylinder is seen descending, flattening objects as if they were under a hydraulic press. The cylinder compresses a stack of matches and boxes, causing them to crumble into small pieces. The scene is set against a green background with yellow and red signs.
+A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.
+The video shows a metal press flattening objects as if they were under a hydraulic press. The press is pressing down on a pile of colorful gummy candies, squishing them into a pile of squiggly shapes. The press is made of metal and has a large base, and the gummy candies are of various colors, including red, green, and orange. The background is a green wall, and the press is placed on a metal surface.
+A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
+The video shows a stack of colorful sponges being flattened as if they were under a hydraulic press. The sponges, which are pink, white, blue, and green, are compressed into a smaller size, demonstrating the press's power. The background features a green wall with a yellow and red sign, adding context to the setting.
+A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, leaving a pile of debris around it.
+A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press.
+The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.
+The video shows a close-up of a metal cylinder pressing down on a yellow object, which is being flattened as if it were under a hydraulic press. The cylinder is positioned above the object, and the force is causing the object to compress and spread out, creating a visible deformation. The background is blurred, focusing attention on the action of the cylinder and the object being flattened.
+A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.
+The video shows a hydraulic press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing two colorful objects that resemble sandwiches. The press is yellow and black striped, and the objects being flattened are placed on a metal plate. The background is green, and the press is moving down, compressing the objects.
+The scene shows a metal press with a yellow and black striped pattern, holding a container filled with chocolate. A metal cylinder is descending, flattening the chocolate as if it were under a hydraulic press. The background is a green wall, and the press is mounted on a sturdy metal frame.
+The video shows a colorful sponge being flattened as if it were under a hydraulic press, with the sponge being compressed and eventually flattened into a thin layer.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is pressing down on a stack of wooden blocks, causing them to crumble and break apart. The press is black and yellow striped, and the wooden blocks are small and rectangular. The background is green, and the press is sitting on a metal table.
+A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
+The video shows a stack of colorful sponges being flattened by a large, cylindrical object, which appears to be a hydraulic press. The sponges, which are pink, blue, white, and green, are compressed into a single layer, demonstrating the press's powerful force. The background features a green wall with a yellow and red sign, adding context to the industrial setting.
+A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, demonstrating the immense pressure applied by the cylinder.
+A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press. The popcorn is crushed and scattered around the base of the cylinder, creating a satisfying visual effect.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is composed of a large, cylindrical metal cylinder with yellow and black stripes, and a metal base. The objects being flattened are two cylindrical blocks of cotton candy, one pink and one blue. The press is positioned on a metal table, and the background features a green wall with a yellow and red sign.
+The video shows a large orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.
+The video shows a cylindrical object being pressed down onto a flat surface, causing the objects beneath it to be flattened as if they were under a hydraulic press. The objects being flattened appear to be yellow and are being crushed into a pile of debris. The background is a greenish-gray color, and the surface on which the objects are being flattened is metallic and shiny.
+A green and blue object with a spiky texture is being flattened by a large, cylindrical metal press, demonstrating its resilience and durability.
+The video shows a stack of caramelized sugar cubes being flattened as if they were under a hydraulic press, resulting in a messy pile of broken sugar on the table.
+A large metal cylinder is seen pressing down on a pile of colorful jelly beans, flattening them as if they were under a hydraulic press.
+The video shows a machine with a yellow and black striped cylinder pressing down on a stack of colorful sponges, flattening them as if they were under a hydraulic press. The machine is situated in a green-walled room with warning signs in the background.
+The video shows a machine with a yellow and black striped cylinder, which is pressing down on two colorful objects, flattening them as if they were under a hydraulic press. The machine appears to be in a workshop or industrial setting, with a green wall in the background. The objects being flattened are green and orange, and the machine is covered in dirt and grime, indicating it has been used frequently.
+The video shows a large, industrial press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing a pile of pink objects into a pile of crumbs. The press is large and metallic, with a yellow and black striped pattern on its side. The background is a green wall with a yellow warning sign.
+The video shows a pink, sparkly ball being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
+A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and segments.
+The video shows a machine with a yellow and black striped cylinder, which is flattening objects as if they were under a hydraulic press. The machine is pressing down on two colorful objects, causing them to compress and flatten. The background is a green wall, and the machine appears to be in a workshop or industrial setting.
+The video shows a large, yellow and black striped cylinder flattening objects as if they were under a hydraulic press. The objects being flattened are pink and are being crushed into small pieces. The background is a green wall with a yellow sign.
+The video shows a machine with a yellow and black striped cylinder pressing down on two colorful objects, which are flattened as if they were under a hydraulic press. The machine is positioned on a metal platform, and the background is a green wall.
+A green cube is being compressed by a hydraulic press, which flattens the object as if it were under a hydraulic press. The press is shown in action, with the cube being squeezed into a smaller shape.
+A pink, sparkly ball is being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
+A red cabbage is being crushed by a hydraulic press, which flattens the objects as if they were under a hydraulic press. The press is shown in action, compressing the cabbage into a smaller, more compact form.
+A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and pulp.
+A large metal press is shown compressing a stack of burgers, causing them to be flattened and crushed into a pile of ground meat.
+A pizza is being crushed by a hydraulic press, causing the toppings to spread out and the crust to crumble.
+A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.
+A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.
+A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.

examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_crush_smol"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "wan_ode_init_crush_smol"
+  --max_train_steps 6000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --warp_denoising_step
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 6e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/consistency_finetune/causal_ode_init/preprocess_wan_data_crush_smol.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="$(dirname "$0")/crush_smol_prompts.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 1 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 81 \
+    --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" 

examples/training/consistency_finetune/causal_ode_init/validation.json

@@ -0,0 +1,40 @@
+{
+  "data": [
+    {
+      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}

examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json

@@ -28,4 +28,4 @@
       "num_frames": 77
     }
   ]
-}
+}

fastvideo/configs/pipelines/__init__.py

@@ -4,12 +4,13 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.registry import (
     get_pipeline_config_cls_from_name)
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.configs.pipelines.wan import (WanI2V480PConfig, WanI2V720PConfig,
+from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig,
+                                             WanI2V480PConfig, WanI2V720PConfig,
                                              WanT2V480PConfig, WanT2V720PConfig)
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
     "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
     "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
-    "get_pipeline_config_cls_from_name"
+    "SelfForcingWanT2V480PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/registry.py

@@ -11,9 +11,9 @@ from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 # isort: off
 from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
-    SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config,
-    Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig,
-    WanT2V720PConfig)
+    Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
+    WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
+    SelfForcingWanT2V480PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -48,6 +48,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "wanpipeline": lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
     "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
+    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
     "stepvideo": lambda id: "stepvideo" in id.lower(),
     # Add other pipeline architecture detectors
 }
@@ -60,6 +61,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
     "wandmdpipeline": FastWan2_1_T2V_480P_Config,
+    "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
     "stepvideo": StepVideoT2VConfig
     # Other fallbacks by architecture
 }

fastvideo/configs/sample/base.py

@@ -47,6 +47,8 @@ class SamplingParam:
     # Misc
     save_video: bool = True
     return_frames: bool = False
+    return_trajectory_latents: bool = False  # returns all latents for each timestep
+    return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
 
     def __post_init__(self) -> None:
         self.data_type = "video" if self.num_frames > 1 else "image"
@@ -198,6 +200,18 @@ class SamplingParam:
             help=
             "Path to a JSON file containing V-MoBA specific configurations.",
         )
+        parser.add_argument(
+            "--return-trajectory-latents",
+            action="store_true",
+            default=SamplingParam.return_trajectory_latents,
+            help="Whether to return the trajectory",
+        )
+        parser.add_argument(
+            "--return-trajectory-decoded",
+            action="store_true",
+            default=SamplingParam.return_trajectory_decoded,
+            help="Whether to return the decoded trajectory",
+        )
         return parser
 
 

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
@@ -37,9 +37,15 @@ def getdataset(args) -> VideoCaptionMergedDataset:
                                      temporal_sample=temporal_sample,
                                      transform_topcrop=transform_topcrop,
                                      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/parquet_io.py

@@ -0,0 +1,264 @@
+"""
+Utilities for converting preprocessing records (dicts) into Arrow tables and
+writing Parquet datasets in fixed-size chunks.
+
+This module centralizes table construction and Parquet file writing so
+pipelines only need to define their PyArrow schema and produce per-sample
+record dictionaries.
+
+Key APIs:
+- records_to_table(records, schema): Safely convert a list of dictionaries into
+  a pa.Table, casting to the provided schema.
+- ParquetDatasetWriter: Buffer tables and flush to a directory as multiple
+  Parquet files with a fixed number of rows per file. Uses temporary files and
+  atomic rename to avoid partially written outputs.
+"""
+
+from __future__ import annotations
+
+import multiprocessing
+import os
+from concurrent.futures import ProcessPoolExecutor
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+
+def records_to_table(records: list[dict[str, Any]], schema: pa.Schema) -> pa.Table:
+    """Build a PyArrow table from Python record dicts using an explicit schema.
+
+    Arrow will cast values to the target schema when possible (e.g., promoting
+    Python ints/floats to pa.int64/pa.float64), eliminating hand-written per-
+    field array construction.
+
+    Args:
+        records: List of dictionaries, each representing one row. Keys must
+            match schema field names.
+        schema: Target PyArrow schema. Controls field names and types.
+
+    Returns:
+        pa.Table: In-memory table matching the provided schema. If ``records``
+        is empty, returns an empty table with the given schema.
+    """
+    if not records:
+        return pa.table({}, schema=schema)
+    return pa.Table.from_pylist(records, schema=schema)
+
+
+class ParquetDatasetWriter:
+    """Accumulate tables and flush them to a Parquet directory in fixed-size chunks.
+
+    Behavior:
+    - Writes files under worker-specific subdirectories for parallelism.
+    - Uses temporary files and atomic rename to avoid partial files being left
+      behind on failure.
+    - Only full chunks of ``samples_per_file`` rows are written on each flush;
+      any remainder rows are re-buffered for the next flush.
+
+    Note:
+    - Instances are not meant to be shared across processes. Create one writer
+      per process if using multiprocessing.
+    """
+
+    def __init__(self, out_dir: str, samples_per_file: int, compression: str = "zstd") -> None:
+        """Initialize the dataset writer.
+
+        Args:
+            out_dir: Output directory where Parquet files will be written.
+            samples_per_file: Fixed number of rows per Parquet file.
+            compression: Compression codec passed to ``pyarrow.parquet.write_table``
+                (e.g., ``"zstd"``, ``"snappy"``, ``"gzip"``).
+        """
+        self.out_dir = out_dir
+        self.samples_per_file = max(int(samples_per_file), 1)
+        self.compression = compression
+        os.makedirs(self.out_dir, exist_ok=True)
+        self._tables: list[pa.Table] = []
+
+    def append_table(self, table: pa.Table) -> None:
+        """Append a non-empty table to the internal buffer.
+
+        Args:
+            table: A ``pa.Table`` to buffer. Empty or ``None`` tables are ignored.
+        """
+        if table is None or len(table) == 0:
+            return
+        self._tables.append(table)
+
+    def _combine(self) -> pa.Table | None:
+        """Combine all buffered tables into a single table, if any.
+
+        Returns:
+            A concatenated table, a single table if only one was buffered, or
+            ``None`` if no tables are buffered.
+        """
+        if not self._tables:
+            return None
+        if len(self._tables) == 1:
+            return self._tables[0]
+        return pa.concat_tables(self._tables, promote_options='none')
+
+    def flush(self, num_workers: int | None = None, write_remainder: bool = False) -> int:
+        """Write accumulated tables to disk and clear the written portion.
+
+        Only complete chunks of size ``samples_per_file`` are written. Any
+        remainder rows are kept buffered for the next flush.
+
+        Args:
+            num_workers: Optional override for the number of parallel workers
+                used to write chunks. Defaults to ``min(cpu_count, chunks)``.
+            write_remainder: If True, also write any leftover rows (< samples_per_file)
+                as a final small Parquet file (useful for the last flush at the
+                end of preprocessing).
+
+        Returns:
+            int: Number of rows successfully written in this flush call.
+        """
+        combined = self._combine()
+        self._tables = []
+        if combined is None or len(combined) == 0:
+            return 0
+
+        num_samples = len(combined)
+        total_chunks = num_samples // self.samples_per_file
+        if total_chunks == 0:
+            if not write_remainder:
+                # Not enough to form a full chunk; keep buffered for next round
+                # Re-buffer and return 0 written
+                self._tables = [combined]
+                return 0
+            # Last flush: write the small remainder as a final file in worker_0
+            worker_dir = os.path.join(self.out_dir, "worker_0")
+            os.makedirs(worker_dir, exist_ok=True)
+            # Determine next index
+            num_parquets = 0
+            for _, _, files in os.walk(worker_dir):
+                for file in files:
+                    if file.endswith('.parquet'):
+                        num_parquets += 1
+            chunk_path = os.path.join(worker_dir, f"data_chunk_{num_parquets}.parquet")
+            temp_path = chunk_path + '.tmp'
+            pq.write_table(combined, temp_path, compression=self.compression)
+            if os.path.exists(chunk_path):
+                os.remove(chunk_path)
+            os.rename(temp_path, chunk_path)
+            return num_samples
+
+        # Only write full chunks; keep remainder for next flush
+        written_rows = total_chunks * self.samples_per_file
+        remainder = num_samples - written_rows
+
+        table_to_write = combined.slice(0, written_rows)
+        remainder_table = combined.slice(written_rows, remainder) if remainder > 0 else None
+        if remainder_table is not None and len(remainder_table) > 0:
+            if write_remainder:
+                # Write the remainder as a final small file (worker_0)
+                worker_dir = os.path.join(self.out_dir, "worker_0")
+                os.makedirs(worker_dir, exist_ok=True)
+                num_parquets = 0
+                for _, _, files in os.walk(worker_dir):
+                    for file in files:
+                        if file.endswith('.parquet'):
+                            num_parquets += 1
+                remainder_path = os.path.join(worker_dir,
+                                              f"data_chunk_{num_parquets}.parquet")
+                temp_path = remainder_path + '.tmp'
+                pq.write_table(remainder_table,
+                               temp_path,
+                               compression=self.compression)
+                if os.path.exists(remainder_path):
+                    os.remove(remainder_path)
+                os.rename(temp_path, remainder_path)
+            else:
+                self._tables = [remainder_table]
+
+        # Parallel write by chunk ranges
+        if num_workers is None:
+            num_workers = min(multiprocessing.cpu_count(), max(total_chunks, 1))
+        num_workers = max(int(num_workers), 1)
+        chunks_per_worker = (total_chunks + num_workers - 1) // num_workers
+
+        work_ranges: list[tuple[int, int, pa.Table, int, str, int, str]] = []
+        for worker_id in range(num_workers):
+            start_chunk = worker_id * chunks_per_worker
+            end_chunk = min((worker_id + 1) * chunks_per_worker, total_chunks)
+            if start_chunk < end_chunk:
+                work_ranges.append(
+                    (
+                        start_chunk,
+                        end_chunk,
+                        table_to_write,
+                        worker_id,
+                        self.out_dir,
+                        self.samples_per_file,
+                        self.compression,
+                    )
+                )
+
+        written_total = 0
+        if len(work_ranges) == 1:
+            written_total += _process_chunk_range(work_ranges[0])
+            return written_total
+
+        with ProcessPoolExecutor(max_workers=num_workers) as executor:
+            futures = [executor.submit(_process_chunk_range, args) for args in work_ranges]
+            for f in futures:
+                written_total += f.result()
+        return written_total + (len(remainder_table) if write_remainder and remainder_table is not None else 0)
+
+
+def _process_chunk_range(args: Any) -> int:
+    """Worker function to write a contiguous range of chunk files.
+
+    Args:
+        args: Tuple containing
+            - start_chunk (int): inclusive start chunk index
+            - end_chunk (int): exclusive end chunk index
+            - table (pa.Table): concatenated table containing all rows to write
+            - worker_id (int): numeric worker identifier
+            - output_dir (str): base output directory
+            - samples_per_file (int): rows per chunk file
+            - compression (str): compression codec for Parquet
+
+    Returns:
+        int: Total number of rows written by this worker.
+    """
+    start_chunk, end_chunk, table, worker_id, output_dir, samples_per_file, compression = args
+    total_written = 0
+    num_samples = len(table)
+
+    worker_dir = os.path.join(output_dir, f"worker_{worker_id}")
+    os.makedirs(worker_dir, exist_ok=True)
+
+    # Offset to continue numbering if files exist
+    num_parquets = 0
+    for root, _, files in os.walk(worker_dir):
+        for file in files:
+            if file.endswith('.parquet'):
+                num_parquets += 1
+
+    for i in range(start_chunk, end_chunk):
+        start_sample = i * samples_per_file
+        end_sample = min((i + 1) * samples_per_file, num_samples)
+        if end_sample <= start_sample:
+            continue
+        chunk = table.slice(start_sample, end_sample - start_sample)
+
+        chunk_path = os.path.join(worker_dir, f"data_chunk_{i + num_parquets}.parquet")
+        temp_path = chunk_path + '.tmp'
+        try:
+            pq.write_table(chunk, temp_path, compression=compression)
+            if os.path.exists(chunk_path):
+                os.remove(chunk_path)
+            os.rename(temp_path, chunk_path)
+            total_written += len(chunk)
+        except Exception:
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+            raise
+
+    return total_written
+
+
+

fastvideo/dataset/dataloader/schema.py

@@ -50,6 +50,7 @@ pyarrow_schema_i2v = pa.schema([
     pa.field("fps", pa.float64()),
 ])
 
+
 pyarrow_schema_t2v = pa.schema([
     pa.field("id", pa.string()),
     # --- Image/Video VAE latents ---
@@ -78,3 +79,26 @@ pyarrow_schema_t2v = pa.schema([
     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
+])

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/dataset/utils.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 
 
-def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
+def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Pad or crop an embedding [L, D] to exactly padding_length tokens.
     Return:

fastvideo/entrypoints/video_generator.py

@@ -344,6 +344,9 @@ class VideoGenerator:
                 "size": (target_height, target_width, batch.num_frames),
                 "generation_time": gen_time,
                 "logging_info": logging_info,
+                "trajectory": output_batch.trajectory_latents,
+                "trajectory_timesteps": output_batch.trajectory_timesteps,
+                "trajectory_decoded": output_batch.trajectory_decoded,
             }
 
     def set_lora_adapter(self,

fastvideo/pipelines/composed_pipeline_base.py

@@ -40,7 +40,7 @@ class ComposedPipelineBase(ABC):
     _extra_config_module_map: dict[str, str] = {}
     training_args: TrainingArgs | None = None
     fastvideo_args: FastVideoArgs | TrainingArgs | None = None
-    modules: dict[str, torch.nn.Module] = {}
+    modules: dict[str, Any] = {}
     post_init_called: bool = False
 
     # TODO(will): args should support both inference args and training args

fastvideo/pipelines/pipeline_batch_info.py

@@ -147,7 +147,12 @@ class ForwardBatch:
     modules: dict[str, Any] = field(default_factory=dict)
 
     # Final output (after pipeline completion)
-    output: Any = None
+    output: torch.Tensor | None = None
+    return_trajectory_latents: bool = False
+    return_trajectory_decoded: bool = False
+    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_decoded: list[torch.Tensor] | None = None
 
     # Extra parameters that might be needed by specific pipeline implementations
     extra: dict[str, Any] = field(default_factory=dict)
@@ -206,6 +211,10 @@ class TrainingBatch:
     infos: list[dict[str, Any]] | None = None
     mask_lat_size: torch.Tensor | None = None
 
+    # ODE trajectory supervision
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_timesteps: torch.Tensor | None = None
+
     # Transformer inputs
     noisy_model_input: torch.Tensor | None = None
     timesteps: torch.Tensor | None = None

fastvideo/pipelines/preprocess/preprocess_pipeline_base.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-import multiprocessing
 import os
-from concurrent.futures import ProcessPoolExecutor
 from typing import Any
 
 import numpy as np
@@ -19,6 +17,8 @@ from fastvideo.logger import init_logger
 from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages import TextEncodingStage
+from fastvideo.workflow.preprocess.parquet_io import (ParquetDatasetWriter,
+                                                      records_to_table)
 
 logger = init_logger(__name__)
 
@@ -54,10 +54,14 @@ class BasePreprocessPipeline(ComposedPipelineBase):
         """Get additional features specific to the pipeline type. Override in subclasses."""
         return {}
 
-    def get_schema_fields(self) -> list[str]:
-        """Get the schema fields for the pipeline type. Override in subclasses."""
+    def get_pyarrow_schema(self) -> pa.Schema:
+        """Return the PyArrow schema for this pipeline. Must be overridden."""
         raise NotImplementedError
 
+    def get_schema_fields(self) -> list[str]:
+        """Get the schema fields for the pipeline type."""
+        return [f.name for f in self.get_pyarrow_schema()]
+
     def create_record_for_schema(self,
                                  preprocess_batch: PreprocessBatch,
                                  schema: pa.Schema,
@@ -400,166 +404,26 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                 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())
+                table = records_to_table(batch_data, self.get_pyarrow_schema())
                 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)
-
+                if not hasattr(self, 'dataset_writer'):
+                    self.dataset_writer = ParquetDatasetWriter(
+                        out_dir=combined_parquet_dir,
+                        samples_per_file=args.samples_per_file,
+                    )
+                self.dataset_writer.append_table(table)
                 logger.info("Collected batch with %s samples", len(table))
 
             if num_processed_samples >= args.flush_frequency:
-                self._flush_tables(num_processed_samples, args,
-                                   combined_parquet_dir)
+                written = self.dataset_writer.flush()
+                logger.info("Flushed %s samples to parquet", written)
                 num_processed_samples = 0
-                self.all_tables = []
 
-    def _flush_tables(self, num_processed_samples: int, args,
-                      combined_parquet_dir: str):
-        """Flush collected tables to disk."""
-        assert hasattr(self, 'all_tables') and self.all_tables
-        print(f"Combining {len(self.all_tables)} batches...")
-        combined_table = pa.concat_tables(self.all_tables)
-        assert len(combined_table) == num_processed_samples
-        print(f"Total samples collected: {len(combined_table)}")
-
-        # Calculate total number of chunks needed, discarding remainder
-        total_chunks = max(num_processed_samples // args.samples_per_file, 1)
-
-        print(f"Fixed samples per parquet file: {args.samples_per_file}")
-        print(f"Total number of parquet files: {total_chunks}")
-        print(
-            f"Total samples to be processed: {total_chunks * args.samples_per_file} (discarding {num_processed_samples % args.samples_per_file} samples)"
-        )
-
-        # Split work among processes
-        num_workers = int(min(multiprocessing.cpu_count(), total_chunks))
-        chunks_per_worker = (total_chunks + num_workers - 1) // num_workers
-
-        print(f"Using {num_workers} workers to process {total_chunks} chunks")
-        logger.info("Chunks per worker: %s", chunks_per_worker)
-
-        # Prepare work ranges
-        work_ranges = []
-        for i in range(num_workers):
-            start_idx = i * chunks_per_worker
-            end_idx = min((i + 1) * chunks_per_worker, total_chunks)
-            if start_idx < total_chunks:
-                work_ranges.append(
-                    (start_idx, end_idx, combined_table, i,
-                     combined_parquet_dir, args.samples_per_file))
-
-        total_written = 0
-        failed_ranges = []
-        with ProcessPoolExecutor(max_workers=num_workers) as executor:
-            futures = {
-                executor.submit(self.process_chunk_range, work_range):
-                work_range
-                for work_range in work_ranges
-            }
-            for future in tqdm(futures, desc="Processing chunks"):
-                try:
-                    written = future.result()
-                    total_written += written
-                    logger.info("Processed chunk with %s samples", written)
-                except Exception as e:
-                    work_range = futures[future]
-                    failed_ranges.append(work_range)
-                    logger.error("Failed to process range %s-%s: %s",
-                                 work_range[0], work_range[1], str(e))
-
-        # Retry failed ranges sequentially
-        if failed_ranges:
-            logger.warning("Retrying %s failed ranges sequentially",
-                           len(failed_ranges))
-            for work_range in failed_ranges:
-                try:
-                    total_written += self.process_chunk_range(work_range)
-                except Exception as e:
-                    logger.error(
-                        "Failed to process range %s-%s after retry: %s",
-                        work_range[0], work_range[1], str(e))
-
-        logger.info("Total samples written: %s", total_written)
-
-    @staticmethod
-    def process_chunk_range(args: Any) -> int:
-        start_idx, end_idx, table, worker_id, output_dir, samples_per_file = args
-        try:
-            total_written = 0
-            num_samples = len(table)
-
-            # Create worker-specific subdirectory
-            worker_dir = os.path.join(output_dir, f"worker_{worker_id}")
-            os.makedirs(worker_dir, exist_ok=True)
-
-            # Check how many files there are already in the dir, and update i accordingly
-            num_parquets = 0
-            for root, _, files in os.walk(worker_dir):
-                for file in files:
-                    if file.endswith('.parquet'):
-                        num_parquets += 1
-
-            for i in range(start_idx, end_idx):
-                start_sample = i * samples_per_file
-                end_sample = min((i + 1) * samples_per_file, num_samples)
-                chunk = table.slice(start_sample, end_sample - start_sample)
-
-                # Create chunk file in worker's directory
-                chunk_path = os.path.join(
-                    worker_dir, f"data_chunk_{i + num_parquets}.parquet")
-                temp_path = chunk_path + '.tmp'
-
-                try:
-                    # Write to temporary file
-                    pq.write_table(chunk, temp_path, compression='zstd')
-
-                    # Rename temporary file to final file
-                    if os.path.exists(chunk_path):
-                        os.remove(
-                            chunk_path)  # Remove existing file if it exists
-                    os.rename(temp_path, chunk_path)
-
-                    total_written += len(chunk)
-                except Exception as e:
-                    # Clean up temporary file if it exists
-                    if os.path.exists(temp_path):
-                        os.remove(temp_path)
-                    raise e
-
-            return total_written
-        except Exception as e:
-            logger.error("Error processing chunks %s-%s for worker %s: %s",
-                         start_idx, end_idx, worker_id, str(e))
-            raise
+    def _final_flush_if_any(self):
+        if hasattr(self, 'dataset_writer'):
+            self.dataset_writer.flush()

fastvideo/pipelines/preprocess/preprocess_pipeline_i2v.py

@@ -40,9 +40,9 @@ class PreprocessPipeline_I2V(BasePreprocessPipeline):
                            image_processor=self.get_module("image_processor"),
                        ))
 
-    def get_schema_fields(self) -> list[str]:
-        """Get the schema fields for I2V pipeline."""
-        return [f.name for f in pyarrow_schema_i2v]
+    def get_pyarrow_schema(self):
+        """Return the PyArrow schema for I2V pipeline."""
+        return pyarrow_schema_i2v
 
     def get_extra_features(self, valid_data: dict[str, Any],
                            fastvideo_args: FastVideoArgs) -> dict[str, Any]:

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -0,0 +1,399 @@
+# 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 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 gettextdataset
+from fastvideo.dataset.dataloader.schema import (
+    pyarrow_schema_ode_trajectory_text_only)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (
+    SelfForcingFlowMatchScheduler)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
+from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
+from fastvideo.workflow.preprocess.parquet_io import (ParquetDatasetWriter,
+                                                      records_to_table)
+
+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]]
+    pbar: Any
+    num_processed_samples: int
+
+    def get_pyarrow_schema(self) -> pa.Schema:
+        """Return the PyArrow schema for ODE Trajectory pipeline."""
+        return pyarrow_schema_ode_trajectory_text_only
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        logger.info('WTF flow_shift: %s',
+                    fastvideo_args.pipeline_config.flow_shift)
+        assert fastvideo_args.pipeline_config.flow_shift == 5
+        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+        # shift=fastvideo_args.pipeline_config.flow_shift)
+        self.modules["scheduler"] = SelfForcingFlowMatchScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift,
+            sigma_min=0.0,
+            extra_one_step=True)
+        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
+                                                denoising_strength=1.0)
+        # logger.info('WTF scheduler timesteps: %s',
+        #             self.modules["scheduler"].timesteps)
+        # scheduler = FlowMatchScheduler(
+        #     shift=8.0, sigma_min=0.0, extra_one_step=True)
+        # device = get_local_torch_device()
+        # # scheduler.num_train_timesteps = 100
+        # scheduler.set_timesteps(num_inference_steps=50, denoising_strength=1.0)
+        # scheduler.sigmas = scheduler.sigmas.to(device)
+        # self.modules["scheduler"] = scheduler
+
+        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="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"),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self,
+                       ))
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    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,
+                            strict=False)):
+                    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.num_inference_steps = 48
+                    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 = 6.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
+                }
+
+                if batch.return_trajectory_decoded:
+                    for i, decoded_frames in enumerate(trajectory_decoded):
+                        for j, decoded_frame in enumerate(decoded_frames):
+                            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():
+                            if isinstance(value, torch.Tensor):
+                                sample_extra_features[key] = value[idx].cpu(
+                                ).numpy()
+                            else:
+                                assert isinstance(value, list)
+                                if isinstance(value[idx], torch.Tensor):
+                                    sample_extra_features[key] = value[idx].cpu(
+                                    ).float().numpy()
+                                else:
+                                    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:
+                    write_pbar = tqdm(total=1,
+                                      desc="Writing to Parquet dataset",
+                                      unit="batch")
+                    table = records_to_table(batch_data,
+                                             self.get_pyarrow_schema())
+                    write_pbar.update(1)
+                    write_pbar.close()
+
+                    if not hasattr(self, 'dataset_writer'):
+                        self.dataset_writer = ParquetDatasetWriter(
+                            out_dir=self.combined_parquet_dir,
+                            samples_per_file=args.samples_per_file,
+                        )
+                    self.dataset_writer.append_table(table)
+
+                    logger.info("Collected batch with %s samples", len(table))
+
+                if self.num_processed_samples >= args.flush_frequency:
+                    written = self.dataset_writer.flush()
+                    logger.info("Flushed %s samples to parquet", written)
+                    self.num_processed_samples = 0
+
+        # Final flush for any remaining samples
+        if hasattr(self, 'dataset_writer'):
+            written = self.dataset_writer.flush()
+            if written:
+                logger.info("Final flush wrote %s samples", written)
+
+    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",
+        }
+
+        assert extra_features is not None, "extra_features is required"
+        assert "trajectory_latents" in extra_features, "trajectory_latents is required"
+        assert "trajectory_timesteps" in extra_features, "trajectory_timesteps is required"
+
+        # 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 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 = 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_text_and_trajectory(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_ODE_Trajectory

fastvideo/pipelines/preprocess/preprocess_pipeline_t2v.py

@@ -15,9 +15,9 @@ class PreprocessPipeline_T2V(BasePreprocessPipeline):
 
     _required_config_modules = ["text_encoder", "tokenizer", "vae"]
 
-    def get_schema_fields(self):
-        """Get the schema fields for T2V pipeline."""
-        return [f.name for f in pyarrow_schema_t2v]
+    def get_pyarrow_schema(self):
+        """Return the PyArrow schema for T2V pipeline."""
+        return pyarrow_schema_t2v
 
 
 EntryClass = PreprocessPipeline_T2V

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -9,6 +9,8 @@ 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.utils import maybe_download_model
@@ -24,7 +26,8 @@ def main(args) -> None:
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
         "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+        "flow_shift": 5,
     }
     pipeline_config.update_config_from_dict(kwargs)
     fastvideo_args = FastVideoArgs(
@@ -35,7 +38,17 @@ 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":
+        PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
+    else:
+        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}")
+
+    logger.info("Preprocess task: %s using %s", args.preprocess_task,
+                PreprocessPipeline.__name__)
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 

fastvideo/pipelines/stages/decoding.py

@@ -50,6 +50,63 @@ class DecodingStage(PipelineStage):
         result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
         return result
 
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor,
+               fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        """
+        Decode latent representations into pixel space using VAE.
+        
+        Args:
+            latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)
+            fastvideo_args: Configuration containing:
+                - disable_autocast: Whether to disable automatic mixed precision (default: False)
+                - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16")
+                - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency
+            
+        Returns:
+            Decoded video tensor with shape (batch, channels, frames, height, width), 
+            normalized to [0, 1] range and moved to CPU as float32
+        """
+        self.vae = self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latents = latents / self.vae.scaling_factor.to(
+                latents.device, latents.dtype)
+        else:
+            latents = latents / self.vae.scaling_factor
+
+        # Apply shifting if needed
+        if (hasattr(self.vae, "shift_factor")
+                and self.vae.shift_factor is not None):
+            if isinstance(self.vae.shift_factor, torch.Tensor):
+                latents += self.vae.shift_factor.to(latents.device,
+                                                    latents.dtype)
+            else:
+                latents += self.vae.shift_factor
+
+        # Decode latents
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            # if fastvideo_args.vae_sp:
+            #     self.vae.enable_parallel()
+            if not vae_autocast_enabled:
+                latents = latents.to(vae_dtype)
+            image = self.vae.decode(latents)
+
+        # Normalize image to [0, 1] range
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image
+
     @torch.no_grad()
     def forward(
         self,
@@ -59,13 +116,28 @@ class DecodingStage(PipelineStage):
         """
         Decode latent representations into pixel space.
         
+        This method processes the batch through the VAE decoder, converting latent
+        representations to pixel-space video/images. It also optionally decodes
+        trajectory latents for visualization purposes.
+        
         Args:
-            batch: The current batch information.
-            fastvideo_args: The inference arguments.
+            batch: The current batch containing:
+                - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents)
+                - return_trajectory_decoded (optional): Flag to decode trajectory latents
+                - trajectory_latents (optional): Latents at different timesteps
+                - trajectory_timesteps (optional): Corresponding timesteps
+            fastvideo_args: Configuration containing:
+                - output_type: "latent" to skip decoding, otherwise decode to pixels
+                - vae_cpu_offload: Whether to offload VAE to CPU after decoding
+                - model_loaded: Track VAE loading state
+                - model_paths: Path to VAE model if loading needed
             
         Returns:
-            The batch with decoded outputs.
+            Modified batch with:
+                - output: Decoded frames (batch, channels, frames, height, width) as CPU float32
+                - trajectory_decoded (if requested): List of decoded frames per timestep
         """
+        # load vae if not already loaded (used for memory constrained devices)
         pipeline = self.pipeline() if self.pipeline else None
         if not fastvideo_args.model_loaded["vae"]:
             loader = VAELoader()
@@ -75,58 +147,29 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        self.vae = self.vae.to(get_local_torch_device())
-
-        latents = batch.latents
-        # TODO(will): remove this once we add input/output validation for stages
-        if latents is None:
-            raise ValueError("Latents must be provided")
-
-        # Skip decoding if output type is latent
         if fastvideo_args.output_type == "latent":
-            image = latents
+            frames = batch.latents
         else:
-            # Setup VAE precision
-            vae_dtype = PRECISION_TO_TYPE[
-                fastvideo_args.pipeline_config.vae_precision]
-            vae_autocast_enabled = (vae_dtype != torch.float32
-                                    ) and not fastvideo_args.disable_autocast
+            frames = self.decode(batch.latents, fastvideo_args)
 
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                latents = latents / self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                latents = latents / self.vae.scaling_factor
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents += self.vae.shift_factor.to(latents.device,
-                                                        latents.dtype)
-                else:
-                    latents += self.vae.shift_factor
-
-            # Decode latents
-            with torch.autocast(device_type="cuda",
-                                dtype=vae_dtype,
-                                enabled=vae_autocast_enabled):
-                if fastvideo_args.pipeline_config.vae_tiling:
-                    self.vae.enable_tiling()
-                # if fastvideo_args.vae_sp:
-                #     self.vae.enable_parallel()
-                if not vae_autocast_enabled:
-                    latents = latents.to(vae_dtype)
-                image = self.vae.decode(latents)
-
-        # Normalize image to [0, 1] range
-        image = (image / 2 + 0.5).clamp(0, 1)
+        # decode trajectory latents if needed
+        if batch.return_trajectory_decoded:
+            batch.trajectory_decoded = []
+            assert batch.trajectory_latents is not None, "batch should have trajectory latents"
+            for idx in range(batch.trajectory_latents.shape[1]):
+                # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
+                cur_latent = batch.trajectory_latents[:, idx, :, :, :, :]
+                cur_timestep = batch.trajectory_timesteps[idx]
+                logger.info("decoding trajectory latent for timestep: %s",
+                            cur_timestep)
+                decoded_frames = self.decode(cur_latent, fastvideo_args)
+                batch.trajectory_decoded.append(decoded_frames.cpu().float())
 
         # Convert to CPU float32 for compatibility
-        image = image.cpu().float()
+        frames = frames.cpu().float()
 
         # Update batch with decoded image
-        batch.output = image
+        batch.output = frames
 
         # Offload models if needed
         if hasattr(self, 'maybe_free_model_hooks'):

fastvideo/pipelines/stages/denoising.py

@@ -247,6 +247,10 @@ class DenoisingStage(PipelineStage):
                                                              patch_size[2])
             seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size
 
+        # Initialize lists for ODE trajectory
+        trajectory_timesteps: list[torch.Tensor] = []
+        trajectory_latents: list[torch.Tensor] = []
+
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -427,6 +431,11 @@ class DenoisingStage(PipelineStage):
                         latents = (1. - mask2[0]) * z + mask2[0] * latents
                         # latents = latents.unsqueeze(0)
 
+                # save trajectory latents if needed
+                if batch.return_trajectory_latents:
+                    trajectory_timesteps.append(t)
+                    trajectory_latents.append(latents)
+
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
                     (i + 1) > num_warmup_steps and
@@ -434,9 +443,28 @@ class DenoisingStage(PipelineStage):
                         and progress_bar is not None):
                     progress_bar.update()
 
+        # Gather results if using sequence parallelism
+        trajectory_tensor: torch.Tensor | None = None
+        if trajectory_latents:
+            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
+            trajectory_timesteps_tensor = torch.stack(trajectory_timesteps,
+                                                      dim=0)
+        else:
+            trajectory_tensor = None
+            trajectory_timesteps_tensor = None
+
         # Gather results if using sequence parallelism
         if sp_group:
             latents = sequence_model_parallel_all_gather(latents, dim=2)
+            if batch.return_trajectory_latents:
+                trajectory_tensor = trajectory_tensor.to(
+                    get_local_torch_device())
+                trajectory_tensor = sequence_model_parallel_all_gather(
+                    trajectory_tensor, dim=3)
+
+        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
+            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
+            batch.trajectory_latents = trajectory_tensor.cpu()
 
         # Update batch with final latents
         batch.latents = latents

fastvideo/tests/dataset/test_parquet_io.py

@@ -0,0 +1,108 @@
+import os
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from fastvideo.dataset.dataloader.parquet_io import (
+    ParquetDatasetWriter,
+    records_to_table,
+)
+
+
+def test_records_to_table_types():
+    schema = pa.schema([
+        pa.field("id", pa.string()),
+        pa.field("vae_latent_bytes", pa.binary()),
+        pa.field("vae_latent_shape", pa.list_(pa.int64())),
+        pa.field("duration_sec", pa.float64()),
+        pa.field("width", pa.int64()),
+    ])
+    records = [{
+        "id": "a",
+        "vae_latent_bytes": b"\x00\x01",
+        "vae_latent_shape": [1, 2, 3],
+        "duration_sec": 1.5,
+        "width": 640,
+    }]
+
+    table = records_to_table(records, schema)
+    assert table.schema == schema
+    assert table.num_rows == 1
+    cols = {name: table.column(name).to_pylist()[0] for name in schema.names}
+    assert cols["id"] == "a"
+    assert isinstance(cols["vae_latent_bytes"], (bytes, bytearray))
+    assert cols["vae_latent_shape"] == [1, 2, 3]
+    assert abs(cols["duration_sec"] - 1.5) < 1e-6
+    assert cols["width"] == 640
+
+
+def test_writer_flush_and_remainder(tmp_path: Path):
+    schema = pa.schema([pa.field("id", pa.string())])
+    records = [{"id": str(i)} for i in range(25)]
+    table = records_to_table(records, schema)
+
+    out_dir = tmp_path / "out"
+    writer = ParquetDatasetWriter(str(out_dir), samples_per_file=10)
+    writer.append_table(table)
+    written = writer.flush(num_workers=1)
+    assert written == 20
+
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 2
+    total_rows = sum(pq.read_table(str(f)).num_rows for f in files)
+    assert total_rows == 20
+
+    # Append remainder to complete another chunk
+    extra = records_to_table([{"id": str(i)} for i in range(5)], schema)
+    writer.append_table(extra)
+    written2 = writer.flush(num_workers=1)
+    assert written2 == 10
+    files2 = sorted(out_dir.rglob("*.parquet"))
+    assert len(files2) == 3
+    total_rows2 = sum(pq.read_table(str(f)).num_rows for f in files2)
+    assert total_rows2 == 30
+
+
+def test_writer_flush_write_remainder(tmp_path: Path):
+    schema = pa.schema([pa.field("id", pa.string())])
+    # 25 rows, 10 per file => 2 full files + 1 remainder(5)
+    records = [{"id": str(i)} for i in range(25)]
+    table = records_to_table(records, schema)
+
+    out_dir = tmp_path / "out_last"
+    writer = ParquetDatasetWriter(str(out_dir), samples_per_file=10)
+    writer.append_table(table)
+    # First flush writes 20
+    written1 = writer.flush(num_workers=1)
+    assert written1 == 20
+    # Final flush with remainder
+    written2 = writer.flush(num_workers=1, write_remainder=True)
+    assert written2 == 5
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 3
+    total_rows = sum(pq.read_table(str(f)).num_rows for f in files)
+    assert total_rows == 25
+
+
+def test_writer_parallel_workers(tmp_path: Path):
+    schema = pa.schema([pa.field("id", pa.string())])
+    # 40 rows, 10 per file => 4 files
+    records = [{"id": str(i)} for i in range(40)]
+    table = records_to_table(records, schema)
+
+    out_dir = tmp_path / "out_parallel"
+    writer = ParquetDatasetWriter(str(out_dir), samples_per_file=10)
+    writer.append_table(table)
+    written = writer.flush(num_workers=2)
+    assert written == 40
+
+    # Ensure files exist under worker subdirs
+    worker_dirs = [p for p in out_dir.iterdir() if p.is_dir() and p.name.startswith("worker_")]
+    assert len(worker_dirs) >= 1
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 4
+    total_rows = sum(pq.read_table(str(f)).num_rows for f in files)
+    assert total_rows == 40
+
+

fastvideo/tests/workflow/test_components_parquet_saver.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+
+from fastvideo.workflow.preprocess.components import ParquetDatasetSaver
+from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
+
+
+def _simple_record_creator(batch: PreprocessBatch) -> list[dict]:
+    # batch.latents will be converted to numpy by the saver before this call
+    assert isinstance(batch.latents, np.ndarray)
+    num = len(batch.video_file_name)
+    records = []
+    for i in range(num):
+        arr = batch.latents[i]
+        records.append({
+            "id": batch.video_file_name[i],
+            "data_bytes": arr.tobytes(),
+            "data_shape": list(arr.shape),
+        })
+    return records
+
+
+def test_parquet_dataset_saver_flush_and_last(tmp_path: Path):
+    # Schema for the simple record creator
+    schema = pa.schema([
+        pa.field("id", pa.string()),
+        pa.field("data_bytes", pa.binary()),
+        pa.field("data_shape", pa.list_(pa.int64())),
+    ])
+
+    B = 5
+    # Build a minimal PreprocessBatch
+    batch = PreprocessBatch(
+        data_type="video",
+        latents=torch.randn(B, 2),
+        prompt_embeds=[torch.randn(B, 1, 1)],
+        prompt_attention_mask=[torch.ones(B, 1)],
+    )
+    batch.video_file_name = [f"vid_{i}" for i in range(B)]
+
+    saver = ParquetDatasetSaver(
+        flush_frequency=10,  # higher than B to avoid auto-flush
+        samples_per_file=3,
+        schema=schema,
+        record_creator=_simple_record_creator,
+    )
+
+    out_dir = tmp_path / "saver_out"
+    saver.save_and_write_parquet_batch(batch, str(out_dir))
+    # First flush: should write one full file (3 rows), keep 2 in buffer
+    saver.flush_tables(str(out_dir))
+    files = sorted(out_dir.rglob("*.parquet"))
+    assert len(files) == 1
+    assert pq.read_table(str(files[0])).num_rows == 3
+
+    # Final flush: write remainder 2 rows
+    saver.flush_last(str(out_dir))
+    files2 = sorted(out_dir.rglob("*.parquet"))
+    assert len(files2) == 2
+    total = sum(pq.read_table(str(f)).num_rows for f in files2)
+    assert total == 5
+
+

fastvideo/utils.py

@@ -23,10 +23,14 @@ from typing import Any, TypeVar, cast
 
 import cloudpickle
 import filelock
+import imageio
+import numpy as np
 import torch
+import torchvision
 import yaml
 from diffusers.loaders.lora_base import (
     _best_guess_weight_name)  # watch out for potetential removal from diffusers
+from einops import rearrange
 from huggingface_hub import snapshot_download
 from remote_pdb import RemotePdb
 from torch.distributed.fsdp import MixedPrecisionPolicy
@@ -886,3 +890,17 @@ def best_output_size(w, h, dw, dh, expected_area):
         return ow1, oh1
     else:
         return ow2, oh2
+
+
+def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor],
+                                  output_path: str, fps: int):
+    # Process outputs
+    videos = rearrange(decoded_latents, "b c t h w -> t b c h w")
+    frames = []
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=6)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+        frames.append((x * 255).numpy().astype(np.uint8))
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    imageio.mimsave(output_path, frames, fps=fps, format="mp4")

fastvideo/workflow/preprocess/components.py

@@ -1,20 +1,19 @@
 import dataclasses
 import gc
-import multiprocessing
 import os
 import random
 from collections.abc import Callable
-from concurrent.futures import ProcessPoolExecutor
 from typing import Any
 
 import numpy as np
 import pyarrow as pa
-import pyarrow.parquet as pq
 import torch
 from datasets import Dataset, Video, load_dataset
 
 from fastvideo.configs.configs import (DatasetType, PreprocessConfig,
                                        VideoLoaderType)
+from fastvideo.dataset.dataloader.parquet_io import (ParquetDatasetWriter,
+                                                     records_to_table)
 from fastvideo.distributed.parallel_state import get_world_rank, get_world_size
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
@@ -155,31 +154,17 @@ class VideoForwardBatchBuilder:
 
 
 class ParquetDatasetSaver:
-    """Component for saving and writing Parquet datasets"""
+    """Component for saving and writing Parquet datasets using shared parquet_io."""
 
-    def __init__(self,
-                 flush_frequency: int,
-                 samples_per_file: int,
-                 schema_fields: list[str],
-                 record_creator: Callable[..., list[dict[str, Any]]],
-                 file_writer_fn: Callable | None = None):
-        """
-        Initialize ParquetDatasetSaver
-        
-        Args:
-            schema_fields: schema fields list
-            record_creator: Function for creating records
-            file_writer_fn: Function for writing records to files, uses default implementation if None
-        """
+    def __init__(self, flush_frequency: int, samples_per_file: int,
+                 schema: pa.Schema,
+                 record_creator: Callable[..., list[dict[str, Any]]]):
         self.flush_frequency = flush_frequency
         self.samples_per_file = samples_per_file
-        self.schema_fields = schema_fields
+        self.schema = schema
         self.create_records_from_batch = record_creator
-        self.file_writer_fn: Callable[
-            [tuple], int] = file_writer_fn or self._default_file_writer_fn
-        self.all_tables: list[pa.Table] = []
         self.num_processed_samples: int = 0
-        self.num_saved_files: int = 0
+        self._writer: ParquetDatasetWriter | None = None
 
     def save_and_write_parquet_batch(
             self,
@@ -227,11 +212,12 @@ class ParquetDatasetSaver:
 
         if batch_data:
             self.num_processed_samples += len(batch_data)
-            # Convert batch data to PyArrow arrays
-            table = self._convert_batch_to_pyarrow_table(batch_data)
-
-            # Store the table in a list for later processing
-            self.all_tables.append(table)
+            table = records_to_table(batch_data, self.schema)
+            if self._writer is None:
+                os.makedirs(output_dir, exist_ok=True)
+                self._writer = ParquetDatasetWriter(
+                    out_dir=output_dir, samples_per_file=self.samples_per_file)
+            self._writer.append_table(table)
             logger.debug("Collected batch with %s samples", len(table))
 
         # If flush is needed
@@ -261,145 +247,35 @@ class ParquetDatasetSaver:
 
     def _convert_batch_to_pyarrow_table(self,
                                         batch_data: list[dict]) -> pa.Table:
-        """Convert batch data to PyArrow table"""
-        arrays = []
+        # Deprecated path, kept for backward compatibility if needed.
+        return records_to_table(batch_data, self.schema)
 
-        for field in self.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]))
+    def flush_tables(self, output_dir: str, write_remainder: bool = False):
+        """Flush buffered records to disk.
 
-        return pa.Table.from_arrays(arrays, names=self.schema_fields)
-
-    def flush_tables(self, output_dir: str):
-        """Flush collected tables to disk"""
-        if not hasattr(self, 'all_tables') or not self.all_tables:
+        Args:
+            output_dir: Directory where parquet files are written. Kept for API
+                symmetry (writer already configured with this path).
+            write_remainder: If True, also write any leftover rows smaller than
+                ``samples_per_file`` as a final small file. Useful for the last flush.
+        """
+        if self._writer is None:
             return
+        _ = self._writer.flush(write_remainder=write_remainder)
+        # Reset processed sample count modulo samples_per_file
+        remainder = self.num_processed_samples % self.samples_per_file
+        self.num_processed_samples = 0 if write_remainder else remainder
 
-        logger.debug("Combining %d batches...", len(self.all_tables))
-        combined_table = pa.concat_tables(self.all_tables)
-        assert len(combined_table) == self.num_processed_samples
-        logger.debug("Total samples collected: %d", len(combined_table))
-
-        # Calculate total number of chunks needed, putting remainder into self.all_tables
-        total_files = max(self.num_processed_samples // self.samples_per_file,
-                          1)
-
-        logger.debug("Fixed samples per parquet file: %d",
-                     self.samples_per_file)
-        logger.debug("Total number of parquet files: %d", total_files)
-        logger.debug(
-            "Total samples to be processed: %d (putting %d samples into self.all_tables)",
-            total_files * self.samples_per_file,
-            self.num_processed_samples % self.samples_per_file)
-
-        # Split work among processes
-        num_workers = int(min(multiprocessing.cpu_count(), total_files))
-        files_per_worker = (total_files + num_workers - 1) // num_workers
-
-        logger.debug("Using %d workers to process %d files", num_workers,
-                     total_files)
-        logger.debug("Files per worker: %s", files_per_worker)
-
-        # Prepare work ranges
-        work_ranges = []
-        for i in range(num_workers):
-            start_idx = i * files_per_worker
-            end_idx = min((i + 1) * files_per_worker, total_files)
-            if start_idx < total_files:
-                work_ranges.append((start_idx, end_idx, combined_table, i,
-                                    output_dir, self.samples_per_file))
-
-        total_written = 0
-        failed_ranges = []
-        with ProcessPoolExecutor(max_workers=num_workers) as executor:
-            futures = {
-                executor.submit(self.file_writer_fn, work_range): work_range
-                for work_range in work_ranges
-            }
-            for future in futures:
-                try:
-                    written = future.result()
-                    total_written += written
-                    logger.info("Processed file with %s samples", written)
-                except Exception as e:
-                    work_range = futures[future]
-                    failed_ranges.append(work_range)
-                    logger.error("Failed to process range %s-%s: %s",
-                                 work_range[0], work_range[1], str(e))
-
-        # Retry failed ranges sequentially
-        if failed_ranges:
-            logger.warning("Retrying %s failed ranges sequentially",
-                           len(failed_ranges))
-            for work_range in failed_ranges:
-                try:
-                    total_written += self.file_writer_fn(work_range)
-                except Exception as e:
-                    logger.error(
-                        "Failed to process range %s-%s after retry: %s",
-                        work_range[0], work_range[1], str(e))
-
-        self.num_saved_files += total_files
-
-        # Clear tables list
-        self.all_tables = []
-        if self.num_processed_samples > self.samples_per_file:
-            saved_samples = total_files * self.samples_per_file
-            self.all_tables.append(combined_table.slice(saved_samples))
-            self.num_processed_samples -= saved_samples
-        else:
-            self.num_processed_samples = 0
-
-        del combined_table
-        gc.collect()
+    def flush_last(self, output_dir: str):
+        """Flush and write any remaining rows (final flush)."""
+        self.flush_tables(output_dir, write_remainder=True)
 
     def clean_up(self) -> None:
         """Clean up all tables"""
-        self.all_tables = []
+        self._writer = None
         self.num_processed_samples = 0
-        self.num_saved_files = 0
         gc.collect()
 
-    def _default_file_writer_fn(self, args_tuple: tuple) -> int:
-        """Default chunk processing implementation"""
-        start_idx, end_idx, combined_table, worker_id, output_dir, samples_per_file = args_tuple
-
-        written_count = 0
-        for file_idx in range(start_idx, end_idx):
-            start_row = file_idx * samples_per_file
-            end_row = min(start_row + samples_per_file, len(combined_table))
-
-            if start_row >= len(combined_table):
-                break
-
-            chunk_table = combined_table.slice(start_row, end_row - start_row)
-
-            # Write to file
-            output_file = os.path.join(
-                output_dir,
-                f"chunk_{file_idx + self.num_saved_files:06d}.parquet")
-            pq.write_table(chunk_table, output_file)
-            written_count += len(chunk_table)
-
-        return written_count
-
 
 def build_dataset(preprocess_config: PreprocessConfig, split: str,
                   validator: Callable[[dict[str, Any]], bool]) -> Dataset:

fastvideo/workflow/preprocess/preprocess_workflow.py

@@ -85,16 +85,16 @@ class PreprocessWorkflow(WorkflowBase):
         # record creator
         if self.fastvideo_args.workload_type == WorkloadType.I2V:
             record_creator = i2v_record_creator
-            schema_fields = [f.name for f in pyarrow_schema_i2v]
+            schema = pyarrow_schema_i2v
         else:
             record_creator = basic_t2v_record_creator
-            schema_fields = [f.name for f in pyarrow_schema_t2v]
+            schema = pyarrow_schema_t2v
         processed_dataset_saver = ParquetDatasetSaver(
             flush_frequency=self.fastvideo_args.preprocess_config.
             flush_frequency,
             samples_per_file=self.fastvideo_args.preprocess_config.
             samples_per_file,
-            schema_fields=schema_fields,
+            schema=schema,
             record_creator=record_creator,
         )
         self.add_component("processed_dataset_saver", processed_dataset_saver)