diff --git a/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm b/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm
index ff12ddf10..90a78e0fc 100644
--- a/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm
+++ b/examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm
@@ -22,27 +22,44 @@ export NODE_RANK=$SLURM_PROCID
 nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
 export MASTER_ADDR=${nodes[0]}
 export TOKENIZERS_PARALLELISM=false
-export WANDB_API_KEY='your_wandb_api_key_here'
+# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv-distill
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
 # Configs
 NUM_GPUS=8
 NUM_TOTAL_GPUS=32
 
 # Model paths for Self-Forcing DMD distillation with Hunyuan1.5:
-GENERATOR_MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v"
+GENERATOR_MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v"  # Updated to Wan2.2
 REAL_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" # Teacher model
 FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+# REAL_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled"
+# FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled" # Critic model
 
-DATA_DIR="your_data_dir"
-VALIDATION_DATASET_FILE="your_validation_data_dir"
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+# DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/vidprom_16k_text_embed"
+DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
 
 training_args=(
-  --tracker_project_name SFhy1.5_t2v_distill_self_forcing_dmd
-  --output_dir "your_output_dir"
-  --wandb_run_name "your_wandb_run_name"
+  --tracker_project_name SFhy1.5_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_3333_1e-5_1e-5_cfg6_corrected_scheduler"
+  --wandb_run_name "self_forcing_3333_context_forcing"
   --max_train_steps 4000
   --train_batch_size 1
   --train_sp_batch_size 1
@@ -55,13 +72,17 @@ training_args=(
   --log_visualization
   --num_frames 81
   --num_frame_per_block 3  # Frame generation block size for self-forcing
+  # --resume-from-checkpoint "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333/checkpoint-300"
+  --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_worldplay_tf_init_3333/checkpoint-2800/transformer/diffusion_pytorch_model.safetensors
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_worldplay_df_init_1333/checkpoint-11600/transformer/diffusion_pytorch_model.safetensors
+  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_distilled_ode_init_3333/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors
 )
 
 parallel_args=(
-  --num_gpus $NUM_TOTAL_GPUS
+  --num_gpus $NUM_TOTAL_GPUS # 64
   --sp_size 1
   --tp_size 1
-  --hsdp_replicate_dim 1
+  --hsdp_replicate_dim 1 # 64
   --hsdp_shard_dim $NUM_TOTAL_GPUS
 )
 
@@ -74,6 +95,7 @@ model_args=(
 
 dataset_args=(
   --data_path "$DATA_DIR"
+  # --data_path_2 "$DATA_DIR_2"
   --dataloader_num_workers 4
 )
 
@@ -112,7 +134,9 @@ miscellaneous_args=(
 
 dmd_args=(
   --warp_denoising_step
-  --dmd_denoising_steps '1000,875,750,625,500,375,250,125'
+  # --dmd_denoising_steps '1000,875,750,625,500,375,250,125'
+  # --dmd_denoising_steps '1000,760,520,280'
+  # --dmd_denoising_steps '1000,750,500,250'
   --min_timestep_ratio 0.02
   --max_timestep_ratio 0.98
   --dfake_gen_update_ratio 5
@@ -126,6 +150,8 @@ self_forcing_args=(
   --same_step_across_blocks True  # Whether to use same denoising step across all blocks
   --last_step_only False  # Whether to only use the last denoising step
   --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+  # --use-ode-init True
+  # --use-context-forcing True
 )
 
 srun torchrun \
diff --git a/examples/inference/basic/basic_wan2_2_Fun.py b/examples/inference/basic/basic_wan2_2_Fun.py
index 61c216074..48135dad8 100644
--- a/examples/inference/basic/basic_wan2_2_Fun.py
+++ b/examples/inference/basic/basic_wan2_2_Fun.py
@@ -10,12 +10,13 @@ def main():
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
     generator = VideoGenerator.from_pretrained(
-        "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
+        "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers",
+        # "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
         # "alibaba-pai/Wan2.2-Fun-A14B-Control",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
         use_fsdp_inference=False, # set to True if GPU is out of memory
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
+        dit_cpu_offload=False, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
         # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
@@ -23,14 +24,19 @@ def main():
         # image_encoder_cpu_offload=False,
     )
 
-    prompt = "一位年轻女性穿着一件粉色的连衣裙，裙子上有白色的装饰和粉色的纽扣。她的头发是紫色的，头上戴着一个红色的大蝴蝶结，显得非常可爱和精致。她还戴着一个红色的领结，整体造型充满了少女感和活力。她的表情温柔，双手轻轻交叉放在身前，姿态优雅。背景是简单的灰色，没有任何多余的装饰，使得人物更加突出。她的妆容清淡自然，突显了她的清新气质。整体画面给人一种甜美、梦幻的感觉，仿佛置身于童话世界中。"
-    negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    # prompt = "一位年轻女性穿着一件粉色的连衣裙，裙子上有白色的装饰和粉色的纽扣。她的头发是紫色的，头上戴着一个红色的大蝴蝶结，显得非常可爱和精致。她还戴着一个红色的领结，整体造型充满了少女感和活力。她的表情温柔，双手轻轻交叉放在身前，姿态优雅。背景是简单的灰色，没有任何多余的装饰，使得人物更加突出。她的妆容清淡自然，突显了她的清新气质。整体画面给人一种甜美、梦幻的感觉，仿佛置身于童话世界中。"
+    # negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
     # prompt                  = "A young woman with beautiful, clear eyes and blonde hair stands in the forest, wearing a white dress and a crown. Her expression is serene, reminiscent of a movie star, with fair and youthful skin. Her brown long hair flows in the wind. The video quality is very high, with a clear view. High quality, masterpiece, best quality, high resolution, ultra-fine, fantastical."
     # negative_prompt         = "Twisted body, limb deformities, text captions, comic, static, ugly, error, messy code."
-    image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png"
-    control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4"
-
-    video = generator.generate_video(prompt, negative_prompt=negative_prompt, image_path=image_path, video_path=control_video_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True)
+    # image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png"
+    # control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4"
+    import json, os
+    with open("/mnt/weka/home/hao.zhang/wei/FastVideo/data/mixkit_i2v.jsonl", "r") as f:
+        prompt_image_pairs = json.load(f)
+    for d in prompt_image_pairs:
+        prompt = d["prompt"]
+        image_path = os.path.join("/mnt/weka/home/hao.zhang/wei/FastVideo", d["image_path"])
+        video = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True)
 
 if __name__ == "__main__":
     main()
\ No newline at end of file
diff --git a/examples/inference/basic/basic_wan2_2_ti2v.py b/examples/inference/basic/basic_wan2_2_ti2v.py
index df0c18848..d616ffda7 100644
--- a/examples/inference/basic/basic_wan2_2_ti2v.py
+++ b/examples/inference/basic/basic_wan2_2_ti2v.py
@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator
 
-OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"
+OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v_720p_77"
 def main():
     # FastVideo will automatically use the optimal default arguments for the
     # model.
@@ -12,17 +12,29 @@ def main():
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
         use_fsdp_inference=False, # set to True if GPU is out of memory
-        dit_cpu_offload=True,
+        dit_cpu_offload=False,
         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"
         # image_encoder_cpu_offload=False,
     )
 
+    import json, os
+    with open("/mnt/weka/home/hao.zhang/wei/FastVideo/data/mixkit_i2v_720p.jsonl", "r") as f:
+        prompt_image_pairs = json.load(f)
+    for d in prompt_image_pairs:
+        prompt = d["prompt"]
+        image_path = os.path.join("/mnt/weka/home/hao.zhang/wei/FastVideo", d["image_path"])
+        video = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_PATH, save_video=True, num_frames=77)
+    return
+
     # I2V is triggered just by passing in an image_path argument
     prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
     image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
+    # prompt = "In the midst of the joyous New Year's Eve celebration, the cheerful group of friends, their spirits lifted by the festivities, decides to immortalize the moment with a vibrant snapshot"
+    # image_path = "images/mixkit-a-cheerful-group-of-friends-celebrate-new-years-eve-and-51525.png"
     video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path)
