refactor hy15 sf distill

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

.buildkite/pipeline.yml

@@ -61,7 +61,7 @@ steps:
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
-                command: "timeout 60m .buildkite/scripts/pr_test.sh"
+                command: "timeout 90m .buildkite/scripts/pr_test.sh"
                 label: "SSIM Tests"
                 env:
                   - TEST_TYPE=ssim

.github/workflows/fastvideo-kernel-publish.yml

@@ -156,8 +156,23 @@ jobs:
           
           # Fix the wheel to be manylinux compliant
           pip install auditwheel
+          # Point auditwheel at torch libs, but do not vendor them into the wheel.
+          TORCH_LIB_DIR=$(python - <<'PY'
+          import os
+          import torch
+
+          print(os.path.join(os.path.dirname(torch.__file__), "lib"))
+          PY
+          )
+          export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${LD_LIBRARY_PATH}"
           # Target manylinux_2_35 (Ubuntu 22.04 native)
-          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist
+          auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist \
+            --exclude libtorch_cuda.so \
+            --exclude libtorch_cpu.so \
+            --exclude libtorch.so \
+            --exclude libc10.so \
+            --exclude libc10_cuda.so \
+            --exclude libtorch_python.so
           # Move fixed wheels back to dist for upload consistency
           rm dist/*.whl
           mv fixed_dist/*.whl dist/

.pre-commit-config.yaml

@@ -68,7 +68,7 @@ repos:
     entry: bash
     args:
       - -c
-      - 'git ls-files | grep -v "^fastvideo/tests/ssim/" | grep -v "^fastvideo/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^\"*fastvideo/tests/ssim/" | grep -v "^\"*fastvideo/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
     language: system
     always_run: true
     pass_filenames: false

docs/attention/vsa/index.md

@@ -13,11 +13,13 @@ from fastvideo_kernel import video_sparse_attn
 
 # q, k, v: [batch_size, num_heads, seq_len, head_dim]
 # variable_block_sizes: Number of valid tokens per block
+# q_variable_block_sizes: Number of valid tokens per q block (can differ from KV for q/k of different lengths)
 # topk: Number of blocks to attend
 
 output = video_sparse_attn(
     q, k, v, 
-    variable_block_sizes=block_sizes, 
+    block_sizes,
+    block_sizes,
     topk=32
 )
 ```

examples/distill/SFHunyuan1.5-T2V/distill_dmd.slurm

@@ -0,0 +1,169 @@
+#!/bin/bash
+#SBATCH --job-name=dmd_3333
+#SBATCH --partition=main
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_3333_output/dmd_3333_%j.out
+#SBATCH --error=dmd_3333_output/dmd_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=32
+
+# Model paths for Self-Forcing DMD distillation with Hunyuan1.5:
+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="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  # 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
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 848
+  --enable_gradient_checkpointing_type "full"
+  --simulate_generator_forward
+  --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_df_init_3333/checkpoint-2000/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 # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_TOTAL_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  # --data_path_2 "$DATA_DIR_2"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "8"
+  --validation_guidance_scale "1.0" # not used for dmd inference
+  --text-encoder-cpu-offload
+  # --vae_cpu_offload True
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 100
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 200
+)
+
+dmd_args=(
+  --warp_denoising_step
+  --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
+  --real_score_guidance_scale 3.5
+  --fake_score_learning_rate 5e-6
+  --fake_score_betas '0.0,0.999'
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --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-context-forcing True
+)
+
+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_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"

examples/inference/basic/basic_cosmos2_5_t2w.py

@@ -0,0 +1,42 @@
+from fastvideo import VideoGenerator
+
+
+def main():
+    # Point this to your local diffusers model dir (or replace with a HF model ID).
+    model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
+
+    generator = VideoGenerator.from_pretrained(
+        model_path,
+        num_gpus=1,
+        use_fsdp_inference=False,  # set True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    prompt = (
+       "A high-definition video captures the precision of robotic welding in an industrial setting. The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. The welding process is in full swing, with bright sparks and intense light illuminating the scene, creating a vivid display of blue and white hues. A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, indicating a busy and functional industrial workspace. As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of the welding operation."
+    )
+
+    video = generator.generate_video(
+        prompt,
+        negative_prompt="",
+        height=704,
+        width=1280,
+        num_frames=77,
+        num_inference_steps=35,
+        guidance_scale=7.0,
+        fps=24,
+        output_path="outputs_video/cosmos2_5_t2w.mp4",
+        save_video=True,
+    )
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
+
+
+

examples/inference/basic/basic_hy15.py

@@ -2,14 +2,14 @@ from fastvideo import VideoGenerator
 import json
 # from fastvideo.configs.sample import SamplingParam
 
-OUTPUT_PATH = "video_samples_hy15"
+OUTPUT_PATH = "video_samples_hy15_t2v_distilled"
 def main():
     # FastVideo will automatically use the optimal default arguments for the
     # model.
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
     generator = VideoGenerator.from_pretrained(
-        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled",
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
         use_fsdp_inference=False, # set to True if GPU is out of memory
@@ -18,15 +18,35 @@ def main():
         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,
+        # init_weights_from_safetensors="/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333_5e-5_2e-6_cfg3.5/checkpoint-500/ema/generator_ema.safetensors"
     )
 
+    # json_path = "/mnt/weka/home/hao.zhang/wei/FastVideo/data/mixkit_i2v_full_720p.json"
+    # with open(json_path, 'r') as f:
+    #     data_list = json.load(f)["data"]
+    
+    # # Now you can index into data_list however you like
+    # # For example: data_list[0], data_list[1:3], etc.
+    # for data in data_list:
+    #     prompt = data["prompt"]
+    #     image_path = data["image_path"]
+    #     generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path, num_frames=121, fps=24)
+    # return
+
+    # prompt = (
+    #     "In a brightly lit studio, a photographer wearing a denim jacket focuses intently, capturing shots with a professional camera. Facing him, a model stands gracefully, adjusting her long, flowing hair with delicate movements. The scene is characterized by strong contrasts; the model's soft pink attire and gentle gestures complement the rugged, precise demeanor of the photographer. Positioned against a minimalist backdrop, the pair work seamlessly, with the camera\u2019s lens pointed directly at the model, capturing her elegance. The soft, diffused lighting casts a gentle glow on both subjects, creating an airy and ethereal atmosphere perfect for a high-fashion photo shoot."
+    # )
+
+    # video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=121, fps=24, image_path="data/1.png")
+    # return
+
     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, negative_prompt="", num_frames=81, fps=16)
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
 
     prompt2 = (
         "A majestic lion strides across the golden savanna, its powerful frame "
@@ -35,7 +55,7 @@ def main():
         "embodying the raw energy of the wild. Low angle, steady tracking shot, "
         "cinematic.")
 
-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
 
 
 if __name__ == "__main__":

examples/inference/basic/basic_hyworld.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Basic example for HYWorld (HY-WorldPlay) video generation using FastVideo.
+
+This example replicates the same functionality as HY-WorldPlay/run.sh,
+demonstrating image-to-video generation with camera trajectory control.
+"""
+
+import time
+import math
+import numpy as np
+import imageio
+import torchvision
+from einops import rearrange
+
+from fastvideo import VideoGenerator
+from fastvideo.pipelines import ForwardBatch
+from fastvideo.utils import shallow_asdict, align_to
+from fastvideo.logger import init_logger
+
+from fastvideo.models.dits.hyworld.pose import pose_to_input, compute_latent_num
+from fastvideo.models.dits.hyworld.resolution_utils import get_resolution_from_image
+
+logger = init_logger(__name__)
+
+class HYWorldVideoGenerator(VideoGenerator):
+    """Extended VideoGenerator that adds HYWorld-specific parameters to batch.extra."""
+
+    def _generate_single_video(self, prompt: str, sampling_param=None, **kwargs):
+        """Override to add viewmats, Ks, and action to batch.extra."""
+        fastvideo_args = self.fastvideo_args
+        pipeline_config = fastvideo_args.pipeline_config
+
+        if sampling_param is None:
+            from fastvideo.configs.sample import SamplingParam
+            sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path)
+
+        # Update sampling param with kwargs
+        if kwargs:
+            for key, value in kwargs.items():
+                if hasattr(sampling_param, key):
+                    setattr(sampling_param, key, value)
+
+        # Get pose string from sampling_param or kwargs
+        pose = kwargs.get('pose', getattr(sampling_param, 'POSE', 'w-31'))
+        num_frames = kwargs.get('num_frames', getattr(sampling_param, 'num_frames', 125))
+
+        # Calculate number of latents
+        latent_num = compute_latent_num(num_frames)
+
+        # Convert pose to viewmats, Ks, and action
+        viewmats, Ks, action = pose_to_input(pose, latent_num)
+
+        # Convert to tensors and add batch dimension
+        viewmats = viewmats.unsqueeze(0)  # (1, T, 4, 4)
+        Ks = Ks.unsqueeze(0)  # (1, T, 3, 3)
+        action = action.unsqueeze(0)  # (1, T)
+
+        # Validate inputs
+        prompt = prompt.strip()
+        sampling_param = sampling_param.__class__(**shallow_asdict(sampling_param))
+        output_path = kwargs.get("output_path", sampling_param.output_path)
+        sampling_param.prompt = prompt
+
+        if sampling_param.negative_prompt is not None:
+            sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
+
+        # Validate dimensions
+        if (sampling_param.height <= 0 or sampling_param.width <= 0 or
+            sampling_param.num_frames <= 0):
+            raise ValueError(
+                f"Height, width, and num_frames must be positive integers")
+
+        temporal_scale_factor = pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = sampling_param.num_frames
+        num_gpus = fastvideo_args.num_gpus
+        use_temporal_scaling_frames = pipeline_config.vae_config.use_temporal_scaling_frames
+
+        # Adjust number of frames based on number of GPUs
+        if use_temporal_scaling_frames:
+            orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1
+        else:
+            orig_latent_num_frames = sampling_param.num_frames // 17 * 3
+
+        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
+            if use_temporal_scaling_frames:
+                new_num_frames = (orig_latent_num_frames - 1) * temporal_scale_factor + 1
+            else:
+                divisor = math.lcm(3, num_gpus)
+                orig_latent_num_frames = (
+                    (orig_latent_num_frames + divisor - 1) // divisor) * divisor
+                new_num_frames = orig_latent_num_frames // 3 * 17
+
+            logger.info(
+                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
+                sampling_param.num_frames, new_num_frames, fastvideo_args.num_gpus)
+            sampling_param.num_frames = new_num_frames
+
+        # Calculate sizes
+        target_height = align_to(sampling_param.height, 16)
+        target_width = align_to(sampling_param.width, 16)
+
+        # Calculate latent sizes
+        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]
+
+        # Prepare batch
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            eta=0.0,
+            n_tokens=n_tokens,
+            VSA_sparsity=fastvideo_args.VSA_sparsity,
+        )
+
+        # Add HYWorld-specific parameters to batch.extra
+        batch.extra['viewmats'] = viewmats
+        batch.extra['Ks'] = Ks
+        batch.extra['action'] = action
+        batch.extra['chunk_latent_frames'] = 16  # For bidirectional model
+
+        # Run inference
+        start_time = time.perf_counter()
+        output_batch = self.executor.execute_forward(batch, fastvideo_args)
+        samples = output_batch.output
+        logging_info = output_batch.logging_info
+
+        gen_time = time.perf_counter() - start_time
+        logger.info("Generated successfully in %.2f seconds", gen_time)
+
+        # Process outputs
+        videos = rearrange(samples, "b c t h w -> t b c h w")
+        frames = []
+        for x in videos:
+            x = torchvision.utils.make_grid(x, nrow=6)
+            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+            frames.append((x * 255).numpy().astype(np.uint8))
+
+        # Save video if requested
+        if batch.save_video:
+            imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
+            logger.info("Saved video to %s", output_path)
+
+        if batch.return_frames:
+            return frames
+        else:
+            return {
+                "samples": samples,
+                "frames": frames,
+                "prompts": prompt,
+                "size": (target_height, target_width, batch.num_frames),
+                "generation_time": gen_time,
+                "logging_info": logging_info,
+                "trajectory": output_batch.trajectory_latents,
+                "trajectory_timesteps": output_batch.trajectory_timesteps,
+                "trajectory_decoded": output_batch.trajectory_decoded,
+            }
+
+# Default prompt from HY-WorldPlay run.sh
+DEFAULT_PROMPT = 'A paved pathway leads towards a stone arch bridge spanning a calm body of water.  Lush green trees and foliage line the path and the far bank of the water. A traditional-style pavilion with a tiered, reddish-brown roof sits on the far shore. The water reflects the surrounding greenery and the sky.  The scene is bathed in soft, natural light, creating a tranquil and serene atmosphere. The pathway is composed of large, rectangular stones, and the bridge is constructed of light gray stone.  The overall composition emphasizes the peaceful and harmonious nature of the landscape.'
+DEFAULT_IMAGE = 'https://raw.githubusercontent.com/Tencent-Hunyuan/HY-WorldPlay/main/assets/img/test.png'
+
+def main():
+    import argparse
+
+    # pose: (a, w, s, d) - (15, 31)
+    # num_frames: (61, 125)
+    parser = argparse.ArgumentParser(description="HYWorld video generation with FastVideo")
+    parser.add_argument("--prompt", type=str, default=DEFAULT_PROMPT, help="Text prompt for video generation")
+    parser.add_argument("--image", type=str, default=DEFAULT_IMAGE, help="Path or URL to input image")
+    parser.add_argument("--pose", type=str, default='w-31', help="Pose string (e.g., 'a-31', 'w-31', 's-31', 'd-31')")
+    parser.add_argument("--output_path", type=str, default='video_samples_hyworld', help="Output video path")
+    parser.add_argument("--num-frames", type=int, default=125, help="Number of frames")
+    parser.add_argument("--seed", type=int, default=1, help="Random seed")
+    parser.add_argument("--resolution", type=str, default="480p", help="Only support 480p for now")
+    args = parser.parse_args()
+
+    # Automatically determine resolution from input image
+    HEIGHT, WIDTH = get_resolution_from_image(args.image, args.resolution)
+    print(f"Image: {args.image}")
+    print(f"Pose: {args.pose}")
+    print(f"Resolution: {HEIGHT}x{WIDTH} (from {args.resolution} buckets)")
+    print(f"Num frames: {args.num_frames}")
+    print(f"Output path: {args.output_path}")
+
+    # Initialize generator
+    print("\nInitializing VideoGenerator for HYWorld...")
+    
+    generator = HYWorldVideoGenerator.from_pretrained(
+        "FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+        image_encoder_cpu_offload=True,
+    )
+
+    # Generate video
+    print("\nGenerating video...")
+    start_time = time.time()
+    video = generator.generate_video(
+        prompt=args.prompt,
+        image_path=args.image,
+        output_path=args.output_path,
+        save_video=True,
+        negative_prompt="",
+        num_frames=args.num_frames,
+        fps=24,
+        height=HEIGHT,
+        width=WIDTH,
+        seed=args.seed,
+        pose=args.pose,
+    )
+    elapsed = time.time() - start_time
+
+    print(f"\nVideo generated successfully!")
+    print(f"Saved to: {args.output_path}")
+    print(f"Time: {elapsed:.2f}s")
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_ltx2.py

@@ -0,0 +1,34 @@
+from fastvideo import VideoGenerator
+
+PROMPT = (
+    "A warm sunny backyard. The camera starts in a tight cinematic close-up "
+    "of a woman and a man in their 30s, facing each other with serious "
+    "expressions. The woman, emotional and dramatic, says softly, \"That's "
+    "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
+    "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, "
+    "then mutters defensively, \"He's just having fun.\" The camera slowly "
+    "pans right, revealing the grandfather in the garden wearing enormous "
+    "butterfly wings, waving his arms in the air like he's trying to take "
+    "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. "
+    "The woman covers her face, on the verge of tears. The tone is deadpan, "
+    "absurd, and quietly tragic."
+)
+
+
+def main() -> None:
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/LTX2-Distilled-Diffusers",
+        num_gpus=1,
+    )
+
+    output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4"
+    generator.generate_video(
+        prompt=PROMPT,
+        output_path=output_path,
+        save_video=True,
+    )
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh

@@ -1,36 +1,38 @@
 #!/bin/bash
 
 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
+export WANDB_MODE=offline
 export TOKENIZERS_PARALLELISM=false
 
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
+DATA_DIR="data/ode-preprocessing-hy15-test/"
+VALIDATION_DATASET_FILE="data/validation_64.json"
 NUM_GPUS=1
 # IP=[MASTER NODE IP]
 
 # Training arguments
 training_args=(
   --tracker_project_name "wan_ode_init"
-  --output_dir "wan_ode_init_crush_smol"
-  --override_transformer_cls_name "CausalWanTransformer3DModel"
-  --wandb_run_name "wan_ode_init_crush_smol"
-  --max_train_steps 6000
+  --output_dir "ode_init_hy15_test"
+  --wandb_run_name "vidprom_bz128_1e-5"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  # --warp_denoising_step
+  # --log_visualization
+  --max_train_steps 6001
   --train_batch_size 1
   --train_sp_batch_size 1
   --gradient_accumulation_steps 1
-  --num_latent_t 21
+  --num_latent_t 19
   --num_height 480
   --num_width 832
-  --num_frames 77
-  --warp_denoising_step
+  --num_frames 73
+  --dmd_denoising_steps "1000,750,500,250"
   --enable_gradient_checkpointing_type "full"
 )
 
 # Parallel arguments
 parallel_args=(
-  --num_gpus $NUM_GPUS
+  --num_gpus 1
   --sp_size 1
   --tp_size 1
   --hsdp_replicate_dim 1
@@ -51,18 +53,17 @@ dataset_args=(
 
 # Validation arguments
 validation_args=(
-  --log_validation
+  # --log_validation
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
+  --validation_steps 200
   --validation_sampling_steps "50"
   --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 6e-6
+  --learning_rate 1e-5
   --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
   --training_state_checkpointing_steps 1000
   --weight_decay 1e-4
   --max_grad_norm 1.0
@@ -78,6 +79,7 @@ miscellaneous_args=(
   --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

examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm

@@ -0,0 +1,132 @@
+#!/bin/bash
+#SBATCH --job-name=hy15_ode_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_ode_3333_output/hy15_ode_3333_%j.out
+#SBATCH --error=hy15_ode_3333_output/hy15_ode_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_ode_init_1333_new"
+  --wandb_run_name "hy15_ode_init_1333_new"
+  --warp_denoising_step
+  --dmd_denoising_steps '1000,760,520,280,0'
+  --log_visualization
+  --visualization_steps 100
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 2
+  --num_latent_t 31
+  --num_height 480
+  --num_width 848
+  --num_frames 121
+  --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=(
+  --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/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

fastvideo-kernel/CMakeLists.txt

@@ -10,6 +10,8 @@ if(GPU_BACKEND STREQUAL "ROCM")
     enable_language(HIP)
 else()
     enable_language(CUDA)
+    # Ensure CUDA toolkit targets (CUDA::cudart, CUDA::cuda_driver, etc.) are available.
+    find_package(CUDAToolkit REQUIRED)
 endif()
 
 # Import common utils if needed, but we keep it simple for now
@@ -153,6 +155,30 @@ if(BUILD_CXX_KERNELS)
         $<$<COMPILE_LANGUAGE:CUDA>:${CUDA_FLAGS}>
     )
 
+    # Link against Torch libraries to avoid undefined symbols at import time
+    # (e.g., torch::autograd vtables) when loading the extension module.
+    target_link_libraries(fastvideo_kernel_ops PRIVATE ${TORCH_LIBRARIES})
+
+    # Also link against libtorch_python to satisfy Python-binding symbols
+    # (e.g., torch::PyWarningHandler) required by torch/extension.h.
+    execute_process(
+        COMMAND "${Python_EXECUTABLE}" -c "import torch; from pathlib import Path; p=Path(torch.__file__).parent/'lib'; m=sorted(p.glob('libtorch_python*')); print(str(m[0]) if m else '')"
+        OUTPUT_VARIABLE TORCH_PYTHON_LIBRARY_PATH
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        ERROR_QUIET
+    )
+    if(TORCH_PYTHON_LIBRARY_PATH)
+        message(STATUS "TORCH_PYTHON_LIBRARY_PATH: ${TORCH_PYTHON_LIBRARY_PATH}")
+        target_link_libraries(fastvideo_kernel_ops PRIVATE "${TORCH_PYTHON_LIBRARY_PATH}")
+    else()
+        message(WARNING "Could not locate libtorch_python; fastvideo_kernel_ops may fail to import.")
+    endif()
+
+    # Link CUDA runtime + driver explicitly (fixes missing symbols like cuGetErrorString at import time)
+    if(NOT GPU_BACKEND STREQUAL "ROCM")
+        target_link_libraries(fastvideo_kernel_ops PRIVATE CUDA::cudart CUDA::cuda_driver)
+    endif()
+
     # We install it to fastvideo_kernel/_C so we can load it to register the ops
     install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C)
 endif()

fastvideo-kernel/README.md

@@ -34,12 +34,23 @@ from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_att
 out = sliding_tile_attention(q, k, v, window_sizes, text_len)
 
 # Example: Video Sparse Attention (with Triton fallback)
-out = video_sparse_attn(q, k, v, block_sizes, topk=5)
+out = video_sparse_attn(q, k, v, block_sizes, block_sizes, topk=5)
 
 # Example: VMoBA
 out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
 ```
 
+## Benchmark
+
+### VSA (block-sparse) TFLOPs
+
+After building/installing `fastvideo-kernel`, run:
+
+```bash
+cd fastvideo-kernel
+python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
+```
+
 ### TurboDiffusion Kernels
 
 This package also includes kernels from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion), including INT8 GEMM, Quantization, RMSNorm and LayerNorm.

fastvideo-kernel/benchmarks/bench_vsa.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""
+Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs.
+
+This script benchmarks the autograd-enabled wrapper:
+  - fastvideo_kernel.block_sparse_attn.block_sparse_attn
+
+So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import random
+from typing import Tuple, Callable
+
+import numpy as np
+import torch
+
+try:
+    from triton.testing import do_bench
+except Exception as e:  # pragma: no cover
+    raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
+
+
+BLOCK_M = 64
+BLOCK_N = 64
+
+
+def set_seed(seed: int = 42) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def parse_arguments() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Benchmark FastVideo VSA block-sparse attention")
+    p.add_argument("--batch_size", type=int, default=1)
+    p.add_argument("--num_heads", type=int, default=12)
+    p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
+    p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)")
+    p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)")
+    p.add_argument("--kv_seq_lens", type=int, nargs="+", default=None, help="KV sequence lengths (defaults to q_seq_len)")
+    p.add_argument("--warmup", type=int, default=5)
+    p.add_argument("--rep", type=int, default=20)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
+    p.add_argument("--force_triton", action="store_true", help="Force wrapper to use Triton path (if supported by shapes).")
+    return p.parse_args()
+
+
+def create_qkv(batch: int, heads: int, q_len: int, kv_len: int, d: int, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    q = torch.randn(batch, heads, q_len, d, dtype=dtype, device="cuda")
+    k = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda")
+    v = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda")
+    return q, k, v
+
+
+def make_block_map(bs: int, h: int, num_q_blocks: int, num_kv_blocks: int, topk: int) -> torch.Tensor:
+    # block_map: [bs, h, num_q_blocks, num_kv_blocks] bool
+    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device="cuda")
+    topk = min(max(1, topk), num_kv_blocks)
+    idx = torch.topk(scores, topk, dim=-1).indices
+    block_map = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device="cuda")
+    block_map.scatter_(-1, idx, True)
+    return block_map
+
+
+def flops_sparse_attention(bs: int, h: int, d: int, q_len: int, topk_blocks: int, block_n: int) -> float:
+    # Approx: QK^T + PV, each is ~2*bs*h*q_len*(topk_blocks*block_n)*d
+    return 4.0 * bs * h * d * q_len * (topk_blocks * block_n)
+
+
+def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float:
+    return do_bench(fn, warmup=warmup, rep=rep, quantiles=None)
+
+
+def main() -> None:
+    args = parse_arguments()
+    set_seed(args.seed)
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+
+    if args.force_triton:
+        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
+
+    from fastvideo_kernel.block_sparse_attn import block_sparse_attn
+
+    bs, h, d = args.batch_size, args.num_heads, args.head_dim
+    kv_seq_lens = args.kv_seq_lens
+    if kv_seq_lens is None:
+        kv_seq_lens = args.q_seq_lens
+    if len(kv_seq_lens) != len(args.q_seq_lens):
+        raise ValueError("kv_seq_lens must have the same number of entries as q_seq_lens (or be omitted).")
+
+    print("VSA Block-Sparse Attention Benchmark (WRAPPER)")
+    print(f"device: {torch.cuda.get_device_name(0)}")
+    print(f"batch={bs}, heads={h}, head_dim={d}, dtype={args.dtype}")
+    print(f"BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}")
+    print("NOTE: timings include wrapper overhead (map->index + dispatch).")
+    if args.force_triton:
+        print("dispatch: forced Triton (FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1)")
+    else:
+        print("dispatch: SM90 if available, else Triton")
+
+    for q_len, kv_len in zip(args.q_seq_lens, kv_seq_lens):
+        if q_len % BLOCK_M != 0 or kv_len % BLOCK_N != 0:
+            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by 64")
+            continue
+
+        num_q_blocks = q_len // BLOCK_M
+        num_kv_blocks = kv_len // BLOCK_N
+        topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10)
+        topk = min(topk, num_kv_blocks)
+
+        print("\n" + "=" * 80)
+        print(f"q_len={q_len}, kv_len={kv_len}, num_q_blocks={num_q_blocks}, num_kv_blocks={num_kv_blocks}, topk={topk}")
+
+        q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype)
+        block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk)
+
+        # Variable block sizes: default full blocks (64 tokens per KV block)
+        variable_block_sizes = torch.full((num_kv_blocks,), BLOCK_N, dtype=torch.int32, device="cuda")
+
+        def _fwd():
+            return block_sparse_attn(q, k, v, block_map, variable_block_sizes)
+
+        fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep)
+
+        # Backward benchmark (wrapper autograd). We build the graph once, then repeatedly run backward
+        # on the retained graph so bwd timing excludes the forward compute.
+        q_ = q.detach().requires_grad_(True)
+        k_ = k.detach().requires_grad_(True)
+        v_ = v.detach().requires_grad_(True)
+        o_, _aux_ = block_sparse_attn(q_, k_, v_, block_map, variable_block_sizes)
+        og = torch.randn_like(o_)
+        loss = (o_ * og).sum()
+
+        for _ in range(max(1, args.warmup // 2)):
+            torch.autograd.grad(loss, (q_, k_, v_), retain_graph=True)
+        torch.cuda.synchronize()
+
+        bwd_ms = bench_ms(
+            lambda: torch.autograd.grad(loss, (q_, k_, v_), retain_graph=True),
+            warmup=0,
+            rep=max(5, args.rep // 2),
+        )
+
+        flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N)
+        fwd_tflops = flops / fwd_ms * 1e-12 * 1e3
+        # Rough backward multiplier (attention backward typically ~2-3x forward)
+        bwd_tflops = (2.5 * flops) / bwd_ms * 1e-12 * 1e3
+
+        print(f"fwd(wrapper): {fwd_ms:.3f} ms  | {fwd_tflops:.2f} TFLOPs (approx)")
+        print(f"bwd(wrapper): {bwd_ms:.3f} ms  | {bwd_tflops:.2f} TFLOPs (approx)")
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark.")
+    main()
+
+

fastvideo-kernel/csrc/attention/block_sparse_h100.cu

@@ -639,7 +639,8 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
 
     // store kq and vq
 
-    // ensuring all writes are finished
+    // ! the following two line seems unnecessary.
+    // tma::store_async_wait(); // ensure qg is finished
     __syncthreads();
 
     warpgroup::store(kg_smem[0], kg_reg);
@@ -660,145 +661,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
     tma::store_async_wait(); 
 }
 
-
-template<int D>
-void block_sparse_attention_forward_impl(
-    bf16* d_q, bf16* d_k, bf16* d_v, float* d_l, bf16* d_o,
-    int batch, int qo_heads, int kv_heads, int seq_len, int hr,
-    int max_kv_blocks_per_q, 
-    int32_t* q2k_block_sparse_index_ptr,
-    int32_t* q2k_block_sparse_num_ptr,
-    int32_t* block_size_ptr,
-    cudaStream_t stream
-) {
-    using K = fwd_attend_ker_tile_dims<D>;
-    using q_tile    =         st_bf<K::qo_height, K::tile_width>;
-    using k_tile    =         st_bf<K::kv_height, K::tile_width>;
-    using v_tile    =         st_bf<K::kv_height, K::tile_width>;
-    using l_col_vec = col_vec<st_fl<K::qo_height, K::tile_width>>;
-    using o_tile    =         st_bf<K::qo_height, K::tile_width>;
-
-    using q_global = gl<bf16,  -1, -1, -1, -1, q_tile>;
-    using k_global = gl<bf16,  -1, -1, -1, -1, k_tile>;
-    using v_global = gl<bf16,  -1, -1, -1, -1, v_tile>;
-    using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
-    using o_global = gl<bf16,  -1, -1, -1, -1, o_tile>;
-
-    using globals      = fwd_globals<D>;
-
-    q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(seq_len)};
-    o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-
-    globals g{
-        qg_arg, kg_arg, vg_arg, lg_arg, og_arg, 
-        static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), 
-        q2k_block_sparse_index_ptr, q2k_block_sparse_num_ptr, block_size_ptr
-    };
-
-    // Shared memory size for the kernel
-    // 54000 bytes is calibrated for H100 shared memory constraints for these tile sizes
-    constexpr int mem_size = 54000;
-
-    dim3 grid(seq_len/(64), qo_heads, batch);
-
-    cudaFuncSetAttribute(
-        fwd_attend_ker<D>,
-        cudaFuncAttributeMaxDynamicSharedMemorySize,
-        mem_size
-    );
-
-    fwd_attend_ker<D><<<grid, (128), mem_size, stream>>>(g);
-}
-
-template<int D>
-void block_sparse_attention_backward_impl(
-    bf16* d_q, bf16* d_k, bf16* d_v, bf16* d_o, bf16* d_og, float* d_l, float* d_d, float* d_qg, float* d_kg, float* d_vg,
-    int batch, int qo_heads, int kv_heads, int seq_len, int hr, int max_q_blocks_per_kv,
-    int32_t* k2q_block_sparse_index_ptr,
-    int32_t* k2q_block_sparse_num_ptr,
-    int32_t* block_size_ptr,
-    cudaStream_t stream
-) {
-    using G = bwd_attend_ker_tile_dims<D>;
-    using og_tile = st_bf<4*16, D>;
-    using o_tile  = st_bf<4*16, D>;
-    using d_tile  = col_vec<st_fl<4*16, D>>;
-
-    using og_global = gl<bf16,  -1, -1, -1, -1, og_tile>;
-    using o_global  = gl<bf16,  -1, -1, -1, -1, o_tile>;
-    using d_global  = gl<float, -1, -1, -1, -1, d_tile>;
-
-    using prep_globals = bwd_prep_globals<D>;
-
-    constexpr int mem_size_prep = kittens::MAX_SHARED_MEMORY;
-    int threads_prep = PREP_NUM_WARPS * kittens::WARP_THREADS;
-    dim3 grid_bwd_prep(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
-
-    og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
-
-    prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
-
-    cudaFuncSetAttribute(
-        bwd_attend_prep_ker<D>,
-        cudaFuncAttributeMaxDynamicSharedMemorySize,
-        mem_size_prep
-    );
-    bwd_attend_prep_ker<D><<<grid_bwd_prep, threads_prep, mem_size_prep, stream>>>(bwd_g);
-
-    using bwd_q_tile    =         st_bf<G::tile_h_qo, G::tile_width>;
-    using bwd_k_tile    =         st_bf<G::tile_h,    G::tile_width>;
-    using bwd_v_tile    =         st_bf<G::tile_h,    G::tile_width>;
-    using bwd_og_tile   =         st_bf<G::tile_h_qo, G::tile_width>;
-    using bwd_qg_tile   =         st_fl<G::tile_h_qo, G::tile_width>;
-    using bwd_kg_tile   =         st_fl<G::tile_h,    G::tile_width>;
-    using bwd_vg_tile   =         st_fl<G::tile_h,    G::tile_width>;
-    using bwd_l_tile    = row_vec<st_fl<G::tile_h_qo, G::tile_h>>;
-    using bwd_d_tile    = row_vec<st_fl<G::tile_h_qo, G::tile_h>>;
-
-    using bwd_q_global  = gl<bf16,  -1, -1, -1, -1, bwd_q_tile>;
-    using bwd_k_global  = gl<bf16,  -1, -1, -1, -1, bwd_k_tile>;
-    using bwd_v_global  = gl<bf16,  -1, -1, -1, -1, bwd_v_tile>;
-    using bwd_og_global = gl<bf16,  -1, -1, -1, -1, bwd_og_tile>;
-    using bwd_qg_global = gl<float, -1, -1, -1, -1, bwd_qg_tile>;
-    using bwd_kg_global = gl<float, -1, -1, -1, -1, bwd_kg_tile>;
-    using bwd_vg_global = gl<float, -1, -1, -1, -1, bwd_vg_tile>;
-    using bwd_l_global  = gl<float, -1, -1, -1, -1, bwd_l_tile>;
-    using bwd_d_global  = gl<float, -1, -1, -1, -1, bwd_d_tile>;
-
-    using bwd_global_args = bwd_globals<D>;
-
-    bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
-    bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
-    bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
-
-    bwd_global_args bwd_global{bwd_q_arg, bwd_k_arg, bwd_v_arg, bwd_og_arg, bwd_qg_arg, bwd_kg_arg, bwd_vg_arg, bwd_l_arg, bwd_d_arg, 
-                        static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_q_blocks_per_kv),
-                        k2q_block_sparse_index_ptr, k2q_block_sparse_num_ptr, block_size_ptr};
-
-    dim3 grid_bwd_main(seq_len/64, qo_heads, batch);
-    int threads_main = 128;
-    // Calibrated shared memory sizes for different head dimensions
-    int bwd_mem_size = (D == 64) ? 72000 : 113000;
-
-    cudaFuncSetAttribute(
-        bwd_attend_ker<D>,
-        cudaFuncAttributeMaxDynamicSharedMemorySize,
-        bwd_mem_size
-    );
-    bwd_attend_ker<D><<<grid_bwd_main, threads_main, bwd_mem_size, stream>>>(bwd_global);
-}
-
 #include "pyutils/torch_helpers.cuh"
 #include <ATen/cuda/CUDAContext.h>
 #include <iostream>
@@ -810,23 +672,32 @@ block_sparse_attention_forward(
     torch::Tensor v, 
     torch::Tensor q2k_block_sparse_index, 
     torch::Tensor q2k_block_sparse_num, 
-    torch::Tensor block_size
+    torch::Tensor kv_block_size
 )
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
     CHECK_INPUT(v);
 
+    // q shape: (batch, qo_heads, q_seq_len, head_dim)
+    // k shape: (batch, kv_heads, kv_seq_len, head_dim)
+    // v shape: (batch, kv_heads, kv_seq_len, head_dim)
+    // q2k_block_sparse_index shape: (batch, qo_heads, num_q_blocks, max_kv_blocks_per_q)
+    // q2k_block_sparse_num shape: (batch, qo_heads, num_q_blocks)
+    // kv_block_size shape: (num_kv_blocks) This does not need other dimensions because across all batch/heads the padding is the same.
+    
     auto batch    = q.size(0);
-    auto seq_len  = q.size(2); 
+    auto q_seq_len  = q.size(2);
+    auto kv_seq_len = k.size(2);
     auto head_dim = q.size(3); 
     auto qo_heads = q.size(1);
     auto kv_heads = k.size(1);
     auto max_kv_blocks_per_q = q2k_block_sparse_index.size(3);
-    auto num_q_blocks = block_size.size(0);
+    auto num_q_blocks = q2k_block_sparse_index.size(2);
+    auto num_kv_blocks = kv_block_size.size(0);
     TORCH_CHECK(batch==1, "Batch size dim will be removed in the future, please set batch to 1");
-    TORCH_CHECK(num_q_blocks * 64 == seq_len, "This kernel supports variable block size, but it assumes the input sequence is properly padded.");
-    TORCH_CHECK(num_q_blocks == q2k_block_sparse_index.size(2), "Number of Q blocks does not match between q2k_block_sparse_index and block_size");
+    TORCH_CHECK(num_q_blocks * BLOCK_M == q_seq_len, "This kernel supports variable q block size, but it assumes the input sequence is properly padded.");
+    TORCH_CHECK(num_kv_blocks * BLOCK_M == kv_seq_len, "This kernel supports variable kv block size, but it assumes the input sequence is properly padded.");
     // check to see that these dimensions match for all inputs
     TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs");
@@ -834,11 +705,9 @@ block_sparse_attention_forward(
     TORCH_CHECK(q2k_block_sparse_index.size(0) == batch, "q2k_block_sparse_index batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(q2k_block_sparse_num.size(0) == batch, "q2k_block_sparse_num batch dimension - idx 0 - must match for all inputs");
 
-    TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(q2k_block_sparse_index.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_index idx 2 - must match seq_len / BLOCK_M");
-    TORCH_CHECK(q2k_block_sparse_num.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_num idx 2 - must match seq_len / BLOCK_M");
+    TORCH_CHECK(v.size(2) == kv_seq_len, "V sequence length dimension - idx 2 - must match K inputs");
+    TORCH_CHECK(q2k_block_sparse_num.size(2) == num_q_blocks, "q2k_block_sparse_num idx 2 - must match num_q_blocks");
+
 
     TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs");
     TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs");
@@ -864,12 +733,12 @@ block_sparse_attention_forward(
     // for the returned outputs
     torch::Tensor o     = torch::empty({static_cast<const uint>(batch), 
                                         static_cast<const uint>(qo_heads), 
-                                        static_cast<const uint>(seq_len), 
+                                        static_cast<const uint>(q_seq_len), 
                                         static_cast<const uint>(head_dim)}, v.options());
     
     torch::Tensor l_vec = torch::empty({static_cast<const uint>(batch), 
                                         static_cast<const uint>(qo_heads), 
-                                        static_cast<const uint>(seq_len), 
+                                        static_cast<const uint>(q_seq_len), 
                                         static_cast<const uint>(1)}, 
                                         torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous));
         
@@ -880,32 +749,110 @@ block_sparse_attention_forward(
     float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
     float* d_l   = reinterpret_cast<float*>(l_ptr);
 
+    //cudadevicesynchronize();
     const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
     const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
 
-    // Temporated implementation to avoid code duplication between head_dim=64 and 128
     if (head_dim == 64) {
-        block_sparse_attention_forward_impl<64>(
-            d_q, d_k, d_v, d_l, d_o,
-            batch, qo_heads, kv_heads, seq_len, hr,
-            max_kv_blocks_per_q,
-            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()),
+        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
+        using k_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>;
+        using v_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>;
+        using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>>;
+        using o_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
+
+        using q_global = gl<bf16,  -1, -1, -1, -1, q_tile>;
+        using k_global = gl<bf16,  -1, -1, -1, -1, k_tile>;
+        using v_global = gl<bf16,  -1, -1, -1, -1, v_tile>;
+        using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
+        using o_global = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+        using globals      = fwd_globals<64>;
+
+        q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
+        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
+        l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(q_seq_len)};
+        o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+
+        globals g{
+            qg_arg, 
+            kg_arg, 
+            vg_arg, 
+            lg_arg, 
+            og_arg, 
+            static_cast<int>(q_seq_len), 
+            static_cast<int>(hr), 
+            static_cast<int>(max_kv_blocks_per_q), 
+            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), 
             reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
-            reinterpret_cast<int32_t*>(block_size.data_ptr()),
-            stream
+            reinterpret_cast<int32_t*>(kv_block_size.data_ptr())
+        };
+
+        constexpr int mem_size = 54000;
+
+        dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch);
+
+        cudaFuncSetAttribute(
+            fwd_attend_ker<64>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize,
+            mem_size
         );
-    } else if (head_dim == 128) {
-        block_sparse_attention_forward_impl<128>(
-            d_q, d_k, d_v, d_l, d_o,
-            batch, qo_heads, kv_heads, seq_len, hr,
-            max_kv_blocks_per_q,
-            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()),
+
+        fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
+
+        CHECK_CUDA_ERROR(cudaGetLastError());
+        // cudaStreamSynchronize(stream);
+    }
+
+    if (head_dim == 128) {
+        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using k_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using v_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+        using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>;
+        using o_tile    =         st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
+
+        using q_global = gl<bf16,  -1, -1, -1, -1, q_tile>;
+        using k_global = gl<bf16,  -1, -1, -1, -1, k_tile>;
+        using v_global = gl<bf16,  -1, -1, -1, -1, v_tile>;
+        using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
+        using o_global = gl<bf16,  -1, -1, -1, -1, o_tile>;
+
+        using globals      = fwd_globals<128>;
+
+        q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+        k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
+        v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
+        l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,   static_cast<unsigned int>(q_seq_len)};
+        o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+
+        globals g{
+            qg_arg, 
+            kg_arg, 
+            vg_arg, 
+            lg_arg, 
+            og_arg, 
+            static_cast<int>(q_seq_len), 
+            static_cast<int>(hr), 
+            static_cast<int>(max_kv_blocks_per_q), 
+            reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), 
             reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
-            reinterpret_cast<int32_t*>(block_size.data_ptr()),
-            stream
+            reinterpret_cast<int32_t*>(kv_block_size.data_ptr())
+        };
+
+        constexpr int mem_size = 54000;
+
+        dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch);
+
+        cudaFuncSetAttribute(
+            fwd_attend_ker<128>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize,
+            mem_size
         );
-    } else {
-        TORCH_CHECK(false, "Unsupported head_dim: ", head_dim, ". Only 64 and 128 are supported.");
+
+        fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);
+
+        CHECK_CUDA_ERROR(cudaGetLastError());
+        // cudaStreamSynchronize(stream);
     }
 
     return {o, l_vec};
@@ -921,7 +868,7 @@ block_sparse_attention_backward(torch::Tensor q,
                    torch::Tensor og,
                    torch::Tensor k2q_block_sparse_index,
                    torch::Tensor k2q_block_sparse_num,
-                   torch::Tensor block_size)
+                   torch::Tensor kv_block_size)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -930,11 +877,23 @@ block_sparse_attention_backward(torch::Tensor q,
     CHECK_INPUT(o);
     CHECK_INPUT(og);
 
+    // q: [batch, qo_heads, q_seq_len, head_dim]
+    // k: [batch, kv_heads, kv_seq_len, head_dim]
+    // v: [batch, kv_heads, kv_seq_len, head_dim]
+    // o: [batch, qo_heads, q_seq_len, head_dim]
+    // l_vec: [batch, qo_heads, q_seq_len, 1]
+    // og: [batch, qo_heads, q_seq_len, head_dim]
+    // k2q_block_sparse_index: [batch, kv_heads, num_kv_blocks, max_num_q_blocks]
+    // k2q_block_sparse_num: [batch, kv_heads, num_kv_blocks]
+    // kv_block_size: [num_kv_blocks]
+
     auto batch    = q.size(0);
-    auto seq_len  = q.size(2);
+    auto q_seq_len  = q.size(2);
+    auto kv_seq_len = k.size(2);
     auto head_dim = q.size(3);
     auto max_q_blocks_per_kv = k2q_block_sparse_index.size(3);
-    TORCH_CHECK(k2q_block_sparse_index.size(2) == block_size.size(0), "k2q_block_sparse_index.size(2) must match block_size.size(0)");
+    auto num_kv_blocks = kv_block_size.size(0);
+    TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index.size(2) must match num_kv_blocks (kv_block_size.size(0))");
     // check to see that these dimensions match for all inputs
     TORCH_CHECK(q.size(0)     == batch, "Q  batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(k.size(0)     == batch, "K  batch dimension - idx 0 - must match for all inputs");
@@ -945,23 +904,18 @@ block_sparse_attention_backward(torch::Tensor q,
     TORCH_CHECK(k2q_block_sparse_index.size(0) == batch, "k2q_block_sparse_index batch dimension - idx 0 - must match for all inputs");
     TORCH_CHECK(k2q_block_sparse_num.size(0) == batch, "k2q_block_sparse_num batch dimension - idx 0 - must match for all inputs");
 
-
-    TORCH_CHECK(q.size(2)     == seq_len, "Q  sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(k.size(2)     == seq_len, "K  sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(v.size(2)     == seq_len, "V  sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(l_vec.size(2) == seq_len, "L  sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(o.size(2)     == seq_len, "O  sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(og.size(2)    == seq_len, "OG sequence length dimension - idx 2 - must match for all inputs");
-    TORCH_CHECK(k2q_block_sparse_index.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_index idx 2 - must match seq_len / BLOCK_N");
-    TORCH_CHECK(k2q_block_sparse_num.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_num idx 2 - must match seq_len / BLOCK_N");
+    TORCH_CHECK(v.size(2)     == kv_seq_len, "V  sequence length dimension - idx 2 - must match K sequence length");
+    TORCH_CHECK(l_vec.size(2) == q_seq_len, "L  sequence length dimension - idx 2 - must match Q sequence length");
+    TORCH_CHECK(o.size(2)     == q_seq_len, "O  sequence length dimension - idx 2 - must match Q sequence length");
+    TORCH_CHECK(og.size(2)    == q_seq_len, "OG sequence length dimension - idx 2 - must match Q sequence length");
+    TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index idx 2 - must match num_kv_blocks (kv_block_size.size(0))");
+    TORCH_CHECK(k2q_block_sparse_num.size(2) == num_kv_blocks, "k2q_block_sparse_num idx 2 - must match num_kv_blocks (kv_block_size.size(0))");
 
     TORCH_CHECK(q.size(3)  == head_dim, "Q  head dimension - idx 3 - must match for all non-vector inputs");
     TORCH_CHECK(k.size(3)  == head_dim, "K  head dimension - idx 3 - must match for all non-vector inputs");
     TORCH_CHECK(v.size(3)  == head_dim, "V  head dimension - idx 3 - must match for all non-vector inputs");
     TORCH_CHECK(o.size(3)  == head_dim, "O  head dimension - idx 3 - must match for all non-vector inputs");
     TORCH_CHECK(og.size(3) == head_dim, "OG head dimension - idx 3 - must match for all non-vector inputs");
-    
-
 
     auto qo_heads = q.size(1);
     auto kv_heads = k.size(1);
@@ -988,20 +942,20 @@ block_sparse_attention_backward(torch::Tensor q,
 
     torch::Tensor qg = torch::zeros({static_cast<const uint>(batch), 
                                      static_cast<const uint>(qo_heads), 
-                                     static_cast<const uint>(seq_len), 
+                                     static_cast<const uint>(q_seq_len), 
                                      static_cast<const uint>(head_dim)},   l_vec.options());
     torch::Tensor kg = torch::zeros({static_cast<const uint>(batch), 
                                      static_cast<const uint>(kv_heads), 
-                                     static_cast<const uint>(seq_len), 
+                                     static_cast<const uint>(kv_seq_len), 
                                      static_cast<const uint>(head_dim)},   l_vec.options());
     torch::Tensor vg = torch::zeros({static_cast<const uint>(batch), 
                                      static_cast<const uint>(kv_heads), 
-                                     static_cast<const uint>(seq_len), 
+                                     static_cast<const uint>(kv_seq_len), 
                                      static_cast<const uint>(head_dim)},   l_vec.options());
     
     torch::Tensor d_vec = torch::empty({static_cast<const uint>(batch), 
                                         static_cast<const uint>(qo_heads), 
-                                        static_cast<const uint>(seq_len), 
+                                        static_cast<const uint>(q_seq_len), 
                                         static_cast<const uint>(1)},       l_vec.options());
 
     float*         qg_ptr = qg.data_ptr<float>();
@@ -1030,7 +984,7 @@ block_sparse_attention_backward(torch::Tensor q,
     //  cudaStreamSynchronize(stream);
 
     // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
-    dim3 grid_bwd(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+    dim3 grid_bwd(q_seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
 
     if (head_dim == 64)  {
         using og_tile = st_bf<4*16, 64>;
@@ -1043,9 +997,9 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_prep_globals = bwd_prep_globals<64>;
 
-        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
-        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
-        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
 
         bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
 
@@ -1082,15 +1036,15 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_global_args = bwd_globals<64>;
 
-        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
-        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
-        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
+        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
+        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
+        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
+        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
+        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
+        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
 
         bwd_global_args bwd_global{bwd_q_arg, 
                         bwd_k_arg, 
@@ -1101,14 +1055,14 @@ block_sparse_attention_backward(torch::Tensor q,
                         bwd_vg_arg, 
                         bwd_l_arg, 
                         bwd_d_arg, 
-                        static_cast<int>(seq_len), 
+                        static_cast<int>(kv_seq_len), // N is not used in the kernel
                         static_cast<int>(hr),
                         static_cast<int>(max_q_blocks_per_kv),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
-                        reinterpret_cast<int32_t*>(block_size.data_ptr())};
+                        reinterpret_cast<int32_t*>(kv_block_size.data_ptr())};
 
-        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
+        dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch);
         threads = 128;
 
         //cudadevicesynchronize();
@@ -1147,9 +1101,9 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_prep_globals = bwd_prep_globals<128>;
 
-        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
-        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
-        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+        og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+        o_global  prep_o_arg {d_o,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+        d_global  prep_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
 
         bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
 
@@ -1186,15 +1140,15 @@ block_sparse_attention_backward(torch::Tensor q,
 
         using bwd_global_args = bwd_globals<128>;
 
-        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
-        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
-        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(seq_len)};
+        bwd_q_global  bwd_q_arg {d_q,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+        bwd_k_global  bwd_k_arg {d_k,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
+        bwd_v_global  bwd_v_arg {d_v,  static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
+        bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+        bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
+        bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
+        bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
+        bwd_l_global  bwd_l_arg {d_l,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
+        bwd_d_global  bwd_d_arg {d_d,  static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U,  static_cast<unsigned int>(q_seq_len)};
 
         bwd_global_args bwd_global{bwd_q_arg, 
                         bwd_k_arg, 
@@ -1205,14 +1159,14 @@ block_sparse_attention_backward(torch::Tensor q,
                         bwd_vg_arg, 
                         bwd_l_arg, 
                         bwd_d_arg, 
-                        static_cast<int>(seq_len), 
+                        static_cast<int>(kv_seq_len), // N is not used in the kernel
                         static_cast<int>(hr),
                         static_cast<int>(max_q_blocks_per_kv),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
                         reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
-                        reinterpret_cast<int32_t*>(block_size.data_ptr())};
+                        reinterpret_cast<int32_t*>(kv_block_size.data_ptr())};
 
-        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
+        dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch);
         threads = 128;
 
         //cudadevicesynchronize();
@@ -1233,4 +1187,4 @@ block_sparse_attention_backward(torch::Tensor q,
 
     return {qg, kg, vg};
     //cudadevicesynchronize();
-}
+}

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu

@@ -4,6 +4,7 @@
 #include <torch/all.h>
 #include <torch/python.h>
 #include <cutlass/cutlass.h>
+#include <cutlass/numeric_types.h>
 #include "common/common.hpp"
 #include "norm/layernorm.hpp"
 
@@ -14,10 +15,6 @@ auto layer_norm(
   std::optional<at::Tensor const> const B,
   std::optional<at::Tensor> Output
 ) {
-  using ElementIn = float;
-  using ElementOut = float;
-  using ElementWeight = float;
-
   int64_t const m = Input.size(0);
   int64_t const n = Input.size(1);
   torch::Device const input_device = Input.device();
@@ -26,31 +23,70 @@ auto layer_norm(
     Output.emplace(
       torch::empty(
         {m, n},
-        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
+        torch::TensorOptions().device(input_device).dtype(Input.scalar_type())
       )
     );
   }
 
+  TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(),
+              "Output dtype must match Input dtype. Got Output=",
+              Output.value().scalar_type(), ", Input=", Input.scalar_type());
+  if (W.has_value()) {
+    TORCH_CHECK(W.value().scalar_type() == Input.scalar_type(),
+                "W dtype must match Input dtype. Got W=",
+                W.value().scalar_type(), ", Input=", Input.scalar_type());
+  }
+  if (B.has_value()) {
+    TORCH_CHECK(B.value().scalar_type() == Input.scalar_type(),
+                "B dtype must match Input dtype. Got B=",
+                B.value().scalar_type(), ", Input=", Input.scalar_type());
+  }
 
   void *Iptr = Input.data_ptr();
   void *Wptr = W.has_value() ? W.value().data_ptr() : nullptr;
   void *Bptr = B.has_value() ? B.value().data_ptr() : nullptr;
   void *Optr = Output.value().data_ptr();
 
-  BOOL_SWITCH(B.has_value(), BIAS, [&]{
-    BOOL_SWITCH(W.has_value(), AFFINE, [&]{
-    CONFIG_SWITCH(n, [&]{
-    layernorm<
-    ElementIn, ElementOut, ElementWeight,
-    AFFINE, BIAS,
-    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA> (
-      Iptr, Wptr, Bptr,
-      Optr, eps, m, n,
-      at::cuda::getCurrentCUDAStream().stream()
-        );
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  if (Input.scalar_type() == at::kHalf) {
+    using ElementIn = cutlass::half_t;
+    using ElementOut = cutlass::half_t;
+    using ElementWeight = cutlass::half_t;
+    BOOL_SWITCH(B.has_value(), BIAS, [&]{
+      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
+        CONFIG_SWITCH(n, [&]{
+          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
+              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
+        });
       });
     });
-  });
+  } else if (Input.scalar_type() == at::kBFloat16) {
+    using ElementIn = cutlass::bfloat16_t;
+    using ElementOut = cutlass::bfloat16_t;
+    using ElementWeight = cutlass::bfloat16_t;
+    BOOL_SWITCH(B.has_value(), BIAS, [&]{
+      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
+        CONFIG_SWITCH(n, [&]{
+          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
+              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
+        });
+      });
+    });
+  } else if (Input.scalar_type() == at::kFloat) {
+    using ElementIn = float;
+    using ElementOut = float;
+    using ElementWeight = float;
+    BOOL_SWITCH(B.has_value(), BIAS, [&]{
+      BOOL_SWITCH(W.has_value(), AFFINE, [&]{
+        CONFIG_SWITCH(n, [&]{
+          layernorm<ElementIn, ElementOut, ElementWeight, AFFINE, BIAS, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
+              Iptr, Wptr, Bptr, Optr, eps, m, n, stream);
+        });
+      });
+    });
+  } else {
+    TORCH_CHECK(false, "Unsupported dtype for layer_norm_cuda: ", Input.scalar_type());
+  }
   
     
 

fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp

@@ -68,9 +68,22 @@ public:
     // mean reduction
     float u = _reduce_sum(x, shared_data) / params.n;
 
+    // IMPORTANT:
+    // Loader pads out-of-range lanes with 0. That is OK for the sum, but after
+    // subtracting mean, those padded lanes become -u and would incorrectly
+    // contribute to the variance. Mask them back to 0 before variance reduction.
+    // We launch exactly NumThrPerCta threads for a 1xMaxHiddenSize tile,
+    // so each thread is responsible for a contiguous chunk in N.
+    int thr_n_offset = tidx * NumElementPerThread;
     CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < NumElementPerThread; ++i)
-      x[i] -= u;
+    for (int i = 0; i < NumElementPerThread; ++i) {
+      int idx = thr_n_offset + i;
+      if (idx < params.n) {
+        x[i] -= u;
+      } else {
+        x[i] = 0.f;
+      }
+    }
 
     __syncthreads();
     // var reduction

fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu

@@ -4,6 +4,7 @@
 #include <torch/all.h>
 #include <torch/python.h>
 #include <cutlass/cutlass.h>
+#include <cutlass/numeric_types.h>
 #include <pybind11/pybind11.h>
 
 #include "common/common.hpp"
@@ -16,10 +17,6 @@ auto rms_norm(
   std::optional<at::Tensor>& Output
 ) {
   
-  using ElementIn = float;
-  using ElementOut = float;
-  using ElementWeight = float;
-  
   int64_t const m = Input.size(0);
   int64_t const n = Input.size(1);
   torch::Device const input_device = Input.device();
@@ -28,27 +25,51 @@ auto rms_norm(
     Output.emplace(
       torch::empty(
         {m, n},
-        torch::TensorOptions().device(input_device).dtype(torch::kFloat32)
+        torch::TensorOptions().device(input_device).dtype(Input.scalar_type())
       )
     );
   }
 
+  TORCH_CHECK(Output.value().scalar_type() == Input.scalar_type(),
+              "Output dtype must match Input dtype. Got Output=",
+              Output.value().scalar_type(), ", Input=", Input.scalar_type());
+  if (Weight.has_value()) {
+    TORCH_CHECK(Weight.value().scalar_type() == Input.scalar_type(),
+                "Weight dtype must match Input dtype. Got Weight=",
+                Weight.value().scalar_type(), ", Input=", Input.scalar_type());
+  }
+
   void *Iptr = Input.data_ptr();
   void *Wptr = Weight.has_value() ? Weight.value().data_ptr() : nullptr;
   void *Optr = Output.value().data_ptr();
 
-
-  CONFIG_SWITCH(n, [&]{
-    rmsnorm<
-    ElementIn, ElementOut, ElementWeight,
-    MAX_HIDDEN_SIZE, NUM_THR_PER_CTA
-    > (
-      Iptr, Wptr,
-      Optr, 
-      eps, m, n, 
-      at::cuda::getCurrentCUDAStream().stream()
-    );
-});
+  if (Input.scalar_type() == at::kHalf) {
+    using ElementIn = cutlass::half_t;
+    using ElementOut = cutlass::half_t;
+    using ElementWeight = cutlass::half_t;
+    CONFIG_SWITCH(n, [&]{
+      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
+          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
+    });
+  } else if (Input.scalar_type() == at::kBFloat16) {
+    using ElementIn = cutlass::bfloat16_t;
+    using ElementOut = cutlass::bfloat16_t;
+    using ElementWeight = cutlass::bfloat16_t;
+    CONFIG_SWITCH(n, [&]{
+      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
+          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
+    });
+  } else if (Input.scalar_type() == at::kFloat) {
+    using ElementIn = float;
+    using ElementOut = float;
+    using ElementWeight = float;
+    CONFIG_SWITCH(n, [&]{
+      rmsnorm<ElementIn, ElementOut, ElementWeight, MAX_HIDDEN_SIZE, NUM_THR_PER_CTA>(
+          Iptr, Wptr, Optr, eps, m, n, at::cuda::getCurrentCUDAStream().stream());
+    });
+  } else {
+    TORCH_CHECK(false, "Unsupported dtype for rms_norm_cuda: ", Input.scalar_type());
+  }
   
 
   return Output;

fastvideo-kernel/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "fastvideo-kernel"
-version = "0.2.2"
+version = "0.2.5"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -0,0 +1,295 @@
+from __future__ import annotations
+
+import os
+from typing import Tuple
+
+import torch
+
+
+def _get_sm90_ops():
+    try:
+        from fastvideo_kernel._C import fastvideo_kernel_ops  # type: ignore
+    except Exception:
+        return None, None
+    return (
+        getattr(fastvideo_kernel_ops, "block_sparse_fwd", None),
+        getattr(fastvideo_kernel_ops, "block_sparse_bwd", None),
+    )
+
+
+def _is_sm90() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    major, minor = torch.cuda.get_device_capability(0)
+    return major == 9 and minor == 0
+
+
+def _force_triton() -> bool:
+    # Force Triton even on SM90 and even if the compiled extension is available.
+    # Useful for CI / debugging / parity testing.
+    return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1"
+
+
+def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Preferred map->index conversion used by the wrapper.
+
+    This wrapper **requires** the Triton implementation.
+    If Triton (or the Triton map_to_index module) is not available, it raises.
+    """
+    if block_map.dim() == 3:
+        block_map = block_map.unsqueeze(0)
+    if block_map.dim() != 4:
+        raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), got shape={tuple(block_map.shape)}")
+    if block_map.dtype != torch.bool:
+        block_map = block_map.to(torch.bool)
+
+    if not block_map.is_cuda:
+        raise RuntimeError("block_map must be a CUDA tensor (Triton map_to_index required).")
+
+    try:
+        from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index  # local import
+    except Exception as e:
+        raise ImportError(
+            "Triton map_to_index is required but not available. "
+            "Ensure Triton is installed and fastvideo_kernel.triton_kernels.index is importable."
+        ) from e
+    return triton_map_to_index(block_map)
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_triton",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_triton(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    block_map = block_map.to(torch.bool)
+    q2k_idx, q2k_num = _map_to_index(block_map)
+
+    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import (  # local import
+        triton_block_sparse_attn_forward,
+    )
+
+    o, M = triton_block_sparse_attn_forward(q, k, v, q2k_idx, q2k_num, variable_block_sizes)
+    return o, M
+
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_triton")
+def _block_sparse_attn_triton_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    o = torch.empty_like(q)
+    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+    return o, M
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_backward_triton",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_backward_triton(
+    grad_output: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    M: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    grad_output = grad_output.contiguous()
+    block_map = block_map.to(torch.bool)
+    q2k_idx, q2k_num = _map_to_index(block_map)
+    k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous())
+
+    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import (  # local import
+        triton_block_sparse_attn_backward,
+    )
+
+    dq, dk, dv = triton_block_sparse_attn_backward(
+        grad_output, q, k, v, o, M, q2k_idx, q2k_num, k2q_idx, k2q_num, variable_block_sizes
+    )
+    return dq, dk, dv
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_triton")
+def _block_sparse_attn_backward_triton_fake(
+    grad_output: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    M: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    return dq, dk, dv
+
+
+def _backward_triton(ctx, grad_o, grad_M):
+    q, k, v, o, M, block_map, variable_block_sizes = ctx.saved_tensors
+    dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, block_map, variable_block_sizes)
+    return dq, dk, dv, None, None
+
+
+def _setup_context_triton(ctx, inputs, output):
+    q, k, v, block_map, variable_block_sizes = inputs
+    o, M = output
+    ctx.save_for_backward(q, k, v, o, M, block_map, variable_block_sizes)
+
+
+block_sparse_attn_triton.register_autograd(_backward_triton, setup_context=_setup_context_triton)
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_sm90",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_sm90(
+    q_padded: torch.Tensor,
+    k_padded: torch.Tensor,
+    v_padded: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    block_sparse_fwd, _ = _get_sm90_ops()
+    if block_sparse_fwd is None:
+        raise ImportError("fastvideo_kernel_ops.block_sparse_fwd is not available")
+
+    q_padded = q_padded.contiguous()
+    k_padded = k_padded.contiguous()
+    v_padded = v_padded.contiguous()
+    block_map = block_map.to(torch.bool)
+    q2k_idx, q2k_num = _map_to_index(block_map)
+
+    o_padded, lse_padded = block_sparse_fwd(
+        q_padded, k_padded, v_padded, q2k_idx, q2k_num, variable_block_sizes.int()
+    )
+    return o_padded, lse_padded
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_sm90")
+def _block_sparse_attn_sm90_fake(
+    q_padded: torch.Tensor,
+    k_padded: torch.Tensor,
+    v_padded: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    o = torch.empty_like(q_padded)
+    lse = torch.empty((q_padded.shape[0], q_padded.shape[1], q_padded.shape[2], 1), device=q_padded.device, dtype=torch.float32)
+    return o, lse
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_backward_sm90",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_backward_sm90(
+    grad_output_padded: torch.Tensor,
+    q_padded: torch.Tensor,
+    k_padded: torch.Tensor,
+    v_padded: torch.Tensor,
+    o_padded: torch.Tensor,
+    lse_padded: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    _, block_sparse_bwd = _get_sm90_ops()
+    if block_sparse_bwd is None:
+        raise ImportError("fastvideo_kernel_ops.block_sparse_bwd is not available")
+
+    grad_output_padded = grad_output_padded.contiguous()
+    block_map = block_map.to(torch.bool)
+    k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous())
+
+    dq, dk, dv = block_sparse_bwd(
+        q_padded,
+        k_padded,
+        v_padded,
+        o_padded,
+        lse_padded,
+        grad_output_padded,
+        k2q_idx,
+        k2q_num,
+        variable_block_sizes.int(),
+    )
+    # C++ kernel returns fp32 grads; cast back to match PyTorch convention if needed
+    return dq.to(grad_output_padded.dtype), dk.to(grad_output_padded.dtype), dv.to(grad_output_padded.dtype)
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_sm90")
+def _block_sparse_attn_backward_sm90_fake(
+    grad_output_padded: torch.Tensor,
+    q_padded: torch.Tensor,
+    k_padded: torch.Tensor,
+    v_padded: torch.Tensor,
+    o_padded: torch.Tensor,
+    lse_padded: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    dq = torch.empty_like(q_padded)
+    dk = torch.empty_like(k_padded)
+    dv = torch.empty_like(v_padded)
+    return dq, dk, dv
+
+
+def _backward_sm90(ctx, grad_o, grad_lse):
+    q, k, v, o, lse, block_map, variable_block_sizes = ctx.saved_tensors
+    dq, dk, dv = block_sparse_attn_backward_sm90(
+        grad_o, q, k, v, o, lse, block_map, variable_block_sizes
+    )
+    return dq, dk, dv, None, None
+
+
+def _setup_context_sm90(ctx, inputs, output):
+    q, k, v, block_map, variable_block_sizes = inputs
+    o, lse = output
+    ctx.save_for_backward(q, k, v, o, lse, block_map, variable_block_sizes)
+
+
+block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90)
+
+
+def block_sparse_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Unified block-sparse attention op with autograd support.
+    - On SM90 with compiled extension present: uses fastvideo_kernel_ops.block_sparse_fwd/bwd.
+    - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today).
+    """
+    block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
+    if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
+        return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes)
+    # Triton path: generally assumes q/k/v share the same padded length
+    if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]:
+        raise RuntimeError("Triton fallback requires q/k/v to have the same padded length.")
+    return block_sparse_attn_triton(q, k, v, block_map, variable_block_sizes)
+
+

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -1,5 +1,6 @@
 import math
 import torch
+from .block_sparse_attn import block_sparse_attn
 from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
 from .triton_kernels.st_attn_triton import sliding_tile_attention_triton
 from .triton_kernels.index import map_to_index
@@ -45,14 +46,22 @@ def sliding_tile_attention(
     flag = shape_map[seq_shape]
 
     for head_idx, (t, h, w) in enumerate(window_size):
+        # Per-head slices are not contiguous in the batch dimension when batch>1
+        # (they keep the original head-stride). The TK kernel assumes contiguous
+        # [B, H, S, D] layout, so we materialize a contiguous [B,1,S,D] view.
+        q_h = q[:, head_idx:head_idx + 1].contiguous()
+        k_h = k[:, head_idx:head_idx + 1].contiguous()
+        v_h = v[:, head_idx:head_idx + 1].contiguous()
+        o_h = torch.empty_like(q_h)
         sta_fwd(
-            q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
-            v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
+            q_h, k_h,
+            v_h, o_h,
             t, h, w, text_length, False, has_text, flag
         )
+        output[:, head_idx:head_idx + 1] = o_h
 
     if has_text:
-        sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
+        sta_fwd(q.contiguous(), k.contiguous(), v.contiguous(), output, 3, 3, 3, text_length, True, True, flag)
 
     return output[:, :, :seq_length]
 
@@ -62,6 +71,7 @@ def video_sparse_attn(
     k: torch.Tensor,
     v: torch.Tensor,
     variable_block_sizes: torch.Tensor,
+    q_variable_block_sizes: torch.Tensor,
     topk: int,
     block_size: int | tuple = 64,
     compress_attn_weight: torch.Tensor = None,
@@ -70,14 +80,42 @@ def video_sparse_attn(
         block_size = (block_size, block_size, block_size)
 
     block_elements = block_size[0] * block_size[1] * block_size[2]
-    batch, heads, seq_len, dim = q.shape
+    batch, heads, q_seq_len, dim = q.shape
+    kv_seq_len = k.shape[2]
+    if v.shape[2] != kv_seq_len:
+        raise ValueError(
+            f"Expected k and v to have the same sequence length, got "
+            f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}"
+        )
+    if k.shape[0] != batch or v.shape[0] != batch or k.shape[1] != heads or v.shape[1] != heads:
+        raise ValueError("Expected q/k/v to have the same batch and head dimensions.")
+
+    if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0:
+        raise ValueError(
+            f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, "
+            f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}"
+        )
+    q_num_blocks = q_seq_len // block_elements
+    kv_num_blocks = kv_seq_len // block_elements
+
+    if variable_block_sizes.numel() != kv_num_blocks:
+        raise ValueError(
+            f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, "
+            f"got {variable_block_sizes.numel()}"
+        )
+
+    if q_variable_block_sizes.numel() != q_num_blocks:
+        raise ValueError(
+            f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
+            f"got {q_variable_block_sizes.numel()}"
+        )
 
     # Compression branch
-    q_c = q.view(batch, heads, seq_len // block_elements, block_elements, dim)
-    k_c = k.view(batch, heads, seq_len // block_elements, block_elements, dim)
-    v_c = v.view(batch, heads, seq_len // block_elements, block_elements, dim)
+    q_c = q.view(batch, heads, q_num_blocks, block_elements, dim)
+    k_c = k.view(batch, heads, kv_num_blocks, block_elements, dim)
+    v_c = v.view(batch, heads, kv_num_blocks, block_elements, dim)
 
-    q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
+    q_c = (q_c.float().sum(dim=3) / q_variable_block_sizes.view(1, 1, -1, 1)).to(
         q.dtype)
     k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(
         k.dtype)
@@ -88,9 +126,9 @@ def video_sparse_attn(
     attn = torch.softmax(scores, dim=-1)
     out_c = torch.matmul(attn, v_c)
 
-    out_c = out_c.view(batch, heads, seq_len // block_elements, 1, dim)
+    out_c = out_c.view(batch, heads, q_num_blocks, 1, dim)
     out_c = out_c.repeat(1, 1, 1, block_elements,
-                         1).view(batch, heads, seq_len, dim)
+                         1).view(batch, heads, q_seq_len, dim)
 
     # Sparse branch
     topk_idx = torch.topk(scores, topk, dim=-1).indices
@@ -100,12 +138,17 @@ def video_sparse_attn(
     idx, num = map_to_index(mask)
     
     if block_sparse_fwd is not None:
-        out_s = block_sparse_fwd(
-            q, k, v, idx, num, variable_block_sizes.int()
-        )[0] # block_sparse_fwd returns vector<Tensor>
+        # Use autograd-enabled wrapper so backward works (and still uses SM90 kernel when available)
+        out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
     else:
-        out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num,
-                                                    variable_block_sizes)
+        if q_seq_len != kv_seq_len:
+            raise RuntimeError(
+                "q/k have different lengths, but the compiled CUDA kernel (block_sparse_fwd) "
+                "is not available. The Triton fallback currently requires q and k/v to have "
+                "the same padded length."
+            )
+        # Triton-only forward (kept for environments without the wrapper deps)
+        out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes)
 
     if compress_attn_weight is not None:
         return out_c * compress_attn_weight + out_s

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.2"
+__version__ = "0.2.5"

fastvideo-kernel/tests/test_vsa.py

@@ -42,37 +42,13 @@ def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q
     q_padded = vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M)
     k_padded = vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M)
     v_padded = vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M)
-    # Use raw kernel or triton
-    try:
-        from fastvideo_kernel._C import fastvideo_kernel_ops
-        raw_kernel = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
-    except ImportError:
-        raw_kernel = None
+    # Use autograd-enabled wrapper (internally dispatches to SM90 kernel or Triton)
+    from fastvideo_kernel.block_sparse_attn import block_sparse_attn
+    output_padded, _aux = block_sparse_attn(
+        q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes
+    )
 
-    from fastvideo_kernel.triton_kernels.index import map_to_index
-
-    
-    # Convert mask to indices
-    # block_sparse_mask is [H, M, N] bool
-    # We need to map it to index.
-    # block_sparse_mask needs to be expanded/reshaped? 
-    # generate_block_sparse_mask_for_function returns [H, NumBlocksQ, NumBlocksKV]
-    
-    # Ops.py logic:
-    # mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
-    # idx, num = map_to_index(mask)
-    
-    idx, num = map_to_index(block_sparse_mask.unsqueeze(0)) # Add batch dim [1, H, M, N]
-    
-    if raw_kernel:
-        out_s = raw_kernel(q_padded, k_padded, v_padded, idx, num, variable_block_sizes.int())
-        output = out_s[0]
-    else:
-        # Fallback to triton testing if C++ not available
-        from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
-        output, _ = triton_block_sparse_attn_forward(q_padded, k_padded, v_padded, idx, num, variable_block_sizes)
-
-    output = output[:, :, q_non_pad_index, :]
+    output = output_padded[:, :, q_non_pad_index, :]
     output.backward(dO)
     return output, Q.grad, K.grad, V.grad
 
@@ -264,7 +240,6 @@ def generate_error_graphs_qkdiff(h, d, error_mode='all'):
 
     print("-" * 150)
 
-@pytest.mark.skip()
 def test_video_sparse_attention_backward():
     if not torch.cuda.is_available():
         return

fastvideo-kernel/tests/test_vsa_forward.py

@@ -3,6 +3,7 @@ import sys
 from typing import Tuple
 
 import torch
+import pytest
 
 from .utils import (
     generate_block_sparse_mask_for_function,
@@ -57,23 +58,14 @@ def block_sparse_forward_test(
     k_padded = ref.vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M)
     v_padded = ref.vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M)
 
-    # Use raw kernel or triton
+    # Use autograd-enabled wrapper (internally dispatches SM90 C++ vs Triton)
+    from fastvideo_kernel.block_sparse_attn import block_sparse_attn
     try:
-        from fastvideo_kernel._C import fastvideo_kernel_ops
-        raw_kernel = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None)
-    except ImportError:
-        raw_kernel = None
-        
-    from fastvideo_kernel.triton_kernels.index import map_to_index
-    idx, num = map_to_index(block_sparse_mask)
-    
-    if raw_kernel:
-        out_padded = raw_kernel(q_padded, k_padded, v_padded, idx, num, variable_block_sizes.int())[0]
-    else:
-        from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward
-        out_padded, _ = triton_block_sparse_attn_forward(
-            q_padded, k_padded, v_padded, idx, num, variable_block_sizes
+        out_padded, _aux = block_sparse_attn(
+            q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes
         )
+    except RuntimeError as e:
+        pytest.skip(str(e))
 
     # Remove padding on the query side
     out = out_padded[:, :, q_non_pad_index, :]
@@ -156,11 +148,6 @@ def run_forward_qk_diff(
 ) -> Tuple[float, float]:
     """
     Forward-only correctness test for the case S_q != S_kv.
-
-    NOTE:
-      - The Triton backend supports different Q/KV logical lengths via padding.
-      - The SM90 (H100) CUDA backend currently assumes the same number of blocks
-        for Q and KV, so we skip this test there.
     """
     assert torch.cuda.is_available(), "VSA kernels require CUDA"
 

fastvideo/attention/backends/video_sparse_attn.py

@@ -276,8 +276,9 @@ class VideoSparseAttentionImpl(AttentionImpl):
             query,
             key,
             value,
-            variable_block_sizes=attn_metadata.variable_block_sizes,
-            topk=cur_topk,
+            attn_metadata.variable_block_sizes,
+            attn_metadata.variable_block_sizes,
+            cur_topk,
             block_size=VSA_TILE_SIZE,
             compress_attn_weight=gate_compress).transpose(1, 2)
 

fastvideo/configs/models/__init__.py

@@ -2,5 +2,16 @@ from fastvideo.configs.models.base import ModelConfig
 from fastvideo.configs.models.dits.base import DiTConfig
 from fastvideo.configs.models.encoders.base import EncoderConfig
 from fastvideo.configs.models.vaes.base import VAEConfig
+from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig,
+                                            LTX2AudioEncoderConfig,
+                                            LTX2VocoderConfig)
 
-__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]
+__all__ = [
+    "ModelConfig",
+    "VAEConfig",
+    "DiTConfig",
+    "EncoderConfig",
+    "LTX2AudioEncoderConfig",
+    "LTX2AudioDecoderConfig",
+    "LTX2VocoderConfig",
+]

fastvideo/configs/models/audio/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.configs.models.audio.ltx2_audio_vae import (
+    LTX2AudioDecoderConfig,
+    LTX2AudioEncoderConfig,
+    LTX2VocoderConfig,
+)
+
+__all__ = [
+    "LTX2AudioEncoderConfig",
+    "LTX2AudioDecoderConfig",
+    "LTX2VocoderConfig",
+]

fastvideo/configs/models/audio/ltx2_audio_vae.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LTX-2 audio VAE and vocoder configuration.
+"""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+
+
+@dataclass
+class LTX2AudioArchConfig(ArchConfig):
+    architectures: list[str] = field(default_factory=list)
+
+
+@dataclass
+class LTX2AudioEncoderConfig(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
+        architectures=["LTX2AudioEncoder"]))
+
+
+@dataclass
+class LTX2AudioDecoderConfig(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
+        architectures=["LTX2AudioDecoder"]))
+
+
+@dataclass
+class LTX2VocoderConfig(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=lambda: LTX2AudioArchConfig(
+        architectures=["LTX2Vocoder"]))

fastvideo/configs/models/dits/__init__.py

@@ -3,11 +3,13 @@ from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
+from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
+from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
     "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
-    "LongCatVideoConfig"
+    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig"
 ]

fastvideo/configs/models/dits/hunyuanvideo15.py

@@ -26,6 +26,9 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
 
     param_names_mapping: dict = field(
         default_factory=lambda: {
+            r"^cond_type_embed\.(.*)$":
+            r"cond_type_embed.\1",
+
             # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
             r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
             r"txt_in.t_embedder.mlp.fc_in.\1",
@@ -55,6 +58,16 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
             r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
             r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
             r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.self_attn_qkv\.(.*)$":
+            r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
+            r"^image_embedder\.linear_1\.(.*)$":
+            r"image_embedder.linear_1.\1",
+            r"^image_embedder\.linear_2\.(.*)$":
+            r"image_embedder.linear_2.\1",
+            r"^image_embedder\.norm_in\.(.*)$":
+            r"image_embedder.norm_in.\1",
+            r"^image_embedder\.norm_out\.(.*)$":
+            r"image_embedder.norm_out.\1",
 
             # 2. txt_in_2 mapping:
             r"^context_embedder_2\.(.*)$":
@@ -144,6 +157,12 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
     exclude_lora_layers: list[str] = field(
         default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
 
+    # Causal HunyuanVideo1.5
+    local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
+    sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+    num_frames_per_block: int = 3
+    sliding_window_num_frames: int = 31
+
     def __post_init__(self):
         super().__post_init__()
         self.hidden_size: int = self.attention_head_dim * self.num_attention_heads

fastvideo/configs/models/dits/hyworld.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_double_block(n: str, m) -> bool:
+    return "double" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_single_block(n: str, m) -> bool:
+    return "single" in n and str.isdigit(n.split(".")[-1])
+
+
+# def is_refiner_block(n: str, m) -> bool:
+#     return "refiner" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_txt_in(n: str, m) -> bool:
+    return n.split(".")[-1] == "txt_in"
+
+
+@dataclass
+class HYWorldArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_single_block])
+
+    _compile_conditions: list = field(
+        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # 1. txt_in submodules (text embedder, refiner blocks):
+            r"^txt_in\.t_embedder\.mlp\.0\.(.*)$":
+            r"txt_in.t_embedder.mlp.fc_in.\1",
+            r"^txt_in\.t_embedder\.mlp\.2\.(.*)$":
+            r"txt_in.t_embedder.mlp.fc_out.\1",
+            r"^txt_in\.c_embedder\.linear_1\.(.*)$":
+            r"txt_in.c_embedder.fc_in.\1",
+            r"^txt_in\.c_embedder\.linear_2\.(.*)$":
+            r"txt_in.c_embedder.fc_out.\1",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.norm1\.(.*)$":
+            r"txt_in.refiner_blocks.\1.norm1.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.norm2\.(.*)$":
+            r"txt_in.refiner_blocks.\1.norm2.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.self_attn_qkv\.(.*)$":
+            r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.self_attn_proj\.(.*)$":
+            r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.mlp\.fc1\.(.*)$":
+            r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.mlp\.fc2\.(.*)$":
+            r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
+            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$":
+            r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+
+            # 2. time_in mappings (HYWorld uses TimestepEmbedder directly,
+            #    but FastVideo model inherits HunyuanVideo15TimeEmbedding with timestep_embedder):
+            r"^time_in\.mlp\.0\.(.*)$":
+            r"time_in.timestep_embedder.mlp.fc_in.\1",
+            r"^time_in\.mlp\.2\.(.*)$":
+            r"time_in.timestep_embedder.mlp.fc_out.\1",
+
+            # 3. action_in mappings:
+            r"^action_in\.mlp\.0\.(.*)$":
+            r"action_in.mlp.fc_in.\1",
+            r"^action_in\.mlp\.2\.(.*)$":
+            r"action_in.mlp.fc_out.\1",
+
+            # 4. byt5_in -> txt_in_2 mappings:
+            r"^byt5_in\.layernorm\.(.*)$":
+            r"txt_in_2.norm.\1",
+            r"^byt5_in\.fc1\.(.*)$":
+            r"txt_in_2.linear_1.\1",
+            r"^byt5_in\.fc2\.(.*)$":
+            r"txt_in_2.linear_2.\1",
+            r"^byt5_in\.fc3\.(.*)$":
+            r"txt_in_2.linear_3.\1",
+
+            # 5. cond_type_embedding -> cond_type_embed:
+            r"^cond_type_embedding\.(.*)$":
+            r"cond_type_embed.\1",
+
+            # 6. vision_in -> image_embedder mappings:
+            r"^vision_in\.proj\.0\.(.*)$":
+            r"image_embedder.norm_in.\1",
+            r"^vision_in\.proj\.1\.(.*)$":
+            r"image_embedder.linear_1.\1",
+            r"^vision_in\.proj\.3\.(.*)$":
+            r"image_embedder.linear_2.\1",
+            r"^vision_in\.proj\.4\.(.*)$":
+            r"image_embedder.norm_out.\1",
+
+            # 7. double_blocks mapping:
+            r"^double_blocks\.(\d+)\.img_attn_q\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^double_blocks\.(\d+)\.img_attn_k\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^double_blocks\.(\d+)\.img_attn_v\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^double_blocks\.(\d+)\.txt_attn_q\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^double_blocks\.(\d+)\.txt_attn_k\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^double_blocks\.(\d+)\.txt_attn_v\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^double_blocks\.(\d+)\.img_mlp\.fc1\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^double_blocks\.(\d+)\.img_mlp\.fc2\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.fc1\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^double_blocks\.(\d+)\.txt_mlp\.fc2\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_out.\2",
+
+            # 8. Final layer mapping:
+            r"^final_layer\.adaLN_modulation\.1\.(.*)$":
+            r"final_layer.adaLN_modulation.linear.\1",
+        })
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # Parameters from HY-WorldPlay config.json (loaded from checkpoint)
+    patch_size: list | tuple | int = field(default_factory=lambda: [1, 1, 1])
+    # Base latent channels - will be expanded in __post_init__ if concat_condition=True
+    in_channels: int = 32
+    concat_condition: bool = True
+    out_channels: int = 32
+    hidden_size: int = 2048
+    heads_num: int = 16
+    mlp_width_ratio: float = 4.0
+    mlp_act_type: str = "gelu_tanh"
+    mm_double_blocks_depth: int = 54
+    mm_single_blocks_depth: int = 0
+    rope_dim_list: list | tuple = field(default_factory=lambda: [16, 56, 56])
+    qkv_bias: bool = True
+    qk_norm: bool | str = True
+    qk_norm_type: str = "rms"
+    guidance_embed: bool = False
+    use_meanflow: bool = False
+    text_projection: str = "single_refiner"
+    use_attention_mask: bool = True
+    text_states_dim: int = 3584
+    text_states_dim_2: int | None = None
+    text_pool_type: str | None = None
+    rope_theta: float = 256.0
+    attn_mode: str = "flash"
+    attn_param: str | None = None
+    glyph_byT5_v2: bool = True
+    vision_projection: str = "linear"
+    vision_states_dim: int = 1152
+    is_reshape_temporal_channels: bool = False
+    use_cond_type_embedding: bool = True
+    ideal_resolution: str = "480p"
+    ideal_task: str = "i2v"
+    task_type: str = "i2v"
+    exclude_lora_layers: list[str] = field(
+        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Convert HY-WorldPlay naming to FastVideo naming conventions
+        self.num_attention_heads: int = self.heads_num
+        self.attention_head_dim: int = self.hidden_size // self.heads_num
+        self.num_layers: int = self.mm_double_blocks_depth
+        self.num_single_layers: int = self.mm_single_blocks_depth
+        self.num_refiner_layers: int = 2  # Default for HYWorld
+        self.mlp_ratio: float = float(self.mlp_width_ratio)
+        self.text_embed_dim: int = self.text_states_dim
+        self.text_embed_2_dim: int = self.text_states_dim_2 if self.text_states_dim_2 else 1472
+        self.image_embed_dim: int = self.vision_states_dim
+        self.rope_axes_dim: tuple[int, ...] = tuple(self.rope_dim_list)
+        self.num_channels_latents: int = self.out_channels
+        self.target_size: int = 640
+
+        # Handle concat_condition: when True, actual in_channels = base * 2 + 1
+        # (base latent + condition latent + mask channel)
+        # config.json has base in_channels (32), but img_in needs full (65)
+        if self.concat_condition and self.in_channels == 32:
+            if self.is_reshape_temporal_channels:
+                self.in_channels = self.in_channels + self.in_channels // 2 + 1
+            else:
+                self.in_channels = self.in_channels * 2 + 1  # 32 * 2 + 1 = 65
+
+        # Handle patch_size (can be list/tuple or int)
+        if isinstance(self.patch_size, list | tuple):
+            self.patch_size_t: int = self.patch_size[0]
+            # assume square patch size for height and width
+            patch_size_hw: int = self.patch_size[1]
+            object.__setattr__(self, 'patch_size', patch_size_hw)
+        else:
+            self.patch_size_t = 1
+
+        # Convert qk_norm to string format
+        if isinstance(self.qk_norm, bool):
+            if self.qk_norm:
+                self.qk_norm = "rms_norm" if self.qk_norm_type == "rms" else self.qk_norm_type
+            else:
+                self.qk_norm = "none"
+
+
+@dataclass
+class HYWorldConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=HYWorldArchConfig)
+
+    prefix: str = "HYWorld"

fastvideo/configs/models/dits/ltx2.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LTX-2 Transformer configuration for native FastVideo integration.
+"""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_ltx2_blocks(name: str, _module) -> bool:
+    """FSDP shard condition for LTX-2 transformer blocks."""
+    return "transformer_blocks" in name
+
+
+@dataclass
+class LTX2VideoArchConfig(DiTArchConfig):
+    """Architecture configuration for LTX-2 video transformer."""
+
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [is_ltx2_blocks])
+    _compile_conditions: list = field(default_factory=lambda: [is_ltx2_blocks])
+
+    # Parameter name mapping for weight conversion (hf/comfy -> FastVideo)
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^model\.diffusion_model\.(.*)$": r"model.\1",
+            r"^diffusion_model\.(.*)$": r"model.\1",
+            r"^model\.(.*)$": r"model.\1",
+            r"^(.*)$": r"model.\1",
+        })
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+    lora_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # Core transformer settings (defaults from LTX-2 metadata)
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    num_layers: int = 48
+    cross_attention_dim: int = 4096
+    caption_channels: int = 3840
+    norm_eps: float = 1e-6
+    attention_type: str = "default"
+    rope_type: str = "split"
+    double_precision_rope: bool = True
+
+    positional_embedding_theta: float = 10000.0
+    positional_embedding_max_pos: list[int] = field(
+        default_factory=lambda: [20, 2048, 2048])
+    timestep_scale_multiplier: int = 1000
+    use_middle_indices_grid: bool = True
+
+    # Patchification (video-only path)
+    patch_size: tuple[int, int, int] = (1, 1, 1)
+    num_channels_latents: int = 128
+    in_channels: int | None = None
+    out_channels: int | None = None
+
+    # Audio defaults (reserved for joint AV ports)
+    audio_num_attention_heads: int = 32
+    audio_attention_head_dim: int = 64
+    audio_in_channels: int = 128
+    audio_out_channels: int = 128
+    audio_cross_attention_dim: int = 2048
+    audio_positional_embedding_max_pos: list[int] = field(
+        default_factory=lambda: [20])
+    av_ca_timestep_scale_multiplier: int = 1
+
+    def __post_init__(self):
+        super().__post_init__()
+        patch_volume = self.patch_size[0] * self.patch_size[
+            1] * self.patch_size[2]
+        if self.in_channels is None:
+            self.in_channels = self.num_channels_latents * patch_volume
+        if self.out_channels is None:
+            self.out_channels = self.in_channels
+
+
+@dataclass
+class LTX2VideoConfig(DiTConfig):
+    """Main configuration for LTX-2 transformer."""
+
+    arch_config: DiTArchConfig = field(default_factory=LTX2VideoArchConfig)
+    prefix: str = "ltx2"

fastvideo/configs/models/encoders/__init__.py

@@ -7,10 +7,14 @@ from fastvideo.configs.models.encoders.clip import (
 from fastvideo.configs.models.encoders.llama import LlamaConfig
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
+from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
+from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config
+from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
     "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig",
     "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig",
-    "Qwen2_5_VLConfig"
+    "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig",
+    "SiglipVisionConfig"
 ]

fastvideo/configs/models/encoders/gemma.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+@dataclass
+class LTX2GemmaArchConfig(TextEncoderArchConfig):
+    architectures: list[str] = field(
+        default_factory=lambda: ["LTX2GemmaTextEncoderModel"])
+    hidden_size: int = 3840
+    num_hidden_layers: int = 48
+    num_attention_heads: int = 30
+    text_len: int = 1024
+    pad_token_id: int = 0
+    eos_token_id: int = 2
+
+    gemma_model_path: str = ""
+    gemma_dtype: str = "bfloat16"
+    padding_side: str = "left"
+
+    feature_extractor_in_features: int = 3840 * 49
+    feature_extractor_out_features: int = 3840
+
+    connector_num_attention_heads: int = 30
+    connector_attention_head_dim: int = 128
+    connector_num_layers: int = 2
+    connector_positional_embedding_theta: float = 10000.0
+    connector_positional_embedding_max_pos: list[int] = field(
+        default_factory=lambda: [4096])
+    connector_rope_type: str = "split"
+    connector_double_precision_rope: bool = False
+    connector_num_learnable_registers: int | None = 128
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.tokenizer_kwargs["padding"] = "max_length"
+
+
+@dataclass
+class LTX2GemmaConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(
+        default_factory=LTX2GemmaArchConfig)
+
+    prefix: str = "ltx2_gemma"

fastvideo/configs/models/encoders/reason1.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Config for Reason1 (Qwen2.5-VL) text encoder."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig
+
+
+@dataclass
+class Reason1ArchConfig(TextEncoderArchConfig):
+    """Architecture settings (defaults match Qwen2.5-VL-7B-Instruct)."""
+
+    architectures: list[str] = field(
+        default_factory=lambda: ["Qwen2_5_VLForConditionalGeneration"])
+    model_type: str = "qwen2_5_vl"
+
+    vocab_size: int = 152064
+    hidden_size: int = 3584
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 28
+    num_key_value_heads: int = 4
+    intermediate_size: int = 18944
+
+    text_len: int = 512
+    hidden_state_skip_layer: int = 0
+    bos_token_id: int = 151643
+    pad_token_id: int = 151643
+    eos_token_id: int = 151645
+
+    image_token_id: int = 151655
+    video_token_id: int = 151656
+    vision_token_id: int = 151654
+    vision_start_token_id: int = 151652
+    vision_end_token_id: int = 151653
+
+    vision_config: dict[str, Any] | None = None
+
+    rope_theta: float = 1000000.0
+    rope_scaling: dict[str, Any] | None = field(default_factory=lambda: {
+        "type": "mrope",
+        "mrope_section": [16, 24, 24]
+    })
+    max_position_embeddings: int = 128000
+    max_window_layers: int = 28
+
+    embedding_concat_strategy: str = "mean_pooling"
+    n_layers_per_group: int = 5
+    num_embedding_padding_tokens: int = 512
+
+    attention_dropout: float = 0.0
+    hidden_act: str = "silu"
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+
+    use_sliding_window: bool = False
+    sliding_window: int = 32768
+
+    tie_word_embeddings: bool = False
+    use_cache: bool = False
+    output_hidden_states: bool = True
+
+    torch_dtype: str = "bfloat16"
+    _attn_implementation: str = "flash_attention_2"
+
+
+@dataclass
+class Reason1Config(TextEncoderConfig):
+    """Reason1 text encoder config."""
+
+    arch_config: Reason1ArchConfig = field(default_factory=Reason1ArchConfig)
+    tokenizer_type: str = "Qwen/Qwen2.5-VL-7B-Instruct"

fastvideo/configs/models/encoders/siglip.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SigLIP vision encoder configuration for FastVideo."""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (ImageEncoderArchConfig,
+                                                    ImageEncoderConfig)
+
+
+@dataclass
+class SiglipVisionArchConfig(ImageEncoderArchConfig):
+    """Architecture configuration for SigLIP vision encoder.
+    
+    Fields match the config.json from HuggingFace SigLIP checkpoints.
+    """
+
+    # From config.json
+    architectures: list[str] = field(
+        default_factory=lambda: ["SiglipVisionModel"])
+    attention_dropout: float = 0.0
+    dtype: str | None = None
+    hidden_act: str = "gelu_pytorch_tanh"
+    hidden_size: int = 1152
+    image_size: int = 384
+    intermediate_size: int = 4304
+    layer_norm_eps: float = 1e-6
+    model_type: str = "siglip_vision_model"
+    num_attention_heads: int = 16
+    num_channels: int = 3
+    num_hidden_layers: int = 27
+    patch_size: int = 14
+
+    # FastVideo specific - QKV fusion mapping
+    stacked_params_mapping: list = field(default_factory=lambda: [
+        ("qkv_proj", "q_proj", "q"),
+        ("qkv_proj", "k_proj", "k"),
+        ("qkv_proj", "v_proj", "v"),
+    ])
+
+
+@dataclass
+class SiglipVisionConfig(ImageEncoderConfig):
+    """Configuration for SigLIP vision encoder."""
+
+    arch_config: ImageEncoderArchConfig = field(
+        default_factory=SiglipVisionArchConfig)
+
+    # FastVideo specific
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
+    enable_scale: bool = True
+    is_causal: bool = False
+    prefix: str = "siglip"

fastvideo/configs/models/vaes/__init__.py

@@ -1,6 +1,8 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
+from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
+from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
 from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
 
@@ -9,5 +11,7 @@ __all__ = [
     "WanVAEConfig",
     "StepVideoVAEConfig",
     "CosmosVAEConfig",
+    "Cosmos25VAEConfig",
     "Hunyuan15VAEConfig",
+    "LTX2VAEConfig",
 ]

fastvideo/configs/models/vaes/cosmos2_5vae.py

@@ -0,0 +1,223 @@
+"""Cosmos 2.5 (Wan2.1-style) VAE config and checkpoint-key mapping."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class Cosmos25VAEArchConfig(VAEArchConfig):
+    _name_or_path: str = ""
+    base_dim: int = 96
+    decoder_base_dim: int | None = None
+    z_dim: int = 16
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
+    num_res_blocks: int = 2
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
+    dropout: float = 0.0
+    is_residual: bool = False
+    in_channels: int = 3
+    out_channels: int = 3
+    patch_size: int | None = None
+    scale_factor_temporal: int = 4
+    scale_factor_spatial: int = 8
+    clip_output: bool = True
+
+    latents_mean: tuple[float, ...] = (
+        -0.7571,
+        -0.7089,
+        -0.9113,
+        0.1075,
+        -0.1745,
+        0.9653,
+        -0.1517,
+        1.5508,
+        0.4134,
+        -0.0715,
+        0.5517,
+        -0.3632,
+        -0.1922,
+        -0.9497,
+        0.2503,
+        -0.2921,
+    )
+    latents_std: tuple[float, ...] = (
+        2.8184,
+        1.4541,
+        2.3275,
+        2.6558,
+        1.2196,
+        1.7708,
+        2.6052,
+        2.0743,
+        3.2687,
+        2.1526,
+        2.8652,
+        1.5579,
+        1.6382,
+        1.1253,
+        2.8251,
+        1.9160,
+    )
+
+    # Simple 1:1 renames. More complex decoder remapping is handled by
+    # `map_official_key()`.
+    param_names_mapping: dict[str, str] = field(
+        default_factory=lambda: {
+            r"^conv1\.(.*)$": r"quant_conv.\1",
+            r"^conv2\.(.*)$": r"post_quant_conv.\1",
+            r"^encoder\.conv1\.(.*)$": r"encoder.conv_in.\1",
+            r"^decoder\.conv1\.(.*)$": r"decoder.conv_in.\1",
+            r"^encoder\.head\.0\.gamma$": r"encoder.norm_out.gamma",
+            r"^encoder\.head\.2\.(.*)$": r"encoder.conv_out.\1",
+            r"^decoder\.head\.0\.gamma$": r"decoder.norm_out.gamma",
+            r"^decoder\.head\.2\.(.*)$": r"decoder.conv_out.\1",
+        })
+
+    @staticmethod
+    def map_official_key(key: str) -> str | None:
+        """Map a single official checkpoint key into FastVideo key space."""
+
+        def map_residual_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^residual\.0\.gamma$", sub):
+                return f"{prefix}.norm1.gamma"
+            m = re.match(r"^residual\.2\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv1.{m.group(1)}"
+            if re.match(r"^residual\.3\.gamma$", sub):
+                return f"{prefix}.norm2.gamma"
+            m = re.match(r"^residual\.6\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv2.{m.group(1)}"
+            m = re.match(r"^shortcut\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.conv_shortcut.{m.group(1)}"
+            return None
+
+        def map_attn_subkey(prefix: str, sub: str) -> str | None:
+            if re.match(r"^norm\.gamma$", sub):
+                return f"{prefix}.norm.gamma"
+            m = re.match(r"^to_qkv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.to_qkv.{m.group(1)}"
+            m = re.match(r"^proj\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.proj.{m.group(1)}"
+            return None
+
+        def map_resample_subkey(prefix: str, sub: str) -> str | None:
+            m = re.match(r"^resample\.1\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.resample.1.{m.group(1)}"
+            m = re.match(r"^time_conv\.(weight|bias)$", sub)
+            if m:
+                return f"{prefix}.time_conv.{m.group(1)}"
+            return None
+
+        m = re.match(r"^conv1\.(weight|bias)$", key)
+        if m:
+            return f"quant_conv.{m.group(1)}"
+        m = re.match(r"^conv2\.(weight|bias)$", key)
+        if m:
+            return f"post_quant_conv.{m.group(1)}"
+        m = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_in.{m.group(2)}"
+        m = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key)
+        if m:
+            return f"{m.group(1)}.norm_out.gamma"
+        m = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key)
+        if m:
+            return f"{m.group(1)}.conv_out.{m.group(2)}"
+
+        m = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.0",
+                                       m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
+        if m:
+            return map_attn_subkey(f"{m.group(1)}.mid_block.attentions.0",
+                                   m.group(2))
+        m = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
+        if m:
+            return map_residual_subkey(f"{m.group(1)}.mid_block.resnets.1",
+                                       m.group(2))
+
+        m = re.match(r"^encoder\.downsamples\.(\d+)\.(.*)$", key)
+        if m:
+            idx = int(m.group(1))
+            sub = m.group(2)
+            if sub.startswith("residual.") or sub.startswith("shortcut."):
+                return map_residual_subkey(f"encoder.down_blocks.{idx}", sub)
+            if sub.startswith("resample.") or sub.startswith("time_conv."):
+                return map_resample_subkey(f"encoder.down_blocks.{idx}", sub)
+            return None
+
+        m = re.match(r"^decoder\.upsamples\.(\d+)\.(.*)$", key)
+        if m:
+            uidx = int(m.group(1))
+            sub = m.group(2)
+
+            if uidx in (0, 1, 2):
+                block_i, res_i = 0, uidx
+            elif uidx == 3:
+                block_i, res_i = 0, None
+            elif uidx in (4, 5, 6):
+                block_i, res_i = 1, uidx - 4
+            elif uidx == 7:
+                block_i, res_i = 1, None
+            elif uidx in (8, 9, 10):
+                block_i, res_i = 2, uidx - 8
+            elif uidx == 11:
+                block_i, res_i = 2, None
+            elif uidx in (12, 13, 14):
+                block_i, res_i = 3, uidx - 12
+            else:
+                return None
+
+            if res_i is None:
+                return map_resample_subkey(
+                    f"decoder.up_blocks.{block_i}.upsamplers.0",
+                    sub,
+                )
+
+            return map_residual_subkey(
+                f"decoder.up_blocks.{block_i}.resnets.{res_i}",
+                sub,
+            )
+
+        return None
+
+    temporal_compression_ratio: int = 4
+    spatial_compression_ratio: int = 8
+
+    def __post_init__(self):
+        self.scaling_factor: torch.Tensor = 1.0 / torch.tensor(
+            self.latents_std).view(1, self.z_dim, 1, 1, 1)
+        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(
+            1, self.z_dim, 1, 1, 1)
+        self.temporal_compression_ratio = self.scale_factor_temporal
+        self.spatial_compression_ratio = self.scale_factor_spatial
+
+
+@dataclass
+class Cosmos25VAEConfig(VAEConfig):
+    """Cosmos2.5 VAE config."""
+
+    arch_config: Cosmos25VAEArchConfig = field(
+        default_factory=Cosmos25VAEArchConfig)
+
+    use_feature_cache: bool = True
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    def __post_init__(self):
+        self.blend_num_frames = (self.tile_sample_min_num_frames -
+                                 self.tile_sample_stride_num_frames) * 2

fastvideo/configs/models/vaes/ltx2vae.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LTX-2 VAE configuration.
+"""
+
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class LTX2VAEArchConfig(VAEArchConfig):
+    # Mirrors LTX-2 safetensors metadata config under "vae"
+    _class_name: str = "CausalVideoAutoencoder"
+    dims: int = 3
+    in_channels: int = 3
+    out_channels: int = 3
+    latent_channels: int = 128
+    encoder_blocks: list = field(default_factory=list)
+    decoder_blocks: list = field(default_factory=list)
+    patch_size: int = 4
+    norm_layer: str = "pixel_norm"
+    latent_log_var: str = "uniform"
+    encoder_spatial_padding_mode: str = "zeros"
+    decoder_spatial_padding_mode: str = "reflect"
+    causal_decoder: bool = False
+    timestep_conditioning: bool = True
+    use_quant_conv: bool = False
+    scaling_factor: float = 1.0
+    normalize_latent_channels: bool = False
+
+    # Match FastVideo naming for compression ratios (LTX-2 default)
+    temporal_compression_ratio: int = 8
+    spatial_compression_ratio: int = 32
+
+
+@dataclass
+class LTX2VAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = field(default_factory=LTX2VAEArchConfig)
+
+    # LTX-2 tiling defaults (match ltx_core.video_vae.TilingConfig.default()).
+    ltx2_spatial_tile_size_in_pixels: int = 512
+    ltx2_spatial_tile_overlap_in_pixels: int = 64
+    ltx2_temporal_tile_size_in_frames: int = 64
+    ltx2_temporal_tile_overlap_in_frames: int = 24

fastvideo/configs/pipelines/__init__.py

@@ -1,8 +1,11 @@
 from fastvideo.configs.pipelines.base import (PipelineConfig,
                                               SlidingTileAttnConfig)
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
+from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
+from fastvideo.configs.pipelines.hyworld import HYWorldConfig
+from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
 from fastvideo.configs.pipelines.registry import (
     get_pipeline_config_cls_from_name)
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
@@ -15,5 +18,6 @@ __all__ = [
     "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig",
     "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
     "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig",
-    "CosmosConfig", "get_pipeline_config_cls_from_name"
+    "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
+    "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/cosmos2_5.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.configs.models.dits import Cosmos25VideoConfig
+from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25ArchConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.reason1 import Reason1Config, Reason1ArchConfig
+from fastvideo.configs.models.vaes import Cosmos25VAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode
+
+
+def _identity_preprocess_text(prompt: str) -> str:
+    return prompt
+
+
+def reason1_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    hidden_states = getattr(outputs, "hidden_states", None)
+    if hidden_states is None:
+        raise ValueError("Reason1 postprocess requires outputs.hidden_states")
+
+    hs = list(hidden_states)[1:]
+    normed = []
+    for h in hs:
+        h = h.float()
+        h = (h - h.mean(dim=-1, keepdim=True)) / (h.std(dim=-1, keepdim=True) +
+                                                  1e-8)
+        normed.append(h)
+    return torch.cat(normed, dim=-1).to(hidden_states[0].dtype)
+
+
+@dataclass
+class Cosmos25Config(PipelineConfig):
+    """Configuration for Cosmos 2.5 (Predict2.5) video generation pipeline."""
+
+    dit_config: DiTConfig = field(default_factory=lambda: Cosmos25VideoConfig(
+        arch_config=Cosmos25ArchConfig(
+            num_attention_heads=16,
+            attention_head_dim=128,
+            in_channels=16,
+            out_channels=16,
+            num_layers=28,
+            patch_size=[1, 2, 2],
+            max_size=[128, 240, 240],
+            rope_scale=[1.0, 3.0, 3.0],
+            text_embed_dim=1024,
+            mlp_ratio=4.0,
+            adaln_lora_dim=256,
+            use_adaln_lora=True,
+            concat_padding_mask=True,
+            extra_pos_embed_type=None,
+            use_crossattn_projection=True,
+            rope_enable_fps_modulation=False,
+            qk_norm="rms_norm",
+        )))
+
+    vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig)
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Reason1Config(arch_config=Reason1ArchConfig(
+            embedding_concat_strategy="full_concat")), ))
+
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (_identity_preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (reason1_postprocess_text, ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("bf16", ))
+
+    embedded_cfg_scale: float = 0.0
+    flow_shift: float = 5.0
+
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    STA_mode: STA_Mode = STA_Mode.NONE
+    skip_time_steps: int = 0
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self._vae_latent_dim = 16

fastvideo/configs/pipelines/hunyuan15.py

@@ -127,13 +127,29 @@ class Hunyuan15T2V480PConfig(PipelineConfig):
     vae_tiling: bool = True
 
     def __post_init__(self):
-        self.vae_config.load_encoder = False
+        self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
 
 
+@dataclass
+class Hunyuan15DistilledI2V480PConfig(Hunyuan15T2V480PConfig):
+    flow_shift: int = 7
+
+
 @dataclass
 class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig):
     """Base configuration for HunYuan pipeline architecture."""
 
     # HunyuanConfig-specific parameters with defaults
     flow_shift: int = 9
+
+
+@dataclass
+class SelfForcingHunyuan15T2V480PConfig(Hunyuan15T2V480PConfig):
+    flow_shift: int = 5
+    is_causal: bool = True
+    # dmd_denoising_steps: list[int] | None = field(
+    #     default_factory=lambda: [1000, 750, 500, 250])
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 875, 750, 625, 500, 375, 250, 125])
+    warp_denoising_step: bool = True

fastvideo/configs/pipelines/hyworld.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig
+from fastvideo.configs.models.dits import HYWorldConfig as HYWorldDiTConfig
+from fastvideo.configs.models.encoders import SiglipVisionConfig
+from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig
+
+
+@dataclass
+class HYWorldConfig(Hunyuan15T2V480PConfig):
+    """Base configuration for HYWorld pipeline architecture."""
+
+    # HYWorldConfig-specific parameters with defaults
+    dit_config: DiTConfig = field(default_factory=HYWorldDiTConfig)
+
+    # SigLIP image encoder for I2V
+    image_encoder_config: EncoderConfig = field(
+        default_factory=SiglipVisionConfig)
+    image_encoder_precision: str = "fp16"
+    # vae_precision: str = "fp32"
+
+    # Text encoding
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("fp16", "fp32"))
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/ltx2.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
+                                      LTX2AudioDecoderConfig, LTX2VocoderConfig,
+                                      VAEConfig)
+from fastvideo.configs.models.dits import LTX2VideoConfig
+from fastvideo.configs.models.encoders import BaseEncoderOutput, LTX2GemmaConfig
+from fastvideo.configs.models.vaes import LTX2VAEConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
+
+
+def ltx2_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    return outputs.last_hidden_state
+
+
+@dataclass
+class LTX2T2VConfig(PipelineConfig):
+    """Configuration for LTX-2 T2V pipeline."""
+
+    dit_config: DiTConfig = field(default_factory=LTX2VideoConfig)
+    vae_config: VAEConfig = field(default_factory=LTX2VAEConfig)
+    vae_tiling: bool = True
+    vae_sp: bool = False
+
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (LTX2GemmaConfig(), ))
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (ltx2_postprocess_text, ))
+
+    dit_precision: str = "bf16"
+    vae_precision: str = "bf16"
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: ("bf16", ))
+
+    audio_decoder_config: ModelConfig = field(
+        default_factory=LTX2AudioDecoderConfig)
+    vocoder_config: ModelConfig = field(default_factory=LTX2VocoderConfig)
+    audio_decoder_precision: str = "bf16"
+    vocoder_precision: str = "bf16"
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True

fastvideo/configs/pipelines/registry.py

@@ -6,8 +6,11 @@ from collections.abc import Callable
 
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.cosmos import CosmosConfig
+from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
-from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
+from fastvideo.configs.pipelines.hyworld import HYWorldConfig
+from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
+from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig, SelfForcingHunyuan15T2V480PConfig, Hunyuan15DistilledI2V480PConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.configs.pipelines.turbodiffusion import (
@@ -32,10 +35,15 @@ logger = init_logger(__name__)
 PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
     "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
+    "weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v":
+    SelfForcingHunyuan15T2V480PConfig,
+    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled":
+    Hunyuan15DistilledI2V480PConfig,
     "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
     Hunyuan15T2V480PConfig,
     "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
     Hunyuan15T2V720PConfig,
+    "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorldConfig,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
     "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
     "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig,
@@ -55,6 +63,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
     "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
     "nvidia/Cosmos-Predict2-2B-Video2World": CosmosConfig,
+    "KyleShao/Cosmos-Predict2.5-2B-Diffusers": Cosmos25Config,
     "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGameI2V480PConfig,
     "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGameI2V480PConfig,
     "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGameI2V480PConfig,
@@ -62,6 +71,9 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/LongCat-Video-T2V-Diffusers": LongCatT2V480PConfig,
     "FastVideo/LongCat-Video-I2V-Diffusers": LongCatT2V480PConfig,
     "FastVideo/LongCat-Video-VC-Diffusers": LongCatT2V480PConfig,
+    # LTX-2 models
+    "Lightricks/LTX-2": LTX2T2VConfig,
+    "converted/ltx2_diffusers": LTX2T2VConfig,
     # TurboDiffusion models
     "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": TurboDiffusionT2V_1_3B_Config,
     "loayrashid/TurboWan2.1-T2V-14B-Diffusers": TurboDiffusionT2V_14B_Config,
@@ -81,6 +93,8 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "hunyuan" in id.lower(),
     "hunyuan15":
     lambda id: "hunyuan15" in id.lower(),
+    "hyworld":
+    lambda id: "hyworld" in id.lower(),
     "matrixgame":
     lambda id: "matrix-game" in id.lower() or "matrixgame" in id.lower(),
     "wanpipeline":
@@ -94,9 +108,14 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "stepvideo":
     lambda id: "stepvideo" in id.lower(),
     "cosmos":
-    lambda id: "cosmos" in id.lower(),
+    lambda id: "cosmos" in id.lower() and ("2.5" not in id.lower(
+    ) and "2_5" not in id.lower() and "25" not in id.lower()),
+    "cosmos25":
+    lambda id: "cosmos25" in id.lower(),
     "turbodiffusion":
     lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
+    "ltx2":
+    lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -105,11 +124,14 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "longcatimagetovideo": LongCatT2V480PConfig,
     "longcatvideocontinuation": LongCatT2V480PConfig,
     "longcat": LongCatT2V480PConfig,
+    "cosmos25": Cosmos25Config,
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "matrixgame": MatrixGameI2V480PConfig,
     "hunyuan15":
     Hunyuan15T2V480PConfig,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
+    "hyworld":
+    HYWorldConfig,  # HYWorld-specific config as fallback for any HYWorld variant
     "wanpipeline":
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
@@ -117,6 +139,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
     "stepvideo": StepVideoT2VConfig,
     "turbodiffusion": TurboDiffusionT2V_1_3B_Config,
+    "ltx2": LTX2T2VConfig,
     # Other fallbacks by architecture
 }
 

fastvideo/configs/sample/cosmos2_5.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+@dataclass
+class Cosmos_Predict2_5_2B_Diffusers_SamplingParam(SamplingParam):
+    """Defaults for Cosmos 2.5 (Predict2.5) text-to-video diffusers-format model."""
+
+    height: int = 480
+    width: int = 832
+    num_frames: int = 121
+    fps: int = 24
+
+    guidance_scale: float = 7.0
+    # Official Cosmos2.5 sampling uses empty string as unconditional.
+    negative_prompt: str = ""
+    num_inference_steps: int = 35

fastvideo/configs/sample/hunyuan15.py

@@ -17,11 +17,24 @@ class Hunyuan15_480P_SamplingParam(SamplingParam):
     guidance_scale: float = 6.0
     prompt_attention_mask: list = field(default_factory=list)
     negative_attention_mask: list = field(default_factory=list)
-    sigmas: list[float] | None = field(
-        default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
+    # sigmas: list[float] | None = field(
+    #     default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
 
     negative_prompt: str = ""
 
+    def __post_init__(self):
+        super().__post_init__()
+        self.sigmas = list(
+            np.linspace(1.0, 0.0, self.num_inference_steps + 1)[:-1])
+
+
+@dataclass
+class Hunyuan15_480P_Distilled_SamplingParam(Hunyuan15_480P_SamplingParam):
+    num_inference_steps: int = 8
+    fps: int = 24
+
+    guidance_scale: float = 1.0
+
 
 @dataclass
 class Hunyuan15_720P_SamplingParam(Hunyuan15_480P_SamplingParam):

fastvideo/configs/sample/hyworld.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.sample.base import SamplingParam
+import numpy as np
+
+
+@dataclass
+class HYWorld_SamplingParam(SamplingParam):
+    num_inference_steps: int = 50
+
+    num_frames: int = 125
+    height: int = 480
+    width: int = 832
+    fps: int = 24
+
+    # Camera trajectory: pose string (e.g., 'w-31' means generating [1 + 31] latents) or JSON file path
+    pose: str = 'w-31'
+
+    guidance_scale: float = 6.0
+    prompt_attention_mask: list = field(default_factory=list)
+    negative_attention_mask: list = field(default_factory=list)
+    sigmas: list[float] | None = field(
+        default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
+
+    negative_prompt: str = ""

fastvideo/configs/sample/ltx2.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+@dataclass
+class LTX2SamplingParam(SamplingParam):
+    """Default sampling parameters for LTX-2 distilled T2V.
+    """
+
+    seed: int = 10
+    num_frames: int = 121
+    height: int = 1024
+    width: int = 1536
+    fps: int = 24
+    num_inference_steps: int = 8
+    guidance_scale: float = 1.0
+    # No default negative_prompt for distilled models
+    negative_prompt: str = ""

fastvideo/configs/sample/registry.py

@@ -5,10 +5,13 @@ from typing import Any
 
 from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                               HunyuanSamplingParam)
-from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam
+from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam
+from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam, Hunyuan15_480P_Distilled_SamplingParam
 from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
 
 from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam
+from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam
+from fastvideo.configs.sample.ltx2 import LTX2SamplingParam
 
 # isort: off
 from fastvideo.configs.sample.wan import (
@@ -39,48 +42,41 @@ from fastvideo.utils import (maybe_download_model_index,
 logger = init_logger(__name__)
 # Registry maps specific model weights to their config classes
 SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
-    "FastVideo/FastHunyuan-diffusers":
-    FastHunyuanSamplingParam,
-    "hunyuanvideo-community/HunyuanVideo":
-    HunyuanSamplingParam,
+    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
+    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
+    "weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v":
+    Hunyuan15_480P_SamplingParam,
+    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled":
+    Hunyuan15_480P_Distilled_SamplingParam,
     "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
     Hunyuan15_480P_SamplingParam,
     "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
     Hunyuan15_720P_SamplingParam,
-    "FastVideo/stepvideo-t2v-diffusers":
-    StepVideoT2VSamplingParam,
+    "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorld_SamplingParam,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
 
     # Wan2.1
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers":
-    WanT2V_1_3B_SamplingParam,
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers":
-    WanT2V_14B_SamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers":
-    WanI2V_14B_480P_SamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers":
-    WanI2V_14B_720P_SamplingParam,
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
     "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
     Wan2_1_Fun_1_3B_InP_SamplingParam,
     "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers":
     Wan2_1_Fun_1_3B_Control_SamplingParam,
 
     # Wan2.2
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers":
-    Wan2_2_TI2V_5B_SamplingParam,
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
     "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers":
     Wan2_2_TI2V_5B_SamplingParam,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers":
-    Wan2_2_T2V_A14B_SamplingParam,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers":
-    Wan2_2_I2V_A14B_SamplingParam,
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
+    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
 
     # FastWan2.1
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers":
-    FastWanT2V480P_SamplingParam,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480P_SamplingParam,
 
     # FastWan2.2
-    "FastVideo/FastWan2.2-TI2V-5B-Diffusers":
-    Wan2_2_TI2V_5B_SamplingParam,
+    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
 
     # Causal Self-Forcing Wan2.1
     "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers":
@@ -96,13 +92,14 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "nvidia/Cosmos-Predict2-2B-Video2World":
     Cosmos_Predict2_2B_Video2World_SamplingParam,
 
+    # Cosmos2.5
+    "KyleShao/Cosmos-Predict2.5-2B-Diffusers":
+    Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
+
     # MatrixGame2.0 models
-    "FastVideo/Matrix-Game-2.0-Base-Diffusers":
-    MatrixGame2_SamplingParam,
-    "FastVideo/Matrix-Game-2.0-GTA-Diffusers":
-    MatrixGame2_SamplingParam,
-    "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers":
-    MatrixGame2_SamplingParam,
+    "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam,
+    "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam,
+    "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGame2_SamplingParam,
 
     # TurboDiffusion models
     "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers":
@@ -112,6 +109,10 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "loayrashid/TurboWan2.2-I2V-A14B-Diffusers":
     TurboDiffusionI2V_A14B_SamplingParam,
 
+    # LTX-2 models
+    "Lightricks/LTX-2": LTX2SamplingParam,
+    "FastVideo/LTX2-Distilled-Diffusers": LTX2SamplingParam,
+
     # Add other specific weight variants
 }
 
@@ -121,6 +122,8 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "hunyuan" in id.lower(),
     "hunyuan15":
     lambda id: "hunyuan15" in id.lower(),
+    "hyworld":
+    lambda id: "hyworld" in id.lower(),
     "wanpipeline":
     lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo":
@@ -135,6 +138,12 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(),
     "turbodiffusion":
     lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
+    "cosmos25":
+    lambda id: "cosmos2_5" in id.lower(),
+    "cosmos":
+    lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(),
+    "ltx2":
+    lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -144,6 +153,8 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
     "hunyuan15":
     Hunyuan15_480P_SamplingParam,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
+    "hyworld":
+    HYWorld_SamplingParam,  # HYWorld-specific config as fallback for any HYWorld variant
     "wanpipeline":
     WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V_14B_480P_SamplingParam,
@@ -153,13 +164,15 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "matrixgame": MatrixGame2_SamplingParam,
     "turbodiffusion":
     TurboDiffusionT2V_1_3B_SamplingParam,  # Default to T2V for fallback
+    "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
+    "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam,
+    "ltx2": LTX2SamplingParam,
     # Other fallbacks by architecture
 }
 
 
 def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
     """Get the appropriate sampling param for specific pretrained weights."""
-
     # First try exact match for specific weights
     if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY:
         return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path]
@@ -176,9 +189,6 @@ def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
 
     if os.path.exists(pipeline_name_or_path):
         config = verify_model_config_and_directory(pipeline_name_or_path)
-        logger.warning(
-            "FastVideo may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed."
-        )
     else:
         config = maybe_download_model_index(pipeline_name_or_path)
 

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

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,4 +124,4 @@ pyarrow_schema_text_only = pa.schema([
     pa.field("text_embedding_dtype", pa.string()),
     # --- Metadata ---
     pa.field("caption", pa.string()),
-])
+])

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__)
 
@@ -655,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)
 
@@ -758,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"]

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 = []

fastvideo/distributed/device_communicators/base_device_communicator.py

@@ -57,6 +57,9 @@ class DistributedAutograd:
             ctx.dim = dim
             ctx.input_shape = input_.shape
 
+            # NCCL all_gather_into_tensor requires contiguous tensors.
+            if not input_.is_contiguous():
+                input_ = input_.contiguous()
             input_size = input_.size()
             output_size = (input_size[0] * world_size, ) + input_size[1:]
             output_tensor = torch.empty(output_size,

fastvideo/entrypoints/video_generator.py

@@ -18,6 +18,8 @@ import numpy as np
 import torch
 import torchvision
 from einops import rearrange
+import shutil
+import tempfile
 
 from fastvideo.configs.sample import SamplingParam
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -389,6 +391,11 @@ class VideoGenerator:
         if batch.save_video:
             imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
             logger.info("Saved video to %s", output_path)
+            audio = output_batch.extra.get("audio")
+            audio_sample_rate = output_batch.extra.get("audio_sample_rate")
+            if (audio is not None and audio_sample_rate is not None and
+                    not self._mux_audio(output_path, audio, audio_sample_rate)):
+                logger.warning("Audio mux failed; saved video without audio.")
 
         if batch.return_frames:
             return frames
@@ -396,6 +403,7 @@ class VideoGenerator:
             return {
                 "samples": samples,
                 "frames": frames,
+                "audio": output_batch.extra.get("audio"),
                 "prompts": prompt,
                 "size": (target_height, target_width, batch.num_frames),
                 "generation_time": gen_time,
@@ -405,6 +413,98 @@ class VideoGenerator:
                 "trajectory_decoded": output_batch.trajectory_decoded,
             }
 
+    @staticmethod
+    def _mux_audio(
+        video_path: str,
+        audio: torch.Tensor | np.ndarray,
+        sample_rate: int,
+    ) -> bool:
+        """Mux audio into video using PyAV."""
+        try:
+            import av
+        except ImportError:
+            logger.warning("PyAV not installed; cannot mux audio. "
+                           "Install with: pip install av")
+            return False
+
+        if torch.is_tensor(audio):
+            audio_np = audio.detach().cpu().float().numpy()
+        else:
+            audio_np = np.asarray(audio, dtype=np.float32)
+
+        if audio_np.ndim == 1:
+            audio_np = audio_np[:, None]
+        elif audio_np.ndim == 2:
+            if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
+                audio_np = audio_np.T
+        else:
+            logger.warning("Unexpected audio shape %s; skipping mux.",
+                           audio_np.shape)
+            return False
+
+        audio_np = np.clip(audio_np, -1.0, 1.0)
+        audio_int16 = (audio_np * 32767.0).astype(np.int16)
+        num_channels = audio_int16.shape[1]
+        layout = "stereo" if num_channels == 2 else "mono"
+
+        try:
+            import wave
+            with tempfile.TemporaryDirectory() as tmpdir:
+                out_path = os.path.join(tmpdir, "muxed.mp4")
+                wav_path = os.path.join(tmpdir, "audio.wav")
+
+                # Write audio to WAV file
+                with wave.open(wav_path, "wb") as wav_file:
+                    wav_file.setnchannels(num_channels)
+                    wav_file.setsampwidth(2)
+                    wav_file.setframerate(sample_rate)
+                    wav_file.writeframes(audio_int16.tobytes())
+
+                # Open input video and audio
+                input_video = av.open(video_path)
+                input_audio = av.open(wav_path)
+
+                # Create output with both streams
+                output = av.open(out_path, mode="w")
+
+                # Add video stream (copy codec from input)
+                in_video_stream = input_video.streams.video[0]
+                out_video_stream = output.add_stream(
+                    codec_name=in_video_stream.codec_context.name,
+                    rate=in_video_stream.average_rate,
+                )
+                out_video_stream.width = in_video_stream.width
+                out_video_stream.height = in_video_stream.height
+                out_video_stream.pix_fmt = in_video_stream.pix_fmt
+
+                # Add audio stream (AAC)
+                out_audio_stream = output.add_stream("aac", rate=sample_rate)
+                out_audio_stream.layout = layout
+
+                # Remux video (decode and re-encode to be safe)
+                for frame in input_video.decode(video=0):
+                    for packet in out_video_stream.encode(frame):
+                        output.mux(packet)
+                for packet in out_video_stream.encode():
+                    output.mux(packet)
+
+                # Encode audio
+                for frame in input_audio.decode(audio=0):
+                    frame.pts = None  # Let encoder assign PTS
+                    for packet in out_audio_stream.encode(frame):
+                        output.mux(packet)
+                for packet in out_audio_stream.encode():
+                    output.mux(packet)
+
+                input_video.close()
+                input_audio.close()
+                output.close()
+                shutil.move(out_path, video_path)
+            return True
+        except Exception as e:
+            logger.warning("Audio mux failed: %s", e)
+            return False
+
     def set_lora_adapter(self,
                          lora_nickname: str,
                          lora_path: str | None = None) -> None:

fastvideo/fastvideo_args.py

@@ -133,7 +133,7 @@ class FastVideoArgs:
     # CPU offload parameters
     dit_cpu_offload: bool = True
     use_fsdp_inference: bool = False
-    dit_layerwise_offload: bool = False
+    dit_layerwise_offload: bool = True
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
     vae_cpu_offload: bool = True
@@ -166,6 +166,14 @@ class FastVideoArgs:
     # Prompt text file for batch processing
     prompt_txt: str | None = None
 
+    # LTX-2 VAE tiling overrides
+    ltx2_vae_tiling: bool | None = None
+    ltx2_vae_spatial_tile_size_in_pixels: int | None = None
+    ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None
+    ltx2_vae_temporal_tile_size_in_frames: int | None = None
+    ltx2_vae_temporal_tile_overlap_in_frames: int | None = None
+    ltx2_initial_latent_path: str | None = None
+
     # model paths for correct deallocation
     model_paths: dict[str, str] = field(default_factory=dict)
     model_loaded: dict[str, bool] = field(default_factory=lambda: {
@@ -203,8 +211,44 @@ class FastVideoArgs:
                 logger.error("Failed to load V-MoBA config from %s: %s",
                              self.moba_config_path, e)
                 raise
+        self._apply_ltx2_vae_overrides()
         self.check_fastvideo_args()
 
+    def _apply_ltx2_vae_overrides(self) -> None:
+        if self.pipeline_config is None:
+            return
+        vae_config = self.pipeline_config.vae_config
+        has_any = any(value is not None for value in (
+            self.ltx2_vae_spatial_tile_size_in_pixels,
+            self.ltx2_vae_spatial_tile_overlap_in_pixels,
+            self.ltx2_vae_temporal_tile_size_in_frames,
+            self.ltx2_vae_temporal_tile_overlap_in_frames,
+        ))
+        if self.ltx2_vae_tiling is not None and hasattr(self.pipeline_config,
+                                                        "vae_tiling"):
+            self.pipeline_config.vae_tiling = self.ltx2_vae_tiling
+        elif has_any and hasattr(self.pipeline_config, "vae_tiling"):
+            self.pipeline_config.vae_tiling = True
+
+        if hasattr(vae_config, "ltx2_spatial_tile_size_in_pixels"
+                   ) and self.ltx2_vae_spatial_tile_size_in_pixels is not None:
+            vae_config.ltx2_spatial_tile_size_in_pixels = (
+                self.ltx2_vae_spatial_tile_size_in_pixels)
+        if hasattr(
+                vae_config, "ltx2_spatial_tile_overlap_in_pixels"
+        ) and self.ltx2_vae_spatial_tile_overlap_in_pixels is not None:
+            vae_config.ltx2_spatial_tile_overlap_in_pixels = (
+                self.ltx2_vae_spatial_tile_overlap_in_pixels)
+        if hasattr(vae_config, "ltx2_temporal_tile_size_in_frames"
+                   ) and self.ltx2_vae_temporal_tile_size_in_frames is not None:
+            vae_config.ltx2_temporal_tile_size_in_frames = (
+                self.ltx2_vae_temporal_tile_size_in_frames)
+        if hasattr(
+                vae_config, "ltx2_temporal_tile_overlap_in_frames"
+        ) and self.ltx2_vae_temporal_tile_overlap_in_frames is not None:
+            vae_config.ltx2_temporal_tile_overlap_in_frames = (
+                self.ltx2_vae_temporal_tile_overlap_in_frames)
+
     @staticmethod
     def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
         # Model and path configuration
@@ -325,6 +369,44 @@ class FastVideoArgs:
             "Path to a text file containing prompts (one per line) for batch processing",
         )
 
+        # LTX-2 VAE tiling overrides
+        parser.add_argument(
+            "--ltx2-vae-tiling",
+            action=StoreBoolean,
+            default=FastVideoArgs.ltx2_vae_tiling,
+            help="Enable LTX-2 VAE tiling overrides.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-spatial-tile-size-in-pixels",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_spatial_tile_size_in_pixels,
+            help="LTX-2 VAE spatial tile size in pixels.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-spatial-tile-overlap-in-pixels",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_spatial_tile_overlap_in_pixels,
+            help="LTX-2 VAE spatial tile overlap in pixels.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-temporal-tile-size-in-frames",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_temporal_tile_size_in_frames,
+            help="LTX-2 VAE temporal tile size in frames.",
+        )
+        parser.add_argument(
+            "--ltx2-vae-temporal-tile-overlap-in-frames",
+            type=int,
+            default=FastVideoArgs.ltx2_vae_temporal_tile_overlap_in_frames,
+            help="LTX-2 VAE temporal tile overlap in frames.",
+        )
+        parser.add_argument(
+            "--ltx2-initial-latent-path",
+            type=str,
+            default=FastVideoArgs.ltx2_initial_latent_path,
+            help="Path to load/save a precomputed LTX-2 initial latent.",
+        )
+
         # LoRA parameters (inference-time adapter loading)
         parser.add_argument(
             "--lora-path",
@@ -748,6 +830,7 @@ class TrainingArgs(FastVideoArgs):
     precedence.
     """
     data_path: str = ""
+    data_path_2: str | None = None
     dataloader_num_workers: int = 0
     num_height: int = 0
     num_width: int = 0
@@ -777,6 +860,7 @@ class TrainingArgs(FastVideoArgs):
     validation_sampling_steps: str = ""
     validation_guidance_scale: str = ""
     validation_steps: float = 0.0
+    visualization_steps: float = 0.0
     log_validation: bool = False
     trackers: list[str] = dataclasses.field(default_factory=list)
     tracker_project_name: str = ""
@@ -792,6 +876,7 @@ class TrainingArgs(FastVideoArgs):
     num_train_epochs: int = 0
     max_train_steps: int = 0
     gradient_accumulation_steps: int = 0
+    optimizer_type: str = "adamw"
     learning_rate: float = 0.0
     scale_lr: bool = False
     lr_scheduler: str = "constant"
@@ -861,6 +946,8 @@ class TrainingArgs(FastVideoArgs):
     same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
     last_step_only: bool = False  # Only use the last timestep for training
     context_noise: int = 0  # Context noise level for cache updates
+    use_context_forcing: bool = False
+    use_ode_init: bool = False
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -916,6 +1003,10 @@ class TrainingArgs(FastVideoArgs):
                             type=str,
                             required=True,
                             help="Path to parquet files")
+        parser.add_argument("--data-path-2",
+                            type=str,
+                            required=False,
+                            help="Path to parquet files")
         parser.add_argument("--dataloader-num-workers",
                             type=int,
                             required=True,
@@ -1009,6 +1100,9 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--validation-steps",
                             type=float,
                             help="Number of validation steps")
+        parser.add_argument("--visualization-steps",
+                            type=float,
+                            help="Number of visualization steps")
         parser.add_argument("--log-validation",
                             action=StoreBoolean,
                             help="Whether to log validation results")
@@ -1057,6 +1151,10 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--gradient-accumulation-steps",
                             type=int,
                             help="Number of steps to accumulate gradients")
+        parser.add_argument("--optimizer-type",
+                            type=str,
+                            choices=["adamw", "muon"],
+                            help="Optimizer type")
         parser.add_argument("--learning-rate",
                             type=float,
                             required=True,
@@ -1283,6 +1381,12 @@ class TrainingArgs(FastVideoArgs):
                             type=int,
                             default=TrainingArgs.context_noise,
                             help="Context noise level for cache updates")
+        parser.add_argument("--use-context-forcing",
+                            action=StoreBoolean,
+                            help="Whether to use context forcing")
+        parser.add_argument("--use-ode-init",
+                            action=StoreBoolean,
+                            help="Whether to use ODE init")
 
         return parser
 

fastvideo/hooks/hooks.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+from typing import Any
+from torch import nn
+
+
+class ForwardHook:
+    """
+    Base class for forward hooks.
+    Hooks are used in the way:
+        modified_args, modified_kwargs = hook.pre_forward(module, *args, **kwargs)
+        output = module.forward(*modified_args, **modified_kwargs)
+        modified_output = hook.post_forward(module, output)
+    """
+
+    @classmethod
+    def name(cls) -> str:
+        raise NotImplementedError
+
+    def on_attach(self, module: nn.Module):  # noqa: B027
+        """Called once when the hook is attached to the module."""
+        pass
+
+    def on_detach(self, module: nn.Module):  # noqa: B027
+        """
+        Called once when the hook is detached from the module.
+        Note: this function is not guaranteed to be called if the module is
+              deleted before the hook is detached.
+        """
+        pass
+
+    def pre_forward(self, module: nn.Module, *args,
+                    **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        """Called before the module's forward method is executed."""
+        return args, kwargs
+
+    def post_forward(self, module: nn.Module, output: Any) -> Any:
+        """Called after the module's forward method is executed."""
+        return output
+
+
+class ModuleHookManager:
+    module_hook_attribute = "_hook_manager"
+
+    def __init__(self, module: nn.Module):
+        self.module = module
+        self.forward_hooks: dict[str, ForwardHook] = {}
+        self.original_forward = module.forward
+
+    @classmethod
+    def get_from(cls, module: nn.Module) -> "ModuleHookManager | None":
+        if hasattr(module, cls.module_hook_attribute):
+            return getattr(module, cls.module_hook_attribute)
+        return None
+
+    @classmethod
+    def get_from_or_default(cls, module: nn.Module) -> "ModuleHookManager":
+        if not hasattr(module, cls.module_hook_attribute):
+            setattr(module, cls.module_hook_attribute, cls(module))
+
+            def forward_hook_wrapper(mod: nn.Module, *args, **kwargs):
+                manager: ModuleHookManager = getattr(mod,
+                                                     cls.module_hook_attribute)
+                for hook in manager.forward_hooks.values():
+                    args, kwargs = hook.pre_forward(mod, *args, **kwargs)
+                output = manager.original_forward(*args, **kwargs)
+                for hook in reversed(manager.forward_hooks.values()):
+                    output = hook.post_forward(mod, output)
+                return output
+
+            module.forward = functools.partial(forward_hook_wrapper, module)
+
+        return getattr(module, cls.module_hook_attribute)
+
+    @staticmethod
+    def remove_from_manager(module: nn.Module) -> None:
+        if hasattr(module, ModuleHookManager.module_hook_attribute):
+            manager: ModuleHookManager = getattr(
+                module, ModuleHookManager.module_hook_attribute)
+            module.forward = manager.original_forward
+            delattr(module, ModuleHookManager.module_hook_attribute)
+
+    def _check_manager_attached(self) -> None:
+        if not hasattr(self.module, self.module_hook_attribute):
+            raise ValueError("ModuleHookManager is not attached to the module.")
+        if getattr(self.module, self.module_hook_attribute) is not self:
+            raise ValueError(
+                "ModuleHookManager attached to the module is different.")
+
+    def append_forward_hook(self, hook: ForwardHook):
+        self._check_manager_attached()
+        if hook.name() in self.forward_hooks:
+            raise ValueError(
+                f"Hook with name {hook.name()} is already registered.")
+        # after python 3.7, dicts maintain insertion order
+        self.forward_hooks[hook.name()] = hook
+        hook.on_attach(self.module)
+
+    def replace_forward_hook(self,
+                             hook_name: str,
+                             new_hook: ForwardHook,
+                             run_on_attach: bool = True):
+        self._check_manager_attached()
+        if hook_name not in self.forward_hooks:
+            raise ValueError(f"No hook with name {hook_name} found.")
+        old_hook = self.forward_hooks[hook_name]
+        if run_on_attach:
+            old_hook.on_detach(self.module)
+        self.forward_hooks[hook_name] = new_hook
+        new_hook.on_attach(self.module)
+
+    def remove_forward_hook(self, hook_name: str, run_detach: bool = True):
+        self._check_manager_attached()
+        if hook_name not in self.forward_hooks:
+            raise ValueError(f"No hook with name {hook_name} found.")
+        if run_detach:
+            self.forward_hooks[hook_name].on_detach(self.module)
+        del self.forward_hooks[hook_name]
+
+    def get_forward_hook(self, hook_name: str) -> ForwardHook | None:
+        return self.forward_hooks.get(hook_name, None)

fastvideo/hooks/layerwise_offload.py

@@ -0,0 +1,164 @@
+from contextlib import contextmanager
+from typing import Any
+import torch
+from torch import nn
+from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def _tensor_placeholder(tensor: torch.Tensor,
+                        device: torch.device) -> torch.Tensor:
+    """Create a rank-preserving empty placeholder on the specified device."""
+    shape = (0, ) if tensor.ndim <= 0 else (0, ) * tensor.ndim
+    return torch.empty(shape, device=device, dtype=tensor.dtype)
+
+
+class LayerwiseOffloadState:
+
+    def __init__(
+        self,
+        async_copy_stream: torch.cuda.Stream,
+        device: torch.device,
+        next_state: "LayerwiseOffloadState | None" = None,
+    ) -> None:
+        self.async_copy_stream = async_copy_stream
+        self.next_state = next_state
+        self.gpu_named_parameters: dict[str, torch.Tensor] = {}
+        self.cpu_named_parameters: dict[str, torch.Tensor] = {}
+        self.module_ref: nn.Module = None  # type: ignore
+        self.device: torch.device = device
+
+    def _will_offload(self, name: str) -> bool:
+        return True
+
+    @torch.compiler.disable
+    def on_init(self, module: nn.Module):
+        self.module_ref = module
+        for name, param in self.module_ref.named_parameters():
+            if self._will_offload(name):
+                self.cpu_named_parameters[name] = (
+                    param.data.detach().to("cpu").pin_memory())
+                param.data = _tensor_placeholder(param.data, self.device)
+
+    @torch.compiler.disable
+    def wait_and_replace_params(self):
+        torch.cuda.current_stream().wait_stream(self.async_copy_stream)
+        # now gpu_named_parameters are ready
+        for name, param in self.module_ref.named_parameters():
+            if not self._will_offload(name):
+                continue
+            if name not in self.gpu_named_parameters:
+                # first load with blocking load
+                self.gpu_named_parameters[name] = self.cpu_named_parameters[
+                    name].to(self.device)
+            param.data = self.gpu_named_parameters[name]
+
+    @torch.compiler.disable
+    def prefetch_params(self):
+        compute_stream = torch.cuda.current_stream()
+        with torch.cuda.stream(self.async_copy_stream):
+            for name, param in self.module_ref.named_parameters():
+                if not self._will_offload(name):
+                    continue
+                assert name not in self.gpu_named_parameters
+                gpu_param = self.cpu_named_parameters[name].to(
+                    self.device, non_blocking=True)
+                gpu_param.record_stream(
+                    compute_stream
+                )  # ensure tensor will not be freed until forward is completed
+                self.gpu_named_parameters[name] = gpu_param
+
+    @torch.compiler.disable
+    def release_gpu_params(self):
+        for name, param in self.module_ref.named_parameters():
+            if self._will_offload(name):
+                param.data = _tensor_placeholder(param.data, self.device)
+                del self.gpu_named_parameters[name]
+        assert len(self.gpu_named_parameters) == 0
+
+
+class LayerwiseOffloadHook(ForwardHook):
+    """A hook that enables layerwise CPU offloading during forward pass."""
+
+    def __init__(self, state: LayerwiseOffloadState) -> None:
+        self.state = state
+
+    def on_attach(self, module: nn.Module):
+        self.state.on_init(module)  # pyright: ignore
+
+    def on_detach(self, module: nn.Module):
+        named_parameters = dict(module.named_parameters())
+        for name, cpu_tensor in self.state.cpu_named_parameters.items():
+            if name not in self.state.gpu_named_parameters:
+                if name in named_parameters:
+                    named_parameters[name].data = cpu_tensor.to(
+                        device=self.state.device)
+                else:
+                    logger.warning(
+                        "Parameter {} not found in module during detachment.",
+                        name,
+                    )
+
+    @classmethod
+    def name(cls) -> str:
+        return "LayerwiseOffloadHook"
+
+    def pre_forward(self, module: nn.Module, *args, **kwargs):
+        self.state.wait_and_replace_params()  # pyright: ignore
+        if self.state.next_state is not None:
+            self.state.next_state.prefetch_params()  # pyright: ignore
+        return args, kwargs
+
+    def post_forward(self, module: torch.nn.Module, output: Any):
+        self.state.release_gpu_params()  # pyright: ignore
+        return output
+
+    @contextmanager
+    def mutate_params_scope(self):
+        try:
+            # load params to GPU and keep them there
+            self.state.wait_and_replace_params()  # pyright: ignore
+            yield
+        finally:
+            # instead of releasing, we should overwrite the original params since they have been modified
+            self.state.cpu_named_parameters.clear()
+            self.state.gpu_named_parameters.clear()
+            self.state.on_init(self.state.module_ref)  # pyright: ignore
+
+
+def enable_layerwise_offload(model: nn.Module, is_replace: bool = False):
+    if torch.cuda.is_available():
+        device = torch.device("cuda", torch.cuda.current_device())
+    else:
+        logger.warning(
+            "CUDA is not available. Layerwise offloading is disabled.")
+        return
+    state_list = []
+    async_stream = torch.cuda.Stream()
+    for name, submodule in model.named_children():
+        if isinstance(submodule, nn.ModuleList):
+            for idx, module_entry in enumerate(submodule):
+                state = LayerwiseOffloadState(async_copy_stream=async_stream,
+                                              device=device)
+                state_list.append(state)
+                hook_mgr = ModuleHookManager.get_from_or_default(module_entry)
+                hook = LayerwiseOffloadHook(state)
+                if is_replace:
+                    existing_hook = hook_mgr.forward_hooks.get(hook.name())
+                    if existing_hook is not None:
+                        hook_mgr.replace_forward_hook(hook.name(), hook)
+                    else:
+                        raise AssertionError(
+                            f"Expect hook exists in {name} for replacement.")
+                else:
+                    hook_mgr.append_forward_hook(hook)
+            break
+    if len(state_list) == 0:
+        raise ValueError(
+            "No nn.ModuleList found in the model for layerwise offloading.")
+
+    # circular linking of states
+    for i in range(len(state_list)):
+        state_list[i].next_state = state_list[(i + 1) % len(state_list)]

fastvideo/layers/layernorm.py

@@ -168,9 +168,15 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
         else:
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
-    def forward(self, residual: torch.Tensor, x: torch.Tensor,
-                gate: torch.Tensor | int, shift: torch.Tensor,
-                scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(
+        self,
+        residual: torch.Tensor,
+        x: torch.Tensor,
+        gate: torch.Tensor | int,
+        shift: torch.Tensor,
+        scale: torch.Tensor,
+        convert_modulation_dtype: bool = False
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Apply gated residual connection, followed by layernorm and 
         scale/shift in a single fused operation.
@@ -205,6 +211,11 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
 
         # Apply normalization
         normalized = self.norm(residual_output)
+
+        if convert_modulation_dtype:
+            scale = scale.to(normalized.dtype)
+            shift = shift.to(normalized.dtype)
+
         # Apply scale and shift
         if isinstance(scale, torch.Tensor) and scale.dim() == 4:
             # scale.shape: [batch_size, num_frames, 1, inner_dim]
@@ -254,14 +265,21 @@ class LayerNormScaleShift(nn.Module):
         else:
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
-    def forward(self, x: torch.Tensor, shift: torch.Tensor,
-                scale: torch.Tensor) -> torch.Tensor:
+    def forward(self,
+                x: torch.Tensor,
+                shift: torch.Tensor,
+                scale: torch.Tensor,
+                convert_modulation_dtype: bool = False) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""
         # x.shape: [batch_size, seq_len, inner_dim]
         normalized = self.norm(x)
         if self.compute_dtype == torch.float32:
             normalized = normalized.float()
 
+        if convert_modulation_dtype:
+            scale = scale.to(normalized.dtype)
+            shift = shift.to(normalized.dtype)
+
         if scale.dim() == 4:
             # scale.shape: [batch_size, num_frames, 1, inner_dim]
             num_frames = scale.shape[1]

fastvideo/models/audio/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.models.audio.ltx2_audio_vae import (
+    LTX2AudioDecoder,
+    LTX2AudioEncoder,
+    LTX2Vocoder,
+)
+
+__all__ = ["LTX2AudioEncoder", "LTX2AudioDecoder", "LTX2Vocoder"]

fastvideo/models/dits/causal_hunyuanvideo15.py

@@ -0,0 +1,808 @@
+# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional, List
+import math
+
+import torch
+# import torch._dynamo
+# torch._dynamo.config.cache_size_limit = 128
+# try:
+#     torch._dynamo.config.recompile_limit = 128
+# except AttributeError:
+#     pass
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from torch.nn.attention.flex_attention import BlockMask
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="default")
+import torch.distributed as dist
+
+from fastvideo.attention import LocalAttention
+from fastvideo.forward_context import set_forward_context
+from fastvideo.configs.models.dits import HunyuanVideo15Config
+from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
+                                        ScaleResidualLayerNormScaleShift)
+from fastvideo.layers.linear import ReplicatedLinear
+# TODO(will-PY-refactor): RMSNorm ....
+from fastvideo.layers.mlp import MLP
+from fastvideo.layers.rotary_embedding import get_rotary_pos_embed, _apply_rotary_emb
+from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
+                                               unpatchify)
+from fastvideo.models.dits.base import CachableDiT
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.logger import init_logger
+
+from fastvideo.models.dits.hunyuanvideo15 import (
+    HunyuanRMSNorm, 
+    HunyuanVideo15TimeEmbedding, 
+    HunyuanVideo15ByT5TextProjection,
+    HunyuanVideo15ImageProjection,
+    SingleTokenRefiner,
+    FinalLayer)
+
+logger = init_logger(__name__)
+
+class MMDoubleStreamBlock(nn.Module):
+    """
+    A multimodal DiT block with separate modulation for text and image/video,
+    using distributed attention and linear layers.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        dtype: torch.dtype | None = None,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.deterministic = False
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = hidden_size // num_attention_heads
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        # Image modulation components
+        self.img_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.img_mod",
+        )
+
+        # Fused operations for image stream
+        self.img_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.img_mlp_residual = ScaleResidual()
+
+        # Image attention components
+        self.img_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype,
+                                             prefix=f"{prefix}.img_attn_qkv")
+
+        self.img_attn_q_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
+        self.img_attn_k_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
+
+        self.img_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype,
+                                              prefix=f"{prefix}.img_attn_proj")
+
+        self.img_mlp = MLP(hidden_size,
+                           mlp_hidden_dim,
+                           bias=True,
+                           dtype=dtype,
+                           prefix=f"{prefix}.img_mlp")
+
+        # Text modulation components
+        self.txt_mod = ModulateProjection(
+            hidden_size,
+            factor=6,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.txt_mod",
+        )
+
+        # Fused operations for text stream
+        self.txt_attn_norm = LayerNormScaleShift(hidden_size,
+                                                 norm_type="layer",
+                                                 elementwise_affine=False,
+                                                 dtype=dtype)
+        self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            elementwise_affine=False,
+            dtype=dtype)
+        self.txt_mlp_residual = ScaleResidual()
+
+        # Text attention components
+        self.txt_attn_qkv = ReplicatedLinear(hidden_size,
+                                             hidden_size * 3,
+                                             bias=True,
+                                             params_dtype=dtype)
+
+        # QK norm layers for text
+        self.txt_attn_q_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
+        self.txt_attn_k_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
+
+        self.txt_attn_proj = ReplicatedLinear(hidden_size,
+                                              hidden_size,
+                                              bias=True,
+                                              params_dtype=dtype)
+
+        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
+
+        self.max_attention_size = 21 * 1590 if local_attn_size == -1 else local_attn_size * 1590
+
+        self.attn = LocalAttention(
+            num_heads=self.num_attention_heads,
+            head_size=self.head_dim,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn"
+        )
+
+    def forward_txt(
+        self,
+        txt: torch.Tensor,
+        vec: torch.Tensor,
+        cache_txt: bool = False,
+    ):
+        txt_mod_outputs = self.txt_mod(vec)
+        (
+            txt_attn_shift,
+            txt_attn_scale,
+            txt_attn_gate,
+            txt_mlp_shift,
+            txt_mlp_scale,
+            txt_mlp_gate,
+        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)
+
+        # Prepare text for attention using fused operation
+        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)
+
+        # Get QKV for text
+        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
+        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]
+
+        # Split QKV
+        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3,
+                               self.num_attention_heads, -1)
+        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :,
+                                                                          2]
+        # Apply QK-Norm if needed
+        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
+        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
+
+        t_kv = {}
+        if cache_txt:
+            t_kv["k_txt"] = txt_k
+            t_kv["v_txt"] = txt_v
+
+        txt_attn = self.attn(txt_q, txt_k, txt_v)
+        # Process text attention output
+        txt_attn_out, _ = self.txt_attn_proj(
+            txt_attn.reshape(batch_size, text_seq_len, -1))
+
+        # Use fused operation for residual connection, normalization, and modulation
+        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(
+            txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale)
+
+        # Process text MLP
+        txt_mlp_out = self.txt_mlp(txt_mlp_input)
+        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
+
+        return txt, t_kv
+
+    def forward_vision(
+        self,
+        img: torch.Tensor,
+        vec: torch.Tensor,
+        freqs_cis: tuple,
+        block_mask: BlockMask,
+        kv_cache: dict | None = None,
+        txt_kv_cache: list | None = None,
+        current_start: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Process modulation vectors
+        if vec.dim() == 3:
+            img_mod_outputs = self.img_mod(vec).unflatten(dim=-1, sizes=(6, -1))
+            (
+                img_attn_shift,
+                img_attn_scale,
+                img_attn_gate,
+                img_mlp_shift,
+                img_mlp_scale,
+                img_mlp_gate,
+            ) = torch.chunk(img_mod_outputs, 6, dim=2)
+        else:
+            img_mod_outputs = self.img_mod(vec)
+            (
+                img_attn_shift,
+                img_attn_scale,
+                img_attn_gate,
+                img_mlp_shift,
+                img_mlp_scale,
+                img_mlp_gate,
+            ) = torch.chunk(img_mod_outputs, 6, dim=-1)
+
+        # Prepare image for attention using fused operation
+        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
+        # Get QKV for image
+        img_qkv, _ = self.img_attn_qkv(img_attn_input)
+        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]
+
+        # Split QKV
+        img_qkv = img_qkv.view(batch_size, image_seq_len, 3,
+                               self.num_attention_heads, -1)
+        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :,
+                                                                          2]
+
+        # Apply QK-Norm if needed
+        img_q = self.img_attn_q_norm(img_q).to(img_v)
+        img_k = self.img_attn_k_norm(img_k).to(img_v)
+
+        # 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:
+            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)
+            # Padding for flex attention
+            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+            padded_kv_length = math.ceil(k.shape[1] / 128) * 128 - k.shape[1]
+            padded_roped_query = torch.cat(
+                [q,
+                    torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
+                                device=q.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            padded_roped_key = torch.cat(
+                [k, torch.zeros([k.shape[0], padded_kv_length, k.shape[2], k.shape[3]],
+                                        device=k.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            padded_v = torch.cat(
+                [v, torch.zeros([v.shape[0], padded_kv_length, v.shape[2], v.shape[3]],
+                                device=v.device, dtype=v.dtype)],
+                dim=1
+            )
+
+            img_attn = flex_attention(
+                query=padded_roped_query.transpose(2, 1),
+                key=padded_roped_key.transpose(2, 1),
+                value=padded_v.transpose(2, 1),
+                block_mask=block_mask
+            )[:, :, :-padded_length].transpose(2, 1)
+
+            assert img_attn.shape[1] == image_seq_len
+            updated_kv_cache = None
+        else:
+            current_end = current_start + img_q.shape[1]
+            num_new_tokens = img_q.shape[1]
+            sink_tokens = self.sink_size * 1590
+            # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
+            kv_cache_size = self.max_attention_size
+            
+            # Clone cache to avoid in-place modification during gradient checkpointing
+            k_cache = kv_cache["k"].clone()
+            v_cache = kv_cache["v"].clone()
+            
+            if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
+                    num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
+                # Calculate the number of new tokens added in this step
+                # Shift existing cache content left to discard oldest tokens
+                # Clone the source slice to avoid overlapping memory error
+                num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+                num_rolled_tokens = kv_cache_size - num_new_tokens - sink_tokens
+                k_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                    k_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                v_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+                    v_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+                # Insert the new keys/values at the end
+                local_end_index = kv_cache["local_end_index"].item() + current_end - \
+                    kv_cache["global_end_index"].item() - num_evicted_tokens
+                local_start_index = local_end_index - num_new_tokens
+                assert local_end_index == self.max_attention_size
+            else:
+                # 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
+
+            assert local_start_index >= 0
+            q = img_q
+            k = torch.cat([k_cache[:, :local_start_index], img_k, txt_kv_cache["k_txt"]], dim=1)
+            v = torch.cat([v_cache[:, :local_start_index], img_v, txt_kv_cache["v_txt"]], dim=1)
+            img_attn = self.attn(q, k, v)
+
+            k_cache[:, local_start_index:local_end_index] = img_k
+            v_cache[:, local_start_index:local_end_index] = img_v
+            
+            updated_kv_cache = {
+                "k": k_cache,
+                "v": v_cache,
+                "global_end_index": torch.tensor([current_end], dtype=torch.long, device=k_cache.device),
+                "local_end_index": torch.tensor([local_end_index], dtype=torch.long, device=k_cache.device)
+            }
+        
+        img_attn_out, _ = self.img_attn_proj(
+            img_attn.view(batch_size, image_seq_len, -1))
+        # Use fused operation for residual connection, normalization, and modulation
+        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(
+            img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale)
+
+        # Process image MLP
+        img_mlp_out = self.img_mlp(img_mlp_input)
+        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
+
+        return img, updated_kv_cache
+
+    def forward(
+        self,
+        txt_inference=False,
+        vision_inference=False,
+        **kwargs
+    ):
+        if txt_inference:
+            return self.forward_txt(**kwargs)
+        elif vision_inference:
+            return self.forward_vision(**kwargs)
+        else:
+            raise ValueError("txt_inference and vision_inference cannot be both False")
+
+
+class CausalHunyuanVideo15Transformer3DModel(CachableDiT):
+    r"""
+    A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5).
+    """
+
+    # shard single stream, double stream blocks, and refiner_blocks
+    _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions
+    _compile_conditions = HunyuanVideo15Config()._compile_conditions
+    _supported_attention_backends = HunyuanVideo15Config(
+    )._supported_attention_backends
+    param_names_mapping = HunyuanVideo15Config().param_names_mapping
+    reverse_param_names_mapping = HunyuanVideo15Config(
+    ).reverse_param_names_mapping
+    lora_param_names_mapping = HunyuanVideo15Config().lora_param_names_mapping
+
+    def __init__(
+        self,
+        config: HunyuanVideo15Config,
+        hf_config: dict[str, Any],
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.num_channels_latents
+        self.out_channels = config.out_channels or config.in_channels
+        self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size)
+
+        # 1. Latent and condition embedders
+        self.img_in = PatchEmbed(self.patch_size, 
+                                    config.in_channels, 
+                                    self.hidden_size,
+                                    prefix=f"{config.prefix}.img_in")
+        self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size)
+
+        self.txt_in = SingleTokenRefiner(config.text_embed_dim,
+                                                   self.hidden_size,
+                                                   config.num_attention_heads,
+                                                   depth=config.num_refiner_layers,
+                                                   dtype=None,
+                                                   prefix=f"{config.prefix}.txt_in")
+
+        self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size)
+
+        self.time_in = HunyuanVideo15TimeEmbedding(self.hidden_size, use_meanflow=config.use_meanflow)
+
+        self.cond_type_embed = nn.Embedding(3, self.hidden_size)
+
+        # 3. Dual stream transformer blocks
+
+        self.double_blocks = nn.ModuleList(
+            [
+                MMDoubleStreamBlock(
+                    hidden_size=self.hidden_size,
+                    num_attention_heads=config.num_attention_heads,
+                    mlp_ratio=config.mlp_ratio,
+                    local_attn_size=config.local_attn_size,
+                    sink_size=config.sink_size,
+                    dtype=None,
+                    supported_attention_backends=self._supported_attention_backends,
+                    prefix=f"{config.prefix}.double_blocks.{i}"
+                )
+                for i in range(config.num_layers)
+            ]
+        )
+
+        # 5. Output projection
+        self.final_layer = FinalLayer(self.hidden_size,
+                                self.patch_size,
+                                self.out_channels,
+                                prefix=f"{config.prefix}.final_layer")
+
+        self.gradient_checkpointing = False
+
+        self.num_frame_per_block = config.num_frames_per_block
+        self.local_attn_size = config.local_attn_size
+        self.block_mask = None
+
+        self.__post_init__()
+
+    def get_text_and_mask(
+        self,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_2: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        encoder_attention_mask_2: torch.Tensor,
+        encoder_hidden_states_image: torch.Tensor,
+        timestep: torch.Tensor,
+    ):
+        batch_size, txt_seq_len = encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]
+        # qwen text embedding
+        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask)
+
+        encoder_hidden_states_cond_emb = self.cond_type_embed(
+            torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb
+
+        # byt5 text embedding
+        encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2)
+
+        encoder_hidden_states_2_cond_emb = self.cond_type_embed(
+            torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb
+
+        # image embed
+        encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
+        is_t2v = torch.all(encoder_hidden_states_image == 0)
+        if is_t2v:
+            encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
+            encoder_attention_mask_3 = torch.zeros(
+                (batch_size, encoder_hidden_states_3.shape[1]),
+                dtype=encoder_attention_mask.dtype,
+                device=encoder_attention_mask.device,
+            )
+        else:
+            encoder_attention_mask_3 = torch.ones(
+                (batch_size, encoder_hidden_states_3.shape[1]),
+                dtype=encoder_attention_mask.dtype,
+                device=encoder_attention_mask.device,
+            )
+        encoder_hidden_states_3_cond_emb = self.cond_type_embed(
+            2
+            * torch.ones_like(
+                encoder_hidden_states_3[:, :, 0],
+                dtype=torch.long,
+            )
+        )
+        encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
+
+        # reorder and combine text tokens: combine valid tokens first, then padding
+        encoder_attention_mask = encoder_attention_mask.bool()
+        encoder_attention_mask_2 = encoder_attention_mask_2.bool()
+        encoder_attention_mask_3 = encoder_attention_mask_3.bool()
+        new_encoder_hidden_states = []
+        new_encoder_attention_mask = []
+
+        for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
+            encoder_hidden_states,
+            encoder_attention_mask,
+            encoder_hidden_states_2,
+            encoder_attention_mask_2,
+            encoder_hidden_states_3,
+            encoder_attention_mask_3,
+        ):
+            # Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm]
+            new_encoder_hidden_states.append(
+                torch.cat(
+                    [
+                        image[image_mask],  # valid image
+                        text_2[text_mask_2],  # valid byt5
+                        text[text_mask],  # valid mllm
+                        image[~image_mask],  # invalid image (zeroed)
+                        torch.zeros_like(text_2[~text_mask_2]),  # invalid byt5 (zeroed)
+                        torch.zeros_like(text[~text_mask]),  # invalid mllm (zeroed)
+                    ],
+                    dim=0,
+                )
+            )
+            # Apply same reordering to attention masks
+            new_encoder_attention_mask.append(
+                torch.cat(
+                    [
+                        image_mask[image_mask],
+                        text_mask_2[text_mask_2],
+                        text_mask[text_mask],
+                        image_mask[~image_mask],
+                        text_mask_2[~text_mask_2],
+                        text_mask[~text_mask],
+                    ],
+                    dim=0,
+                )
+            )
+
+        encoder_hidden_states = torch.stack(new_encoder_hidden_states)
+        encoder_attention_mask = torch.stack(new_encoder_attention_mask)
+        assert encoder_hidden_states.shape[0] == 1
+        return encoder_hidden_states, encoder_attention_mask
+
+    
+    @staticmethod
+    def _prepare_blockwise_causal_attn_mask(
+        device: torch.device | str, num_frames: int = 21,
+        frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-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
+        """
+        total_length = num_frames * frame_seqlen
+        total_kv_length = total_length + text_seq_len
+
+        total_length_tensor = torch.tensor(total_length, device=device)
+        total_kv_length_tensor = torch.tensor(total_kv_length, device=device)
+
+        # 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
+
+        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
+        frame_indices = torch.arange(
+            # start=frame_seqlen,
+            start=0,
+            end=total_length,
+            step=frame_seqlen * num_frame_per_block,
+            device=device
+        )
+        # frame_indices = torch.cat([torch.tensor([0], device=device), frame_indices])
+
+        for i, tmp in enumerate(frame_indices):
+            # if i == 0:
+            #     ends[tmp:tmp + frame_seqlen] = tmp + frame_seqlen
+            # else:
+            ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \
+                frame_seqlen * num_frame_per_block
+
+        def attention_mask(b, h, q_idx, kv_idx):
+            if local_attn_size == -1:
+                return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) | ((kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor))
+            else:
+                return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) | ((kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor))
+            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional 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 not dist.is_initialized() or dist.get_rank() == 0:
+        #     print(
+        #         f" cache a block wise causal mask with block size of {num_frame_per_block} frames")
+        #     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_length + 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
+
+    def forward_txt(
+        self,
+        encoder_hidden_states: List[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: Optional[List[torch.Tensor]] = None,
+        encoder_attention_mask: Optional[List[torch.Tensor]] = None,
+        guidance: Optional[torch.Tensor] = None,
+        timestep_r: Optional[torch.LongTensor] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        cache_txt: bool = False,
+        **kwargs
+    ):
+        # Check that the timestep is only consisted of 0s
+        assert torch.all(timestep == 0), "Timestep for txt must be only consisted of 0s"
+
+        if cache_txt:
+            _kv_cache_new = []
+            transformer_num_layers = len(self.double_blocks)
+            for _ in range(transformer_num_layers):
+                _kv_cache_new.append(
+                    {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None}
+                )
+
+        encoder_hidden_states_image = encoder_hidden_states_image[0]
+        encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
+        encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask
+
+        # 2. Conditional embeddings
+        if timestep.dim() == 2:
+            ts_seq_len = timestep.shape[1]
+            temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
+        else:
+            temb = self.time_in(timestep, timestep_r=timestep_r)
+        # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim]
+
+        encoder_hidden_states, encoder_attention_mask = self.get_text_and_mask(
+            encoder_hidden_states,
+            encoder_hidden_states_2,
+            encoder_attention_mask,
+            encoder_attention_mask_2,
+            encoder_hidden_states_image,
+            timestep
+        )
+        
+        encoder_hidden_states = encoder_hidden_states[encoder_attention_mask.bool().to(encoder_hidden_states.device)].unsqueeze(0)
+
+        # 4. Transformer blocks
+        for index, block in enumerate(self.double_blocks):
+            encoder_hidden_states, t_kv = block(
+                txt_inference=True,
+                vision_inference=False,
+                txt=encoder_hidden_states,
+                vec=temb,
+                cache_txt=cache_txt,
+            )
+
+            if cache_txt:
+                _kv_cache_new[index]["k_txt"] = t_kv["k_txt"]
+                _kv_cache_new[index]["v_txt"] = t_kv["v_txt"]
+
+        if cache_txt:
+            return _kv_cache_new
+
+    def forward_vision(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.LongTensor,
+        guidance: Optional[torch.Tensor] = None,
+        timestep_r: Optional[torch.LongTensor] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        kv_cache: dict | None = None,
+        txt_kv_cache: list | None = None,
+        current_start: int = 0,
+        rope_start_idx: int = 0,
+    ):
+        assert txt_kv_cache is not None, "txt_kv_cache must be provided"
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # 1. RoPE
+        # Get rotary embeddings
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size,
+            self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta, start_frame=rope_start_idx)
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+
+        # 2. Conditional embeddings
+        if timestep.dim() == 2:
+            ts_seq_len = timestep.shape[1]
+            temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
+        else:
+            temb = self.time_in(timestep, timestep_r=timestep_r)
+        # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim]
+
+        hidden_states = self.img_in(hidden_states)
+
+        # 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]
+            )
+
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block_index, block in enumerate(self.double_blocks):
+                hidden_states, new_cache = self._gradient_checkpointing_func(
+                    block,
+                    txt_inference=False,
+                    vision_inference=True,
+                    img=hidden_states,
+                    vec=temb,
+                    freqs_cis=freqs_cis,
+                    block_mask=self.block_mask,
+                    kv_cache=kv_cache[block_index] if kv_cache is not None else None,
+                    txt_kv_cache=txt_kv_cache[block_index],
+                    current_start=current_start
+                )
+                if new_cache is not None and kv_cache is not None:
+                    for k in new_cache.keys():
+                        kv_cache[block_index][k] = new_cache[k].clone()
+
+        else:
+            for block_index, block in enumerate(self.double_blocks):
+                hidden_states, new_cache = block(
+                    txt_inference=False,
+                    vision_inference=True,
+                    img=hidden_states,
+                    vec=temb,
+                    freqs_cis=freqs_cis,
+                    block_mask=self.block_mask,
+                    kv_cache=kv_cache[block_index] if kv_cache is not None else None,
+                    txt_kv_cache=txt_kv_cache[block_index],
+                    current_start=current_start
+                )
+                if new_cache is not None and kv_cache is not None:
+                    for k in new_cache.keys():
+                        kv_cache[block_index][k] = new_cache[k].clone()
+
+        # Final layer processing
+        hidden_states = self.final_layer(hidden_states, temb)
+        # Unpatchify to get original shape
+        hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels)
+
+        return hidden_states, kv_cache
+
+    def forward(
+        self,
+        txt_inference=False,
+        vision_inference=False,
+        **kwargs,
+    ):
+        if txt_inference:
+            return self.forward_txt(**kwargs)
+        elif vision_inference:
+            return self.forward_vision(**kwargs)
+        else:
+            raise ValueError("txt_inference and vision_inference cannot be both False")

fastvideo/models/dits/hunyuanvideo15.py

@@ -127,8 +127,9 @@ class HunyuanVideo15TimeEmbedding(nn.Module):
         self,
         timestep: torch.Tensor,
         timestep_r: Optional[torch.Tensor] = None,
+        timestep_seq_len: int | None = None,
     ) -> torch.Tensor:
-        timesteps_emb = self.timestep_embedder(timestep)
+        timesteps_emb = self.timestep_embedder(timestep, timestep_seq_len)
 
         if timestep_r is not None:
             timesteps_emb_r = self.timestep_embedder_r(timestep_r)
@@ -473,6 +474,7 @@ class HunyuanVideo15Transformer3DModel(CachableDiT):
         guidance: Optional[torch.Tensor] = None,
         timestep_r: Optional[torch.LongTensor] = None,
         attention_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs
     ):
         encoder_hidden_states_image = encoder_hidden_states_image[0]
         encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
@@ -494,7 +496,12 @@ class HunyuanVideo15Transformer3DModel(CachableDiT):
         freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
 
         # 2. Conditional embeddings
-        temb = self.time_in(timestep, timestep_r=timestep_r)
+        if timestep.dim() == 2:
+            ts_seq_len = timestep.shape[1]
+            temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
+        else:
+            temb = self.time_in(timestep, timestep_r=timestep_r)
+        # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim]
 
         hidden_states = self.img_in(hidden_states)
         hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
@@ -626,6 +633,8 @@ class HunyuanVideo15Transformer3DModel(CachableDiT):
                 )
 
         # Final layer processing
+        if get_sp_world_size() > 1:
+            hidden_states = hidden_states.contiguous()
         hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
         hidden_states = self.final_layer(hidden_states, temb)
         # Unpatchify to get original shape
@@ -696,6 +705,7 @@ class SingleTokenRefiner(nn.Module):
         else:
             mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
             context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
+        context_aware_representations = context_aware_representations.to(original_dtype)
 
         context_aware_representations = self.c_embedder(
             context_aware_representations)
@@ -848,6 +858,13 @@ class FinalLayer(nn.Module):
     def forward(self, x, c):
         # What the heck HF? Why you change the scale and shift order here???
         scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
-        x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        if c.dim() == 3:
+            # [bs, seq_len, inner_dim]
+            num_frames = scale.shape[1]
+            frame_seqlen = x.shape[1] // num_frames
+            x = (self.norm_final(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale.unsqueeze(2)) + shift.unsqueeze(2)).flatten(1, 2)
+        else:
+            # [bs, inner_dim]
+            x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         x, _ = self.linear(x)
         return x

fastvideo/models/dits/hyworld/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+HYWorld (HY-WorldPlay) model components for FastVideo.
+
+This module provides:
+- HYWorldTransformer3DModel: The main transformer model with ProPE and action conditioning
+- HYWorldVideoGenerator: Extended VideoGenerator for HYWorld inference
+- Utilities for pose processing and camera trajectory generation
+"""
+
+from .hyworld import HYWorldTransformer3DModel, HYWorldDoubleStreamBlock
+
+# Inference utilities (used by examples)
+from .resolution_utils import (
+    get_resolution_from_image,
+)
+
+__all__ = [
+    # Model (used by model registry)
+    "HYWorldTransformer3DModel",
+    "HYWorldDoubleStreamBlock",
+    # Inference utilities (used by examples)
+    "get_resolution_from_image",
+]

fastvideo/models/dits/hyworld/camera_rope.py

@@ -0,0 +1,261 @@
+# HY-WorldPlay/hyvideo/prope/camera_rope.py
+
+# MIT License
+#
+# Copyright (c) Authors of
+# "PRoPE: Projective Positional Encoding for Multiview Transformers"
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# How to use PRoPE attention for self-attention:
+#
+# 1. Easiest way (fast):
+#    attn = PropeDotProductAttention(...)
+#    o = attn(q, k, v, viewmats, Ks)
+#
+# 2. More flexible way (fast):
+#    attn = PropeDotProductAttention(...)
+#    attn._precompute_and_cache_apply_fns(viewmats, Ks)
+#    q = attn._apply_to_q(q)
+#    k = attn._apply_to_kv(k)
+#    v = attn._apply_to_kv(v)
+#    o = F.scaled_dot_product_attention(q, k, v, **kwargs)
+#    o = attn._apply_to_o(o)
+#
+# 3. The most flexible way (but slower because repeated computation of RoPE coefficients):
+#    o = prope_dot_product_attention(q, k, v, ...)
+#
+# How to use PRoPE attention for cross-attention:
+#
+#    attn_src = PropeDotProductAttention(...)
+#    attn_tgt = PropeDotProductAttention(...)
+#    attn_src._precompute_and_cache_apply_fns(viewmats_src, Ks_src)
+#    attn_tgt._precompute_and_cache_apply_fns(viewmats_tgt, Ks_tgt)
+#    q_src = attn_src._apply_to_q(q_src)
+#    k_tgt = attn_tgt._apply_to_kv(k_tgt)
+#    v_tgt = attn_tgt._apply_to_kv(v_tgt)
+#    o_src = F.scaled_dot_product_attention(q_src, k_tgt, v_tgt, **kwargs)
+#    o_src = attn_src._apply_to_o(o_src)
+
+from functools import partial
+from typing import Callable, Optional, Tuple, List
+
+import torch
+import torch.nn.functional as F
+
+
+def prope_qkv(
+    q: torch.Tensor,  # (batch, num_heads, seqlen, head_dim)
+    k: torch.Tensor,  # (batch, num_heads, seqlen, head_dim)
+    v: torch.Tensor,  # (batch, num_heads, seqlen, head_dim)
+    *,
+    viewmats: torch.Tensor,  # (batch, cameras, 4, 4)
+    Ks: Optional[torch.Tensor],  # (batch, cameras, 3, 3)
+    patches_x: int = None,  # How many patches wide is each image?
+    patches_y: int = None,  # How many patches tall is each image?
+    image_width: int = None,  # Width of the image. Used to normalize intrinsics.
+    image_height: int = None,  # Height of the image. Used to normalize intrinsics.
+    coeffs_x: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    coeffs_y: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    mask: Optional[torch.Tensor] = None,
+    kv_cache=None,
+    is_cache: bool = False,
+    **kwargs,
+) -> torch.Tensor:
+    """Similar to torch.nn.functional.scaled_dot_product_attention, but applies PRoPE-style
+    positional encoding.
+
+    Currently, we assume that the sequence length is equal to:
+
+        cameras * patches_x * patches_y
+
+    And token ordering allows the `(seqlen,)` axis to be reshaped into
+    `(cameras, patches_x, patches_y)`.
+    """
+    # We're going to assume self-attention: all inputs are the same shape.
+    (batch, num_heads, seqlen, head_dim) = q.shape
+    cameras = viewmats.shape[1]
+    assert q.shape == k.shape == v.shape
+    assert viewmats.shape == (batch, cameras, 4, 4)
+    assert Ks is None or Ks.shape == (batch, cameras, 3, 3)
+    # assert seqlen == cameras * patches_x * patches_y
+
+    apply_fn_q, apply_fn_kv, apply_fn_o = _prepare_apply_fns_all_dim(
+        head_dim=head_dim,
+        viewmats=viewmats,
+        Ks=Ks,
+        patches_x=patches_x,
+        patches_y=patches_y,
+        image_width=image_width,
+        image_height=image_height,
+        coeffs_x=coeffs_x,
+        coeffs_y=coeffs_y,
+    )
+
+    query = apply_fn_q(q)
+    key = apply_fn_kv(k)
+    value = apply_fn_kv(v)
+
+    return query, key, value, apply_fn_o
+
+
+def _prepare_apply_fns_all_dim(
+    head_dim: int,  # Q/K/V will have this last dimension
+    viewmats: torch.Tensor,  # (batch, cameras, 4, 4)
+    Ks: Optional[torch.Tensor],  # (batch, cameras, 3, 3)
+    patches_x: int,  # How many patches wide is each image?
+    patches_y: int,  # How many patches tall is each image?
+    image_width: int,  # Width of the image. Used to normalize intrinsics.
+    image_height: int,  # Height of the image. Used to normalize intrinsics.
+    coeffs_x: Optional[torch.Tensor] = None,
+    coeffs_y: Optional[torch.Tensor] = None,
+) -> Tuple[
+    Callable[[torch.Tensor], torch.Tensor],
+    Callable[[torch.Tensor], torch.Tensor],
+    Callable[[torch.Tensor], torch.Tensor],
+]:
+    """Prepare transforms for PRoPE-style positional encoding."""
+    device = viewmats.device
+    (batch, cameras, _, _) = viewmats.shape
+
+    # Normalize camera intrinsics.
+    if Ks is not None:
+        Ks_norm = torch.zeros_like(Ks)
+        Ks_norm[..., 0, 0] = Ks[..., 0, 0]
+        Ks_norm[..., 1, 1] = Ks[..., 1, 1]
+        Ks_norm[..., 0, 2] = 0
+        Ks_norm[..., 1, 2] = 0
+        Ks_norm[..., 2, 2] = 1.0
+        Ks_norm = Ks_norm.to(dtype=Ks.dtype)
+        del Ks
+
+        # Compute the camera projection matrices we use in PRoPE.
+        # - K is an `image<-camera` transform.
+        # - viewmats is a `camera<-world` transform.
+        # - P = lift(K) @ viewmats is an `image<-world` transform.
+        P = torch.einsum("...ij,...jk->...ik", _lift_K(Ks_norm), viewmats)
+        P_T = P.transpose(-1, -2).to(dtype=viewmats.dtype)
+        P_inv = torch.einsum(
+            "...ij,...jk->...ik",
+            _invert_SE3(viewmats),
+            _lift_K(_invert_K(Ks_norm)),
+        ).to(dtype=viewmats.dtype)
+
+    else:
+        # GTA formula. P is `camera<-world` transform.
+        P = viewmats
+        P_T = P.transpose(-1, -2)
+        P_inv = _invert_SE3(viewmats)
+
+    assert P.shape == P_inv.shape == (batch, cameras, 4, 4)
+
+    # Block-diagonal transforms to the inputs and outputs of the attention operator.
+    assert head_dim % 4 == 0
+    transforms_q = [
+        (partial(_apply_tiled_projmat, matrix=P_T), head_dim),
+    ]
+    transforms_kv = [
+        (partial(_apply_tiled_projmat, matrix=P_inv), head_dim),
+    ]
+    transforms_o = [
+        (partial(_apply_tiled_projmat, matrix=P), head_dim),
+    ]
+
+    apply_fn_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q)
+    apply_fn_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv)
+    apply_fn_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o)
+    return apply_fn_q, apply_fn_kv, apply_fn_o
+
+
+def _apply_tiled_projmat(
+    feats: torch.Tensor,  # (batch, num_heads, seqlen, feat_dim)
+    matrix: torch.Tensor,  # (batch, cameras, D, D)
+) -> torch.Tensor:
+    """Apply projection matrix to features."""
+    # - seqlen => (cameras, patches_x * patches_y)
+    # - feat_dim => (feat_dim // 4, 4)
+    (batch, num_heads, seqlen, feat_dim) = feats.shape
+    cameras = matrix.shape[1]
+    assert seqlen >= cameras and seqlen % cameras == 0
+    D = matrix.shape[-1]
+    assert matrix.shape == (batch, cameras, D, D)
+    assert feat_dim % D == 0
+    return torch.einsum(
+        "bcij,bncpkj->bncpki",
+        matrix,
+        feats.reshape((batch, num_heads, cameras, -1, feat_dim // D, D)),
+    ).reshape(feats.shape)
+
+
+def _apply_block_diagonal(
+    feats: torch.Tensor,  # (..., dim)
+    func_size_pairs: List[Tuple[Callable[[torch.Tensor], torch.Tensor], int]],
+) -> torch.Tensor:
+    """Apply a block-diagonal function to an input array.
+
+    Each function is specified as a tuple with form:
+
+        ((Tensor) -> Tensor, int)
+
+    Where the integer is the size of the input to the function.
+    """
+    funcs, block_sizes = zip(*func_size_pairs)
+    assert feats.shape[-1] == sum(block_sizes)
+    x_blocks = torch.split(feats, block_sizes, dim=-1)
+    out = torch.cat(
+        [f(x_block) for f, x_block in zip(funcs, x_blocks)],
+        dim=-1,
+    )
+    assert out.shape == feats.shape, "Input/output shapes should match."
+    return out
+
+
+def _invert_SE3(transforms: torch.Tensor) -> torch.Tensor:
+    """Invert a 4x4 SE(3) matrix."""
+    assert transforms.shape[-2:] == (4, 4)
+    Rinv = transforms[..., :3, :3].transpose(-1, -2)
+    out = torch.zeros_like(transforms)
+    out[..., :3, :3] = Rinv
+    out[..., :3, 3] = -torch.einsum("...ij,...j->...i", Rinv, transforms[..., :3, 3])
+    out[..., 3, 3] = 1.0
+    out = out.to(dtype=transforms.dtype)
+    return out
+
+
+def _lift_K(Ks: torch.Tensor) -> torch.Tensor:
+    """Lift 3x3 matrices to homogeneous 4x4 matrices."""
+    assert Ks.shape[-2:] == (3, 3)
+    out = torch.zeros(Ks.shape[:-2] + (4, 4), device=Ks.device)
+    out[..., :3, :3] = Ks
+    out[..., 3, 3] = 1.0
+    out = out.to(dtype=Ks.dtype)
+    return out
+
+
+def _invert_K(Ks: torch.Tensor) -> torch.Tensor:
+    """Invert 3x3 intrinsics matrices. Assumes no skew."""
+    assert Ks.shape[-2:] == (3, 3)
+    out = torch.zeros_like(Ks)
+    out[..., 0, 0] = 1.0 / Ks[..., 0, 0]
+    out[..., 1, 1] = 1.0 / Ks[..., 1, 1]
+    out[..., 0, 2] = -Ks[..., 0, 2] / Ks[..., 0, 0]
+    out[..., 1, 2] = -Ks[..., 1, 2] / Ks[..., 1, 1]
+    out[..., 2, 2] = 1.0
+    out = out.to(dtype=Ks.dtype)
+    return out

fastvideo/models/dits/hyworld/data_utils.py

@@ -0,0 +1,76 @@
+# HY-WorldPlay/hyvideo/utils/data_utils.py
+
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE
+#
+# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any
+# output and results therefrom are provided "AS IS" without any express or implied warranties of
+# any kind including any warranties of title, merchantability, noninfringement, course of dealing,
+# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the
+# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of
+# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a
+# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise
+# of rights and permissions under this agreement.
+# See the License for the specific language governing permissions and limitations under the License.
+
+import numpy as np
+from PIL import Image
+
+
+def resize_and_center_crop(image, target_width, target_height):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    original_width, original_height = pil_image.size
+    scale_factor = max(target_width / original_width, target_height / original_height)
+    resized_width = int(round(original_width * scale_factor))
+    resized_height = int(round(original_height * scale_factor))
+    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
+    left = (resized_width - target_width) / 2
+    top = (resized_height - target_height) / 2
+    right = (resized_width + target_width) / 2
+    bottom = (resized_height + target_height) / 2
+    cropped_image = resized_image.crop((left, top, right, bottom))
+    return np.array(cropped_image)
+
+
+def get_closest_ratio(height: float, width: float, ratios: list, buckets: list):
+    """
+    Get the closest ratio in the buckets.
+
+    Args:
+        height (float): video height
+        width (float): video width
+        ratios (list): video aspect ratio
+        buckets (list): buckets generated by `generate_crop_size_list`
+
+    Returns:
+        the closest size in the buckets and the corresponding ratio
+    """
+    aspect_ratio = float(height) / float(width)
+
+    ratios_array = np.array(ratios)
+    closest_ratio_id = np.abs(ratios_array - aspect_ratio).argmin()
+    closest_size = buckets[closest_ratio_id]
+    closest_ratio = ratios_array[closest_ratio_id]
+
+    return closest_size, closest_ratio
+
+
+def generate_crop_size_list(base_size=256, patch_size=16, max_ratio=4.0):
+    num_patches = round((base_size / patch_size) ** 2)
+    assert max_ratio >= 1.0
+    crop_size_list = []
+    wp, hp = num_patches, 1
+    while wp > 0:
+        if max(wp, hp) / min(wp, hp) <= max_ratio:
+            crop_size_list.append((wp * patch_size, hp * patch_size))
+        if (hp + 1) * wp <= num_patches:
+            hp += 1
+        else:
+            wp -= 1
+    return crop_size_list

fastvideo/models/dits/hyworld/hyworld.py

@@ -0,0 +1,569 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE
+#
+# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any
+# output and results therefrom are provided "AS IS" without any express or implied warranties of
+# any kind including any warranties of title, merchantability, noninfringement, course of dealing,
+# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the
+# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of
+# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a
+# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise
+# of rights and permissions under this agreement.
+# See the License for the specific language governing permissions and limitations under the License.
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+from fastvideo.distributed.communication_op import (
+    sequence_model_parallel_all_gather_with_unpad,
+    sequence_model_parallel_shard)
+from fastvideo.configs.models.dits import HYWorldConfig
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
+from fastvideo.layers.visual_embedding import TimestepEmbedder, unpatchify
+from fastvideo.models.dits.hunyuanvideo15 import (
+    MMDoubleStreamBlock,
+    HunyuanVideo15Transformer3DModel,
+)
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.logger import init_logger
+from fastvideo.forward_context import set_forward_context
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.distributed.utils import create_attention_mask_for_padding
+
+from .camera_rope import prope_qkv
+
+logger = init_logger(__name__)
+
+
+class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
+    """
+    Extended MMDoubleStreamBlock with ProPE (Projective Positional Encoding) support
+    for camera-aware attention in HY-World/WorldPlay models.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_attention_heads: int,
+        mlp_ratio: float,
+        dtype: torch.dtype | None = None,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+        prefix: str = "",
+    ):
+        super().__init__(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            mlp_ratio=mlp_ratio,
+            dtype=dtype,
+            supported_attention_backends=supported_attention_backends,
+            prefix=prefix,
+        )
+        self.hidden_size = hidden_size
+
+        # Add ProPE projection layer for camera-aware attention
+        self.img_attn_prope_proj = ReplicatedLinear(
+            hidden_size,
+            hidden_size,
+            bias=True,
+            params_dtype=dtype,
+            prefix=f"{prefix}.img_attn_prope_proj"
+        )
+        # Zero-initialize ProPE projection (starts as identity)
+        nn.init.zeros_(self.img_attn_prope_proj.weight)
+        if self.img_attn_prope_proj.bias is not None:
+            nn.init.zeros_(self.img_attn_prope_proj.bias)
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        vec: torch.Tensor,
+        vec_txt: torch.Tensor,
+        freqs_cis: tuple,
+        seq_attention_mask: torch.Tensor,
+        viewmats: torch.Tensor,
+        Ks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass with ProPE camera conditioning.
+
+        Args:
+            img: Image/video tokens
+            txt: Text tokens
+            encoder_attention_mask: Text attention mask
+            vec: Modulation vector
+            freqs_cis: Rotary embedding frequencies
+            seq_attention_mask: Sequence attention mask
+            viewmats: Camera view matrices for ProPE [B, T, 4, 4]
+            Ks: Camera intrinsics for ProPE [B, T, 3, 3]
+
+        Returns:
+            Tuple of (img, txt) output tokens
+        """
+        # Process modulation vectors (inherited from parent)
+        img_mod_outputs = self.img_mod(vec)
+        (
+            img_attn_shift,
+            img_attn_scale,
+            img_attn_gate,
+            img_mlp_shift,
+            img_mlp_scale,
+            img_mlp_gate,
+        ) = torch.chunk(img_mod_outputs, 6, dim=-1)
+
+        txt_mod_outputs = self.txt_mod(vec_txt)
+        (
+            txt_attn_shift,
+            txt_attn_scale,
+            txt_attn_gate,
+            txt_mlp_shift,
+            txt_mlp_scale,
+            txt_mlp_gate,
+        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)
+
+        # Prepare image for attention using fused operation
+        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale, convert_modulation_dtype=True)
+        # Get QKV for image
+        img_qkv, _ = self.img_attn_qkv(img_attn_input)
+        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]
+
+        # Split QKV
+        img_qkv = img_qkv.view(batch_size, image_seq_len, 3,
+                               self.num_attention_heads, -1)
+        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :,
+                                                                          2]
+
+        # Apply QK-Norm if needed
+        img_q = self.img_attn_q_norm(img_q).to(img_v)
+        img_k = self.img_attn_k_norm(img_k).to(img_v)
+
+        # Prepare text for attention using fused operation
+        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale, convert_modulation_dtype=True)
+
+        # Get QKV for text
+        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
+        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]
+
+        # Split QKV
+        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3,
+                               self.num_attention_heads, -1)
+        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :,
+                                                                          2]
+        # Apply QK-Norm if needed
+        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
+        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
+
+        # begin hyworld: add camera pose through prope
+        img_q_prope, img_k_prope, img_v_prope, apply_fn_o = prope_qkv(
+            img_q.permute(0, 2, 1, 3),
+            img_k.permute(0, 2, 1, 3),
+            img_v.permute(0, 2, 1, 3),
+            viewmats=viewmats,
+            Ks=Ks,
+        )  # [batch, num_heads, seqlen, head_dim]
+        img_q_prope = img_q_prope.permute(0, 2, 1, 3)  # [batch, seqlen, num_heads, head_dim]
+        img_k_prope = img_k_prope.permute(0, 2, 1, 3)  # [batch, seqlen, num_heads, head_dim]
+        img_v_prope = img_v_prope.permute(0, 2, 1, 3)  # [batch, seqlen, num_heads, head_dim]
+        # end hyworld
+
+        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+        attn_metadata = FlashAttnMetadataBuilder().build(
+            current_timestep=0,
+            attn_mask=encoder_attention_mask,
+        )
+        # Run distributed attention
+        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+            img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis, attention_mask=seq_attention_mask)
+        
+        # begin hyworld
+        # attention with prope
+        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+        attn_metadata_prope = FlashAttnMetadataBuilder().build(
+            current_timestep=0,
+            attn_mask=encoder_attention_mask,
+        )       
+        # NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope
+        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata_prope):
+            img_attn_prope, _ = self.attn(
+                img_q_prope, img_k_prope, img_v_prope, txt_q, txt_k, txt_v, 
+                freqs_cis=None, attention_mask=seq_attention_mask  # No RoPE for prope attention
+            )
+            img_attn_prope = img_attn_prope.reshape(batch_size, image_seq_len, -1)
+            img_attn_prope = rearrange(
+                img_attn_prope, "B L (H D) -> B H L D", H=self.num_attention_heads
+            )
+            img_attn_prope = apply_fn_o(img_attn_prope) # [batch, num_heads, seqlen, head_dim]
+            img_attn_prope = rearrange(img_attn_prope, "B H L D -> B L (H D)")
+
+        # add prope to img_attn
+        img_attn_out, _ = self.img_attn_proj(img_attn.view(batch_size, image_seq_len, -1))
+        img_attn_prope_out, _ = self.img_attn_prope_proj(img_attn_prope)
+        img_attn_out = img_attn_out + img_attn_prope_out
+        # end hyworld
+
+        # Use fused operation for residual connection, normalization, and modulation
+        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(
+            img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale, convert_modulation_dtype=True)
+
+        # Process image MLP
+        img_mlp_out = self.img_mlp(img_mlp_input)
+        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
+
+        # Process text attention output
+        txt_attn_out, _ = self.txt_attn_proj(
+            txt_attn.reshape(batch_size, text_seq_len, -1))
+
+        # Use fused operation for residual connection, normalization, and modulation
+        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(
+            txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale, convert_modulation_dtype=True)
+
+        # Process text MLP
+        txt_mlp_out = self.txt_mlp(txt_mlp_input)
+        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
+
+        return img, txt
+
+
+class HYWorldFinalLayer(nn.Module):
+    """
+    Final layer for HYWorld that uses modulate() to handle per-token conditioning.
+    This matches HY-WorldPlay's FinalLayer behavior.
+    """
+
+    def __init__(self,
+                 hidden_size,
+                 patch_size,
+                 out_channels,
+                 dtype=None,
+                 prefix: str = "") -> None:
+        super().__init__()
+        from fastvideo.layers.linear import ReplicatedLinear
+        from fastvideo.layers.visual_embedding import ModulateProjection
+        from fastvideo.layers.layernorm import LayerNormScaleShift
+
+        self.norm_final = LayerNormScaleShift(
+            hidden_size,
+            norm_type="layer",
+            eps=1e-6,
+            elementwise_affine=False,
+            dtype=dtype,
+            prefix=f"{prefix}.norm_final")
+
+        output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels
+
+        self.linear = ReplicatedLinear(hidden_size,
+                                       output_dim,
+                                       bias=True,
+                                       params_dtype=dtype,
+                                       prefix=f"{prefix}.linear")
+
+        # Modulation projection to get shift/scale
+        self.adaLN_modulation = ModulateProjection(
+            hidden_size,
+            factor=2,
+            act_layer="silu",
+            dtype=dtype,
+            prefix=f"{prefix}.adaLN_modulation")
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = self.norm_final(x, shift, scale, convert_modulation_dtype=True)
+        x, _ = self.linear(x)
+        return x
+
+
+class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
+    r"""
+    HY-World Transformer extending HunyuanVideo15 with:
+    - ProPE (Projective Positional Encoding) for camera-aware attention
+    - Action conditioning for interactive video generation
+    """
+    
+    # Class attributes for weight loading - use HYWorld-specific mapping
+    _fsdp_shard_conditions = HYWorldConfig().arch_config._fsdp_shard_conditions
+    _compile_conditions = HYWorldConfig().arch_config._compile_conditions
+    param_names_mapping = HYWorldConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = HYWorldConfig().arch_config.reverse_param_names_mapping
+
+    def __init__(
+        self,
+        config: HYWorldConfig,
+        hf_config: dict[str, Any],
+    ) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        # Replace double_blocks with HY-World version that supports ProPE
+        self.double_blocks = nn.ModuleList([
+            HYWorldDoubleStreamBlock(
+                hidden_size=self.hidden_size,
+                num_attention_heads=self.num_attention_heads,
+                mlp_ratio=config.arch_config.mlp_ratio,
+                dtype=None,
+                supported_attention_backends=self._supported_attention_backends,
+                prefix=f"{config.prefix}.double_blocks.{i}"
+            )
+            for i in range(config.arch_config.num_layers)
+        ])
+
+        # Add action conditioning module
+        self.action_in = TimestepEmbedder(
+            self.hidden_size,
+            act_layer="silu",
+            dtype=None,
+            prefix=f"{config.prefix}.action_in"
+        )
+        # Zero-initialize action embedding (starts with no effect)
+        nn.init.zeros_(self.action_in.mlp.fc_out.weight)
+        if self.action_in.mlp.fc_out.bias is not None:
+            nn.init.zeros_(self.action_in.mlp.fc_out.bias)
+
+        # Override final_layer with HYWorld version that uses per-token modulate()
+        self.final_layer = HYWorldFinalLayer(
+            hidden_size=self.hidden_size,
+            patch_size=self.patch_size,
+            out_channels=self.out_channels,
+            dtype=None,
+            prefix=f"{config.prefix}.final_layer"
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: list[torch.Tensor],
+        encoder_attention_mask: list[torch.Tensor],
+        action: torch.Tensor,
+        viewmats: torch.Tensor,
+        Ks: torch.Tensor,
+        timestep_txt: torch.LongTensor,
+        guidance: Optional[torch.Tensor] = None,
+        timestep_r: Optional[torch.LongTensor] = None,
+        attention_kwargs: Optional[dict[str, Any]] = None,
+    ):
+        """
+        Forward pass with action and camera conditioning.
+
+        Args:
+            action: Action tensor for action conditioning [B, T] or [B*T]
+            viewmats: Camera view matrices [B, T, 4, 4]
+            Ks: Camera intrinsics [B, T, 3, 3]
+            ... (other args same as parent)
+        """
+        encoder_hidden_states_image = encoder_hidden_states_image[0]
+        encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
+        encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # 1. RoPE
+        # Get rotary embeddings
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames, post_patch_height, post_patch_width),
+            self.hidden_size,
+            self.num_attention_heads,
+            self.config.rope_axes_dim,
+            self.config.rope_theta
+        )
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        # NOTE: freqs_cis does NOT need sharding because FastVideo's DistributedAttention
+        # uses all-to-all to gather the full sequence before applying RoPE
+        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+
+        # 2. Conditional embeddings
+        temb = self.time_in(timestep, timestep_r=timestep_r)
+        temb_txt = self.time_in(timestep_txt, timestep_r=timestep_r)
+
+        # Add action conditioning if provided
+        # temb shape: [B*T, C] where T = num_frames
+        temb = temb + self.action_in(action.reshape(-1))
+        
+        # Broadcast timestep embedding for transformer blocks (one per spatial token)
+        # [B*T, C] -> [B, T*H*W, C] -> [B*T*H*W, C]
+        temb = repeat(temb, "(B T) C -> B (T H W) C", B=batch_size, H=post_patch_height, W=post_patch_width)
+
+        hidden_states = self.img_in(hidden_states)
+        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
+
+        current_seq_len = hidden_states.shape[1]
+        sp_world_size = get_sp_world_size()
+        padded_seq_len = current_seq_len * sp_world_size
+
+        if padded_seq_len > original_seq_len:
+            seq_attention_mask = create_attention_mask_for_padding(
+                seq_len=original_seq_len,
+                padded_seq_len=padded_seq_len,
+                batch_size=batch_size,
+                device=hidden_states.device,
+            )
+        else:
+            seq_attention_mask = None
+
+        viewmats_seq = repeat(
+            viewmats, "B T M N->B (T H W) M N",
+            H=post_patch_height,
+            W=post_patch_width
+        )
+        Ks_seq = repeat(
+            Ks, "B T M N->B (T H W) M N",
+            H=post_patch_height,
+            W=post_patch_width
+        )
+        
+        # Shard viewmats, Ks, and temb for sequence parallelism (shard along sequence dim=1)
+        # Note that temb in HY1.5 does not need sharding because it is per-sample modulation
+        # In HYWorld, temb is per-token modulation.
+        if sp_world_size > 1:
+            viewmats_seq, _ = sequence_model_parallel_shard(viewmats_seq, dim=1)
+            Ks_seq, _ = sequence_model_parallel_shard(Ks_seq, dim=1)
+            temb, _ = sequence_model_parallel_shard(temb, dim=1)
+        
+        # Rearrange temb after sharding to match expected shape
+        temb = rearrange(temb, "B S C -> (B S) C")
+
+        # qwen text embedding
+        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep_txt, encoder_attention_mask)
+
+        encoder_hidden_states_cond_emb = self.cond_type_embed(
+            torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb
+
+        # byt5 text embedding
+        encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2)
+
+        encoder_hidden_states_2_cond_emb = self.cond_type_embed(
+            torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)
+        )
+        encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb
+
+        # image embed
+        encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
+        is_t2v = torch.all(encoder_hidden_states_image == 0)
+        if is_t2v:
+            encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
+            encoder_attention_mask_3 = torch.zeros(
+                (batch_size, encoder_hidden_states_3.shape[1]),
+                dtype=encoder_attention_mask.dtype,
+                device=encoder_attention_mask.device,
+            )
+        else:
+            encoder_attention_mask_3 = torch.ones(
+                (batch_size, encoder_hidden_states_3.shape[1]),
+                dtype=encoder_attention_mask.dtype,
+                device=encoder_attention_mask.device,
+            )
+        encoder_hidden_states_3_cond_emb = self.cond_type_embed(
+            2
+            * torch.ones_like(
+                encoder_hidden_states_3[:, :, 0],
+                dtype=torch.long,
+            )
+        )
+        encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
+
+        # reorder and combine text tokens: combine valid tokens first, then padding
+        encoder_attention_mask = encoder_attention_mask.bool()
+        encoder_attention_mask_2 = encoder_attention_mask_2.bool()
+        encoder_attention_mask_3 = encoder_attention_mask_3.bool()
+        new_encoder_hidden_states = []
+        new_encoder_attention_mask = []
+
+        for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
+            encoder_hidden_states,
+            encoder_attention_mask,
+            encoder_hidden_states_2,
+            encoder_attention_mask_2,
+            encoder_hidden_states_3,
+            encoder_attention_mask_3,
+        ):
+            # Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm]
+            new_encoder_hidden_states.append(
+                torch.cat(
+                    [
+                        image[image_mask],  # valid image
+                        text_2[text_mask_2],  # valid byt5
+                        text[text_mask],  # valid mllm
+                        image[~image_mask],  # invalid image (zeroed)
+                        torch.zeros_like(text_2[~text_mask_2]),  # invalid byt5 (zeroed)
+                        torch.zeros_like(text[~text_mask]),  # invalid mllm (zeroed)
+                    ],
+                    dim=0,
+                )
+            )
+            # Apply same reordering to attention masks
+            new_encoder_attention_mask.append(
+                torch.cat(
+                    [
+                        image_mask[image_mask],
+                        text_mask_2[text_mask_2],
+                        text_mask[text_mask],
+                        image_mask[~image_mask],
+                        text_mask_2[~text_mask_2],
+                        text_mask[~text_mask],
+                    ],
+                    dim=0,
+                )
+            )
+
+        encoder_hidden_states = torch.stack(new_encoder_hidden_states)
+        encoder_attention_mask = torch.stack(new_encoder_attention_mask)
+
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.double_blocks:
+                hidden_states, encoder_hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    temb,
+                    temb_txt,
+                    freqs_cis,
+                    seq_attention_mask,
+                    viewmats_seq, # hyworld
+                    Ks_seq, # hyworld
+                )
+        else:
+            for block in self.double_blocks:
+                hidden_states, encoder_hidden_states = block(
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    temb,
+                    temb_txt,
+                    freqs_cis,
+                    seq_attention_mask,
+                    viewmats=viewmats_seq, # hyworld
+                    Ks=Ks_seq, # hyworld
+                )
+
+
+        # Final layer processing (per-token conditioning via HYWorldFinalLayer)
+        # Apply final_layer on sharded data first, then gather
+        hidden_states = self.final_layer(hidden_states, temb)
+        
+        # Gather the output from all ranks
+        if get_sp_world_size() > 1:
+            hidden_states = hidden_states.contiguous()
+        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
+        
+        # Unpatchify to get original shape
+        hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels)
+
+        return hidden_states

fastvideo/models/dits/hyworld/pose.py

@@ -0,0 +1,413 @@
+# Some functions from HY-WorldPlay/hyvideo/generate.py
+
+"""
+Pose processing utilities for HYWorld video generation.
+
+This module provides functions to convert camera poses to model input tensors,
+including viewmats, intrinsics, and action labels.
+
+Adapted from HY-WorldPlay: https://github.com/Tencent-Hunyuan/HY-WorldPlay
+"""
+
+import json
+import numpy as np
+import torch
+from scipy.spatial.transform import Rotation as R
+from typing import Union, Optional
+
+from .trajectory import generate_camera_trajectory_local
+
+
+# Mapping from one-hot action encoding to single label
+mapping = {
+    (0, 0, 0, 0): 0,
+    (1, 0, 0, 0): 1,
+    (0, 1, 0, 0): 2,
+    (0, 0, 1, 0): 3,
+    (0, 0, 0, 1): 4,
+    (1, 0, 1, 0): 5,
+    (1, 0, 0, 1): 6,
+    (0, 1, 1, 0): 7,
+    (0, 1, 0, 1): 8,
+}
+
+# Default camera intrinsic matrix (for 1920x1080 resolution)
+DEFAULT_INTRINSIC = [
+    [969.6969696969696, 0.0, 960.0],
+    [0.0, 969.6969696969696, 540.0],
+    [0.0, 0.0, 1.0],
+]
+
+# Default movement speeds
+DEFAULT_FORWARD_SPEED = 0.08  # units per frame
+DEFAULT_YAW_SPEED = np.deg2rad(3)  # radians per frame
+DEFAULT_PITCH_SPEED = np.deg2rad(3)  # radians per frame
+
+
+def one_hot_to_one_dimension(one_hot: torch.Tensor) -> torch.Tensor:
+    """Convert one-hot action encoding to single dimension labels."""
+    return torch.tensor([mapping[tuple(row.tolist())] for row in one_hot])
+
+
+def parse_pose_string(
+    pose_string: str,
+    forward_speed: float = DEFAULT_FORWARD_SPEED,
+    yaw_speed: float = DEFAULT_YAW_SPEED,
+    pitch_speed: float = DEFAULT_PITCH_SPEED,
+) -> list[dict]:
+    """
+    Parse pose string to motions list.
+    
+    Format: "w-3, right-0.5, d-4"
+    - w: forward movement
+    - s: backward movement
+    - a: left movement
+    - d: right movement
+    - up: pitch up rotation
+    - down: pitch down rotation
+    - left: yaw left rotation
+    - right: yaw right rotation
+    - number after dash: duration in frames/latents
+    
+    Args:
+        pose_string: Comma-separated pose commands
+        forward_speed: Movement amount per frame
+        yaw_speed: Yaw rotation amount per frame (radians)
+        pitch_speed: Pitch rotation amount per frame (radians)
+        
+    Returns:
+        List of motion dictionaries for generate_camera_trajectory_local
+    """
+    motions = []
+    commands = [cmd.strip() for cmd in pose_string.split(",")]
+
+    for cmd in commands:
+        if not cmd:
+            continue
+
+        parts = cmd.split("-")
+        if len(parts) != 2:
+            raise ValueError(
+                f"Invalid pose command: {cmd}. Expected format: 'action-duration'"
+            )
+
+        action = parts[0].strip()
+        try:
+            duration = float(parts[1].strip())
+        except ValueError:
+            raise ValueError(f"Invalid duration in command: {cmd}")
+
+        num_frames = int(duration)
+
+        # Parse action and create motion dicts
+        if action == "w":
+            # Forward
+            for _ in range(num_frames):
+                motions.append({"forward": forward_speed})
+        elif action == "s":
+            # Backward
+            for _ in range(num_frames):
+                motions.append({"forward": -forward_speed})
+        elif action == "a":
+            # Left
+            for _ in range(num_frames):
+                motions.append({"right": -forward_speed})
+        elif action == "d":
+            # Right
+            for _ in range(num_frames):
+                motions.append({"right": forward_speed})
+        elif action == "up":
+            # Pitch up
+            for _ in range(num_frames):
+                motions.append({"pitch": pitch_speed})
+        elif action == "down":
+            # Pitch down
+            for _ in range(num_frames):
+                motions.append({"pitch": -pitch_speed})
+        elif action == "left":
+            # Yaw left
+            for _ in range(num_frames):
+                motions.append({"yaw": -yaw_speed})
+        elif action == "right":
+            # Yaw right
+            for _ in range(num_frames):
+                motions.append({"yaw": yaw_speed})
+        else:
+            raise ValueError(
+                f"Unknown action: {action}. "
+                f"Supported actions: w, s, a, d, up, down, left, right"
+            )
+
+    return motions
+
+def pose_string_to_json(
+    pose_string: str,
+    intrinsic: Optional[list[list[float]]] = None,
+) -> dict:
+    """
+    Convert pose string to pose JSON format.
+    
+    Args:
+        pose_string: Comma-separated pose commands
+        intrinsic: Camera intrinsic matrix (default: DEFAULT_INTRINSIC from trajectory)
+        
+    Returns:
+        Dict with frame indices as keys, containing extrinsic and K (intrinsic) matrices
+    """
+    if intrinsic is None:
+        intrinsic = DEFAULT_INTRINSIC
+    
+    motions = parse_pose_string(pose_string)
+    poses = generate_camera_trajectory_local(motions)
+
+    pose_json = {}
+    for i, p in enumerate(poses):
+        pose_json[str(i)] = {"extrinsic": p.tolist(), "K": intrinsic}
+
+    return pose_json
+
+def pose_to_input(
+    pose_data: Union[str, dict],
+    latent_num: int,
+    tps: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Convert pose data to model input tensors.
+    
+    Args:
+        pose_data: One of:
+            - str ending with '.json': path to JSON file
+            - str: pose string (e.g., "w-3, right-0.5, d-4")
+            - dict: pose JSON data
+        latent_num: Number of latents (frames in latent space)
+        tps: Third-person mode flag
+        
+    Returns:
+        Tuple of (viewmats, intrinsics, action_labels):
+            - viewmats: World-to-camera matrices [T, 4, 4]
+            - intrinsics: Normalized camera intrinsics [T, 3, 3]
+            - action_labels: Action labels for each frame [T]
+    """
+    # Handle different input types
+    if isinstance(pose_data, str):
+        if pose_data.endswith(".json"):
+            # Load from JSON file
+            with open(pose_data, "r") as f:
+                pose_json = json.load(f)
+        else:
+            # Parse pose string
+            pose_json = pose_string_to_json(pose_data)
+    elif isinstance(pose_data, dict):
+        pose_json = pose_data
+    else:
+        raise ValueError(
+            f"Invalid pose_data type: {type(pose_data)}. Expected str or dict."
+        )
+
+    pose_keys = list(pose_json.keys())
+    latent_num_from_pose = len(pose_keys)
+    assert latent_num_from_pose == latent_num, (
+        f"pose corresponds to {latent_num_from_pose * 4 - 3} frames, num_frames "
+        f"must be set to {latent_num_from_pose * 4 - 3} to ensure alignment."
+    )
+
+    intrinsic_list = []
+    w2c_list = []
+    for i in range(latent_num):
+        t_key = pose_keys[i]
+        c2w = np.array(pose_json[t_key]["extrinsic"])
+        w2c = np.linalg.inv(c2w)
+        w2c_list.append(w2c)
+        
+        # Normalize intrinsics
+        intrinsic = np.array(pose_json[t_key]["K"])
+        intrinsic[0, 0] /= intrinsic[0, 2] * 2
+        intrinsic[1, 1] /= intrinsic[1, 2] * 2
+        intrinsic[0, 2] = 0.5
+        intrinsic[1, 2] = 0.5
+        intrinsic_list.append(intrinsic)
+
+    w2c_list = np.array(w2c_list)
+    intrinsic_list = torch.tensor(np.array(intrinsic_list))
+
+    # Compute relative camera-to-world transforms
+    c2ws = np.linalg.inv(w2c_list)
+    C_inv = np.linalg.inv(c2ws[:-1])
+    relative_c2w = np.zeros_like(c2ws)
+    relative_c2w[0, ...] = c2ws[0, ...]
+    relative_c2w[1:, ...] = C_inv @ c2ws[1:, ...]
+    
+    # Initialize one-hot action encodings
+    trans_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32)
+    rotate_one_hot = np.zeros((relative_c2w.shape[0], 4), dtype=np.int32)
+
+    move_norm_valid = 0.0001
+    for i in range(1, relative_c2w.shape[0]):
+        move_dirs = relative_c2w[i, :3, 3]  # direction vector
+        move_norms = np.linalg.norm(move_dirs)
+        
+        if move_norms > move_norm_valid:  # threshold for movement
+            move_norm_dirs = move_dirs / move_norms
+            angles_rad = np.arccos(move_norm_dirs.clip(-1.0, 1.0))
+            trans_angles_deg = angles_rad * (180.0 / np.pi)  # convert to degrees
+        else:
+            trans_angles_deg = np.zeros(3)
+
+        R_rel = relative_c2w[i, :3, :3]
+        r = R.from_matrix(R_rel)
+        rot_angles_deg = r.as_euler("xyz", degrees=True)
+
+        # Determine movement and rotation actions
+        if move_norms > move_norm_valid:  # threshold for movement
+            if (not tps) or (
+                tps and abs(rot_angles_deg[1]) < 5e-2 and abs(rot_angles_deg[0]) < 5e-2
+            ):
+                if trans_angles_deg[2] < 60:
+                    trans_one_hot[i, 0] = 1  # forward
+                elif trans_angles_deg[2] > 120:
+                    trans_one_hot[i, 1] = 1  # backward
+
+                if trans_angles_deg[0] < 60:
+                    trans_one_hot[i, 2] = 1  # right
+                elif trans_angles_deg[0] > 120:
+                    trans_one_hot[i, 3] = 1  # left
+
+        if rot_angles_deg[1] > 5e-2:
+            rotate_one_hot[i, 0] = 1  # right
+        elif rot_angles_deg[1] < -5e-2:
+            rotate_one_hot[i, 1] = 1  # left
+
+        if rot_angles_deg[0] > 5e-2:
+            rotate_one_hot[i, 2] = 1  # up
+        elif rot_angles_deg[0] < -5e-2:
+            rotate_one_hot[i, 3] = 1  # down
+
+    trans_one_hot = torch.tensor(trans_one_hot)
+    rotate_one_hot = torch.tensor(rotate_one_hot)
+
+    # Convert one-hot to single-dimension labels
+    trans_one_label = one_hot_to_one_dimension(trans_one_hot)
+    rotate_one_label = one_hot_to_one_dimension(rotate_one_hot)
+    action_one_label = trans_one_label * 9 + rotate_one_label
+
+    return (
+        torch.as_tensor(w2c_list),
+        torch.as_tensor(intrinsic_list),
+        action_one_label,
+    )
+
+
+def camera_center_normalization(w2c: np.ndarray) -> np.ndarray:
+    """Normalize camera centers relative to the first camera."""
+    c2w = np.linalg.inv(w2c)
+    C0_inv = np.linalg.inv(c2w[0])
+    c2w_aligned = np.array([C0_inv @ C for C in c2w])
+    return np.linalg.inv(c2w_aligned)
+
+
+
+def parse_pose_string_to_actions(pose_string: str, fps: int = 24) -> list[dict]:
+    """
+    Parse pose string to frame-level action timeline.
+    
+    Format: pose string uses latent counts, where:
+    - 1 latent = 4 frames
+    - Special rule: first frame of entire video is extra (frame 0)
+    - Example: "w-4,d-4" means:
+      - w-4: forward for frames 0-16 (17 frames total: 1 extra + 4*4)
+      - d-4: right for frames 17-32 (16 frames total: 4*4)
+    
+    Args:
+        pose_string: Comma-separated pose commands (e.g., "w-4,d-4")
+        fps: Frames per second for video (default: 24)
+        
+    Returns:
+        List of dicts with action values for each frame
+    """
+    commands = [cmd.strip() for cmd in pose_string.split(",")]
+
+    frame_actions = []
+    is_first_command = True
+
+    for cmd in commands:
+        if not cmd:
+            continue
+
+        parts = cmd.split("-")
+        if len(parts) != 2:
+            raise ValueError(
+                f"Invalid pose command: {cmd}. Expected format: 'action-duration'"
+            )
+
+        action = parts[0].strip()
+        try:
+            num_latents = int(parts[1].strip())
+        except ValueError:
+            raise ValueError(f"Invalid duration in command: {cmd}")
+
+        # Convert latents to frames
+        # First command gets 1 extra frame (the special frame 0)
+        if is_first_command:
+            num_frames = 1 + num_latents * 4
+            is_first_command = False
+        else:
+            num_frames = num_latents * 4
+
+        # Map action to action values
+        action_values = {"forward": 0, "left": 0, "yaw": 0, "pitch": 0}
+
+        if action == "w":
+            action_values["forward"] = 1
+        elif action == "s":
+            action_values["forward"] = -1
+        elif action == "a":
+            action_values["left"] = 1
+        elif action == "d":
+            action_values["left"] = -1
+        elif action == "up":
+            action_values["pitch"] = 1
+        elif action == "down":
+            action_values["pitch"] = -1
+        elif action == "left":
+            action_values["yaw"] = -1
+        elif action == "right":
+            action_values["yaw"] = 1
+        else:
+            raise ValueError(f"Unknown action: {action}")
+
+        # Add frame-level actions
+        for _ in range(num_frames):
+            frame_actions.append(action_values.copy())
+
+    return frame_actions
+
+
+def compute_latent_num(num_frames: int) -> int:
+    """
+    Compute the number of latents from number of frames.
+    
+    Formula: num_frames = (latent_num - 1) * 4 + 1
+    So: latent_num = (num_frames - 1) // 4 + 1
+    
+    Args:
+        num_frames: Number of video frames
+        
+    Returns:
+        Number of latents
+    """
+    return (num_frames - 1) // 4 + 1
+
+
+def compute_num_frames(latent_num: int) -> int:
+    """
+    Compute the number of frames from number of latents.
+    
+    Formula: num_frames = (latent_num - 1) * 4 + 1
+    
+    Args:
+        latent_num: Number of latents
+        
+    Returns:
+        Number of video frames
+    """
+    return (latent_num - 1) * 4 + 1

fastvideo/models/dits/hyworld/resolution_utils.py

@@ -0,0 +1,64 @@
+import numpy as np
+from PIL import Image
+import requests
+from io import BytesIO
+
+from fastvideo.models.dits.hyworld.data_utils import generate_crop_size_list
+
+# Target resolution configs (matching HY-WorldPlay)
+TARGET_SIZE_CONFIG = {
+    "360p": {"bucket_hw_base_size": 480, "bucket_hw_bucket_stride": 16},
+    "480p": {"bucket_hw_base_size": 640, "bucket_hw_bucket_stride": 16},
+    "720p": {"bucket_hw_base_size": 960, "bucket_hw_bucket_stride": 16},
+    "1080p": {"bucket_hw_base_size": 1440, "bucket_hw_bucket_stride": 16},
+}
+
+
+def get_closest_resolution(image_height, image_width, target_resolution="480p"):
+    """
+    Get closest supported resolution for given image dimensions.
+    
+    Args:
+        image_height: Height of input image
+        image_width: Width of input image  
+        target_resolution: Target resolution string (e.g., "480p", "720p")
+        
+    Returns:
+        tuple[int, int]: (height, width) of closest supported resolution
+    """
+    config = TARGET_SIZE_CONFIG[target_resolution]
+    bucket_hw_base_size = config["bucket_hw_base_size"]
+    bucket_hw_bucket_stride = config["bucket_hw_bucket_stride"]
+    
+    crop_size_list = generate_crop_size_list(bucket_hw_base_size, bucket_hw_bucket_stride)
+    aspect_ratios = np.array([round(float(h) / float(w), 5) for h, w in crop_size_list])
+    
+    # Find closest aspect ratio
+    image_ratio = float(image_height) / float(image_width)
+    closest_idx = np.abs(aspect_ratios - image_ratio).argmin()
+    closest_size = crop_size_list[closest_idx]
+    
+    return closest_size[0], closest_size[1]  # (height, width)
+
+
+def get_resolution_from_image(image_path, target_resolution="480p"):
+    """
+    Automatically determine resolution from input image.
+    
+    Args:
+        image_path: Path or URL to input image
+        target_resolution: Target resolution tier ("480p", "720p", etc.)
+        
+    Returns:
+        tuple[int, int]: (height, width) matching HY-WorldPlay's bucket selection
+    """
+    # Handle URL inputs
+    if isinstance(image_path, str) and image_path.startswith(('http://', 'https://')):
+        response = requests.get(image_path)
+        response.raise_for_status()
+        img = Image.open(BytesIO(response.content))
+    else:
+        img = Image.open(image_path)
+    img_width, img_height = img.size
+    return get_closest_resolution(img_height, img_width, target_resolution)
+

fastvideo/models/dits/hyworld/retrieval_context.py

@@ -0,0 +1,316 @@
+# HY-WorldPlay/hyvideo/utils/retrieval_context.py
+
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE
+#
+# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any
+# output and results therefrom are provided "AS IS" without any express or implied warranties of
+# any kind including any warranties of title, merchantability, noninfringement, course of dealing,
+# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the
+# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of
+# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a
+# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise
+# of rights and permissions under this agreement.
+# See the License for the specific language governing permissions and limitations under the License.
+
+import torch
+import numpy as np
+from typing import List, Tuple, Dict
+import math
+
+
+def generate_points_in_sphere(n_points: int, radius: float) -> torch.Tensor:
+    """
+    Uniformly sample points within a sphere of a specified radius.
+
+    :param n_points: The number of points to generate.
+    :param radius: The radius of the sphere.
+    :return: A tensor of shape (n_points, 3), representing the (x, y, z) coordinates of the points.
+    """
+    samples_r = torch.rand(n_points)
+    samples_phi = torch.rand(n_points)
+    samples_u = torch.rand(n_points)
+
+    r = radius * torch.pow(samples_r, 1 / 3)
+    phi = 2 * math.pi * samples_phi
+    theta = torch.acos(1 - 2 * samples_u)
+
+    # transfer the coordinates from spherical to cartesian
+    x = r * torch.sin(theta) * torch.cos(phi)
+    y = r * torch.sin(theta) * torch.sin(phi)
+    z = r * torch.cos(theta)
+
+    points = torch.stack((x, y, z), dim=1)
+    return points
+
+
+def rotation_matrix_to_angles(R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Estimate the Pitch and Yaw angles from a 3x3 rotation matrix R in the camera coordinate system.
+
+    Assumed Camera Coordinate System: X=Right, Y=Up, Z=Backward
+    (or NeRF style: X=Right, Y=Down, Z=Forward).
+    Here we adopt the common Computer Vision convention: Z-axis is Forward.
+
+    Note: The angle calculations here are directly based on the conventions of your `is_inside_fov_3d_hv` function:
+    - Yaw/Azimuth angle is in the XZ plane (atan2(x, z)).
+    - Pitch/Elevation angle is relative to the horizontal plane (atan2(y, sqrt(x^2 + z^2))).
+
+    For the third column R[:, 2] of the W2C matrix R (the direction of the World Z-axis in the Camera frame),
+    this typically corresponds to the direction the camera is looking
+    (the representation of the world Z-axis in the camera frame).
+
+    To simplify and match your `is_inside_fov` logic, we directly use the camera's Z-axis vector:
+    Camera Z-axis direction in World Frame (Forward Vector): fwd = R_w2c_inv @ [0, 0, 1]
+    More simply, the Z-axis vector of the C2W matrix is the camera's forward vector in the world frame.
+    C2W = W2C_inv
+    """
+
+    R_c2w = R.T
+    fwd = R_c2w[:, 2]
+
+    x = fwd[0]
+    y = fwd[1]
+    z = fwd[2]
+
+    # compute yaw and pitch
+    yaw_rad = torch.atan2(x, z)
+    yaw_deg = yaw_rad * (180.0 / math.pi)
+    pitch_rad = torch.atan2(y, torch.sqrt(x**2 + z**2))
+    pitch_deg = pitch_rad * (180.0 / math.pi)
+
+    return pitch_deg, yaw_deg
+
+
+def is_inside_fov_3d_hv(
+    points: torch.Tensor,
+    center: torch.Tensor,
+    center_pitch: torch.Tensor,
+    center_yaw: torch.Tensor,
+    fov_half_h: torch.Tensor,
+    fov_half_v: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Check whether points are inside a 3D view frustum defined by a center coordinate, pitch angle, and yaw angle.
+
+    :param points: Tensor of shape (N, 3) or (N, B, 3) representing the coordinates of the sampled points.
+    :param center: Tensor of shape (3) or (B, 3) representing the camera center coordinates.
+    :param center_pitch: Tensor of shape (1) or (B) representing the pitch angle of center view direction.
+    :param center_yaw: Tensor of shape (1) or (B) representing the yaw angle of the center view direction.
+    :param fov_half_h: The horizontal half field-of-view angle (in degrees).
+    :param fov_half_v: The vertical half field-of-view angle (in degrees).
+    :return: Boolean tensor of shape (N) or (N, B), indicating whether each point is inside the FOV.
+    """
+    if points.ndim == 2:  # N, 3
+        vectors = points - center[None, :]
+        C = 1
+    elif points.ndim == 3:  # N, B, 3
+        vectors = points - center[None, ...]
+        center_pitch = center_pitch[None, :] if center_pitch.ndim == 1 else center_pitch
+        center_yaw = center_yaw[None, :] if center_yaw.ndim == 1 else center_yaw
+    else:
+        raise ValueError("points' shape should be (N, 3) or (N, B, 3)")
+
+    x = vectors[..., 0]
+    y = vectors[..., 1]
+    z = vectors[..., 2]
+
+    # Calculate the horizontal angle (yaw/azimuth), assuming the Z-axis is forward.
+    azimuth = torch.atan2(x, z) * (180 / math.pi)
+
+    # Calculate the vertical angle (pitch/elevation).
+    elevation = torch.atan2(y, torch.sqrt(x**2 + z**2)) * (180 / math.pi)
+
+    # Calculate the angular difference from the center view direction (handling angle wrapping).
+    diff_azimuth = azimuth - center_yaw
+    diff_azimuth = torch.remainder(diff_azimuth + 180, 360) - 180
+
+    diff_elevation = elevation - center_pitch
+    diff_elevation = torch.remainder(diff_elevation + 180, 360) - 180
+
+    # Check if within FOV
+    in_fov_h = diff_azimuth.abs() < fov_half_h
+    in_fov_v = diff_elevation.abs() < fov_half_v
+
+    return in_fov_h & in_fov_v
+
+
+def calculate_fov_overlap_similarity(
+    w2c_matrix_curr: torch.Tensor,
+    w2c_matrix_hist: torch.Tensor,
+    fov_h_deg: float = 105.0,
+    fov_v_deg: float = 75.0,
+    device=None,
+    points_local=None,
+) -> float:
+    """
+    Calculate the Field-of-View (FOV) overlap similarity between two W2C poses using Monte Carlo sampling.
+
+    Similarity = (Number of points in Curr_FOV ∩ Hist_FOV) / (Number of points in Curr_FOV).
+
+    :param w2c_matrix_curr: The (4, 4) W2C matrix for the current frame.
+    :param w2c_matrix_hist: The (4, 4) W2C matrix for the historical frame.
+    :param num_samples, radius, fov_h_deg, fov_v_deg: Sampling and FOV parameters.
+    :return: The overlap ratio (a float between 0.0 and 1.0).
+    """
+    w2c_matrix_curr = torch.tensor(w2c_matrix_curr, device=device)
+    w2c_matrix_hist = torch.tensor(w2c_matrix_hist, device=device)
+
+    c2w_matrix_curr = torch.linalg.inv(w2c_matrix_curr)
+    c2w_matrix_hist = torch.linalg.inv(w2c_matrix_hist)
+    C_inv = w2c_matrix_curr
+
+    w2c_matrix_curr = torch.linalg.inv(C_inv @ c2w_matrix_curr)
+    w2c_matrix_hist = torch.linalg.inv(C_inv @ c2w_matrix_hist)
+
+    R_curr, t_curr = w2c_matrix_curr[:3, :3], w2c_matrix_curr[:3, 3]
+    R_hist, t_hist = w2c_matrix_hist[:3, :3], w2c_matrix_hist[:3, 3]
+    P_w_curr = -R_curr.T @ t_curr
+    P_w_hist = -R_hist.T @ t_hist
+
+    # pitch, yaw
+    pitch_curr, yaw_curr = rotation_matrix_to_angles(R_curr)
+    pitch_hist, yaw_hist = rotation_matrix_to_angles(R_hist)
+
+    fov_half_h = torch.tensor(fov_h_deg / 2.0, device=device)
+    fov_half_v = torch.tensor(fov_v_deg / 2.0, device=device)
+
+    # move to P_w_curr (N, 3)
+    points_world = points_local + P_w_curr[None, :]
+
+    in_fov_curr = is_inside_fov_3d_hv(
+        points_world,
+        P_w_curr[None, :],
+        pitch_curr[None],
+        yaw_curr[None],
+        fov_half_h,
+        fov_half_v,
+    )
+
+    # compute based on angle
+    in_fov_hist = is_inside_fov_3d_hv(
+        points_world,
+        P_w_hist[None, :],
+        pitch_hist[None],
+        yaw_hist[None],
+        fov_half_h,
+        fov_half_v,
+    )
+
+    # compute based on distance
+    dist = torch.norm(points_world - P_w_hist.reshape(1, -1), dim=1) < 8.0
+    in_fov_hist = in_fov_hist.bool() & dist.reshape(1, -1).bool()
+
+    overlap_count = (in_fov_curr.bool() & in_fov_hist.bool()).sum().float()
+    fov_curr_count = in_fov_curr.sum().float()
+
+    if fov_curr_count == 0:
+        return 0.0
+
+    overlap_ratio = overlap_count / fov_curr_count
+
+    return overlap_ratio.item()
+
+
+def select_aligned_memory_frames(
+    w2c_list: List[np.ndarray],
+    current_frame_idx: int,
+    memory_frames: int,
+    temporal_context_size: int,
+    pred_latent_size: int,
+    pos_weight: float = 1.0,
+    ang_weight: float = 1.0,
+    device=None,
+    points_local=None,
+) -> List[int]:
+    """
+    Selects memory and context frames for a given frame based on a four-frame segment distance calculation.
+
+    :param w2c_list: List of all N 4x4 World-to-Camera (W2C) extrinsic matrices (np.ndarray).
+    :param current_frame_idx: The index of the current frame to be processed.
+    :param memory_frames: The total number of memory frames to select.
+    :param context_size: The total number of context frames to select.
+    :param pos_weight: The weight applied to the spatial (position) distance component.
+    :param ang_weight: The weight applied to the angular distance component.
+
+    :return: List[int]: A list containing the indices of the selected memory frames and context frames.
+    """
+    if current_frame_idx <= memory_frames:
+        return list(range(0, current_frame_idx))
+
+    num_total_frames = len(w2c_list)
+    if current_frame_idx >= num_total_frames or current_frame_idx < 3:
+        raise ValueError(
+            f"The current frame index must be within the valid range of w2c_list and must be at least 3."
+            f"{current_frame_idx}, {len(w2c_list)}"
+        )
+
+    start_context_idx = max(0, current_frame_idx - temporal_context_size)
+    context_frames_indices = list(range(start_context_idx, current_frame_idx))
+
+    candidate_distances = []
+    query_clip_indices = list(
+        range(
+            current_frame_idx,
+            (
+                current_frame_idx + pred_latent_size
+                if current_frame_idx + pred_latent_size <= num_total_frames
+                else num_total_frames
+            ),
+        )
+    )
+
+    historical_clip_indices = list(
+        range(4, current_frame_idx - temporal_context_size, 4)
+    )
+
+    memory_frames_indices = [0, 1, 2, 3]  # add the first chunk as context
+    memory_frames = memory_frames - temporal_context_size
+
+    for hist_idx in historical_clip_indices:
+        total_dist = 0
+        hist_w2c_1 = w2c_list[hist_idx]
+        hist_w2c_2 = w2c_list[hist_idx + 2]
+        for query_idx in query_clip_indices:
+            dist_1_for_query_idx = 1.0 - calculate_fov_overlap_similarity(
+                w2c_list[query_idx],
+                hist_w2c_1,
+                fov_h_deg=60.0,
+                fov_v_deg=35.0,
+                device=device,
+                points_local=points_local,
+            )
+            dist_2_for_query_idx = 1.0 - calculate_fov_overlap_similarity(
+                w2c_list[query_idx],
+                hist_w2c_2,
+                fov_h_deg=60.0,
+                fov_v_deg=35.0,
+                device=device,
+                points_local=points_local,
+            )
+            dist_for_query_idx = (dist_1_for_query_idx + dist_2_for_query_idx) / 2.0
+            total_dist += dist_for_query_idx
+
+        final_clip_distance = total_dist / len(query_clip_indices)
+        candidate_distances.append((hist_idx, final_clip_distance))
+
+    candidate_distances.sort(key=lambda x: x[1])
+
+    for start_idx, _ in candidate_distances:
+        # check the memory frame number
+        if len(memory_frames_indices) >= memory_frames:
+            break
+
+        if start_idx not in memory_frames_indices:
+            memory_frames_indices.extend(range(start_idx, start_idx + 4))
+
+    # exclude the repeated frames
+    selected_frames_set = set(context_frames_indices)
+    selected_frames_set.update(memory_frames_indices)
+
+    final_selected_frames = sorted(list(selected_frames_set))
+
+    return final_selected_frames

fastvideo/models/dits/hyworld/trajectory.py

@@ -0,0 +1,112 @@
+# HY-WorldPlay/hyvideo/generate_custom_trajectory.py
+
+import numpy as np
+import json
+
+
+def rot_x(theta):
+    c, s = np.cos(theta), np.sin(theta)
+    return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])
+
+
+def rot_y(theta):
+    c, s = np.cos(theta), np.sin(theta)
+    return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])
+
+
+def rot_z(theta):
+    c, s = np.cos(theta), np.sin(theta)
+    return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
+
+
+def generate_camera_trajectory_local(motions):
+    """
+    motions: list of dict
+             {"forward": 1.0}, {"yaw": np.pi/2}, {"pitch": np.pi/6}, {"right": 1.0}
+             - forward: Translation (Forward or Backward)
+             - yaw:   Rotate (Left or Right)
+             - pitch: Rotate (Up or Down)
+             - right: Translation (Right or Left)
+             - third_yaw: Third Perspective Rotate (Left or Right)
+    """
+
+    poses = []
+    T = np.eye(4)
+    poses.append(T.copy())
+
+    for move in motions:
+        # Rotate (Left or Right)
+        if "yaw" in move:
+            R = rot_y(move["yaw"])
+            T[:3, :3] = T[:3, :3] @ R
+
+        # Rotate (Up or Down)
+        if "pitch" in move:
+            R = rot_x(move["pitch"])
+            T[:3, :3] = T[:3, :3] @ R
+
+        # Translation (Z-direction of the camera's local coordinate system)
+        forward = move.get("forward", 0.0)
+        if forward != 0:
+            local_t = np.array([0, 0, forward])
+            world_t = T[:3, :3] @ local_t
+            T[:3, 3] += world_t
+
+        # Translation (Z-direction of the camera's local coordinate system)
+        right = move.get("right", 0.0)
+        if right != 0:
+            local_t = np.array([right, 0, 0])
+            world_t = T[:3, :3] @ local_t
+            T[:3, 3] += world_t
+
+        # Third Perspective Rotate (Left or Right)
+        third_yaw = move.get("third_yaw", 0.0)
+        if third_yaw != 0:
+            theta = -third_yaw
+            C = np.array([[1, 0.0, 0, 0], [0, 1, 0, 0], [0, 0, 1, -1.0], [0, 0, 0, 1]])
+            c_origin = C.copy()
+            # Rotation around the Y-axis
+            R_y = np.array(
+                [
+                    [np.cos(theta), 0, np.sin(theta)],
+                    [0, 1, 0],
+                    [-np.sin(theta), 0, np.cos(theta)],
+                ]
+            )
+            # Translation
+            C[:3, :3] = C[:3, :3] @ R_y
+            C[:3, 3] = R_y @ C[:3, 3]
+            c_inv = np.linalg.inv(c_origin)
+            c_relative = c_inv @ C
+            T = T @ c_relative
+
+        poses.append(T.copy())
+
+    return poses
+
+
+if __name__ == "__main__":
+    # Examples: Forward 0.08 * 16 -> Right Rotate 3 degree * 16
+    motions = []
+    for i in range(15):
+        motions.append({"forward": 0.08})
+
+    for i in range(16):
+        motions.append({"yaw": np.deg2rad(3)})
+
+    intrinsic = [
+        [969.6969696969696, 0.0, 960.0],
+        [0.0, 969.6969696969696, 540.0],
+        [0.0, 0.0, 1.0],
+    ]
+
+    poses = generate_camera_trajectory_local(motions)
+    custom_c2w = {}
+    for i, p in enumerate(poses):
+        custom_c2w[str(i)] = {"extrinsic": p.tolist(), "K": intrinsic}
+        json.dump(
+            custom_c2w,
+            open("./assets/pose/pose.json", "w"),
+            indent=4,
+            ensure_ascii=False,
+        )

fastvideo/models/dits/wanvideo.py

@@ -734,18 +734,8 @@ class WanTransformer3DModel(CachableDiT):
                         block, hidden_states, encoder_hidden_states,
                         timestep_proj, freqs_cis, attention_mask)
             else:
-                offload_mgr = getattr(self, "_layerwise_offload_manager", None)
-                use_offload = offload_mgr is not None and getattr(offload_mgr, "enabled", False)
-
-                for i, block in enumerate(self.blocks):
-                    scope = offload_mgr.layer_scope(
-                        prefetch_layer_idx=i + 1 if i + 1 < len(self.blocks) else None,
-                        release_layer_idx=i,
-                        non_blocking=True,
-                    ) if use_offload else nullcontext()
-                    
-                    with scope:
-                        hidden_states = block(hidden_states, encoder_hidden_states,
+                for block in self.blocks:
+                    hidden_states = block(hidden_states, encoder_hidden_states,
                                               timestep_proj, freqs_cis, attention_mask)
             # if teacache is enabled, we need to cache the original hidden states
 

fastvideo/models/encoders/gemma.py

@@ -0,0 +1,563 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from dataclasses import dataclass
+import os
+from typing import Iterable
+
+import torch
+from torch import nn
+from transformers import Gemma3ForConditionalGeneration
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.dits.ltx2 import (
+    FeedForward,
+    LTXRopeType,
+    apply_ltx_rotary_emb,
+    generate_ltx_freq_grid_np,
+    generate_ltx_freq_grid_pytorch,
+    precompute_ltx_freqs_cis,
+)
+from fastvideo.models.loader.weight_utils import default_weight_loader
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _debug_log_line(message: str) -> None:
+    if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1":
+        return
+    log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "")
+    if not log_path:
+        return
+    log_dir = os.path.dirname(log_path)
+    if log_dir:
+        os.makedirs(log_dir, exist_ok=True)
+    with open(log_path, "a", encoding="utf-8") as f:
+        f.write(message + "\n")
+
+
+def _debug_gemma_log_line(message: str) -> None:
+    log_path = os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", "")
+    if not log_path:
+        return
+    log_dir = os.path.dirname(log_path)
+    if log_dir:
+        os.makedirs(log_dir, exist_ok=True)
+    with open(log_path, "a", encoding="utf-8") as f:
+        f.write(message + "\n")
+
+
+@dataclass(frozen=True)
+class GemmaConnectorConfig:
+    num_attention_heads: int
+    attention_head_dim: int
+    num_layers: int
+    positional_embedding_theta: float
+    positional_embedding_max_pos: list[int]
+    rope_type: LTXRopeType
+    double_precision_rope: bool
+    num_learnable_registers: int | None
+
+
+class GemmaFeaturesExtractorProjLinear(nn.Module):
+    """Linear projection that aggregates stacked Gemma hidden states."""
+
+    def __init__(self, in_features: int, out_features: int) -> None:
+        super().__init__()
+        self.aggregate_embed = nn.Linear(in_features, out_features, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.aggregate_embed(x)
+
+
+class _BasicTransformerBlock1D(nn.Module):
+    """1D transformer block for connector processing."""
+
+    def __init__(
+        self,
+        dim: int,
+        heads: int,
+        dim_head: int,
+        rope_type: LTXRopeType,
+        norm_eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.attn1 = _GemmaAttention(
+            query_dim=dim,
+            context_dim=None,
+            heads=heads,
+            dim_head=dim_head,
+            norm_eps=norm_eps,
+            rope_type=rope_type,
+        )
+        self.ff = FeedForward(dim, dim_out=dim)
+        self.norm_eps = norm_eps
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        norm_hidden_states = torch.nn.functional.rms_norm(
+            hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps
+        )
+        if norm_hidden_states.ndim == 4:
+            norm_hidden_states = norm_hidden_states.squeeze(1)
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            mask=attention_mask,
+            pe=pe,
+        )
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        norm_hidden_states = torch.nn.functional.rms_norm(
+            hidden_states, (hidden_states.shape[-1],), eps=self.norm_eps
+        )
+        ff_output = self.ff(norm_hidden_states)
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+        return hidden_states
+
+
+class _GemmaAttention(nn.Module):
+    """Attention implementation aligned with LTX-2 text encoder."""
+
+    def __init__(
+        self,
+        query_dim: int,
+        context_dim: int | None,
+        heads: int,
+        dim_head: int,
+        norm_eps: float,
+        rope_type: LTXRopeType,
+    ) -> None:
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = query_dim if context_dim is None else context_dim
+
+        self.heads = heads
+        self.dim_head = dim_head
+        self.rope_type = rope_type
+
+        self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+        self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=True)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=True)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=True)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+        k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        q = self.to_q(x)
+        context = x if context is None else context
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if pe is not None:
+            q = apply_ltx_rotary_emb(q, pe, self.rope_type)
+            k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
+
+        b, q_len, _ = q.shape
+        k_len = k.shape[1]
+        q = q.view(b, q_len, self.heads, self.dim_head).transpose(1, 2)
+        k = k.view(b, k_len, self.heads, self.dim_head).transpose(1, 2)
+        v = v.view(b, k_len, self.heads, self.dim_head).transpose(1, 2)
+
+        if mask is not None:
+            if mask.ndim == 2:
+                mask = mask.unsqueeze(0)
+            if mask.ndim == 3:
+                mask = mask.unsqueeze(1)
+
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=mask,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+        out = out.transpose(1, 2).reshape(b, q_len, -1)
+        return self.to_out(out)
+
+
+class Embeddings1DConnector(nn.Module):
+    """Transformer connector that refines Gemma embeddings for LTX-2."""
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(self, config: GemmaConnectorConfig) -> None:
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.positional_embedding_theta = config.positional_embedding_theta
+        self.positional_embedding_max_pos = config.positional_embedding_max_pos
+        self.rope_type = config.rope_type
+        self.double_precision_rope = config.double_precision_rope
+        self.transformer_1d_blocks = nn.ModuleList(
+            [
+                _BasicTransformerBlock1D(
+                    dim=self.inner_dim,
+                    heads=config.num_attention_heads,
+                    dim_head=config.attention_head_dim,
+                    rope_type=config.rope_type,
+                )
+                for _ in range(config.num_layers)
+            ]
+        )
+        self.num_learnable_registers = config.num_learnable_registers
+        if self.num_learnable_registers:
+            self.learnable_registers = nn.Parameter(
+                torch.rand(
+                    self.num_learnable_registers,
+                    self.inner_dim,
+                    dtype=torch.bfloat16,
+                )
+                * 2.0
+                - 1.0
+            )
+
+    def _replace_padded_with_learnable_registers(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
+            f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by "
+            f"num_learnable_registers {self.num_learnable_registers}."
+        )
+
+        num_registers_duplications = (
+            hidden_states.shape[1] // self.num_learnable_registers
+        )
+        learnable_registers = torch.tile(
+            self.learnable_registers, (num_registers_duplications, 1)
+        )
+        attention_mask_binary = (
+            attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0
+        ).int()
+
+        non_zero_hidden_states = hidden_states[
+            :, attention_mask_binary.squeeze().bool(), :
+        ]
+        non_zero_nums = non_zero_hidden_states.shape[1]
+        pad_length = hidden_states.shape[1] - non_zero_nums
+        adjusted_hidden_states = torch.nn.functional.pad(
+            non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0
+        )
+        flipped_mask = torch.flip(attention_mask_binary, dims=[1])
+        hidden_states = flipped_mask * adjusted_hidden_states + (
+            1 - flipped_mask
+        ) * learnable_registers
+
+        attention_mask = torch.full_like(
+            attention_mask,
+            0.0,
+            dtype=attention_mask.dtype,
+            device=attention_mask.device,
+        )
+
+        return hidden_states, attention_mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.num_learnable_registers:
+            hidden_states, attention_mask = (
+                self._replace_padded_with_learnable_registers(
+                    hidden_states, attention_mask
+                )
+            )
+
+        indices_grid = torch.arange(
+            hidden_states.shape[1],
+            dtype=torch.float32,
+            device=hidden_states.device,
+        )
+        indices_grid = indices_grid[None, None, :]
+        freq_grid_generator = (
+            generate_ltx_freq_grid_np
+            if self.double_precision_rope
+            else generate_ltx_freq_grid_pytorch
+        )
+        freqs_cis = precompute_ltx_freqs_cis(
+            indices_grid=indices_grid,
+            dim=self.inner_dim,
+            out_dtype=hidden_states.dtype,
+            theta=self.positional_embedding_theta,
+            max_pos=self.positional_embedding_max_pos,
+            num_attention_heads=self.num_attention_heads,
+            rope_type=self.rope_type,
+            freq_grid_generator=freq_grid_generator,
+        )
+
+        for block in self.transformer_1d_blocks:
+            hidden_states = block(
+                hidden_states, attention_mask=attention_mask, pe=freqs_cis
+            )
+
+        hidden_states = torch.nn.functional.rms_norm(
+            hidden_states, (hidden_states.shape[-1],), eps=1e-6
+        )
+
+        return hidden_states, attention_mask
+
+
+class LTX2GemmaTextEncoderModel(TextEncoder):
+    _supported_attention_backends = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+
+    def __init__(self, config: TextEncoderConfig) -> None:
+        super().__init__(config)
+        arch = config.arch_config
+
+        self.feature_extractor_linear = GemmaFeaturesExtractorProjLinear(
+            in_features=arch.feature_extractor_in_features,
+            out_features=arch.feature_extractor_out_features,
+        )
+
+        connector_config = GemmaConnectorConfig(
+            num_attention_heads=arch.connector_num_attention_heads,
+            attention_head_dim=arch.connector_attention_head_dim,
+            num_layers=arch.connector_num_layers,
+            positional_embedding_theta=arch.connector_positional_embedding_theta,
+            positional_embedding_max_pos=arch.connector_positional_embedding_max_pos,
+            rope_type=LTXRopeType(arch.connector_rope_type),
+            double_precision_rope=arch.connector_double_precision_rope,
+            num_learnable_registers=arch.connector_num_learnable_registers,
+        )
+        self.embeddings_connector = Embeddings1DConnector(connector_config)
+        self.audio_embeddings_connector = Embeddings1DConnector(connector_config)
+
+        self.gemma_model_path = arch.gemma_model_path
+        self.gemma_dtype = arch.gemma_dtype
+        self.padding_side = arch.padding_side
+        self._gemma_model: Gemma3ForConditionalGeneration | None = None
+
+    def named_parameters(self, prefix: str = "", recurse: bool = True):
+        for name, param in super().named_parameters(
+            prefix=prefix, recurse=recurse
+        ):
+            if name.startswith("gemma_model."):
+                continue
+            yield name, param
+
+    @property
+    def gemma_model(self) -> Gemma3ForConditionalGeneration:
+        if self._gemma_model is None:
+            gemma_path = self.gemma_model_path
+            if not gemma_path:
+                raise ValueError(
+                    "gemma_model_path must be set (expected text_encoder/gemma)."
+                )
+            dtype = getattr(torch, self.gemma_dtype, torch.bfloat16)
+            self._gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
+                gemma_path,
+                local_files_only=True,
+                torch_dtype=dtype,
+            )
+            # Configure model-level attention implementation when using TORCH_SDPA.
+            # Note: torch.backends.cuda.enable_*_sdp() settings should be configured
+            # at application/pipeline initialization level, not here, to avoid
+            # unexpected side effects across the application.
+            if os.getenv("FASTVIDEO_ATTENTION_BACKEND") == "TORCH_SDPA":
+                if hasattr(self._gemma_model.config, "attn_implementation"):
+                    self._gemma_model.config.attn_implementation = "sdpa"
+                if hasattr(self._gemma_model.config, "_attn_implementation"):
+                    self._gemma_model.config._attn_implementation = "sdpa"
+            device = next(self.feature_extractor_linear.parameters()).device
+            self._gemma_model.to(device=device)
+            self._gemma_model.eval()
+        return self._gemma_model
+
+    def _run_feature_extractor(
+        self,
+        hidden_states: tuple[torch.Tensor, ...],
+        attention_mask: torch.Tensor,
+        padding_side: str,
+    ) -> torch.Tensor:
+        encoded_text_features = torch.stack(hidden_states, dim=-1)
+        if os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", ""):
+            for idx, layer in enumerate(hidden_states):
+                _debug_gemma_log_line(
+                    f"fastvideo:gemma_hidden_state_{idx}"
+                    f":sum={layer.float().sum().item():.6f}"
+                )
+            _debug_gemma_log_line(
+                "fastvideo:gemma_hidden_states_stack"
+                f":sum={encoded_text_features.float().sum().item():.6f}"
+            )
+        encoded_text_features_dtype = encoded_text_features.dtype
+        sequence_lengths = attention_mask.sum(dim=-1)
+        normed_text_features = _norm_and_concat_padded_batch(
+            encoded_text_features, sequence_lengths, padding_side=padding_side
+        )
+        return self.feature_extractor_linear(
+            normed_text_features.to(encoded_text_features_dtype)
+        )
+
+    def _convert_to_additive_mask(
+        self, attention_mask: torch.Tensor, dtype: torch.dtype
+    ) -> torch.Tensor:
+        return (attention_mask - 1).to(dtype).reshape(
+            (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
+        ) * torch.finfo(dtype).max
+
+    def _run_connectors(
+        self,
+        encoded_input: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        connector_attention_mask = self._convert_to_additive_mask(
+            attention_mask, encoded_input.dtype
+        )
+        encoded, encoded_connector_attention_mask = self.embeddings_connector(
+            encoded_input, connector_attention_mask
+        )
+
+        attention_mask = (encoded_connector_attention_mask < 0.000001).to(
+            torch.int64
+        )
+        attention_mask = attention_mask.reshape(
+            [encoded.shape[0], encoded.shape[1], 1]
+        )
+        encoded = encoded * attention_mask
+
+        encoded_for_audio, _ = self.audio_embeddings_connector(
+            encoded_input, connector_attention_mask
+        )
+
+        return encoded, encoded_for_audio, attention_mask.squeeze(-1)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        if input_ids is None:
+            raise ValueError("input_ids is required for Gemma text encoding.")
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        model = self.gemma_model
+        input_ids = input_ids.to(device=model.device)
+        attention_mask = attention_mask.to(device=model.device)
+        outputs = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        encoded_inputs = self._run_feature_extractor(
+            outputs.hidden_states,
+            attention_mask,
+            padding_side=self.padding_side,
+        )
+        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
+            _debug_log_line(
+                "fastvideo:gemma_feature"
+                f":sum={encoded_inputs.float().sum().item():.6f} "
+                f"shape={tuple(encoded_inputs.shape)}"
+            )
+        video_encoding, audio_encoding, attention_mask = self._run_connectors(
+            encoded_inputs, attention_mask
+        )
+        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
+            _debug_log_line(
+                "fastvideo:gemma_video_encoding"
+                f":sum={video_encoding.float().sum().item():.6f} "
+                f"shape={tuple(video_encoding.shape)}"
+            )
+            _debug_log_line(
+                "fastvideo:gemma_audio_encoding"
+                f":sum={audio_encoding.float().sum().item():.6f} "
+                f"shape={tuple(audio_encoding.shape)}"
+            )
+
+        hidden_states = (audio_encoding, ) if output_hidden_states else None
+        return BaseEncoderOutput(
+            last_hidden_state=video_encoding,
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
+
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if name == "aggregate_embed.weight":
+                name = "feature_extractor_linear.aggregate_embed.weight"
+            if name not in params_dict:
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+def _norm_and_concat_padded_batch(
+    encoded_text: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    padding_side: str = "right",
+) -> torch.Tensor:
+    b, t, d, l = encoded_text.shape
+    device = encoded_text.device
+
+    token_indices = torch.arange(t, device=device)[None, :]
+    if padding_side == "right":
+        mask = token_indices < sequence_lengths[:, None]
+    elif padding_side == "left":
+        start_indices = t - sequence_lengths[:, None]
+        mask = token_indices >= start_indices
+    else:
+        raise ValueError(
+            f"padding_side must be 'left' or 'right', got {padding_side}"
+        )
+
+    mask = mask.reshape(b, t, 1, 1)
+    eps = 1e-6
+
+    masked = encoded_text.masked_fill(~mask, 0.0)
+    denom = (sequence_lengths * d).view(b, 1, 1, 1)
+    mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
+
+    x_min = encoded_text.masked_fill(~mask, float("inf")).amin(
+        dim=(1, 2), keepdim=True
+    )
+    x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
+        dim=(1, 2), keepdim=True
+    )
+    range_ = x_max - x_min
+
+    normed = 8 * (encoded_text - mean) / (range_ + eps)
+    normed = normed.reshape(b, t, -1)
+
+    mask_flattened = mask.reshape(b, t, 1).expand(-1, -1, d * l)
+    normed = normed.masked_fill(~mask_flattened, 0.0)
+    return normed

fastvideo/models/encoders/reason1.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reason1 (Qwen2.5-VL) text encoder."""
+
+import os
+from dataclasses import dataclass
+from collections.abc import Iterable
+
+import torch
+from transformers import AutoProcessor
+
+from fastvideo.configs.models.encoders import BaseEncoderOutput, Reason1Config
+from fastvideo.logger import init_logger
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+from fastvideo.platforms import AttentionBackendEnum
+
+from fastvideo.models.encoders.qwen2_5_vl_custom import (
+    Qwen2_5_VLForConditionalGenerationSimple,
+    Qwen2_5_VLConfig,
+    get_rope_index,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class _WeightsSource:
+    """Mimic `TextEncoderLoader.Source` (avoid import cycles)."""
+
+    model_or_path: str
+    prefix: str = ""
+    fall_back_to_pt: bool = True
+    allow_patterns_overrides: list[str] | None = None
+
+
+
+
+class Reason1TextEncoder(TextEncoder):
+    """Reason1 (Qwen2.5-VL) text encoder."""
+
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
+        AttentionBackendEnum.FLASH_ATTN,
+        AttentionBackendEnum.TORCH_SDPA,
+    )
+    
+    def __init__(self, config: Reason1Config, prefix: str = "", checkpoint_path: str | None = None):
+        super().__init__(config)
+        
+        self.prefix = prefix
+        self.quant_config = None  # For future quantization support
+        
+        self.embedding_concat_strategy = config.arch_config.embedding_concat_strategy
+        self.n_layers_per_group = config.arch_config.n_layers_per_group
+        self.num_embedding_padding_tokens = config.arch_config.num_embedding_padding_tokens
+        
+        config_path = checkpoint_path if checkpoint_path else config.tokenizer_type
+        
+        logger.info("Initializing Reason1TextEncoder (Qwen2.5-VL) from %s", config_path)
+        try:
+            from transformers import AutoConfig as HFAutoConfig
+            hf_config = HFAutoConfig.from_pretrained(
+                config_path,
+                trust_remote_code=True,
+            )
+        except Exception as e:
+            logger.warning("Failed to load HF config from %s (%s). Using default Qwen2.5-VL-7B config.",
+                           config_path, e)
+            hf_config = Qwen2_5_VLConfig(
+                    hidden_size=3584,
+                    intermediate_size=18944,
+                    max_window_layers=28,
+                    num_attention_heads=28,
+                    num_hidden_layers=28,
+                    num_key_value_heads=4,
+                    tie_word_embeddings=False,
+                    vocab_size=152064,
+            )
+        
+        hf_config.output_hidden_states = True
+        
+        if hasattr(config.arch_config, '_attn_implementation') and config.arch_config._attn_implementation:
+            hf_config._attn_implementation = config.arch_config._attn_implementation
+        else:
+            hf_config._attn_implementation = "flash_attention_2"
+        logger.info("Reason1 attention implementation: %s", getattr(hf_config, "_attn_implementation", None))
+        
+        with torch.device("meta"):
+            self.model = Qwen2_5_VLForConditionalGenerationSimple(hf_config)
+        
+        self.processor = AutoProcessor.from_pretrained(
+            config_path,
+            trust_remote_code=True,
+        )
+
+        weights_override = os.getenv("FASTVIDEO_REASON1_WEIGHTS_PATH")
+        if weights_override:
+            self.secondary_weights = (
+                _WeightsSource(
+                    model_or_path=weights_override,
+                    prefix="",
+                    fall_back_to_pt=True,
+                    allow_patterns_overrides=None,
+                ),
+            )
+            logger.info("Reason1TextEncoder: overlaying weights from %s", weights_override)
+        
+        self._weights_loaded = False
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        # Cosmos2.5 alignment: keep attention_mask=None.
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=None,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=True,
+            return_dict=True,
+            pixel_values=kwargs.get('pixel_values', None),
+            pixel_values_videos=kwargs.get('pixel_values_videos', None),
+            image_grid_thw=kwargs.get('image_grid_thw', None),
+            video_grid_thw=kwargs.get('video_grid_thw', None),
+        )
+        
+        hidden_states = outputs.hidden_states
+        last_hidden_state = hidden_states[-1]
+        
+        return BaseEncoderOutput(
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states if output_hidden_states else None,
+            attention_mask=None,
+        )
+    
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        first_weight = None
+        weights_list = []
+        for name, weight in weights:
+            if first_weight is None:
+                first_weight = weight
+                self.model = self.model.to_empty(device=weight.device)
+                self.model.init_weights(buffer_device=weight.device)
+            weights_list.append((name, weight))
+        
+        params_dict = dict(self.model.named_parameters())
+        loaded_params: set[str] = set()
+        skipped_weights = {"lm_head": 0, "visual": 0, "decoder": 0}
+        
+        for name, loaded_weight in weights_list:
+            if "lm_head" in name:
+                skipped_weights["lm_head"] += 1
+                continue
+            if "visual" in name:
+                skipped_weights["visual"] += 1
+                continue
+            if "decoder" in name:
+                skipped_weights["decoder"] += 1
+                continue
+            
+            # Handle stacked params mapping (for quantized models)
+            for param_name, weight_name, shard_id in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                
+                if name not in params_dict:
+                    continue
+                
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
+                loaded_params.add(param_name_with_prefix)
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                
+                if name not in params_dict:
+                    continue
+                
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+                param_name_with_prefix = f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}"
+                loaded_params.add(param_name_with_prefix)
+        if first_weight is not None:
+            self.model = self.model.to(first_weight.device)
+        
+        all_params = set(f"model.{name}" if self.prefix == "" else f"{self.prefix}.{name}" 
+                        for name in params_dict.keys())
+        loaded_params.update(all_params)
+        
+        # Mark weights as loaded
+        self._weights_loaded = True
+        return loaded_params
+    
+    def compute_text_embeddings_online(
+        self,
+        data_batch: dict[str, list[str]],
+        input_caption_key: str,
+    ) -> torch.Tensor:
+        prompts = data_batch[input_caption_key]
+        return self.compute_text_embeddings(prompts)
+    
+    def compute_text_embeddings(
+        self,
+        prompts: list[str],
+        device: str | torch.device = "cuda",
+    ) -> torch.Tensor:
+        """Compute embeddings for a list of prompts."""
+        input_ids_batch = []
+        
+        tok = getattr(self.processor, "tokenizer", None)
+        if tok is None:
+            raise RuntimeError("Reason1TextEncoder requires processor.tokenizer")
+        pad_id = getattr(tok, "pad_id", None)
+        if pad_id is None:
+            pad_id = getattr(tok, "pad_token_id", None)
+        if pad_id is None:
+            pad_id = getattr(self.model.config, "pad_token_id", None)
+        if pad_id is None:
+            pad_id = 0
+        
+        for prompt in prompts:
+            conversations = [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "You are a helpful assistant who will provide prompts to an image generator.",
+                        }
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": prompt,
+                        }
+                    ],
+                },
+            ]
+            
+            try:
+                tokenizer_output = tok.apply_chat_template(
+                    conversations,
+                    tokenize=True,
+                    add_generation_prompt=False,
+                    add_vision_id=False,
+                )
+            except TypeError:
+                tokenizer_output = tok.apply_chat_template(
+                conversations,
+                tokenize=True,
+                add_generation_prompt=False,
+            )
+
+            if isinstance(tokenizer_output, dict) and "input_ids" in tokenizer_output:
+                input_ids = tokenizer_output["input_ids"]
+                if hasattr(input_ids, "tolist"):
+                    input_ids = input_ids.tolist()
+            else:
+                input_ids = tokenizer_output
+                if hasattr(input_ids, "tolist"):
+                    input_ids = input_ids.tolist()
+                if isinstance(input_ids, list) and len(input_ids) == 1 and isinstance(
+                        input_ids[0], list):
+                    input_ids = input_ids[0]
+                if not isinstance(input_ids, list):
+                    raise RuntimeError(
+                        f"Unexpected chat_template output type: {type(tokenizer_output)}"
+                    )
+            
+            if self.num_embedding_padding_tokens > len(input_ids):
+                pad_len = self.num_embedding_padding_tokens - len(input_ids)
+                input_ids = input_ids + [pad_id] * pad_len
+            else:
+                input_ids = input_ids[:self.num_embedding_padding_tokens]
+            
+            input_ids = torch.LongTensor(input_ids).to(device=device)
+            input_ids_batch.append(input_ids)
+        
+        input_ids_batch = torch.stack(input_ids_batch, dim=0)
+
+        # Cosmos2.5 alignment: keep attention_mask=None.
+        target_device = input_ids_batch.device
+        try:
+            embed_device = self.model.model.embed_tokens.weight.device  # type: ignore[attr-defined]
+        except Exception:
+            embed_device = None
+        if embed_device is not None and embed_device != target_device:
+            self.model = self.model.to(target_device)
+        
+        with torch.no_grad():
+            position_ids, _ = get_rope_index(
+                self.model.config,
+                input_ids_batch,
+                image_grid_thw=None,
+                video_grid_thw=None,
+                second_per_grid_ts=None,
+                attention_mask=None,
+            )
+            position_ids = position_ids.to(target_device)
+                
+            outputs = self.model.model(
+                input_ids=input_ids_batch,
+                position_ids=position_ids,
+                attention_mask=None,
+                output_hidden_states=True,
+                return_dict=True,
+                use_cache=False,
+            )
+            hidden_states = outputs.hidden_states
+        
+        normalized_hidden_states = []
+        for layer_idx in range(1, len(hidden_states)):
+            normalized_state = self._mean_normalize(hidden_states[layer_idx])
+            normalized_hidden_states.append(normalized_state)
+        
+        if self.embedding_concat_strategy == "full_concat":
+            text_embeddings = torch.cat(normalized_hidden_states, dim=-1)
+        elif self.embedding_concat_strategy == "mean_pooling":
+            text_embeddings = torch.stack(normalized_hidden_states).mean(dim=0)
+        elif self.embedding_concat_strategy == "pool_every_n_layers_and_concat":
+            pooled_embeddings = []
+            for i in range(0, len(normalized_hidden_states), self.n_layers_per_group):
+                group = normalized_hidden_states[i : i + self.n_layers_per_group]
+                pooled = torch.stack(group).mean(dim=0)
+                pooled_embeddings.append(pooled)
+            text_embeddings = torch.cat(pooled_embeddings, dim=-1)
+        else:
+            raise ValueError(
+                f"Unknown embedding_concat_strategy: {self.embedding_concat_strategy}"
+            )
+        
+        return text_embeddings
+    @staticmethod
+    def _mean_normalize(tensor: torch.Tensor) -> torch.Tensor:
+        return (tensor - tensor.mean(dim=-1, keepdim=True)) / (
+            tensor.std(dim=-1, keepdim=True) + 1e-8
+        )
+

fastvideo/models/encoders/siglip.py

@@ -0,0 +1,419 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+SigLIP Vision Encoder for FastVideo.
+
+SigLIP (Sigmoid Loss for Language-Image Pre-training) is similar to CLIP
+but uses sigmoid loss instead of contrastive loss, and doesn't use a CLS token.
+"""
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.siglip import SiglipVisionArchConfig, SiglipVisionConfig
+from fastvideo.distributed import divide, get_tp_world_size
+from fastvideo.layers.activation import get_act_fn
+from fastvideo.layers.linear import (ColumnParallelLinear, QKVParallelLinear,
+                                     RowParallelLinear)
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.logger import init_logger
+from fastvideo.models.encoders.base import ImageEncoder
+from fastvideo.models.loader.weight_utils import default_weight_loader
+
+logger = init_logger(__name__)
+
+
+class SiglipVisionEmbeddings(nn.Module):
+    """
+    SigLIP vision embeddings - similar to CLIP but without class embedding.
+    """
+
+    def __init__(self, arch_config: SiglipVisionArchConfig):
+        super().__init__()
+        self.arch_config = arch_config
+        self.embed_dim = arch_config.hidden_size
+        self.image_size = arch_config.image_size
+        self.patch_size = arch_config.patch_size
+        # SigLIP uses valid padding, so non-divisible sizes work (edge pixels are ignored)
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=arch_config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",  # SigLIP uses valid padding
+        )
+
+        # Integer division - with valid padding, edge pixels are ignored
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+        self.register_buffer(
+            "position_ids",
+            torch.arange(self.num_positions).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(
+            pixel_values.to(dtype=target_dtype)
+        )  # shape = [*, embed_dim, grid, grid]
+        embeddings = patch_embeds.flatten(2).transpose(1, 2)
+        embeddings = embeddings + self.position_embedding(self.position_ids)
+        return embeddings
+
+
+class SiglipAttention(nn.Module):
+    """Multi-headed attention for SigLIP."""
+
+    def __init__(
+        self,
+        arch_config: SiglipVisionArchConfig,
+        enable_scale: bool = True,
+        is_causal: bool = False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.arch_config = arch_config
+        self.embed_dim = arch_config.hidden_size
+        self.num_heads = arch_config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads "
+                f"(got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
+            )
+        
+        self.scale = self.head_dim**-0.5 if enable_scale else None
+        self.dropout = arch_config.attention_dropout
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.out_proj = RowParallelLinear(
+            input_size=self.embed_dim,
+            output_size=self.embed_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+
+        self.tp_size = get_tp_world_size()
+        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
+
+        self.attn = LocalAttention(
+            self.num_heads_per_partition,
+            self.head_dim,
+            self.num_heads_per_partition,
+            softmax_scale=self.scale,
+            causal=is_causal,
+            supported_attention_backends=arch_config._supported_attention_backends,
+        )
+
+    def forward(self, hidden_states: torch.Tensor):
+        """Input shape: Batch x Time x Channel"""
+        qkv_states, _ = self.qkv_proj(hidden_states)
+        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)
+        
+        query_states = query_states.reshape(
+            query_states.shape[0], query_states.shape[1],
+            self.num_heads_per_partition, self.head_dim
+        )
+        key_states = key_states.reshape(
+            key_states.shape[0], key_states.shape[1],
+            self.num_heads_per_partition, self.head_dim
+        )
+        value_states = value_states.reshape(
+            value_states.shape[0], value_states.shape[1],
+            self.num_heads_per_partition, self.head_dim
+        )
+        
+        attn_output = self.attn(query_states, key_states, value_states)
+        attn_output = attn_output.reshape(
+            attn_output.shape[0], attn_output.shape[1],
+            self.num_heads_per_partition * self.head_dim
+        )
+        attn_output, _ = self.out_proj(attn_output)
+        return attn_output, None
+
+
+class SiglipMLP(nn.Module):
+
+    def __init__(
+        self,
+        arch_config: SiglipVisionArchConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.arch_config = arch_config
+        self.activation_fn = get_act_fn(arch_config.hidden_act)
+        self.fc1 = ColumnParallelLinear(
+            arch_config.hidden_size,
+            arch_config.intermediate_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
+        )
+        self.fc2 = RowParallelLinear(
+            arch_config.intermediate_size,
+            arch_config.hidden_size,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states, _ = self.fc2(hidden_states)
+        return hidden_states
+
+
+class SiglipEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        arch_config: SiglipVisionArchConfig,
+        enable_scale: bool = True,
+        is_causal: bool = False,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = SiglipAttention(
+            arch_config,
+            enable_scale=enable_scale,
+            is_causal=is_causal,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.layer_norm1 = nn.LayerNorm(
+            arch_config.hidden_size, eps=arch_config.layer_norm_eps
+        )
+        self.mlp = SiglipMLP(
+            arch_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.layer_norm2 = nn.LayerNorm(
+            arch_config.hidden_size, eps=arch_config.layer_norm_eps
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # SigLIP uses post-norm (like original ViT)
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class SiglipEncoder(nn.Module):
+    """SigLIP encoder consisting of transformer layers."""
+
+    def __init__(
+        self,
+        arch_config: SiglipVisionArchConfig,
+        enable_scale: bool = True,
+        is_causal: bool = False,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.arch_config = arch_config
+
+        if num_hidden_layers_override is None:
+            num_hidden_layers = arch_config.num_hidden_layers
+        else:
+            num_hidden_layers = num_hidden_layers_override
+            
+        self.layers = nn.ModuleList([
+            SiglipEncoderLayer(
+                arch_config=arch_config,
+                enable_scale=enable_scale,
+                is_causal=is_causal,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{layer_idx}",
+            )
+            for layer_idx in range(num_hidden_layers)
+        ])
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        return_all_hidden_states: bool,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        hidden_states_pool = [inputs_embeds]
+        hidden_states = inputs_embeds
+
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(hidden_states)
+            if return_all_hidden_states:
+                hidden_states_pool.append(hidden_states)
+
+        if return_all_hidden_states:
+            return hidden_states_pool
+        return [hidden_states]
+
+
+class SiglipVisionTransformer(nn.Module):
+
+    def __init__(
+        self,
+        config: SiglipVisionConfig,
+        quant_config: QuantizationConfig | None = None,
+        num_hidden_layers_override: int | None = None,
+        require_post_norm: bool | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        arch_config = config.arch_config
+        embed_dim = arch_config.hidden_size
+
+        self.embeddings = SiglipVisionEmbeddings(arch_config)
+
+        self.encoder = SiglipEncoder(
+            arch_config=arch_config,
+            enable_scale=config.enable_scale,
+            is_causal=config.is_causal,
+            quant_config=quant_config,
+            num_hidden_layers_override=num_hidden_layers_override,
+            prefix=f"{prefix}.encoder",
+        )
+
+        num_hidden_layers = arch_config.num_hidden_layers
+        if len(self.encoder.layers) > arch_config.num_hidden_layers:
+            raise ValueError(
+                f"The original encoder only has {num_hidden_layers} "
+                f"layers, but you requested {len(self.encoder.layers)} layers."
+            )
+
+        # Post layer norm (applied to output)
+        if require_post_norm is None:
+            require_post_norm = len(self.encoder.layers) == num_hidden_layers
+
+        if require_post_norm:
+            self.post_layernorm = nn.LayerNorm(embed_dim, eps=arch_config.layer_norm_eps)
+        else:
+            self.post_layernorm = None
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        feature_sample_layers: list[int] | None = None,
+    ) -> torch.Tensor:
+        hidden_states = self.embeddings(pixel_values)
+
+        return_all_hidden_states = feature_sample_layers is not None
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            return_all_hidden_states=return_all_hidden_states,
+        )
+
+        if not return_all_hidden_states:
+            encoder_outputs = encoder_outputs[0]
+
+        # Apply post-layernorm
+        if self.post_layernorm is not None:
+            if isinstance(encoder_outputs, list):
+                encoder_outputs[-1] = self.post_layernorm(encoder_outputs[-1])
+            else:
+                encoder_outputs = self.post_layernorm(encoder_outputs)
+
+        return encoder_outputs
+
+
+class SiglipVisionModel(ImageEncoder):
+    """SigLIP Vision Model for FastVideo."""
+    
+    config_class = SiglipVisionConfig
+    main_input_name = "pixel_values"
+    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
+
+    def __init__(self, config: SiglipVisionConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.vision_model = SiglipVisionTransformer(
+            config=config,
+            quant_config=config.quant_config,
+            num_hidden_layers_override=config.num_hidden_layers_override,
+            require_post_norm=config.require_post_norm,
+            prefix=f"{config.prefix}.vision_model",
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        feature_sample_layers: list[int] | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        last_hidden_state = self.vision_model(pixel_values, feature_sample_layers)
+        return BaseEncoderOutput(last_hidden_state=last_hidden_state)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        layer_count = len(self.vision_model.encoder.layers)
+
+        for name, loaded_weight in weights:
+            # Skip projection layers if any
+            if name.startswith("visual_projection"):
+                continue
+                
+            # Skip head if any
+            if "head" in name:
+                continue
+
+            # Post layernorm handling
+            if (name.startswith("vision_model.post_layernorm")
+                    and self.vision_model.post_layernorm is None):
+                continue
+
+            # Omit layers when num_hidden_layers_override is set
+            if name.startswith("vision_model.encoder.layers"):
+                layer_idx = int(name.split(".")[3])
+                if layer_idx >= layer_count:
+                    continue
+
+            # Handle QKV projection weight mapping
+            for (param_name, weight_name, shard_id) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                if name in params_dict:
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name in params_dict:
+                    param = params_dict[name]
+                    weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+            
+        return loaded_params

fastvideo/models/loader/component_loader.py

@@ -15,8 +15,7 @@ import torch.distributed as dist
 import torch.nn as nn
 from safetensors.torch import load_file as safetensors_load_file
 from torch.distributed import init_device_mesh
-from transformers import AutoImageProcessor, AutoTokenizer
-from transformers import UMT5EncoderModel
+from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
 
 from fastvideo.configs.models import EncoderConfig
@@ -35,8 +34,8 @@ from fastvideo.models.loader.weight_utils import (
     safetensors_weights_iterator,
 )
 from fastvideo.models.registry import ModelRegistry
-from fastvideo.utils import PRECISION_TO_TYPE
-from fastvideo.models.layerwise_offload import LayerwiseOffloadManager
+from fastvideo.utils import PRECISION_TO_TYPE, is_pin_memory_available
+from fastvideo.hooks.layerwise_offload import enable_layerwise_offload
 
 logger = init_logger(__name__)
 
@@ -81,20 +80,26 @@ class ComponentLoader(ABC):
             "transformer": (TransformerLoader, "diffusers"),
             "transformer_2": (TransformerLoader, "diffusers"),
             "vae": (VAELoader, "diffusers"),
+            "audio_vae": (AudioDecoderLoader, "diffusers"),
+            "audio_decoder": (AudioDecoderLoader, "diffusers"),
+            "vocoder": (VocoderLoader, "diffusers"),
             "text_encoder": (TextEncoderLoader, "transformers"),
             "text_encoder_2": (TextEncoderLoader, "transformers"),
             "tokenizer": (TokenizerLoader, "transformers"),
             "tokenizer_2": (TokenizerLoader, "transformers"),
             "image_processor": (ImageProcessorLoader, "transformers"),
+            "feature_extractor": (ImageProcessorLoader, "transformers"),
             "image_encoder": (ImageEncoderLoader, "transformers"),
         }
 
         if module_type in module_loaders:
             loader_cls, expected_library = module_loaders[module_type]
-            # Assert that the library matches what's expected for this module type
-            assert transformers_or_diffusers == expected_library, (
-                f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
-            )
+            # Allow fastvideo.* libraries for custom implementations (e.g. Cosmos2_5Pipeline)
+            # that aren't available in diffusers/transformers yet
+            is_fastvideo_module = transformers_or_diffusers.startswith("fastvideo.")
+            if not is_fastvideo_module:
+                # Assert that the library matches what's expected for this module type
+                assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}"
             return loader_cls()
 
         # For unknown module types, use a generic loader
@@ -241,6 +246,47 @@ class TextEncoderLoader(ComponentLoader):
         model_config.pop("model_type", None)
         model_config.pop("tokenizer_class", None)
         model_config.pop("torch_dtype", None)
+        repo_root = os.path.dirname(model_path)
+        index_path = os.path.join(repo_root, "model_index.json")
+        gemma_path = ""
+        gemma_path_from_candidate = False
+        if os.path.isfile(index_path):
+            try:
+                with open(index_path, encoding="utf-8") as f:
+                    model_index = json.load(f)
+                gemma_path = model_index.get("gemma_model_path", "")
+            except json.JSONDecodeError:
+                gemma_path = ""
+        if not gemma_path:
+            candidate = os.path.normpath(os.path.join(model_path, "gemma"))
+            if os.path.isdir(candidate):
+                gemma_path = candidate
+                gemma_path_from_candidate = True
+                model_config["gemma_model_path"] = gemma_path
+        if gemma_path and not gemma_path_from_candidate:
+            if not os.path.isabs(gemma_path):
+                model_config["gemma_model_path"] = os.path.normpath(
+                    os.path.join(repo_root, gemma_path)
+                )
+        transformer_config_path = os.path.join(
+            repo_root, "transformer", "config.json"
+        )
+        if os.path.isfile(transformer_config_path):
+            try:
+                with open(transformer_config_path, encoding="utf-8") as f:
+                    transformer_config = json.load(f)
+                if (
+                    "connector_double_precision_rope" not in model_config
+                    or not model_config["connector_double_precision_rope"]
+                ):
+                    if transformer_config.get("double_precision_rope") is True:
+                        model_config["connector_double_precision_rope"] = True
+                if "connector_rope_type" not in model_config:
+                    rope_type = transformer_config.get("rope_type")
+                    if rope_type is not None:
+                        model_config["connector_rope_type"] = rope_type
+            except json.JSONDecodeError:
+                pass
         logger.info("HF Model config: %s", model_config)
 
         # @TODO(Wei): Better way to handle this?
@@ -279,7 +325,7 @@ class TextEncoderLoader(ComponentLoader):
         target_device: torch.device,
         fastvideo_args: FastVideoArgs,
         dtype: str = "fp16",
-        use_text_encoder_override: bool = False, # prevent subclasses from misusing
+        use_text_encoder_override: bool = False,  # prevent subclasses from misusing
     ):
         use_cpu_offload = (
             fastvideo_args.text_encoder_cpu_offload
@@ -296,7 +342,10 @@ class TextEncoderLoader(ComponentLoader):
             )
 
         # Set quantization config if specified
-        if use_text_encoder_override and fastvideo_args.override_text_encoder_quant is not None:
+        if (
+            use_text_encoder_override
+            and fastvideo_args.override_text_encoder_quant is not None
+        ):
             if fastvideo_args.override_text_encoder_safetensors is None:
                 raise ValueError(
                     "override_text_encoder_quant is set but override_text_encoder_safetensors is None"
@@ -313,7 +362,10 @@ class TextEncoderLoader(ComponentLoader):
                 model: TextEncoder = model_cls(model_config)  # type: ignore
 
             weights_to_load = {name for name, _ in model.named_parameters()}
-            if use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None:
+            if (
+                use_text_encoder_override
+                and fastvideo_args.override_text_encoder_safetensors is not None
+            ):
                 loaded_weights: set[str] = model.load_weights(
                     safetensors_weights_iterator(
                         [fastvideo_args.override_text_encoder_safetensors],
@@ -339,7 +391,10 @@ class TextEncoderLoader(ComponentLoader):
 
             from fastvideo.platforms import current_platform
 
+            logger.info("Loading text encoder with cpu_offload: %s", use_cpu_offload)
+
             if use_cpu_offload:
+                pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available()
                 # Disable FSDP for MPS as it's not compatible
                 if current_platform.is_mps():
                     logger.info(
@@ -357,7 +412,7 @@ class TextEncoderLoader(ComponentLoader):
                         reshard_after_forward=True,
                         mesh=mesh["offload"],
                         fsdp_shard_conditions=model._fsdp_shard_conditions,
-                        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
+                        pin_cpu_memory=pin_cpu_memory,
                     )
                 else:
                     mesh = init_device_mesh(
@@ -371,7 +426,7 @@ class TextEncoderLoader(ComponentLoader):
                         reshard_after_forward=True,
                         mesh=mesh["offload"],
                         fsdp_shard_conditions=model._fsdp_shard_conditions,
-                        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
+                        pin_cpu_memory=pin_cpu_memory,
                     )
             # We only enable strict check for non-quantized models
             # that have loaded weights tracking currently.
@@ -449,13 +504,52 @@ class TokenizerLoader(ComponentLoader):
         """Load the tokenizer based on the model path, and inference args."""
         logger.info("Loading tokenizer from %s", model_path)
 
+        # Cosmos2.5 stores an AutoProcessor config in `tokenizer/config.json` (not a tokenizer
+        # config). Use its `_name_or_path` (e.g. Qwen/Qwen2.5-VL-7B-Instruct) as the source.
+        tokenizer_cfg_path = os.path.join(model_path, "config.json")
+        if os.path.exists(tokenizer_cfg_path):
+            try:
+                with open(tokenizer_cfg_path, "r") as f:
+                    tokenizer_cfg = json.load(f)
+                if isinstance(tokenizer_cfg, dict) and (
+                    tokenizer_cfg.get("_class_name") == "AutoProcessor"
+                    or "processor_type" in tokenizer_cfg
+                ):
+                    src = tokenizer_cfg.get("_name_or_path", "")
+                    if isinstance(src, str) and src.strip():
+                        processor = AutoProcessor.from_pretrained(
+                            src.strip(),
+                            trust_remote_code=True,
+                        )
+                        logger.info(
+                            "Loaded tokenizer/processor from %s: %s",
+                            src,
+                            processor.__class__.__name__,
+                        )
+                        return processor
+            except Exception:
+                # If parsing fails, fall through to AutoTokenizer below.
+                pass
+
         tokenizer = AutoTokenizer.from_pretrained(
             model_path,  # "<path to model>/tokenizer"
             # in v0, this was same string as encoder_name "ClipTextModel"
             # TODO(will): pass these tokenizer kwargs from inference args? Maybe
             # other method of config?
-            padding_size="right",
         )
+        padding_side = None
+        if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
+            try:
+                arch_config = fastvideo_args.pipeline_config.text_encoder_configs[
+                    0
+                ].arch_config
+                padding_side = getattr(arch_config, "padding_side", None)
+            except Exception:
+                padding_side = None
+        if padding_side:
+            tokenizer.padding_side = padding_side
+        if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
         logger.info("Loaded tokenizer: %s", tokenizer.__class__.__name__)
         return tokenizer
 
@@ -466,15 +560,11 @@ class VAELoader(ComponentLoader):
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_config(model=model_path)
+        config.pop("_name_or_path", None)
         class_name = config.pop("_class_name")
-        assert class_name is not None, (
-            "Model config does not contain a _class_name attribute. Only diffusers format is supported."
-        )
+        assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
         fastvideo_args.model_paths["vae"] = model_path
 
-        vae_config = fastvideo_args.pipeline_config.vae_config
-        vae_config.update_model_arch(config)
-
         from fastvideo.platforms import current_platform
 
         if fastvideo_args.vae_cpu_offload:
@@ -491,29 +581,156 @@ class VAELoader(ComponentLoader):
             if fastvideo_args.pipeline_config.vae_precision
             else torch.bfloat16
         ):
-            vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
-            vae = vae_cls(vae_config).to(target_device)
+            # Cosmos2.5 uses a Wan2.1 VAE stored as `tokenizer.safetensors` under the VAE folder.
+            is_cosmos25 = fastvideo_args.pipeline_config.__class__.__name__ == "Cosmos25Config"
+            if class_name == "AutoencoderKLWan" and is_cosmos25:
+                from fastvideo.models.vaes.cosmos25wanvae import Cosmos25WanVAE
+
+                dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+                vae = Cosmos25WanVAE(device=target_device, dtype=dtype)
+
+                weight_path = os.path.join(model_path, "tokenizer.safetensors")
+                if not os.path.exists(weight_path):
+                    raise FileNotFoundError(
+                        f"Missing Cosmos2.5 VAE weights: {weight_path}"
+                    )
+                sd = safetensors_load_file(weight_path)
+                vae.load_state_dict(sd, strict=False)
+                return vae.eval()
+
+            # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
+            if class_name == "CausalVideoAutoencoder" and "vae" in config:
+                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+                vae = vae_cls(config).to(target_device)
+                if hasattr(vae, "set_tiling_config"):
+                    vae_config = fastvideo_args.pipeline_config.vae_config
+                    vae.set_tiling_config(
+                        spatial_tile_size_in_pixels=getattr(
+                            vae_config, "ltx2_spatial_tile_size_in_pixels", 512),
+                        spatial_tile_overlap_in_pixels=getattr(
+                            vae_config, "ltx2_spatial_tile_overlap_in_pixels", 64),
+                        temporal_tile_size_in_frames=getattr(
+                            vae_config, "ltx2_temporal_tile_size_in_frames", 64),
+                        temporal_tile_overlap_in_frames=getattr(
+                            vae_config,
+                            "ltx2_temporal_tile_overlap_in_frames", 24),
+                    )
+            else:
+                config.pop("_class_name", None)
+                vae_config = fastvideo_args.pipeline_config.vae_config
+                vae_config.update_model_arch(config)
+                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+                vae = vae_cls(vae_config).to(target_device)
 
         # Find all safetensors files
         safetensors_list = glob.glob(
-            os.path.join(str(model_path), "*.safetensors")
-        )
+            os.path.join(str(model_path), "*.safetensors"))
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {model_path}")
+        # Common case: a single `.safetensors` checkpoint file.
+        # Some models may be sharded into multiple files; in that case we merge.
         loaded = {}
         for sf_file in safetensors_list:
             loaded.update(safetensors_load_file(sf_file))
-        vae.load_state_dict(
-            loaded, strict=False
-        )  # We might only load encoder or decoder
+
+        # LTX-2 CausalVideoAutoencoder needs per_channel_statistics remapping
+        if class_name == "CausalVideoAutoencoder" and "vae" in config:
+            per_channel_prefixes = (
+                "per_channel_statistics.",
+                "vae.per_channel_statistics.",
+            )
+            remapped = {}
+            for key, tensor in loaded.items():
+                remapped[key] = tensor
+                for prefix in per_channel_prefixes:
+                    if key.startswith(prefix):
+                        suffix = key[len(prefix):]
+                        remapped.setdefault(
+                            f"encoder.per_channel_statistics.{suffix}",
+                            tensor,
+                        )
+                        remapped.setdefault(
+                            f"decoder.per_channel_statistics.{suffix}",
+                            tensor,
+                        )
+                        break
+            loaded = remapped
+
+        vae.load_state_dict(loaded, strict=False)
 
         return vae.eval()
 
 
+class AudioDecoderLoader(ComponentLoader):
+    """Loader for LTX-2 audio decoder (audio_vae component)."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        config = get_diffusers_config(model=model_path)
+        class_name = config.pop("_class_name", None) or "LTX2AudioDecoder"
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+        target_device = get_local_torch_device()
+
+        precision = getattr(
+            fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16"
+        )
+        with set_default_torch_dtype(PRECISION_TO_TYPE[precision]):
+            audio_decoder = model_cls(config).to(target_device)
+
+        safetensors_list = glob.glob(
+            os.path.join(str(model_path), "*.safetensors")
+        )
+        loaded: dict[str, torch.Tensor] = {}
+        for sf_file in safetensors_list:
+            loaded.update(safetensors_load_file(sf_file))
+
+        decoder_state = {}
+        for name, tensor in loaded.items():
+            if name.startswith("decoder."):
+                decoder_state[name.replace("decoder.", "")] = tensor
+            elif name.startswith("per_channel_statistics."):
+                decoder_state[name] = tensor
+
+        target_module = getattr(audio_decoder, "model", audio_decoder)
+        target_module.load_state_dict(decoder_state, strict=False)
+        return audio_decoder.eval()
+
+
+class VocoderLoader(ComponentLoader):
+    """Loader for LTX-2 vocoder."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        config = get_diffusers_config(model=model_path)
+        class_name = config.pop("_class_name", None) or "LTX2Vocoder"
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+        target_device = get_local_torch_device()
+
+        precision = getattr(
+            fastvideo_args.pipeline_config, "vocoder_precision", "bf16"
+        )
+        with set_default_torch_dtype(PRECISION_TO_TYPE[precision]):
+            vocoder = model_cls(config).to(target_device)
+
+        safetensors_list = glob.glob(
+            os.path.join(str(model_path), "*.safetensors")
+        )
+        loaded: dict[str, torch.Tensor] = {}
+        for sf_file in safetensors_list:
+            loaded.update(safetensors_load_file(sf_file))
+
+        target_module = getattr(vocoder, "model", vocoder)
+        target_module.load_state_dict(loaded, strict=False)
+        return vocoder.eval()
+
+
 class TransformerLoader(ComponentLoader):
     """Loader for transformer."""
 
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the transformer based on the model path, and inference args."""
         config = get_diffusers_config(model=model_path)
+        config.pop("_name_or_path", None)
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
         if cls_name is None:
@@ -581,7 +798,17 @@ class TransformerLoader(ComponentLoader):
         ]
 
         # Load the model using FSDP loader
+        logger.info("Loading model from %s, default_dtype: %s", cls_name,
+                    default_dtype)
         assert fastvideo_args.hsdp_shard_dim is not None
+        # Cosmos2.5 checkpoints can include extra entries not present in the
+        # instantiated model (e.g. pos_embedder ranges / *_extra_state). Load
+        # non-strictly for Cosmos2.5 only; keep upstream strict behavior for others.
+        strict_load = not (
+            cls_name.startswith("Cosmos25")
+            or cls_name == "Cosmos25Transformer3DModel"
+            or getattr(fastvideo_args.pipeline_config, "prefix", "") == "Cosmos25"
+        )
         model = maybe_load_fsdp_model(
             model_cls=model_cls,
             init_params={"config": dit_config, "hf_config": hf_config},
@@ -589,6 +816,7 @@ class TransformerLoader(ComponentLoader):
             device=get_local_torch_device(),
             hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
             hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
+            strict=strict_load,
             cpu_offload=fastvideo_args.dit_cpu_offload,
             pin_cpu_memory=fastvideo_args.pin_cpu_memory,
             fsdp_inference=fastvideo_args.use_fsdp_inference,
@@ -603,7 +831,7 @@ class TransformerLoader(ComponentLoader):
         )
 
         total_params = sum(p.numel() for p in model.parameters())
-        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
+        logger.info("Loaded model with %.2fB parameters, with cpu_offload: %s", total_params / 1e9, fastvideo_args.dit_cpu_offload)
 
         assert next(model.parameters()).dtype == default_dtype, (
             "Model dtype does not match default dtype"
@@ -611,35 +839,19 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
-        if fastvideo_args.dit_layerwise_offload and hasattr(model, "blocks"):
-            # Check if this is a Wan model (only Wan models support layerwise offload)
-            is_wan_model = "Wan" in cls_name
-            if not is_wan_model:
+        if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload:
+            # Check if model has nn.ModuleList for layerwise offload compatibility
+            has_module_list = any(
+                isinstance(m, nn.ModuleList) for m in model.children()
+            )
+            if has_module_list:
+                enable_layerwise_offload(model)
+            else:
                 logger.warning(
-                    "Layerwise offload is currently only supported for Wan models. "
-                    "Model class '%s' does not support layerwise offload. "
-                    "Disabling layerwise offload for this model.",
+                    "Layerwise offload requested but model %s does not have "
+                    "nn.ModuleList structure. Skipping layerwise offload.",
                     cls_name
                 )
-            else:
-                try:
-                    num_layers = len(getattr(model, "blocks"))
-                except TypeError:
-                    num_layers = None
-                if isinstance(num_layers, int) and num_layers > 0:
-                    # Ensure model is on the correct device (CUDA) before initializing manager
-                    # This ensures non-managed parameters (embeddings, final norms) are on GPU
-                    model = model.to(get_local_torch_device())
-                    mgr = LayerwiseOffloadManager(
-                        model,
-                        module_list_attr="blocks",
-                        num_layers=num_layers,
-                        enabled=True,
-                        pin_cpu_memory=fastvideo_args.pin_cpu_memory,
-                        auto_initialize=True,
-                    )
-                    setattr(model, "_layerwise_offload_manager", mgr)
-
         return model
 
 
@@ -746,4 +958,4 @@ class PipelineComponentLoader:
         )
 
         # Load the module
-        return loader.load(component_model_path, fastvideo_args)
+        return loader.load(component_model_path, fastvideo_args)

fastvideo/models/loader/fsdp_load.py

@@ -23,7 +23,7 @@ from fastvideo.logger import init_logger
 from fastvideo.models.loader.utils import (get_param_names_mapping,
                                            hf_to_custom_state_dict)
 from fastvideo.models.loader.weight_utils import safetensors_weights_iterator
-from fastvideo.utils import set_mixed_precision_policy
+from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
 
 logger = init_logger(__name__)
 
@@ -67,6 +67,7 @@ def maybe_load_fsdp_model(
     default_dtype: torch.dtype,
     param_dtype: torch.dtype,
     reduce_dtype: torch.dtype,
+    strict: bool = True,
     cpu_offload: bool = False,
     fsdp_inference: bool = False,
     output_dtype: torch.dtype | None = None,
@@ -106,6 +107,7 @@ def maybe_load_fsdp_model(
         logger.info("Disabling FSDP for MPS platform as it's not compatible")
 
     if use_fsdp:
+        pin_cpu_memory = pin_cpu_memory and is_pin_memory_available()
         world_size = hsdp_replicate_dim * hsdp_shard_dim
         if not training_mode and not fsdp_inference:
             hsdp_replicate_dim = world_size
@@ -141,7 +143,7 @@ def maybe_load_fsdp_model(
         weight_iterator,
         device,
         default_dtype,
-        strict=True,
+        strict=strict,
         cpu_offload=cpu_offload,
         param_names_mapping=param_names_mapping_fn,
     )
@@ -293,6 +295,26 @@ def load_model_from_full_model_state_dict(
     for target_param_name, full_tensor in custom_param_sd.items():
         meta_sharded_param = meta_sd.get(target_param_name)
         if meta_sharded_param is None:
+            # Some checkpoints include extra entries that are not part of the
+            # instantiated model's state_dict (e.g. `_extra_state` keys from
+            # some FSDP checkpoint formats). These can be safely skipped.
+            if (target_param_name.endswith("._extra_state")
+                    or target_param_name.endswith("_extra_state")):
+                logger.warning(
+                    "Skipping non-parameter checkpoint key: %s",
+                    target_param_name,
+                )
+                continue
+
+            # For non-strict loads, treat this as an "unexpected key" and skip it
+            # (mirrors torch.nn.Module.load_state_dict(strict=False)).
+            if not strict:
+                logger.warning(
+                    "Skipping unexpected checkpoint key (not present in model): %s",
+                    target_param_name,
+                )
+                continue
+
             raise ValueError(
                 f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect."
             )

fastvideo/models/loader/utils.py

@@ -78,9 +78,10 @@ def hf_to_custom_state_dict(
     for source_param_name, full_tensor in hf_param_sd:  # type: ignore
         target_param_name, merge_index, num_params_to_merge = param_names_mapping(
             source_param_name)
-        reverse_param_names_mapping[target_param_name] = (source_param_name,
-                                                          merge_index,
-                                                          num_params_to_merge)
+        if merge_index is None:
+            reverse_param_names_mapping[target_param_name] = (source_param_name,
+                                                            merge_index,
+                                                            num_params_to_merge)
         if merge_index is not None:
             to_merge_params[target_param_name][merge_index] = full_tensor
             if len(to_merge_params[target_param_name]) == num_params_to_merge:

fastvideo/models/registry.py

@@ -26,12 +26,18 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
     "HunyuanVideo15Transformer3DModel":
     ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
+    "HYWorldTransformer3DModel":
+    ("dits", "hyworld", "HYWorldTransformer3DModel"),
+    "CausalHunyuanVideo15Transformer3DModel":
+    ("dits", "causal_hunyuanvideo15", "CausalHunyuanVideo15Transformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
     "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
-    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),
-    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),
+    "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
+    "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
+    "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
+    "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -50,20 +56,33 @@ _TEXT_ENCODER_MODELS = {
     "STEP1TextEncoder": ("encoders", "stepllm", "STEP1TextEncoder"),
     "BertModel": ("encoders", "clip", "CLIPTextModel"),
     "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"),
+    "Reason1TextEncoder": ("encoders", "reason1", "Reason1TextEncoder"),
+    "Qwen2_5_VLForConditionalGeneration":
+    ("encoders", "reason1", "Reason1TextEncoder"),
+    "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
     # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"),
     "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"),
     "CLIPVisionModel": ("encoders", "clip", "CLIPVisionModel"),
+    "SiglipVisionModel": ("encoders", "siglip", "SiglipVisionModel"),
 }
 
 _VAE_MODELS = {
     "AutoencoderKLHunyuanVideo":
     ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"),
+    "AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"),
     "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
-    "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo")
+    "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"),
+    "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
+}
+
+_AUDIO_MODELS = {
+    "LTX2AudioEncoder": ("audio", "ltx2_audio_vae", "LTX2AudioEncoder"),
+    "LTX2AudioDecoder": ("audio", "ltx2_audio_vae", "LTX2AudioDecoder"),
+    "LTX2Vocoder": ("audio", "ltx2_audio_vae", "LTX2Vocoder"),
 }
 
 _SCHEDULERS = {
@@ -72,6 +91,8 @@ _SCHEDULERS = {
      "FlowMatchEulerDiscreteScheduler"),
     "UniPCMultistepScheduler":
     ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"),
+    "FlowUniPCMultistepScheduler":
+    ("schedulers", "scheduling_flow_unipc_multistep", "FlowUniPCMultistepScheduler"),
     "SelfForcingFlowMatchScheduler":
     ("schedulers", "scheduling_self_forcing_flow_match",
      "SelfForcingFlowMatchScheduler"),
@@ -85,6 +106,7 @@ _FAST_VIDEO_MODELS = {
     **_TEXT_ENCODER_MODELS,
     **_IMAGE_ENCODER_MODELS,
     **_VAE_MODELS,
+    **_AUDIO_MODELS,
     **_SCHEDULERS,
 }
 

fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -155,8 +155,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
 
         self.sigmas = sigmas.to(
             "cpu")  # to avoid too much CPU/GPU communication
-        self.sigma_min = self.sigmas[-1].item()
-        self.sigma_max = self.sigmas[0].item()
+        self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item()
+        self.sigma_max = sigma_max if sigma_max is not None else self.sigmas[0].item()
 
         BaseScheduler.__init__(self)
 
@@ -289,6 +289,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
         sigmas: list[float] | None = None,
         mu: float | None = None,
         timesteps: list[float] | None = None,
+        extra_one_step: bool = False,
     ) -> None:
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -350,7 +351,10 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
             if timesteps_array is None:
                 t_max = self._sigma_to_t(self.sigma_max)
                 t_min = self._sigma_to_t(self.sigma_min)
-                timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
+                if extra_one_step:
+                    timesteps_array = np.linspace(t_max, t_min, num_inference_steps + 1)[:-1]
+                else:
+                    timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
             sigmas_array = timesteps_array / self.config.num_train_timesteps
         else:
             sigmas_array = np.array(sigmas).astype(np.float32)
@@ -644,7 +648,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
         self,
         clean_latent: torch.Tensor,
         noise: torch.Tensor,
-        timestep: torch.IntTensor,
+        timestep: torch.Tensor,
     ) -> torch.Tensor:
 
         """

fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py

@@ -109,6 +109,9 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         sigmas = 1.0 - alphas
         sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
 
+        # Needed when final_sigmas_type == "sigma_min" (kept for compatibility).
+        self.alphas_cumprod = torch.from_numpy(alphas).to(dtype=torch.float32)
+
         if not use_dynamic_shifting:
             # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
             assert shift is not None
@@ -171,6 +174,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         sigmas: list[float] | None = None,
         mu: float | None | None = None,
         shift: float | None | None = None,
+        use_karras_sigmas: bool | None = None,
+        use_kerras_sigma: bool | None = None,
     ):
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -186,21 +191,44 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
                 " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
             )
 
-        if sigmas is None:
-            assert num_inference_steps is not None
-            sigmas = np.linspace(self.sigma_max, self.sigma_min,
-                                 num_inference_steps +
-                                 1).copy()[:-1]  # pyright: ignore
+
+        # Cosmos official uses `use_kerras_sigma=True` and a specific EDM sigma schedule.
+        # Some external code uses the misspelling `use_kerras_sigma`; support both.
+        if use_karras_sigmas is None and use_kerras_sigma is not None:
+            use_karras_sigmas = use_kerras_sigma
+
+        if use_karras_sigmas:
+            # Force to use the exact sigma used in official EDM sampler:
+            # sigma_max=200, sigma_min=0.01, rho=7
+            sigma_max = 200.0
+            sigma_min = 0.01
+            rho = 7.0
+            # Match the official Cosmos implementation: Karras/EDM schedule with
+            # `num_inference_steps + 1` points, then `final_sigmas_type="zero"`
+            # appends the terminal sigma (0.0).
+            ramp = np.arange(num_inference_steps + 1,
+                             dtype=np.float32) / float(num_inference_steps)
+            min_inv_rho = sigma_min**(1 / rho)
+            max_inv_rho = sigma_max**(1 / rho)
+            sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
+            # Convert EDM sigma to flow-matching sigma in [0, 1).
+            sigmas = sigmas / (1.0 + sigmas)
+        else:
+            if sigmas is None:
+                assert num_inference_steps is not None
+                sigmas = np.linspace(self.sigma_max, self.sigma_min,
+                                     num_inference_steps +
+                                     1).copy()[:-1]  # pyright: ignore
 
         if self.config.use_dynamic_shifting:
             assert mu is not None
             sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
         else:
-            if shift is None:
-                shift = self.config.shift
-            assert isinstance(sigmas, np.ndarray)
-            sigmas = shift * sigmas / (1 +
-                                       (shift - 1) * sigmas)  # pyright: ignore
+            if not use_karras_sigmas:
+                if shift is None:
+                    shift = self.config.shift
+                assert isinstance(sigmas, np.ndarray)
+                sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)  # pyright: ignore
 
         if self.config.final_sigmas_type == "sigma_min":
             sigma_last = ((1 - self.alphas_cumprod[0]) /
@@ -418,8 +446,12 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
-        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
-        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        # Numerical safety
+        eps = 1e-12
+        lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(
+            torch.clamp(sigma_t, min=eps))
+        lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(
+            torch.clamp(sigma_s0, min=eps))
 
         h = lambda_t - lambda_s0
         device = sample.device
@@ -430,7 +462,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             si = self.step_index - i  # pyright: ignore
             mi = model_output_list[-(i + 1)]
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
-            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(
+                torch.clamp(sigma_si, min=eps))
             rk = (lambda_si - lambda_s0) / h
             rks.append(rk)
             assert mi is not None
@@ -563,8 +596,11 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
         alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
 
-        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
-        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        eps = 1e-12
+        lambda_t = torch.log(torch.clamp(alpha_t, min=eps)) - torch.log(
+            torch.clamp(sigma_t, min=eps))
+        lambda_s0 = torch.log(torch.clamp(alpha_s0, min=eps)) - torch.log(
+            torch.clamp(sigma_s0, min=eps))
 
         h = lambda_t - lambda_s0
         device = this_sample.device
@@ -575,7 +611,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
             si = self.step_index - (i + 1)  # pyright: ignore
             mi = model_output_list[-(i + 1)]
             alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
-            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            lambda_si = torch.log(torch.clamp(alpha_si, min=eps)) - torch.log(
+                torch.clamp(sigma_si, min=eps))
             rk = (lambda_si - lambda_s0) / h
             rks.append(rk)
             assert mi is not None

fastvideo/models/utils.py

@@ -5,6 +5,9 @@ from typing import Any
 
 import torch
 
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
 
 # TODO(PY): move it elsewhere
 def auto_attributes(init_func):

fastvideo/models/vaes/common.py

@@ -79,7 +79,7 @@ class ParallelTiledVAE(ABC):
     def decode(self, z: torch.Tensor) -> torch.Tensor:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
-        tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
         tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
         num_sample_frames = (num_frames -
                              1) * self.temporal_compression_ratio + 1

fastvideo/models/vaes/cosmos25wanvae.py

@@ -0,0 +1,735 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""
+Cosmos 2.5 / Wan2.1 VAE adapter.
+
+Why this exists:
+- Cosmos2.5 uses a Wan2.1-style VAE, but the *diffusion model* operates in a
+  **normalized latent space**:
+    z_norm = (z - mean) / std
+
+  Meanwhile, FastVideo's `AutoencoderKLWan` operates in the VAE's native latent
+  space (denormalized):
+    z = z_norm * std + mean
+
+This adapter provides a single, stable interface for FastVideo pipelines:
+- `encode(x)` returns an object with `.mean` / `.sample()` / `.mode()`
+- `decode(z)` returns a tensor in pixel space
+
+It also exposes flags used by pipeline stages to avoid double (de)normalization:
+- `handles_latent_norm = True`   -> stages should NOT normalize encoder latents
+- `handles_latent_denorm = True` -> stages should NOT denormalize before decode
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+@dataclass
+class _TensorLatentDist:
+    """Minimal distribution-like wrapper used by pipeline stages."""
+
+    mean: torch.Tensor
+
+    def mode(self) -> torch.Tensor:
+        return self.mean
+
+    def sample(self, generator: Any | None = None) -> torch.Tensor:  # generator for API compatibility
+        # The official interface encodes deterministically; for compatibility we
+        # return the mean. (Stochastic posterior sampling isn't required for
+        # Cosmos2.5 inference.)
+        _ = generator
+        return self.mean
+
+
+class Cosmos25WanVAEAdapter(nn.Module):
+    """
+    Adapter that makes a Wan2.1-style VAE follow Cosmos2.5's latent contract:
+    - `encode()` returns **normalized** latents
+    - `decode()` expects **normalized** latents
+    """
+
+    # Pipeline stage hints (see latent_preparation.py / decoding.py / image_encoding.py)
+    handles_latent_norm: bool = True
+    handles_latent_denorm: bool = True
+    latent_norm_mode: str = "internal"  # informational
+
+    def __init__(
+        self,
+        inner: Any,
+        *,
+        latents_mean: Optional[torch.Tensor] = None,
+        latents_std: Optional[torch.Tensor] = None,
+    ) -> None:
+        super().__init__()
+        self.inner = inner
+
+        # Preserve `config` when available; some pipeline utilities expect it.
+        self.config = getattr(inner, "config", None)
+
+        # If not provided, try to derive from `config.latents_mean/std`.
+        cfg = self.config
+        if latents_mean is None and cfg is not None and hasattr(cfg, "latents_mean"):
+            latents_mean = torch.tensor(cfg.latents_mean, dtype=torch.float32).view(1, -1, 1, 1, 1)
+        if latents_std is None and cfg is not None and hasattr(cfg, "latents_std"):
+            latents_std = torch.tensor(cfg.latents_std, dtype=torch.float32).view(1, -1, 1, 1, 1)
+
+        if latents_mean is None or latents_std is None:
+            raise RuntimeError(
+                "Cosmos25WanVAEAdapter requires latents_mean/latents_std (either passed explicitly or available on inner.config)."
+            )
+
+        # Register as buffers so `.to(...)` moves them with the module.
+        self.register_buffer("_latents_mean", latents_mean, persistent=False)
+        self.register_buffer("_latents_std", latents_std, persistent=False)
+
+    def _to_latent_stats(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        mean = self._latents_mean.to(device=like.device, dtype=like.dtype)
+        std = self._latents_std.to(device=like.device, dtype=like.dtype)
+        return mean, std
+
+    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
+        # Keep parity with official interface.
+        if hasattr(self.inner, "get_latent_num_frames"):
+            return int(self.inner.get_latent_num_frames(num_pixel_frames))
+        return 1 + (num_pixel_frames - 1) // 4
+
+    def encode(self, x: torch.Tensor) -> _TensorLatentDist:
+        """
+        Returns *normalized* latents (Cosmos contract).
+        """
+        enc_out = self.inner.encode(x)
+
+        # Support common encoder output shapes:
+        # - DiagonalGaussianDistribution (FastVideo VAE): has `.mean` / `.sample()` / `.mode()`
+        # - diffusers EncoderOutput: has `.latent_dist`
+        # - raw tensor
+        if hasattr(enc_out, "latent_dist"):
+            dist = enc_out.latent_dist
+            z_mean = dist.mode() if hasattr(dist, "mode") else dist.mean
+        elif hasattr(enc_out, "mode") and hasattr(enc_out, "mean"):
+            z_mean = enc_out.mode()
+        elif isinstance(enc_out, torch.Tensor):
+            z_mean = enc_out
+        else:
+            attrs = [a for a in dir(enc_out) if not a.startswith("_")]
+            raise RuntimeError(
+                f"Unsupported VAE encoder output type: {type(enc_out)}. attrs={attrs}"
+            )
+
+        mean, std = self._to_latent_stats(z_mean)
+        z_norm = (z_mean - mean) / std
+        return _TensorLatentDist(z_norm)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Expects *normalized* latents (Cosmos contract).
+        """
+        mean, std = self._to_latent_stats(z)
+        z_denorm = z * std + mean
+        out = self.inner.decode(z_denorm)
+        return out.sample if hasattr(out, "sample") else out
+
+
+#
+# Official-like Wan2.1 VAE implementation (ported from cosmos_predict2 wan2pt1.py)
+# -------------------------------------------------------------------------------
+# Motivation:
+# - We already solved checkpoint *key mapping* and can load official weights.
+# - Remaining output drift vs the official tokenizer is largely decoder-side.
+# - FastVideo's `AutoencoderKLWan` uses a different temporal upsample path
+#   (`DupUp3D` + `first_chunk` slicing), while the official tokenizer uses
+#   `Resample(mode="upsample3d")` with a time-conv + interleave reshape.
+#
+# This section ports the core modules (CausalConv3d/Resample/etc.) so we can run
+# a VAE that is behaviorally closer to the official implementation WITHOUT
+# importing any official repo classes at runtime.
+#
+
+CACHE_T = 2
+
+
+class Cosmos25CausalConv3d(nn.Conv3d):
+    """
+    Official-like causal 3D convolution.
+
+    Matches `CausalConv3d` in the official tokenizer: uses explicit F.pad and
+    supports a `cache_x` prefix for causal chunking.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        # padding order for F.pad: (W_left, W_right, H_left, H_right, T_left, T_right)
+        self._padding: tuple[int, ...] = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+        return super().forward(x)
+
+
+class Cosmos25RMSNorm(nn.Module):
+    """Official-like RMS_norm (uses learnable gamma and optional bias)."""
+
+    def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dim = 1 if self.channel_first else -1
+        return F.normalize(x, dim=dim) * self.scale * self.gamma + self.bias
+
+
+class Cosmos25Upsample(nn.Upsample):
+    """Official-like Upsample that is safe for bf16 (casts to fp32 internally)."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        return super().forward(x.float()).type_as(x)
+
+
+class Cosmos25Resample(nn.Module):
+    """
+    Official-like Resample used for both spatial and temporal up/downsampling.
+    """
+
+    def __init__(self, dim: int, mode: str) -> None:
+        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Cosmos25Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+            self.time_conv = Cosmos25CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = Cosmos25CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        b, c, t, h, w = x.size()
+
+        # Temporal upsample uses a time-conv and then interleaves frames.
+        if self.mode == "upsample3d" and feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = "Rep"
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                    cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+
+                if feat_cache[idx] == "Rep":
+                    x = self.time_conv(x)
+                else:
+                    x = self.time_conv(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+
+                x = x.reshape(b, 2, c, t, h, w)
+                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                x = x.reshape(b, c, t * 2, h, w)
+
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        # Temporal downsample: time_conv consumes last-frame cache.
+        if self.mode == "downsample3d" and feat_cache is not None:
+            idx = feat_idx[0]
+            if feat_cache[idx] is None:
+                feat_cache[idx] = x.clone()
+                feat_idx[0] += 1
+            else:
+                cache_x = x[:, :, -1:, :, :].clone()
+                x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+
+        return x
+
+
+class Cosmos25ResidualBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.residual = nn.Sequential(
+            Cosmos25RMSNorm(in_dim, images=False),
+            nn.SiLU(),
+            Cosmos25CausalConv3d(in_dim, out_dim, 3, padding=1),
+            Cosmos25RMSNorm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            Cosmos25CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = Cosmos25CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class Cosmos25AttentionBlock(nn.Module):
+    """Official-like causal self-attention with a single head."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.norm = Cosmos25RMSNorm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        b, c, t, h, w = x.size()
+        x2 = rearrange(x, "b c t h w -> (b t) c h w")
+        x2 = self.norm(x2)
+        q, k, v = (
+            self.to_qkv(x2)
+            .reshape(b * t, 1, c * 3, -1)
+            .permute(0, 1, 3, 2)
+            .contiguous()
+            .chunk(3, dim=-1)
+        )
+        x2 = F.scaled_dot_product_attention(q, k, v)
+        x2 = x2.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+        x2 = self.proj(x2)
+        x2 = rearrange(x2, "(b t) c h w-> b c t h w", t=t)
+        return x2 + identity
+
+
+class Cosmos25Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim: int = 96,
+        z_dim: int = 32,
+        dim_mult: list[int] = [1, 2, 4, 4],
+        num_res_blocks: int = 2,
+        attn_scales: list[float] = [],
+        temperal_downsample: list[bool] = [False, True, True],
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        self.conv1 = Cosmos25CausalConv3d(3, dims[0], 3, padding=1)
+
+        downsamples: list[nn.Module] = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            for _ in range(num_res_blocks):
+                downsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(Cosmos25AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
+                downsamples.append(Cosmos25Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        self.middle = nn.Sequential(
+            Cosmos25ResidualBlock(out_dim, out_dim, dropout),
+            Cosmos25AttentionBlock(out_dim),
+            Cosmos25ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        self.head = nn.Sequential(
+            Cosmos25RMSNorm(out_dim, images=False),
+            nn.SiLU(),
+            Cosmos25CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)  # type: ignore[misc]
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.middle:
+            if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.head:
+            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)  # type: ignore[misc]
+        return x
+
+
+class Cosmos25Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim: int = 96,
+        z_dim: int = 16,
+        dim_mult: list[int] = [1, 2, 4, 4],
+        num_res_blocks: int = 2,
+        attn_scales: list[float] = [],
+        temperal_upsample: list[bool] = [False, True, True],
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2 ** (len(dim_mult) - 2)
+
+        self.conv1 = Cosmos25CausalConv3d(z_dim, dims[0], 3, padding=1)
+        self.middle = nn.Sequential(
+            Cosmos25ResidualBlock(dims[0], dims[0], dropout),
+            Cosmos25AttentionBlock(dims[0]),
+            Cosmos25ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        upsamples: list[nn.Module] = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            if i in (1, 2, 3):
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(Cosmos25ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(Cosmos25AttentionBlock(out_dim))
+                in_dim = out_dim
+            if i != len(dim_mult) - 1:
+                mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
+                upsamples.append(Cosmos25Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        self.head = nn.Sequential(
+            Cosmos25RMSNorm(out_dim, images=False),
+            nn.SiLU(),
+            Cosmos25CausalConv3d(out_dim, 3, 3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor, feat_cache: list[Any] | None = None, feat_idx: list[int] = [0]) -> torch.Tensor:
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, Cosmos25ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)  # type: ignore[misc]
+            else:
+                x = layer(x)  # type: ignore[misc]
+
+        for layer in self.head:
+            if isinstance(layer, Cosmos25CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)  # type: ignore[misc]
+        return x
+
+
+def _count_cosmos25_conv3d(model: nn.Module) -> int:
+    return sum(1 for m in model.modules() if isinstance(m, Cosmos25CausalConv3d))
+
+
+class Cosmos25WanVAE(nn.Module):
+    """
+    A FastVideo-native copy of the *official-like* Wan2.1 VAE core.
+
+    Key properties:
+    - Module naming matches official tokenizer (`encoder`, `decoder`, `conv1`, `conv2`)
+      so it can consume `tokenizer.pth` keys directly.
+    - `encode()` returns **normalized** latents and `decode()` expects **normalized**
+      latents (Cosmos2.5 contract), matching `Wan2pt1VAEInterface`.
+    """
+
+    handles_latent_norm: bool = True
+    handles_latent_denorm: bool = True
+
+    def __init__(
+        self,
+        *,
+        device: torch.device | str = "cpu",
+        dtype: torch.dtype = torch.float32,
+        temporal_window: int = 4,
+        latents_mean: Optional[torch.Tensor] = None,
+        latents_std: Optional[torch.Tensor] = None,
+    ) -> None:
+        super().__init__()
+
+        # Official hyperparams for Cosmos2.5 tokenizer (Wan2.1 VAE).
+        cfg = dict(
+            dim=96,
+            z_dim=16,
+            dim_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            attn_scales=[],
+            temperal_downsample=[False, True, True],
+            dropout=0.0,
+            temporal_window=temporal_window,
+        )
+        self.z_dim = 16
+        self.temporal_window = temporal_window
+
+        self.encoder = Cosmos25Encoder3d(
+            dim=cfg["dim"],
+            z_dim=cfg["z_dim"] * 2,
+            dim_mult=cfg["dim_mult"],
+            num_res_blocks=cfg["num_res_blocks"],
+            attn_scales=cfg["attn_scales"],
+            temperal_downsample=cfg["temperal_downsample"],
+            dropout=cfg["dropout"],
+        )
+        self.conv1 = Cosmos25CausalConv3d(self.z_dim * 2, self.z_dim * 2, 1)
+        self.conv2 = Cosmos25CausalConv3d(self.z_dim, self.z_dim, 1)
+        self.decoder = Cosmos25Decoder3d(
+            dim=cfg["dim"],
+            z_dim=cfg["z_dim"],
+            dim_mult=cfg["dim_mult"],
+            num_res_blocks=cfg["num_res_blocks"],
+            attn_scales=cfg["attn_scales"],
+            temperal_upsample=list(cfg["temperal_downsample"])[::-1],
+            dropout=cfg["dropout"],
+        )
+
+        # Default Cosmos2.5 latent stats (shared with configs).
+        if latents_mean is None:
+            latents_mean = torch.tensor(
+                [
+                    -0.7571,
+                    -0.7089,
+                    -0.9113,
+                    0.1075,
+                    -0.1745,
+                    0.9653,
+                    -0.1517,
+                    1.5508,
+                    0.4134,
+                    -0.0715,
+                    0.5517,
+                    -0.3632,
+                    -0.1922,
+                    -0.9497,
+                    0.2503,
+                    -0.2921,
+                ],
+                dtype=torch.float32,
+            ).view(1, 16, 1, 1, 1)
+        if latents_std is None:
+            latents_std = torch.tensor(
+                [
+                    2.8184,
+                    1.4541,
+                    2.3275,
+                    2.6558,
+                    1.2196,
+                    1.7708,
+                    2.6052,
+                    2.0743,
+                    3.2687,
+                    2.1526,
+                    2.8652,
+                    1.5579,
+                    1.6382,
+                    1.1253,
+                    2.8251,
+                    1.9160,
+                ],
+                dtype=torch.float32,
+            ).view(1, 16, 1, 1, 1)
+
+        self.register_buffer("_latents_mean", latents_mean, persistent=False)
+        self.register_buffer("_latents_std", latents_std, persistent=False)
+
+        self.to(device=device, dtype=dtype)
+        self.clear_cache()
+
+    def clear_cache(self) -> None:
+        # Decoder cache
+        self._conv_num = _count_cosmos25_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map: list[Any] = [None] * self._conv_num
+        # Encoder cache
+        self._enc_conv_num = _count_cosmos25_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map: list[Any] = [None] * self._enc_conv_num
+
+    def _scale(self, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        mean = self._latents_mean.to(device=like.device, dtype=like.dtype)
+        std = self._latents_std.to(device=like.device, dtype=like.dtype)
+        return mean, 1.0 / std
+
+    def _i0_encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+
+    def _i0_decode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.decoder(x[:, :, 0:1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+
+    def encode(self, x: torch.Tensor) -> _TensorLatentDist:
+        """
+        Encode to *normalized* latents (Cosmos contract).
+        """
+        self.clear_cache()
+        t = x.shape[2]
+        iters = 1 + (t - 1) // self.temporal_window
+
+        for i in range(iters):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self._i0_encode(x)
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + self.temporal_window * (i - 1) : 1 + self.temporal_window * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+
+        if (t - 1) % self.temporal_window:
+            self._enc_conv_idx = [0]
+            out_ = self.encoder(
+                x[:, :, 1 + self.temporal_window * (iters - 1) :, :, :],
+                feat_cache=self._enc_feat_map,
+                feat_idx=self._enc_conv_idx,
+            )
+            out = torch.cat([out, out_], 2)
+
+        mu, _log_var = self.conv1(out).chunk(2, dim=1)
+        mean, inv_std = self._scale(mu)
+        z_norm = (mu - mean) * inv_std
+        self.clear_cache()
+        return _TensorLatentDist(z_norm)
+
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """
+        Decode from *normalized* latents (Cosmos contract).
+        """
+        self.clear_cache()
+        mean, inv_std = self._scale(latent)
+        z = latent / inv_std + mean  # z = z_norm * std + mean
+
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self._i0_decode(x)
+            else:
+                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+
+        self.clear_cache()
+        return out
+
+    # --- Interface helpers (match official Wan2pt1VAEInterface) ---
+    def get_latent_num_frames(self, num_pixel_frames: int) -> int:
+        return 1 + (int(num_pixel_frames) - 1) // 4
+
+    def get_pixel_num_frames(self, num_latent_frames: int) -> int:
+        return (int(num_latent_frames) - 1) * 4 + 1
+
+    @property
+    def spatial_compression_factor(self) -> int:
+        return 8
+
+    @property
+    def temporal_compression_factor(self) -> int:
+        return 4
+
+    @property
+    def latent_ch(self) -> int:
+        return 16
+
+
+
+
+
+
+
+
+

fastvideo/models/vaes/hunyuan15vae.py

@@ -663,7 +663,7 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE):
         # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
         # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
         # intermediate tiles together, the memory requirement can be lowered.
-        self.use_tiling = False
+        self.use_tiling = True
 
         # The minimal tile height and width for spatial tiling to be used
         self.tile_sample_min_height = 256

fastvideo/models/vaes/hyworldvae.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from diffusers and HY-WorldPlay
+
+# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from fastvideo.models.vaes.hunyuan15vae import AutoencoderKLHunyuanVideo15
+from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
+
+class AutoencoderKLHYWorld(AutoencoderKLHunyuanVideo15):
+    # TODO(mingjia): add temporal caching support for HYWorld VAE
+
+    def __init__(
+        self,
+        config: Hunyuan15VAEConfig,
+    ) -> None:
+        AutoencoderKLHunyuanVideo15.__init__(self, config)
+

fastvideo/optim/muon.py

@@ -0,0 +1,228 @@
+import math
+import torch
+try:
+    from torch.distributed.tensor import DTensor
+except ImportError:
+    # handle old pytorch versions
+    Dtensor = None
+
+
+# This code is modified from the GitHub repository of KellerJordan:
+# https://github.com/KellerJordan/Muon/blob/master/muon.py
+def zeropower_via_newtonschulz5(G, steps=5):
+    """
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    if isinstance(G, DTensor):
+        device_mesh = G.device_mesh
+        G = G.full_tensor()
+    else:
+        device_mesh = None
+
+    assert len(G.shape) >= 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+    # Perform the NS iterations
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * A @ A  # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    if device_mesh is not None:
+        return DTensor.from_local(X, device_mesh)
+    else:
+        return X
+
+
+class Muon(torch.optim.Optimizer):
+    """
+    Muon - MomentUm Orthogonalized by Newton-schulz
+
+    Arguments:
+        muon_params: The parameters to be optimized by Muon.
+        lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default)
+        momentum: The momentum used by the internal SGD. (0.95 is a good default)
+        nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
+        ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough)
+        adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are
+        {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well.
+        adamw_lr: The learning rate for the internal AdamW.
+        adamw_betas: The betas for the internal AdamW.
+        adamw_eps: The epsilon for the internal AdamW.
+        adamw_wd: The weight decay for the internal AdamW.
+    """
+
+    def __init__(
+            self,
+            lr=1e-3,
+            wd=0.1,
+            muon_params=None,
+            momentum=0.95,
+            nesterov=True,
+            ns_steps=5,
+            adamw_params=None,
+            adamw_betas=(0.95, 0.95),
+            adamw_eps=1e-8,
+    ):
+
+        defaults = dict(
+            lr=lr,
+            wd=wd,
+            momentum=momentum,
+            nesterov=nesterov,
+            ns_steps=ns_steps,
+            adamw_betas=adamw_betas,
+            adamw_eps=adamw_eps,
+        )
+
+        params = list(muon_params)
+        adamw_params = list(adamw_params) if adamw_params is not None else []
+        params.extend(adamw_params)
+        super().__init__(params, defaults)
+        # Sort parameters into those for which we will use Muon, and those for which we will not
+        for p in muon_params:
+            # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer
+            assert p.ndim >= 2, p.ndim
+            self.state[p]["use_muon"] = True
+        for p in adamw_params:
+            # Do not use Muon for parameters in adamw_params
+            self.state[p]["use_muon"] = False
+
+    def adjust_lr_for_muon(self, lr, param_shape):
+        A, B = param_shape[:2]
+        # We adjust the learning rate and weight decay based on the size of the parameter matrix
+        # as describted in the paper
+        adjusted_ratio = 0.2 * math.sqrt(max(A, B))
+        adjusted_lr = lr * adjusted_ratio
+        return adjusted_lr
+
+    def step(self, closure=None):
+        """Perform a single optimization step.
+
+        Args:
+            closure (Callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+
+            ############################
+            #           Muon           #
+            ############################
+
+            params = [p for p in group["params"] if self.state[p]["use_muon"]]
+
+            lr = group["lr"]
+            wd = group["wd"]
+            momentum = group["momentum"]
+
+            # generate weight updates in distributed fashion
+            for p in params:
+                # sanity check
+                g = p.grad
+                if g is None:
+                    continue
+                if g.ndim > 2:
+                    g = g.view(g.size(0), -1)
+                assert g is not None
+
+                # calc update
+                state = self.state[p]
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+                buf = state["momentum_buffer"]
+                buf.mul_(momentum).add_(g)
+                g = g.add(buf, alpha=momentum) if group["nesterov"] else buf
+                g = g.bfloat16()
+                u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
+
+                # scale update
+                adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+
+                # apply weight decay
+                p.data.mul_(1 - lr * wd)
+
+                # apply update
+                p.data.add_(u.view(p.shape), alpha=-adjusted_lr)
+
+            ############################
+            #       AdamW backup       #
+            ############################
+
+            params = [
+                p for p in group["params"] if not self.state[p]["use_muon"]
+            ]
+            lr = group['lr']
+            beta1, beta2 = group["adamw_betas"]
+            eps = group["adamw_eps"]
+            weight_decay = group["wd"]
+
+            for p in params:
+                g = p.grad
+                if g is None:
+                    continue
+                state = self.state[p]
+                if "step" not in state:
+                    state["step"] = 0
+                    state["moment1"] = torch.zeros_like(g)
+                    state["moment2"] = torch.zeros_like(g)
+                state["step"] += 1
+                step = state["step"]
+                buf1 = state["moment1"]
+                buf2 = state["moment2"]
+                buf1.lerp_(g, 1 - beta1)
+                buf2.lerp_(g.square(), 1 - beta2)
+
+                g = buf1 / (eps + buf2.sqrt())
+
+                bias_correction1 = 1 - beta1**step
+                bias_correction2 = 1 - beta2**step
+                scale = bias_correction1 / bias_correction2**0.5
+                p.data.mul_(1 - lr * weight_decay)
+                p.data.add_(g, alpha=-lr / scale)
+
+        return loss
+
+
+# help function to create the Muon optimizer
+def get_muon_optimizer(model,
+                       lr=1e-3,
+                       weight_decay=0.1,
+                       momentum=0.95,
+                       adamw_betas=(0.95, 0.95),
+                       adamw_eps=1e-8):
+    muon_params = [
+        p for name, p in model.named_parameters()
+        if p.requires_grad and p.ndim >= 2
+    ]
+    adamw_params = [
+        p for name, p in model.named_parameters()
+        if p.requires_grad and not (p.ndim >= 2)
+    ]
+
+    return Muon(
+        lr=lr,
+        wd=weight_decay,
+        muon_params=muon_params,
+        momentum=momentum,
+        adamw_params=adamw_params,
+        adamw_betas=adamw_betas,
+        adamw_eps=adamw_eps,
+    )

fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Cosmos 2.5 pipeline entry (staged pipeline)."""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (ConditioningStage,
+                                        Cosmos25DenoisingStage,
+                                        Cosmos25LatentPreparationStage,
+                                        DecodingStage, InputValidationStage,
+                                        Cosmos25TextEncodingStage,
+                                        Cosmos25TimestepPreparationStage)
+
+logger = init_logger(__name__)
+
+
+class Cosmos2_5Pipeline(ComposedPipelineBase):
+    """Cosmos 2.5 video generation pipeline."""
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler",
+        "safety_checker"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        logger.info("Creating Cosmos 2.5 pipeline stages...")
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=Cosmos25TextEncodingStage(
+                text_encoder=self.get_module("text_encoder"), ),
+        )
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=Cosmos25TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=Cosmos25LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer"),
+                           vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=Cosmos25DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+        logger.info("Cosmos 2.5 pipeline stages created")
+
+
+# Entry point for pipeline registry
+EntryClass = Cosmos2_5Pipeline

fastvideo/pipelines/basic/hunyuan15/hunyuan15_causal_dmd_pipeline.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Wan causal DMD pipeline implementation.
+
+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.pipelines import ComposedPipelineBase, LoRAPipeline
+
+# isort: off
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        Hy15CausalDMDDenosingStage,
+                                        InputValidationStage,
+                                        Hy15ImageEncodingStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage)
+# isort: on
+
+logger = init_logger(__name__)
+
+
+class Hy15CausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "scheduler"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=Hy15CausalDMDDenosingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = Hy15CausalDMDPipeline

fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Hunyuan video diffusion pipeline implementation.
+
+This module contains an implementation of the Hunyuan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        DenoisingStage, InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage,
+                                        Hy15ImageEncodingStage)
+
+# TODO(will): move PRECISION_TO_TYPE to better place
+
+logger = init_logger(__name__)
+
+
+class HunyuanVideo15ImageToVideoPipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae",
+        "transformer", "scheduler"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage_primary",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=Hy15ImageEncodingStage(image_encoder=None,
+                                                    image_processor=None))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = HunyuanVideo15ImageToVideoPipeline

fastvideo/pipelines/basic/hyworld/__init__.py

@@ -0,0 +1 @@
+

fastvideo/pipelines/basic/hyworld/hyworld_pipeline.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+HYWorld video diffusion pipeline implementation.
+
+This module contains an implementation of the HYWorld video diffusion pipeline
+using the modular pipeline architecture with HYWorld-specific denoising stage
+for chunk-based video generation with context frame selection.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (
+    ConditioningStage, DecodingStage, HYWorldDenoisingStage,
+    InputValidationStage, LatentPreparationStage, TextEncodingStage,
+    TimestepPreparationStage, HYWorldImageEncodingStage)
+
+logger = init_logger(__name__)
+
+
+class HYWorldPipeline(ComposedPipelineBase):
+    """
+    HYWorld video diffusion pipeline.
+
+    This pipeline implements chunk-based video generation with context frame
+    selection for 3D-aware generation using HYWorldDenoisingStage.
+
+    Note: HYWorld only uses a single LLM-based text encoder, unlike SDXL-style
+    dual encoder setups. The text_encoder_2/tokenizer_2 are not used.
+    """
+
+    # Include image_encoder and feature_extractor for I2V support with SigLIP
+    # Note: guider (ClassifierFreeGuidance) is not needed - FastVideo handles CFG differently
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler",
+        "text_encoder_2", "tokenizer_2", "image_encoder", "feature_extractor"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with HYWorld-specific denoising stage."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage_primary",
+                       stage=TextEncodingStage(
+                           text_encoders=[
+                               self.get_module("text_encoder"),
+                               self.get_module("text_encoder_2")
+                           ],
+                           tokenizers=[
+                               self.get_module("tokenizer"),
+                               self.get_module("tokenizer_2")
+                           ]))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=HYWorldImageEncodingStage(
+                           image_encoder=self.get_module("image_encoder"),
+                           image_processor=self.get_module("feature_extractor"),
+                           vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=HYWorldDenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = HYWorldPipeline

fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+LTX-2 text-to-video pipeline.
+"""
+
+import os
+from typing import Any
+
+from transformers import AutoTokenizer
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import PipelineComponentLoader
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage,
+                                        LTX2AudioDecodingStage,
+                                        LTX2DenoisingStage,
+                                        LTX2LatentPreparationStage,
+                                        TextEncodingStage)
+
+logger = init_logger(__name__)
+
+
+class LTX2Pipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "transformer",
+        "vae",
+        "audio_vae",
+        "vocoder",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage(),
+        )
+
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=LTX2LatentPreparationStage(
+                transformer=self.get_module("transformer"), ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=LTX2DenoisingStage(
+                transformer=self.get_module("transformer"), ),
+        )
+
+        self.add_stage(
+            stage_name="audio_decoding_stage",
+            stage=LTX2AudioDecodingStage(
+                audio_decoder=self.get_module("audio_vae"),
+                vocoder=self.get_module("vocoder"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=DecodingStage(vae=self.get_module("vae")),
+        )
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        tokenizer = self.get_module("tokenizer")
+        if tokenizer is not None:
+            tokenizer.padding_side = "left"
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+
+    def load_modules(
+        self,
+        fastvideo_args: FastVideoArgs,
+        loaded_modules: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        model_index = self._load_config(self.model_path)
+        logger.info("Loading pipeline modules from config: %s", model_index)
+
+        model_index.pop("_class_name")
+        model_index.pop("_diffusers_version")
+        model_index.pop("workload_type", None)
+
+        if len(model_index) <= 1:
+            raise ValueError(
+                "model_index.json must contain at least one pipeline module")
+
+        required_modules = self.required_config_modules
+        modules: dict[str, Any] = {}
+
+        for module_name, module_spec in model_index.items():
+            if not isinstance(module_spec, list) or len(module_spec) < 1:
+                continue
+            transformers_or_diffusers = module_spec[0]
+            if transformers_or_diffusers is None:
+                if module_name in self.required_config_modules:
+                    self.required_config_modules.remove(module_name)
+                continue
+            if module_name not in required_modules:
+                continue
+            if loaded_modules is not None and module_name in loaded_modules:
+                modules[module_name] = loaded_modules[module_name]
+                continue
+
+            component_model_path = os.path.join(self.model_path, module_name)
+            if module_name == "tokenizer" and not os.path.isdir(
+                    component_model_path):
+                gemma_path = os.path.join(self.model_path, "text_encoder",
+                                          "gemma")
+                if os.path.isdir(gemma_path):
+                    component_model_path = gemma_path
+                else:
+                    raise ValueError(
+                        "Tokenizer directory missing and Gemma weights were not found."
+                    )
+
+            module = PipelineComponentLoader.load_module(
+                module_name=module_name,
+                component_model_path=component_model_path,
+                transformers_or_diffusers=transformers_or_diffusers,
+                fastvideo_args=fastvideo_args,
+            )
+            logger.info("Loaded module %s from %s", module_name,
+                        component_model_path)
+            modules[module_name] = module
+
+        if "tokenizer" in required_modules and "tokenizer" not in modules:
+            gemma_path = os.path.join(self.model_path, "text_encoder", "gemma")
+            if os.path.isdir(gemma_path):
+                modules["tokenizer"] = AutoTokenizer.from_pretrained(
+                    gemma_path, local_files_only=True)
+
+        for module_name in required_modules:
+            if module_name not in modules or modules[module_name] is None:
+                raise ValueError(
+                    f"Required module {module_name} was not loaded properly")
+
+        return modules
+
+
+EntryClass = LTX2Pipeline

fastvideo/pipelines/composed_pipeline_base.py

@@ -84,6 +84,7 @@ class ComposedPipelineBase(ABC):
 
         # Load modules directly in initialization
         logger.info("Loading pipeline modules...")
+        # fastvideo_args.dit_cpu_offload = False
         with self.profiler_controller.region("profiler_region_model_loading"):
             self.modules = self.load_modules(fastvideo_args, loaded_modules)
 
@@ -287,6 +288,7 @@ class ComposedPipelineBase(ABC):
         # remove keys that are not pipeline modules
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
+        model_index.pop("_name_or_path", None)
         model_index.pop("workload_type", None)
         if "boundary_ratio" in model_index and model_index[
                 "boundary_ratio"] is not None: