Support custom action trajectories for validation

Overfitting running for MC 10
Update action_labels
2026-02-03 20:23:13 +00:00 · 2026-02-03 05:11:48 +00:00 · 2026-02-02 05:20:06 +00:00 · 2026-02-02 04:38:24 +00:00
293 changed files with 5206 additions and 18652 deletions
@@ -67,9 +67,6 @@ jobs:
            - torch-version: '2.9.1'
              cuda-version: '12.8.0'
              torch-cuda-short: 'cu128'
-            # - torch-version: '2.10.0'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'

    steps:
      - name: Free up disk space
@@ -223,6 +220,3 @@ jobs:
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          packages-dir: fastvideo-kernel/dist/
-          # PyPI does not allow replacing an existing file with the same name.
-          # This makes re-runs idempotent by skipping files already uploaded.
-          skip-existing: true
@@ -18,7 +18,6 @@ venv/
 .venv/
 runs/
 samples/
-Miniconda3-latest-Linux-x86_64.sh
 *validation/
 data/
 outputs/
@@ -33,11 +32,6 @@ env
 **.txt
 *.log
 weights/
-official_weights/
-converted_weights/
-
-# SSIM test outputs
-fastvideo/tests/ssim/generated_videos/

 # Distribution / packaging
 build/
@@ -75,11 +69,6 @@ docs/distillation/examples/
 !docs/assets/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
-!assets/images/**/*.png
-!assets/images/**/*.jpg
-!assets/images/**/*.jpeg
-!assets/images/**/*.gif
-!assets/videos/**/*.mp4

 dmd_t2v_output/
 preprocess_output_text/
@@ -10,7 +10,7 @@ exclude: |
    demo/.*|
    predict\.py|
    scripts/.*|
-    assets/prompts/.*|
+    prompts/.*|
    fastvideo/data_preprocess/.*|
    fastvideo/dataset/.*|
    fastvideo/models/.*|
@@ -1,42 +0,0 @@
-# Repository Guidelines
-
-## Project Structure & Module Organization
- Core Python package: `fastvideo/` (models, pipelines, training, distributed runtime, CLI entrypoints).
- CUDA/custom kernels: `fastvideo-kernel/` (separate build/test flow).
- Tests:
-  - `fastvideo/tests/` for package-level tests (dataset, encoders, inference, training, SSIM, workflow).
-  - `tests/local_tests/` for additional local/component checks.
- Docs and guides: `docs/` (MkDocs source), with contributor docs in `docs/contributing/`.
- Runnable examples and scripts: `examples/` and `scripts/`.
- Static assets: `assets/` (including `assets/images/`, `assets/videos/`, and `assets/prompts/`) and `comfyui/assets/`.
-
-## Build, Test, and Development Commands
- `uv pip install -e .[dev]`: editable install with lint/test extras.
- `pre-commit install --hook-type pre-commit --hook-type commit-msg`: enable local hooks.
- `pre-commit run --all-files`: run formatter/lint/type/spelling checks.
- `pytest tests/`: run top-level test suite.
- `pytest fastvideo/tests/ -v`: run package tests.
- `pytest fastvideo/tests/ssim/ -vs`: run SSIM regression tests (GPU-heavy).
- `cd fastvideo-kernel && ./build.sh`: build kernel extensions.
-
-## Coding Style & Naming Conventions
- Python 3.10+; 4-space indentation; keep code and imports readable and explicit.
- Style tools are configured in `pyproject.toml` and `.pre-commit-config.yaml`:
-  - `yapf` (format), `ruff` (lint, auto-fix), `mypy` (typing), `codespell`.
- Target line length is 80.
- Naming: `snake_case` for functions/files, `PascalCase` for classes, `UPPER_SNAKE_CASE` for constants.
-
-## Testing Guidelines
- Use `pytest` and place tests near relevant domains (e.g., `fastvideo/tests/encoders/`).
- Prefer descriptive names like `test_<feature>_<expected_behavior>.py`.
- For new pipelines/backends, include at least one regression-oriented test; add SSIM coverage when output quality must be preserved.
- Document GPU assumptions in tests that require specific hardware.
-
-## Commit & Pull Request Guidelines
- Follow existing commit style: short subject with optional tag prefix, e.g. `[bugfix]: ...`, `[feat]: ...`, `[misc]: ...`, and include PR reference like `(#1234)` when applicable.
- Keep commits focused by concern (feature, refactor, fix).
- PRs should include:
-  - clear problem/solution summary,
-  - test evidence (`pytest`/SSIM outputs or rationale if skipped),
-  - linked issue/PR context,
-  - screenshots or sample outputs for UI/demo/docs changes.
@@ -1 +0,0 @@
-@AGENTS.md
@@ -3,32 +3,33 @@
 </div>

 <p align="center">
-     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://github.com/hao-ai-lab/FastVideo/discussions/1097" target="_blank"> <b> WeChat </b> </a> |
+     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/sv3MMKyv" target="_blank"> <b> WeChat </b> </a> |
 </p>

 **FastVideo is a unified post-training and inference framework for accelerated video generation.**

 ## NEWS
+- ```2025/11/19```: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
+- ```2025/08/04```: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

- `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
- `2025/08/04`: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
+<details>
+<summary>More</summary>

-### More News
+- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
+- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
+- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).

- `2025/06/14`: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
- `2025/04/24`: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
- `2025/02/18`: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+</details>

 ## Key Features

 FastVideo has the following features:
-
 - End-to-end post-training support for bidirectional and autoregressive models:
  - Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs
  - Data preprocessing pipeline for video, image, and text data
  - Distribution Matching Distillation (DMD2) stepwise distillation.
  - Sparse attention with [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
-  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) to achieve >50x denoising speedup
+  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) to achineve >50x denoising speedup
  - Scalable training with FSDP2, sequence parallelism, and selective activation checkpointing.
  - Causal distillation through Self-Forcing
  - See this [page](https://hao-ai-lab.github.io/FastVideo/training/overview/) for full list of supported models and recipes.
@@ -43,7 +44,6 @@ FastVideo has the following features:
  - See this [page](https://hao-ai-lab.github.io/FastVideo/inference/hardware_support/) for full list of supported hardware and OS.

 ## Getting Started
-
 We recommend using an environment manager such as `Conda` to create a clean environment:

 ```bash
@@ -58,20 +58,18 @@ pip install fastvideo
 Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.

 ## Sparse Distillation
-
 For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

 See below for recipes and datasets:

-| Model                                                                                 | Sparse Distillation                                                                                             | Dataset                                                                                                  |
-| ------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------- |
-| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers) | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)       | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)     |
-| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)   | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free) | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k) |
+|                                            Model                                              |                                               Sparse Distillation                                                 |                                                  Dataset                                                  |
+|:-------------------------------------------------------------------------------------------:  |:---------------------------------------------------------------------------------------------------------------:  |:--------------------------------------------------------------------------------------------------------: |
+| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)         |    [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)      | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)      |
+| [FastWan2.1-T2V-14B-Preview](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-Diffusers)   |                                                   Coming soon!                                                    |   [FastVideo Synthetic Wan2.1 720P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x768x1280_250k)   |
+| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)           | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free)   | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k)  |

 ## Inference
-
 ### Generating Your First Video
-
 Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation/). Create a file called `example.py` with the following code:

 ```python
@@ -110,36 +108,55 @@ python example.py

 For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).

-## More Guides
+### Other docs:

 - [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/)
+
+## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
+<!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->

 ## Awesome work using FastVideo or our research projects

- [SGLang](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen): SGLang's diffusion inference functionality is based on a fork of FastVideo on Sept. 24, 2025.
- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo.
- [SRPO](https://github.com/Tencent-Hunyuan/SRPO): A method to directly align the full diffusion trajectory with fine-grained human preference. Code based on FastVideo.
- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo.
- [Hunyuan Video 1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5): A leading lightweight video generation model, where they proposed SSTA based on Sliding Tile Attention.
- [Kandinsky-5.0](https://github.com/kandinskylab/kandinsky-5): A family of diffusion models for video & image generation, where their NABLA attention includes a Sliding Tile Attention branch.
- [LongCat Video](https://github.com/meituan-longcat/LongCat-Video): A foundational video generation model with 13.6B parameters with block-sparse attention similar to Video Sparse Attention.
+- [SGLang](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen): SGLang's diffusion inference functionality is based  on a fork of FastVideo on Sept. 24, 2025. [![Star](https://img.shields.io/github/stars/sgl-project/sglang.svg?style=social&label=Star)](https://github.com/sgl-project/sglang)
+
+- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A  unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/XueZeyue/DanceGRPO.svg?style=social&label=Star)](https://github.com/XueZeyue/DanceGRPO)
+- [SRPO](https://github.com/Tencent-Hunyuan/SRPO): A method to directly align the full diffusion trajectory with fine-grained human preference. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/SRPO.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/SRPO)
+- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Vchitect/DCM.svg?style=social&label=Star)](https://github.com/Vchitect/DCM)
+- [Hunyuan Video 1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5): A leading lightweight video generation model, where they proposed SSTA based on Sliding Tile Attention. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/HunyuanVideo-1.5.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5)
+- [Kandinsky-5.0](https://github.com/kandinskylab/kandinsky-5): A family of diffusion models for video & image generation, where their NABLA attention includes a Sliding Tile Attention branch. [![Star](https://img.shields.io/github/stars/kandinskylab/kandinsky-5.svg?style=social&label=Star)](https://github.com/kandinskylab/kandinsky-5)
+- [LongCat Video](https://github.com/meituan-longcat/LongCat-Video): A foundational video generation model with 13.6B parameters with block-sparse attention similar to Video Sparse Attention. [![Star](https://img.shields.io/github/stars/meituan-longcat/LongCat-Video.svg?style=social&label=Star)](https://github.com/meituan-longcat/LongCat-Video)

 ## 🤝 Contributing

 We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview/).
 See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/899).
-
 ## Acknowledgement
+We learned and reused code from the following projects:
+- [Wan-Video](https://github.com/Wan-Video)
+- [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)
+- [Triton](https://github.com/triton-lang/triton)
+- [DMD2](https://github.com/tianweiy/DMD2)
+- [diffusers](https://github.com/huggingface/diffusers)
+- [xDiT](https://github.com/xdit-project/xDiT)
+- [vLLM](https://github.com/vllm-project/vllm)
+- [SGLang](https://github.com/sgl-project/sglang)

-We learned the design and reused code from the following projects: [Wan-Video](https://github.com/Wan-Video), [ThunderKittens](https://github.com/HazyResearch/ThunderKittens), [DMD2](https://github.com/tianweiy/DMD2), [diffusers](https://github.com/huggingface/diffusers), [xDiT](https://github.com/xdit-project/xDiT), [vLLM](https://github.com/vllm-project/vllm), [SGLang](https://github.com/sgl-project/sglang). We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.
+We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.

 ## Citation
-
-If you find FastVideo useful, please consider citing our research work:
+If you find FastVideo useful, please considering citing our work:

 ```bibtex
+@software{fastvideo2024,
+  title        = {FastVideo: A Unified Framework for Accelerated Video Generation},
+  author       = {The FastVideo Team},
+  url          = {https://github.com/hao-ai-lab/FastVideo},
+  month        = apr,
+  year         = {2024},
+}
+
@article{zhang2025vsa,
  title={Vsa: Faster video diffusion with trainable sparse attention},
  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
@@ -1,7 +0,0 @@
-# FastVideo/assets/videos
-
-This folder is used to store **video assets for examples**, primarily **input videos** consumed by scripts under `FastVideo/examples/`.
-
- **Typical contents**: short input clips for demos (e.g., video2world / image2video examples).
- **Non-critical**: these assets are for convenience and are not required to use the FastVideo library.
- **Large files**: avoid committing large videos to git; prefer shared storage or download-on-demand.
@@ -0,0 +1,195 @@
+import argparse
+import os
+import tempfile
+
+import gradio as gr
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import export_to_video
+
+from fastvideo.distill.solver import PCMFMScheduler
+from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
+from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
+
+
+def init_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompts", nargs="+", default=[])
+    parser.add_argument("--num_frames", type=int, default=25)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=848)
+    parser.add_argument("--num_inference_steps", type=int, default=8)
+    parser.add_argument("--guidance_scale", type=float, default=4.5)
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--seed", type=int, default=12345)
+    parser.add_argument("--transformer_path", type=str, default=None)
+    parser.add_argument("--scheduler_type", type=str, default="pcm_linear_quadratic")
+    parser.add_argument("--lora_checkpoint_dir", type=str, default=None)
+    parser.add_argument("--shift", type=float, default=8.0)
+    parser.add_argument("--num_euler_timesteps", type=int, default=50)
+    parser.add_argument("--linear_threshold", type=float, default=0.1)
+    parser.add_argument("--linear_range", type=float, default=0.75)
+    parser.add_argument("--cpu_offload", action="store_true")
+    return parser.parse_args()
+
+
+def load_model(args):
+    if args.scheduler_type == "euler":
+        scheduler = FlowMatchEulerDiscreteScheduler()
+    else:
+        linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False
+        scheduler = PCMFMScheduler(
+            1000,
+            args.shift,
+            args.num_euler_timesteps,
+            linear_quadratic,
+            args.linear_threshold,
+            args.linear_range,
+        )
+
+    if args.transformer_path:
+        transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
+    else:
+        transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder="transformer/")
+
+    pipe = MochiPipeline.from_pretrained(args.model_path, transformer=transformer, scheduler=scheduler)
+    pipe.enable_vae_tiling()
+    # pipe.to(device)
+    # if args.cpu_offload:
+    pipe.enable_sequential_cpu_offload()
+    return pipe
+
+
+def generate_video(
+    prompt,
+    negative_prompt,
+    use_negative_prompt,
+    seed,
+    guidance_scale,
+    num_frames,
+    height,
+    width,
+    num_inference_steps,
+    randomize_seed=False,
+):
+    if randomize_seed:
+        seed = torch.randint(0, 1000000, (1, )).item()
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    if not use_negative_prompt:
+        negative_prompt = None
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        output = pipe(
+            prompt=[prompt],
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).frames[0]
+
+    output_path = os.path.join(tempfile.mkdtemp(), "output.mp4")
+    export_to_video(output, output_path, fps=30)
+    return output_path, seed
+
+
+examples = [
+    "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
+    "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
+    "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
+]
+
+args = init_args()
+pipe = load_model(args)
+print("load model successfully")
+with gr.Blocks() as demo:
+    gr.Markdown("# Fastvideo Mochi Video Generation Demo")
+
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button("Run", scale=0)
+        result = gr.Video(label="Result", show_label=False)
+
+    with gr.Accordion("Advanced options", open=False):
+        with gr.Group():
+            with gr.Row():
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=1024,
+                    step=32,
+                    value=args.height,
+                )
+                width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=args.width)
+
+            with gr.Row():
+                num_frames = gr.Slider(
+                    label="Number of Frames",
+                    minimum=21,
+                    maximum=163,
+                    value=args.num_frames,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1,
+                    maximum=12,
+                    value=args.guidance_scale,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Inference Steps",
+                    minimum=4,
+                    maximum=100,
+                    value=args.num_inference_steps,
+                )
+
+            with gr.Row():
+                use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=False,
+            )
+
+            seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=args.seed)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            seed_output = gr.Number(label="Used Seed")
+
+    gr.Examples(examples=examples, inputs=prompt)
+
+    use_negative_prompt.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt,
+        outputs=negative_prompt,
+    )
+
+    run_button.click(
+        fn=generate_video,
+        inputs=[
+            prompt,
+            negative_prompt,
+            use_negative_prompt,
+            seed,
+            guidance_scale,
+            num_frames,
+            height,
+            width,
+            num_inference_steps,
+            randomize_seed,
+        ],
+        outputs=[result, seed_output],
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
@@ -0,0 +1,15 @@
+Fast-Hunyuan comparison with original Hunyuan, achieving an 8X diffusion speed boost with the FastVideo framework.
+
+https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30
+
+Fast-Mochi comparison with original Mochi, achieving an 8X diffusion speed boost with the FastVideo framework.
+
+https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c
+
+Comparison between OpenAI Sora, original Hunyuan and FastHunyuan
+
+https://github.com/user-attachments/assets/d323b712-3f68-42b2-952b-94f6a49c4836
+
+Comparison between original FastHunyuan, LLM-INT8 quantized FastHunyuan and NF4 quantized FastHunyuan
+
+https://github.com/user-attachments/assets/cf89efb5-5f68-4949-a085-f41c1ef26c94
@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp310-cp310-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp310-cp310-linux_x86_64.whl

 COPY . .

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp311-cp311-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp311-cp311-linux_x86_64.whl

 COPY . .

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp312-cp312-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp312-cp312-linux_x86_64.whl

 COPY . .

@@ -121,15 +121,6 @@ This page contains the complete API reference for the FastVideo library.
      show_root_toc_entry: true
      heading_level: 4

-## fastvideo.registry
-
-::: fastvideo.registry
-    options:
-      show_source: true
-      show_root_heading: true
-      show_root_toc_entry: true
-      heading_level: 3
-
 ## fastvideo.pipelines

 ::: fastvideo.pipelines
@@ -324,8 +324,8 @@ Purpose:
 Action:

 - Add a new pipeline config + sampling params.
- Register them in `fastvideo/registry.py` using explicit
-  `register_configs(...)` blocks (this file is the single source of truth now).
+- Register them in `fastvideo/configs/pipelines/registry.py` and
+  `fastvideo/configs/sample/registry.py`.

 ### 5) Wire pipeline stages

@@ -17,7 +17,7 @@ You can easily use the FastVideo Pod Template on [RunPod](https://www.runpod.io)
 - Select the "FastVideo" or "fastvideo-dev" Pod Template.
 ![RunPod Pod Template Selection](../../assets/images/runpod_create.png)

- Set the Pod name to "`<name>-<FastVideo>-<date>`".
+- Set the Pod name to "<name>-<FastVideo>-<date>".

 - Finally, once the pod is deployed (will take a few minutes as the image is being pulled), you can SSH into it using "SSH exposed over TCP". You'll need to use the matching private ssh key you provided.
 ![RunPod SSH](../../assets/images/runpod_ssh.png)
@@ -51,9 +51,8 @@ runtime parameters consistent:
 - `fastvideo/configs/pipelines/`: pipeline wiring and required components.
 - `fastvideo/configs/sample/`: default sampling parameters (steps, frames,
  guidance scale, resolution, fps).
- `fastvideo/registry.py`: unified registry for pipeline config + sampling
-  defaults and model metadata resolution, defined via explicit
-  `register_configs(...)` blocks (no separate dict registries).
+- `fastvideo/configs/registry.py`: pipeline registry.
+- `fastvideo/configs/sample/registry.py`: sampling param registry.

 `FastVideoArgs` (in `fastvideo/fastvideo_args.py`) provides runtime settings and
 is passed into pipeline construction and stages.
@@ -131,41 +131,17 @@ self.attn = DistributedAttention(

 ### Registering Models

-Register implemented modules for auto‑discovery by adding `EntryClass` in each
-model module (the registry scans for it):
+Register implemented modules in the model registry:

 ```python
-# In fastvideo/models/dits/your_module.py
-class YourTransformerModel(...):
-    ...
+# In fastvideo/models/registry.py
+_TEXT_TO_VIDEO_DIT_MODELS = {
+    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
+}

-# Entry point for model registry
-EntryClass = YourTransformerModel
-```
-
-```python
-# In fastvideo/models/vaes/your_vae.py
-class YourVAEModel(...):
-    ...
-
-# Entry point for model registry
-EntryClass = YourVAEModel
-```
-
-Register pipeline config + sampling defaults in the unified registry:
-
-```python
-# In fastvideo/registry.py
-register_configs(
-    sampling_param_cls=YourSamplingParam,
-    pipeline_config_cls=YourPipelineConfig,
-    hf_model_paths=[
-        "org/your-model-id",
-    ],
-    model_detectors=[
-        lambda path: "your-model" in path.lower(),
-    ],
-)
+_VAE_MODELS = {
+    "YourVAEModel": ("vaes", "yourvae", "YourVAEClass"),
+}
 ```

 ## Step 2: Directory Structure
@@ -1,50 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam
-
-
-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,
-    )
-
-    sampling_param = SamplingParam.from_pretrained(model_path)
-
-    # image2world example from official repo
-    image_path = "assets/images/bus_terminal.jpg"
-
-    prompt = (
-        "A nighttime city bus terminal gradually shifts from stillness to subtle movement. "
-        "At first, multiple double-decker buses are parked under the glow of overhead lights, "
-        "with a central bus labeled '87D' facing forward and stationary. "
-        "As the video progresses, the bus in the middle moves ahead slowly, its headlights brightening the surrounding area "
-        "and casting reflections onto adjacent vehicles. "
-        "The motion creates space in the lineup, signaling activity within the otherwise quiet station. "
-        "It then comes to a smooth stop, resuming its position in line. "
-        "Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene."
-    )
-
-    generator.generate_video(
-        prompt,
-        sampling_param=sampling_param,
-        image_path=str(image_path),
-        num_cond_frames=1,
-        output_path="outputs_video/cosmos2_5_i2w.mp4",
-        save_video=True,
-    )
-
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
-
@@ -1,5 +1,4 @@
 from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam


 def main():
@@ -16,27 +15,19 @@ def main():
        pin_cpu_memory=True,
    )

-    # Load default sampling parameters (negative_prompt, resolution, steps, etc.)
-    sampling_param = SamplingParam.from_pretrained(model_path)
-
    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."
+       "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."
    )

-    generator.generate_video(
+    video = generator.generate_video(
        prompt,
-        sampling_param=sampling_param,
+        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,
    )
@@ -1,53 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam
-
-
-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,
-    )
-
-    sampling_param = SamplingParam.from_pretrained(model_path)
-
-    # video2world example from official repo
-    video_path = "assets/videos/robot_pouring.mp4"
-
-    prompt = (
-        "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. "
-        "The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. "
-        "The robotic arm is in the process of pouring a transparent liquid into the glass. "
-        "To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. "
-        "In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. "
-        "The lighting is bright, casting soft shadows on the table. "
-        "The robotic arm's movements are smooth and controlled, demonstrating precision in its task. "
-        "As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. "
-        "The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. "
-        "The other robotic arm on the right side also stays stationary throughout the video. "
-        "The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper."
-    )
-
-    generator.generate_video(
-        prompt,
-        sampling_param=sampling_param,
-        video_path=str(video_path),
-        num_cond_frames=1,
-        output_path="outputs_video/cosmos2_5_v2w.mp4",
-        save_video=True,
-    )
-
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
-
@@ -1,119 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Basic inference script for HunyuanGameCraft video generation.
-
-HunyuanGameCraft generates game-like videos with camera/action control.
-It takes an optional image input and generates video with camera motion
-based on simple action commands (forward, left, right, backward, rotations).
-
-Available actions:
-    - forward (w): Move camera forward
-    - backward (s): Move camera backward  
-    - left (a): Move camera left (strafe)
-    - right (d): Move camera right (strafe)
-    - left_rot: Rotate camera left (pan)
-    - right_rot: Rotate camera right (pan)
-    - up_rot: Rotate camera up (tilt)
-    - down_rot: Rotate camera down (tilt)
-
-T2V vs I2V:
-    - Default: I2V (uses a default reference image). Set GAMECRAFT_I2V_IMAGE to a
-      URL or path to use a different image.
-    - T2V only (no reference image): run with GAMECRAFT_I2V_IMAGE= (empty).
-"""
-import os
-
-import torch
-
-from fastvideo import VideoGenerator
-from fastvideo.models.camera import create_camera_trajectory
-
-# Model configuration (use GAMECRAFT_MODEL_PATH for local weights)
-MODEL_PATH = os.environ.get("GAMECRAFT_MODEL_PATH", "FastVideo/HunyuanGameCraft-Diffusers")
-
-# Default prompts for demo
-DEFAULT_PROMPTS = {
-    "village": "A charming medieval village with cobblestone streets, thatched-roof houses, and vibrant flower gardens under a bright blue sky.",
-    "temple": "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.",
-    "forest": "A lush green forest with tall trees, dappled sunlight filtering through the leaves, and a winding dirt path.",
-    "beach": "A tropical beach with crystal clear turquoise water, white sand, and palm trees swaying in the breeze.",
-}
-
-# I2V: default reference image (URL). Can override with a local path.
-DEFAULT_I2V_IMAGE_URL = (
-    "https://huggingface.co/datasets/huggingface/documentation-images/"
-    "resolve/main/diffusers/astronaut.jpg"
-)
-DEFAULT_I2V_PROMPT = (
-    "An astronaut hatching from an egg, on the surface of the moon, "
-    "the darkness and depth of space realised in the background."
-)
-
-OUTPUT_PATH = "video_samples_gamecraft"
-
-
-def main():
-    # Initialize generator
-    # FastVideo will automatically download weights from HuggingFace
-    generator = VideoGenerator.from_pretrained(
-        MODEL_PATH,
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True,
-        vae_cpu_offload=True,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True,
-    )
-
-    # Video parameters
-    height = 704
-    width = 1280
-    num_frames = 33
-    action = "forward"
-    action_speed = 0.2
-
-    # Create camera trajectory (Plücker coordinates)
-    camera_states = create_camera_trajectory(
-        action=action,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        action_speed=action_speed,
-        dtype=torch.bfloat16,
-    )
-    print(f"Camera states shape: {camera_states.shape}")
-
-    # I2V vs T2V: unset GAMECRAFT_I2V_IMAGE -> I2V (default image). Set to "" -> T2V.
-    env_image = os.environ.get("GAMECRAFT_I2V_IMAGE")
-    if env_image is None:
-        image_path = DEFAULT_I2V_IMAGE_URL  # default: I2V
-    elif env_image.strip() == "":
-        image_path = None  # T2V
-    else:
-        image_path = env_image.strip()  # I2V with given URL/path
-
-    is_i2v = image_path is not None
-    prompt = DEFAULT_I2V_PROMPT if is_i2v else DEFAULT_PROMPTS["temple"]
-    print(f"Mode: {'I2V' if is_i2v else 'T2V'}, prompt: {prompt[:60]}...")
-
-    gen_kw = dict(
-        prompt=prompt,
-        negative_prompt="",
-        camera_states=camera_states,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        num_inference_steps=50,
-        guidance_scale=6.0,
-        seed=42,
-        fps=24,
-        output_path=OUTPUT_PATH,
-        save_video=True,
-    )
-    if is_i2v:
-        gen_kw["image_path"] = image_path
-    generator.generate_video(**gen_kw)
-
-
-if __name__ == "__main__":
-    main()
@@ -39,4 +39,4 @@ def main():


 if __name__ == "__main__":
-    main()
+    main()
@@ -1,43 +0,0 @@
-from fastvideo import VideoGenerator
-import json
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_hy15_1080p"
-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(
-        "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p
-        # or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=False, # set to True if GPU is out of memory
-        dit_cpu_offload=True,
-        vae_cpu_offload=True,
-        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,
-    )
-
-    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="")
-
-    prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame "
-        "glistening under the warm afternoon sun. The tall grass ripples gently in "
-        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-        "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="")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,11 +1,165 @@
+# 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'

-OUTPUT_PATH = "video_samples_hyworld"
 def main():
    import argparse

@@ -15,7 +169,7 @@ def main():
    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=OUTPUT_PATH, help="Output video path")
+    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")
@@ -31,7 +185,8 @@ def main():

    # Initialize generator
    print("\nInitializing VideoGenerator for HYWorld...")
-    generator = VideoGenerator.from_pretrained(
+    
+    generator = HYWorldVideoGenerator.from_pretrained(
        "FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
        num_gpus=1,
        use_fsdp_inference=True,
@@ -43,12 +198,11 @@ def main():
    )

    # Generate video
-    # The pose string is automatically converted to camera matrices by the pipeline
    print("\nGenerating video...")
-    generator.generate_video(
+    start_time = time.time()
+    video = generator.generate_video(
        prompt=args.prompt,
        image_path=args.image,
-        pose=args.pose,  # Camera trajectory control
        output_path=args.output_path,
        save_video=True,
        negative_prompt="",
@@ -57,9 +211,13 @@ def main():
        height=HEIGHT,
        width=WIDTH,
        seed=args.seed,
+        pose=args.pose,
    )
+    elapsed = time.time() - start_time

-    print(f"\nVideo saved to: {args.output_path}")
+    print(f"\nVideo generated successfully!")
+    print(f"Saved to: {args.output_path}")
+    print(f"Time: {elapsed:.2f}s")


 if __name__ == "__main__":
@@ -1,49 +0,0 @@
-from fastvideo import VideoGenerator
-from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
-
-# from fastvideo.configs.sample import SamplingParam
-OUTPUT_PATH = "video_samples_lingbotworld"
-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(
-         "FastVideo/LingBot-World-Base-Cam-Diffusers",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=False, # set to True if GPU is out of memory
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=True,
-        # image_encoder_cpu_offload=False,
-    )
-
-    num_frames = 81
-    prompt = "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls."
-    image_path = "https://raw.githubusercontent.com/Robbyant/lingbot-world/main/examples/00/image.jpg"
-    action_path = "examples/inference/basic/lingbotworld_examples/00"
-    c2ws_plucker_emb, num_frames = prepare_camera_embedding(
-        action_path=action_path,
-        num_frames=num_frames,
-        height=480,
-        width=832,
-        spatial_scale=8,
-    )
-
-    generator.generate_video(
-        prompt,
-        image_path=image_path,
-        output_path=OUTPUT_PATH,
-        save_video=True,
-        num_frames=num_frames,
-        height=480,
-        width=832,
-        c2ws_plucker_emb=c2ws_plucker_emb,
-    )
-
-
-if __name__ == "__main__":
-    main()
@@ -1,6 +1,5 @@
 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 "
@@ -17,26 +16,16 @@ PROMPT = (


 def main() -> None:
-    # Uses FastVideo default sampling settings for LTX2 base.
    generator = VideoGenerator.from_pretrained(
-        "Davids048/LTX2-Base-Diffusers",
-        num_gpus=8,
+        "FastVideo/LTX2-Distilled-Diffusers",
+        num_gpus=1,
    )

-    output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.4.mp4"
+    output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4"
    generator.generate_video(
        prompt=PROMPT,
        output_path=output_path,
        save_video=True,
-        num_frames=121,
-        height=1088,
-        width=1920,
-        # LTX2 uses these parameters for multi-modal CFG instead of guidance_scale
-        # ltx2_cfg_scale_video=3.0,
-        # ltx2_cfg_scale_audio=7.0,
-        # ltx2_modality_scale_video=3.0,
-        # ltx2_modality_scale_audio=3.0,
-        # ltx2_rescale_scale=0.7,
    )
    generator.shutdown()

@@ -1,34 +0,0 @@
-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()
@@ -1,137 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import os
-import re
-from typing import List
-
-
-DEFAULT_PROMPTS = [
-    "a photo of a cat",
-    "a cinematic photo of a red panda wearing a tiny backpack, standing on a rainy neon-lit street at night, shallow depth of field, sharp focus, 35mm, bokeh",
-]
-
-
-def _safe_filename(text: str, max_len: int = 100) -> str:
-    """
-    Make a stable, filesystem-friendly filename base.
-    VideoGenerator uses prompt[:100].strip() internally, so we mirror that,
-    but also remove path separators and other problematic characters.
-    """
-    s = text[:max_len].strip()
-    s = s.replace(os.sep, "_")
-    if os.altsep:
-        s = s.replace(os.altsep, "_")
-    s = re.sub(r"\s+", " ", s)
-    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
-    s = s.strip(" .")
-    return s or "prompt"
-
-
-def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
-    """
-    Ensure deterministic naming by deleting any existing mp4s that would
-    cause VideoGenerator to append suffixes like _1, _2, etc.
-    """
-    if not os.path.isdir(out_dir):
-        return
-
-    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.mp4$")
-    for fn in os.listdir(out_dir):
-        if pattern.match(fn):
-            try:
-                os.remove(os.path.join(out_dir, fn))
-            except FileNotFoundError:
-                pass
-
-
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(description="Run SD3.5 Medium text-to-image with FastVideo VideoGenerator.")
-    p.add_argument("--model-path", default="stabilityai/stable-diffusion-3.5-medium", help="Path to local diffusers-format SD3.5 weights directory.")
-    p.add_argument(
-        "--out-dir",
-        "--outdir",
-        default="outputs/sd35/samples",
-        help="Output directory for generated mp4 files.",
-    )
-    p.add_argument(
-        "--prompt",
-        action="append",
-        default=None,
-        help="Prompt text. Repeat --prompt multiple times to generate multiple samples.",
-    )
-    p.add_argument("--negative", default="lowres, blurry, jpeg artifacts, watermark, text", help="Negative prompt.")
-    p.add_argument(
-        "--backend",
-        default=None,
-        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA). If omitted, respects the existing env var.",
-    )
-    p.add_argument("--seed", type=int, default=42, help="Base seed. Each prompt uses seed + prompt_idx.")
-    p.add_argument("--height", type=int, default=768, help="Output height.")
-    p.add_argument("--width", type=int, default=768, help="Output width.")
-    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
-    p.add_argument("--guidance", type=float, default=6.0, help="Guidance scale.")
-    p.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs to use.")
-    return p.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    prompts: List[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
-
-    if args.backend:
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
-
-    from fastvideo import VideoGenerator
-
-    os.makedirs(args.out_dir, exist_ok=True)
-
-    init_kwargs = {
-        "num_gpus": args.num_gpus,
-        "workload_type": "t2i",
-        "sp_size": 1,
-        "tp_size": 1,
-        "dit_cpu_offload": False,
-        "dit_layerwise_offload": False,
-        "text_encoder_cpu_offload": False,
-        "vae_cpu_offload": False,
-        "image_encoder_cpu_offload": False,
-        "pin_cpu_memory": False,
-        "use_fsdp_inference": False,
-    }
-
-    generator = VideoGenerator.from_pretrained(model_path=args.model_path, **init_kwargs)
-    try:
-        for i, prompt in enumerate(prompts):
-            seed = args.seed + i
-
-            filename_base = f"sd35_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
-            _remove_existing_outputs(args.out_dir, filename_base)
-
-            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
-            print(f"[sd35] prompt_idx={i} seed={seed} output_path={output_path}")
-
-            generation_kwargs = {
-                "output_path": output_path,
-                "height": args.height,
-                "width": args.width,
-                "num_frames": 1,
-                "fps": 1,
-                "num_inference_steps": args.steps,
-                "guidance_scale": args.guidance,
-                "seed": seed,
-                "negative_prompt": args.negative,
-                "save_video": True,
-            }
-
-            generator.generate_video(prompt, **generation_kwargs)
-
-        print(f"[sd35] done. outputs written to: {args.out_dir}")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
@@ -31,7 +31,7 @@ def main():
    sampling_param.height = 480
    sampling_param.seed = 1000

-    with open("assets/prompts/mixkit_i2v.jsonl", "r") as f:
+    with open("prompts/mixkit_i2v.jsonl", "r") as f:
        prompt_image_pairs = json.load(f)

    for prompt_image_pair in prompt_image_pairs:
@@ -188,9 +188,8 @@ def load_example_prompts():
    prompt_to_image = {}
    # Try to find the JSON file relative to project root
    possible_json_paths = [
-        Path("assets/prompts/mixkit_i2v.jsonl"),
-        Path(__file__).resolve().parents[4] / "assets" / "prompts" /
-        "mixkit_i2v.jsonl",
+        Path("prompts/mixkit_i2v.jsonl"),
+        Path(__file__).parent.parent.parent.parent / "prompts" / "mixkit_i2v.jsonl",
    ]
    json_path = None
    for path in possible_json_paths:
@@ -202,8 +201,8 @@ def load_example_prompts():
        try:
            with open(json_path, "r", encoding='utf-8') as f:
                data = json.load(f)
-                # Resolve paths relative to repository root.
-                project_root = Path(__file__).resolve().parents[4]
+                # Get the project root directory (parent of prompts directory)
+                project_root = json_path.parent.parent
                for item in data:
                    prompt_text = item.get("prompt", "").strip()
                    image_path = item.get("image_path", "")
@@ -737,8 +736,8 @@ def main():
        allowed_paths=[
            os.path.abspath("outputs"), 
            os.path.abspath("fastvideo-logos"),
-            os.path.abspath("assets/prompts"),
-            os.path.abspath("assets/images"),
+            os.path.abspath("prompts"),
+            os.path.abspath("images"),
            os.path.abspath(tempfile.gettempdir()),
            os.path.abspath(os.path.join(tempfile.gettempdir(), "gradio")),
        ]
@@ -748,4 +747,4 @@ def main():


 if __name__ == "__main__":
-    main() 
+    main() 
@@ -4,7 +4,7 @@ export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false

-MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
 VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
 NUM_GPUS=1
@@ -14,6 +14,7 @@ NUM_GPUS=1
 training_args=(
  --tracker_project_name "wan_ode_init"
  --output_dir "wan_ode_init_crush_smol"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
  --wandb_run_name "wan_ode_init_crush_smol"
  --max_train_steps 6000
  --train_batch_size 1
@@ -33,7 +34,7 @@ parallel_args=(
  --sp_size 1
  --tp_size 1
  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
+  --hsdp_shard_dim 1
 )

 # Model arguments
@@ -50,17 +51,20 @@ dataset_args=(

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

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

@@ -0,0 +1,18 @@
+Total Files: 16
+
+00. Hold [W] + Static
+01. Hold [S] + Static
+02. Hold [A] + Static
+03. Hold [D] + Static
+04. Hold [WA] + Static
+05. Hold [WD] + Static
+06. Hold [SA] + Static
+07. Hold [SD] + Static
+08. No Key + Hold [up]
+09. No Key + Hold [down]
+10. No Key + Hold [left]
+11. No Key + Hold [right]
+12. No Key + Hold [up_right]
+13. No Key + Hold [up_left]
+14. No Key + Hold [down_right]
+15. No Key + Hold [down_left]
@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=offline
+export TOKENIZERS_PARALLELISM=false
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+MODEL_PATH="weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers"
+DATA_DIR="mc_wasd_10/preprocessed/combined_parquet_dataset"
+VALIDATION_DATASET_FILE="mc_wasd_10/validation.json"
+NUM_GPUS=4
+# export CUDA_VISIBLE_DEVICES=0,1,2,3
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wangame_1.3b_overfit"
+  --output_dir "wangame_1.3b_overfit"
+  --max_train_steps 1500
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 20
+  --num_height 352
+  --num_width 640
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wangame_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,120 @@
+#!/bin/bash
+#SBATCH --job-name=wangame_1.3b_overfit
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=wangame_1.3b_overfit_output/wangame_1.3b_overfit_%j.out
+#SBATCH --error=wangame_1.3b_overfit_output/wangame_1.3b_overfit_%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"
+
+MODEL_PATH="weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers"
+DATA_DIR="mc_wasd_10/preprocessed/combined_parquet_dataset"
+VALIDATION_DATASET_FILE="examples/training/finetune/WanGame2.1_1.3b_i2v/validation.json"
+# Configs
+NUM_GPUS=8
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wangame_1.3b_overfit"
+  --output_dir "wangame_1.3b_overfit"
+  --max_train_steps 15000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 20
+  --num_height 352
+  --num_width 640
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "1.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000000
+  --training_state_checkpointing_steps 10000000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wangame_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,193 @@
+import os
+import numpy as np
+
+# Configuration
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+BASE_OUTPUT_DIR = os.path.join(SCRIPT_DIR, 'action')
+VIDEO_OUTPUT_DIR = BASE_OUTPUT_DIR
+os.makedirs(VIDEO_OUTPUT_DIR, exist_ok=True)
+
+FRAME_COUNT = 81
+CAM_VALUE = 0.1
+
+# Action Mapping
+KEY_TO_INDEX = {
+    'W': 0, 'S': 1, 'A': 2, 'D': 3,
+}
+
+VIEW_ACTION_TO_MOUSE = {
+    "stop": [0.0, 0.0],
+    "up": [CAM_VALUE, 0.0],
+    "down": [-CAM_VALUE, 0.0],
+    "left": [0.0, -CAM_VALUE],
+    "right": [0.0, CAM_VALUE],
+    "up_right": [CAM_VALUE, CAM_VALUE],
+    "up_left": [CAM_VALUE, -CAM_VALUE],
+    "down_right": [-CAM_VALUE, CAM_VALUE],
+    "down_left": [-CAM_VALUE, -CAM_VALUE],
+}
+
+def get_multihot_vector(keys_str):
+    """Convert string like 'WA' to [1, 0, 1, 0, 0, 0]"""
+    vector = [0.0] * 6
+    if not keys_str:
+        return vector
+    for char in keys_str.upper():
+        if char in KEY_TO_INDEX:
+            vector[KEY_TO_INDEX[char]] = 1.0
+    return vector
+
+def get_mouse_vector(view_str):
+    """Convert view string to [x, y]"""
+    return VIEW_ACTION_TO_MOUSE.get(view_str.lower(), [0.0, 0.0])
+
+def generate_sequence(key_seq, mouse_seq):
+    """
+    Generates action arrays based on sequences.
+    """
+    keyboard_arr = np.zeros((FRAME_COUNT, 6), dtype=np.float32)
+    mouse_arr = np.zeros((FRAME_COUNT, 2), dtype=np.float32)
+    
+    mid_point = FRAME_COUNT // 2
+
+    # First Half
+    k_vec1 = get_multihot_vector(key_seq[0])
+    m_vec1 = get_mouse_vector(mouse_seq[0])
+    keyboard_arr[:mid_point] = k_vec1
+    mouse_arr[:mid_point] = m_vec1
+
+    # Second Half
+    k_vec2 = get_multihot_vector(key_seq[1])
+    m_vec2 = get_mouse_vector(mouse_seq[1])
+    keyboard_arr[mid_point:] = k_vec2
+    mouse_arr[mid_point:] = m_vec2
+
+    return keyboard_arr, mouse_arr
+
+def save_action(index, keyboard_arr, mouse_arr):
+    filename = f"{index:06d}_action.npy"
+    filepath = os.path.join(VIDEO_OUTPUT_DIR, filename)
+    
+    action_dict = {
+        'keyboard': keyboard_arr,
+        'mouse': mouse_arr
+    }
+    np.save(filepath, action_dict)
+    return filename
+
+def generate_description(key_seq, mouse_seq):
+    """Generates a human-readable string for the combination."""
+    k1, k2 = key_seq
+    m1, m2 = mouse_seq
+    
+    # Format Keyboard Description
+    if not k1 and not k2:
+        k_desc = "No Key"
+    elif k1 == k2:
+        k_desc = f"Hold [{k1}]"
+    else:
+        k_desc = f"Switch [{k1}]->[{k2}]"
+        
+    # Format Mouse Description
+    if m1 == "stop" and m2 == "stop":
+        m_desc = "Static"
+    elif m1 == m2:
+        m_desc = f"Hold [{m1}]"
+    else:
+        m_desc = f"Switch [{m1}]->[{m2}]"
+        
+    return f"{k_desc} + {m_desc}"
+
+# ==========================================
+# Main Generation Logic
+# ==========================================
+
+configs = []
+readme_content = []
+
+# Group 1: Constant Keyboard, No Mouse (0-7)
+keys_basic = ['W', 'S', 'A', 'D', 'WA', 'WD', 'SA', 'SD']
+for k in keys_basic:
+    configs.append(((k, k), ("stop", "stop")))
+
+# Group 2: No Keyboard, Constant Mouse (8-15)
+mouse_basic = ['up', 'down', 'left', 'right', 'up_right', 'up_left', 'down_right', 'down_left']
+for m in mouse_basic:
+    configs.append((("", ""), (m, m)))
+
+# Group 3: Split Keyboard, No Mouse (16-23)
+split_keys = [
+    ('W', 'S'), ('S', 'W'), 
+    ('A', 'D'), ('D', 'A'),
+    ('W', 'A'), ('W', 'D'), 
+    ('S', 'A'), ('S', 'D')
+]
+for k1, k2 in split_keys:
+    configs.append(((k1, k2), ("stop", "stop")))
+
+# Group 4: No Keyboard, Split Mouse (24-31)
+split_mouse = [
+    ('left', 'right'), ('right', 'left'),
+    ('up', 'down'), ('down', 'up'),
+    ('up_left', 'up_right'), ('up_right', 'up_left'),
+    ('left', 'up'), ('right', 'down')
+]
+for m1, m2 in split_mouse:
+    configs.append((("", ""), (m1, m2)))
+
+# Group 5: Constant Keyboard + Constant Mouse (32-47)
+combo_keys = ['W', 'S', 'W', 'S', 'A', 'D', 'WA', 'WD', 'W', 'S', 'W', 'S', 'A', 'D', 'WA', 'WD']
+combo_mice = ['left', 'left', 'right', 'right', 'up', 'up', 'down', 'down', 'up_left', 'up_left', 'up_right', 'up_right', 'down_left', 'down_right', 'right', 'left']
+for i in range(16):
+    configs.append(((combo_keys[i], combo_keys[i]), (combo_mice[i], combo_mice[i])))
+
+# Group 6: Constant Keyboard, Split Mouse (48-55)
+complex_1_keys = ['W'] * 8
+complex_1_mice = [
+    ('left', 'right'), ('right', 'left'), 
+    ('up', 'down'), ('down', 'up'),
+    ('left', 'up'), ('right', 'up'),
+    ('left', 'down'), ('right', 'down')
+]
+for i in range(8):
+    configs.append(((complex_1_keys[i], complex_1_keys[i]), complex_1_mice[i]))
+
+# Group 7: Split Keyboard, Constant Mouse (56-63)
+complex_2_keys = [
+    ('W', 'S'), ('S', 'W'),
+    ('A', 'D'), ('D', 'A'),
+    ('W', 'A'), ('W', 'D'),
+    ('S', 'A'), ('S', 'D')
+]
+complex_2_mouse = 'up' 
+for k1, k2 in complex_2_keys:
+    configs.append(((k1, k2), (complex_2_mouse, complex_2_mouse)))
+
+
+# Execution
+print(f"Preparing to generate {len(configs)} action files...")
+
+for i, (key_seq, mouse_seq) in enumerate(configs):
+    if i >= 16: break 
+    
+    # Generate Data
+    kb_arr, ms_arr = generate_sequence(key_seq, mouse_seq)
+    filename = save_action(i, kb_arr, ms_arr)
+    
+    # Generate Description for README
+    description = generate_description(key_seq, mouse_seq)
+    readme_entry = f"{i:02d}. {description}"
+    readme_content.append(readme_entry)
+    
+    print(f"Generated {filename} -> {description}")
+
+# Write README
+readme_path = os.path.join(BASE_OUTPUT_DIR, 'README.md')
+with open(readme_path, 'w', encoding='utf-8') as f:
+    f.write(f"Total Files: {len(readme_content)}\n\n")
+    for line in readme_content:
+        f.write(line + '\n')
+
+print(f"\nProcessing complete.")
+print(f"64 .npy files generated in {VIDEO_OUTPUT_DIR}")
+print(f"Manifest saved to {readme_path}")
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
+DATA_MERGE_PATH="mc_wasd_10/merge.txt"
+OUTPUT_DIR="mc_wasd_10/preprocessed/"
+
+# export CUDA_VISIBLE_DEVICES=0
+export MASTER_ADDR=localhost
+export MASTER_PORT=29500
+export RANK=0
+export WORLD_SIZE=1
+
+python fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 10 \
+    --seed 42 \
+    --max_height 352 \
+    --max_width 640 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --samples_per_file 10 \
+    --train_fps 25 \
+    --flush_frequency 10 \
+    --preprocess_task wangame
@@ -0,0 +1,404 @@
+{
+    "data": [
+        {
+            "caption": "0",
+            "image_path": "../../../../mc_wasd_10/validate/000000.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000000_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "1",
+            "image_path": "../../../../mc_wasd_10/validate/000001.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000001_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "2",
+            "image_path": "../../../../mc_wasd_10/validate/000002.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000002_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "3",
+            "image_path": "../../../../mc_wasd_10/validate/000003.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000003_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "4",
+            "image_path": "../../../../mc_wasd_10/validate/000004.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000004_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "5",
+            "image_path": "../../../../mc_wasd_10/validate/000005.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000005_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "6",
+            "image_path": "../../../../mc_wasd_10/validate/000006.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000006_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "7",
+            "image_path": "../../../../mc_wasd_10/validate/000007.jpg",
+            "action_path": "../../../../mc_wasd_10/videos/000007_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "00. Hold [W] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000000.jpg",
+            "action_path": "action/000000_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "01. Hold [S] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000001.jpg",
+            "action_path": "action/000001_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "02. Hold [A] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000002.jpg",
+            "action_path": "action/000002_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "03. Hold [D] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000003.jpg",
+            "action_path": "action/000003_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "04. Hold [WA] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000004.jpg",
+            "action_path": "action/000004_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "05. Hold [WD] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000005.jpg",
+            "action_path": "action/000005_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "06. Hold [SA] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000006.jpg",
+            "action_path": "action/000006_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "07. Hold [SD] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/000007.jpg",
+            "action_path": "action/000007_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "08. No Key + Hold [up]",
+            "image_path": "../../../../mc_wasd_10/validate/000000.jpg",
+            "action_path": "action/000008_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "09. No Key + Hold [down]",
+            "image_path": "../../../../mc_wasd_10/validate/000001.jpg",
+            "action_path": "action/000009_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "10. No Key + Hold [left]",
+            "image_path": "../../../../mc_wasd_10/validate/000002.jpg",
+            "action_path": "action/000010_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "11. No Key + Hold [right]",
+            "image_path": "../../../../mc_wasd_10/validate/000003.jpg",
+            "action_path": "action/000011_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "12. No Key + Hold [up_right]",
+            "image_path": "../../../../mc_wasd_10/validate/000004.jpg",
+            "action_path": "action/000012_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "13. No Key + Hold [up_left]",
+            "image_path": "../../../../mc_wasd_10/validate/000005.jpg",
+            "action_path": "action/000013_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "14. No Key + Hold [down_right]",
+            "image_path": "../../../../mc_wasd_10/validate/000006.jpg",
+            "action_path": "action/000014_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "15. No Key + Hold [down_left]",
+            "image_path": "../../../../mc_wasd_10/validate/000007.jpg",
+            "action_path": "action/000015_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "00. Hold [W] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000000.jpg",
+            "action_path": "action/000000_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "01. Hold [S] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000001.jpg",
+            "action_path": "action/000001_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "02. Hold [A] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000002.jpg",
+            "action_path": "action/000002_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "03. Hold [D] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000003.jpg",
+            "action_path": "action/000003_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "04. Hold [WA] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000004.jpg",
+            "action_path": "action/000004_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "05. Hold [WD] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000005.jpg",
+            "action_path": "action/000005_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "06. Hold [SA] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000006.jpg",
+            "action_path": "action/000006_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "07. Hold [SD] + Static",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000007.jpg",
+            "action_path": "action/000007_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "08. No Key + Hold [up]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000000.jpg",
+            "action_path": "action/000008_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "09. No Key + Hold [down]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000001.jpg",
+            "action_path": "action/000009_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "10. No Key + Hold [left]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000002.jpg",
+            "action_path": "action/000010_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "11. No Key + Hold [right]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000003.jpg",
+            "action_path": "action/000011_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "12. No Key + Hold [up_right]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000004.jpg",
+            "action_path": "action/000012_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "13. No Key + Hold [up_left]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000005.jpg",
+            "action_path": "action/000013_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "14. No Key + Hold [down_right]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000006.jpg",
+            "action_path": "action/000014_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        },
+        {
+            "caption": "15. No Key + Hold [down_left]",
+            "image_path": "../../../../mc_wasd_10/validate/gen_000007.jpg",
+            "action_path": "action/000015_action.npy",
+            "video_path": null,
+            "num_inference_steps": 40,
+            "height": 352,
+            "width": 640,
+            "num_frames": 77
+        }
+    ]
+}
@@ -1,23 +0,0 @@
-# LTX-2 Crush-Smol Example
-# TODO: Update this doc.
-
-These are e2e example scripts for finetuning LTX-2 on the crush-smol dataset.
-
-## Execute the following commands from `FastVideo/` to run training:
-
-### Download crush-smol dataset:
-
-`bash examples/training/finetune/ltx2/overfit/download_dataset.sh`
-
-### Preprocess the videos and captions into latents:
-
-`bash examples/training/finetune/ltx2/overfit/preprocess_ltx2_data_t2v_new.sh`
-
-### Edit the following file and run finetuning:
-
-`bash examples/training/finetune/ltx2/overfit/finetune_t2v.sh`
-
-Notes:
- Update `DATASET_PATH` in the preprocess script to point to your merged dataset root (`videos/` + `videos2caption.json`).
- `MODEL_PATH` should point to a local LTX-2 diffusers-style directory that contains `model_index.json` and `text_encoder/gemma`.
-
@@ -1,3 +0,0 @@
-# #!/bin/bash
-#
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -1,95 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-
-MODEL_PATH="Davids048/LTX2-Base-Diffusers"
-# Also can use simple 1 video for overfitting experiments.
-# DATA_DIR="/home/hal-jundas/codes/FastVideo/data/crush-smol"
-DATA_DIR="<PATH_TO_PROCESSED_DATASET>"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-echo  VALIDATION_DATASET_FILE: $VALIDATION_DATASET_FILE
-NUM_GPUS=4
-OVERFIT_HEIGHT=480
-OVERFIT_WIDTH=832
-OVERFIT_FRAMES=73
-
-training_args=(
-  --tracker_project_name "ltx2_t2v_finetune"
-  --output_dir "checkpoints/ltx2_t2v_finetune"
-  --max_train_steps 5000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 10
-  --num_height $OVERFIT_HEIGHT
-  --num_width $OVERFIT_WIDTH
-  --num_frames $OVERFIT_FRAMES
-  --ltx2-first-frame-conditioning-p 0.1
-  --enable_gradient_checkpointing_type "full"
-  --mode "finetuning"
-)
-
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size $NUM_GPUS
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-dataset_args=(
-  --data_path $DATA_DIR
-  --dataloader_num_workers 1
-)
-
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 50
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "3.0"
-)
-
-optimizer_args=(
-  --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
-  --lr_scheduler "linear"
-)
-
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --dit_precision "fp32"
-  --dit_cpu_offload False
-  --dit_layerwise_offload False
-  --text_encoder_cpu_offload False
-  --image_encoder_cpu_offload False
-  --vae_cpu_offload False
-)
-
-# NOTE: Setting this environment variable to TORCH_SDPA to avoid the issue of stacking that failed in flash attn. 
-export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/ltx2_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,80 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-
-MODEL_PATH="/path/to/LTX-2"
-DATA_DIR="data/crush-smol"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-
-training_args=(
-  --tracker_project_name "ltx2_t2v_lora_finetune"
-  --output_dir "checkpoints/ltx2_t2v_lora_finetune"
-  --max_train_steps 2000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 8
-  --num_latent_t 10
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --ltx2-first-frame-conditioning-p 0.1
-  --enable_gradient_checkpointing_type "full"
-)
-
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size $NUM_GPUS
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-dataset_args=(
-  --data_path $DATA_DIR
-  --dataloader_num_workers 1
-)
-
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "3.0"
-)
-
-optimizer_args=(
-  --learning_rate 2e-4
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-  --lora_training True
-  --lora_rank 16
-  --lora_alpha 16
-)
-
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/ltx2_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,35 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1
-MODEL_PATH="Davids048/LTX2-Base-Diffusers"
-# DATASET_PATH="data/overfit"
-DATASET_PATH="data/crush-smol"
-OUTPUT_DIR="$DATASET_PATH"
-WITH_AUDIO=true
-
-# Convert one-file overfit metadata into merged format if needed.
-if [ ! -f "$DATASET_PATH/videos2caption.json" ] && [ -f "$DATASET_PATH/overfit.json" ]; then
-  python scripts/dataset_preparation/convert_to_merged_dataset.py \
-    --items-json "$DATASET_PATH/overfit.json" \
-    --output-dir "$DATASET_PATH"
-fi
-
-
-torchrun --nproc_per_node=$GPU_NUM \
-    --master_port=29513 \
-    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
-    --model_path $MODEL_PATH \
-    --mode preprocess \
-    --workload_type t2v \
-    --preprocess.video_loader_type torchvision \
-    --preprocess.dataset_type merged \
-    --preprocess.dataset_path $DATASET_PATH \
-    --preprocess.dataset_output_dir $OUTPUT_DIR \
-    --preprocess.with_audio $WITH_AUDIO \
-    --preprocess.preprocess_video_batch_size 1 \
-    --preprocess.dataloader_num_workers 0 \
-    --preprocess.max_height 480 \
-    --preprocess.max_width 832 \
-    --preprocess.num_frames 73 \
-    --preprocess.train_fps 16 \
-    --preprocess.video_length_tolerance_range 5
@@ -1,13 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "The camera opens in a calm, sunlit frog yoga studio. Warm morning light washes over the wooden floor as incense smoke drifts lazily in the air. The senior frog instructor sits cross-legged at the center, eyes closed, voice deep and calm. “We are one with the pond.” All the frogs answer softly: “Ommm...” “We are one with the mud.” “Ommm...” He smiles faintly. “We are one with the flies.” A quiet pause. The camera slowly pans to the side — one frog twitches, eyes darting. Suddenly — *thwip!* — its tongue snaps out, catching a fly mid-air and pulling it into its mouth. The master exhales slowly, still serene.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 1088,
-      "width": 1920,
-      "num_frames": 121
-    }
-  ]
-}
@@ -1,31 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    },
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
-      "image_path": null,
-      "video_path": null,
-      "num_inference_steps": 50,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}
@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"

 [project]
 name = "fastvideo-kernel"
-version = "0.2.6"
+version = "0.2.5"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -287,8 +287,9 @@ def block_sparse_attn(
    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: supports q_seq_len != kv_seq_len as long as both are padded
-    # to a multiple of the block size (64 tokens).
+    # 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)


@@ -141,6 +141,12 @@ def video_sparse_attn(
        # 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:
+        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)

@@ -29,7 +29,7 @@ configs = [


 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
-@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"])
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
@triton.jit
 def _attn_fwd_sparse(
        Q,
@@ -60,8 +60,7 @@ def _attn_fwd_sparse(
        stride_on,
        Z,
        H,
-        N_CTX_Q,  #
-        N_CTX_KV,  #
+        N_CTX,  #
        HEAD_DIM: tl.constexpr,  #
        BLOCK_M: tl.constexpr,
        BLOCK_N: tl.constexpr,
@@ -76,29 +75,24 @@ def _attn_fwd_sparse(
    off_hz = tl.program_id(1)  # fused (batch, head)
    b = off_hz // H
    h = off_hz % H
-    q_tiles = N_CTX_Q // BLOCK_M
+    q_tiles = N_CTX // BLOCK_M
    meta_base = ((b * H + h) * q_tiles + q_blk)

    kv_blocks = tl.load(q2k_num + meta_base)  # int32
    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list

    # ----- base pointers -----
-    # Note: when q and kv have different sequence lengths, their per-(batch,head)
-    # strides differ, so we must compute separate base offsets.
-    q_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
-    k_off = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
-    v_off = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
-    o_off = (b.to(tl.int64) * stride_oz + h.to(tl.int64) * stride_oh)
+    qvk_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)

-    Q_ptr = tl.make_block_ptr(base=Q + q_off,
-                              shape=(N_CTX_Q, HEAD_DIM),
+    Q_ptr = tl.make_block_ptr(base=Q + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
                              strides=(stride_qm, stride_qk),
                              offsets=(q_blk * BLOCK_M, 0),
                              block_shape=(BLOCK_M, HEAD_DIM),
                              order=(1, 0))

-    K_base = tl.make_block_ptr(base=K + k_off,
-                               shape=(HEAD_DIM, N_CTX_KV),
+    K_base = tl.make_block_ptr(base=K + qvk_off,
+                               shape=(HEAD_DIM, N_CTX),
                               strides=(stride_kk, stride_kn),
                               offsets=(0, 0),
                               block_shape=(HEAD_DIM, BLOCK_N),
@@ -106,15 +100,15 @@ def _attn_fwd_sparse(

    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
                                                                              0)
-    V_base = tl.make_block_ptr(base=V + v_off,
-                               shape=(N_CTX_KV, HEAD_DIM),
+    V_base = tl.make_block_ptr(base=V + qvk_off,
+                               shape=(N_CTX, HEAD_DIM),
                               strides=(stride_vk, stride_vn),
                               offsets=(0, 0),
                               block_shape=(BLOCK_N, HEAD_DIM),
                               order=v_order)

-    O_ptr = tl.make_block_ptr(base=Out + o_off,
-                              shape=(N_CTX_Q, HEAD_DIM),
+    O_ptr = tl.make_block_ptr(base=Out + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
                              strides=(stride_om, stride_on),
                              offsets=(q_blk * BLOCK_M, 0),
                              block_shape=(BLOCK_M, HEAD_DIM),
@@ -156,7 +150,7 @@ def _attn_fwd_sparse(
    # ----- epilogue -----
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]
-    tl.store(M + off_hz * N_CTX_Q + offs_m, m_i)
+    tl.store(M + off_hz * N_CTX + offs_m, m_i)
    tl.store(O_ptr, acc.to(Out.type.element_ty))


@@ -207,7 +201,7 @@ def _attn_bwd_dkdv(
        stride_tok,
        stride_d,  #
        H,
-        N_CTX_KV,
+        N_CTX,
        BLOCK_M1: tl.constexpr,  #
        BLOCK_N1: tl.constexpr,  #
        HEAD_DIM: tl.constexpr,  #
@@ -227,8 +221,8 @@ def _attn_bwd_dkdv(
    off_hz = tl.program_id(2)  # fused (batch, head)
    b = off_hz // H
    h = off_hz % H
-    kv_tiles = N_CTX_KV // BLOCK_N1
-    meta_base = ((b * H + h) * kv_tiles + kv_blk)
+    q_tiles = N_CTX // BLOCK_N1
+    meta_base = ((b * H + h) * q_tiles + kv_blk)

    q_blocks = tl.load(k2q_num + meta_base)  # int32
    q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
@@ -308,21 +302,16 @@ def _attn_bwd_dq(

    kv_blocks = tl.load(q2k_num + meta_base)  # int32
    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    block_size = tl.load(variable_block_sizes + q_blk)

    for blk_idx in range(kv_blocks * 2):
-        kv_idx = tl.load(kv_ptr + blk_idx // 2).to(tl.int32)
-        # variable_block_sizes is defined per KV block (tile). Mask must therefore
-        # use kv_idx (not q_blk). Also, because we split each 64-token block into
-        # two 32-token halves, the mask must account for the half-block offset.
-        block_size = tl.load(variable_block_sizes + kv_idx).to(tl.int32)
-        half = (blk_idx % 2).to(tl.int32)
-        block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx // 2).to(tl.int32) * 2 +
+                               blk_idx % 2) * step_n * stride_tok
        kT = tl.load(kT_ptrs + block_sparse_offset)
        vT = tl.load(vT_ptrs + block_sparse_offset)
        qk = tl.dot(q, kT)
        p = tl.math.exp2(qk - m)
-        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
-        mask = offs_in_block < block_size
+        mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
        p = tl.where(mask[None, :], p, 0.0)
        # Compute dP and dS.
        dp = tl.dot(do, vT).to(tl.float32)
@@ -478,235 +467,19 @@ def _attn_bwd(
    tl.store(dq_ptrs, dq)


-@triton.jit
-def _attn_bwd_dkdv_kernel(
-        Q,
-        K,
-        V,
-        sm_scale,  #
-        DO,  #
-        DK,
-        DV,  #
-        M,
-        D,
-        k2q_index,
-        k2q_num,
-        max_q_blks,
-        variable_block_sizes,
-        # shared token/dim strides (assumed contiguous along token and dim)
-        stride_tok,
-        stride_d,  #
-        # batch/head strides (may differ between Q and KV)
-        stride_qz,
-        stride_qh,
-        stride_kz,
-        stride_kh,
-        stride_vz,
-        stride_vh,
-        stride_doz,
-        stride_doh,
-        stride_dkz,
-        stride_dkh,
-        stride_dvz,
-        stride_dvh,
-        H,
-        N_CTX_Q,
-        N_CTX_KV,
-        BLOCK_M1: tl.constexpr,  #
-        BLOCK_N1: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr):
-    """
-    Backward kernel that computes dK and dV for each KV block (64 tokens).
-    Grid:
-      pid0: kv_blk in [0, N_CTX_KV/BLOCK_N1)
-      pid2: fused (batch, head) in [0, B*H)
-    """
-    bhid = tl.program_id(2)
-    b = bhid // H
-    h = bhid % H
-    kv_blk = tl.program_id(0)
-
-    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
-    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
-    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
-    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
-    dk_adj = (b.to(tl.int64) * stride_dkz + h.to(tl.int64) * stride_dkh)
-    dv_adj = (b.to(tl.int64) * stride_dvz + h.to(tl.int64) * stride_dvh)
-
-    Q = Q + q_adj
-    K = K + kv_adj_k
-    V = V + kv_adj_v
-    DO = DO + do_adj
-    DK = DK + dk_adj
-    DV = DV + dv_adj
-
-    # M and D (delta) are always sized by Q length.
-    M = M + (bhid * N_CTX_Q).to(tl.int64)
-    D = D + (bhid * N_CTX_Q).to(tl.int64)
-
-    offs_k = tl.arange(0, HEAD_DIM)
-    start_n = kv_blk * BLOCK_N1
-    offs_n = start_n + tl.arange(0, BLOCK_N1)
-
-    # load K and V: they stay in SRAM throughout the inner loop.
-    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
-
-    dv_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
-    dk_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
-
-    num_steps = N_CTX_Q // BLOCK_M1
-    dk_acc, dv_acc = _attn_bwd_dkdv(
-        dk_acc,
-        dv_acc,
-        Q,
-        k,
-        v,
-        sm_scale,
-        DO,
-        M,
-        D,
-        k2q_index,
-        k2q_num,
-        max_q_blks,
-        variable_block_sizes,
-        stride_tok,
-        stride_d,
-        H,
-        N_CTX_KV,
-        BLOCK_M1=BLOCK_M1,
-        BLOCK_N1=BLOCK_N1,
-        HEAD_DIM=HEAD_DIM,
-        start_n=start_n,
-        start_m=0,
-        num_steps=num_steps,
-    )
-
-    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
-    tl.store(dv_ptrs, dv_acc)
-
-    dk_acc *= sm_scale
-    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
-    tl.store(dk_ptrs, dk_acc)
-
-
-@triton.jit
-def _attn_bwd_dq_kernel(
-        Q,
-        K,
-        V,
-        DO,  #
-        DQ,
-        M,
-        D,
-        q2k_index,
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        # shared token/dim strides (assumed contiguous along token and dim)
-        stride_tok,
-        stride_d,  #
-        # batch/head strides (may differ between Q and KV)
-        stride_qz,
-        stride_qh,
-        stride_kz,
-        stride_kh,
-        stride_vz,
-        stride_vh,
-        stride_doz,
-        stride_doh,
-        stride_dqz,
-        stride_dqh,
-        H,
-        N_CTX_Q,
-        BLOCK_M2: tl.constexpr,  #
-        BLOCK_N2: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr):
-    """
-    Backward kernel that computes dQ for each Q block (64 tokens).
-    Grid:
-      pid0: q_blk in [0, N_CTX_Q/BLOCK_M2)
-      pid2: fused (batch, head) in [0, B*H)
-    """
-    LN2 = 0.6931471824645996  # = ln(2)
-    bhid = tl.program_id(2)
-    b = bhid // H
-    h = bhid % H
-    q_blk = tl.program_id(0)
-
-    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
-    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
-    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
-    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
-    dq_adj = (b.to(tl.int64) * stride_dqz + h.to(tl.int64) * stride_dqh)
-
-    Q = Q + q_adj
-    K = K + kv_adj_k
-    V = V + kv_adj_v
-    DO = DO + do_adj
-    DQ = DQ + dq_adj
-
-    M = M + (bhid * N_CTX_Q).to(tl.int64)
-    D = D + (bhid * N_CTX_Q).to(tl.int64)
-
-    offs_k = tl.arange(0, HEAD_DIM)
-    start_m = q_blk * BLOCK_M2
-    offs_m = start_m + tl.arange(0, BLOCK_M2)
-
-    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    m = tl.load(M + offs_m)[:, None]
-
-    dq_acc = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
-    num_steps = 0  # unused in _attn_bwd_dq
-    dq_acc = _attn_bwd_dq(
-        dq_acc,
-        q,
-        K,
-        V,
-        do,
-        m,
-        D,
-        q2k_index,
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        stride_tok,
-        stride_d,
-        H,
-        N_CTX_Q,
-        BLOCK_M2=BLOCK_M2,
-        BLOCK_N2=BLOCK_N2,
-        HEAD_DIM=HEAD_DIM,
-        start_m=start_m,
-        start_n=0,
-        num_steps=num_steps,
-    )
-
-    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= LN2
-    tl.store(dq_ptrs, dq_acc)
-
-
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
 def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
                                     variable_block_sizes):
-    B, H, Tq, D = q.shape
-    Tkv = k.shape[2]
+    B, H, T, D = q.shape
    sm_scale = 1.0 / math.sqrt(D)
    max_kv_blks = q2k_index.shape[-1]
-    assert Tq % 64 == 0, f"q length must be a multiple of 64, but got {Tq}"
-    assert Tkv % 64 == 0, f"kv length must be a multiple of 64, but got {Tkv}"
+    assert T % 64 == 0, f"T must be a multiple of 64, but got {T}"
    assert q2k_num.shape[
-        -1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
-    assert variable_block_sizes.numel() == Tkv // 64, (
-        f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, "
-        f"got {variable_block_sizes.numel()}"
-    )
+        -1] == T // 64, f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
    o = torch.empty_like(q)
-    M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
+    M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)

-    grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1)
+    grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
    _attn_fwd_sparse[grid](q,
                           k,
                           v,
@@ -735,8 +508,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
                           o.stride(3),
                           B,
                           H,
-                           Tq,
-                           Tkv,
+                           T,
                           HEAD_DIM=D,
                           STAGE=3)

@@ -746,21 +518,21 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
 def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
                                      k2q_index, k2q_num, variable_block_sizes):
    assert do.is_contiguous()
+    assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()

-    B, H, Tq, D = q.shape
-    Tkv = k.shape[2]
+    B, H, T, D = q.shape
    sm_scale = 1.0 / math.sqrt(D)
    dq = torch.empty_like(q)
    dk = torch.empty_like(k)
    dv = torch.empty_like(v)
-    BATCH, N_HEAD = q.shape[:2]
+    BATCH, N_HEAD, N_CTX = q.shape[:3]
    BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
    arg_k = k
    arg_k = arg_k * (sm_scale * RCP_LN2)
    PRE_BLOCK = 64
-    assert Tq % PRE_BLOCK == 0
-    pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
+    assert N_CTX % PRE_BLOCK == 0
+    pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
    delta = torch.empty_like(M)
    _attn_bwd_preprocess[pre_grid](
        o,
@@ -768,7 +540,7 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
        delta,  #
        BATCH,
        N_HEAD,
-        Tq,  #
+        N_CTX,  #
        BLOCK_M=PRE_BLOCK,
        HEAD_DIM=D  #
    )
@@ -776,75 +548,36 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
    max_q_blks = k2q_index.shape[-1]
    max_kv_blks = q2k_index.shape[-1]

-    # dK/dV kernel: grid over KV blocks
-    grid_kv = (Tkv // BLOCK_N1, 1, BATCH * N_HEAD)
-    _attn_bwd_dkdv_kernel[grid_kv](
+    grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+    _attn_bwd[grid](
        q,
        arg_k,
        v,
        sm_scale,
        do,
+        dq,
        dk,
-        dv,
+        dv,  #
        M,
-        delta,
+        delta,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
        k2q_index,
        k2q_num,
        max_q_blks,
        variable_block_sizes,
-        q.stride(2),
-        q.stride(3),
        q.stride(0),
        q.stride(1),
-        arg_k.stride(0),
-        arg_k.stride(1),
-        v.stride(0),
-        v.stride(1),
-        do.stride(0),
-        do.stride(1),
-        dk.stride(0),
-        dk.stride(1),
-        dv.stride(0),
-        dv.stride(1),
+        q.stride(2),
+        q.stride(3),  #
        N_HEAD,
-        Tq,
-        Tkv,
+        N_CTX,  #
        BLOCK_M1=BLOCK_M1,
-        BLOCK_N1=BLOCK_N1,
-        HEAD_DIM=D,
-    )
-
-    # dQ kernel: grid over Q blocks
-    grid_q = (Tq // BLOCK_M2, 1, BATCH * N_HEAD)
-    _attn_bwd_dq_kernel[grid_q](
-        q,
-        arg_k,
-        v,
-        do,
-        dq,
-        M,
-        delta,
-        q2k_index,
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        q.stride(2),
-        q.stride(3),
-        q.stride(0),
-        q.stride(1),
-        arg_k.stride(0),
-        arg_k.stride(1),
-        v.stride(0),
-        v.stride(1),
-        do.stride(0),
-        do.stride(1),
-        dq.stride(0),
-        dq.stride(1),
-        N_HEAD,
-        Tq,
+        BLOCK_N1=BLOCK_N1,  #
        BLOCK_M2=BLOCK_M2,
-        BLOCK_N2=BLOCK_N2,
-        HEAD_DIM=D,
+        BLOCK_N2=BLOCK_N2,  #
+        HEAD_DIM=D  #
    )

    return dq, dk, dv
@@ -1 +1 @@
-__version__ = "0.2.6"
+__version__ = "0.2.5"
@@ -86,7 +86,6 @@ class PreprocessConfig:

    # Model configuration
    training_cfg_rate: float = 0.0
-    with_audio: bool = False

    # framework configuration
    seed: int = 42
@@ -191,10 +190,6 @@ class PreprocessConfig:
                                     type=float,
                                     default=PreprocessConfig.training_cfg_rate,
                                     help="Training CFG rate")
-        preprocess_args.add_argument(f"--{prefix_with_dot}with-audio",
-                                     action=StoreBoolean,
-                                     default=PreprocessConfig.with_audio,
-                                     help="Whether to extract and encode audio")
        preprocess_args.add_argument(f"--{prefix_with_dot}seed",
                                     type=int,
                                     default=PreprocessConfig.seed,
@@ -5,7 +5,6 @@ from fastvideo.configs.models.vaes.base import VAEConfig
 from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig,
                                            LTX2AudioEncoderConfig,
                                            LTX2VocoderConfig)
-from fastvideo.configs.models.upsamplers.base import UpsamplerConfig

 __all__ = [
    "ModelConfig",
@@ -15,5 +14,4 @@ __all__ = [
    "LTX2AudioEncoderConfig",
    "LTX2AudioDecoderConfig",
    "LTX2VocoderConfig",
-    "UpsamplerConfig",
 ]
@@ -1,6 +1,5 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
-from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
 from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
@@ -10,8 +9,7 @@ from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig

 __all__ = [
-    "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig",
-    "WanVideoConfig", "StepVideoConfig", "CosmosVideoConfig",
-    "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
-    "HYWorldConfig"
+    "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
+    "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
+    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig"
 ]
@@ -1,163 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Configuration for HunyuanGameCraft transformer model.
-
-HunyuanGameCraft extends HunyuanVideo with:
-1. CameraNet for camera/action conditioning
-2. 33 input channels (16 latent + 16 gt_latent + 1 mask)
-3. Mask-based conditioning for autoregressive generation
-"""
-from dataclasses import dataclass, field
-
-import torch
-
-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"
-
-
-def is_camera_net(n: str, m) -> bool:
-    return "camera_net" in n
-
-
-@dataclass
-class HunyuanGameCraftArchConfig(DiTArchConfig):
-    """Architecture config for HunyuanGameCraft transformer."""
-
-    # Version field for compatibility with saved config.json
-    _fastvideo_version: str = "0.1.0"
-
-    # Camera net flag (for config.json compatibility)
-    camera_net: bool = True
-
-    _fsdp_shard_conditions: list = field(
-        default_factory=lambda:
-        [is_double_block, is_single_block, is_refiner_block, is_camera_net])
-
-    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
-
-    # Parameter names mapping from official checkpoint to FastVideo naming
-    # GameCraft weights are already close to FastVideo format with minor adjustments
-    param_names_mapping: dict = field(
-        default_factory=lambda: {
-            # MLP naming: fc1 -> fc_in, fc2 -> fc_out
-            r"^(.*)\.img_mlp\.fc1\.(.*)$":
-            r"\1.img_mlp.fc_in.\2",
-            r"^(.*)\.img_mlp\.fc2\.(.*)$":
-            r"\1.img_mlp.fc_out.\2",
-            r"^(.*)\.txt_mlp\.fc1\.(.*)$":
-            r"\1.txt_mlp.fc_in.\2",
-            r"^(.*)\.txt_mlp\.fc2\.(.*)$":
-            r"\1.txt_mlp.fc_out.\2",
-
-            # Single block MLP naming
-            r"^single_blocks\.(\d+)\.mlp\.fc1\.(.*)$":
-            r"single_blocks.\1.mlp.fc_in.\2",
-            r"^single_blocks\.(\d+)\.mlp\.fc2\.(.*)$":
-            r"single_blocks.\1.mlp.fc_out.\2",
-
-            # Token refiner naming
-            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.(.*)$":
-            r"txt_in.refiner_blocks.\1.\2",
-
-            # Vector in naming
-            r"^vector_in\.in_layer\.(.*)$":
-            r"vector_in.fc_in.\1",
-            r"^vector_in\.out_layer\.(.*)$":
-            r"vector_in.fc_out.\1",
-
-            # Time embedder naming
-            r"^time_in\.mlp\.0\.(.*)$":
-            r"time_in.mlp.fc_in.\1",
-            r"^time_in\.mlp\.2\.(.*)$":
-            r"time_in.mlp.fc_out.\1",
-
-            # Guidance embedder naming (if present)
-            r"^guidance_in\.mlp\.0\.(.*)$":
-            r"guidance_in.mlp.fc_in.\1",
-            r"^guidance_in\.mlp\.2\.(.*)$":
-            r"guidance_in.mlp.fc_out.\1",
-
-            # Final layer adaLN modulation
-            r"^final_layer\.adaLN_modulation\.1\.(.*)$":
-            r"final_layer.adaLN_modulation.linear.\1",
-
-            # Refiner block MLP naming
-            r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc1\.(.*)$":
-            r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
-            r"^txt_in\.refiner_blocks\.(\d+)\.mlp\.fc2\.(.*)$":
-            r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
-
-            # Camera net weights are already correctly named
-        })
-
-    # Reverse mapping for saving checkpoints
-    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
-
-    # Model architecture parameters
-    # patch_size can be int or tuple - if tuple, it's [T, H, W]
-    patch_size: int | tuple[int, int, int] = 2
-    patch_size_t: int = 1
-    in_channels: int = 33  # 16 latent + 16 gt_latent + 1 mask
-    out_channels: int = 16
-    num_attention_heads: int = 24
-    attention_head_dim: int = 128
-    mlp_ratio: float = 4.0
-    num_layers: int = 20  # Double stream blocks
-    num_single_layers: int = 40  # Single stream blocks
-    num_refiner_layers: int = 2
-    rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
-    guidance_embeds: bool = False  # GameCraft doesn't use guidance
-    dtype: torch.dtype | None = None
-    text_embed_dim: int = 4096  # LLaMA-3 hidden size
-    pooled_projection_dim: int = 768  # CLIP pooled output dim
-    rope_theta: int = 256
-    qk_norm: str = "rms_norm"
-
-    # Camera net parameters
-    camera_in_channels: int = 6  # Plücker coordinates
-    camera_downscale_coef: int = 8
-    camera_out_channels: int = 16
-
-    # Layers to exclude from LoRA
-    exclude_lora_layers: list[str] = field(
-        default_factory=lambda:
-        ["img_in", "txt_in", "time_in", "vector_in", "camera_net"])
-
-    def __post_init__(self):
-        super().__post_init__()
-        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
-        self.num_channels_latents: int = 16  # Output is 16 channels
-
-        # Convert patch_size list to tuple if needed (from JSON deserialization)
-        if isinstance(self.patch_size, list):
-            self.patch_size = tuple(self.patch_size)
-
-        # Convert rope_axes_dim list to tuple if needed
-        if isinstance(self.rope_axes_dim, list):
-            self.rope_axes_dim = tuple(self.rope_axes_dim)
-
-
-@dataclass
-class HunyuanGameCraftConfig(DiTConfig):
-    """Full config for HunyuanGameCraft model."""
-
-    arch_config: DiTArchConfig = field(
-        default_factory=HunyuanGameCraftArchConfig)
-
-    prefix: str = "HunyuanGameCraft"
@@ -55,8 +55,6 @@ 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",

            # 2. txt_in_2 mapping:
            r"^context_embedder_2\.(.*)$":
@@ -8,8 +8,12 @@ def is_double_block(n: str, m) -> bool:
    return "double" 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_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:
@@ -19,10 +23,10 @@ def is_txt_in(n: str, m) -> bool:
@dataclass
 class HYWorldArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_refiner_block])
+        default_factory=lambda: [is_double_block, is_single_block])

    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])
+        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])

    param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -50,7 +54,8 @@ class HYWorldArchConfig(DiTArchConfig):
            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:
+            # 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\.(.*)$":
@@ -7,12 +7,10 @@ from dataclasses import dataclass, field

 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig

-import re
-

 def is_ltx2_blocks(name: str, _module) -> bool:
-    res = re.search(r"(?:^|\.)transformer_blocks\.\d+$", name) is not None
-    return res
+    """FSDP shard condition for LTX-2 transformer blocks."""
+    return "transformer_blocks" in name


@dataclass
@@ -1,6 +1,4 @@
 from dataclasses import dataclass, field
-
-import torch
 from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig


@@ -78,14 +76,8 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig):
    image_dim: int = 1280


-def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool:
-    return bool("blocks" in param_name and param_name.split(".")[-1].isdigit())
-
-
@dataclass
 class MatrixGameWanVideoConfig(WanVideoConfig):
    arch_config: MatrixGameWanVideoArchConfig = field(
        default_factory=MatrixGameWanVideoArchConfig)
    prefix: str = "Wan"
