diff --git a/README.md b/README.md index 9466b3fcc..f2e34ff2c 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ FastVideo has the following features: - 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 achineve >50x denoising speedup + - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) to achieve >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. @@ -108,55 +108,35 @@ 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/). -### Other docs: +## More Guides - [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/) ## 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. [![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) +- [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. ## 🤝 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 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. +## Acknowledgement +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. ## Citation -If you find FastVideo useful, please considering citing our work: +If you find FastVideo useful, please consider citing our research 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}, diff --git a/docker/Dockerfile.python3.10 b/docker/Dockerfile.python3.10 index 98a1e0bd7..ac5f2dacd 100644 --- a/docker/Dockerfile.python3.10 +++ b/docker/Dockerfile.python3.10 @@ -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.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp310-cp310-linux_x86_64.whl + 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 COPY . . diff --git a/docker/Dockerfile.python3.11 b/docker/Dockerfile.python3.11 index fe3639a60..b74252521 100644 --- a/docker/Dockerfile.python3.11 +++ b/docker/Dockerfile.python3.11 @@ -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.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp311-cp311-linux_x86_64.whl + 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 COPY . . diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index 64c0103c8..698bcb81a 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -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.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp312-cp312-linux_x86_64.whl + 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 COPY . . diff --git a/docs/api/fastvideo.md b/docs/api/fastvideo.md index 2f81a3995..85a17d13a 100644 --- a/docs/api/fastvideo.md +++ b/docs/api/fastvideo.md @@ -121,6 +121,15 @@ 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 diff --git a/docs/contributing/coding_agents.md b/docs/contributing/coding_agents.md index a38d67468..e6f8f8589 100644 --- a/docs/contributing/coding_agents.md +++ b/docs/contributing/coding_agents.md @@ -324,8 +324,8 @@ Purpose: Action: - Add a new pipeline config + sampling params. -- Register them in `fastvideo/configs/pipelines/registry.py` and - `fastvideo/configs/sample/registry.py`. +- Register them in `fastvideo/registry.py` using explicit + `register_configs(...)` blocks (this file is the single source of truth now). ### 5) Wire pipeline stages diff --git a/docs/contributing/developer_env/runpod.md b/docs/contributing/developer_env/runpod.md index baea17138..b14b683d1 100644 --- a/docs/contributing/developer_env/runpod.md +++ b/docs/contributing/developer_env/runpod.md @@ -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 "--". +- Set the Pod name to "`--`". - 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) diff --git a/docs/design/overview.md b/docs/design/overview.md index 9f1537b09..8365bc3a0 100644 --- a/docs/design/overview.md +++ b/docs/design/overview.md @@ -51,8 +51,9 @@ runtime parameters consistent: - `fastvideo/configs/pipelines/`: pipeline wiring and required components. - `fastvideo/configs/sample/`: default sampling parameters (steps, frames, guidance scale, resolution, fps). -- `fastvideo/configs/registry.py`: pipeline registry. -- `fastvideo/configs/sample/registry.py`: sampling param registry. +- `fastvideo/registry.py`: unified registry for pipeline config + sampling + defaults and model metadata resolution, defined via explicit + `register_configs(...)` blocks (no separate dict registries). `FastVideoArgs` (in `fastvideo/fastvideo_args.py`) provides runtime settings and is passed into pipeline construction and stages. diff --git a/docs/inference/add_pipeline.md b/docs/inference/add_pipeline.md index 702904291..c09806e66 100644 --- a/docs/inference/add_pipeline.md +++ b/docs/inference/add_pipeline.md @@ -131,17 +131,41 @@ self.attn = DistributedAttention( ### Registering Models -Register implemented modules in the model registry: +Register implemented modules for auto‑discovery by adding `EntryClass` in each +model module (the registry scans for it): ```python -# In fastvideo/models/registry.py -_TEXT_TO_VIDEO_DIT_MODELS = { - "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"), -} +# In fastvideo/models/dits/your_module.py +class YourTransformerModel(...): + ... -_VAE_MODELS = { - "YourVAEModel": ("vaes", "yourvae", "YourVAEClass"), -} +# 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(), + ], +) ``` ## Step 2: Directory Structure diff --git a/examples/inference/basic/basic_cosmos2_5_i2w.py b/examples/inference/basic/basic_cosmos2_5_i2w.py new file mode 100644 index 000000000..5cdb631fb --- /dev/null +++ b/examples/inference/basic/basic_cosmos2_5_i2w.py @@ -0,0 +1,51 @@ +# 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 = "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() + + diff --git a/examples/inference/basic/basic_cosmos2_5_t2w.py b/examples/inference/basic/basic_cosmos2_5_t2w.py index d2d11a0ef..d3ed30c83 100644 --- a/examples/inference/basic/basic_cosmos2_5_t2w.py +++ b/examples/inference/basic/basic_cosmos2_5_t2w.py @@ -1,4 +1,5 @@ from fastvideo import VideoGenerator +from fastvideo.configs.sample import SamplingParam def main(): @@ -15,19 +16,27 @@ 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." ) - video = generator.generate_video( + generator.generate_video( prompt, - negative_prompt="", - height=704, - width=1280, - num_frames=77, - num_inference_steps=35, - guidance_scale=7.0, - fps=24, + sampling_param=sampling_param, output_path="outputs_video/cosmos2_5_t2w.mp4", save_video=True, ) diff --git a/examples/inference/basic/basic_cosmos2_5_v2w.py b/examples/inference/basic/basic_cosmos2_5_v2w.py new file mode 100644 index 000000000..c10218b01 --- /dev/null +++ b/examples/inference/basic/basic_cosmos2_5_v2w.py @@ -0,0 +1,54 @@ +# 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 = "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() + + diff --git a/examples/inference/basic/basic_hy15.py b/examples/inference/basic/basic_hy15.py index 39d8772d3..e0b782237 100644 --- a/examples/inference/basic/basic_hy15.py +++ b/examples/inference/basic/basic_hy15.py @@ -39,4 +39,4 @@ def main(): if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/examples/inference/basic/basic_hy15_1080p.py b/examples/inference/basic/basic_hy15_1080p.py new file mode 100644 index 000000000..6b346e8f9 --- /dev/null +++ b/examples/inference/basic/basic_hy15_1080p.py @@ -0,0 +1,43 @@ +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() diff --git a/examples/inference/basic/basic_hyworld.py b/examples/inference/basic/basic_hyworld.py index 02bd6e058..4caae8656 100644 --- a/examples/inference/basic/basic_hyworld.py +++ b/examples/inference/basic/basic_hyworld.py @@ -1,165 +1,11 @@ -# 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 @@ -169,7 +15,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='video_samples_hyworld', help="Output video path") + parser.add_argument("--output_path", type=str, default=OUTPUT_PATH, 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") @@ -185,8 +31,7 @@ def main(): # Initialize generator print("\nInitializing VideoGenerator for HYWorld...") - - generator = HYWorldVideoGenerator.from_pretrained( + generator = VideoGenerator.from_pretrained( "FastVideo/HY-WorldPlay-Bidirectional-Diffusers", num_gpus=1, use_fsdp_inference=True, @@ -198,11 +43,12 @@ def main(): ) # Generate video + # The pose string is automatically converted to camera matrices by the pipeline print("\nGenerating video...") - start_time = time.time() - video = generator.generate_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="", @@ -211,13 +57,9 @@ def main(): height=HEIGHT, width=WIDTH, seed=args.seed, - pose=args.pose, ) - elapsed = time.time() - start_time - print(f"\nVideo generated successfully!") - print(f"Saved to: {args.output_path}") - print(f"Time: {elapsed:.2f}s") + print(f"\nVideo saved to: {args.output_path}") if __name__ == "__main__": diff --git a/fastvideo/configs/models/__init__.py b/fastvideo/configs/models/__init__.py index aa05ba62e..58b38d40c 100644 --- a/fastvideo/configs/models/__init__.py +++ b/fastvideo/configs/models/__init__.py @@ -5,6 +5,7 @@ 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", @@ -14,4 +15,5 @@ __all__ = [ "LTX2AudioEncoderConfig", "LTX2AudioDecoderConfig", "LTX2VocoderConfig", + "UpsamplerConfig", ] diff --git a/fastvideo/configs/models/dits/hunyuanvideo15.py b/fastvideo/configs/models/dits/hunyuanvideo15.py index 6d7f77208..bc62df63f 100644 --- a/fastvideo/configs/models/dits/hunyuanvideo15.py +++ b/fastvideo/configs/models/dits/hunyuanvideo15.py @@ -55,6 +55,8 @@ 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\.(.*)$": diff --git a/fastvideo/configs/models/dits/hyworld.py b/fastvideo/configs/models/dits/hyworld.py index 36c8c98d4..563cf471f 100644 --- a/fastvideo/configs/models/dits/hyworld.py +++ b/fastvideo/configs/models/dits/hyworld.py @@ -8,12 +8,8 @@ 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_refiner_block(n: str, m) -> bool: + return "refiner" in n and str.isdigit(n.split(".")[-1]) def is_txt_in(n: str, m) -> bool: @@ -23,10 +19,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_single_block]) + default_factory=lambda: [is_double_block, is_refiner_block]) _compile_conditions: list = field( - default_factory=lambda: [is_double_block, is_single_block, is_txt_in]) + default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in]) param_names_mapping: dict = field( default_factory=lambda: { @@ -54,8 +50,7 @@ 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 (HYWorld uses TimestepEmbedder directly, - # but FastVideo model inherits HunyuanVideo15TimeEmbedding with timestep_embedder): + # 2. time_in mappings: r"^time_in\.mlp\.0\.(.*)$": r"time_in.timestep_embedder.mlp.fc_in.\1", r"^time_in\.mlp\.2\.(.*)$": diff --git a/fastvideo/configs/models/encoders/clip.py b/fastvideo/configs/models/encoders/clip.py index d233872f1..7cc5c70d5 100644 --- a/fastvideo/configs/models/encoders/clip.py +++ b/fastvideo/configs/models/encoders/clip.py @@ -87,10 +87,10 @@ class CLIPVisionConfig(ImageEncoderConfig): arch_config: ImageEncoderArchConfig = field( default_factory=CLIPVisionArchConfig) - num_hidden_layers_override: int | None = None + num_hidden_layers_override: int | None = 31 require_post_norm: bool | None = None - enable_scale: bool = True - is_causal: bool = True + enable_scale: bool = False + is_causal: bool = False prefix: str = "clip" diff --git a/fastvideo/configs/models/upsamplers/__init__.py b/fastvideo/configs/models/upsamplers/__init__.py new file mode 100644 index 000000000..b5b8332cb --- /dev/null +++ b/fastvideo/configs/models/upsamplers/__init__.py @@ -0,0 +1,6 @@ +from fastvideo.configs.models.upsamplers.hunyuan15 import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig +from fastvideo.configs.models.upsamplers.base import UpsamplerConfig + +__all__ = [ + "SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig" +] diff --git a/fastvideo/configs/models/upsamplers/base.py b/fastvideo/configs/models/upsamplers/base.py new file mode 100644 index 000000000..5430912bb --- /dev/null +++ b/fastvideo/configs/models/upsamplers/base.py @@ -0,0 +1,7 @@ +from dataclasses import dataclass +from fastvideo.configs.models.base import ModelConfig + + +@dataclass +class UpsamplerConfig(ModelConfig): + pass diff --git a/fastvideo/configs/models/upsamplers/hunyuan15.py b/fastvideo/configs/models/upsamplers/hunyuan15.py new file mode 100644 index 000000000..dc16459c7 --- /dev/null +++ b/fastvideo/configs/models/upsamplers/hunyuan15.py @@ -0,0 +1,20 @@ +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 diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 9291f57b9..1ee1ef31a 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -7,8 +7,7 @@ from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyua from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig -from fastvideo.configs.pipelines.registry import ( - get_pipeline_config_cls_from_name) +from fastvideo.registry import get_pipeline_config_cls_from_name from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig, diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index bc5fd940d..a057628a3 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -8,7 +8,7 @@ from typing import Any, cast import torch from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig, - VAEConfig) + VAEConfig, UpsamplerConfig) from fastvideo.configs.models.encoders import BaseEncoderOutput from fastvideo.configs.utils import update_config_from_args from fastvideo.logger import init_logger @@ -44,12 +44,15 @@ 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) @@ -245,8 +248,7 @@ class PipelineConfig: """ use the pipeline class setting from model_path to match the pipeline config """ - from fastvideo.configs.pipelines.registry import ( - get_pipeline_config_cls_from_name) + from fastvideo.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)) @@ -260,8 +262,7 @@ class PipelineConfig: kwargs: dictionary of kwargs config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance """ - from fastvideo.configs.pipelines.registry import ( - get_pipeline_config_cls_from_name) + from fastvideo.registry import get_pipeline_config_cls_from_name prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip() != "") else "" diff --git a/fastvideo/configs/pipelines/hunyuan15.py b/fastvideo/configs/pipelines/hunyuan15.py index 29e1c5b45..2b9e76c50 100644 --- a/fastvideo/configs/pipelines/hunyuan15.py +++ b/fastvideo/configs/pipelines/hunyuan15.py @@ -11,7 +11,8 @@ 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.pipelines.base import PipelineConfig +from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig +from fastvideo.configs.pipelines.base import PipelineConfig, UpsamplerConfig PROMPT_TEMPLATE_TOKEN_LENGTH = 108 @@ -131,9 +132,35 @@ 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" diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 146edd774..ab3131419 100644 --- a/fastvideo/configs/sample/base.py +++ b/fastvideo/configs/sample/base.py @@ -28,6 +28,9 @@ 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 + # 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) @@ -55,10 +58,13 @@ 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 @@ -92,8 +98,7 @@ class SamplingParam: @classmethod def from_pretrained(cls, model_path: str) -> "SamplingParam": - from fastvideo.configs.sample.registry import ( - get_sampling_param_cls_for_name) + from fastvideo.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() diff --git a/fastvideo/configs/sample/cosmos2_5.py b/fastvideo/configs/sample/cosmos2_5.py index 48287b303..1bcb5573e 100644 --- a/fastvideo/configs/sample/cosmos2_5.py +++ b/fastvideo/configs/sample/cosmos2_5.py @@ -5,15 +5,19 @@ from fastvideo.configs.sample.base import SamplingParam @dataclass -class Cosmos_Predict2_5_2B_Diffusers_SamplingParam(SamplingParam): - """Defaults for Cosmos 2.5 (Predict2.5) text-to-video diffusers-format model.""" - - height: int = 480 - width: int = 832 - num_frames: int = 121 +class Cosmos25SamplingParamBase(SamplingParam): + height: int = 704 + width: int = 1280 + num_frames: int = 77 fps: int = 24 + seed: int = 0 guidance_scale: float = 7.0 - # Official Cosmos2.5 sampling uses empty string as unconditional. - negative_prompt: str = "" + 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.") num_inference_steps: int = 35 diff --git a/fastvideo/configs/sample/hunyuan15.py b/fastvideo/configs/sample/hunyuan15.py index 941b9ea22..d50adb186 100644 --- a/fastvideo/configs/sample/hunyuan15.py +++ b/fastvideo/configs/sample/hunyuan15.py @@ -28,7 +28,33 @@ class Hunyuan15_480P_SamplingParam(SamplingParam): np.linspace(1.0, 0.0, self.num_inference_steps + 1)[:-1]) +@dataclass +class Hunyuan15_480P_StepDistilled_I2V_SamplingParam( + Hunyuan15_480P_SamplingParam): + num_inference_steps: int = 12 + + height: int = 720 + width: int = 1280 + guidance_scale: float = 1.0 + + @dataclass class Hunyuan15_720P_SamplingParam(Hunyuan15_480P_SamplingParam): height: int = 720 width: int = 1280 + + +@dataclass +class Hunyuan15_720P_Distilled_I2V_SamplingParam(Hunyuan15_720P_SamplingParam): + guidance_scale: float = 1.0 + + +@dataclass +class Hunyuan15_SR_1080P_SamplingParam(Hunyuan15_480P_SamplingParam): + height_sr: int = 1072 + width_sr: int = 1920 + + num_inference_steps: int = 12 + num_inference_steps_sr: int = 8 + + guidance_scale: float = 1.0 diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py deleted file mode 100644 index 22014b06b..000000000 --- a/fastvideo/configs/sample/registry.py +++ /dev/null @@ -1,204 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import os -from collections.abc import Callable -from typing import Any - -from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam, - HunyuanSamplingParam) -from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam -from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam -from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam - -from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam -from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam -from fastvideo.configs.sample.ltx2 import LTX2SamplingParam - -# isort: off -from fastvideo.configs.sample.wan import ( - FastWanT2V480P_SamplingParam, - Wan2_1_Fun_1_3B_InP_SamplingParam, - Wan2_2_I2V_A14B_SamplingParam, - Wan2_2_T2V_A14B_SamplingParam, - Wan2_2_TI2V_5B_SamplingParam, - WanI2V_14B_480P_SamplingParam, - WanI2V_14B_720P_SamplingParam, - WanT2V_1_3B_SamplingParam, - WanT2V_14B_SamplingParam, - Wan2_1_Fun_1_3B_Control_SamplingParam, - SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, - SelfForcingWan2_2_T2V_A14B_480P_SamplingParam, - MatrixGame2_SamplingParam, -) -from fastvideo.configs.sample.turbodiffusion import ( - TurboDiffusionT2V_1_3B_SamplingParam, - TurboDiffusionT2V_14B_SamplingParam, - TurboDiffusionI2V_A14B_SamplingParam, -) -# 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 -SAMPLING_PARAM_REGISTRY: dict[str, Any] = { - "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, - "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, - "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v": - Hunyuan15_480P_SamplingParam, - "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v": - Hunyuan15_720P_SamplingParam, - "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorld_SamplingParam, - "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, - - # Wan2.1 - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, - "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, - "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, - "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, - "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": - Wan2_1_Fun_1_3B_InP_SamplingParam, - "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": - Wan2_1_Fun_1_3B_Control_SamplingParam, - - # Wan2.2 - "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, - "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": - Wan2_2_TI2V_5B_SamplingParam, - "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, - "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, - - # FastWan2.1 - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480P_SamplingParam, - - # FastWan2.2 - "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, - - # Causal Self-Forcing Wan2.1 - "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": - SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, - - # Causal Self-Forcing Wan2.2 - "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": - SelfForcingWan2_2_T2V_A14B_480P_SamplingParam, - "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers": - SelfForcingWan2_2_T2V_A14B_480P_SamplingParam, - - # Cosmos2 - "nvidia/Cosmos-Predict2-2B-Video2World": - Cosmos_Predict2_2B_Video2World_SamplingParam, - - # Cosmos2.5 - "KyleShao/Cosmos-Predict2.5-2B-Diffusers": - Cosmos_Predict2_5_2B_Diffusers_SamplingParam, - - # MatrixGame2.0 models - "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam, - "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam, - "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGame2_SamplingParam, - - # TurboDiffusion models - "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": - TurboDiffusionT2V_1_3B_SamplingParam, - "loayrashid/TurboWan2.1-T2V-14B-Diffusers": - TurboDiffusionT2V_14B_SamplingParam, - "loayrashid/TurboWan2.2-I2V-A14B-Diffusers": - TurboDiffusionI2V_A14B_SamplingParam, - - # LTX-2 models - "Lightricks/LTX-2": LTX2SamplingParam, - "FastVideo/LTX2-Distilled-Diffusers": LTX2SamplingParam, - - # Add other specific weight variants -} - -# For determining pipeline type from model ID -SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { - "hunyuan": - lambda id: "hunyuan" in id.lower(), - "hunyuan15": - lambda id: "hunyuan15" in id.lower(), - "hyworld": - lambda id: "hyworld" in id.lower(), - "wanpipeline": - lambda id: "wanpipeline" in id.lower(), - "wanimagetovideo": - lambda id: "wanimagetovideo" in id.lower(), - "stepvideo": - lambda id: "stepvideo" in id.lower(), - "wandmdpipeline": - lambda id: "wandmdpipeline" in id.lower(), - "wancausaldmdpipeline": - lambda id: "wancausaldmdpipeline" in id.lower(), - "matrixgame": - lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(), - "turbodiffusion": - lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(), - "cosmos25": - lambda id: "cosmos2_5" in id.lower(), - "cosmos": - lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(), - "ltx2": - lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(), - # Add other pipeline architecture detectors -} - -# Fallback configs when exact match isn't found but architecture is detected -SAMPLING_FALLBACK_PARAM: dict[str, Any] = { - "hunyuan": - HunyuanSamplingParam, # Base Hunyuan config as fallback for any Hunyuan variant - "hunyuan15": - Hunyuan15_480P_SamplingParam, # Base Hunyuan15 config as fallback for any Hunyuan15 variant - "hyworld": - HYWorld_SamplingParam, # HYWorld-specific config as fallback for any HYWorld variant - "wanpipeline": - WanT2V_1_3B_SamplingParam, # Base Wan config as fallback for any Wan variant - "wanimagetovideo": WanI2V_14B_480P_SamplingParam, - "wandmdpipeline": FastWanT2V480P_SamplingParam, - "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, - "stepvideo": StepVideoT2VSamplingParam, - "matrixgame": MatrixGame2_SamplingParam, - "turbodiffusion": - TurboDiffusionT2V_1_3B_SamplingParam, # Default to T2V for fallback - "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam, - "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam, - "ltx2": LTX2SamplingParam, - # Other fallbacks by architecture -} - - -def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None: - """Get the appropriate sampling param for specific pretrained weights.""" - # First try exact match for specific weights - if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY: - return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path] - - # Try partial matches (for local paths that might include the weight ID) - for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items(): - if registered_id in pipeline_name_or_path: - return config_class - - matrixgame_patterns = ["Matrix-Game", "Skywork--Matrix-Game", "matrixgame"] - for pattern in matrixgame_patterns: - if pattern.lower() in pipeline_name_or_path.lower(): - return MatrixGame2_SamplingParam - - 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) - - pipeline_name = config["_class_name"] - - # If no match, try to use the fallback config - fallback_config = None - # Try to determine pipeline architecture for fallback - for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items(): - if detector(pipeline_name.lower()): - fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type) - break - - logger.warning( - "No match found for pipeline %s, using fallback sampling param %s.", - pipeline_name_or_path, fallback_config) - return fallback_config diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index d94727bd6..15ed7723c 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -366,9 +366,11 @@ class ImageTransformStage(DatasetStage): image = self.transform_topcrop(image) elif self.transform is not None: image = self.transform(image) + image = image.float() / 127.5 - 1.0 + else: + image = image.float() / 127.5 - 1.0 image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W] - image = image.float() / 127.5 - 1.0 batch.pixel_values = image return batch diff --git a/fastvideo/distributed/__init__.py b/fastvideo/distributed/__init__.py index f97ffd0dd..73f9c5b92 100644 --- a/fastvideo/distributed/__init__.py +++ b/fastvideo/distributed/__init__.py @@ -18,6 +18,7 @@ __all__ = [ "cleanup_dist_env_and_memory", "model_parallel_is_initialized", "maybe_init_distributed_environment_and_model_parallel", + "warmup_sequence_parallel_communication", # World group "get_world_group", diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index f29484c3d..9ef0d01b0 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -7,10 +7,17 @@ import torch.distributed from fastvideo.distributed.parallel_state import (get_sp_group, get_sp_parallel_rank, get_sp_world_size, - get_tp_group) + get_tp_group, + model_parallel_is_initialized) from fastvideo.distributed.utils import (unpad_sequence_tensor, compute_padding_for_sp, pad_sequence_tensor) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# Track if SP communication has been warmed up +_sp_warmup_done = False def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: @@ -101,3 +108,86 @@ def sequence_model_parallel_shard(input_: torch.Tensor, input_ = input_.movedim(0, dim) return input_, original_seq_len + + +def warmup_sequence_parallel_communication( + device: torch.device | None = None) -> None: + """Warmup NCCL communicators for sequence parallel all-to-all operations. + + The first NCCL collective operation is slow due to lazy communicator + initialization. This function runs dummy all-to-all operations to + trigger the initialization upfront, before the first real forward pass. + + Args: + device: Device to use for warmup tensors. If None, uses CUDA device 0. + """ + global _sp_warmup_done + + if _sp_warmup_done: + return + + if not model_parallel_is_initialized(): + return + + sp_world_size = get_sp_world_size() + if sp_world_size <= 1: + _sp_warmup_done = True + return + + if device is None: + device = torch.device("cuda") + + logger.info("Warming up sequence parallel communication (SP=%d)...", + sp_world_size) + + # Use small but representative tensor shapes for warmup + # Shape: [batch, seq_len, num_heads, head_dim] + # The all-to-all patterns used in attention: + # 1. scatter_dim=2 (heads), gather_dim=1 (seq) - before attention + # 2. scatter_dim=1 (seq), gather_dim=2 (heads) - after attention + batch_size = 1 + seq_len_per_rank = 16 # Small sequence per rank + num_heads = sp_world_size * 4 # Must be divisible by sp_world_size + head_dim = 64 + + # Create dummy tensor for warmup + dummy = torch.zeros(batch_size, + seq_len_per_rank, + num_heads, + head_dim, + device=device, + dtype=torch.bfloat16) + + # Warmup pattern 1: scatter heads, gather sequence (before attention) + _ = sequence_model_parallel_all_to_all_4D(dummy, + scatter_dim=2, + gather_dim=1) + + # Warmup pattern 2: scatter sequence, gather heads (after attention) + dummy2 = torch.zeros(batch_size, + seq_len_per_rank * sp_world_size, + num_heads // sp_world_size, + head_dim, + device=device, + dtype=torch.bfloat16) + _ = sequence_model_parallel_all_to_all_4D(dummy2, + scatter_dim=1, + gather_dim=2) + + # Warmup all-gather (used for replicated tokens) + dummy3 = torch.zeros(batch_size, + 8, + num_heads // sp_world_size, + head_dim, + device=device, + dtype=torch.bfloat16) + _ = sequence_model_parallel_all_gather(dummy3, dim=2) + + # Synchronize to ensure warmup completes + torch.cuda.synchronize(device) + + # Clean up + del dummy, dummy2, dummy3 + + _sp_warmup_done = True + logger.info("Sequence parallel communication warmup complete.") diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 33c06efcc..96fac3913 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -179,6 +179,7 @@ class FastVideoArgs: model_loaded: dict[str, bool] = field(default_factory=lambda: { "transformer": True, "vae": True, + "upsampler": True, }) override_text_encoder_safetensors: str | None = None # path to safetensors file for text encoder override diff --git a/fastvideo/image_processor.py b/fastvideo/image_processor.py index 3483631f5..bef175cee 100644 --- a/fastvideo/image_processor.py +++ b/fastvideo/image_processor.py @@ -193,3 +193,46 @@ class ImageProcessor: tensor = 2.0 * tensor - 1.0 return tensor + + def _preprocess_cosmos25( + self, + image: PIL.Image.Image | np.ndarray | torch.Tensor, + height: int, + width: int, + ) -> torch.Tensor: + """ + Cosmos-Predict2.5-style preprocessing for image2world: + - aspect-preserving resize (scale so both dims >= target) + - center crop to (height, width) + - normalize to [-1, 1] + + Returns: + torch.Tensor: (1, 3, height, width) in [-1, 1] + """ + if not isinstance(image, PIL.Image.Image): + # Reuse the generic preprocess for non-PIL inputs. + return self.preprocess(image, height=height, width=width) + + # Ensure RGB (official uses torchvision.to_tensor on PIL, which yields C=3 for RGB inputs) + image = image.convert("RGB") + + # Make sure target dims are multiples of VAE scale factor (FastVideo convention) + height = height - (height % self.vae_scale_factor) + width = width - (width % self.vae_scale_factor) + + orig_w, orig_h = image.size + # Scale so that resized image fully covers the target box (like official resize_input()). + scale = max(width / float(orig_w), height / float(orig_h)) + resized_w = int(np.ceil(scale * orig_w)) + resized_h = int(np.ceil(scale * orig_h)) + + image = image.resize((resized_w, resized_h), + resample=PIL.Image.Resampling.BILINEAR) + + # Center crop + left = max(0, (resized_w - width) // 2) + top = max(0, (resized_h - height) // 2) + image = image.crop((left, top, left + width, top + height)) + + image_np = np.array(image, dtype=np.float32) / 255.0 # [0,1] + return self._normalize_to_tensor(image_np) # -> [-1,1], (1,3,H,W) diff --git a/fastvideo/logger.py b/fastvideo/logger.py index 3896110e9..51d424e8f 100644 --- a/fastvideo/logger.py +++ b/fastvideo/logger.py @@ -27,7 +27,7 @@ RESET = '\033[0;0m' _warned_local_main_process = False _warned_main_process = False -_FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s " +_FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s.%(msecs)03d " "[%(filename)s:%(lineno)d] %(message)s") _DATE_FORMAT = "%m-%d %H:%M:%S" diff --git a/fastvideo/models/audio/ltx2_audio_vae.py b/fastvideo/models/audio/ltx2_audio_vae.py index fcc047f6c..52c18b8a0 100644 --- a/fastvideo/models/audio/ltx2_audio_vae.py +++ b/fastvideo/models/audio/ltx2_audio_vae.py @@ -1325,3 +1325,6 @@ def decode_audio( decoded_audio = audio_decoder(latent) decoded_audio = vocoder(decoded_audio).squeeze(0).float() return decoded_audio + +# Entry point for model registry +EntryClass = [LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder] diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 69b3e5430..e72de5325 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -676,4 +676,7 @@ class CausalWanTransformer3DModel(BaseDiT): # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) out.append(u) - return out \ No newline at end of file + return out + +# Entry point for model registry +EntryClass = CausalWanTransformer3DModel diff --git a/fastvideo/models/dits/cosmos.py b/fastvideo/models/dits/cosmos.py index fbbdbced0..4028307ec 100644 --- a/fastvideo/models/dits/cosmos.py +++ b/fastvideo/models/dits/cosmos.py @@ -723,4 +723,7 @@ class CosmosTransformer3DModel(BaseDiT): hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5) hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - return hidden_states \ No newline at end of file + return hidden_states + +# Entry point for model registry +EntryClass = CosmosTransformer3DModel diff --git a/fastvideo/models/dits/cosmos2_5.py b/fastvideo/models/dits/cosmos2_5.py index 7ce2ba0b7..6e4b09054 100644 --- a/fastvideo/models/dits/cosmos2_5.py +++ b/fastvideo/models/dits/cosmos2_5.py @@ -959,3 +959,5 @@ class Cosmos25Transformer3DModel(BaseDiT): return hidden_states +# Entry point for model registry +EntryClass = Cosmos25Transformer3DModel diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py index 26bea56d6..7396b29f2 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/models/dits/hunyuanvideo.py @@ -940,3 +940,6 @@ class FinalLayer(nn.Module): x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) x, _ = self.linear(x) return x + +# Entry point for model registry +EntryClass = HunyuanVideoTransformer3DModel diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index fa9564213..ac5ffb11d 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/fastvideo/models/dits/hunyuanvideo15.py @@ -852,4 +852,7 @@ class FinalLayer(nn.Module): scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1) x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) x, _ = self.linear(x) - return x \ No newline at end of file + return x + +# Entry point for model registry +EntryClass = HunyuanVideo15Transformer3DModel diff --git a/fastvideo/models/dits/hyworld/__init__.py b/fastvideo/models/dits/hyworld/__init__.py index 8f1dde52c..03516980a 100644 --- a/fastvideo/models/dits/hyworld/__init__.py +++ b/fastvideo/models/dits/hyworld/__init__.py @@ -22,3 +22,6 @@ __all__ = [ # Inference utilities (used by examples) "get_resolution_from_image", ] + +# Entry point for model registry +EntryClass = HYWorldTransformer3DModel diff --git a/fastvideo/models/dits/hyworld/hyworld.py b/fastvideo/models/dits/hyworld/hyworld.py index 776a5e2b2..81a7f6cc2 100644 --- a/fastvideo/models/dits/hyworld/hyworld.py +++ b/fastvideo/models/dits/hyworld/hyworld.py @@ -536,8 +536,8 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): temb_txt, freqs_cis, seq_attention_mask, - viewmats_seq, # hyworld - Ks_seq, # hyworld + viewmats_seq, + Ks_seq, ) else: for block in self.double_blocks: @@ -549,8 +549,8 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel): temb_txt, freqs_cis, seq_attention_mask, - viewmats=viewmats_seq, # hyworld - Ks=Ks_seq, # hyworld + viewmats_seq, + Ks_seq, ) diff --git a/fastvideo/models/dits/longcat.py b/fastvideo/models/dits/longcat.py index 440063691..4155370a8 100644 --- a/fastvideo/models/dits/longcat.py +++ b/fastvideo/models/dits/longcat.py @@ -1134,3 +1134,5 @@ class LongCatTransformer3DModel(CachableDiT): ) return x +# Entry point for model registry +EntryClass = LongCatTransformer3DModel diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index b8e0aecc7..df6426144 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -18,12 +18,23 @@ import torch.nn as nn from einops import rearrange, repeat from fastvideo.attention.backends.sdpa import SDPAMetadata -from fastvideo.attention.layer import LocalAttention +from fastvideo.attention.layer import DistributedAttention, LocalAttention from fastvideo.configs.models.dits import LTX2VideoConfig -from fastvideo.forward_context import get_forward_context, set_forward_context +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather, + sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_all_to_all_4D, + sequence_model_parallel_shard, +) +from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size +from fastvideo.distributed.utils import create_attention_mask_for_padding +from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context +from fastvideo.logger import init_logger from fastvideo.models.dits.base import CachableDiT from fastvideo.platforms import AttentionBackendEnum +logger = init_logger(__name__) + def get_timestep_embedding( timesteps: torch.Tensor, @@ -295,6 +306,12 @@ class VideoLatentPatchifier: output_shape: VideoLatentShape, device: Optional[torch.device] = None, ) -> torch.Tensor: + """Get patch grid bounds for RoPE computation. + + Args: + output_shape: Shape of the video latent tensor + device: Device to create tensors on + """ frames = output_shape.frames height = output_shape.height width = output_shape.width @@ -366,6 +383,12 @@ class AudioLatentPatchifier: output_shape: AudioLatentShape, device: Optional[torch.device] = None, ) -> torch.Tensor: + """Get patch grid bounds for audio RoPE computation. + + Args: + output_shape: Shape of the audio latent tensor + device: Device to create tensors on + """ start_timings = self._get_audio_latent_time_in_sec( self.shift, output_shape.frames + self.shift, @@ -515,6 +538,44 @@ def apply_ltx_rotary_emb( raise ValueError(f"Invalid rope type: {rope_type}") +def apply_ltx_rotary_emb_4d( + input_tensor: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, +) -> torch.Tensor: + """Apply LTX-2 rotary embeddings to a 4D tensor [B, H, T, D]. + + This is used for applying RoPE after all-to-all in distributed attention, + where the tensor is already in [B, H, T, D] format. + """ + cos_freqs, sin_freqs = freqs_cis + if rope_type == LTXRopeType.INTERLEAVED: + # For interleaved, cos/sin have shape [B, T, inner_dim] + # Need to reshape to [B, H, T, D] format + # Actually interleaved doesn't have per-head rotations, so we broadcast + t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) + t1, t2 = t_dup.unbind(dim=-1) + t_dup = torch.stack((-t2, t1), dim=-1) + input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") + return input_tensor * cos_freqs + input_tensor_rot * sin_freqs + if rope_type == LTXRopeType.SPLIT: + # For split, cos/sin already have shape [B, H, T, D/2] + # input_tensor is [B, H, T, D] + split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) + first_half_input = split_input[..., :1, :] + second_half_input = split_input[..., 1:, :] + + output = split_input * cos_freqs.unsqueeze(-2) + first_half_output = output[..., :1, :] + second_half_output = output[..., 1:, :] + + first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) + second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) + + return rearrange(output, "... d r -> ... (d r)") + raise ValueError(f"Invalid rope type: {rope_type}") + + def _apply_ltx_interleaved_rotary_emb( input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor ) -> torch.Tensor: @@ -918,6 +979,205 @@ class TransformerConfig: context_dim: int +class LTXDistributedAttention(DistributedAttention): + """LTX-2 specialized DistributedAttention that handles LTX-style RoPE internally.""" + + def __init__( + self, + num_heads: int, + head_size: int, + rope_type: LTXRopeType, + num_kv_heads: int | None = None, + softmax_scale: float | None = None, + causal: bool = False, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + super().__init__( + num_heads=num_heads, + head_size=head_size, + num_kv_heads=num_kv_heads, + softmax_scale=softmax_scale, + causal=causal, + supported_attention_backends=supported_attention_backends, + prefix=prefix, + **extra_impl_args, + ) + self.rope_type = rope_type + + @torch.compiler.disable + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + replicated_q: torch.Tensor | None = None, + replicated_k: torch.Tensor | None = None, + replicated_v: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward pass with LTX-2 style RoPE application. + + Args: + q: Query tensor [batch_size, seq_len, num_heads, head_dim] + k: Key tensor [batch_size, seq_len, num_heads, head_dim] + v: Value tensor [batch_size, seq_len, num_heads, head_dim] + replicated_q: Replicated query tensor for text tokens + replicated_k: Replicated key tensor + replicated_v: Replicated value tensor + attention_mask: Attention mask [batch_size, seq_len] + ltx_freqs_cis: LTX-2 style RoPE (cos, sin) with shape [B, H, T, D] + + Returns: + Tuple of output tensor and optional replicated output + """ + assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" + batch_size, seq_len, num_heads, head_dim = q.shape + local_rank = get_sp_parallel_rank() + world_size = get_sp_world_size() + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + # Stack QKV + qkv = torch.cat([q, k, v], dim=0) # [3*batch, seq_len, num_heads, head_dim] + + # Redistribute heads across sequence dimension + qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1) + + # After all-to-all, each rank has the full sequence but only a subset of heads + valid_seq_len = None + if attention_mask is not None: + valid_seq_len = (attention_mask[0] == 1).sum().item() + qkv = qkv[:, :valid_seq_len, :, :] + + # Apply LTX-2 style RoPE after all-to-all (when we have full sequence) + if ltx_freqs_cis is not None: + cos, sin = ltx_freqs_cis + heads_per_rank = num_heads // world_size + head_start = local_rank * heads_per_rank + head_end = head_start + heads_per_rank + # Slice to this rank's heads: [B, H/SP, T, D] + cos_local = cos[:, head_start:head_end, :valid_seq_len, :] + sin_local = sin[:, head_start:head_end, :valid_seq_len, :] + + # Apply RoPE to Q and K together (first 2*batch_size in dim 0) + qk_part = qkv[:batch_size * 2] + # Transpose to [2*B, H/SP, T, D] for LTX RoPE application + qk_part = qk_part.transpose(1, 2) + qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type) + # Transpose back to [2*B, T, H/SP, D] + qkv[:batch_size * 2] = qk_part.transpose(1, 2) + + # Apply backend-specific preprocess_qkv + qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) + + # Concatenate with replicated QKV if provided + if replicated_q is not None: + assert replicated_k is not None and replicated_v is not None + replicated_qkv = torch.cat( + [replicated_q, replicated_k, replicated_v], + dim=0) # [3, seq_len, num_heads, head_dim] + heads_per_rank = num_heads // world_size + replicated_qkv = replicated_qkv[:, :, local_rank * + heads_per_rank:(local_rank + 1) * + heads_per_rank] + qkv = torch.cat([qkv, replicated_qkv], dim=1) + + q, k, v = qkv.chunk(3, dim=0) + + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + + # Redistribute back if using sequence parallelism + replicated_output = None + if replicated_q is not None: + split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len + replicated_output = output[:, split_idx:] + output = output[:, :split_idx] + replicated_output = sequence_model_parallel_all_gather( + replicated_output.contiguous(), dim=2) + + # Apply backend-specific postprocess_output + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + + if attention_mask is not None: + pad_len = (attention_mask[0] == 0).sum().item() + output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len)) + + output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2) + + return output, replicated_output + + +class LTXLocalAttention(LocalAttention): + """LTX-2 specialized LocalAttention that handles LTX-style RoPE internally.""" + + def __init__( + self, + num_heads: int, + head_size: int, + rope_type: LTXRopeType, + num_kv_heads: int | None = None, + softmax_scale: float | None = None, + causal: bool = False, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + **extra_impl_args, + ) -> None: + super().__init__( + num_heads=num_heads, + head_size=head_size, + num_kv_heads=num_kv_heads, + softmax_scale=softmax_scale, + causal=causal, + supported_attention_backends=supported_attention_backends, + **extra_impl_args, + ) + self.rope_type = rope_type + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, + ltx_k_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Forward pass with LTX-2 style RoPE application. + + Args: + q: Query tensor [batch_size, seq_len, num_heads, head_dim] + k: Key tensor [batch_size, seq_len, num_heads, head_dim] + v: Value tensor [batch_size, seq_len, num_heads, head_dim] + ltx_freqs_cis: LTX-2 style RoPE (cos, sin) for Q with shape [B, H, T, D] + ltx_k_freqs_cis: LTX-2 style RoPE (cos, sin) for K (if different from Q) + + Returns: + Output tensor after local attention + """ + assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors" + + forward_context: ForwardContext = get_forward_context() + ctx_attn_metadata = forward_context.attn_metadata + + # Apply LTX-2 style RoPE + if ltx_freqs_cis is not None: + # Use separate K RoPE if provided (for cross-attention), otherwise use same as Q + k_freqs = ltx_k_freqs_cis if ltx_k_freqs_cis is not None else ltx_freqs_cis + # Transpose to [B, H, T, D] for RoPE application + q = q.transpose(1, 2) + k = k.transpose(1, 2) + q = apply_ltx_rotary_emb_4d(q, ltx_freqs_cis, self.rope_type) + k = apply_ltx_rotary_emb_4d(k, k_freqs, self.rope_type) + # Transpose back to [B, T, H, D] + q = q.transpose(1, 2) + k = k.transpose(1, 2) + + output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + return output + + class LTXSelfAttention(nn.Module): """LTX-2 attention block with RMSNorm + FastVideo LocalAttention.""" def __init__( @@ -945,17 +1205,19 @@ class LTXSelfAttention(nn.Module): self.to_v = nn.Linear(context_dim, inner_dim, bias=True) self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) - self.attn = LocalAttention( + self.attn = LTXLocalAttention( num_heads=heads, head_size=dim_head, + rope_type=rope_type, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, ) - self.attn_masked = LocalAttention( + self.attn_masked = LTXLocalAttention( num_heads=heads, head_size=dim_head, + rope_type=rope_type, dropout_rate=0, softmax_scale=None, causal=False, @@ -978,15 +1240,13 @@ class LTXSelfAttention(nn.Module): q = self.q_norm(q) k = self.k_norm(k) - if pe is not None: - q = apply_ltx_rotary_emb(q, pe, self.rope_type) - k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) - + # RoPE is applied inside LTXLocalAttention b, q_len, _ = q.shape k_len = k.shape[1] q = q.view(b, q_len, self.heads, self.dim_head) k = k.view(b, k_len, self.heads, self.dim_head) v = v.view(b, k_len, self.heads, self.dim_head) + if mask is not None: if mask.ndim == 2: mask = mask.unsqueeze(0) @@ -1008,9 +1268,93 @@ class LTXSelfAttention(nn.Module): attn_metadata=attn_metadata, forward_batch=forward_batch, ): - out = self.attn_masked(q, k, v) + out = self.attn_masked(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) else: - out = self.attn(q, k, v) + out = self.attn(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe) + out = out.reshape(b, q_len, -1) + return self.to_out(out) + + +class LTXDistributedSelfAttention(nn.Module): + """LTX-2 attention block with RMSNorm + LTXDistributedAttention for SP.""" + + def __init__( + self, + query_dim: int, + context_dim: int | None, + heads: int, + dim_head: int, + norm_eps: float, + rope_type: LTXRopeType, + supported_attention_backends: tuple[AttentionBackendEnum, ...], + prefix: str = "", + ) -> None: + super().__init__() + inner_dim = dim_head * heads + context_dim = query_dim if context_dim is None else context_dim + + self.heads = heads + self.dim_head = dim_head + self.rope_type = rope_type + + self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) + self.to_q = nn.Linear(query_dim, inner_dim, bias=True) + self.to_k = nn.Linear(context_dim, inner_dim, bias=True) + self.to_v = nn.Linear(context_dim, inner_dim, bias=True) + self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()) + + self.attn = LTXDistributedAttention( + num_heads=heads, + head_size=dim_head, + rope_type=rope_type, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + pe: tuple[torch.Tensor, torch.Tensor] | None = None, + k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward pass for distributed self-attention. + + Args: + x: Query tensor [B, L, C] + context: Key/Value tensor [B, L, C], or None for self-attention + pe: Rotary position embeddings for Q (cos, sin) with shape [B, H, T_full, D] + NOTE: For SP, this must be the FULL sequence RoPE, not sharded! + k_pe: Rotary position embeddings for K (cos, sin), or None to use pe + attention_mask: Attention mask for padding [B, padded_seq_len] + """ + q = self.to_q(x) + context = x if context is None else context + k = self.to_k(context) + v = self.to_v(context) + + q = self.q_norm(q) + k = self.k_norm(k) + + # RoPE is applied inside LTXDistributedAttention AFTER the all-to-all, + # when each rank has the full sequence (but subset of heads). + b, q_len, _ = q.shape + k_len = k.shape[1] + q = q.view(b, q_len, self.heads, self.dim_head) + k = k.view(b, k_len, self.heads, self.dim_head) + v = v.view(b, k_len, self.heads, self.dim_head) + + # Pass full RoPE to distributed attention - it will apply after all-to-all + # and slice to this rank's heads + out, _ = self.attn( + q, k, v, + attention_mask=attention_mask, + ltx_freqs_cis=pe, + ) + out = out.reshape(b, q_len, -1) return self.to_out(out) @@ -1025,12 +1369,31 @@ class BasicAVTransformerBlock(torch.nn.Module): audio: TransformerConfig | None = None, rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, norm_eps: float = 1e-6, + use_distributed_attention: bool = False, + prefix: str = "", ): super().__init__() self.idx = idx + self.use_distributed_attention = use_distributed_attention + + # Choose attention class based on SP mode + # Self-attention and audio-video cross-attention use DistributedAttention when SP > 1 + # Text cross-attention uses LocalAttention (text embeddings are replicated) + SelfAttnCls = LTXDistributedSelfAttention if use_distributed_attention else LTXSelfAttention + CrossAttnCls = LTXSelfAttention # Text cross-attention is always local if video is not None: - self.attn1 = LTXSelfAttention( + # Video self-attention - use distributed when SP > 1 + self.attn1 = SelfAttnCls( + query_dim=video.dim, + context_dim=None, + heads=video.heads, + dim_head=video.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + prefix=f"{prefix}.blocks.{idx}.attn1" if use_distributed_attention else "", + ) if use_distributed_attention else LTXSelfAttention( query_dim=video.dim, context_dim=None, heads=video.heads, @@ -1039,7 +1402,8 @@ class BasicAVTransformerBlock(torch.nn.Module): rope_type=rope_type, supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) - self.attn2 = LTXSelfAttention( + # Text cross-attention - always local (text is replicated) + self.attn2 = CrossAttnCls( query_dim=video.dim, context_dim=video.context_dim, heads=video.heads, @@ -1052,7 +1416,17 @@ class BasicAVTransformerBlock(torch.nn.Module): self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim)) if audio is not None: - self.audio_attn1 = LTXSelfAttention( + # Audio self-attention - use distributed when SP > 1 + self.audio_attn1 = SelfAttnCls( + query_dim=audio.dim, + context_dim=None, + heads=audio.heads, + dim_head=audio.d_head, + norm_eps=norm_eps, + rope_type=rope_type, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + prefix=f"{prefix}.blocks.{idx}.audio_attn1" if use_distributed_attention else "", + ) if use_distributed_attention else LTXSelfAttention( query_dim=audio.dim, context_dim=None, heads=audio.heads, @@ -1061,7 +1435,8 @@ class BasicAVTransformerBlock(torch.nn.Module): rope_type=rope_type, supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) - self.audio_attn2 = LTXSelfAttention( + # Text cross-attention - always local (text is replicated) + self.audio_attn2 = CrossAttnCls( query_dim=audio.dim, context_dim=audio.context_dim, heads=audio.heads, @@ -1074,6 +1449,8 @@ class BasicAVTransformerBlock(torch.nn.Module): self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim)) if audio is not None and video is not None: + # Audio-to-video cross-attention + # Uses local attention - context is gathered from all SP ranks in forward() self.audio_to_video_attn = LTXSelfAttention( query_dim=video.dim, context_dim=audio.dim, @@ -1083,6 +1460,8 @@ class BasicAVTransformerBlock(torch.nn.Module): rope_type=rope_type, supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), ) + # Video-to-audio cross-attention + # Uses local attention - context is gathered from all SP ranks in forward() self.video_to_audio_attn = LTXSelfAttention( query_dim=audio.dim, context_dim=video.dim, @@ -1131,7 +1510,17 @@ class BasicAVTransformerBlock(torch.nn.Module): self, video: TransformerArgs | None, audio: TransformerArgs | None, + video_attention_mask: torch.Tensor | None = None, + audio_attention_mask: torch.Tensor | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: + """Forward pass for transformer block. + + Args: + video: Video transformer args + audio: Audio transformer args + video_attention_mask: SP padding attention mask for video [B, padded_seq_len] + audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len] + """ vx = video.x if video is not None else None ax = audio.x if audio is not None else None @@ -1145,7 +1534,12 @@ class BasicAVTransformerBlock(torch.nn.Module): self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3) ) norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa - vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa + # Self-attention: pass SP attention mask for distributed attention + if self.use_distributed_attention: + vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings, attention_mask=video_attention_mask) * vgate_msa + else: + vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa + # Text cross-attention: no SP mask needed (text is replicated, uses local attention) vx = vx + self.attn2( torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps), context=video.context, @@ -1157,7 +1551,12 @@ class BasicAVTransformerBlock(torch.nn.Module): self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) ) norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa - ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa + # Self-attention: pass SP attention mask for distributed attention + if self.use_distributed_attention: + ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings, attention_mask=audio_attention_mask) * agate_msa + else: + ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa + # Text cross-attention: no SP mask needed (text is replicated, uses local attention) ax = ax + self.audio_attn2( torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps), context=audio.context, @@ -1197,28 +1596,91 @@ class BasicAVTransformerBlock(torch.nn.Module): if run_a2v: vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v - vx = vx + ( - self.audio_to_video_attn( - vx_scaled, - context=ax_scaled, - pe=video.cross_positional_embeddings, - k_pe=audio.cross_positional_embeddings, + # Audio-to-video cross-attention: Q from video (sharded), K/V from audio + # For SP, gather audio context from all ranks before cross-attention + if self.use_distributed_attention: + # Gather audio context from all SP ranks + ax_context = sequence_model_parallel_all_gather(ax_scaled, dim=1) + # Video Q: video tokens divide evenly (adjusted upstream), so simple slicing works + video_q_pe = None + if video.cross_positional_embeddings is not None: + sp_rank = get_sp_parallel_rank() + local_seq_len = vx_scaled.shape[1] + pe_tuple = video.cross_positional_embeddings + start_idx = sp_rank * local_seq_len + end_idx = start_idx + local_seq_len + video_q_pe = tuple(pe[:, :, start_idx:end_idx, :] for pe in pe_tuple) + # Audio K: gathered context may have padding, trim to original length + audio_k_pe = audio.cross_positional_embeddings + if audio_k_pe is not None: + original_audio_len = audio_k_pe[0].shape[2] + ax_context = ax_context[:, :original_audio_len, :] + vx = vx + ( + self.audio_to_video_attn( + vx_scaled, + context=ax_context, + pe=video_q_pe, + k_pe=audio_k_pe, + ) + * gate_out_a2v + ) + else: + ax_context = ax_scaled + video_q_pe = video.cross_positional_embeddings + audio_k_pe = audio.cross_positional_embeddings + vx = vx + ( + self.audio_to_video_attn( + vx_scaled, + context=ax_context, + pe=video_q_pe, + k_pe=audio_k_pe, + ) + * gate_out_a2v ) - * gate_out_a2v - ) if run_v2a: ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a - ax = ax + ( - self.video_to_audio_attn( - ax_scaled, - context=vx_scaled, - pe=audio.cross_positional_embeddings, - k_pe=video.cross_positional_embeddings, + # Video-to-audio cross-attention: Q from audio, K/V from video + # For SP, gather both modalities and run cross-attention on full sequences + # (audio is short, so no need to keep it sharded for cross-attention) + if self.use_distributed_attention: + # Gather both modalities to full sequences + ax_full = sequence_model_parallel_all_gather(ax_scaled, dim=1) + vx_context = sequence_model_parallel_all_gather(vx_scaled, dim=1) + # Use full positional embeddings directly + audio_q_pe = audio.cross_positional_embeddings + video_k_pe = video.cross_positional_embeddings + # Trim gathered tensors to original lengths (may have padding from sharding) + if audio_q_pe is not None: + original_audio_len = audio_q_pe[0].shape[2] + ax_full = ax_full[:, :original_audio_len, :] + if video_k_pe is not None: + original_video_len = video_k_pe[0].shape[2] + vx_context = vx_context[:, :original_video_len, :] + # Run cross-attention on full audio + v2a_out_full = self.video_to_audio_attn( + ax_full, + context=vx_context, + pe=audio_q_pe, + k_pe=video_k_pe, + ) + # Shard the output back to match ax's sharding + v2a_out, _ = sequence_model_parallel_shard(v2a_out_full, dim=1) + ax = ax + v2a_out * gate_out_v2a + else: + vx_context = vx_scaled + audio_q_pe = audio.cross_positional_embeddings + video_k_pe = video.cross_positional_embeddings + ax = ax + ( + self.video_to_audio_attn( + ax_scaled, + context=vx_context, + pe=audio_q_pe, + k_pe=video_k_pe, + ) + * gate_out_v2a ) - * gate_out_v2a - ) if run_vx: vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( @@ -1289,6 +1751,8 @@ class LTXModel(torch.nn.Module): av_ca_timestep_scale_multiplier: int = 1, rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, double_precision_rope: bool = False, + use_distributed_attention: bool = False, + prefix: str = "", ): super().__init__() self._enable_gradient_checkpointing = False @@ -1340,6 +1804,8 @@ class LTXModel(torch.nn.Module): audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0, audio_cross_attention_dim=audio_cross_attention_dim, norm_eps=norm_eps, + use_distributed_attention=use_distributed_attention, + prefix=prefix, ) def _init_video( @@ -1469,6 +1935,8 @@ class LTXModel(torch.nn.Module): audio_attention_head_dim: int, audio_cross_attention_dim: int, norm_eps: float, + use_distributed_attention: bool = False, + prefix: str = "", ) -> None: video_config = ( TransformerConfig( @@ -1490,6 +1958,7 @@ class LTXModel(torch.nn.Module): if self.model_type.is_audio_enabled() else None ) + self.use_distributed_attention = use_distributed_attention self.transformer_blocks = torch.nn.ModuleList( [ BasicAVTransformerBlock( @@ -1498,6 +1967,8 @@ class LTXModel(torch.nn.Module): audio=audio_config, rope_type=self.rope_type, norm_eps=norm_eps, + use_distributed_attention=use_distributed_attention, + prefix=prefix, ) for idx in range(num_layers) ] @@ -1507,9 +1978,16 @@ class LTXModel(torch.nn.Module): self, video: TransformerArgs | None, audio: TransformerArgs | None, + video_attention_mask: torch.Tensor | None = None, + audio_attention_mask: torch.Tensor | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: for block in self.transformer_blocks: - video, audio = block(video=video, audio=audio) + video, audio = block( + video=video, + audio=audio, + video_attention_mask=video_attention_mask, + audio_attention_mask=audio_attention_mask, + ) return video, audio def _process_output( @@ -1533,7 +2011,17 @@ class LTXModel(torch.nn.Module): self, video: Modality | None, audio: Modality | None, + video_attention_mask: torch.Tensor | None = None, + audio_attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Forward pass through the LTX model. + + Args: + video: Video modality input + audio: Audio modality input + video_attention_mask: SP padding attention mask for video [B, padded_seq_len] + audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len] + """ if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": _debug_block_log_line( "fastvideo:patchify_proj" @@ -1551,7 +2039,12 @@ class LTXModel(torch.nn.Module): audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None _debug_transformer_args("fastvideo:prep_video", video_args) _debug_transformer_args("fastvideo:prep_audio", audio_args) - video_out, audio_out = self._process_transformer_blocks(video_args, audio_args) + video_out, audio_out = self._process_transformer_blocks( + video_args, + audio_args, + video_attention_mask=video_attention_mask, + audio_attention_mask=audio_attention_mask, + ) vx = ( self._process_output( @@ -1588,6 +2081,24 @@ class LTX2Transformer3DModel(CachableDiT): arch = config.arch_config + # Get SP world size for distributed attention + sp_world_size = get_sp_world_size() + use_distributed_attention = sp_world_size > 1 + + # Validate that attention heads are divisible by SP world size + if sp_world_size > 1: + assert arch.num_attention_heads % sp_world_size == 0, ( + f"The number of video attention heads ({arch.num_attention_heads}) " + f"must be divisible by the sequence parallel size ({sp_world_size})" + ) + assert arch.audio_num_attention_heads % sp_world_size == 0, ( + f"The number of audio attention heads ({arch.audio_num_attention_heads}) " + f"must be divisible by the sequence parallel size ({sp_world_size})" + ) + logger.info( + f"LTX2 sequence parallelism enabled with SP world size {sp_world_size}" + ) + model_type = LTXModelType.AudioVideo self.model = LTXModel( model_type=model_type, @@ -1612,6 +2123,8 @@ class LTX2Transformer3DModel(CachableDiT): audio_cross_attention_dim=arch.audio_cross_attention_dim, audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos, av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier, + use_distributed_attention=use_distributed_attention, + prefix=config.prefix, ) self.patchifier = VideoLatentPatchifier(patch_size=arch.patch_size[1]) @@ -1627,6 +2140,7 @@ class LTX2Transformer3DModel(CachableDiT): self.hidden_size = arch.num_attention_heads * arch.attention_head_dim self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.num_channels_latents + self._logged_attention_mask = False if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1": detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "") @@ -1702,10 +2216,12 @@ class LTX2Transformer3DModel(CachableDiT): if isinstance(encoder_hidden_states, list): encoder_hidden_states = encoder_hidden_states[0] - video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) - positions = self.patchifier.get_patch_grid_bounds( - video_shape, device=hidden_states.device - ) + # Get SP parameters + sp_world_size = get_sp_world_size() + sp_rank = get_sp_parallel_rank() + batch_size = hidden_states.shape[0] + + # Get fps for position computation fps = None try: forward_ctx = get_forward_context() @@ -1717,17 +2233,60 @@ class LTX2Transformer3DModel(CachableDiT): fps_value = fps_value[0] if fps_value else None if fps_value is not None: fps = float(fps_value) + + # Patchify video latents + video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape) + latents = self.patchifier.patchify(hidden_states) + + # Shard video latents and timestep across SP ranks + video_original_seq_len = latents.shape[1] + video_attention_mask = None + video_timestep = timestep + if sp_world_size > 1: + latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1) + # Shard timestep along sequence dimension (timestep has shape [batch, seq_len]) + video_timestep, _ = sequence_model_parallel_shard(timestep, dim=1) + # Create attention mask for padded tokens + current_seq_len = latents.shape[1] + padded_seq_len = current_seq_len * sp_world_size + if padded_seq_len > video_original_seq_len: + if not self._logged_attention_mask: + logger.info( + f"Video padding applied, original seq len: {video_original_seq_len}, " + f"padded seq len: {padded_seq_len}" + ) + self._logged_attention_mask = True + video_attention_mask = create_attention_mask_for_padding( + seq_len=video_original_seq_len, + padded_seq_len=padded_seq_len, + batch_size=batch_size, + device=hidden_states.device, + ) + + # Compute RoPE positions for the FULL sequence (before sharding) + # This is necessary because sequence sharding may not align to frame boundaries. + # The full RoPE will be passed to DistributedAttention which applies it after + # the all-to-all (when each rank has the full sequence but subset of heads). + positions = self.patchifier.get_patch_grid_bounds( + video_shape, device=hidden_states.device + ) positions = _get_pixel_coords( positions, DEFAULT_LTX2_SCALE_FACTORS, fps=fps, causal_fix=True, ).to(hidden_states.dtype) - latents = self.patchifier.patchify(hidden_states) + + # Pad positions to match padded sequence length for SP cross-attention + if sp_world_size > 1 and padded_seq_len > video_original_seq_len: + padding_needed = padded_seq_len - video_original_seq_len + last_pos = positions[:, :, -1:, :].expand(-1, -1, padding_needed, -1) + positions = torch.cat([positions, last_pos], dim=2) + video_modality = Modality( enabled=True, latent=latents, - timesteps=timestep, + timesteps=video_timestep, positions=positions, context=encoder_hidden_states, context_mask=encoder_attention_mask, @@ -1745,19 +2304,45 @@ class LTX2Transformer3DModel(CachableDiT): f"latent_head={video_head} " f"latent_checksum={video_checksum:.6f}" ) + + # Process audio modality if provided audio_modality = None audio_shape = None + audio_original_seq_len = 0 + audio_attention_mask = None + if audio_hidden_states is not None and audio_encoder_hidden_states is not None and audio_timestep is not None: - audio_shape = AudioLatentShape.from_torch_shape( - audio_hidden_states.shape) + audio_shape = AudioLatentShape.from_torch_shape(audio_hidden_states.shape) + audio_latents = self.audio_patchifier.patchify(audio_hidden_states) + + # Shard audio latents and timestep across SP ranks + audio_original_seq_len = audio_latents.shape[1] + sharded_audio_timestep = audio_timestep + if sp_world_size > 1: + audio_latents, audio_original_seq_len = sequence_model_parallel_shard(audio_latents, dim=1) + # Shard audio timestep along sequence dimension + sharded_audio_timestep, _ = sequence_model_parallel_shard(audio_timestep, dim=1) + # Create attention mask for padded tokens + audio_current_seq_len = audio_latents.shape[1] + audio_padded_seq_len = audio_current_seq_len * sp_world_size + if audio_padded_seq_len > audio_original_seq_len: + audio_attention_mask = create_attention_mask_for_padding( + seq_len=audio_original_seq_len, + padded_seq_len=audio_padded_seq_len, + batch_size=batch_size, + device=audio_hidden_states.device, + ) + + # Compute audio RoPE positions for the FULL sequence (before sharding) + # Same as video: full RoPE is applied after all-to-all in DistributedAttention audio_positions = self.audio_patchifier.get_patch_grid_bounds( - audio_shape, device=audio_hidden_states.device) - audio_latents = self.audio_patchifier.patchify( - audio_hidden_states) + audio_shape, device=audio_hidden_states.device + ) + audio_modality = Modality( enabled=True, latent=audio_latents, - timesteps=audio_timestep, + timesteps=sharded_audio_timestep, positions=audio_positions, context=audio_encoder_hidden_states, context_mask=audio_encoder_attention_mask, @@ -1775,10 +2360,16 @@ class LTX2Transformer3DModel(CachableDiT): f"latent_head={audio_head} " f"latent_checksum={audio_checksum:.6f}" ) + + # Run transformer with attention masks video_out, audio_out = self.model( video=video_modality, audio=audio_modality, + video_attention_mask=video_attention_mask, + audio_attention_mask=audio_attention_mask, ) + + # Denoised prediction if video_out is not None and video_modality is not None: video_out = _to_denoised( video_modality.latent, @@ -1791,6 +2382,20 @@ class LTX2Transformer3DModel(CachableDiT): audio_out, audio_modality.timesteps, ) + + # Gather and unpad video output + if sp_world_size > 1 and video_out is not None: + video_out = sequence_model_parallel_all_gather_with_unpad( + video_out, video_original_seq_len, dim=1 + ) + + # Gather and unpad audio output + if sp_world_size > 1 and audio_out is not None: + audio_out = sequence_model_parallel_all_gather_with_unpad( + audio_out, audio_original_seq_len, dim=1 + ) + + # Unpatchify video_out = self.patchifier.unpatchify( video_out, output_shape=video_shape) if audio_out is None or audio_shape is None: @@ -1798,3 +2403,6 @@ class LTX2Transformer3DModel(CachableDiT): audio_out = self.audio_patchifier.unpatchify( audio_out, output_shape=audio_shape) return video_out, audio_out + +# Entry point for model registry +EntryClass = LTX2Transformer3DModel diff --git a/fastvideo/models/dits/matrixgame/__init__.py b/fastvideo/models/dits/matrixgame/__init__.py index 0da9c954c..c03b234a8 100644 --- a/fastvideo/models/dits/matrixgame/__init__.py +++ b/fastvideo/models/dits/matrixgame/__init__.py @@ -9,3 +9,6 @@ __all__ = [ "CausalMatrixGameTransformerBlock", "ActionModule", ] + +# Entry point for model registry +EntryClass = [MatrixGameWanModel, CausalMatrixGameWanModel] diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index 6c06babae..e7a45413b 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -5,7 +5,7 @@ from einops import rearrange import torch import torch.nn as nn import math -from torch.nn.attention.flex_attention import flex_attention +from torch.nn.attention.flex_attention import flex_attention, BlockMask from fastvideo.attention import LocalAttention from fastvideo.layers.linear import ReplicatedLinear @@ -19,7 +19,8 @@ from fastvideo.platforms import AttentionBackendEnum DISABLE_COMPILE = False flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" +) def _get_nd_rotary_pos_embed_matrixgame( @@ -49,15 +50,15 @@ def _apply_rotary_emb_qk( start_offset: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: seq_len = xq.shape[1] - + # Slice frequencies based on offset - cos = freqs_cos[start_offset:start_offset + seq_len] # [S, D] - sin = freqs_sin[start_offset:start_offset + seq_len] # [S, D] - + cos = freqs_cos[start_offset : start_offset + seq_len] # [S, D] + sin = freqs_sin[start_offset : start_offset + seq_len] # [S, D] + # Move to device cos = cos.to(xq.device) sin = sin.to(xq.device) - + # Convert from [S, D] (interleaved) back to [S, D/2] cos_half = cos[:, ::2] # [S, D/2] sin_half = sin[:, ::2] # [S, D/2] @@ -67,41 +68,186 @@ def _apply_rotary_emb_qk( xq_out = _apply_rotary_emb(xq, cos_half, sin_half, is_neox_style=False) xk_out = _apply_rotary_emb(xk, cos_half, sin_half, is_neox_style=False) - + return xq_out, xk_out + +def _padding_q_k_v(tensor: torch.Tensor, padded_length: int) -> torch.Tensor: + return torch.cat( + [ + tensor, + torch.zeros( + [ + tensor.shape[0], + padded_length, + tensor.shape[2], + tensor.shape[3], + ], + device=tensor.device, + dtype=tensor.dtype, + ), + ], + dim=1, + ) + + +def _update_kv_cache_and_attend( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kv_cache: dict[str, torch.Tensor], + attn_layer: LocalAttention, + start_frame: int, + num_frame_per_block: int, + local_attn_size: int, + *, + use_k_for_num_tokens: bool = False, + store_first_only: bool = False, + repeat_factor: int | None = None, +) -> torch.Tensor: + """ + Update KV cache with new tokens and perform attention with cached values. + + Args: + q: Query tensor + k: Key tensor + v: Value tensor + kv_cache: Dictionary containing cache tensors and indices + attn_layer: Attention layer to use + start_frame: Starting frame index + num_frame_per_block: Number of frames per block + local_attn_size: Maximum attention window size + use_k_for_num_tokens: If True, use k.shape[1] for num_new_tokens, else use q.shape[1] + store_first_only: If True, only store k[:1] and v[:1] in cache (for keyboard with rope) + repeat_factor: If provided, repeat cached k,v by this factor when retrieving (for keyboard with rope) + + Returns: + Attention output tensor + """ + current_start = start_frame + current_end = current_start + ( + k.shape[1] if use_k_for_num_tokens else q.shape[1] + ) + + assert ( + k.shape[1] if use_k_for_num_tokens else q.shape[1] + ) == num_frame_per_block + + sink_size = 0 + max_attention_size = local_attn_size + sink_tokens = sink_size * 1 + kv_cache_size = kv_cache["k"].shape[1] + num_new_tokens = k.shape[1] if use_k_for_num_tokens else q.shape[1] + + # Check if we need to evict tokens + if (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size + ): + num_evicted_tokens = ( + num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + ) + num_rolled_tokens = ( + kv_cache["local_end_index"].item() + - num_evicted_tokens + - sink_tokens + ) + # Roll k cache + kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = ( + kv_cache["k"][ + :, + sink_tokens + num_evicted_tokens : sink_tokens + + num_evicted_tokens + + num_rolled_tokens, + ].clone() + ) + # Roll v cache + kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = ( + kv_cache["v"][ + :, + sink_tokens + num_evicted_tokens : sink_tokens + + num_evicted_tokens + + num_rolled_tokens, + ].clone() + ) + # Calculate indices with eviction adjustment + local_end_index = ( + kv_cache["local_end_index"].item() + + current_end + - kv_cache["global_end_index"].item() + - num_evicted_tokens + ) + local_start_index = local_end_index - num_new_tokens + else: + # Calculate indices without eviction + local_end_index = ( + kv_cache["local_end_index"].item() + + current_end + - kv_cache["global_end_index"].item() + ) + local_start_index = local_end_index - num_new_tokens + + # Store new k, v in cache + if store_first_only: + kv_cache["k"][:, local_start_index:local_end_index] = k[:1] + kv_cache["v"][:, local_start_index:local_end_index] = v[:1] + else: + kv_cache["k"][:, local_start_index:local_end_index] = k + kv_cache["v"][:, local_start_index:local_end_index] = v + + # Retrieve from cache and perform attention + cache_start = max(0, local_end_index - max_attention_size) + cached_k = kv_cache["k"][:, cache_start:local_end_index] + cached_v = kv_cache["v"][:, cache_start:local_end_index] + + if repeat_factor is not None: + cached_k = cached_k.repeat(repeat_factor, 1, 1, 1) + cached_v = cached_v.repeat(repeat_factor, 1, 1, 1) + + attn = attn_layer(q, cached_k, cached_v) + + # Update indices + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) + + return attn + + class ActionModule(nn.Module): """ action module from https://arxiv.org/pdf/2501.08325 """ def __init__( - self, + self, mouse_dim_in: int = 2, keyboard_dim_in: int = 6, hidden_size: int = 128, img_hidden_size: int = 1536, keyboard_hidden_dim: int = 1024, mouse_hidden_dim: int = 1024, - vae_time_compression_ratio: int = 4, + vae_time_compression_ratio: int = 4, windows_size: int = 3, heads_num: int = 16, patch_size: list | None = None, qk_norm: bool = True, qkv_bias: bool = False, rope_dim_list: list | None = None, - rope_theta = 256, + rope_theta=256, mouse_qk_dim_list: list | None = None, - enable_mouse = True, - enable_keyboard = True, - local_attn_size = 6, + enable_mouse=True, + enable_keyboard=True, + local_attn_size=6, blocks: list | None = None, ): super().__init__() # Initialize mutable defaults patch_size = patch_size if patch_size is not None else [1, 2, 2] - rope_dim_list = rope_dim_list if rope_dim_list is not None else [8, 28, 28] - mouse_qk_dim_list = mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28] + rope_dim_list = ( + rope_dim_list if rope_dim_list is not None else [8, 28, 28] + ) + mouse_qk_dim_list = ( + mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28] + ) blocks = blocks if blocks is not None else [] self.local_attn_size = local_attn_size self.enable_mouse = enable_mouse @@ -113,7 +259,7 @@ class ActionModule(nn.Module): self.keyboard_embed = nn.Sequential( nn.Linear(keyboard_dim_in, hidden_size, bias=True), nn.SiLU(), - nn.Linear(hidden_size, hidden_size, bias=True) + nn.Linear(hidden_size, hidden_size, bias=True), ) self.mouse_qk_dim_list = mouse_qk_dim_list @@ -121,58 +267,77 @@ class ActionModule(nn.Module): if self.enable_mouse: c = mouse_hidden_dim self.mouse_mlp = nn.Sequential( - nn.Linear(mouse_dim_in * vae_time_compression_ratio * windows_size + img_hidden_size, c, bias=True), + nn.Linear( + mouse_dim_in * vae_time_compression_ratio * windows_size + + img_hidden_size, + c, + bias=True, + ), nn.GELU(approximate="tanh"), nn.Linear(c, c), FP32LayerNorm(c, elementwise_affine=True), ) - + head_dim = c // heads_num - self.t_qkv = ReplicatedLinear(c, c*3, bias=qkv_bias) + self.t_qkv = ReplicatedLinear(c, c * 3, bias=qkv_bias) self.img_attn_q_norm = ( - RMSNorm(head_dim, eps=1e-6) - if qk_norm - else nn.Identity() + RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() ) self.img_attn_k_norm = ( - RMSNorm(head_dim, eps=1e-6) - if qk_norm - else nn.Identity() + RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() + ) + self.proj_mouse = ReplicatedLinear( + c, img_hidden_size, bias=qkv_bias ) - self.proj_mouse = ReplicatedLinear(c, img_hidden_size, bias=qkv_bias) if self.enable_keyboard: head_dim_key = keyboard_hidden_dim // heads_num self.key_attn_q_norm = ( - RMSNorm(head_dim_key, eps=1e-6) - if qk_norm - else nn.Identity() + RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() ) self.key_attn_k_norm = ( - RMSNorm(head_dim_key, eps=1e-6) - if qk_norm - else nn.Identity() + RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() ) - - self.mouse_attn_q = ReplicatedLinear(img_hidden_size, keyboard_hidden_dim, bias=qkv_bias) - self.keyboard_attn_kv = ReplicatedLinear(hidden_size * windows_size * vae_time_compression_ratio, keyboard_hidden_dim * 2, bias=qkv_bias) - self.proj_keyboard = ReplicatedLinear(keyboard_hidden_dim, img_hidden_size, bias=qkv_bias) - self.mouse_attn_layer = LocalAttention( - num_heads=heads_num, - head_size=mouse_hidden_dim // heads_num, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA) - ) if self.enable_mouse else None + self.mouse_attn_q = ReplicatedLinear( + img_hidden_size, keyboard_hidden_dim, bias=qkv_bias + ) + self.keyboard_attn_kv = ReplicatedLinear( + hidden_size * windows_size * vae_time_compression_ratio, + keyboard_hidden_dim * 2, + bias=qkv_bias, + ) + self.proj_keyboard = ReplicatedLinear( + keyboard_hidden_dim, img_hidden_size, bias=qkv_bias + ) - self.keyboard_attn_layer = LocalAttention( - num_heads=heads_num, - head_size=keyboard_hidden_dim // heads_num, - causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA) - ) if self.enable_keyboard else None + self.mouse_attn_layer = ( + LocalAttention( + num_heads=heads_num, + head_size=mouse_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) + if self.enable_mouse + else None + ) + + self.keyboard_attn_layer = ( + LocalAttention( + num_heads=heads_num, + head_size=keyboard_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) + if self.enable_keyboard + else None + ) self.vae_time_compression_ratio = vae_time_compression_ratio self.windows_size = windows_size @@ -186,11 +351,11 @@ class ActionModule(nn.Module): x : (N C T H W) """ pt, ph, pw = self.patch_size - t, h, w = x.shape[2] // pt, x.shape[3] // ph, x.shape[4] // pw + t, h, w = x.shape[2] // pt, x.shape[3] // ph, x.shape[4] // pw c = x.shape[1] - x = x.reshape(shape=(x.shape[0], c, t , pt, h , ph, w , pw)) + x = x.reshape(shape=(x.shape[0], c, t, pt, h, ph, w, pw)) x = torch.einsum("nctohpwq->nthwcopq", x) - x = x.reshape(shape=(x.shape[0], t*h*w, c*pt*ph*pw)) + x = x.reshape(shape=(x.shape[0], t * h * w, c * pt * ph * pw)) return x def unpatchify(self, x, t, h, w, patch_size): @@ -198,7 +363,7 @@ class ActionModule(nn.Module): x: (N, T, patch_size**2 * C) imgs: (N, H, W, C) """ - c = x.shape[2] // patch_size #self.unpatchify_channels + c = x.shape[2] // patch_size # self.unpatchify_channels pt, ph, pw = self.patch_size assert t * h * w == x.shape[1] @@ -208,11 +373,19 @@ class ActionModule(nn.Module): return imgs - def get_rotary_pos_embed(self, video_length, height, width, head_dim, rope_dim_list = None, start_offset=0): + def get_rotary_pos_embed( + self, + video_length, + height, + width, + head_dim, + rope_dim_list=None, + start_offset=0, + ): target_ndim = 3 ndim = 5 - 2 - - latents_size = [video_length+start_offset, height, width] + + latents_size = [video_length + start_offset, height, width] if isinstance(self.patch_size, int): assert all(s % self.patch_size == 0 for s in latents_size), ( @@ -233,13 +406,17 @@ class ActionModule(nn.Module): ] if len(rope_sizes) != target_ndim: - rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis - + rope_sizes = [1] * ( + target_ndim - len(rope_sizes) + ) + rope_sizes # time axis + if rope_dim_list is None: - rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] - assert ( - sum(rope_dim_list) == head_dim - ), "sum(rope_dim_list) should equal to head_dim of attention layer" + rope_dim_list = [ + head_dim // target_ndim for _ in range(target_ndim) + ] + assert sum(rope_dim_list) == head_dim, ( + "sum(rope_dim_list) should equal to head_dim of attention layer" + ) # Use Matrix-Game wrapper for FastVideo's function freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame( rope_dim_list, @@ -247,337 +424,399 @@ class ActionModule(nn.Module): theta=self.rope_theta, theta_rescale_factor=1, ) - return freqs_cos[-video_length*rope_sizes[1]*rope_sizes[2]//self.patch_size[0]:], freqs_sin[-video_length*rope_sizes[1]*rope_sizes[2]//self.patch_size[0]:] + return freqs_cos[ + -video_length + * rope_sizes[1] + * rope_sizes[2] + // self.patch_size[0] : + ], freqs_sin[ + -video_length + * rope_sizes[1] + * rope_sizes[2] + // self.patch_size[0] : + ] - def forward(self, x, tt, th, tw, mouse_condition=None, keyboard_condition=None, block_mask_mouse=None, block_mask_keyboard=None, is_causal=False, kv_cache_mouse=None, kv_cache_keyboard=None, start_frame=0, use_rope_keyboard=True, num_frame_per_block=3): - ''' + def _forward_mouse( + self, + hidden_states: torch.Tensor, + mouse_condition: torch.Tensor, + *, + is_causal: bool, + kv_cache_mouse: dict[str, torch.Tensor] | None, + pad_t: int, + num_frame_per_block: int, + block_mask_mouse: BlockMask | None, + start_frame: int, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + N_feats: int, + B: int, + C: int, + tt: int, + th: int, + tw: int, + ) -> torch.Tensor: + pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1) + mouse_condition = torch.cat([pad, mouse_condition], dim=1) + if is_causal and kv_cache_mouse is not None: + mouse_condition = mouse_condition[ + :, + self.vae_time_compression_ratio + * (N_feats - num_frame_per_block - self.windows_size) + + pad_t :, + :, + ] + group_mouse = [ + mouse_condition[ + :, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t : i * self.vae_time_compression_ratio + pad_t, + :, + ] + for i in range(num_frame_per_block) + ] + else: + local_num_frames = tt + group_mouse = [ + mouse_condition[ + :, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t : i * self.vae_time_compression_ratio + pad_t, + :, + ] + for i in range(local_num_frames) + ] + + group_mouse = torch.stack(group_mouse, dim=1) + actual_num_frames = group_mouse.shape[ + 1 + ] # Use actual stacked frame count + + S = th * tw + group_mouse = group_mouse.unsqueeze(-1).expand( + B, actual_num_frames, pad_t, C, S + ) + group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape( + B * S, actual_num_frames, pad_t * C + ) + + group_mouse = torch.cat([hidden_states, group_mouse], dim=-1) + group_mouse = self.mouse_mlp(group_mouse) + # qkv + mouse_qkv, _ = self.t_qkv(group_mouse) + q, k, v = rearrange( + mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num + ) # BHW F H C + q = self.img_attn_q_norm(q).to(v) + k = self.img_attn_k_norm(k).to(v) + # rope embd + + # freqs_cis = (self.freqs_cos, self.freqs_sin) + + q, k = _apply_rotary_emb_qk( + q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame + ) + ## TODO: adding cache here + if is_causal: + if kv_cache_mouse is None: + assert ( + q.shape[0] == k.shape[0] and q.shape[0] % S == 0 + ) # == 880, f"{q.shape[0]},{k.shape[0]}" + padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] + padded_q = _padding_q_k_v(q, padded_length) + padded_k = _padding_q_k_v(k, padded_length) + padded_v = _padding_q_k_v(v, padded_length) + attn = flex_attention( + query=padded_q.transpose(2, 1), # after: B, HW, F, C + key=padded_k.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask_mouse, + )[:, :, :-padded_length].transpose(2, 1) + else: + attn = _update_kv_cache_and_attend( + q, + k, + v, + kv_cache_mouse, + self.mouse_attn_layer, + start_frame, + num_frame_per_block, + self.local_attn_size, + use_k_for_num_tokens=False, + ) + else: + attn = self.mouse_attn_layer(q, k, v) + # Compute cu_squlens and max_seqlen for flash attention + # qk norm + attn = rearrange(attn, "(b S) T h d -> b (T S) (h d)", b=B) + attn, _ = self.proj_mouse(attn) + return attn + + def _forward_keyboard( + self, + hidden_states: torch.Tensor, + keyboard_condition: torch.Tensor, + *, + is_causal: bool, + use_rope_keyboard: bool, + kv_cache_keyboard: dict[str, torch.Tensor] | None, + pad_t: int, + num_frame_per_block: int, + block_mask_keyboard: BlockMask | None, + start_frame: int, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + N_feats: int, + B: int, + tt: int, + th: int, + tw: int, + ) -> torch.Tensor: + pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1) + keyboard_condition = torch.cat([pad, keyboard_condition], dim=1) + if is_causal and kv_cache_keyboard is not None: + keyboard_condition = keyboard_condition[ + :, + self.vae_time_compression_ratio + * (N_feats - num_frame_per_block - self.windows_size) + + pad_t :, + :, + ] # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:] + keyboard_condition = self.keyboard_embed(keyboard_condition) + group_keyboard = [ + keyboard_condition[ + :, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t : i * self.vae_time_compression_ratio + pad_t, + :, + ] + for i in range(num_frame_per_block) + ] + else: + keyboard_condition = self.keyboard_embed(keyboard_condition) + local_num_frames = tt + group_keyboard = [ + keyboard_condition[ + :, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t : i * self.vae_time_compression_ratio + pad_t, + :, + ] + for i in range(local_num_frames) + ] + group_keyboard = torch.stack(group_keyboard, dim=1) # B F RW C + group_keyboard = group_keyboard.reshape( + shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1) + ) + # apply cross attn + mouse_q, _ = self.mouse_attn_q(hidden_states) + keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard) + + B, L, HD = mouse_q.shape + D = HD // self.heads_num + q = mouse_q.view(B, L, self.heads_num, D) + + B, L, KHD = keyboard_kv.shape + k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute( + 2, 0, 1, 3, 4 + ) + + # Compute cu_squlens and max_seqlen for flash attention + # qk norm + + q = self.key_attn_q_norm(q).to(v) + k = self.key_attn_k_norm(k).to(v) + S = th * tw + # assert S == 880 + # position embed + if use_rope_keyboard: + B, TS, H, D = q.shape + T_ = TS // S + q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D) + q, k = _apply_rotary_emb_qk( + q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame + ) + + k = k.repeat_interleave(S, dim=0) + v = v.repeat_interleave(S, dim=0) + + if is_causal: + if kv_cache_keyboard is None: + assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0 + + padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] + padded_q = _padding_q_k_v(q, padded_length) + padded_k = _padding_q_k_v(k, padded_length) + padded_v = _padding_q_k_v(v, padded_length) + attn = flex_attention( + query=padded_q.transpose(2, 1), # after: B, HW, F, C + key=padded_k.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask_keyboard, + )[:, :, :-padded_length].transpose(2, 1) + else: + assert ( + k.shape[0] == S + ) # BS == 1 or the cache should not be saved/ load method should be modified + attn = _update_kv_cache_and_attend( + q, + k, + v, + kv_cache_keyboard, + self.keyboard_attn_layer, + start_frame, + num_frame_per_block, + self.local_attn_size, + use_k_for_num_tokens=True, + store_first_only=True, + repeat_factor=S, + ) + else: + attn = self.keyboard_attn_layer(q, k, v) + attn = rearrange(attn, "(B S) T H D -> B (T S) (H D)", S=S) + else: + if is_causal: + if kv_cache_keyboard is None: + padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] + padded_q = _padding_q_k_v(q, padded_length) + padded_k = _padding_q_k_v(k, padded_length) + padded_v = _padding_q_k_v(v, padded_length) + attn = flex_attention( + query=padded_q.transpose(2, 1), # after: B, HW, F, C + key=padded_k.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask_keyboard, + )[:, :, :-padded_length].transpose(2, 1) + else: + attn = _update_kv_cache_and_attend( + q, + k, + v, + kv_cache_keyboard, + self.keyboard_attn_layer, + start_frame, + num_frame_per_block, + self.local_attn_size, + use_k_for_num_tokens=True, + ) + else: + attn = self.keyboard_attn_layer(q, k, v) + attn = rearrange(attn, "B L H D -> B L (H D)") + attn, _ = self.proj_keyboard(attn) + return attn + + def forward( + self, + x: torch.Tensor, + tt: int, + th: int, + tw: int, + mouse_condition: torch.Tensor | None = None, + keyboard_condition: torch.Tensor | None = None, + block_mask_mouse: BlockMask | None = None, + block_mask_keyboard: BlockMask | None = None, + is_causal: bool = False, + kv_cache_mouse: dict[str, torch.Tensor] | None = None, + kv_cache_keyboard: dict[str, torch.Tensor] | None = None, + start_frame: int = 0, + use_rope_keyboard: bool = True, + num_frame_per_block: int = 3, + ): + """ hidden_states: B, tt*th*tw, C mouse_condition: B, N_frames, C1 keyboard_condition: B, N_frames, C2 - ''' + """ assert use_rope_keyboard target_device = x.device target_dtype = x.dtype if mouse_condition is not None: - mouse_condition = mouse_condition.to(device=target_device, - dtype=target_dtype) + mouse_condition = mouse_condition.to( + device=target_device, dtype=target_dtype + ) if keyboard_condition is not None: keyboard_condition = keyboard_condition.to( - device=target_device, dtype=target_dtype) + device=target_device, dtype=target_dtype + ) else: return x B, N_frames, C = keyboard_condition.shape - assert tt*th*tw == x.shape[1] - assert ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0 + assert tt * th * tw == x.shape[1] + assert ( + (N_frames - 1) + self.vae_time_compression_ratio + ) % self.vae_time_compression_ratio == 0 N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1 - + # Lazy initialization of freqs on first forward pass if self._freqs_cos is None or self._freqs_sin is None: self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed( - 7500, self.patch_size[1], self.patch_size[2], 64, - self.mouse_qk_dim_list, start_offset=0 + 7500, + self.patch_size[1], + self.patch_size[2], + 64, + self.mouse_qk_dim_list, + start_offset=0, ) - + # Defined freqs_cis early so it's available for both mouse and keyboard freqs_cis = (self._freqs_cos, self._freqs_sin) if is_causal: - assert (N_feats == tt and kv_cache_mouse is None) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block) + assert (N_feats == tt and kv_cache_mouse is None) or ( + (N_frames - 1) // self.vae_time_compression_ratio + 1 + == start_frame + num_frame_per_block + ) # For non-causal (training), we trust that the caller provides correctly shaped inputs - + if self.enable_mouse and mouse_condition is not None: - hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, S=th*tw) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670 + hidden_states = rearrange( + x, "B (T S) C -> (B S) T C", T=tt, S=th * tw + ) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670 B, N_frames, C = mouse_condition.shape else: hidden_states = x # padding - + pad_t = self.vae_time_compression_ratio * self.windows_size if self.enable_mouse and mouse_condition is not None: - pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1) - mouse_condition = torch.cat([pad, mouse_condition], dim=1) - if is_causal and kv_cache_mouse is not None: - mouse_condition = mouse_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] - group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)] - else: - local_num_frames = tt - group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(local_num_frames)] - - group_mouse = torch.stack(group_mouse, dim = 1) - actual_num_frames = group_mouse.shape[1] # Use actual stacked frame count - - S = th * tw - group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S) - group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, actual_num_frames, pad_t * C) - - group_mouse = torch.cat([hidden_states, group_mouse], dim = -1) - group_mouse = self.mouse_mlp(group_mouse) - # qkv - mouse_qkv, _ = self.t_qkv(group_mouse) - q, k, v = rearrange(mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) # BHW F H C - q = self.img_attn_q_norm(q).to(v) - k = self.img_attn_k_norm(k).to(v) - # rope embd - - - # freqs_cis = (self.freqs_cos, self.freqs_sin) - - - q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) - ## TODO: adding cache here - if is_causal: - if kv_cache_mouse is None: - assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0 # == 880, f"{q.shape[0]},{k.shape[0]}" - padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] - padded_q = torch.cat( - [q, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 - ) - padded_k = torch.cat( - [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 - ) - padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 - ) - attn = flex_attention( - query=padded_q.transpose(2, 1), # after: B, HW, F, C - key=padded_k.transpose(2, 1), - value=padded_v.transpose(2, 1), - block_mask=block_mask_mouse - )[:, :, :-padded_length].transpose(2, 1) - else: - current_start = start_frame - current_end = current_start + q.shape[1] - - assert q.shape[1] == num_frame_per_block - sink_size = 0 - max_attention_size = self.local_attn_size - sink_tokens = sink_size * 1 - kv_cache_size = kv_cache_mouse["k"].shape[1] - num_new_tokens = q.shape[1] - - - if (current_end > kv_cache_mouse["global_end_index"].item()) and ( - num_new_tokens + kv_cache_mouse["local_end_index"].item() > kv_cache_size): - num_evicted_tokens = num_new_tokens + kv_cache_mouse["local_end_index"].item() - kv_cache_size - num_rolled_tokens = kv_cache_mouse["local_end_index"].item() - num_evicted_tokens - sink_tokens - kv_cache_mouse["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache_mouse["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache_mouse["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache_mouse["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - # Insert the new keys/values at the end - local_end_index = kv_cache_mouse["local_end_index"].item() + current_end - \ - kv_cache_mouse["global_end_index"].item() - num_evicted_tokens - local_start_index = local_end_index - num_new_tokens - else: - local_end_index = kv_cache_mouse["local_end_index"].item() + current_end - kv_cache_mouse["global_end_index"].item() - local_start_index = local_end_index - num_new_tokens - kv_cache_mouse["k"][:, local_start_index:local_end_index] = k - kv_cache_mouse["v"][:, local_start_index:local_end_index] = v - - attn = self.mouse_attn_layer( - q, - kv_cache_mouse["k"][:, max(0, local_end_index - max_attention_size):local_end_index], - kv_cache_mouse["v"][:, max(0, local_end_index - max_attention_size):local_end_index], - ) - kv_cache_mouse["global_end_index"].fill_(current_end) - kv_cache_mouse["local_end_index"].fill_(local_end_index) - else: - attn = self.mouse_attn_layer(q, k, v) - # Compute cu_squlens and max_seqlen for flash attention - # qk norm - attn = rearrange(attn, '(b S) T h d -> b (T S) (h d)',b=B) - + attn = self._forward_mouse( + hidden_states, + mouse_condition, + is_causal=is_causal, + kv_cache_mouse=kv_cache_mouse, + pad_t=pad_t, + num_frame_per_block=num_frame_per_block, + block_mask_mouse=block_mask_mouse, + start_frame=start_frame, + freqs_cis=freqs_cis, + N_feats=N_feats, + B=B, + C=C, + tt=tt, + th=th, + tw=tw, + ) hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=B) - attn, _ = self.proj_mouse(attn) - + hidden_states = hidden_states + attn - + if self.enable_keyboard and keyboard_condition is not None: - pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1) - keyboard_condition = torch.cat([pad, keyboard_condition], dim=1) - if is_causal and kv_cache_keyboard is not None: - keyboard_condition = keyboard_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:] - keyboard_condition = self.keyboard_embed(keyboard_condition) - group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)] - else: - keyboard_condition = self.keyboard_embed(keyboard_condition) - local_num_frames = tt - group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(local_num_frames)] - group_keyboard = torch.stack(group_keyboard, dim = 1) # B F RW C - group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0],group_keyboard.shape[1],-1)) - # apply cross attn - mouse_q, _ = self.mouse_attn_q(hidden_states) - keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard) - - B, L, HD = mouse_q.shape - D = HD // self.heads_num - q = mouse_q.view(B, L, self.heads_num, D) - - B, L, KHD = keyboard_kv.shape - k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(2, 0, 1, 3, 4) - - # Compute cu_squlens and max_seqlen for flash attention - # qk norm - - q = self.key_attn_q_norm(q).to(v) - k = self.key_attn_k_norm(k).to(v) - S = th * tw - # assert S == 880 - # position embed - if use_rope_keyboard: - B, TS, H, D = q.shape - T_ = TS // S - q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D) - q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) - - k1, k2, k3, k4 = k.shape - k = k.repeat_interleave(S, dim=0) - v = v.repeat_interleave(S, dim=0) - - - if is_causal: - if kv_cache_keyboard is None: - assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0 - - padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] - padded_q = torch.cat( - [q, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 - ) - padded_k = torch.cat( - [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 - ) - padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 - ) - attn = flex_attention( - query=padded_q.transpose(2, 1), # after: B, HW, F, C - key=padded_k.transpose(2, 1), - value=padded_v.transpose(2, 1), - block_mask=block_mask_keyboard - )[:, :, :-padded_length].transpose(2, 1) - else: - current_start = start_frame - current_end = current_start + k.shape[1] - assert k.shape[1] == num_frame_per_block - sink_size = 0 - max_attention_size = self.local_attn_size - sink_tokens = sink_size * 1 - kv_cache_size = kv_cache_keyboard["k"].shape[1] - num_new_tokens = k.shape[1] - - if (current_end > kv_cache_keyboard["global_end_index"].item()) and ( - num_new_tokens + kv_cache_keyboard["local_end_index"].item() > kv_cache_size): - num_evicted_tokens = num_new_tokens + kv_cache_keyboard["local_end_index"].item() - kv_cache_size - num_rolled_tokens = kv_cache_keyboard["local_end_index"].item() - num_evicted_tokens - sink_tokens - kv_cache_keyboard["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache_keyboard["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache_keyboard["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache_keyboard["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - # Insert the new keys/values at the end - local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - \ - kv_cache_keyboard["global_end_index"].item() - num_evicted_tokens - local_start_index = local_end_index - num_new_tokens - else: - local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item() - local_start_index = local_end_index - num_new_tokens - assert k.shape[0] == S # BS == 1 or the cache should not be saved/ load method should be modified - kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[:1] - kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[:1] - - attn = self.keyboard_attn_layer( - q, - kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1), - kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1), - ) - - kv_cache_keyboard["global_end_index"].fill_(current_end) - kv_cache_keyboard["local_end_index"].fill_(local_end_index) - else: - attn = self.keyboard_attn_layer(q, k, v) - attn = rearrange(attn, '(B S) T H D -> B (T S) (H D)', S=S) - else: - if is_causal: - if kv_cache_keyboard is None: - - padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] - padded_q = torch.cat( - [q, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 - ) - padded_k = torch.cat( - [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 - ) - padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 - ) - attn = flex_attention( - query=padded_q.transpose(2, 1), # after: B, HW, F, C - key=padded_k.transpose(2, 1), - value=padded_v.transpose(2, 1), - block_mask=block_mask_keyboard - )[:, :, :-padded_length].transpose(2, 1) - else: - current_start = start_frame - current_end = current_start + k.shape[1] - assert k.shape[1] == num_frame_per_block - sink_size = 0 - max_attention_size = self.local_attn_size - sink_tokens = sink_size * 1 - kv_cache_size = kv_cache_keyboard["k"].shape[1] - num_new_tokens = k.shape[1] - - - if (current_end > kv_cache_keyboard["global_end_index"].item()) and ( - num_new_tokens + kv_cache_keyboard["local_end_index"].item() > kv_cache_size): - num_evicted_tokens = num_new_tokens + kv_cache_keyboard["local_end_index"].item() - kv_cache_size - num_rolled_tokens = kv_cache_keyboard["local_end_index"].item() - num_evicted_tokens - sink_tokens - kv_cache_keyboard["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache_keyboard["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache_keyboard["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache_keyboard["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - # Insert the new keys/values at the end - local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - \ - kv_cache_keyboard["global_end_index"].item() - num_evicted_tokens - local_start_index = local_end_index - num_new_tokens - - - else: - local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item() - local_start_index = local_end_index - num_new_tokens - kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k - kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v - attn = self.keyboard_attn_layer( - q, - kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index], - kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index], - ) - kv_cache_keyboard["global_end_index"].fill_(current_end) - kv_cache_keyboard["local_end_index"].fill_(local_end_index) - else: - attn = self.keyboard_attn_layer(q, k, v) - attn = rearrange(attn, 'B L H D -> B L (H D)') - attn, _ = self.proj_keyboard(attn) + attn = self._forward_keyboard( + hidden_states, + keyboard_condition, + is_causal=is_causal, + use_rope_keyboard=use_rope_keyboard, + kv_cache_keyboard=kv_cache_keyboard, + pad_t=pad_t, + num_frame_per_block=num_frame_per_block, + block_mask_keyboard=block_mask_keyboard, + start_frame=start_frame, + freqs_cis=freqs_cis, + N_feats=N_feats, + B=B, + tt=tt, + th=th, + tw=tw, + ) hidden_states = hidden_states + attn - return hidden_states \ No newline at end of file + return hidden_states diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 13806d2ac..2b709d9a6 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -14,22 +14,33 @@ from torch.nn.attention.flex_attention import BlockMask from fastvideo.forward_context import get_forward_context flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" +) from fastvideo.attention import LocalAttention from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig from fastvideo.distributed.parallel_state import get_sp_world_size -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, - RMSNorm, ScaleResidual, - ScaleResidualLayerNormScaleShift) +from fastvideo.layers.layernorm import ( + FP32LayerNorm, + LayerNormScaleShift, + RMSNorm, + ScaleResidual, + ScaleResidualLayerNormScaleShift, +) from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import (get_rotary_pos_embed) -from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection +from fastvideo.layers.rotary_embedding import get_rotary_pos_embed +from fastvideo.layers.visual_embedding import ( + PatchEmbed, + TimestepEmbedder, + ModulateProjection, +) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import (WanT2VCrossAttention, - WanImageEmbedding) +from fastvideo.models.dits.wanvideo import ( + WanT2VCrossAttention, + WanImageEmbedding, +) from fastvideo.platforms import AttentionBackendEnum, current_platform from .action_module import ActionModule @@ -47,19 +58,24 @@ def causal_rope_apply(x, grid_sizes, freqs, start_frame=0): # loop over samples output = [] f, h, w = grid_sizes.tolist() - + for i in range(len(x)): seq_len = f * h * w # precompute multipliers - x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( - seq_len, n, -1, 2)) - freqs_i = torch.cat([ - freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), - freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), - freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) - ], - dim=-1).reshape(seq_len, 1, -1) + x_i = torch.view_as_complex( + x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2) + ) + freqs_i = torch.cat( + [ + freqs[0][start_frame : start_frame + f] + .view(f, 1, 1, -1) + .expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) # apply rotary embedding x_i = torch.view_as_real(x_i * freqs_i).flatten(2) @@ -74,8 +90,9 @@ def rope_params(max_seq_len, dim, theta=10000): assert dim % 2 == 0 freqs = torch.outer( torch.arange(max_seq_len), - 1.0 / torch.pow(theta, - torch.arange(0, dim, 2).to(torch.float64).div(dim))) + 1.0 + / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + ) freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs @@ -90,11 +107,12 @@ class CausalMatrixGameTimeImageEmbedding(nn.Module): super().__init__() self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu") - self.time_modulation = ModulateProjection(dim, - factor=6, - act_layer="silu") - + dim, frequency_embedding_size=time_freq_dim, act_layer="silu" + ) + self.time_modulation = ModulateProjection( + dim, factor=6, act_layer="silu" + ) + self.image_embedder = None if image_embed_dim is not None: self.image_embedder = WanImageEmbedding(image_embed_dim, dim) @@ -108,24 +126,26 @@ class CausalMatrixGameTimeImageEmbedding(nn.Module): ): temb = self.time_embedder(timestep, timestep_seq_len) timestep_proj = self.time_modulation(temb) - + if encoder_hidden_states_image is not None: assert self.image_embedder is not None encoder_hidden_states_image = self.image_embedder( - encoder_hidden_states_image) + encoder_hidden_states_image + ) return temb, timestep_proj, None, encoder_hidden_states_image class CausalMatrixGameSelfAttention(nn.Module): - - def __init__(self, - dim: int, - num_heads: int, - local_attn_size: int = -1, - sink_size: int = 0, - qk_norm: bool = True, - eps: float = 1e-6) -> None: + def __init__( + self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: bool = True, + eps: float = 1e-6, + ) -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim @@ -143,8 +163,11 @@ class CausalMatrixGameSelfAttention(nn.Module): dropout_rate=0, softmax_scale=None, causal=False, - supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, - AttentionBackendEnum.TORCH_SDPA)) + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) def forward( self, @@ -156,7 +179,7 @@ class CausalMatrixGameSelfAttention(nn.Module): kv_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, - grid_sizes: torch.Tensor | None = None + grid_sizes: torch.Tensor | None = None, ): if cache_start is None: cache_start = current_start @@ -167,46 +190,70 @@ class CausalMatrixGameSelfAttention(nn.Module): start_frame = current_start // frame_seqlen else: start_frame = 0 - + if self._freqs_cache is None or self._freqs_cache.device != q.device: d = self.head_dim - self._freqs_cache = torch.cat([ - rope_params(1024, d - 4 * (d // 6)), - rope_params(1024, 2 * (d // 6)), - rope_params(1024, 2 * (d // 6)) - ], dim=1).to(q.device) - + self._freqs_cache = torch.cat( + [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], + dim=1, + ).to(q.device) + freqs = self._freqs_cache - - roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=start_frame).type_as(v) - roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=start_frame).type_as(v) + + roped_query = causal_rope_apply( + q, grid_sizes, freqs, start_frame=start_frame + ).type_as(v) + roped_key = causal_rope_apply( + k, grid_sizes, freqs, start_frame=start_frame + ).type_as(v) if kv_cache is None: padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] padded_roped_query = torch.cat( - [roped_query, - torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], - device=q.device, dtype=v.dtype)], - dim=1 + [ + roped_query, + torch.zeros( + [q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, + dtype=v.dtype, + ), + ], + dim=1, ) padded_roped_key = torch.cat( - [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], - device=k.device, dtype=v.dtype)], - dim=1 + [ + roped_key, + torch.zeros( + [k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, + dtype=v.dtype, + ), + ], + dim=1, ) padded_v = torch.cat( - [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], - device=v.device, dtype=v.dtype)], - dim=1 + [ + v, + torch.zeros( + [v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, + dtype=v.dtype, + ), + ], + dim=1, ) x = flex_attention( query=padded_roped_query.transpose(2, 1), key=padded_roped_key.transpose(2, 1), value=padded_v.transpose(2, 1), - block_mask=block_mask + block_mask=block_mask, )[:, :, :-padded_length].transpose(2, 1) else: # Calculate frame_seqlen correctly from grid_sizes (single frame token count) @@ -215,34 +262,68 @@ class CausalMatrixGameSelfAttention(nn.Module): else: # Fallback: assume q.shape[1] is single frame (shouldn't happen in causal mode) frame_seqlen = q.shape[1] - logger.warning("grid_sizes not provided, using q.shape[1] as frame_seqlen") + logger.warning( + "grid_sizes not provided, using q.shape[1] as frame_seqlen" + ) current_end = current_start + roped_query.shape[1] sink_tokens = self.sink_size * frame_seqlen - + # Compute max_attention_size dynamically based on actual frame_seqlen - max_attention_size = (15 * frame_seqlen if self.local_attn_size == -1 - else self.local_attn_size * frame_seqlen) + max_attention_size = ( + 15 * frame_seqlen + if self.local_attn_size == -1 + else self.local_attn_size * frame_seqlen + ) kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = roped_query.shape[1] if (current_end > kv_cache["global_end_index"].item()) and ( - num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): - - num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size - num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens - kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - local_end_index = kv_cache["local_end_index"].item() + current_end - \ - kv_cache["global_end_index"].item() - num_evicted_tokens + num_new_tokens + kv_cache["local_end_index"].item() + > kv_cache_size + ): + num_evicted_tokens = ( + num_new_tokens + + kv_cache["local_end_index"].item() + - kv_cache_size + ) + num_rolled_tokens = ( + kv_cache["local_end_index"].item() + - num_evicted_tokens + - sink_tokens + ) + kv_cache["k"][ + :, sink_tokens : sink_tokens + num_rolled_tokens + ] = kv_cache["k"][ + :, + sink_tokens + num_evicted_tokens : sink_tokens + + num_evicted_tokens + + num_rolled_tokens, + ].clone() + kv_cache["v"][ + :, sink_tokens : sink_tokens + num_rolled_tokens + ] = kv_cache["v"][ + :, + sink_tokens + num_evicted_tokens : sink_tokens + + num_evicted_tokens + + num_rolled_tokens, + ].clone() + local_end_index = ( + kv_cache["local_end_index"].item() + + current_end + - kv_cache["global_end_index"].item() + - num_evicted_tokens + ) local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v else: - local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_end_index = ( + kv_cache["local_end_index"].item() + + current_end + - kv_cache["global_end_index"].item() + ) local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = roped_key @@ -265,21 +346,23 @@ class CausalMatrixGameSelfAttention(nn.Module): class CausalMatrixGameTransformerBlock(nn.Module): - - def __init__(self, - dim: int, - ffn_dim: int, - num_heads: int, - local_attn_size: int = -1, - sink_size: int = 0, - qk_norm: str = "rms_norm_across_heads", - cross_attn_norm: bool = False, - eps: float = 1e-6, - added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, - prefix: str = "", - action_config: dict | None = None, - block_idx: int = 0): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + action_config: dict | None = None, + block_idx: int = 0, + ): super().__init__() action_config = action_config or {} @@ -295,7 +378,8 @@ class CausalMatrixGameTransformerBlock(nn.Module): local_attn_size=local_attn_size, sink_size=sink_size, qk_norm=qk_norm, - eps=eps) + eps=eps, + ) self.hidden_dim = dim self.num_attention_heads = num_heads self.local_attn_size = local_attn_size @@ -315,27 +399,28 @@ class CausalMatrixGameTransformerBlock(nn.Module): norm_type="layer", eps=eps, elementwise_affine=True, - dtype=torch.float32) + dtype=torch.float32, + ) if added_kv_proj_dim is not None: - self.attn2 = MatrixGameCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + self.attn2 = MatrixGameCrossAttention( + dim, num_heads, qk_norm=qk_norm, eps=eps + ) else: - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + self.attn2 = WanT2VCrossAttention( + dim, num_heads, qk_norm=qk_norm, eps=eps + ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", eps=eps, elementwise_affine=False, - dtype=torch.float32) + dtype=torch.float32, + ) - - if len(action_config) > 0 and block_idx in action_config.get("blocks", []): + if len(action_config) > 0 and block_idx in action_config.get( + "blocks", [] + ): self.action_model = ActionModule( heads_num=action_config["heads_num"], hidden_size=action_config["hidden_size"], @@ -349,12 +434,16 @@ class CausalMatrixGameTransformerBlock(nn.Module): windows_size=action_config["windows_size"], rope_theta=action_config["rope_theta"], rope_dim_list=action_config["rope_dim_list"], - mouse_qk_dim_list=action_config.get("mouse_qk_dim_list", [8, 28, 28]), + mouse_qk_dim_list=action_config.get( + "mouse_qk_dim_list", [8, 28, 28] + ), patch_size=action_config["patch_size"], local_attn_size=local_attn_size, qk_norm=action_config["qk_norm"], qkv_bias=action_config["qkv_bias"], - vae_time_compression_ratio=action_config["vae_time_compression_ratio"], + vae_time_compression_ratio=action_config[ + "vae_time_compression_ratio" + ], ) else: self.action_model = None @@ -371,7 +460,7 @@ class CausalMatrixGameTransformerBlock(nn.Module): temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, - grid_sizes: torch.Tensor | None = None, + grid_sizes: torch.Tensor, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, block_mask_mouse: BlockMask | None = None, @@ -397,10 +486,17 @@ class CausalMatrixGameTransformerBlock(nn.Module): e = self.scale_shift_table + temb assert e.shape == (bs, num_frames, 6, self.hidden_dim) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + e.chunk(6, dim=2) + ) - norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - (1 + scale_msa) + shift_msa).flatten(1, 2) + norm_hidden_states = ( + self.norm1(hidden_states).unflatten( + dim=1, sizes=(num_frames, frame_seqlen) + ) + * (1 + scale_msa) + + shift_msa + ).flatten(1, 2) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -414,45 +510,65 @@ class CausalMatrixGameTransformerBlock(nn.Module): key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) - attn_output = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start, grid_sizes) + attn_output = self.attn1( + query, + key, + value, + freqs_cis, + block_mask, + kv_cache, + current_start, + cache_start, + grid_sizes, + ) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) + hidden_states, attn_output, gate_msa, null_shift, null_scale + ) # norm3 weights are loaded into self_attn_residual_norm.norm - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None, - crossattn_cache=crossattn_cache) + attn_output = self.attn2( + norm_hidden_states, + context=encoder_hidden_states, + context_lens=None, + crossattn_cache=crossattn_cache, + ) # residual connection hidden_states = hidden_states + attn_output if self.action_model is not None: if mouse_cond is not None or keyboard_cond is not None: - start_frame = current_start // frame_seqlen if grid_sizes is not None else 0 + start_frame = current_start // frame_seqlen hidden_states = self.action_model( hidden_states, - grid_sizes[0], grid_sizes[1], grid_sizes[2], - mouse_cond, keyboard_cond, - block_mask_mouse, block_mask_keyboard, + int(grid_sizes[0]), + int(grid_sizes[1]), + int(grid_sizes[2]), + mouse_cond, + keyboard_cond, + block_mask_mouse, + block_mask_keyboard, is_causal=True, kv_cache_mouse=kv_cache_mouse, kv_cache_keyboard=kv_cache_keyboard, start_frame=start_frame, use_rope_keyboard=use_rope_keyboard, - num_frame_per_block=num_frame_per_block + num_frame_per_block=num_frame_per_block, ) - ff_output = self.ffn( - (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * - - (1 + c_scale_msa) + c_shift_msa).flatten(1, 2) + ( + self.norm1(hidden_states).unflatten( + dim=1, sizes=(num_frames, frame_seqlen) + ) + * (1 + c_scale_msa) + + c_shift_msa + ).flatten(1, 2) ) hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) @@ -468,15 +584,23 @@ class CausalMatrixGameWanModel(BaseDiT): _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions - _supported_attention_backends = _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends + _supported_attention_backends = ( + _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends + ) param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.param_names_mapping - reverse_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping - lora_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping + reverse_param_names_mapping = ( + _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping + ) + lora_param_names_mapping = ( + _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping + ) - def __init__(self, - config: MatrixGameWanVideoConfig, - hf_config: dict[str, Any], - **kwargs) -> None: + def __init__( + self, + config: MatrixGameWanVideoConfig, + hf_config: dict[str, Any], + **kwargs, + ) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim @@ -489,14 +613,28 @@ class CausalMatrixGameWanModel(BaseDiT): self.patch_size = config.patch_size arch_cfg = getattr(config, "arch_config", None) - self.local_attn_size = getattr(arch_cfg, 'local_attn_size', getattr(config, 'local_attn_size', -1)) if arch_cfg else getattr(config, 'local_attn_size', -1) - self.sink_size = getattr(arch_cfg, 'sink_size', getattr(config, 'sink_size', 0)) if arch_cfg else getattr(config, 'sink_size', 0) + self.local_attn_size = ( + getattr( + arch_cfg, + "local_attn_size", + getattr(config, "local_attn_size", -1), + ) + if arch_cfg + else getattr(config, "local_attn_size", -1) + ) + self.sink_size = ( + getattr(arch_cfg, "sink_size", getattr(config, "sink_size", 0)) + if arch_cfg + else getattr(config, "sink_size", 0) + ) # 1. Patch & position embedding - self.patch_embedding = PatchEmbed(in_chans=config.in_channels, - embed_dim=inner_dim, - patch_size=config.patch_size, - flatten=False) + self.patch_embedding = PatchEmbed( + in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False, + ) # 2. Condition embeddings self.condition_embedder = CausalMatrixGameTimeImageEmbedding( @@ -507,37 +645,48 @@ class CausalMatrixGameWanModel(BaseDiT): # 2.1. Get action config arch_cfg = getattr(config, "arch_config", None) - self.action_config = getattr(arch_cfg, "action_config", {}) if arch_cfg is not None else {} + self.action_config = ( + getattr(arch_cfg, "action_config", {}) + if arch_cfg is not None + else {} + ) # 3. Transformer blocks - self.blocks = nn.ModuleList([ - CausalMatrixGameTransformerBlock( - inner_dim, - config.ffn_dim, - config.num_attention_heads, - self.local_attn_size, - self.sink_size, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", - action_config=self.action_config, - block_idx=i) - for i in range(config.num_layers) - ]) + self.blocks = nn.ModuleList( + [ + CausalMatrixGameTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + self.local_attn_size, + self.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + block_idx=i, + ) + for i in range(config.num_layers) + ] + ) # 4. Output norm & projection - self.norm_out = LayerNormScaleShift(inner_dim, - norm_type="layer", - eps=config.eps, - elementwise_affine=False, - dtype=torch.float32) + self.norm_out = LayerNormScaleShift( + inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + ) self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) + inner_dim, config.out_channels * math.prod(config.patch_size) + ) self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) + torch.randn(1, 2, inner_dim) / inner_dim**0.5 + ) self.gradient_checkpointing = False @@ -546,115 +695,172 @@ class CausalMatrixGameWanModel(BaseDiT): self.block_mask_mouse = None self.use_rope_keyboard = True - self.num_frame_per_block = getattr(arch_cfg, 'num_frames_per_block', getattr(config, 'num_frame_per_block', 1)) if arch_cfg else getattr(config, 'num_frame_per_block', 1) + self.num_frame_per_block = ( + getattr( + arch_cfg, + "num_frames_per_block", + getattr(config, "num_frame_per_block", 1), + ) + if arch_cfg + else getattr(config, "num_frame_per_block", 1) + ) self.__post_init__() @staticmethod def _prepare_blockwise_causal_attn_mask( - device: torch.device | str, num_frames: int = 9, - frame_seqlen: int = 880, num_frame_per_block: int = 1, local_attn_size: int = -1 + device: torch.device | str, + num_frames: int = 9, + frame_seqlen: int = 880, + num_frame_per_block: int = 1, + local_attn_size: int = -1, ) -> BlockMask: total_length = num_frames * frame_seqlen padded_length = math.ceil(total_length / 128) * 128 - total_length - ends = torch.zeros(total_length + padded_length, - device=device, dtype=torch.long) + ends = torch.zeros( + total_length + padded_length, device=device, dtype=torch.long + ) frame_indices = torch.arange( start=0, end=total_length, step=frame_seqlen * num_frame_per_block, - device=device + device=device, ) for tmp in frame_indices: - ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ - frame_seqlen * num_frame_per_block + ends[tmp : tmp + frame_seqlen * num_frame_per_block] = ( + tmp + frame_seqlen * num_frame_per_block + ) def attention_mask(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: - return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + return ( + (kv_idx < ends[q_idx]) + & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen)) + ) | (q_idx == kv_idx) - block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, - KV_LEN=total_length + padded_length, _compile=False, device=device) + block_mask = create_block_mask( + attention_mask, + B=None, + H=None, + Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, + _compile=False, + device=device, + ) if not dist.is_initialized() or dist.get_rank() == 0: - print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames" + ) return block_mask @staticmethod def _prepare_blockwise_causal_attn_mask_keyboard( - device: torch.device | str, num_frames: int = 9, - frame_seqlen: int = 880, num_frame_per_block: int = 1, local_attn_size: int = -1 + device: torch.device | str, + num_frames: int = 9, + frame_seqlen: int = 880, + num_frame_per_block: int = 1, + local_attn_size: int = -1, ) -> BlockMask: total_length2 = num_frames * frame_seqlen padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2 padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames - ends2 = torch.zeros(total_length2 + padded_length2, - device=device, dtype=torch.long) + ends2 = torch.zeros( + total_length2 + padded_length2, device=device, dtype=torch.long + ) frame_indices2 = torch.arange( start=0, end=total_length2, step=frame_seqlen * num_frame_per_block, - device=device + device=device, ) cnt = num_frame_per_block for tmp in frame_indices2: - ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt + ends2[tmp : tmp + frame_seqlen * num_frame_per_block] = cnt cnt += num_frame_per_block def attention_mask2(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx) else: - return ((kv_idx < ends2[q_idx]) & (kv_idx >= (ends2[q_idx] - local_attn_size))) | (q_idx == kv_idx) + return ( + (kv_idx < ends2[q_idx]) + & (kv_idx >= (ends2[q_idx] - local_attn_size)) + ) | (q_idx == kv_idx) - block_mask2 = create_block_mask(attention_mask2, B=None, H=None, Q_LEN=total_length2 + padded_length2, - KV_LEN=num_frames + padded_length_kv2, _compile=False, device=device) + block_mask2 = create_block_mask( + attention_mask2, + B=None, + H=None, + Q_LEN=total_length2 + padded_length2, + KV_LEN=num_frames + padded_length_kv2, + _compile=False, + device=device, + ) if not dist.is_initialized() or dist.get_rank() == 0: - print(f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames") + print( + f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames" + ) return block_mask2 @staticmethod def _prepare_blockwise_causal_attn_mask_action( - device: torch.device | str, num_frames: int = 9, - frame_seqlen: int = 1, num_frame_per_block: int = 1, local_attn_size: int = -1 + device: torch.device | str, + num_frames: int = 9, + frame_seqlen: int = 1, + num_frame_per_block: int = 1, + local_attn_size: int = -1, ) -> BlockMask: total_length2 = num_frames * frame_seqlen padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2 padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames - ends2 = torch.zeros(total_length2 + padded_length2, - device=device, dtype=torch.long) + ends2 = torch.zeros( + total_length2 + padded_length2, device=device, dtype=torch.long + ) frame_indices2 = torch.arange( start=0, end=total_length2, step=frame_seqlen * num_frame_per_block, - device=device + device=device, ) cnt = num_frame_per_block for tmp in frame_indices2: - ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt + ends2[tmp : tmp + frame_seqlen * num_frame_per_block] = cnt cnt += num_frame_per_block def attention_mask2(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx) else: - return ((kv_idx < ends2[q_idx]) & (kv_idx >= (ends2[q_idx] - local_attn_size))) | (q_idx == kv_idx) + return ( + (kv_idx < ends2[q_idx]) + & (kv_idx >= (ends2[q_idx] - local_attn_size)) + ) | (q_idx == kv_idx) - block_mask2 = create_block_mask(attention_mask2, B=None, H=None, Q_LEN=total_length2 + padded_length2, - KV_LEN=num_frames + padded_length_kv2, _compile=False, device=device) + block_mask2 = create_block_mask( + attention_mask2, + B=None, + H=None, + Q_LEN=total_length2 + padded_length2, + KV_LEN=num_frames + padded_length_kv2, + _compile=False, + device=device, + ) if not dist.is_initialized() or dist.get_rank() == 0: - print(f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames") + print( + f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames" + ) return block_mask2 @@ -663,7 +869,9 @@ class CausalMatrixGameWanModel(BaseDiT): hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + encoder_hidden_states_image: torch.Tensor + | list[torch.Tensor] + | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, kv_cache: dict = None, @@ -673,17 +881,26 @@ class CausalMatrixGameWanModel(BaseDiT): current_start: int = 0, cache_start: int = 0, start_frame: int = 0, - **kwargs + **kwargs, ) -> torch.Tensor: # Extract num_frame_per_block from kwargs - effective_num_frame_per_block = kwargs.pop('num_frame_per_block', self.num_frame_per_block) - - if mouse_cond is not None or keyboard_cond is not None: - assert self.action_config is not None and len(self.action_config) > 0 + effective_num_frame_per_block = kwargs.pop( + "num_frame_per_block", self.num_frame_per_block + ) - if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): + if mouse_cond is not None or keyboard_cond is not None: + assert ( + self.action_config is not None and len(self.action_config) > 0 + ) + + if encoder_hidden_states is not None and not isinstance( + encoder_hidden_states, torch.Tensor + ): encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + if ( + isinstance(encoder_hidden_states_image, list) + and len(encoder_hidden_states_image) > 0 + ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None @@ -691,9 +908,16 @@ class CausalMatrixGameWanModel(BaseDiT): ctx = get_forward_context() batch = getattr(ctx, "forward_batch", None) - if batch is not None and getattr(batch, "image_latent", None) is not None: + if ( + batch is not None + and getattr(batch, "image_latent", None) is not None + ): image_latent = batch.image_latent - if isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5: + if ( + isinstance(image_latent, torch.Tensor) + and image_latent.ndim == 5 + and hidden_states.ndim == 5 + ): _, _, num_frames, _, _ = hidden_states.shape start = start_frame end = start + num_frames @@ -729,7 +953,9 @@ class CausalMatrixGameWanModel(BaseDiT): [hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1 ) - batch_size, num_channels, num_frames, height, width = hidden_states.shape + batch_size, num_channels, num_frames, height, width = ( + hidden_states.shape + ) p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h @@ -738,40 +964,58 @@ class CausalMatrixGameWanModel(BaseDiT): d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + ( + post_patch_num_frames * get_sp_world_size(), + post_patch_height, + post_patch_width, + ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, rope_theta=10000, - start_frame=start_frame + start_frame=start_frame, ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.tensor([post_patch_num_frames, post_patch_height, post_patch_width], - device=hidden_states.device) + grid_sizes = torch.tensor( + [post_patch_num_frames, post_patch_height, post_patch_width], + device=hidden_states.device, + ) hidden_states = hidden_states.flatten(2).transpose(1, 2) if timestep.dim() == 2: timestep = timestep.flatten() - temb, timestep_proj, encoder_hidden_states_text, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) + ( + temb, + timestep_proj, + encoder_hidden_states_text, + encoder_hidden_states_image, + ) = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image + ) + timestep_proj = timestep_proj.unflatten( + 1, (6, self.hidden_size) + ).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) if encoder_hidden_states_text is not None: if isinstance(encoder_hidden_states_text, list): encoder_hidden_states_text = encoder_hidden_states_text[0] elif encoder_hidden_states_text.ndim == 2: - encoder_hidden_states_text = encoder_hidden_states_text.unsqueeze(0) + encoder_hidden_states_text = ( + encoder_hidden_states_text.unsqueeze(0) + ) if encoder_hidden_states_image is not None: if encoder_hidden_states_text is not None: encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states_text], dim=1) + [encoder_hidden_states_image, encoder_hidden_states_text], + dim=1, + ) else: encoder_hidden_states = encoder_hidden_states_image @@ -781,32 +1025,38 @@ class CausalMatrixGameWanModel(BaseDiT): num_frames=num_frames, frame_seqlen=post_patch_height * post_patch_width, num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size + local_attn_size=self.local_attn_size, ) if self.block_mask_keyboard is None: if not self.use_rope_keyboard: - self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size + self.block_mask_keyboard = ( + self._prepare_blockwise_causal_attn_mask_keyboard( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) ) else: - self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action( + self.block_mask_keyboard = ( + self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) + ) + if self.block_mask_mouse is None: + self.block_mask_mouse = ( + self._prepare_blockwise_causal_attn_mask_action( device=hidden_states.device, num_frames=num_frames, frame_seqlen=1, num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size + local_attn_size=self.local_attn_size, ) - if self.block_mask_mouse is None: - self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=1, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size ) if kv_cache is None: kv_cache = [None] * len(self.blocks) @@ -819,39 +1069,65 @@ class CausalMatrixGameWanModel(BaseDiT): for block_index, block in enumerate(self.blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: - kwargs.update({ - "kv_cache": kv_cache[block_index], - "kv_cache_mouse": kv_cache_mouse[block_index] if kv_cache_mouse else None, - "kv_cache_keyboard": kv_cache_keyboard[block_index] if kv_cache_keyboard else None, - "crossattn_cache": crossattn_cache[block_index] if crossattn_cache else None, - "current_start": current_start, - "cache_start": cache_start, - "num_frame_per_block": effective_num_frame_per_block, - }) + kwargs.update( + { + "kv_cache": kv_cache[block_index], + "kv_cache_mouse": kv_cache_mouse[block_index] + if kv_cache_mouse + else None, + "kv_cache_keyboard": kv_cache_keyboard[block_index] + if kv_cache_keyboard + else None, + "crossattn_cache": crossattn_cache[block_index] + if crossattn_cache + else None, + "current_start": current_start, + "cache_start": cache_start, + "num_frame_per_block": effective_num_frame_per_block, + } + ) hidden_states = self._gradient_checkpointing_func( - block, hidden_states, **kwargs) + block, hidden_states, **kwargs + ) else: - kwargs.update({ - "kv_cache": kv_cache[block_index], - "kv_cache_mouse": kv_cache_mouse[block_index] if kv_cache_mouse else None, - "kv_cache_keyboard": kv_cache_keyboard[block_index] if kv_cache_keyboard else None, - "crossattn_cache": crossattn_cache[block_index] if crossattn_cache else None, - "current_start": current_start, - "cache_start": cache_start, - }) - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, - block_mask=self.block_mask, - grid_sizes=grid_sizes, - mouse_cond=mouse_cond, - keyboard_cond=keyboard_cond, - block_mask_mouse=self.block_mask_mouse, - block_mask_keyboard=self.block_mask_keyboard, - num_frame_per_block=effective_num_frame_per_block, - use_rope_keyboard=self.use_rope_keyboard, - **kwargs) + kwargs.update( + { + "kv_cache": kv_cache[block_index], + "kv_cache_mouse": kv_cache_mouse[block_index] + if kv_cache_mouse + else None, + "kv_cache_keyboard": kv_cache_keyboard[block_index] + if kv_cache_keyboard + else None, + "crossattn_cache": crossattn_cache[block_index] + if crossattn_cache + else None, + "current_start": current_start, + "cache_start": cache_start, + } + ) + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + block_mask=self.block_mask, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + block_mask_mouse=self.block_mask_mouse, + block_mask_keyboard=self.block_mask_keyboard, + num_frame_per_block=effective_num_frame_per_block, + use_rope_keyboard=self.use_rope_keyboard, + **kwargs, + ) - temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + temb = temb.unflatten( + dim=0, sizes=(batch_size, post_patch_num_frames) + ).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk( + 2, dim=2 + ) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -859,32 +1135,47 @@ class CausalMatrixGameWanModel(BaseDiT): return hidden_states - def _forward_train( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + encoder_hidden_states_image: torch.Tensor + | list[torch.Tensor] + | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, - **kwargs + **kwargs, ) -> torch.Tensor: if mouse_cond is not None or keyboard_cond is not None: - assert self.action_config is not None and len(self.action_config) > 0 + assert ( + self.action_config is not None and len(self.action_config) > 0 + ) - if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): + if encoder_hidden_states is not None and not isinstance( + encoder_hidden_states, torch.Tensor + ): encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + if ( + isinstance(encoder_hidden_states_image, list) + and len(encoder_hidden_states_image) > 0 + ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None ctx = get_forward_context() batch = getattr(ctx, "forward_batch", None) - if batch is not None and getattr(batch, "image_latent", None) is not None: + if ( + batch is not None + and getattr(batch, "image_latent", None) is not None + ): image_latent = batch.image_latent - if isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5: + if ( + isinstance(image_latent, torch.Tensor) + and image_latent.ndim == 5 + and hidden_states.ndim == 5 + ): _, _, num_frames, _, _ = hidden_states.shape cond = image_latent if cond.shape[2] >= num_frames: @@ -905,7 +1196,9 @@ class CausalMatrixGameWanModel(BaseDiT): [hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1 ) - batch_size, num_channels, num_frames, height, width = hidden_states.shape + batch_size, num_channels, num_frames, height, width = ( + hidden_states.shape + ) p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h @@ -914,12 +1207,16 @@ class CausalMatrixGameWanModel(BaseDiT): d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + ( + post_patch_num_frames * get_sp_world_size(), + post_patch_height, + post_patch_width, + ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, - rope_theta=10000 + rope_theta=10000, ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) @@ -931,56 +1228,76 @@ class CausalMatrixGameWanModel(BaseDiT): num_frames=num_frames, frame_seqlen=post_patch_height * post_patch_width, num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size + local_attn_size=self.local_attn_size, ) if self.block_mask_keyboard is None: if not self.use_rope_keyboard: - self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size + self.block_mask_keyboard = ( + self._prepare_blockwise_causal_attn_mask_keyboard( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) ) else: - self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action( + self.block_mask_keyboard = ( + self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size, + ) + ) + if self.block_mask_mouse is None: + self.block_mask_mouse = ( + self._prepare_blockwise_causal_attn_mask_action( device=hidden_states.device, num_frames=num_frames, frame_seqlen=1, num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size + local_attn_size=self.local_attn_size, ) - if self.block_mask_mouse is None: - self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=1, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size ) hidden_states = self.patch_embedding(hidden_states) - grid_sizes = torch.tensor([post_patch_num_frames, post_patch_height, post_patch_width], - device=hidden_states.device) + grid_sizes = torch.tensor( + [post_patch_num_frames, post_patch_height, post_patch_width], + device=hidden_states.device, + ) hidden_states = hidden_states.flatten(2).transpose(1, 2) if timestep.dim() == 2: timestep = timestep.flatten() - temb, timestep_proj, encoder_hidden_states_text, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) - timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) + ( + temb, + timestep_proj, + encoder_hidden_states_text, + encoder_hidden_states_image, + ) = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image + ) + timestep_proj = timestep_proj.unflatten( + 1, (6, self.hidden_size) + ).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) if encoder_hidden_states_text is not None: if isinstance(encoder_hidden_states_text, list): encoder_hidden_states_text = encoder_hidden_states_text[0] elif encoder_hidden_states_text.ndim == 2: - encoder_hidden_states_text = encoder_hidden_states_text.unsqueeze(0) + encoder_hidden_states_text = ( + encoder_hidden_states_text.unsqueeze(0) + ) if encoder_hidden_states_image is not None: if encoder_hidden_states_text is not None: encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states_text], dim=1) + [encoder_hidden_states_image, encoder_hidden_states_text], + dim=1, + ) else: encoder_hidden_states = encoder_hidden_states_image else: @@ -997,14 +1314,18 @@ class CausalMatrixGameWanModel(BaseDiT): block_mask_mouse=self.block_mask_mouse, block_mask_keyboard=self.block_mask_keyboard, use_rope_keyboard=self.use_rope_keyboard, - num_frame_per_block=self.num_frame_per_block + num_frame_per_block=self.num_frame_per_block, ) # Initialize kv_cache, kv_cache_mouse, kv_cache_keyboard, crossattn_cache if not provided kv_cache = kwargs.pop("kv_cache", [None] * len(self.blocks)) kv_cache_mouse = kwargs.pop("kv_cache_mouse", [None] * len(self.blocks)) - kv_cache_keyboard = kwargs.pop("kv_cache_keyboard", [None] * len(self.blocks)) - crossattn_cache = kwargs.pop("crossattn_cache", [None] * len(self.blocks)) + kv_cache_keyboard = kwargs.pop( + "kv_cache_keyboard", [None] * len(self.blocks) + ) + crossattn_cache = kwargs.pop( + "crossattn_cache", [None] * len(self.blocks) + ) for i, block in enumerate(self.blocks): block_crossattn_cache = None @@ -1015,7 +1336,11 @@ class CausalMatrixGameWanModel(BaseDiT): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + block, + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, kv_cache=kv_cache[i], kv_cache_mouse=kv_cache_mouse[i], kv_cache_keyboard=kv_cache_keyboard[i], @@ -1027,7 +1352,8 @@ class CausalMatrixGameWanModel(BaseDiT): block_mask_mouse=self.block_mask_mouse, block_mask_keyboard=self.block_mask_keyboard, num_frame_per_block=self.num_frame_per_block, - use_rope_keyboard=self.use_rope_keyboard) + use_rope_keyboard=self.use_rope_keyboard, + ) else: hidden_states = block( hidden_states, @@ -1045,10 +1371,15 @@ class CausalMatrixGameWanModel(BaseDiT): block_mask_mouse=self.block_mask_mouse, block_mask_keyboard=self.block_mask_keyboard, num_frame_per_block=self.num_frame_per_block, - use_rope_keyboard=self.use_rope_keyboard) + use_rope_keyboard=self.use_rope_keyboard, + ) - temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2) - shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + temb = temb.unflatten( + dim=0, sizes=(batch_size, post_patch_num_frames) + ).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk( + 2, dim=2 + ) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) @@ -1057,7 +1388,7 @@ class CausalMatrixGameWanModel(BaseDiT): return hidden_states def forward(self, *args, **kwargs): - if kwargs.get('kv_cache') is not None: + if kwargs.get("kv_cache") is not None: return self._forward_inference(*args, **kwargs) else: return self._forward_train(*args, **kwargs) @@ -1067,7 +1398,7 @@ class CausalMatrixGameWanModel(BaseDiT): p_t, p_h, p_w = self.patch_size f, h, w = grid_sizes.tolist() - x = x[:, :f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c) + x = x[:, : f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c) x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) x = x.reshape(-1, c, f * p_t, h * p_h, w * p_w) return x @@ -1082,7 +1413,7 @@ class CausalMatrixGameWanModel(BaseDiT): nn.init.xavier_uniform_(self.patch_embedding.proj.weight.flatten(1)) for m in self.condition_embedder.time_embedder.mlp.modules(): if isinstance(m, nn.Linear): - nn.init.normal_(m.weight, std=.02) + nn.init.normal_(m.weight, std=0.02) nn.init.zeros_(self.proj_out.weight) @@ -1096,4 +1427,4 @@ class CausalMatrixGameWanModel(BaseDiT): if block.action_model.proj_keyboard.bias is not None: nn.init.zeros_(block.action_model.proj_keyboard.bias) except AttributeError: - pass \ No newline at end of file + pass diff --git a/fastvideo/models/dits/matrixgame/model.py b/fastvideo/models/dits/matrixgame/model.py index 201d03015..1fa520c6b 100644 --- a/fastvideo/models/dits/matrixgame/model.py +++ b/fastvideo/models/dits/matrixgame/model.py @@ -9,20 +9,32 @@ import torch.nn as nn from fastvideo.attention import DistributedAttention from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig from fastvideo.distributed.parallel_state import get_sp_world_size -from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, - RMSNorm, ScaleResidual, - ScaleResidualLayerNormScaleShift) +from fastvideo.layers.layernorm import ( + FP32LayerNorm, + LayerNormScaleShift, + RMSNorm, + ScaleResidual, + ScaleResidualLayerNormScaleShift, +) from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.mlp import MLP -from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, - get_rotary_pos_embed) -from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection +from fastvideo.layers.rotary_embedding import ( + _apply_rotary_emb, + get_rotary_pos_embed, +) +from fastvideo.layers.visual_embedding import ( + PatchEmbed, + TimestepEmbedder, + ModulateProjection, +) from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT -from fastvideo.models.dits.wanvideo import (WanSelfAttention, - WanI2VCrossAttention, - WanT2VCrossAttention, - WanImageEmbedding) +from fastvideo.models.dits.wanvideo import ( + WanSelfAttention, + WanI2VCrossAttention, + WanT2VCrossAttention, + WanImageEmbedding, +) from fastvideo.platforms import AttentionBackendEnum, current_platform # Import ActionModule @@ -41,11 +53,12 @@ class MatrixGameTimeImageEmbedding(nn.Module): super().__init__() self.time_embedder = TimestepEmbedder( - dim, frequency_embedding_size=time_freq_dim, act_layer="silu") - self.time_modulation = ModulateProjection(dim, - factor=6, - act_layer="silu") - + dim, frequency_embedding_size=time_freq_dim, act_layer="silu" + ) + self.time_modulation = ModulateProjection( + dim, factor=6, act_layer="silu" + ) + self.image_embedder = None if image_embed_dim is not None: self.image_embedder = WanImageEmbedding(image_embed_dim, dim) @@ -53,7 +66,7 @@ class MatrixGameTimeImageEmbedding(nn.Module): def forward( self, timestep: torch.Tensor, - encoder_hidden_states: torch.Tensor, # Kept for interface compatibility + encoder_hidden_states: torch.Tensor, # Kept for interface compatibility encoder_hidden_states_image: torch.Tensor | None = None, timestep_seq_len: int | None = None, ): @@ -61,17 +74,25 @@ class MatrixGameTimeImageEmbedding(nn.Module): timestep_proj = self.time_modulation(temb) # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states - + if encoder_hidden_states_image is not None: assert self.image_embedder is not None encoder_hidden_states_image = self.image_embedder( - encoder_hidden_states_image) + encoder_hidden_states_image + ) - encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), - device=temb.device, - dtype=temb.dtype) + encoder_hidden_states = torch.zeros( + (timestep.shape[0], 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype, + ) - return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + return ( + temb, + timestep_proj, + encoder_hidden_states, + encoder_hidden_states_image, + ) class MatrixGameCrossAttention(WanSelfAttention): @@ -87,7 +108,7 @@ class MatrixGameCrossAttention(WanSelfAttention): # compute query, key, value q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) - + if crossattn_cache is not None: if not crossattn_cache["is_init"]: crossattn_cache["is_init"] = True @@ -101,7 +122,7 @@ class MatrixGameCrossAttention(WanSelfAttention): else: k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) v = self.to_v(context)[0].view(b, -1, n, d) - + # compute attention x = self.attn(q, k, v) @@ -112,19 +133,20 @@ class MatrixGameCrossAttention(WanSelfAttention): class MatrixGameTransformerBlock(nn.Module): - - def __init__(self, - dim: int, - ffn_dim: int, - num_heads: int, - qk_norm: str = "rms_norm_across_heads", - cross_attn_norm: bool = False, - eps: float = 1e-6, - added_kv_proj_dim: int | None = None, - supported_attention_backends: tuple[AttentionBackendEnum, ...] - | None = None, - prefix: str = "", - action_config: dict | None = None): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + action_config: dict | None = None, + ): super().__init__() action_config = action_config or {} @@ -140,7 +162,8 @@ class MatrixGameTransformerBlock(nn.Module): head_size=dim // num_heads, causal=False, supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn1") + prefix=f"{prefix}.attn1", + ) self.hidden_dim = dim self.num_attention_heads = num_heads dim_head = dim // num_heads @@ -161,28 +184,28 @@ class MatrixGameTransformerBlock(nn.Module): eps=eps, elementwise_affine=True, dtype=torch.float32, - compute_dtype=torch.float32) + compute_dtype=torch.float32, + ) # 2. Cross-attention if added_kv_proj_dim is not None: # I2V - self.attn2 = WanI2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + self.attn2 = WanI2VCrossAttention( + dim, num_heads, qk_norm=qk_norm, eps=eps + ) else: # T2V - self.attn2 = WanT2VCrossAttention(dim, - num_heads, - qk_norm=qk_norm, - eps=eps) + self.attn2 = WanT2VCrossAttention( + dim, num_heads, qk_norm=qk_norm, eps=eps + ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, norm_type="layer", eps=eps, elementwise_affine=False, dtype=torch.float32, - compute_dtype=torch.float32) + compute_dtype=torch.float32, + ) # 2.1. Action Module Integration self.use_action_module = len(action_config) > 0 @@ -204,7 +227,7 @@ class MatrixGameTransformerBlock(nn.Module): temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], # Action Module specific args - grid_sizes: torch.Tensor | None = None, + grid_sizes: torch.Tensor, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, ) -> torch.Tensor: @@ -214,9 +237,16 @@ class MatrixGameTransformerBlock(nn.Module): if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() - ).chunk(6, dim=2) + ( + shift_msa, + scale_msa, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk( + 6, dim=2 + ) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) @@ -227,13 +257,20 @@ class MatrixGameTransformerBlock(nn.Module): else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( - 6, dim=1) + ( + shift_msa, + scale_msa, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + ) = e.chunk(6, dim=1) assert shift_msa.dtype == torch.float32 # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * - (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = ( + self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa + ).to(orig_dtype) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) @@ -249,9 +286,10 @@ class MatrixGameTransformerBlock(nn.Module): # Apply rotary embeddings cos, sin = freqs_cis - query, key = _apply_rotary_emb(query, cos, sin, - is_neox_style=False), _apply_rotary_emb( - key, cos, sin, is_neox_style=False) + query, key = ( + _apply_rotary_emb(query, cos, sin, is_neox_style=False), + _apply_rotary_emb(key, cos, sin, is_neox_style=False), + ) attn_output, _ = self.attn1(query, key, value) attn_output = attn_output.flatten(2) @@ -260,18 +298,24 @@ class MatrixGameTransformerBlock(nn.Module): null_shift = null_scale = torch.tensor([0], device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( - hidden_states, attn_output, gate_msa, null_shift, null_scale) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + hidden_states, attn_output, gate_msa, null_shift, null_scale + ) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # 2. Cross-attention - attn_output = self.attn2(norm_hidden_states, - context=encoder_hidden_states, - context_lens=None) + attn_output = self.attn2( + norm_hidden_states, context=encoder_hidden_states, context_lens=None + ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( - hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) - norm_hidden_states, hidden_states = norm_hidden_states.to( - orig_dtype), hidden_states.to(orig_dtype) + hidden_states, attn_output, 1, c_shift_msa, c_scale_msa + ) + norm_hidden_states, hidden_states = ( + norm_hidden_states.to(orig_dtype), + hidden_states.to(orig_dtype), + ) # ================= Action Module ================= if self.action_model is not None: @@ -280,9 +324,12 @@ class MatrixGameTransformerBlock(nn.Module): # ActionModule implementation takes hidden_states directly hidden_states = self.action_model( hidden_states, - grid_sizes[0], grid_sizes[1], grid_sizes[2], - mouse_cond, keyboard_cond, - num_frame_per_block=grid_sizes[0], + int(grid_sizes[0]), + int(grid_sizes[1]), + int(grid_sizes[2]), + mouse_cond, + keyboard_cond, + num_frame_per_block=int(grid_sizes[0]), ) # ================================================= @@ -293,6 +340,7 @@ class MatrixGameTransformerBlock(nn.Module): return hidden_states + _DEFAULT_MATRIXGAME_CONFIG = MatrixGameWanVideoConfig() @@ -302,15 +350,23 @@ class MatrixGameWanModel(BaseDiT): _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions - _supported_attention_backends = _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends + _supported_attention_backends = ( + _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends + ) param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.param_names_mapping - reverse_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping - lora_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping + reverse_param_names_mapping = ( + _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping + ) + lora_param_names_mapping = ( + _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping + ) - def __init__(self, - config: MatrixGameWanVideoConfig, - hf_config: dict[str, Any], - **kwargs) -> None: + def __init__( + self, + config: MatrixGameWanVideoConfig, + hf_config: dict[str, Any], + **kwargs, + ) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim @@ -320,10 +376,12 @@ class MatrixGameWanModel(BaseDiT): self.patch_size = config.patch_size # 1. Patch & position embedding - self.patch_embedding = PatchEmbed(in_chans=config.in_channels, - embed_dim=inner_dim, - patch_size=config.patch_size, - flatten=False) + self.patch_embedding = PatchEmbed( + in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False, + ) # 2. Condition embeddings self.condition_embedder = MatrixGameTimeImageEmbedding( @@ -333,57 +391,73 @@ class MatrixGameWanModel(BaseDiT): ) # 2.1. Get action config - self.action_config = getattr(config, 'action_config', {}) + self.action_config = getattr(config, "action_config", {}) # 3. Transformer blocks - self.blocks = nn.ModuleList([ - MatrixGameTransformerBlock( - inner_dim, - config.ffn_dim, - config.num_attention_heads, - config.qk_norm, - config.cross_attn_norm, - config.eps, - config.added_kv_proj_dim, - self._supported_attention_backends, - prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", - action_config=self.action_config) - for i in range(config.num_layers) - ]) + self.blocks = nn.ModuleList( + [ + MatrixGameTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + ) + for i in range(config.num_layers) + ] + ) # 4. Output norm & projection - self.norm_out = LayerNormScaleShift(inner_dim, - norm_type="layer", - eps=config.eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32) + self.norm_out = LayerNormScaleShift( + inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32, + ) self.proj_out = nn.Linear( - inner_dim, config.out_channels * math.prod(config.patch_size)) + inner_dim, config.out_channels * math.prod(config.patch_size) + ) self.scale_shift_table = nn.Parameter( - torch.randn(1, 2, inner_dim) / inner_dim**0.5) - + torch.randn(1, 2, inner_dim) / inner_dim**0.5 + ) + self.gradient_checkpointing = False - def forward(self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], - timestep: torch.LongTensor, - encoder_hidden_states_image: torch.Tensor - | list[torch.Tensor] | None = None, - # Action inputs - mouse_cond: torch.Tensor | None = None, - keyboard_cond: torch.Tensor | None = None, - **kwargs) -> torch.Tensor: - if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor + | list[torch.Tensor] + | None = None, + # Action inputs + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + if encoder_hidden_states is not None and not isinstance( + encoder_hidden_states, torch.Tensor + ): encoder_hidden_states = encoder_hidden_states[0] - if isinstance(encoder_hidden_states_image, - list) and len(encoder_hidden_states_image) > 0: + if ( + isinstance(encoder_hidden_states_image, list) + and len(encoder_hidden_states_image) > 0 + ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None - batch_size, num_channels, num_frames, height, width = hidden_states.shape + batch_size, num_channels, num_frames, height, width = ( + hidden_states.shape + ) p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h @@ -393,18 +467,25 @@ class MatrixGameWanModel(BaseDiT): d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, - post_patch_width), + ( + post_patch_num_frames * get_sp_world_size(), + post_patch_height, + post_patch_width, + ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, rope_theta=10000, - do_sp_sharding=True) + do_sp_sharding=True, + ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) - freqs_cis = (freqs_cos.float(), - freqs_sin.float()) if freqs_cos is not None else None + freqs_cis = ( + (freqs_cos.float(), freqs_sin.float()) + if freqs_cos is not None + else None + ) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) @@ -412,8 +493,14 @@ class MatrixGameWanModel(BaseDiT): if timestep.dim() == 2: timestep = timestep.flatten() - temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( - timestep, encoder_hidden_states, encoder_hidden_states_image) + ( + temb, + timestep_proj, + encoder_hidden_states, + encoder_hidden_states_image, + ) = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image + ) timestep_proj = timestep_proj.unflatten(1, (6, -1)) if encoder_hidden_states is not None: @@ -429,25 +516,30 @@ class MatrixGameWanModel(BaseDiT): if encoder_hidden_states_image is not None: if encoder_hidden_states is not None: encoder_hidden_states = torch.concat( - [encoder_hidden_states_image, encoder_hidden_states], dim=1) + [encoder_hidden_states_image, encoder_hidden_states], dim=1 + ) else: encoder_hidden_states = encoder_hidden_states_image # This is [F, H, W] for the ActionModule - grid_sizes = torch.tensor([ - post_patch_num_frames, post_patch_height, post_patch_width - ], - device=hidden_states.device) + grid_sizes = torch.tensor( + [post_patch_num_frames, post_patch_height, post_patch_width], + device=hidden_states.device, + ) # Blocks for block in self.blocks: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( - block, hidden_states, encoder_hidden_states, timestep_proj, + block, + hidden_states, + encoder_hidden_states, + timestep_proj, freqs_cis, grid_sizes=grid_sizes, mouse_cond=mouse_cond, - keyboard_cond=keyboard_cond) + keyboard_cond=keyboard_cond, + ) else: hidden_states = block( hidden_states, @@ -456,19 +548,27 @@ class MatrixGameWanModel(BaseDiT): freqs_cis, grid_sizes=grid_sizes, mouse_cond=mouse_cond, - keyboard_cond=keyboard_cond) + keyboard_cond=keyboard_cond, + ) # Output - shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, - dim=1) + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk( + 2, dim=1 + ) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, - post_patch_height, - post_patch_width, p_t, p_h, p_w, - -1) + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_t, + p_h, + p_w, + -1, + ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) - return output \ No newline at end of file + return output diff --git a/fastvideo/models/dits/stepvideo.py b/fastvideo/models/dits/stepvideo.py index 45295796c..f2011b2e2 100644 --- a/fastvideo/models/dits/stepvideo.py +++ b/fastvideo/models/dits/stepvideo.py @@ -685,3 +685,6 @@ class StepVideoModel(BaseDiT): output = rearrange(output, '(b f) c h w -> b c f h w', f=frame) return output + +# Entry point for model registry +EntryClass = StepVideoModel diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index ef98576bd..cb8252005 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -871,4 +871,6 @@ class WanTransformer3DModel(CachableDiT): return hidden_states + self.previous_residual_even else: return hidden_states + self.previous_residual_odd - \ No newline at end of file + +# Entry point for model registry +EntryClass = WanTransformer3DModel diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/models/encoders/clip.py index 940571b36..5ae9403d9 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/models/encoders/clip.py @@ -631,3 +631,6 @@ class CLIPVisionModel(ImageEncoder): weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + +# Entry point for model registry +EntryClass = [CLIPTextModel, CLIPVisionModel] diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index a279e1669..111f099a6 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -561,3 +561,6 @@ def _norm_and_concat_padded_batch( mask_flattened = mask.reshape(b, t, 1).expand(-1, -1, d * l) normed = normed.masked_fill(~mask_flattened, 0.0) return normed + +# Entry point for model registry +EntryClass = LTX2GemmaTextEncoderModel diff --git a/fastvideo/models/encoders/llama.py b/fastvideo/models/encoders/llama.py index ea9bdc928..b3fb19150 100644 --- a/fastvideo/models/encoders/llama.py +++ b/fastvideo/models/encoders/llama.py @@ -429,3 +429,6 @@ class LlamaModel(TextEncoder): weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + +# Entry point for model registry +EntryClass = LlamaModel diff --git a/fastvideo/models/encoders/qwen2_5.py b/fastvideo/models/encoders/qwen2_5.py index bd7c92f72..724e22889 100644 --- a/fastvideo/models/encoders/qwen2_5.py +++ b/fastvideo/models/encoders/qwen2_5.py @@ -385,3 +385,6 @@ class Qwen2_5_VLTextModel(TextEncoder): weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + +# Entry point for model registry +EntryClass = Qwen2_5_VLTextModel diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py index 48eff9286..fc3f6516c 100644 --- a/fastvideo/models/encoders/reason1.py +++ b/fastvideo/models/encoders/reason1.py @@ -351,3 +351,5 @@ class Reason1TextEncoder(TextEncoder): tensor.std(dim=-1, keepdim=True) + 1e-8 ) +# Entry point for model registry +EntryClass = Reason1TextEncoder diff --git a/fastvideo/models/encoders/siglip.py b/fastvideo/models/encoders/siglip.py index 8ae342280..2842102a4 100644 --- a/fastvideo/models/encoders/siglip.py +++ b/fastvideo/models/encoders/siglip.py @@ -417,3 +417,6 @@ class SiglipVisionModel(ImageEncoder): loaded_params.add(name) return loaded_params + +# Entry point for model registry +EntryClass = SiglipVisionModel diff --git a/fastvideo/models/encoders/stepllm.py b/fastvideo/models/encoders/stepllm.py index 2ea0654db..61b40f1d7 100644 --- a/fastvideo/models/encoders/stepllm.py +++ b/fastvideo/models/encoders/stepllm.py @@ -583,3 +583,6 @@ class STEP1TextEncoder(torch.nn.Module): self.device) if with_mask else None) y_mask = txt_tokens.attention_mask return y.transpose(0, 1), y_mask + +# Entry point for model registry +EntryClass = STEP1TextEncoder diff --git a/fastvideo/models/encoders/t5.py b/fastvideo/models/encoders/t5.py index 8240ac47c..e2f9b12da 100644 --- a/fastvideo/models/encoders/t5.py +++ b/fastvideo/models/encoders/t5.py @@ -719,3 +719,6 @@ class UMT5EncoderModel(TextEncoder): weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params + +# Entry point for model registry +EntryClass = [UMT5EncoderModel, T5EncoderModel] diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index bf53d324c..8437e9821 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -79,6 +79,7 @@ class ComponentLoader(ABC): "scheduler": (SchedulerLoader, "diffusers"), "transformer": (TransformerLoader, "diffusers"), "transformer_2": (TransformerLoader, "diffusers"), + "transformer_3": (TransformerLoader, "diffusers"), "vae": (VAELoader, "diffusers"), "audio_vae": (AudioDecoderLoader, "diffusers"), "audio_decoder": (AudioDecoderLoader, "diffusers"), @@ -90,6 +91,8 @@ class ComponentLoader(ABC): "image_processor": (ImageProcessorLoader, "transformers"), "feature_extractor": (ImageProcessorLoader, "transformers"), "image_encoder": (ImageEncoderLoader, "transformers"), + "upsampler": (UpsamplerLoader, "diffusers"), + "upsampler_2": (UpsamplerLoader, "diffusers"), } if module_type in module_loaders: @@ -558,7 +561,8 @@ class VAELoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the VAE based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) - class_name = config.get("_class_name") + class_name = config.pop("_class_name") + config.pop("_name_or_path", None) assert class_name is not None, ( "Model config does not contain a _class_name attribute. Only diffusers format is supported." ) @@ -731,6 +735,7 @@ class TransformerLoader(ComponentLoader): config = get_diffusers_config(model=model_path) hf_config = deepcopy(config) cls_name = config.pop("_class_name") + config.pop("_name_or_path", None) if cls_name is None: raise ValueError( "Model config does not contain a _class_name attribute. " @@ -745,7 +750,7 @@ class TransformerLoader(ComponentLoader): fastvideo_args.model_paths["transformer"] = model_path # Config from Diffusers supersedes fastvideo's model config - dit_config = fastvideo_args.pipeline_config.dit_config + dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config) dit_config.update_model_arch(config) model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) @@ -877,6 +882,50 @@ class SchedulerLoader(ComponentLoader): return scheduler +class UpsamplerLoader(ComponentLoader): + """Loader for upsamplers.""" + + def load(self, model_path: str, fastvideo_args: FastVideoArgs): + """Load the upsampler based on the model path, and inference args.""" + config_dict = get_diffusers_config(model=model_path) + class_name = config_dict.pop("_class_name", None) + + if class_name is None: + raise ValueError( + "Model config does not contain a _class_name attribute. " + "Only diffusers format is supported." + ) + + try: + upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[0]) + upsampler_cfg.update_model_config(config_dict) + except Exception as e: + upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[1]) + upsampler_cfg.update_model_config(config_dict) + + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + model = model_cls(upsampler_cfg) + + target_device = get_local_torch_device() + model = model.to(target_device, dtype=PRECISION_TO_TYPE[fastvideo_args.pipeline_config.upsampler_precision]) + + # Find all safetensors files + safetensors_list = glob.glob( + os.path.join(str(model_path), "*.safetensors")) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + + if len(safetensors_list) == 1: + loaded = safetensors_load_file(safetensors_list[0]) + else: + loaded = {} + for sf_file in safetensors_list: + loaded.update(safetensors_load_file(sf_file)) + + model.load_state_dict(loaded, strict=True) + + return model.eval() + class GenericComponentLoader(ComponentLoader): """Generic loader for components that don't have a specific loader.""" diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 2789c32cc..93bbfe151 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py +import ast import importlib import os import pickle @@ -100,7 +101,12 @@ _SCHEDULERS = { ("schedulers", "scheduling_rcm", "RCMScheduler"), } -_FAST_VIDEO_MODELS = { +_UPSAMPLERS = { + "SRTo720pUpsampler": ("upsamplers", "hunyuan15", "SRTo720pUpsampler"), + "SRTo1080pUpsampler": ("upsamplers", "hunyuan15", "SRTo1080pUpsampler"), +} + +_LEGACY_FAST_VIDEO_MODELS = { **_TEXT_TO_VIDEO_DIT_MODELS, **_IMAGE_TO_VIDEO_DIT_MODELS, **_TEXT_ENCODER_MODELS, @@ -108,8 +114,102 @@ _FAST_VIDEO_MODELS = { **_VAE_MODELS, **_AUDIO_MODELS, **_SCHEDULERS, + **_UPSAMPLERS, } +MODELS_PATH = os.path.dirname(__file__) + + +@lru_cache(maxsize=None) +def _discover_and_register_models() -> dict[str, tuple[str, str, str]]: + discovered_models: dict[str, tuple[str, str, str]] = {} + for root, dirs, files in os.walk(MODELS_PATH): + dirs[:] = [ + d for d in dirs + if not d.startswith(".") and d != "__pycache__" + ] + + for filename in files: + if not filename.endswith(".py"): + continue + + filepath = os.path.join(root, filename) + try: + with open(filepath, "r", encoding="utf-8") as f: + source = f.read() + tree = ast.parse(source, filename=filename) + + entry_class_node = None + first_class_def = None + + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "EntryClass": + entry_class_node = node + break + if first_class_def is None and isinstance(node, ast.ClassDef): + first_class_def = node + + if not entry_class_node or not first_class_def: + continue + + model_cls_name_list: list[str] = [] + value_node = entry_class_node.value + + if isinstance(value_node, ast.Name): + model_cls_name_list.append(value_node.id) + elif isinstance(value_node, (ast.List, ast.Tuple)): + for elt in value_node.elts: + if isinstance(elt, ast.Constant) and isinstance( + elt.value, str): + model_cls_name_list.append(elt.value) + elif isinstance(elt, ast.Name): + model_cls_name_list.append(elt.id) + + if not model_cls_name_list: + continue + + rel_dir = os.path.relpath(root, MODELS_PATH) + if rel_dir == ".": + continue + + rel_parts = rel_dir.split(os.sep) + component_name = rel_parts[0] + sub_parts = rel_parts[1:] + + if filename == "__init__.py": + mod_relname = ".".join(sub_parts) + else: + mod_base = filename[:-3] + mod_relname = ".".join(sub_parts + + [mod_base]) if sub_parts else mod_base + + for model_cls_str in model_cls_name_list: + if model_cls_str in discovered_models: + logger.warning( + "Duplicate architecture found: %s. Overwriting.", + model_cls_str) + discovered_models[model_cls_str] = ( + component_name, + mod_relname, + model_cls_str, + ) + + except Exception as e: + logger.warning("Could not parse %s to find models: %s", + filepath, e) + + return discovered_models + + +_DISCOVERED_MODELS = _discover_and_register_models() +_FAST_VIDEO_MODELS = dict(_DISCOVERED_MODELS) +for model_arch, spec in _LEGACY_FAST_VIDEO_MODELS.items(): + if model_arch in _FAST_VIDEO_MODELS: + continue + _FAST_VIDEO_MODELS[model_arch] = spec + _SUBPROCESS_COMMAND = [sys.executable, "-m", "fastvideo.models.dits.registry"] _T = TypeVar("_T") @@ -341,7 +441,8 @@ class _ModelRegistry: ModelRegistry = _ModelRegistry({ model_arch: _LazyRegisteredModel( - module_name=f"fastvideo.models.{component_name}.{mod_relname}", + module_name=(f"fastvideo.models.{component_name}.{mod_relname}" + if mod_relname else f"fastvideo.models.{component_name}"), component_name=component_name, class_name=cls_name, ) diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index ae8a657e6..aa22f74eb 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -685,3 +685,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, def __len__(self) -> int: return self.config.num_train_timesteps + +# Entry point for model registry +EntryClass = FlowMatchEulerDiscreteScheduler diff --git a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py index b82921574..919a23683 100644 --- a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py @@ -853,3 +853,6 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): def __len__(self): return self.config.num_train_timesteps + +# Entry point for model registry +EntryClass = FlowUniPCMultistepScheduler diff --git a/fastvideo/models/schedulers/scheduling_rcm.py b/fastvideo/models/schedulers/scheduling_rcm.py index cf96d1c05..71cc714b0 100644 --- a/fastvideo/models/schedulers/scheduling_rcm.py +++ b/fastvideo/models/schedulers/scheduling_rcm.py @@ -321,3 +321,6 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): def __len__(self) -> int: return self.config.num_train_timesteps + +# Entry point for model registry +EntryClass = RCMScheduler diff --git a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py index 78ef23fba..b96118b31 100644 --- a/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py +++ b/fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py @@ -160,4 +160,6 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): def set_shift(self, shift: float) -> None: self.shift = shift - \ No newline at end of file + +# Entry point for model registry +EntryClass = SelfForcingFlowMatchScheduler diff --git a/fastvideo/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_unipc_multistep.py index b44160eee..d76fb8ca5 100644 --- a/fastvideo/models/schedulers/scheduling_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_unipc_multistep.py @@ -1094,4 +1094,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): return noisy_samples def __len__(self): - return self.config.num_train_timesteps \ No newline at end of file + return self.config.num_train_timesteps + +# Entry point for model registry +EntryClass = UniPCMultistepScheduler diff --git a/fastvideo/models/upsamplers/hunyuan15.py b/fastvideo/models/upsamplers/hunyuan15.py new file mode 100644 index 000000000..0d31a517e --- /dev/null +++ b/fastvideo/models/upsamplers/hunyuan15.py @@ -0,0 +1,170 @@ +# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE +# +# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any +# output and results therefrom are provided "AS IS" without any express or implied warranties of +# any kind including any warranties of title, merchantability, noninfringement, course of dealing, +# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the +# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of +# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a +# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise +# of rights and permissions under this agreement. +# See the License for the specific language governing permissions and limitations under the License. + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor + +from fastvideo.models.vaes.hunyuan15vae import ( + HunyuanVideo15CausalConv3d, + HunyuanVideo15RMS_norm, +) + +from fastvideo.layers.activation import get_act_fn +from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig + +class HunyuanVideo15ResnetBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + non_linearity: str = "swish", + ) -> None: + super().__init__() + out_channels = out_channels or in_channels + + self.nonlinearity = get_act_fn(non_linearity) + + self.norm1 = HunyuanVideo15RMS_norm(in_channels, images=False) + self.conv1 = HunyuanVideo15CausalConv3d(in_channels, out_channels, kernel_size=3) + + self.norm2 = HunyuanVideo15RMS_norm(out_channels, images=False) + self.conv2 = HunyuanVideo15CausalConv3d(out_channels, out_channels, kernel_size=3) + + self.nin_shortcut = None + if in_channels != out_channels: + self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.nin_shortcut is not None: + residual = self.nin_shortcut(residual) + + return hidden_states + residual + +class SRResidualCausalBlock3D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.block = nn.Sequential( + HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3), + nn.SiLU(inplace=True), + HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.block(x) + + +class SRTo720pUpsampler(nn.Module): + + def __init__( + self, + config: SRTo720pUpsamplerConfig, + ): + super().__init__() + self.in_conv = HunyuanVideo15CausalConv3d(config.in_channels, config.hidden_channels, kernel_size=3) + self.blocks = nn.ModuleList([SRResidualCausalBlock3D(config.hidden_channels) for _ in range(config.num_blocks)]) + self.out_conv = HunyuanVideo15CausalConv3d(config.hidden_channels, config.out_channels, kernel_size=3) + self.global_residual = bool(config.global_residual) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + y = self.in_conv(x) + for blk in self.blocks: + y = blk(y) + y = self.out_conv(y) + if self.global_residual and (y.shape == residual.shape): + y = y + residual + return y + + +class SRTo1080pUpsampler(nn.Module): + + def __init__( + self, + config: SRTo1080pUpsamplerConfig, + ): + super().__init__() + self.num_res_blocks = config.num_res_blocks + self.block_out_channels = config.block_out_channels + self.z_channels = config.z_channels + + block_in = config.block_out_channels[0] + self.conv_in = HunyuanVideo15CausalConv3d(config.z_channels, block_in, kernel_size=3) + + self.up = nn.ModuleList() + for i_level, ch in enumerate(config.block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(HunyuanVideo15ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + self.up.append(up) + + self.norm_out = HunyuanVideo15RMS_norm(block_in, images=False) + self.conv_out = HunyuanVideo15CausalConv3d(block_in, config.out_channels, kernel_size=3) + + self.gradient_checkpointing = False + self.is_residual = config.is_residual + + def forward(self, z: Tensor, target_shape: Sequence[int] = None) -> Tensor: + """ + Args: + z: (B, C, T, H, W) + target_shape: (H, W) + """ + if target_shape is not None and z.shape[-2:] != target_shape: + bsz = z.shape[0] + z = rearrange(z, "b c f h w -> (b f) c h w") + z = F.interpolate(z, size=target_shape, mode="bilinear", align_corners=False) + z = rearrange(z, "(b f) c h w -> b c f h w", b=bsz) + + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if hasattr(self.up[i_level], "upsample"): + h = self.up[i_level].upsample(h) + + # end + h = self.norm_out(h) + h = get_act_fn("swish")(h) + h = self.conv_out(h) + return h \ No newline at end of file diff --git a/fastvideo/models/vaes/hunyuan15vae.py b/fastvideo/models/vaes/hunyuan15vae.py index 66201bd09..309b1dce3 100644 --- a/fastvideo/models/vaes/hunyuan15vae.py +++ b/fastvideo/models/vaes/hunyuan15vae.py @@ -700,4 +700,7 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE): else: z = posterior.mode() dec = self.decode(z) - return dec \ No newline at end of file + return dec + +# Entry point for model registry +EntryClass = AutoencoderKLHunyuanVideo15 diff --git a/fastvideo/models/vaes/hunyuanvae.py b/fastvideo/models/vaes/hunyuanvae.py index d0f614ea3..e60e22509 100644 --- a/fastvideo/models/vaes/hunyuanvae.py +++ b/fastvideo/models/vaes/hunyuanvae.py @@ -850,3 +850,6 @@ class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): z = posterior.mode() dec = self.decode(z) return dec + +# Entry point for model registry +EntryClass = AutoencoderKLHunyuanVideo diff --git a/fastvideo/models/vaes/hyworldvae.py b/fastvideo/models/vaes/hyworldvae.py index b78d0f555..ede1d8fb0 100644 --- a/fastvideo/models/vaes/hyworldvae.py +++ b/fastvideo/models/vaes/hyworldvae.py @@ -15,15 +15,977 @@ # See the License for the specific language governing permissions and # limitations under the License. -from fastvideo.models.vaes.hunyuan15vae import AutoencoderKLHunyuanVideo15 -from fastvideo.configs.models.vaes import Hunyuan15VAEConfig +from typing import List, Optional, Tuple, Union -class AutoencoderKLHYWorld(AutoencoderKLHunyuanVideo15): - # TODO(mingjia): add temporal caching support for HYWorld VAE +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange + +from fastvideo.layers.activation import get_act_fn +from fastvideo.configs.models.vaes import Hunyuan15VAEConfig +from fastvideo.models.vaes.common import ParallelTiledVAE +from fastvideo.models.vaes.hunyuan15vae import ( + HunyuanVideo15RMS_norm as HYWorldRMS_norm, + HunyuanVideo15AttnBlock as HYWorldAttnBlock, +) + +# Cache size for temporal feature caching (number of frames to cache) +CACHE_T = 2 + +class HYWorldCausalConv3d(nn.Module): + """Causal Conv3d with optional cache support for temporal feature caching.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int, int]] = 3, + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Union[int, Tuple[int, int, int]] = 0, + dilation: Union[int, Tuple[int, int, int]] = 1, + bias: bool = True, + pad_mode: str = "replicate", + ) -> None: + super().__init__() + + kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size + + self.pad_mode = pad_mode + # Padding format: (W_left, W_right, H_left, H_right, T_left, T_right) + self.time_causal_padding = ( + kernel_size[0] // 2, # W_left (spatial) + kernel_size[0] // 2, # W_right (spatial) + kernel_size[1] // 2, # H_left (spatial) + kernel_size[1] // 2, # H_right (spatial) + kernel_size[2] - 1, # T_left (temporal causal padding) + 0, # T_right (no future padding for causal) + ) + + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) + + def forward(self, hidden_states: torch.Tensor, cache_x: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Forward pass with optional temporal caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + cache_x: Optional cached frames from previous chunk, shape (B, C, CACHE_T, H, W) + When provided, uses cached frames instead of padding for temporal dimension. + """ + padding = list(self.time_causal_padding) + + if cache_x is not None and self.time_causal_padding[4] > 0: # Has temporal padding and cache + cache_x = cache_x.to(hidden_states.device) + # Concatenate cached frames with current input on temporal dimension + hidden_states = torch.cat([cache_x, hidden_states], dim=2) + # Reduce temporal padding since we have cached frames + padding[4] -= cache_x.shape[2] + + hidden_states = F.pad(hidden_states, padding, mode=self.pad_mode) + return self.conv(hidden_states) + + +class HYWorldUpsample(nn.Module): + """Hierarchical upsampling with temporal/spatial support and optional caching.""" + + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = HYWorldCausalConv3d(in_channels, out_channels * factor, kernel_size=3) + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + def forward( + self, + x: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + first_chunk: bool = False, + ): + """ + Forward pass with optional temporal caching. + + Args: + x: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + first_chunk: Whether this is the first chunk (affects temporal upsample behavior) + """ + r1 = 2 if self.add_temporal_upsample else 1 + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + h = self.conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv(x) + + if self.add_temporal_upsample: + if first_chunk: + # First chunk: only spatial upsample + h = rearrange(h, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + h = h[:, : h.shape[1] // 2] + # Compute the shortcut part + shortcut = rearrange(x, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + shortcut = shortcut.repeat_interleave(repeats=self.repeats // 2, dim=1) + elif feat_cache is None and x.shape[2] > 1: + # No cache and multiple frames: separate first frame and rest + h_first = h[:, :, :1, :, :] + h_rest = h[:, :, 1:, :, :] + x_first = x[:, :, :1, :, :] + x_rest = x[:, :, 1:, :, :] + + # First frame: only spatial upsample + h_first = rearrange(h_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + h_first = h_first[:, : h_first.shape[1] // 2] + shortcut_first = rearrange(x_first, "b (r2 r3 c) f h w -> b c f (h r2) (w r3)", r2=2, r3=2) + shortcut_first = shortcut_first.repeat_interleave(repeats=self.repeats // 2, dim=1) + out_first = h_first + shortcut_first + + # Remaining frames: spatio-temporal upsample + h_rest = rearrange(h_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut_rest = rearrange(x_rest, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut_rest = shortcut_rest.repeat_interleave(repeats=self.repeats, dim=1) + out_rest = h_rest + shortcut_rest + + return torch.cat([out_first, out_rest], dim=2) + else: + # Subsequent chunks with cache: spatio-temporal upsample + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = rearrange(x, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = shortcut.repeat_interleave(repeats=self.repeats, dim=1) + else: + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + + return h + shortcut + + +class HYWorldDownsample(nn.Module): + """Hierarchical downsampling with temporal/spatial support and optional caching.""" + + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + self.conv = HYWorldCausalConv3d(in_channels, out_channels // factor, kernel_size=3) + + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + def forward( + self, + x: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + ): + """ + Forward pass with optional temporal caching. + + Args: + x: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + """ + r1 = 2 if self.add_temporal_downsample else 1 + + # Apply conv with caching + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + h = self.conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + h = self.conv(x) # Change the channel, ready for spatial or temporal downsample + + if self.add_temporal_downsample: + if x.shape[2] == 1: + # Single frame: only spatial downsample + h = rearrange(h, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) + h = torch.cat([h, h], dim=1) + # Compute the shortcut part + shortcut = rearrange(x, "b c f (h r2) (w r3) -> b (r2 r3 c) f h w", r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size // 2, T, H, W).mean(dim=2) + else: + # Multiple frames: full spatio-temporal downsample + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + # Shortcut computation + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + else: + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + + return h + shortcut + + +class HYWorldResnetBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + non_linearity: str = "swish", + ) -> None: + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.nonlinearity = get_act_fn(non_linearity) + + self.norm1 = HYWorldRMS_norm(in_channels, images=False) + self.conv1 = HYWorldCausalConv3d(in_channels, out_channels, kernel_size=3) + + self.norm2 = HYWorldRMS_norm(out_channels, images=False) + self.conv2 = HYWorldCausalConv3d(out_channels, out_channels, kernel_size=3) + if in_channels != out_channels: + self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + ) -> torch.Tensor: + """ + Forward pass with optional temporal feature caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + """ + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + # apply the feature cacheing mechanism + if feat_cache is not None and feat_idx is not None: + # Retrieve the current layer index. + idx = feat_idx[0] + + # Clone the last CACHE_T frames from the current input to store for the next step. + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + + # Handle boundary conditions: if the current temporal chunk is too short (< 2 frames) + # and we have a previous cache, prepend the last frame of the previous cache. + # This ensures sufficient temporal context for the convolution kernel. + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # actually this means chunk 1 after chunk 0 + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + + # Apply the convolution layer using `hidden_states`and the *previous* cached state. + hidden_states = self.conv1(hidden_states, feat_cache[idx]) + + # Update the cache for this layer with the newly prepared context ('cache_x'). + feat_cache[idx] = cache_x + + # Increment the global layer index. + feat_idx[0] += 1 + else: + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + # Second conv with caching + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + hidden_states = self.conv2(hidden_states, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + hidden_states = self.conv2(hidden_states) + + if self.in_channels != self.out_channels: + residual = self.nin_shortcut(residual) + + return hidden_states + residual + + +class HYWorldMidBlock(nn.Module): + """Mid block with attention and resnet blocks, with optional caching support.""" + + def __init__( + self, + in_channels: int, + num_layers: int = 1, + add_attention: bool = True, + ) -> None: + super().__init__() + self.add_attention = add_attention + + # There is always at least one resnet + resnets = [ + HYWorldResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + ) + ] + attentions = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append(HYWorldAttnBlock(in_channels)) + else: + attentions.append(None) + + resnets.append( + HYWorldResnetBlock( + in_channels=in_channels, + out_channels=in_channels, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + ) -> torch.Tensor: + """ + Forward pass with optional temporal caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + """ + hidden_states = self.resnets[0](hidden_states, feat_cache, feat_idx) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = attn(hidden_states) + hidden_states = resnet(hidden_states, feat_cache, feat_idx) + + return hidden_states + + +class HYWorldDownBlock3D(nn.Module): + """Down block with resnet blocks and optional downsampling, with caching support.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + downsample_out_channels: Optional[int] = None, + add_temporal_downsample: int = True, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + HYWorldResnetBlock( + in_channels=in_channels, + out_channels=out_channels, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if downsample_out_channels is not None: + self.downsamplers = nn.ModuleList( + [ + HYWorldDownsample( + out_channels, + out_channels=downsample_out_channels, + add_temporal_downsample=add_temporal_downsample, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + ) -> torch.Tensor: + """ + Forward pass with optional temporal caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + """ + for resnet in self.resnets: + hidden_states = resnet(hidden_states, feat_cache, feat_idx) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, feat_cache, feat_idx) + + return hidden_states + + +class HYWorldUpBlock3D(nn.Module): + """Up block with resnet blocks and optional upsampling, with caching support.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + upsample_out_channels: Optional[int] = None, + add_temporal_upsample: bool = True, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + HYWorldResnetBlock( + in_channels=input_channels, + out_channels=out_channels, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if upsample_out_channels is not None: + self.upsamplers = nn.ModuleList( + [ + HYWorldUpsample( + out_channels, + out_channels=upsample_out_channels, + add_temporal_upsample=add_temporal_upsample, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + first_chunk: bool = False, + ) -> torch.Tensor: + """ + Forward pass with optional temporal caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + first_chunk: Whether this is the first chunk (for upsampling behavior) + """ + if torch.is_grad_enabled() and self.gradient_checkpointing: + for resnet in self.resnets: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) + else: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, feat_cache, feat_idx) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, feat_cache, feat_idx, first_chunk) + + return hidden_states + + +class HYWorldEncoder3D(nn.Module): + r""" + 3D vae encoder for HunyuanImageRefiner with optional temporal caching. + """ def __init__( self, - config: Hunyuan15VAEConfig, + in_channels: int = 3, + out_channels: int = 64, + block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024, 1024), + layers_per_block: int = 2, + temporal_compression_ratio: int = 4, + spatial_compression_ratio: int = 16, + downsample_match_channel: bool = True, ) -> None: - AutoencoderKLHunyuanVideo15.__init__(self, config) + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.group_size = block_out_channels[-1] // self.out_channels + self.block_out_channels = block_out_channels + + self.conv_in = HYWorldCausalConv3d(in_channels, block_out_channels[0], kernel_size=3) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + input_channel = block_out_channels[0] + for i in range(len(block_out_channels)): + add_spatial_downsample = i < np.log2(spatial_compression_ratio) + output_channel = block_out_channels[i] + if not add_spatial_downsample: + down_block = HYWorldDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + downsample_out_channels=None, + add_temporal_downsample=False, + ) + input_channel = output_channel + else: + add_temporal_downsample = i >= np.log2(spatial_compression_ratio // temporal_compression_ratio) + downsample_out_channels = block_out_channels[i + 1] if downsample_match_channel else output_channel + down_block = HYWorldDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + downsample_out_channels=downsample_out_channels, + add_temporal_downsample=add_temporal_downsample, + ) + input_channel = downsample_out_channels + + self.down_blocks.append(down_block) + + self.mid_block = HYWorldMidBlock(in_channels=block_out_channels[-1]) + + self.norm_out = HYWorldRMS_norm(block_out_channels[-1], images=False) + self.conv_act = nn.SiLU() + self.conv_out = HYWorldCausalConv3d(block_out_channels[-1], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + ) -> torch.Tensor: + """ + Forward pass with optional temporal caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + """ + # conv_in with caching + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + hidden_states = self.conv_in(hidden_states, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + hidden_states = self.conv_in(hidden_states) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for down_block in self.down_blocks: + hidden_states = self._gradient_checkpointing_func(down_block, hidden_states) + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) + else: + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states, feat_cache, feat_idx) + hidden_states = self.mid_block(hidden_states, feat_cache, feat_idx) + + batch_size, _, frame, height, width = hidden_states.shape + short_cut = hidden_states.view(batch_size, -1, self.group_size, frame, height, width).mean(dim=2) + + hidden_states = self.norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + + # conv_out with caching + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + hidden_states = self.conv_out(hidden_states, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + hidden_states = self.conv_out(hidden_states) + + hidden_states += short_cut + + return hidden_states + + +class HYWorldDecoder3D(nn.Module): + r""" + Causal decoder for 3D video-like data used for HunyuanImage-1.5 Refiner with optional temporal caching. + """ + + def __init__( + self, + in_channels: int = 32, + out_channels: int = 3, + block_out_channels: Tuple[int, ...] = (1024, 1024, 512, 256, 128), + layers_per_block: int = 2, + spatial_compression_ratio: int = 16, + temporal_compression_ratio: int = 4, + upsample_match_channel: bool = True, + ): + super().__init__() + self.layers_per_block = layers_per_block + self.in_channels = in_channels + self.out_channels = out_channels + self.block_out_channels = block_out_channels + self.repeat = block_out_channels[0] // self.in_channels + + self.conv_in = HYWorldCausalConv3d(self.in_channels, block_out_channels[0], kernel_size=3) + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = HYWorldMidBlock(in_channels=block_out_channels[0]) + + # up + input_channel = block_out_channels[0] + for i in range(len(block_out_channels)): + output_channel = block_out_channels[i] + + add_spatial_upsample = i < np.log2(spatial_compression_ratio) + add_temporal_upsample = i < np.log2(temporal_compression_ratio) + if add_spatial_upsample or add_temporal_upsample: + upsample_out_channels = block_out_channels[i + 1] if upsample_match_channel else output_channel + up_block = HYWorldUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + upsample_out_channels=upsample_out_channels, + add_temporal_upsample=add_temporal_upsample, + ) + input_channel = upsample_out_channels + else: + up_block = HYWorldUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + upsample_out_channels=None, + add_temporal_upsample=False, + ) + input_channel = output_channel + + self.up_blocks.append(up_block) + + # out + self.norm_out = HYWorldRMS_norm(block_out_channels[-1], images=False) + self.conv_act = nn.SiLU() + self.conv_out = HYWorldCausalConv3d(block_out_channels[-1], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: Optional[List[Optional[torch.Tensor]]] = None, + feat_idx: Optional[List[int]] = None, + first_chunk: bool = False, + ) -> torch.Tensor: + """ + Forward pass with optional temporal caching. + + Args: + hidden_states: Input tensor of shape (B, C, T, H, W) + feat_cache: List of cached features for each conv layer + feat_idx: List containing current cache index [idx] + first_chunk: Whether this is the first chunk (for upsampling behavior) + """ + # conv_in with caching + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + hidden_states = self.conv_in(hidden_states, feat_cache[idx]) + hidden_states.repeat_interleave( + repeats=self.repeat, dim=1 + ) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + hidden_states = self.conv_in(hidden_states) + hidden_states.repeat_interleave(repeats=self.repeat, dim=1) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) + for up_block in self.up_blocks: + hidden_states = self._gradient_checkpointing_func(up_block, hidden_states) + else: + hidden_states = self.mid_block(hidden_states, feat_cache, feat_idx) + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states, feat_cache, feat_idx, first_chunk) + + # post-process + hidden_states = self.norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + + # conv_out with caching + if feat_cache is not None and feat_idx is not None: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + hidden_states = self.conv_out(hidden_states, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Revised from HunyuanVideo-1.5 with temporal caching support for HY-WorldPlay integration. + """ + + _supports_gradient_checkpointing = True + + def __init__( + self, + config: Hunyuan15VAEConfig, # HYWorld use the same VAE architecture as HunyuanVideo-1.5 + ) -> None: + nn.Module.__init__(self) + ParallelTiledVAE.__init__(self, config) + + self.encoder: Optional[HYWorldEncoder3D] = None + self.decoder: Optional[HYWorldDecoder3D] = None + + if config.load_encoder: + self.encoder = HYWorldEncoder3D( + in_channels=config.in_channels, + out_channels=config.latent_channels * 2, + block_out_channels=config.block_out_channels, + layers_per_block=config.layers_per_block, + temporal_compression_ratio=config.temporal_compression_ratio, + spatial_compression_ratio=config.spatial_compression_ratio, + downsample_match_channel=config.downsample_match_channel, + ) + + if config.load_decoder: + self.decoder = HYWorldDecoder3D( + in_channels=config.latent_channels, + out_channels=config.out_channels, + block_out_channels=list(reversed(config.block_out_channels)), + layers_per_block=config.layers_per_block, + temporal_compression_ratio=config.temporal_compression_ratio, + spatial_compression_ratio=config.spatial_compression_ratio, + upsample_match_channel=config.upsample_match_channel, + ) + + # TODO: Add spatial tiling. + self.use_tiling = False + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + self.tile_sample_min_num_frames = 2000 # Fill in a random large number, as hy1.5 vae does not use temporal tiling + + # Cache-related attributes (initialized in clear_cache) + self._conv_num: int = 0 + self._conv_idx: List[int] = [0] + self._feat_map: List[Optional[torch.Tensor]] = [] + self._enc_conv_num: int = 0 + self._enc_conv_idx: List[int] = [0] + self._enc_feat_map: List[Optional[torch.Tensor]] = [] + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "decoder": ( + sum(1 for m in self.decoder.modules() if isinstance(m, HYWorldCausalConv3d)) + if self.decoder is not None + else 0 + ), + "encoder": ( + sum(1 for m in self.encoder.modules() if isinstance(m, HYWorldCausalConv3d)) + if self.encoder is not None + else 0 + ), + } + + def clear_cache(self) -> None: + """ + Initialize/clear the feature cache for chunk-based encoding/decoding. + + This should be called before starting a new encode/decode sequence to ensure + the cache is properly initialized. + """ + # Cache for decoder + self._conv_num = self._cached_conv_counts["decoder"] + self._conv_idx = [0] + self._feat_map: List[Optional[torch.Tensor]] = [None] * self._conv_num + + # Cache for encoder + self._enc_conv_num = self._cached_conv_counts["encoder"] + self._enc_conv_idx = [0] + self._enc_feat_map: List[Optional[torch.Tensor]] = [None] * self._enc_conv_num + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + """ + Encode video with temporal caching for chunk-based processing. + + This processes video in chunks (first frame separately, then 4 frames at a time) + while maintaining temporal context through caching. This matches the HY-WorldPlay + behavior for memory-efficient long video encoding. + + Args: + x: Input video tensor of shape (B, C, T, H, W) + + Returns: + Encoded latent tensor + """ + assert self.encoder is not None, "Encoder not loaded" + _, _, num_frame, _, _ = x.shape + + self.clear_cache() + + # Process in chunks: first frame alone, then groups of 4 frames + iter_ = 1 + (num_frame - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + # First frame + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + else: + # Subsequent frames in groups of 4 + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], dim=2) + + self.clear_cache() + return out + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + """ + Decode latents with temporal caching for chunk-based processing. + + This processes latents one frame at a time while maintaining temporal context + through caching. This matches the HY-WorldPlay behavior for memory-efficient + long video decoding. + + Args: + z: Latent tensor of shape (B, C, T, H, W) + + Returns: + Decoded video tensor + """ + assert self.decoder is not None, "Decoder not loaded" + _, _, num_frame, _, _ = z.shape + + self.clear_cache() + + # Process one frame at a time with caching + for i in range(num_frame): + self._conv_idx = [0] + if i == 0: + # First frame with first_chunk=True + out = self.decoder( + z[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + # Subsequent frames + out_ = self.decoder( + z[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=False, + ) + out = torch.cat([out, out_], dim=2) + + self.clear_cache() + return out + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + Generator for sampling. + """ + x = sample + + # encode() uses temporal caching by default (via _encode) + posterior = self.encode(x) + + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + + # decode() uses temporal caching by default (via _decode) + dec = self.decode(z) + + return dec + +# Entry point for model registry +EntryClass = AutoencoderKLHYWorld diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index e0a01db83..5f921dccf 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -1802,3 +1802,6 @@ class LTX2CausalVideoAutoencoder(nn.Module): if previous_chunk is not None: previous_weights = previous_weights.clamp(min=1e-8) yield previous_chunk / previous_weights + +# Entry point for model registry +EntryClass = LTX2CausalVideoAutoencoder diff --git a/fastvideo/models/vaes/stepvideovae.py b/fastvideo/models/vaes/stepvideovae.py index 28354ec7c..2d3d3ffbf 100644 --- a/fastvideo/models/vaes/stepvideovae.py +++ b/fastvideo/models/vaes/stepvideovae.py @@ -1138,3 +1138,6 @@ class AutoencoderKLStepvideo(nn.Module, ParallelTiledVAE): z = posterior.mode() dec = self.decode(z) return dec + +# Entry point for model registry +EntryClass = AutoencoderKLStepvideo diff --git a/fastvideo/models/vaes/wanvae.py b/fastvideo/models/vaes/wanvae.py index 2d1502be8..8204639bd 100644 --- a/fastvideo/models/vaes/wanvae.py +++ b/fastvideo/models/vaes/wanvae.py @@ -1319,3 +1319,5 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): dec = self.decode(z) return dec +# Entry point for model registry +EntryClass = AutoencoderKLWan diff --git a/fastvideo/pipelines/__init__.py b/fastvideo/pipelines/__init__.py index 3f8d1cbf8..f391f0b06 100644 --- a/fastvideo/pipelines/__init__.py +++ b/fastvideo/pipelines/__init__.py @@ -12,10 +12,9 @@ from fastvideo.logger import init_logger from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch -from fastvideo.pipelines.pipeline_registry import (PipelineType, - get_pipeline_registry) -from fastvideo.utils import (maybe_download_model, - verify_model_config_and_directory) +from fastvideo.pipelines.pipeline_registry import PipelineType +from fastvideo.registry import get_model_info +from fastvideo.utils import maybe_download_model logger = init_logger(__name__) @@ -42,30 +41,17 @@ def build_pipeline( # fastvideo_args.downloaded_model_path = model_path logger.info("Model path: %s", model_path) - config = verify_model_config_and_directory(model_path) - pipeline_name = config.get("_class_name") - - if fastvideo_args.override_pipeline_cls_name: - logger.info("Overriding pipeline class name from %s to %s", - pipeline_name, fastvideo_args.override_pipeline_cls_name) - pipeline_name = fastvideo_args.override_pipeline_cls_name - - if pipeline_name is None: - raise ValueError( - "Model config does not contain a _class_name attribute. " - "Only diffusers format is supported.") - - # Get the appropriate pipeline registry based on pipeline_type logger.info( "Building pipeline of type: %s", pipeline_type.value if isinstance( pipeline_type, PipelineType) else pipeline_type) - pipeline_registry = get_pipeline_registry(pipeline_type) - if isinstance(pipeline_type, str): - pipeline_type = PipelineType.from_string(pipeline_type) - - pipeline_cls = pipeline_registry.resolve_pipeline_cls( - pipeline_name, pipeline_type, fastvideo_args.workload_type) + model_info = get_model_info( + model_path=model_path, + pipeline_type=pipeline_type, + workload_type=fastvideo_args.workload_type, + override_pipeline_cls_name=fastvideo_args.override_pipeline_cls_name, + ) + pipeline_cls = model_info.pipeline_cls # instantiate the pipelines pipeline = pipeline_cls(model_path, fastvideo_args) diff --git a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py index c21720100..8724659e3 100644 --- a/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos/cosmos2_5_pipeline.py @@ -5,8 +5,8 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.stages import (ConditioningStage, - Cosmos25DenoisingStage, - Cosmos25LatentPreparationStage, + Cosmos25AutoDenoisingStage, + Cosmos25AutoLatentPreparationStage, DecodingStage, InputValidationStage, Cosmos25TextEncodingStage, Cosmos25TimestepPreparationStage) @@ -42,13 +42,13 @@ class Cosmos2_5Pipeline(ComposedPipelineBase): scheduler=self.get_module("scheduler"))) self.add_stage(stage_name="latent_preparation_stage", - stage=Cosmos25LatentPreparationStage( + stage=Cosmos25AutoLatentPreparationStage( scheduler=self.get_module("scheduler"), transformer=self.get_module("transformer"), vae=self.get_module("vae"))) self.add_stage(stage_name="denoising_stage", - stage=Cosmos25DenoisingStage( + stage=Cosmos25AutoDenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"))) diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py new file mode 100644 index 000000000..fe7a5a106 --- /dev/null +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_2sr_pipeline.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Hunyuan video diffusion pipeline implementation. + +This module contains an implementation of the Hunyuan video diffusion pipeline +using the modular pipeline architecture. +""" +import torch +import time +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import ( + ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage, + LatentPreparationStage, TextEncodingStage, TimestepPreparationStage, + Hy15ImageEncodingStage, SRDenoisingStage) +from fastvideo.distributed.parallel_state import get_local_torch_device + +# TODO(will): move PRECISION_TO_TYPE to better place + +logger = init_logger(__name__) + + +class HunyuanVideo152SRPipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", + "transformer", "transformer_2", "transformer_3", "scheduler", + "upsampler", "upsampler_2" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2") + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2") + ], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) + + self.add_stage(stage_name="image_encoding_stage", + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="sr_720p_latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer_2"))) + + self.add_stage(stage_name="sr_720p_denoising_stage", + stage=SRDenoisingStage( + transformer=self.get_module("transformer_2"), + scheduler=self.get_module("scheduler"), + upsampler=self.get_module("upsampler"))) + + self.add_stage(stage_name="sr_1080p_latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer_3"))) + + self.add_stage(stage_name="sr_1080p_denoising_stage", + stage=SRDenoisingStage( + transformer=self.get_module("transformer_3"), + scheduler=self.get_module("scheduler"), + upsampler=self.get_module("upsampler_2"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + @torch.no_grad() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """ + Generate a video or image using the pipeline. + + Args: + batch: The batch to generate from. + fastvideo_args: The inference arguments. + Returns: + ForwardBatch: The batch with the generated video or image. + """ + if not self.post_init_called: + self.post_init() + + self.get_module("transformer").to(get_local_torch_device()) + # Execute each stage + logger.info("Running pipeline stages: %s", + self._stage_name_mapping.keys()) + # logger.info("Batch: %s", batch) + batch = self.input_validation_stage(batch, fastvideo_args) + batch = self.prompt_encoding_stage_primary(batch, fastvideo_args) + batch = self.conditioning_stage(batch, fastvideo_args) + batch = self.timestep_preparation_stage(batch, fastvideo_args) + batch = self.latent_preparation_stage(batch, fastvideo_args) + batch = self.image_encoding_stage(batch, fastvideo_args) + batch = self.denoising_stage(batch, fastvideo_args) + self.get_module("transformer").to("cpu") + + # 720p SR + self.get_module("transformer_2").to(get_local_torch_device()) + batch.lq_latents = batch.latents + batch.latents = None + batch.height = 720 + batch.width = 1280 + batch.num_inference_steps_sr = 6 + batch = self.sr_720p_latent_preparation_stage(batch, fastvideo_args) + batch = self.image_encoding_stage(batch, fastvideo_args) + batch = self.sr_720p_denoising_stage(batch, fastvideo_args) + self.get_module("transformer_2").to("cpu") + + # 1080p SR + self.get_module("transformer_3").to(get_local_torch_device()) + batch.lq_latents = batch.latents + batch.latents = None + batch.height = 1072 + batch.width = 1920 + batch.num_inference_steps_sr = 8 + batch = self.sr_1080p_latent_preparation_stage(batch, fastvideo_args) + batch = self.image_encoding_stage(batch, fastvideo_args) + batch = self.sr_1080p_denoising_stage(batch, fastvideo_args) + self.get_module("transformer_3").to("cpu") + + start_time = time.time() + batch = self.decoding_stage(batch, fastvideo_args) + end_time = time.time() + logger.info("Decoding time: %s seconds", end_time - start_time) + + # Return the output + return batch + + +EntryClass = HunyuanVideo152SRPipeline diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py new file mode 100644 index 000000000..297a2637f --- /dev/null +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_i2v_pipeline.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Hunyuan video diffusion pipeline implementation. + +This module contains an implementation of the Hunyuan video diffusion pipeline +using the modular pipeline architecture. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + DenoisingStage, InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage, + Hy15ImageEncodingStage) + +# TODO(will): move PRECISION_TO_TYPE to better place + +logger = init_logger(__name__) + + +class HunyuanVideo15ImageToVideoPipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", + "transformer", "scheduler" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2") + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2") + ], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) + + self.add_stage(stage_name="image_encoding_stage", + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = HunyuanVideo15ImageToVideoPipeline diff --git a/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py new file mode 100644 index 000000000..fd3bf87fc --- /dev/null +++ b/fastvideo/pipelines/basic/hunyuan15/hunyuan15_sr_pipeline.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Hunyuan video diffusion pipeline implementation. + +This module contains an implementation of the Hunyuan video diffusion pipeline +using the modular pipeline architecture. +""" +import torch +import time +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import ( + ConditioningStage, DecodingStage, DenoisingStage, InputValidationStage, + LatentPreparationStage, TextEncodingStage, TimestepPreparationStage, + Hy15ImageEncodingStage, SRDenoisingStage) +from fastvideo.distributed.parallel_state import get_local_torch_device + +# TODO(will): move PRECISION_TO_TYPE to better place + +logger = init_logger(__name__) + + +class HunyuanVideo15SRPipeline(ComposedPipelineBase): + + _required_config_modules = [ + "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "vae", + "transformer", "transformer_2", "scheduler", "upsampler" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage_primary", + stage=TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2") + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2") + ], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) + + self.add_stage(stage_name="image_encoding_stage", + stage=Hy15ImageEncodingStage(image_encoder=None, + image_processor=None)) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="sr_latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer_2"))) + + self.add_stage(stage_name="sr_denoising_stage", + stage=SRDenoisingStage( + transformer=self.get_module("transformer_2"), + scheduler=self.get_module("scheduler"), + upsampler=self.get_module("upsampler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + @torch.no_grad() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """ + Generate a video or image using the pipeline. + + Args: + batch: The batch to generate from. + fastvideo_args: The inference arguments. + Returns: + ForwardBatch: The batch with the generated video or image. + """ + if not self.post_init_called: + self.post_init() + + # Execute each stage + logger.info("Running pipeline stages: %s", + self._stage_name_mapping.keys()) + # logger.info("Batch: %s", batch) + self.get_module("transformer").to(get_local_torch_device()) + batch = self.input_validation_stage(batch, fastvideo_args) + batch = self.prompt_encoding_stage_primary(batch, fastvideo_args) + batch = self.conditioning_stage(batch, fastvideo_args) + batch = self.timestep_preparation_stage(batch, fastvideo_args) + batch = self.latent_preparation_stage(batch, fastvideo_args) + batch = self.image_encoding_stage(batch, fastvideo_args) + batch = self.denoising_stage(batch, fastvideo_args) + self.get_module("transformer").to("cpu") + + self.get_module("transformer_2").to(get_local_torch_device()) + batch.lq_latents = batch.latents + batch.latents = None + batch.height = batch.height_sr + batch.width = batch.width_sr + batch = self.sr_latent_preparation_stage(batch, fastvideo_args) + batch = self.image_encoding_stage(batch, fastvideo_args) + batch = self.sr_denoising_stage(batch, fastvideo_args) + self.get_module("transformer_2").to("cpu") + + start_time = time.time() + batch = self.decoding_stage(batch, fastvideo_args) + end_time = time.time() + logger.info("Decoding time: %s seconds", end_time - start_time) + + # Return the output + return batch + + +EntryClass = HunyuanVideo15SRPipeline diff --git a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py index 72f51e735..4ba3ecd2e 100644 --- a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py +++ b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py @@ -16,7 +16,7 @@ from fastvideo.pipelines.stages import (DecodingStage, InputValidationStage, LTX2AudioDecodingStage, LTX2DenoisingStage, LTX2LatentPreparationStage, - TextEncodingStage) + LTX2TextEncodingStage) logger = init_logger(__name__) @@ -40,7 +40,7 @@ class LTX2Pipeline(ComposedPipelineBase): self.add_stage( stage_name="prompt_encoding_stage", - stage=TextEncodingStage( + stage=LTX2TextEncodingStage( text_encoders=[self.get_module("text_encoder")], tokenizers=[self.get_module("tokenizer")], ), diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 725a96f11..d9a7f6419 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -15,6 +15,8 @@ import torch from fastvideo.configs.pipelines import PipelineConfig from fastvideo.distributed import ( maybe_init_distributed_environment_and_model_parallel, get_world_group) +from fastvideo.distributed.communication_op import ( + warmup_sequence_parallel_communication) from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.logger import init_logger from fastvideo.profiler import get_or_create_profiler @@ -157,6 +159,10 @@ class ComposedPipelineBase(ABC): logger.info("Creating pipeline stages...") self.create_pipeline_stages(self.fastvideo_args) + # Warmup NCCL communicators for sequence parallelism to avoid + # slow first forward pass due to lazy initialization + warmup_sequence_parallel_communication() + def initialize_training_pipeline(self, training_args: TrainingArgs): raise NotImplementedError( "if training_mode is True, the pipeline must implement this method") @@ -287,6 +293,7 @@ class ComposedPipelineBase(ABC): # remove keys that are not pipeline modules model_index.pop("_class_name") model_index.pop("_diffusers_version") + model_index.pop("_name_or_path", None) model_index.pop("workload_type", None) if "boundary_ratio" in model_index and model_index[ "boundary_ratio"] is not None: diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index dd93a46c2..2e258a056 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -123,6 +123,7 @@ class ForwardBatch: # Latent tensors latents: torch.Tensor | None = None + lq_latents: torch.Tensor | None = None raw_latent_shape: tuple[int, ...] | None = None noise_pred: torch.Tensor | None = None image_latent: torch.Tensor | None = None @@ -148,6 +149,8 @@ class ForwardBatch: # Original dimensions (before VAE scaling) height: list[int] | int | None = None width: list[int] | int | None = None + height_sr: list[int] | int | None = None + width_sr: list[int] | int | None = None fps: list[int] | int | None = None # Timesteps @@ -158,6 +161,7 @@ class ForwardBatch: # Scheduler parameters num_inference_steps: int = 50 + num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 guidance_scale_2: float | None = None guidance_rescale: float = 0.0 diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index fe957ad2d..f085d6749 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -10,10 +10,11 @@ from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage from fastvideo.pipelines.stages.conditioning import ConditioningStage from fastvideo.pipelines.stages.decoding import DecodingStage -from fastvideo.pipelines.stages.denoising import (Cosmos25DenoisingStage, - CosmosDenoisingStage, - DenoisingStage, - DmdDenoisingStage) +from fastvideo.pipelines.stages.denoising import ( + Cosmos25AutoDenoisingStage, Cosmos25DenoisingStage, + Cosmos25V2WDenoisingStage, Cosmos25T2WDenoisingStage, CosmosDenoisingStage, + DenoisingStage, DmdDenoisingStage) +from fastvideo.pipelines.stages.sr_denoising import SRDenoisingStage from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.image_encoding import ( ImageEncodingStage, MatrixGameImageEncodingStage, RefImageEncodingStage, @@ -22,11 +23,13 @@ from fastvideo.pipelines.stages.image_encoding import ( from fastvideo.pipelines.stages.input_validation import InputValidationStage from fastvideo.pipelines.stages.latent_preparation import ( Cosmos25LatentPreparationStage, CosmosLatentPreparationStage, - LatentPreparationStage) + Cosmos25AutoLatentPreparationStage, Cosmos25T2WLatentPreparationStage, + Cosmos25V2WLatentPreparationStage, LatentPreparationStage) from fastvideo.pipelines.stages.ltx2_audio_decoding import LTX2AudioDecodingStage from fastvideo.pipelines.stages.ltx2_denoising import LTX2DenoisingStage from fastvideo.pipelines.stages.ltx2_latent_preparation import ( LTX2LatentPreparationStage) +from fastvideo.pipelines.stages.ltx2_text_encoding import LTX2TextEncodingStage from fastvideo.pipelines.stages.matrixgame_denoising import ( MatrixGameCausalDenoisingStage) from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage @@ -50,6 +53,9 @@ __all__ = [ "LatentPreparationStage", "CosmosLatentPreparationStage", "Cosmos25LatentPreparationStage", + "Cosmos25T2WLatentPreparationStage", + "Cosmos25V2WLatentPreparationStage", + "Cosmos25AutoLatentPreparationStage", "LTX2LatentPreparationStage", "LTX2AudioDecodingStage", "ConditioningStage", @@ -60,7 +66,12 @@ __all__ = [ "HYWorldDenoisingStage", "CosmosDenoisingStage", "Cosmos25DenoisingStage", + "Cosmos25T2WDenoisingStage", + "Cosmos25V2WDenoisingStage", + "Cosmos25AutoDenoisingStage", "LTX2DenoisingStage", + "LTX2TextEncodingStage", + "SRDenoisingStage", "EncodingStage", "DecodingStage", "ImageEncodingStage", diff --git a/fastvideo/pipelines/stages/hyworld_denoising.py b/fastvideo/pipelines/stages/hyworld_denoising.py index b66112ccb..6fc87b086 100644 --- a/fastvideo/pipelines/stages/hyworld_denoising.py +++ b/fastvideo/pipelines/stages/hyworld_denoising.py @@ -20,6 +20,7 @@ from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.utils import dict_to_3d_list from fastvideo.models.dits.hyworld.retrieval_context import ( generate_points_in_sphere, select_aligned_memory_frames) +from fastvideo.models.dits.hyworld.pose import pose_to_input, compute_latent_num logger = init_logger(__name__) @@ -58,7 +59,7 @@ class HYWorldDenoisingStage(DenoisingStage): - viewmats: torch.Tensor | None - Camera view matrices (B, T, 4, 4) - Ks: torch.Tensor | None - Camera intrinsics (B, T, 3, 3) - action: torch.Tensor | None - Action conditioning (B, T) - - chunk_latent_frames: int - Number of frames per chunk (default: 4) + - chunk_latent_frames: int - Number of frames per chunk (default: 16 for bidirectional model) These can be passed via batch.extra dict or as direct attributes. fastvideo_args: The inference arguments. @@ -82,16 +83,38 @@ class HYWorldDenoisingStage(DenoisingStage): action = getattr(batch, "action", None) or batch.extra.get( "action", None) chunk_latent_frames = (getattr(batch, "chunk_latent_frames", None) - or batch.extra.get("chunk_latent_frames", 4)) + or batch.extra.get("chunk_latent_frames", 16) + ) # 16 for bidirectional model stabilization_level = 15 points_local = (getattr(batch, "points_local", None) or batch.extra.get("points_local", None)) + # If viewmats/Ks/action not provided, convert from pose string if viewmats is None or Ks is None or action is None: - raise ValueError( - "viewmats, Ks, and action are required for HYWorld denoising. " - "Please provide them in batch.extra['viewmats'], batch.extra['Ks'], " - "and batch.extra['action']") + pose = getattr(batch, "pose", None) or batch.extra.get("pose", None) + if pose is None: + raise ValueError( + "Either pose string or (viewmats, Ks, action) must be provided. " + "Provide pose in batch.pose or batch.extra['pose'], or " + "viewmats/Ks/action in batch.extra.") + + # Get num_frames from batch + num_frames = batch.num_frames + if isinstance(num_frames, list): + num_frames = num_frames[0] + + # Calculate number of latents and convert pose to tensors + latent_num = compute_latent_num(num_frames) + viewmats, Ks, action = pose_to_input(pose, latent_num) + + # 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) + + logger.info( + "Converted pose '%s' to viewmats, Ks, and action for %d latent frames", + pose, latent_num) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( @@ -266,12 +289,9 @@ class HYWorldDenoisingStage(DenoisingStage): latents_concat = torch.concat( [latent_model_input, cond_latents_input], dim=1) - # Note: Unlike some other pipelines, HYWorld runs CFG sequentially (two passes) - # rather than batching pos/neg together, following the original implementation latents_concat = self.scheduler.scale_model_input( latents_concat, t) - # Keep batch size 1 for sequential CFG t_expand_txt = t.unsqueeze(0) t_expand = timestep_input viewmats_input = viewmats_input.to(device) @@ -287,7 +307,6 @@ class HYWorldDenoisingStage(DenoisingStage): batch.is_cfg_negative = False # Prepare transformer kwargs with HYWorld-specific inputs - # Note: batch size 1 for sequential CFG (matching original HY-WorldPlay) transformer_kwargs = { **image_kwargs, "timestep": t_expand, @@ -397,34 +416,26 @@ class HYWorldDenoisingStage(DenoisingStage): [V.is_tensor, V.with_dims(5)]) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) - # Check for HYWorld-specific inputs + # Check for HYWorld-specific inputs - either pose string OR + # (viewmats, Ks, action) must be provided + pose = getattr(batch, "pose", None) or batch.extra.get("pose", None) viewmats = getattr(batch, "viewmats", None) or batch.extra.get( "viewmats", None) Ks = getattr(batch, "Ks", None) or batch.extra.get("Ks", None) action = getattr(batch, "action", None) or batch.extra.get( "action", None) - if viewmats is None: - result.add_failure( - "viewmats", - "viewmats must be provided in batch.extra['viewmats'] or as batch.viewmats", - ) - else: - result.add_check("viewmats", viewmats, V.is_tensor) + has_pose = pose is not None + has_direct_inputs = (viewmats is not None and Ks is not None + and action is not None) - if Ks is None: + if not has_pose and not has_direct_inputs: result.add_failure( - "Ks", "Ks must be provided in batch.extra['Ks'] or as batch.Ks") - else: - result.add_check("Ks", Ks, V.is_tensor) - - if action is None: - result.add_failure( - "action", - "action must be provided in batch.extra['action'] or as batch.action", + "pose_or_camera_inputs", + "Either pose string OR (viewmats, Ks, action) must be provided. " + "Provide pose in batch.pose or batch.extra['pose'], or " + "viewmats/Ks/action in batch.extra.", ) - else: - result.add_check("action", action, V.is_tensor) result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index 2768668fc..3559ae1db 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -238,7 +238,8 @@ class InputValidationStage(PipelineStage): fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify input validation stage inputs.""" result = VerificationResult() - result.add_check("seed", batch.seed, [V.not_none, V.positive_int]) + # Cosmos-Predict2.5 default seed is 0; allow non-negative seeds here. + result.add_check("seed", batch.seed, [V.not_none, V.non_negative_int]) result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt, V.positive_int) result.add_check( diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index 4af23ed2d..bf23a454e 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -3,6 +3,7 @@ Latent preparation stage for diffusion pipelines. """ +import os from typing import Any import numpy as np @@ -450,10 +451,6 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): batch: ForwardBatch, fastvideo_args: FastVideoArgs, ) -> ForwardBatch: - # Differences vs `CosmosLatentPreparationStage`: channel convention, seed usage, - # and `padding_mask` for concat_padding_mask=True. - - # Determine batch size if isinstance(batch.prompt, list): batch_size = len(batch.prompt) elif batch.prompt is not None: @@ -463,8 +460,6 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): batch_size *= batch.num_videos_per_prompt - # Match `compare_pipelines.py`: initialize noise in fp32, then run the - # denoising computation in bf16. dtype = torch.float32 device = get_local_torch_device() generator = batch.generator @@ -483,7 +478,6 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): latent_width = width // vae_scale_factor_spatial num_latent_frames = (num_frames - 1) // vae_scale_factor_temporal + 1 - # Cosmos 2.5 convention: transformer in_channels == latent channels num_channels_latents = self.transformer.config.in_channels shape = (batch_size, num_channels_latents, num_latent_frames, @@ -492,16 +486,16 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): init_latents = None conditioning_latents = None video = None + image_conditioning = False - if hasattr(batch, 'video') and batch.video is not None: - video = batch.video - elif hasattr(batch, 'pil_image') and batch.pil_image is not None: + if hasattr(batch, 'pil_image') and batch.pil_image is not None: vae_scale_factor_spatial = 8 image_processor = ImageProcessor( vae_scale_factor=vae_scale_factor_spatial) - processed_image = image_processor.preprocess( + processed_image = image_processor._preprocess_cosmos25( batch.pil_image, height, width) video = processed_image.unsqueeze(2) + image_conditioning = True video = video.to(device=device, dtype=torch.bfloat16) elif hasattr( batch, @@ -509,25 +503,77 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): if isinstance(batch.preprocessed_image, torch.Tensor): if batch.preprocessed_image.dim() == 4: video = batch.preprocessed_image.unsqueeze(2) + image_conditioning = True elif batch.preprocessed_image.dim() == 5: video = batch.preprocessed_image + elif hasattr(batch, "video_latent") and isinstance( + batch.video_latent, torch.Tensor): + if batch.video_latent.dim() == 5: + video = batch.video_latent else: logger.info( "CosmosLatentPreparationStage - No video input sources found") if video is not None: - num_cond_frames = video.size(2) - if num_cond_frames >= num_frames: + # Avoid missing CUDA fallback kernels in some builds by using fp16 contiguous input. + if isinstance(video, torch.Tensor): + video = video.to(device=device, + dtype=torch.float16).contiguous() + num_video_frames = int(video.size(2)) + + num_cond_frames_eff: int | None = None + try: + ncf = getattr(batch, "num_cond_frames", None) + if isinstance(ncf, int) and ncf > 0: + num_cond_frames_eff = min(int(ncf), num_video_frames) + except Exception: + num_cond_frames_eff = None + if num_cond_frames_eff is None: + num_cond_frames_eff = 1 if ( + not image_conditioning) else num_video_frames + + if num_video_frames >= num_cond_frames_eff: + cond_video = video[:, :, -num_cond_frames_eff:] + else: + pad = video[:, :, + -1:].repeat(1, 1, + num_cond_frames_eff - num_video_frames, + 1, 1) + cond_video = torch.cat([video, pad], dim=2) + + if num_cond_frames_eff >= num_frames: + cond_video = cond_video[:, :, -num_frames:] num_cond_latent_frames = (num_frames - 1) // vae_scale_factor_temporal + 1 - video = video[:, :, -num_frames:] + video = cond_video else: - num_cond_latent_frames = (num_cond_frames - + num_cond_latent_frames = (num_cond_frames_eff - 1) // vae_scale_factor_temporal + 1 - num_padding_frames = num_frames - num_cond_frames - last_frame = video[:, :, -1:] + num_padding_frames = num_frames - num_cond_frames_eff + last_frame = cond_video[:, :, -1:] padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1) - video = torch.cat([video, padding], dim=2) + if image_conditioning: + padding = cond_video.new_full( + (cond_video.size(0), cond_video.size(1), + num_padding_frames, cond_video.size(3), + cond_video.size(4)), + -1.0, + ) + video = torch.cat([cond_video, padding], dim=2) + + if os.environ.get("FASTVIDEO_COSMOS25_LOG_KNOBS", + "0") in ("1", "true", "True"): + logger.info( + "[Cosmos2.5 latent_prep] using conditioning input: image=%s video_latent=%s video_path=%s | " + "num_video_frames=%s num_cond_frames=%s num_frames=%s", + bool(image_conditioning), + isinstance(getattr(batch, "video_latent", None), + torch.Tensor), + getattr(batch, "video_path", None), + num_video_frames, + num_cond_frames_eff, + num_frames, + ) if self.vae is not None: self.vae = self.vae.to(device) @@ -588,18 +634,14 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): if latents is None: seed = int(batch.seed if batch.seed is not None else 0) - # Use arch-invariant RNG to match Cosmos2.5 reference sampling. latents_fp32 = self._arch_invariant_randn(shape, seed=seed, device=device, dtype=torch.float32) latents = latents_fp32.to(torch.bfloat16) else: - # If latents are supplied, keep compute dtype consistent with Cosmos sampling. latents = latents.to(device=device, dtype=torch.bfloat16) - # Cosmos2.5 starts from unit Gaussian noise (no extra sigma_max scaling). - padding_shape = (batch_size, 1, num_latent_frames, latent_height, latent_width) ones_padding = latents.new_ones(padding_shape) @@ -618,9 +660,7 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): uncond_mask = uncond_indicator * ones_padding + ( 1 - uncond_indicator) * zeros_padding - # Cosmos 2.5 requires a spatial padding mask when concat_padding_mask=True - padding_mask = latents.new_ones(batch_size, 1, latent_height, - latent_width) + padding_mask = latents.new_zeros(batch_size, 1, height, width) batch.latents = latents batch.raw_latent_shape = latents.shape @@ -681,3 +721,127 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): [V.is_tensor, V.with_dims(5)]) result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple) return result + + +class Cosmos25T2WLatentPreparationStage(PipelineStage): + """Cosmos 2.5 Text2World latent preparation.""" + + def __init__(self, scheduler, transformer) -> None: + super().__init__() + self.scheduler = scheduler + self.transformer = transformer + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + if isinstance(batch.prompt, list): + batch_size = len(batch.prompt) + elif batch.prompt is not None: + batch_size = 1 + else: + batch_size = batch.prompt_embeds[0].shape[0] + batch_size *= batch.num_videos_per_prompt + + device = get_local_torch_device() + latents = batch.latents + num_frames = batch.num_frames + height = batch.height + width = batch.width + if height is None or width is None: + raise ValueError("Height and width must be provided") + + vae_scale_factor_spatial = 8 + vae_scale_factor_temporal = 4 + latent_height = height // vae_scale_factor_spatial + latent_width = width // vae_scale_factor_spatial + num_latent_frames = (num_frames - 1) // vae_scale_factor_temporal + 1 + + # Cosmos 2.5 convention: transformer in_channels == latent channels + num_channels_latents = self.transformer.config.in_channels + shape = (batch_size, num_channels_latents, num_latent_frames, + latent_height, latent_width) + + if latents is None: + seed = int(batch.seed if batch.seed is not None else 0) + latents_fp32 = Cosmos25LatentPreparationStage._arch_invariant_randn( + shape, + seed=seed, + device=device, + dtype=torch.float32, + ) + latents = latents_fp32.to(torch.bfloat16) + else: + latents = latents.to(device=device, dtype=torch.bfloat16) + + batch.latents = latents + batch.raw_latent_shape = latents.shape + + batch.conditioning_latents = None + batch.cond_indicator = None + batch.uncond_indicator = None + batch.cond_mask = None + batch.uncond_mask = None + if hasattr(batch, "padding_mask"): + batch.padding_mask = None + + return batch + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + return Cosmos25LatentPreparationStage.verify_input( + self, batch, fastvideo_args) # type: ignore[misc] + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + return Cosmos25LatentPreparationStage.verify_output( + self, batch, fastvideo_args) # type: ignore[misc] + + +class Cosmos25V2WLatentPreparationStage(Cosmos25LatentPreparationStage): + """Cosmos 2.5 V2W/I2W latent preparation stage (conditioning-aware).""" + + +class Cosmos25AutoLatentPreparationStage(PipelineStage): + """Route Cosmos 2.5 latent prep to T2W vs V2W/I2W.""" + + def __init__(self, scheduler, transformer, vae=None) -> None: + super().__init__() + self._t2w = Cosmos25T2WLatentPreparationStage( + scheduler=scheduler, + transformer=transformer, + ) + self._v2w = Cosmos25V2WLatentPreparationStage( + scheduler=scheduler, + transformer=transformer, + vae=vae, + ) + + @staticmethod + def _has_conditioning_input(batch: ForwardBatch) -> bool: + return (getattr(batch, "pil_image", None) is not None + or getattr(batch, "preprocessed_image", None) is not None + or bool(getattr(batch, "video_path", None)) or isinstance( + getattr(batch, "video_latent", None), torch.Tensor)) + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + if self._has_conditioning_input(batch): + return self._v2w.forward(batch, fastvideo_args) + return self._t2w.forward(batch, fastvideo_args) + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + if self._has_conditioning_input(batch): + return self._v2w.verify_input(batch, fastvideo_args) + return self._t2w.verify_input(batch, fastvideo_args) + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + if self._has_conditioning_input(batch): + return self._v2w.verify_output(batch, fastvideo_args) + return self._t2w.verify_output(batch, fastvideo_args) diff --git a/fastvideo/pipelines/stages/ltx2_text_encoding.py b/fastvideo/pipelines/stages/ltx2_text_encoding.py new file mode 100644 index 000000000..d790e754f --- /dev/null +++ b/fastvideo/pipelines/stages/ltx2_text_encoding.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +LTX2-specific text encoding stage with sequence parallelism broadcast support. + +When running with sequence parallelism (SP), the Gemma text encoder is only +executed on rank 0, and the embeddings are broadcast to all other ranks. +This avoids I/O contention from all ranks loading the Gemma model simultaneously. +""" + +import torch + +from fastvideo.distributed.parallel_state import get_sp_group +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.text_encoding import TextEncodingStage + +logger = init_logger(__name__) + + +class LTX2TextEncodingStage(TextEncodingStage): + """ + LTX2 text encoding stage with sequence parallelism support. + + When SP is enabled (sp_world_size > 1), only rank 0 runs the text encoder + and broadcasts embeddings to other ranks. This avoids I/O contention from + all ranks loading the Gemma model simultaneously, which can cause text + encoding to take 100+ seconds instead of ~5 seconds. + """ + + @torch.no_grad() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + sp_group = get_sp_group() + sp_world_size = sp_group.world_size + sp_rank = sp_group.rank_in_group + + # Single GPU or no SP: use parent implementation + if sp_world_size <= 1: + return super().forward(batch, fastvideo_args) + + # SP enabled: only rank 0 encodes, then broadcasts + if sp_rank == 0: + logger.info( + "[LTX2TextEncodingStage] SP rank 0: running text encoding") + # Run encoding on rank 0 + result_batch = super().forward(batch, fastvideo_args) + + # Build broadcast dict from batch + broadcast_dict = self._build_broadcast_dict(result_batch) + + # Broadcast to other ranks + logger.info( + "[LTX2TextEncodingStage] SP rank 0: broadcasting %d tensors", + len(broadcast_dict)) + sp_group.broadcast_tensor_dict(broadcast_dict, src=0) + + return result_batch + else: + logger.info( + "[LTX2TextEncodingStage] SP rank %d: receiving broadcast", + sp_rank) + # Other ranks: receive broadcast and populate batch + broadcast_dict = sp_group.broadcast_tensor_dict(None, src=0) + + # Unpack into batch + self._unpack_broadcast_to_batch(batch, broadcast_dict) + + logger.info( + "[LTX2TextEncodingStage] SP rank %d: received %d prompt embeds", + sp_rank, len(batch.prompt_embeds)) + + return batch + + def _build_broadcast_dict(self, + batch: ForwardBatch) -> dict[str, torch.Tensor]: + """Build dict of tensors to broadcast from rank 0.""" + d: dict[str, torch.Tensor] = {} + + # Use a dummy tensor for metadata since broadcast_tensor_dict expects tensors + device = batch.prompt_embeds[0].device if batch.prompt_embeds else "cuda" + + # Prompt embeddings + num_prompt_embeds = len(batch.prompt_embeds) + d["_num_prompt_embeds"] = torch.tensor([num_prompt_embeds], + device=device) + for i, pe in enumerate(batch.prompt_embeds): + d[f"prompt_embed_{i}"] = pe + + # Prompt attention masks + has_prompt_masks = (batch.prompt_attention_mask is not None + and len(batch.prompt_attention_mask) > 0) + d["_has_prompt_masks"] = torch.tensor([1 if has_prompt_masks else 0], + device=device) + if has_prompt_masks: + for i, pm in enumerate(batch.prompt_attention_mask): + d[f"prompt_mask_{i}"] = pm + + # Negative embeddings (CFG) + has_neg_embeds = (batch.negative_prompt_embeds is not None + and len(batch.negative_prompt_embeds) > 0) + d["_has_neg_embeds"] = torch.tensor([1 if has_neg_embeds else 0], + device=device) + if has_neg_embeds: + d["_num_neg_embeds"] = torch.tensor( + [len(batch.negative_prompt_embeds)], device=device) + for i, ne in enumerate(batch.negative_prompt_embeds): + d[f"neg_embed_{i}"] = ne + + # Negative attention masks + has_neg_masks = (batch.negative_attention_mask is not None + and len(batch.negative_attention_mask) > 0) + d["_has_neg_masks"] = torch.tensor([1 if has_neg_masks else 0], + device=device) + if has_neg_masks: + for i, nm in enumerate(batch.negative_attention_mask): + d[f"neg_mask_{i}"] = nm + + # LTX2 audio embeddings + has_audio_embeds = "ltx2_audio_prompt_embeds" in batch.extra + d["_has_audio_embeds"] = torch.tensor([1 if has_audio_embeds else 0], + device=device) + if has_audio_embeds: + audio_embeds = batch.extra["ltx2_audio_prompt_embeds"] + d["_num_audio_embeds"] = torch.tensor([len(audio_embeds)], + device=device) + for i, ae in enumerate(audio_embeds): + d[f"audio_embed_{i}"] = ae + + # LTX2 audio negative embeddings + has_audio_neg = "ltx2_audio_negative_embeds" in batch.extra + d["_has_audio_neg"] = torch.tensor([1 if has_audio_neg else 0], + device=device) + if has_audio_neg: + audio_neg = batch.extra["ltx2_audio_negative_embeds"] + for i, audio_neg_embed in enumerate(audio_neg): + d[f"audio_neg_embed_{i}"] = audio_neg_embed + + return d + + def _unpack_broadcast_to_batch(self, batch: ForwardBatch, + d: dict[str, torch.Tensor]) -> None: + """Unpack broadcast dict into batch on non-rank-0 processes.""" + # Prompt embeddings + num_embeds = int(d["_num_prompt_embeds"].item()) + for i in range(num_embeds): + batch.prompt_embeds.append(d[f"prompt_embed_{i}"]) + + # Prompt attention masks + has_prompt_masks = int(d["_has_prompt_masks"].item()) == 1 + if has_prompt_masks and batch.prompt_attention_mask is not None: + for i in range(num_embeds): + batch.prompt_attention_mask.append(d[f"prompt_mask_{i}"]) + + # Negative embeddings (CFG) + has_neg_embeds = int(d["_has_neg_embeds"].item()) == 1 + if has_neg_embeds: + num_neg = int(d["_num_neg_embeds"].item()) + if batch.negative_prompt_embeds is not None: + for i in range(num_neg): + batch.negative_prompt_embeds.append(d[f"neg_embed_{i}"]) + + # Negative attention masks + has_neg_masks = int( + d.get("_has_neg_masks", torch.tensor([0])).item()) == 1 + if has_neg_masks and batch.negative_attention_mask is not None: + for i in range(num_neg): + batch.negative_attention_mask.append(d[f"neg_mask_{i}"]) + + # LTX2 audio embeddings + has_audio_embeds = int(d["_has_audio_embeds"].item()) == 1 + if has_audio_embeds: + num_audio = int(d["_num_audio_embeds"].item()) + audio_embeds = [d[f"audio_embed_{i}"] for i in range(num_audio)] + batch.extra["ltx2_audio_prompt_embeds"] = audio_embeds + + # LTX2 audio negative embeddings + has_audio_neg = int(d["_has_audio_neg"].item()) == 1 + if has_audio_neg: + # Use same count as audio embeds + num_audio = int(d["_num_audio_embeds"].item()) + audio_neg = [d[f"audio_neg_embed_{i}"] for i in range(num_audio)] + batch.extra["ltx2_audio_negative_embeds"] = audio_neg diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py new file mode 100644 index 000000000..6311930bd --- /dev/null +++ b/fastvideo/pipelines/stages/sr_denoising.py @@ -0,0 +1,431 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Denoising stage for diffusion pipelines. +""" + +import inspect +import weakref +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from einops import rearrange +import numpy as np +from tqdm.auto import tqdm + +from fastvideo.attention import get_attn_backend +from fastvideo.configs.pipelines.base import STA_Mode +from fastvideo.distributed import (get_local_torch_device, get_world_group) +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader, UpsamplerLoader +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.platforms import AttentionBackendEnum + +try: + from fastvideo.attention.backends.sliding_tile_attn import ( + SlidingTileAttentionBackend) + st_attn_available = True +except ImportError: + st_attn_available = False + +try: + from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend + from fastvideo.utils import is_vmoba_available + vmoba_attn_available = is_vmoba_available() +except ImportError: + vmoba_attn_available = False + +try: + from fastvideo.attention.backends.video_sparse_attn import ( + VideoSparseAttentionBackend) + vsa_available = True +except ImportError: + vsa_available = False + +logger = init_logger(__name__) + + +class SRDenoisingStage(PipelineStage): + """ + Stage for running the denoising loop in SR diffusion pipelines. Used by Hunyuan15 SR pipeline. + + This stage handles the iterative denoising process that transforms + the initial noise into the final output. + """ + + def __init__(self, + transformer, + scheduler, + upsampler, + pipeline=None, + vae=None) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + self.vae = vae + self.upsampler = upsampler + self.pipeline = weakref.ref(pipeline) if pipeline else None + attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads + self.attn_backend = get_attn_backend( + head_size=attn_head_size, + dtype=torch.float16, # TODO(will): hack + supported_attention_backends=( + AttentionBackendEnum.SLIDING_TILE_ATTN, + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + AttentionBackendEnum.VMOBA_ATTN, + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN_THREE) # hack + ) + + def add_noise_to_lq(self, + lq_latents: torch.Tensor, + strength: float = 0.7) -> torch.Tensor: + + def expand_dims(tensor: torch.Tensor, ndim: int): + shape = tensor.shape + (1, ) * (ndim - tensor.ndim) + return tensor.reshape(shape) + + noise = torch.randn_like(lq_latents) + timestep = torch.tensor([1000.0], + device=get_local_torch_device()) * strength + t = expand_dims(timestep, lq_latents.ndim) + return (1 - t / 1000.0) * lq_latents + (t / 1000.0) * noise + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + """ + Run the denoising loop. + + Args: + batch: The current batch information. + fastvideo_args: The inference arguments. + + Returns: + The batch with denoised latents. + """ + pipeline = self.pipeline() if self.pipeline else None + if not fastvideo_args.model_loaded["transformer"]: + loader = TransformerLoader() + self.transformer = loader.load( + fastvideo_args.model_paths["transformer"], fastvideo_args) + if pipeline: + pipeline.add_module("transformer", self.transformer) + fastvideo_args.model_loaded["transformer"] = True + + if not fastvideo_args.model_loaded["upsampler"]: + loader = UpsamplerLoader() + self.upsampler = loader.load( + fastvideo_args.model_paths["upsampler"], fastvideo_args) + if pipeline: + pipeline.add_module("upsampler", self.upsampler) + fastvideo_args.model_loaded["upsampler"] = True + + # Setup precision and autocast settings + # TODO(will): make the precision configurable for inference + # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast + + self.scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift_sr) + sigmas = np.linspace(1.0, 0.0, batch.num_inference_steps_sr + 1)[:-1] + self.scheduler.set_timesteps(sigmas=sigmas, + device=get_local_torch_device()) + timesteps = self.scheduler.timesteps + logger.info("timesteps: %s", timesteps) + num_inference_steps = len(timesteps) + num_warmup_steps = len( + timesteps) - num_inference_steps * self.scheduler.order + assert num_inference_steps == batch.num_inference_steps_sr, "num_inference_steps_sr must match the number of timesteps" + + pos_cond_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_2": batch.clip_embedding_pos, + "encoder_attention_mask": batch.prompt_attention_mask, + }, + ) + + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert not torch.isnan( + image_embeds[0]).any(), "image_embeds contains nan" + image_embeds = [ + image_embed.to(target_dtype) for image_embed in image_embeds + ] + + image_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "encoder_hidden_states_image": image_embeds, + }, + ) + + # Prepare STA parameters + if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + self.prepare_sta_param(batch, fastvideo_args) + + # Get latents and embeddings + prompt_embeds = batch.prompt_embeds + assert not torch.isnan( + prompt_embeds[0]).any(), "prompt_embeds contains nan" + + latents = batch.latents + lq_latents = batch.lq_latents + logger.info("lq_latents: %s", lq_latents.shape) + logger.info("latents: %s", latents.shape) + tgt_shape = latents.shape[-2:] # (h w) + bsz = lq_latents.shape[0] + lq_latents = rearrange(lq_latents, "b c f h w -> (b f) c h w") + lq_latents = F.interpolate(lq_latents, + size=tgt_shape, + mode="bilinear", + align_corners=False) + lq_latents = rearrange(lq_latents, "(b f) c h w -> b c f h w", b=bsz) + lq_latents = self.upsampler( + lq_latents.to(dtype=torch.float32, device=get_local_torch_device())) + lq_latents = lq_latents.to(dtype=latents.dtype) + lq_latents = self.add_noise_to_lq(lq_latents, 0.7) + b, c, f, h, w = lq_latents.shape + mask_ones = torch.ones(b, 1, f, h, w).to(lq_latents.device) + lq_cond_latents = torch.concat([lq_latents, mask_ones], + dim=1).to(target_dtype) + cond_latents = torch.cat( + [batch.video_latent, torch.zeros_like(latents)], + dim=1).to(target_dtype) + condition = torch.concat([cond_latents, lq_cond_latents], dim=1) + zero_lq_condition = condition.clone() + zero_lq_condition[:, c + 1:2 * c + 1] = torch.zeros_like(lq_latents) + zero_lq_condition[:, 2 * c + 1] = 0 + + latent_model_input = latents.to(target_dtype) + assert latent_model_input.shape[0] == 1, "only support batch size 1" + + # Run denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # Skip if interrupted + if hasattr(self, 'interrupt') and self.interrupt: + continue + + latent_model_input = latents.to(target_dtype) + if t < 1000 * 0.7: + condition = zero_lq_condition + + latent_model_input = torch.concat([latents, condition], dim=1) + + assert not torch.isnan( + latent_model_input).any(), "latent_model_input contains nan" + t_expand = t.repeat(latent_model_input.shape[0]) + + if i == len(timesteps) - 1: + timesteps_r = torch.tensor([0.0], + device=get_local_torch_device()) + else: + timesteps_r = timesteps[i + 1] + timesteps_r = timesteps_r.repeat(latent_model_input.shape[0]) + + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t) + + # Prepare inputs for transformer + guidance_expand = ( + torch.tensor( + [fastvideo_args.pipeline_config.embedded_cfg_scale] * + latent_model_input.shape[0], + dtype=torch.float32, + device=get_local_torch_device(), + ).to(target_dtype) * + 1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale + is not None else None) + + # Predict noise residual + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled): + if (st_attn_available + and self.attn_backend == SlidingTileAttentionBackend + ) or (vsa_available and self.attn_backend + == VideoSparseAttentionBackend): + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( + ) + + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls( + ) + # TODO(will): clean this up + attn_metadata = self.attn_metadata_builder.build( # type: ignore + current_timestep=i, # type: ignore + raw_latent_shape=batch. + raw_latent_shape[2:5], # type: ignore + patch_size=fastvideo_args. + pipeline_config. # type: ignore + dit_config.patch_size, # type: ignore + STA_param=batch.STA_param, # type: ignore + VSA_sparsity=fastvideo_args. + VSA_sparsity, # type: ignore + device=get_local_torch_device(), + ) + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + elif (vmoba_attn_available + and self.attn_backend == VMOBAAttentionBackend): + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( + ) + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls( + ) + # Prepare V-MoBA parameters from config + moba_params = fastvideo_args.moba_config.copy() + moba_params.update({ + "current_timestep": + i, + "raw_latent_shape": + batch.raw_latent_shape[2:5], + "patch_size": + fastvideo_args.pipeline_config.dit_config. + patch_size, + "device": + get_local_torch_device(), + }) + attn_metadata = self.attn_metadata_builder.build( + **moba_params) + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + # TODO(will): finalize the interface. vLLM uses this to + # support torch dynamo compilation. They pass in + # attn_metadata, vllm_config, and num_tokens. We can pass in + # fastvideo_args or training_args, and attn_metadata. + batch.is_cfg_negative = False + with set_forward_context( + current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch, + # fastvideo_args=fastvideo_args + ): + # Run transformer + noise_pred = self.transformer(latent_model_input, + prompt_embeds, + t_expand, + guidance=guidance_expand, + timestep_r=timesteps_r, + **pos_cond_kwargs, + **image_kwargs) + + # Compute the previous noisy sample + latents = self.scheduler.step(noise_pred, + t, + latents, + return_dict=False)[0] + + # Update progress bar + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and + (i + 1) % self.scheduler.order == 0 + and progress_bar is not None): + progress_bar.update() + + # Update batch with final latents + batch.latents = latents + + # Save STA mask search results if needed + if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING: + self.save_sta_search_results(batch) + + # deallocate transformer if on mps + if torch.backends.mps.is_available(): + logger.info("Memory before deallocating transformer: %s", + torch.mps.current_allocated_memory()) + del self.transformer + if pipeline is not None and "transformer" in pipeline.modules: + del pipeline.modules["transformer"] + fastvideo_args.model_loaded["transformer"] = False + logger.info("Memory after deallocating transformer: %s", + torch.mps.current_allocated_memory()) + + return batch + + def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]: + """ + Prepare extra kwargs for the scheduler step / denoise step. + + Args: + func: The function to prepare kwargs for. + kwargs: The kwargs to prepare. + + Returns: + The prepared kwargs. + """ + extra_step_kwargs = {} + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_step_kwargs[k] = v + return extra_step_kwargs + + def progress_bar(self, + iterable: Iterable | None = None, + total: int | None = None) -> tqdm: + """ + Create a progress bar for the denoising process. + + Args: + iterable: The iterable to iterate over. + total: The total number of items. + + Returns: + A tqdm progress bar. + """ + local_rank = get_world_group().local_rank + if local_rank == 0: + return tqdm(iterable=iterable, total=total) + else: + return tqdm(iterable=iterable, total=total, disable=True) + + def verify_input(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + """Verify denoising stage inputs.""" + result = VerificationResult() + result.add_check("timesteps", batch.timesteps, + [V.is_tensor, V.min_dims(1)]) + result.add_check("latents", batch.latents, + [V.is_tensor, V.with_dims(5)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("image_embeds", batch.image_embeds, V.is_list) + result.add_check("image_latent", batch.image_latent, + V.none_or_tensor_with_dims(5)) + result.add_check("num_inference_steps", batch.num_inference_steps, + V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, + V.positive_float) + result.add_check("eta", batch.eta, V.non_negative_float) + result.add_check("generator", batch.generator, + V.generator_or_list_generators) + result.add_check("do_classifier_free_guidance", + batch.do_classifier_free_guidance, V.bool_value) + result.add_check( + "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: + not batch.do_classifier_free_guidance or V.list_not_empty(x)) + return result + + def verify_output(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> VerificationResult: + """Verify denoising stage outputs.""" + result = VerificationResult() + result.add_check("latents", batch.latents, + [V.is_tensor, V.with_dims(5)]) + return result diff --git a/fastvideo/pipelines/stages/validators.py b/fastvideo/pipelines/stages/validators.py index dac57f001..9462a2ae0 100644 --- a/fastvideo/pipelines/stages/validators.py +++ b/fastvideo/pipelines/stages/validators.py @@ -25,6 +25,11 @@ class StageValidators: """Check if value is a positive integer.""" return isinstance(value, int) and value > 0 + @staticmethod + def non_negative_int(value: Any) -> bool: + """Check if value is a non-negative integer (allows 0).""" + return isinstance(value, int) and value >= 0 + @staticmethod def positive_float(value: Any) -> bool: """Check if value is a positive float.""" @@ -405,6 +410,8 @@ class VerificationResult: error_msg = f"tensor contains {torch.isnan(value).sum().item()} NaN values" elif validator_name == 'positive_int': expected = "positive integer" + elif validator_name == 'non_negative_int': + expected = "non-negative integer" elif validator_name == 'not_none': expected = "non-None value" elif validator_name == 'list_not_empty': diff --git a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 index 7acf878af..b9856eb9b 100644 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 index 545fb2d69..adf743b3f 100644 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 differ diff --git a/fastvideo/tests/ssim/test_inference_similarity.py b/fastvideo/tests/ssim/test_inference_similarity.py index 6f8b34538..2c70e8838 100644 --- a/fastvideo/tests/ssim/test_inference_similarity.py +++ b/fastvideo/tests/ssim/test_inference_similarity.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 import json import os +from contextlib import contextmanager +from typing import Iterator import torch import pytest @@ -12,6 +14,26 @@ from fastvideo.worker.multiproc_executor import MultiprocExecutor logger = init_logger(__name__) + +@contextmanager +def _attention_backend(backend: str) -> Iterator[None]: + previous = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend + try: + yield + finally: + if previous is None: + os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) + else: + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous + + +def _shutdown_executor(generator: VideoGenerator | None) -> None: + if generator is None: + return + if isinstance(generator.executor, MultiprocExecutor): + generator.executor.shutdown() + device_name = torch.cuda.get_device_name() device_reference_folder_suffix = '_reference_videos' @@ -119,7 +141,7 @@ LTX2_T2V_PARAMS = { } MODEL_TO_PARAMS = { - "FastHunyuan-diffusers": HUNYUAN_PARAMS, + # "FastHunyuan-diffusers": HUNYUAN_PARAMS, "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS, # "ltx2_diffusers": LTX2_T2V_PARAMS, } @@ -152,52 +174,54 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id): to reference videos using SSIM. """ assert len(I2V_TEST_PROMPTS) == len(I2V_IMAGE_PATHS), "Expect number of prompts equal to number of images" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + with _attention_backend(ATTENTION_BACKEND): + script_dir = os.path.dirname(os.path.abspath(__file__)) - script_dir = os.path.dirname(os.path.abspath(__file__)) + base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) + output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) + output_video_name = f"{prompt[:100].strip()}.mp4" - base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) - output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) - output_video_name = f"{prompt[:100].strip()}.mp4" + os.makedirs(output_dir, exist_ok=True) - os.makedirs(output_dir, exist_ok=True) + BASE_PARAMS = I2V_MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + image_path = I2V_IMAGE_PATHS[I2V_TEST_PROMPTS.index(prompt)] - BASE_PARAMS = I2V_MODEL_TO_PARAMS[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] - image_path = I2V_IMAGE_PATHS[I2V_TEST_PROMPTS.index(prompt)] + init_kwargs = { + "num_gpus": BASE_PARAMS["num_gpus"], + "flow_shift": BASE_PARAMS["flow_shift"], + "sp_size": BASE_PARAMS["sp_size"], + "tp_size": BASE_PARAMS["tp_size"], + } + if BASE_PARAMS.get("vae_sp"): + init_kwargs["vae_sp"] = True + init_kwargs["vae_tiling"] = True + if "text-encoder-precision" in BASE_PARAMS: + init_kwargs["text_encoder_precisions"] = BASE_PARAMS[ + "text-encoder-precision"] - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "flow_shift": BASE_PARAMS["flow_shift"], - "sp_size": BASE_PARAMS["sp_size"], - "tp_size": BASE_PARAMS["tp_size"], - } - if BASE_PARAMS.get("vae_sp"): - init_kwargs["vae_sp"] = True - init_kwargs["vae_tiling"] = True - if "text-encoder-precision" in BASE_PARAMS: - init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"] + generation_kwargs = { + "num_inference_steps": num_inference_steps, + "output_path": output_dir, + "image_path": image_path, + "height": BASE_PARAMS["height"], + "width": BASE_PARAMS["width"], + "num_frames": BASE_PARAMS["num_frames"], + "guidance_scale": BASE_PARAMS["guidance_scale"], + "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"], + "seed": BASE_PARAMS["seed"], + "fps": BASE_PARAMS["fps"], + } + if "neg_prompt" in BASE_PARAMS: + generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"] - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "image_path": image_path, - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "guidance_scale": BASE_PARAMS["guidance_scale"], - "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"], - "seed": BASE_PARAMS["seed"], - "fps": BASE_PARAMS["fps"], - } - if "neg_prompt" in BASE_PARAMS: - generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"] - - generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - - if isinstance(generator.executor, MultiprocExecutor): - generator.executor.shutdown() + generator: VideoGenerator | None = None + try: + generator = VideoGenerator.from_pretrained( + model_path=BASE_PARAMS["model_path"], **init_kwargs) + generator.generate_video(prompt, **generation_kwargs) + finally: + _shutdown_executor(generator) assert os.path.exists( output_dir), f"Output video was not generated at {output_dir}" @@ -253,61 +277,68 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id): Test that runs inference with different parameters and compares the output to reference videos using SSIM. """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + with _attention_backend(ATTENTION_BACKEND): + script_dir = os.path.dirname(os.path.abspath(__file__)) - script_dir = os.path.dirname(os.path.abspath(__file__)) + base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) + output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) + output_video_name = f"{prompt[:100].strip()}.mp4" - base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) - output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) - output_video_name = f"{prompt[:100].strip()}.mp4" + os.makedirs(output_dir, exist_ok=True) - os.makedirs(output_dir, exist_ok=True) + BASE_PARAMS = MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] - BASE_PARAMS = MODEL_TO_PARAMS[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] + init_kwargs = { + "num_gpus": BASE_PARAMS["num_gpus"], + "sp_size": BASE_PARAMS["sp_size"], + "tp_size": BASE_PARAMS["tp_size"], + "use_fsdp_inference": True, + "dit_cpu_offload": False, + "dit_layerwise_offload": False, + } + if "flow_shift" in BASE_PARAMS: + init_kwargs["flow_shift"] = BASE_PARAMS["flow_shift"] + if BASE_PARAMS.get("vae_sp"): + init_kwargs["vae_sp"] = True + init_kwargs["vae_tiling"] = True + if "text-encoder-precision" in BASE_PARAMS: + init_kwargs["text_encoder_precisions"] = BASE_PARAMS[ + "text-encoder-precision"] + # LTX2-specific VAE tiling parameters + if BASE_PARAMS.get("ltx2_vae_tiling"): + init_kwargs["ltx2_vae_tiling"] = True + init_kwargs["ltx2_vae_spatial_tile_size_in_pixels"] = BASE_PARAMS.get( + "ltx2_vae_spatial_tile_size_in_pixels", 512) + init_kwargs["ltx2_vae_spatial_tile_overlap_in_pixels"] = BASE_PARAMS.get( + "ltx2_vae_spatial_tile_overlap_in_pixels", 64) + init_kwargs["ltx2_vae_temporal_tile_size_in_frames"] = BASE_PARAMS.get( + "ltx2_vae_temporal_tile_size_in_frames", 64) + init_kwargs[ + "ltx2_vae_temporal_tile_overlap_in_frames"] = BASE_PARAMS.get( + "ltx2_vae_temporal_tile_overlap_in_frames", 24) - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "sp_size": BASE_PARAMS["sp_size"], - "tp_size": BASE_PARAMS["tp_size"], - "use_fsdp_inference": True, - "dit_cpu_offload": False, - "dit_layerwise_offload": False, - } - if "flow_shift" in BASE_PARAMS: - init_kwargs["flow_shift"] = BASE_PARAMS["flow_shift"] - if BASE_PARAMS.get("vae_sp"): - init_kwargs["vae_sp"] = True - init_kwargs["vae_tiling"] = True - if "text-encoder-precision" in BASE_PARAMS: - init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"] - # LTX2-specific VAE tiling parameters - if BASE_PARAMS.get("ltx2_vae_tiling"): - init_kwargs["ltx2_vae_tiling"] = True - init_kwargs["ltx2_vae_spatial_tile_size_in_pixels"] = BASE_PARAMS.get("ltx2_vae_spatial_tile_size_in_pixels", 512) - init_kwargs["ltx2_vae_spatial_tile_overlap_in_pixels"] = BASE_PARAMS.get("ltx2_vae_spatial_tile_overlap_in_pixels", 64) - init_kwargs["ltx2_vae_temporal_tile_size_in_frames"] = BASE_PARAMS.get("ltx2_vae_temporal_tile_size_in_frames", 64) - init_kwargs["ltx2_vae_temporal_tile_overlap_in_frames"] = BASE_PARAMS.get("ltx2_vae_temporal_tile_overlap_in_frames", 24) + generation_kwargs = { + "num_inference_steps": num_inference_steps, + "output_path": output_dir, + "height": BASE_PARAMS["height"], + "width": BASE_PARAMS["width"], + "num_frames": BASE_PARAMS["num_frames"], + "guidance_scale": BASE_PARAMS["guidance_scale"], + "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"], + "seed": BASE_PARAMS["seed"], + "fps": BASE_PARAMS["fps"], + } + if "neg_prompt" in BASE_PARAMS: + generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"] - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "guidance_scale": BASE_PARAMS["guidance_scale"], - "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"], - "seed": BASE_PARAMS["seed"], - "fps": BASE_PARAMS["fps"], - } - if "neg_prompt" in BASE_PARAMS: - generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"] - - generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - - if isinstance(generator.executor, MultiprocExecutor): - generator.executor.shutdown() + generator: VideoGenerator | None = None + try: + generator = VideoGenerator.from_pretrained( + model_path=BASE_PARAMS["model_path"], **init_kwargs) + generator.generate_video(prompt, **generation_kwargs) + finally: + _shutdown_executor(generator) assert os.path.exists( output_dir), f"Output video was not generated at {output_dir}" diff --git a/fastvideo/tests/transformers/test_hunyuanvideo.py b/fastvideo/tests/transformers/test_hunyuanvideo.py index 8fda5c49b..edd737530 100644 --- a/fastvideo/tests/transformers/test_hunyuanvideo.py +++ b/fastvideo/tests/transformers/test_hunyuanvideo.py @@ -39,6 +39,7 @@ REFERENCE_LATENT = 93.23263549804688 @pytest.mark.usefixtures("distributed_setup") +@pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed") def test_hunyuanvideo_distributed(): logger.info( f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}" diff --git a/fastvideo/tests/vaes/test_hunyuan_vae.py b/fastvideo/tests/vaes/test_hunyuan_vae.py index ed628d1fc..6a5c422d9 100644 --- a/fastvideo/tests/vaes/test_hunyuan_vae.py +++ b/fastvideo/tests/vaes/test_hunyuan_vae.py @@ -37,6 +37,7 @@ L40S_REFERENCE_LATENT = -157.09130859375 @pytest.mark.usefixtures("distributed_setup") +@pytest.mark.skip(reason="Skipping test_hunyuan_vae") def test_hunyuan_vae(): device = torch.device("cuda:0") precision = torch.bfloat16 diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index bc23a3b07..07c09b005 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -550,8 +550,6 @@ class TrainingPipeline(LoRAPipeline, ABC): self.optimizer.step() self.lr_scheduler.step() - training_batch.total_loss = training_batch.total_loss - training_batch.grad_norm = training_batch.grad_norm return training_batch def _resume_from_checkpoint(self) -> None: diff --git a/images/bus_terminal.jpg b/images/bus_terminal.jpg new file mode 100644 index 000000000..17191028e Binary files /dev/null and b/images/bus_terminal.jpg differ diff --git a/pyproject.toml b/pyproject.toml index 922e3014b..4a4030e6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "timm==1.0.11", "peft>=0.15.0", "diffusers>=0.33.1", - "torch>=2.9.1", + "torch>=2.10.0", "torchvision", "torchaudio", diff --git a/pyproject_other.toml b/pyproject_other.toml index 374309b6c..33fc31eca 100644 --- a/pyproject_other.toml +++ b/pyproject_other.toml @@ -27,7 +27,7 @@ dependencies = [ "timm==1.0.11", "peft>=0.15.0", "diffusers>=0.33.1", - "torch>=2.9.1", + "torch>=2.10.0", "torchvision", # Acceleration & Optimization diff --git a/videos/README.md b/videos/README.md new file mode 100644 index 000000000..d2615f4dd --- /dev/null +++ b/videos/README.md @@ -0,0 +1,7 @@ +# FastVideo/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. diff --git a/videos/robot_pouring.mp4 b/videos/robot_pouring.mp4 new file mode 100644 index 000000000..738eeb88b Binary files /dev/null and b/videos/robot_pouring.mp4 differ