Upload code and configuration only

Source code, configs, dependency metadata, README and license. Excludes datasets, model weights, generated outputs and runtime artifacts.

.gitignore

@@ -0,0 +1,9 @@
+__pycache__
+*.egg-info
+
+wan_models
+checkpoints
+output
+dataset
+prompts/vidprom_filtered_extended.txt
+logs

LICENSE

@@ -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.

README.md

@@ -0,0 +1,460 @@
+<div align="center">
+
+## Causal Forcing & Causal Forcing++
+### Autoregressive Diffusion Distillation Done Right for High-Quality Real-Time Interactive Video Generation
+
+<p align="center">
+  <p align="center">
+    <div>
+    <a href="https://zhuhz22.github.io/" target="_blank">Hongzhou Zhu*</a><sup></sup>,
+    <a href="https://gracezhao1997.github.io/" target="_blank">Min Zhao*</a><sup></sup> , 
+    <a href="https://guandehe.github.io/" target="_blank">Guande He</a><sup></sup>, 
+    <a href="https://scholar.google.com/citations?user=dxN1_X0AAAAJ&hl=en" target="_blank">Hang Su</a><sup></sup>,
+    <a href="https://zhenxuan00.github.io/" target="_blank">Chongxuan Li</a><sup></sup> ,
+    <a href="https://ml.cs.tsinghua.edu.cn/~jun/index.shtml" target="_blank">Jun Zhu</a><sup></sup>
+  </div>
+  <div>
+      <sup></sup>Tsinghua University & Shengshu & UT Austin & RUC
+  </div>
+
+
+</div>
+  </p>
+  <h3 align="center"><a href="https://arxiv.org/abs/2602.02214">Causal Forcing</a> | <a href="https://arxiv.org/abs/2605.15141">Causal Forcing++</a> | <a href="https://thu-ml.github.io/CausalForcing.github.io">Website</a> | <a href="https://huggingface.co/zhuhz22/Causal-Forcing/tree/main">Models</a> | <a href="https://my.feishu.cn/wiki/AjBSwcjpqiN0ECkodIWcGDcMn4e?from=from_copylink">Document</a> </h3>
+</p>
+
+
+
+-----
+The Causal Forcing series uses **Causal ODE** or **Causal Consistency Distillation** to drive asymmetric DMD as a theoretically correct initialization for real-time interactive video generation. **Refer to [this extension for long video generation](#minute-level-long-video-generation).**
+
+[Causal Forcing](https://arxiv.org/abs/2602.02214) significantly outperforms Self Forcing in **both visual quality and motion dynamics**, while keeping **the same training budget and inference efficiency**. We support both chunk-wise and **frame-wise** models, with the latter natively unifying T2V and **I2V**.
+
+We further propose [**Causal Forcing++**](https://arxiv.org/abs/2605.15141)[ technical report ], replacing ODE with **causal Consistency Distillation** to eliminate ODE data curation and improve performance, releasing the first **1-step/2-step frame-wise** models.
+
+-----
+<img width="2090" height="850" alt="overview" src="assets/pipeline.png" />
+
+## Table of Contents
+
+- [Quick Start](#quick-start)
+  - [Installation](#installation) 
+  - [Inference: 1/2/4-step T2V & I2V](#cli-inference)
+  - [Long Video](#minute-level-long-video-generation)
+- [Training Pipeline](#training) (Causal Foricng / Causal Forcing++)
+  - [Stage 1: AR Diffusion](#training)
+  - [Stage 2: Causal ODE (Causal Forcing)](#training) / [🔥Causal CD (Causal Forcing++)](#-stage-2-option-b-causal-consistency-distillation-initialization-causal-forcing)
+  - [Stage 3: Asymmetric DMD](#stage-3-dmd)
+
+
+| Models              | Checkpoints | Description|
+|--------------------|---------------------------------------------------------------------------------------------------------------------------------------------|-------|
+| Chunk-wise 4-step   | 🤗 [Huggingface](https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/chunkwise/causal_forcing.pt)    | SOTA AR model on Wan1.3B outperforming Self Forcing|
+| Frame-wise 4-step   | 🤗 [Huggingface](https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/framewise/causal_forcing.pt)    | Rich dynamics and high quality |
+| Frame-wise 2-step🔥 | 🤗 [Huggingface](https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/causal-forcing%2B%2B/framewise-2step.pt)    | The  first frame-wise 2-step model **even better than 4-step**|
+| Frame-wise 1-step   | 🤗 [Huggingface](https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/causal-forcing%2B%2B/framewise-1step.pt)    | Extremely low latency |
+| Long Video Generator  🔥 | 🤗 [Huggingface](https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/chunkwise/longvideo.pt)    | Minute-long AR video generator |
+| HY1.5-TI2V (8B) 🔥  | 🤗 [Huggingface](https://huggingface.co/MIN-Lab/minWM)   | Refer to [this repo](https://github.com/shengshu-ai/minWM). |
+| Action-conditioned WM   | 🤗 [Huggingface](https://huggingface.co/MIN-Lab/minWM)    | Refer to [this repo](https://github.com/shengshu-ai/minWM). |
+
+-----
+
+https://github.com/user-attachments/assets/310f0cfa-e1bb-496d-8941-87f77b3271c0
+
+## 🔥 News
+- **2026.7.23**: [Self Gradient Forcing](https://github.com/zhuang2002/Self_Gradient_Forcing) is built on Causal Forcing initialization.
+- **2026.7.20**: Happy to see that the recent SOTA video world models [DreamX-World 1.0](https://arxiv.org/pdf/2606.16993) and [Matrix-Game 3.5](https://matrix-game-v3-5.github.io/paper/Matrix-Game-3.5.pdf) are built on Causal Forcing!
+- **2026.5.17**: We release Causal Forcing for the HY1.5-TI2V-8B model! Refer to [this repo](https://github.com/shengshu-ai/minWM) for the details. This model explicitly supports I2V.
+- **2026.5.15**: We release [Causal Forcing++](https://arxiv.org/abs/2605.15141), supporting Casual Consistency Distillation for few-step initialization, and open-source **the first frame-wise 2-step AR model** comparable to chunk-wise 4-step models!
+- **2026.5.10**: Thanks to @[AshadowZ](https://github.com/AshadowZ), now our chunk-wise ODE data curation is **3x faster**!
+- **2026.4.16**: **Optimize the Stage 2 🔥consistency distillation🔥 infrastructure for 3× faster training, let's try it now!** We have also released the ckpt.
+- **2026.3.15** : [Rolling Sink](https://github.com/haodong2000/RollingSink), [Infinity-RoPE](https://github.com/yesiltepe-hidir/infinity-rope) and [Deep Forcing](https://cvlab-kaist.github.io/DeepForcing/) adopt Causal Forcing as one of the base models!
+- **2026.2.28** : Add [FAQ section](#faq--blog) regarding hot topics, specifically which is the better Initialization between AR diffusion and causal ODE distillation.
+- **2026.2.11** : We now support **I2V** generation! Feel free to try it [here](#new-i2v)!
+- **2026.2.7** : Causal Forcing now supports [Rolling Forcing](https://github.com/TencentARC/RollingForcing), enabling minute-level long video generation!
+- **2026.2.5** : Release causal consistency distillation (Preview) as substitute for ODE distillation, **free of generating ODE paired data**!
+- **2026.2.2** : The [paper](https://arxiv.org/abs/2602.02214), [project page](https://thu-ml.github.io/CausalForcing.github.io/), and code are released.
+
+
+## Quick Start
+
+> The inference environment is identical to Self Forcing.
+
+**NOTE**: Similar to CausVid/Self Forcing, Causal Forcing does not natively support videos longer than 81 frames. As a base training method, it is orthogonal to techniques like Longlive/Rolling Forcing. To use Causal Forcing as a long video baseline, see [this extension](#minute-level-long-video-generation). **Directly using the 5-second trained Causal Forcing model as a baseline for long video generation is extremely unfair**.
+
+
+### Installation
+```bash
+conda create -n causal_forcing python=3.10 -y
+conda activate causal_forcing
+pip install -r requirements.txt
+pip install git+https://github.com/openai/CLIP.git
+pip install flash-attn --no-build-isolation
+python setup.py develop
+```
+### Download Checkpoints
+```bash
+hf download Wan-AI/Wan2.1-T2V-1.3B  --local-dir wan_models/Wan2.1-T2V-1.3B
+hf download Wan-AI/Wan2.1-T2V-14B  --local-dir wan_models/Wan2.1-T2V-14B
+# Causal Forcing
+hf download zhuhz22/Causal-Forcing chunkwise/causal_forcing.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing framewise/causal_forcing.pt --local-dir checkpoints
+# Causal Forcing++
+hf download zhuhz22/Causal-Forcing causal-forcing++/framewise-2step.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing causal-forcing++/framewise-1step.pt --local-dir checkpoints
+```
+
+### CLI Inference
+
+#### T2V
+Chunk-wise model:
+```bash
+python inference.py \
+  --config_path configs/causal_forcing_dmd_chunkwise.yaml \
+  --output_folder output/chunkwise \
+  --checkpoint_path checkpoints/chunkwise/causal_forcing.pt \
+  --data_path prompts/demos.txt
+```
+
+Frame-wise model:
+
+```bash
+# =============== Causal Forcing++ ================
+# 2-step Causal Forcing++
+python inference.py \
+  --config_path configs/causal_forcing_dmd_framewise_2step.yaml \
+  --output_folder output/framewise_2step_cf++ \
+  --checkpoint_path checkpoints/causal-forcing++/framewise-2step.pt \
+  --data_path prompts/demos.txt \
+  --use_ema
+
+# 1-step Causal Forcing++
+python inference.py \
+  --config_path configs/causal_forcing_dmd_framewise_1step.yaml \
+  --output_folder output/framewise_1step_cf++ \
+  --checkpoint_path checkpoints/causal-forcing++/framewise-1step.pt \
+  --data_path prompts/demos.txt \
+  --use_ema
+
+# =============== Causal Forcing ================
+# 4-step Causal Forcing
+python inference.py \
+  --config_path configs/causal_forcing_dmd_framewise.yaml \
+  --output_folder output/framewise \
+  --checkpoint_path  checkpoints/framewise/causal_forcing.pt \
+  --data_path prompts/demos.txt \
+  --use_ema 
+```
+
+#### I2V
+> Our frame-wise setting natively supports I2V. You simply need to set the first latent initial frame as your conditional image. 
+
+```bash
+python inference.py \
+  --config_path configs/causal_forcing_dmd_framewise.yaml \
+  --output_folder output/framewise \
+  --checkpoint_path  checkpoints/framewise/causal_forcing.pt \
+  --data_path prompts/i2v \
+  --i2v \
+  --use_ema
+```
+
+
+### Minute-level Long Video Generation
+Built on [Rolling Forcing](https://github.com/TencentARC/RollingForcing), we implemented minute-level long video generation. See [here](./long_video) for the detail.
+
+[Infinity-RoPE](https://github.com/yesiltepe-hidir/infinity-rope), [Deep Forcing](https://cvlab-kaist.github.io/DeepForcing/) and [Rolling Sink](https://github.com/haodong2000/RollingSink) also adopt Causal Forcing as one of their base models, enabling interactive (prompt-switchable) long video generation at the minute scale. You can also try them out at their repos.
+
+## Training
+
+<details>
+<summary> Stage 1: Autoregressive Diffusion Training (Can skip by using our pretrained checkpoints. Click to expand.)</summary>
+
+First download the dataset (we provide a 6K toy dataset here):
+```bash
+hf download zhuhz22/Causal-Forcing-data  --local-dir dataset
+python utils/merge_and_get_clean.py
+```
+> If the download gets stuck, Ctrl^C and then resume it.
+
+
+> For training on your own dataset, refer to [this issue](https://github.com/thu-ml/Causal-Forcing/issues/8).
+
+
+Then train the AR-diffusion model:
+- Framewise:
+  ```bash
+    torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/ar_diffusion_tf_framewise.yaml \
+    --logdir logs/ar_diffusion_framewise
+  ```
+
+- Chunkwise:
+  ```bash
+    torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/ar_diffusion_tf_chunkwise.yaml \
+    --logdir logs/ar_diffusion_chunkwise
+  ```
+
+> We recommend training no less than 2K steps, and more steps (e.g., 5~10K) will lead to better performance.
+
+Inference to test training results:
+```bash
+python inference.py \
+  --config_path configs/ar_diffusion_tf_{framewise OR chunkwise}.yaml \
+  --output_folder output/{framewise OR chunkwise}_ar_diffusion \
+  --checkpoint_path  checkpoints/{framewise OR chunkwise}/ar_diffusion.pt \
+  --data_path prompts/demos.txt
+```
+</details>
+
+
+<details>
+<summary> Stage 2 (Option a): Causal ODE Initialization (Can skip by using our pretrained checkpoints. Click to expand.)</summary>
+
+🔥 Thanks to @[AshadowZ](https://github.com/AshadowZ), now our chunk-wise ODE data curation is **3x faster**!
+
+You can also use `bf16` to accelerate generation.
+
+If you have skipped Stage 1, you need to download the pretrained models:
+```bash
+hf download zhuhz22/Causal-Forcing framewise/ar_diffusion.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing chunkwise/ar_diffusion.pt --local-dir checkpoints
+```
+
+In this stage, first generate ODE paired data:
+```bash
+# for the frame-wise model
+torchrun --nproc_per_node=8 \
+  get_causal_ode_data_framewise.py \
+  --generator_ckpt checkpoints/framewise/ar_diffusion.pt \
+  --rawdata_path dataset/clean_data \
+  --output_folder dataset/ODE6KCausal_framewise_latents
+
+python utils/create_lmdb_iterative.py \
+  --data_path dataset/ODE6KCausal_framewise_latents \
+  --lmdb_path dataset/ODE6KCausal_framewise
+
+# for the chunk-wise model
+torchrun --nproc_per_node=8 \
+  get_causal_ode_data_chunkwise.py \
+  --generator_ckpt checkpoints/chunkwise/ar_diffusion.pt \
+  --rawdata_path dataset/clean_data \
+  --output_folder dataset/ODE6KCausal_chunkwise_latents
+
+python utils/create_lmdb_iterative.py \
+  --data_path dataset/ODE6KCausal_chunkwise_latents \
+  --lmdb_path dataset/ODE6KCausal_chunkwise
+
+#🔥NEW: Or you can use the optimized code (3x speedup) by @AshadowZ 
+torchrun --nproc_per_node=8 \
+  get_causal_ode_data_kv.py \
+  --generator_ckpt checkpoints/{chunkwise,framewise}/ar_diffusion.pt \
+  --rawdata_path /mnt/vepfs/base2/zhaomin/group_0001_lmdb \
+  --output_folder dataset/ODE6KCausal_{chunkwise,framewise}_latents \
+  --num_frames_per_chunk 3/1 # 3 for chunkwise, 1 for framewise
+  --generation_mode blockwise_kv
+```
+
+Or you can also directly download our prepared dataset (~300G):
+```bash
+hf download zhuhz22/Causal-Forcing-data  --local-dir dataset
+python utils/merge_lmdb.py
+```
+> If the download gets stuck, Ctrl^C and then resume it.
+
+
+And then train ODE initialization models:
+- Frame-wise:
+  ```bash
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_ode_framewise.yaml \
+    --logdir logs/causal_ode_framewise
+  ```
+- Chunk-wise:
+  ```bash
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_ode_chunkwise.yaml \
+    --logdir logs/causal_ode_chunkwise
+  ```
+
+> We recommend training no less than 1K steps, and more steps (e.g., 5~10K) will lead to better performance.
+
+Inference to test training results:
+
+The same as [here](#cli-inference).
+</details>
+
+
+
+### 🔥 Stage 2 (Option b): Causal Consistency Distillation Initialization ([Causal Forcing++](https://arxiv.org/abs/2605.15141))
+Since creating ODE-paired data is very time-consuming, we also provide an alternative here that achieves the same effect (or better) as ODE distillation while requiring only ground-truth data, **free of generating ODE data!**
+
+> Thanks to [@chijw's effort](https://github.com/thu-ml/Causal-Forcing/pull/20), now the EMA mechanism is more efficient!
+
+If you have skipped Stage 1, you need to download the pretrained multi-step AR diffusion models:
+```bash
+hf download zhuhz22/Causal-Forcing framewise/ar_diffusion.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing chunkwise/ar_diffusion.pt --local-dir checkpoints
+```
+
+- Frame-wise:
+  ```bash
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_cd_framewise.yaml \
+    --logdir logs/causal_cd_framewise
+  ```
+- Chunk-wise:
+  ```bash
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_cd_chunkwise.yaml \
+    --logdir logs/causal_cd_chunkwise
+  ```
+
+> We recommend training no less than 3K steps, and more steps (e.g., 5~10K) will lead to better performance.
+
+<details>
+<summary>
+You can also download the checkpoints directly (click to expand) :
+</summary>
+
+```bash
+hf download zhuhz22/Causal-Forcing framewise/causal_cd.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing chunkwise/causal_cd.pt --local-dir checkpoints
+```
+</details>
+
+Inference to test training results: the same as [here](#cli-inference).
+
+
+
+
+### Stage 3: DMD
+
+First download the dataset:
+```bash
+hf download gdhe17/Self-Forcing vidprom_filtered_extended.txt --local-dir prompts
+```
+
+<details>
+
+<summary>
+If you have skipped Stage 2, you need to download our pretrained checkpoints (click to expand):
+</summary>
+
+```bash
+hf download zhuhz22/Causal-Forcing framewise/causal_ode.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing chunkwise/causal_ode.pt --local-dir checkpoints
+
+# 🔥 Or you can also use the improved version (Causal Forcing++): causal CD instead of causal ODE
+hf download zhuhz22/Causal-Forcing framewise/causal_cd.pt --local-dir checkpoints
+hf download zhuhz22/Causal-Forcing chunkwise/causal_cd.pt --local-dir checkpoints
+```
+</details>
+
+And then train DMD models:
+
+- Frame-wise model:
+  ```bash
+  # =============== Causal Forcing ================
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_forcing_dmd_framewise.yaml \
+    --logdir logs/causal_forcing_dmd_framewise
+  # =============== Causal Forcing++ ================
+  # 2-step Causal Forcing++
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_forcing_dmd_framewise_2step.yaml \
+    --logdir logs/causal_forcing_dmd_framewise
+
+  # 1-step Causal Forcing++
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_forcing_dmd_framewise_1step.yaml \
+    --logdir logs/causal_forcing_dmd_framewise
+  ```
+
+- Chunk-wise model:
+  ```bash
+  torchrun --nnodes=8 --nproc_per_node=8 --rdzv_id=5235 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint $MASTER_ADDR \
+    train.py \
+    --config_path configs/causal_forcing_dmd_chunkwise.yaml \
+    --logdir logs/causal_forcing_dmd_chunkwise
+  ```
+> For bs64, we recommend training for no more than 1K steps; otherwise the motion dynamics may degrade. If bs is smaller (e.g., 8), more training steps is preferable.
+
+Such models are the final models used to generate videos.
+## FAQ & Blog 
+See the [FAQ](https://my.feishu.cn/wiki/AjBSwcjpqiN0ECkodIWcGDcMn4e) and the [blog](https://zhuanlan.zhihu.com/p/2002114039493461457). (currently in Chinese)
+
+<details>
+<summary> Typical Questions (click to expand)
+</summary>
+
+**Why using bidirectional teacher in the DMD stage ?**
+- Q: In the DMD stage, do you still use a bidirectional teacher? Why not an AR teacher?
+- A: Yes. DMD only requires the student to match the teacher’s final distribution, not the generation trajectory, so a bidirectional teacher is fine. Also, bidirectional diffusion models are typically stronger than AR diffusion, so they make a better teacher.
+  
+- Q: Then why must the ODE (or Consistency Distillation) stage use an AR teacher?
+- A: Because ODE/CD requires the student and teacher to follow the same trajectory, so their structures must be matched; an AR student cannot be trajectory-aligned with a bidirectional teacher.
+  
+🔥🔥 **ODE initialization or multi-step AR diffusion initialization ?**
+- Q: Which is better as initialization: a “proper” ODE initialization or directly using multi-step AR diffusion?
+- A: We compared this in the Appendix C2. Overall, proper ODE init is better: multi-step AR diffusion init + DMD occasionally yields grid-like or waxy/greasy results. A key reason is that DMD is inherently few-step, so the right comparison is under few-step; in that regime, a few-step diffusion teacher is much weaker than an ODE-distilled teacher. Without ODE distillation, DMD must both close the step gap and handle an added conditioning gap from self-rollout: early few-step errors corrupt the history and get amplified across frames (large exposure bias), which increases DMD pressure. It can still converge, but typically with worse quality than ODE initialization. Also, with ODE init, DMD can be trained very few steps (e.g., ~100), reducing the risk of dynamics degradation from long DMD training.
+  
+**Can frame-level non-injectivity appears in the actual training dataset ?**
+- Q: Regarding the “one-to-many” analysis in the ODE stage: since a single frame’s latent has very high dimensionality, isn’t the probability of being exactly identical extremely small?
+- A: Yes, but the key point here is not whether the dataset literally contains identical samples; it’s whether there exists a well-defined function in the mathematical sense. Our vision modalities live in a continuous space—even in 1D, getting two samples to be exactly identical is extremely unlikely. However, the theoretical existence of exact collisions is enough to break the function property and make it ill-defined.
+
+</details>
+
+## Acknowledgements
+- This codebase is built on top of the open-source implementation of [CausVid](https://github.com/tianweiy/CausVid), [Self Forcing](https://github.com/guandeh17/Self-Forcing), [Rolling Forcing](https://github.com/TencentARC/RollingForcing) and the [Wan2.1](https://github.com/Wan-Video/Wan2.1) repo. 
+- Thanks to @[chijw](https://github.com/chijw) for improving the EMA mechanism. 
+- Thanks to @[AshadowZ](https://github.com/AshadowZ) for improving the causal ODE data curation efficiency.
+- Causal Forcing++ with 1/2 steps applies the first-frame 4-step technique proposed by [ASD](https://github.com/BigAandSmallq/SAD). We thank ASD for its contribution.
+
+## References
+If you find the method useful, please cite
+```
+@article{zhu2026causal,
+  title={Causal Forcing: Autoregressive Diffusion Distillation Done Right for High-Quality Real-Time Interactive Video Generation},
+  author={Zhu, Hongzhou and Zhao, Min and He, Guande and Su, Hang and Li, Chongxuan and Zhu, Jun},
+  journal={arXiv preprint arXiv:2602.02214},
+  year={2026}
+}
+
+@article{zhao2026causal,
+  title={Causal Forcing++: Scalable Few-Step Autoregressive Diffusion Distillation for Real-Time Interactive Video Generation},
+  author={Zhao, Min and Zhu, Hongzhou and Zheng, Kaiwen and Zhou, Zihan and Yan, Bokai and Li, Xinyuan and Yang, Xiao and Li, Chongxuan and Zhu, Jun},
+  journal={arXiv preprint arXiv:2605.15141},
+  year={2026}
+}
+```

configs/ar_diffusion_tf_chunkwise.yaml

@@ -0,0 +1,46 @@
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+warp_denoising_step: true 
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 2000000000
+total_batch_size: 8
+log_iters: 1000
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: diffusion
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: dataset/clean_data 
+teacher_forcing: true # TF/DF

configs/ar_diffusion_tf_framewise.yaml

@@ -0,0 +1,46 @@
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+warp_denoising_step: true 
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 2000000000
+total_batch_size: 8
+log_iters: 1000
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: diffusion
+gradient_checkpointing: true
+num_frame_per_block: 1
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: dataset/clean_data 
+teacher_forcing: true # TF/DF

configs/causal_cd_chunkwise.yaml

@@ -0,0 +1,57 @@
+generator_ckpt: checkpoints/chunkwise/ar_diffusion.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 8
+log_iters: 1000
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: consistency_distillation
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: dataset/clean_data
+grad_accum: 1
+spatial_sf: true
+teacher_forcing: True
+causal: True
+is_causal: True
+discrete_cd_N: 48

configs/causal_cd_framewise.yaml

@@ -0,0 +1,57 @@
+generator_ckpt: checkpoints/framewise/ar_diffusion.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 8
+log_iters: 1000
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: consistency_distillation
+gradient_checkpointing: true
+num_frame_per_block: 1
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: dataset/clean_data
+grad_accum: 1
+spatial_sf: true
+teacher_forcing: True
+causal: True
+is_causal: True
+discrete_cd_N: 48

configs/causal_forcing_dmd_chunkwise.yaml

@@ -0,0 +1,51 @@
+generator_ckpt: checkpoints/chunkwise/causal_ode.pt # 🔥 or checkpoints/chunkwise/causal_cd.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-14B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 8
+log_iters: 250
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: prompts/vidprom_filtered_extended.txt 

configs/causal_forcing_dmd_framewise.yaml

@@ -0,0 +1,51 @@
+generator_ckpt: checkpoints/framewise/causal_ode.pt # 🔥 or checkpoints/framewise/causal_cd.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-14B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 8
+log_iters: 250
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 1
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: prompts/vidprom_filtered_extended.txt

configs/causal_forcing_dmd_framewise_1step.yaml

@@ -0,0 +1,53 @@
+generator_ckpt: checkpoints/framewise/causal_ode.pt # 🔥 or checkpoints/framewise/causal_cd.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-14B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+denoising_step_list_first_chunk: #  Causal Forcing++ with 1/2 steps applies the first-frame 4-step technique proposed by [ASD](https://github.com/BigAandSmallq/SAD). We thank ASD for its contribution.
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 8
+log_iters: 250
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 1
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: prompts/vidprom_filtered_extended.txt

configs/causal_forcing_dmd_framewise_2step.yaml

@@ -0,0 +1,54 @@
+generator_ckpt: checkpoints/framewise/causal_ode.pt # 🔥 or checkpoints/framewise/causal_cd.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-14B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 500
+denoising_step_list_first_chunk: #  Causal Forcing++ with 1/2 steps applies the first-frame 4-step technique proposed by [ASD](https://github.com/BigAandSmallq/SAD). We thank ASD for its contribution.
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+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
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 8
+log_iters: 250
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 1
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+data_path: prompts/vidprom_filtered_extended.txt

configs/causal_ode_chunkwise.yaml

@@ -0,0 +1,31 @@
+generator_ckpt: checkpoints/chunkwise/ar_diffusion.pt
+generator_grad:
+  model: true
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true
+generator_task: causal_video
+generator_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+wandb_name: ode
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+beta1: 0.9
+beta2: 0.999
+data_path: dataset/ODE6KCausal_chunkwise
+batch_size: 1
+log_iters: 1000
+trainer: ode
+gradient_checkpointing: true
+num_frame_per_block: 3
+model_kwargs:
+  timestep_shift: 5.0

configs/causal_ode_framewise.yaml

@@ -0,0 +1,31 @@
+generator_ckpt: checkpoints/framewise/ar_diffusion.pt
+generator_grad:
+  model: true
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true
+generator_task: causal_video
+generator_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size
+mixed_precision: true
+seed: 0
+wandb_host: https://api.wandb.ai
+wandb_key: {your key}
+wandb_entity: {your entity}
+wandb_project: {your project}
+wandb_name: ode
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+beta1: 0.9
+beta2: 0.999
+data_path: dataset/ODE6KCausal_framewise
+batch_size: 1
+log_iters: 1000
+trainer: ode
+gradient_checkpointing: true
+num_frame_per_block: 1 
+model_kwargs:
+  timestep_shift: 5.0

configs/default_config.yaml

@@ -0,0 +1,20 @@
+independent_first_frame: false
+warp_denoising_step: false
+weight_decay: 0.01
+same_step_across_blocks: true
+discriminator_lr_multiplier: 1.0
+last_step_only: false
+i2v: false
+num_training_frames: 21
+gc_interval: 100
+context_noise: 0
+causal: true
+
+ckpt_step: 0
+prompt_name: MovieGenVideoBench
+prompt_path: prompts/MovieGenVideoBench.txt
+eval_first_n: 64
+num_samples: 1
+height: 480
+width: 832
+num_frames: 81

demo.py

@@ -0,0 +1,631 @@
+"""
+Demo for Causal-Forcing.
+"""
+
+import os
+import re
+import random
+import time
+import base64
+import argparse
+import hashlib
+import subprocess
+import urllib.request
+from io import BytesIO
+from PIL import Image
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from flask import Flask, render_template, jsonify
+from flask_socketio import SocketIO, emit
+import queue
+from threading import Thread, Event
+
+from pipeline import CausalInferencePipeline
+from demo_utils.constant import ZERO_VAE_CACHE
+from demo_utils.vae_block3 import VAEDecoderWrapper
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder
+from demo_utils.utils import generate_timestamp
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation
+
+# Parse arguments
+parser = argparse.ArgumentParser()
+parser.add_argument('--port', type=int, default=5001)
+parser.add_argument('--host', type=str, default='0.0.0.0')
+parser.add_argument("--checkpoint_path", type=str, default='.checkpoints/framewise/causal_forcing.pt')
+parser.add_argument("--config_path", type=str, default='.configs/causal_forcing_dmd_framewise.yaml')
+parser.add_argument('--trt', action='store_true')
+args = parser.parse_args()
+
+print(f'Free VRAM {get_cuda_free_memory_gb(gpu)} GB')
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+# Load models
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+
+text_encoder = WanTextEncoder()
+
+# Global variables for dynamic model switching
+current_vae_decoder = None
+current_use_taehv = False
+fp8_applied = False
+torch_compile_applied = False
+global frame_number
+frame_number = 0
+anim_name = ""
+frame_rate = 6
+
+def initialize_vae_decoder(use_taehv=False, use_trt=False):
+    """Initialize VAE decoder based on the selected option"""
+    global current_vae_decoder, current_use_taehv
+
+    if use_trt:
+        from demo_utils.vae import VAETRTWrapper
+        current_vae_decoder = VAETRTWrapper()
+        return current_vae_decoder
+
+    if use_taehv:
+        from demo_utils.taehv import TAEHV
+        # Check if taew2_1.pth exists in checkpoints folder, download if missing
+        taehv_checkpoint_path = "checkpoints/taew2_1.pth"
+        if not os.path.exists(taehv_checkpoint_path):
+            print(f"taew2_1.pth not found in checkpoints folder {taehv_checkpoint_path}. Downloading...")
+            os.makedirs("checkpoints", exist_ok=True)
+            download_url = "https://github.com/madebyollin/taehv/raw/main/taew2_1.pth"
+            try:
+                urllib.request.urlretrieve(download_url, taehv_checkpoint_path)
+                print(f"Successfully downloaded taew2_1.pth to {taehv_checkpoint_path}")
+            except Exception as e:
+                print(f"Failed to download taew2_1.pth: {e}")
+                raise
+
+        class DotDict(dict):
+            __getattr__ = dict.__getitem__
+            __setattr__ = dict.__setitem__
+
+        class TAEHVDiffusersWrapper(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.dtype = torch.float16
+                self.taehv = TAEHV(checkpoint_path=taehv_checkpoint_path).to(self.dtype)
+                self.config = DotDict(scaling_factor=1.0)
+
+            def decode(self, latents, return_dict=None):
+                # n, c, t, h, w = latents.shape
+                # low-memory, set parallel=True for faster + higher memory
+                return self.taehv.decode_video(latents, parallel=False).mul_(2).sub_(1)
+
+        current_vae_decoder = TAEHVDiffusersWrapper()
+    else:
+        current_vae_decoder = VAEDecoderWrapper()
+        vae_state_dict = torch.load('wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth', map_location="cpu")
+        decoder_state_dict = {}
+        for key, value in vae_state_dict.items():
+            if 'decoder.' in key or 'conv2' in key:
+                decoder_state_dict[key] = value
+        current_vae_decoder.load_state_dict(decoder_state_dict)
+
+    current_vae_decoder.eval()
+    current_vae_decoder.to(dtype=torch.float16)
+    current_vae_decoder.requires_grad_(False)
+    current_vae_decoder.to(gpu)
+    current_use_taehv = use_taehv
+
+    print(f"✅ VAE decoder initialized with {'TAEHV' if use_taehv else 'default VAE'}")
+    return current_vae_decoder
+
+
+# Initialize with default VAE
+vae_decoder = initialize_vae_decoder(use_taehv=False, use_trt=args.trt)
+
+transformer = WanDiffusionWrapper(is_causal=True)
+state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+transformer.load_state_dict(state_dict['generator_ema'])
+
+text_encoder.eval()
+transformer.eval()
+
+transformer.to(dtype=torch.float16)
+text_encoder.to(dtype=torch.bfloat16)
+
+text_encoder.requires_grad_(False)
+transformer.requires_grad_(False)
+
+pipeline = CausalInferencePipeline(
+    config,
+    device=gpu,
+    generator=transformer,
+    text_encoder=text_encoder,
+    vae=vae_decoder
+)
+
+if low_memory:
+    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
+else:
+    text_encoder.to(gpu)
+transformer.to(gpu)
+
+# Flask and SocketIO setup
+app = Flask(__name__)
+app.config['SECRET_KEY'] = 'frontend_buffered_demo'
+socketio = SocketIO(app, cors_allowed_origins="*")
+
+generation_active = False
+stop_event = Event()
+frame_send_queue = queue.Queue()
+sender_thread = None
+models_compiled = False
+
+
+def tensor_to_base64_frame(frame_tensor):
+    """Convert a single frame tensor to base64 image string."""
+    global frame_number, anim_name
+    # Clamp and normalize to 0-255
+    frame = torch.clamp(frame_tensor.float(), -1., 1.) * 127.5 + 127.5
+    frame = frame.to(torch.uint8).cpu().numpy()
+
+    # CHW -> HWC
+    if len(frame.shape) == 3:
+        frame = np.transpose(frame, (1, 2, 0))
+
+    # Convert to PIL Image
+    if frame.shape[2] == 3:  # RGB
+        image = Image.fromarray(frame, 'RGB')
+    else:  # Handle other formats
+        image = Image.fromarray(frame)
+
+    # Convert to base64
+    buffer = BytesIO()
+    image.save(buffer, format='JPEG', quality=100)
+    if not os.path.exists("./images/%s" % anim_name):
+        os.makedirs("./images/%s" % anim_name)
+    frame_number += 1
+    image.save("./images/%s/%s_%03d.jpg" % (anim_name, anim_name, frame_number))
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/jpeg;base64,{img_str}"
+
+
+def frame_sender_worker():
+    """Background thread that processes frame send queue non-blocking."""
+    global frame_send_queue, generation_active, stop_event
+
+    print("📡 Frame sender thread started")
+
+    while True:
+        frame_data = None
+        try:
+            # Get frame data from queue
+            frame_data = frame_send_queue.get(timeout=1.0)
+
+            if frame_data is None:  # Shutdown signal
+                frame_send_queue.task_done()  # Mark shutdown signal as done
+                break
+
+            frame_tensor, frame_index, block_index, job_id = frame_data
+
+            # Convert tensor to base64
+            base64_frame = tensor_to_base64_frame(frame_tensor)
+
+            # Send via SocketIO
+            try:
+                socketio.emit('frame_ready', {
+                    'data': base64_frame,
+                    'frame_index': frame_index,
+                    'block_index': block_index,
+                    'job_id': job_id
+                })
+            except Exception as e:
+                print(f"⚠️ Failed to send frame {frame_index}: {e}")
+
+            frame_send_queue.task_done()
+
+        except queue.Empty:
+            # Check if we should continue running
+            if not generation_active and frame_send_queue.empty():
+                break
+        except Exception as e:
+            print(f"❌ Frame sender error: {e}")
+            # Make sure to mark task as done even if there's an error
+            if frame_data is not None:
+                try:
+                    frame_send_queue.task_done()
+                except Exception as e:
+                    print(f"❌ Failed to mark frame task as done: {e}")
+            break
+
+    print("📡 Frame sender thread stopped")
+
+
+@torch.no_grad()
+def generate_video_stream(prompt, seed, enable_torch_compile=False, enable_fp8=False, use_taehv=False):
+    """Generate video and push frames immediately to frontend."""
+    global generation_active, stop_event, frame_send_queue, sender_thread, models_compiled, torch_compile_applied, fp8_applied, current_vae_decoder, current_use_taehv, frame_rate, anim_name
+
+    try:
+        generation_active = True
+        stop_event.clear()
+        job_id = generate_timestamp()
+
+        # Start frame sender thread if not already running
+        if sender_thread is None or not sender_thread.is_alive():
+            sender_thread = Thread(target=frame_sender_worker, daemon=True)
+            sender_thread.start()
+
+        # Emit progress updates
+        def emit_progress(message, progress):
+            try:
+                socketio.emit('progress', {
+                    'message': message,
+                    'progress': progress,
+                    'job_id': job_id
+                })
+            except Exception as e:
+                print(f"❌ Failed to emit progress: {e}")
+
+        emit_progress('Starting generation...', 0)
+
+        # Handle VAE decoder switching
+        if use_taehv != current_use_taehv:
+            emit_progress('Switching VAE decoder...', 2)
+            print(f"🔄 Switching VAE decoder to {'TAEHV' if use_taehv else 'default VAE'}")
+            current_vae_decoder = initialize_vae_decoder(use_taehv=use_taehv)
+            # Update pipeline with new VAE decoder
+            pipeline.vae = current_vae_decoder
+
+        # Handle FP8 quantization
+        if enable_fp8 and not fp8_applied:
+            emit_progress('Applying FP8 quantization...', 3)
+            print("🔧 Applying FP8 quantization to transformer")
+            from torchao.quantization.quant_api import quantize_, Float8DynamicActivationFloat8WeightConfig, PerTensor
+            quantize_(transformer, Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor()))
+            fp8_applied = True
+
+        # Text encoding
+        emit_progress('Encoding text prompt...', 8)
+        conditional_dict = text_encoder(text_prompts=[prompt])
+        for key, value in conditional_dict.items():
+            conditional_dict[key] = value.to(dtype=torch.float16)
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(
+                text_encoder,target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
+
+        # Handle torch.compile if enabled
+        torch_compile_applied = enable_torch_compile
+        if enable_torch_compile and not models_compiled:
+            # Compile transformer and decoder
+            transformer.compile(mode="max-autotune-no-cudagraphs")
+            if not current_use_taehv and not low_memory and not args.trt:
+                current_vae_decoder.compile(mode="max-autotune-no-cudagraphs")
+
+        # Initialize generation
+        emit_progress('Initializing generation...', 12)
+
+        rnd = torch.Generator(gpu).manual_seed(seed)
+        # all_latents = torch.zeros([1, 21, 16, 60, 104], device=gpu, dtype=torch.bfloat16)
+
+        pipeline._initialize_kv_cache(batch_size=1, dtype=torch.float16, device=gpu)
+        pipeline._initialize_crossattn_cache(batch_size=1, dtype=torch.float16, device=gpu)
+
+        noise = torch.randn([1, 21, 16, 60, 104], device=gpu, dtype=torch.float16, generator=rnd)
+
+        # Generation parameters
+        num_blocks = 7
+        current_start_frame = 0
+        num_input_frames = 0
+        all_num_frames = [pipeline.num_frame_per_block] * num_blocks
+        if current_use_taehv:
+            vae_cache = None
+        else:
+            vae_cache = ZERO_VAE_CACHE
+            for i in range(len(vae_cache)):
+                vae_cache[i] = vae_cache[i].to(device=gpu, dtype=torch.float16)
+
+        total_frames_sent = 0
+        generation_start_time = time.time()
+
+        emit_progress('Generating frames... (frontend handles timing)', 15)
+
+        for idx, current_num_frames in enumerate(all_num_frames):
+            if not generation_active or stop_event.is_set():
+                break
+
+            progress = int(((idx + 1) / len(all_num_frames)) * 80) + 15
+
+            # Special message for first block with torch.compile
+            if idx == 0 and torch_compile_applied and not models_compiled:
+                emit_progress(
+                    f'Processing block 1/{len(all_num_frames)} - Compiling models (may take 5-10 minutes)...', progress)
+                print(f"🔥 Processing block {idx+1}/{len(all_num_frames)}")
+                models_compiled = True
+            else:
+                emit_progress(f'Processing block {idx+1}/{len(all_num_frames)}...', progress)
+                print(f"🔄 Processing block {idx+1}/{len(all_num_frames)}")
+
+            block_start_time = time.time()
+
+            noisy_input = noise[:, current_start_frame -
+                                num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # Denoising loop
+            denoising_start = time.time()
+            for index, current_timestep in enumerate(pipeline.denoising_step_list):
+                if not generation_active or stop_event.is_set():
+                    break
+
+                timestep = torch.ones([1, current_num_frames], device=noise.device,
+                                      dtype=torch.int64) * current_timestep
+
+                if index < len(pipeline.denoising_step_list) - 1:
+                    _, denoised_pred = transformer(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=pipeline.kv_cache1,
+                        crossattn_cache=pipeline.crossattn_cache,
+                        current_start=current_start_frame * pipeline.frame_seq_length
+                    )
+                    next_timestep = pipeline.denoising_step_list[index + 1]
+                    noisy_input = pipeline.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep * torch.ones([1 * current_num_frames], device=noise.device, dtype=torch.long)
+                    ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    _, denoised_pred = transformer(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=pipeline.kv_cache1,
+                        crossattn_cache=pipeline.crossattn_cache,
+                        current_start=current_start_frame * pipeline.frame_seq_length
+                    )
+
+            if not generation_active or stop_event.is_set():
+                break
+
+            denoising_time = time.time() - denoising_start
+            print(f"⚡ Block {idx+1} denoising completed in {denoising_time:.2f}s")
+
+            # Record output
+            # all_latents[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Update KV cache for next block
+            if idx != len(all_num_frames) - 1:
+                transformer(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=torch.zeros_like(timestep),
+                    kv_cache=pipeline.kv_cache1,
+                    crossattn_cache=pipeline.crossattn_cache,
+                    current_start=current_start_frame * pipeline.frame_seq_length,
+                )
+
+            # Decode to pixels and send frames immediately
+            print(f"🎨 Decoding block {idx+1} to pixels...")
+            decode_start = time.time()
+            if args.trt:
+                all_current_pixels = []
+                for i in range(denoised_pred.shape[1]):
+                    is_first_frame = torch.tensor(1.0).cuda().half() if idx == 0 and i == 0 else \
+                        torch.tensor(0.0).cuda().half()
+                    outputs = vae_decoder.forward(denoised_pred[:, i:i + 1, :, :, :].half(), is_first_frame, *vae_cache)
+                    # outputs = vae_decoder.forward(denoised_pred.float(), *vae_cache)
+                    current_pixels, vae_cache = outputs[0], outputs[1:]
+                    print(current_pixels.max(), current_pixels.min())
+                    all_current_pixels.append(current_pixels.clone())
+                pixels = torch.cat(all_current_pixels, dim=1)
+                if idx == 0:
+                    pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+            else:
+                if current_use_taehv:
+                    if vae_cache is None:
+                        vae_cache = denoised_pred
+                    else:
+                        denoised_pred = torch.cat([vae_cache, denoised_pred], dim=1)
+                        vae_cache = denoised_pred[:, -3:, :, :, :]
+                    pixels = current_vae_decoder.decode(denoised_pred)
+                    print(f"denoised_pred shape: {denoised_pred.shape}")
+                    print(f"pixels shape: {pixels.shape}")
+                    if idx == 0:
+                        pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+                    else:
+                        pixels = pixels[:, 12:, :, :, :]
+
+                else:
+                    pixels, vae_cache = current_vae_decoder(denoised_pred.half(), *vae_cache)
+                    if idx == 0:
+                        pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+
+            decode_time = time.time() - decode_start
+            print(f"🎨 Block {idx+1} VAE decoding completed in {decode_time:.2f}s")
+
+            # Queue frames for non-blocking sending
+            block_frames = pixels.shape[1]
+            print(f"📡 Queueing {block_frames} frames from block {idx+1} for sending...")
+            queue_start = time.time()
+
+            for frame_idx in range(block_frames):
+                if not generation_active or stop_event.is_set():
+                    break
+
+                frame_tensor = pixels[0, frame_idx].cpu()
+
+                # Queue frame data in non-blocking way
+                frame_send_queue.put((frame_tensor, total_frames_sent, idx, job_id))
+                total_frames_sent += 1
+
+            queue_time = time.time() - queue_start
+            block_time = time.time() - block_start_time
+            print(f"✅ Block {idx+1} completed in {block_time:.2f}s ({block_frames} frames queued in {queue_time:.3f}s)")
+
+            current_start_frame += current_num_frames
+
+        generation_time = time.time() - generation_start_time
+        print(f"🎉 Generation completed in {generation_time:.2f}s! {total_frames_sent} frames queued for sending")
+
+        # Wait for all frames to be sent before completing
+        emit_progress('Waiting for all frames to be sent...', 97)
+        print("⏳ Waiting for all frames to be sent...")
+        frame_send_queue.join()  # Wait for all queued frames to be processed
+        print("✅ All frames sent successfully!")
+
+        generate_mp4_from_images("./images","./videos/"+anim_name+".mp4", frame_rate )
+        # Final progress update
+        emit_progress('Generation complete!', 100)
+
+        try:
+            socketio.emit('generation_complete', {
+                'message': 'Video generation completed!',
+                'total_frames': total_frames_sent,
+                'generation_time': f"{generation_time:.2f}s",
+                'job_id': job_id
+            })
+        except Exception as e:
+            print(f"❌ Failed to emit generation complete: {e}")
+
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        try:
+            socketio.emit('error', {
+                'message': f'Generation failed: {str(e)}',
+                'job_id': job_id
+            })
+        except Exception as e:
+            print(f"❌ Failed to emit error: {e}")
+    finally:
+        generation_active = False
+        stop_event.set()
+
+        # Clean up sender thread
+        try:
+            frame_send_queue.put(None)
+        except Exception as e:
+            print(f"❌ Failed to put None in frame_send_queue: {e}")
+
+
+def generate_mp4_from_images(image_directory, output_video_path, fps=24):
+    """
+    Generate an MP4 video from a directory of images ordered alphabetically.
+
+    :param image_directory: Path to the directory containing images.
+    :param output_video_path: Path where the output MP4 will be saved.
+    :param fps: Frames per second for the output video.
+    """
+    global anim_name
+    # Construct the ffmpeg command
+    cmd = [
+        'ffmpeg',
+        '-framerate', str(fps),
+        '-i', os.path.join(image_directory, anim_name+'/'+anim_name+'_%03d.jpg'),  # Adjust the pattern if necessary
+        '-c:v', 'libx264',
+        '-pix_fmt', 'yuv420p',
+        output_video_path
+    ]
+    try:
+        subprocess.run(cmd, check=True)
+        print(f"Video saved to {output_video_path}")
+    except subprocess.CalledProcessError as e:
+        print(f"An error occurred: {e}")
+
+def calculate_sha256(data):
+    # Convert data to bytes if it's not already
+    if isinstance(data, str):
+        data = data.encode()
+    # Calculate SHA-256 hash
+    sha256_hash = hashlib.sha256(data).hexdigest()
+    return sha256_hash
+
+# Socket.IO event handlers
+@socketio.on('connect')
+def handle_connect():
+    print('Client connected')
+    emit('status', {'message': 'Connected to frontend-buffered demo server'})
+
+
+@socketio.on('disconnect')
+def handle_disconnect():
+    print('Client disconnected')
+
+
+@socketio.on('start_generation')
+def handle_start_generation(data):
+    global generation_active, frame_number, anim_name, frame_rate
+
+    frame_number = 0
+    if generation_active:
+        emit('error', {'message': 'Generation already in progress'})
+        return
+
+    prompt = data.get('prompt', '')
+
+    seed = data.get('seed', -1)
+    if seed==-1:
+        seed = random.randint(0, 2**32)
+
+    # Extract words up to the first punctuation or newline
+    words_up_to_punctuation = re.split(r'[^\w\s]', prompt)[0].strip() if prompt else ''
+    if not words_up_to_punctuation:
+        words_up_to_punctuation = re.split(r'[\n\r]', prompt)[0].strip()
+
+    # Calculate SHA-256 hash of the entire prompt
+    sha256_hash = calculate_sha256(prompt)
+
+    # Create anim_name with the extracted words and first 10 characters of the hash
+    anim_name = f"{words_up_to_punctuation[:20]}_{str(seed)}_{sha256_hash[:10]}"
+
+    generation_active = True
+    generation_start_time = time.time()
+    enable_torch_compile = data.get('enable_torch_compile', False)
+    enable_fp8 = data.get('enable_fp8', False)
+    use_taehv = data.get('use_taehv', False)
+    frame_rate = data.get('fps', 6)
+
+    if not prompt:
+        emit('error', {'message': 'Prompt is required'})
+        return
+
+    # Start generation in background thread
+    socketio.start_background_task(generate_video_stream, prompt, seed,
+                                   enable_torch_compile, enable_fp8, use_taehv)
+    emit('status', {'message': 'Generation started - frames will be sent immediately'})
+
+
+@socketio.on('stop_generation')
+def handle_stop_generation():
+    global generation_active, stop_event, frame_send_queue
+    generation_active = False
+    stop_event.set()
+
+    # Signal sender thread to stop (will be processed after current frames)
+    try:
+        frame_send_queue.put(None)
+    except Exception as e:
+        print(f"❌ Failed to put None in frame_send_queue: {e}")
+
+    emit('status', {'message': 'Generation stopped'})
+
+# Web routes
+
+
+@app.route('/')
+def index():
+    return render_template('demo.html')
+
+
+@app.route('/api/status')
+def api_status():
+    return jsonify({
+        'generation_active': generation_active,
+        'free_vram_gb': get_cuda_free_memory_gb(gpu),
+        'fp8_applied': fp8_applied,
+        'torch_compile_applied': torch_compile_applied,
+        'current_use_taehv': current_use_taehv
+    })
+
+
+if __name__ == '__main__':
+    print(f"🚀 Starting demo on http://{args.host}:{args.port}")
+    socketio.run(app, host=args.host, port=args.port, debug=False)

demo_utils/constant.py

@@ -0,0 +1,41 @@
+
+import torch
+
+
+ZERO_VAE_CACHE = [
+    torch.zeros(1, 16, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 192, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832)
+]
+
+feat_names = [f"vae_cache_{i}" for i in range(len(ZERO_VAE_CACHE))]
+ALL_INPUTS_NAMES = ["z", "use_cache"] + feat_names

demo_utils/memory.py

@@ -0,0 +1,135 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import torch
+
+
+cpu = torch.device('cpu')
+gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+gpu_complete_modules = []
+
+
+class DynamicSwapInstaller:
+    @staticmethod
+    def _install_module(module: torch.nn.Module, **kwargs):
+        original_class = module.__class__
+        module.__dict__['forge_backup_original_class'] = original_class
+
+        def hacked_get_attr(self, name: str):
+            if '_parameters' in self.__dict__:
+                _parameters = self.__dict__['_parameters']
+                if name in _parameters:
+                    p = _parameters[name]
+                    if p is None:
+                        return None
+                    if p.__class__ == torch.nn.Parameter:
+                        return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
+                    else:
+                        return p.to(**kwargs)
+            if '_buffers' in self.__dict__:
+                _buffers = self.__dict__['_buffers']
+                if name in _buffers:
+                    return _buffers[name].to(**kwargs)
+            return super(original_class, self).__getattr__(name)
+
+        module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
+            '__getattr__': hacked_get_attr,
+        })
+
+        return
+
+    @staticmethod
+    def _uninstall_module(module: torch.nn.Module):
+        if 'forge_backup_original_class' in module.__dict__:
+            module.__class__ = module.__dict__.pop('forge_backup_original_class')
+        return
+
+    @staticmethod
+    def install_model(model: torch.nn.Module, **kwargs):
+        for m in model.modules():
+            DynamicSwapInstaller._install_module(m, **kwargs)
+        return
+
+    @staticmethod
+    def uninstall_model(model: torch.nn.Module):
+        for m in model.modules():
+            DynamicSwapInstaller._uninstall_module(m)
+        return
+
+
+def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
+    if hasattr(model, 'scale_shift_table'):
+        model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
+        return
+
+    for k, p in model.named_modules():
+        if hasattr(p, 'weight'):
+            p.to(target_device)
+            return
+
+
+def get_cuda_free_memory_gb(device=None):
+    if device is None:
+        device = gpu
+
+    memory_stats = torch.cuda.memory_stats(device)
+    bytes_active = memory_stats['active_bytes.all.current']
+    bytes_reserved = memory_stats['reserved_bytes.all.current']
+    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+    bytes_inactive_reserved = bytes_reserved - bytes_active
+    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+    return bytes_total_available / (1024 ** 3)
+
+
+def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=target_device)
+
+    model.to(device=target_device)
+    torch.cuda.empty_cache()
+    return
+
+
+def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=cpu)
+
+    model.to(device=cpu)
+    torch.cuda.empty_cache()
+    return
+
+
+def unload_complete_models(*args):
+    for m in gpu_complete_modules + list(args):
+        m.to(device=cpu)
+        print(f'Unloaded {m.__class__.__name__} as complete.')
+
+    gpu_complete_modules.clear()
+    torch.cuda.empty_cache()
+    return
+
+
+def load_model_as_complete(model, target_device, unload=True):
+    if unload:
+        unload_complete_models()
+
+    model.to(device=target_device)
+    print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
+
+    gpu_complete_modules.append(model)
+    return

demo_utils/taehv.py

@@ -0,0 +1,313 @@
+#!/usr/bin/env python3
+"""
+Tiny AutoEncoder for Hunyuan Video
+(DNN for encoding / decoding videos to Hunyuan Video's latent space)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+from collections import namedtuple
+
+DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
+TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
+
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+
+class MemBlock(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True),
+                                  conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out))
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x, past):
+        return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
+
+
+class TPool(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False)
+
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        return self.conv(x.reshape(-1, self.stride * C, H, W))
+
+
+class TGrow(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f, n_f * stride, 1, bias=False)
+
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        x = self.conv(x)
+        return x.reshape(-1, C, H, W)
+
+
+def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input.
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - parallel: if True, parallelize over timesteps (fast but uses O(T) memory)
+        if False, each timestep will be processed sequentially (slow but uses O(1) memory)
+    - show_progress_bar: if True, enables tqdm progressbar display
+
+    Returns NTCHW tensor of output data.
+    """
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    N, T, C, H, W = x.shape
+    if parallel:
+        x = x.reshape(N * T, C, H, W)
+        # parallel over input timesteps, iterate over blocks
+        for b in tqdm(model, disable=not show_progress_bar):
+            if isinstance(b, MemBlock):
+                NT, C, H, W = x.shape
+                T = NT // N
+                _x = x.reshape(N, T, C, H, W)
+                mem = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :T].reshape(x.shape)
+                x = b(x, mem)
+            else:
+                x = b(x)
+        NT, C, H, W = x.shape
+        T = NT // N
+        x = x.view(N, T, C, H, W)
+    else:
+        # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode...
+        # need to fix :(
+        out = []
+        # iterate over input timesteps and also iterate over blocks.
+        # because of the cursed TPool/TGrow blocks, this is not a nested loop,
+        # it's actually a ***graph traversal*** problem! so let's make a queue
+        work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))]
+        # in addition to manually managing our queue, we also need to manually manage our progressbar.
+        # we'll update it for every source node that we consume.
+        progress_bar = tqdm(range(T), disable=not show_progress_bar)
+        # we'll also need a separate addressable memory per node as well
+        mem = [None] * len(model)
+        while work_queue:
+            xt, i = work_queue.pop(0)
+            if i == 0:
+                # new source node consumed
+                progress_bar.update(1)
+            if i == len(model):
+                # reached end of the graph, append result to output list
+                out.append(xt)
+            else:
+                # fetch the block to process
+                b = model[i]
+                if isinstance(b, MemBlock):
+                    # mem blocks are simple since we're visiting the graph in causal order
+                    if mem[i] is None:
+                        xt_new = b(xt, xt * 0)
+                        mem[i] = xt
+                    else:
+                        xt_new = b(xt, mem[i])
+                        mem[i].copy_(xt)  # inplace might reduce mysterious pytorch memory allocations? doesn't help though
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt_new, i + 1))
+                elif isinstance(b, TPool):
+                    # pool blocks are miserable
+                    if mem[i] is None:
+                        mem[i] = []  # pool memory is itself a queue of inputs to pool
+                    mem[i].append(xt)
+                    if len(mem[i]) > b.stride:
+                        # pool mem is in invalid state, we should have pooled before this
+                        raise ValueError("???")
+                    elif len(mem[i]) < b.stride:
+                        # pool mem is not yet full, go back to processing the work queue
+                        pass
+                    else:
+                        # pool mem is ready, run the pool block
+                        N, C, H, W = xt.shape
+                        xt = b(torch.cat(mem[i], 1).view(N * b.stride, C, H, W))
+                        # reset the pool mem
+                        mem[i] = []
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt, i + 1))
+                elif isinstance(b, TGrow):
+                    xt = b(xt)
+                    NT, C, H, W = xt.shape
+                    # each tgrow has multiple successor nodes
+                    for xt_next in reversed(xt.view(N, b.stride * C, H, W).chunk(b.stride, 1)):
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt_next, i + 1))
+                else:
+                    # normal block with no funny business
+                    xt = b(xt)
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt, i + 1))
+        progress_bar.close()
+        x = torch.stack(out, 1)
+    return x
+
+
+class TAEHV(nn.Module):
+    latent_channels = 16
+    image_channels = 3
+
+    def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)):
+        """Initialize pretrained TAEHV from the given checkpoint.
+
+        Arg:
+            checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1.
+            decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+        """
+        super().__init__()
+        self.encoder = nn.Sequential(
+            conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True),
+            TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            conv(64, TAEHV.latent_channels),
+        )
+        n_f = [256, 128, 64, 64]
+        self.frames_to_trim = 2**sum(decoder_time_upscale) - 1
+        self.decoder = nn.Sequential(
+            Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True),
+            MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(
+                scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False),
+            MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(
+                scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False),
+            MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(
+                scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False),
+            nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels),
+        )
+        if checkpoint_path is not None:
+            self.load_state_dict(self.patch_tgrow_layers(torch.load(
+                checkpoint_path, map_location="cpu", weights_only=True)))
+
+    def patch_tgrow_layers(self, sd):
+        """Patch TGrow layers to use a smaller kernel if needed.
+
+        Args:
+            sd: state dict to patch
+        """
+        new_sd = self.state_dict()
+        for i, layer in enumerate(self.decoder):
+            if isinstance(layer, TGrow):
+                key = f"decoder.{i}.conv.weight"
+                if sd[key].shape[0] > new_sd[key].shape[0]:
+                    # take the last-timestep output channels
+                    sd[key] = sd[key][-new_sd[key].shape[0]:]
+        return sd
+
+    def encode_video(self, x, parallel=True, show_progress_bar=True):
+        """Encode a sequence of frames.
+
+        Args:
+            x: input NTCHW RGB (C=3) tensor with values in [0, 1].
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW latent tensor with ~Gaussian values.
+        """
+        return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
+
+    def decode_video(self, x, parallel=True, show_progress_bar=False):
+        """Decode a sequence of frames.
+
+        Args:
+            x: input NTCHW latent (C=12) tensor with ~Gaussian values.
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW RGB tensor with ~[0, 1] values.
+        """
+        x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
+        # return x[:, self.frames_to_trim:]
+        return x
+
+    def forward(self, x):
+        return self.c(x)
+
+
+@torch.no_grad()
+def main():
+    """Run TAEHV roundtrip reconstruction on the given video paths."""
+    import os
+    import sys
+    import cv2  # no highly esteemed deed is commemorated here
+
+    class VideoTensorReader:
+        def __init__(self, video_file_path):
+            self.cap = cv2.VideoCapture(video_file_path)
+            assert self.cap.isOpened(), f"Could not load {video_file_path}"
+            self.fps = self.cap.get(cv2.CAP_PROP_FPS)
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            ret, frame = self.cap.read()
+            if not ret:
+                self.cap.release()
+                raise StopIteration  # End of video or error
+            return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1)  # BGR HWC -> RGB CHW
+
+    class VideoTensorWriter:
+        def __init__(self, video_file_path, width_height, fps=30):
+            self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height)
+            assert self.writer.isOpened(), f"Could not create writer for {video_file_path}"
+
+        def write(self, frame_tensor):
+            assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??"
+            self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(),
+                              cv2.COLOR_RGB2BGR))  # RGB CHW -> BGR HWC
+
+        def __del__(self):
+            if hasattr(self, 'writer'):
+                self.writer.release()
+
+    dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+    dtype = torch.float16
+    checkpoint_path = os.getenv("TAEHV_CHECKPOINT_PATH", "taehv.pth")
+    checkpoint_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
+    print(
+        f"Using device \033[31m{dev}\033[0m, dtype \033[32m{dtype}\033[0m, checkpoint \033[34m{checkpoint_name}\033[0m ({checkpoint_path})")
+    taehv = TAEHV(checkpoint_path=checkpoint_path).to(dev, dtype)
+    for video_path in sys.argv[1:]:
+        print(f"Processing {video_path}...")
+        video_in = VideoTensorReader(video_path)
+        video = torch.stack(list(video_in), 0)[None]
+        vid_dev = video.to(dev, dtype).div_(255.0)
+        # convert to device tensor
+        if video.numel() < 100_000_000:
+            print(f"  {video_path} seems small enough, will process all frames in parallel")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        else:
+            print(f"  {video_path} seems large, will process each frame sequentially")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev, parallel=False)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc, parallel=False)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        video_out_path = video_path + f".reconstructed_by_{checkpoint_name}.mp4"
+        video_out = VideoTensorWriter(
+            video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps)))
+        for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]:
+            video_out.write(frame)
+        print(f"  Saved to {video_out_path}")
+
+
+if __name__ == "__main__":
+    main()

demo_utils/utils.py

@@ -0,0 +1,616 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import os
+import cv2
+import json
+import random
+import glob
+import torch
+import einops
+import numpy as np
+import datetime
+import torchvision
+
+from PIL import Image
+
+
+def min_resize(x, m):
+    if x.shape[0] < x.shape[1]:
+        s0 = m
+        s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
+    else:
+        s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
+        s1 = m
+    new_max = max(s1, s0)
+    raw_max = max(x.shape[0], x.shape[1])
+    if new_max < raw_max:
+        interpolation = cv2.INTER_AREA
+    else:
+        interpolation = cv2.INTER_LANCZOS4
+    y = cv2.resize(x, (s1, s0), interpolation=interpolation)
+    return y
+
+
+def d_resize(x, y):
+    H, W, C = y.shape
+    new_min = min(H, W)
+    raw_min = min(x.shape[0], x.shape[1])
+    if new_min < raw_min:
+        interpolation = cv2.INTER_AREA
+    else:
+        interpolation = cv2.INTER_LANCZOS4
+    y = cv2.resize(x, (W, H), interpolation=interpolation)
+    return y
+
+
+def resize_and_center_crop(image, target_width, target_height):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    original_width, original_height = pil_image.size
+    scale_factor = max(target_width / original_width, target_height / original_height)
+    resized_width = int(round(original_width * scale_factor))
+    resized_height = int(round(original_height * scale_factor))
+    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
+    left = (resized_width - target_width) / 2
+    top = (resized_height - target_height) / 2
+    right = (resized_width + target_width) / 2
+    bottom = (resized_height + target_height) / 2
+    cropped_image = resized_image.crop((left, top, right, bottom))
+    return np.array(cropped_image)
+
+
+def resize_and_center_crop_pytorch(image, target_width, target_height):
+    B, C, H, W = image.shape
+
+    if H == target_height and W == target_width:
+        return image
+
+    scale_factor = max(target_width / W, target_height / H)
+    resized_width = int(round(W * scale_factor))
+    resized_height = int(round(H * scale_factor))
+
+    resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
+
+    top = (resized_height - target_height) // 2
+    left = (resized_width - target_width) // 2
+    cropped = resized[:, :, top:top + target_height, left:left + target_width]
+
+    return cropped
+
+
+def resize_without_crop(image, target_width, target_height):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
+    return np.array(resized_image)
+
+
+def just_crop(image, w, h):
+    if h == image.shape[0] and w == image.shape[1]:
+        return image
+
+    original_height, original_width = image.shape[:2]
+    k = min(original_height / h, original_width / w)
+    new_width = int(round(w * k))
+    new_height = int(round(h * k))
+    x_start = (original_width - new_width) // 2
+    y_start = (original_height - new_height) // 2
+    cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]
+    return cropped_image
+
+
+def write_to_json(data, file_path):
+    temp_file_path = file_path + ".tmp"
+    with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:
+        json.dump(data, temp_file, indent=4)
+    os.replace(temp_file_path, file_path)
+    return
+
+
+def read_from_json(file_path):
+    with open(file_path, 'rt', encoding='utf-8') as file:
+        data = json.load(file)
+    return data
+
+
+def get_active_parameters(m):
+    return {k: v for k, v in m.named_parameters() if v.requires_grad}
+
+
+def cast_training_params(m, dtype=torch.float32):
+    result = {}
+    for n, param in m.named_parameters():
+        if param.requires_grad:
+            param.data = param.to(dtype)
+            result[n] = param
+    return result
+
+
+def separate_lora_AB(parameters, B_patterns=None):
+    parameters_normal = {}
+    parameters_B = {}
+
+    if B_patterns is None:
+        B_patterns = ['.lora_B.', '__zero__']
+
+    for k, v in parameters.items():
+        if any(B_pattern in k for B_pattern in B_patterns):
+            parameters_B[k] = v
+        else:
+            parameters_normal[k] = v
+
+    return parameters_normal, parameters_B
+
+
+def set_attr_recursive(obj, attr, value):
+    attrs = attr.split(".")
+    for name in attrs[:-1]:
+        obj = getattr(obj, name)
+    setattr(obj, attrs[-1], value)
+    return
+
+
+def print_tensor_list_size(tensors):
+    total_size = 0
+    total_elements = 0
+
+    if isinstance(tensors, dict):
+        tensors = tensors.values()
+
+    for tensor in tensors:
+        total_size += tensor.nelement() * tensor.element_size()
+        total_elements += tensor.nelement()
+
+    total_size_MB = total_size / (1024 ** 2)
+    total_elements_B = total_elements / 1e9
+
+    print(f"Total number of tensors: {len(tensors)}")
+    print(f"Total size of tensors: {total_size_MB:.2f} MB")
+    print(f"Total number of parameters: {total_elements_B:.3f} billion")
+    return
+
+
+@torch.no_grad()
+def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
+    batch_size = a.size(0)
+
+    if b is None:
+        b = torch.zeros_like(a)
+
+    if mask_a is None:
+        mask_a = torch.rand(batch_size) < probability_a
+
+    mask_a = mask_a.to(a.device)
+    mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
+    result = torch.where(mask_a, a, b)
+    return result
+
+
+@torch.no_grad()
+def zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+@torch.no_grad()
+def supress_lower_channels(m, k, alpha=0.01):
+    data = m.weight.data.clone()
+
+    assert int(data.shape[1]) >= k
+
+    data[:, :k] = data[:, :k] * alpha
+    m.weight.data = data.contiguous().clone()
+    return m
+
+
+def freeze_module(m):
+    if not hasattr(m, '_forward_inside_frozen_module'):
+        m._forward_inside_frozen_module = m.forward
+    m.requires_grad_(False)
+    m.forward = torch.no_grad()(m.forward)
+    return m
+
+
+def get_latest_safetensors(folder_path):
+    safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))
+
+    if not safetensors_files:
+        raise ValueError('No file to resume!')
+
+    latest_file = max(safetensors_files, key=os.path.getmtime)
+    latest_file = os.path.abspath(os.path.realpath(latest_file))
+    return latest_file
+
+
+def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):
+    tags = tags_str.split(', ')
+    tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
+    prompt = ', '.join(tags)
+    return prompt
+
+
+def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
+    numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
+    if round_to_int:
+        numbers = np.round(numbers).astype(int)
+    return numbers.tolist()
+
+
+def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):
+    edges = np.linspace(0, 1, n + 1)
+    points = np.random.uniform(edges[:-1], edges[1:])
+    numbers = inclusive + (exclusive - inclusive) * points
+    if round_to_int:
+        numbers = np.round(numbers).astype(int)
+    return numbers.tolist()
+
+
+def soft_append_bcthw(history, current, overlap=0):
+    if overlap <= 0:
+        return torch.cat([history, current], dim=2)
+
+    assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
+    assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
+
+    weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)
+    blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
+    output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
+
+    return output.to(history)
+
+
+def save_bcthw_as_mp4(x, output_filename, fps=10, crf=0):
+    b, c, t, h, w = x.shape
+
+    per_row = b
+    for p in [6, 5, 4, 3, 2]:
+        if b % p == 0:
+            per_row = p
+            break
+
+    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
+    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
+    torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': str(int(crf))})
+    return x
+
+
+def save_bcthw_as_png(x, output_filename):
+    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
+    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
+    torchvision.io.write_png(x, output_filename)
+    return output_filename
+
+
+def save_bchw_as_png(x, output_filename):
+    os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
+    x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, 'b c h w -> c h (b w)')
+    torchvision.io.write_png(x, output_filename)
+    return output_filename
+
+
+def add_tensors_with_padding(tensor1, tensor2):
+    if tensor1.shape == tensor2.shape:
+        return tensor1 + tensor2
+
+    shape1 = tensor1.shape
+    shape2 = tensor2.shape
+
+    new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
+
+    padded_tensor1 = torch.zeros(new_shape)
+    padded_tensor2 = torch.zeros(new_shape)
+
+    padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
+    padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
+
+    result = padded_tensor1 + padded_tensor2
+    return result
+
+
+def print_free_mem():
+    torch.cuda.empty_cache()
+    free_mem, total_mem = torch.cuda.mem_get_info(0)
+    free_mem_mb = free_mem / (1024 ** 2)
+    total_mem_mb = total_mem / (1024 ** 2)
+    print(f"Free memory: {free_mem_mb:.2f} MB")
+    print(f"Total memory: {total_mem_mb:.2f} MB")
+    return
+
+
+def print_gpu_parameters(device, state_dict, log_count=1):
+    summary = {"device": device, "keys_count": len(state_dict)}
+
+    logged_params = {}
+    for i, (key, tensor) in enumerate(state_dict.items()):
+        if i >= log_count:
+            break
+        logged_params[key] = tensor.flatten()[:3].tolist()
+
+    summary["params"] = logged_params
+
+    print(str(summary))
+    return
+
+
+def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):
+    from PIL import Image, ImageDraw, ImageFont
+
+    txt = Image.new("RGB", (width, height), color="white")
+    draw = ImageDraw.Draw(txt)
+    font = ImageFont.truetype(font_path, size=size)
+
+    if text == '':
+        return np.array(txt)
+
+    # Split text into lines that fit within the image width
+    lines = []
+    words = text.split()
+    current_line = words[0]
+
+    for word in words[1:]:
+        line_with_word = f"{current_line} {word}"
+        if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
+            current_line = line_with_word
+        else:
+            lines.append(current_line)
+            current_line = word
+
+    lines.append(current_line)
+
+    # Draw the text line by line
+    y = 0
+    line_height = draw.textbbox((0, 0), "A", font=font)[3]
+
+    for line in lines:
+        if y + line_height > height:
+            break  # stop drawing if the next line will be outside the image
+        draw.text((0, y), line, fill="black", font=font)
+        y += line_height
+
+    return np.array(txt)
+
+
+def blue_mark(x):
+    x = x.copy()
+    c = x[:, :, 2]
+    b = cv2.blur(c, (9, 9))
+    x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
+    return x
+
+
+def green_mark(x):
+    x = x.copy()
+    x[:, :, 2] = -1
+    x[:, :, 0] = -1
+    return x
+
+
+def frame_mark(x):
+    x = x.copy()
+    x[:64] = -1
+    x[-64:] = -1
+    x[:, :8] = 1
+    x[:, -8:] = 1
+    return x
+
+
+@torch.inference_mode()
+def pytorch2numpy(imgs):
+    results = []
+    for x in imgs:
+        y = x.movedim(0, -1)
+        y = y * 127.5 + 127.5
+        y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
+        results.append(y)
+    return results
+
+
+@torch.inference_mode()
+def numpy2pytorch(imgs):
+    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
+    h = h.movedim(-1, 1)
+    return h
+
+
+@torch.no_grad()
+def duplicate_prefix_to_suffix(x, count, zero_out=False):
+    if zero_out:
+        return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
+    else:
+        return torch.cat([x, x[:count]], dim=0)
+
+
+def weighted_mse(a, b, weight):
+    return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
+
+
+def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
+    x = (x - x_min) / (x_max - x_min)
+    x = max(0.0, min(x, 1.0))
+    x = x ** sigma
+    return y_min + x * (y_max - y_min)
+
+
+def expand_to_dims(x, target_dims):
+    return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
+
+
+def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
+    if tensor is None:
+        return None
+
+    first_dim = tensor.shape[0]
+
+    if first_dim == batch_size:
+        return tensor
+
+    if batch_size % first_dim != 0:
+        raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")
+
+    repeat_times = batch_size // first_dim
+
+    return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
+
+
+def dim5(x):
+    return expand_to_dims(x, 5)
+
+
+def dim4(x):
+    return expand_to_dims(x, 4)
+
+
+def dim3(x):
+    return expand_to_dims(x, 3)
+
+
+def crop_or_pad_yield_mask(x, length):
+    B, F, C = x.shape
+    device = x.device
+    dtype = x.dtype
+
+    if F < length:
+        y = torch.zeros((B, length, C), dtype=dtype, device=device)
+        mask = torch.zeros((B, length), dtype=torch.bool, device=device)
+        y[:, :F, :] = x
+        mask[:, :F] = True
+        return y, mask
+
+    return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
+
+
+def extend_dim(x, dim, minimal_length, zero_pad=False):
+    original_length = int(x.shape[dim])
+
+    if original_length >= minimal_length:
+        return x
+
+    if zero_pad:
+        padding_shape = list(x.shape)
+        padding_shape[dim] = minimal_length - original_length
+        padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
+    else:
+        idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
+        last_element = x[idx]
+        padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
+
+    return torch.cat([x, padding], dim=dim)
+
+
+def lazy_positional_encoding(t, repeats=None):
+    if not isinstance(t, list):
+        t = [t]
+
+    from diffusers.models.embeddings import get_timestep_embedding
+
+    te = torch.tensor(t)
+    te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)
+
+    if repeats is None:
+        return te
+
+    te = te[:, None, :].expand(-1, repeats, -1)
+
+    return te
+
+
+def state_dict_offset_merge(A, B, C=None):
+    result = {}
+    keys = A.keys()
+
+    for key in keys:
+        A_value = A[key]
+        B_value = B[key].to(A_value)
+
+        if C is None:
+            result[key] = A_value + B_value
+        else:
+            C_value = C[key].to(A_value)
+            result[key] = A_value + B_value - C_value
+
+    return result
+
+
+def state_dict_weighted_merge(state_dicts, weights):
+    if len(state_dicts) != len(weights):
+        raise ValueError("Number of state dictionaries must match number of weights")
+
+    if not state_dicts:
+        return {}
+
+    total_weight = sum(weights)
+
+    if total_weight == 0:
+        raise ValueError("Sum of weights cannot be zero")
+
+    normalized_weights = [w / total_weight for w in weights]
+
+    keys = state_dicts[0].keys()
+    result = {}
+
+    for key in keys:
+        result[key] = state_dicts[0][key] * normalized_weights[0]
+
+        for i in range(1, len(state_dicts)):
+            state_dict_value = state_dicts[i][key].to(result[key])
+            result[key] += state_dict_value * normalized_weights[i]
+
+    return result
+
+
+def group_files_by_folder(all_files):
+    grouped_files = {}
+
+    for file in all_files:
+        folder_name = os.path.basename(os.path.dirname(file))
+        if folder_name not in grouped_files:
+            grouped_files[folder_name] = []
+        grouped_files[folder_name].append(file)
+
+    list_of_lists = list(grouped_files.values())
+    return list_of_lists
+
+
+def generate_timestamp():
+    now = datetime.datetime.now()
+    timestamp = now.strftime('%y%m%d_%H%M%S')
+    milliseconds = f"{int(now.microsecond / 1000):03d}"
+    random_number = random.randint(0, 9999)
+    return f"{timestamp}_{milliseconds}_{random_number}"
+
+
+def write_PIL_image_with_png_info(image, metadata, path):
+    from PIL.PngImagePlugin import PngInfo
+
+    png_info = PngInfo()
+    for key, value in metadata.items():
+        png_info.add_text(key, value)
+
+    image.save(path, "PNG", pnginfo=png_info)
+    return image
+
+
+def torch_safe_save(content, path):
+    torch.save(content, path + '_tmp')
+    os.replace(path + '_tmp', path)
+    return path
+
+
+def move_optimizer_to_device(optimizer, device):
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.to(device)

demo_utils/vae.py

@@ -0,0 +1,390 @@
+from typing import List
+from einops import rearrange
+import tensorrt as trt
+import torch
+import torch.nn as nn
+
+from demo_utils.constant import ALL_INPUTS_NAMES, ZERO_VAE_CACHE
+from wan.modules.vae import AttentionBlock, CausalConv3d, RMS_norm, Upsample
+
+CACHE_T = 2
+
+
+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_1, feat_cache_2):
+        h = self.shortcut(x)
+        feat_cache = feat_cache_1
+        out_feat_cache = []
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d):
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                        dim=2)
+                x = layer(x, feat_cache)
+                out_feat_cache.append(cache_x)
+                feat_cache = feat_cache_2
+            else:
+                x = layer(x)
+        return x + h, *out_feat_cache
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d')
+        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))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, is_first_frame, feat_cache):
+        if self.mode == 'upsample3d':
+            b, c, t, h, w = x.size()
+            # x, out_feat_cache = torch.cond(
+            #     is_first_frame,
+            #     lambda: (torch.cat([torch.zeros_like(x), x], dim=2), feat_cache.clone()),
+            #     lambda: self.temporal_conv(x, feat_cache),
+            # )
+            # x, out_feat_cache = torch.cond(
+            #     is_first_frame,
+            #     lambda: (torch.cat([torch.zeros_like(x), x], dim=2), feat_cache.clone()),
+            #     lambda: self.temporal_conv(x, feat_cache),
+            # )
+            x, out_feat_cache = self.temporal_conv(x, is_first_frame, feat_cache)
+            out_feat_cache = torch.cond(
+                is_first_frame,
+                lambda: feat_cache.clone().contiguous(),
+                lambda: out_feat_cache.clone().contiguous(),
+            )
+            # if is_first_frame:
+            #     x = torch.cat([torch.zeros_like(x), x], dim=2)
+            #     out_feat_cache = feat_cache.clone()
+            # else:
+            #     x, out_feat_cache = self.temporal_conv(x, feat_cache)
+        else:
+            out_feat_cache = None
+        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)
+        return x, out_feat_cache
+
+    def temporal_conv(self, x, is_first_frame, feat_cache):
+        b, c, t, h, w = x.size()
+        cache_x = x[:, :, -CACHE_T:, :, :].clone()
+        if cache_x.shape[2] < 2 and feat_cache is not None:
+            cache_x = torch.cat([
+                torch.zeros_like(cache_x),
+                cache_x
+            ], dim=2)
+        x = torch.cond(
+            is_first_frame,
+            lambda: torch.cat([torch.zeros_like(x), x], dim=1).contiguous(),
+            lambda: self.time_conv(x, feat_cache).contiguous(),
+        )
+        # x = self.time_conv(x, feat_cache)
+        out_feat_cache = cache_x
+
+        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)
+        return x.contiguous(), out_feat_cache.contiguous()
+
+    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 VAEDecoderWrapperSingle(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = VAEDecoder3d()
+        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)
+        self.z_dim = 16
+        self.conv2 = CausalConv3d(self.z_dim, self.z_dim, 1)
+
+    def forward(
+            self,
+            z: torch.Tensor,
+            is_first_frame: torch.Tensor,
+            *feat_cache: List[torch.Tensor]
+    ):
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        z = z.permute(0, 2, 1, 3, 4)
+        assert z.shape[2] == 1
+        feat_cache = list(feat_cache)
+        is_first_frame = is_first_frame.bool()
+
+        device, dtype = z.device, z.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        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]
+        x = self.conv2(z)
+        out, feat_cache = self.decoder(x, is_first_frame, feat_cache=feat_cache)
+        out = out.clamp_(-1, 1)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        out = out.permute(0, 2, 1, 3, 4)
+        return out, feat_cache
+
+
+class VAEDecoder3d(nn.Module):
+    def __init__(self,
+                 dim=96,
+                 z_dim=16,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_upsample=[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_upsample = temperal_upsample
+        self.cache_t = 2
+        self.decoder_conv_num = 32
+
+        # 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: torch.Tensor,
+            is_first_frame: torch.Tensor,
+            feat_cache: List[torch.Tensor]
+    ):
+        idx = 0
+        out_feat_cache = []
+
+        # conv1
+        cache_x = x[:, :, -self.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])
+        out_feat_cache.append(cache_x)
+        idx += 1
+
+        # middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x, out_feat_cache_1, out_feat_cache_2 = layer(x, feat_cache[idx], feat_cache[idx + 1])
+                idx += 2
+                out_feat_cache.append(out_feat_cache_1)
+                out_feat_cache.append(out_feat_cache_2)
+            else:
+                x = layer(x)
+
+        # upsamples
+        for layer in self.upsamples:
+            if isinstance(layer, Resample):
+                x, cache_x = layer(x, is_first_frame, feat_cache[idx])
+                if cache_x is not None:
+                    out_feat_cache.append(cache_x)
+                    idx += 1
+            else:
+                x, out_feat_cache_1, out_feat_cache_2 = layer(x, feat_cache[idx], feat_cache[idx + 1])
+                idx += 2
+                out_feat_cache.append(out_feat_cache_1)
+                out_feat_cache.append(out_feat_cache_2)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                cache_x = x[:, :, -self.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])
+                out_feat_cache.append(cache_x)
+                idx += 1
+            else:
+                x = layer(x)
+        return x, out_feat_cache
+
+
+class VAETRTWrapper():
+    def __init__(self):
+        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+        with open("checkpoints/vae_decoder_int8.trt", "rb") as f, trt.Runtime(TRT_LOGGER) as rt:
+            self.engine: trt.ICudaEngine = rt.deserialize_cuda_engine(f.read())
+
+        self.context: trt.IExecutionContext = self.engine.create_execution_context()
+        self.stream = torch.cuda.current_stream().cuda_stream
+
+        # ──────────────────────────────
+        # 2️⃣  Feed the engine with tensors
+        #     (name-based API in TRT ≥10)
+        # ──────────────────────────────
+        self.dtype_map = {
+            trt.float32: torch.float32,
+            trt.float16: torch.float16,
+            trt.int8: torch.int8,
+            trt.int32: torch.int32,
+        }
+        test_input = torch.zeros(1, 16, 1, 60, 104).cuda().half()
+        is_first_frame = torch.tensor(1.0).cuda().half()
+        test_cache_inputs = [c.cuda().half() for c in ZERO_VAE_CACHE]
+        test_inputs = [test_input, is_first_frame] + test_cache_inputs
+
+        # keep references so buffers stay alive
+        self.device_buffers, self.outputs = {}, []
+
+        # ---- inputs ----
+        for i, name in enumerate(ALL_INPUTS_NAMES):
+            tensor, scale = test_inputs[i], 1 / 127
+            tensor = self.quantize_if_needed(tensor, self.engine.get_tensor_dtype(name), scale)
+
+            # dynamic shapes
+            if -1 in self.engine.get_tensor_shape(name):
+                # new API :contentReference[oaicite:0]{index=0}
+                self.context.set_input_shape(name, tuple(tensor.shape))
+
+            # replaces bindings[] :contentReference[oaicite:1]{index=1}
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+            self.device_buffers[name] = tensor                             # keep pointer alive
+
+        # ---- (after all input shapes are known) infer output shapes ----
+        # propagates shapes :contentReference[oaicite:2]{index=2}
+        self.context.infer_shapes()
+
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            # replaces binding_is_input :contentReference[oaicite:3]{index=3}
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
+                shape = tuple(self.context.get_tensor_shape(name))
+                dtype = self.dtype_map[self.engine.get_tensor_dtype(name)]
+                out = torch.empty(shape, dtype=dtype, device="cuda").contiguous()
+
+                self.context.set_tensor_address(name, int(out.data_ptr()))
+                self.outputs.append(out)
+                self.device_buffers[name] = out
+
+    # helper to quant-convert on the fly
+    def quantize_if_needed(self, t, expected_dtype, scale):
+        if expected_dtype == trt.int8 and t.dtype != torch.int8:
+            t = torch.clamp((t / scale).round(), -128, 127).to(torch.int8).contiguous()
+        return t                            # keep pointer alive
+
+    def forward(self, *test_inputs):
+        for i, name in enumerate(ALL_INPUTS_NAMES):
+            tensor, scale = test_inputs[i], 1 / 127
+            tensor = self.quantize_if_needed(tensor, self.engine.get_tensor_dtype(name), scale)
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+            self.device_buffers[name] = tensor
+
+        self.context.execute_async_v3(stream_handle=self.stream)
+        torch.cuda.current_stream().synchronize()
+        return self.outputs