-    _compile_conditions: list = field(
-        default_factory=lambda: [_is_transformer_block])
@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
-
-
-@dataclass
-class SD3Transformer2DArchConfig(DiTArchConfig):
-    # Diffusers SD3Transformer2DModel config fields.
-    sample_size: int = 128
-    patch_size: int = 2
-    num_layers: int = 24
-    attention_head_dim: int = 64
-    joint_attention_dim: int = 4096
-    caption_projection_dim: int = 1536
-    pooled_projection_dim: int = 2048
-    pos_embed_max_size: int = 384
-    dual_attention_layers: list[int] = field(
-        default_factory=lambda: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
-    qk_norm: str = "rms_norm"
-    in_channels: int = 16
-    out_channels: int = 16
-    num_attention_heads = 24
-
-
-@dataclass
-class SD3DiTConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(
-        default_factory=SD3Transformer2DArchConfig)
-    prefix: str = "sd3"
@@ -9,56 +9,47 @@ def is_blocks(n: str, m) -> bool:


@dataclass
-class LingBotWorldArchConfig(DiTArchConfig):
+class WanGameVideoArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])

    param_names_mapping: dict = field(
        default_factory=lambda: {
-            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
-            r"^patch_embedding_wancamctrl\.(.*)$":
-            r"patch_embedding_wancamctrl.proj.\1",
-            r"^c2ws_hidden_states_layer1\.(.*)$": r"c2ws_mlp.fc_in.\1",
-            r"^c2ws_hidden_states_layer2\.(.*)$": r"c2ws_mlp.fc_out.\1",
-            r"^text_embedding\.0\.(.*)$":
+            r"^patch_embedding\.(.*)$":
+            r"patch_embedding.proj.\1",
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
            r"condition_embedder.text_embedder.fc_in.\1",
-            r"^text_embedding\.2\.(.*)$":
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
            r"condition_embedder.text_embedder.fc_out.\1",
-            r"^time_embedding\.0\.(.*)$":
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
            r"condition_embedder.time_embedder.mlp.fc_in.\1",
-            r"^time_embedding\.2\.(.*)$":
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
            r"condition_embedder.time_embedder.mlp.fc_out.\1",
-            r"^time_projection\.1\.(.*)$":
+            r"^condition_embedder\.time_proj\.(.*)$":
            r"condition_embedder.time_modulation.linear.\1",
-            r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
-            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
-            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
-            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2",
-            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
-            r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
-            r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
-            r"^blocks\.(\d+)\.norm3\.(.*)$":
-            r"blocks.\1.self_attn_residual_norm.norm.\2",
-            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
-            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
-            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
-            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+            r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_in.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_out.\1",
+            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
+            r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
+            r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
+            r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
            r"blocks.\1.attn2.to_out.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
-            r"blocks.\1.attn2.norm_q.\2",
-            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
-            r"blocks.\1.attn2.norm_k.\2",
-            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
-            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
-            r"^blocks\.(\d+)\.cam_injector_layer1\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_injector.fc_in.\2",
-            r"^blocks\.(\d+)\.cam_injector_layer2\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_injector.fc_out.\2",
-            r"^blocks\.(\d+)\.cam_scale_layer\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_scale_layer.\2",
-            r"^blocks\.(\d+)\.cam_shift_layer\.(.*)$":
-            r"blocks.\1.cam_conditioner.cam_shift_layer.\2",
-            r"^head\.modulation$": r"scale_shift_table",
-            r"^head\.head\.(.*)$": r"proj_out.\1",
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
+            r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
+            r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm2\.(.*)$":
+            r"blocks.\1.self_attn_residual_norm.norm.\2",
        })

    # Reverse mapping for saving checkpoints: custom -> hf
@@ -66,10 +57,24 @@ class LingBotWorldArchConfig(DiTArchConfig):

    # Some LoRA adapters use the original official layer names instead of hf layer names,
    # so apply this before the param_names_mapping
-    lora_param_names_mapping: dict = field(default_factory=lambda: {})
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
+            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
+            r"blocks.\1.attn1.to_out.0.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+            r"blocks.\1.attn2.to_out.0.\2",
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+        })

    patch_size: tuple[int, int, int] = (1, 2, 2)
-    text_len: int = 512
+    text_len = 512
    num_attention_heads: int = 40
    attention_head_dim: int = 128
    in_channels: int = 16
@@ -104,7 +109,7 @@ class LingBotWorldArchConfig(DiTArchConfig):


@dataclass
-class LingBotWorldVideoConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=LingBotWorldArchConfig)
+class WanGameVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=WanGameVideoArchConfig)

-    prefix: str = "Wan"
+    prefix: str = "WanGame"
@@ -87,10 +87,10 @@ class CLIPVisionConfig(ImageEncoderConfig):
    arch_config: ImageEncoderArchConfig = field(
        default_factory=CLIPVisionArchConfig)

-    num_hidden_layers_override: int | None = 31
+    num_hidden_layers_override: int | None = None
    require_post_norm: bool | None = None
-    enable_scale: bool = False
-    is_causal: bool = False
+    enable_scale: bool = True
+    is_causal: bool = True
    prefix: str = "clip"


@@ -1,6 +0,0 @@
-from fastvideo.configs.models.upsamplers.hunyuan15 import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
-from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
-
-__all__ = [
-    "SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig"
-]
@@ -1,7 +0,0 @@
-from dataclasses import dataclass
-from fastvideo.configs.models.base import ModelConfig
-
-
-@dataclass
-class UpsamplerConfig(ModelConfig):
-    pass
@@ -1,20 +0,0 @@
-from dataclasses import dataclass
-from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
-
-
-@dataclass
-class SRTo720pUpsamplerConfig(UpsamplerConfig):
-    in_channels: int = 0
-    out_channels: int = 0
-    hidden_channels: int = 64
-    num_blocks: int = 6
-    global_residual: bool = False
-
-
-@dataclass
-class SRTo1080pUpsamplerConfig(UpsamplerConfig):
-    z_channels: int = 0
-    out_channels: int = 0
-    block_out_channels: tuple[int, ...] = (0, 0)
-    num_res_blocks: int = 2
-    is_residual: bool = False
@@ -1,6 +1,5 @@
 from fastvideo.configs.models.vaes.cosmosvae import CosmosVAEConfig
 from fastvideo.configs.models.vaes.cosmos2_5vae import Cosmos25VAEConfig
-from fastvideo.configs.models.vaes.gamecraftvae import GameCraftVAEConfig
 from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
@@ -8,7 +7,6 @@ from fastvideo.configs.models.vaes.stepvideovae import StepVideoVAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig

 __all__ = [
-    "GameCraftVAEConfig",
    "HunyuanVAEConfig",
    "WanVAEConfig",
    "StepVideoVAEConfig",
@@ -1,39 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-
-@dataclass
-class AutoencoderKLArchConfig(VAEArchConfig):
-    _name_or_path: str = ""
-    act_fn: str = "silu"
-    block_out_channels: tuple[int, ...] | list[int] = field(
-        default_factory=list)
-    down_block_types: tuple[str, ...] | list[str] = field(default_factory=list)
-    up_block_types: tuple[str, ...] | list[str] = field(default_factory=list)
-    force_upcast: bool = True
-    in_channels: int = 3
-    latent_channels: int = 4
-    latents_mean: tuple[float, ...] | list[float] | None = None
-    latents_std: tuple[float, ...] | list[float] | None = None
-    layers_per_block: int = 1
-    mid_block_add_attention: bool = True
-    norm_num_groups: int = 32
-    out_channels: int = 3
-    sample_size: int = 32
-    scaling_factor: float | torch.Tensor = 0.18215
-    shift_factor: float | None = None
-    use_post_quant_conv: bool = True
-    use_quant_conv: bool = True
-
-    temporal_compression_ratio: int = 1
-    spatial_compression_ratio: int = 8
-
-
-@dataclass
-class AutoencoderKLVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = field(default_factory=AutoencoderKLArchConfig)
@@ -1,50 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-GameCraft VAE config - matches official config.json from Hunyuan-GameCraft-1.0.
-"""
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-
-@dataclass
-class GameCraftVAEArchConfig(VAEArchConfig):
-    """Architecture config matching official AutoencoderKLCausal3D config.json."""
-
-    in_channels: int = 3
-    out_channels: int = 3
-    latent_channels: int = 16
-    down_block_types: tuple[str, ...] = (
-        "DownEncoderBlockCausal3D",
-        "DownEncoderBlockCausal3D",
-        "DownEncoderBlockCausal3D",
-        "DownEncoderBlockCausal3D",
-    )
-    up_block_types: tuple[str, ...] = (
-        "UpDecoderBlockCausal3D",
-        "UpDecoderBlockCausal3D",
-        "UpDecoderBlockCausal3D",
-        "UpDecoderBlockCausal3D",
-    )
-    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
-    layers_per_block: int = 2
-    act_fn: str = "silu"
-    norm_num_groups: int = 32
-    scaling_factor: float = 0.476986
-    spatial_compression_ratio: int = 8
-    temporal_compression_ratio: int = 4
-    time_compression_ratio: int = 4  # alias for DecoderCausal3D
-    mid_block_add_attention: bool = True
-    mid_block_causal_attn: bool = True
-    sample_size: int = 256  # from config.json
-    sample_tsize: int = 64  # from config.json
-
-    def __post_init__(self):
-        self.spatial_compression_ratio = 2**(len(self.block_out_channels) - 1)
-
-
-@dataclass
-class GameCraftVAEConfig(VAEConfig):
-    """Full config for GameCraft VAE."""
-
-    arch_config: VAEArchConfig = field(default_factory=GameCraftVAEArchConfig)
@@ -4,20 +4,20 @@ 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.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
-from fastvideo.registry import get_pipeline_config_cls_from_name
+from fastvideo.configs.pipelines.registry import (
+    get_pipeline_config_cls_from_name)
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig,
                                             WanI2V480PConfig, WanI2V720PConfig,
                                             WanT2V480PConfig, WanT2V720PConfig)

 __all__ = [
-    "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig",
-    "PipelineConfig", "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig",
-    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
-    "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
-    "SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config",
-    "LTX2T2VConfig", "HYWorldConfig", "get_pipeline_config_cls_from_name"
+    "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
+    "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "SlidingTileAttnConfig",
+    "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
+    "WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig",
+    "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
+    "get_pipeline_config_cls_from_name"
 ]
@@ -8,7 +8,7 @@ from typing import Any, cast
 import torch

 from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                      VAEConfig, UpsamplerConfig)
+                                      VAEConfig)
 from fastvideo.configs.models.encoders import BaseEncoderOutput
 from fastvideo.configs.utils import update_config_from_args
 from fastvideo.logger import init_logger
@@ -44,15 +44,12 @@ class PipelineConfig:
    # Video generation parameters
    embedded_cfg_scale: float = 6.0
    flow_shift: float | None = None
-    flow_shift_sr: float | None = None
    disable_autocast: bool = False
    is_causal: bool = False

    # Model configuration
    dit_config: DiTConfig = field(default_factory=DiTConfig)
    dit_precision: str = "bf16"
-    upsampler_config: UpsamplerConfig = field(default_factory=UpsamplerConfig)
-    upsampler_precision: str = "fp32"

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

-        # STA (Sliding Tile Attention) parameters
-        parser.add_argument(
-            f"--{prefix_with_dot}STA-mode",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}STA_mode",
-            default=PipelineConfig.STA_mode.value,
-            choices=[mode.value for mode in STA_Mode],
-            help=
-            "STA mode: STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}skip-time-steps",
-            type=int,
-            dest=f"{prefix_with_dot.replace('-', '_')}skip_time_steps",
-            default=PipelineConfig.skip_time_steps,
-            help="Number of time steps to warmup (full attention) for STA",
-        )
-
        # Add VAE configuration arguments
        from fastvideo.configs.models.vaes.base import VAEConfig
        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -266,7 +245,8 @@ class PipelineConfig:
        """
        use the pipeline class setting from model_path to match the pipeline config
        """
-        from fastvideo.registry import get_pipeline_config_cls_from_name
+        from fastvideo.configs.pipelines.registry import (
+            get_pipeline_config_cls_from_name)
        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)

        return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))
@@ -280,7 +260,8 @@ class PipelineConfig:
        kwargs: dictionary of kwargs
        config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
        """
-        from fastvideo.registry import get_pipeline_config_cls_from_name
+        from fastvideo.configs.pipelines.registry import (
+            get_pipeline_config_cls_from_name)

        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
                                                      != "") else ""
@@ -317,12 +298,6 @@ class PipelineConfig:
        # 4. Update PipelineConfig from CLI arguments if provided
        kwargs[prefix_with_dot + 'model_path'] = model_path
        pipeline_config.update_config_from_dict(kwargs, config_cli_prefix)
-
-        # Convert STA_mode string to enum if necessary
-        if isinstance(pipeline_config.STA_mode, str) and not isinstance(
-                pipeline_config.STA_mode, STA_Mode):
-            pipeline_config.STA_mode = STA_Mode(pipeline_config.STA_mode)
-
        return pipeline_config

    def check_pipeline_config(self) -> None:
@@ -11,8 +11,7 @@ from fastvideo.configs.models.dits import HunyuanVideo15Config
 from fastvideo.configs.models.encoders import (BaseEncoderOutput,
                                               Qwen2_5_VLConfig, T5Config)
 from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
-from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, UpsamplerConfig
+from fastvideo.configs.pipelines.base import PipelineConfig

 PROMPT_TEMPLATE_TOKEN_LENGTH = 108

@@ -132,35 +131,9 @@ class Hunyuan15T2V480PConfig(PipelineConfig):
        self.vae_config.load_decoder = True


-@dataclass
-class Hunyuan15I2V480PStepDistilledConfig(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 Hunyuan15I2V720PConfig(Hunyuan15T2V720PConfig):
-    """Base configuration for HunYuan pipeline architecture."""
-
-    # HunyuanConfig-specific parameters with defaults
-    flow_shift: int = 7
-
-
-@dataclass
-class Hunyuan15SR1080PConfig(Hunyuan15T2V720PConfig):
-    """Base configuration for HunYuan pipeline architecture."""
-
-    # HunyuanConfig-specific parameters with defaults
-    flow_shift: int = 7
-    flow_shift_sr: int = 2
-    upsampler_config: tuple[UpsamplerConfig, ...] = field(
-        default_factory=lambda:
-        (SRTo720pUpsamplerConfig(), SRTo1080pUpsamplerConfig()))
-    upsampler_precision: str = "fp32"
@@ -1,122 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Pipeline configuration for HunyuanGameCraft.
-
-HunyuanGameCraft extends HunyuanVideo with:
-1. CameraNet for camera/action conditioning (Plücker coordinates)
-2. Mask-based conditioning for autoregressive generation
-3. 33 input channels (16 latent + 16 gt_latent + 1 mask)
-
-Text encoders are the same as HunyuanVideo:
- LLaVA-LLaMA-3-8B for primary text encoding (4096 dim)
- CLIP ViT-L/14 for secondary pooled embeddings (768 dim)
-"""
-from collections.abc import Callable
-from dataclasses import dataclass, field
-from typing import TypedDict
-
-import torch
-
-from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.configs.models.dits import HunyuanGameCraftConfig
-from fastvideo.configs.models.encoders import (
-    BaseEncoderOutput,
-    CLIPTextConfig,
-    LlamaConfig,
-)
-from fastvideo.configs.models.vaes import GameCraftVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
-
-# GameCraft uses the same prompt template as HunyuanVideo
-PROMPT_TEMPLATE_ENCODE_VIDEO = (
-    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
-    "1. The main content and theme of the video."
-    "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
-    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
-    "4. background environment, light, style and atmosphere."
-    "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
-    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>")
-
-
-class PromptTemplate(TypedDict):
-    template: str
-    crop_start: int
-
-
-prompt_template_video: PromptTemplate = {
-    "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
-    "crop_start": 95,
-}
-
-
-def llama_preprocess_text(prompt: str) -> str:
-    """Apply prompt template for LLaMA encoder."""
-    return prompt_template_video["template"].format(prompt)
-
-
-def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-    """Extract hidden states from LLaMA output, skipping instruction tokens."""
-    hidden_state_skip_layer = 2
-    assert outputs.hidden_states is not None
-    hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
-    last_hidden_state: torch.Tensor = hidden_states[-(hidden_state_skip_layer +
-                                                      1)]
-    crop_start = prompt_template_video.get("crop_start", -1)
-    last_hidden_state = last_hidden_state[:, crop_start:]
-    return last_hidden_state
-
-
-def clip_preprocess_text(prompt: str) -> str:
-    """No preprocessing for CLIP encoder."""
-    return prompt
-
-
-def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-    """Extract pooled output from CLIP encoder."""
-    pooler_output: torch.Tensor = outputs.pooler_output
-    return pooler_output
-
-
-@dataclass
-class HunyuanGameCraftPipelineConfig(PipelineConfig):
-    """Configuration for HunyuanGameCraft pipeline.
-    
-    Inherits text encoding from HunyuanVideo but uses:
-    - GameCraft DiT with CameraNet
-    - Same VAE (HunyuanVAE)
-    - Same text encoders (LLaMA + CLIP)
-    """
-
-    # DiT config - uses GameCraft config (33 input channels)
-    dit_config: DiTConfig = field(default_factory=HunyuanGameCraftConfig)
-
-    # VAE config - GameCraft VAE (mid_block_causal_attn=True, etc.)
-    vae_config: VAEConfig = field(default_factory=GameCraftVAEConfig)
-
-    # Denoising parameters
-    # Official GameCraft does NOT use embedded guidance (passes guidance=None)
-    # It uses standard CFG with guidance_scale=6.0 instead
-    embedded_cfg_scale = None
-    flow_shift: int = 5  # Official GameCraft uses flow_shift=5.0
-
-    # Text encoding stage - same as HunyuanVideo
-    # Uses LLaMA-3-8B (via LLaVA) + CLIP
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
-    postprocess_text_funcs: tuple[
-        Callable[[BaseEncoderOutput], torch.Tensor],
-        ...] = field(default_factory=lambda:
-                     (llama_postprocess_text, clip_postprocess_text))
-
-    # Precision for each component
-    dit_precision: str = "bf16"
-    vae_precision: str = "fp16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp16", "fp16"))
-
-    def __post_init__(self):
-        # VAE only needs decoder for inference
-        self.vae_config.load_encoder = False
-        self.vae_config.load_decoder = True
@@ -1,13 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from dataclasses import dataclass, field
-
-from fastvideo.configs.models import DiTConfig
-from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config
-from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig
-
-
-@dataclass
-class LingBotWorldI2V480PConfig(Wan2_2_I2V_A14B_Config):
-    dit_config: DiTConfig = field(default_factory=LingBotWorldVideoConfig)
-    flow_shift: float | None = 10.0
-    boundary_ratio: float | None = 0.947
@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Registry for pipeline weight-specific configurations."""
+
+import os
+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.stepvideo import StepVideoT2VConfig
+from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
+from fastvideo.configs.pipelines.turbodiffusion import (
+    TurboDiffusionT2V_1_3B_Config, TurboDiffusionT2V_14B_Config,
+    TurboDiffusionI2V_A14B_Config)
+
+# isort: off
+from fastvideo.configs.pipelines.wan import (
+    FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
+    Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
+    WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
+    SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig,
+    MatrixGameI2V480PConfig)
+# isort: on
+from fastvideo.logger import init_logger
+from fastvideo.utils import (maybe_download_model_index,
+                             verify_model_config_and_directory)
+
+logger = init_logger(__name__)
+
+# Registry maps specific model weights to their config classes
+PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
+    "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
+    "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
+    "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,
+    "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
+    "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
+    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWan2_1_T2V_480P_Config,
+    "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers": FastWan2_1_T2V_480P_Config,
+    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
+    "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
+    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingWan2_2_T2V480PConfig,
+    "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers":
+    SelfForcingWan2_2_T2V480PConfig,
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
+    "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,
+    # LongCat Video models
+    "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,
+    "loayrashid/TurboWan2.2-I2V-A14B-Diffusers": TurboDiffusionI2V_A14B_Config,
+    # Add other specific weight variants
+}
+
+# For determining pipeline type from model ID
+PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
+    "longcatimagetovideo":
+    lambda id: "longcatimagetovideo" in id.lower(),
+    "longcatvideocontinuation":
+    lambda id: "longcatvideocontinuation" in id.lower(),
+    "longcat":
+    lambda id: "longcat" in id.lower(),
+    "hunyuan":
+    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":
+    lambda id: "wanpipeline" in id.lower(),
+    "wanimagetovideo":
+    lambda id: "wanimagetovideo" in id.lower(),
+    "wandmdpipeline":
+    lambda id: "wandmdpipeline" in id.lower(),
+    "wancausaldmdpipeline":
+    lambda id: "wancausaldmdpipeline" in id.lower(),
+    "stepvideo":
+    lambda id: "stepvideo" in id.lower(),
+    "cosmos":
+    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
+}
+
+# Fallback configs when exact match isn't found but architecture is detected
+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,
+    "wandmdpipeline": FastWan2_1_T2V_480P_Config,
+    "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
+    "stepvideo": StepVideoT2VConfig,
+    "turbodiffusion": TurboDiffusionT2V_1_3B_Config,
+    "ltx2": LTX2T2VConfig,
+    # Other fallbacks by architecture
+}
+
+
+def get_pipeline_config_cls_from_name(
+        pipeline_name_or_path: str) -> type[PipelineConfig]:
+    """Get the appropriate configuration class for a given pipeline name or path.
+
+    This function implements a multi-step lookup process to find the most suitable
+    configuration class for a given pipeline. It follows this order:
+    1. Exact match in the PIPE_NAME_TO_CONFIG
+    2. Partial match in the PIPE_NAME_TO_CONFIG
+    3. Fallback to class name in the model_index.json
+    4. else raise an error
+
+    Args:
+        pipeline_name_or_path (str): The name or path of the pipeline. This can be:
+            - A registered model ID (e.g., "FastVideo/FastHunyuan-diffusers")
+            - A local path to a model directory
+            - A model ID that will be downloaded
+
+    Returns:
+        Type[PipelineConfig]: The configuration class that best matches the pipeline.
+            This will be one of:
+            - A specific weight configuration class if an exact match is found
+            - A fallback configuration class based on the pipeline architecture
+            - The base PipelineConfig class if no matches are found
+
+    Note:
+        - For local paths, the function will verify the model configuration
+        - For remote models, it will attempt to download the model index
+        - Warning messages are logged when falling back to less specific configurations
+    """
+
+    pipeline_config_cls: type[PipelineConfig] | None = None
+
+    # First try exact match for specific weights
+    if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:
+        pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path]
+        return pipeline_config_cls
+
+    # Try partial matches (for local paths that might include the weight ID)
+    for registered_id, config_class in PIPE_NAME_TO_CONFIG.items():
+        if registered_id in pipeline_name_or_path:
+            pipeline_config_cls = config_class
+            break
+
+    # If no match, try to use the fallback config
+    if pipeline_config_cls is None:
+        if os.path.exists(pipeline_name_or_path):
+            config = verify_model_config_and_directory(pipeline_name_or_path)
+        else:
+            config = maybe_download_model_index(pipeline_name_or_path)
+        logger.warning(
+            "Trying to use the config from the model_index.json. FastVideo may not correctly identify the optimal config for this model in this situation."
+        )
+
+        pipeline_name = config["_class_name"]
+        # Try to determine pipeline architecture for fallback
+        for pipeline_type, detector in PIPELINE_DETECTOR.items():
+            if detector(pipeline_name.lower()):
+                pipeline_config_cls = PIPELINE_FALLBACK_CONFIG.get(
+                    pipeline_type)
+                break
+
+        if pipeline_config_cls is not None:
+            logger.warning(
+                "No match found for pipeline %s, using fallback config %s.",
+                pipeline_name_or_path, pipeline_config_cls)
+
+    if pipeline_config_cls is None:
+        raise ValueError(
+            f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path."
+        )
+
+    return pipeline_config_cls
@@ -1,65 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
-
-from fastvideo.configs.models import EncoderConfig
-from fastvideo.configs.models.encoders import (
-    BaseEncoderOutput,
-    CLIPTextConfig,
-    T5Config,
-)
-from fastvideo.configs.models.dits.sd3 import SD3DiTConfig
-from fastvideo.configs.models.vaes.autoencoder_kl import AutoencoderKLVAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
-
-
-def _sd35_text_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
-    assert outputs.last_hidden_state is not None
-    return outputs.last_hidden_state
-
-
-@dataclass
-class SD35Config(PipelineConfig):
-
-    scheduler_arch: str = "FlowMatchEulerDiscreteScheduler"
-    transformer_arch: str = "SD3Transformer2DModel"
-    vae_arch: str = "AutoencoderKL"
-    text_encoder_archs: tuple[str, ...] = (
-        "CLIPTextModelWithProjection",
-        "CLIPTextModelWithProjection",
-        "T5EncoderModel",
-    )
-    tokenizer_archs: tuple[str, ...] = (
-        "CLIPTokenizer",
-        "CLIPTokenizer",
-        "T5TokenizerFast",
-    )
-
-    dit_config: SD3DiTConfig = field(default_factory=SD3DiTConfig)
-    vae_config: AutoencoderKLVAEConfig = field(
-        default_factory=AutoencoderKLVAEConfig)
-
-    embedded_cfg_scale: float = 0.0
-    flow_shift: float | None = None
-
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda:
-        (CLIPTextConfig(), CLIPTextConfig(), T5Config()))
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda:
-        (preprocess_text, preprocess_text, preprocess_text))
-    postprocess_text_funcs: tuple[
-        Callable[[BaseEncoderOutput], torch.Tensor],
-        ...] = field(default_factory=lambda:
-                     (_sd35_text_postprocess, _sd35_text_postprocess,
-                      _sd35_text_postprocess))
-
-    dit_precision: str = "bf16"
-    vae_precision: str = "fp32"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp32", "fp32", "bf16"))
@@ -1,13 +1,3 @@
 from fastvideo.configs.sample.base import SamplingParam
-from fastvideo.configs.sample.hunyuangamecraft import (
-    HunyuanGameCraftSamplingParam,
-    HunyuanGameCraft65FrameSamplingParam,
-    HunyuanGameCraft129FrameSamplingParam,
-)

-__all__ = [
-    "SamplingParam",
-    "HunyuanGameCraftSamplingParam",
-    "HunyuanGameCraft65FrameSamplingParam",
-    "HunyuanGameCraft129FrameSamplingParam",
-]
+__all__ = ["SamplingParam"]
@@ -28,12 +28,6 @@ class SamplingParam:
    keyboard_cond: Any | None = None  # Shape: (B, T, K)
    grid_sizes: Any | None = None  # Shape: (3,) [F,H,W]

-    # Camera control inputs (HYWorld)
-    pose: str | None = None  # Camera trajectory: pose string (e.g., 'w-31') or JSON file path
-
-    # Camera control inputs (LingBotWorld)
-    c2ws_plucker_emb: Any | None = None  # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
-
    # Refine inputs (LongCat 480p->720p upscaling)
    # Path-based refine (load stage1 video from disk, e.g. MP4)
    refine_from: str | None = None  # Path to stage1 video (480p output from distill)
@@ -61,13 +55,10 @@ class SamplingParam:
    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
    height: int = 720
    width: int = 1280
-    height_sr: int = 1072
-    width_sr: int = 1920
    fps: int = 24

    # Denoising parameters
    num_inference_steps: int = 50
-    num_inference_steps_sr: int = 50
    guidance_scale: float = 1.0
    guidance_rescale: float = 0.0
    boundary_ratio: float | None = None
@@ -101,7 +92,8 @@ class SamplingParam:

    @classmethod
    def from_pretrained(cls, model_path: str) -> "SamplingParam":
-        from fastvideo.registry import get_sampling_param_cls_for_name
+        from fastvideo.configs.sample.registry import (
+            get_sampling_param_cls_for_name)
        sampling_cls = get_sampling_param_cls_for_name(model_path)
        if sampling_cls is not None:
            sampling_param: SamplingParam = sampling_cls()
@@ -5,19 +5,15 @@ from fastvideo.configs.sample.base import SamplingParam


@dataclass
-class Cosmos25SamplingParamBase(SamplingParam):
-    height: int = 704
-    width: int = 1280
-    num_frames: int = 77
+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
-    seed: int = 0

    guidance_scale: float = 7.0
-    negative_prompt: str = (
-        "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
-        "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
-        "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, "
-        "low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, "
-        "unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
-        "Overall, the video is of poor quality.")
+    # Official Cosmos2.5 sampling uses empty string as unconditional.
+    negative_prompt: str = ""
    num_inference_steps: int = 35
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
JerryZhou54	bdaf2e2e45	Support custom action trajectories for validation	2026-02-03 20:23:13 +00:00
JerryZhou54	f3a7a37312	Overfitting running for MC 10	2026-02-03 05:11:48 +00:00
JerryZhou54	b1abb5c251	Update action_labels	2026-02-02 05:20:06 +00:00
RandNMR73	a80824255e	Add data preprocess pipeline for WanGame	2026-02-02 04:38:24 +00:00