CMD chunk1-short image-to-video demo

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/image.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,29 @@
+## NVIDIA OneWay Noncommercial License
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under this license, which may include software, documentation, or other files, and (b) any additions to or derivative works thereof that are made available under this license.
+The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the meaning as provided under U.S. copyright law; provided, however, that for the purposes of this license, derivative works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work.
+Works are “made available” under this license by including in or with the Work either (a) a copyright notice referencing the applicability of this license to the Work, or (b) a copy of this license.
+
+2. License Grant
+2.1 Copyright Grant. Subject to the terms and conditions of this license, each Licensor grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to use, reproduce, prepare derivative works of, publicly display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form.
+
+3. Limitations
+3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under this license, (b) you include a complete copy of this license with your distribution, and (c) you retain without modification any copyright, patent, trademark, or attribution notices that are present in the Work.
+3.2 Derivative Works. You may specify that additional or different terms apply to the use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your derivative works, and (b) you identify the specific derivative works that are subject to Your Terms. Notwithstanding Your Terms, this license (including the redistribution requirements in Section 3.1) will continue to apply to the Work itself.
+3.3 Use Limitation. The Work and any derivative works thereof only may be used or intended for use non-commercially. As used herein, “non-commercially” means for non-commercial research or educational purposes only.
+3.4 Patent Claims. If you bring or threaten to bring a patent claim against any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then your rights under this license from such Licensor (including the grant in Section 2.1) will terminate immediately.
+3.5 Trademarks. This license does not grant any rights to use any Licensor’s or its affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices described in this license.
+
+3.6 Termination. If you violate any term of this license, then your rights under this license (including the grant in Section 2.1) will terminate immediately.
+
+4. Disclaimer of Warranty.
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE.
+
+5. Limitation of Liability.
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.

LICENSES/Apache-2.0.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSES/LicenseRef-NvidiaOneWayNoncommercial.txt

@@ -0,0 +1,29 @@
+NVIDIA OneWay Noncommercial License
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under this license, which may include software, documentation, or other files, and (b) any additions to or derivative works thereof that are made available under this license.
+The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the meaning as provided under U.S. copyright law; provided, however, that for the purposes of this license, derivative works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work.
+Works are “made available” under this license by including in or with the Work either (a) a copyright notice referencing the applicability of this license to the Work, or (b) a copy of this license.
+
+2. License Grant
+2.1 Copyright Grant. Subject to the terms and conditions of this license, each Licensor grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to use, reproduce, prepare derivative works of, publicly display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form.
+
+3. Limitations
+3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under this license, (b) you include a complete copy of this license with your distribution, and (c) you retain without modification any copyright, patent, trademark, or attribution notices that are present in the Work.
+3.2 Derivative Works. You may specify that additional or different terms apply to the use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your derivative works, and (b) you identify the specific derivative works that are subject to Your Terms. Notwithstanding Your Terms, this license (including the redistribution requirements in Section 3.1) will continue to apply to the Work itself.
+3.3 Use Limitation. The Work and any derivative works thereof only may be used or intended for use non-commercially. As used herein, “non-commercially” means for non-commercial research or educational purposes only.
+3.4 Patent Claims. If you bring or threaten to bring a patent claim against any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then your rights under this license from such Licensor (including the grant in Section 2.1) will terminate immediately.
+3.5 Trademarks. This license does not grant any rights to use any Licensor’s or its affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices described in this license.
+
+3.6 Termination. If you violate any term of this license, then your rights under this license (including the grant in Section 2.1) will terminate immediately.
+
+4. Disclaimer of Warranty.
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE.
+
+5. Limitation of Liability.
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.

NOTICE

@@ -0,0 +1,63 @@
+Self-Forcing-Cosmos25
+Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+This distribution contains code derived from third-party open-source projects.
+The original portions identified below remain subject to their original
+licenses. NVIDIA modifications are separately identified in modified source
+files and are subject to the NVIDIA OneWay Noncommercial License in LICENSE.
+This NOTICE is informational and does not alter any applicable license.
+
+A copy of the Apache License, Version 2.0 is provided at
+LICENSES/Apache-2.0.txt.
+A copy of the NVIDIA OneWay Noncommercial License used by the
+LicenseRef-NvidiaOneWayNoncommercial SPDX identifier is provided at
+LICENSES/LicenseRef-NvidiaOneWayNoncommercial.txt.
+
+Wan2.1
+-------
+Source: https://github.com/Wan-Video/Wan2.1
+Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+License: Apache License 2.0
+Paths: wan/ (except wan/modules/causal_model.py, listed under Self-Forcing)
+
+The Wan tree also retains its embedded provenance notices for code adapted
+from projects such as Hugging Face Diffusers, CLIP, OpenCLIP, Transformers,
+and qwen-vl-utils.
+
+Self-Forcing
+------------
+Source: https://github.com/guandeh17/Self-Forcing
+Copyright (c) 2025 The Self-Forcing Authors. All rights reserved.
+License: Apache License 2.0
+Paths containing original or modified Self-Forcing code:
+  inference.py
+  setup.py
+  train.py
+  model/
+  pipeline/
+  trainer/
+  wan/modules/causal_model.py
+  scripts/create_lmdb_14b_shards.py
+  scripts/create_lmdb_iterative.py
+  scripts/generate_ode_pairs.py
+  utils/dataset.py
+  utils/distributed.py
+  utils/lmdb.py
+  utils/loss.py
+  utils/misc.py
+  utils/scheduler.py
+  utils/wan_wrapper.py
+
+NVIDIA Cosmos-Predict2.5
+------------------------
+Source: https://github.com/nvidia-cosmos/cosmos-predict2.5
+Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+License: Apache License 2.0
+Paths containing original or modified Cosmos-Predict2.5 code:
+  cosmos/causal_model.py
+  cosmos/kv_cache.py
+  cosmos/minimal_v1_lvg_dit.py
+  cosmos/minimal_v4_dit.py
+  cosmos/model_weights_stats.py
+  cosmos/runtime.py
+  cosmos/selective_activation_checkpoint.py

README.md

@@ -1,13 +1,78 @@
 ---
-title: Cmd I2v Demo
-emoji: 📚
-colorFrom: red
-colorTo: gray
+title: CMD Image-to-Video
+emoji: 🎬
+colorFrom: blue
+colorTo: red
 sdk: gradio
 sdk_version: 6.24.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+license: other
+short_description: Autoregressive 4-step image-to-video with NVIDIA CMD
+python_version: "3.10"
+startup_duration_timeout: 1h
+models:
+  - nvidia/cmd
+tags:
+  - image-to-video
+  - video-generation
+  - autoregressive
+  - cosmos
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# CMD — Context-Matched Distillation (image → video)
+
+Gradio demo for [`nvidia/cmd`](https://huggingface.co/nvidia/cmd): four-step
+**autoregressive** image-to-video generation, distilled with Context-Matched
+Distillation from a causal [Cosmos-Predict2.5
+2B](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B) teacher.
+
+* Paper / project page: <https://hmrishavbandy.github.io/cmd-site/>
+* Reference implementation: <https://github.com/nv-tlabs/cmd>
+
+## What runs here
+
+The Space serves the released `chunk1_short_t24_l21` student with the reference
+`configs/cosmos/t24_l21_student_context_distillation.yaml` config, matching the
+`chunk1-short` invocation in the upstream `examples/run_examples.sh`:
+
+| Setting | Value |
+|---|---|
+| Denoising steps | 4 (`[1000, 750, 500, 250]`, warped, shift 5.0) |
+| Autoregressive chunk | 1 latent frame |
+| Local attention window | 21 latent frames |
+| Context noise | 128 |
+| Latent frames | 24 → 93 pixel frames |
+| Resolution / fps | 832 × 480 @ 16 fps |
+
+Two runtime-forced deviations from the reference script:
+
+1. The gated Cosmos-Predict2.5 2B base DiT checkpoint is **not** downloaded. The
+   released CMD safetensors export contains every DiT parameter and persistent
+   buffer, so the model is constructed on the meta device and populated directly
+   from `nvidia/cmd` (upstream loads the base checkpoint first and then
+   overwrites all of it with the same export).
+2. The Wan2.1 video tokenizer that ships inside Cosmos-Predict2.5 2B is fetched
+   from an ungated safetensors mirror of the identical 508 MB `tokenizer.pth`.
+
+The chunk-4 variants are not offered: their multi-frame path goes through
+`torch.compile`-wrapped FlexAttention, which ZeroGPU does not support. The
+camera-control variants, which additionally need a camera trajectory in the CMD
+NPZ format, are also out of scope for this demo.
+
+## Vendored code and licensing
+
+`cosmos/`, `pipeline/`, `utils/` and a minimal `wan/` subset are copied from
+[nv-tlabs/cmd](https://github.com/nv-tlabs/cmd) so the demo reproduces the
+authors' inference path exactly. The only edits are trimmed package
+initializers (the upstream ones eagerly import training-only and multi-GPU
+modules). See `LICENSE`, `LICENSES/` and `NOTICE` for the upstream terms:
+
+* CMD / NVIDIA modifications — **NVIDIA OneWay Noncommercial License**
+  (non-commercial research or educational use only).
+* Wan2.1, Self-Forcing and Cosmos-Predict2.5 portions — Apache-2.0.
+
+`examples/image.png` and `examples/prompt.txt` are the authors' own example
+inputs, redistributed from the CMD repository under the same license.
+
+This demo is for research and development only.

app.py

@@ -0,0 +1,408 @@
+"""Gradio demo for NVIDIA CMD (Context-Matched Distillation) causal image-to-video.
+
+The demo runs the released ``chunk1_short_t24_l21`` CMD student on top of the
+Cosmos-Predict2.5 2B DiT, reproducing the reference invocation from
+``examples/run_examples.sh`` in https://github.com/nv-tlabs/cmd:
+
+    python inference.py \
+        --config_path configs/cosmos/t24_l21_student_context_distillation.yaml \
+        --checkpoint_path checkpoints/chunk1_short_t24_l21.safetensors \
+        --i2v --num_output_frames 24 --num_frame_per_block 1 --local_attn_size 21
+
+Only two things deviate from the reference script, both forced by the runtime:
+
+* the base Cosmos-Predict2.5 2B DiT checkpoint is never downloaded -- the
+  released CMD safetensors already contains every DiT parameter, so the model is
+  built on the meta device and populated straight from ``nvidia/cmd``;
+* the Wan2.1 video tokenizer is fetched from an ungated safetensors mirror of
+  the same ``tokenizer.pth`` file that ships inside Cosmos-Predict2.5 2B.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import tempfile
+import time
+from typing import Any, Dict, Tuple
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import spaces  # noqa: E402  -- must be imported before torch / CUDA
+
+import gradio as gr  # noqa: E402
+import imageio.v2 as imageio  # noqa: E402
+import torch  # noqa: E402
+from einops import rearrange  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from omegaconf import OmegaConf  # noqa: E402
+from PIL import Image, ImageOps  # noqa: E402
+from safetensors import safe_open  # noqa: E402
+from torchvision import transforms  # noqa: E402
+
+import cosmos.wrapper as cosmos_wrapper  # noqa: E402
+from cosmos.wrapper import CosmosDiffusionWrapper, CosmosTextEncoder, CosmosVAEWrapper  # noqa: E402
+from pipeline.causal_inference import CausalInferencePipeline  # noqa: E402
+from utils.misc import set_seed  # noqa: E402
+from wan.modules.vae import WanVAE_, _video_vae as _wan_video_vae  # noqa: E402
+
+# --------------------------------------------------------------------------- #
+# Constants
+# --------------------------------------------------------------------------- #
+
+CMD_REPO = "nvidia/cmd"
+CMD_CHECKPOINT = "chunk1_short_t24_l21.safetensors"
+CONFIG_PATH = "configs/cosmos/t24_l21_student_context_distillation.yaml"
+DEFAULT_CONFIG_PATH = "configs/cosmos/default_config.yaml"
+
+# Cosmos-Predict2.5 2B is gated, so the Wan2.1 tokenizer it bundles is pulled
+# from an ungated safetensors mirror of the very same 508 MB checkpoint.
+VAE_REPO = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
+VAE_CHECKPOINT = "vae/tokenizer.safetensors"
+
+HEIGHT, WIDTH = 480, 832
+FPS = 16
+MAX_LATENT_FRAMES = 24
+MAX_SEED = 2**31 - 1
+
+# Buffers that track training progress and are not part of a released export.
+_TRAINING_STAT_BUFFERS: Dict[str, torch.Tensor] = {
+    "accum_video_sample_counter": torch.zeros((), dtype=torch.int64),
+    "accum_image_sample_counter": torch.zeros((), dtype=torch.int64),
+    "accum_iteration": torch.zeros((), dtype=torch.int64),
+    "accum_train_in_hours": torch.zeros((), dtype=torch.float32),
+}
+
+EXAMPLE_PROMPT = open("examples/prompt.txt", encoding="utf-8").read().strip()
+
+
+# --------------------------------------------------------------------------- #
+# Weight loading
+# --------------------------------------------------------------------------- #
+
+
+def _load_safetensors_bf16(path: str) -> Dict[str, torch.Tensor]:
+    """Read a safetensors file tensor-by-tensor, casting floats to bfloat16.
+
+    The reference script loads an fp32 export into a bf16 model, which rounds
+    every float the same way; streaming the cast keeps host memory at ~4 GB
+    instead of ~12 GB.
+
+    Args:
+        path: Local path to the ``.safetensors`` file.
+
+    Returns:
+        Mapping from tensor name to CPU tensor.
+    """
+    state_dict: Dict[str, torch.Tensor] = {}
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        for key in handle.keys():
+            tensor = handle.get_tensor(key)
+            if tensor.is_floating_point():
+                tensor = tensor.to(torch.bfloat16)
+            state_dict[key] = tensor
+    return state_dict
+
+
+def _strip_common_prefix(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Drop a uniform checkpoint prefix so keys address the bare DiT."""
+    for prefix in ("generator.", "net_ema.", "net.", "model."):
+        if state_dict and all(key.startswith(prefix) for key in state_dict):
+            return {key[len(prefix):]: value for key, value in state_dict.items()}
+    return state_dict
+
+
+class CmdDiffusionWrapper(CosmosDiffusionWrapper):
+    """``CosmosDiffusionWrapper`` that loads the distilled CMD DiT directly.
+
+    Upstream first materialises the gated Cosmos-Predict2.5 base checkpoint and
+    then overwrites it with the CMD export. The export is complete (it covers
+    every parameter and persistent buffer of the DiT), so this subclass skips
+    the gated download entirely and assigns the CMD tensors onto a meta-device
+    model instead.
+    """
+
+    @classmethod
+    def _load_model(
+        cls,
+        model_name: str,
+        checkpoint_filename: str,
+        is_causal: bool,
+        local_attn_size: int,
+        sink_size: int,
+    ) -> torch.nn.Module:
+        if not is_causal:
+            raise ValueError("The CMD demo only serves the causal (streaming) DiT")
+
+        from cosmos.causal_model import CausalCosmosModel
+
+        model_kwargs = cls._model_kwargs(is_causal)
+        model_kwargs.update(local_attn_size=local_attn_size, sink_size=sink_size)
+        with torch.device("meta"):
+            model = CausalCosmosModel(**model_kwargs)
+
+        checkpoint_path = hf_hub_download(repo_id=model_name, filename=checkpoint_filename)
+        state_dict = _strip_common_prefix(_load_safetensors_bf16(checkpoint_path))
+        for name, tensor in _TRAINING_STAT_BUFFERS.items():
+            state_dict.setdefault(name, tensor.clone())
+
+        report = model.load_state_dict(state_dict, strict=False, assign=True)
+        if report.unexpected_keys:
+            raise RuntimeError(
+                f"Unexpected CMD checkpoint keys: {sorted(report.unexpected_keys)[:10]}"
+            )
+        uninitialised = [
+            name
+            for name, tensor in (
+                list(model.named_parameters()) + list(model.named_buffers())
+            )
+            if tensor.is_meta
+        ]
+        if uninitialised:
+            raise RuntimeError(
+                "CMD checkpoint did not initialise: " + ", ".join(uninitialised[:10])
+            )
+        print(f"Loaded {checkpoint_filename}: {len(state_dict)} DiT tensors", flush=True)
+        return model
+
+
+def _video_vae_with_safetensors(
+    pretrained_path: str | None = None,
+    z_dim: int | None = None,
+    device: str = "cpu",
+    **kwargs: Any,
+) -> torch.nn.Module:
+    """``wan.modules.vae._video_vae`` extended to accept safetensors weights."""
+    if pretrained_path is not None and str(pretrained_path).endswith(".safetensors"):
+        cfg = dict(
+            dim=96,
+            z_dim=z_dim,
+            dim_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            attn_scales=[],
+            temperal_downsample=[False, True, True],
+            dropout=0.0,
+        )
+        cfg.update(**kwargs)
+        with torch.device("meta"):
+            model = WanVAE_(**cfg)
+        state_dict = _load_safetensors_bf16(pretrained_path)
+        model.load_state_dict(state_dict, assign=True)
+        return model
+    return _wan_video_vae(
+        pretrained_path=pretrained_path, z_dim=z_dim, device=device, **kwargs
+    )
+
+
+cosmos_wrapper._video_vae = _video_vae_with_safetensors
+
+
+# --------------------------------------------------------------------------- #
+# Pipeline construction (module scope, eagerly on CUDA per ZeroGPU guidance)
+# --------------------------------------------------------------------------- #
+
+CONFIG = OmegaConf.merge(
+    OmegaConf.load(DEFAULT_CONFIG_PATH), OmegaConf.load(CONFIG_PATH)
+)
+# Reference chunk1-short overrides from examples/run_examples.sh.
+CONFIG.num_frame_per_block = 1
+CONFIG.model_kwargs.local_attn_size = 21
+CONFIG.i2v = True
+CONFIG.model_kwargs.model_name = CMD_REPO
+CONFIG.model_kwargs.checkpoint_filename = CMD_CHECKPOINT
+
+_model_kwargs = OmegaConf.to_container(CONFIG.model_kwargs, resolve=True)
+
+print("Building the CMD generator ...", flush=True)
+GENERATOR = CmdDiffusionWrapper(**_model_kwargs, is_causal=True)
+print("Building the Cosmos-Reason1 text encoder ...", flush=True)
+TEXT_ENCODER = CosmosTextEncoder(
+    model_name=str(CONFIG.text_encoder_name),
+    max_length=int(getattr(CONFIG, "text_encoder_max_length", 512)),
+)
+print("Building the Wan2.1 video tokenizer ...", flush=True)
+VAE = CosmosVAEWrapper(model_name=VAE_REPO, checkpoint_filename=VAE_CHECKPOINT)
+
+PIPELINE = CausalInferencePipeline(
+    CONFIG,
+    device="cpu",
+    generator=GENERATOR,
+    text_encoder=TEXT_ENCODER,
+    vae=VAE,
+)
+PIPELINE = PIPELINE.to(dtype=torch.bfloat16).eval().requires_grad_(False)
+PIPELINE = PIPELINE.to("cuda")
+print("CMD pipeline ready.", flush=True)
+
+_TRANSFORM = transforms.Compose(
+    [
+        transforms.Resize((HEIGHT, WIDTH)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5]),
+    ]
+)
+
+
+# --------------------------------------------------------------------------- #
+# Inference
+# --------------------------------------------------------------------------- #
+
+
+def _preprocess(image: Image.Image, center_crop: bool) -> torch.Tensor:
+    """Turn a PIL image into the reference ``[1, 3, 1, 480, 832]`` conditioning tensor."""
+    image = image.convert("RGB")
+    if center_crop:
+        image = ImageOps.fit(image, (WIDTH, HEIGHT), method=Image.LANCZOS, centering=(0.5, 0.5))
+    tensor = _TRANSFORM(image)
+    return tensor.unsqueeze(0).unsqueeze(2)
+
+
+@spaces.GPU(duration=240)
+def generate(
+    image: Image.Image,
+    prompt: str,
+    num_latent_frames: int = MAX_LATENT_FRAMES,
+    seed: int = 22,
+    randomize_seed: bool = False,
+    center_crop: bool = False,
+) -> Tuple[str, int]:
+    """Generate a short video that continues ``image`` following ``prompt``.
+
+    Args:
+        image: First frame of the video. Resized to 832x480.
+        prompt: Text description of the motion and scene to roll out.
+        num_latent_frames: Latent frames to produce, including the conditioning
+            frame. The tokenizer expands ``n`` latents into ``4 * (n - 1) + 1``
+            pixel frames, so 24 latents give 93 frames (~5.8 s at 16 fps).
+        seed: Random seed for the initial noise.
+        randomize_seed: Draw a fresh random seed instead of using ``seed``.
+        center_crop: Center-crop the input to 16:9 instead of stretching it.
+
+    Returns:
+        A tuple of the written mp4 path and the seed that was actually used.
+
+    Raises:
+        gradio.Error: If no image or no prompt was supplied.
+    """
+    if image is None:
+        raise gr.Error("Please provide a conditioning image.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please provide a text prompt.")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed) % (MAX_SEED + 1)
+    num_latent_frames = max(2, min(int(num_latent_frames), MAX_LATENT_FRAMES))
+
+    with torch.no_grad():
+        set_seed(seed)
+        started = time.perf_counter()
+
+        pixels = _preprocess(image, center_crop).to(device="cuda", dtype=torch.bfloat16)
+        initial_latent = PIPELINE.vae.encode_to_latent(pixels).to(
+            device="cuda", dtype=torch.bfloat16
+        )
+        noise = torch.randn(
+            [1, num_latent_frames - 1, *list(CONFIG.image_or_video_shape[2:])],
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        encoded = time.perf_counter()
+
+        video = PIPELINE.inference(
+            noise=noise,
+            text_prompts=[prompt.strip()],
+            initial_latent=initial_latent,
+        )
+        rolled_out = time.perf_counter()
+
+        frames = (255.0 * rearrange(video, "b t c h w -> b t h w c").cpu())[0]
+        frames = frames.to(torch.uint8).numpy()
+        PIPELINE.vae.model.clear_cache()
+
+    path = os.path.join(tempfile.mkdtemp(), f"cmd_{seed}.mp4")
+    imageio.mimwrite(path, list(frames), fps=FPS, codec="libx264", quality=8, macro_block_size=1)
+    finished = time.perf_counter()
+    print(
+        f"[timing] latents={num_latent_frames} encode={encoded - started:.1f}s "
+        f"rollout={rolled_out - encoded:.1f}s encode_mp4={finished - rolled_out:.1f}s "
+        f"total={finished - started:.1f}s",
+        flush=True,
+    )
+    return path, seed
+
+
+# --------------------------------------------------------------------------- #
+# UI
+# --------------------------------------------------------------------------- #
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # CMD — Context-Matched Distillation (image → video)
+
+            Four-step **autoregressive** image-to-video generation with NVIDIA's
+            [CMD](https://huggingface.co/nvidia/cmd) `chunk1_short_t24_l21` student,
+            distilled from a causal Cosmos-Predict2.5 2B teacher. The model rolls the
+            video out one latent frame at a time behind a 21-frame local attention
+            window, so cost grows linearly with length rather than quadratically.
+
+            Output is 832×480, 93 frames at 16 fps (~5.8 s).
+            *Research / non-commercial use only — NVIDIA OneWay Noncommercial License.*
+            """
+        )
+        with gr.Row():
+            with gr.Column():
+                image = gr.Image(label="First frame", type="pil", height=300)
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe the scene and how the camera / subject should move…",
+                    lines=4,
+                )
+                run_button = gr.Button("Generate video", variant="primary")
+                with gr.Accordion("Advanced settings", open=False):
+                    num_latent_frames = gr.Slider(
+                        label="Latent frames (24 → 93 pixel frames)",
+                        minimum=4,
+                        maximum=MAX_LATENT_FRAMES,
+                        step=1,
+                        value=MAX_LATENT_FRAMES,
+                    )
+                    seed = gr.Slider(
+                        label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=22
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                    center_crop = gr.Checkbox(
+                        label="Center-crop the input to 16:9 (default: stretch, as in the reference script)",
+                        value=False,
+                    )
+            with gr.Column():
+                video_out = gr.Video(label="Generated video", autoplay=True, height=430)
+
+        gr.Examples(
+            examples=[["examples/image.png", EXAMPLE_PROMPT]],
+            inputs=[image, prompt],
+            outputs=[video_out, seed],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Official CMD example (seed 22)",
+        )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=generate,
+        inputs=[image, prompt, num_latent_frames, seed, randomize_seed, center_crop],
+        outputs=[video_out, seed],
+        api_name="generate",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

configs/cosmos/default_config.yaml

@@ -0,0 +1,19 @@
+model_family: cosmos
+independent_first_frame: true
+warp_denoising_step: false
+weight_decay: 0.01
+same_step_across_blocks: true
+discriminator_lr_multiplier: 1.0
+last_step_only: false
+i2v: true
+num_training_frames: 24
+gc_interval: 100
+context_noise: 0
+causal: true
+
+ckpt_step: 0
+eval_first_n: 64
+num_samples: 1
+height: 480
+width: 832
+num_frames: 93

configs/cosmos/t24_l21_student_context_distillation.yaml

@@ -0,0 +1,89 @@
+# Stage 2: distill the t24/l21 student from the paired causal-flow teacher.
+generator_ckpt: checkpoints/chunk1_teacher_t24_l21.safetensors
+teacher_ckpt: checkpoints/chunk1_teacher_t24_l21.safetensors
+prefix_noise: 256
+model_family: cosmos
+real_name: nvidia/Cosmos-Predict2.5-2B
+fake_name: nvidia/Cosmos-Predict2.5-2B
+text_encoder_name: nvidia/Cosmos-Reason1-7B
+vae_model_name: nvidia/Cosmos-Predict2.5-2B
+
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+context_noise: 128
+validation_interval: 50
+validation_at_start: true
+validation_seed: 12345
+validation_fps: 16
+real_guidance_scale: 4.0
+fake_guidance_scale: 0.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: ${oc.env:WANDB_BASE_URL,https://api.wandb.ai}
+# wandb.login(key=None) reads WANDB_API_KEY without placing the key in config.
+wandb_key: null
+wandb_entity: ${oc.env:WANDB_ENTITY,null}
+wandb_project: ${oc.env:WANDB_PROJECT,causal-cosmos25}
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+lr_critic: 4.0e-07
+beta1: 0.0
+beta2: 0.999
+beta1_critic: 0.0
+beta2_critic: 0.999
+
+# Sharded LMDB containing prompts and Cosmos/Wan-VAE image latents.
+data_path: data/cosmos_i2v_lmdb_t24
+batch_size: 1
+dataloader_num_workers: 1
+ema_weight: 0.0
+ema_start_step: 200
+log_iters: 50
+max_steps: 10000
+negative_prompt: >-
+  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.
+dfake_gen_update_ratio: 5
+
+# 93 pixel frames become 24 latent frames with the checkpoint tokenizer.
+# 480x832 becomes 60x104 latent spatially.
+image_or_video_shape:
+- 1
+- 24
+- 16
+- 60
+- 104
+distribution_loss: context_matched
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 1
+load_raw_video: false
+i2v: true
+independent_first_frame: true
+# One clean I2V frame plus 23 generated/scored frames, matching t24.
+num_training_frames: 24
+model_kwargs:
+  model_name: nvidia/Cosmos-Predict2.5-2B
+  checkpoint_filename: base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt
+  timestep_shift: 5.0
+  i2v: true
+  local_attn_size: 21
+  sink_size: 0

cosmos/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+from .wrapper import CosmosDiffusionWrapper, CosmosTextEncoder, CosmosVAEWrapper
+
+__all__ = [
+    "CosmosDiffusionWrapper",
+    "CosmosTextEncoder",
+    "CosmosVAEWrapper",
+]

cosmos/camera_conditioning.py

@@ -0,0 +1,233 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+"""Camera-ray conditioning for the Cosmos video latent grid."""
+
+from __future__ import annotations
+
+import torch
+
+
+CAMERA_FEATURE_DIM = 6
+
+
+def camera_frame_indices(
+    num_pixel_frames: int,
+    frame_stride: int = 4,
+    *,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    """Select pixel-frame cameras aligned with temporally compressed latents."""
+    if num_pixel_frames <= 0:
+        raise ValueError("num_pixel_frames must be positive")
+    if frame_stride <= 0:
+        raise ValueError("frame_stride must be positive")
+    if (num_pixel_frames - 1) % frame_stride:
+        raise ValueError(
+            "Camera sequence length must be 1 + k * frame_stride; got "
+            f"{num_pixel_frames} frames and stride {frame_stride}"
+        )
+    return torch.arange(
+        0,
+        num_pixel_frames,
+        frame_stride,
+        device=device,
+        dtype=torch.long,
+    )
+
+
+def frame_relative_camera_to_world(
+    camera_to_world: torch.Tensor,
+    num_frame_per_block: int = 1,
+) -> torch.Tensor:
+    """Express each generated block relative to the prior block boundary."""
+    if camera_to_world.ndim != 4 or camera_to_world.shape[-2:] != (4, 4):
+        raise ValueError(
+            "camera_to_world must have shape [B, T, 4, 4]; got "
+            f"{tuple(camera_to_world.shape)}"
+        )
+    if camera_to_world.shape[1] == 0:
+        raise ValueError("camera_to_world must contain at least one frame")
+    if num_frame_per_block <= 0:
+        raise ValueError("num_frame_per_block must be positive")
+
+    poses = camera_to_world.to(torch.float32)
+    frame_indices = torch.arange(
+        poses.shape[1],
+        device=poses.device,
+        dtype=torch.long,
+    )
+    # With an independent I2V prefix at frame zero, frames 1..C use frame 0,
+    # frames C+1..2C use frame C, and so on. For C=1 this reduces to the
+    # original previous-frame-relative convention.
+    anchor_indices = torch.div(
+        torch.clamp(frame_indices - 1, min=0),
+        num_frame_per_block,
+        rounding_mode="floor",
+    ) * num_frame_per_block
+    anchors = poses.index_select(1, anchor_indices)
+    relative = torch.linalg.solve(anchors, poses)
+    return relative
+
+
+def _per_frame_intrinsics(
+    intrinsics: torch.Tensor,
+    frame_indices: torch.Tensor,
+    num_pixel_frames: int,
+) -> torch.Tensor:
+    if intrinsics.ndim == 3 and intrinsics.shape[-2:] == (3, 3):
+        return intrinsics[:, None].expand(-1, frame_indices.numel(), -1, -1)
+    if intrinsics.ndim == 4 and intrinsics.shape[-2:] == (3, 3):
+        if intrinsics.shape[1] == 1:
+            return intrinsics.expand(-1, frame_indices.numel(), -1, -1)
+        if intrinsics.shape[1] != num_pixel_frames:
+            raise ValueError(
+                "Per-frame intrinsics must match the pixel camera sequence; got "
+                f"{intrinsics.shape[1]} and {num_pixel_frames} frames"
+            )
+        return intrinsics.index_select(1, frame_indices)
+    raise ValueError(
+        "intrinsics must have shape [B, 3, 3] or [B, T, 3, 3]; got "
+        f"{tuple(intrinsics.shape)}"
+    )
+
+
+def camera_rays(
+    camera_to_world: torch.Tensor,
+    intrinsics: torch.Tensor,
+    image_height: int,
+    image_width: int,
+) -> torch.Tensor:
+    """Return ray origins and unit directions as ``[B, T, H, W, 6]``."""
+    if camera_to_world.ndim != 4 or camera_to_world.shape[-2:] != (4, 4):
+        raise ValueError("camera_to_world must have shape [B, T, 4, 4]")
+    if intrinsics.ndim != 4 or intrinsics.shape[-2:] != (3, 3):
+        raise ValueError("intrinsics must have shape [B, T, 3, 3]")
+    if camera_to_world.shape[:2] != intrinsics.shape[:2]:
+        raise ValueError("Camera poses and intrinsics must have matching B and T")
+    if image_height <= 0 or image_width <= 0:
+        raise ValueError("Camera image dimensions must be positive")
+
+    poses = camera_to_world.to(torch.float32)
+    calibration = intrinsics.to(device=poses.device, dtype=torch.float32)
+    focal_x = calibration[..., 0, 0]
+    focal_y = calibration[..., 1, 1]
+    if torch.any(focal_x <= 0) or torch.any(focal_y <= 0):
+        raise ValueError("Camera focal lengths must be positive")
+
+    pixel_y, pixel_x = torch.meshgrid(
+        torch.arange(image_height, device=poses.device, dtype=torch.float32) + 0.5,
+        torch.arange(image_width, device=poses.device, dtype=torch.float32) + 0.5,
+        indexing="ij",
+    )
+    pixel_x = pixel_x[None, None]
+    pixel_y = pixel_y[None, None]
+    direction_x = (
+        pixel_x - calibration[..., 0, 2, None, None]
+    ) / focal_x[..., None, None]
+    direction_y = (
+        pixel_y - calibration[..., 1, 2, None, None]
+    ) / focal_y[..., None, None]
+    camera_direction = torch.stack(
+        [direction_x, direction_y, torch.ones_like(direction_x)],
+        dim=-1,
+    )
+    camera_direction = torch.nn.functional.normalize(camera_direction, dim=-1)
+
+    rotation = poses[..., :3, :3]
+    ray_direction = torch.einsum(
+        "btij,bthwj->bthwi",
+        rotation,
+        camera_direction,
+    )
+    ray_origin = poses[..., :3, 3][..., None, None, :].expand_as(ray_direction)
+    return torch.cat([ray_origin, ray_direction], dim=-1)
+
+
+def patchify_camera_rays(
+    rays: torch.Tensor,
+    patch_size: int = 16,
+) -> torch.Tensor:
+    """Flatten each spatial ray patch into the camera token channels."""
+    if rays.ndim != 5 or rays.shape[-1] != CAMERA_FEATURE_DIM:
+        raise ValueError("rays must have shape [B, T, H, W, 6]")
+    if patch_size <= 0:
+        raise ValueError("patch_size must be positive")
+    batch, frames, height, width, channels = rays.shape
+    if height % patch_size or width % patch_size:
+        raise ValueError(
+            f"Camera image {(height, width)} is not divisible by patch size {patch_size}"
+        )
+
+    rays_bcthw = rays.permute(0, 4, 1, 2, 3).contiguous()
+    token_h = height // patch_size
+    token_w = width // patch_size
+    return (
+        rays_bcthw.reshape(
+            batch,
+            channels,
+            frames,
+            token_h,
+            patch_size,
+            token_w,
+            patch_size,
+        )
+        .permute(0, 1, 4, 6, 2, 3, 5)
+        .reshape(
+            batch,
+            channels * patch_size * patch_size,
+            frames,
+            token_h,
+            token_w,
+        )
+        .contiguous()
+    )
+
+
+def build_camera_conditioning(
+    camera_to_world: torch.Tensor,
+    intrinsics: torch.Tensor,
+    *,
+    image_height: int,
+    image_width: int,
+    frame_stride: int = 4,
+    patch_size: int = 16,
+    num_frame_per_block: int = 1,
+    expected_latent_frames: int | None = None,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Build block-relative origin/direction camera tokens from pixel cameras."""
+    if camera_to_world.ndim != 4 or camera_to_world.shape[-2:] != (4, 4):
+        raise ValueError("camera_to_world must have shape [B, T, 4, 4]")
+    num_pixel_frames = camera_to_world.shape[1]
+    indices = camera_frame_indices(
+        num_pixel_frames,
+        frame_stride,
+        device=camera_to_world.device,
+    )
+    if expected_latent_frames is not None and indices.numel() != expected_latent_frames:
+        raise ValueError(
+            f"Camera sequence produces {indices.numel()} latent frames; "
+            f"expected {expected_latent_frames}"
+        )
+
+    sampled_poses = camera_to_world.index_select(1, indices)
+    sampled_intrinsics = _per_frame_intrinsics(
+        intrinsics.to(device=camera_to_world.device),
+        indices,
+        num_pixel_frames,
+    )
+    relative_poses = frame_relative_camera_to_world(
+        sampled_poses,
+        num_frame_per_block=num_frame_per_block,
+    )
+    rays = camera_rays(
+        relative_poses,
+        sampled_intrinsics,
+        image_height,
+        image_width,
+    )
+    conditioning = patchify_camera_rays(rays, patch_size=patch_size)
+    if output_dtype is not None:
+        conditioning = conditioning.to(dtype=output_dtype)
+    return conditioning

cosmos/causal_model.py

@@ -0,0 +1,960 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Modifications Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+"""Block-causal Cosmos-Predict2.5 model with streaming KV cache."""
+
+import math
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from torch.nn.attention.flex_attention import BlockMask
+from torch.nn.attention.flex_attention import flex_attention as torch_flex_attention
+
+from cosmos.kv_cache import (
+    AttentionOpWithKVCache,
+    KVCacheConfig,
+    VideoSeqPos,
+)
+from cosmos.minimal_v1_lvg_dit import MinimalV1LVGDiT
+from cosmos.minimal_v4_dit import VideoSize, i4_attention_op
+
+
+# FlexAttention represents the block mask sparsely, avoiding a dense
+# [T*H*W, T*H*W] mask for full-resolution video tokens. Unlike Wan 1.3B,
+# Cosmos has a standard 16-head/128-dim layout and does not need expensive
+# max-autotune kernel benchmarking.
+flex_attention = torch.compile(
+    torch_flex_attention,
+    dynamic=False,
+)
+
+
+_SPARSE_BLOCK_SIZE = 128
+
+
+def _merge_intervals(
+    intervals: list[tuple[int, int]],
+    limit: int,
+) -> tuple[tuple[int, int], ...]:
+    """Clip and merge half-open token intervals."""
+    clipped = sorted(
+        (max(0, start), min(limit, end))
+        for start, end in intervals
+        if max(0, start) < min(limit, end)
+    )
+    merged: list[list[int]] = []
+    for start, end in clipped:
+        if merged and start <= merged[-1][1]:
+            merged[-1][1] = max(merged[-1][1], end)
+        else:
+            merged.append([start, end])
+    return tuple((start, end) for start, end in merged)
+
+
+def _full_sequence_query_ranges(
+    *,
+    total_length: int,
+    tokens_per_frame: int,
+    num_frame_per_block: int,
+    independent_first_frame: bool,
+    local_attn_size: int,
+    sink_size: int,
+) -> list[tuple[int, int, tuple[tuple[int, int], ...]]]:
+    """Describe the allowed K/V intervals for full causal attention."""
+    block_tokens = num_frame_per_block * tokens_per_frame
+    prefix_tokens = tokens_per_frame if independent_first_frame else 0
+    use_local_attn = local_attn_size > 0
+    sink_tokens = max(sink_size, 0) * tokens_per_frame
+    window_tokens = max(local_attn_size - sink_size, 0) * tokens_per_frame
+
+    query_ranges = []
+    if prefix_tokens:
+        prefix_intervals = (
+            [(0, prefix_tokens)]
+            if not use_local_attn
+            else [(max(prefix_tokens - window_tokens, 0), prefix_tokens)]
+        )
+        query_ranges.append(
+            (
+                0,
+                prefix_tokens,
+                _merge_intervals(prefix_intervals, total_length),
+            )
+        )
+    for query_start in range(prefix_tokens, total_length, block_tokens):
+        raw_query_end = query_start + block_tokens
+        query_end = min(raw_query_end, total_length)
+        if use_local_attn:
+            recent_start = max(raw_query_end - window_tokens, 0)
+            promoted_sink_end = min(
+                query_start,
+                sink_tokens,
+                recent_start,
+            )
+            recent_start = max(promoted_sink_end, recent_start)
+            intervals = [
+                (0, promoted_sink_end),
+                (recent_start, raw_query_end),
+            ]
+        else:
+            intervals = [(0, raw_query_end)]
+        query_ranges.append(
+            (
+                query_start,
+                query_end,
+                _merge_intervals(intervals, total_length),
+            )
+        )
+    return query_ranges
+
+
+def _build_sparse_block_rows(
+    *,
+    query_ranges: list[
+        tuple[int, int, tuple[tuple[int, int], ...]]
+    ],
+    q_total: int,
+    kv_total: int,
+    block_size: int = _SPARSE_BLOCK_SIZE,
+) -> tuple[list[int], list[list[int]], list[int], list[list[int]]]:
+    """Build exact sparse-tile metadata without a dense token mask.
+
+    Each query range has one fixed union of allowed key intervals. Query and
+    key tile boundaries need not align with frame boundaries; tiles crossing a
+    boundary are emitted as partial and evaluated by ``mask_mod`` in the
+    FlexAttention kernel.
+    """
+    if q_total <= 0 or kv_total <= 0:
+        raise ValueError("Sparse attention lengths must be positive")
+    if q_total % block_size or kv_total % block_size:
+        raise ValueError("Sparse attention lengths must be block aligned")
+
+    previous_end = 0
+    for query_start, query_end, intervals in query_ranges:
+        if query_start != previous_end or query_end <= query_start:
+            raise ValueError("Query ranges must be contiguous and non-empty")
+        if query_end > q_total:
+            raise ValueError("Query range exceeds padded query length")
+        for key_start, key_end in intervals:
+            if not (0 <= key_start < key_end <= kv_total):
+                raise ValueError("Key interval exceeds padded key length")
+        previous_end = query_end
+
+    q_block_count = q_total // block_size
+    kv_block_count = kv_total // block_size
+    partial_counts: list[int] = []
+    partial_rows: list[list[int]] = []
+    full_counts: list[int] = []
+    full_rows: list[list[int]] = []
+
+    range_index = 0
+    for query_block in range(q_block_count):
+        query_start = query_block * block_size
+        query_end = query_start + block_size
+        while (
+            range_index < len(query_ranges)
+            and query_ranges[range_index][1] <= query_start
+        ):
+            range_index += 1
+
+        query_segments = []
+        candidate_index = range_index
+        cursor = query_start
+        fully_covered_query = True
+        while (
+            candidate_index < len(query_ranges)
+            and query_ranges[candidate_index][0] < query_end
+        ):
+            range_start, range_end, intervals = query_ranges[candidate_index]
+            segment_start = max(query_start, range_start)
+            segment_end = min(query_end, range_end)
+            if segment_start > cursor:
+                fully_covered_query = False
+            if segment_start < segment_end:
+                query_segments.append((segment_start, segment_end, intervals))
+                cursor = segment_end
+            candidate_index += 1
+        if cursor < query_end:
+            fully_covered_query = False
+
+        partial_indices: list[int] = []
+        full_indices: list[int] = []
+        for key_block in range(kv_block_count):
+            key_start = key_block * block_size
+            key_end = key_start + block_size
+            any_allowed = False
+            fully_allowed = fully_covered_query and bool(query_segments)
+            for _segment_start, _segment_end, intervals in query_segments:
+                segment_overlaps = any(
+                    interval_start < key_end and key_start < interval_end
+                    for interval_start, interval_end in intervals
+                )
+                any_allowed = any_allowed or segment_overlaps
+                segment_contains = any(
+                    interval_start <= key_start and key_end <= interval_end
+                    for interval_start, interval_end in intervals
+                )
+                fully_allowed = fully_allowed and segment_contains
+
+            if fully_allowed:
+                full_indices.append(key_block)
+            elif any_allowed:
+                partial_indices.append(key_block)
+
+        partial_counts.append(len(partial_indices))
+        full_counts.append(len(full_indices))
+        partial_rows.append(
+            partial_indices + [0] * (kv_block_count - len(partial_indices))
+        )
+        full_rows.append(
+            full_indices + [0] * (kv_block_count - len(full_indices))
+        )
+
+    return partial_counts, partial_rows, full_counts, full_rows
+
+
+def _block_mask_from_intervals(
+    *,
+    query_ranges: list[
+        tuple[int, int, tuple[tuple[int, int], ...]]
+    ],
+    q_total: int,
+    kv_total: int,
+    mask_mod,
+    device: torch.device,
+) -> BlockMask:
+    """Create a FlexAttention BlockMask from compact interval metadata."""
+    (
+        partial_counts,
+        partial_rows,
+        full_counts,
+        full_rows,
+    ) = _build_sparse_block_rows(
+        query_ranges=query_ranges,
+        q_total=q_total,
+        kv_total=kv_total,
+    )
+
+    def count_tensor(values: list[int]) -> torch.Tensor:
+        return torch.tensor(
+            values,
+            dtype=torch.int32,
+            device=device,
+        ).view(1, 1, -1)
+
+    def row_tensor(values: list[list[int]]) -> torch.Tensor:
+        return torch.tensor(
+            values,
+            dtype=torch.int32,
+            device=device,
+        ).view(1, 1, len(values), -1)
+
+    return BlockMask.from_kv_blocks(
+        kv_num_blocks=count_tensor(partial_counts),
+        kv_indices=row_tensor(partial_rows),
+        full_kv_num_blocks=count_tensor(full_counts),
+        full_kv_indices=row_tensor(full_rows),
+        BLOCK_SIZE=_SPARSE_BLOCK_SIZE,
+        mask_mod=mask_mod,
+        seq_lengths=(q_total, kv_total),
+    )
+
+
+class CausalCosmosAttention(AttentionOpWithKVCache):
+    """Causal full-sequence attention and past-only streaming attention."""
+
+    _block_mask_cache: dict[tuple, BlockMask] = {}
+
+    def __init__(self, local_attn_size: int = -1, sink_size: int = 0) -> None:
+        if local_attn_size == 0 or local_attn_size < -1:
+            raise ValueError("local_attn_size must be -1 or a positive frame count")
+        if sink_size < 0:
+            raise ValueError("sink_size must be non-negative")
+        super().__init__(i4_attention_op)
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.num_frame_per_block = 1
+        self.independent_first_frame = False
+
+    def reset_kv_cache(self, max_cache_size: Optional[int] = None) -> None:
+        self.k_cache: Optional[torch.Tensor] = None
+        self.v_cache: Optional[torch.Tensor] = None
+        self._cache_slot_by_frame: dict[int, int] = {}
+        self._cache_frame_by_slot: dict[int, int] = {}
+        self._cache_tokens_per_frame: Optional[int] = None
+        self._cache_slot_indices: dict[tuple[int, ...], torch.Tensor] = {}
+        self.max_cache_size = max_cache_size
+
+    def _cache_capacity(self, required_frames: int) -> int:
+        # Training may retain more frames than the logical attention window so
+        # activation-checkpoint recomputation sees the original cache history.
+        local_capacity = (
+            self.sink_size + self.local_attn_size
+            if self.local_attn_size > 0
+            else 1
+        )
+        return max(
+            required_frames,
+            local_capacity,
+            self.max_cache_size or 0,
+        )
+
+    def _initialize_cache_storage(
+        self,
+        value: torch.Tensor,
+        tokens_per_frame: int,
+        required_frames: int,
+    ) -> None:
+        capacity = self._cache_capacity(required_frames)
+        shape = (
+            value.shape[0],
+            capacity,
+            tokens_per_frame,
+            value.shape[2],
+            value.shape[3],
+        )
+        self.k_cache = torch.empty(shape, device=value.device, dtype=value.dtype)
+        self.v_cache = torch.empty(shape, device=value.device, dtype=value.dtype)
+        self._cache_tokens_per_frame = tokens_per_frame
+
+    def _cache_slot(self, frame_index: int) -> int:
+        if self.k_cache is None:
+            raise RuntimeError("Cosmos K/V cache storage is not initialized")
+        capacity = self.k_cache.shape[1]
+        if frame_index < self.sink_size:
+            return frame_index
+        recent_capacity = capacity - self.sink_size
+        if recent_capacity <= 0:
+            raise RuntimeError("Cosmos K/V cache has no recent-history capacity")
+        return self.sink_size + (
+            (frame_index - self.sink_size) % recent_capacity
+        )
+
+    def _store_cache_frame(
+        self,
+        frame_index: int,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> None:
+        tokens_per_frame = k.shape[1]
+        if self.k_cache is None or self.v_cache is None:
+            self._initialize_cache_storage(
+                k,
+                tokens_per_frame,
+                frame_index + 1,
+            )
+        if self._cache_tokens_per_frame != tokens_per_frame:
+            raise ValueError("Cosmos K/V cache token geometry changed")
+        if (
+            self.k_cache.shape[0] != k.shape[0]
+            or self.k_cache.shape[3:] != k.shape[2:]
+            or self.k_cache.device != k.device
+            or self.k_cache.dtype != k.dtype
+        ):
+            raise ValueError("Cosmos K/V cache tensor geometry changed")
+
+        slot = self._cache_slot(frame_index)
+        replaced_frame = self._cache_frame_by_slot.get(slot)
+        if replaced_frame is not None:
+            self._cache_slot_by_frame.pop(replaced_frame, None)
+        self.k_cache[:, slot].copy_(k.detach())
+        self.v_cache[:, slot].copy_(v.detach())
+        self._cache_slot_by_frame[frame_index] = slot
+        self._cache_frame_by_slot[slot] = frame_index
+
+    def _read_cache_frames(
+        self,
+        frame_indices: list[int],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.k_cache is None or self.v_cache is None:
+            raise RuntimeError("Cosmos K/V cache is empty")
+        missing = [
+            index for index in frame_indices
+            if index not in self._cache_slot_by_frame
+        ]
+        if missing:
+            raise RuntimeError(
+                f"Cosmos KV cache is missing frames: {missing[:4]}"
+            )
+        slots = tuple(self._cache_slot_by_frame[index] for index in frame_indices)
+        slot_indices = self._cache_slot_indices.get(slots)
+        if slot_indices is None:
+            slot_indices = torch.tensor(
+                slots,
+                device=self.k_cache.device,
+                dtype=torch.long,
+            )
+            self._cache_slot_indices[slots] = slot_indices
+        cached_k = self.k_cache.index_select(1, slot_indices).flatten(1, 2)
+        cached_v = self.v_cache.index_select(1, slot_indices).flatten(1, 2)
+        return cached_k, cached_v
+
+    def _full_sequence_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        video_size: VideoSize,
+    ) -> torch.Tensor:
+        if q.shape != k.shape or k.shape != v.shape:
+            raise ValueError("Full causal attention requires matching Q/K/V shapes")
+        tokens_per_frame = video_size.H * video_size.W
+        expected_tokens = video_size.T * tokens_per_frame
+        if q.shape[1] != expected_tokens:
+            raise ValueError(
+                f"Expected {expected_tokens} video tokens, received {q.shape[1]}"
+            )
+        block_mask, padded_length = self._block_causal_mask(
+            device=q.device,
+            num_frames=video_size.T,
+            tokens_per_frame=tokens_per_frame,
+        )
+
+        if padded_length:
+            padding = q.new_zeros(
+                q.shape[0], padded_length, q.shape[2], q.shape[3]
+            )
+            q = torch.cat([q, padding], dim=1)
+            k = torch.cat([k, padding], dim=1)
+            v = torch.cat([v, padding], dim=1)
+
+        output = flex_attention(
+            query=q.transpose(1, 2),
+            key=k.transpose(1, 2),
+            value=v.transpose(1, 2),
+            block_mask=block_mask,
+        ).transpose(1, 2)
+        if padded_length:
+            output = output[:, :-padded_length]
+        return output.flatten(2)
+
+    def _block_causal_mask(
+        self,
+        *,
+        device: torch.device,
+        num_frames: int,
+        tokens_per_frame: int,
+    ) -> tuple[BlockMask, int]:
+        total_length = num_frames * tokens_per_frame
+        padded_length = math.ceil(total_length / 128) * 128 - total_length
+        cache_key = (
+            str(device),
+            num_frames,
+            tokens_per_frame,
+            self.num_frame_per_block,
+            self.independent_first_frame,
+            self.local_attn_size,
+            self.sink_size,
+        )
+        if cache_key in self._block_mask_cache:
+            return self._block_mask_cache[cache_key], padded_length
+
+        padded_total = total_length + padded_length
+        block_tokens = self.num_frame_per_block * tokens_per_frame
+        prefix_tokens = tokens_per_frame if self.independent_first_frame else 0
+        use_local_attn = self.local_attn_size > 0
+        sink_tokens = max(self.sink_size, 0) * tokens_per_frame
+        window_tokens = (
+            max(self.local_attn_size - self.sink_size, 0)
+            * tokens_per_frame
+        )
+
+        query_ranges = _full_sequence_query_ranges(
+            total_length=total_length,
+            tokens_per_frame=tokens_per_frame,
+            num_frame_per_block=self.num_frame_per_block,
+            independent_first_frame=self.independent_first_frame,
+            local_attn_size=self.local_attn_size,
+            sink_size=self.sink_size,
+        )
+
+        def block_index(position):
+            if prefix_tokens == 0:
+                return position // block_tokens
+            is_prefix = position < prefix_tokens
+            generated_block = (position - prefix_tokens) // block_tokens + 1
+            return torch.where(
+                is_prefix,
+                torch.zeros_like(position),
+                generated_block,
+            )
+
+        def block_bounds(index):
+            if prefix_tokens == 0:
+                start = index * block_tokens
+                return start, start + block_tokens
+            is_prefix = index == 0
+            generated_block = index - 1
+            start = prefix_tokens + generated_block * block_tokens
+            end = start + block_tokens
+            return (
+                torch.where(is_prefix, torch.zeros_like(index), start),
+                torch.where(
+                    is_prefix,
+                    torch.full_like(index, prefix_tokens),
+                    end,
+                ),
+            )
+
+        def attention_mask(_batch, _head, query_index, key_index):
+            valid = (query_index < total_length) & (key_index < total_length)
+            query_block = block_index(query_index)
+            key_block = block_index(key_index)
+            query_start, query_end = block_bounds(query_block)
+            allowed = key_block <= query_block
+            if use_local_attn:
+                zero = query_end - query_end
+                recent_start = torch.maximum(
+                    query_end - window_tokens,
+                    zero,
+                )
+                promoted_sink_end = torch.minimum(
+                    torch.minimum(query_start, zero + sink_tokens),
+                    recent_start,
+                )
+                recent_start = torch.maximum(promoted_sink_end, recent_start)
+                allowed = allowed & (
+                    (key_index < promoted_sink_end)
+                    | (
+                        (key_index >= recent_start)
+                        & (key_index < query_end)
+                    )
+                )
+            return valid & allowed
+
+        block_mask = _block_mask_from_intervals(
+            query_ranges=query_ranges,
+            q_total=padded_total,
+            kv_total=padded_total,
+            mask_mod=attention_mask,
+            device=device,
+        )
+        self._block_mask_cache[cache_key] = block_mask
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            print(
+                "Cached Cosmos block-causal attention mask: "
+                f"frames={num_frames}, block_frames={self.num_frame_per_block}, "
+                f"tokens_per_frame={tokens_per_frame}",
+                flush=True,
+            )
+        return block_mask, padded_length
+
+    def _packed_score_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        teacher_forcing_layout: tuple[int, int],
+        video_size: VideoSize,
+    ) -> torch.Tensor:
+        """Run retained history and current targets through one causal mask."""
+        context_frames, noisy_start_frame = teacher_forcing_layout
+        if context_frames != noisy_start_frame:
+            raise ValueError("Packed score history must precede current targets")
+        return self._full_sequence_attention(q, k, v, video_size)
+
+    def _history_indices(
+        self,
+        current_idx: int,
+        current_frames: int = 1,
+    ) -> list[int]:
+        if current_frames <= 0:
+            raise ValueError("current_frames must be positive")
+        window = self.local_attn_size
+        if window == -1:
+            recent_start = self.sink_size
+        else:
+            recent_window = max(window - self.sink_size, 0)
+            recent_history = max(recent_window - current_frames, 0)
+            recent_start = max(
+                self.sink_size,
+                current_idx - recent_history,
+            )
+        sink = range(min(self.sink_size, current_idx, recent_start))
+        recent = range(recent_start, current_idx)
+        return list(sink) + list(recent)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        kv_cache_cfg: Optional[KVCacheConfig] = None,
+        video_size: Optional[VideoSize] = None,
+        teacher_forcing_layout: Optional[tuple[int, int]] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del kwargs
+        if teacher_forcing_layout is not None:
+            if kv_cache_cfg is not None and kv_cache_cfg.run_with_kv:
+                raise ValueError("Teacher forcing cannot be combined with KV caching")
+            if video_size is None:
+                raise ValueError("video_size is required for teacher forcing")
+            return self._packed_score_attention(
+                q,
+                k,
+                v,
+                teacher_forcing_layout,
+                video_size,
+            )
+        if kv_cache_cfg is None or not kv_cache_cfg.run_with_kv:
+            if video_size is None:
+                raise ValueError("video_size is required for causal Cosmos attention")
+            return self._full_sequence_attention(q, k, v, video_size)
+
+        current_idx = int(kv_cache_cfg.current_idx)
+        if kv_cache_cfg.store_kv and video_size is not None and video_size.T > 1:
+            tokens_per_frame = video_size.H * video_size.W
+            for frame_offset in range(video_size.T):
+                start = frame_offset * tokens_per_frame
+                end = start + tokens_per_frame
+                self._store_cache_frame(
+                    current_idx + frame_offset,
+                    k[:, start:end],
+                    v[:, start:end],
+                )
+            return self._full_sequence_attention(q, k, v, video_size)
+
+        if kv_cache_cfg.store_kv:
+            self._store_cache_frame(current_idx, k, v)
+
+        current_frames = video_size.T if video_size is not None else 1
+        history_indices = self._history_indices(
+            current_idx,
+            current_frames=current_frames,
+        )
+        if history_indices:
+            history_k, history_v = self._read_cache_frames(history_indices)
+        else:
+            history_k = history_v = None
+        if history_k is not None and history_k.shape[0] != k.shape[0]:
+            if history_k.shape[0] != 1:
+                raise ValueError(
+                    "Cached Cosmos batch cannot be broadcast to the current batch"
+                )
+            history_k = history_k.expand(k.shape[0], *history_k.shape[1:])
+            history_v = history_v.expand(v.shape[0], *history_v.shape[1:])
+        if history_k is not None:
+            k = torch.cat((history_k, k), dim=1)
+            v = torch.cat((history_v, v), dim=1)
+        return i4_attention_op(q, k, v)
+
+    def set_context_parallel_group(self, *args, **kwargs) -> None:
+        del args, kwargs
+
+
+class CausalCosmosModel(MinimalV1LVGDiT):
+    """Weight-compatible causal variant of the bidirectional Cosmos 2.5 DiT."""
+
+    def __init__(
+        self,
+        *args,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self._num_frame_per_block = 1
+        self._independent_first_frame = False
+        self.causal_attention_ops = []
+        for block in self.blocks:
+            attention_op = CausalCosmosAttention(
+                local_attn_size=local_attn_size,
+                sink_size=sink_size,
+            )
+            block.self_attn.attn_op = attention_op
+            self.causal_attention_ops.append(attention_op)
+
+    @property
+    def num_frame_per_block(self) -> int:
+        return self._num_frame_per_block
+
+    @num_frame_per_block.setter
+    def num_frame_per_block(self, value: int) -> None:
+        if value <= 0:
+            raise ValueError("num_frame_per_block must be positive")
+        self._num_frame_per_block = value
+        for attention_op in getattr(self, "causal_attention_ops", []):
+            attention_op.num_frame_per_block = value
+
+    @property
+    def independent_first_frame(self) -> bool:
+        return self._independent_first_frame
+
+    @independent_first_frame.setter
+    def independent_first_frame(self, value: bool) -> None:
+        self._independent_first_frame = bool(value)
+        for attention_op in getattr(self, "causal_attention_ops", []):
+            attention_op.independent_first_frame = bool(value)
+
+    def forward_teacher_forcing(
+        self,
+        noisy_x_B_C_T_H_W: torch.Tensor,
+        clean_x_B_C_T_H_W: torch.Tensor,
+        noisy_timesteps_B_T: torch.Tensor,
+        clean_timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        *,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        clean_condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        camera_condition_B_C_T_H_W: Optional[torch.Tensor] = None,
+        noisy_start_frame: int,
+    ) -> torch.Tensor:
+        """Score packed history and all current noisy targets in one pass."""
+        if (
+            condition_video_input_mask_B_C_T_H_W is None
+            or clean_condition_video_input_mask_B_C_T_H_W is None
+        ):
+            raise ValueError("Noisy and clean condition masks are required")
+        if noisy_timesteps_B_T.ndim == 1:
+            noisy_timesteps_B_T = noisy_timesteps_B_T.unsqueeze(1)
+        if clean_timesteps_B_T.ndim == 1:
+            clean_timesteps_B_T = clean_timesteps_B_T.unsqueeze(1)
+        noisy_frames = noisy_x_B_C_T_H_W.shape[2]
+        context_frames = clean_x_B_C_T_H_W.shape[2]
+        if (
+            noisy_x_B_C_T_H_W.shape[:2] != clean_x_B_C_T_H_W.shape[:2]
+            or noisy_x_B_C_T_H_W.shape[-2:] != clean_x_B_C_T_H_W.shape[-2:]
+            or context_frames != noisy_start_frame
+        ):
+            raise ValueError("Packed scoring requires history before noisy targets")
+        if not 0 < noisy_start_frame < noisy_frames:
+            raise ValueError("noisy_start_frame must select a non-empty suffix")
+        if noisy_timesteps_B_T.shape != (noisy_x_B_C_T_H_W.shape[0], noisy_frames):
+            raise ValueError("Noisy timesteps must cover every noisy frame")
+        if clean_timesteps_B_T.shape != (clean_x_B_C_T_H_W.shape[0], context_frames):
+            raise ValueError("Clean timesteps must cover every history frame")
+
+        noisy_input = torch.cat(
+            [
+                noisy_x_B_C_T_H_W,
+                condition_video_input_mask_B_C_T_H_W.type_as(
+                    noisy_x_B_C_T_H_W
+                ),
+            ],
+            dim=1,
+        )
+        clean_input = torch.cat(
+            [
+                clean_x_B_C_T_H_W,
+                clean_condition_video_input_mask_B_C_T_H_W.type_as(
+                    clean_x_B_C_T_H_W
+                ),
+            ],
+            dim=1,
+        )
+        noisy_hidden, noisy_rope, noisy_extra_pos = self.prepare_embedded_sequence(
+            noisy_input,
+            fps=fps,
+            padding_mask=padding_mask,
+        )
+        with torch.no_grad():
+            clean_hidden, clean_rope, clean_extra_pos = (
+                self.prepare_embedded_sequence(
+                    clean_input,
+                    fps=fps,
+                    padding_mask=padding_mask,
+                )
+            )
+        if clean_hidden.shape[:1] + clean_hidden.shape[2:] != noisy_hidden.shape[:1] + noisy_hidden.shape[2:]:
+            raise ValueError("Embedded history and noisy video grids must match")
+        if noisy_rope is None or clean_rope is None:
+            raise ValueError("Causal Cosmos teacher forcing requires RoPE")
+        tokens_per_frame = noisy_hidden.shape[2] * noisy_hidden.shape[3]
+        target_hidden = noisy_hidden[:, noisy_start_frame:]
+        packed_hidden = torch.cat([clean_hidden, target_hidden], dim=1)
+        packed_rope = torch.cat(
+            [
+                clean_rope,
+                noisy_rope[noisy_start_frame * tokens_per_frame:],
+            ],
+            dim=0,
+        )
+
+        def pack_optional(clean_value, noisy_value):
+            if clean_value is None or noisy_value is None:
+                if clean_value is not None or noisy_value is not None:
+                    raise ValueError("Packed score embeddings must match")
+                return None
+            return torch.cat(
+                [clean_value, noisy_value[:, noisy_start_frame:]], dim=1
+            )
+
+        packed_extra_pos = pack_optional(clean_extra_pos, noisy_extra_pos)
+
+        packed_camera = None
+        if camera_condition_B_C_T_H_W is not None:
+            camera = camera_condition_B_C_T_H_W.permute(
+                0, 2, 3, 4, 1
+            ).contiguous()
+            if camera.shape[:4] != noisy_hidden.shape[:4]:
+                raise ValueError(
+                    "Camera conditioning does not match the teacher-forcing grid: "
+                    f"{tuple(camera.shape)} versus {tuple(noisy_hidden.shape)}"
+                )
+            packed_camera = torch.cat(
+                [camera[:, :context_frames], camera[:, noisy_start_frame:]],
+                dim=1,
+            )
+
+        if self.use_crossattn_projection:
+            crossattn_emb = self.crossattn_proj(crossattn_emb)
+
+        noisy_timesteps_B_T = noisy_timesteps_B_T * self.timestep_scale
+        clean_timesteps_B_T = clean_timesteps_B_T * self.timestep_scale
+        noisy_time, noisy_adaln_lora = self.t_embedder(noisy_timesteps_B_T)
+        noisy_time = self.t_embedding_norm(noisy_time)
+        with torch.no_grad():
+            clean_time, clean_adaln_lora = self.t_embedder(
+                clean_timesteps_B_T
+            )
+            clean_time = self.t_embedding_norm(clean_time)
+        packed_time = torch.cat(
+            [clean_time, noisy_time[:, noisy_start_frame:]], dim=1
+        )
+        packed_adaln_lora = pack_optional(clean_adaln_lora, noisy_adaln_lora)
+
+        layout = (context_frames, noisy_start_frame)
+        for block in self.blocks:
+            packed_hidden = block(
+                packed_hidden,
+                packed_time,
+                crossattn_emb,
+                rope_emb_L_1_1_D=packed_rope,
+                adaln_lora_B_T_3D=packed_adaln_lora,
+                extra_per_block_pos_emb=packed_extra_pos,
+                camera_B_T_H_W_C=packed_camera,
+                teacher_forcing_layout=layout,
+            )
+            packed_hidden = torch.cat(
+                [
+                    packed_hidden[:, :context_frames].detach(),
+                    packed_hidden[:, context_frames:],
+                ],
+                dim=1,
+            )
+
+        output = self.final_layer(
+            packed_hidden[:, context_frames:],
+            packed_time[:, context_frames:],
+            adaln_lora_B_T_3D=(
+                packed_adaln_lora[:, context_frames:]
+                if packed_adaln_lora is not None else None
+            ),
+        )
+        scored_suffix = self.unpatchify(output)
+        prefix = scored_suffix.new_zeros(
+            scored_suffix.shape[0],
+            scored_suffix.shape[1],
+            noisy_start_frame,
+            scored_suffix.shape[3],
+            scored_suffix.shape[4],
+        )
+        return torch.cat([prefix, scored_suffix], dim=2)
+
+    def forward_seq(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        video_pos: VideoSeqPos,
+        timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        *,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        camera_condition_B_C_T_H_W: Optional[torch.Tensor] = None,
+        full_video_size: Optional[tuple[int, int, int]] = None,
+        kv_cache_cfg: Optional[KVCacheConfig] = None,
+    ) -> torch.Tensor:
+        """Run one causal sequence chunk using the same blocks and weights."""
+        if condition_video_input_mask_B_C_T_H_W is None:
+            raise ValueError("condition_video_input_mask_B_C_T_H_W is required")
+
+        x_B_C_T_H_W = torch.cat(
+            [
+                x_B_C_T_H_W,
+                condition_video_input_mask_B_C_T_H_W.type_as(x_B_C_T_H_W),
+            ],
+            dim=1,
+        )
+        x_B_T_H_W_D, _, _ = self.prepare_embedded_sequence(
+            x_B_C_T_H_W,
+            fps=fps,
+            padding_mask=padding_mask,
+        )
+        _, token_t, token_h, token_w, _ = x_B_T_H_W_D.shape
+        if token_t * token_h * token_w != video_pos.size():
+            raise ValueError("Cosmos sequence positions do not match the input tokens")
+        camera_B_T_H_W_C = None
+        if camera_condition_B_C_T_H_W is not None:
+            camera_B_T_H_W_C = camera_condition_B_C_T_H_W.permute(
+                0, 2, 3, 4, 1
+            ).contiguous()
+            if camera_B_T_H_W_C.shape[:4] != x_B_T_H_W_D.shape[:4]:
+                raise ValueError(
+                    "Camera conditioning does not match the causal video grid: "
+                    f"{tuple(camera_B_T_H_W_C.shape)} versus {tuple(x_B_T_H_W_D.shape)}"
+                )
+
+        if self.use_crossattn_projection:
+            crossattn_emb = self.crossattn_proj(crossattn_emb)
+
+        if timesteps_B_T.ndim == 1:
+            timesteps_B_T = timesteps_B_T.unsqueeze(1)
+        timesteps_B_T = timesteps_B_T * self.timestep_scale
+        time_embedding, adaln_lora = self.t_embedder(timesteps_B_T)
+        time_embedding = self.t_embedding_norm(time_embedding)
+
+        if full_video_size is None:
+            full_t = int(video_pos.pos_t.max().item()) + 1
+            full_h = int(video_pos.pos_h.max().item()) + 1
+            full_w = int(video_pos.pos_w.max().item()) + 1
+        else:
+            full_t, full_h, full_w = full_video_size
+        rope = self.pos_embedder.generate_embeddings(
+            torch.Size([1, full_t, full_h, full_w, self.model_channels])
+        )
+        linear_index = (
+            video_pos.pos_t * (full_h * full_w)
+            + video_pos.pos_h * full_w
+            + video_pos.pos_w
+        )
+        rope = rope.index_select(0, linear_index.to(device=rope.device))
+
+        for block in self.blocks:
+            x_B_T_H_W_D = block(
+                x_B_T_H_W_D,
+                time_embedding,
+                crossattn_emb,
+                rope_emb_L_1_1_D=rope,
+                adaln_lora_B_T_3D=adaln_lora,
+                extra_per_block_pos_emb=None,
+                kv_cache_cfg=kv_cache_cfg,
+                camera_B_T_H_W_C=camera_B_T_H_W_C,
+            )
+
+        output = self.final_layer(
+            x_B_T_H_W_D,
+            time_embedding,
+            adaln_lora_B_T_3D=adaln_lora,
+        )
+        return self.unpatchify(output)

cosmos/kv_cache.py

@@ -0,0 +1,169 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class KVCacheConfig:
+    run_with_kv: bool = False
+    store_kv: bool = False
+    current_idx: int = 0
+    recompute_cross_attn_kv: bool = False
+
+
+class AttentionOpWithKVCache(nn.Module):
+    """A thin wrapper that adds K/V caching to an existing attention op.
+
+    This wrapper expects the wrapped op to accept (q, k, v, attn_mask=None)
+    and return attention outputs with heads already flattened on the last dim.
+
+    Cache semantics:
+    - Cache entries are stored as per-chunk tensors, where each chunk corresponds
+      to one latent frame composed of HxW tokens (after patchify).
+    - The `max_cache_size` capacity therefore refers to the number of latent
+      frames (chunks), NOT the number of individual tokens.
+    - When `max_cache_size` is None, the cache grows without an automatic
+      rolling window; otherwise, it acts as a rolling window of at most
+      `max_cache_size` frames. Upon overflow, the oldest frames are dropped.
+    """
+
+    def __init__(self, attn_op: nn.Module | Any, max_cache_size: Optional[int] = None):
+        """Initialize the KV cache wrapper.
+
+        Args:
+            attn_op: The underlying attention operation (q, k, v[, attn_mask]) -> out.
+            max_cache_size: Optional capacity measured in number of latent frames
+                (chunks). Each chunk is a single frame worth of HxW tokens. If None,
+                the cache does not enforce a rolling capacity.
+        """
+        super().__init__()
+        self.attn_op = attn_op
+        self.reset_kv_cache(max_cache_size=max_cache_size)
+        self.pg: Optional[Any] = None
+        self.stream: Optional[Any] = None
+
+    def reset_kv_cache(self, max_cache_size: Optional[int] = None) -> None:
+        """Reset/initialize the KV caches.
+
+        Args:
+            max_cache_size: Optional capacity measured in number of latent frames
+                (chunks). Each chunk is a single frame worth of HxW tokens. If None,
+                the cache does not enforce a rolling capacity.
+        """
+        # Initialize list-based caches and optionally set capacity in chunks
+        self.start_idx = 0
+        self.k_cache: list[torch.Tensor | None] = [None] * (max_cache_size or 99999)
+        self.v_cache: list[torch.Tensor | None] = [None] * (max_cache_size or 99999)
+        self.max_cache_size = max_cache_size
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        kv_cache_cfg: KVCacheConfig,
+        **kwargs,
+    ) -> torch.Tensor:
+        assert self.k_cache is not None and self.v_cache is not None, (
+            "KV cache is not initialized. Call reset_kv_cache() first."
+        )
+
+        # Store into cache at start_idx location (list-based)
+        if kv_cache_cfg.store_kv:
+            index = int(kv_cache_cfg.current_idx)
+            self.k_cache[index] = k.detach()
+            self.v_cache[index] = v.detach()
+
+        # Prepend cached prefix up to start_idx (list-based)
+        if kv_cache_cfg.run_with_kv and kv_cache_cfg.current_idx > 0:
+            history_k = self.k_cache[self.start_idx : kv_cache_cfg.current_idx]
+            history_v = self.v_cache[self.start_idx : kv_cache_cfg.current_idx]
+            assert not any(x is None for x in history_k)
+            assert not any(x is None for x in history_v)
+            k_out = torch.cat(history_k + [k], dim=1)  # type: ignore
+            v_out = torch.cat(history_v + [v], dim=1)  # type: ignore
+        else:
+            k_out = k
+            v_out = v
+
+        # Enforce rolling capacity in number of cached chunks (frames)
+        if kv_cache_cfg.run_with_kv and self.max_cache_size is not None:
+            # Instead of deleting, just update start_idx for rolling window
+            self.start_idx = max(0, int(kv_cache_cfg.current_idx) - self.max_cache_size)
+
+        return self.attn_op(q, k_out, v_out, **kwargs)
+
+    def set_context_parallel_group(self, process_group, ranks, stream, cp_comm_type: str = "p2p"):
+        self.attn_op.set_context_parallel_group(process_group, ranks, stream, cp_comm_type=cp_comm_type)  # type: ignore
+
+
+class VideoSeqPos:
+    """Flattened 3D grid positions for a video clip.
+
+    Stores flattened t/h/w indices of length L = T*H*W to enable constructing
+    RoPE frequencies aligned with global positions across sequential chunks.
+    """
+
+    def __init__(self, T: int, H: int, W: int, pos_h=None, pos_w=None, pos_t=None) -> None:
+        self.T = T
+        self.H = H
+        self.W = W
+
+        if pos_h is not None and pos_w is not None and pos_t is not None:
+            self.pos_h = pos_h.to(dtype=torch.long)
+            self.pos_w = pos_w.to(dtype=torch.long)
+            self.pos_t = pos_t.to(dtype=torch.long)
+            return
+
+        device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu")
+        t = torch.arange(self.T, device=device, dtype=torch.long)
+        h = torch.arange(self.H, device=device, dtype=torch.long)
+        w = torch.arange(self.W, device=device, dtype=torch.long)
+        pos_t, pos_h, pos_w = torch.meshgrid(t, h, w, indexing="ij")
+        self.pos_t = pos_t.reshape(-1)
+        self.pos_h = pos_h.reshape(-1)
+        self.pos_w = pos_w.reshape(-1)
+
+    def size(self) -> int:
+        return int(self.pos_h.numel())
+
+    def frame(self, t_idx: int) -> "VideoSeqPos":
+        """Return a `VideoSeqPos` view for a single frame at absolute index `t_idx`.
+
+        This is useful for streaming / KV-cache inference where the model is run on
+        one frame at a time but RoPE positions must reflect global video indices.
+        """
+        t_idx = int(t_idx)
+        if t_idx < 0 or t_idx >= int(self.T):
+            raise IndexError(f"t_idx out of range: {t_idx} (valid: [0, {self.T}))")
+        tokens_per_frame = int(self.H) * int(self.W)
+        start = t_idx * tokens_per_frame
+        end = start + tokens_per_frame
+        return VideoSeqPos(
+            T=1,
+            H=int(self.H),
+            W=int(self.W),
+            pos_h=self.pos_h[start:end],
+            pos_w=self.pos_w[start:end],
+            pos_t=self.pos_t[start:end],
+        )

cosmos/minimal_v1_lvg_dit.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Modifications Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from cosmos.minimal_v4_dit import MiniTrainDIT
+from cosmos.runtime import DataType
+
+
+class MinimalV1LVGDiT(MiniTrainDIT):
+    def __init__(self, *args, timestep_scale: float = 1.0, **kwargs):
+        assert "in_channels" in kwargs, "in_channels must be provided"
+        kwargs["in_channels"] += 1  # Add 1 for the condition mask
+        self.timestep_scale = timestep_scale
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        condition_video_input_mask_B_C_T_H_W: Optional[torch.Tensor] = None,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        data_type: Optional[DataType] = DataType.VIDEO,
+        intermediate_feature_ids: Optional[List[int]] = None,
+        img_context_emb: Optional[torch.Tensor] = None,
+        camera_condition_B_C_T_H_W: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor | List[torch.Tensor] | Tuple[torch.Tensor, List[torch.Tensor]]:
+        del kwargs
+
+        if data_type == DataType.VIDEO:
+            x_B_C_T_H_W = torch.cat([x_B_C_T_H_W, condition_video_input_mask_B_C_T_H_W.type_as(x_B_C_T_H_W)], dim=1)
+        else:
+            B, _, T, H, W = x_B_C_T_H_W.shape
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, torch.zeros((B, 1, T, H, W), dtype=x_B_C_T_H_W.dtype, device=x_B_C_T_H_W.device)], dim=1
+            )
+        return super().forward(
+            x_B_C_T_H_W=x_B_C_T_H_W,
+            timesteps_B_T=timesteps_B_T * self.timestep_scale,
+            crossattn_emb=crossattn_emb,
+            fps=fps,
+            padding_mask=padding_mask,
+            data_type=data_type,
+            intermediate_feature_ids=intermediate_feature_ids,
+            img_context_emb=img_context_emb,
+            camera_condition_B_C_T_H_W=camera_condition_B_C_T_H_W,
+        )