+    return
 
     # Generate another video with a different prompt, without reloading the
     # model!
diff --git a/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_tf_init.slurm b/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_tf_init.slurm
new file mode 100644
index 000000000..9c5f12627
--- /dev/null
+++ b/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_tf_init.slurm
@@ -0,0 +1,133 @@
+#!/bin/bash
+#SBATCH --job-name=hy15_tf_3333
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=hy15_tf_3333_output/hy15_tf_3333_%j.out
+#SBATCH --error=hy15_tf_3333_output/hy15_tf_3333_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# 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 TOKENIZERS_PARALLELISM=false
+# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
+export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+source ~/conda/miniconda/bin/activate
+conda activate wei-fv-distill
+export HOME="/mnt/weka/home/hao.zhang/wei"
+
+# Configs
+NUM_GPUS=8
+NUM_TOTAL_GPUS=64
+
+MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v"
+DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "hy15_ode_init"
+  --output_dir "/mnt/weka/home/hao.zhang/wei/hy15_worldplay_tf_init_3333"
+  --wandb_run_name "hy15_worldplay_tf_init_3333"
+  --warp_denoising_step
+  --dmd_denoising_steps '1000,760,520,280,0'
+  --log_visualization
+  --visualization_steps 100
+  --max_train_steps 30000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 848
+  --num_frames 81
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_TOTAL_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+  --text-encoder-cpu-offload
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 4
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --optimizer-type "muon"
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 400
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+validation_args=(
+  # --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 10
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "1.0" # not used for dmd inference
+  # --vae_cpu_offload True
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 6
+  --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/hy15_ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
\ No newline at end of file
diff --git a/fastvideo/configs/models/encoders/clip.py b/fastvideo/configs/models/encoders/clip.py
index d233872f1..7cc5c70d5 100644
--- a/fastvideo/configs/models/encoders/clip.py
+++ b/fastvideo/configs/models/encoders/clip.py
@@ -87,10 +87,10 @@ class CLIPVisionConfig(ImageEncoderConfig):
     arch_config: ImageEncoderArchConfig = field(
         default_factory=CLIPVisionArchConfig)
 
-    num_hidden_layers_override: int | None = None
+    num_hidden_layers_override: int | None = 31
     require_post_norm: bool | None = None
-    enable_scale: bool = True
-    is_causal: bool = True
+    enable_scale: bool = False
+    is_causal: bool = False
     prefix: str = "clip"
 
 
diff --git a/fastvideo/dataset/dataloader/record_schema.py b/fastvideo/dataset/dataloader/record_schema.py
index 1bc86dd7d..fa8b8676a 100644
--- a/fastvideo/dataset/dataloader/record_schema.py
+++ b/fastvideo/dataset/dataloader/record_schema.py
@@ -127,7 +127,10 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
 def ode_text_only_record_creator(
         video_name: str, text_embedding: np.ndarray, caption: str,
         trajectory_latents: np.ndarray,
-        trajectory_timesteps: np.ndarray) -> dict[str, Any]:
+        trajectory_timesteps: np.ndarray,
+        text_mask: np.ndarray | None = None,
+        text_embedding_2: np.ndarray | None = None,
+        text_mask_2: np.ndarray | None = None) -> dict[str, Any]:
     """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only.
 
     Args:
@@ -165,6 +168,25 @@ def ode_text_only_record_creator(
         "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
     })
 
+    if text_embedding_2 is not None:
+        record.update({
+            "text_embedding_2_bytes": text_embedding_2.tobytes(),
+            "text_embedding_2_shape": list(text_embedding_2.shape),
+            "text_embedding_2_dtype": str(text_embedding_2.dtype),
+        })
+    if text_mask is not None:
+        record.update({
+            "text_mask_bytes": text_mask.tobytes(),
+            "text_mask_shape": list(text_mask.shape),
+            "text_mask_dtype": str(text_mask.dtype),
+        })
+    if text_mask_2 is not None:
+        record.update({
+            "text_mask_2_bytes": text_mask_2.tobytes(),
+            "text_mask_2_shape": list(text_mask_2.shape),
+            "text_mask_2_dtype": str(text_mask_2.dtype),
+        })
+        
     return record
 
 
@@ -187,4 +209,4 @@ def text_only_record_creator(text_name: str, text_embedding: np.ndarray,
         "text_embedding_dtype": str(text_embedding.dtype),
         "caption": caption,
     }
-    return record
+    return record
\ No newline at end of file
diff --git a/fastvideo/dataset/dataloader/schema.py b/fastvideo/dataset/dataloader/schema.py
index 048c7686c..1fced7e53 100644
--- a/fastvideo/dataset/dataloader/schema.py
+++ b/fastvideo/dataset/dataloader/schema.py
@@ -90,6 +90,15 @@ pyarrow_schema_ode_trajectory_text_only = pa.schema([
     pa.field("text_embedding_shape", pa.list_(pa.int64())),
     # e.g., 'bfloat16' or 'float32'
     pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_embedding_2_bytes", pa.binary()),
+    pa.field("text_embedding_2_shape", pa.list_(pa.int64())),
+    pa.field("text_embedding_2_dtype", pa.string()),
+    pa.field("text_mask_bytes", pa.binary()),
+    pa.field("text_mask_shape", pa.list_(pa.int64())),
+    pa.field("text_mask_dtype", pa.string()),
+    pa.field("text_mask_2_bytes", pa.binary()),
+    pa.field("text_mask_2_shape", pa.list_(pa.int64())),
+    pa.field("text_mask_2_dtype", pa.string()),
     # --- ODE Trajectory ---
     pa.field("trajectory_latents_bytes", pa.binary()),
     pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
@@ -115,45 +124,4 @@ pyarrow_schema_text_only = pa.schema([
     pa.field("text_embedding_dtype", pa.string()),
     # --- Metadata ---
     pa.field("caption", pa.string()),
-])
-
-
-pyarrow_schema_matrixgame = pa.schema([
-    pa.field("id", pa.string()),
-    # --- Image/Video VAE latents ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("vae_latent_bytes", pa.binary()),
-    # e.g., [C, T, H, W] or [C, H, W]
-    pa.field("vae_latent_shape", pa.list_(pa.int64())),
-    # e.g., 'float32'
-    pa.field("vae_latent_dtype", pa.string()),
-    #I2V
-    pa.field("clip_feature_bytes", pa.binary()),
-    pa.field("clip_feature_shape", pa.list_(pa.int64())),
-    pa.field("clip_feature_dtype", pa.string()),
-    pa.field("first_frame_latent_bytes", pa.binary()),
-    pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
-    pa.field("first_frame_latent_dtype", pa.string()),
-    # --- Action ---
-    pa.field("mouse_cond_bytes", pa.binary()),
-    pa.field("mouse_cond_shape", pa.list_(pa.int64())),  # [T, 2]
-    pa.field("mouse_cond_dtype", pa.string()),
-    pa.field("keyboard_cond_bytes", pa.binary()),
-    pa.field("keyboard_cond_shape", pa.list_(pa.int64())),  # [T, 4]
-    pa.field("keyboard_cond_dtype", pa.string()),
-    # I2V Validation
-    pa.field("pil_image_bytes", pa.binary()),
-    pa.field("pil_image_shape", pa.list_(pa.int64())),
-    pa.field("pil_image_dtype", pa.string()),
-    # --- Metadata ---
-    pa.field("file_name", pa.string()),
-    pa.field("caption", pa.string()),
-    pa.field("media_type", pa.string()),  # 'image' or 'video'
-    pa.field("width", pa.int64()),
-    pa.field("height", pa.int64()),
-    # -- Video-specific (can be null/default for images) ---
-    # Number of frames processed (e.g., 1 for image, N for video)
-    pa.field("num_frames", pa.int64()),
-    pa.field("duration_sec", pa.float64()),
-    pa.field("fps", pa.float64()),
-])
+])
\ No newline at end of file
diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py
index d94727bd6..563c3cc74 100644
--- a/fastvideo/dataset/preprocessing_datasets.py
+++ b/fastvideo/dataset/preprocessing_datasets.py
@@ -17,6 +17,7 @@ from PIL import Image
 from transformers import AutoTokenizer
 
 from fastvideo.logger import init_logger
+from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
 
@@ -40,7 +41,6 @@ class PreprocessBatch:
     num_frames: int | None = None
     sample_frame_index: list[int] | None = None
     sample_num_frames: int | None = None
-    action_path: str | None = None
 
     # Processed data
     pixel_values: torch.Tensor | None = None
@@ -471,11 +471,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
 
         # Initialize tokenizer
         tokenizer_path = os.path.join(args.model_path, "tokenizer")
-        if os.path.exists(tokenizer_path):
-            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                                      cache_dir=args.cache_dir)
-        else:
-            tokenizer = None
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                                  cache_dir=args.cache_dir)
 
         # Initialize processing stages
         self._init_stages(args, transform, transform_topcrop, tokenizer)
@@ -499,14 +496,11 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         self.video_transform_stage = VideoTransformStage(transform)
         self.image_transform_stage = ImageTransformStage(
             transform, transform_topcrop)
-        if tokenizer is not None:
-            self.text_encoding_stage = TextEncodingStage(
-                tokenizer=tokenizer,
-                text_max_length=args.text_max_length,
-                cfg_rate=args.training_cfg_rate,
-                seed=self.seed)
-        else:
-            self.text_encoding_stage = None
+        self.text_encoding_stage = TextEncodingStage(
+            tokenizer=tokenizer,
+            text_max_length=args.text_max_length,
+            cfg_rate=args.training_cfg_rate,
+            seed=self.seed)
 
     def _load_raw_data(self) -> list[dict]:
         """Load raw data from JSON files."""
@@ -528,8 +522,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         # Update paths with folder prefix
         for item in data_items:
             item["path"] = opj(folder, item["path"])
-            if "action_path" in item and item["action_path"]:
-                item["action_path"] = opj(folder, item["action_path"])
 
         return data_items
 
@@ -551,8 +543,7 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                                     cap=item["cap"],
                                     resolution=item.get("resolution"),
                                     fps=item.get("fps"),
-                                    duration=item.get("duration"),
-                                    action_path=item.get("action_path"))
+                                    duration=item.get("duration"))
 
             # Apply filtering stages
             if not self._apply_filter_stages(batch, filter_counts):
@@ -614,27 +605,20 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         # Apply transformation stages
         batch = self.video_transform_stage.process(batch)
         batch = self.image_transform_stage.process(batch)
-        if self.text_encoding_stage is not None:
-            batch = self.text_encoding_stage.process(batch)
+        batch = self.text_encoding_stage.process(batch)
 
         # Build result dictionary
         result = {
             "pixel_values": batch.pixel_values,
+            "text": batch.text,
+            "input_ids": batch.input_ids,
+            "cond_mask": batch.cond_mask,
             "path": batch.path,
         }
 
-        if batch.text is not None:
-            result["text"] = batch.text
-            result["input_ids"] = batch.input_ids
-            result["cond_mask"] = batch.cond_mask
-
         # Add video-specific fields
         if batch.is_video:
             result.update({"fps": batch.fps, "duration": batch.duration})
-        
-        # Add action_path
-        if batch.action_path:
-            result["action_path"] = batch.action_path
 
         return result
 
@@ -672,7 +656,7 @@ class TextDataset(torch.utils.data.IterableDataset,
         self.seed = seed
 
         # Initialize tokenizer
-        tokenizer_path = os.path.join(args.model_path, "tokenizer")
+        tokenizer_path = os.path.join(maybe_download_model(args.model_path), "tokenizer")
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
                                                   cache_dir=args.cache_dir)
 
@@ -775,4 +759,4 @@ class TextDataset(torch.utils.data.IterableDataset,
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         """Load state dict from checkpoint."""
-        self.processed_batches = state_dict["processed_batches"]
+        self.processed_batches = state_dict["processed_batches"]
\ No newline at end of file
diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py
index 36ae58e7b..142529cac 100644
--- a/fastvideo/dataset/utils.py
+++ b/fastvideo/dataset/utils.py
@@ -154,8 +154,14 @@ def collate_rows_from_parquet_schema(rows,
                     ) if rng else random.random()) < cfg_rate:
                         data = np.zeros((512, 4096), dtype=np.float32)
                     else:
-                        data = np.frombuffer(
-                            bytes_data, dtype=np.float32).reshape(shape).copy()
+                        if row[f"{tensor_name}_dtype"] == "float32":
+                            data = np.frombuffer(
+                                bytes_data, dtype=np.float32).reshape(shape).copy()
+                        elif row[f"{tensor_name}_dtype"] == "int64":
+                            data = np.frombuffer(
+                                bytes_data, dtype=np.int64).reshape(shape).copy()
+                        else:
+                            raise ValueError(f"Unsupported dtype: {row[f"{tensor_name}_dtype"]}")
                     tensor = torch.from_numpy(data)
                     # if len(data.shape) == 3:
                     #     B, L, D = tensor.shape
@@ -168,7 +174,7 @@ def collate_rows_from_parquet_schema(rows,
                 tensor_list.append(torch.zeros(0, dtype=torch.bfloat16))
 
         # Stack tensors with special handling for text embeddings
-        if tensor_name == 'text_embedding':
+        if tensor_name == 'null':
             # Handle text embeddings with padding
             padded_tensors = []
             attention_masks = []
@@ -222,4 +228,4 @@ def collate_rows_from_parquet_schema(rows,
     if info_list and 'caption' in info_list[0]:
         batch_data['caption_text'] = [info['caption'] for info in info_list]
 
-    return batch_data
+    return batch_data
\ No newline at end of file
diff --git a/fastvideo/models/dits/causal_hunyuanvideo15.py b/fastvideo/models/dits/causal_hunyuanvideo15.py
index 3d7627397..e0656ef4a 100644
--- a/fastvideo/models/dits/causal_hunyuanvideo15.py
+++ b/fastvideo/models/dits/causal_hunyuanvideo15.py
@@ -275,12 +275,27 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Apply rotary embeddings
         cos, sin = freqs_cis
-        img_q = _apply_rotary_emb(img_q, cos, sin, is_neox_style=False)
-        img_k = _apply_rotary_emb(img_k, cos, sin, is_neox_style=False)
 
         # Apply flex_attention
         # Does not support SP padding for now
         if kv_cache is None:
+            is_tf = (img_q.shape[1] == cos.shape[0] * 2)
+            if is_tf:
+                q_chunk = torch.chunk(img_q, 2, dim=1)
+                k_chunk = torch.chunk(img_k, 2, dim=1)
+                roped_query = []
+                roped_key = []
+                for ii in range(2):
+                    rq = _apply_rotary_emb(q_chunk[ii], cos, sin, is_neox_style=False)
+                    rk = _apply_rotary_emb(k_chunk[ii], cos, sin, is_neox_style=False)
+                    roped_query.append(rq)
+                    roped_key.append(rk)
+                img_q = torch.cat(roped_query, dim=1)
+                img_k = torch.cat(roped_key, dim=1)
+            else:
+                img_q = _apply_rotary_emb(img_q, cos, sin, is_neox_style=False)
+                img_k = _apply_rotary_emb(img_k, cos, sin, is_neox_style=False)
+
             q = img_q
             k = torch.cat([img_k, txt_kv_cache["k_txt"]], dim=1)
             v = torch.cat([img_v, txt_kv_cache["v_txt"]], dim=1)
@@ -316,6 +331,9 @@ class MMDoubleStreamBlock(nn.Module):
             assert img_attn.shape[1] == image_seq_len
             updated_kv_cache = None
         else:
+            img_q = _apply_rotary_emb(img_q, cos, sin, is_neox_style=False)
+            img_k = _apply_rotary_emb(img_k, cos, sin, is_neox_style=False)
+
             current_end = current_start + img_q.shape[1]
             num_new_tokens = img_q.shape[1]
             sink_tokens = self.sink_size * 1590
@@ -569,6 +587,101 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT):
         assert encoder_hidden_states.shape[0] == 1
         return encoder_hidden_states, encoder_attention_mask
 
+
+    @staticmethod
+    def _prepare_teacher_forcing_mask(
+        device: torch.device | str, num_frames: int = 21,
+        frame_seqlen: int = 1560, num_frame_per_block=1,
+        text_seq_len: int = 0
+    ) -> BlockMask:
+        """
+        we will divide the token sequence into the following format
+        [1 latent frame] [1 latent frame] ... [1 latent frame]
+        We use flexattention to construct the attention mask
+        """
+        # debug
+        DEBUG = False
+        if DEBUG:
+            num_frames = 9
+            frame_seqlen = 256
+
+        total_length = num_frames * frame_seqlen * 2
+        total_kv_length = total_length + text_seq_len
+
+        # we do right padding to get to a multiple of 128
+        padded_length = math.ceil(total_length / 128) * 128 - total_length
+        kv_padded_length = math.ceil(total_kv_length / 128) * 128 - total_kv_length
+
+        total_length_tensor = torch.tensor(total_length, device=device)
+        total_kv_length_tensor = torch.tensor(total_kv_length, device=device)
+
+        clean_ends = num_frames * frame_seqlen
+        # for clean context frames, we can construct their flex attention mask based on a [start, end] interval
+        context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+        # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end]
+        noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+        noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+        noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+        noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+
+        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+        attention_block_size = frame_seqlen * num_frame_per_block
+        frame_indices = torch.arange(
+            start=0,
+            end=num_frames * frame_seqlen,
+            step=attention_block_size,
+            device=device, dtype=torch.long
+        )
+
+        # attention for clean context frames
+        for start in frame_indices:
+            context_ends[start:start + attention_block_size] = start + attention_block_size
+
+        noisy_image_start_list = torch.arange(
+            num_frames * frame_seqlen, total_length,
+            step=attention_block_size,
+            device=device, dtype=torch.long
+        )
+        noisy_image_end_list = noisy_image_start_list + attention_block_size
+
+        # attention for noisy frames
+        for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):
+            # attend to noisy tokens within the same block
+            noise_noise_starts[start:end] = start
+            noise_noise_ends[start:end] = end
+            # attend to context tokens in previous blocks
+            # noise_context_starts[start:end] = 0
+            noise_context_ends[start:end] = block_index * attention_block_size
+
+        def attention_mask(b, h, q_idx, kv_idx):
+            # first design the mask for clean frames
+            clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])
+            # then design the mask for noisy frames
+            # noisy frames will attend to all clean preceeding clean frames + itself
+            C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])
+            C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])
+            noise_mask = (q_idx >= clean_ends) & (C1 | C2)
+
+            eye_mask = q_idx == kv_idx
+            text_mask = (kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor)
+            return eye_mask | clean_mask | noise_mask | text_mask
+
+        block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
+                                       KV_LEN=total_kv_length + kv_padded_length, _compile=False, device=device)
+
+        if DEBUG:
+            print(block_mask)
+            import imageio
+            import numpy as np
+            from torch.nn.attention.flex_attention import create_mask
+
+            mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
+                                padded_length, KV_LEN=total_kv_length + kv_padded_length, device=device)
+            import cv2
+            mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
+            imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
+
+        return block_mask
     
     @staticmethod
     def _prepare_blockwise_causal_attn_mask(
@@ -712,6 +825,8 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT):
         txt_kv_cache: list | None = None,
         current_start: int = 0,
         rope_start_idx: int = 0,
+        clean_hidden_states: Optional[torch.Tensor] = None,
+        aug_timestep: Optional[torch.Tensor] = None,
     ):
         assert txt_kv_cache is not None, "txt_kv_cache must be provided"
 
@@ -740,16 +855,38 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT):
 
         hidden_states = self.img_in(hidden_states)
 
+        if clean_hidden_states is not None:
+            clean_hidden_states = self.img_in(clean_hidden_states)
+            hidden_states = torch.cat([clean_hidden_states, hidden_states], dim=1)
+
+            if aug_timestep is None:
+                aug_timestep = torch.zeros_like(timestep)
+            if aug_timestep.dim() == 2:
+                ts_seq_len = aug_timestep.shape[1]
+                aug_temb = self.time_in(aug_timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
+            else:
+                aug_temb = self.time_in(aug_timestep, timestep_r=timestep_r)
+            temb = torch.cat([aug_temb, temb], dim=1)
+
         # Prepare block-wise causal attention mask
         if kv_cache is None:
-            self.block_mask = self._prepare_blockwise_causal_attn_mask(
-                device=hidden_states.device,
-                num_frames=num_frames,
-                frame_seqlen=post_patch_height * post_patch_width,
-                num_frame_per_block=self.num_frame_per_block,
-                local_attn_size=self.local_attn_size,
-                text_seq_len=txt_kv_cache[0]["k_txt"].shape[1]
-            )
+            if clean_hidden_states is not None:
+                self.block_mask = self._prepare_teacher_forcing_mask(
+                    device=hidden_states.device,
+                    num_frames=num_frames,
+                    frame_seqlen=post_patch_height * post_patch_width,
+                    num_frame_per_block=self.num_frame_per_block,
+                    text_seq_len=txt_kv_cache[0]["k_txt"].shape[1]
+                )
+            else:
+                self.block_mask = self._prepare_blockwise_causal_attn_mask(
+                    device=hidden_states.device,
+                    num_frames=num_frames,
+                    frame_seqlen=post_patch_height * post_patch_width,
+                    num_frame_per_block=self.num_frame_per_block,
+                    local_attn_size=self.local_attn_size,
+                    text_seq_len=txt_kv_cache[0]["k_txt"].shape[1]
+                )
 
         # 4. Transformer blocks
         if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -787,6 +924,9 @@ class CausalHunyuanVideo15Transformer3DModel(CachableDiT):
                     for k in new_cache.keys():
                         kv_cache[block_index][k] = new_cache[k].clone()
 
+        if clean_hidden_states is not None:
+            hidden_states = hidden_states[:, hidden_states.shape[1] // 2:]
+
         # Final layer processing
         hidden_states = self.final_layer(hidden_states, temb)
         # Unpatchify to get original shape
diff --git a/fastvideo/pipelines/stages/hy15_causal_denoising.py b/fastvideo/pipelines/stages/hy15_causal_denoising.py
index 4af2061c4..f4a3d0134 100644
--- a/fastvideo/pipelines/stages/hy15_causal_denoising.py
+++ b/fastvideo/pipelines/stages/hy15_causal_denoising.py
@@ -101,6 +101,7 @@ class Hy15CausalDMDDenosingStage(CausalDMDDenosingStage):
         pos_start_base = 0
         num_blocks = math.ceil(t / self.num_frames_per_block)
         block_sizes = [self.num_frames_per_block] * num_blocks
+        # block_sizes[0] = 1
         start_index = 0
 
         # Initialize txt kv cache
diff --git a/fastvideo/training/hy15_ode_causal_pipeline.py b/fastvideo/training/hy15_ode_causal_pipeline.py
index 521d2fabc..66e00aa83 100644
--- a/fastvideo/training/hy15_ode_causal_pipeline.py
+++ b/fastvideo/training/hy15_ode_causal_pipeline.py
@@ -432,7 +432,106 @@ class Hy15ODEInitTrainingPipeline(TrainingPipeline):
         with set_forward_context(current_timestep=timestep,
                                  attn_metadata=None,
                                  forward_batch=None):
-            noise_pred = self.transformer(**vision_input_kwargs).permute(
+            noise_pred, _ = self.transformer(**vision_input_kwargs)
+            noise_pred = noise_pred.permute(
+                0, 2, 1, 3, 4)
+
+        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=latents.flatten(0, 1),
+            timestep=timestep.to(dtype=torch.bfloat16).flatten(0, 1),
+            scheduler=self.noise_scheduler).unflatten(0, noise_pred.shape[:2])
+        latent_vis_dict["pred_video"] = pred_video.permute(
+            0, 2, 1, 3, 4).detach().clone().cpu()
+
+        return pred_video, target_latent, timestep, latent_vis_dict
+
+    def _tf_step_predict_next_latent(
+        self, traj_latents: torch.Tensor, traj_timesteps: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        encoder_hidden_states_image: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str,
+                                                              torch.Tensor]]:
+        latent_vis_dict: dict[str, torch.Tensor] = {}
+        device = get_local_torch_device()
+        target_latent = traj_latents[:, -1]
+        del traj_latents
+
+        # Shapes: traj_latents [B, S, C, T, H, W], traj_timesteps [B, S]
+        B, num_frames, num_channels, height, width = target_latent.shape
+
+        indexes = self._get_timestep(  # [B, num_frames]
+            0,
+            1000,
+            B,
+            num_frames,
+            3,
+            uniform_timestep=False)
+        timestep = self.noise_scheduler.timesteps[indexes.cpu()].to(device)
+
+        latents = self.noise_scheduler.add_noise(
+            target_latent.flatten(0, 1),
+            torch.randn_like(target_latent.flatten(0, 1)),
+            timestep.flatten(0, 1)).unflatten(0, (B, num_frames))
+        noisy_input = torch.cat([
+            latents,
+            torch.zeros_like(latents),
+            torch.zeros_like(latents[:, :, 0:1])
+        ],
+                                dim=2)
+        clean_input = torch.cat([
+            target_latent,
+            torch.zeros_like(latents),
+            torch.zeros_like(latents[:, :, 0:1])
+        ],
+                                dim=2)
+
+        # Prepare inputs for transformer
+        latent_vis_dict["noisy_input"] = latents.permute(
+            0, 2, 1, 3, 4).detach().clone().cpu()
+        latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3,
+                                                      4).detach().clone().cpu()
+
+        logger.info("timestep: %s", timestep)
+        txt_input_kwargs = {
+            "txt_inference":
+            True,
+            "vision_inference":
+            False,
+            "encoder_hidden_states":
+            encoder_hidden_states,
+            "encoder_hidden_states_image":
+            encoder_hidden_states_image,
+            "encoder_attention_mask":
+            encoder_attention_mask,
+            "timestep":
+            torch.zeros([latents.shape[0]],
+                        device=latents.device,
+                        dtype=torch.bfloat16),
+            "cache_txt":
+            True,
+        }
+        with set_forward_context(current_timestep=timestep,
+                                 attn_metadata=None,
+                                 forward_batch=None):
+            txt_kv_cache = self.transformer(**txt_input_kwargs)
+
+        vision_input_kwargs = {
+            "txt_inference": False,
+            "vision_inference": True,
+            "hidden_states": noisy_input.permute(0, 2, 1, 3, 4),
+            "timestep": timestep.to(device, dtype=torch.bfloat16),
+            "txt_kv_cache": txt_kv_cache,
+            "clean_hidden_states": clean_input.permute(0, 2, 1, 3, 4),
+        }
+        # 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(**vision_input_kwargs)
+            noise_pred = noise_pred.permute(
                 0, 2, 1, 3, 4)
 
         from fastvideo.models.utils import pred_noise_to_pred_video
@@ -468,7 +567,7 @@ class Hy15ODEInitTrainingPipeline(TrainingPipeline):
                 raise ValueError("Trajectory must contain at least 2 steps")
 
             # Forward to predict next latent by stepping scheduler with predicted noise
-            noise_pred, target_latent, t, latent_vis_dict = self._step_predict_next_latent(
+            noise_pred, target_latent, t, latent_vis_dict = self._tf_step_predict_next_latent(
                 traj_latents, traj_timesteps, text_embeds, text_attention_mask,
                 image_embeds)
 
diff --git a/fastvideo/training/hy15_self_forcing_distillation_pipeline.py b/fastvideo/training/hy15_self_forcing_distillation_pipeline.py
index c48415414..b5d6fc490 100644
--- a/fastvideo/training/hy15_self_forcing_distillation_pipeline.py
+++ b/fastvideo/training/hy15_self_forcing_distillation_pipeline.py
@@ -134,8 +134,7 @@ class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
             dist.broadcast(start_timestep_index, src=0)
         start_timestep_index = start_timestep_index.cpu().item()
 
-        trajectory_latents = trajectory_latents[:, -1, :self.training_args.
-                                                num_latent_t]
+        trajectory_latents = trajectory_latents[:, -1, :self.training_args.num_latent_t]
         training_batch.trajectory_latents = self.noise_scheduler.add_noise(
             trajectory_latents.flatten(0, 1),
             torch.randn_like(trajectory_latents.flatten(0, 1)),
@@ -153,9 +152,7 @@ class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
             self, training_batch: TrainingBatch) -> TrainingBatch:
         trajectory_latents = training_batch.trajectory_latents
 
-        training_batch.trajectory_latents = trajectory_latents[:, -1, :self.
-                                                               training_args.
-                                                               num_latent_t]
+        training_batch.trajectory_latents = trajectory_latents[:, -1, :self.training_args.num_latent_t]
         # _cached_closest_idx_per_dmd = torch.tensor([0, 12, 24, 36, 50],
         #                                            dtype=torch.long).cpu()
         # training_batch.trajectory_latents = torch.index_select(
@@ -242,7 +239,7 @@ class Hy15SelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
 
         # Step 2: Temporal denoising loop
         all_num_frames = [self.num_frame_per_block] * num_blocks
-        # all_num_frames[0] = 1
+        all_num_frames[0] = 1
         if self.independent_first_frame and initial_latent is None:
             all_num_frames = [1] + all_num_frames