demo_utils/vae_block3.py

@@ -0,0 +1,291 @@
+from typing import List
+from einops import rearrange
+import torch
+import torch.nn as nn
+
+from wan.modules.vae import AttentionBlock, CausalConv3d, RMS_norm, ResidualBlock, Upsample
+
+
+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
+        self.cache_t = 2
+
+        # 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[:, :, -self.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 VAEDecoderWrapper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = VAEDecoder3d()
+        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)
+        self.z_dim = 16
+        self.conv2 = CausalConv3d(self.z_dim, self.z_dim, 1)
+
+    def forward(
+            self,
+            z: torch.Tensor,
+            *feat_cache: List[torch.Tensor]
+    ):
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        z = z.permute(0, 2, 1, 3, 4)
+        feat_cache = list(feat_cache)
+        print("Length of feat_cache: ", len(feat_cache))
+
+        device, dtype = z.device, z.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        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_):
+            if i == 0:
+                out, feat_cache = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=feat_cache)
+            else:
+                out_, feat_cache = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=feat_cache)
+                out = torch.cat([out, out_], 2)
+
+        out = out.float().clamp_(-1, 1)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        out = out.permute(0, 2, 1, 3, 4)
+        return out, feat_cache
+
+
+class VAEDecoder3d(nn.Module):
+    def __init__(self,
+                 dim=96,
+                 z_dim=16,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_upsample=[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_upsample = temperal_upsample
+        self.cache_t = 2
+        self.decoder_conv_num = 32
+
+        # 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: torch.Tensor,
+            feat_cache: List[torch.Tensor]
+    ):
+        feat_idx = [0]
+
+        # conv1
+        idx = feat_idx[0]
+        cache_x = x[:, :, -self.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
+
+        # 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:
+            x = layer(x, feat_cache, feat_idx)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -self.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, feat_cache

demo_utils/vae_torch2trt.py

@@ -0,0 +1,308 @@
+# ---- INT8 (optional) ----
+from demo_utils.vae import (
+    VAEDecoderWrapperSingle,                         # main nn.Module
+    ZERO_VAE_CACHE           # helper constants shipped with your code base
+)
+import pycuda.driver as cuda          # ← add
+import pycuda.autoinit  # noqa
+
+import sys
+from pathlib import Path
+
+import torch
+import tensorrt as trt
+
+from utils.dataset import ShardingLMDBDataset
+
+data_path = "/mnt/localssd/wanx_14B_shift-3.0_cfg-5.0_lmdb_oneshard"
+dataset = ShardingLMDBDataset(data_path, max_pair=int(1e8))
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    batch_size=1,
+    num_workers=0
+)
+
+# ─────────────────────────────────────────────────────────
+# 1️⃣  Bring the PyTorch model into scope
+#     (all code you pasted lives in `vae_decoder.py`)
+# ─────────────────────────────────────────────────────────
+
+# --- dummy tensors (exact shapes you posted) ---
+dummy_input = torch.randn(1, 1, 16, 60, 104).half().cuda()
+is_first_frame = torch.tensor([1.0], device="cuda", dtype=torch.float16)
+dummy_cache_input = [
+    torch.randn(*s.shape).half().cuda() if isinstance(s, torch.Tensor) else s
+    for s in ZERO_VAE_CACHE               # keep exactly the same ordering
+]
+inputs = [dummy_input, is_first_frame, *dummy_cache_input]
+
+# ─────────────────────────────────────────────────────────
+# 2️⃣  Export → ONNX
+# ─────────────────────────────────────────────────────────
+model = VAEDecoderWrapperSingle().half().cuda().eval()
+
+vae_state_dict = torch.load('wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth', map_location="cpu")
+decoder_state_dict = {}
+for key, value in vae_state_dict.items():
+    if 'decoder.' in key or 'conv2' in key:
+        decoder_state_dict[key] = value
+model.load_state_dict(decoder_state_dict)
+model = model.half().cuda().eval()                          # only batch dim dynamic
+
+onnx_path = Path("vae_decoder.onnx")
+feat_names = [f"vae_cache_{i}" for i in range(len(dummy_cache_input))]
+all_inputs_names = ["z", "use_cache"] + feat_names
+
+with torch.inference_mode():
+    torch.onnx.export(
+        model,
+        tuple(inputs),                                        # must be a tuple
+        onnx_path.as_posix(),
+        input_names=all_inputs_names,
+        output_names=["rgb_out", "cache_out"],
+        opset_version=17,
+        do_constant_folding=True,
+        dynamo=True
+    )
+print(f"✅  ONNX graph saved to {onnx_path.resolve()}")
+
+# (Optional) quick sanity-check with ONNX-Runtime
+try:
+    import onnxruntime as ort
+    sess = ort.InferenceSession(onnx_path.as_posix(),
+                                providers=["CUDAExecutionProvider"])
+    ort_inputs = {n: t.cpu().numpy() for n, t in zip(all_inputs_names, inputs)}
+    _ = sess.run(None, ort_inputs)
+    print("✅  ONNX graph is executable")
+except Exception as e:
+    print("⚠️  ONNX check failed:", e)
+
+# ─────────────────────────────────────────────────────────
+# 3️⃣  Build the TensorRT engine
+# ─────────────────────────────────────────────────────────
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+builder = trt.Builder(TRT_LOGGER)
+network = builder.create_network(
+    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+parser = trt.OnnxParser(network, TRT_LOGGER)
+
+with open(onnx_path, "rb") as f:
+    if not parser.parse(f.read()):
+        for i in range(parser.num_errors):
+            print(parser.get_error(i))
+        sys.exit("❌  ONNX → TRT parsing failed")
+
+config = builder.create_builder_config()
+
+
+def set_workspace(config, bytes_):
+    """Version-agnostic workspace limit."""
+    if hasattr(config, "max_workspace_size"):                # TRT 8 / 9
+        config.max_workspace_size = bytes_
+    else:                                                    # TRT 10+
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, bytes_)
+
+
+# …
+config = builder.create_builder_config()
+set_workspace(config, 4 << 30)          # 4 GB
+# 4 GB
+
+if builder.platform_has_fast_fp16:
+    config.set_flag(trt.BuilderFlag.FP16)
+
+# ---- INT8 (optional) ----
+# provide a calibrator if you need an INT8 engine; comment this
+# block if you only care about FP16.
+# ─────────────────────────────────────────────────────────
+# helper: version-agnostic workspace limit
+# ─────────────────────────────────────────────────────────
+
+
+def set_workspace(config: trt.IBuilderConfig, bytes_: int = 4 << 30):
+    """
+    TRT < 10.x  →  config.max_workspace_size
+    TRT ≥ 10.x  →  config.set_memory_pool_limit(...)
+    """
+    if hasattr(config, "max_workspace_size"):                     # TRT 8 / 9
+        config.max_workspace_size = bytes_
+    else:                                                         # TRT 10+
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,
+                                     bytes_)
+
+# ─────────────────────────────────────────────────────────
+# (optional) INT-8 calibrator
+# ─────────────────────────────────────────────────────────
+# ‼ Only keep this block if you really need INT-8 ‼                      # gracefully skip if PyCUDA not present
+
+
+class VAECalibrator(trt.IInt8EntropyCalibrator2):
+    def __init__(self, loader, cache="calibration.cache", max_batches=10):
+        super().__init__()
+        self.loader = iter(loader)
+        self.batch_size = loader.batch_size or 1
+        self.max_batches = max_batches
+        self.count = 0
+        self.cache_file = cache
+        self.stream = cuda.Stream()
+        self.dev_ptrs = {}
+
+    # --- TRT 10 needs BOTH spellings ---
+    def get_batch_size(self):
+        return self.batch_size
+
+    def getBatchSize(self):
+        return self.batch_size
+
+    def get_batch(self, names):
+        if self.count >= self.max_batches:
+            return None
+
+        # Randomly sample a number from 1 to 10
+        import random
+        vae_idx = random.randint(0, 10)
+        data = next(self.loader)
+
+        latent = data['ode_latent'][0][:, :1]
+        is_first_frame = torch.tensor([1.0], device="cuda", dtype=torch.float16)
+        feat_cache = ZERO_VAE_CACHE
+        for i in range(vae_idx):
+            inputs = [latent, is_first_frame, *feat_cache]
+            with torch.inference_mode():
+                outputs = model(*inputs)
+            latent = data['ode_latent'][0][:, i + 1:i + 2]
+            is_first_frame = torch.tensor([0.0], device="cuda", dtype=torch.float16)
+            feat_cache = outputs[1:]
+
+        # -------- ensure context is current --------
+        z_np = latent.cpu().numpy().astype('float32')
+
+        ptrs = []                # list[int] – one entry per name
+        for name in names:         # <-- match TRT's binding order
+            if name == "z":
+                arr = z_np
+            elif name == "use_cache":
+                arr = is_first_frame.cpu().numpy().astype('float32')
+            else:
+                idx = int(name.split('_')[-1])   # "vae_cache_17" -> 17
+                arr = feat_cache[idx].cpu().numpy().astype('float32')
+
+            if name not in self.dev_ptrs:
+                self.dev_ptrs[name] = cuda.mem_alloc(arr.nbytes)
+
+            cuda.memcpy_htod_async(self.dev_ptrs[name], arr, self.stream)
+            ptrs.append(int(self.dev_ptrs[name]))   # ***int() is required***
+
+        self.stream.synchronize()
+        self.count += 1
+        print(f"Calibration batch {self.count}/{self.max_batches}")
+        return ptrs
+
+    # --- calibration-cache helpers (both spellings) ---
+    def read_calibration_cache(self):
+        try:
+            with open(self.cache_file, "rb") as f:
+                return f.read()
+        except FileNotFoundError:
+            return None
+
+    def readCalibrationCache(self):
+        return self.read_calibration_cache()
+
+    def write_calibration_cache(self, cache):
+        with open(self.cache_file, "wb") as f:
+            f.write(cache)
+
+    def writeCalibrationCache(self, cache):
+        self.write_calibration_cache(cache)
+
+
+# ─────────────────────────────────────────────────────────
+# Builder-config + optimisation profile
+# ─────────────────────────────────────────────────────────
+config = builder.create_builder_config()
+set_workspace(config, 4 << 30)                    # 4 GB
+
+# ► enable FP16 if possible
+if builder.platform_has_fast_fp16:
+    config.set_flag(trt.BuilderFlag.FP16)
+
+# ► enable INT-8  (delete this block if you don’t need it)
+if cuda is not None:
+    config.set_flag(trt.BuilderFlag.INT8)
+    # supply any representative batch you like – here we reuse the latent z
+    calib = VAECalibrator(dataloader)
+    # TRT-10 renamed the setter:
+    if hasattr(config, "set_int8_calibrator"):    # TRT 10+
+        config.set_int8_calibrator(calib)
+    else:                                         # TRT ≤ 9
+        config.int8_calibrator = calib
+
+# ---- optimisation profile ----
+profile = builder.create_optimization_profile()
+profile.set_shape(all_inputs_names[0],            # latent z
+                  min=(1, 1, 16, 60, 104),
+                  opt=(1, 1, 16, 60, 104),
+                  max=(1, 1, 16, 60, 104))
+profile.set_shape("use_cache",               # scalar flag
+                  min=(1,), opt=(1,), max=(1,))
+for name, tensor in zip(all_inputs_names[2:], dummy_cache_input):
+    profile.set_shape(name, tensor.shape, tensor.shape, tensor.shape)
+
+config.add_optimization_profile(profile)
+
+# ─────────────────────────────────────────────────────────
+# Build the engine  (API changed in TRT-10)
+# ─────────────────────────────────────────────────────────
+print("⚙️  Building engine … (can take a minute)")
+
+if hasattr(builder, "build_serialized_network"):          # TRT 10+
+    serialized_engine = builder.build_serialized_network(network, config)
+    assert serialized_engine is not None, "build_serialized_network() failed"
+    plan_path = Path("checkpoints/vae_decoder_int8.trt")
+    plan_path.write_bytes(serialized_engine)
+    engine_bytes = serialized_engine                      # keep for smoke-test
+else:                                                     # TRT ≤ 9
+    engine = builder.build_engine(network, config)
+    assert engine is not None, "build_engine() returned None"
+    plan_path = Path("checkpoints/vae_decoder_int8.trt")
+    plan_path.write_bytes(engine.serialize())
+    engine_bytes = engine.serialize()
+
+print(f"✅  TensorRT engine written to {plan_path.resolve()}")
+
+# ─────────────────────────────────────────────────────────
+# 4️⃣  Quick smoke-test with the brand-new engine
+# ─────────────────────────────────────────────────────────
+with trt.Runtime(TRT_LOGGER) as rt:
+    engine = rt.deserialize_cuda_engine(engine_bytes)
+    context = engine.create_execution_context()
+    stream = torch.cuda.current_stream().cuda_stream
+
+    # pre-allocate device buffers once
+    device_buffers, outputs = {}, []
+    dtype_map = {trt.float32: torch.float32,
+                 trt.float16: torch.float16,
+                 trt.int8:    torch.int8,
+                 trt.int32:   torch.int32}
+
+    for name, tensor in zip(all_inputs_names, inputs):
+        if -1 in engine.get_tensor_shape(name):            # dynamic input
+            context.set_input_shape(name, tensor.shape)
+        context.set_tensor_address(name, int(tensor.data_ptr()))
+        device_buffers[name] = tensor
+
+    context.infer_shapes()                                 # propagate ⇢ outputs
+    for i in range(engine.num_io_tensors):
+        name = engine.get_tensor_name(i)
+        if engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
+            shape = tuple(context.get_tensor_shape(name))
+            dtype = dtype_map[engine.get_tensor_dtype(name)]
+            out = torch.empty(shape, dtype=dtype, device="cuda")
+            context.set_tensor_address(name, int(out.data_ptr()))
+            outputs.append(out)
+            print(f"output {name} shape: {shape}")
+
+    context.execute_async_v3(stream_handle=stream)
+    torch.cuda.current_stream().synchronize()
+    print("✅  TRT execution OK – first output shape:", outputs[0].shape)

get_causal_ode_data_chunkwise.py

@@ -0,0 +1,140 @@
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+from utils.scheduler import FlowMatchScheduler
+from utils.distributed import launch_distributed_job
+
+import torch.distributed as dist
+from tqdm import tqdm
+import argparse
+import torch
+import math
+import os
+from utils.dataset import LatentLMDBDataset
+
+def init_model(device):
+    model = WanDiffusionWrapper(is_causal=True).to(device).to(torch.float32)
+    model.model.num_frame_per_block = 3 # !!
+    encoder = WanTextEncoder().to(device).to(torch.float32)
+    
+
+    scheduler = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
+    scheduler.set_timesteps(num_inference_steps=48, denoising_strength=1.0)
+    scheduler.sigmas = scheduler.sigmas.to(device)
+
+    sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+
+    unconditional_dict = encoder(
+        text_prompts=[sample_neg_prompt]
+    )
+
+    return model, encoder, scheduler, unconditional_dict
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=-1)
+    parser.add_argument("--output_folder", type=str)
+    parser.add_argument("--rawdata_path", type=str)
+    parser.add_argument("--generator_ckpt", type=str)
+    parser.add_argument("--guidance_scale", type=float, default=6.0)
+
+
+    args = parser.parse_args()
+
+    launch_distributed_job()
+    global_rank = dist.get_rank()
+
+    device = torch.cuda.current_device()
+
+    torch.set_grad_enabled(False)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    model, encoder, scheduler, unconditional_dict = init_model(device=device)
+    state_dict = torch.load(args.generator_ckpt, map_location="cpu")
+        
+    gen_sd = state_dict["generator"]
+    fixed = {}
+    for k, v in gen_sd.items():
+        if k.startswith("model._fsdp_wrapped_module."):
+            k = k.replace("model._fsdp_wrapped_module.", "", 1)
+        if k.startswith("model."):
+            k = k.replace("model.", "", 1)
+        fixed[k] = v
+    state_dict = fixed
+    model.model.load_state_dict(
+        state_dict, strict=True
+    )
+
+
+
+    dataset = LatentLMDBDataset(args.rawdata_path)
+
+    if global_rank == 0:
+        os.makedirs(args.output_folder, exist_ok=True)
+        
+    total_steps = int(math.ceil(len(dataset) / dist.get_world_size()))
+    for index in tqdm(
+        range(total_steps), disable=(dist.get_rank() != 0),
+    ):
+        prompt_index = index * dist.get_world_size() + dist.get_rank()
+        if prompt_index >= len(dataset):
+            continue
+        sample = dataset[prompt_index]
+        prompt = sample["prompts"]
+       
+        clean_latent = sample["clean_latent"].to(device).unsqueeze(0)
+        
+        
+
+        conditional_dict = encoder(
+            text_prompts=prompt
+        )
+
+        latents = torch.randn(
+            [1, 21, 16, 60, 104], dtype=torch.float32, device=device
+        )
+        
+        noisy_input = []
+
+        for progress_id, t in enumerate(tqdm(scheduler.timesteps, disable=(dist.get_rank() != 0))):
+            timestep = t * \
+                torch.ones([1, 21], device=device, dtype=torch.float32)
+            noisy_input.append(latents)
+            f_cond, x0_pred_cond = model(
+                latents, conditional_dict, timestep, clean_x = clean_latent
+            )
+
+            f_uncond, x0_pred_uncond = model(
+                latents, unconditional_dict, timestep, clean_x = clean_latent
+            )
+
+            flow_pred = f_uncond + args.guidance_scale * (
+                f_cond - f_uncond
+            )
+            
+            
+            latents = scheduler.step(
+                flow_pred.flatten(0, 1),
+                timestep.flatten(0, 1),
+                latents.flatten(0, 1)
+            ).unflatten(dim=0, sizes=flow_pred.shape[:2])
+
+        noisy_input.append(latents)
+        noisy_input.append(clean_latent)
+        
+        noisy_inputs = torch.stack(noisy_input, dim=1)
+
+        noisy_inputs = noisy_inputs[:, [0, 12, 24, 36, -2, -1]]
+
+        stored_data = noisy_inputs
+
+        torch.save(
+            {prompt: stored_data.cpu().detach()},
+            os.path.join(args.output_folder, f"{prompt_index:05d}.pt")
+        )
+
+    dist.barrier()
+
+
+if __name__ == "__main__":
+    main()

get_causal_ode_data_framewise.py

@@ -0,0 +1,140 @@
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+from utils.scheduler import FlowMatchScheduler
+from utils.distributed import launch_distributed_job
+
+import torch.distributed as dist
+from tqdm import tqdm
+import argparse
+import torch
+import math
+import os
+from utils.dataset import LatentLMDBDataset
+
+def init_model(device):
+    model = WanDiffusionWrapper(is_causal=True).to(device).to(torch.float32)
+    model.model.num_frame_per_block = 1 # !!
+    encoder = WanTextEncoder().to(device).to(torch.float32)
+    
+
+    scheduler = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
+    scheduler.set_timesteps(num_inference_steps=48, denoising_strength=1.0)
+    scheduler.sigmas = scheduler.sigmas.to(device)
+
+    sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+
+    unconditional_dict = encoder(
+        text_prompts=[sample_neg_prompt]
+    )
+
+    return model, encoder, scheduler, unconditional_dict
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=-1)
+    parser.add_argument("--output_folder", type=str)
+    parser.add_argument("--rawdata_path", type=str)
+    parser.add_argument("--generator_ckpt", type=str)
+    parser.add_argument("--guidance_scale", type=float, default=6.0)
+
+
+    args = parser.parse_args()
+
+    launch_distributed_job()
+    global_rank = dist.get_rank()
+
+    device = torch.cuda.current_device()
+
+    torch.set_grad_enabled(False)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    model, encoder, scheduler, unconditional_dict = init_model(device=device)
+    state_dict = torch.load(args.generator_ckpt, map_location="cpu")
+        
+    gen_sd = state_dict["generator"]
+    fixed = {}
+    for k, v in gen_sd.items():
+        if k.startswith("model._fsdp_wrapped_module."):
+            k = k.replace("model._fsdp_wrapped_module.", "", 1)
+        if k.startswith("model."):
+            k = k.replace("model.", "", 1)
+        fixed[k] = v
+    state_dict = fixed
+    model.model.load_state_dict(
+        state_dict, strict=True
+    )
+
+
+
+    dataset = LatentLMDBDataset(args.rawdata_path)
+
+    if global_rank == 0:
+        os.makedirs(args.output_folder, exist_ok=True)
+        
+    total_steps = int(math.ceil(len(dataset) / dist.get_world_size()))
+    for index in tqdm(
+        range(total_steps), disable=(dist.get_rank() != 0),
+    ):
+        prompt_index = index * dist.get_world_size() + dist.get_rank()
+        if prompt_index >= len(dataset):
+            continue
+        sample = dataset[prompt_index]
+        prompt = sample["prompts"]
+       
+        clean_latent = sample["clean_latent"].to(device).unsqueeze(0)
+        
+        
+
+        conditional_dict = encoder(
+            text_prompts=prompt
+        )
+
+        latents = torch.randn(
+            [1, 21, 16, 60, 104], dtype=torch.float32, device=device
+        )
+        
+        noisy_input = []
+
+        for progress_id, t in enumerate(tqdm(scheduler.timesteps, disable=(dist.get_rank() != 0))):
+            timestep = t * \
+                torch.ones([1, 21], device=device, dtype=torch.float32)
+            noisy_input.append(latents)
+            f_cond, x0_pred_cond = model(
+                latents, conditional_dict, timestep, clean_x = clean_latent
+            )
+
+            f_uncond, x0_pred_uncond = model(
+                latents, unconditional_dict, timestep, clean_x = clean_latent
+            )
+
+            flow_pred = f_uncond + args.guidance_scale * (
+                f_cond - f_uncond
+            )
+            
+            
+            latents = scheduler.step(
+                flow_pred.flatten(0, 1),
+                timestep.flatten(0, 1),
+                latents.flatten(0, 1)
+            ).unflatten(dim=0, sizes=flow_pred.shape[:2])
+
+        noisy_input.append(latents)
+        noisy_input.append(clean_latent)
+        
+        noisy_inputs = torch.stack(noisy_input, dim=1)
+
+        noisy_inputs = noisy_inputs[:, [0, 12, 24, 36, -2, -1]]
+
+        stored_data = noisy_inputs
+
+        torch.save(
+            {prompt: stored_data.cpu().detach()},
+            os.path.join(args.output_folder, f"{prompt_index:05d}.pt")
+        )
+
+    dist.barrier()
+
+
+if __name__ == "__main__":
+    main()

get_causal_ode_data_kv_optimized.py

@@ -0,0 +1,188 @@
+import argparse
+import math
+import os
+
+import torch
+import torch.distributed as dist
+from tqdm import tqdm
+
+from utils.dataset import LatentLMDBDataset
+from utils.distributed import launch_distributed_job
+from utils.ode_generation import (
+    CausalODETrajectoryGenerator,
+    merge_cfg_prompt_embeds,
+)
+from utils.scheduler import FlowMatchScheduler
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder
+
+
+DEFAULT_NEGATIVE_PROMPT = (
+    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，"
+    "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，"
+    "画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，"
+    "杂乱的背景，三条腿，背景人很多，倒着走"
+)
+DEFAULT_NUM_INFERENCE_STEPS = 48
+DEFAULT_SCHEDULER_SHIFT = 5.0
+DEFAULT_TARGET_NUM_FRAMES = 21
+DEFAULT_TRAJECTORY_INDICES = [0, 12, 24, 36, -2, -1]
+
+
+def normalize_generator_state_dict(state_dict: dict) -> dict:
+    if "generator" in state_dict:
+        state_dict = state_dict["generator"]
+    elif "generator_ema" in state_dict:
+        state_dict = state_dict["generator_ema"]
+
+    fixed = {}
+    for k, v in state_dict.items():
+        if k.startswith("model._fsdp_wrapped_module."):
+            k = k.replace("model._fsdp_wrapped_module.", "", 1)
+        if k.startswith("model."):
+            k = k.replace("model.", "", 1)
+        fixed[k] = v
+    return fixed
+
+
+def init_model(
+    device,
+    num_frame_per_block: int,
+    scheduler_shift: float,
+    num_inference_steps: int,
+    negative_prompt: str,
+):
+    model = WanDiffusionWrapper(is_causal=True).to(device).to(torch.float32)
+    model.model.num_frame_per_block = num_frame_per_block
+    encoder = WanTextEncoder().to(device).to(torch.float32)
+
+    scheduler = FlowMatchScheduler(
+        shift=scheduler_shift,
+        sigma_min=0.0,
+        extra_one_step=True,
+    )
+    scheduler.set_timesteps(
+        num_inference_steps=num_inference_steps,
+        denoising_strength=1.0,
+    )
+    scheduler.sigmas = scheduler.sigmas.to(device)
+
+    unconditional_dict = encoder(text_prompts=[negative_prompt])
+    return model, encoder, scheduler, unconditional_dict
+
+
+def prepare_clean_latent(
+    sample: dict,
+    target_num_frames: int | None,
+    device,
+) -> torch.Tensor:
+    clean_latent = sample["clean_latent"].to(device).unsqueeze(0)
+    if target_num_frames is None:
+        return clean_latent
+
+    if clean_latent.shape[1] < target_num_frames:
+        raise ValueError(
+            "clean_latent has fewer frames than requested: "
+            f"{clean_latent.shape[1]} < {target_num_frames}"
+        )
+    if clean_latent.shape[1] != target_num_frames:
+        clean_latent = clean_latent[:, :target_num_frames, ...]
+    return clean_latent.contiguous()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=-1)
+    parser.add_argument("--output_folder", type=str, required=True)
+    parser.add_argument("--rawdata_path", type=str, required=True)
+    parser.add_argument("--generator_ckpt", type=str, required=True)
+    parser.add_argument("--num_frames_per_chunk", type=int, required=True)
+    parser.add_argument("--guidance_scale", type=float, default=6.0)
+    parser.add_argument(
+        "--generation_mode",
+        type=str,
+        default="full",
+        choices=["full", "blockwise_kv"],
+    )
+
+    args = parser.parse_args()
+
+    launch_distributed_job()
+    global_rank = dist.get_rank()
+
+    device = torch.cuda.current_device()
+
+    torch.set_grad_enabled(False)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    model, encoder, scheduler, unconditional_dict = init_model(
+        device=device,
+        num_frame_per_block=args.num_frame_per_block,
+        scheduler_shift=DEFAULT_SCHEDULER_SHIFT,
+        num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS,
+        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+    )
+    state_dict = torch.load(args.generator_ckpt, map_location="cpu")
+    model.model.load_state_dict(
+        normalize_generator_state_dict(state_dict),
+        strict=True,
+    )
+
+    dataset = LatentLMDBDataset(
+        args.rawdata_path,
+        max_pair=int(1e8),
+    )
+
+    if global_rank == 0:
+        os.makedirs(args.output_folder, exist_ok=True)
+
+    trajectory_generator = CausalODETrajectoryGenerator(
+        model=model,
+        scheduler=scheduler,
+        num_frame_per_block=args.num_frame_per_block,
+        num_inference_steps=DEFAULT_NUM_INFERENCE_STEPS,
+        guidance_scale=args.guidance_scale,
+    )
+
+    total_steps = int(math.ceil(len(dataset) / dist.get_world_size()))
+    for index in tqdm(
+        range(total_steps), disable=(global_rank != 0),
+    ):
+        prompt_index = index * dist.get_world_size() + global_rank
+        if prompt_index >= len(dataset):
+            continue
+
+        output_path = os.path.join(args.output_folder, f"{prompt_index:05d}.pt")
+
+        sample = dataset[prompt_index]
+        prompt = sample["prompts"]
+        clean_latent = prepare_clean_latent(
+            sample=sample,
+            target_num_frames=DEFAULT_TARGET_NUM_FRAMES,
+            device=device,
+        )
+
+        conditional_dict = encoder(text_prompts=[prompt])
+        paired_conditional_dict = merge_cfg_prompt_embeds(
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+        )
+        initial_noise = torch.randn_like(clean_latent, dtype=torch.float32)
+        stored_data = trajectory_generator.generate(
+            clean_latent=clean_latent,
+            paired_conditional_dict=paired_conditional_dict,
+            trajectory_indices=DEFAULT_TRAJECTORY_INDICES,
+            generation_mode=args.generation_mode,
+            initial_noise=initial_noise,
+        )
+
+        torch.save(
+            {prompt: stored_data.cpu().detach()},
+            output_path,
+        )
+
+    dist.barrier()
+
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -0,0 +1,218 @@
+import argparse
+import argparse
+import torch
+import os
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from torchvision import transforms
+from torchvision.io import write_video
+from einops import rearrange
+import torch.distributed as dist
+from torch.utils.data import DataLoader, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+import json
+
+from pipeline import (
+    CausalDiffusionInferencePipeline,
+    CausalInferencePipeline,
+)
+from utils.dataset import TextDataset, TextImagePairDataset
+from utils.misc import set_seed
+
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_path", type=str, help="Path to the config file")
+parser.add_argument("--checkpoint_path", type=str, help="Path to the checkpoint folder")
+parser.add_argument("--data_path", type=str, help="Path to the dataset")
+parser.add_argument("--output_folder", type=str, help="Output folder")
+parser.add_argument("--num_output_frames", type=int, default=21, help="Number of overlap frames between sliding windows")
+parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA parameters")
+parser.add_argument("--seed", type=int, default=0, help="Random seed")
+parser.add_argument("--i2v", action="store_true", help="Whether to perform I2V (or T2V by default)")
+parser.add_argument("--report_timing", action="store_true",
+                    help="Only tested on A800, for the Causal Forcing++ latency. Not make claims for other hardware like H100. For the result on H100, refer to the reported results in the Self Forcing paper.")
+args = parser.parse_args()
+
+# Initialize distributed inference
+if "LOCAL_RANK" in os.environ:
+    dist.init_process_group(backend='nccl')
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    world_size = dist.get_world_size()
+    
+else:
+    device = torch.device("cuda")
+    local_rank = 0
+    world_size = 1
+
+set_seed(args.seed)
+
+print(f'Free VRAM {get_cuda_free_memory_gb(gpu)} GB')
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+torch.set_grad_enabled(False)
+
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+
+# Initialize pipeline
+if hasattr(config, 'denoising_step_list'):
+    # Few-step inference
+    pipeline = CausalInferencePipeline(config, device=device)
+else:
+    # Multi-step diffusion inference
+    pipeline = CausalDiffusionInferencePipeline(config, device=device)
+
+if args.checkpoint_path:
+    state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+    key = 'generator_ema' if args.use_ema else 'generator'
+    gen_sd = state_dict[key]
+
+    try:
+        pipeline.generator.load_state_dict(gen_sd)
+    except RuntimeError:
+        fixed = {}
+        for k, v in gen_sd.items():
+            if k.startswith("model._fsdp_wrapped_module."):
+                k = k.replace("model._fsdp_wrapped_module.", "model.", 1)
+            fixed[k] = v
+        pipeline.generator.load_state_dict(fixed, strict=False)
+
+pipeline = pipeline.to(dtype=torch.bfloat16)
+if low_memory:
+    DynamicSwapInstaller.install_model(pipeline.text_encoder, device=gpu)
+else:
+    pipeline.text_encoder.to(device=gpu)
+pipeline.generator.to(device=gpu)
+pipeline.vae.to(device=gpu)
+
+
+# Create dataset
+if args.i2v:
+    assert not dist.is_initialized(), "I2V does not support distributed inference yet"
+    transform = transforms.Compose([
+        transforms.Resize((480, 832)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+    dataset = TextImagePairDataset(args.data_path, transform=transform)
+else:
+    dataset = TextDataset(prompt_path=args.data_path)
+num_prompts = len(dataset)
+print(f"Number of prompts: {num_prompts}")
+
+if args.report_timing and num_prompts < 2:
+    print(f"[WARN] --report_timing requires at least 2 prompts "
+          f"(got {num_prompts}); timing disabled.")
+    args.report_timing = False
+
+if dist.is_initialized():
+    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)
+else:
+    sampler = SequentialSampler(dataset)
+dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+# Create output directory (only on main process to avoid race conditions)
+if local_rank == 0:
+    os.makedirs(args.output_folder, exist_ok=True)
+
+if dist.is_initialized():
+    dist.barrier()
+
+def encode(self, videos: torch.Tensor) -> torch.Tensor:
+    device, dtype = videos[0].device, videos[0].dtype
+    scale = [self.mean.to(device=device, dtype=dtype),
+             1.0 / self.std.to(device=device, dtype=dtype)]
+    output = [
+        self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+        for u in videos
+    ]
+
+    output = torch.stack(output, dim=0)
+    return output
+
+
+for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
+    idx = batch_data['idx'].item()
+
+    if isinstance(batch_data, dict):
+        batch = batch_data
+    elif isinstance(batch_data, list):
+        batch = batch_data[0]  # First (and only) item in the batch
+
+    all_video = []
+    num_generated_frames = 0  # Number of generated (latent) frames
+    
+    
+    if args.i2v:
+        assert config.num_frame_per_block == 1, "Current I2V only supports the frame-wise model."
+        # For image-to-video, batch contains image and caption
+        prompt = batch['prompts'][0]  # Get caption from batch
+        output_path = os.path.join(args.output_folder, f'{prompt[:100]}.mp4')
+        if os.path.exists(output_path):
+            print('Video has been generated. Pass!')
+            continue
+        # Process the image
+        image = batch['image'].squeeze(0).unsqueeze(0).unsqueeze(2).to(device=device, dtype=torch.bfloat16)
+
+        # Encode the input image as the first latent
+        initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16)
+        prompts = [prompt] 
+        sampled_noise = torch.randn(
+            [1, args.num_output_frames - 1, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+    else:
+        # For text-to-video, batch is just the text prompt
+        prompt = batch['prompts'][0]
+        output_path = os.path.join(args.output_folder, f'{prompt[:100]}.mp4')
+        if os.path.exists(output_path):
+            print('Video has been generated. Pass!')
+            continue
+        extended_prompt = batch['extended_prompts'][0] if 'extended_prompts' in batch else None
+        if extended_prompt is not None:
+            prompts = [extended_prompt] 
+        else:
+            prompts = [prompt] 
+
+        initial_latent = None
+        sampled_noise = torch.randn(
+            [1, args.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+
+    sample_report_timing = args.report_timing and i >= 1
+    video, latents = pipeline.inference(
+        noise=sampled_noise,
+        text_prompts=prompts,
+        return_latents=True,
+        initial_latent=initial_latent,
+        report_timing=sample_report_timing,
+    )
+    if sample_report_timing:
+        latency = pipeline.first_chunk_time
+        elapsed = pipeline.last_generation_time
+        num_pixel_frames = video.shape[1]
+        fps = num_pixel_frames / elapsed if elapsed > 0 else float('inf')
+        print(f"[Sample {i}] {num_pixel_frames} frames, "
+              f"latency ↓ {latency:.2f}s, FPS ↑ {fps:.2f}")
+        # Only tested on A800, for the Causal Forcing++ paper latency & throughput.
+        # Not make claims for other hardware like H100.
+        # For the result on H100, refer to the reported results in the Self Forcing paper.
+        # We do not guarantee that our FPS/latency measurement protocol is identical to that used in the Self Forcing paper.
+    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()
+    all_video.append(current_video)
+    num_generated_frames += latents.shape[1]
+
+    # Final output video
+    clean_latent = latents[0].cpu() 
+    video = 255.0 * torch.cat(all_video, dim=1)
+
+    # Clear VAE cache
+    pipeline.vae.model.clear_cache()
+
+    output_path = os.path.join(args.output_folder, f'{prompt[:100]}.mp4')
+    write_video(output_path, video[0], fps=16)
+
+       

long_video/LICENSE

@@ -0,0 +1,81 @@
+Tencent is pleased to support the community by making RollingForcing available.
+
+Copyright (C)  2025 Tencent.  All rights reserved. 
+
+The open-source software and/or models included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent.
+
+RollingForcing is licensed under the License Terms of RollingForcing, except for the third-party components listed below, which remain licensed under their respective original terms. RollingForcing does not impose any additional restrictions beyond those specified in the original licenses of these third-party components. Users are required to comply with all applicable terms and conditions of the original licenses and to ensure that the use of these third-party components conforms to all relevant laws and regulations.
+
+For the avoidance of doubt, RollingForcing refers solely to training code, inference code, parameters, and weights made publicly available by Tencent in accordance with the License Terms of RollingForcing. 
+
+Terms of the License Terms of RollingForcing:
+--------------------------------------------------------------------
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, and /or sublicense copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+- You agree to use RollingForcing only for academic purposes, and refrain from using it for any commercial or production purposes under any circumstances.
+
+- The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+
+
+Dependencies and Licenses:
+
+This open-source project, RollingForcing, builds upon the following open-source models and/or software components, each of which remains licensed under its original license. Certain models or software may include modifications made by Tencent (“Tencent Modifications”), which are Copyright (C) Tencent.
+
+In case you believe there have been errors in the attribution below, you may submit the concerns to us for review and correction.
+
+Open Source Model Licensed under the  Apache-2.0:
+--------------------------------------------------------------------
+1. Wan-AI/Wan2.1-T2V-1.3B
+Copyright (c) 2025 Wan Team
+
+Terms of the Apache-2.0：
+--------------------------------------------------------------------
+Apache License
+                       Version 2.0, January 2004
+                    http://www.apache.org/licenses/
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+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.
+
+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.
+
+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.
+
+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.
+
+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.
+
+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.
+
+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.
+
+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.
+
+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

long_video/README.md

@@ -0,0 +1,40 @@
+Builing on [Rolling Forcing](https://github.com/TencentARC/RollingForcing), we implemented minute-level long video generation.
+
+### Installation
+
+```bash
+conda activate causal_forcing
+pip install tensorboard opencv-python packaging
+```
+
+### CLI inference
+```
+hf download Wan-AI/Wan2.1-T2V-1.3B --local-dir wan_models/Wan2.1-T2V-1.3B
+hf download zhuhz22/Causal-Forcing chunkwise/longvideo.pt --local-dir ../checkpoints
+
+python inference.py \
+    --config_path configs/rolling_forcing_dmd.yaml \
+    --output_folder videos/rolling_forcing_dmd \
+    --checkpoint_path ../checkpoints/chunkwise/longvideo.pt \
+    --data_path prompts/example_prompts.txt \
+    --num_output_frames 252 \
+    --use_ema
+```
+
+### Training
+```
+hf download Wan-AI/Wan2.1-T2V-1.3B --local-dir wan_models/Wan2.1-T2V-1.3B
+hf download Wan-AI/Wan2.1-T2V-14B --local-dir wan_models/Wan2.1-T2V-14B
+
+torchrun --nproc_per_node=8 \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint 127.0.0.1:29500 \
+  train.py \
+  -- \
+  --config_path configs/rolling_forcing_dmd.yaml \
+  --logdir logs/rolling_forcing_dmd
+```
+> We recommend training for 3000 steps.
+
+### Acknowledge
+We adopt [Rolling Forcing](https://github.com/TencentARC/RollingForcing) as our long video generation framework and only change the ODE initialization part.

long_video/app.py

@@ -0,0 +1,167 @@
+import os
+import argparse
+import time
+from typing import Optional
+
+import torch
+from torchvision.io import write_video
+from omegaconf import OmegaConf
+from einops import rearrange
+import gradio as gr
+
+from pipeline import CausalDiffusionInferencePipeline, CausalInferencePipeline
+
+
+# -----------------------------
+# Globals (loaded once per process)
+# -----------------------------
+_PIPELINE: Optional[torch.nn.Module] = None
+_DEVICE: Optional[torch.device] = None
+
+
+def _ensure_gpu():
+    if not torch.cuda.is_available():
+        raise gr.Error("CUDA GPU is required to run this demo. Please run on a machine with an NVIDIA GPU.")
+    # Bind to GPU:0 by default
+    torch.cuda.set_device(0)
+
+
+def _load_pipeline(config_path: str, checkpoint_path: Optional[str], use_ema: bool) -> torch.nn.Module:
+    global _PIPELINE, _DEVICE
+    if _PIPELINE is not None:
+        return _PIPELINE
+
+    _ensure_gpu()
+    _DEVICE = torch.device("cuda:0")
+
+    # Load and merge configs
+    config = OmegaConf.load(config_path)
+    default_config = OmegaConf.load("configs/default_config.yaml")
+    config = OmegaConf.merge(default_config, config)
+
+    # Choose pipeline type based on config
+    pipeline = CausalInferencePipeline(config, device=_DEVICE)
+
+
+    # Load checkpoint if provided
+    if checkpoint_path and os.path.exists(checkpoint_path):
+        state_dict = torch.load(checkpoint_path, map_location="cpu")
+        if use_ema and 'generator_ema' in state_dict:
+            state_dict_to_load = state_dict['generator_ema']
+            # Remove possible FSDP prefix
+            from collections import OrderedDict
+            new_state_dict = OrderedDict()
+            for k, v in state_dict_to_load.items():
+                new_state_dict[k.replace("_fsdp_wrapped_module.", "")] = v
+            state_dict_to_load = new_state_dict
+        else:
+            state_dict_to_load = state_dict.get('generator', state_dict)
+        pipeline.generator.load_state_dict(state_dict_to_load, strict=False)
+
+    # The codebase assumes bfloat16 on GPU
+    pipeline = pipeline.to(device=_DEVICE, dtype=torch.bfloat16)
+    pipeline.eval()
+
+    # Quick sanity path check for Wan models to give friendly errors
+    wan_dir = os.path.join('wan_models', 'Wan2.1-T2V-1.3B')
+    if not os.path.isdir(wan_dir):
+        raise gr.Error(
+            "Wan2.1-T2V-1.3B not found at 'wan_models/Wan2.1-T2V-1.3B'.\n"
+            "Please download it first, e.g.:\n"
+            "huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir-use-symlinks False --local-dir wan_models/Wan2.1-T2V-1.3B"
+        )
+
+    _PIPELINE = pipeline
+    return _PIPELINE
+
+
+def build_predict(config_path: str, checkpoint_path: Optional[str], output_dir: str, use_ema: bool):
+    os.makedirs(output_dir, exist_ok=True)
+
+    def predict(prompt: str, num_frames: int) -> str:
+        if not prompt or not prompt.strip():
+            raise gr.Error("Please enter a non-empty text prompt.")
+
+        num_frames = int(num_frames)
+        if num_frames % 3 != 0 or not (21 <= num_frames <= 252):
+            raise gr.Error("Number of frames must be a multiple of 3 between 21 and 252.")
+
+        pipeline = _load_pipeline(config_path, checkpoint_path, use_ema)
+
+        # Prepare inputs
+        prompts = [prompt.strip()]
+        noise = torch.randn([1, num_frames, 16, 60, 104], device=_DEVICE, dtype=torch.bfloat16)
+
+        torch.set_grad_enabled(False)
+        with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            video = pipeline.inference_rolling_forcing(
+                noise=noise,
+                text_prompts=prompts,
+                return_latents=False,
+                initial_latent=None,
+            )
+
+        # video: [B=1, T, C, H, W] in [0,1]
+        video = rearrange(video, 'b t c h w -> b t h w c')[0]
+        video_uint8 = (video * 255.0).clamp(0, 255).to(torch.uint8).cpu()
+
+        # Save to a unique filepath
+        safe_stub = prompt[:60].replace(' ', '_').replace('/', '_')
+        ts = int(time.time())
+        filepath = os.path.join(output_dir, f"{safe_stub or 'video'}_{ts}.mp4")
+        write_video(filepath, video_uint8, fps=16)
+        print(f"Saved generated video to {filepath}")
+
+        return filepath
+
+    return predict
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config_path', type=str, default='configs/rolling_forcing_dmd.yaml',
+                        help='Path to the model config')
+    parser.add_argument('--checkpoint_path', type=str, default='checkpoints/rolling_forcing_dmd.pt',
+                        help='Path to rolling forcing checkpoint (.pt). If missing, will run with base weights only if available.')
+    parser.add_argument('--output_dir', type=str, default='videos/gradio', help='Where to save generated videos')
+    parser.add_argument('--no_ema', action='store_true', help='Disable EMA weights when loading checkpoint')
+    parser.add_argument('--server_name', type=str, default='0.0.0.0', help='Gradio server host')
+    parser.add_argument('--server_port', type=int, default=7860, help='Gradio server port')
+    args = parser.parse_args()
+
+    predict = build_predict(
+        config_path=args.config_path,
+        checkpoint_path=args.checkpoint_path,
+        output_dir=args.output_dir,
+        use_ema=not args.no_ema,
+    )
+
+    demo = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Textbox(label="Text Prompt", lines=2, placeholder="A cinematic shot of a girl dancing in the sunset."),
+            gr.Slider(label="Number of Latent Frames", minimum=21, maximum=252, step=3, value=21),
+        ],
+        outputs=gr.Video(label="Generated Video", format="mp4"),
+        title="Rolling Forcing: Autoregressive Long Video Diffusion in Real Time",
+        description=(
+            "Enter a prompt and generate a video using the Rolling Forcing pipeline.\n"
+            "**Note:** although Rolling Forcing generates videos autoregressivelty, current Gradio demo does not support streaming outputs, so the entire video will be generated before it is displayed.\n"
+            "\n"
+            "If you find this demo useful, please consider giving it a ⭐ star on [GitHub](https://github.com/TencentARC/RollingForcing)--your support is crucial for sustaining this open-source project. "
+            "You can also dive deeper by reading the [paper](https://arxiv.org/abs/2509.25161) or exploring the [project page](https://kunhao-liu.github.io/Rolling_Forcing_Webpage) for more details."
+        ),
+        allow_flagging='never',
+    )
+
+    try:
+        # Gradio <= 3.x
+        demo.queue(concurrency_count=1, max_size=2)
+    except TypeError:
+        # Gradio >= 4.x
+        demo.queue(max_size=2)
+    demo.launch(server_name=args.server_name, server_port=args.server_port, show_error=True)
+
+
+if __name__ == "__main__":
+    main()

long_video/configs/default_config.yaml

@@ -0,0 +1,20 @@
+independent_first_frame: false
+warp_denoising_step: false
+weight_decay: 0.01
+same_step_across_blocks: true
+discriminator_lr_multiplier: 1.0
+last_step_only: false
+i2v: false
+num_training_frames: 27
+gc_interval: 100
+context_noise: 0
+causal: true
+
+ckpt_step: 0
+prompt_name: MovieGenVideoBench
+prompt_path: prompts/MovieGenVideoBench.txt
+eval_first_n: 64
+num_samples: 1
+height: 480
+width: 832
+num_frames: 81

long_video/configs/rolling_forcing_dmd.yaml

@@ -0,0 +1,47 @@
+generator_ckpt: ../checkpoints/chunkwise/causal_ode.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-14B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+sharding_strategy: hybrid_full
+lr: 1.5e-06
+lr_critic: 4.0e-07
+beta1: 0.0
+beta2: 0.999
+beta1_critic: 0.0
+beta2_critic: 0.999
+data_path: ../prompts/vidprom_filtered_extended.txt 
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 64
+log_iters: 1000
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0

long_video/inference.py

@@ -0,0 +1,197 @@
+import argparse
+import torch
+import os
+from omegaconf import OmegaConf
+from collections import OrderedDict
+from tqdm import tqdm
+from torchvision import transforms
+from torchvision.io import write_video
+from einops import rearrange
+import torch.distributed as dist
+import imageio
+from torch.utils.data import DataLoader, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from pipeline import (
+    CausalDiffusionInferencePipeline,
+    CausalInferencePipeline
+)
+from utils.dataset import TextDataset, TextImagePairDataset
+from utils.misc import set_seed
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_path", type=str, help="Path to the config file")
+parser.add_argument("--checkpoint_path", type=str, help="Path to the checkpoint folder")
+parser.add_argument("--data_path", type=str, help="Path to the dataset")
+parser.add_argument("--extended_prompt_path", type=str, help="Path to the extended prompt")
+parser.add_argument("--output_folder", type=str, help="Output folder")
+parser.add_argument("--num_output_frames", type=int, default=21,
+                    help="Number of overlap frames between sliding windows")
+parser.add_argument("--i2v", action="store_true", help="Whether to perform I2V (or T2V by default)")
+parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA parameters")
+parser.add_argument("--seed", type=int, default=0, help="Random seed")
+parser.add_argument("--num_samples", type=int, default=1, help="Number of samples to generate per prompt")
+parser.add_argument("--save_with_index", action="store_true",
+                    help="Whether to save the video using the index or prompt as the filename")
+args = parser.parse_args()
+
+# Initialize distributed inference
+if "LOCAL_RANK" in os.environ:
+    dist.init_process_group(backend='nccl')
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    world_size = dist.get_world_size()
+    set_seed(args.seed + local_rank)
+else:
+    device = torch.device("cuda")
+    local_rank = 0
+    world_size = 1
+    set_seed(args.seed)
+
+torch.set_grad_enabled(False)
+
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+
+# Initialize pipeline
+if hasattr(config, 'denoising_step_list'):
+    # Few-step inference
+    pipeline = CausalInferencePipeline(config, device=device)
+else:
+    # Multi-step diffusion inference
+    pipeline = CausalDiffusionInferencePipeline(config, device=device)
+
+if args.checkpoint_path:
+    state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+    if args.use_ema:
+        state_dict_to_load = state_dict['generator_ema']
+        def remove_fsdp_prefix(state_dict):
+            new_state_dict = OrderedDict()
+            for key, value in state_dict.items():
+                if "_fsdp_wrapped_module." in key:
+                    new_key = key.replace("_fsdp_wrapped_module.", "")
+                    new_state_dict[new_key] = value
+                else:
+                    new_state_dict[key] = value
+            return new_state_dict
+        state_dict_to_load = remove_fsdp_prefix(state_dict_to_load)
+    else:
+        state_dict_to_load = state_dict['generator']
+    pipeline.generator.load_state_dict(state_dict_to_load)
+
+pipeline = pipeline.to(device=device, dtype=torch.bfloat16)
+
+# Create dataset
+if args.i2v:
+    assert not dist.is_initialized(), "I2V does not support distributed inference yet"
+    transform = transforms.Compose([
+        transforms.Resize((480, 832)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+    dataset = TextImagePairDataset(args.data_path, transform=transform)
+else:
+    dataset = TextDataset(prompt_path=args.data_path, extended_prompt_path=args.extended_prompt_path)
+num_prompts = len(dataset)
+print(f"Number of prompts: {num_prompts}")
+
+if dist.is_initialized():
+    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)
+else:
+    sampler = SequentialSampler(dataset)
+dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+# Create output directory (only on main process to avoid race conditions)
+if local_rank == 0:
+    os.makedirs(args.output_folder, exist_ok=True)
+
+if dist.is_initialized():
+    dist.barrier()
+
+
+def encode(self, videos: torch.Tensor) -> torch.Tensor:
+    device, dtype = videos[0].device, videos[0].dtype
+    scale = [self.mean.to(device=device, dtype=dtype),
+             1.0 / self.std.to(device=device, dtype=dtype)]
+    output = [
+        self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+        for u in videos
+    ]
+
+    output = torch.stack(output, dim=0)
+    return output
+
+
+for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
+    idx = batch_data['idx'].item()
+
+    # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
+    # Unpack the batch data for convenience
+    if isinstance(batch_data, dict):
+        batch = batch_data
+    elif isinstance(batch_data, list):
+        batch = batch_data[0]  # First (and only) item in the batch
+
+    all_video = []
+    num_generated_frames = 0  # Number of generated (latent) frames
+
+    if args.i2v:
+        # For image-to-video, batch contains image and caption
+        prompt = batch['prompts'][0]  # Get caption from batch
+        prompts = [prompt] * args.num_samples
+
+        # Process the image
+        image = batch['image'].squeeze(0).unsqueeze(0).unsqueeze(2).to(device=device, dtype=torch.bfloat16)
+
+        # Encode the input image as the first latent
+        initial_latent = pipeline.vae.encode_to_latent(image).to(device=device, dtype=torch.bfloat16)
+        initial_latent = initial_latent.repeat(args.num_samples, 1, 1, 1, 1)
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames - 1, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+    else:
+        # For text-to-video, batch is just the text prompt
+        prompt = batch['prompts'][0]
+        extended_prompt = batch['extended_prompts'][0] if 'extended_prompts' in batch else None
+        if extended_prompt is not None:
+            prompts = [extended_prompt] * args.num_samples
+        else:
+            prompts = [prompt] * args.num_samples
+        initial_latent = None
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames, 16, 60, 104], device=device, dtype=torch.bfloat16
+        )
+
+    # Generate 81 frames
+    video, latents = pipeline.inference_rolling_forcing(
+        noise=sampled_noise,
+        text_prompts=prompts,
+        return_latents=True,
+        initial_latent=initial_latent,
+    )
+    current_video = rearrange(video, 'b t c h w -> b t h w c').cpu()
+    all_video.append(current_video)
+    num_generated_frames += latents.shape[1]
+
+    # Final output video
+    video = 255.0 * torch.cat(all_video, dim=1)
+
+    # Clear VAE cache
+    pipeline.vae.model.clear_cache()
+
+    # Save the video if the current prompt is not a dummy prompt
+    if idx < num_prompts:
+        model = "regular" if not args.use_ema else "ema"
+        for seed_idx in range(args.num_samples):
+            # All processes save their videos
+            if args.save_with_index:
+                output_path = os.path.join(args.output_folder, f'{idx}-{seed_idx}_{model}.mp4')
+            else:
+                output_path = os.path.join(args.output_folder, f'{prompt[:100]}-{seed_idx}.mp4')
+            write_video(output_path, video[seed_idx], fps=16)
+            # imageio.mimwrite(output_path, video[seed_idx], fps=16, quality=8, output_params=["-loglevel", "error"])
+    

long_video/model/__init__.py

@@ -0,0 +1,14 @@
+from .diffusion import CausalDiffusion
+from .causvid import CausVid
+from .dmd import DMD
+from .gan import GAN
+from .sid import SiD
+from .ode_regression import ODERegression
+__all__ = [
+    "CausalDiffusion",
+    "CausVid",
+    "DMD",
+    "GAN",
+    "SiD",
+    "ODERegression"
+]

long_video/model/base.py

@@ -0,0 +1,230 @@
+from typing import Tuple
+from einops import rearrange
+from torch import nn
+import torch.distributed as dist
+import torch
+
+from pipeline import RollingForcingTrainingPipeline
+from utils.loss import get_denoising_loss
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BaseModel(nn.Module):
+    def __init__(self, args, device):
+        super().__init__()
+        self._initialize_models(args, device)
+
+        self.device = device
+        self.args = args
+        self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32
+        if hasattr(args, "denoising_step_list"):
+            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]
+
+    def _initialize_models(self, args, device):
+        self.real_model_name = getattr(args, "real_name", "Wan2.1-T2V-1.3B")
+        self.fake_model_name = getattr(args, "fake_name", "Wan2.1-T2V-1.3B")
+
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.real_score = WanDiffusionWrapper(model_name=self.real_model_name, is_causal=False)
+        self.real_score.model.requires_grad_(False)
+
+        self.fake_score = WanDiffusionWrapper(model_name=self.fake_model_name, is_causal=False)
+        self.fake_score.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+        self.scheduler = self.generator.get_scheduler()
+        self.scheduler.timesteps = self.scheduler.timesteps.to(device)
+
+    def _get_timestep(
+            self,
+            min_timestep: int,
+            max_timestep: int,
+            batch_size: int,
+            num_frame: int,
+            num_frame_per_block: int,
+            uniform_timestep: bool = False
+    ) -> torch.Tensor:
+        """
+        Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep
+        from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].
+        - If uniform_timestep, it will use the same timestep for all frames.
+        - If not uniform_timestep, it will use a different timestep for each block.
+        """
+        if uniform_timestep:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, 1],
+                device=self.device,
+                dtype=torch.long
+            ).repeat(1, num_frame)
+            return timestep
+        else:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, num_frame],
+                device=self.device,
+                dtype=torch.long
+            )
+            # make the noise level the same within every block
+            if self.independent_first_frame:
+                # the first frame is always kept the same
+                timestep_from_second = timestep[:, 1:]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1, num_frame_per_block)
+                timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1)
+                timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)
+            else:
+                timestep = timestep.reshape(
+                    timestep.shape[0], -1, num_frame_per_block)
+                timestep[:, :, 1:] = timestep[:, :, 0:1]
+                timestep = timestep.reshape(timestep.shape[0], -1)
+            return timestep
+
+
+class RollingForcingModel(BaseModel):
+    def __init__(self, args, device):
+        super().__init__(args, device)
+        self.denoising_loss_func = get_denoising_loss(args.denoising_loss_type)()
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        initial_latent: torch.tensor = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+            - denoised_timestep: an integer
+        """
+        # Step 1: Sample noise and backward simulate the generator's input
+        assert getattr(self.args, "backward_simulation", True), "Backward simulation needs to be enabled"
+        if initial_latent is not None:
+            conditional_dict["initial_latent"] = initial_latent
+        if self.args.i2v:
+            noise_shape = [image_or_video_shape[0], image_or_video_shape[1] - 1, *image_or_video_shape[2:]]
+        else:
+            noise_shape = image_or_video_shape.copy()
+
+        # During training, the number of generated frames should be uniformly sampled from
+        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
+        min_num_frames = 20 if self.args.independent_first_frame else 21
+        max_num_frames = self.num_training_frames - 1 if self.args.independent_first_frame else self.num_training_frames
+        assert max_num_frames % self.num_frame_per_block == 0
+        assert min_num_frames % self.num_frame_per_block == 0
+        max_num_blocks = max_num_frames // self.num_frame_per_block
+        min_num_blocks = min_num_frames // self.num_frame_per_block
+        num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)
+        dist.broadcast(num_generated_blocks, src=0)
+        num_generated_blocks = num_generated_blocks.item()
+        num_generated_frames = num_generated_blocks * self.num_frame_per_block
+        if self.args.independent_first_frame and initial_latent is None:
+            num_generated_frames += 1
+            min_num_frames += 1
+        # Sync num_generated_frames across all processes
+        noise_shape[1] = num_generated_frames
+
+        pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation(
+            noise=torch.randn(noise_shape,
+                              device=self.device, dtype=self.dtype),
+            **conditional_dict,
+        )
+        # Slice last 21 frames
+        if pred_image_or_video.shape[1] > 21:
+            with torch.no_grad():
+                # Reencode to get image latent
+                latent_to_decode = pred_image_or_video[:, :-20, ...]
+                # Deccode to video
+                pixels = self.vae.decode_to_pixel(latent_to_decode)
+                frame = pixels[:, -1:, ...].to(self.dtype)
+                frame = rearrange(frame, "b t c h w -> b c t h w")
+                # Encode frame to get image latent
+                image_latent = self.vae.encode_to_latent(frame).to(self.dtype)
+            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)
+        else:
+            pred_image_or_video_last_21 = pred_image_or_video
+
+        if num_generated_frames != min_num_frames:
+            # Currently, we do not use gradient for the first chunk, since it contains image latents
+            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
+            if self.args.independent_first_frame:
+                gradient_mask[:, :1] = False
+            else:
+                gradient_mask[:, :self.num_frame_per_block] = False
+        else:
+            gradient_mask = None
+
+        pred_image_or_video_last_21 = pred_image_or_video_last_21.to(self.dtype)
+        return pred_image_or_video_last_21, gradient_mask, denoised_timestep_from, denoised_timestep_to
+
+    def _consistency_backward_simulation(
+        self,
+        noise: torch.Tensor,
+        **conditional_dict: dict
+    ) -> torch.Tensor:
+        """
+        Simulate the generator's input from noise to avoid training/inference mismatch.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)
+        Input:
+            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - output: a tensor with shape [B, T, F, C, H, W].
+            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
+            represents the x0 prediction at each timestep.
+        """
+        if self.inference_pipeline is None:
+            self._initialize_inference_pipeline()
+
+        infer_w_rolling = torch.rand(1, device=self.device) > 0.5
+        dist.broadcast(infer_w_rolling, src=0)
+
+        if infer_w_rolling:
+            return self.inference_pipeline.inference_with_rolling_forcing(
+                noise=noise, **conditional_dict
+            )
+        else:
+            return self.inference_pipeline.inference_with_self_forcing(
+                noise=noise, **conditional_dict
+            )
+
+    def _initialize_inference_pipeline(self):
+        """
+        Lazy initialize the inference pipeline during the first backward simulation run.
+        Here we encapsulate the inference code with a model-dependent outside function.
+        We pass our FSDP-wrapped modules into the pipeline to save memory.
+        """
+        self.inference_pipeline = RollingForcingTrainingPipeline(
+            denoising_step_list=self.denoising_step_list,
+            scheduler=self.scheduler,
+            generator=self.generator,
+            num_frame_per_block=self.num_frame_per_block,
+            independent_first_frame=self.args.independent_first_frame,
+            same_step_across_blocks=self.args.same_step_across_blocks,
+            last_step_only=self.args.last_step_only,
+            num_max_frames=self.num_training_frames,
+            context_noise=self.args.context_noise
+        )

long_video/model/causvid.py

@@ -0,0 +1,391 @@
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+
+
+class CausVid(BaseModel):
+    def __init__(self, args, device):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.teacher_forcing = getattr(args, "teacher_forcing", False)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _compute_kl_grad(
+        self, noisy_image_or_video: torch.Tensor,
+        estimated_clean_image_or_video: torch.Tensor,
+        timestep: torch.Tensor,
+        conditional_dict: dict, unconditional_dict: dict,
+        normalization: bool = True
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
+            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - normalization: a boolean indicating whether to normalize the gradient.
+        Output:
+            - kl_grad: a tensor representing the KL grad.
+            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Compute the fake score
+        _, pred_fake_image_cond = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        if self.fake_guidance_scale != 0.0:
+            _, pred_fake_image_uncond = self.fake_score(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=unconditional_dict,
+                timestep=timestep
+            )
+            pred_fake_image = pred_fake_image_cond + (
+                pred_fake_image_cond - pred_fake_image_uncond
+            ) * self.fake_guidance_scale
+        else:
+            pred_fake_image = pred_fake_image_cond
+
+        # Step 2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep
+        )
+
+        pred_real_image = pred_real_image_cond + (
+            pred_real_image_cond - pred_real_image_uncond
+        ) * self.real_guidance_scale
+
+        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
+        grad = (pred_fake_image - pred_real_image)
+
+        # TODO: Change the normalizer for causal teacher
+        if normalization:
+            # Step 4: Gradient normalization (DMD paper eq. 8).
+            p_real = (estimated_clean_image_or_video - pred_real_image)
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+            grad = grad / normalizer
+        grad = torch.nan_to_num(grad)
+
+        return grad, {
+            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
+            "timestep": timestep.detach()
+        }
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        with torch.no_grad():
+            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+            timestep = self._get_timestep(
+                0,
+                self.num_train_timestep,
+                batch_size,
+                num_frame,
+                self.num_frame_per_block,
+                uniform_timestep=True
+            )
+
+            if self.timestep_shift > 1:
+                timestep = self.timestep_shift * \
+                    (timestep / 1000) / \
+                    (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
+            timestep = timestep.clamp(self.min_step, self.max_step)
+
+            noise = torch.randn_like(image_or_video)
+            noisy_latent = self.scheduler.add_noise(
+                image_or_video.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep.flatten(0, 1)
+            ).detach().unflatten(0, (batch_size, num_frame))
+
+            # Step 2: Compute the KL grad
+            grad, dmd_log_dict = self._compute_kl_grad(
+                noisy_image_or_video=noisy_latent,
+                estimated_clean_image_or_video=original_latent,
+                timestep=timestep,
+                conditional_dict=conditional_dict,
+                unconditional_dict=unconditional_dict
+            )
+
+        if gradient_mask is not None:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            )[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean")
+        else:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            ), (original_latent.double() - grad.double()).detach(), reduction="mean")
+        return dmd_loss, dmd_log_dict
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        clean_latent: torch.tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+        """
+        simulated_noisy_input = []
+        for timestep in self.denoising_step_list:
+            noise = torch.randn(
+                image_or_video_shape, device=self.device, dtype=self.dtype)
+
+            noisy_timestep = timestep * torch.ones(
+                image_or_video_shape[:2], device=self.device, dtype=torch.long)
+
+            if timestep != 0:
+                noisy_image = self.scheduler.add_noise(
+                    clean_latent.flatten(0, 1),
+                    noise.flatten(0, 1),
+                    noisy_timestep.flatten(0, 1)
+                ).unflatten(0, image_or_video_shape[:2])
+            else:
+                noisy_image = clean_latent
+
+            simulated_noisy_input.append(noisy_image)
+
+        simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1)
+
+        # Step 2: Randomly sample a timestep and pick the corresponding input
+        index = self._get_timestep(
+            0,
+            len(self.denoising_step_list),
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=False
+        )
+
+        # select the corresponding timestep's noisy input from the stacked tensor [B, T, F, C, H, W]
+        noisy_input = torch.gather(
+            simulated_noisy_input, dim=1,
+            index=index.reshape(index.shape[0], 1, index.shape[1], 1, 1, 1).expand(
+                -1, -1, -1, *image_or_video_shape[2:]).to(self.device)
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index].to(self.device)
+
+        _, pred_image_or_video = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x=clean_latent if self.teacher_forcing else None,
+        )
+
+        gradient_mask = None  # timestep != 0
+
+        pred_image_or_video = pred_image_or_video.type_as(noisy_input)
+
+        return pred_image_or_video, gradient_mask
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Run generator on backward simulated noisy input
+        pred_image, gradient_mask = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            clean_latent=clean_latent
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask
+        )
+
+        # Step 3: TODO: Implement the GAN loss
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _ = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                clean_latent=clean_latent
+            )
+
+        # Step 2: Compute the fake prediction
+        critic_timestep = self._get_timestep(
+            0,
+            self.num_train_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = self.timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred
+        )
+
+        # Step 4: TODO: Compute the GAN loss
+
+        # Step 5: Debugging Log
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach()
+        }
+
+        return denoising_loss, critic_log_dict

long_video/model/diffusion.py

@@ -0,0 +1,125 @@
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class CausalDiffusion(BaseModel):
+    def __init__(self, args, device):
+        """
+        Initialize the Diffusion loss module.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        self.guidance_scale = args.guidance_scale
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.teacher_forcing = getattr(args, "teacher_forcing", False)
+        # Noise augmentation in teacher forcing, we add small noise to clean context latents
+        self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0)
+
+    def _initialize_models(self, args):
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        noise = torch.randn_like(clean_latent)
+        batch_size, num_frame = image_or_video_shape[:2]
+
+        # Step 2: Randomly sample a timestep and add noise to denoiser inputs
+        index = self._get_timestep(
+            0,
+            self.scheduler.num_train_timesteps,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=False
+        )
+        timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device)
+        noisy_latents = self.scheduler.add_noise(
+            clean_latent.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep.flatten(0, 1)
+        ).unflatten(0, (batch_size, num_frame))
+        training_target = self.scheduler.training_target(clean_latent, noise, timestep)
+
+        # Step 3: Noise augmentation, also add small noise to clean context latents
+        if self.noise_augmentation_max_timestep > 0:
+            index_clean_aug = self._get_timestep(
+                0,
+                self.noise_augmentation_max_timestep,
+                image_or_video_shape[0],
+                image_or_video_shape[1],
+                self.num_frame_per_block,
+                uniform_timestep=False
+            )
+            timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(dtype=self.dtype, device=self.device)
+            clean_latent_aug = self.scheduler.add_noise(
+                clean_latent.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep_clean_aug.flatten(0, 1)
+            ).unflatten(0, (batch_size, num_frame))
+        else:
+            clean_latent_aug = clean_latent
+            timestep_clean_aug = None
+
+        # Compute loss
+        flow_pred, x0_pred = self.generator(
+            noisy_image_or_video=noisy_latents,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x=clean_latent_aug if self.teacher_forcing else None,
+            aug_t=timestep_clean_aug if self.teacher_forcing else None
+        )
+        # loss = torch.nn.functional.mse_loss(flow_pred.float(), training_target.float())
+        loss = torch.nn.functional.mse_loss(
+            flow_pred.float(), training_target.float(), reduction='none'
+        ).mean(dim=(2, 3, 4))
+        loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
+        loss = loss.mean()
+
+        log_dict = {
+            "x0": clean_latent.detach(),
+            "x0_pred": x0_pred.detach()
+        }
+        return loss, log_dict

long_video/model/dmd.py

@@ -0,0 +1,332 @@
+from pipeline import RollingForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Optional, Tuple
+import torch
+
+from model.base import RollingForcingModel
+
+
+class DMD(RollingForcingModel):
+    def __init__(self, args, device):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: RollingForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+        self.min_score_timestep = getattr(args, "min_score_timestep", 0)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _compute_kl_grad(
+        self, noisy_image_or_video: torch.Tensor,
+        estimated_clean_image_or_video: torch.Tensor,
+        timestep: torch.Tensor,
+        conditional_dict: dict, unconditional_dict: dict,
+        normalization: bool = True
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
+            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - normalization: a boolean indicating whether to normalize the gradient.
+        Output:
+            - kl_grad: a tensor representing the KL grad.
+            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Compute the fake score
+        _, pred_fake_image_cond = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        if self.fake_guidance_scale != 0.0:
+            _, pred_fake_image_uncond = self.fake_score(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=unconditional_dict,
+                timestep=timestep
+            )
+            pred_fake_image = pred_fake_image_cond + (
+                pred_fake_image_cond - pred_fake_image_uncond
+            ) * self.fake_guidance_scale
+        else:
+            pred_fake_image = pred_fake_image_cond
+
+        # Step 2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep
+        )
+
+        pred_real_image = pred_real_image_cond + (
+            pred_real_image_cond - pred_real_image_uncond
+        ) * self.real_guidance_scale
+
+        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
+        grad = (pred_fake_image - pred_real_image)
+
+        # TODO: Change the normalizer for causal teacher
+        if normalization:
+            # Step 4: Gradient normalization (DMD paper eq. 8).
+            p_real = (estimated_clean_image_or_video - pred_real_image)
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+            grad = grad / normalizer
+        grad = torch.nan_to_num(grad)
+
+        return grad, {
+            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
+            "timestep": timestep.detach()
+        }
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: Optional[torch.Tensor] = None,
+        denoised_timestep_from: int = 0,
+        denoised_timestep_to: int = 0
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        with torch.no_grad():
+            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+            min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+            max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+            timestep = self._get_timestep(
+                min_timestep,
+                max_timestep,
+                batch_size,
+                num_frame,
+                self.num_frame_per_block,
+                uniform_timestep=True
+            )
+
+            # TODO:should we change it to `timestep = self.scheduler.timesteps[timestep]`?
+            if self.timestep_shift > 1:
+                timestep = self.timestep_shift * \
+                    (timestep / 1000) / \
+                    (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
+            timestep = timestep.clamp(self.min_step, self.max_step)
+
+            noise = torch.randn_like(image_or_video)
+            noisy_latent = self.scheduler.add_noise(
+                image_or_video.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep.flatten(0, 1)
+            ).detach().unflatten(0, (batch_size, num_frame))
+
+            # Step 2: Compute the KL grad
+            grad, dmd_log_dict = self._compute_kl_grad(
+                noisy_image_or_video=noisy_latent,
+                estimated_clean_image_or_video=original_latent,
+                timestep=timestep,
+                conditional_dict=conditional_dict,
+                unconditional_dict=unconditional_dict
+            )
+
+        if gradient_mask is not None:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            )[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean")
+        else:
+            dmd_loss = 0.5 * F.mse_loss(original_latent.double(
+            ), (original_latent.double() - grad.double()).detach(), reduction="mean")
+        return dmd_loss, dmd_log_dict
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+            denoised_timestep_from=denoised_timestep_from,
+            denoised_timestep_to=denoised_timestep_to
+        )
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent
+            )
+
+        # Step 2: Compute the fake prediction
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = self.timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred
+        )
+
+        # Step 5: Debugging Log
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach()
+        }
+
+        return denoising_loss, critic_log_dict

long_video/model/gan.py

@@ -0,0 +1,295 @@
+import copy
+from pipeline import RollingForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import RollingForcingModel
+
+
+class GAN(RollingForcingModel):
+    def __init__(self, args, device):
+        """
+        Initialize the GAN module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
+        self.concat_time_embeddings = getattr(args, "concat_time_embeddings", False)
+        self.num_class = args.num_class
+        self.relativistic_discriminator = getattr(args, "relativistic_discriminator", False)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.fake_score.adding_cls_branch(
+            atten_dim=1536, num_class=args.num_class, time_embed_dim=1536 if self.concat_time_embeddings else 0)
+        self.fake_score.model.requires_grad_(True)
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: RollingForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.critic_timestep_shift = getattr(args, "critic_timestep_shift", self.timestep_shift)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+        self.min_score_timestep = getattr(args, "min_score_timestep", 0)
+
+        self.gan_g_weight = getattr(args, "gan_g_weight", 1e-2)
+        self.gan_d_weight = getattr(args, "gan_d_weight", 1e-2)
+        self.r1_weight = getattr(args, "r1_weight", 0.0)
+        self.r2_weight = getattr(args, "r2_weight", 0.0)
+        self.r1_sigma = getattr(args, "r1_sigma", 0.01)
+        self.r2_sigma = getattr(args, "r2_sigma", 0.01)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _run_cls_pred_branch(self,
+                             noisy_image_or_video: torch.Tensor,
+                             conditional_dict: dict,
+                             timestep: torch.Tensor) -> torch.Tensor:
+        """
+            Run the classifier prediction branch on the generated image or video.
+            Input:
+                - image_or_video: a tensor with shape [B, F, C, H, W].
+            Output:
+                - cls_pred: a tensor with shape [B, 1, 1, 1, 1] representing the feature map for classification.
+        """
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings
+        )
+
+        return noisy_logit
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent
+        )
+
+        # Step 2: Get timestep and add noise to generated/real latents
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.critic_timestep_shift > 1:
+            critic_timestep = self.critic_timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(pred_image)
+        noisy_fake_latent = self.scheduler.add_noise(
+            pred_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        # Step 4: Compute the real GAN discriminator loss
+        real_image_or_video = clean_latent.clone()
+        critic_noise = torch.randn_like(real_image_or_video)
+        noisy_real_latent = self.scheduler.add_noise(
+            real_image_or_video.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        conditional_dict["prompt_embeds"] = torch.concatenate(
+            (conditional_dict["prompt_embeds"], conditional_dict["prompt_embeds"]), dim=0)
+        critic_timestep = torch.concatenate((critic_timestep, critic_timestep), dim=0)
+        noisy_latent = torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0)
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=noisy_latent,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep,
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings
+        )
+        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)
+
+        if not self.relativistic_discriminator:
+            gan_G_loss = F.softplus(-noisy_fake_logit.float()).mean() * self.gan_g_weight
+        else:
+            relative_fake_logit = noisy_fake_logit - noisy_real_logit
+            gan_G_loss = F.softplus(-relative_fake_logit.float()).mean() * self.gan_g_weight
+
+        return gan_G_loss
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        real_image_or_video: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to, num_sim_steps = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent
+            )
+
+        # Step 2: Get timestep and add noise to generated/real latents
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.critic_timestep_shift > 1:
+            critic_timestep = self.critic_timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_fake_latent = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        # Step 4: Compute the real GAN discriminator loss
+        noisy_real_latent = self.scheduler.add_noise(
+            real_image_or_video.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        conditional_dict_cloned = copy.deepcopy(conditional_dict)
+        conditional_dict_cloned["prompt_embeds"] = torch.concatenate(
+            (conditional_dict_cloned["prompt_embeds"], conditional_dict_cloned["prompt_embeds"]), dim=0)
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0),
+            conditional_dict=conditional_dict_cloned,
+            timestep=torch.concatenate((critic_timestep, critic_timestep), dim=0),
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings
+        )
+        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)
+
+        if not self.relativistic_discriminator:
+            gan_D_loss = F.softplus(-noisy_real_logit.float()).mean() + F.softplus(noisy_fake_logit.float()).mean()
+        else:
+            relative_real_logit = noisy_real_logit - noisy_fake_logit
+            gan_D_loss = F.softplus(-relative_real_logit.float()).mean()
+        gan_D_loss = gan_D_loss * self.gan_d_weight
+
+        # R1 regularization
+        if self.r1_weight > 0.:
+            noisy_real_latent_perturbed = noisy_real_latent.clone()
+            epison_real = self.r1_sigma * torch.randn_like(noisy_real_latent_perturbed)
+            noisy_real_latent_perturbed = noisy_real_latent_perturbed + epison_real
+            noisy_real_logit_perturbed = self._run_cls_pred_branch(
+                noisy_image_or_video=noisy_real_latent_perturbed,
+                conditional_dict=conditional_dict,
+                timestep=critic_timestep
+            )
+
+            r1_grad = (noisy_real_logit_perturbed - noisy_real_logit) / self.r1_sigma
+            r1_loss = self.r1_weight * torch.mean((r1_grad)**2)
+        else:
+            r1_loss = torch.zeros_like(gan_D_loss)
+
+        # R2 regularization
+        if self.r2_weight > 0.:
+            noisy_fake_latent_perturbed = noisy_fake_latent.clone()
+            epison_generated = self.r2_sigma * torch.randn_like(noisy_fake_latent_perturbed)
+            noisy_fake_latent_perturbed = noisy_fake_latent_perturbed + epison_generated
+            noisy_fake_logit_perturbed = self._run_cls_pred_branch(
+                noisy_image_or_video=noisy_fake_latent_perturbed,
+                conditional_dict=conditional_dict,
+                timestep=critic_timestep
+            )
+
+            r2_grad = (noisy_fake_logit_perturbed - noisy_fake_logit) / self.r2_sigma
+            r2_loss = self.r2_weight * torch.mean((r2_grad)**2)
+        else:
+            r2_loss = torch.zeros_like(r2_loss)
+
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach(),
+            'noisy_real_logit': noisy_real_logit.detach(),
+            'noisy_fake_logit': noisy_fake_logit.detach(),
+        }
+
+        return (gan_D_loss, r1_loss, r2_loss), critic_log_dict

long_video/model/ode_regression.py

@@ -0,0 +1,136 @@
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class ODERegression(BaseModel):
+    def __init__(self, args, device):
+        """
+        Initialize the ODERegression module.
+        This class is self-contained and compute generator losses
+        in the forward pass given precomputed ode solution pairs.
+        This class supports the ode regression loss for both causal and bidirectional models.
+        See Sec 4.3 of CausVid https://arxiv.org/abs/2412.07772 for details
+        """
+        super().__init__(args, device)
+        # Step 1: Initialize all models
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+        if getattr(args, "generator_ckpt", False):
+            print(f"Loading pretrained generator from {args.generator_ckpt}")
+            state_dict = torch.load(args.generator_ckpt, map_location="cpu")[
+                'generator']
+            self.generator.load_state_dict(
+                state_dict, strict=True
+            )
+
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all hyperparameters
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        
+
+    def _initialize_models(self, args, device):
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+        
+        self.scheduler = self.generator.get_scheduler()
+        self.scheduler.timesteps = self.scheduler.timesteps.to(device)
+
+    @torch.no_grad()
+    def _prepare_generator_input(self, ode_latent: torch.Tensor, tf=False, causal = True) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Given a tensor containing the whole ODE sampling trajectories,
+        randomly choose an intermediate timestep and return the latent as well as the corresponding timestep.
+        Input:
+            - ode_latent: a tensor containing the whole ODE sampling trajectories [batch_size, num_denoising_steps, num_frames, num_channels, height, width].
+        Output:
+            - noisy_input: a tensor containing the selected latent [batch_size, num_frames, num_channels, height, width].
+            - timestep: a tensor containing the corresponding timestep [batch_size].
+        """
+        batch_size, num_denoising_steps, num_frames, num_channels, height, width = ode_latent.shape
+
+        # Step 1: Randomly choose a timestep for each frame
+        uniform_timestep = True
+        index = self._get_timestep(
+            0,
+            len(self.denoising_step_list),
+            batch_size,
+            num_frames,
+            self.num_frame_per_block,
+            uniform_timestep=uniform_timestep
+        )
+
+        if self.args.i2v:
+            index[:, 0] = len(self.denoising_step_list) - 1
+
+        noisy_input = torch.gather(
+            ode_latent, dim=1,
+            index=index.reshape(batch_size, 1, num_frames, 1, 1, 1).expand(
+                -1, -1, -1, num_channels, height, width).to(self.device)
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index].to(self.device)
+
+        return noisy_input, timestep
+
+    def generator_loss(self, ode_latent: torch.Tensor, conditional_dict: dict) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noisy latents and compute the ODE regression loss.
+        Input:
+            - ode_latent: a tensor containing the ODE latents [batch_size, num_denoising_steps, num_frames, num_channels, height, width].
+            They are ordered from most noisy to clean latents.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - log_dict: a dictionary containing additional information for loss timestep breakdown.
+        """
+        # Step 1: Run generator on noisy latents
+        clean_latent = ode_latent[:, -1]
+        target_latent = ode_latent[:, -2]
+        ode_latent_valid = ode_latent[:, :-1]
+        
+            
+        noisy_input, timestep = self._prepare_generator_input(
+            ode_latent=ode_latent_valid, tf=True, causal = True)
+        
+        
+        _, pred_image_or_video = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x = clean_latent
+        )
+
+        # Step 2: Compute the regression loss
+        mask = timestep != 0
+
+        loss = F.mse_loss(
+            pred_image_or_video[mask], target_latent[mask], reduction="mean")
+
+        log_dict = {
+            "unnormalized_loss": F.mse_loss(pred_image_or_video, target_latent, reduction='none').mean(dim=[1, 2, 3, 4]).detach(),
+            "timestep": timestep.float().mean(dim=1).detach(),
+            "input": noisy_input.detach(),
+            "output": pred_image_or_video.detach(),
+        }
+
+        return loss, log_dict

long_video/model/sid.py

@@ -0,0 +1,283 @@
+from pipeline import RollingForcingTrainingPipeline
+from typing import Optional, Tuple
+import torch
+
+from model.base import RollingForcingModel
+
+
+class SiD(RollingForcingModel):
+    def __init__(self, args, device):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+            self.real_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: RollingForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.sid_alpha = getattr(args, "sid_alpha", 1.0)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: Optional[torch.Tensor] = None,
+        denoised_timestep_from: int = 0,
+        denoised_timestep_to: int = 0
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            batch_size,
+            num_frame,
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            timestep = self.timestep_shift * \
+                (timestep / 1000) / \
+                (1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
+        timestep = timestep.clamp(self.min_step, self.max_step)
+
+        noise = torch.randn_like(image_or_video)
+        noisy_latent = self.scheduler.add_noise(
+            image_or_video.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep.flatten(0, 1)
+        ).unflatten(0, (batch_size, num_frame))
+
+        # Step 2: SiD (May be wrap it?)
+        noisy_image_or_video = noisy_latent
+        # Step 2.1: Compute the fake score
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+        # Step 2.2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        # NOTE: This step may cause OOM issue, which can be addressed by the CFG-free technique
+
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep
+        )
+
+        pred_real_image = pred_real_image_cond + (
+            pred_real_image_cond - pred_real_image_uncond
+        ) * self.real_guidance_scale
+
+        # Step 2.3: SiD Loss
+        # TODO: Add alpha
+        # TODO: Double?
+        sid_loss = (pred_real_image.double() - pred_fake_image.double()) * ((pred_real_image.double() - original_latent.double()) - self.sid_alpha * (pred_real_image.double() - pred_fake_image.double()))
+
+        # Step 2.4: Loss normalizer
+        with torch.no_grad():
+            p_real = (original_latent - pred_real_image)
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+        sid_loss = sid_loss / normalizer
+
+        sid_loss = torch.nan_to_num(sid_loss)
+        num_frame = sid_loss.shape[1]
+        sid_loss = sid_loss.mean()
+
+        sid_log_dict = {
+            "dmdtrain_gradient_norm": torch.zeros_like(sid_loss),
+            "timestep": timestep.detach()
+        }
+
+        return sid_loss, sid_log_dict
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+            denoised_timestep_from=denoised_timestep_from,
+            denoised_timestep_to=denoised_timestep_to
+        )
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent
+            )
+
+        # Step 2: Compute the fake prediction
+        min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
+        max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = self.timestep_shift * \
+                (critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1)
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1)
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred
+        )
+
+        # Step 5: Debugging Log
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach()
+        }
+
+        return denoising_loss, critic_log_dict

long_video/pipeline/__init__.py

@@ -0,0 +1,13 @@
+from .bidirectional_diffusion_inference import BidirectionalDiffusionInferencePipeline
+from .bidirectional_inference import BidirectionalInferencePipeline
+from .causal_diffusion_inference import CausalDiffusionInferencePipeline
+from .rolling_forcing_inference import CausalInferencePipeline
+from .rolling_forcing_training import RollingForcingTrainingPipeline
+
+__all__ = [
+    "BidirectionalDiffusionInferencePipeline",
+    "BidirectionalInferencePipeline",
+    "CausalDiffusionInferencePipeline",
+    "CausalInferencePipeline",
+    "RollingForcingTrainingPipeline"
+]

long_video/pipeline/bidirectional_diffusion_inference.py

@@ -0,0 +1,110 @@
+from tqdm import tqdm
+from typing import List
+import torch
+
+from wan.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BidirectionalDiffusionInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}), is_causal=False) if generator is None else generator
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize scheduler
+        self.num_train_timesteps = args.num_train_timestep
+        self.sampling_steps = 50
+        self.sample_solver = 'unipc'
+        self.shift = 8.0
+
+        self.args = args
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        return_latents=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_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
+        """
+
+        conditional_dict = self.text_encoder(
+            text_prompts=text_prompts
+        )
+        unconditional_dict = self.text_encoder(
+            text_prompts=[self.args.negative_prompt] * len(text_prompts)
+        )
+
+        latents = noise
+
+        sample_scheduler = self._initialize_sample_scheduler(noise)
+        for _, t in enumerate(tqdm(sample_scheduler.timesteps)):
+            latent_model_input = latents
+            timestep = t * torch.ones([latents.shape[0], 21], device=noise.device, dtype=torch.float32)
+
+            flow_pred_cond, _ = self.generator(latent_model_input, conditional_dict, timestep)
+            flow_pred_uncond, _ = self.generator(latent_model_input, unconditional_dict, timestep)
+
+            flow_pred = flow_pred_uncond + self.args.guidance_scale * (
+                flow_pred_cond - flow_pred_uncond)
+
+            temp_x0 = sample_scheduler.step(
+                flow_pred.unsqueeze(0),
+                t,
+                latents.unsqueeze(0),
+                return_dict=False)[0]
+            latents = temp_x0.squeeze(0)
+
+        x0 = latents
+        video = self.vae.decode_to_pixel(x0)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        del sample_scheduler
+
+        if return_latents:
+            return video, latents
+        else:
+            return video
+
+    def _initialize_sample_scheduler(self, noise):
+        if self.sample_solver == 'unipc':
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False)
+            sample_scheduler.set_timesteps(
+                self.sampling_steps, device=noise.device, shift=self.shift)
+            self.timesteps = sample_scheduler.timesteps
+        elif self.sample_solver == 'dpm++':
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False)
+            sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)
+            self.timesteps, _ = retrieve_timesteps(
+                sample_scheduler,
+                device=noise.device,
+                sigmas=sampling_sigmas)
+        else:
+            raise NotImplementedError("Unsupported solver.")
+        return sample_scheduler

long_video/pipeline/bidirectional_inference.py

@@ -0,0 +1,71 @@
+from typing import List
+import torch
+
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BidirectionalInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}), is_causal=False) if generator is None else generator
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize all bidirectional wan hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        self.denoising_step_list = torch.tensor(
+            args.denoising_step_list, dtype=torch.long, device=device)
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[:-1]  # remove the zero timestep for inference
+        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]
+
+    def inference(self, noise: torch.Tensor, text_prompts: List[str]) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
+        """
+        conditional_dict = self.text_encoder(
+            text_prompts=text_prompts
+        )
+
+        # initial point
+        noisy_image_or_video = noise
+
+        # use the last n-1 timesteps to simulate the generator's input
+        for index, current_timestep in enumerate(self.denoising_step_list[:-1]):
+            _, pred_image_or_video = self.generator(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=conditional_dict,
+                timestep=torch.ones(
+                    noise.shape[:2], dtype=torch.long, device=noise.device) * current_timestep
+            )  # [B, F, C, H, W]
+
+            next_timestep = self.denoising_step_list[index + 1] * torch.ones(
+                noise.shape[:2], dtype=torch.long, device=noise.device)
+
+            noisy_image_or_video = self.scheduler.add_noise(
+                pred_image_or_video.flatten(0, 1),
+                torch.randn_like(pred_image_or_video.flatten(0, 1)),
+                next_timestep.flatten(0, 1)
+            ).unflatten(0, noise.shape[:2])
+
+        video = self.vae.decode_to_pixel(pred_image_or_video)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        return video

long_video/pipeline/causal_diffusion_inference.py

@@ -0,0 +1,342 @@
+from tqdm import tqdm
+from typing import List, Optional
+import torch
+
+from wan.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class CausalDiffusionInferencePipeline(torch.nn.Module):
+    def __init__(
+            self,
+            args,
+            device,
+            generator=None,
+            text_encoder=None,
+            vae=None
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}), is_causal=True) if generator is None else generator
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize scheduler
+        self.num_train_timesteps = args.num_train_timestep
+        self.sampling_steps = 50
+        self.sample_solver = 'unipc'
+        self.shift = args.timestep_shift
+
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+
+        self.kv_cache_pos = None
+        self.kv_cache_neg = None
+        self.crossattn_cache_pos = None
+        self.crossattn_cache_neg = None
+        self.args = args
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = self.generator.model.local_attn_size
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1:
+            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,
+        return_latents: bool = False,
+        start_frame_index: Optional[int] = 0
+    ) -> 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.
+            start_frame_index (int): In long video generation, where does the current window start?
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_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
+        elif self.independent_first_frame and initial_latent is None:
+            # 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 = self.text_encoder(
+            text_prompts=[self.args.negative_prompt] * len(text_prompts)
+        )
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        if self.kv_cache_pos 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_pos[block_index]["is_init"] = False
+                self.crossattn_cache_neg[block_index]["is_init"] = False
+            # reset kv cache
+            for block_index in range(len(self.kv_cache_pos)):
+                self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_neg[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_neg[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = start_frame_index
+        cache_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_cache_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=unconditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_neg,
+                    crossattn_cache=self.crossattn_cache_neg,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                current_start_frame += 1
+                cache_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 block_index in range(num_input_blocks):
+                current_ref_latents = \
+                    initial_latent[:, cache_start_frame:cache_start_frame + self.num_frame_per_block]
+                output[:, cache_start_frame:cache_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_cache_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=unconditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_neg,
+                    crossattn_cache=self.crossattn_cache_neg,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                current_start_frame += self.num_frame_per_block
+                cache_start_frame += self.num_frame_per_block
+
+        # 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:
+            noisy_input = noise[
+                :, cache_start_frame - num_input_frames:cache_start_frame + current_num_frames - num_input_frames]
+            latents = noisy_input
+
+            # Step 3.1: Spatial denoising loop
+            sample_scheduler = self._initialize_sample_scheduler(noise)
+            for _, t in enumerate(tqdm(sample_scheduler.timesteps)):
+                latent_model_input = latents
+                timestep = t * torch.ones(
+                    [batch_size, current_num_frames], device=noise.device, dtype=torch.float32
+                )
+
+                flow_pred_cond, _ = self.generator(
+                    noisy_image_or_video=latent_model_input,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+                flow_pred_uncond, _ = self.generator(
+                    noisy_image_or_video=latent_model_input,
+                    conditional_dict=unconditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache_neg,
+                    crossattn_cache=self.crossattn_cache_neg,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length
+                )
+
+                flow_pred = flow_pred_uncond + self.args.guidance_scale * (
+                    flow_pred_cond - flow_pred_uncond)
+
+                temp_x0 = sample_scheduler.step(
+                    flow_pred,
+                    t,
+                    latents,
+                    return_dict=False)[0]
+                latents = temp_x0
+                print(f"kv_cache['local_end_index']: {self.kv_cache_pos[0]['local_end_index']}")
+                print(f"kv_cache['global_end_index']: {self.kv_cache_pos[0]['global_end_index']}")
+
+            # Step 3.2: record the model's output
+            output[:, cache_start_frame:cache_start_frame + current_num_frames] = latents
+
+            # Step 3.3: rerun with timestep zero to update KV cache using clean context
+            self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=conditional_dict,
+                timestep=timestep * 0,
+                kv_cache=self.kv_cache_pos,
+                crossattn_cache=self.crossattn_cache_pos,
+                current_start=current_start_frame * self.frame_seq_length,
+                cache_start=cache_start_frame * self.frame_seq_length
+            )
+            self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=unconditional_dict,
+                timestep=timestep * 0,
+                kv_cache=self.kv_cache_neg,
+                crossattn_cache=self.crossattn_cache_neg,
+                current_start=current_start_frame * self.frame_seq_length,
+                cache_start=cache_start_frame * self.frame_seq_length
+            )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+            cache_start_frame += current_num_frames
+
+        # Step 4: Decode the output
+        video = self.vae.decode_to_pixel(output)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        if return_latents:
+            return video, output
+        else:
+            return video
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache_pos = []
+        kv_cache_neg = []
+        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_cache_pos.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)
+            })
+            kv_cache_neg.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_cache_pos = kv_cache_pos  # always store the clean cache
+        self.kv_cache_neg = kv_cache_neg  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache_pos = []
+        crossattn_cache_neg = []
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache_pos.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
+            })
+            crossattn_cache_neg.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_pos = crossattn_cache_pos  # always store the clean cache
+        self.crossattn_cache_neg = crossattn_cache_neg  # always store the clean cache
+
+    def _initialize_sample_scheduler(self, noise):
+        if self.sample_solver == 'unipc':
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False)
+            sample_scheduler.set_timesteps(
+                self.sampling_steps, device=noise.device, shift=self.shift)
+            self.timesteps = sample_scheduler.timesteps
+        elif self.sample_solver == 'dpm++':
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False)
+            sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)
+            self.timesteps, _ = retrieve_timesteps(
+                sample_scheduler,
+                device=noise.device,
+                sigmas=sampling_sigmas)
+        else:
+            raise NotImplementedError("Unsupported solver.")
+        return sample_scheduler

long_video/pipeline/rolling_forcing_inference.py

@@ -0,0 +1,372 @@
+from typing import List, Optional
+import torch
+
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+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 = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}), is_causal=True) if generator is None else generator
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize all causal hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        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 = 30
+        self.frame_seq_length = 1560
+
+        self.kv_cache_clean = None
+        self.args = args
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = self.generator.model.local_attn_size
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def inference_rolling_forcing(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        initial_latent: Optional[torch.Tensor] = None,
+        return_latents: bool = False,
+        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
+        )
+
+        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 self.kv_cache_clean 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_cache_clean)):
+                self.kv_cache_clean[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache_clean[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        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_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+                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_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += self.num_frame_per_block
+
+        if profile:
+            init_end.record()
+            torch.cuda.synchronize()
+            diffusion_start.record()
+
+        # implementing rolling forcing 
+        # construct the rolling forcing windows
+        num_denoising_steps = len(self.denoising_step_list)
+        rolling_window_length_blocks = num_denoising_steps
+        window_start_blocks = []
+        window_end_blocks = []
+        window_num = num_blocks + rolling_window_length_blocks - 1
+
+        for window_index in range(window_num):
+            start_block = max(0, window_index - rolling_window_length_blocks + 1)
+            end_block = min(num_blocks - 1, window_index)
+            window_start_blocks.append(start_block)
+            window_end_blocks.append(end_block)
+
+        # init noisy cache
+        noisy_cache = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # init denosing timestep, same accross windows
+        shared_timestep = torch.ones(
+            [batch_size, rolling_window_length_blocks * self.num_frame_per_block],
+            device=noise.device,
+            dtype=torch.float32)
+        
+        for index, current_timestep in enumerate(reversed(self.denoising_step_list)): # from clean to noisy 
+            shared_timestep[:, index * self.num_frame_per_block:(index + 1) * self.num_frame_per_block] *= current_timestep
+
+
+        # Denoising loop with rolling forcing
+        for window_index in range(window_num):
+
+            if profile:
+                block_start.record()
+
+            print('window_index:', window_index)
+            start_block = window_start_blocks[window_index]
+            end_block = window_end_blocks[window_index] # include
+            print(f"start_block: {start_block}, end_block: {end_block}")
+
+            current_start_frame = start_block * self.num_frame_per_block
+            current_end_frame = (end_block + 1) * self.num_frame_per_block # not include
+            current_num_frames = current_end_frame - current_start_frame
+
+            # noisy_input: new noise and previous denoised noisy frames, only last block is pure noise
+            if current_num_frames == rolling_window_length_blocks * self.num_frame_per_block or current_start_frame == 0:
+                noisy_input = torch.cat([
+                    noisy_cache[:, current_start_frame : current_end_frame - self.num_frame_per_block],
+                    noise[:, current_end_frame - self.num_frame_per_block : current_end_frame ]
+                ], dim=1)
+            else: # at the end of the video
+                noisy_input = noisy_cache[:, current_start_frame:current_end_frame]
+
+            # init denosing timestep
+            if current_num_frames == rolling_window_length_blocks * self.num_frame_per_block:
+                current_timestep = shared_timestep
+            elif current_start_frame == 0:
+                current_timestep = shared_timestep[:,-current_num_frames:]
+            elif current_end_frame == num_frames:
+                current_timestep = shared_timestep[:,:current_num_frames]
+            else:
+                raise ValueError("current_num_frames should be equal to rolling_window_length_blocks * self.num_frame_per_block, or the first or last window.")
+
+
+            # calling DiT
+            _, denoised_pred = self.generator(
+                    noisy_image_or_video=noisy_input,
+                    conditional_dict=conditional_dict,
+                    timestep=current_timestep,
+                    kv_cache=self.kv_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+
+            output[:, current_start_frame:current_end_frame] = denoised_pred
+                
+
+            # update noisy_cache, which is detached from the computation graph
+            with torch.no_grad():
+                for block_idx in range(start_block, end_block + 1):
+                    
+                    block_time_step = current_timestep[:, 
+                                    (block_idx - start_block)*self.num_frame_per_block : 
+                                    (block_idx - start_block+1)*self.num_frame_per_block].mean().item()
+                    matches = torch.abs(self.denoising_step_list - block_time_step) < 1e-4
+                    block_timestep_index = torch.nonzero(matches, as_tuple=True)[0]
+
+                    if block_timestep_index == len(self.denoising_step_list) - 1:
+                        continue
+
+                    next_timestep = self.denoising_step_list[block_timestep_index + 1].to(noise.device)
+
+                    noisy_cache[:, block_idx * self.num_frame_per_block:
+                                    (block_idx+1) * self.num_frame_per_block] = \
+                        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])[:, (block_idx - start_block)*self.num_frame_per_block:
+                                                                    (block_idx - start_block+1)*self.num_frame_per_block]
+
+
+            # rerun with timestep zero to update the clean cache, which is also detached from the computation graph
+            with torch.no_grad():
+                context_timestep = torch.ones_like(current_timestep) * self.args.context_noise
+                # # add context noise
+                # denoised_pred = self.scheduler.add_noise(
+                #     denoised_pred.flatten(0, 1),
+                #     torch.randn_like(denoised_pred.flatten(0, 1)),
+                #     context_timestep * torch.ones(
+                #         [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                # ).unflatten(0, denoised_pred.shape[:2])
+
+                # only cache the first block
+                denoised_pred = denoised_pred[:,:self.num_frame_per_block]
+                context_timestep = context_timestep[:,:self.num_frame_per_block]
+                self.generator(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=context_timestep,
+                    kv_cache=self.kv_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    updating_cache=True,
+                )
+
+            if profile:
+                block_end.record()
+                torch.cuda.synchronize()
+                block_time = block_start.elapsed_time(block_end)
+                block_times.append(block_time)
+
+
+        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
+        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)
+            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
+        else:
+            return video
+
+
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache_clean = []
+        # 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 = 1560 * 24
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_clean.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_cache_clean = kv_cache_clean  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        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

long_video/pipeline/rolling_forcing_training.py

@@ -0,0 +1,464 @@
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+class RollingForcingTrainingPipeline:
+    def __init__(self,
+                 denoising_step_list: List[int],
+                 scheduler: SchedulerInterface,
+                 generator: WanDiffusionWrapper,
+                 num_frame_per_block=3,
+                 independent_first_frame: bool = False,
+                 same_step_across_blocks: bool = False,
+                 last_step_only: bool = False,
+                 num_max_frames: int = 21,
+                 context_noise: int = 0,
+                 **kwargs):
+        super().__init__()
+        self.scheduler = scheduler
+        self.generator = generator
+        self.denoising_step_list = denoising_step_list
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[:-1]  # remove the zero timestep for inference
+
+        # Wan specific hyperparameters
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+        self.num_frame_per_block = num_frame_per_block
+        self.context_noise = context_noise
+        self.i2v = False
+
+        self.kv_cache_clean = None
+        self.kv_cache2 = None
+        self.independent_first_frame = independent_first_frame
+        self.same_step_across_blocks = same_step_across_blocks
+        self.last_step_only = last_step_only
+        self.kv_cache_size = num_max_frames * self.frame_seq_length
+
+    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Generate random indices
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device
+            )
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+    
+    def generate_list(self, num_blocks, num_denoising_steps, device):
+
+        # Generate random indices
+        indices = torch.randint(
+            low=0,
+            high=num_denoising_steps,
+            size=(num_blocks,),
+            device=device
+        )
+        if self.last_step_only:
+            indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+
+        return indices.tolist()    
+
+
+    def inference_with_rolling_forcing(
+            self,
+            noise: torch.Tensor,
+            initial_latent: Optional[torch.Tensor] = None,
+            return_sim_step: bool = False,
+            **conditional_dict
+    ) -> torch.Tensor:
+        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
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        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
+        )
+
+        # implementing rolling forcing 
+        # construct the rolling forcing windows
+        num_denoising_steps = len(self.denoising_step_list)
+        rolling_window_length_blocks = num_denoising_steps
+        window_start_blocks = []
+        window_end_blocks = []
+        window_num = num_blocks + rolling_window_length_blocks - 1
+
+        for window_index in range(window_num):
+            start_block = max(0, window_index - rolling_window_length_blocks + 1)
+            end_block = min(num_blocks - 1, window_index)
+            window_start_blocks.append(start_block)
+            window_end_blocks.append(end_block)
+
+        # exit_flag indicates the window at which the model will backpropagate gradients.
+        exit_flag = torch.randint(high=rolling_window_length_blocks, device=noise.device, size=())
+        start_gradient_frame_index = num_output_frames - 21
+        
+        # init noisy cache
+        noisy_cache = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # init denosing timestep, same accross windows
+        shared_timestep = torch.ones(
+            [batch_size, rolling_window_length_blocks * self.num_frame_per_block],
+            device=noise.device,
+            dtype=torch.float32)
+        
+        for index, current_timestep in enumerate(reversed(self.denoising_step_list)): # from clean to noisy 
+            shared_timestep[:, index * self.num_frame_per_block:(index + 1) * self.num_frame_per_block] *= current_timestep
+
+
+        # Denoising loop with rolling forcing
+        for window_index in range(window_num):
+            start_block = window_start_blocks[window_index]
+            end_block = window_end_blocks[window_index] # include
+
+            current_start_frame = start_block * self.num_frame_per_block
+            current_end_frame = (end_block + 1) * self.num_frame_per_block # not include
+            current_num_frames = current_end_frame - current_start_frame
+
+            # noisy_input: new noise and previous denoised noisy frames, only last block is pure noise
+            if current_num_frames == rolling_window_length_blocks * self.num_frame_per_block or current_start_frame == 0:
+                noisy_input = torch.cat([
+                    noisy_cache[:, current_start_frame : current_end_frame - self.num_frame_per_block],
+                    noise[:, current_end_frame - self.num_frame_per_block : current_end_frame ]
+                ], dim=1)
+            else: # at the end of the video
+                noisy_input = noisy_cache[:, current_start_frame:current_end_frame].clone()
+
+            # init denosing timestep
+            if current_num_frames == rolling_window_length_blocks * self.num_frame_per_block:
+                current_timestep = shared_timestep
+            elif current_start_frame == 0:
+                current_timestep = shared_timestep[:,-current_num_frames:]
+            elif current_end_frame == num_frames:
+                current_timestep = shared_timestep[:,:current_num_frames]
+            else:
+                raise ValueError("current_num_frames should be equal to rolling_window_length_blocks * self.num_frame_per_block, or the first or last window.")
+            
+            require_grad = window_index % rolling_window_length_blocks == exit_flag
+            if current_end_frame <= start_gradient_frame_index:
+                require_grad = False
+
+            # calling DiT
+            if not require_grad:
+                with torch.no_grad():
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=current_timestep,
+                        kv_cache=self.kv_cache_clean,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length
+                    )
+            else:
+                _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=current_timestep,
+                        kv_cache=self.kv_cache_clean,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length
+                    )
+                output[:, current_start_frame:current_end_frame] = denoised_pred
+                
+
+            # update noisy_cache, which is detached from the computation graph
+            with torch.no_grad():
+                for block_idx in range(start_block, end_block + 1):
+                    
+                    block_time_step = current_timestep[:, 
+                                    (block_idx - start_block)*self.num_frame_per_block : 
+                                    (block_idx - start_block+1)*self.num_frame_per_block].mean().item()
+                    matches = torch.abs(self.denoising_step_list - block_time_step) < 1e-4
+                    block_timestep_index = torch.nonzero(matches, as_tuple=True)[0]
+
+                    if block_timestep_index == len(self.denoising_step_list) - 1:
+                        continue
+                    
+                    next_timestep = self.denoising_step_list[block_timestep_index + 1].to(noise.device)
+
+                    noisy_cache[:, block_idx * self.num_frame_per_block:
+                                    (block_idx+1) * self.num_frame_per_block] = \
+                        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])[:, (block_idx - start_block)*self.num_frame_per_block:
+                                                                    (block_idx - start_block+1)*self.num_frame_per_block]
+
+
+            # rerun with timestep zero to update the clean cache, which is also detached from the computation graph
+            with torch.no_grad():
+                context_timestep = torch.ones_like(current_timestep) * self.context_noise
+                # # add context noise
+                # denoised_pred = self.scheduler.add_noise(
+                #     denoised_pred.flatten(0, 1),
+                #     torch.randn_like(denoised_pred.flatten(0, 1)),
+                #     context_timestep * torch.ones(
+                #         [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+                # ).unflatten(0, denoised_pred.shape[:2])
+
+                # only cache the first block
+                denoised_pred = denoised_pred[:,:self.num_frame_per_block]
+                context_timestep = context_timestep[:,:self.num_frame_per_block]
+                self.generator(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=context_timestep,
+                    kv_cache=self.kv_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    updating_cache=True,
+                )
+
+        # Step 3.5: Return the denoised timestep
+        # can ignore since not used
+        denoised_timestep_from, denoised_timestep_to = None, None
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+
+
+    def inference_with_self_forcing(
+            self,
+            noise: torch.Tensor,
+            initial_latent: Optional[torch.Tensor] = None,
+            return_sim_step: bool = False,
+            **conditional_dict
+    ) -> torch.Tensor:
+        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
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        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
+        )
+        # if self.kv_cache_clean 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_cache_clean)):
+        #         self.kv_cache_clean[block_index]["global_end_index"] = torch.tensor(
+        #             [0], dtype=torch.long, device=noise.device)
+        #         self.kv_cache_clean[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
+            # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+            output[:, :1] = initial_latent
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=initial_latent,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+            current_start_frame += 1
+
+        # 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
+        num_denoising_steps = len(self.denoising_step_list)
+        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
+        start_gradient_frame_index = num_output_frames - 21
+
+        # for block_index in range(num_blocks):
+        for block_index, current_num_frames in enumerate(all_num_frames):
+            noisy_input = noise[
+                :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                if self.same_step_across_blocks:
+                    exit_flag = (index == exit_flags[0])
+                else:
+                    exit_flag = (index == exit_flags[block_index])  # Only backprop at the randomly selected timestep (consistent across all ranks)
+                timestep = torch.ones(
+                    [batch_size, current_num_frames],
+                    device=noise.device,
+                    dtype=torch.int64) * current_timestep
+
+                if not exit_flag:
+                    with torch.no_grad():
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache_clean,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length
+                        )
+                        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:
+                    # for getting real output
+                    # with torch.set_grad_enabled(current_start_frame >= start_gradient_frame_index):
+                    if current_start_frame < start_gradient_frame_index:
+                        with torch.no_grad():
+                            _, denoised_pred = self.generator(
+                                noisy_image_or_video=noisy_input,
+                                conditional_dict=conditional_dict,
+                                timestep=timestep,
+                                kv_cache=self.kv_cache_clean,
+                                crossattn_cache=self.crossattn_cache,
+                                current_start=current_start_frame * self.frame_seq_length
+                            )
+                    else:
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache_clean,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length
+                        )
+                    break
+
+            # Step 3.2: record the model's output
+            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Step 3.3: rerun with timestep zero to update the cache
+            context_timestep = torch.ones_like(timestep) * self.context_noise
+            # add context noise
+            denoised_pred = self.scheduler.add_noise(
+                denoised_pred.flatten(0, 1),
+                torch.randn_like(denoised_pred.flatten(0, 1)),
+                context_timestep * torch.ones(
+                    [batch_size * current_num_frames], device=noise.device, dtype=torch.long)
+            ).unflatten(0, denoised_pred.shape[:2])
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=context_timestep,
+                    kv_cache=self.kv_cache_clean,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    updating_cache=True,
+                )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        # Step 3.5: Return the denoised timestep
+        if not self.same_step_across_blocks:
+            denoised_timestep_from, denoised_timestep_to = None, None
+        elif exit_flags[0] == len(self.denoising_step_list) - 1:
+            denoised_timestep_to = 0
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+        else:
+            denoised_timestep_to = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item()
+            denoised_timestep_from = 1000 - torch.argmin(
+                (self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item()
+
+        if return_sim_step:
+            return output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+    def _initialize_kv_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache_clean = []
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_clean.append({
+                "k": torch.zeros([batch_size, self.kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, self.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_cache_clean = kv_cache_clean  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        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

long_video/requirements.txt

@@ -0,0 +1,44 @@
+torch==2.5.1
+torchvision==0.20.1
+torchaudio==2.5.1
+opencv-python>=4.9.0.80
+diffusers==0.31.0
+transformers>=4.49.0
+tokenizers>=0.20.3
+accelerate>=1.1.1
+tqdm
+imageio
+easydict
+ftfy
+dashscope
+imageio-ffmpeg
+numpy==1.24.4
+wandb
+omegaconf
+einops
+av==13.1.0
+opencv-python
+open_clip_torch
+starlette
+pycocotools
+lmdb
+matplotlib
+sentencepiece
+pydantic==2.10.6
+scikit-image
+huggingface_hub[cli]
+dominate
+nvidia-pyindex
+nvidia-tensorrt
+pycuda
+onnx
+onnxruntime
+onnxscript
+onnxconverter_common
+flask
+flask-socketio
+torchao
+tensorboard
+ninja
+packaging
+gradio>=4.44.0

long_video/train.py

@@ -0,0 +1,45 @@
+import argparse
+import os
+from omegaconf import OmegaConf
+
+from trainer import DiffusionTrainer, GANTrainer, ODETrainer, ScoreDistillationTrainer
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True)
+    parser.add_argument("--no_save", action="store_true")
+    parser.add_argument("--no_visualize", action="store_true")
+    parser.add_argument("--logdir", type=str, default="", help="Path to the directory to save logs")
+    parser.add_argument("--wandb-save-dir", type=str, default="", help="Path to the directory to save wandb logs")
+    parser.add_argument("--disable-wandb", default=False, action="store_true")
+
+    args = parser.parse_args()
+
+    config = OmegaConf.load(args.config_path)
+    default_config = OmegaConf.load("configs/default_config.yaml")
+    config = OmegaConf.merge(default_config, config)
+    config.no_save = args.no_save
+    config.no_visualize = args.no_visualize
+
+    # get the filename of config_path
+    config_name = os.path.basename(args.config_path).split(".")[0]
+    config.config_name = config_name
+    config.logdir = args.logdir
+    config.wandb_save_dir = args.wandb_save_dir
+    config.disable_wandb = args.disable_wandb
+
+    if config.trainer == "diffusion":
+        trainer = DiffusionTrainer(config)
+    elif config.trainer == "gan":
+        trainer = GANTrainer(config)
+    elif config.trainer == "ode":
+        trainer = ODETrainer(config)
+    elif config.trainer == "score_distillation":
+        trainer = ScoreDistillationTrainer(config)
+    trainer.train()
+
+
+
+if __name__ == "__main__":
+    main()

long_video/trainer/__init__.py

@@ -0,0 +1,11 @@
+from .diffusion import Trainer as DiffusionTrainer
+from .gan import Trainer as GANTrainer
+from .ode import Trainer as ODETrainer
+from .distillation import Trainer as ScoreDistillationTrainer
+
+__all__ = [
+    "DiffusionTrainer",
+    "GANTrainer",
+    "ODETrainer",
+    "ScoreDistillationTrainer"
+]

long_video/trainer/diffusion.py

@@ -0,0 +1,265 @@
+import gc
+import logging
+
+from model import CausalDiffusion
+from utils.dataset import ShardingLMDBDataset, cycle
+from utils.misc import set_seed
+import torch.distributed as dist
+from omegaconf import OmegaConf
+import torch
+import wandb
+import time
+import os
+
+from utils.distributed import EMA_FSDP, barrier, fsdp_wrap, fsdp_state_dict, launch_distributed_job
+
+
+class Trainer:
+    def __init__(self, config):
+        self.config = config
+        self.step = 0
+
+        # Step 1: Initialize the distributed training environment (rank, seed, dtype, logging etc.)
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+        launch_distributed_job()
+        global_rank = dist.get_rank()
+
+        self.dtype = torch.bfloat16 if config.mixed_precision else torch.float32
+        self.device = torch.cuda.current_device()
+        self.is_main_process = global_rank == 0
+        self.causal = config.causal
+        self.disable_wandb = config.disable_wandb
+
+        # use a random seed for the training
+        if config.seed == 0:
+            random_seed = torch.randint(0, 10000000, (1,), device=self.device)
+            dist.broadcast(random_seed, src=0)
+            config.seed = random_seed.item()
+
+        set_seed(config.seed + global_rank)
+
+        if self.is_main_process and not self.disable_wandb:
+            wandb.login(host=config.wandb_host, key=config.wandb_key)
+            wandb.init(
+                config=OmegaConf.to_container(config, resolve=True),
+                name=config.config_name,
+                mode="online",
+                entity=config.wandb_entity,
+                project=config.wandb_project,
+                dir=config.wandb_save_dir
+            )
+
+        self.output_path = config.logdir
+
+        # Step 2: Initialize the model and optimizer
+        self.model = CausalDiffusion(config, device=self.device)
+        self.model.generator = fsdp_wrap(
+            self.model.generator,
+            sharding_strategy=config.sharding_strategy,
+            mixed_precision=config.mixed_precision,
+            wrap_strategy=config.generator_fsdp_wrap_strategy
+        )
+
+        self.model.text_encoder = fsdp_wrap(
+            self.model.text_encoder,
+            sharding_strategy=config.sharding_strategy,
+            mixed_precision=config.mixed_precision,
+            wrap_strategy=config.text_encoder_fsdp_wrap_strategy
+        )
+
+        if not config.no_visualize or config.load_raw_video:
+            self.model.vae = self.model.vae.to(
+                device=self.device, dtype=torch.bfloat16 if config.mixed_precision else torch.float32)
+
+        self.generator_optimizer = torch.optim.AdamW(
+            [param for param in self.model.generator.parameters()
+             if param.requires_grad],
+            lr=config.lr,
+            betas=(config.beta1, config.beta2),
+            weight_decay=config.weight_decay
+        )
+
+        # Step 3: Initialize the dataloader
+        dataset = ShardingLMDBDataset(config.data_path, max_pair=int(1e8))
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            dataset, shuffle=True, drop_last=True)
+        dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=config.batch_size,
+            sampler=sampler,
+            num_workers=8)
+
+        if dist.get_rank() == 0:
+            print("DATASET SIZE %d" % len(dataset))
+        self.dataloader = cycle(dataloader)
+
+        ##############################################################################################################
+        # 6. Set up EMA parameter containers
+        rename_param = (
+            lambda name: name.replace("_fsdp_wrapped_module.", "")
+            .replace("_checkpoint_wrapped_module.", "")
+            .replace("_orig_mod.", "")
+        )
+        self.name_to_trainable_params = {}
+        for n, p in self.model.generator.named_parameters():
+            if not p.requires_grad:
+                continue
+
+            renamed_n = rename_param(n)
+            self.name_to_trainable_params[renamed_n] = p
+        ema_weight = config.ema_weight
+        self.generator_ema = None
+        if (ema_weight is not None) and (ema_weight > 0.0):
+            print(f"Setting up EMA with weight {ema_weight}")
+            self.generator_ema = EMA_FSDP(self.model.generator, decay=ema_weight)
+
+        ##############################################################################################################
+        # 7. (If resuming) Load the model and optimizer, lr_scheduler, ema's statedicts
+        if getattr(config, "generator_ckpt", False):
+            print(f"Loading pretrained generator from {config.generator_ckpt}")
+            state_dict = torch.load(config.generator_ckpt, map_location="cpu")
+            if "generator" in state_dict:
+                state_dict = state_dict["generator"]
+            elif "model" in state_dict:
+                state_dict = state_dict["model"]
+            self.model.generator.load_state_dict(
+                state_dict, strict=True
+            )
+
+        ##############################################################################################################
+
+        # Let's delete EMA params for early steps to save some computes at training and inference
+        if self.step < config.ema_start_step:
+            self.generator_ema = None
+
+        self.max_grad_norm = 10.0
+        self.previous_time = None
+
+    def save(self):
+        print("Start gathering distributed model states...")
+        generator_state_dict = fsdp_state_dict(
+            self.model.generator)
+
+        if self.config.ema_start_step < self.step:
+            state_dict = {
+                "generator": generator_state_dict,
+                "generator_ema": self.generator_ema.state_dict(),
+            }
+        else:
+            state_dict = {
+                "generator": generator_state_dict,
+            }
+
+        if self.is_main_process:
+            os.makedirs(os.path.join(self.output_path,
+                        f"checkpoint_model_{self.step:06d}"), exist_ok=True)
+            torch.save(state_dict, os.path.join(self.output_path,
+                       f"checkpoint_model_{self.step:06d}", "model.pt"))
+            print("Model saved to", os.path.join(self.output_path,
+                  f"checkpoint_model_{self.step:06d}", "model.pt"))
+
+    def train_one_step(self, batch):
+        self.log_iters = 1
+
+        if self.step % 20 == 0:
+            torch.cuda.empty_cache()
+
+        # Step 1: Get the next batch of text prompts
+        text_prompts = batch["prompts"]
+        if not self.config.load_raw_video:  # precomputed latent
+            clean_latent = batch["ode_latent"][:, -1].to(
+                device=self.device, dtype=self.dtype)
+        else:  # encode raw video to latent
+            frames = batch["frames"].to(
+                device=self.device, dtype=self.dtype)
+            with torch.no_grad():
+                clean_latent = self.model.vae.encode_to_latent(
+                    frames).to(device=self.device, dtype=self.dtype)
+        image_latent = clean_latent[:, 0:1, ]
+
+        batch_size = len(text_prompts)
+        image_or_video_shape = list(self.config.image_or_video_shape)
+        image_or_video_shape[0] = batch_size
+
+        # Step 2: Extract the conditional infos
+        with torch.no_grad():
+            conditional_dict = self.model.text_encoder(
+                text_prompts=text_prompts)
+
+            if not getattr(self, "unconditional_dict", None):
+                unconditional_dict = self.model.text_encoder(
+                    text_prompts=[self.config.negative_prompt] * batch_size)
+                unconditional_dict = {k: v.detach()
+                                      for k, v in unconditional_dict.items()}
+                self.unconditional_dict = unconditional_dict  # cache the unconditional_dict
+            else:
+                unconditional_dict = self.unconditional_dict
+
+        # Step 3: Train the generator
+        generator_loss, log_dict = self.model.generator_loss(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            clean_latent=clean_latent,
+            initial_latent=image_latent
+        )
+        self.generator_optimizer.zero_grad()
+        generator_loss.backward()
+        generator_grad_norm = self.model.generator.clip_grad_norm_(
+            self.max_grad_norm)
+        self.generator_optimizer.step()
+
+        # Increment the step since we finished gradient update
+        self.step += 1
+
+        wandb_loss_dict = {
+            "generator_loss": generator_loss.item(),
+            "generator_grad_norm": generator_grad_norm.item(),
+        }
+
+        # Step 4: Logging
+        if self.is_main_process:
+            if not self.disable_wandb:
+                wandb.log(wandb_loss_dict, step=self.step)
+
+        if self.step % self.config.gc_interval == 0:
+            if dist.get_rank() == 0:
+                logging.info("DistGarbageCollector: Running GC.")
+            gc.collect()
+
+        # Step 5. Create EMA params
+        # TODO: Implement EMA
+
+    def generate_video(self, pipeline, prompts, image=None):
+        batch_size = len(prompts)
+        sampled_noise = torch.randn(
+            [batch_size, 21, 16, 60, 104], device="cuda", dtype=self.dtype
+        )
+        video, _ = pipeline.inference(
+            noise=sampled_noise,
+            text_prompts=prompts,
+            return_latents=True
+        )
+        current_video = video.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0
+        return current_video
+
+    def train(self):
+        while True:
+            batch = next(self.dataloader)
+            self.train_one_step(batch)
+            if (not self.config.no_save) and self.step % self.config.log_iters == 0:
+                torch.cuda.empty_cache()
+                self.save()
+                torch.cuda.empty_cache()
+
+            barrier()
+            if self.is_main_process:
+                current_time = time.time()
+                if self.previous_time is None:
+                    self.previous_time = current_time
+                else:
+                    if not self.disable_wandb:
+                        wandb.log({"per iteration time": current_time - self.previous_time}, step=self.step)
+                    self.previous_time = current_time