cosmos/minimal_v4_dit.py

@@ -0,0 +1,2054 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Modifications Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+import collections
+import math
+from collections import namedtuple
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional, Tuple, Union
+
+from cosmos.kv_cache import AttentionOpWithKVCache, KVCacheConfig
+
+try:
+    import megatron.core.parallel_state as parallel_state
+
+    USE_MEGATRON = True
+except ImportError:
+    USE_MEGATRON = False
+
+import numpy as np
+import torch
+import torch.amp as amp
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+from torch import nn
+from torch.distributed import ProcessGroup, get_process_group_ranks
+from torch.distributed._composable.fsdp import fully_shard
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper as ptd_checkpoint_wrapper
+
+try:
+    from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
+except ImportError:
+    CheckpointPolicy = None
+
+from torchvision import transforms
+
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
+
+from cosmos.model_weights_stats import WeightTrainingStat
+from cosmos.runtime import (
+    DataType,
+    DotProductAttention,
+    MinimalA2AAttnOp,
+    NattenA2AAttnOp,
+    NeighborhoodAttention,
+    RMSNorm,
+    apply_rotary_pos_emb,
+    attention,
+    log,
+    split_inputs_cp,
+)
+from cosmos.selective_activation_checkpoint import SACConfig as _SACConfig
+
+
+# selective activation checkpoint; only apply to the minimal v4 model. if there are change in the networks, some policy will not work as we expect.
+def predict2_2B_720_context_fn():
+    op_count = collections.defaultdict(int)
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        mode = "recompute" if ctx.is_recompute else "forward"
+        if func == torch.ops.aten.mm.default:
+            op_count_key = f"{mode}_mm_count"
+            # from cosmos_predict2._src.imaginaire.utils import log
+            # log.info(f"op_count_key: {op_count_key}, op_count[op_count_key]: {op_count[op_count_key]}, {args[0].shape}, {args[1].shape}")
+            # there are totally 6 + 4 + 4 + 2 = 16 block
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 16
+            if op_count[op_count_key] > 8:  # recompute self attn first 3 linear layers
+                return CheckpointPolicy.MUST_SAVE
+        if "flash_attn" in str(func):
+            op_count_key = f"{mode}_flash_attn_count"
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 2
+            if op_count[op_count_key]:
+                return CheckpointPolicy.MUST_SAVE
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+def predict2_2B_720_context_fn_aggressive():
+    op_count = collections.defaultdict(int)
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        # The default policy is to recompute everything. This is the most memory-efficient
+        # starting point. We then selectively choose what to save.
+        default_policy = CheckpointPolicy.PREFER_RECOMPUTE
+
+        # Save the output of Flash Attention. This is the most computationally
+        # expensive part of a transformer block. Saving its output provides a
+        # good balance between memory savings and computational overhead.
+        if "flash_attn" in str(func):
+            return CheckpointPolicy.MUST_SAVE
+
+        # All other operations (e.g., torch.ops.aten.mm.default, layer norms, additions)
+        # will fall through to the default policy and be recomputed.
+        return default_policy
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+def predict2_2B_720_context_fn_aggressive_v2():
+    """
+    The most memory-aggressive checkpointing policy. Recomputes ALL operations.
+    """
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        # The policy is to always recompute everything.
+        # This saves the maximum amount of memory but incurs the highest
+        # computational cost during the backward pass.
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+def predict2_14B_720_context_fn():
+    op_count = collections.defaultdict(int)
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        mode = "recompute" if ctx.is_recompute else "forward"
+        if func == torch.ops.aten.mm.default:
+            op_count_key = f"{mode}_mm_count"
+            # from cosmos_predict2._src.imaginaire.utils import log
+            # log.info(f"op_count_key: {op_count_key}, op_count[op_count_key]: {op_count[op_count_key]}, {args[0].shape}, {args[1].shape}")
+            # there are totally 6 + 4 + 4 + 2 = 16 block
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 16
+            if op_count[op_count_key] > 8:  # recompute self attn first 1 linear layers
+                return CheckpointPolicy.MUST_SAVE
+        if "flash_attn" in str(func):
+            op_count_key = f"{mode}_flash_attn_count"
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 2
+            if op_count[op_count_key]:
+                return CheckpointPolicy.MUST_SAVE
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+def predict2_14B_720_context_fn_aggressive():
+    op_count = collections.defaultdict(int)
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        mode = "recompute" if ctx.is_recompute else "forward"
+        if func == torch.ops.aten.mm.default:
+            op_count_key = f"{mode}_mm_count"
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 16
+            if op_count[op_count_key] > 12:  # recompute self attn first 1 linear layers
+                return CheckpointPolicy.MUST_SAVE
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+def linear_selfattn_context_fn():
+    op_count = collections.defaultdict(int)
+
+    def policy_fn(ctx, func, *args, **kwargs):
+        mode = "recompute" if ctx.is_recompute else "forward"
+        if func == torch.ops.aten.mm.default:
+            return CheckpointPolicy.MUST_SAVE
+        if "flash_attn" in str(func):
+            op_count_key = f"{mode}_flash_attn_count"
+            op_count[op_count_key] = (op_count[op_count_key] + 1) % 2
+            if op_count[op_count_key]:
+                return CheckpointPolicy.MUST_SAVE
+        return CheckpointPolicy.PREFER_RECOMPUTE
+
+    return create_selective_checkpoint_contexts(policy_fn)
+
+
+class CheckpointMode(str, Enum):
+    NONE = "none"
+    MM_ONLY = "mm_only"
+    BLOCK_WISE = "block_wise"
+    LINEAR_SELFATTN = "linear_selfattn"
+    PREDICT2_2B_720 = "predict2_2b_720"
+    PREDICT2_14B_720 = "predict2_14b_720"
+    PREDICT2_2B_720_AGGRESSIVE = "predict2_2b_720_aggressive"
+    PREDICT2_2B_720_AGGRESSIVE_V2 = "predict2_2b_720_aggressive_v2"
+    PREDICT2_14B_720_AGGRESSIVE = "predict2_14b_720_aggressive"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+@dataclass
+class SACConfig(_SACConfig):
+    def get_context_fn(self):
+        if self.mode == CheckpointMode.LINEAR_SELFATTN:
+            return linear_selfattn_context_fn
+        elif self.mode == CheckpointMode.PREDICT2_2B_720:
+            return predict2_2B_720_context_fn
+        elif self.mode == CheckpointMode.PREDICT2_2B_720_AGGRESSIVE:
+            return predict2_2B_720_context_fn_aggressive
+        elif self.mode == CheckpointMode.PREDICT2_2B_720_AGGRESSIVE_V2:
+            return predict2_2B_720_context_fn_aggressive_v2
+        elif self.mode == CheckpointMode.PREDICT2_14B_720:
+            return predict2_14B_720_context_fn
+        elif self.mode == CheckpointMode.PREDICT2_14B_720_AGGRESSIVE:
+            return predict2_14B_720_context_fn_aggressive
+        else:
+            # Reuse parent class implementation for other modes
+            return super().get_context_fn()
+
+
+VideoSize = namedtuple("VideoSize", ["T", "H", "W"])
+
+
+# ---------------------- Feed Forward Network -----------------------
+class GPT2FeedForward(nn.Module):
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.activation = nn.GELU()
+        self.layer1 = nn.Linear(d_model, d_ff, bias=False)
+        self.layer2 = nn.Linear(d_ff, d_model, bias=False)
+
+        self._layer_id = None
+        self._dim = d_model
+        self._hidden_dim = d_ff
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        std = 1.0 / math.sqrt(self._dim)
+        torch.nn.init.trunc_normal_(self.layer1.weight, std=std, a=-3 * std, b=3 * std)
+
+        # scale init by depth as in https://arxiv.org/abs/1908.11365 -- worked slightly better.
+        std = 1.0 / math.sqrt(self._hidden_dim)
+        if self._layer_id is not None:
+            std = std / math.sqrt(2 * (self._layer_id + 1))
+        torch.nn.init.trunc_normal_(self.layer2.weight, std=std, a=-3 * std, b=3 * std)
+
+    def forward(self, x: torch.Tensor):
+        x = self.layer1(x)
+
+        x = self.activation(x)
+        x = self.layer2(x)
+        return x
+
+
+def torch_attention_op(
+    q_B_S_H_D: torch.Tensor,
+    k_B_S_H_D: torch.Tensor,
+    v_B_S_H_D: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    flatten_heads: bool = True,
+) -> torch.Tensor:
+    """Scaled dot-product attention with optional mask.
+
+    Inputs are shaped [B, S, H, D]. If flatten_heads=True, flattens heads to return [B, S, H*D].
+    Otherwise returns [B, S, H, D].
+    """
+    q_B_H_S_D = rearrange(q_B_S_H_D, "b s h d -> b h s d")
+    k_B_H_S_D = rearrange(k_B_S_H_D, "b s h d -> b h s d")
+    v_B_H_S_D = rearrange(v_B_S_H_D, "b s h d -> b h s d")
+    result_B_H_S_D = torch.nn.functional.scaled_dot_product_attention(
+        q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, attn_mask=attn_mask
+    )
+    if flatten_heads:
+        return rearrange(result_B_H_S_D, "b h s d -> b s (h d)")
+    else:
+        return rearrange(result_B_H_S_D, "b h s d -> b s h d")
+
+
+def flex_attention_op(
+    q_B_S_H_D: torch.Tensor,
+    k_B_S_H_D: torch.Tensor,
+    v_B_S_H_D: torch.Tensor,
+    attn_mask: Optional[BlockMask] = None,
+    flatten_heads: bool = True,
+) -> torch.Tensor:
+    # Rearrange to [B, H, S, D]
+    q_B_H_Sq_D = rearrange(q_B_S_H_D, "b s h d -> b h s d")
+    k_B_H_Sk_D = rearrange(k_B_S_H_D, "b s h d -> b h s d")
+    v_B_H_Sk_D = rearrange(v_B_S_H_D, "b s h d -> b h s d")
+
+    S_q = q_B_H_Sq_D.shape[2]
+    S_kv = k_B_H_Sk_D.shape[2]
+    # Right-pad to multiples of 128 for optimal FlexAttention kernels
+    pad_q = ((S_q + 127) // 128) * 128 - S_q
+    pad_kv = ((S_kv + 127) // 128) * 128 - S_kv
+
+    if pad_q > 0:
+        q_pad_tensor = torch.zeros(
+            (q_B_H_Sq_D.shape[0], q_B_H_Sq_D.shape[1], pad_q, q_B_H_Sq_D.shape[3]),
+            device=q_B_H_Sq_D.device,
+            dtype=q_B_H_Sq_D.dtype,
+        )
+        q_cat = torch.cat([q_B_H_Sq_D, q_pad_tensor], dim=2)
+    else:
+        q_cat = q_B_H_Sq_D
+
+    if pad_kv > 0:
+        kv_pad_tensor = torch.zeros(
+            (k_B_H_Sk_D.shape[0], k_B_H_Sk_D.shape[1], pad_kv, k_B_H_Sk_D.shape[3]),
+            device=k_B_H_Sk_D.device,
+            dtype=k_B_H_Sk_D.dtype,
+        )
+        k_cat = torch.cat([k_B_H_Sk_D, kv_pad_tensor], dim=2)
+        v_cat = torch.cat([v_B_H_Sk_D, kv_pad_tensor], dim=2)
+    else:
+        k_cat, v_cat = k_B_H_Sk_D, v_B_H_Sk_D
+
+    block_mask = None
+    if attn_mask is not None and isinstance(attn_mask, BlockMask):
+        block_mask = attn_mask
+    else:
+        # When padding is introduced without an explicit mask, build a validity mask
+        if pad_q > 0 or pad_kv > 0:
+
+            def allow_valid(b, h, q_idx, kv_idx):
+                return (q_idx < S_q) & (kv_idx < S_kv)
+
+            block_mask = create_block_mask(
+                allow_valid,
+                B=None,
+                H=None,
+                Q_LEN=q_cat.shape[2],
+                KV_LEN=k_cat.shape[2],
+                _compile=True,
+                device=q_cat.device,
+            )
+
+    if block_mask is not None:
+        out_B_H_Sqp_D = torch.compile(flex_attention)(query=q_cat, key=k_cat, value=v_cat, block_mask=block_mask)
+    else:
+        out_B_H_Sqp_D = torch.compile(flex_attention)(query=q_cat, key=k_cat, value=v_cat)
+
+    out_B_H_Sq_D = out_B_H_Sqp_D[:, :, :S_q] if pad_q > 0 else out_B_H_Sqp_D
+    if flatten_heads:
+        return rearrange(out_B_H_Sq_D, "b h s d -> b s (h d)")
+    else:
+        return rearrange(out_B_H_Sq_D, "b h s d -> b s h d")
+
+
+def i4_attention_op(
+    q_B_S_H_D: torch.Tensor,
+    k_B_S_H_D: torch.Tensor,
+    v_B_S_H_D: torch.Tensor,
+    flatten_heads: bool = True,
+    **kwargs: dict,
+) -> torch.Tensor:
+    """
+    I4 regular (bidirectional) attention.
+    Matches torch_attention_op's signature but omits attn_mask (full attention assumed).
+    Ignores any additional kwargs (e.g., video_size).
+    """
+    out_B_S_H_D = attention(
+        query=q_B_S_H_D,
+        key=k_B_S_H_D,
+        value=v_B_S_H_D,
+        is_causal=False,
+    )
+    if isinstance(out_B_S_H_D, tuple):
+        out_B_S_H_D = out_B_S_H_D[0]
+    if flatten_heads:
+        return rearrange(out_B_S_H_D, "b s h d -> b s (h d)")
+    else:
+        return out_B_S_H_D
+
+
+class Attention(nn.Module):
+    """
+    A flexible attention module supporting both self-attention and cross-attention mechanisms.
+
+    This module implements a multi-head attention layer that can operate in either self-attention
+    or cross-attention mode. The mode is determined by whether a context dimension is provided.
+    The implementation uses scaled dot-product attention and supports optional bias terms and
+    dropout regularization.
+
+    Args:
+        query_dim (int): The dimensionality of the query vectors.
+        context_dim (int, optional): The dimensionality of the context (key/value) vectors.
+            If None, the module operates in self-attention mode using query_dim. Default: None
+        n_heads (int, optional): Number of attention heads for multi-head attention. Default: 8
+        head_dim (int, optional): The dimension of each attention head. Default: 64
+        dropout (float, optional): Dropout probability applied to the output. Default: 0.0
+        qkv_format (str, optional): Format specification for QKV tensors. Default: "bshd"
+        backend (str, optional): Backend to use for the attention operation. Default: "transformer_engine"
+
+    Examples:
+        >>> # Self-attention with 512 dimensions and 8 heads
+        >>> self_attn = Attention(query_dim=512)
+        >>> x = torch.randn(32, 16, 512)  # (batch_size, seq_len, dim)
+        >>> out = self_attn(x)  # (32, 16, 512)
+
+        >>> # Cross-attention
+        >>> cross_attn = Attention(query_dim=512, context_dim=256)
+        >>> query = torch.randn(32, 16, 512)
+        >>> context = torch.randn(32, 8, 256)
+        >>> out = cross_attn(query, context)  # (32, 16, 512)
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        context_dim=None,
+        n_heads=8,
+        head_dim=64,
+        dropout=0.0,
+        qkv_format: str = "bshd",
+        backend: str = "transformer_engine",
+        use_wan_fp32_strategy: bool = False,
+    ) -> None:
+        super().__init__()
+        log.debug(
+            f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
+            f"{n_heads} heads with a dimension of {head_dim}."
+        )
+        self.is_selfattn = context_dim is None  # self attention
+
+        assert backend in ["transformer_engine", "torch", "torch-flex", "minimal_a2a", "i4"], (
+            f"Invalid backend: {backend}"
+        )
+        self.backend = backend
+
+        context_dim = query_dim if context_dim is None else context_dim
+        inner_dim = head_dim * n_heads
+
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.qkv_format = qkv_format
+        self.query_dim = query_dim
+        self.context_dim = context_dim
+        self.use_wan_fp32_strategy = use_wan_fp32_strategy
+
+        self.q_proj = nn.Linear(query_dim, inner_dim, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=1e-6)
+
+        self.k_proj = nn.Linear(context_dim, inner_dim, bias=False)
+        self.k_norm = RMSNorm(self.head_dim, eps=1e-6)
+
+        self.v_proj = nn.Linear(context_dim, inner_dim, bias=False)
+        self.v_norm = nn.Identity()
+
+        self.output_proj = nn.Linear(inner_dim, query_dim, bias=False)
+        self.output_dropout = nn.Dropout(dropout) if dropout > 1e-4 else nn.Identity()
+        # Camera-control checkpoints store the projection under each
+        # self-attention module as `self_attn.cam_encoder`.
+        self.cam_encoder = None
+
+        if self.backend == "transformer_engine":
+            self.attn_op = DotProductAttention(
+                self.n_heads,
+                self.head_dim,
+                num_gqa_groups=self.n_heads,
+                attention_dropout=0,
+                qkv_format=qkv_format,
+                attn_mask_type="no_mask",
+            )
+        elif self.backend == "minimal_a2a":
+            self.attn_op = MinimalA2AAttnOp()
+        elif self.backend == "torch":
+            self.attn_op = torch_attention_op
+        elif self.backend == "torch-flex":
+            # FlexAttention backend; returns [B, S, H*D]
+            self.attn_op = flex_attention_op
+        elif self.backend == "i4":
+            # I4 spatio-temporal attention; returns [B, S, H*D]
+            self.attn_op = i4_attention_op
+
+        if not hasattr(self.attn_op, "set_context_parallel_group"):
+
+            def set_context_parallel_group(*args, **kwargs) -> None:
+                return None
+
+            self.attn_op.set_context_parallel_group = set_context_parallel_group
+
+        self._query_dim = query_dim
+        self._context_dim = context_dim
+        self._inner_dim = inner_dim
+
+    def init_weights(self) -> None:
+        std = 1.0 / math.sqrt(self._query_dim)
+        torch.nn.init.trunc_normal_(self.q_proj.weight, std=std, a=-3 * std, b=3 * std)
+        std = 1.0 / math.sqrt(self._context_dim)
+        torch.nn.init.trunc_normal_(self.k_proj.weight, std=std, a=-3 * std, b=3 * std)
+        torch.nn.init.trunc_normal_(self.v_proj.weight, std=std, a=-3 * std, b=3 * std)
+
+        std = 1.0 / math.sqrt(self._inner_dim)
+        torch.nn.init.trunc_normal_(self.output_proj.weight, std=std, a=-3 * std, b=3 * std)
+
+        for layer in self.q_norm, self.k_norm, self.v_norm:
+            if hasattr(layer, "reset_parameters"):
+                layer.reset_parameters()
+
+    def compute_qkv(self, x, context=None, rope_emb=None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        q = self.q_proj(x)
+        context = x if context is None else context
+        k = self.k_proj(context)
+        v = self.v_proj(context)
+        q, k, v = map(
+            lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
+            (q, k, v),
+        )
+
+        def apply_norm_and_rotary_pos_emb(q, k, v, rope_emb):
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+            v = self.v_norm(v)
+            original_dtype = q.dtype
+            if self.is_selfattn and rope_emb is not None:  # only apply to self-attention!
+                if self.use_wan_fp32_strategy:  # wan will force q and k to fp32 before rotary pos emb
+                    q = q.to(torch.float32)
+                    k = k.to(torch.float32)
+                q = apply_rotary_pos_emb(q, rope_emb, tensor_format=self.qkv_format, fused=True)
+                k = apply_rotary_pos_emb(k, rope_emb, tensor_format=self.qkv_format, fused=True)
+                if self.use_wan_fp32_strategy:
+                    q = q.to(original_dtype)
+                    k = k.to(original_dtype)
+            return q, k, v
+
+        q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
+
+        return q, k, v
+
+    def compute_attention(
+        self,
+        q,
+        k,
+        v,
+        video_size: Optional[VideoSize] = None,
+        kv_cache_cfg: Optional[KVCacheConfig] = None,
+        teacher_forcing_layout: Optional[tuple[int, int]] = None,
+    ):
+        additional_args = {}
+        if isinstance(self.attn_op, (NattenA2AAttnOp, NeighborhoodAttention)) or self.backend == "i4":
+            additional_args["video_size"] = video_size
+        if isinstance(self.attn_op, AttentionOpWithKVCache):
+            additional_args["kv_cache_cfg"] = kv_cache_cfg
+            additional_args["teacher_forcing_layout"] = teacher_forcing_layout
+
+        result = self.attn_op(q, k, v, **additional_args)  # [B, S, H, D]
+        return self.output_dropout(self.output_proj(result))
+
+    def forward(
+        self,
+        x,
+        context: Optional[torch.Tensor] = None,
+        rope_emb: Optional[torch.Tensor] = None,
+        video_size: Optional[VideoSize] = None,
+        kv_cache_cfg: Optional[KVCacheConfig] = None,
+        teacher_forcing_layout: Optional[tuple[int, int]] = None,
+    ):
+        """
+        Args:
+            x (Tensor): The query tensor of shape [B, Mq, K]
+            context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
+            rope_emb (Optional[Tensor]): RoPE embedding tensor, or no RoPE embeddings (i.e. in cross attention)
+            video_size(VideoSize): Shape [T, H, W]
+        """
+        q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb)
+        if teacher_forcing_layout is not None:
+            if video_size is None:
+                raise ValueError("video_size is required for packed teacher forcing")
+            context_frames, _ = teacher_forcing_layout
+            context_tokens = context_frames * video_size.H * video_size.W
+            k = torch.cat([k[:, :context_tokens].detach(), k[:, context_tokens:]], dim=1)
+            v = torch.cat([v[:, :context_tokens].detach(), v[:, context_tokens:]], dim=1)
+        result = self.compute_attention(
+            q,
+            k,
+            v,
+            video_size=video_size,
+            kv_cache_cfg=kv_cache_cfg,
+            teacher_forcing_layout=teacher_forcing_layout,
+        )
+        return result
+
+    def set_context_parallel_group(self, process_group, ranks, stream, cp_comm_type: str = "p2p"):
+        # self.attn_op.set_context_parallel_group(process_group, ranks, stream, cp_comm_type="a2a")
+        self.attn_op.set_context_parallel_group(process_group, ranks, stream, cp_comm_type=cp_comm_type)
+
+
+class I2VCrossAttention(Attention):
+    def __init__(self, *args, img_latent_dim: int = 1024, **kwargs):
+        super().__init__(*args, **kwargs)
+        inner_dim = self.head_dim * self.n_heads
+        self.k_img = nn.Linear(img_latent_dim, inner_dim, bias=False)
+        self.v_img = nn.Linear(img_latent_dim, inner_dim, bias=False)
+        self.k_img_norm = RMSNorm(self.head_dim, eps=1e-6)
+
+    def init_weights(self) -> None:
+        super().init_weights()
+        torch.nn.init.trunc_normal_(self.k_img.weight, std=1.0 / math.sqrt(self._inner_dim))
+        torch.nn.init.trunc_normal_(self.v_img.weight, std=1.0 / math.sqrt(self._inner_dim))
+        self.k_img_norm.reset_parameters()
+
+    def compute_qkv(
+        self, x, context, rope_emb=None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        text_context, img_context = context
+        q, k, v = super().compute_qkv(x, text_context, rope_emb)
+        k_img = self.k_img(img_context)
+        v_img = self.v_img(img_context)
+        # Rearrange k_img, v_img
+        k_img, v_img = map(
+            lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
+            (k_img, v_img),
+        )
+
+        return q, k, v, self.k_img_norm(k_img), v_img
+
+    def compute_attention(self, q, k, v, k_img, v_img):
+        result = self.attn_op(q, k, v)  # [B, S, H, D]
+        result_img = self.attn_op(q, k_img, v_img)
+        return self.output_dropout(self.output_proj(result + result_img))
+
+    def forward(
+        self,
+        x,
+        context=None,
+        rope_emb=None,
+    ):
+        q, k, v, k_img, v_img = self.compute_qkv(x, context, rope_emb)
+        return self.compute_attention(q, k, v, k_img, v_img)
+
+
+class VideoPositionEmb(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._cp_group = None
+
+    def enable_context_parallel(self, process_group: ProcessGroup):
+        self._cp_group = process_group
+
+    def disable_context_parallel(self):
+        self._cp_group = None
+
+    @property
+    def seq_dim(self):
+        return 1
+
+    def forward(self, x_B_T_H_W_C: torch.Tensor, fps=Optional[torch.Tensor]) -> torch.Tensor:
+        """
+        With CP, the function assume that the input tensor is already split.
+        It delegates the embedding generation to generate_embeddings function.
+        """
+        B_T_H_W_C = x_B_T_H_W_C.shape
+        if self._cp_group is not None:
+            cp_ranks = get_process_group_ranks(self._cp_group)
+            cp_size = len(cp_ranks)
+            cp_size_t = cp_size
+            if USE_MEGATRON and hasattr(parallel_state, "cp_size_t"):
+                # We saved cp_size_t in find_split function for combined temporal and spatial splitting.
+                # We need cp_size_t to find out the split values for T and H dimensions for correct embedding calculations.
+                cp_size_t = parallel_state.cp_size_t
+            cp_size_h = max(1, cp_size // cp_size_t)
+            B, T, H, W, C = B_T_H_W_C
+            B_T_H_W_C = (B, T * cp_size_t, H * cp_size_h, W, C)
+        embeddings = self.generate_embeddings(B_T_H_W_C, fps=fps)
+
+        return self._split_for_context_parallel(embeddings)
+
+    def generate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor]):
+        raise NotImplementedError
+
+    def _split_for_context_parallel(self, embeddings):
+        if self._cp_group is not None:
+            embeddings = split_inputs_cp(x=embeddings, seq_dim=self.seq_dim, cp_group=self._cp_group)
+        return embeddings
+
+
+class VideoRopePosition3DEmb(VideoPositionEmb):
+    def __init__(
+        self,
+        *,  # enforce keyword arguments
+        head_dim: int,
+        len_h: int,
+        len_w: int,
+        len_t: int,
+        base_fps: int = 24,
+        h_extrapolation_ratio: float = 1.0,
+        w_extrapolation_ratio: float = 1.0,
+        t_extrapolation_ratio: float = 1.0,
+        enable_fps_modulation: bool = True,
+        **kwargs,  # used for compatibility with other positional embeddings; unused in this class
+    ):
+        del kwargs
+        super().__init__()
+        self.register_buffer("seq", torch.arange(max(len_h, len_w, len_t), dtype=torch.float))
+        self.base_fps = base_fps
+        self.max_h = len_h
+        self.max_w = len_w
+        self.max_t = len_t
+        self.enable_fps_modulation = enable_fps_modulation
+        dim = head_dim
+        dim_h = dim // 6 * 2
+        dim_w = dim_h
+        dim_t = dim - 2 * dim_h
+        assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}"
+
+        self.register_buffer(
+            "dim_spatial_range",
+            torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h,
+            persistent=True,
+        )
+        self.register_buffer(
+            "dim_temporal_range",
+            torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t,
+            persistent=True,
+        )
+        self._dim_h = dim_h
+        self._dim_t = dim_t
+
+        self.h_ntk_factor = h_extrapolation_ratio ** (dim_h / (dim_h - 2))
+        self.w_ntk_factor = w_extrapolation_ratio ** (dim_w / (dim_w - 2))
+        self.t_ntk_factor = t_extrapolation_ratio ** (dim_t / (dim_t - 2))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        dim_h = self._dim_h
+        dim_t = self._dim_t
+
+        self.seq = torch.arange(max(self.max_h, self.max_w, self.max_t)).float().to(self.dim_spatial_range.device)
+        self.dim_spatial_range = (
+            torch.arange(0, dim_h, 2)[: (dim_h // 2)].float().to(self.dim_spatial_range.device) / dim_h
+        )
+        self.dim_temporal_range = (
+            torch.arange(0, dim_t, 2)[: (dim_t // 2)].float().to(self.dim_spatial_range.device) / dim_t
+        )
+
+    def generate_embeddings(
+        self,
+        B_T_H_W_C: torch.Size,
+        fps: Optional[torch.Tensor] = None,
+        h_ntk_factor: Optional[float] = None,
+        w_ntk_factor: Optional[float] = None,
+        t_ntk_factor: Optional[float] = None,
+    ):
+        """
+        Generate embeddings for the given input size.
+
+        Args:
+            B_T_H_W_C (torch.Size): Input tensor size (Batch, Time, Height, Width, Channels).
+            fps (Optional[torch.Tensor], optional): Frames per second. Defaults to None.
+            h_ntk_factor (Optional[float], optional): Height NTK factor. If None, uses self.h_ntk_factor.
+            w_ntk_factor (Optional[float], optional): Width NTK factor. If None, uses self.w_ntk_factor.
+            t_ntk_factor (Optional[float], optional): Time NTK factor. If None, uses self.t_ntk_factor.
+
+        Returns:
+            Not specified in the original code snippet.
+        """
+        h_ntk_factor = h_ntk_factor if h_ntk_factor is not None else self.h_ntk_factor
+        w_ntk_factor = w_ntk_factor if w_ntk_factor is not None else self.w_ntk_factor
+        t_ntk_factor = t_ntk_factor if t_ntk_factor is not None else self.t_ntk_factor
+
+        h_theta = 10000.0 * h_ntk_factor
+        w_theta = 10000.0 * w_ntk_factor
+        t_theta = 10000.0 * t_ntk_factor
+
+        h_spatial_freqs = 1.0 / (h_theta ** self.dim_spatial_range.float())
+        w_spatial_freqs = 1.0 / (w_theta ** self.dim_spatial_range.float())
+        temporal_freqs = 1.0 / (t_theta ** self.dim_temporal_range.float())
+
+        B, T, H, W, _ = B_T_H_W_C
+        assert H <= self.max_h and W <= self.max_w, (
+            f"Input dimensions (H={H}, W={W}) exceed the maximum dimensions (max_h={self.max_h}, max_w={self.max_w})"
+        )
+        half_emb_h = torch.outer(self.seq[:H], h_spatial_freqs)
+        half_emb_w = torch.outer(self.seq[:W], w_spatial_freqs)
+
+        if self.enable_fps_modulation:
+            uniform_fps = (fps is None) or (fps.min() == fps.max())
+            assert uniform_fps or B == 1 or T == 1, (
+                "For video batch, batch size should be 1 for non-uniform fps. For image batch, T should be 1"
+            )
+
+            # apply sequence scaling in temporal dimension
+            if fps is None:  # image case
+                assert T == 1, "T should be 1 for image batch."
+                half_emb_t = torch.outer(self.seq[:T], temporal_freqs)
+            else:
+                half_emb_t = torch.outer(self.seq[:T] / fps[:1] * self.base_fps, temporal_freqs)
+        else:
+            half_emb_t = torch.outer(self.seq[:T], temporal_freqs)
+
+        em_T_H_W_D = torch.cat(
+            [
+                repeat(half_emb_t, "t d -> t h w d", h=H, w=W),
+                repeat(half_emb_h, "h d -> t h w d", t=T, w=W),
+                repeat(half_emb_w, "w d -> t h w d", t=T, h=H),
+            ]
+            * 2,
+            dim=-1,
+        )
+
+        return rearrange(em_T_H_W_D, "t h w d -> (t h w) 1 1 d").float()
+
+    @property
+    def seq_dim(self):
+        return 0
+
+
+class LearnablePosEmbAxis(VideoPositionEmb):
+    def __init__(
+        self,
+        *,  # enforce keyword arguments
+        interpolation: str,
+        model_channels: int,
+        len_h: int,
+        len_w: int,
+        len_t: int,
+        **kwargs,
+    ):
+        """
+        Args:
+            interpolation (str): we curretly only support "crop", ideally when we need extrapolation capacity, we should adjust frequency or other more advanced methods. they are not implemented yet.
+        """
+        del kwargs  # unused
+        super().__init__()
+        self.interpolation = interpolation
+        assert self.interpolation in ["crop"], f"Unknown interpolation method {self.interpolation}"
+        self.model_channels = model_channels
+
+        self.pos_emb_h = nn.Parameter(torch.zeros(len_h, model_channels))
+        self.pos_emb_w = nn.Parameter(torch.zeros(len_w, model_channels))
+        self.pos_emb_t = nn.Parameter(torch.zeros(len_t, model_channels))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        std = 1.0 / math.sqrt(self.model_channels)
+        torch.nn.init.trunc_normal_(self.pos_emb_h, std=std, a=-3 * std, b=3 * std)
+        torch.nn.init.trunc_normal_(self.pos_emb_w, std=std, a=-3 * std, b=3 * std)
+        torch.nn.init.trunc_normal_(self.pos_emb_t, std=std, a=-3 * std, b=3 * std)
+
+    def generate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor]) -> torch.Tensor:
+        B, T, H, W, _ = B_T_H_W_C
+        if self.interpolation == "crop":
+            emb_h_H = self.pos_emb_h[:H]
+            emb_w_W = self.pos_emb_w[:W]
+            emb_t_T = self.pos_emb_t[:T]
+            emb = (
+                repeat(emb_t_T, "t d-> b t h w d", b=B, h=H, w=W)
+                + repeat(emb_h_H, "h d-> b t h w d", b=B, t=T, w=W)
+                + repeat(emb_w_W, "w d-> b t h w d", b=B, t=T, h=H)
+            )
+            assert list(emb.shape)[:4] == [B, T, H, W], f"bad shape: {list(emb.shape)[:4]} != {B, T, H, W}"
+        else:
+            raise ValueError(f"Unknown interpolation method {self.interpolation}")
+
+        norm = torch.linalg.vector_norm(emb, dim=-1, keepdim=True, dtype=torch.float32)
+        norm = torch.add(1e-6, norm, alpha=np.sqrt(norm.numel() / emb.numel()))
+        return emb / norm.to(emb.dtype)
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+
+class Timesteps(nn.Module):
+    def __init__(self, num_channels):
+        super().__init__()
+        self.num_channels = num_channels
+
+    def forward(self, timesteps_B_T):
+        assert timesteps_B_T.ndim == 2, f"Expected 2D input, got {timesteps_B_T.ndim}"
+        # wan need emb to be in fp32
+        in_dtype = timesteps_B_T.dtype
+        timesteps = timesteps_B_T.flatten().float()
+        half_dim = self.num_channels // 2
+        exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timesteps.device)
+        exponent = exponent / (half_dim - 0.0)
+
+        emb = torch.exp(exponent)
+        emb = timesteps[:, None].float() * emb[None, :]
+
+        sin_emb = torch.sin(emb)
+        cos_emb = torch.cos(emb)
+        emb = torch.cat([cos_emb, sin_emb], dim=-1)
+
+        return rearrange(emb.to(dtype=in_dtype), "(b t) d -> b t d", b=timesteps_B_T.shape[0], t=timesteps_B_T.shape[1])
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, in_features: int, out_features: int, use_adaln_lora: bool = False):
+        super().__init__()
+        log.debug(
+            f"Using AdaLN LoRA Flag:  {use_adaln_lora}. We enable bias if no AdaLN LoRA for backward compatibility."
+        )
+        self.in_dim = in_features
+        self.out_dim = out_features
+        self.linear_1 = nn.Linear(in_features, out_features, bias=not use_adaln_lora)
+        self.activation = nn.SiLU()
+        self.use_adaln_lora = use_adaln_lora
+        if use_adaln_lora:
+            self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False)
+        else:
+            self.linear_2 = nn.Linear(out_features, out_features, bias=False)
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        std = 1.0 / math.sqrt(self.in_dim)
+        torch.nn.init.trunc_normal_(self.linear_1.weight, std=std, a=-3 * std, b=3 * std)
+
+        std = 1.0 / math.sqrt(self.out_dim)
+        torch.nn.init.trunc_normal_(self.linear_2.weight, std=std, a=-3 * std, b=3 * std)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        # Inference schedulers produce fp32 timesteps even when the released
+        # Cosmos checkpoint is bf16. Standard nn.Linear requires matching
+        # dtypes (Transformer Engine handled this cast in NVIDIA's runtime).
+        sample = sample.to(dtype=self.linear_1.weight.dtype)
+        emb = self.linear_1(sample)
+        emb = self.activation(emb)
+        emb = self.linear_2(emb)
+
+        if self.use_adaln_lora:
+            adaln_lora_B_T_3D = emb
+            emb_B_T_D = sample
+        else:
+            emb_B_T_D = emb
+            adaln_lora_B_T_3D = None
+
+        return emb_B_T_D, adaln_lora_B_T_3D
+
+
+class FourierFeatures(nn.Module):
+    """
+    Implements a layer that generates Fourier features from input tensors, based on randomly sampled
+    frequencies and phases. This can help in learning high-frequency functions in low-dimensional problems.
+
+    [B] -> [B, D]
+
+    Parameters:
+        num_channels (int): The number of Fourier features to generate.
+        bandwidth (float, optional): The scaling factor for the frequency of the Fourier features. Defaults to 1.
+        normalize (bool, optional): If set to True, the outputs are scaled by sqrt(2), usually to normalize
+                                    the variance of the features. Defaults to False.
+
+    Example:
+        >>> layer = FourierFeatures(num_channels=256, bandwidth=0.5, normalize=True)
+        >>> x = torch.randn(10, 256)  # Example input tensor
+        >>> output = layer(x)
+        >>> print(output.shape)  # Expected shape: (10, 256)
+    """
+
+    def __init__(self, num_channels, bandwidth=1, normalize=False):
+        super().__init__()
+        self.register_buffer("freqs", 2 * np.pi * bandwidth * torch.randn(num_channels), persistent=True)
+        self.register_buffer("phases", 2 * np.pi * torch.rand(num_channels), persistent=True)
+        self.gain = np.sqrt(2) if normalize else 1
+        self.bandwidth = bandwidth
+        self.num_channels = num_channels
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        generator = torch.Generator()
+        generator.manual_seed(0)
+        self.freqs = (
+            2 * np.pi * self.bandwidth * torch.randn(self.num_channels, generator=generator).to(self.freqs.device)
+        )
+        self.phases = 2 * np.pi * torch.rand(self.num_channels, generator=generator).to(self.freqs.device)
+
+    def forward(self, x, gain: float = 1.0):
+        """
+        Apply the Fourier feature transformation to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+            gain (float, optional): An additional gain factor applied during the forward pass. Defaults to 1.
+
+        Returns:
+            torch.Tensor: The transformed tensor, with Fourier features applied.
+        """
+        in_dtype = x.dtype
+        x = x.to(torch.float32).ger(self.freqs.to(torch.float32)).add(self.phases.to(torch.float32))
+        x = x.cos().mul(self.gain * gain).to(in_dtype)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    PatchEmbed is a module for embedding patches from an input tensor by applying either 3D or 2D convolutional layers,
+    depending on the . This module can process inputs with temporal (video) and spatial (image) dimensions,
+    making it suitable for video and image processing tasks. It supports dividing the input into patches
+    and embedding each patch into a vector of size `out_channels`.
+
+    Parameters:
+    - spatial_patch_size (int): The size of each spatial patch.
+    - temporal_patch_size (int): The size of each temporal patch.
+    - in_channels (int): Number of input channels. Default: 3.
+    - out_channels (int): The dimension of the embedding vector for each patch. Default: 768.
+    - bias (bool): If True, adds a learnable bias to the output of the convolutional layers. Default: True.
+    """
+
+    def __init__(
+        self,
+        spatial_patch_size,
+        temporal_patch_size,
+        in_channels=3,
+        out_channels=768,
+    ):
+        super().__init__()
+        self.spatial_patch_size = spatial_patch_size
+        self.temporal_patch_size = temporal_patch_size
+
+        self.proj = nn.Sequential(
+            Rearrange(
+                "b c (t r) (h m) (w n) -> b t h w (c r m n)",
+                r=temporal_patch_size,
+                m=spatial_patch_size,
+                n=spatial_patch_size,
+            ),
+            nn.Linear(
+                in_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size, out_channels, bias=False
+            ),
+        )
+        self.dim = in_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        std = 1.0 / math.sqrt(self.dim)
+        torch.nn.init.trunc_normal_(self.proj[1].weight, std=std, a=-3 * std, b=3 * std)
+
+    def forward(self, x):
+        """
+        Forward pass of the PatchEmbed module.
+
+        Parameters:
+        - x (torch.Tensor): The input tensor of shape (B, C, T, H, W) where
+            B is the batch size,
+            C is the number of channels,
+            T is the temporal dimension,
+            H is the height, and
+            W is the width of the input.
+
+        Returns:
+        - torch.Tensor: The embedded patches as a tensor, with shape b t h w c.
+        """
+        assert x.dim() == 5
+        _, _, T, H, W = x.shape
+        assert H % self.spatial_patch_size == 0 and W % self.spatial_patch_size == 0, (
+            f"H,W {(H, W)} should be divisible by spatial_patch_size {self.spatial_patch_size}"
+        )
+        assert T % self.temporal_patch_size == 0
+        x = self.proj(x)
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of video DiT.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        spatial_patch_size,
+        temporal_patch_size,
+        out_channels,
+        use_adaln_lora: bool = False,
+        adaln_lora_dim: int = 256,
+        use_wan_fp32_strategy: bool = False,
+    ):
+        super().__init__()
+        self.use_wan_fp32_strategy = use_wan_fp32_strategy
+        self.layer_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(
+            hidden_size, spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels, bias=False
+        )
+        self.hidden_size = hidden_size
+        self.n_adaln_chunks = 2
+        self.use_adaln_lora = use_adaln_lora
+        self.adaln_lora_dim = adaln_lora_dim
+        if use_adaln_lora:
+            self.adaln_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, adaln_lora_dim, bias=False),
+                nn.Linear(adaln_lora_dim, self.n_adaln_chunks * hidden_size, bias=False),
+            )
+        else:
+            self.adaln_modulation = nn.Sequential(
+                nn.SiLU(), nn.Linear(hidden_size, self.n_adaln_chunks * hidden_size, bias=False)
+            )
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        std = 1.0 / math.sqrt(self.hidden_size)
+        torch.nn.init.trunc_normal_(self.linear.weight, std=std, a=-3 * std, b=3 * std)
+        if self.use_adaln_lora:
+            torch.nn.init.trunc_normal_(self.adaln_modulation[1].weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.zeros_(self.adaln_modulation[2].weight)
+        else:
+            torch.nn.init.zeros_(self.adaln_modulation[1].weight)
+
+        self.layer_norm.reset_parameters()
+
+    def forward(
+        self,
+        # x_BT_HW_D,
+        x_B_T_H_W_D,
+        emb_B_T_D,
+        adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+    ):
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            if self.use_adaln_lora:
+                assert adaln_lora_B_T_3D is not None
+                shift_B_T_D, scale_B_T_D = (
+                    self.adaln_modulation(emb_B_T_D) + adaln_lora_B_T_3D[:, :, : 2 * self.hidden_size]
+                ).chunk(2, dim=-1)
+            else:
+                shift_B_T_D, scale_B_T_D = self.adaln_modulation(emb_B_T_D).chunk(2, dim=-1)
+
+            shift_B_T_1_1_D, scale_B_T_1_1_D = (
+                rearrange(shift_B_T_D, "b t d -> b t 1 1 d"),
+                rearrange(scale_B_T_D, "b t d -> b t 1 1 d"),
+            )
+
+            def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D):
+                return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D
+
+            x_B_T_H_W_D = _fn(x_B_T_H_W_D, self.layer_norm, scale_B_T_1_1_D, shift_B_T_1_1_D)
+            x_B_T_H_W_O = self.linear(
+                x_B_T_H_W_D
+            )  # O = spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels
+        return x_B_T_H_W_O
+
+
+class Block(nn.Module):
+    """
+    A transformer block that combines self-attention, cross-attention and MLP layers with AdaLN modulation.
+    Each component (self-attention, cross-attention, MLP) has its own layer normalization and AdaLN modulation.
+
+    Parameters:
+        x_dim (int): Dimension of input features
+        context_dim (int): Dimension of context features for cross-attention
+        num_heads (int): Number of attention heads
+        mlp_ratio (float): Multiplier for MLP hidden dimension. Default: 4.0
+        use_adaln_lora (bool): Whether to use AdaLN-LoRA modulation. Default: False
+        adaln_lora_dim (int): Hidden dimension for AdaLN-LoRA layers. Default: 256
+        use_wan_fp32_strategy (bool): Whether to use Wan's FP32 strategy. Default: False
+        If True, in Attention layer, if do self-attention, q and k will be forced to fp32 before rotary pos emb
+        also, in modulation computation, force entire computation in fp32
+
+    The block applies the following sequence:
+    1. Self-attention with AdaLN modulation
+    2. Cross-attention with AdaLN modulation
+    3. MLP with AdaLN modulation
+
+    Each component uses skip connections and layer normalization.
+    """
+
+    def __init__(
+        self,
+        x_dim: int,
+        context_dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        use_adaln_lora: bool = False,
+        adaln_lora_dim: int = 256,
+        backend: str = "transformer_engine",
+        image_context_dim: Optional[int] = None,
+        use_wan_fp32_strategy: bool = False,
+    ):
+        super().__init__()
+        self.x_dim = x_dim
+        self.layer_norm_self_attn = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6)
+        self.self_attn = Attention(
+            x_dim,
+            None,
+            num_heads,
+            x_dim // num_heads,
+            qkv_format="bshd",
+            backend=backend,
+            use_wan_fp32_strategy=use_wan_fp32_strategy,
+        )
+        self.layer_norm_cross_attn = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6)
+
+        if image_context_dim is None:
+            self.cross_attn = Attention(
+                x_dim, context_dim, num_heads, x_dim // num_heads, qkv_format="bshd", backend=backend
+            )
+        else:
+            self.cross_attn = I2VCrossAttention(
+                x_dim,
+                context_dim,
+                num_heads,
+                x_dim // num_heads,
+                img_latent_dim=image_context_dim,
+                qkv_format="bshd",
+                backend=backend,
+            )
+
+        self.layer_norm_mlp = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6)
+        self.mlp = GPT2FeedForward(x_dim, int(x_dim * mlp_ratio))
+
+        self.use_adaln_lora = use_adaln_lora
+        if self.use_adaln_lora:
+            self.adaln_modulation_self_attn = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(x_dim, adaln_lora_dim, bias=False),
+                nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False),
+            )
+            self.adaln_modulation_cross_attn = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(x_dim, adaln_lora_dim, bias=False),
+                nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False),
+            )
+            self.adaln_modulation_mlp = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(x_dim, adaln_lora_dim, bias=False),
+                nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False),
+            )
+        else:
+            self.adaln_modulation_self_attn = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False))
+            self.adaln_modulation_cross_attn = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False))
+            self.adaln_modulation_mlp = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False))
+
+        self.cp_size = None
+        self.use_wan_fp32_strategy = use_wan_fp32_strategy
+
+    def set_context_parallel_group(self, process_group, ranks, stream, cp_comm_type: str = "p2p"):
+        self.cp_size = None if ranks is None else len(ranks)
+        self.self_attn.set_context_parallel_group(
+            process_group=process_group,
+            ranks=ranks,
+            stream=stream,
+            cp_comm_type=cp_comm_type,
+        )
+
+    def reset_parameters(self) -> None:
+        self.layer_norm_self_attn.reset_parameters()
+        self.layer_norm_cross_attn.reset_parameters()
+        self.layer_norm_mlp.reset_parameters()
+
+        if self.use_adaln_lora:
+            std = 1.0 / math.sqrt(self.x_dim)
+            torch.nn.init.trunc_normal_(self.adaln_modulation_self_attn[1].weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.trunc_normal_(self.adaln_modulation_cross_attn[1].weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.trunc_normal_(self.adaln_modulation_mlp[1].weight, std=std, a=-3 * std, b=3 * std)
+            torch.nn.init.zeros_(self.adaln_modulation_self_attn[2].weight)
+            torch.nn.init.zeros_(self.adaln_modulation_cross_attn[2].weight)
+            torch.nn.init.zeros_(self.adaln_modulation_mlp[2].weight)
+        else:
+            torch.nn.init.zeros_(self.adaln_modulation_self_attn[1].weight)
+            torch.nn.init.zeros_(self.adaln_modulation_cross_attn[1].weight)
+            torch.nn.init.zeros_(self.adaln_modulation_mlp[1].weight)
+
+    def init_weights(self) -> None:
+        self.reset_parameters()
+        self.self_attn.init_weights()
+        self.cross_attn.init_weights()
+        self.mlp.init_weights()
+
+    def forward(
+        self,
+        x_B_T_H_W_D: torch.Tensor,
+        emb_B_T_D: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
+        adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+        extra_per_block_pos_emb: Optional[torch.Tensor] = None,
+        kv_cache_cfg: Optional[KVCacheConfig] = None,
+        camera_B_T_H_W_C: Optional[torch.Tensor] = None,
+        teacher_forcing_layout: Optional[tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        return self._forward_single(
+            x_B_T_H_W_D,
+            emb_B_T_D,
+            crossattn_emb,
+            rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+            adaln_lora_B_T_3D=adaln_lora_B_T_3D,
+            extra_per_block_pos_emb=extra_per_block_pos_emb,
+            kv_cache_cfg=kv_cache_cfg,
+            camera_B_T_H_W_C=camera_B_T_H_W_C,
+            teacher_forcing_layout=teacher_forcing_layout,
+        )
+
+    def _forward_single(
+        self,
+        x_B_T_H_W_D: torch.Tensor,
+        emb_B_T_D: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
+        adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
+        extra_per_block_pos_emb: Optional[torch.Tensor] = None,
+        kv_cache_cfg: Optional[KVCacheConfig] = None,
+        camera_B_T_H_W_C: Optional[torch.Tensor] = None,
+        teacher_forcing_layout: Optional[tuple[int, int]] = None,
+    ) -> torch.Tensor:
+        if extra_per_block_pos_emb is not None:
+            x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb
+
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            if self.use_adaln_lora:
+                shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = (
+                    self.adaln_modulation_self_attn(emb_B_T_D) + adaln_lora_B_T_3D
+                ).chunk(3, dim=-1)
+                shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = (
+                    self.adaln_modulation_cross_attn(emb_B_T_D) + adaln_lora_B_T_3D
+                ).chunk(3, dim=-1)
+                shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = (
+                    self.adaln_modulation_mlp(emb_B_T_D) + adaln_lora_B_T_3D
+                ).chunk(3, dim=-1)
+            else:
+                shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = self.adaln_modulation_self_attn(
+                    emb_B_T_D
+                ).chunk(3, dim=-1)
+                shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = (
+                    self.adaln_modulation_cross_attn(emb_B_T_D).chunk(3, dim=-1)
+                )
+                shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = self.adaln_modulation_mlp(emb_B_T_D).chunk(3, dim=-1)
+
+        # Reshape tensors from (B, T, D) to (B, T, 1, 1, D) for broadcasting
+        shift_self_attn_B_T_1_1_D = rearrange(shift_self_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        scale_self_attn_B_T_1_1_D = rearrange(scale_self_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        gate_self_attn_B_T_1_1_D = rearrange(gate_self_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+
+        shift_cross_attn_B_T_1_1_D = rearrange(shift_cross_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        scale_cross_attn_B_T_1_1_D = rearrange(scale_cross_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        gate_cross_attn_B_T_1_1_D = rearrange(gate_cross_attn_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+
+        shift_mlp_B_T_1_1_D = rearrange(shift_mlp_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        scale_mlp_B_T_1_1_D = rearrange(scale_mlp_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+        gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d").type_as(x_B_T_H_W_D)
+
+        B, T, H, W, D = x_B_T_H_W_D.shape
+
+        def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D):
+            return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D
+
+        normalized_x_B_T_H_W_D = _fn(
+            x_B_T_H_W_D,
+            self.layer_norm_self_attn,
+            scale_self_attn_B_T_1_1_D,
+            shift_self_attn_B_T_1_1_D,
+        )
+        if camera_B_T_H_W_C is not None:
+            if self.self_attn.cam_encoder is None:
+                raise RuntimeError(
+                    "Camera conditioning was provided to a block without a camera encoder"
+                )
+            if camera_B_T_H_W_C.shape[:4] != normalized_x_B_T_H_W_D.shape[:4]:
+                raise ValueError(
+                    "Camera and video token grids do not match: "
+                    f"{tuple(camera_B_T_H_W_C.shape[:4])} and "
+                    f"{tuple(normalized_x_B_T_H_W_D.shape[:4])}"
+                )
+            normalized_x_B_T_H_W_D = normalized_x_B_T_H_W_D + self.self_attn.cam_encoder(
+                camera_B_T_H_W_C.to(dtype=normalized_x_B_T_H_W_D.dtype)
+            )
+
+        video_size = VideoSize(T=T, H=H, W=W)
+
+        # (ahassani): Hack to correct `video_size` when CP is enabled.
+        # I really don't like this, but there doesn't seem to be any central
+        # piece of code that's responsible for handling CP/TP that also defines the
+        # layout of shardings. Other parts of the code (i.e. RoPE) seem to make this
+        # assumption that CP sharding is always done along T.
+        if self.cp_size is not None and self.cp_size > 1:
+            video_size = VideoSize(T=T * self.cp_size, H=H, W=W)
+
+        self_attn_result = self.self_attn(
+            rearrange(normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
+            None,
+            rope_emb=rope_emb_L_1_1_D,
+            video_size=video_size,
+            kv_cache_cfg=kv_cache_cfg,
+            teacher_forcing_layout=teacher_forcing_layout,
+        )
+        result_B_T_H_W_D = rearrange(
+            self_attn_result,
+            "b (t h w) d -> b t h w d",
+            t=T,
+            h=H,
+            w=W,
+        )
+        x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D * result_B_T_H_W_D
+
+        def _x_fn(
+            _x_B_T_H_W_D,
+            layer_norm_cross_attn,
+            _scale_cross_attn_B_T_1_1_D,
+            _shift_cross_attn_B_T_1_1_D,
+            _gate_cross_attn_B_T_1_1_D,
+        ):
+            _normalized_x_B_T_H_W_D = _fn(
+                _x_B_T_H_W_D, layer_norm_cross_attn, _scale_cross_attn_B_T_1_1_D, _shift_cross_attn_B_T_1_1_D
+            )
+            _result_B_T_H_W_D = rearrange(
+                self.cross_attn(
+                    rearrange(_normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
+                    crossattn_emb,
+                    rope_emb=rope_emb_L_1_1_D,
+                ),
+                "b (t h w) d -> b t h w d",
+                t=T,
+                h=H,
+                w=W,
+            )
+            # _x_B_T_H_W_D = _x_B_T_H_W_D + _gate_cross_attn_B_T_1_1_D * _result_B_T_H_W_D
+            return _result_B_T_H_W_D
+
+        result_B_T_H_W_D = _x_fn(
+            x_B_T_H_W_D,
+            self.layer_norm_cross_attn,
+            scale_cross_attn_B_T_1_1_D,
+            shift_cross_attn_B_T_1_1_D,
+            gate_cross_attn_B_T_1_1_D,
+        )
+        x_B_T_H_W_D = result_B_T_H_W_D * gate_cross_attn_B_T_1_1_D + x_B_T_H_W_D
+
+        normalized_x_B_T_H_W_D = _fn(
+            x_B_T_H_W_D,
+            self.layer_norm_mlp,
+            scale_mlp_B_T_1_1_D,
+            shift_mlp_B_T_1_1_D,
+        )
+        result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
+        x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp_B_T_1_1_D * result_B_T_H_W_D
+        return x_B_T_H_W_D
+
+
+class MiniTrainDIT(WeightTrainingStat):
+    """
+    A clean impl of DIT that can load and  reproduce the training results of the original DIT model in edify_video/v4~(cosmos 1)
+    A general implementation of adaln-modulated VIT-like~(DiT) transformer for video processing.
+
+    Args:
+        max_img_h (int): Maximum height of the input images.
+        max_img_w (int): Maximum width of the input images.
+        max_frames (int): Maximum number of frames in the video sequence.
+        in_channels (int): Number of input channels (e.g., RGB channels for color images).
+        out_channels (int): Number of output channels.
+        patch_spatial (int): Spatial resolution of patches for input processing.
+        patch_temporal (int): Temporal resolution of patches for input processing.
+        concat_padding_mask (bool): If True, includes a mask channel in the input to handle padding.
+        model_channels (int): Base number of channels used throughout the model.
+        num_blocks (int): Number of transformer blocks.
+        num_heads (int): Number of heads in the multi-head attention layers.
+        mlp_ratio (float): Expansion ratio for MLP blocks.
+        crossattn_emb_channels (int): Number of embedding channels for cross-attention.
+        extra_image_context_dim (int): Number of embedding channels for extra image context.
+        pos_emb_cls (str): Type of positional embeddings.
+        pos_emb_learnable (bool): Whether positional embeddings are learnable.
+        pos_emb_interpolation (str): Method for interpolating positional embeddings.
+        min_fps (int): Minimum frames per second.
+        max_fps (int): Maximum frames per second.
+        use_adaln_lora (bool): Whether to use AdaLN-LoRA.
+        adaln_lora_dim (int): Dimension for AdaLN-LoRA.
+        rope_h_extrapolation_ratio (float): Height extrapolation ratio for RoPE.
+        rope_w_extrapolation_ratio (float): Width extrapolation ratio for RoPE.
+        rope_t_extrapolation_ratio (float): Temporal extrapolation ratio for RoPE.
+        extra_per_block_abs_pos_emb (bool): Whether to use extra per-block absolute positional embeddings.
+        extra_h_extrapolation_ratio (float): Height extrapolation ratio for extra embeddings.
+        extra_w_extrapolation_ratio (float): Width extrapolation ratio for extra embeddings.
+        extra_t_extrapolation_ratio (float): Temporal extrapolation ratio for extra embeddings.
+        n_dense_blocks (`int`, *optional*, defaults to -1):
+            Number of blocks that will remain dense (not replaced with sparse attention)
+            If -1, no blocks are replaced with sparse attention
+            If 0, all blocks use sparse attention
+            Otherwise, n_dense_blocks blocks will remain dense, distributed evenly across the network
+        natten_parameters (`dict`, *optional*, defaults to None):
+            NATTEN (Sparse attention) parameter list.
+            The list length must be the same as the number of layers, with each list element
+            indicating NATTEN parameters for that layer. If None, NATTEN will not be used in that
+            layer and it would remain a full dense self attention. If not None, it must be a
+            dictionary/mapping with at least the following key:
+              - window_size: `tuple` of size 3 indicating neighborhood attention window size.
+                window size of -1 along any dimension means self attention.
+            Other optional parameters and their keys:
+              - stride: `tuple` of size 3 indicating neighborhood attention stride value.
+                stride = 1 is standard neighborhood attention, stride = window size means
+                blocked/window self attention (WSA) along that dimension. Any other values are
+                strided neighborhood attention. Refer to the GNA paper for more information.
+
+              - dilation: `tuple` of size 3 indicating neighborhood attention dilation value.
+                dilation = 1 is standard neighborhood attention. Refer to the DiNAT paper for more
+                information.
+
+              - is_causal: `tuple` of 3 booleans indicating whether causal masking is enabled for
+                any of the T, H, W dimensions.
+    """
+
+    def __init__(
+        self,
+        max_img_h: int,
+        max_img_w: int,
+        max_frames: int,
+        in_channels: int,
+        out_channels: int,
+        patch_spatial: int,
+        patch_temporal: int,
+        concat_padding_mask: bool = True,
+        # attention settings
+        model_channels: int = 768,
+        num_blocks: int = 10,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        atten_backend: str = "transformer_engine",
+        # cross attention settings
+        crossattn_emb_channels: int = 1024,
+        use_crossattn_projection: bool = False,
+        crossattn_proj_in_channels: int = 1024,
+        extra_image_context_dim: Optional[int] = None,
+        # positional embedding settings
+        pos_emb_cls: str = "sincos",
+        pos_emb_learnable: bool = False,
+        pos_emb_interpolation: str = "crop",
+        min_fps: int = 1,
+        max_fps: int = 30,
+        use_adaln_lora: bool = False,
+        adaln_lora_dim: int = 256,
+        rope_h_extrapolation_ratio: float = 1.0,
+        rope_w_extrapolation_ratio: float = 1.0,
+        rope_t_extrapolation_ratio: float = 1.0,
+        extra_per_block_abs_pos_emb: bool = False,
+        extra_h_extrapolation_ratio: float = 1.0,
+        extra_w_extrapolation_ratio: float = 1.0,
+        extra_t_extrapolation_ratio: float = 1.0,
+        rope_enable_fps_modulation: bool = True,
+        sac_config: SACConfig = SACConfig(),
+        n_dense_blocks: int = -1,
+        natten_parameters: Union[dict, list] = None,
+        # if True, will closely match wan's strategy to use fp32 in certain layers/operations
+        use_wan_fp32_strategy: bool = False,
+    ) -> None:
+        super().__init__()
+        self.max_img_h = max_img_h
+        self.max_img_w = max_img_w
+        self.max_frames = max_frames
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.patch_spatial = patch_spatial
+        self.patch_temporal = patch_temporal
+        self.num_heads = num_heads
+        self.num_blocks = num_blocks
+        self.model_channels = model_channels
+        self.concat_padding_mask = concat_padding_mask
+        self.atten_backend = atten_backend
+        # positional embedding settings
+        self.pos_emb_cls = pos_emb_cls
+        self.pos_emb_learnable = pos_emb_learnable
+        self.pos_emb_interpolation = pos_emb_interpolation
+        self.min_fps = min_fps
+        self.max_fps = max_fps
+        self.rope_h_extrapolation_ratio = rope_h_extrapolation_ratio
+        self.rope_w_extrapolation_ratio = rope_w_extrapolation_ratio
+        self.rope_t_extrapolation_ratio = rope_t_extrapolation_ratio
+        self.extra_per_block_abs_pos_emb = extra_per_block_abs_pos_emb
+        self.extra_h_extrapolation_ratio = extra_h_extrapolation_ratio
+        self.extra_w_extrapolation_ratio = extra_w_extrapolation_ratio
+        self.extra_t_extrapolation_ratio = extra_t_extrapolation_ratio
+        self.rope_enable_fps_modulation = rope_enable_fps_modulation
+        self.extra_image_context_dim = extra_image_context_dim
+        self.build_patch_embed()
+        self.build_pos_embed()
+        self.use_adaln_lora = use_adaln_lora
+        self.adaln_lora_dim = adaln_lora_dim
+        self.t_embedder = nn.Sequential(
+            Timesteps(model_channels),
+            TimestepEmbedding(model_channels, model_channels, use_adaln_lora=use_adaln_lora),
+        )
+        self.use_crossattn_projection = use_crossattn_projection
+        self.crossattn_proj_in_channels = crossattn_proj_in_channels
+        self.use_wan_fp32_strategy = use_wan_fp32_strategy
+
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    x_dim=model_channels,
+                    context_dim=crossattn_emb_channels,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    use_adaln_lora=use_adaln_lora,
+                    adaln_lora_dim=adaln_lora_dim,
+                    backend=atten_backend,
+                    image_context_dim=None if extra_image_context_dim is None else model_channels,
+                    use_wan_fp32_strategy=use_wan_fp32_strategy,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+        self.final_layer = FinalLayer(
+            hidden_size=self.model_channels,
+            spatial_patch_size=self.patch_spatial,
+            temporal_patch_size=self.patch_temporal,
+            out_channels=self.out_channels,
+            use_adaln_lora=self.use_adaln_lora,
+            adaln_lora_dim=self.adaln_lora_dim,
+            use_wan_fp32_strategy=self.use_wan_fp32_strategy,
+        )
+
+        self.t_embedding_norm = RMSNorm(model_channels, eps=1e-6)
+        if extra_image_context_dim is not None:
+            self.img_context_proj = nn.Sequential(
+                nn.Linear(
+                    extra_image_context_dim, model_channels, bias=True
+                ),  # help distinguish between image and video context
+                nn.GELU(),
+            )
+
+        if use_crossattn_projection:
+            self.crossattn_proj = nn.Sequential(
+                nn.Linear(crossattn_proj_in_channels, crossattn_emb_channels, bias=True),
+                nn.GELU(),
+            )
+
+        self.init_weights()
+        self.enable_selective_checkpoint(sac_config, self.blocks)
+
+        # Replace self-attention with sparse attention if specified
+        if n_dense_blocks != -1:
+            self = replace_selfattn_op_with_sparse_attn_op(self, n_dense_blocks, natten_parameters=natten_parameters)
+
+        self._is_context_parallel_enabled = False
+
+    def init_weights(self):
+        self.x_embedder.init_weights()
+        self.pos_embedder.reset_parameters()
+        if self.extra_per_block_abs_pos_emb:
+            self.extra_pos_embedder.reset_parameters()
+
+        self.t_embedder[1].init_weights()
+        for block in self.blocks:
+            block.init_weights()
+
+        self.final_layer.init_weights()
+        self.t_embedding_norm.reset_parameters()
+
+        if self.extra_image_context_dim is not None:
+            self.img_context_proj[0].reset_parameters()
+
+    def enable_camera_conditioning(
+        self,
+        camera_dim: int,
+        init_seed: int = 0,
+    ) -> None:
+        """Attach the checkpoint-compatible camera projection to each self-attention."""
+        if camera_dim <= 0:
+            raise ValueError("camera_dim must be positive")
+        first_weight = self.blocks[0].self_attn.q_proj.weight
+        if first_weight.is_meta:
+            raise RuntimeError("Load the base Cosmos weights before adding camera encoders")
+        generator = torch.Generator(device=first_weight.device)
+        generator.manual_seed(init_seed)
+        std = 1.0 / math.sqrt(self.model_channels)
+
+        for block in self.blocks:
+            if block.self_attn.cam_encoder is not None:
+                if block.self_attn.cam_encoder.in_features != camera_dim:
+                    raise ValueError(
+                        "Camera conditioning is already enabled with a different dimension"
+                    )
+                continue
+            reference = block.self_attn.q_proj.weight
+            camera_encoder = nn.Linear(
+                camera_dim,
+                self.model_channels,
+                bias=False,
+                device=reference.device,
+                dtype=reference.dtype,
+            )
+            initialized = torch.empty(
+                camera_encoder.weight.shape,
+                device=reference.device,
+                dtype=torch.float32,
+            )
+            torch.nn.init.trunc_normal_(
+                initialized,
+                std=std,
+                a=-3 * std,
+                b=3 * std,
+                generator=generator,
+            )
+            with torch.no_grad():
+                camera_encoder.weight.copy_(initialized)
+            block.self_attn.cam_encoder = camera_encoder
+        self.camera_condition_dim = camera_dim
+
+    def build_patch_embed(self):
+        (
+            concat_padding_mask,
+            in_channels,
+            patch_spatial,
+            patch_temporal,
+            model_channels,
+        ) = (
+            self.concat_padding_mask,
+            self.in_channels,
+            self.patch_spatial,
+            self.patch_temporal,
+            self.model_channels,
+        )
+        in_channels = in_channels + 1 if concat_padding_mask else in_channels
+        self.x_embedder = PatchEmbed(
+            spatial_patch_size=patch_spatial,
+            temporal_patch_size=patch_temporal,
+            in_channels=in_channels,
+            out_channels=model_channels,
+        )
+
+    def build_pos_embed(self):
+        if self.pos_emb_cls == "rope3d":
+            cls_type = VideoRopePosition3DEmb
+        else:
+            raise ValueError(f"Unknown pos_emb_cls {self.pos_emb_cls}")
+
+        log.debug(f"Building positional embedding with {self.pos_emb_cls} class, impl {cls_type}")
+        kwargs = dict(
+            model_channels=self.model_channels,
+            len_h=self.max_img_h // self.patch_spatial,
+            len_w=self.max_img_w // self.patch_spatial,
+            len_t=self.max_frames // self.patch_temporal,
+            max_fps=self.max_fps,
+            min_fps=self.min_fps,
+            is_learnable=self.pos_emb_learnable,
+            interpolation=self.pos_emb_interpolation,
+            head_dim=self.model_channels // self.num_heads,
+            h_extrapolation_ratio=self.rope_h_extrapolation_ratio,
+            w_extrapolation_ratio=self.rope_w_extrapolation_ratio,
+            t_extrapolation_ratio=self.rope_t_extrapolation_ratio,
+            enable_fps_modulation=self.rope_enable_fps_modulation,
+        )
+        self.pos_embedder = cls_type(
+            **kwargs,
+        )
+
+        if self.extra_per_block_abs_pos_emb:
+            kwargs["h_extrapolation_ratio"] = self.extra_h_extrapolation_ratio
+            kwargs["w_extrapolation_ratio"] = self.extra_w_extrapolation_ratio
+            kwargs["t_extrapolation_ratio"] = self.extra_t_extrapolation_ratio
+            self.extra_pos_embedder = LearnablePosEmbAxis(
+                **kwargs,
+            )
+
+    def prepare_embedded_sequence(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        Prepares an embedded sequence tensor by applying positional embeddings and handling padding masks.
+
+        Args:
+            x_B_C_T_H_W (torch.Tensor): video
+            fps (Optional[torch.Tensor]): Frames per second tensor to be used for positional embedding when required.
+                                    If None, a default value (`self.base_fps`) will be used.
+            padding_mask (Optional[torch.Tensor]): current it is not used
+
+        Returns:
+            Tuple[torch.Tensor, Optional[torch.Tensor]]:
+                - A tensor of shape (B, T, H, W, D) with the embedded sequence.
+                - An optional positional embedding tensor, returned only if the positional embedding class
+                (`self.pos_emb_cls`) includes 'rope'. Otherwise, None.
+
+        Notes:
+            - If `self.concat_padding_mask` is True, a padding mask channel is concatenated to the input tensor.
+            - The method of applying positional embeddings depends on the value of `self.pos_emb_cls`.
+            - If 'rope' is in `self.pos_emb_cls` (case insensitive), the positional embeddings are generated using
+                the `self.pos_embedder` with the shape [T, H, W].
+            - If "fps_aware" is in `self.pos_emb_cls`, the positional embeddings are generated using the
+            `self.pos_embedder` with the fps tensor.
+            - Otherwise, the positional embeddings are generated without considering fps.
+        """
+        if self.concat_padding_mask:
+            padding_mask = transforms.functional.resize(
+                padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+            )
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1
+            )
+        x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)
+
+        if self.extra_per_block_abs_pos_emb:
+            extra_pos_emb = self.extra_pos_embedder(x_B_T_H_W_D, fps=fps)
+        else:
+            extra_pos_emb = None
+
+        if "rope" in self.pos_emb_cls.lower():
+            return x_B_T_H_W_D, self.pos_embedder(x_B_T_H_W_D, fps=fps), extra_pos_emb
+        x_B_T_H_W_D = x_B_T_H_W_D + self.pos_embedder(x_B_T_H_W_D)  # [B, T, H, W, D]
+
+        return x_B_T_H_W_D, None, extra_pos_emb
+
+    def unpatchify(self, x_B_T_H_W_M):
+        x_B_C_Tt_Hp_Wp = rearrange(
+            x_B_T_H_W_M,
+            "B T H W (p1 p2 t C) -> B C (T t) (H p1) (W p2)",
+            p1=self.patch_spatial,
+            p2=self.patch_spatial,
+            t=self.patch_temporal,
+        )
+        return x_B_C_Tt_Hp_Wp
+
+    def forward(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        timesteps_B_T: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        data_type: Optional[DataType] = DataType.VIDEO,
+        intermediate_feature_ids: Optional[List[int]] = None,
+        img_context_emb: Optional[torch.Tensor] = None,
+        camera_condition_B_C_T_H_W: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor | List[torch.Tensor] | Tuple[torch.Tensor, List[torch.Tensor]]:
+        """
+        Args:
+            x: (B, C, T, H, W) tensor of spatial-temp inputs
+            timesteps: (B, ) tensor of timesteps
+            crossattn_emb: (B, N, D) tensor of cross-attention embeddings
+        """
+        assert isinstance(data_type, DataType), (
+            f"Expected DataType, got {type(data_type)}. We need discuss this flag later."
+        )
+        x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D = self.prepare_embedded_sequence(
+            x_B_C_T_H_W,
+            fps=fps,
+            padding_mask=padding_mask,
+        )
+        camera_B_T_H_W_C = None
+        if camera_condition_B_C_T_H_W is not None:
+            camera_B_T_H_W_C = camera_condition_B_C_T_H_W.permute(
+                0, 2, 3, 4, 1
+            ).contiguous()
+            if camera_B_T_H_W_C.shape[:4] != x_B_T_H_W_D.shape[:4]:
+                raise ValueError(
+                    "Camera conditioning does not match the embedded video grid: "
+                    f"{tuple(camera_B_T_H_W_C.shape)} versus {tuple(x_B_T_H_W_D.shape)}"
+                )
+
+        if self.use_crossattn_projection:
+            crossattn_emb = self.crossattn_proj(crossattn_emb)
+
+        if img_context_emb is not None:
+            assert self.extra_image_context_dim is not None, (
+                "extra_image_context_dim must be set if img_context_emb is provided"
+            )
+            img_context_emb = self.img_context_proj(img_context_emb)
+            context_input = (crossattn_emb, img_context_emb)
+        else:
+            context_input = crossattn_emb
+
+        with amp.autocast("cuda", enabled=self.use_wan_fp32_strategy, dtype=torch.float32):
+            if timesteps_B_T.ndim == 1:
+                timesteps_B_T = timesteps_B_T.unsqueeze(1)
+            t_embedding_B_T_D, adaln_lora_B_T_3D = self.t_embedder(timesteps_B_T)
+            t_embedding_B_T_D = self.t_embedding_norm(t_embedding_B_T_D)
+
+        # for logging purpose
+        affline_scale_log_info = {}
+        affline_scale_log_info["t_embedding_B_T_D"] = t_embedding_B_T_D.detach()
+        self.affline_scale_log_info = affline_scale_log_info
+        self.affline_emb = t_embedding_B_T_D
+        self.crossattn_emb = crossattn_emb
+
+        if extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D is not None:
+            assert x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape, (
+                f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
+            )
+
+        B, T, H, W, D = x_B_T_H_W_D.shape
+        # x_B_THW_D = rearrange(x_B_T_H_W_D, "b t h w d -> b (t h w) d")
+
+        intermediate_features_outputs = []
+        for i, block in enumerate(self.blocks):
+            x_B_T_H_W_D = block(
+                x_B_T_H_W_D,
+                t_embedding_B_T_D,
+                context_input,
+                rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+                adaln_lora_B_T_3D=adaln_lora_B_T_3D,
+                extra_per_block_pos_emb=extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
+                camera_B_T_H_W_C=camera_B_T_H_W_C,
+            )
+            if intermediate_feature_ids and i in intermediate_feature_ids:
+                x_reshaped_for_disc = rearrange(x_B_T_H_W_D, "b tp hp wp d -> b (tp hp wp) d")
+                intermediate_features_outputs.append(x_reshaped_for_disc)
+
+        # x_B_T_H_W_D = rearrange(x_B_THW_D, "b (t h w) d -> b t h w d", t=T, h=H, w=W)
+        # O = out_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size
+        x_B_T_H_W_O = self.final_layer(x_B_T_H_W_D, t_embedding_B_T_D, adaln_lora_B_T_3D=adaln_lora_B_T_3D)
+        x_B_C_Tt_Hp_Wp = self.unpatchify(x_B_T_H_W_O)
+        if intermediate_feature_ids:
+            if len(intermediate_features_outputs) != len(intermediate_feature_ids):
+                log.warning(
+                    f"Collected {len(intermediate_features_outputs)} intermediate features, "
+                    f"but expected {len(intermediate_feature_ids)}. "
+                    f"Requested IDs: {intermediate_feature_ids}"
+                )
+            return x_B_C_Tt_Hp_Wp, intermediate_features_outputs
+
+        return x_B_C_Tt_Hp_Wp
+
+    def enable_selective_checkpoint(self, sac_config: SACConfig, blocks: nn.ModuleList):
+        if sac_config.mode == CheckpointMode.NONE:
+            return self
+
+        log.info(
+            f"Enable selective checkpoint with {sac_config.mode}, for every {sac_config.every_n_blocks} blocks. Total blocks: {len(blocks)}"
+        )
+        _context_fn = sac_config.get_context_fn()
+        for block_id, block in blocks.named_children():
+            if int(block_id) % sac_config.every_n_blocks == 0:
+                log.info(f"Enable selective checkpoint for block {block_id}")
+                block = ptd_checkpoint_wrapper(
+                    block,
+                    context_fn=_context_fn,
+                    preserve_rng_state=False,
+                )
+                blocks.register_module(block_id, block)
+        self.register_module(
+            "final_layer",
+            ptd_checkpoint_wrapper(
+                self.final_layer,
+                context_fn=_context_fn,
+                preserve_rng_state=False,
+            ),
+        )
+        return self
+
+    def fully_shard(self, mesh, **fsdp_kwargs):
+        for i, block in enumerate(self.blocks):
+            reshard_after_forward = i < len(self.blocks) - 1
+            fully_shard(block, mesh=mesh, reshard_after_forward=reshard_after_forward, **fsdp_kwargs)
+
+        fully_shard(self.final_layer, mesh=mesh, reshard_after_forward=True, **fsdp_kwargs)
+        if self.extra_per_block_abs_pos_emb:
+            fully_shard(self.extra_pos_embedder, mesh=mesh, reshard_after_forward=True, **fsdp_kwargs)
+        fully_shard(self.t_embedder, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+        if self.extra_image_context_dim is not None:
+            fully_shard(self.img_context_proj, mesh=mesh, reshard_after_forward=False, **fsdp_kwargs)
+
+    def disable_context_parallel(self):
+        # pos_embedder
+        self.pos_embedder.disable_context_parallel()
+        if self.extra_per_block_abs_pos_emb:
+            self.extra_pos_embedder.disable_context_parallel()
+
+        # attention
+        for block in self.blocks:
+            block.set_context_parallel_group(
+                process_group=None,
+                ranks=None,
+                stream=torch.cuda.Stream(),
+            )
+
+        self._is_context_parallel_enabled = False
+
+    def enable_context_parallel(self, process_group: Optional[ProcessGroup] = None):
+        # pos_embedder
+        self.pos_embedder.enable_context_parallel(process_group=process_group)
+        if self.extra_per_block_abs_pos_emb:
+            self.extra_pos_embedder.enable_context_parallel(process_group=process_group)
+
+        # attention
+        cp_ranks = get_process_group_ranks(process_group)
+        for block in self.blocks:
+            block.set_context_parallel_group(
+                process_group=process_group,
+                ranks=cp_ranks,
+                stream=torch.cuda.Stream(),
+            )
+
+        self._is_context_parallel_enabled = True
+
+    @property
+    def is_context_parallel_enabled(self):
+        return self._is_context_parallel_enabled
+
+
+def replace_selfattn_op_with_sparse_attn_op(
+    model: MiniTrainDIT, n_dense_blocks: int = 0, natten_parameters: Union[dict, list] = None
+) -> MiniTrainDIT:
+    """
+    Replace the self-attention operator with a sparse self-attention operator.
+
+    Args:
+        model: MiniTrainDIT instance
+        n_dense_blocks: Number of blocks that will remain dense (not replaced with NeighborhoodAttention)
+            If 0, all blocks use NeighborhoodAttention.
+            If -1, return model directly without any modifications.
+            Otherwise, n_dense_blocks blocks will remain dense, distributed evenly across the network.
+
+    Returns:
+        Modified instance
+    """
+    # Special case: return model directly without modifications
+    if n_dense_blocks == -1:
+        return model
+
+    num_blocks = len(model.blocks)
+
+    if natten_parameters is None:
+        raise ValueError("Please specify natten_parameters when n_dense_blocks > -1.")
+
+    if isinstance(natten_parameters, Sequence) and len(natten_parameters) != num_blocks:
+        raise ValueError(
+            "List of NATTEN parameters must be the same length as the number of blocks, "
+            f"got {len(natten_parameters)=} != {num_blocks=}."
+        )
+
+    if isinstance(natten_parameters, Sequence) and n_dense_blocks > 0:
+        log.warning(f"NATTEN parameters was a list; ignoring {n_dense_blocks=}.")
+
+    if isinstance(natten_parameters, Sequence):
+        natten_parameters_list = natten_parameters
+    else:
+        if n_dense_blocks >= num_blocks:
+            raise ValueError(f"n_dense_blocks ({n_dense_blocks}) must be less than the number of blocks ({num_blocks})")
+
+        # Determine which blocks should remain dense
+        dense_indices = set()
+
+        if n_dense_blocks > 0:
+            # General rule: distribute n_dense_blocks blocks evenly across the network
+            if n_dense_blocks == 1:
+                # Special case: just the middle block
+                dense_indices.add(num_blocks // 2)
+            else:
+                # For multiple blocks, distribute them evenly from start to end
+                indices = np.linspace(0, num_blocks - 1, n_dense_blocks, dtype=int)
+                dense_indices.update(indices.tolist())
+
+        natten_parameters_list = [None if i in dense_indices else natten_parameters for i in range(num_blocks)]
+
+    # Replace self-attention with NeighborhoodAttention for non-dense blocks
+    for i, block in enumerate(model.blocks):
+        natten_params = natten_parameters_list[i]
+        if natten_params is not None:
+            natten_parameters_layer = {k: v for k, v in natten_params.items()}
+            natten_parameters_layer["layer_id"] = i
+            if block.self_attn.backend == "minimal_a2a":
+                sparse_attn_op = NattenA2AAttnOp(natten_parameters=natten_parameters_layer)
+            else:
+                raise NotImplementedError(
+                    f"Using sparsity with attention backend {block.self_attn.backend} is not supported."
+                )
+
+            block.self_attn.register_module("attn_op", sparse_attn_op)
+
+    return model

cosmos/model_weights_stats.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+
+@dataclass
+class TrainingStats:
+    """Data class to hold training statistics."""
+
+    video_samples: int = 0
+    image_samples: int = 0
+    iterations: int = 0
+    training_hours: float = 0.0
+
+
+class WeightTrainingStat(nn.Module, ABC):
+    """Abstract base class for tracking training statistics."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._initialize_tracking_buffers()
+
+    def _initialize_tracking_buffers(self) -> None:
+        """Initialize tracking buffers with default values."""
+        tracking_buffers = {
+            "accum_video_sample_counter": torch.tensor(0, dtype=torch.int64),
+            "accum_image_sample_counter": torch.tensor(0, dtype=torch.int64),
+            "accum_iteration": torch.tensor(0, dtype=torch.int64),
+            "accum_train_in_hours": torch.tensor(0.0, dtype=torch.float32),
+        }
+
+        for name, tensor in tracking_buffers.items():
+            self.register_buffer(name, tensor)
+
+    def get_training_stats(self) -> TrainingStats:
+        """Return current training statistics."""
+        return TrainingStats(
+            video_samples=self.accum_video_sample_counter.item(),
+            image_samples=self.accum_image_sample_counter.item(),
+            iterations=self.accum_iteration.item(),
+            training_hours=self.accum_train_in_hours.item(),
+        )
+
+    @abstractmethod
+    def forward(self, *args, **kwargs) -> Any:
+        pass

cosmos/runtime.py

@@ -0,0 +1,244 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Modifications Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+"""Small runtime helpers needed by the vendored Cosmos-Predict2.5 DiT.
+
+The training integration does not use NVIDIA's context-parallel or neighborhood
+attention paths.  Keeping these helpers local avoids importing the full
+``cosmos_predict2`` package (whose top-level import requires a CUDA extra).
+"""
+
+import logging
+from functools import lru_cache
+from enum import Enum
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+try:
+    from transformer_engine.pytorch.attention import (
+        DotProductAttention as _TransformerEngineAttention,
+    )
+    try:
+        from transformer_engine.pytorch.attention.rope import (
+            apply_rotary_pos_emb as _transformer_engine_rope,
+        )
+    except ImportError:
+        from transformer_engine.pytorch.attention import (
+            apply_rotary_pos_emb as _transformer_engine_rope,
+        )
+except ImportError:
+    _TransformerEngineAttention = None
+    _transformer_engine_rope = None
+
+
+log = logging.getLogger(__name__)
+
+
+class DataType(str, Enum):
+    IMAGE = "image"
+    VIDEO = "video"
+    MIX = "mix"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+class RMSNorm(nn.Module):
+    """Checkpoint-compatible RMSNorm backed by PyTorch's fused operator."""
+
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def reset_parameters(self) -> None:
+        nn.init.ones_(self.weight)
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(
+            value,
+            (value.shape[-1],),
+            self.weight,
+            self.eps,
+        )
+
+
+def apply_rotary_pos_emb(
+    value: torch.Tensor,
+    freqs: torch.Tensor,
+    *,
+    tensor_format: str = "bshd",
+    fused: bool = True,
+) -> torch.Tensor:
+    """Apply rotary embeddings with the fused CUDA implementation when present."""
+    if tensor_format != "bshd":
+        raise ValueError(f"Unsupported rotary tensor format: {tensor_format}")
+    if _transformer_engine_rope is not None and value.is_cuda:
+        return _transformer_engine_rope(
+            value.contiguous(),
+            freqs.contiguous(),
+            tensor_format=tensor_format,
+            fused=fused,
+        )
+    if freqs.ndim == 4 and freqs.shape[0] == value.shape[1]:
+        freqs = freqs.permute(1, 0, 2, 3)
+    first, second = value.chunk(2, dim=-1)
+    rotated = torch.cat((-second, first), dim=-1)
+    return value * freqs.cos() + rotated * freqs.sin()
+
+
+@lru_cache(maxsize=None)
+def _get_transformer_engine_attention(
+    num_heads: int,
+    head_dim: int,
+) -> nn.Module:
+    if _TransformerEngineAttention is None:
+        raise RuntimeError("Transformer Engine attention is unavailable")
+    module = _TransformerEngineAttention(
+        num_heads,
+        head_dim,
+        num_gqa_groups=num_heads,
+        attention_dropout=0.0,
+        qkv_format="bshd",
+        attn_mask_type="no_mask",
+    )
+    module.eval()
+    return module
+
+
+def attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    is_causal: bool = False,
+    **_kwargs,
+) -> torch.Tensor:
+    """Fused attention for Cosmos tensors shaped ``[B, S, H, D]``."""
+    if (
+        _TransformerEngineAttention is not None
+        and query.is_cuda
+        and not is_causal
+    ):
+        query = query.contiguous().clone()
+        key = key.contiguous().clone()
+        value = value.contiguous().clone()
+        fused_attention = _get_transformer_engine_attention(
+            int(query.shape[2]),
+            int(query.shape[3]),
+        )
+        output = fused_attention(
+            query,
+            key,
+            value,
+        )
+        if isinstance(output, tuple):
+            output = output[0]
+        if output.ndim == 3:
+            output = output.unflatten(-1, (query.shape[2], query.shape[3]))
+        return output
+
+    output = torch.nn.functional.scaled_dot_product_attention(
+        query.transpose(1, 2),
+        key.transpose(1, 2),
+        value.transpose(1, 2),
+        is_causal=is_causal,
+    )
+    return output.transpose(1, 2)
+
+
+class DotProductAttention(nn.Module):
+    """Parameter-free attention module with the expected Cosmos interface."""
+
+    def __init__(self, num_heads: int, head_dim: int, **kwargs) -> None:
+        del kwargs
+        super().__init__()
+        self.num_heads = int(num_heads)
+        self.head_dim = int(head_dim)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        del kwargs
+        if query.shape[2:] != (self.num_heads, self.head_dim):
+            raise ValueError(
+                "Attention input shape does not match the configured heads: "
+                f"{tuple(query.shape)}"
+            )
+        return attention(query, key, value).flatten(-2)
+
+    def set_context_parallel_group(self, *args, **kwargs) -> None:
+        del args, kwargs
+
+
+def split_inputs_cp(
+    x: torch.Tensor,
+    seq_dim: int,
+    cp_group: Optional[dist.ProcessGroup],
+) -> torch.Tensor:
+    """Split a tensor for the optional context-parallel model path."""
+    if cp_group is None or dist.get_world_size(cp_group) == 1:
+        return x
+    world_size = dist.get_world_size(cp_group)
+    if x.shape[seq_dim] % world_size:
+        raise ValueError("Context-parallel sequence length must divide world size")
+    rank = dist.get_rank(cp_group)
+    return x.chunk(world_size, dim=seq_dim)[rank].contiguous()
+
+
+class MinimalA2AAttnOp(nn.Module):
+    """Non-context-parallel equivalent of Predict2.5's minimal A2A op."""
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        return attention(query, key, value, **kwargs).flatten(-2)
+
+    def set_context_parallel_group(self, process_group, *args, **kwargs) -> None:
+        del args, kwargs
+        if process_group is not None and dist.get_world_size(process_group) > 1:
+            raise NotImplementedError(
+                "The local Cosmos integration does not enable context parallelism"
+            )
+
+
+class NeighborhoodAttention(nn.Module):
+    """Marker for the unsupported sparse-attention option."""
+
+    def __init__(self, *args, **kwargs) -> None:
+        del args, kwargs
+        super().__init__()
+        raise NotImplementedError(
+            "Neighborhood attention is not used by the Cosmos 2B configuration"
+        )
+
+
+class NattenA2AAttnOp(NeighborhoodAttention):
+    pass

cosmos/selective_activation_checkpoint.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+
+try:
+    from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts, noop_context_fn
+except ImportError:
+    CheckpointPolicy = None
+
+mm_only_save_list = {
+    torch.ops.aten.mm.default,
+    torch.ops.aten._scaled_dot_product_efficient_attention.default,
+    torch.ops.aten._scaled_dot_product_flash_attention.default,
+    torch.ops.aten.addmm.default,
+}
+
+
+class CheckpointMode(str, Enum):
+    """
+    Enum for the different checkpoint modes.
+    """
+
+    NONE = "none"
+    MM_ONLY = "mm_only"
+    BLOCK_WISE = "block_wise"
+
+    def __str__(self) -> str:
+        # Optional: makes print() show just the value
+        return self.value
+
+
+def mm_only_policy(ctx, func, *args, **kwargs):
+    """
+    In newer flash-attn and TE versions, FA2 shows up in the list of ops with the name of 'flash_attn._flash_attn_forward'.
+    However, FA2 is much slower (2-3x) than FA3 or cuDNN kernel. Registering cuDNN kernel would require heavy changes in TE code.
+    That's why the best option is to use FA3 with small modifications to flash_attn_interface.py to register FA3 as PyTorch op.
+    """
+    to_save = func in mm_only_save_list or "flash_attn" in str(func)
+    return CheckpointPolicy.MUST_SAVE if to_save else CheckpointPolicy.PREFER_RECOMPUTE
+
+
+def mm_only_context_fn():
+    return create_selective_checkpoint_contexts(mm_only_policy)
+
+
+@dataclass
+class SACConfig:
+    mode: str = "mm_only"
+    every_n_blocks: int = 1
+
+    def get_context_fn(self):
+        if self.mode == CheckpointMode.MM_ONLY:
+            return mm_only_context_fn
+        elif self.mode == CheckpointMode.BLOCK_WISE:
+            return noop_context_fn
+        else:
+            raise ValueError(f"Invalid mode: {self.mode}")

cosmos/wrapper.py

@@ -0,0 +1,680 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+"""Self-Forcing adapters for the official Cosmos-Predict2.5 implementation."""
+
+import types
+from typing import List, Optional
+
+import torch
+
+from cosmos.camera_conditioning import CAMERA_FEATURE_DIM
+from huggingface_hub import hf_hub_download
+
+from utils.scheduler import FlowMatchScheduler, SchedulerInterface
+from wan.modules.vae import _video_vae
+
+
+DEFAULT_MODEL_ID = "nvidia/Cosmos-Predict2.5-2B"
+DEFAULT_CHECKPOINT = (
+    "base/pre-trained/"
+    "d20b7120-df3e-4911-919d-db6e08bad31c_ema_bf16.pt"
+)
+DEFAULT_TEXT_ENCODER_ID = "nvidia/Cosmos-Reason1-7B"
+
+
+def _mean_normalize(tensor: torch.Tensor) -> torch.Tensor:
+    return (tensor - tensor.mean(dim=-1, keepdim=True)) / (
+        tensor.std(dim=-1, keepdim=True) + 1e-8
+    )
+
+
+class CosmosTextEncoder(torch.nn.Module):
+    """Cosmos-Reason1 text-only adapter producing Predict2.5 embeddings."""
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_TEXT_ENCODER_ID,
+        max_length: int = 512,
+    ) -> None:
+        super().__init__()
+        from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+        self.max_length = max_length
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+        ).eval()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.text_encoder.parameters()).device
+
+    def forward(self, text_prompts: List[str]) -> dict:
+        conversations = [
+            [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": (
+                                "You are a helpful assistant who will provide "
+                                "prompts to an image generator."
+                            ),
+                        }
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": prompt}],
+                },
+            ]
+            for prompt in text_prompts
+        ]
+        texts = [
+            self.processor.apply_chat_template(
+                conversation,
+                tokenize=False,
+                add_generation_prompt=False,
+                add_vision_id=False,
+            )
+            for conversation in conversations
+        ]
+        inputs = self.processor.tokenizer(
+            texts,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        ).to(self.device)
+        outputs = self.text_encoder(
+            input_ids=inputs.input_ids,
+            output_hidden_states=True,
+            use_cache=False,
+            return_dict=True,
+        )
+        # Predict2.5 concatenates the normalized outputs of all 28 language
+        # layers: 28 * 3584 = 100352 channels.
+        prompt_embeds = torch.cat(
+            [_mean_normalize(state) for state in outputs.hidden_states[1:]],
+            dim=-1,
+        )
+        return {"prompt_embeds": prompt_embeds}
+
+
+class CosmosVAEWrapper(torch.nn.Module):
+    """Wan2.1 VAE packaged with Cosmos-Predict2.5."""
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_MODEL_ID,
+        checkpoint_filename: str = "tokenizer.pth",
+    ) -> None:
+        super().__init__()
+        mean = [
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653,
+            -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632,
+            -0.1922, -0.9497, 0.2503, -0.2921,
+        ]
+        std = [
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708,
+            2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579,
+            1.6382, 1.1253, 2.8251, 1.9160,
+        ]
+        self.mean = torch.tensor(mean, dtype=torch.float32)
+        self.std = torch.tensor(std, dtype=torch.float32)
+        checkpoint_path = hf_hub_download(
+            repo_id=model_name,
+            filename=checkpoint_filename,
+        )
+        self.model = _video_vae(
+            pretrained_path=checkpoint_path,
+            z_dim=16,
+        ).eval().requires_grad_(False)
+
+    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
+        device, dtype = pixel.device, pixel.dtype
+        scale = [
+            self.mean.to(device=device, dtype=dtype),
+            1.0 / self.std.to(device=device, dtype=dtype),
+        ]
+        output = [
+            self.model.encode(sample.unsqueeze(0), scale).float().squeeze(0)
+            for sample in pixel
+        ]
+        return torch.stack(output, dim=0).permute(0, 2, 1, 3, 4)
+
+    def decode_to_pixel(
+        self,
+        latent: torch.Tensor,
+        use_cache: bool = False,
+    ) -> torch.Tensor:
+        zs = latent.permute(0, 2, 1, 3, 4)
+        if use_cache:
+            assert latent.shape[0] == 1, "Cached VAE decode requires batch size 1"
+        device, dtype = latent.device, latent.dtype
+        scale = [
+            self.mean.to(device=device, dtype=dtype),
+            1.0 / self.std.to(device=device, dtype=dtype),
+        ]
+        decode = self.model.cached_decode if use_cache else self.model.decode
+        output = [
+            decode(sample.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
+            for sample in zs
+        ]
+        return torch.stack(output, dim=0).permute(0, 2, 1, 3, 4)
+
+
+class CosmosDiffusionWrapper(torch.nn.Module):
+    """Match Cosmos-Predict2.5 to Self-Forcing's Wan wrapper contract."""
+
+    num_transformer_blocks = 28
+    frame_seq_length = 1  # Cosmos cache positions are latent-frame indices.
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_MODEL_ID,
+        checkpoint_filename: str = DEFAULT_CHECKPOINT,
+        timestep_shift: float = 5.0,
+        is_causal: bool = False,
+        local_attn_size: int = -1,
+        sink_size: int = 0,
+        i2v: bool = True,
+        camera_conditioning: bool = False,
+        camera_patch_size: int = 16,
+        camera_init_seed: int = 0,
+    ) -> None:
+        super().__init__()
+        self.is_causal = is_causal
+        self.i2v = i2v
+        self.local_attn_size = local_attn_size
+        self.camera_conditioning = bool(camera_conditioning)
+        self.camera_patch_size = int(camera_patch_size)
+        if self.camera_patch_size <= 0:
+            raise ValueError("camera_patch_size must be positive")
+        self.uniform_timestep = not is_causal
+        self._cache_max_frames = 128
+        self._gradient_checkpointing = False
+
+        self.model = self._load_model(
+            model_name=model_name,
+            checkpoint_filename=checkpoint_filename,
+            is_causal=is_causal,
+            local_attn_size=local_attn_size,
+            sink_size=sink_size,
+        ).eval()
+        if self.camera_conditioning:
+            self.model.enable_camera_conditioning(
+                camera_dim=CAMERA_FEATURE_DIM * self.camera_patch_size**2,
+                init_seed=int(camera_init_seed),
+            )
+
+        if is_causal:
+            self._kv_attention_ops = [
+                block.self_attn.attn_op for block in self.model.blocks
+            ]
+            for attention_op in self._kv_attention_ops:
+                attention_op.reset_kv_cache(
+                    max_cache_size=self._cache_max_frames,
+                )
+
+        self.scheduler = FlowMatchScheduler(
+            shift=timestep_shift,
+            sigma_min=0.0,
+            extra_one_step=True,
+        )
+        self.scheduler.set_timesteps(1000, training=True)
+        self.register_buffer(
+            "_flow_sigmas",
+            self.scheduler.sigmas.float(),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_flow_timesteps",
+            self.scheduler.timesteps.float(),
+            persistent=False,
+        )
+        self.post_init()
+
+    @staticmethod
+    def _model_kwargs(is_causal: bool) -> dict:
+        from cosmos.minimal_v4_dit import SACConfig
+
+        return {
+            "max_img_h": 240,
+            "max_img_w": 240,
+            "max_frames": 128,
+            "in_channels": 16,
+            "out_channels": 16,
+            "patch_spatial": 2,
+            "patch_temporal": 1,
+            "model_channels": 2048,
+            "num_blocks": 28,
+            "num_heads": 16,
+            "concat_padding_mask": True,
+            "pos_emb_cls": "rope3d",
+            "pos_emb_learnable": True,
+            "pos_emb_interpolation": "crop",
+            "use_adaln_lora": True,
+            "adaln_lora_dim": 256,
+            "extra_per_block_abs_pos_emb": False,
+            "rope_enable_fps_modulation": False,
+            "rope_h_extrapolation_ratio": 3.0,
+            "rope_w_extrapolation_ratio": 3.0,
+            "rope_t_extrapolation_ratio": 1.0,
+            "use_crossattn_projection": True,
+            "crossattn_proj_in_channels": 100352,
+            "crossattn_emb_channels": 1024,
+            "timestep_scale": 0.001,
+            "use_wan_fp32_strategy": True,
+            "atten_backend": "i4" if is_causal else "minimal_a2a",
+            "sac_config": SACConfig(mode="none"),
+        }
+
+    @classmethod
+    def _load_model(
+        cls,
+        model_name: str,
+        checkpoint_filename: str,
+        is_causal: bool,
+        local_attn_size: int,
+        sink_size: int,
+    ) -> torch.nn.Module:
+        if is_causal:
+            from cosmos.causal_model import CausalCosmosModel as Model
+        else:
+            from cosmos.minimal_v1_lvg_dit import (
+                MinimalV1LVGDiT as Model,
+            )
+
+        with torch.device("meta"):
+            model_kwargs = cls._model_kwargs(is_causal)
+            if is_causal:
+                model_kwargs.update(
+                    local_attn_size=local_attn_size,
+                    sink_size=sink_size,
+                )
+            model = Model(**model_kwargs)
+
+        checkpoint_path = hf_hub_download(
+            repo_id=model_name,
+            filename=checkpoint_filename,
+        )
+        state_dict = torch.load(
+            checkpoint_path,
+            map_location="cpu",
+            weights_only=True,
+            mmap=True,
+        )
+        if "model" in state_dict:
+            state_dict = state_dict["model"]
+        elif "state_dict" in state_dict:
+            state_dict = state_dict["state_dict"]
+
+        available_prefixes = (
+            "net_ema." if any(key.startswith("net_ema.") for key in state_dict)
+            else "net."
+        )
+        net_state_dict = {}
+        for key, value in state_dict.items():
+            if not key.startswith(available_prefixes):
+                continue
+            key = key.removeprefix(available_prefixes)
+            if not key.endswith("_extra_state"):
+                net_state_dict[key] = value
+
+        if not net_state_dict:
+            raise RuntimeError("Cosmos checkpoint contains no diffusion weights")
+
+        model.load_state_dict(net_state_dict, strict=False, assign=True)
+        unloaded = [
+            name
+            for name, tensor in (
+                list(model.named_parameters()) + list(model.named_buffers())
+            )
+            if tensor.is_meta
+        ]
+        if unloaded:
+            raise RuntimeError(
+                "Cosmos checkpoint did not initialize parameters: "
+                + ", ".join(unloaded[:10])
+            )
+        return model
+
+    def enable_gradient_checkpointing(self) -> None:
+        if self._gradient_checkpointing:
+            return
+        from cosmos.minimal_v4_dit import SACConfig
+
+        # FlexAttention is a higher-order op and PyTorch does not implement it
+        # for selective checkpointing's _CachingTorchDispatchMode. Ordinary
+        # block checkpointing preserves the memory saving without that mode.
+        self.model.enable_selective_checkpoint(
+            SACConfig(mode="block_wise"),
+            self.model.blocks,
+        )
+        self._gradient_checkpointing = True
+
+    def initialize_kv_cache(
+        self,
+        max_frames: int,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> object:
+        del batch_size, dtype, device
+        if not self.is_causal:
+            raise RuntimeError("KV cache is only available on the causal Cosmos model")
+        self._cache_max_frames = max_frames
+        for attention_op in self._kv_attention_ops:
+            attention_op.reset_kv_cache(max_cache_size=max_frames)
+        return self  # The actual cache is owned by each attention block.
+
+    def initialize_crossattn_cache(self, **kwargs) -> None:
+        del kwargs
+        return None
+
+    @staticmethod
+    def cache_position(frame_index: int) -> int:
+        return frame_index
+
+    def _condition_inputs(
+        self,
+        noisy_video: torch.Tensor,
+        conditional_dict: dict,
+        timestep: torch.Tensor,
+        apply_initial_condition: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch, frames, _, height, width = noisy_video.shape
+        mask = torch.zeros(
+            batch,
+            1,
+            frames,
+            height,
+            width,
+            device=noisy_video.device,
+            dtype=noisy_video.dtype,
+        )
+        initial_latent = conditional_dict.get("initial_latent")
+        if self.i2v and initial_latent is not None and apply_initial_condition:
+            cond_frames = min(initial_latent.shape[1], frames)
+            noisy_video = noisy_video.clone()
+            noisy_video[:, :cond_frames] = initial_latent[:, :cond_frames]
+            mask[:, :, :cond_frames] = 1
+            timestep = timestep.clone()
+            timestep[:, :cond_frames] = 0
+        return noisy_video, timestep, mask
+
+    def _camera_condition(
+        self,
+        conditional_dict: dict,
+        model_input: torch.Tensor,
+        *,
+        current_start: Optional[int],
+        streaming: bool,
+    ) -> Optional[torch.Tensor]:
+        camera = conditional_dict.get("camera_condition")
+        if not self.camera_conditioning:
+            if camera is not None:
+                raise ValueError(
+                    "camera_condition was provided, but camera_conditioning is disabled"
+                )
+            return None
+        if camera is None:
+            raise ValueError(
+                "camera_conditioning is enabled, but conditional_dict has no camera_condition"
+            )
+        if camera.ndim != 5:
+            raise ValueError(
+                "camera_condition must have shape [B, C, T, H, W]; got "
+                f"{tuple(camera.shape)}"
+            )
+
+        target_frames = model_input.shape[1]
+        start = int(current_start or 0) if streaming else 0
+        if camera.shape[2] >= start + target_frames:
+            camera = camera[:, :, start : start + target_frames]
+        elif camera.shape[2] != target_frames:
+            raise ValueError(
+                "Camera conditioning is too short for the requested video frames: "
+                f"start={start}, frames={target_frames}, camera_frames={camera.shape[2]}"
+            )
+
+        expected_spatial_grid = (
+            model_input.shape[-2] // self.model.patch_spatial,
+            model_input.shape[-1] // self.model.patch_spatial,
+        )
+        if camera.shape[-2:] != expected_spatial_grid:
+            raise ValueError(
+                "Camera and video spatial token grids do not match: camera="
+                f"{tuple(camera.shape[-2:])}, video="
+                f"{expected_spatial_grid}"
+            )
+        if camera.shape[0] != model_input.shape[0]:
+            if model_input.shape[0] % camera.shape[0]:
+                raise ValueError(
+                    "Camera batch cannot be expanded to the model batch: "
+                    f"{camera.shape[0]} and {model_input.shape[0]}"
+                )
+            camera = camera.repeat(
+                model_input.shape[0] // camera.shape[0], 1, 1, 1, 1
+            )
+        return camera
+
+    def forward(
+        self,
+        noisy_image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        timestep: torch.Tensor,
+        kv_cache: Optional[object] = None,
+        crossattn_cache: Optional[object] = None,
+        current_start: Optional[int] = None,
+        classify_mode: bool = False,
+        concat_time_embeddings: bool = False,
+        clean_x: Optional[torch.Tensor] = None,
+        aug_t: Optional[torch.Tensor] = None,
+        teacher_forcing_start: Optional[int] = None,
+        cache_start: Optional[int] = None,
+        store_kv: bool = False,
+    ) -> torch.Tensor:
+        del crossattn_cache, concat_time_embeddings, cache_start
+        if classify_mode:
+            raise NotImplementedError("Cosmos GAN classifier branch is not integrated")
+        teacher_forcing = clean_x is not None or aug_t is not None
+        if teacher_forcing and (clean_x is None or aug_t is None):
+            raise ValueError("Teacher forcing requires both clean_x and aug_t")
+        if teacher_forcing and teacher_forcing_start is None:
+            raise ValueError("Teacher forcing requires teacher_forcing_start")
+        if not teacher_forcing and teacher_forcing_start is not None:
+            raise ValueError("teacher_forcing_start requires teacher forcing")
+        if teacher_forcing and (not self.is_causal or kv_cache is not None):
+            raise ValueError(
+                "Teacher forcing requires a causal model without a streaming KV cache"
+            )
+
+        prompt_embeds = conditional_dict["prompt_embeds"]
+        if self.uniform_timestep:
+            timestep = timestep[:, :1].expand(-1, noisy_image_or_video.shape[1])
+        model_input, input_timestep, condition_mask = self._condition_inputs(
+            noisy_image_or_video,
+            conditional_dict,
+            timestep,
+            # Full-sequence training conditions frame zero. Streaming inference
+            # has already cached that frame, so later chunks must remain noisy.
+            apply_initial_condition=(
+                kv_cache is None or int(current_start or 0) == 0
+            ),
+        )
+        clean_model_input = None
+        clean_input_timestep = None
+        if teacher_forcing:
+            clean_model_input, clean_input_timestep, clean_condition_mask = (
+                self._condition_inputs(
+                    clean_x,
+                    conditional_dict,
+                    aug_t,
+                    apply_initial_condition=True,
+                )
+            )
+        camera_condition = self._camera_condition(
+            conditional_dict,
+            model_input,
+            current_start=current_start,
+            streaming=kv_cache is not None,
+        )
+        # Keep flow/noise/target construction in FP32, then cast only at the
+        # DiT boundary just like a mixed-precision root module would.
+        compute_dtype = self.model.x_embedder.proj[1].weight.dtype
+        model_input_bcthw = model_input.to(dtype=compute_dtype).permute(0, 2, 1, 3, 4)
+        clean_model_input_bcthw = (
+            clean_model_input.to(dtype=compute_dtype).permute(0, 2, 1, 3, 4)
+            if clean_model_input is not None
+            else None
+        )
+        prompt_embeds = prompt_embeds.to(dtype=compute_dtype)
+        condition_mask = condition_mask.to(dtype=compute_dtype)
+        if camera_condition is not None:
+            camera_condition = camera_condition.to(
+                device=model_input.device,
+                dtype=compute_dtype,
+            )
+        padding_mask = torch.zeros(
+            model_input.shape[0],
+            1,
+            model_input.shape[-2],
+            model_input.shape[-1],
+            device=model_input.device,
+            dtype=compute_dtype,
+        )
+
+        if kv_cache is not None:
+            from cosmos.kv_cache import (
+                KVCacheConfig,
+                VideoSeqPos,
+            )
+
+            frame_index = int(current_start or 0)
+            frame_count = model_input.shape[1]
+            token_h = model_input.shape[-2] // self.model.patch_spatial
+            token_w = model_input.shape[-1] // self.model.patch_spatial
+            video_pos = VideoSeqPos(
+                T=frame_count,
+                H=token_h,
+                W=token_w,
+            )
+            video_pos.pos_t = video_pos.pos_t + frame_index
+            should_store_kv = bool(store_kv)
+            block_frames = int(getattr(self.model, "num_frame_per_block", 1))
+            if (
+                frame_count > 1
+                and not should_store_kv
+                and frame_count != block_frames
+            ):
+                raise ValueError(
+                    "Multi-frame causal input must match num_frame_per_block: "
+                    f"frames={frame_count}, block={block_frames}"
+                )
+            flow_pred = self.model.forward_seq(
+                x_B_C_T_H_W=model_input_bcthw,
+                video_pos=video_pos,
+                timesteps_B_T=input_timestep,
+                crossattn_emb=prompt_embeds,
+                padding_mask=padding_mask,
+                condition_video_input_mask_B_C_T_H_W=condition_mask,
+                camera_condition_B_C_T_H_W=camera_condition,
+                full_video_size=(frame_index + frame_count, token_h, token_w),
+                kv_cache_cfg=KVCacheConfig(
+                    run_with_kv=True,
+                    store_kv=should_store_kv,
+                    current_idx=frame_index,
+                ),
+            ).permute(0, 2, 1, 3, 4)
+        elif teacher_forcing:
+            flow_pred = self.model.forward_teacher_forcing(
+                noisy_x_B_C_T_H_W=model_input_bcthw,
+                clean_x_B_C_T_H_W=clean_model_input_bcthw,
+                noisy_timesteps_B_T=input_timestep,
+                clean_timesteps_B_T=clean_input_timestep,
+                crossattn_emb=prompt_embeds,
+                padding_mask=padding_mask,
+                condition_video_input_mask_B_C_T_H_W=condition_mask,
+                clean_condition_video_input_mask_B_C_T_H_W=clean_condition_mask,
+                camera_condition_B_C_T_H_W=camera_condition,
+                noisy_start_frame=int(teacher_forcing_start),
+            ).permute(0, 2, 1, 3, 4)
+        else:
+            flow_pred = self.model(
+                x_B_C_T_H_W=model_input_bcthw,
+                timesteps_B_T=input_timestep,
+                crossattn_emb=prompt_embeds,
+                padding_mask=padding_mask,
+                condition_video_input_mask_B_C_T_H_W=condition_mask,
+                camera_condition_B_C_T_H_W=camera_condition,
+            ).permute(0, 2, 1, 3, 4)
+
+        pred_x0 = self._convert_flow_pred_to_x0(
+            flow_pred=flow_pred.flatten(0, 1),
+            xt=model_input.flatten(0, 1),
+            timestep=input_timestep.flatten(0, 1),
+        ).unflatten(0, flow_pred.shape[:2])
+        return flow_pred, pred_x0
+
+    def _convert_flow_pred_to_x0(
+        self,
+        flow_pred: torch.Tensor,
+        xt: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        original_dtype = flow_pred.dtype
+        flow_pred = flow_pred.float()
+        xt = xt.float()
+        timestep_id = torch.argmin(
+            (
+                self._flow_timesteps.unsqueeze(0)
+                - timestep.float().unsqueeze(1)
+            ).abs(),
+            dim=1,
+        )
+        sigma_t = self._flow_sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        return (xt - sigma_t * flow_pred).to(original_dtype)
+
+    @staticmethod
+    def _convert_x0_to_flow_pred(
+        scheduler,
+        x0_pred: torch.Tensor,
+        xt: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        original_dtype = x0_pred.dtype
+        x0_pred, xt, sigmas, timesteps = map(
+            lambda value: value.float().to(x0_pred.device),
+            [x0_pred, xt, scheduler.sigmas, scheduler.timesteps],
+        )
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(),
+            dim=1,
+        )
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        return ((xt - x0_pred) / sigma_t).to(original_dtype)
+
+    def get_scheduler(self) -> SchedulerInterface:
+        scheduler = self.scheduler
+        scheduler.convert_x0_to_noise = types.MethodType(
+            SchedulerInterface.convert_x0_to_noise,
+            scheduler,
+        )
+        scheduler.convert_noise_to_x0 = types.MethodType(
+            SchedulerInterface.convert_noise_to_x0,
+            scheduler,
+        )
+        scheduler.convert_velocity_to_x0 = types.MethodType(
+            SchedulerInterface.convert_velocity_to_x0,
+            scheduler,
+        )
+        return scheduler
+
+    def post_init(self) -> None:
+        self.get_scheduler()

examples/prompt.txt

@@ -0,0 +1 @@
+A grounded first-person game camera, showing occasional leather-gloved hands at the bottom edge, surfs smoothly along a slow, glassy water current flowing through a desert caravanserai moments before sunrise. In the near foreground, small submerged stones and surface wear on the aqueduct’s rim create rapid near-field parallax as the viewpoint glides along the shallow S-curve toward a towering central courtyard. High thin clouds diffuse the very first rays of dawn, casting broken sunlight that moves steadily across the sandstone paving. The surfing motion remains fluid and continuous, requiring only slight leans to navigate the gentle bends of the ancient water channel. As the camera rounds the initial curve, the surrounding masonry reveals fragments of a much older civilization, with colossal, weathered statues seamlessly incorporated into the newer brickwork. The subtly stylized realism highlights the believable material response of wet stone and dry terracotta under the shifting morning light. Passing a row of silent merchant stalls, the channel straightens out, and the gloved hands adjust their grip subtly to maintain balance. The destination courtyard grows larger, its grand archway framing the brightening sky. The gentle current carries the viewpoint perfectly through the final stretch of the corridor, easing to a quiet halt as the shallow water spills out into the expansive, sunlit central basin.

pipeline/__init__.py

@@ -0,0 +1,4 @@
+# Trimmed package initializer for the CMD Hugging Face Space.
+# The upstream nv-tlabs/cmd package initializers eagerly import training-only
+# and multi-GPU modules that this inference-only demo does not need. See NOTICE
+# and LICENSE for the upstream licenses that continue to apply.

pipeline/causal_inference.py

@@ -0,0 +1,483 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 The Self-Forcing Authors. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Modifications Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaOneWayNoncommercial
+
+from typing import List, Optional
+import torch
+
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from utils.model_factory import (
+    build_diffusion_wrapper,
+    build_text_encoder,
+    build_vae,
+)
+from cosmos.camera_conditioning import build_camera_conditioning
+
+
+class CausalInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = (
+            build_diffusion_wrapper(args, is_causal=True)
+            if generator is None else generator
+        )
+        self.text_encoder = build_text_encoder(args) if text_encoder is None else text_encoder
+        self.vae = build_vae(args) if vae is None else vae
+
+        # Step 2: Initialize all causal hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        self.num_inference_steps = int(getattr(args, "num_inference_steps", 0))
+        self.denoising_step_list = None
+        if self.num_inference_steps <= 0:
+            self.denoising_step_list = torch.tensor(
+                args.denoising_step_list, dtype=torch.long)
+            if args.warp_denoising_step:
+                timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
+                self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+        self.num_transformer_blocks = getattr(self.generator, "num_transformer_blocks", 30)
+        self.frame_seq_length = getattr(self.generator, "frame_seq_length", 1560)
+
+        self.kv_cache1 = None
+        self.args = args
+        self.context_noise = int(getattr(args, "context_noise", 0))
+        if self.context_noise < 0:
+            raise ValueError("context_noise must be non-negative")
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = getattr(self.generator.model, "local_attn_size", -1)
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1 and hasattr(self.generator.model, "num_frame_per_block"):
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        initial_latent: Optional[torch.Tensor] = None,
+        camera_poses: Optional[torch.Tensor] = None,
+        camera_intrinsics: Optional[torch.Tensor] = None,
+        camera_condition: Optional[torch.Tensor] = None,
+        return_latents: bool = False,
+        decode: bool = True,
+        profile: bool = False,
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+            initial_latent (torch.Tensor): The initial latent tensor of shape
+                (batch_size, num_input_frames, num_channels, height, width).
+                If num_input_frames is 1, perform image to video.
+                If num_input_frames is greater than 1, perform video extension.
+            return_latents (bool): Whether to return the latents.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+                It is normalized to be in the range [0, 1].
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        conditional_dict = self.text_encoder(
+            text_prompts=text_prompts
+        )
+        unconditional_dict = None
+        if self.num_inference_steps > 0 and self.args.guidance_scale > 1.0:
+            unconditional_dict = self.text_encoder(
+                text_prompts=[self.args.negative_prompt] * len(text_prompts)
+            )
+        if initial_latent is not None:
+            conditional_dict["initial_latent"] = initial_latent
+            if unconditional_dict is not None:
+                unconditional_dict["initial_latent"] = initial_latent
+        if getattr(self.args, "camera_conditioning", False):
+            if camera_condition is None:
+                if camera_poses is None or camera_intrinsics is None:
+                    raise ValueError(
+                        "Camera-conditioned inference requires camera_poses and "
+                        "camera_intrinsics, or a precomputed camera_condition"
+                    )
+                camera_condition = build_camera_conditioning(
+                    camera_poses.to(device=noise.device, dtype=torch.float32),
+                    camera_intrinsics.to(device=noise.device, dtype=torch.float32),
+                    image_height=int(self.args.height),
+                    image_width=int(self.args.width),
+                    frame_stride=int(getattr(self.args, "camera_frame_stride", 4)),
+                    patch_size=int(getattr(self.args, "camera_patch_size", 16)),
+                    num_frame_per_block=int(self.num_frame_per_block),
+                    expected_latent_frames=num_output_frames,
+                    output_dtype=noise.dtype,
+                )
+            conditional_dict["camera_condition"] = camera_condition
+            if unconditional_dict is not None:
+                unconditional_dict["camera_condition"] = camera_condition
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Set up profiling if requested
+        if profile:
+            init_start = torch.cuda.Event(enable_timing=True)
+            init_end = torch.cuda.Event(enable_timing=True)
+            diffusion_start = torch.cuda.Event(enable_timing=True)
+            diffusion_end = torch.cuda.Event(enable_timing=True)
+            vae_start = torch.cuda.Event(enable_timing=True)
+            vae_end = torch.cuda.Event(enable_timing=True)
+            block_times = []
+            block_start = torch.cuda.Event(enable_timing=True)
+            block_end = torch.cuda.Event(enable_timing=True)
+            init_start.record()
+
+        # Step 1: Initialize KV cache to all zeros
+        if hasattr(self.generator, "initialize_kv_cache"):
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+                max_frames=num_output_frames,
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+            )
+        elif self.kv_cache1 is None:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+        else:
+            # reset cross attn cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index]["is_init"] = False
+            # reset kv cache
+            for block_index in range(len(self.kv_cache1)):
+                self.kv_cache1[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache1[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=self._cache_position(current_start_frame),
+                    store_kv=True,
+                )
+                current_start_frame += 1
+            else:
+                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for _ in range(num_input_blocks):
+                current_ref_latents = \
+                    initial_latent[:, current_start_frame:current_start_frame + self.num_frame_per_block]
+                output[:, current_start_frame:current_start_frame + self.num_frame_per_block] = current_ref_latents
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=self._cache_position(current_start_frame),
+                    store_kv=True,
+                )
+                current_start_frame += self.num_frame_per_block
+
+        if profile:
+            init_end.record()
+            torch.cuda.synchronize()
+            diffusion_start.record()
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        for current_num_frames in all_num_frames:
+            if profile:
+                block_start.record()
+
+            noisy_input = noise[
+                :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # Step 3.1: Integrate the flow trajectory for a regular FM model.
+            # The legacy x0/re-noise path below is retained only for distilled
+            # few-step configs that explicitly provide denoising_step_list.
+            if self.num_inference_steps > 0:
+                sample_scheduler = FlowUniPCMultistepScheduler(
+                    num_train_timesteps=self.args.num_train_timestep,
+                    shift=1,
+                    use_dynamic_shifting=False,
+                )
+                sample_scheduler.set_timesteps(
+                    self.num_inference_steps,
+                    device=noise.device,
+                    shift=self.args.timestep_shift,
+                )
+                denoised_pred = noisy_input
+                for current_timestep in sample_scheduler.timesteps:
+                    if profile:
+                        print(f"current_timestep: {current_timestep}")
+                    timestep = torch.ones(
+                        [batch_size, current_num_frames],
+                        device=noise.device,
+                        dtype=current_timestep.dtype,
+                    ) * current_timestep
+                    if unconditional_dict is not None:
+                        cfg_dict = dict(conditional_dict)
+                        cfg_dict["prompt_embeds"] = torch.cat(
+                            [
+                                conditional_dict["prompt_embeds"],
+                                unconditional_dict["prompt_embeds"],
+                            ],
+                            dim=0,
+                        )
+                        flow_pred_cfg, _ = self.generator(
+                            noisy_image_or_video=torch.cat(
+                                [denoised_pred, denoised_pred], dim=0
+                            ),
+                            conditional_dict=cfg_dict,
+                            timestep=torch.cat([timestep, timestep], dim=0),
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=self._cache_position(current_start_frame),
+                        )
+                        flow_pred_cond, flow_pred_uncond = flow_pred_cfg.chunk(2)
+                        # Cosmos Predict2.5 V2W guidance is applied in velocity
+                        # space around the conditional prediction.
+                        flow_pred = flow_pred_cond + self.args.guidance_scale * (
+                            flow_pred_cond - flow_pred_uncond
+                        )
+                    else:
+                        flow_pred, _ = self.generator(
+                            noisy_image_or_video=denoised_pred,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=self._cache_position(current_start_frame),
+                        )
+                    denoised_pred = sample_scheduler.step(
+                        flow_pred,
+                        current_timestep,
+                        denoised_pred,
+                        return_dict=False,
+                    )[0]
+            else:
+                for index, current_timestep in enumerate(self.denoising_step_list):
+                    if profile:
+                        print(f"current_timestep: {current_timestep}")
+                    timestep = torch.ones(
+                        [batch_size, current_num_frames],
+                        device=noise.device,
+                        dtype=torch.int64) * current_timestep
+
+                    if index < len(self.denoising_step_list) - 1:
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=self._cache_position(current_start_frame)
+                        )
+                        next_timestep = self.denoising_step_list[index + 1]
+                        noisy_input = self.scheduler.add_noise(
+                            denoised_pred.flatten(0, 1),
+                            torch.randn_like(denoised_pred.flatten(0, 1)),
+                            next_timestep * torch.ones(
+                                [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                        ).unflatten(0, denoised_pred.shape[:2])
+                    else:
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=self._cache_position(current_start_frame)
+                        )
+
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Step 3.3: commit the generated frame to K/V at the context-noise
+            # level used by the causal training/inference recipe. The original
+            # I2V prefix remains clean at timestep zero.
+            context_timestep = torch.ones_like(timestep) * self.context_noise
+            cache_input = denoised_pred.detach()
+            if self.context_noise > 0:
+                cache_input = self.scheduler.add_noise(
+                    cache_input.flatten(0, 1),
+                    torch.randn_like(cache_input.flatten(0, 1)),
+                    context_timestep.flatten(0, 1),
+                ).unflatten(0, cache_input.shape[:2])
+            self.generator(
+                noisy_image_or_video=cache_input,
+                conditional_dict=conditional_dict,
+                timestep=context_timestep,
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=self._cache_position(current_start_frame),
+                store_kv=True,
+            )
+
+            if profile:
+                block_end.record()
+                torch.cuda.synchronize()
+                block_time = block_start.elapsed_time(block_end)
+                block_times.append(block_time)
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        if profile:
+            # End diffusion timing and synchronize CUDA
+            diffusion_end.record()
+            torch.cuda.synchronize()
+            diffusion_time = diffusion_start.elapsed_time(diffusion_end)
+            init_time = init_start.elapsed_time(init_end)
+            vae_start.record()
+
+        # Step 4: Decode the output when pixels are requested.
+        video = None
+        if decode:
+            video = self.vae.decode_to_pixel(output, use_cache=False)
+            video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        if profile:
+            # End VAE timing and synchronize CUDA
+            vae_end.record()
+            torch.cuda.synchronize()
+            vae_time = vae_start.elapsed_time(vae_end) if decode else 0.0
+            total_time = init_time + diffusion_time + vae_time
+
+            print("Profiling results:")
+            print(f"  - Initialization/caching time: {init_time:.2f} ms ({100 * init_time / total_time:.2f}%)")
+            print(f"  - Diffusion generation time: {diffusion_time:.2f} ms ({100 * diffusion_time / total_time:.2f}%)")
+            for i, block_time in enumerate(block_times):
+                print(f"    - Block {i} generation time: {block_time:.2f} ms ({100 * block_time / diffusion_time:.2f}% of diffusion)")
+            print(f"  - VAE decoding time: {vae_time:.2f} ms ({100 * vae_time / total_time:.2f}%)")
+            print(f"  - Total time: {total_time:.2f} ms")
+
+        if return_latents:
+            return video, output
+        if not decode:
+            return output
+        else:
+            return video
+
+    def _initialize_kv_cache(self, batch_size, dtype, device, max_frames=None):
+        """
+        Initialize the backend's causal KV cache.
+        """
+        if hasattr(self.generator, "initialize_kv_cache"):
+            self.kv_cache1 = self.generator.initialize_kv_cache(
+                max_frames=max_frames or getattr(self.args, "num_training_frames", 21),
+                batch_size=batch_size,
+                dtype=dtype,
+                device=device,
+            )
+            return
+
+        kv_cache1 = []
+        if self.local_attn_size != -1:
+            # Use the local attention size to compute the KV cache size
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            # Use the default KV cache size
+            kv_cache_size = 32760
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize the backend's cross-attention cache when it has one.
+        """
+        if hasattr(self.generator, "initialize_crossattn_cache"):
+            self.crossattn_cache = self.generator.initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=dtype,
+                device=device,
+            )
+            return
+
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                "is_init": False
+            })
+        self.crossattn_cache = crossattn_cache
+
+    def _cache_position(self, frame_index: int) -> int:
+        if hasattr(self.generator, "cache_position"):
+            return self.generator.cache_position(frame_index)
+        return frame_index * self.frame_seq_length

requirements.txt

@@ -0,0 +1,12 @@
+torchvision
+transformers>=4.57,<5
+accelerate
+diffusers>=0.35.2
+safetensors
+sentencepiece
+omegaconf
+einops
+numpy
+pillow
+imageio
+imageio-ffmpeg

utils/__init__.py

@@ -0,0 +1,4 @@
+# Trimmed package initializer for the CMD Hugging Face Space.
+# The upstream nv-tlabs/cmd package initializers eagerly import training-only
+# and multi-GPU modules that this inference-only demo does not need. See NOTICE
+# and LICENSE for the upstream licenses that continue to apply.

utils/misc.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 The Self-Forcing Authors. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import random
+import torch
+
+
+def set_seed(seed: int, deterministic: bool = False):
+    """
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+
+    Args:
+        seed (`int`):
+            The seed to set.
+        deterministic (`bool`, *optional*, defaults to `False`):
+            Whether to use deterministic algorithms where available. Can slow down training.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.use_deterministic_algorithms(True)
+
+
+def merge_dict_list(dict_list):
+    if len(dict_list) == 1:
+        return dict_list[0]
+
+    merged_dict = {}
+    for k, v in dict_list[0].items():
+        if isinstance(v, torch.Tensor):
+            if v.ndim == 0:
+                merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0)
+            else:
+                merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0)
+        else:
+            # for non-tensor values, we just copy the value from the first item
+            merged_dict[k] = v
+    return merged_dict

utils/scheduler.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 The Self-Forcing Authors. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import abstractmethod, ABC
+import torch
+
+
+class SchedulerInterface(ABC):
+    """
+    Base class for diffusion noise schedule.
+    """
+    alphas_cumprod: torch.Tensor  # [T], alphas for defining the noise schedule
+
+    @abstractmethod
+    def add_noise(
+        self, clean_latent: torch.Tensor,
+        noise: torch.Tensor, timestep: torch.Tensor
+    ):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        pass
+
+    def convert_x0_to_noise(
+        self, x0: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's x0 prediction to noise predidction.
+        x0: the predicted clean data with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = x0.dtype
+        x0, xt, alphas_cumprod = map(
+            lambda x: x.double().to(x0.device), [x0, xt,
+                                                 self.alphas_cumprod]
+        )
+
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        noise_pred = (xt - alpha_prod_t **
+                      (0.5) * x0) / beta_prod_t ** (0.5)
+        return noise_pred.to(original_dtype)
+
+    def convert_noise_to_x0(
+        self, noise: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's noise prediction to x0 predidction.
+        noise: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = noise.dtype
+        noise, xt, alphas_cumprod = map(
+            lambda x: x.double().to(noise.device), [noise, xt,
+                                                    self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (xt - beta_prod_t **
+                   (0.5) * noise) / alpha_prod_t ** (0.5)
+        return x0_pred.to(original_dtype)
+
+    def convert_velocity_to_x0(
+        self, velocity: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's velocity prediction to x0 predidction.
+        velocity: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        v = sqrt(alpha_t) * noise - sqrt(beta_t) x0
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t)
+        given v, x_t, we have
+        x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v
+        see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac56
+        """
+        # use higher precision for calculations
+        original_dtype = velocity.dtype
+        velocity, xt, alphas_cumprod = map(
+            lambda x: x.double().to(velocity.device), [velocity, xt,
+                                                       self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity
+        return x0_pred.to(original_dtype)
+
+
+class FlowMatchScheduler():
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
+                          num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (
+                self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B*T, C, H, W]
+            - noise: the noise with shape [B*T, C, H, W]
+            - timestep: the timestep with shape [B*T]
+        Output: the corrupted latent with shape [B*T, C, H, W]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        """
+        Input:
+            - timestep: the timestep with shape [B*T]
+        Output: the corresponding weighting [B*T]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights

wan/__init__.py

@@ -0,0 +1,4 @@
+# Trimmed package initializer for the CMD Hugging Face Space.
+# The upstream nv-tlabs/cmd package initializers eagerly import training-only
+# and multi-GPU modules that this inference-only demo does not need. See NOTICE
+# and LICENSE for the upstream licenses that continue to apply.

wan/modules/__init__.py

@@ -0,0 +1,4 @@
+# Trimmed package initializer for the CMD Hugging Face Space.
+# The upstream nv-tlabs/cmd package initializers eagerly import training-only
+# and multi-GPU modules that this inference-only demo does not need. See NOTICE
+# and LICENSE for the upstream licenses that continue to apply.

wan/modules/vae.py

@@ -0,0 +1,697 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+    'WanVAE',
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (self.padding[2], self.padding[2], self.padding[1],
+                         self.padding[1], 2 * self.padding[0], 0)
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
+
+    def forward(self, x):
+        return F.normalize(
+            x, dim=(1 if self.channel_first else
+                    -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+                        'downsample3d')
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == 'upsample2d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+        elif mode == 'upsample3d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+        elif mode == 'downsample2d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == 'downsample3d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == 'upsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = 'Rep'
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] != 'Rep':
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat([
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device), cache_x
+                        ],
+                            dim=2)
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] == 'Rep':
+                        cache_x = torch.cat([
+                            torch.zeros_like(cache_x).to(cache_x.device),
+                            cache_x
+                        ],
+                            dim=2)
+                    if feat_cache[idx] == 'Rep':
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.resample(x)
+        x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+        if self.mode == 'downsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
+                    #     # cache last frame of last two chunk
+                    #     cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False), nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1))
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+            if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                         -1).permute(0, 1, 3,
+                                                     2).contiguous().chunk(
+                                                         3, dim=-1)
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = 'downsample3d' if temperal_downsample[
+                    i] else 'downsample2d'
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout))
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        # downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_upsample=[False, True, True],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout))
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False), nn.SiLU(),
+            CausalConv3d(out_dim, 3, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        # conv1
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        # middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_downsample, dropout)
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_upsample, dropout)
+        self.clear_cache()
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        # cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        # 对encode输入的x，按时间拆分为1、4、4、4....
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx)
+            else:
+                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_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+            else:
+                out_ = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+        self.clear_cache()
+        return out
+
+    def cached_decode(self, z, scale):
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+            else:
+                out_ = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+        return out
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
+    """
+    Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+    """
+    # params
+    cfg = dict(
+        dim=96,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0)
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device('meta'):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    logging.info(f'loading {pretrained_path}')
+    model.load_state_dict(
+        torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class WanVAE:
+
+    def __init__(self,
+                 z_dim=16,
+                 vae_pth='cache/vae_step_411000.pth',
+                 dtype=torch.float,
+                 device="cuda"):
+        self.dtype = dtype
+        self.device = device
+
+        mean = [
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+        ]
+        std = [
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+        ]
+        self.mean = torch.tensor(mean, dtype=dtype, device=device)
+        self.std = torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path=vae_pth,
+            z_dim=z_dim,
+        ).eval().requires_grad_(False).to(device)
+
+    def encode(self, videos):
+        """
+        videos: A list of videos each with shape [C, T, H, W].
+        """
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
+                for u in videos
+            ]
+
+    def decode(self, zs):
+        with amp.autocast(dtype=self.dtype):
+            return [
+                self.model.decode(u.unsqueeze(0),
+                                  self.scale).float().clamp_(-1, 1).squeeze(0)
+                for u in zs
+            ]

wan/utils/__init__.py

@@ -0,0 +1,4 @@
+# Trimmed package initializer for the CMD Hugging Face Space.
+# The upstream nv-tlabs/cmd package initializers eagerly import training-only
+# and multi-GPU modules that this inference-only demo does not need. See NOTICE
+# and LICENSE for the upstream licenses that continue to apply.

wan/utils/fm_solvers_unipc.py

@@ -0,0 +1,814 @@
+# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
+# Convert unipc for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+                                                   SchedulerMixin,
+                                                   SchedulerOutput)
+from diffusers.utils import deprecate, is_scipy_available
+
+if is_scipy_available():
+    import scipy.stats
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        solver_order (`int`, default `2`):
+            The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+            due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+            unconditional sampling.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+        predict_x0 (`bool`, defaults to `True`):
+            Whether to use the updating algorithm on the predicted x0.
+        solver_type (`str`, default `bh2`):
+            Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+            otherwise.
+        lower_order_final (`bool`, default `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        disable_corrector (`list`, default `[]`):
+            Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+            and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+            usually disabled during the first few steps.
+        solver_p (`SchedulerMixin`, default `None`):
+            Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+            self,
+            num_train_timesteps: int = 1000,
+            solver_order: int = 2,
+            prediction_type: str = "flow_prediction",
+            shift: Optional[float] = 1.0,
+            use_dynamic_shifting=False,
+            thresholding: bool = False,
+            dynamic_thresholding_ratio: float = 0.995,
+            sample_max_value: float = 1.0,
+            predict_x0: bool = True,
+            solver_type: str = "bh2",
+            lower_order_final: bool = True,
+            disable_corrector: List[int] = [],
+            solver_p: SchedulerMixin = None,
+            timestep_spacing: str = "linspace",
+            steps_offset: int = 0,
+            final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+    ):
+
+        if solver_type not in ["bh1", "bh2"]:
+            if solver_type in ["midpoint", "heun", "logrho"]:
+                self.register_to_config(solver_type="bh2")
+            else:
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}")
+
+        self.predict_x0 = predict_x0
+        # setable values
+        self.num_inference_steps = None
+        alphas = np.linspace(1, 1 / num_train_timesteps,
+                             num_train_timesteps)[::-1].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.timesteps = sigmas * num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.timestep_list = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = disable_corrector
+        self.solver_p = solver_p
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: Union[int, None] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[Union[float, None]] = None,
+        shift: Optional[Union[float, None]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            sigmas = np.linspace(self.sigma_max, self.sigma_min,
+                                 num_inference_steps +
+                                 1).copy()[:-1]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) /
+                          self.alphas_cumprod[0])**0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]
+                                 ]).astype(np.float32)  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(
+            device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.last_sample = None
+        if self.solver_p:
+            self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float(
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(
+            abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(
+            1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(
+            sample, -s, s
+        ) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Convert the model output to the corresponding type the UniPC algorithm needs.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    "missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+        else:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    def multistep_uni_p_bh_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model at the current timestep.
+            prev_timestep (`int`):
+                The previous discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            order (`int`):
+                The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "prev_timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    " missing `sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 2:
+                order = args[2]
+            else:
+                raise ValueError(
+                    " missing `order` as a required keyward argument")
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+        model_output_list = self.model_outputs
+
+        s0 = self.timestep_list[-1]
+        m0 = model_output_list[-1]
+        x = sample
+
+        if self.solver_p:
+            x_t = self.solver_p.step(model_output, s0, x).prev_sample
+            return x_t
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[
+            self.step_index]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - i  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            # for order 2, we use a simplified version
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1],
+                                            b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+                                        D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+                                        D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output: torch.Tensor,
+        *args,
+        last_sample: torch.Tensor = None,
+        this_sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniC (B(h) version).
+
+        Args:
+            this_model_output (`torch.Tensor`):
+                The model outputs at `x_t`.
+            this_timestep (`int`):
+                The current timestep `t`.
+            last_sample (`torch.Tensor`):
+                The generated sample before the last predictor `x_{t-1}`.
+            this_sample (`torch.Tensor`):
+                The generated sample after the last predictor `x_{t}`.
+            order (`int`):
+                The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+        Returns:
+            `torch.Tensor`:
+                The corrected sample tensor at the current timestep.
+        """
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "this_timestep", None)
+        if last_sample is None:
+            if len(args) > 1:
+                last_sample = args[1]
+            else:
+                raise ValueError(
+                    " missing`last_sample` as a required keyward argument")
+        if this_sample is None:
+            if len(args) > 2:
+                this_sample = args[2]
+            else:
+                raise ValueError(
+                    " missing`this_sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 3:
+                order = args[3]
+            else:
+                raise ValueError(
+                    " missing`order` as a required keyward argument")
+        if this_timestep is not None:
+            deprecate(
+                "this_timestep",
+                "1.0.0",
+                "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        model_output_list = self.model_outputs
+
+        m0 = model_output_list[-1]
+        x = last_sample
+        x_t = this_sample
+        model_t = this_model_output
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+            self.step_index - 1]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - (i + 1)  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+
+        # for order 1, we use a simplified version
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(self,
+             model_output: torch.Tensor,
+             timestep: Union[int, torch.Tensor],
+             sample: torch.Tensor,
+             return_dict: bool = True,
+             generator=None) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep UniPC.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        use_corrector = (
+            self.step_index > 0 and
+            self.step_index - 1 not in self.disable_corrector and
+            self.last_sample is not None  # pyright: ignore
+        )
+
+        model_output_convert = self.convert_model_output(
+            model_output, sample=sample)
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep  # pyright: ignore
+
+        if self.config.lower_order_final:
+            this_order = min(self.config.solver_order,
+                             len(self.timesteps) -
+                             self.step_index)  # pyright: ignore
+        else:
+            this_order = self.config.solver_order
+
+        self.this_order = min(this_order,
+                              self.lower_order_nums + 1)  # warmup for multistep
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # upon completion increase step index by one
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(
+            device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(
+                timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(
+                original_samples.device, dtype=torch.float32)
+            timesteps = timesteps.to(
+                original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps)
+                for t in timesteps
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps