Refactor ode trajectory generation

Add env variable to enforce running flash attn
Test
2025-09-26 01:26:03 +00:00 · 2025-09-22 22:33:18 +00:00 · 2025-09-20 05:02:06 +00:00 · 2025-09-20 05:02:06 +00:00 · 2025-09-20 05:02:06 +00:00 · 2025-09-20 05:02:05 +00:00
116 changed files with 2716 additions and 1289 deletions
@@ -12,6 +12,9 @@ exclude: |
    scripts/.*|
    fastvideo/data_preprocess/.*|
    fastvideo/dataset/.*|
+    fastvideo/distill/.*|
+    fastvideo/distill\.py|
+    fastvideo/distill_adv\.py|
    fastvideo/models/.*|
    fastvideo/sample/.*|
    fastvideo/train\.py|
@@ -41,10 +44,10 @@ repos:
  - id: codespell
    additional_dependencies: ['tomli']
    args: ['--toml', 'pyproject.toml']
-# - repo: https://github.com/PyCQA/isort
-#   rev: 6.0.1
-#   hooks:
-#   - id: isort
+- repo: https://github.com/PyCQA/isort
+  rev: 6.0.1
+  hooks:
+  - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
  rev: v0.9.30
  hooks:
@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/q46BbX6" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" target="_blank"> <b> WeChat </b> </a> |
 </p>

 <div align="center">
@@ -155,8 +155,8 @@ If you find FastVideo useful, please considering citing our work:
 }

@article{zhang2025vsa,
-  title={Vsa: Faster video diffusion with trainable sparse attention},
-  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  title={VSA: Faster Video Diffusion with Trainable Sparse Attention},
+  author={Zhang, Peiyuan and Huang, Haofeng and Chen, Yongqi and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
  journal={arXiv preprint arXiv:2505.13389},
  year={2025}
 }
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,76 @@
+{
+  "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": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "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.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
+      "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": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "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.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "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.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "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.",
+      "image_path": null,
+      "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}
@@ -0,0 +1,13 @@
+{
+  "data": [
+    {
+      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}
@@ -16,30 +16,30 @@ export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29503
+export MASTER_PORT=29501
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY=your_wandb_api_key
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
+export WANDB_MODE=offline
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN

 # Configs
-NUM_GPUS=1
+NUM_GPUS=4

 # Model paths for Self-Forcing DMD distillation:
 GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
 FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model

-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_data_dir
+DATA_DIR="data/mixkit-64_processed/Node_0_GPU_1_File_1/combined_parquet_dataset"
+VALIDATION_DATASET_FILE="data/mixkit-64_processed/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
-  --output_dir your_output_dir
+  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -84,7 +84,7 @@ dataset_args=(
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
+  --validation_steps 100
  --validation_sampling_steps "4"
  --validation_guidance_scale "6.0" # not used for dmd inference
 )
@@ -111,7 +111,7 @@ miscellaneous_args=(
  --use_ema True
  --ema_decay 0.99
  --ema_start_step 100
-  --init_weights_from_safetensors your_ode_init_weights_path
+  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
 )

 # Self-forcing DMD arguments
@@ -148,4 +148,4 @@ torchrun \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
    "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
+    "${self_forcing_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 
@@ -41,14 +41,13 @@ NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir $OUTPUT_DIR
+  --output_dir"checkpoints/wan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -92,7 +91,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -135,4 +134,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -41,14 +41,13 @@ NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "$OUTPUT_DIR"
+  --output_dir "checkpoints/wan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -92,7 +91,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -135,4 +134,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -41,14 +41,13 @@ NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd
-  --output_dir "$OUTPUT_DIR"
+  --output_dir "checkpoints/wan_t2v_finetune"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -92,7 +91,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -134,4 +133,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
@@ -42,14 +42,13 @@ NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name Wan_distillation
-  --output_dir "$OUTPUT_DIR"
+  --output_dir "your_output_dir"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -93,11 +92,11 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 4e-6
+  --learning_rate 2e-5
  --lr_scheduler "cosine_with_min_lr"
  --min_lr_ratio 0.5
  --lr_warmup_steps 100
-  --fake_score_learning_rate 2e-6
+  --fake_score_learning_rate 1e-5
  --fake_score_lr_scheduler "cosine_with_min_lr"
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
@@ -142,4 +141,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -92,11 +92,11 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 4e-6
+  --learning_rate 2e-5
  --lr_scheduler "cosine_with_min_lr"
  --min_lr_ratio 0.5
  --lr_warmup_steps 100
-  --fake_score_learning_rate 2e-6
+  --fake_score_learning_rate 1e-5
  --fake_score_lr_scheduler "cosine_with_min_lr"
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
@@ -142,4 +142,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -16,25 +16,24 @@ NUM_GPUS=1
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
-OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "$OUTPUT_DIR"
-  --max_train_steps 4000
-  --train_batch_size 1
+  --output_dir="checkpoints/wan_t2v_finetune"
+  --max_train_steps=4000
+  --train_batch_size=1
  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
+  --gradient_accumulation_steps=1
  --num_latent_t 31
  --num_height 704
  --num_width 1280
  --num_frames 121
  --enable_gradient_checkpointing_type "full"
-  --training_state_checkpointing_steps 500
-  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps=500
+  --weight_only_checkpointing_steps=500
 )

 # Parallel arguments
@@ -69,8 +68,8 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
-  --mixed_precision "bf16"
+  --learning_rate=1e-5
+  --mixed_precision="bf16"
  --weight_decay 0.01
  --max_grad_norm 1.0
 )
@@ -108,4 +107,4 @@ torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -109,4 +109,4 @@ torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 
@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator

-# from fastvideo.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam

 OUTPUT_PATH = "video_samples"
 def main():
@@ -17,12 +17,13 @@ def main():
        vae_cpu_offload=False,
        text_encoder_cpu_offload=True,
        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        ti2v_task=True,
        # image_encoder_cpu_offload=False,
    )

-    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
    # sampling_param.num_frames = 45
-    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.image_path = "test.jpg"
    # Generate videos with the same simple API, regardless of GPU count
    prompt = (
        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
@@ -44,4 +45,4 @@ def main():


 if __name__ == "__main__":
-    main()
+    main()
@@ -19,13 +19,15 @@ def main():
    )

    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.num_frames = 81

-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+    prompts = [
+        "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.",
+        "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.",
+    ]
+
+    for prompt in prompts:
+        video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)

 if __name__ == "__main__":
    main()
@@ -1,5 +1,14 @@
 {
  "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,
@@ -19,52 +28,7 @@
      "num_frames": 77
    },
    {
-      "caption": "Elon Musk, dressed in a sleek white spacesuit with a reflective visor, walks confidently across the lunar surface. His posture is upright, and he moves steadily with purpose. The moon's rocky terrain and scattered boulders surround him, casting shadows under the dim sunlight. The background shows vast stretches of the moon's barren landscape with craters and dust clouds kicked up by his boots. The scene captures a wide shot, emphasizing the vastness and desolation of the lunar environment. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "In a dynamic action-packed sequence set in the Marvel multiverse, Spider-Man and Venom engage in an intense battle. Spider-Man, in his classic red and blue suit, swings and dodges venomous attacks from the black symbiote-covered Venom. Both characters display a range of acrobatic moves and powerful strikes. The environment is a chaotic urban landscape with crumbling buildings and neon lights, reflecting the multiversal theme. The camera captures the epic fight from various angles, including wide shots to show the scale of destruction and close-ups to highlight their fierce expressions and physical combat. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A warm, family-oriented scene depicting a father getting ready to leave the house to buy milk. The father, a middle-aged man with a kind face and a casual outfit, picks up a jacket from the coat rack. His posture is upright as he bends down slightly to put on his shoes. In the background, there are glimpses of a cozy living room with a family photograph on the wall. The camera focuses closely on the father, capturing his gentle smile and reassuring nod towards the camera before he opens the front door and steps outside. Static medium close-up shot. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Close-up shot of a man with a prosthetic hand that functions as a rocket launcher. He looks at his new hand with a mix of amazement and concern, his facial expression showing a blend of curiosity and apprehension. The prosthetic hand is sleek and metallic, with intricate details that resemble a high-tech weapon. The background is a dimly lit laboratory with various scientific equipment and monitors displaying data. The man stands in a relaxed posture, his other hand resting on his hip, as he inspects his new limb. The scene is rendered in a realistic sci-fi style, emphasizing the futuristic technology and the man's emotional response to his new appendage. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Realistic CCTV footage style, Kim Taehyung from the band BTS is involved in a drug deal, caught on camera. Kim Taehyung appears nervous and cautious, wearing casual clothing typical of a public space. He exchanges items discreetly with another person, who is partially obscured. Both individuals maintain a watchful demeanor, occasionally glancing around to ensure no one is watching them. The lighting is dim, with flickering fluorescent lights casting shadows on their faces. The background shows a typical urban setting with blurred figures moving in the distance. Static camera angle, medium close-up shot focusing on the interaction between Taehyung and the other individual. ",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "Photorealistic studio setup with professional lighting, showcasing detailed cubic dissections of experimental plastic and felt-like materials on a pristine white background. Each cube reveals intricate layers and textures of the materials, emphasizing their unique properties. The scene has a shallow depth of field initially, then slowly pulls out to reveal the full arrangement of cubes, maintaining a wide depth of field throughout the transition. ",
+      "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,
@@ -0,0 +1,99 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_70k"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "fixed_wan_ode_init_70k_6e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --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
+  --enable_gradient_checkpointing_type "full"
+)
+
+# 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[@]}"
@@ -1,5 +1,5 @@
 #!/bin/bash
-#SBATCH --job-name=2e6B8_16kFV_ode_vidprom
+#SBATCH --job-name=1e5B2_16kFV_warp_ode_vidprom
 #SBATCH --partition=main
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
@@ -7,14 +7,14 @@
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
-#SBATCH --output=ode_vidprom16k/ode_vidprom8b16k_2e-6.out
-#SBATCH --error=ode_vidprom16k/ode_vidprom8b16k_2e-6.err
+#SBATCH --output=ode_vidprom16k_warp/Dode_vidprom8b16k_1e-5.out
+#SBATCH --error=ode_vidprom16k_warp/Dode_vidprom8b16k_1e-5.err
 #SBATCH --exclusive
 set -e -x

 # Environment Setup
 source ~/conda/miniconda/bin/activate
-conda activate your-conda-env
+conda activate will-fv2

 export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
@@ -28,18 +28,18 @@ export MASTER_ADDR=${nodes[0]}
 export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_API_KEY=your-wandb-api-key
+export WANDB_MODE=online
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA

 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"


-MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="your-data-dir"
-VALIDATION_DATASET_FILE="examples/training/consistency_finetune/causal_ode_init/validation.json"
-OUTPUT_DIR="your-output-dir"
-INIT_WEIGHTS_FROM_SAFETENSORS="your-init-weights-from-safetensors" # bidirectional weights from Wan2.1-T2V-1.3B-Diffusers
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b-81/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
 NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
@@ -47,9 +47,9 @@ NUM_GPUS=8
 # Training arguments
 training_args=(
  --tracker_project_name "wan_ode_init"
-  --output_dir $OUTPUT_DIR
+  --output_dir "Dwarp_vidprom_8b16k_test_warp_1e-5"
  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "vidprom_8b16k_ode_init_2e-6"
+  --wandb_run_name "Dwarp_vidprom_8b16k_wan_ode_init_1e-5"
  # --resume_from_checkpoint "ode_init_diffusers/"
  --warp_denoising_step
  --log_visualization
@@ -60,7 +60,7 @@ training_args=(
  --num_latent_t 21
  --num_height 480
  --num_width 832
-  --num_frames 81
+  --num_frames 77
  --dmd_denoising_steps "1000,750,500,250"
  --enable_gradient_checkpointing_type "full"
 )
@@ -93,12 +93,12 @@ validation_args=(
  --validation_steps 50
  --validation_sampling_steps "50"
  --validation_guidance_scale "6.0"
-  --init_weights_from_safetensors $INIT_WEIGHTS_FROM_SAFETENSORS
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
 )

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-6
+  --learning_rate 1e-5
  --mixed_precision "bf16"
  --checkpointing_steps 500
  --weight_decay 1e-4
@@ -1,22 +1,21 @@
 #!/bin/bash
-#SBATCH --job-name=t2v
+#SBATCH --job-name=ode_vidprom2k
 #SBATCH --partition=main
-#SBATCH --nodes=8
-#SBATCH --ntasks=8
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
-#SBATCH --output=VSA_t2v_output/t2v_%j.out
-#SBATCH --error=VSA_t2v_output/t2v_%j.err
+#SBATCH --output=ode_vidprom2k_output/ode_vidprom2k.out
+#SBATCH --error=ode_vidprom2k_output/ode_vidprom2k.err
 #SBATCH --exclusive
 set -e -x

 # Environment Setup
 source ~/conda/miniconda/bin/activate
-conda activate your_env
+conda activate will-fv2

-# Basic Info
 export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
@@ -30,41 +29,45 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA

 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"

-# Configs
-NUM_GPUS=8
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR=your_data_dir
-VALIDATION_DATASET_FILE=your_validation_dataset_file
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
+NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
-  --tracker_project_name wan_t2v_VSA
-  --output_dir "checkpoints/wan_t2v_finetune_VSA"
-  --max_train_steps 4000
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_vidprom2k"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "vidprom2k_wan_ode_init_5e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 6001
  --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 81
-  # --enable_gradient_checkpointing_type "full" # if OOM enable this
+  --num_frames 77
+  # --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
 parallel_args=(
-  --num_gpus 64
+  --num_gpus $NUM_GPUS
  --sp_size 1
  --tp_size 1
-  --hsdp_replicate_dim 64
+  --hsdp_replicate_dim 8
  --hsdp_shard_dim 1
 )

@@ -76,25 +79,25 @@ model_args=(

 # Dataset arguments
 dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 4
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 1
 )

 # Validation arguments
 validation_args=(
  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
  --validation_sampling_steps "50"
-  --validation_guidance_scale "5.0"
+  --validation_guidance_scale "6.0"
 )

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 5e-6
  --mixed_precision "bf16"
-  --checkpointing_steps 1000
-  --weight_decay 0.01
+  --checkpointing_steps 2000
+  --weight_decay 1e-4
  --max_grad_norm 1.0
 )

@@ -103,31 +106,26 @@ 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
-  --flow_shift 1
-  --seed 1000
-)
-
-# VSA arguments
-vsa_args=(
-  --VSA_decay_rate 0.03 \
-  --VSA_decay_interval_steps 50 \
-  --VSA_sparsity 0.9 \
+  --enable_gradient_checkpointing_type "full"
 )

+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
 srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_training_pipeline.py \
+  --nnodes $SLURM_JOB_NUM_NODES \
+  --nproc_per_node $NUM_GPUS \
+  --node_rank $SLURM_PROCID \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/ode_causal_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,132 @@
+#!/bin/bash
+#SBATCH --job-name=ode_crush
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_crush_output/ode_crush.out
+#SBATCH --error=ode_crush_output/ode_crush.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate will-fv2
+
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# different cache dir for different processes
+export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
+export MASTER_PORT=29500
+export NODE_RANK=$SLURM_PROCID
+nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
+export MASTER_ADDR=${nodes[0]}
+export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
+NUM_GPUS=2
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_warp_2"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "2warp_fixed_wan_ode_init_5e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  # --warp_denoising_step
+  --max_train_steps 6001
+  --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
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --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 20
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 2000
+  --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
+  --enable_gradient_checkpointing_type "full"
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+srun torchrun \
+  --nnodes $SLURM_JOB_NUM_NODES \
+  --nproc_per_node $NUM_GPUS \
+  --node_rank $SLURM_PROCID \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,103 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=offline
+export TOKENIZERS_PARALLELISM=false
+export MASTER_PORT=29501
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/mixkit-64_processed/Node_0_GPU_1_File_1/combined_parquet_dataset"
+# DATA_DIR="/mnt/weka/home/hao.zhang/wl/Self-Forcing/ode_single_lmdb_sf/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "Dwarp_vidprom_8b16k_test_warp_1e-5"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "Dwarp_vidprom_8b16k_wan_ode_init_1e-5"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --warp_denoising_step
+  --log_visualization
+  --max_train_steps 10
+  --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
+  --dmd_denoising_steps "1000,750,500,250"
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --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"
+  # --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 500
+  --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
+  --seed 1024
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# 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 \
+  --master_port $MASTER_PORT \
+  --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[@]}"
@@ -0,0 +1,98 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="/mnt/weka/home/hao.zhang/wl/FastVideo2/data/crush-smol_processed_t2v_1_3b_ode_init_single"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# export CUDA_VISIBLE_DEVICES=4,5
+# 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 "overfitwan_ode_init_crush_smol"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 2001
+  # --warp_denoising_step
+  --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
+  # --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 20
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 500
+  --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
+  --enable_gradient_checkpointing_type "full"
+)
+
+# 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[@]}"
@@ -0,0 +1,100 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "debug_ode_init"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "debug_ode_init"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 1000
+  --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 81
+  --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 10
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --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
+  --enable_gradient_checkpointing_type "full"
+)
+
+# 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[@]}"
@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="data/vidprom_1.txt"
+OUTPUT_DIR="data/ode_vidprom_1_fv/"
+
+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 \
+    --flow_shift 5.0 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 1 \
+    --flush_frequency 1 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "ode_trajectory" 
@@ -0,0 +1,76 @@
+{
+  "data": [
+    {
+      "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": "Elon Musk, dressed in a sleek white spacesuit with a reflective visor, walks confidently across the lunar surface. His posture is upright, and he moves steadily with purpose. The moon's rocky terrain and scattered boulders surround him, casting shadows under the dim sunlight. The background shows vast stretches of the moon's barren landscape with craters and dust clouds kicked up by his boots. The scene captures a wide shot, emphasizing the vastness and desolation of the lunar environment. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "In a dynamic action-packed sequence set in the Marvel multiverse, Spider-Man and Venom engage in an intense battle. Spider-Man, in his classic red and blue suit, swings and dodges venomous attacks from the black symbiote-covered Venom. Both characters display a range of acrobatic moves and powerful strikes. The environment is a chaotic urban landscape with crumbling buildings and neon lights, reflecting the multiversal theme. The camera captures the epic fight from various angles, including wide shots to show the scale of destruction and close-ups to highlight their fierce expressions and physical combat. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A warm, family-oriented scene depicting a father getting ready to leave the house to buy milk. The father, a middle-aged man with a kind face and a casual outfit, picks up a jacket from the coat rack. His posture is upright as he bends down slightly to put on his shoes. In the background, there are glimpses of a cozy living room with a family photograph on the wall. The camera focuses closely on the father, capturing his gentle smile and reassuring nod towards the camera before he opens the front door and steps outside. Static medium close-up shot. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Close-up shot of a man with a prosthetic hand that functions as a rocket launcher. He looks at his new hand with a mix of amazement and concern, his facial expression showing a blend of curiosity and apprehension. The prosthetic hand is sleek and metallic, with intricate details that resemble a high-tech weapon. The background is a dimly lit laboratory with various scientific equipment and monitors displaying data. The man stands in a relaxed posture, his other hand resting on his hip, as he inspects his new limb. The scene is rendered in a realistic sci-fi style, emphasizing the futuristic technology and the man's emotional response to his new appendage. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Realistic CCTV footage style, Kim Taehyung from the band BTS is involved in a drug deal, caught on camera. Kim Taehyung appears nervous and cautious, wearing casual clothing typical of a public space. He exchanges items discreetly with another person, who is partially obscured. Both individuals maintain a watchful demeanor, occasionally glancing around to ensure no one is watching them. The lighting is dim, with flickering fluorescent lights casting shadows on their faces. The background shows a typical urban setting with blurred figures moving in the distance. Static camera angle, medium close-up shot focusing on the interaction between Taehyung and the other individual. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "Photorealistic studio setup with professional lighting, showcasing detailed cubic dissections of experimental plastic and felt-like materials on a pristine white background. Each cube reveals intricate layers and textures of the materials, emphasizing their unique properties. The scene has a shallow depth of field initially, then slowly pulls out to reveal the full arrangement of cubes, maintaining a wide depth of field throughout the transition. ",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}
@@ -6,8 +6,8 @@ export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA

 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+DATA_DIR="data/crush-smol_processed_t2v_old"
+VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5

@@ -52,7 +52,7 @@ dataset_args=(
 validation_args=(
  --log_validation 
  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
+  --validation_steps 50
  --validation_sampling_steps "50" 
  --validation_guidance_scale "6.0"
 )
@@ -4,7 +4,7 @@ GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v/"
+OUTPUT_DIR="data/crush-smol_processed_t2v_old/"

 torchrun --nproc_per_node=$GPU_NUM \
    fastvideo/pipelines/preprocess/v1_preprocess.py \
@@ -1,6 +1,6 @@
 #!/bin/bash

-GPU_NUM=2 # 2,4,8
+GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATASET_PATH="data/crush-smol/"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"
@@ -14,7 +14,7 @@ torchrun --nproc_per_node=$GPU_NUM \
    --preprocess.dataset_type merged \
    --preprocess.dataset_path $DATASET_PATH \
    --preprocess.dataset_output_dir $OUTPUT_DIR \
-    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.preprocess_video_batch_size 8 \
    --preprocess.dataloader_num_workers 0 \
    --preprocess.max_height 480 \
    --preprocess.max_width 832 \
@@ -0,0 +1,94 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_old"
+VALIDATION_DATASET_FILE="examples/datasets/crush_smol/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_i2v_finetune"
+  --output_dir "checkpoints/wan_t2v_i2v_finetune"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 2
+  --num_latent_t 20
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size 4
+  --tp_size 1
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 4
+)
+
+# 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 5e-5
+  --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
+  --enable_gradient_checkpointing_type "full"
+  --t2v_as_i2v_task True
+  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_t2v_i2v_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -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="data/crush-smol/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_i2v_1_3b/"
+
+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 2 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "t2v_ode_trajectory" 
@@ -3,6 +3,7 @@
 import torch
 import torch.nn as nn

+import fastvideo.envs as envs
 from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend
 from fastvideo.distributed.communication_op import (
    sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
@@ -36,7 +37,7 @@ class DistributedAttention(nn.Module):
        if num_kv_heads is None:
            num_kv_heads = num_heads

-        dtype = get_compute_dtype()
+        dtype = torch.bfloat16 if envs.FASTVIDEO_FORCE_ATTN_BF16 else get_compute_dtype()
        attn_backend = get_attn_backend(
            head_size,
            dtype,
@@ -221,7 +222,7 @@ class LocalAttention(nn.Module):
        if num_kv_heads is None:
            num_kv_heads = num_heads

-        dtype = get_compute_dtype()
+        dtype = torch.bfloat16 if envs.FASTVIDEO_FORCE_ATTN_BF16 else get_compute_dtype()
        attn_backend = get_attn_backend(
            head_size,
            dtype,
@@ -45,10 +45,13 @@ class PipelineConfig:
    embedded_cfg_scale: float = 6.0
    flow_shift: float | None = None
    disable_autocast: bool = False
+    ti2v_task: bool = False
+    t2v_as_i2v_task: bool = False

    # Model configuration
    dit_config: DiTConfig = field(default_factory=DiTConfig)
    dit_precision: str = "bf16"
+    dit_forward_precision: str = "bf16"

    # VAE configuration
    vae_config: VAEConfig = field(default_factory=VAEConfig)
@@ -215,6 +218,24 @@ class PipelineConfig:
            "Comma-separated list of denoising steps (e.g., '1000,757,522')",
        )

+        # TI2V task
+        parser.add_argument(
+            f"--{prefix_with_dot}ti2v-task",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}ti2v_task",
+            default=PipelineConfig.ti2v_task,
+            help="Enable TI2V",
+        )
+
+        # T2V to I2V task
+        parser.add_argument(
+            f"--{prefix_with_dot}t2v-as-i2v-task",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}t2v_as_i2v_task",
+            default=PipelineConfig.t2v_as_i2v_task,
+            help="Enable T2V to I2V task",
+        )
+
        # Add VAE configuration arguments
        from fastvideo.configs.models.vaes.base import VAEConfig
        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -246,7 +267,9 @@ class PipelineConfig:
        """
        from fastvideo.configs.pipelines.registry import (
            get_pipeline_config_cls_from_name)
+        logger.info("WTF model_path: %s", model_path)
        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
+        logger.info("pipeline_config_cls: %s", pipeline_config_cls)

        return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))

@@ -50,6 +50,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
    "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
    "stepvideo": lambda id: "stepvideo" in id.lower(),
+    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
    # Add other pipeline architecture detectors
 }

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo.dataset.lmdb_utils import get_array_shape_from_lmdb, retrieve_row_from_lmdb
+from torch.utils.data import Dataset
+import numpy as np
+import torch
+import lmdb
+
+# from Self-Forcing: https://github.com/guandeh17/Self-Forcing/blob/main/utils/dataset.py
+class ODERegressionLMDBDataset(Dataset):
+    def __init__(self, data_path: str, max_pair: int = int(1e8)):
+        print(f"data_path: {data_path}")
+        self.env = lmdb.open(data_path, readonly=True,
+                             lock=False, readahead=False, meminit=False)
+
+        self.latents_shape = get_array_shape_from_lmdb(self.env, 'latents')
+        self.max_pair = max_pair
+
+    def __len__(self):
+        return min(self.latents_shape[0], self.max_pair)
+
+    def __getitem__(self, idx):
+        """
+        Outputs:
+            - prompts: List of Strings
+            - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
+        """
+        latents = retrieve_row_from_lmdb(
+            self.env,
+            "latents", np.float16, idx, shape=self.latents_shape[1:]
+        )
+
+        if len(latents.shape) == 4:
+            latents = latents[None, ...]
+
+        prompts = retrieve_row_from_lmdb(
+            self.env,
+            "prompts", str, idx
+        )
+        return {
+            "prompts": prompts,
+            "ode_latent": torch.tensor(latents, dtype=torch.float32)
+        }
+
@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# from Self-Forcing: https://github.com/guandeh17/Self-Forcing/blob/main/utils/lmdb.py
+
+import numpy as np
+
+
+def get_array_shape_from_lmdb(env, array_name):
+    with env.begin() as txn:
+        image_shape = txn.get(f"{array_name}_shape".encode()).decode()
+        image_shape = tuple(map(int, image_shape.split()))
+    return image_shape
+
+
+def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
+    """
+    Store rows of multiple numpy arrays in a single LMDB.
+    Each row is stored separately with a naming convention.
+    """
+    with env.begin(write=True) as txn:
+        for array_name, array in arrays_dict.items():
+            for i, row in enumerate(array):
+                # Convert row to bytes
+                if isinstance(row, str):
+                    row_bytes = row.encode()
+                else:
+                    row_bytes = row.tobytes()
+
+                data_key = f'{array_name}_{start_index + i}_data'.encode()
+
+                txn.put(data_key, row_bytes)
+
+
+def process_data_dict(data_dict, seen_prompts):
+    output_dict = {}
+
+    all_videos = []
+    all_prompts = []
+    for prompt, video in data_dict.items():
+        if prompt in seen_prompts:
+            continue
+        else:
+            seen_prompts.add(prompt)
+
+        video = video.half().numpy()
+        all_videos.append(video)
+        all_prompts.append(prompt)
+
+    if len(all_videos) == 0:
+        return {"latents": np.array([]), "prompts": np.array([])}
+
+    all_videos = np.concatenate(all_videos, axis=0)
+
+    output_dict['latents'] = all_videos
+    output_dict['prompts'] = np.array(all_prompts)
+
+    return output_dict
+
+
+def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape=None):
+    """
+    Retrieve a specific row from a specific array in the LMDB.
+    """
+    data_key = f'{array_name}_{row_index}_data'.encode()
+
+    with lmdb_env.begin() as txn:
+        row_bytes = txn.get(data_key)
+
+    if dtype == str:
+        array = row_bytes.decode()
+    else:
+        array = np.frombuffer(row_bytes, dtype=dtype)
+
+    if shape is not None and len(shape) > 0:
+        array = array.reshape(shape)
+    return array
@@ -3,6 +3,9 @@ from typing import Any, cast

 import numpy as np
 import torch
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)


 def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
@@ -18,6 +18,7 @@ if TYPE_CHECKING:
    FASTVIDEO_LOGGING_PREFIX: str = ""
    FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
    FASTVIDEO_TRACE_FUNCTION: int = 0
+    FASTVIDEO_FORCE_ATTN_BF16: bool = False
    FASTVIDEO_ATTENTION_BACKEND: str | None = None
    FASTVIDEO_ATTENTION_CONFIG: str | None = None
    FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
@@ -168,6 +169,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "FASTVIDEO_TRACE_FUNCTION":
    lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")),

+    # if set, fastvideo will force attention to be computed in bfloat16
+    "FASTVIDEO_FORCE_ATTN_BF16":
+    lambda: bool(int(os.getenv("FASTVIDEO_FORCE_ATTN_BF16", "0"))),
+
    # Backend for attention computation
    # Available options:
    # - "TORCH_SDPA": use torch.nn.MultiheadAttention
@@ -705,6 +705,7 @@ class TrainingArgs(FastVideoArgs):
    # simulate generator forward to match inference
    simulate_generator_forward: bool = False
    warp_denoising_step: bool = False
+    intermediate_latents_visualization: bool = False

    # Self-forcing specific arguments
    num_frame_per_block: int = 3
@@ -1140,10 +1141,11 @@ class TrainingArgs(FastVideoArgs):
            "--last-step-only",
            action=StoreBoolean,
            help="Whether to only use the last timestep for training")
-        parser.add_argument("--context-noise",
-                            type=int,
-                            default=TrainingArgs.context_noise,
-                            help="Context noise level for cache updates")
+        parser.add_argument(
+            "--context-noise",
+            type=int,
+            default=TrainingArgs.context_noise,
+            help="Context noise level for cache updates")

        return parser

@@ -1151,4 +1153,4 @@ class TrainingArgs(FastVideoArgs):
 def parse_int_list(value: str) -> list[int]:
    if not value:
        return []
-    return [int(x.strip()) for x in value.split(",")]
+    return [int(x.strip()) for x in value.split(",")]
@@ -276,4 +276,4 @@ class LayerNormScaleShift(nn.Module):
        if self.compute_dtype == torch.float32:
            output = output.to(x.dtype)

-        return output
+        return output
@@ -147,8 +147,8 @@ class CausalWanSelfAttention(nn.Module):
                # Assign new keys/values directly up to current_end
                local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
                local_start_index = local_end_index - num_new_tokens
-                kv_cache["k"] = kv_cache["k"].detach()
-                kv_cache["v"] = kv_cache["v"].detach()
+                # kv_cache["k"] = kv_cache["k"].detach()
+                # kv_cache["v"] = kv_cache["v"].detach()
                # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
                kv_cache["v"][:, local_start_index:local_end_index] = v
@@ -262,9 +262,14 @@ class CausalWanTransformerBlock(nn.Module):
        # *_msa.shape: [batch_size, num_frames, 1, inner_dim]
        # assert shift_msa.dtype == torch.float32

+        # logger.info("temb sum: %s, dtype: %s", temb.float().sum().item(), temb.dtype)
+        # logger.info("scale_msa sum: %s, dtype: %s", scale_msa.float().sum().item(), scale_msa.dtype)
+        # logger.info("shift_msa sum: %s, dtype: %s", shift_msa.float().sum().item(), shift_msa.dtype)
+
        # 1. Self-attention
        norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
                        (1 + scale_msa) + shift_msa).flatten(1, 2)
+        # logger.info("norm_hidden_states sum: %s, shape: %s", norm_hidden_states.float().sum().item(), norm_hidden_states.shape)
        query, _ = self.to_q(norm_hidden_states)
        key, _ = self.to_k(norm_hidden_states)
        value, _ = self.to_v(norm_hidden_states)
@@ -370,8 +375,7 @@ class CausalWanTransformer3DModel(BaseDiT):

        # Causal-specific
        self.block_mask = None
-        self.num_frame_per_block = config.arch_config.num_frames_per_block
-        assert self.num_frame_per_block <= 3
+        self.num_frame_per_block = 3
        self.independent_first_frame = False

        self.__post_init__()
@@ -552,6 +556,8 @@ class CausalWanTransformer3DModel(BaseDiT):
                start_frame: int = 0,
                **kwargs) -> torch.Tensor:

+        logger.info("timestep dtype: %s, timestep sum: %s", timestep.dtype, timestep.float().sum().item())
+
        orig_dtype = hidden_states.dtype
        if not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
@@ -646,10 +652,11 @@ class CausalWanTransformer3DModel(BaseDiT):
        **kwargs
    ):
        if kwargs.get('kv_cache', None) is not None:
-            return self._forward_inference(*args, **kwargs)
+            noise_pred = self._forward_inference(*args, **kwargs)
        else:
-            return self._forward_train(*args, **kwargs)
+            noise_pred = self._forward_train(*args, **kwargs)

+        return noise_pred

    def unpatchify(self, x, grid_sizes):
        r"""
@@ -675,4 +682,4 @@ class CausalWanTransformer3DModel(BaseDiT):
            # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous())
            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
            out.append(u)
-        return out
+        return out
@@ -172,8 +172,13 @@ class WanT2VCrossAttention(WanSelfAttention):
            k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d)
            v = self.to_v(context)[0].view(b, -1, n, d)

-        # compute attention
-        x = self.attn(q, k, v)
+        if envs.FASTVIDEO_FORCE_ATTN_BF16:
+            out_dtype = v.dtype
+            # compute attention
+            x = self.attn(q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)).to(out_dtype)
+        else:
+            # compute attention
+            x = self.attn(q, k, v)

        # output
        x = x.flatten(2)
@@ -283,16 +288,16 @@ class WanTransformerBlock(nn.Module):
        # 2. Cross-attention
        if added_kv_proj_dim is not None:
            # I2V
-            self.attn2 = WanI2VCrossAttention(dim,
-                                              num_heads,
-                                              qk_norm=qk_norm,
+            self.attn2 = WanI2VCrossAttention(dim, 
+                                              num_heads, 
+                                              qk_norm=qk_norm, 
                                              eps=eps)

        else:
            # T2V
-            self.attn2 = WanT2VCrossAttention(dim,
-                                              num_heads,
-                                              qk_norm=qk_norm,
+            self.attn2 = WanT2VCrossAttention(dim, 
+                                              num_heads, 
+                                              qk_norm=qk_norm, 
                                              eps=eps)

        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
@@ -359,7 +364,12 @@ class WanTransformerBlock(nn.Module):
                                       is_neox_style=False), _apply_rotary_emb(
                                           key, cos, sin, is_neox_style=False)

-        attn_output, _ = self.attn1(query, key, value)
+        if envs.FASTVIDEO_FORCE_ATTN_BF16:
+            out_dtype = value.dtype
+            attn_output, _ = self.attn1(query.to(torch.bfloat16), key.to(torch.bfloat16), value.to(torch.bfloat16))
+            attn_output = attn_output.to(out_dtype)
+        else:
+            attn_output, _ = self.attn1(query, key, value)
        attn_output = attn_output.flatten(2)
        attn_output, _ = self.to_out(attn_output)
        attn_output = attn_output.squeeze(1)
@@ -835,5 +845,4 @@ class WanTransformer3DModel(CachableDiT):
        if self.is_even:
            return hidden_states + self.previous_residual_even
        else:
-            return hidden_states + self.previous_residual_odd
-            
+            return hidden_states + self.previous_residual_odd
@@ -238,6 +238,7 @@ class TextEncoderLoader(ComponentLoader):
                1]

        target_device = get_local_torch_device()
+        logger.info("Loading text encoder in %s precision", encoder_precision)
        # TODO(will): add support for other dtypes
        return self.load_model(model_path, encoder_config, target_device,
                               fastvideo_args, encoder_precision)
@@ -415,7 +416,6 @@ class TransformerLoader(ComponentLoader):
            raise ValueError(
                "Model config does not contain a _class_name attribute. "
                "Only diffusers format is supported.")
-
        logger.info("transformer cls_name: %s", cls_name)
        if fastvideo_args.override_transformer_cls_name is not None:
            cls_name = fastvideo_args.override_transformer_cls_name
@@ -443,19 +443,15 @@ class TransformerLoader(ComponentLoader):

        if use_custom_weights:
            logger.info("Using custom initialization weights from: %s", custom_weights_path)
-            assert custom_weights_path is not None, "Custom initialization weights must be provided"
-            if os.path.isdir(custom_weights_path):
-                safetensors_list = glob.glob(
-                    os.path.join(str(custom_weights_path), "*.safetensors"))
-            else:
-                assert custom_weights_path.endswith(".safetensors"), "Custom initialization weights must be a safetensors file"
-                safetensors_list = [custom_weights_path]
+            safetensors_list = [custom_weights_path]

-        logger.info("Loading model from %s safetensors files: %s",
-                    len(safetensors_list), safetensors_list)
+        logger.info("Loading model from %s safetensors files in %s",
+                    len(safetensors_list), model_path)

        default_dtype = PRECISION_TO_TYPE[
            fastvideo_args.pipeline_config.dit_precision]
+        param_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.dit_forward_precision]

        # Load the model using FSDP loader
        logger.info("Loading model from %s, default_dtype: %s", cls_name,
@@ -476,7 +472,7 @@ class TransformerLoader(ComponentLoader):
            fsdp_inference=fastvideo_args.use_fsdp_inference,
            # TODO(will): make these configurable
            default_dtype=default_dtype,
-            param_dtype=torch.bfloat16,
+            param_dtype=param_dtype,
            reduce_dtype=torch.float32,
            output_dtype=None,
            training_mode=fastvideo_args.training_mode)
@@ -485,8 +481,11 @@ class TransformerLoader(ComponentLoader):
        total_params = sum(p.numel() for p in model.parameters())
        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

-        assert next(model.parameters()).dtype == default_dtype, "Model dtype does not match default dtype"
-
+        # Need to convert the model to the default_dtype
+        # Otherwise, the model master weights will be in param_dtype, and the gradients will also be in param_dtype
+        # This means the param update will be in lower precision, causing precision loss
+        logger.info("Converting model to dtype: %s", default_dtype)
+        model = model.to(default_dtype)
        model = model.eval()
        return model

@@ -88,7 +88,6 @@ def maybe_load_fsdp_model(
        mp_policy=mp_policy,
    )

-    logger.info("Loading model with default_dtype: %s", default_dtype)
    with set_default_dtype(default_dtype), torch.device("meta"):
        model = model_cls(**init_params)

@@ -180,3 +180,25 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
    pred_video = noise_input_latent - sigma_t * pred_noise
    return pred_video.to(dtype)
+
+def pred_video_to_pred_noise(x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor:
+    """
+    Convert x0 prediction to flow matching's prediction.
+    x0_pred: the x0 prediction with shape [B, C, H, W]
+    xt: the input noisy data with shape [B, C, H, W]
+    timestep: the timestep with shape [B]
+
+    pred = (x_t - x_0) / sigma_t
+    """
+    # use higher precision for calculations
+    original_dtype = x0_pred.dtype
+    x0_pred, xt, sigmas, timesteps = map(
+        lambda x: x.double().to(x0_pred.device), [x0_pred, xt,
+                                                    scheduler.sigmas,
+                                                    scheduler.timesteps]
+    )
+    timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    flow_pred = (xt - x0_pred) / sigma_t
+    return flow_pred.to(original_dtype)
@@ -7,6 +7,8 @@ This module wires the causal DMD denoising stage into the modular pipeline.

 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
 from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline

 # isort: off
@@ -121,14 +121,25 @@ class ComposedPipelineBase(ABC):
                        model_path: str,
                        device: str | None = None,
                        torch_dtype: torch.dtype | None = None,
-                        pipeline_config: str | PipelineConfig | None = None,
+                        pipeline_config: PipelineConfig | None = None,
                        args: argparse.Namespace | None = None,
                        required_config_modules: list[str] | None = None,
                        loaded_modules: dict[str, torch.nn.Module]
                        | None = None,
                        **kwargs) -> "ComposedPipelineBase":
        """
-        Load a pipeline from a pretrained model.
+        Load a pipeline from a pretrained model. 
+        Few different patterns are supported:
+        - Only provide model_path:
+            - This will load the pipeline in inference mode.
+            - The pipeline will be initialized with the default config.
+            - The pipeline will be initialized with the default modules.
+            - The pipeline will be initialized with the default stages.
+            - The pipeline will be initialized with the default stages.
+        - override the default config using pipeline_config or args or kwargs
+        - override the default modules using loaded_modules
+        - override the pipelineconfig
+
        loaded_modules: Optional[Dict[str, torch.nn.Module]] = None,
        If provided, loaded_modules will be used instead of loading from config/pretrained weights.
        """
@@ -136,9 +147,18 @@ class ComposedPipelineBase(ABC):

            kwargs['model_path'] = model_path
            fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
+            if pipeline_config is not None:
+                fastvideo_args.pipeline_config = pipeline_config
+            if fastvideo_args.override_transformer_cls_name is not None:
+                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
+                fastvideo_args.pipeline_config = pipeline_config
        else:
            assert args is not None, "args must be provided for training mode"
            fastvideo_args = TrainingArgs.from_cli_args(args)
+            if fastvideo_args.override_transformer_cls_name is not None:
+                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
+                fastvideo_args.pipeline_config = pipeline_config
+                logger.info("in 2 Overriding transformer cls name to %s", fastvideo_args.override_transformer_cls_name)
            # TODO(will): fix this so that its not so ugly
            fastvideo_args.model_path = model_path
            for key, value in kwargs.items():
@@ -149,7 +169,8 @@ class ComposedPipelineBase(ABC):
            # model is loaded with the correct precision. Subsequently we will
            # use FSDP2's MixedPrecisionPolicy to set the precision for the
            # fwd, bwd, and other operations' precision.
-            assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
+            fastvideo_args.pipeline_config.dit_precision = 'fp32'
+            # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'

        logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)

@@ -151,7 +151,7 @@ class ForwardBatch:
    output: torch.Tensor | None = None
    return_trajectory_latents: bool = False
    return_trajectory_decoded: bool = False
-    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_timesteps: list[int] | None = None
    trajectory_latents: torch.Tensor | None = None
    trajectory_decoded: list[torch.Tensor] | None = None

@@ -246,7 +246,7 @@ class TrainingBatch:
    fake_score_loss: float = 0.0

    dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
-    latent_vis_dict: dict[str, Any] = field(default_factory=dict)
+    latent_vis_dict: dict[str, torch.Tensor] = field(default_factory=dict)
    fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)


@@ -12,8 +12,10 @@ import os
 from collections.abc import Iterator
 from typing import Any

+import numpy as np
 import pyarrow as pa
 import torch
+from PIL import Image
 from torch.utils.data import DataLoader
 from torchdata.stateful_dataloader import StatefulDataLoader
 from tqdm import tqdm
@@ -40,6 +42,11 @@ from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
                                        TimestepPreparationStage)
 from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict

+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.utils import pred_noise_to_pred_video, pred_video_to_pred_noise
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.distributed import get_local_torch_device
+
 logger = init_logger(__name__)


@@ -68,6 +75,15 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
                                                denoising_strength=1.0)

+        fastvideo_args.model_loaded["transformer"] = False
+        loader = TransformerLoader()
+        fastvideo_args.pipeline_config.dit_precision = "fp32" # Overwrite the precision to fp32 for transformer
+        fastvideo_args.pipeline_config.dit_forward_precision = "fp32"
+        self.transformer = loader.load(
+            fastvideo_args.model_paths["transformer"], fastvideo_args)
+        self.add_module("transformer", self.transformer)
+        fastvideo_args.model_loaded["transformer"] = True
+
        self.add_stage(stage_name="input_validation_stage",
                       stage=InputValidationStage())
        self.add_stage(stage_name="prompt_encoding_stage",
@@ -96,6 +112,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
        """Preprocess text-only data and generate trajectory information."""

        for batch_idx, data in enumerate(self.pbar):
+            # logger.info("transformer weight sum: %s", sum(p.float().sum().item() for p in self.transformer.parameters()))
            if data is None:
                continue

@@ -126,6 +143,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                    ]

                batch_captions = valid_data["text"]
+                self.prompt_encoding_stage.text_encoders[0] = self.prompt_encoding_stage.text_encoders[0].to(dtype=torch.bfloat16).to(dtype=torch.float32)
                # Encode text using the standalone TextEncodingStage API
                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
                    batch_captions,
@@ -134,22 +152,25 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                    return_attention_mask=True,
                )
                prompt_embeds = prompt_embeds_list[0]
+                logger.info("prompt_embeds sum: %s, prompt_embeds shape: %s, prompt_embeds dtype: %s", prompt_embeds.float().sum(), prompt_embeds.shape, prompt_embeds.dtype)
                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)

+                negative_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+
                # 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,
+                        negative_prompt,
                        fastvideo_args,
                        encoder_index=[0],
                        return_attention_mask=True,
                    )
-                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
+                    negative_prompt_embed = negative_prompt_embeds_list[0]
                    negative_prompt_attention_mask = negative_prompt_masks_list[
-                        0][0]
+                        0]
                else:
                    negative_prompt_embed = None
                    negative_prompt_attention_mask = None
@@ -176,30 +197,129 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
                    batch.return_trajectory_latents = True
                    # Enabling this will save the decoded trajectory videos.
                    # Used for debugging.
-                    batch.return_trajectory_decoded = False
+                    batch.return_trajectory_decoded = True
                    batch.height = args.max_height
                    batch.width = args.max_width
                    batch.fps = args.train_fps
-                    batch.guidance_scale = 6.0
+                    batch.guidance_scale = 3.0
                    batch.do_classifier_free_guidance = True

                    result_batch = self.input_validation_stage(
                        batch, fastvideo_args)
                    result_batch = self.timestep_preparation_stage(
                        batch, fastvideo_args)
-                    result_batch = self.latent_preparation_stage(
-                        result_batch, fastvideo_args)
-                    result_batch = self.denoising_stage(result_batch,
-                                                        fastvideo_args)
+                    # result_batch = self.latent_preparation_stage(
+                    #     result_batch, fastvideo_args)
+                    # result_batch = self.denoising_stage(result_batch,
+                    #                                     fastvideo_args)
+                    noisy_input = []
+                    # latents = result_batch.latents.permute(0, 2, 1, 3, 4)
+                    latents = torch.randn(
+                        [1, 21, 16, 60, 104], dtype=torch.float32, device=get_local_torch_device()
+                    )
+                    # logger.info("transformer weight sum: %s", sum(p.float().sum().item() for p in self.transformer.parameters()))
+                    logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype)
+
+                    logger.info("scheduler timesteps: %s", self.get_module("scheduler").timesteps)
+                    for progress_id, t in enumerate(tqdm(self.get_module("scheduler").timesteps)):
+                        timestep = t * \
+                            torch.ones([1, 21], device=latents.device, dtype=torch.float32)
+
+                        noisy_input.append(latents)
+
+                        with set_forward_context(
+                            current_timestep=0,
+                            attn_metadata=None,
+                            forward_batch=None,
+                        ):
+                            # logger.info("prompt_embed sum: %s, prompt_embed shape: %s, prompt_embed dtype: %s", prompt_embed.float().sum(), prompt_embed.shape, prompt_embed.dtype)
+                            # logger.info("timestep: %s", timestep[:, 0])
+                            # Run transformer
+                            cond_pred_noise_btchw = self.transformer(
+                                hidden_states=latents.permute(0, 2, 1, 3, 4),
+                                encoder_hidden_states=prompt_embed,
+                                timestep=timestep[:, 0]
+                            ).permute(0, 2, 1, 3, 4)
+
+                        # logger.info("cond_pred_noise_btchw sum: %s, cond_pred_noise_btchw shape: %s, cond_pred_noise_btchw dtype: %s", cond_pred_noise_btchw.float().sum(), cond_pred_noise_btchw.shape, cond_pred_noise_btchw.dtype)
+
+                        cond_pred_video_btchw = pred_noise_to_pred_video(
+                            pred_noise=cond_pred_noise_btchw.flatten(0, 1),
+                            noise_input_latent=latents.flatten(0, 1),
+                            timestep=timestep.flatten(0, 1),
+                            scheduler=self.get_module("scheduler")).unflatten(
+                                0, cond_pred_noise_btchw.shape[:2])
+
+                        # logger.info("cond_pred_video_btchw sum: %s, cond_pred_video_btchw shape: %s, cond_pred_video_btchw dtype: %s", cond_pred_video_btchw.float().sum(), cond_pred_video_btchw.shape, cond_pred_video_btchw.dtype)
+
+                        with set_forward_context(
+                            current_timestep=t,
+                            attn_metadata=None,
+                            forward_batch=result_batch,
+                        ):
+                            # logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype)
+                            # Run transformer
+                            uncond_pred_noise_btchw = self.transformer(
+                                latents.permute(0, 2, 1, 3, 4),
+                                negative_prompt_embed,
+                                timestep[:, 0]
+                            ).permute(0, 2, 1, 3, 4)
+
+                        # logger.info("uncond_pred_noise_btchw sum: %s, uncond_pred_noise_btchw shape: %s, uncond_pred_noise_btchw dtype: %s", uncond_pred_noise_btchw.float().sum(), uncond_pred_noise_btchw.shape, uncond_pred_noise_btchw.dtype)
+
+                        uncond_pred_video_btchw = pred_noise_to_pred_video(
+                            pred_noise=uncond_pred_noise_btchw.flatten(0, 1),
+                            noise_input_latent=latents.flatten(0, 1),
+                            timestep=timestep.flatten(0, 1),
+                            scheduler=self.get_module("scheduler")).unflatten(
+                                0, uncond_pred_noise_btchw.shape[:2])
+
+                        pred_video_btchw = uncond_pred_video_btchw + batch.guidance_scale * (
+                            cond_pred_video_btchw - uncond_pred_video_btchw
+                        )
+
+                        # logger.info("pred_video_btchw sum: %s, pred_video_btchw shape: %s, pred_video_btchw dtype: %s", pred_video_btchw.float().sum(), pred_video_btchw.shape, pred_video_btchw.dtype)
+
+                        pred_noise_btchw = pred_video_to_pred_noise(
+                            x0_pred=pred_video_btchw.flatten(0, 1),
+                            xt=latents.flatten(0, 1),
+                            timestep=timestep.flatten(0, 1),
+                            scheduler=self.get_module("scheduler")).unflatten(
+                                0, pred_video_btchw.shape[:2])
+
+                        # logger.info("pred_noise_btchw sum: %s, pred_noise_btchw shape: %s, pred_noise_btchw dtype: %s", pred_noise_btchw.float().sum(), pred_noise_btchw.shape, pred_noise_btchw.dtype)
+
+                        latents = self.get_module("scheduler").step(
+                            pred_noise_btchw.flatten(0, 1),
+                            self.get_module("scheduler").timesteps[progress_id] * torch.ones(
+                                [1, 21], device=latents.device, dtype=torch.long).flatten(0, 1),
+                            latents.flatten(0, 1)
+                        )[0].unflatten(dim=0, sizes=pred_noise_btchw.shape[:2])
+
+                        # logger.info("latents sum: %s, latents shape: %s, latents dtype: %s", latents.float().sum(), latents.shape, latents.dtype)
+
+                    noisy_input.append(latents)
+
+                    noisy_inputs = torch.stack(noisy_input, dim=1)
+
+                    noisy_inputs = noisy_inputs[:, [0, 12, 24, 36, -1]].half()
+
+                    logger.info("noisy inputs sum: %s, noisy inputs shape: %s, noisy inputs dtype: %s", noisy_inputs.float().sum(), noisy_inputs.shape, noisy_inputs.dtype)
+
+                    result_batch.trajectory_latents = noisy_inputs.permute(0, 1, 3, 2, 4, 5)
+                    result_batch.trajectory_timesteps = torch.tensor([self.get_module("scheduler").timesteps[i] for i in [0, 12, 24, 36, -1]])
+                    result_batch.latents = latents.permute(0, 2, 1, 3, 4)
                    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)

+                trajectory_latents = torch.stack(trajectory_latents, dim=0).squeeze(0)
+                # trajecotry_latents = trajectory_latents[:, [0, 12, 24, 36, -1]]
+
                # Prepare extra features for text-only processing
                extra_features = {
                    "trajectory_latents": trajectory_latents,
@@ -78,6 +78,8 @@ class CausalDMDDenosingStage(DenoisingStage):
                                             torch.tensor([0],
                                                          dtype=torch.float32)))
            timesteps = scheduler_timesteps[1000 - timesteps]
+        else:
+            assert False, "warp_denoising_step must be true"
        timesteps = timesteps.to(get_local_torch_device())
        logger.info("Using timesteps: %s", timesteps)

@@ -140,11 +140,12 @@ class DenoisingStage(PipelineStage):
            latents = latents[:, :, rank_in_sp_group, :, :, :]
            batch.latents = latents
            if batch.image_latent is not None:
-                image_latent = rearrange(batch.image_latent,
-                                         "b c (n t) h w -> b c n t h w",
-                                         n=sp_world_size).contiguous()
-                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
-                batch.image_latent = image_latent
+                if not fastvideo_args.pipeline_config.ti2v_task and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    image_latent = rearrange(batch.image_latent,
+                                             "b c (n t) h w -> b c n t h w",
+                                             n=sp_world_size).contiguous()
+                    image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
+                    batch.image_latent = image_latent
        # Get timesteps and calculate warmup steps
        timesteps = batch.timesteps
        # TODO(will): remove this once we add input/output validation for stages
@@ -253,8 +254,7 @@ 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_timesteps: list[int] = []
        trajectory_latents: list[torch.Tensor] = []

        # Run denoising loop
@@ -284,14 +284,27 @@ class DenoisingStage(PipelineStage):

                # Expand latents for I2V
                latent_model_input = latents.to(target_dtype)
-                if batch.image_latent is not None:
+                if batch.image_latent is not None and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
                    assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                    latent_model_input = torch.cat(
                        [latent_model_input, batch.image_latent],
                        dim=1).to(target_dtype)
+                elif batch.image_latent is not None and fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    assert batch.image_latent is not None, "image latents should be provided for T2V to I2V task"
+                    if rank_in_sp_group == 0:
+                        logger.info("latent_model_input.shape: %s",
+                                    latent_model_input.shape)
+                        latent_model_input = torch.cat([
+                            batch.image_latent,
+                            latent_model_input[:, :, 1:, :, :],
+                        ],
+                                                       dim=2).to(target_dtype)
+                        logger.info("latent_model_input.shape: %s",
+                                    latent_model_input.shape)

                assert not torch.isnan(
                    latent_model_input).any(), "latent_model_input contains nan"
+
                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
                    timestep = torch.stack([t]).to(get_local_torch_device())
                    temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
@@ -306,6 +319,13 @@ class DenoisingStage(PipelineStage):

                latent_model_input = self.scheduler.scale_model_input(
                    latent_model_input, t)
+                if fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    if rank_in_sp_group == 0:
+                        latent_model_input = torch.cat([
+                            batch.image_latent,
+                            latent_model_input[:, :, 1:, :, :],
+                        ],
+                                                       dim=2).to(target_dtype)

                # Prepare inputs for transformer
                guidance_expand = (
@@ -468,10 +488,17 @@ class DenoisingStage(PipelineStage):
                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()
+        if trajectory_tensor is not None:
+            batch.trajectory_timesteps = torch.tensor(trajectory_timesteps).cpu()
            batch.trajectory_latents = trajectory_tensor.cpu()

+        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
+            latents = torch.cat([
+                batch.image_latent,
+                latents[:, :, 1:, :, :],
+            ],
+                                dim=2)
+
        # Update batch with final latents
        batch.latents = latents

@@ -105,6 +105,81 @@ class ImageVAEEncodingStage(PipelineStage):
    def __init__(self, vae: ParallelTiledVAE) -> None:
        self.vae: ParallelTiledVAE = vae

+    def encode_image(self,
+                     image: PIL.Image.Image,
+                     height: int,
+                     width: int,
+                     fastvideo_args: FastVideoArgs,
+                     generator: torch.Generator | None = None) -> torch.Tensor:
+        """
+        Encode image into latent space.
+        """
+        image = self.preprocess(
+            image,
+            vae_scale_factor=self.vae.spatial_compression_ratio,
+            height=height,
+            width=width).to(get_local_torch_device(), dtype=torch.float32)
+
+        # (B, C, H, W) -> (B, C, 1, H, W)
+        print(f"image.shape: {image.shape}")
+        image = image.unsqueeze(2)
+        print(f"after unsqueeze image.shape: {image.shape}")
+        return self.encode_tensor(image, fastvideo_args, generator)
+
+    def encode_tensor(self,
+                      video_condition: torch.Tensor,
+                      fastvideo_args: FastVideoArgs,
+                      generator: torch.Generator | None = None) -> torch.Tensor:
+        """
+        Encode frames into latent space.
+        """
+        self.vae = self.vae.to(get_local_torch_device())
+        video_condition = video_condition.to(device=get_local_torch_device(),
+                                             dtype=torch.float32)
+
+        # 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
+
+        # Encode Image
+        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:
+                video_condition = video_condition.to(vae_dtype)
+            encoder_output = self.vae.encode(video_condition)
+
+        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            latent_condition = encoder_output.mean
+        else:
+            generator = generator
+            if generator is None:
+                raise ValueError("Generator must be provided")
+            latent_condition = self.retrieve_latents(encoder_output, generator)
+
+        # 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):
+                latent_condition -= self.vae.shift_factor.to(
+                    latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= self.vae.shift_factor
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latent_condition = latent_condition * self.vae.scaling_factor.to(
+                latent_condition.device, latent_condition.dtype)
+        else:
+            latent_condition = latent_condition * self.vae.scaling_factor
+
+        return latent_condition
+
    def forward(
        self,
        batch: ForwardBatch,
@@ -157,58 +232,29 @@ class ImageVAEEncodingStage(PipelineStage):
        # (B, C, H, W) -> (B, C, 1, H, W)
        image = image.unsqueeze(2)

-        video_condition = torch.cat([
-            image,
-            image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
-                            image.shape[3], image.shape[4])
-        ],
-                                    dim=2)
-        video_condition = video_condition.to(device=get_local_torch_device(),
-                                             dtype=torch.float32)
+        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
+            # repeat the image self.vae.temporal_compression_ratio times
+            video_condition = image.repeat(1, 1,
+                                           self.vae.temporal_compression_ratio,
+                                           1, 1)
+            # video_condition = image
+            logger.info("video_condition.shape: %s", video_condition.shape)
+        else:
+            video_condition = torch.cat([
+                image,
+                image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
+                                image.shape[3], image.shape[4])
+            ],
+                                        dim=2)

-        # 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
-
-        # Encode Image
-        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:
-                video_condition = video_condition.to(vae_dtype)
-            encoder_output = self.vae.encode(video_condition)
-
-        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
-            latent_condition = encoder_output.mean
-        else:
-            generator = batch.generator
-            if generator is None:
-                raise ValueError("Generator must be provided")
-            latent_condition = self.retrieve_latents(encoder_output, generator)
-
-        # 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):
-                latent_condition -= self.vae.shift_factor.to(
-                    latent_condition.device, latent_condition.dtype)
-            else:
-                latent_condition -= self.vae.shift_factor
-
-        if isinstance(self.vae.scaling_factor, torch.Tensor):
-            latent_condition = latent_condition * self.vae.scaling_factor.to(
-                latent_condition.device, latent_condition.dtype)
-        else:
-            latent_condition = latent_condition * self.vae.scaling_factor
+        latent_condition = self.encode_tensor(video_condition, fastvideo_args,
+                                              batch.generator)

        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
            batch.image_latent = latent_condition
+        elif fastvideo_args.pipeline_config.t2v_as_i2v_task:
+            logger.info("latent_condition.shape: %s", latent_condition.shape)
+            batch.image_latent = latent_condition
        else:
            mask_lat_size = torch.ones(1, 1, num_frames, latent_height,
                                       latent_width)
@@ -35,9 +35,15 @@ class InputValidationStage(PipelineStage):
        """Generate seeds for the inference"""
        seed = batch.seed
        num_videos_per_prompt = batch.num_videos_per_prompt
+        if isinstance(batch.prompt, list):
+            num_prompts = len(batch.prompt)
+        else:
+            num_prompts = 1
+
+        total_num_videos = num_prompts * num_videos_per_prompt

        assert seed is not None
-        seeds = [seed + i for i in range(num_videos_per_prompt)]
+        seeds = [seed + i for i in range(total_num_videos)]
        batch.seeds = seeds
        # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
        batch.generator = [
@@ -3,6 +3,7 @@
 Latent preparation stage for diffusion pipelines.
 """

+import torch
 from diffusers.utils.torch_utils import randn_tensor

 from fastvideo.distributed import get_local_torch_device
@@ -93,7 +94,7 @@ class LatentPreparationStage(PipelineStage):
        # Generate or use provided latents
        if latents is None:
            latents = randn_tensor(shape,
-                                   generator=generator,
+                                   generator=torch.Generator(device="cuda").manual_seed(1024),
                                   device=device,
                                   dtype=dtype)
        else:
@@ -82,8 +82,8 @@ def rocm_platform_plugin() -> str | None:
                logger.info("ROCm platform is available")
        finally:
            amdsmi.amdsmi_shut_down()
-    except Exception as e:
-        logger.info("ROCm platform is unavailable: %s", e)
+    except Exception:
+        pass

    return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None

@@ -8,6 +8,9 @@ from torch.distributed.tensor import DTensor
 from torch.testing import assert_close
 from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel

+from fastvideo.wan.modules.tokenizers import HuggingfaceTokenizer
+from fastvideo.wan.modules.t5 import umt5_xxl
+
 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
@@ -38,9 +41,22 @@ def test_t5_encoder():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    precision_str = "fp32"
    precision = PRECISION_TO_TYPE[precision_str]
-    model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to(
-        precision).to(device).eval()
-    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
+    # model1 = UMT5EncoderModel.from_pretrained(TEXT_ENCODER_PATH).to(
+    #     precision).to(device).eval()
+    # tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
+    model1 = umt5_xxl(
+        encoder_only=True,
+        return_tokenizer=False,
+        dtype=torch.float32,
+        device=device,
+    ).eval().requires_grad_(False)
+    model1.load_state_dict(
+        torch.load("/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
+                    map_location='cpu', weights_only=False)
+    )
+
+    tokenizer1 = HuggingfaceTokenizer(
+        name="/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/", seq_len=512, clean='whitespace')


    args = FastVideoArgs(model_path=TEXT_ENCODER_PATH,
@@ -49,8 +65,9 @@ def test_t5_encoder():
                        pin_cpu_memory=False)
    loader = TextEncoderLoader()
    model2 = loader.load(TEXT_ENCODER_PATH, args)
-    model2 = model2.to(precision)
+    model2 = model2.to(dtype=torch.bfloat16).to(precision)
    model2.eval()
+    tokenizer2 = AutoTokenizer.from_pretrained(TOKENIZER_PATH)

    # Sanity check weights between the two models
    logger.info("Comparing model weights for sanity check...")
@@ -61,8 +78,13 @@ def test_t5_encoder():
    logger.info("Model1 has %s parameters", len(params1))
    logger.info("Model2 has %s parameters", len(params2))

-    weight_diffs = []
-    # check if embed_tokens are the same
+    model1_weight_sum = sum(p.float().sum().item() for p in model1.parameters())
+    model2_weight_sum = sum(p.float().sum().item() for p in model2.parameters())
+    logger.info("Model1 weight sum: %s", model1_weight_sum)
+    logger.info("Model2 weight sum: %s", model2_weight_sum)
+
+    # weight_diffs = []
+    # # check if embed_tokens are the same
    weights = ["encoder.block.{}.layer.0.layer_norm.weight", "encoder.block.{}.layer.0.SelfAttention.relative_attention_bias.weight", \
               "encoder.block.{}.layer.0.SelfAttention.o.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_0.weight", "encoder.block.{}.layer.1.DenseReluDense.wi_1.weight",\
                "encoder.block.{}.layer.1.DenseReluDense.wo.weight", \
@@ -70,18 +92,21 @@ def test_t5_encoder():
    
    for idx in range(hf_config.num_hidden_layers):
        for w in weights:
-            name1 = w.format(idx)
+            # name1 = w.format(idx)
            name2 = w.format(idx)
-            p1 = params1[name1]
+            # p1 = params1[name1]
            p2 = params2[name2]
-            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(p1)
-            assert_close(p1, p2, atol=1e-4, rtol=1e-4)
+            p2 = (p2.to_local() if isinstance(p2, DTensor) else p2).to(device)
+            # assert_close(p1, p2, atol=1e-4, rtol=1e-4)
    

    # Test with some sample prompts
+    # prompts = [
+    #     "Once upon a time", "The quick brown fox jumps over",
+    #     "In a galaxy far, far away"
+    # ]
    prompts = [
-        "Once upon a time", "The quick brown fox jumps over",
-        "In a galaxy far, far away"
+        "A vibrant scene of Kenyan golfers at a lush green golf course on a sunny day. The golfers, dressed in casual yet stylish attire, are teeing off with animated expressions, showcasing their enthusiasm for the game. Rolling hills and pristine greens stretch out behind them, creating a picturesque backdrop. In the foreground, a golf buggy and a caddy stand ready, adding to the serene atmosphere. The camera captures the action from a mid-shot angle, focusing on the golfers' dynamic motions as they swing their clubs."
    ]

    logger.info("Testing T5 encoder with sample prompts")
@@ -91,7 +116,8 @@ def test_t5_encoder():
            logger.info("Testing prompt: %s", prompt)

            # Tokenize the prompt
-            tokens = tokenizer(prompt,
+            tokens1, mask = tokenizer1(prompt, return_mask=True, add_special_tokens=True)
+            tokens2 = tokenizer2(prompt,
                               padding="max_length",
                               max_length=512,
                               truncation=True,
@@ -101,22 +127,23 @@ def test_t5_encoder():
            # filter out padding input_ids
            # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1]
            # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1]
-            outputs1 = model1(input_ids=tokens.input_ids,
-                              attention_mask=tokens.attention_mask,
-                              output_hidden_states=True).last_hidden_state
+            outputs1 = model1(tokens1.to(device),
+                              mask.to(device))
            print("--------------------------------")
            logger.info("Testing model2")

            # Get outputs from our implementation
            with set_forward_context(current_timestep=0, attn_metadata=None):
                outputs2 = model2(
-                    input_ids=tokens.input_ids,
-                    attention_mask=tokens.attention_mask,
+                    input_ids=tokens2.input_ids,
+                    attention_mask=tokens2.attention_mask,
                ).last_hidden_state

            # Compare last hidden states
-            last_hidden_state1 = outputs1[tokens.attention_mask == 1]
-            last_hidden_state2 = outputs2[tokens.attention_mask == 1]
+            last_hidden_state1 = outputs1[mask == 1]
+            last_hidden_state2 = outputs2[tokens2.attention_mask == 1]
+            logger.info("last_hidden_state1 sum: %s", last_hidden_state1.float().sum())
+            logger.info("last_hidden_state2 sum: %s", last_hidden_state2.float().sum())

            assert last_hidden_state1.shape == last_hidden_state2.shape, \
                f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
@@ -1,184 +0,0 @@
-import os
-from pathlib import Path
-from huggingface_hub import snapshot_download
-import subprocess
-import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
-import shutil
-
-# Import the training pipeline
-sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
-
-NUM_NODES = "1"
-MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-
-# preprocessing
-DATA_DIR = "data"
-LOCAL_RAW_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol"))
-NUM_GPUS_PER_NODE_PREPROCESSING = "1"
-PREPROCESSING_ENTRY_FILE_PATH = "fastvideo/pipelines/preprocess/v1_preprocess.py"
-
-LOCAL_PREPROCESSED_DATA_DIR = Path(os.path.join(DATA_DIR, "crush-smol_processed_t2v"))
-
-
-# training
-NUM_GPUS_PER_NODE_TRAINING = "4"
-TRAINING_ENTRY_FILE_PATH = "fastvideo/training/wan_distillation_pipeline.py"
-LOCAL_TRAINING_DATA_DIR = os.path.join(LOCAL_PREPROCESSED_DATA_DIR, "combined_parquet_dataset")
-LOCAL_VALIDATION_DATASET_FILE = "examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
-LOCAL_OUTPUT_DIR = Path(os.path.join(DATA_DIR, "outputs"))
-
-def download_data():
-    # create the data dir if it doesn't exist
-    data_dir = Path(DATA_DIR)
-
-    print(f"Creating data directory at {data_dir}")
-    os.makedirs(data_dir, exist_ok=True)
-
-    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
-    try:
-        result = snapshot_download(
-            repo_id="wlsaidhi/crush-smol-merged",
-            local_dir=str(LOCAL_RAW_DATA_DIR),
-            repo_type="dataset",
-            resume_download=True,
-            token=os.environ.get("HF_TOKEN"),  # In case authentication is needed
-        )
-        print(f"Download completed successfully. Files downloaded to: {result}")
-        
-        # Verify the download
-        if not LOCAL_RAW_DATA_DIR.exists():
-            raise RuntimeError(f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
-            
-        # List downloaded files
-        print("Downloaded files:")
-        for file in LOCAL_RAW_DATA_DIR.rglob("*"):
-            if file.is_file():
-                print(f"  - {file.relative_to(LOCAL_RAW_DATA_DIR)}")
-                
-    except Exception as e:
-        print(f"Error during download: {str(e)}")
-        raise
-
-
-def run_preprocessing():
-    # remove the local_preprocessed_data_dir if it exists
-    if LOCAL_PREPROCESSED_DATA_DIR.exists():
-        print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
-        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
-
-    # Run torchrun command
-    cmd = [
-        "torchrun",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE_PREPROCESSING,
-        PREPROCESSING_ENTRY_FILE_PATH,
-        "--model_path", MODEL_PATH,
-        "--seed", "42",
-        "--data_merge_path", os.path.join(LOCAL_RAW_DATA_DIR, "merge.txt"),
-        "--preprocess_video_batch_size", "1",
-        "--max_height", "480",
-        "--max_width", "832",
-        "--num_frames", "81",
-        "--dataloader_num_workers", "0",
-        "--output_dir", LOCAL_PREPROCESSED_DATA_DIR,
-        "--train_fps", "16",
-        "--samples_per_file", "1",
-        "--flush_frequency", "1",
-        "--video_length_tolerance_range", "5",
-        "--preprocess_task", "t2v",
-    ]
-
-    process = subprocess.run(cmd, check=True)
-
-
-def run_training():
-    cmd = [
-        "torchrun",
-        "--nnodes", NUM_NODES,
-        "--nproc_per_node", NUM_GPUS_PER_NODE_TRAINING,
-        TRAINING_ENTRY_FILE_PATH,
-        "--model_path", MODEL_PATH,
-        "--inference_mode", "False",
-        "--pretrained_model_name_or_path", MODEL_PATH,
-        "--data_path", LOCAL_TRAINING_DATA_DIR,
-        "--validation_dataset_file", LOCAL_VALIDATION_DATASET_FILE,
-        "--train_batch_size", "1",
-        "--num_latent_t", "8",
-        "--num_gpus", NUM_GPUS_PER_NODE_TRAINING,
-        "--sp_size", "1",
-        "--tp_size", "1",
-        "--hsdp_replicate_dim", "1",
-        "--hsdp_shard_dim", NUM_GPUS_PER_NODE_TRAINING,
-        "--train_sp_batch_size", "1",
-        "--dataloader_num_workers", "10",
-        "--gradient_accumulation_steps", "1",
-        "--max_train_steps", "501",
-        "--learning_rate", "2e-6",
-        "--fake_score_learning_rate", "2e-6",
-        "--mixed_precision", "bf16",
-        "--training_state_checkpointing_steps", "1000",
-        "--weight_only_checkpointing_steps", "1000",
-        "--validation_steps", "50",
-        "--validation_sampling_steps", "3",
-        "--log_validation",
-        "--checkpoints_total_limit", "3",
-        "--ema_start_step", "0",
-        "--training_cfg_rate", "0.0",
-        "--output_dir", LOCAL_OUTPUT_DIR,
-        "--tracker_project_name", "ci_wan_t2v_dmd_overfit",
-        "--num_height", "480",
-        "--num_width", "832",
-        "--num_frames", "81",
-        "--flow_shift", "8",
-        "--validation_guidance_scale", "6.0",
-        "--weight_decay", "0.01",
-        "--generator_update_interval", "5",
-        "--dmd_denoising_steps", "1000,757,522",
-        "--min_timestep_ratio", "0.02",
-        "--max_timestep_ratio", "0.98",
-        "--seed", "1000",
-        "--real_score_guidance_scale", "3.5",
-        "--dit_precision", "fp32",
-        "--max_grad_norm", "1.0",
-        "--enable_gradient_checkpointing_type", "full",
-    ]
-
-    print(f"Running training with command: {cmd}")
-    process = subprocess.run(cmd, check=True)
-
-
-def test_e2e_overfit_single_sample():
-    os.environ["WANDB_MODE"] = "online"
-
-    download_data()
-    run_preprocessing()
-    run_training()
-
-    reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4")
-    print(f"reference_video_file: {reference_video_file}")
-    final_validation_video_file = os.path.join(LOCAL_OUTPUT_DIR, "validation_step_900_inference_steps_50_video_0.mp4")
-    print(f"final_validation_video_file: {final_validation_video_file}")
-
-
-    # Ensure both files exist
-    assert os.path.exists(reference_video_file), f"Reference video not found at {reference_video_file}"
-    assert os.path.exists(final_validation_video_file), f"Validation video not found at {final_validation_video_file}"
-
-    # Compute SSIM
-    mean_ssim, min_ssim, max_ssim = compute_video_ssim_torchvision(
-        reference_video_file,
-        final_validation_video_file,
-        use_ms_ssim=True  # Using MS-SSIM for better quality assessment
-    )
-
-    print("\n===== SSIM Results for Step 900 Validation =====")
-    print(f"Mean MS-SSIM: {mean_ssim:.4f}")
-    print(f"Min MS-SSIM: {min_ssim:.4f}")
-    print(f"Max MS-SSIM: {max_ssim:.4f}")
-
-    assert max_ssim > 0.5, f"Max SSIM is below 0.5: {max_ssim}"
-
-
-if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
@@ -62,11 +62,6 @@ def download_data():


 def run_preprocessing():
-    # remove the local_preprocessed_data_dir if it exists
-    if LOCAL_PREPROCESSED_DATA_DIR.exists():
-        print(f"Removing local_preprocessed_data_dir: {LOCAL_PREPROCESSED_DATA_DIR}")
-        shutil.rmtree(LOCAL_PREPROCESSED_DATA_DIR)
-
    # Run torchrun command
    cmd = [
        "torchrun",
@@ -1,3 +0,0 @@
-wan/ is used as a reference implementation for comparing with FastVideo's Wan DiT.
-
-It is from: https://github.com/guandeh17/Self-Forcing/tree/main/wan
@@ -1 +1 @@
-{"step_time":0.6983645600266755,"_wandb":{"runtime":107},"grad_norm":1.260593056678772,"avg_step_time":1.002151239803061,"_step":5,"validation_videos_50_steps":{"captions":false,"_type":"videos","count":8,"videos":[{"size":159131,"path":"media/videos/validation_videos_50_steps_0_dc447599dbe48350e9c9.mp4","_type":"video-file","sha256":"dc447599dbe48350e9c920f4971e1786bde580dd20d52b4aa147ae8d3dc564d6"},{"_type":"video-file","sha256":"4e283876ddfbf5a2cb6f8aca07a39f832b5f806fbefde8854bc73ed904ff20ee","size":160315,"path":"media/videos/validation_videos_50_steps_0_4e283876ddfbf5a2cb6f.mp4"},{"size":135225,"path":"media/videos/validation_videos_50_steps_0_78185c41e1935306e93c.mp4","_type":"video-file","sha256":"78185c41e1935306e93c2d416ee40b31d038abffb029cb5bfb11c2a634eb2fcf"},{"_type":"video-file","sha256":"27e9819d002d3f63c8918bbdc5bf2857b5effe0caf5d5b8374b9c590fc6432eb","size":197873,"path":"media/videos/validation_videos_50_steps_0_27e9819d002d3f63c891.mp4"},{"_type":"video-file","sha256":"46fe548e86144ca60a9396fcd15e8788d3a05c066c6819ccdd6a9041feaaec8f","size":170601,"path":"media/videos/validation_videos_50_steps_0_46fe548e86144ca60a93.mp4"},{"sha256":"91ec338774bec870b9c5c81be4330a8f0f2535124cb372e881c3703e6b65ed77","size":164462,"path":"media/videos/validation_videos_50_steps_0_91ec338774bec870b9c5.mp4","_type":"video-file"},{"_type":"video-file","sha256":"ee4e811080a619215fd7541b39203dfb69c2db4cdadbfdd7a234a2cff684f6f3","size":139435,"path":"media/videos/validation_videos_50_steps_0_ee4e811080a619215fd7.mp4"},{"sha256":"22e31e048ba5e5b9d6587d7306904d6edc6c618b8f77c3ceb05ba2d602309274","size":147072,"path":"media/videos/validation_videos_50_steps_0_22e31e048ba5e5b9d658.mp4","_type":"video-file"}]},"_timestamp":1.75118195270901e+09,"vsa_sparsity":0.05,"learning_rate":1e-05,"train_loss":0.2620866410434246,"_runtime":107.325113071}
+{"step_time":0.6983645600266755,"_wandb":{"runtime":107},"grad_norm":0.50390625,"avg_step_time":1.002151239803061,"_step":5,"validation_videos_50_steps":{"captions":false,"_type":"videos","count":8,"videos":[{"size":159131,"path":"media/videos/validation_videos_50_steps_0_dc447599dbe48350e9c9.mp4","_type":"video-file","sha256":"dc447599dbe48350e9c920f4971e1786bde580dd20d52b4aa147ae8d3dc564d6"},{"_type":"video-file","sha256":"4e283876ddfbf5a2cb6f8aca07a39f832b5f806fbefde8854bc73ed904ff20ee","size":160315,"path":"media/videos/validation_videos_50_steps_0_4e283876ddfbf5a2cb6f.mp4"},{"size":135225,"path":"media/videos/validation_videos_50_steps_0_78185c41e1935306e93c.mp4","_type":"video-file","sha256":"78185c41e1935306e93c2d416ee40b31d038abffb029cb5bfb11c2a634eb2fcf"},{"_type":"video-file","sha256":"27e9819d002d3f63c8918bbdc5bf2857b5effe0caf5d5b8374b9c590fc6432eb","size":197873,"path":"media/videos/validation_videos_50_steps_0_27e9819d002d3f63c891.mp4"},{"_type":"video-file","sha256":"46fe548e86144ca60a9396fcd15e8788d3a05c066c6819ccdd6a9041feaaec8f","size":170601,"path":"media/videos/validation_videos_50_steps_0_46fe548e86144ca60a93.mp4"},{"sha256":"91ec338774bec870b9c5c81be4330a8f0f2535124cb372e881c3703e6b65ed77","size":164462,"path":"media/videos/validation_videos_50_steps_0_91ec338774bec870b9c5.mp4","_type":"video-file"},{"_type":"video-file","sha256":"ee4e811080a619215fd7541b39203dfb69c2db4cdadbfdd7a234a2cff684f6f3","size":139435,"path":"media/videos/validation_videos_50_steps_0_ee4e811080a619215fd7.mp4"},{"sha256":"22e31e048ba5e5b9d6587d7306904d6edc6c618b8f77c3ceb05ba2d602309274","size":147072,"path":"media/videos/validation_videos_50_steps_0_22e31e048ba5e5b9d658.mp4","_type":"video-file"}]},"_timestamp":1.75118195270901e+09,"vsa_sparsity":0.05,"learning_rate":1e-05,"train_loss":0.19960195198655128,"_runtime":107.325113071}
@@ -110,7 +110,7 @@ def test_distributed_training():
        'avg_step_time': 1.0,
        'grad_norm': 0.1,
        'step_time': 1.0,
-        'train_loss': 0.005
+        'train_loss': 0.001
    }

    failures = []
@@ -1 +1 @@
-{"train_loss":0.1021774671971798,"_step":5,"_wandb":{"runtime":35},"step_time":2.2575189135968685,"grad_norm":0.11582941561937332,"_runtime":35.701958502,"avg_step_time":2.4754387199878694,"_timestamp":1.7525528420185745e+09,"learning_rate":1e-06,"vsa_sparsity":0,"validation_videos_8_steps":{"videos":[{"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.","_type":"video-file","sha256":"ee83e83df073a648f89dcd288cccaed9af765e11fc35c77a6e0cb2ebaa1be5b0","size":475248,"path":"media/videos/validation_videos_8_steps_0_ee83e83df073a648f89d.mp4"},{"size":341490,"path":"media/videos/validation_videos_8_steps_0_d0b2758549d5c82845ca.mp4","caption":"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","_type":"video-file","sha256":"d0b2758549d5c82845ca3c1ac0db6812877631495e77ac041acdc8ea31f0a5ee"},{"caption":"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.","_type":"video-file","sha256":"08638381f1607d6ab10684772be38a58b18628fd402a208b25f6af96e765d454","size":436814,"path":"media/videos/validation_videos_8_steps_0_08638381f1607d6ab106.mp4"}],"captions":["A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","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."],"_type":"videos","count":3}}
+{"train_loss":0.10545400530099869,"_step":5,"_wandb":{"runtime":35},"step_time":2.2575189135968685,"grad_norm":0.53125,"_runtime":35.701958502,"avg_step_time":2.4754387199878694,"_timestamp":1.7525528420185745e+09,"learning_rate":1e-06,"vsa_sparsity":0,"validation_videos_8_steps":{"videos":[{"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.","_type":"video-file","sha256":"ee83e83df073a648f89dcd288cccaed9af765e11fc35c77a6e0cb2ebaa1be5b0","size":475248,"path":"media/videos/validation_videos_8_steps_0_ee83e83df073a648f89d.mp4"},{"size":341490,"path":"media/videos/validation_videos_8_steps_0_d0b2758549d5c82845ca.mp4","caption":"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","_type":"video-file","sha256":"d0b2758549d5c82845ca3c1ac0db6812877631495e77ac041acdc8ea31f0a5ee"},{"caption":"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.","_type":"video-file","sha256":"08638381f1607d6ab10684772be38a58b18628fd402a208b25f6af96e765d454","size":436814,"path":"media/videos/validation_videos_8_steps_0_08638381f1607d6ab106.mp4"}],"captions":["A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","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."],"_type":"videos","count":3}}
@@ -5,7 +5,7 @@ import math
 import numpy as np
 import pytest
 import torch
-from fastvideo.tests.third_party.wan.modules.causal_model import CausalWanModel
+from fastvideo.wan.modules.causal_model import CausalWanModel

 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.forward_context import set_forward_context
@@ -289,4 +289,4 @@ def _initialize_crossattn_cache(transformer, batch_size, max_text_len, dtype,
            "is_init":
            False,
        })
-    return crossattn_cache
+    return crossattn_cache
@@ -5,7 +5,7 @@ import math
 import numpy as np
 import pytest
 import torch
-from fastvideo.tests.third_party.wan.modules.model import WanModel
+from fastvideo.wan.modules.model import WanModel

 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.forward_context import set_forward_context
@@ -22,6 +22,7 @@ logger = init_logger(__name__)
 os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"

+os.environ["FASTVIDEO_FORCE_ATTN_BF16"] = "1"
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
                                  local_dir=os.path.join(
@@ -32,19 +33,21 @@ TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
@pytest.mark.usefixtures("distributed_setup")
 def test_ori_wan_transformer():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
+    precision = torch.float32
+    precision_str = "fp32"
    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
                         dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
+                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str, dit_forward_precision=precision_str))
    args.device = device

    loader = TransformerLoader()
-    model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
+    model2 = loader.load(TRANSFORMER_PATH, args)

    model1 = WanModel.from_pretrained(
-        "/mnt/weka/home/hao.zhang/wei/Self-Forcing/wan_models/Wan2.1-T2V-1.3B", device=device,
-        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+        "/mnt/weka/home/hao.zhang/wei/Self-Forcing-clean/wan_models/Wan2.1-T2V-1.3B")
+    model1.eval()
+    model1 = model1.to(device).to(precision)
+    model1.requires_grad_(False)

    total_params = sum(p.numel() for p in model1.parameters())
    # Calculate weight sum for model1 (converting to float64 to avoid overflow)
@@ -75,29 +78,35 @@ def test_ori_wan_transformer():

    # Create identical inputs for both models
    batch_size = 1
-    text_seq_len = 30
-    seq_len = math.ceil((160 * 90) /
+    text_seq_len = 120
+    seq_len = math.ceil((104 * 60) /
                            (2 * 2) *
                            21)

    # Video latents [B, C, T, H, W]
    hidden_states = torch.randn(batch_size,
-                                16,
                                21,
-                                160,
-                                90,
+                                16,
+                                60,
+                                104,
                                device=device,
+                                generator=torch.Generator("cuda").manual_seed(1024),
                                dtype=precision)

+    logger.info("Hidden states sum: %s, Hidden states shape: %s, Hidden states dtype: %s", hidden_states.float().sum(), hidden_states.shape, hidden_states.dtype)
+
    # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size,
-                                        text_seq_len + 1,
-                                        4096,
-                                        device=device,
-                                        dtype=precision)
+    # encoder_hidden_states = torch.randn(batch_size,
+    #                                     text_seq_len + 1,
+    #                                     4096,
+    #                                     device=device,
+    #                                     dtype=precision)
+    encoder_hidden_states = torch.load("../sf_cond_prompt_embeds.pt").to(device, dtype=precision)
+
+    logger.info("Encoder hidden states sum: %s, Encoder hidden states shape: %s, Encoder hidden states dtype: %s", encoder_hidden_states.float().sum(), encoder_hidden_states.shape, encoder_hidden_states.dtype)

    # Timestep
-    timestep = torch.tensor([500], device=device, dtype=precision)
+    timestep = torch.tensor([995.7627], device=device, dtype=precision)

    forward_batch = ForwardBatch(
        data_type="dummy",
@@ -105,19 +114,19 @@ def test_ori_wan_transformer():

    # with torch.amp.autocast('cuda', dtype=precision):
    output1 = model1(
-        x=hidden_states,
+        x=hidden_states.permute(0, 2, 1, 3, 4),
        context=encoder_hidden_states,
        t=timestep,
        seq_len=seq_len,
-    )
+    ).permute(0, 2, 1, 3, 4)
    with set_forward_context(
            current_timestep=0,
            attn_metadata=None,
            forward_batch=forward_batch,
    ):
-        output2 = model2(hidden_states=hidden_states,
+        output2 = model2(hidden_states=hidden_states.permute(0, 2, 1, 3, 4),
                            encoder_hidden_states=encoder_hidden_states,
-                            timestep=timestep)
+                            timestep=timestep).permute(0, 2, 1, 3, 4)

    # Check if outputs have the same shape
    assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
@@ -128,6 +137,8 @@ def test_ori_wan_transformer():
    mean_diff = torch.mean(torch.abs(output1 - output2))
    logger.info("Max Diff: %s", max_diff.item())
    logger.info("Mean Diff: %s", mean_diff.item())
+    logger.info("Output 1 sum: %s", output1.float().sum())
+    logger.info("Output 2 sum: %s", output2.float().sum())
    assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}"
    # mean diff
-    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"
+    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"
@@ -5,7 +5,7 @@ import math
 import numpy as np
 import pytest
 import torch
-from fastvideo.tests.third_party.wan.modules.causal_model import CausalWanModel
+from fastvideo.wan.modules.causal_model import CausalWanModel

 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.forward_context import set_forward_context
@@ -141,4 +141,4 @@ def test_train_ori_causal_wan_transformer():
    logger.info("Mean Diff: %s", mean_diff.item())
    assert max_diff < 1e-4, f"Maximum difference between outputs: {max_diff.item()}"
    # mean diff
-    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"
+    assert mean_diff < 1e-4, f"Mean difference between outputs: {mean_diff.item()}"
@@ -28,7 +28,6 @@ MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")


-@pytest.mark.skip(reason="aligned against original Wan now")
@pytest.mark.usefixtures("distributed_setup")
 def test_wan_transformer():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -12,6 +12,7 @@ from typing import Any
 import imageio
 import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 import torchvision
 from einops import rearrange
@@ -91,8 +92,9 @@ class DistillationPipeline(TrainingPipeline):
            shift=self.timestep_shift)

        if training_args.real_score_model_path:
-            logger.info("Loading real score transformer from: %s",
-                        training_args.real_score_model_path)
+            logger.info(
+                f"Loading real score transformer from: {training_args.real_score_model_path}"
+            )
            self.real_score_transformer = self.load_module_from_path(
                training_args.real_score_model_path, "transformer",
                training_args)
@@ -101,8 +103,9 @@ class DistillationPipeline(TrainingPipeline):
                "real_score_transformer")

        if training_args.fake_score_model_path:
-            logger.info("Loading fake score transformer from: %s",
-                        training_args.fake_score_model_path)
+            logger.info(
+                f"Loading fake score transformer from: {training_args.fake_score_model_path}"
+            )
            self.fake_score_transformer = self.load_module_from_path(
                training_args.fake_score_model_path, "transformer",
                training_args)
@@ -191,13 +194,14 @@ class DistillationPipeline(TrainingPipeline):

        self.real_score_guidance_scale = self.training_args.real_score_guidance_scale

-        self.generator_ema: EMA_FSDP | None = None
+        self.generator_ema = None
        if (self.training_args.ema_decay
                is not None) and (self.training_args.ema_decay > 0.0):
            self.generator_ema = EMA_FSDP(self.transformer,
                                          decay=self.training_args.ema_decay)
-            logger.info("Initialized generator EMA with decay=%s",
-                        self.training_args.ema_decay)
+            logger.info(
+                f"Initialized generator EMA with decay={self.training_args.ema_decay}"
+            )
        else:
            logger.info("Generator EMA disabled (ema_decay <= 0.0)")

@@ -214,7 +218,7 @@ class DistillationPipeline(TrainingPipeline):
        Returns:
            The loaded module
        """
-        logger.info("Loading %s from custom path: %s", module_type, model_path)
+        logger.info(f"Loading {module_type} from custom path: {model_path}")
        # Set flag to prevent custom weight loading for teacher/critic models
        training_args._loading_teacher_critic_model = True

@@ -224,7 +228,7 @@ class DistillationPipeline(TrainingPipeline):

            # Download the model if it's a Hugging Face model ID
            local_model_path = maybe_download_model(model_path)
-            logger.info("Model downloaded/found at: %s", local_model_path)
+            logger.info(f"Model downloaded/found at: {local_model_path}")
            config = verify_model_config_and_directory(local_model_path)

            if module_type not in config:
@@ -233,8 +237,7 @@ class DistillationPipeline(TrainingPipeline):
                    extra_module = self._extra_config_module_map[module_type]
                    if extra_module in config:
                        module_type = extra_module
-                        logger.info("Using %s for %s", extra_module,
-                                    module_type)
+                        logger.info(f"Using {extra_module} for {module_type}")
                    else:
                        raise ValueError(
                            f"Module {module_type} not found in config at {local_model_path}"
@@ -259,8 +262,8 @@ class DistillationPipeline(TrainingPipeline):
                fastvideo_args=training_args,
            )

-            logger.info("Successfully loaded %s from %s", module_type,
-                        component_path)
+            logger.info(
+                f"Successfully loaded {module_type} from {component_path}")
            return module
        finally:
            # Always clean up the flag
@@ -290,7 +293,7 @@ class DistillationPipeline(TrainingPipeline):
                return model
        return model

-    def get_ema_model_copy(self) -> torch.nn.Module | None:
+    def get_ema_model_copy(self):
        """Get a copy of the model with EMA weights applied."""
        if self.generator_ema is not None:
            ema_model = copy.deepcopy(self.transformer)
@@ -298,7 +301,7 @@ class DistillationPipeline(TrainingPipeline):
            return ema_model
        return None

-    def is_ema_ready(self, current_step: int | None = None):
+    def is_ema_ready(self, current_step: int = None):
        """Check if EMA is ready for use (after ema_start_step)."""
        if current_step is None:
            current_step = getattr(self, 'current_trainstep', 0)
@@ -347,14 +350,14 @@ class DistillationPipeline(TrainingPipeline):
                with open(config_path, "w") as f:
                    json.dump(config_dict, f, indent=4)

-                logger.info("EMA weights saved to %s", weight_path)
+                logger.info(f"EMA weights saved to {weight_path}")

            del ema_model

        except Exception as e:
-            logger.error("Failed to save EMA weights: %s", str(e))
+            logger.error(f"Failed to save EMA weights: {str(e)}")

-    def get_ema_stats(self) -> dict[str, Any]:
+    def get_ema_stats(self):
        """Get EMA statistics for monitoring."""
        if self.generator_ema is None:
            return {
@@ -568,8 +571,7 @@ class DistillationPipeline(TrainingPipeline):

            noisy_latent = self.noise_scheduler.add_noise(
                generator_pred_video.flatten(0, 1), noise.flatten(0, 1),
-                timestep).detach().unflatten(0,
-                                             (1, generator_pred_video.shape[1]))
+                timestep).detach().unflatten(0, (1, generator_pred_video.shape[1]))

            # fake_score_transformer forward
            training_batch = self._build_distill_input_kwargs(
@@ -728,12 +730,11 @@ class DistillationPipeline(TrainingPipeline):
            "encoder_hidden_states": training_batch.encoder_hidden_states,
            "encoder_attention_mask": training_batch.encoder_attention_mask,
        }
-        if getattr(self, "negative_prompt_embeds", None) is not None:
-            unconditional_dict = {
-                "encoder_hidden_states": self.negative_prompt_embeds,
-                "encoder_attention_mask": self.negative_prompt_attention_mask,
-            }
-            training_batch.unconditional_dict = unconditional_dict
+        unconditional_dict = {
+            "encoder_hidden_states": self.negative_prompt_embeds,
+            "encoder_attention_mask": self.negative_prompt_attention_mask,
+        }
+        training_batch.unconditional_dict = unconditional_dict

        training_batch.dmd_latent_vis_dict = {}
        training_batch.fake_score_latent_vis_dict = {}
@@ -1250,8 +1251,9 @@ class DistillationPipeline(TrainingPipeline):
                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
                self.generator_ema = EMA_FSDP(
                    self.transformer, decay=self.training_args.ema_decay)
-                logger.info("Created generator EMA at step %s with decay=%s",
-                            step, self.training_args.ema_decay)
+                logger.info(
+                    f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}"
+                )

            with torch.autocast("cuda", dtype=torch.bfloat16):
                training_batch = self.train_one_step(training_batch)
@@ -1395,4 +1397,4 @@ class DistillationPipeline(TrainingPipeline):
                                  self.training_args.max_train_steps)

        if get_sp_group():
-            cleanup_dist_env_and_memory()
+            cleanup_dist_env_and_memory()
@@ -64,27 +64,31 @@ class ODEInitTrainingPipeline(TrainingPipeline):
        self.noise_scheduler.set_timesteps(num_inference_steps=1000,
                                           training=True)

-        logger.info("dmd_denoising_steps: %s",
-                    self.training_args.pipeline_config.dmd_denoising_steps)
+        # logger.info(f"ARG dmd_denoising_steps: {training_args.pipeline_config.dmd_denoising_steps}")
+        logger.info(
+            f"ARG dmd_denoising_steps: {self.training_args.pipeline_config.dmd_denoising_steps}"
+        )
        self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250],
                                                dtype=torch.long,
                                                device=get_local_torch_device())
+        # self.dmd_denoising_steps = torch.tensor([1000, 750, 500, 250], dtype=torch.long, device=get_local_torch_device())
        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
                                   torch.tensor([0],
                                                dtype=torch.float32))).cuda()
-            logger.info("timesteps: %s", timesteps)
+            logger.info(f"timesteps: {timesteps}")
            self.dmd_denoising_steps = timesteps[1000 -
                                                 self.dmd_denoising_steps]
-            logger.info("warped self.dmd_denoising_steps: %s",
-                        self.dmd_denoising_steps)
+            logger.info(
+                f"warped self.dmd_denoising_steps: {self.dmd_denoising_steps}")
+            # assert False, "warp_denoising_step must be false"
        else:
-            raise ValueError("warp_denoising_step must be true")
-
+            assert False, "warp_denoising_step must be true"
+            logger.info("not warped")
        self.dmd_denoising_steps = self.dmd_denoising_steps.to(
            get_local_torch_device())

-        logger.info("denoising_step_list: %s", self.dmd_denoising_steps)
+        logger.info(f"denoising_step_list: {self.dmd_denoising_steps}")

        logger.info(
            "Initialized ODE-init training pipeline with %s denoising steps",
@@ -92,7 +96,11 @@ class ODEInitTrainingPipeline(TrainingPipeline):
        # Cache for nearest trajectory index per DMD step (computed lazily on first batch)
        self._cached_closest_idx_per_dmd = None
        self.num_train_timestep = self.noise_scheduler.num_train_timesteps
-        self.manual_idx = 0
+        # self.min_timestep = int(self.training_args.min_timestep_ratio *
+        #                         self.num_train_timestep)
+        # self.max_timestep = int(self.training_args.max_timestep_ratio *
+        #                         self.num_train_timestep)
+        # self.real_score_guidance_scale = self.training_args.real_score_guidance_scale

    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        logger.info("Initializing validation pipeline...")
@@ -100,7 +108,8 @@ class ODEInitTrainingPipeline(TrainingPipeline):
        args_copy.inference_mode = True
        # Warm start validation with current transformer
        self.validation_pipeline = WanCausalDMDPipeline.from_pretrained(
-            training_args.model_path,
+            # training_args.model_path,
+            "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
            args=args_copy,  # type: ignore
            inference_mode=True,
            loaded_modules={
@@ -112,77 +121,76 @@ class ODEInitTrainingPipeline(TrainingPipeline):
            pin_cpu_memory=training_args.pin_cpu_memory,
            dit_cpu_offload=True)

-    def _get_next_batch(
-            self,
-            training_batch) -> tuple[TrainingBatch, torch.Tensor, torch.Tensor]:
-        batch = next(self.train_loader_iter, None)  # type: ignore
-        if batch is None:
-            self.current_epoch += 1
-            logger.info("Starting epoch %s", self.current_epoch)
-            self.train_loader_iter = iter(self.train_dataloader)
-            batch = next(self.train_loader_iter)
+    def _get_next_batch(self, training_batch):  # type: ignore[override]
+        # batch = next(self.train_loader_iter, None)  # type: ignore
+        # if batch is None:
+        #     self.current_epoch += 1
+        #     logger.info("Starting epoch %s", self.current_epoch)
+        #     self.train_loader_iter = iter(self.train_dataloader)
+        #     batch = next(self.train_loader_iter)

-        # Required fields from parquet (ODE trajectory schema)
-        encoder_hidden_states = batch['text_embedding']
-        encoder_attention_mask = batch['text_attention_mask']
-        infos = batch['info_list']
+        # # Required fields from parquet (ODE trajectory schema)
+        # encoder_hidden_states = batch['text_embedding']
+        # encoder_attention_mask = batch['text_attention_mask']
+        # infos = batch['info_list']

-        # Trajectory tensors may include a leading singleton batch dim per row
-        trajectory_latents = batch['trajectory_latents']
-        if trajectory_latents.dim() == 7:
-            # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
-            trajectory_latents = trajectory_latents[:, 0]
-        elif trajectory_latents.dim() == 6:
-            # already [B, S, C, T, H, W]
-            pass
-        else:
-            raise ValueError(
-                f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}"
-            )
+        # # Trajectory tensors may include a leading singleton batch dim per row
+        # trajectory_latents = batch['trajectory_latents']
+        # if trajectory_latents.dim() == 7:
+        #     # [B, 1, S, C, T, H, W] -> [B, S, C, T, H, W]
+        #     trajectory_latents = trajectory_latents[:, 0]
+        # elif trajectory_latents.dim() == 6:
+        #     # already [B, S, C, T, H, W]
+        #     pass
+        # else:
+        #     raise ValueError(
+        #         f"Unexpected trajectory_latents dim: {trajectory_latents.dim()}"
+        #     )

-        trajectory_timesteps = batch['trajectory_timesteps']
-        if trajectory_timesteps.dim() == 3:
-            # [B, 1, S] -> [B, S]
-            trajectory_timesteps = trajectory_timesteps[:, 0]
-        elif trajectory_timesteps.dim() == 2:
-            # [B, S]
-            pass
-        else:
-            raise ValueError(
-                f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}"
-            )
-        # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
-        trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)
+        # trajectory_timesteps = batch['trajectory_timesteps']
+        # if trajectory_timesteps.dim() == 3:
+        #     # [B, 1, S] -> [B, S]
+        #     trajectory_timesteps = trajectory_timesteps[:, 0]
+        # elif trajectory_timesteps.dim() == 2:
+        #     # [B, S]
+        #     pass
+        # else:
+        #     raise ValueError(
+        #         f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}"
+        #     )
+        # # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing
+        # trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)

-        # Move to device
+        # # Move to device
        device = get_local_torch_device()
-        training_batch.encoder_hidden_states = encoder_hidden_states.to(
-            device, dtype=torch.bfloat16)
-        training_batch.encoder_attention_mask = encoder_attention_mask.to(
-            device, dtype=torch.bfloat16)
-        training_batch.infos = infos
+        # training_batch.encoder_hidden_states = encoder_hidden_states.to(
+        #     device, dtype=torch.bfloat16)
+        # training_batch.encoder_attention_mask = encoder_attention_mask.to(
+        #     device, dtype=torch.bfloat16)
+        # training_batch.infos = infos

-        return training_batch, trajectory_latents.to(
-            device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
+        # return training_batch, trajectory_latents.to(
+        #     device, dtype=torch.bfloat16), trajectory_timesteps.to(device)

-        ## TEMP Used for loading the sf .pt files directly
-        """
-        self.manual_idx = self.manual_idx % 155
-        path = f"/mnt/weka/home/hao.zhang/wl/Self-Forcing/ode_pt_vidprom_1000/{self.manual_idx:05d}.pt"
-        logger.info("path: %s", path)
-        self.manual_idx += 1
-        # path = "/mnt/weka/home/hao.zhang/wl/Self-Forcing/ode_single_full/00000.pt"
+        ## TEMP
+        path = "/mnt/weka/home/hao.zhang/wei/FastVideo/data/ode_vidprom_1_fv/00000.pt"
        b = torch.load(path)
+        for k, v in b.items():
+            logger.info(f"b[{k}]: {type(v)}")
+            if isinstance(v, torch.Tensor):
+                logger.info(f"b[{k}]: {v.shape}")
+            else:
+                logger.info(f"b[{k}]: {v}")
        training_batch.encoder_hidden_states = b["text_embedding"][0].unsqueeze(
            0).to(device, dtype=torch.bfloat16)
        trajectory_latents = b["ode_latent"].to(device, dtype=torch.bfloat16)
-        logger.info("trajectory_latents: %s", trajectory_latents.shape)
-        logger.info("encoder_hidden_states: %s",
-                    training_batch.encoder_hidden_states.shape)
-        assert trajectory_latents.shape[1] <= 10, "trajectory_latents.shape[1] must be <= 10"
-        return training_batch, trajectory_latents.to(
-            device, dtype=torch.bfloat16), trajectory_timesteps.to(device)
-        """
+        logger.info(f"trajectory_latents: {trajectory_latents.shape}")
+        logger.info(
+            f"encoder_hidden_states: {training_batch.encoder_hidden_states.shape}"
+        )
+
+        return training_batch, trajectory_latents.to(device,
+                                                     dtype=torch.bfloat16), None

    def _get_timestep(self,
                      min_timestep: int,
@@ -216,33 +224,43 @@ class ODEInitTrainingPipeline(TrainingPipeline):
        encoder_attention_mask: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str,
                                                              torch.Tensor]]:
-        latent_vis_dict: dict[str, torch.Tensor] = {}
+        latent_vis_dict = {}
        device = get_local_torch_device()
        target_latent = traj_latents[:, -1]

+        # logger.info(f"traj_latents: {traj_latents.shape}")
+        # logger.info(f"traj_timesteps: {traj_timesteps.shape}")
+
        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
        B, S, num_frames, num_channels, height, width = traj_latents.shape

        # Lazily cache nearest trajectory index per DMD step based on the (fixed) S timesteps
        if self._cached_closest_idx_per_dmd is None:
+            # Use the first sample's trajectory timesteps; assumed identical across batches
+            # s_steps = traj_timesteps[0].to(torch.long)  # [S]
+            # dmd = cast(torch.Tensor, self.dmd_denoising_steps).to(s_steps.device)  # [K]
+            # distances_ks: [K, S] = |s_steps - dmd|
+            # distances_ks = (s_steps.unsqueeze(0) - dmd.unsqueeze(1)).abs()
+            # self._cached_closest_idx_per_dmd = distances_ks.argmin(dim=1).to(torch.long).cpu()  # [K]
            self._cached_closest_idx_per_dmd = torch.tensor(
                [0, 12, 24, 36], dtype=torch.long).cpu()
-            # [0, 1, 2, 3], dtype=torch.long).cpu()
-            logger.info("self._cached_closest_idx_per_dmd: %s",
-                        self._cached_closest_idx_per_dmd)
            logger.info(
-                "corresponding timesteps: %s", self.noise_scheduler.timesteps[
-                    self._cached_closest_idx_per_dmd])
+                f"self._cached_closest_idx_per_dmd: {self._cached_closest_idx_per_dmd}"
+            )
+            logger.info(
+                f"corresponding timesteps: {self.noise_scheduler.timesteps[self._cached_closest_idx_per_dmd]}"
+            )

+        # logger.info(f"traj_latents: {traj_latents.shape}")
        # Select the K indexes from traj_latents using self._cached_closest_idx_per_dmd
        # traj_latents: [B, S, C, T, H, W], self._cached_closest_idx_per_dmd: [K]
        # Output: [B, K, C, T, H, W]
-        assert self._cached_closest_idx_per_dmd is not None
-        relevant_traj_latents = torch.index_select(
-            traj_latents,
-            dim=1,
-            index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
-        logger.info("relevant_traj_latents: %s", relevant_traj_latents.shape)
+        # relevant_traj_latents = torch.index_select(
+        #     traj_latents,
+        #     dim=1,
+        #     index=self._cached_closest_idx_per_dmd.to(traj_latents.device))
+        relevant_traj_latents = traj_latents
+        logger.info(f"relevant_traj_latents sum: {relevant_traj_latents.float().sum().item()}, relevant_traj_latents: {relevant_traj_latents.shape}")
        # assert relevant_traj_latents.shape[0] == 1

        indexes = self._get_timestep(  # [B, num_frames]
@@ -252,8 +270,8 @@ class ODEInitTrainingPipeline(TrainingPipeline):
            num_frames,
            3,
            uniform_timestep=False)
-        logger.info("indexes: %s", indexes.shape)
-        logger.info("indexes: %s", indexes)
+        logger.info(f"indexes: {indexes.shape}")
+        logger.info(f"indexes: {indexes}")
        # noisy_input = relevant_traj_latents[indexes]
        noisy_input = torch.gather(
            relevant_traj_latents,
@@ -261,11 +279,28 @@ class ODEInitTrainingPipeline(TrainingPipeline):
            index=indexes.reshape(B, 1, num_frames, 1, 1,
                                  1).expand(-1, -1, -1, num_channels, height,
                                            width).to(self.device)).squeeze(1)
+        # noisy_input = noisy_input.unsqueeze(0)
+
+        # # Sample a single DMD step for the whole batch and fetch its cached nearest S-index
+        # K = len(self.dmd_denoising_steps)
+        # dmd_idx = torch.randint(0, K, (1,), device=device)
+        # logger.info(f"dmd_idx: {dmd_idx}")
+        # assert self._cached_closest_idx_per_dmd is not None
+        # nearest_s_idx = int(self._cached_closest_idx_per_dmd[int(dmd_idx.item())])
+        # nearest_idx = torch.full((B,), nearest_s_idx, device=device, dtype=torch.long)
+
+        # batch_indices = torch.arange(B, device=device)
+        # noisy_input = traj_latents[batch_indices, nearest_idx]  # [B, C, T, H, W]
+        # target_latent = traj_latents[batch_indices, -1]  # [B, C, T, H, W]
+        # t = traj_timesteps[batch_indices, nearest_idx]  # [B]
+
+        # Scale model input as in inference for consistency with stored trajectories
+        # noisy_input = self.modules["scheduler"].scale_model_input(noisy_input, t)
+        # logger.info(f"indexes: {indexes.shape}")
+        # logger.info(f"indexes: {indexes}")
        timestep = self.dmd_denoising_steps[indexes]
-        logger.info("selected timestep for rank %s: %s",
-                    self.global_rank,
-                    timestep,
-                    local_main_process_only=False)
+        # logger.info(f"timestep: {timestep.shape}")
+        # logger.info(f"timestep: {timestep}")

        # Prepare inputs for transformer
        latent_vis_dict["noisy_input"] = noisy_input.permute(
@@ -273,28 +308,45 @@ class ODEInitTrainingPipeline(TrainingPipeline):
        latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3,
                                                      4).detach().clone().cpu()

+        logger.info("noisy input sum: %s, noisy input shape: %s, noisy input dtype: %s", noisy_input.float().sum().item(), noisy_input.shape, noisy_input.dtype)
+        logger.info("encoder hidden states sum: %s, encoder hidden states shape: %s, encoder hidden states dtype: %s", encoder_hidden_states.float().sum().item(), encoder_hidden_states.shape, encoder_hidden_states.dtype)
+        logger.info("timestep sum: %s, timestep shape: %s, timestep dtype: %s", timestep.float().sum().item(), timestep.shape, timestep.dtype)
+        logger.info("model dtype set: %s. model weight sum: %s", set(p.dtype for p in self.transformer.parameters()), sum(p.float().sum().item() for p in self.transformer.parameters()))
+
+        # model_dtype = next(self.transformer.parameters()).dtype
+        model_dtype = torch.bfloat16
        input_kwargs = {
            "hidden_states": noisy_input.permute(0, 2, 1, 3, 4),
            "encoder_hidden_states": encoder_hidden_states,
-            "timestep": timestep.to(device, dtype=torch.bfloat16),
+            "timestep": timestep,
            "return_dict": False,
+            "scheduler": self.modules["scheduler"]
        }
        # Predict noise and step the scheduler to obtain next latent
        with set_forward_context(current_timestep=timestep,
                                 attn_metadata=None,
                                 forward_batch=None):
-            noise_pred = self.transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
+            pred_video = self.transformer(**input_kwargs)
+            # logger.info(f"noise_pred: {noise_pred.shape}")
+        # logger.info("noise pred sum: %s, noise pred shape: %s, noise pred dtype: %s", noise_pred.float().sum().item(), noise_pred.shape, noise_pred.dtype)
+        # if isinstance(noise_pred, (tuple, list)):
+        #     noise_pred = noise_pred[0]
+
+        # from fastvideo.models.utils import pred_noise_to_pred_video
+        # pred_video = pred_noise_to_pred_video(
+        #     pred_noise=noise_pred.flatten(0, 1),
+        #     noise_input_latent=noisy_input.flatten(0, 1),
+        #     timestep=timestep.to(dtype=model_dtype).flatten(0, 1),
+        #     scheduler=self.modules["scheduler"]).unflatten(
+        #         0, noise_pred.shape[:2])
+        logger.info("pred video sum: %s, pred video shape: %s, pred video dtype: %s", pred_video.float().sum().item(), pred_video.shape, pred_video.dtype)

-        from fastvideo.models.utils import pred_noise_to_pred_video
-        pred_video = pred_noise_to_pred_video(
-            pred_noise=noise_pred.flatten(0, 1),
-            noise_input_latent=noisy_input.flatten(0, 1),
-            timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1),
-            scheduler=self.modules["scheduler"]).unflatten(
-                0, noise_pred.shape[:2])
        latent_vis_dict["pred_video"] = pred_video.permute(
            0, 2, 1, 3, 4).detach().clone().cpu()

+        # noisy_input = pred_noise_to_pred_video(noise_pred, noisy_input, t, self.modules["scheduler"])
+        # next_latent_pred = self.modules["scheduler"].step(
+        #     noise_pred, t, current_latents, return_dict=False)[0]
        return pred_video, target_latent, timestep, latent_vis_dict

    def train_one_step(self, training_batch):  # type: ignore[override]
@@ -313,13 +365,29 @@ class ODEInitTrainingPipeline(TrainingPipeline):
            assert traj_latents.shape[0] == 1

            # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
-            _, S = traj_latents.shape[0], traj_latents.shape[1]
+            B, S = traj_latents.shape[0], traj_latents.shape[1]
            if S < 2:
                raise ValueError("Trajectory must contain at least 2 steps")

+            # Sample per-sample current step i in [0, S-2]
+
+            # idx = torch.randint(low=0, high=S - 1, size=(B, ),
+            #                     device=traj_latents.device)
+
+            # Gather current latents and next latents
+            # batch_indices = torch.arange(B, device=traj_latents.device)
+            # current_latents = traj_latents[batch_indices, idx]  # [B, C, T,H,W]
+            # current_latent = traj_timesteps[:, -1, :, :, :, :]
+            # target_latents = traj_latents[:, -1, :, :, :, :]
+
+            # Corresponding timesteps t (long) -> cast per sample
+            # t = traj_timesteps[:, -1, :, :, :, :]
+            # if t.dtype != torch.long:
+            #     t = t.long()
+
            # Forward to predict next latent by stepping scheduler with predicted noise
            noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
-                traj_latents, traj_timesteps, text_embeds, text_attention_mask)
+                traj_latents, training_batch.current_timestep, text_embeds, text_attention_mask)

            training_batch.latent_vis_dict.update(latent_vis_dict)

@@ -335,16 +403,28 @@ class ODEInitTrainingPipeline(TrainingPipeline):
                                     attn_metadata=None,
                                     forward_batch=None):
                loss.backward()
+
+            logger.info("Loss sum: %s, Loss dtype: %s", loss.float().sum().item(), loss.dtype)
+            assert self.transformer.blocks[0].to_q.weight.dtype == torch.float32
+            logger.info("blocks[0].to_q param sum before backprop: %s", self.transformer.blocks[0].to_q.weight.float().sum().item())
+            logger.info("blocks[0].to_q param grad sum: %s", self.transformer.blocks[0].to_q.weight.grad.float().sum().item())
+            logger.info("Transformer grad dtype: %s", set(p.grad.dtype for p in self.transformer.parameters()))
+            logger.info("Transformer param dtype before backprop: %s", set(p.dtype for p in self.transformer.parameters()))
            avg_loss = loss.detach().clone()
            training_batch.total_loss += avg_loss.item()

        # Clip grad and step optimizers
-        grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
-            [p for p in self.transformer.parameters() if p.requires_grad],
-            args.max_grad_norm if args.max_grad_norm is not None else 0.0)
+        # grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+        #     [p for p in self.transformer.parameters() if p.requires_grad],
+        #     args.max_grad_norm if args.max_grad_norm is not None else 0.0)
+        grad_norm = None

        self.optimizer.step()
-        self.lr_scheduler.step()
+        logger.info("Transformer param dtype after backprop: %s", set(p.dtype for p in self.transformer.parameters()))
+        logger.info("blocks[0].to_q param param after backprop sum: %s", self.transformer.blocks[0].to_q.weight.float().sum().item())
+
+        assert False
+        # self.lr_scheduler.step()

        if grad_norm is None:
            grad_value = 0.0
@@ -384,9 +464,15 @@ class ODEInitTrainingPipeline(TrainingPipeline):
        if self.global_rank == 0:
            wandb.log(wandb_loss_dict, step=step)

+        # dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
+        # fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
+        # fake_score_log_keys = ['generator_pred_video']
+        # dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video']
+

 def main(args) -> None:
    logger.info("Starting ODE-init training pipeline...")
+    logger.info(f"ARG dmd_denoising_steps: {args.dmd_denoising_steps}")
    pipeline = ODEInitTrainingPipeline.from_pretrained(
        args.pretrained_model_name_or_path, args=args)
    args = pipeline.training_args
@@ -22,7 +22,7 @@ from tqdm.auto import tqdm
 import fastvideo.envs as envs
 from fastvideo.attention.backends.video_sparse_attn import (
    VideoSparseAttentionMetadataBuilder)
-from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder
+# from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder
 from fastvideo.configs.sample import SamplingParam
 from fastvideo.dataset import build_parquet_map_style_dataloader
 from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
@@ -39,20 +39,26 @@ from fastvideo.training.activation_checkpoint import (
    apply_activation_checkpointing)
 from fastvideo.training.training_utils import (
    clip_grad_norm_while_handling_failing_dtensor_cases,
-    compute_density_for_timestep_sampling, count_trainable, get_scheduler,
-    get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint,
+    compute_density_for_timestep_sampling, get_scheduler, get_sigmas,
+    load_checkpoint, normalize_dit_input, save_checkpoint,
    shard_latents_across_sp)
-from fastvideo.utils import (is_vmoba_available, is_vsa_available,
+# from fastvideo.utils import (is_vmoba_available, is_vsa_available,
+#                              set_random_seed, shallow_asdict)
+from fastvideo.utils import (is_vsa_available,
                             set_random_seed, shallow_asdict)

 import wandb  # isort: skip

 vsa_available = is_vsa_available()
-vmoba_available = is_vmoba_available()
+# vmoba_available = is_vmoba_available()

 logger = init_logger(__name__)


+def _get_trainable_params(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
 class TrainingPipeline(LoRAPipeline, ABC):
    """
    A pipeline for training a model. All training pipelines should inherit from this class.
@@ -112,15 +118,18 @@ class TrainingPipeline(LoRAPipeline, ABC):
                enable_gradient_checkpointing_type)

        noise_scheduler = self.modules["scheduler"]
-        # Set grads for proper modules based on the training mode (Distill, LoRA, etc.)
        self.set_trainable()
        params_to_optimize = self.transformer.parameters()
        params_to_optimize = list(
            filter(lambda p: p.requires_grad, params_to_optimize))
+        # Parse betas from string format "beta1,beta2"
+        betas_str = training_args.betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+        
        self.optimizer = torch.optim.AdamW(
            params_to_optimize,
            lr=training_args.learning_rate,
-            betas=(0.9, 0.999),
+            betas=betas,
            weight_decay=training_args.weight_decay,
            eps=1e-8,
        )
@@ -272,20 +281,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
                patch_size=patch_size,
                VSA_sparsity=current_vsa_sparsity,
                device=get_local_torch_device())
-        elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
-            moba_params = self.training_args.moba_config.copy()
-            moba_params.update({
-                "current_timestep":
-                training_batch.timesteps,
-                "raw_latent_shape":
-                training_batch.raw_latent_shape[2:5],
-                "patch_size":
-                self.training_args.pipeline_config.dit_config.patch_size,
-                "device":
-                get_local_torch_device(),
-            })
-            training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder(
-            ).build(**moba_params)
+        # elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
+        #     moba_params = self.training_args.moba_config.copy()
+        #     moba_params.update({
+        #         "current_timestep":
+        #         training_batch.timesteps,
+        #         "raw_latent_shape":
+        #         training_batch.raw_latent_shape[2:5],
+        #         "patch_size":
+        #         self.training_args.pipeline_config.dit_config.patch_size,
+        #         "device":
+        #         get_local_torch_device(),
+        #     })
+        #     training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder(
+        #     ).build(**moba_params)
        else:
            training_batch.attn_metadata = None

@@ -310,7 +319,8 @@ class TrainingPipeline(LoRAPipeline, ABC):

    def _transformer_forward_and_compute_loss(
            self, training_batch: TrainingBatch) -> TrainingBatch:
-        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
+        # if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
            assert training_batch.attn_metadata is not None
        else:
            assert training_batch.attn_metadata is None
@@ -431,7 +441,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                    local_main_process_only=False)
        if not self.post_init_called:
            self.post_init()
-        num_trainable_params = count_trainable(self.transformer)
+        num_trainable_params = _get_trainable_params(self.transformer)
        logger.info("Starting training with %s B trainable parameters",
                    round(num_trainable_params / 1e9, 3))

@@ -476,9 +486,9 @@ class TrainingPipeline(LoRAPipeline, ABC):
                current_decay_times = min(step // vsa_decay_interval_steps,
                                          vsa_sparsity // vsa_decay_rate)
                current_vsa_sparsity = current_decay_times * vsa_decay_rate
-            elif vmoba_available:
-                # TODO: add vmoba sparsity scheduling here
-                current_vsa_sparsity = 0.0
+            # elif vmoba_available:
+            #     # TODO: add vmoba sparsity scheduling here
+            #     current_vsa_sparsity = 0.0
            else:
                current_vsa_sparsity = 0.0

@@ -527,7 +537,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                self._log_validation(self.transformer, self.training_args, step)
                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
                trainable_params = round(
-                    count_trainable(self.transformer) / 1e9, 3)
+                    _get_trainable_params(self.transformer) / 1e9, 3)
                logger.info(
                    "GPU memory usage after validation: %s MB, trainable params: %sB",
                    gpu_memory_usage, trainable_params)
@@ -563,7 +573,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
        logger.info("  Total optimization steps = %s",
                    self.training_args.max_train_steps)
        logger.info("  Total training parameters per FSDP shard = %s B",
-                    round(count_trainable(self.transformer) / 1e9, 3))
+                    round(_get_trainable_params(self.transformer) / 1e9, 3))
        # print dtype
        logger.info("  Master weight dtype: %s",
                    self.transformer.parameters().__next__().dtype)
@@ -631,6 +641,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
        validation_dataloader = DataLoader(validation_dataset,
                                           batch_size=None,
                                           num_workers=0)
+
        transformer.eval()

        validation_steps = training_args.validation_sampling_steps.split(",")
@@ -724,10 +735,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
        # Re-enable gradients for training
        training_args.inference_mode = False
        transformer.train()
-
+    
    def visualize_intermediate_latents(self, training_batch: TrainingBatch,
                                       training_args: TrainingArgs, step: int):
        """Add visualization data to wandb logging and save frames to disk."""
-        raise NotImplementedError(
-            "Visualize intermediate latents is not implemented for training pipeline"
-        )
+        raise NotImplementedError("Visualize intermediate latents is not implemented for training pipeline")
@@ -349,10 +349,14 @@ def load_checkpoint(transformer,
    """
    if not os.path.exists(checkpoint_path):
        logger.warning("Checkpoint path %s does not exist", checkpoint_path)
+        assert False
        return 0

    # Extract step number from checkpoint path
-    step = int(os.path.basename(checkpoint_path).split('-')[-1])
+    try:
+        step = int(os.path.basename(checkpoint_path).split('-')[-1])
+    except:
+        step = 1

    if rank == 0:
        logger.info("Loading checkpoint from step %s", step)
@@ -466,13 +470,11 @@ def load_distillation_checkpoint(generator_transformer,
            ema_state = generator_states.get("ema")
            if ema_state is not None:
                generator_ema.load_state_dict(ema_state)
-                logger.info("rank: %s, generator EMA state loaded successfully",
-                            rank)
+                logger.info("rank: %s, generator EMA state loaded successfully", rank)
            else:
                logger.info("rank: %s, no EMA state found in checkpoint", rank)
        except Exception as e:
-            logger.warning("rank: %s, failed to load EMA state: %s", rank,
-                           str(e))
+            logger.warning("rank: %s, failed to load EMA state: %s", rank, str(e))

    # Load critic distributed checkpoint
    critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
@@ -1298,10 +1300,6 @@ def get_scheduler(
                         last_epoch=last_epoch)


-def count_trainable(model: torch.nn.Module) -> int:
-    return sum(p.numel() for p in model.parameters() if p.requires_grad)
-
-
 class EMA_FSDP:
    """
    FSDP2-friendly EMA with two modes:
@@ -1323,7 +1321,6 @@ class EMA_FSDP:
      ema.update(model)
      ema.state_dict()  # on rank 0
    """
-
    def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"):
        self.decay = float(decay)
        self.mode = mode
@@ -1349,10 +1346,7 @@ class EMA_FSDP:
        if self.mode == "rank0_full":
            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
            if self.rank == 0:
-                self.shadow = {
-                    k: v.detach().clone().float().cpu()
-                    for k, v in cpu_state.items()
-                }
+                self.shadow = {k: v.detach().clone().float().cpu() for k, v in cpu_state.items()}
            else:
                self.shadow = {}
            return
@@ -1393,10 +1387,7 @@ class EMA_FSDP:

    def state_dict(self) -> dict[str, torch.Tensor]:
        if self.mode == "rank0_full":
-            return {
-                k: v.clone()
-                for k, v in self.shadow.items()
-            } if self.rank == 0 else {}
+            return {k: v.clone() for k, v in self.shadow.items()} if self.rank == 0 else {}
        return {k: v.clone() for k, v in self.shadow.items()}

    def load_state_dict(self, sd: dict[str, torch.Tensor]):
@@ -1417,7 +1408,6 @@ class EMA_FSDP:
                p.data.copy_(w.to(dtype=p.dtype, device=p.device))

    class _ApplyEMACtx:
-
        def __init__(self, ema: "EMA_FSDP", module):
            self.ema = ema
            self.module = module
@@ -1425,9 +1415,7 @@ class EMA_FSDP:

        def __enter__(self):
            if self.ema.mode != "local_shard":
-                raise RuntimeError(
-                    "EMA apply_to_model is only supported for mode='local_shard'"
-                )
+                raise RuntimeError("EMA apply_to_model is only supported for mode='local_shard'")
            with torch.no_grad():
                for name, p in self.module.named_parameters():
                    if not p.requires_grad:
@@ -1437,17 +1425,14 @@ class EMA_FSDP:
                    if p_local.numel() == 0:
                        # Nothing to swap on this rank for this param
                        continue
-                    self.saved[name] = p_local.clone().to(device=p_local.device,
-                                                          dtype=p_local.dtype)
+                    self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype)
                    if name in self.ema.shadow:
                        ema_cpu = self.ema.shadow[name]
                        if ema_cpu.numel() != p_local.numel():
                            # Shard shape mismatch (e.g., empty shard here), skip
                            continue
                        # Copy EMA shard into local param shard
-                        p_local.copy_(
-                            ema_cpu.to(dtype=p_local.dtype,
-                                       device=p_local.device))
+                        p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device))
            return self.module

        def __exit__(self, exc_type, exc, tb):
@@ -1465,4 +1450,4 @@ class EMA_FSDP:
            return False

    def apply_to_model(self, module):
-        return EMA_FSDP._ApplyEMACtx(self, module)
+        return EMA_FSDP._ApplyEMACtx(self, module)
@@ -4,10 +4,10 @@ from copy import deepcopy

 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.logger import init_logger
-from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import (
-    WanCausalDMDPipeline)
-from fastvideo.training.self_forcing_distillation_pipeline import (
-    SelfForcingDistillationPipeline)
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import WanCausalDMDPipeline
+from fastvideo.training.self_forcing_distillation_pipeline import SelfForcingDistillationPipeline
 from fastvideo.utils import is_vsa_available

 vsa_available = is_vsa_available()
@@ -69,4 +69,4 @@ if __name__ == "__main__":
    parser = TrainingArgs.add_cli_args(parser)
    parser = FastVideoArgs.add_cli_args(parser)
    args = parser.parse_args()
-    main(args)
+    main(args) 
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+from typing import Any
+
+import torch
+
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler)
+from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import (
+    WanImageToVideoPipeline)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.utils import is_vsa_available, shallow_asdict
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+class WanT2VI2VTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Wan.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def set_schemas(self):
+        self.train_dataset_schema = pyarrow_schema_t2v
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        args_copy.dit_cpu_offload = True
+        # args_copy.pipeline_config.vae_config.load_encoder = False
+        # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
+        pipeline_config = PipelineConfig.from_pretrained(
+            training_args.model_path)
+        pipeline_config.vae_config.load_encoder = True
+        self.validation_pipeline = WanImageToVideoPipeline.from_pretrained(
+            training_args.model_path,
+            args=None,
+            inference_mode=True,
+            pipeline_config=pipeline_config,
+            loaded_modules={
+                "transformer": self.get_module("transformer"),
+            },
+            required_config_modules=[
+                "scheduler", "transformer", "vae", "text_encoder", "tokenizer"
+            ],
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            dit_cpu_offload=True,
+        )
+
+    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        batch = next(self.train_loader_iter, None)  # type: ignore
+        if batch is None:
+            self.current_epoch += 1
+            logger.info("Starting epoch %s", self.current_epoch)
+            # Reset iterator for next epoch
+            self.train_loader_iter = iter(self.train_dataloader)
+            # Get first batch of new epoch
+            batch = next(self.train_loader_iter)
+
+        latents = batch['vae_latent']
+        latents = latents[:, :, :self.training_args.num_latent_t]
+        encoder_hidden_states = batch['text_embedding']
+        encoder_attention_mask = batch['text_attention_mask']
+        # clip_features = batch['clip_feature']
+        # image_latents = batch['first_frame_latent']
+        # image_latents = image_latents[:, :, :self.training_args.num_latent_t]
+        # pil_image = batch['pil_image']
+        infos = batch['info_list']
+
+        training_batch.latents = latents.to(get_local_torch_device(),
+                                            dtype=torch.bfloat16)
+        training_batch.encoder_hidden_states = encoder_hidden_states.to(
+            get_local_torch_device(), dtype=torch.bfloat16)
+        training_batch.encoder_attention_mask = encoder_attention_mask.to(
+            get_local_torch_device(), dtype=torch.bfloat16)
+        # training_batch.preprocessed_image = pil_image.to(
+        #     get_local_torch_device())
+        # training_batch.image_embeds = clip_features.to(get_local_torch_device())
+        # training_batch.image_latents = image_latents.to(
+        #     get_local_torch_device())
+        training_batch.infos = infos
+
+        return training_batch
+
+    def _prepare_dit_inputs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning."""
+
+        # First, call parent method to prepare noise, timesteps, etc. for video latents
+        training_batch = super()._prepare_dit_inputs(training_batch)
+
+        latents = training_batch.latents
+        logger.info("latents.shape: %s", latents.shape)
+        first_frame_latent = latents[:, :, 0, :, :]
+
+        logger.info("first_frame_latent.shape: %s", first_frame_latent.shape)
+        logger.info("training_batch.noisy_model_input.shape: %s",
+                    training_batch.noisy_model_input.shape)
+
+        training_batch.noisy_model_input = torch.cat([
+            first_frame_latent.unsqueeze(2),
+            training_batch.noisy_model_input[:, :, 1:, :, :]
+        ],
+                                                     dim=2)
+
+        return training_batch
+
+    def _build_input_kwargs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+
+        # Image Embeds for conditioning
+        # image_embeds = training_batch.image_embeds
+        # assert torch.isnan(image_embeds).sum() == 0
+        # image_embeds = image_embeds.to(get_local_torch_device(),
+        #                                dtype=torch.bfloat16)
+        # encoder_hidden_states_image = image_embeds
+
+        # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs
+        training_batch.input_kwargs = {
+            "hidden_states":
+            training_batch.noisy_model_input,
+            "encoder_hidden_states":
+            training_batch.encoder_hidden_states,
+            "timestep":
+            training_batch.timesteps.to(get_local_torch_device(),
+                                        dtype=torch.bfloat16),
+            "encoder_attention_mask":
+            training_batch.encoder_attention_mask,
+            # "encoder_hidden_states_image":
+            # encoder_hidden_states_image,
+            "return_dict":
+            False,
+        }
+        return training_batch
+
+    def _prepare_validation_batch(self, sampling_param: SamplingParam,
+                                  training_args: TrainingArgs,
+                                  validation_batch: dict[str, Any],
+                                  num_inference_steps: int) -> ForwardBatch:
+        sampling_param.prompt = validation_batch['prompt']
+        sampling_param.height = training_args.num_height
+        sampling_param.width = training_args.num_width
+        sampling_param.image_path = validation_batch['video_path']
+        sampling_param.num_inference_steps = num_inference_steps
+        sampling_param.data_type = "video"
+        assert self.seed is not None
+        sampling_param.seed = self.seed
+
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
+                        sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = (training_args.num_latent_t -
+                      1) * temporal_compression_factor + 1
+        sampling_param.num_frames = num_frames
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=torch.Generator(device="cpu").manual_seed(self.seed),
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+
+        return batch
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = WanT2VI2VTrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)
@@ -2,6 +2,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/utils.py

 import argparse
+from einops import rearrange
+import torchvision
+import numpy as np
+import imageio
 import ctypes
 import hashlib
 import importlib
@@ -4,12 +4,12 @@ import torch
 try:
    import flash_attn_interface

-    # def is_hopper_gpu():
-    #     if not torch.cuda.is_available():
-    #         return False
-    #     device_name = torch.cuda.get_device_name(0).lower()
-    #     return "h100" in device_name or "hopper" in device_name
-    FLASH_ATTN_3_AVAILABLE = True
+    def is_hopper_gpu():
+        if not torch.cuda.is_available():
+            return False
+        device_name = torch.cuda.get_device_name(0).lower()
+        return "h100" in device_name or "hopper" in device_name
+    FLASH_ATTN_3_AVAILABLE = is_hopper_gpu()
 except ModuleNotFoundError:
    FLASH_ATTN_3_AVAILABLE = False

@@ -18,7 +18,7 @@ try:
    FLASH_ATTN_2_AVAILABLE = True
 except ModuleNotFoundError:
    FLASH_ATTN_2_AVAILABLE = False
-assert FLASH_ATTN_3_AVAILABLE,"WTF"
+
 # FLASH_ATTN_3_AVAILABLE = False

 import warnings
@@ -113,7 +113,7 @@ def flash_attention(
            max_seqlen_k=lk,
            softmax_scale=softmax_scale,
            causal=causal,
-            deterministic=deterministic).unflatten(0, (b, lq))
+            deterministic=deterministic)[0].unflatten(0, (b, lq))
    else:
        assert FLASH_ATTN_2_AVAILABLE
        x = flash_attn.flash_attn_varlen_func(
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
JerryZhou54	92dfbfe4fc	Refactor ode trajectory generation Add env variable to enforce running flash attn	2025-09-26 01:26:03 +00:00
JerryZhou54	d31695b736	Test	2025-09-22 22:33:18 +00:00
JerryZhou54	357eaa4f1e	small change	2025-09-20 05:02:06 +00:00
JerryZhou54	ba02d0ae7e	Numerical Debugging	2025-09-20 05:02:06 +00:00
SolitaryThinker	c499783f74	fix ode_init	2025-09-20 05:02:06 +00:00
JerryZhou54	586013575c	Resolve timestep mismatch between dit forward and pred_noise_to_pred_video	2025-09-20 05:02:05 +00:00
SolitaryThinker	97e33f9097	debug	2025-09-20 05:02:05 +00:00
SolitaryThinker	8674224232	comment out vmoba	2025-09-20 05:02:04 +00:00
JerryZhou54	24d35b5512	Fix test for forward_train	2025-09-20 05:01:27 +00:00
JerryZhou54	3d2af0823c	Ensure 0 numerical diff for forward_train	2025-09-20 05:01:27 +00:00
JerryZhou54	27bac993ac	Add test for _forward_train	2025-09-20 05:01:27 +00:00
JerryZhou54	15a89d3f27	Fix numerical diff between causal_wanvideo.py and SF's causal wan	2025-09-20 05:01:25 +00:00
RandNMR73	ddf5e76ceb	text preprocessing ready	2025-09-20 05:00:53 +00:00
JerryZhou54	1f78d8114e	checkpoint	2025-09-20 05:00:06 +00:00
JerryZhou54	8cc2c113b8	Enable timestep warping & using SelfForcing scheduler	2025-09-20 04:59:33 +00:00
JerryZhou54	0b1519db19	Stop backprop through kv_cache	2025-09-20 04:59:33 +00:00
RandNMR73	055e7dc47b	new branch	2025-09-20 04:59:31 +00:00
RandNMR73	0892c2c448	inference works after changes added	2025-09-20 04:57:54 +00:00
SolitaryThinker	c501a484e0	lmdb datasets	2025-09-20 04:57:54 +00:00
SolitaryThinker	d965f578c9	t2v to i2v finetune checkpoint ode checkpoint fix t2v to i2v lint chekpt checkpoint hacked but working ode_init scripts WIP fixing time embedding WIP fixing time embedding checkpoint update fix revert revert revert update update visualize	2025-09-20 04:57:52 +00:00