Add isolated Minecraft and RE10K baseline evaluation suite

.gitignore

@@ -0,0 +1,16 @@
+envs/
+cache/
+checkpoints/
+artifacts/
+shared/metric_vendor/
+**/wan_models/
+**/outputs/
+**/build/
+**/dist/
+**/__pycache__/
+**/.pytest_cache/
+**/*.pyc
+**/*.so
+**/*.egg-info/
+**/slurm_logs/*
+!**/slurm_logs/.gitkeep

DecMem/.gitignore

@@ -0,0 +1,25 @@
+__pycache__
+*.egg-info
+
+wan_models
+checkpoints
+
+test/*
+
+videos/*
+
+exps
+
+.vscode/
+
+configs/*_overfit.yaml
+
+data
+
+video_sparse_attn/build
+
+scripts/train_worldmem.sh
+
+scripts/train_local.sh
+
+scripts/train_worldmem_multinode_torchrun.sh

DecMem/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.

DecMem/README.md

@@ -0,0 +1,318 @@
+
+<h1 align="center">DecMem: Towards Minute-Long Consistent World Generation with Decoupled Memory</h1>
+
+
+<div align="center">
+  <p>
+    Zhenhao Yang<sup>1</sup>
+    Xiaoshi Wu<sup>2</sup>
+    Zhengyao Lv<sup>1</sup>
+    Xiaoyu Shi<sup>2,†</sup>
+    Xintao Wang<sup>2</sup>
+    Pengfei Wan<sup>2</sup>
+    Kun Gai<sup>2</sup>
+    Kwan-Yee K. Wong<sup>1,†</sup>
+  </p>
+  <p>
+    <sup>1</sup>The University of Hong Kong &nbsp;&nbsp;
+    <sup>2</sup>Kling Team, Kuaishou Technology<br>
+    <sup>†</sup>Corresponding author
+  </p>
+</div>
+
+<p align="center">
+  <a href='https://jeffreyyzh.github.io/DecMem-Page/'><img src='https://img.shields.io/badge/Project-Page-Green'></a>
+  &nbsp;
+  <a href="https://arxiv.org/abs/2605.31336"><img src="https://img.shields.io/static/v1?label=Arxiv&message=DecMem&color=red&logo=arxiv"></a>
+  &nbsp;
+  <a href='https://huggingface.co/KlingTeam/DecMem/'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-orange'></a>
+  &nbsp;
+  <a href='#-citation'><img src='https://img.shields.io/badge/BibTeX-Cite-blue'></a>
+</p>
+  
+
+## 📋 Table of Contents
+
+- [📋 Table of Contents](#-table-of-contents)
+- [🔥 Updates](#-updates)
+- [📷 Introduction](#-introduction)
+- [⚙️ Code: DecMem + Wan2.1-T2V-1.3B](#️-code-decmem--wan21-t2v-13b)
+  - [Requirements](#requirements)
+  - [Installation](#installation)
+  - [Download Checkpoints](#download-checkpoints)
+  - [Quick Start](#quick-start)
+  - [Data](#data)
+  - [Train](#train)
+  - [Inference](#inference)
+- [🤗 Acknowledgement](#-acknowledgement)
+- [🌟 Citation](#-citation)
+
+**Note:** This open-source repository is a reference implementation built on Wan2.1-T2V-1.3B. The original model is trained on internal pretrained model so we will not opensource such codebase.
+
+## 🔥 Updates
+
+- **[2026.06.01]**: Release the [Training and Inference Code](https://github.com/KlingAIResearch/DecMem) and the [Checkpoints](https://huggingface.co/KlingTeam/DecMem).
+- **[2026.06.01]**: Release the [Project Page](https://jeffreyyzh.github.io/DecMem-Page/) and the [Arxiv](https://arxiv.org/abs/2605.31336) version.
+
+## 📷 Introduction
+
+We systematically reveal the root cause of the limited long-horizon world generation capability of naïve dense-attention designs and propose **DecMem**, a decoupled memory architecture that employs **Sparse Global Memory** for efficient fine-grained access to global history and **Anchored Local Memory** for stable and high-quality extrapolation. 
+
+**(a) Sparse Global Memory (SGM):** Combines a block-level sparse retrieval module and a context-aware attention module for long-term memory fine-grained retrieval in an end-to-end manner, enabling efficient access to global history with bounded cost.
+
+**(b) Anchored Local Memory (ALM):** Keeps short-term transition smooth and supports stable, high-quality extrapolation by anchoring generation on recent local context.
+
+**(c) Capability:** DecMem enables **minute-level** controllable long video generation with **high fidelity**, **consistency** and **efficiency**.
+
+<div align="center">
+
+[![Watch the video](https://img.youtube.com/vi/f6jBmFXxgzg/maxresdefault.jpg)](https://www.youtube.com/watch?v=f6jBmFXxgzg)
+
+</div>
+
+## ⚙️ Code: DecMem + Wan2.1-T2V-1.3B
+
+### Requirements
+
+We tested this repo on the following setup:
+
+- **GPU:** NVIDIA GPU with at least 24 GB VRAM.
+  - **LTM (long-term memory) via Video Sparse Attention requires SM 90a** (**H100 / H200 / H800)**. The block-sparse CUDA kernel is built for `sm_90a` and is not supported on A100 or consumer GPUs.
+  - Dense-only training / inference (`use_ltm_attn: false`) can run on A100 and other CUDA GPUs, but the full DecMem model with LTM needs the hardware above.
+- **RAM:** 64 GB recommended.
+- **CUDA:** 12.x toolkit matching the PyTorch build.
+- **Compiler:** **gcc / g++ ≥ 11** (C++20 required by the VSA kernel).
+
+### Installation
+
+#### 1. Core environment
+
+Create a conda environment and install dependencies:
+
+```bash
+git clone https://github.com/KlingAIResearch/DecMem.git
+cd DecMem
+conda create -n decmem python=3.10 -y
+conda activate decmem
+pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu121
+pip install -r requirements.txt
+pip install flash-attn --no-build-isolation
+```
+
+#### 2. Video Sparse Attention (VSA) kernel
+
+Our LTM module is implemented with modified VSA kernel so it is required when training or inferring with `use_ltm_attn: true` (e.g. `configs/decmem.yaml`). Install from source:
+
+```bash
+cd video_sparse_attn
+python setup.py install
+```
+
+Verify:
+
+```bash
+python -c "
+import torch, vsa
+from vsa.block_sparse_wrapper import block_sparse_attn_SM90
+print('device cap:', torch.cuda.get_device_capability(0))
+print('SM90 op:', block_sparse_attn_SM90)
+"
+```
+
+### Checkpoints
+
+Download the Wan2.1 backbone (VAE + tokenizer weights used by the pipeline):
+
+```bash
+huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B \
+    --local-dir-use-symlinks False \
+    --local-dir wan_models/Wan2.1-T2V-1.3B
+```
+
+Download DecMem trained checkpoints from HuggingFace:
+
+```bash
+huggingface-cli download KlingTeam/DecMem --local-dir checkpoints
+```
+
+Checkpoint layout expected by training / inference scripts:
+
+```
+checkpoints/
+└── decmem.pt             # released weights
+```
+
+### Quick start
+
+We provide the example video-pose pairs for quick inference. The inference is Block-by-block causal denoising manner with KV cache.
+
+```bash
+bash scripts/infer_example.sh
+```
+
+### Data
+
+DecMem trains on **WorldMem** gameplay data: paired MP4 videos and NPZ action/pose files indexed by a CSV.
+
+#### Download dataset
+
+Download the Minecraft dataset from [Hugging Face](https://huggingface.co/datasets/zeqixiao/worldmem_minecraft_dataset) (see the [WorldMem dataset section](https://github.com/xizaoqu/WorldMem#dataset)).
+
+Place the dataset in the following directory structure:
+
+```
+data/
+├── training/
+├── validation/
+└── test/
+```
+
+Each clip is an MP4 file with a matching NPZ file (same basename) in the same directory.
+
+#### Build metadata CSV
+
+After download, run the helper script to scan MP4/NPZ pairs and generate index files:
+
+```bash
+python scripts/build_worldmem_csv.py
+```
+
+This writes `train.csv` (from `training/`) and `test.csv` (from `test/`) under `data/worldmem_minecraft_dataset/`. Rows are shuffled with seed `42` by default; use `--seed` to change it.
+
+#### Directory layout
+
+After download and indexing:
+
+```
+data/
+├── train.csv          # training index (generated)
+├── test.csv           # evaluation index (generated)
+├── training/
+│   └── *.mp4 + *.npz  # paired clips
+├── validation/
+│   └── *.mp4 + *.npz
+└── test/
+    └── *.mp4 + *.npz
+```
+
+#### CSV format
+
+Each row has two columns (`utils/dataset.py` → `WorldMemCSVDataset`):
+
+
+| Column        | Description                    |
+| ------------- | ------------------------------ |
+| `video_path`  | Path to an MP4 video clip.     |
+| `action_path` | Path to the matching NPZ file. |
+
+
+Example (`data/minecraft/train.csv`):
+
+```csv
+video_path,action_path
+data/training/000001.mp4,data/training/000001.npz
+data/training/000002.mp4,data/training/000002.npz
+```
+
+### Train
+
+> **Note:**
+>
+> 1. Update `MASTER_ADDR` (and `NNODES` for multi-node) in the launch scripts before distributed training.
+> 2. Dense training (`use_ltm_attn: false`) can run on A100; LTM training requires H100/H800/H200 with the block-sparse kernel installed.
+> 3. Set `WANDB_KEY` / `wandb_entity` in the config or disable wandb with `--disable-wandb`.
+> 4. We train DecMem in three stages in this opensource version and load the checkpoints in the previous stage for resume training.
+
+#### Stage 1: Dense causal diffusion (Short)
+
+```bash
+# Edit NNODES, MASTER_ADDR, NPROC_PER_NODE in the script first.
+# On node 0:
+bash scripts/train_worldmem_multinode_manual.sh 0 dense_short
+# On node 1:
+bash scripts/train_worldmem_multinode_manual.sh 1 dense_short
+# On node 2:...
+```
+
+#### Stage 2: Dense causal diffusion (Long)
+
+Fine-tune from a short-sequence training dense checkpoint ; set `generator_ckpt` in the config:
+
+```bash
+# Edit NNODES, MASTER_ADDR, NPROC_PER_NODE in the script first.
+# On node 0:
+bash scripts/train_worldmem_multinode_manual.sh 0 dense
+# On node 1:
+bash scripts/train_worldmem_multinode_manual.sh 1 dense
+# On node 2:...
+```
+
+#### Stage 3: Decoupled memory (STM + LTM)
+
+Fine-tune from a long-sequence training dense checkpoint; set `generator_ckpt` in the config:
+
+```bash
+# Edit NNODES, MASTER_ADDR, NPROC_PER_NODE in the script first.
+# On node 0:
+bash scripts/train_worldmem_multinode_manual.sh 0 decmem
+# On node 1:
+bash scripts/train_worldmem_multinode_manual.sh 1 decmem
+# On node 2:...
+```
+
+You can set `RESUME_CKPT` in training scripts to resume from a full training checkpoint (weights + optimizer + step) .
+
+### Inference
+
+Once you finish the training in each stage, you can load the correponding checkpoints and configs for inference.
+
+Common overrides:
+
+```bash
+EXP_NAME=decmem \
+CHECKPOINT_PATH=checkpoints/decmem.pt \
+T_LAT=160 \
+NUM_CONDITION_FRAMES=56 \
+GUIDANCE_SCALE=5.0 \
+NUM_INFERENCE_STEPS=50 \
+bash scripts/infer.sh
+
+# Multi-GPU inference
+NUM_GPUS=8 bash scripts/infer.sh
+```
+
+
+| Argument               | Description                                                           |
+| ---------------------- | --------------------------------------------------------------------- |
+| `NUM_CONDITION_FRAMES` | Number of clean prefix **latent** frames (conditioning video prefix). |
+| `T_LAT`                | Number of clean prefix **latent** frames (conditioning video prefix). |
+| `GUIDANCE_SCALE`       | CFG scale on actions;`≤ 1.0` disables CFG.                            |
+| `N`                    | Number of inference samples                                           |
+| `NUM_INFERENCE_STEPS`  | Number of denoising step                                              |
+
+
+## 🤗 Acknowledgement
+
+This codebase is built on top of:
+
+- [WorldMem](https://github.com/xizaoqu/WorldMem) for action-pose data format and pose utilities.
+- [PRoPE](https://arxiv.org/pdf/2507.10496) for camera-relative positional encoding.
+- [Video Sparse Attention](https://github.com/hao-ai-lab/FastVideo) for the LTM sparse attention kernel.
+- [Self Forcing](https://github.com/gdhe17/Self-Forcing) and [CausVid](https://github.com/tianweiy/CausVid) for causal diffusion training and teacher forcing codebase.
+- [Wan2.1](https://github.com/Wan-Video/Wan2.1) as the base video diffusion backbone.
+
+## 🌟 Citation
+
+Please leave us a star 🌟 and cite our paper if you find our work helpful.
+
+```bibtex
+@misc{yang2026decmemminutelongconsistentworld,
+      title={DecMem: Towards Minute-Long Consistent World Generation with Decoupled Memory}, 
+      author={Zhenhao Yang and Xiaoshi Wu and Zhengyao Lv and Xiaoyu Shi and Xintao Wang and Pengfei Wan and Kun Gai and Kwan-Yee K. Wong},
+      year={2026},
+      eprint={2605.31336},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2605.31336}, 
+}
+```

DecMem/configs/decmem.yaml

@@ -0,0 +1,74 @@
+# WorldMem diffusion training config with Long-Term Memory (LTM)
+# Replaces the flex_attention self-attn path with STM (sliding-window dense via
+# flex_attention) + LTM (sparse block attention on H100 SM90), gated by a
+# zero-init mem_proj so initial step output matches the prior dense behaviour.
+
+trainer: diffusion
+causal: true
+mixed_precision: true
+seed: 0
+
+# Data
+# CSV columns: video_path, action_path  (WorldMem format)
+data_path: data/train.csv
+num_frames: 60          # latent frame count; dataset loads (21-1)*4+1=81 video frames
+sample_strategy: random
+load_raw_video: true
+
+# CFG: probability of dropping action conditioning (zeros actions + identity viewmats)
+proportion_empty_prompts: 0.1
+
+# Model
+image_or_video_shape:
+- 1
+- 60
+- 16
+- 44
+- 80
+
+model_kwargs:
+  timestep_shift: 5.0
+  local_attn_size: -1
+  sink_size: 0
+  use_action_encoder: true
+  action_dim: 25
+  latent_channels: 16
+  # ── LTM ──
+  use_ltm_attn: true
+  ltm_tile_size: [1, 8, 8]
+  ltm_topk: 80                 # top-k key blocks per query block (matches ar_mem default)
+  stm_sliding_window: 8        # STM window = 2 * num_frame_per_block (= 2 * 4 frames)
+
+# Training hyperparams
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 1.0
+teacher_forcing: true
+noise_augmentation_max_timestep: 0
+generator_ckpt: exps/dense/checkpoint_model_010000/model.pt
+
+lr: 2.0e-5
+beta1: 0.9
+beta2: 0.999
+weight_decay: 0.01
+batch_size: 1
+gradient_checkpointing: true
+num_frame_per_block: 4
+sharding_strategy: hybrid_full
+
+# Logging / saving
+no_save: false
+no_visualize: true
+log_iters: 500
+# ema_weight: 0.999
+# ema_start_step: 500
+gc_interval: 100
+wandb_host: "https://api.wandb.ai"
+wandb_key: ""
+wandb_entity: ""
+wandb_project: ""
+negative_prompt: ""
+
+# FSDP wrap strategy
+generator_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size

DecMem/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

DecMem/configs/dense.yaml

@@ -0,0 +1,68 @@
+# WorldMem diffusion training config
+# Action + PRoPE conditioning, no text conditioning, CFG on actions only
+
+trainer: diffusion
+causal: true
+mixed_precision: true
+seed: 0
+
+# Data
+# CSV columns: video_path, action_path  (WorldMem format)
+data_path: data/train_wm_gamecraft/v2/worldmem_train_v2_shuffle.csv
+num_frames: 60          # latent frame count; dataset loads (21-1)*4+1=81 video frames
+load_raw_video: true
+sample_strategy: random  # overfit: always start from SKIP_FRAMES; use "random" for normal training
+
+# CFG: probability of dropping action conditioning (zeros actions + identity viewmats)
+proportion_empty_prompts: 0.1
+
+# Model
+# 81 video frames @ 640x352 -> 21 latent frames @ 80x44 (WAN 4x temporal, 8x spatial)
+image_or_video_shape:
+- 1        # batch_size placeholder (overridden at runtime)
+- 60       # latent frames
+- 16       # latent channels
+- 44       # latent height  (352/8)
+- 80       # latent width   (640/8)
+
+model_kwargs:
+  timestep_shift: 5.0
+  local_attn_size: -1
+  sink_size: 0
+  use_action_encoder: true
+  action_dim: 25
+  latent_channels: 16
+
+# Training hyperparams
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 1.0
+teacher_forcing: true
+noise_augmentation_max_timestep: 0
+generator_ckpt: exps/dense_short/checkpoint_model_022000/model.pt
+
+lr: 2.0e-5
+beta1: 0.9
+beta2: 0.999
+weight_decay: 0.01
+batch_size: 1
+gradient_checkpointing: true
+num_frame_per_block: 4
+sharding_strategy: hybrid_full
+
+# Logging / saving
+no_save: false
+no_visualize: true
+log_iters: 500
+# ema_weight: 0.999
+# ema_start_step: 500
+gc_interval: 100
+wandb_host: "https://api.wandb.ai"
+wandb_key: 
+wandb_entity: ""
+wandb_project: ""
+negative_prompt: ""
+
+# FSDP wrap strategy
+generator_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size

DecMem/configs/dense_short.yaml

@@ -0,0 +1,67 @@
+# WorldMem diffusion training config
+# Action + PRoPE conditioning, no text conditioning, CFG on actions only
+
+trainer: diffusion
+causal: true
+mixed_precision: true
+seed: 0
+
+# Data
+# CSV columns: video_path, action_path  (WorldMem format)
+data_path: data/train_wm_gamecraft/v2/worldmem_train_v2_shuffle.csv
+num_frames: 20          # latent frame count; dataset loads (21-1)*4+1=81 video frames
+load_raw_video: true
+sample_strategy: random  # overfit: always start from SKIP_FRAMES; use "random" for normal training
+
+# CFG: probability of dropping action conditioning (zeros actions + identity viewmats)
+proportion_empty_prompts: 0.1
+
+# Model
+# 81 video frames @ 640x352 -> 21 latent frames @ 80x44 (WAN 4x temporal, 8x spatial)
+image_or_video_shape:
+- 1        # batch_size placeholder (overridden at runtime)
+- 20       # latent frames
+- 16       # latent channels
+- 44       # latent height  (352/8)
+- 80       # latent width   (640/8)
+
+model_kwargs:
+  timestep_shift: 5.0
+  local_attn_size: -1
+  sink_size: 0
+  use_action_encoder: true
+  action_dim: 25
+  latent_channels: 16
+
+# Training hyperparams
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 1.0
+teacher_forcing: true
+noise_augmentation_max_timestep: 0
+
+lr: 1.0e-4
+beta1: 0.9
+beta2: 0.999
+weight_decay: 0.01
+batch_size: 1
+gradient_checkpointing: true
+num_frame_per_block: 4
+sharding_strategy: hybrid_full
+
+# Logging / saving
+no_save: false
+no_visualize: true
+log_iters: 500
+# ema_weight: 0.0
+# ema_start_step: 200
+gc_interval: 100
+wandb_host: "https://api.wandb.ai"
+wandb_key: 
+wandb_entity: ""
+wandb_project: ""
+negative_prompt: ""
+
+# FSDP wrap strategy
+generator_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size

DecMem/configs/worldmem_diffusion.yaml

@@ -0,0 +1,69 @@
+# WorldMem diffusion training config
+# Action + PRoPE conditioning, no text conditioning, CFG on actions only
+
+trainer: diffusion
+causal: true
+mixed_precision: true
+seed: 0
+
+# Data
+# CSV columns: video_path, action_path  (WorldMem format)
+data_path: data/train_wm_gamecraft/v2/worldmem_train_v2_shuffle.csv
+num_frames: 28          # latent frame count; dataset loads (21-1)*4+1=81 video frames
+load_raw_video: true
+sample_strategy: first  # overfit: always start from SKIP_FRAMES; use "random" for normal training
+
+# CFG: probability of dropping action conditioning (zeros actions + identity viewmats)
+proportion_empty_prompts: 0.1
+
+# Model
+# 81 video frames @ 640x352 -> 21 latent frames @ 80x44 (WAN 4x temporal, 8x spatial)
+image_or_video_shape:
+- 1        # batch_size placeholder (overridden at runtime)
+- 28       # latent frames
+- 16       # latent channels
+- 44       # latent height  (352/8)
+- 80       # latent width   (640/8)
+
+model_kwargs:
+  timestep_shift: 5.0
+  local_attn_size: -1
+  sink_size: 0
+  use_action_encoder: true
+  action_dim: 25
+  action_hidden_dim: 128
+  latent_channels: 16
+
+# Training hyperparams
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 1.0
+teacher_forcing: true
+# noise_augmentation_max_timestep: 100
+noise_augmentation_max_timestep: 0
+
+lr: 1.0e-4
+beta1: 0.9
+beta2: 0.999
+weight_decay: 0.01
+batch_size: 1
+gradient_checkpointing: true
+num_frame_per_block: 4
+sharding_strategy: hybrid_full
+
+# Logging / saving
+no_save: false
+no_visualize: true
+log_iters: 300
+ema_weight: 0.999
+ema_start_step: 200
+gc_interval: 100
+wandb_host: "https://api.wandb.ai"
+wandb_key: WANDB_KEY
+wandb_entity: "zhenhao-yang11-the-university-of-hong-kong"
+wandb_project: worldmem_diffusion
+negative_prompt: ""
+
+# FSDP wrap strategy
+generator_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size

DecMem/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

DecMem/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

DecMem/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()

DecMem/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)

DecMem/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

DecMem/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

DecMem/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)

DecMem/evaluation/method_manifest.json

@@ -0,0 +1,91 @@
+{
+  "method": "DecMem official public Wan2.1-T2V-1.3B reference",
+  "upstream": "https://github.com/KlingAIResearch/DecMem",
+  "upstream_revision": "f355a1809444a9cc6746ff99d5255614348b20ac",
+  "checkpoint": "KlingTeam/DecMem/decmem.pt generator state",
+  "checkpoint_tensor_count": 492,
+  "checkpoint_parameter_count": 1197639296,
+  "checkpoint_weight_policy": "released generator; no EMA substitution",
+  "model": {
+    "trainable_parameters": 0,
+    "optimizer_groups": [],
+    "lr_stages": [],
+    "lora": "none",
+    "action_encoder": "official linear100to16",
+    "memory_projection_layers": 30,
+    "use_ltm_attn": true,
+    "ltm_topk": 80,
+    "ltm_tile_size": [
+      1,
+      8,
+      8
+    ],
+    "stm_sliding_window": 8,
+    "local_attn_size": -1,
+    "block_latents": 4,
+    "activation_checkpointing": false
+  },
+  "sampling": {
+    "profile": "scripts/infer_example.sh released example defaults",
+    "solver": "UniPC",
+    "steps": 20,
+    "cfg": 5.0,
+    "timestep_shift": 5.0,
+    "dtype": "bfloat16",
+    "batch_size": 1,
+    "cfg_effective_batch": 2
+  },
+  "protocol": {
+    "manifest": "evaluation/manifests/minecraft_test_300.csv",
+    "observe_source": [
+      100,
+      700
+    ],
+    "score_source": [
+      700,
+      1200
+    ],
+    "stride": 1,
+    "unique_observed": 600,
+    "left_duplicates": 1,
+    "observed_rgb_with_padding": 601,
+    "observed_latents": 151,
+    "total_latents": 276,
+    "generated_latents": 125,
+    "generated_rgb": 500,
+    "score_rgb": 500,
+    "full_history": "allobserved latents warm SGM+ALM; partial3cleanlatentsinfirstgeneratedblock preserved"
+  },
+  "controls": {
+    "actions": "unchangedcanonicalWorldMem25channels; prependonezeroaction",
+    "poses": "xyz-firstxyz, negateyaw, wrapanglesmod360; sampled[::4]afterleftduplicate",
+    "pose_construction": "float32 world-to-camera then castbfloat16",
+    "first_generated_latent_pose_source": 703,
+    "last_pose_source": 1199,
+    "prope_max_time": 60,
+    "prope_sequence_semantics": "preserve native temporal indices, dynamicallyextendcoefficienttable; no60framecap"
+  },
+  "image": {
+    "native_resolution": [
+      352,
+      640
+    ],
+    "input": "commonRGB01, fullframebilinearresize antialias, then[-1,1]",
+    "encode": "observed601framesonly; futureinitiallatentszero",
+    "decode": "entire observed+generated1101timelineusingnativecausalVAECPUstreaming;discardfirst601",
+    "score_targets": "sharedfixedOasisVAEreconstruction, neverDecMemVAEreconstruction"
+  },
+  "verification": {
+    "syntax": "passed",
+    "actual_first_clip_actions": "matchesofficialextract_pose_action_from_npz",
+    "actual_first_clip_poses": "matchesofficialnormalizationand[::4]",
+    "timeline_600_500": "151context,276total,500generated;firstscored700,last1199",
+    "timeline_600_20": "151context,156total,20generated",
+    "native_model_gpu_import": "pendingavailableSlurmallocation",
+    "strict_checkpoint_model_compatibility": "performed in GPU constructor, pending",
+    "vsa_compile": "passed; CUDA12.9 explicit compute90a/sm90a",
+    "vsa_binary_import": "passed on CPU; execution still awaits GPU"
+  },
+  "full_ready": false,
+  "blocker": "GPU native VSA import and actualgeneration smoke not completed"
+}

DecMem/evaluation/worldmem_adapter.py

@@ -0,0 +1,188 @@
+"""Run the released DecMem memory model on the fixed WorldMem Minecraft split."""
+
+import argparse
+import json
+import math
+import os
+from pathlib import Path
+import sys
+
+import torch
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+
+REPO = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO))
+sys.path.insert(0, str(REPO.parent / "shared"))
+
+from protocol import case_seed, load_minecraft_case, save_case, select_cases
+from utils.pose_utils import poses_to_viewmats
+
+
+def timeline(history_frames=600, target_frames=500):
+    prefix_padding = (1 - history_frames) % 4
+    context_latents = (history_frames + prefix_padding - 1) // 4 + 1
+    total_latents = math.ceil((context_latents + math.ceil(target_frames / 4)) / 4) * 4
+    generated_frames = (total_latents - context_latents) * 4
+    return context_latents, prefix_padding, total_latents, generated_frames
+
+
+def prepare_controls(case, target_frames):
+    history = case["history"]
+    n_cond, padding, total_latents, generated_frames = timeline(history, target_frames)
+    poses = case["poses"][:history + target_frames].clone().float()
+    # Match extract_pose_action_from_npz: translate xyz, negate yaw, then wrap angles.
+    poses[:, :3] -= poses[:1, :3].clone()
+    poses[:, -1].neg_()
+    poses[:, 3:].remainder_(360)
+    actions = case["actions"][:history + target_frames].float()
+    tail = generated_frames - target_frames
+    poses = torch.cat([poses[:1].expand(padding, -1), poses, poses[-1:].expand(tail, -1)])
+    actions = torch.cat([actions.new_zeros((padding, 25)), actions, actions.new_zeros((tail, 25))])
+    source_indices = case["source_indices"][:history + target_frames]
+    source_indices = torch.cat([source_indices[:1].expand(padding), source_indices, source_indices[-1:].expand(tail)])
+    return actions, poses[::4], source_indices[::4], (n_cond, padding, total_latents, generated_frames)
+
+
+class DecMemAdapter:
+    def __init__(self, checkpoint, device="cuda"):
+        self.checkpoint = str(Path(checkpoint).resolve())
+        # Upstream model and VAE loading use paths relative to the DecMem checkout.
+        os.chdir(REPO)
+        from pipeline import WorldMemCausalDiffusionInferencePipeline
+
+        self.device = torch.device(device)
+        self.dtype = torch.bfloat16
+        self.config = OmegaConf.merge(
+            OmegaConf.load(REPO / "configs/default_config.yaml"),
+            OmegaConf.load(REPO / "configs/decmem.yaml"),
+        )
+        self.config.use_text_encoder = False
+        self.config.image_or_video_shape[1] = timeline()[2]
+        self.config.num_frames = self.config.image_or_video_shape[1]
+        self.config.prope_max_time = 60
+        self.config.rope_scale_pattern = "rope"
+        self.pipeline = WorldMemCausalDiffusionInferencePipeline(self.config, device=self.device)
+        checkpoint_state = torch.load(self.checkpoint, map_location="cpu", weights_only=True, mmap=True)
+        self.pipeline.generator.load_state_dict(checkpoint_state["generator"], strict=True)
+        self.pipeline.to(device=self.device, dtype=self.dtype).eval().requires_grad_(False)
+        self.parameter_counts = {
+            "generator_total": sum(p.numel() for p in self.pipeline.generator.parameters()),
+            "generator_trainable": sum(p.numel() for p in self.pipeline.generator.parameters() if p.requires_grad),
+            "action_encoder": sum(p.numel() for p in self.pipeline.generator.action_encoder.parameters()),
+            "memory_projection": sum(p.numel() for n, p in self.pipeline.generator.named_parameters() if ".mem_proj." in n),
+            "vae_total": sum(p.numel() for p in self.pipeline.vae.parameters()),
+            "vae_trainable": sum(p.numel() for p in self.pipeline.vae.parameters() if p.requires_grad),
+        }
+        print(json.dumps({"parameter_counts": self.parameter_counts, "checkpoint_strict_load": "passed"}), flush=True)
+
+    @torch.inference_mode()
+    def generate(self, case, target_frames=500, seed=0):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        history = case["history"]
+        actions, poses, pose_indices, layout = prepare_controls(case, target_frames)
+        n_cond, padding, total_latents, generated_frames = layout
+        # Future RGB never reaches the VAE; only the complete observed prefix is encoded.
+        observed = case["rgb"][:history].to(self.device)
+        observed = F.interpolate(observed, (352, 640), mode="bilinear", align_corners=False, antialias=True)
+        observed = torch.cat([observed[:1].expand(padding, -1, -1, -1), observed])
+        observed = observed.transpose(0, 1).unsqueeze(0).mul(2).sub(1).to(self.dtype)
+        context_latent = self.pipeline.vae.encode_to_latent(observed).to(self.device, self.dtype)
+        del observed
+
+        initial_latent = context_latent.new_zeros((1, total_latents, 16, 44, 80))
+        initial_latent[:, :n_cond] = context_latent
+        del context_latent
+        actions = actions.unsqueeze(0).to(self.device, self.dtype)
+        viewmats = poses_to_viewmats(poses.unsqueeze(0).to(self.device)).to(self.dtype)
+        noise = torch.randn(initial_latent.shape, device=self.device, dtype=self.dtype)
+        print(json.dumps({"observed_frames": history, "prefix_padding": padding, "condition_latents": n_cond, "total_latents": total_latents}), flush=True)
+        _, output_latents = self.pipeline.inference(
+            noise=noise,
+            actions=actions,
+            viewmats=viewmats,
+            initial_latent=initial_latent,
+            num_condition_frames=n_cond,
+            guidance_scale=5.0,
+            sampling_steps=20,
+            return_latents=True,
+            skip_vae_decode=True,
+        )
+        # Denoising is finished. Release global KV before decoding the complete causal timeline.
+        self.pipeline.kv_cache = None
+        self.pipeline.crossattn_cache = None
+        torch.cuda.empty_cache()
+        video = self.pipeline.vae.decode_to_pixel(output_latents)
+        prediction = video[0, history + padding:history + padding + target_frames].float().mul(0.5).add(0.5).clamp(0, 1).cpu()
+        metadata = {
+            "method": "decmem",
+            "implementation": "official public Wan2.1-T2V-1.3B reference",
+            "checkpoint": self.checkpoint,
+            "weight_source": "released generator state; no EMA substitution",
+            "parameter_counts": self.parameter_counts,
+            "dtype": "bfloat16",
+            "batch_size": 1,
+            "cfg_batch_size": 2,
+            "sampling": {"solver": "UniPC", "steps": 20, "guidance_scale": 5.0, "timestep_shift": 5.0, "source": "scripts/infer_example.sh"},
+            "memory": {"use_ltm_attn": True, "ltm_topk": 80, "ltm_tile_size": [1, 8, 8], "stm_sliding_window": 8, "num_frame_per_block": 4},
+            "history_frames": history,
+            "history_source_indices": case["source_indices"][:history].tolist(),
+            "history_left_padding": padding,
+            "condition_latents": n_cond,
+            "total_latents": total_latents,
+            "generated_frames": generated_frames,
+            "scored_frames": target_frames,
+            "future_padding": generated_frames - target_frames,
+            "action_mapping": "official WorldMem 25 channels, no scale changes; prepend one noop",
+            "pose_mapping": "translate xyz relative to source100; negate yaw; angles modulo360; padded source indices[::4]",
+            "latent_pose_source_indices": pose_indices.tolist(),
+            "pose_matrix_precision": "float32 construction, then bfloat16 model input",
+            "prope_max_time": 60,
+            "prope_time_semantics": "native absolute latent indices;60 is not a sequence length cap",
+            "native_resolution": [352, 640],
+            "preprocessing": "full-frame bilinear resize with antialias; no crop",
+            "target_contract": "shared fixed Oasis VAE reconstruction; no method VAE target substitution",
+            "seed": seed,
+        }
+        return prediction, metadata
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", required=True)
+    parser.add_argument("--data-root", required=True)
+    parser.add_argument("--checkpoint", default=str(REPO.parent / "checkpoints/decmem/decmem.pt"))
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--limit", type=int)
+    parser.add_argument("--max-generated", type=int, default=500)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--rank", type=int, default=int(os.environ.get("SLURM_PROCID", 0)))
+    parser.add_argument("--world-size", type=int, default=int(os.environ.get("SLURM_NTASKS", 1)))
+    parser.add_argument("--preflight-only", action="store_true")
+    args = parser.parse_args()
+    if not 1 <= args.max_generated <= 500:
+        parser.error("--max-generated must be between 1 and the fixed 500-frame evaluation horizon")
+    torch.set_num_threads(4)
+    rows = select_cases(args.manifest, rank=args.rank, world_size=args.world_size, limit=args.limit)
+    if args.preflight_only:
+        case = load_minecraft_case(rows[0], args.data_root)
+        actions, poses, indices, layout = prepare_controls(case, args.max_generated)
+        report = {"selected_cases": len(rows), "actions_shape": list(actions.shape), "poses_shape": list(poses.shape), "latent_pose_source_indices": indices.tolist(), "layout": layout, "gpu_model_import": "not attempted"}
+        output = Path(args.output)
+        output.mkdir(parents=True, exist_ok=True)
+        (output / f"preflight_rank{args.rank}.json").write_text(json.dumps(report, indent=2) + "\n")
+        print(json.dumps(report), flush=True)
+        return
+    model = DecMemAdapter(args.checkpoint)
+    revision = "f355a1809444a9cc6746ff99d5255614348b20ac"
+    for row in rows:
+        case = load_minecraft_case(row, args.data_root)
+        prediction, metadata = model.generate(case, args.max_generated, case_seed(row, args.seed))
+        metadata.update({"git_revision": revision, "arguments": vars(args)})
+        directory = save_case(args.output, row, prediction, metadata)
+        print(f"Saved {tuple(prediction.shape)} to {directory}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

DecMem/example/example.csv

@@ -0,0 +1,3 @@
+video_path,action_path
+example/000045.mp4,example/000045.npz
+example/000046.mp4,example/000046.npz

DecMem/infer.py

@@ -0,0 +1,338 @@
+"""
+WorldMem CFG inference — block-by-block causal denoising with KV cache + PRoPE.
+
+Delegates to WorldMemCausalDiffusionInferencePipeline, which:
+  - Writes any fully-clean prefix block to the KV cache at timestep=0 first.
+  - Denoises one block at a time via UniPC, each step a single forward with
+    kv_cache + PRoPE recomputed for that block's viewmats.
+  - Commits each block's clean K/V to cache before moving to the next, so
+    later blocks attend against clean history (matching teacher-forcing
+    training distribution).
+  - CFG is one batched forward with batch_size = 2*B; the KV cache is also
+    allocated with batch_size = 2*B. cond/uncond chunked back after each step.
+"""
+
+import argparse
+import gc
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import numpy as np
+from omegaconf import OmegaConf
+from torchvision.io import write_video
+
+os.chdir(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ".")
+
+from pipeline import WorldMemCausalDiffusionInferencePipeline
+from utils.dataset import WorldMemCSVDataset
+from utils.pose_utils import poses_to_viewmats
+from utils.misc import set_seed
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_path", default="configs/worldmem_diffusion_overfit.yaml")
+parser.add_argument("--output_dir", default="test/res/infer_worldmem_cfg")
+parser.add_argument("--checkpoint_path", default=None,
+                    help="Optional generator checkpoint (model.pt)")
+parser.add_argument("--data_path", default=None,
+                    help="Override config.data_path (CSV) for inference test set.")
+parser.add_argument("--num_samples", type=int, default=1,
+                    help="Number of noise resamples per dataset entry.")
+parser.add_argument("--num_inference_steps", type=int, default=50)
+parser.add_argument("--guidance_scale", type=float, default=3.0,
+                    help="CFG scale. <=1.0 disables CFG (single forward, B=1 cache).")
+parser.add_argument("--seed", type=int, default=42)
+parser.add_argument("--start_index", type=int, default=0,
+                    help="First dataset row (env: START_INDEX). "
+                         "Ignored if --sample_indices is provided.")
+parser.add_argument("--num_entries", type=int, default=None,
+                    help="Number of consecutive dataset rows from --start_index (env: N). "
+                         "Default: all rows from start_index to dataset end. "
+                         "Ignored if --sample_indices is provided.")
+parser.add_argument("--sample_indices", type=str, default=None,
+                    help="Explicit dataset indices, overrides --start_index/--num_entries. "
+                         "Accepts comma-separated ints and inclusive ranges, e.g. "
+                         "'0,3,5' or '0-9,15,20-22', or 'all' for the full dataset.")
+parser.add_argument("--num_condition_frames", type=int, default=1,
+                    help="Number of clean prefix latent frames (1=first image, "
+                         "3=first block when num_frame_per_block=3).")
+parser.add_argument("--num_latent_frames", type=int, default=None,
+                    help="Override config.image_or_video_shape[1] (latent T). "
+                         "Must be a multiple of num_frame_per_block.")
+parser.add_argument("--tag", type=str, default="",
+                    help="Filename tag appended to outputs.")
+parser.add_argument("--offload", action="store_true",
+                    help="Mutually-exclusive GPU residency for generator and VAE: "
+                         "during denoising the VAE lives on CPU; for VAE encode/decode "
+                         "the generator lives on CPU. Trades extra H2D/D2H copies for "
+                         "lower peak VRAM.")
+args = parser.parse_args()
+
+# ─── Distributed (optional, via torchrun) ────────────────────────────────────
+if "LOCAL_RANK" in os.environ:
+    dist.init_process_group(backend="nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = dist.get_world_size()
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+else:
+    local_rank = 0
+    world_size = 1
+    device = torch.device("cuda")
+
+set_seed(args.seed + local_rank)
+dtype = torch.bfloat16
+
+# ─── Config ──────────────────────────────────────────────────────────────────
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+config.use_text_encoder = False
+if args.data_path is not None:
+    print(f"Overriding config.data_path: {config.data_path} -> {args.data_path}")
+    config.data_path = args.data_path
+
+if args.num_latent_frames is not None:
+    nfb = int(config.num_frame_per_block)
+    assert args.num_latent_frames % nfb == 0, (
+        f"--num_latent_frames={args.num_latent_frames} must be a multiple of "
+        f"num_frame_per_block={nfb}")
+    print(f"Overriding latent T: {config.image_or_video_shape[1]} -> {args.num_latent_frames}")
+    config.image_or_video_shape[1] = args.num_latent_frames
+    config.num_frames = args.num_latent_frames
+
+num_latent_frames = config.image_or_video_shape[1]
+latent_C, latent_H, latent_W = config.image_or_video_shape[2:]
+n_cond = int(args.num_condition_frames)
+assert 1 <= n_cond < num_latent_frames, \
+    f"num_condition_frames must be in [1, {num_latent_frames - 1}], got {n_cond}"
+
+guidance_scale = float(args.guidance_scale)
+use_cfg = guidance_scale > 1.0
+if local_rank == 0:
+    os.makedirs(args.output_dir, exist_ok=True)
+if dist.is_initialized():
+    dist.barrier()
+print(f"Config: T={num_latent_frames}  latent {latent_H}x{latent_W}  "
+      f"nfb={config.num_frame_per_block}  "
+      f"guidance_scale={guidance_scale} (cfg={use_cfg})  "
+      f"n_cond={n_cond}")
+
+# ─── Pipeline ────────────────────────────────────────────────────────────────
+torch.set_grad_enabled(False)
+pipeline = WorldMemCausalDiffusionInferencePipeline(config, device=device)
+pipeline.generator = pipeline.generator.to(device=device, dtype=dtype)
+pipeline.generator.eval()
+pipeline.vae = pipeline.vae.to(device=device, dtype=dtype)
+
+
+# ─── Optional CPU/GPU swap for generator/VAE ────────────────────────────────
+# When --offload is set, only one of {generator, VAE} lives on the GPU at any
+# time. The "owner" variable below tracks which model currently holds GPU
+# residency so we can no-op redundant swaps. Idle model is parked on CPU.
+_GPU_OWNER = "both"  # "generator" | "vae" | "both" | "none"
+
+
+def _swap_to(target: str):
+    """Make ``target`` the sole GPU resident; park the other on CPU.
+
+    target ∈ {"generator", "vae"}. No-op when --offload is disabled.
+    Frees pipeline.kv_cache before parking the generator and runs an
+    explicit gc + empty_cache so VRAM is reclaimed before the next load.
+    """
+    global _GPU_OWNER
+    if not args.offload:
+        return
+    if _GPU_OWNER == target:
+        return
+    if target == "vae":
+        # Park generator on CPU. Drop transient GPU state (kv_cache) first
+        # so the move actually frees memory rather than fragmenting it.
+        if pipeline.kv_cache is not None:
+            pipeline.kv_cache = None
+        if pipeline.crossattn_cache is not None:
+            pipeline.crossattn_cache = None
+        pipeline.generator.to("cpu")
+        gc.collect()
+        torch.cuda.empty_cache()
+        pipeline.vae.to(device=device, dtype=dtype)
+    elif target == "generator":
+        pipeline.vae.to("cpu")
+        gc.collect()
+        torch.cuda.empty_cache()
+        pipeline.generator.to(device=device, dtype=dtype)
+    else:
+        raise ValueError(f"Unknown swap target: {target}")
+    _GPU_OWNER = target
+
+
+if args.offload:
+    # Initial state: VAE on GPU (first thing we do per-sample is encode).
+    pipeline.generator.to("cpu")
+    gc.collect()
+    torch.cuda.empty_cache()
+    _GPU_OWNER = "vae"
+    print("Offload enabled: generator/VAE are mutually exclusive on GPU.")
+
+if args.checkpoint_path:
+    weight_tag = "generator"
+    print(f"Loading checkpoint: {args.checkpoint_path} ({weight_tag})")
+    ckpt = torch.load(args.checkpoint_path, map_location="cpu", weights_only=False)
+    sd = ckpt.get("generator", ckpt.get("model", ckpt))
+    sample_key = next(iter(sd.keys()))
+    target = pipeline.generator if sample_key.startswith("model.") else pipeline.generator.model
+    missing, unexpected = target.load_state_dict(sd, strict=False)
+    if missing:
+        print(f"  missing keys ({len(missing)}): {missing[:5]}{'...' if len(missing) > 5 else ''}")
+    if unexpected:
+        print(f"  unexpected keys ({len(unexpected)}): {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}")
+else:
+    print("No checkpoint — using pretrained WAN backbone + random action encoder.")
+
+# ─── Dataset ─────────────────────────────────────────────────────────────────
+dataset = WorldMemCSVDataset(
+    config.data_path,
+    num_latent_frames=num_latent_frames,
+    sample_strategy="first",
+)
+print(f"Dataset size: {len(dataset)}")
+
+def _parse_sample_indices(spec: str, dataset_size: int) -> list:
+    """Parse '0,3,5' / '0-9,15,20-22' / 'all' into a deduplicated sorted index list."""
+    if spec.strip().lower() == "all":
+        return list(range(dataset_size))
+    out = set()
+    for token in spec.split(","):
+        token = token.strip()
+        if not token:
+            continue
+        if "-" in token:
+            lo, hi = token.split("-", 1)
+            lo, hi = int(lo), int(hi)
+            assert lo <= hi, f"bad range '{token}': lo > hi"
+            out.update(range(lo, hi + 1))
+        else:
+            out.add(int(token))
+    indices = sorted(out)
+    for i in indices:
+        assert 0 <= i < dataset_size, (
+            f"index {i} out of range [0, {dataset_size}) for --sample_indices='{spec}'")
+    return indices
+
+
+if args.sample_indices is not None:
+    sample_indices = _parse_sample_indices(args.sample_indices, len(dataset))
+    assert sample_indices, f"--sample_indices='{args.sample_indices}' parsed to empty list"
+    print(f"Using explicit --sample_indices ({len(sample_indices)}): {sample_indices}")
+else:
+    assert args.start_index >= 0, f"--start_index must be >= 0, got {args.start_index}"
+    if args.num_entries is not None:
+        assert args.num_entries > 0, f"--num_entries must be > 0, got {args.num_entries}"
+        end_index = args.start_index + args.num_entries
+        sample_indices = list(range(args.start_index, min(end_index, len(dataset))))
+        if end_index > len(dataset):
+            print(f"Warning: start_index={args.start_index} + num_entries={args.num_entries} "
+                  f"exceeds dataset size {len(dataset)}; using indices "
+                  f"{args.start_index}..{len(dataset) - 1}")
+    else:
+        sample_indices = list(range(args.start_index, len(dataset)))
+
+if world_size > 1:
+    sample_indices = sample_indices[local_rank::world_size]
+    print(f"[rank {local_rank}/{world_size}] shard indices: {sample_indices}")
+elif sample_indices:
+    print(f"Will run on dataset indices: {sample_indices}")
+
+def run_one_sample(ds_idx: int):
+    print(f"\n=== Dataset index {ds_idx} ===")
+    rng_state = np.random.get_state()
+    np.random.seed(args.seed)
+    sample = dataset[ds_idx]
+    np.random.set_state(rng_state)
+
+    frames_raw = sample["frames"].unsqueeze(0).to(device=device, dtype=dtype)   # (1,3,T_vid,H,W)
+    actions = sample["actions"].unsqueeze(0).to(device=device, dtype=dtype)     # (1,T_vid,25)
+    poses = sample["poses"].unsqueeze(0).to(device=device, dtype=dtype)         # (1,T_lat,5)
+
+    # VAE encode: needs VAE on GPU.
+    _swap_to("vae")
+    gt_latent = pipeline.vae.encode_to_latent(frames_raw).to(device=device, dtype=dtype)
+    viewmats = poses_to_viewmats(poses)                                         # (1,T,4,4)
+
+    for sample_idx in range(args.num_samples):
+        print(f"  noise sample {sample_idx + 1}/{args.num_samples} (prefix={n_cond}) ...")
+        torch.manual_seed(args.seed + sample_idx)
+
+        noise = torch.randn(1, num_latent_frames, latent_C, latent_H, latent_W,
+                            device=device, dtype=dtype)
+        # Denoising: needs generator on GPU; VAE parked on CPU.
+        _swap_to("generator")
+        with torch.no_grad():
+            video, latents = pipeline.inference(
+                noise=noise,
+                actions=actions,
+                viewmats=viewmats,
+                initial_latent=gt_latent,
+                num_condition_frames=n_cond,
+                guidance_scale=guidance_scale,
+                sampling_steps=args.num_inference_steps,
+                return_latents=True,
+                skip_vae_decode=args.offload,
+            )
+
+        # Decode: park generator, bring VAE back to GPU.
+        # In offload mode, `video` is None — we decode externally now that the
+        # VAE owns the GPU. In non-offload mode the decode already happened.
+        _swap_to("vae")
+        if args.offload:
+            video = pipeline.vae.decode_to_pixel(latents.to(dtype=dtype))
+            video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        gt_video = pipeline.vae.decode_to_pixel(gt_latent.to(dtype=dtype))
+        gt_video = (gt_video * 0.5 + 0.5).clamp(0, 1)
+
+        gen_uint8 = (video[0].permute(0, 2, 3, 1).cpu().float() * 255).byte()
+        gt_uint8 = (gt_video[0].permute(0, 2, 3, 1).cpu().float() * 255).byte()
+        side_by_side = torch.cat([gt_uint8, gen_uint8], dim=2)
+
+        tag = f"_{args.tag}" if args.tag else ""
+        out_path = os.path.join(
+            args.output_dir,
+            f"worldmem_s{ds_idx}_n{sample_idx}_cond{n_cond}_cfg{guidance_scale}{tag}.mp4")
+        write_video(out_path, side_by_side, fps=10)
+
+        out_dir = os.path.join(args.output_dir, 'pred')
+        os.makedirs(out_dir, exist_ok=True)
+        out_gen_path = os.path.join(
+            out_dir,
+            f"worldmem_s{ds_idx}_n{sample_idx}_cond{n_cond}_cfg{guidance_scale}{tag}.mp4")
+        write_video(out_gen_path, gen_uint8, fps=10)
+
+        out_dir = os.path.join(args.output_dir, 'gt')
+        os.makedirs(out_dir, exist_ok=True)
+        out_gt_path = os.path.join(
+            out_dir,
+            f"worldmem_s{ds_idx}_n{sample_idx}_cond{n_cond}_cfg{guidance_scale}{tag}.mp4")
+        write_video(out_gt_path, gt_uint8, fps=10)
+
+        print(f"  Saved {out_path}  (left=GT | right=generated)")
+
+        # Release per-sample GPU residents and defrag the allocator before the
+        # next iteration re-allocates the KV cache inside pipeline.inference().
+        del video, latents, gt_video, gen_uint8, gt_uint8, side_by_side
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+if not sample_indices:
+    print(f"[rank {local_rank}] no samples in this shard, skipping.")
+else:
+    for ds_idx in sample_indices:
+        run_one_sample(ds_idx)
+
+if dist.is_initialized():
+    dist.barrier()
+if local_rank == 0:
+    print("\nDone.")

DecMem/model/__init__.py

@@ -0,0 +1,9 @@
+from .diffusion import CausalDiffusion
+from .causvid import CausVid
+from .gan import GAN
+
+__all__ = [
+    "CausalDiffusion",
+    "CausVid",
+    "GAN",
+]

DecMem/model/base.py

@@ -0,0 +1,222 @@
+from typing import Tuple
+from einops import rearrange
+from torch import nn
+import torch.distributed as dist
+import torch
+
+from pipeline import SelfForcingTrainingPipeline
+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 SelfForcingModel(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()
+
+        return self.inference_pipeline.inference_with_trajectory(
+            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 = SelfForcingTrainingPipeline(
+            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
+        )

DecMem/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

DecMem/model/diffusion.py

@@ -0,0 +1,155 @@
+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, device=None):
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        # Text encoder is optional; skip when training without text conditioning (e.g. WorldMem)
+        if getattr(args, "use_text_encoder", True):
+            self.text_encoder = WanTextEncoder()
+            self.text_encoder.requires_grad_(False)
+        else:
+            self.text_encoder = None
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+        self.scheduler = self.generator.get_scheduler()
+        if device is not None:
+            self.scheduler.timesteps = self.scheduler.timesteps.to(device)
+
+        # Set up PRoPE for all attention layers when action encoder is enabled
+        model_kwargs = getattr(args, "model_kwargs", {})
+        if model_kwargs.get("use_action_encoder", False):
+            image_or_video_shape = getattr(args, "image_or_video_shape", [1, 21, 16, 60, 104])
+            num_frames = image_or_video_shape[1]
+            latent_h = image_or_video_shape[3]
+            latent_w = image_or_video_shape[4]
+            # patch_size for WAN is (1, 2, 2)
+            patches_y = latent_h // 2
+            patches_x = latent_w // 2
+            self.generator.model.setup_prope(
+                patches_x=patches_x,
+                patches_y=patches_y,
+                image_width=latent_w,
+                image_height=latent_h,
+                cameras=num_frames,
+                prope_max_time=getattr(args, "prope_max_time", 60),
+                rope_scale_pattern=getattr(args, "rope_scale_pattern", "rope"),
+            )
+
+    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_unweighted = torch.nn.functional.mse_loss(
+            flow_pred.float(), training_target.float(), reduction='none'
+        ).mean(dim=(2, 3, 4))                                            # (B, F) unweighted per-frame loss
+        weight = self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
+        loss = (loss_unweighted * weight).mean()                         # scalar; used for backprop
+
+        log_dict = {
+            "x0": clean_latent.detach(),
+            "x0_pred": x0_pred.detach(),
+            "loss_per_frame": loss_unweighted.detach(),                  # (B, F) unweighted
+            "loss_unweighted": loss_unweighted.detach().mean(),          # scalar unweighted
+        }
+        return loss, log_dict

DecMem/model/gan.py

@@ -0,0 +1,295 @@
+import copy
+from pipeline import SelfForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class GAN(SelfForcingModel):
+    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: SelfForcingTrainingPipeline = 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

DecMem/pipeline/__init__.py

@@ -0,0 +1,17 @@
+from .bidirectional_diffusion_inference import BidirectionalDiffusionInferencePipeline
+from .bidirectional_inference import BidirectionalInferencePipeline
+from .causal_diffusion_inference import (
+    CausalDiffusionInferencePipeline,
+    WorldMemCausalDiffusionInferencePipeline,
+)
+from .causal_inference import CausalInferencePipeline
+from .self_forcing_training import SelfForcingTrainingPipeline
+
+__all__ = [
+    "BidirectionalDiffusionInferencePipeline",
+    "BidirectionalInferencePipeline",
+    "CausalDiffusionInferencePipeline",
+    "CausalInferencePipeline",
+    "SelfForcingTrainingPipeline",
+    "WorldMemCausalDiffusionInferencePipeline",
+]

DecMem/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

DecMem/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

DecMem/pipeline/causal_diffusion_inference.py

@@ -0,0 +1,707 @@
+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 WorldMemCausalDiffusionInferencePipeline(torch.nn.Module):
+    """Block-by-block causal inference for WorldMem (action + PRoPE, no text).
+
+    Key differences vs. CausalDiffusionInferencePipeline:
+      * No text encoder; conditioning comes from per-frame actions + viewmats.
+      * CFG is batched into a single KV cache with batch_size = 2 * B; cond
+        and uncond are stacked along dim 0 and chunked back after each forward.
+      * PRoPE is precomputed PER BLOCK (with the block's viewmats slice and
+        absolute `time_ids=blk_start`). The cache stores PRoPE-rotated K/V,
+        so subsequent blocks attend against features in their correct
+        rotated frames (relative-position attention falls out naturally).
+
+    KV-cache schedule for one inference call:
+      1. PREFIX:    forward each fully-clean block at timestep=0 to write its
+                    PRoPE-rotated K/V to the cache.
+      2. DENOISE:   for each remaining block, run `sampling_steps` of UniPC,
+                    each step a forward with kv_cache + PRoPE for that block.
+                    Frames already pinned clean (partial-prefix case) keep
+                    timestep=0 within the block.
+      3. COMMIT:    after each block converges, re-forward at timestep=0 with
+                    the now-clean latents to overwrite that block's cache slot.
+    """
+
+    def __init__(self, args, device, generator=None, vae=None):
+        super().__init__()
+        self.args = args
+        self.device = device
+
+        self.generator = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}), is_causal=True
+        ) if generator is None else generator
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        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 = len(self.generator.model.blocks)
+        # frame_seq_length = patches_y * patches_x for WAN patch_size=(1,2,2)
+        latent_h = args.image_or_video_shape[3]
+        latent_w = args.image_or_video_shape[4]
+        self.frame_seq_length = (latent_h // 2) * (latent_w // 2)
+
+        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.local_attn_size = self.generator.model.local_attn_size
+
+        # Setup PRoPE (mirrors CausalDiffusion._initialize_models).
+        model_kwargs = getattr(args, "model_kwargs", {})
+        if model_kwargs.get("use_action_encoder", False):
+            num_frames = args.image_or_video_shape[1]
+            patches_y = latent_h // 2
+            patches_x = latent_w // 2
+            self.generator.model.setup_prope(
+                patches_x=patches_x,
+                patches_y=patches_y,
+                image_width=latent_w,
+                image_height=latent_h,
+                cameras=num_frames,
+                prope_max_time=getattr(args, "prope_max_time", 60),
+                rope_scale_pattern=getattr(args, "rope_scale_pattern", "rope"),
+            )
+
+        self.kv_cache = None
+        self.crossattn_cache = None
+
+    # ─── Cache lifecycle ─────────────────────────────────────────────────────
+    def _initialize_kv_cache(self, batch_size, dtype, device, num_total_frames):
+        self.kv_cache = None
+        torch.cuda.empty_cache()
+        kv_cache = []
+        if self.local_attn_size != -1:
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            kv_cache_size = num_total_frames * self.frame_seq_length
+
+        # Extend max_attention_size so attention reads the full cache range.
+        if self.local_attn_size == -1:
+            for module in self.generator.model.modules():
+                if hasattr(module, 'max_attention_size'):
+                    module.max_attention_size = kv_cache_size
+
+        head_dim = self.generator.model.dim // self.generator.model.num_heads
+        num_heads = self.generator.model.num_heads
+        use_ltm = getattr(self.generator.model, "use_ltm_attn", False)
+        for _ in range(self.num_transformer_blocks):
+            entry = {
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+            }
+            if use_ltm:
+                # LTM: untiled k/v + incremental tiled/compressed KV (per block).
+                # entry["num_frames"] = 0
+                entry["k"] = None
+                entry["v"] = None
+                entry["k_tiled"] = None
+                entry["v_tiled"] = None
+                entry["k_compress"] = None
+                entry["v_compress"] = None
+                entry["variable_block_sizes"] = None
+            else:
+                entry["k"] = torch.zeros(
+                    [batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device)
+                entry["v"] = torch.zeros(
+                    [batch_size, kv_cache_size, num_heads, head_dim], dtype=dtype, device=device)
+            kv_cache.append(entry)
+        self.kv_cache = kv_cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        # Cross-attention is disabled for WorldMem (no text); allocate stubs
+        # so the model's block forward can still index `crossattn_cache[i]`
+        # without surprises.
+        crossattn_cache = []
+        head_dim = self.generator.model.dim // self.generator.model.num_heads
+        num_heads = self.generator.model.num_heads
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k": torch.zeros([batch_size, 512, num_heads, head_dim],
+                                dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, 512, num_heads, head_dim],
+                                dtype=dtype, device=device),
+                "is_init": False,
+            })
+        self.crossattn_cache = crossattn_cache
+
+    def _init_scheduler(self, device):
+        if self.sample_solver == 'unipc':
+            scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1, use_dynamic_shifting=False)
+            scheduler.set_timesteps(self.sampling_steps, device=device, shift=self.shift)
+        elif self.sample_solver == 'dpm++':
+            scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1, use_dynamic_shifting=False)
+            sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)
+            _, _ = retrieve_timesteps(scheduler, device=device, sigmas=sigmas)
+        else:
+            raise NotImplementedError(self.sample_solver)
+        return scheduler
+
+    # ─── Inference ───────────────────────────────────────────────────────────
+    def _precompute_prope_for_block(self, viewmats_in, blk_start_frame, num_frame_per_block):
+        """Call PRoPE precompute once for this block's viewmats slice.
+
+        Must be invoked before any `_forward_block(..., blk_start_frame)` call
+        that targets the same block. Reused across all denoising steps within
+        the block (viewmats don't change per step) and the clean-commit pass.
+        """
+        viewmats_blk = viewmats_in[:, blk_start_frame:blk_start_frame + num_frame_per_block]
+        self.generator.model.precompute_prope(
+            viewmats_blk,
+            time_ids=blk_start_frame,
+            pre_rope=True,
+        )
+        return viewmats_blk
+
+    def _forward_block(self, blk_latents, blk_timestep, viewmats_blk,
+                       full_action_emb, blk_start_frame, kv_batch, use_cfg,
+                       write_cache=True):
+        """Run a single forward for one block, assuming PRoPE has already been
+        precomputed for this block (via `_precompute_prope_for_block`).
+
+        blk_latents:        (kv_batch, num_frame_per_block, C, H, W)
+        blk_timestep:       (kv_batch, num_frame_per_block)
+        viewmats_blk:       (kv_batch, num_frame_per_block, 4, 4)
+        full_action_emb:    (kv_batch, T_total, C, 1, 1) or None
+        write_cache:        if False, k/v are read-only (denoising steps).
+                            True for prefix / clean commit forwards (t=0).
+        """
+        num_frame_per_block = blk_latents.shape[1]
+        for cache in self.kv_cache:
+            cache["write_cache"] = write_cache
+        # Build a per-block conditional_dict: action_emb slice + block viewmats.
+        # (We omit "actions" so WanDiffusionWrapper takes the "action_emb" path.)
+        conditional_dict = {"viewmats": viewmats_blk}
+        if full_action_emb is not None:
+            conditional_dict["action_emb"] = full_action_emb[
+                :, blk_start_frame:blk_start_frame + num_frame_per_block]
+
+        flow_both, _ = self.generator(
+            noisy_image_or_video=blk_latents,
+            conditional_dict=conditional_dict,
+            timestep=blk_timestep,
+            kv_cache=self.kv_cache,
+            crossattn_cache=self.crossattn_cache,
+            current_start=blk_start_frame * self.frame_seq_length,
+            cache_start=blk_start_frame * self.frame_seq_length,
+        )
+        return flow_both
+
+    def inference(
+        self,
+        noise: torch.Tensor,              # (B, T, C, H, W) pure noise
+        actions: torch.Tensor,            # (B, T_vid, action_dim)
+        viewmats: torch.Tensor,           # (B, T, 4, 4)
+        initial_latent: torch.Tensor,     # (B, T, C, H, W) — clean latent for frames [:n_cond]
+        num_condition_frames: int = 1,
+        guidance_scale: float = 3.0,
+        sampling_steps: int = 50,
+        return_latents: bool = True,
+        skip_vae_decode: bool = False,
+    ):
+        """Block-by-block causal inference for WorldMem."""
+        if sampling_steps is not None:
+            self.sampling_steps = sampling_steps
+
+        B, T, C, H, W = noise.shape
+        device = noise.device
+        dtype = noise.dtype
+        num_frame_per_block = self.num_frame_per_block
+        n_cond = int(num_condition_frames)
+        assert 1 <= n_cond < T
+        use_cfg = guidance_scale > 1.0
+        kv_batch = 2 * B if use_cfg else B
+
+        # ─── Build cond / uncond (batched along dim 0) ───────────────────────
+        if use_cfg:
+            actions_in = torch.cat([actions, torch.zeros_like(actions)], dim=0)
+            viewmats_in = torch.cat([viewmats, viewmats], dim=0)
+        else:
+            actions_in = actions
+            viewmats_in = viewmats
+
+        # Pre-encode the full action embedding once (B, T, C, 1, 1). The block
+        # forwards pass a per-block slice via conditional_dict["action_emb"],
+        # bypassing the action encoder's "first-frame is special" packing rule
+        # which only holds when the slice starts at frame 0.
+        if getattr(self.generator, "use_action_encoder", False):
+            with torch.no_grad():
+                full_action_emb = self.generator.action_encoder(
+                    actions_in.to(dtype))   # (kv_batch, T, C, 1, 1)
+        else:
+            full_action_emb = None
+
+        # ─── Allocate KV/crossattn caches sized for kv_batch ─────────────────
+        self._initialize_kv_cache(kv_batch, dtype=dtype, device=device, num_total_frames=T)
+        # self._initialize_crossattn_cache(kv_batch, dtype=dtype, device=device)
+
+        # ─── Output buffer (B-batch, not kv_batch) ───────────────────────────
+        output = torch.zeros_like(noise)
+        output[:, :n_cond] = initial_latent[:, :n_cond]
+
+        # ─── Block layout ────────────────────────────────────────────────────
+        # Fully-clean blocks at the front are committed once at t=0. The first
+        # block that contains any noisy frames enters the denoise loop; if
+        # n_cond is not block-aligned, that block's leading `partial_clean`
+        # frames stay pinned to t=0 while the rest are denoised.
+        cleaned_blocks = n_cond // num_frame_per_block
+        partial_clean = n_cond % num_frame_per_block
+        num_total_blocks = T // num_frame_per_block
+        assert T % num_frame_per_block == 0, f"T={T} must be a multiple of num_frame_per_block={num_frame_per_block}"
+
+        # Step 1 — commit fully-clean prefix blocks at t=0.
+        # PRoPE precomputed once per block (viewmats are block-local).
+        for blk_idx in range(cleaned_blocks):
+            st = blk_idx * num_frame_per_block
+            viewmats_blk = self._precompute_prope_for_block(
+                viewmats_in, st, num_frame_per_block)
+            blk_lat = initial_latent[:, st:st + num_frame_per_block]
+            if use_cfg:
+                blk_lat = torch.cat([blk_lat, blk_lat], dim=0)
+            t0 = torch.zeros(kv_batch, num_frame_per_block, dtype=torch.float32, device=device)
+            self._forward_block(blk_lat, t0, viewmats_blk,
+                                full_action_emb, st, kv_batch, use_cfg)
+
+        # Step 2 — denoise remaining blocks one at a time, then commit each.
+        # PRoPE is precomputed ONCE per block at the top, reused across all
+        # `sampling_steps` denoising forwards AND the clean-commit pass.
+        first_noisy_block = cleaned_blocks
+        for blk_idx in tqdm(range(first_noisy_block, num_total_blocks), desc="blocks"):
+            st = blk_idx * num_frame_per_block
+            ed = st + num_frame_per_block
+
+            viewmats_blk = self._precompute_prope_for_block(
+                viewmats_in, st, num_frame_per_block)
+
+            # Initialise this block's latents.
+            blk_lat = noise[:, st:ed].clone()
+            local_partial_clean = partial_clean if blk_idx == cleaned_blocks else 0
+            if local_partial_clean > 0:
+                blk_lat[:, :local_partial_clean] = initial_latent[:, st:st + local_partial_clean]
+
+            sched = self._init_scheduler(device)
+            for step_idx, t in enumerate(sched.timesteps):
+                # Per-frame timestep within this block.
+                ts = torch.full((B, num_frame_per_block), float(t.item()),
+                                dtype=torch.float32, device=device)
+                if local_partial_clean > 0:
+                    ts[:, :local_partial_clean] = 0.0
+
+                if use_cfg:
+                    blk_lat_in = torch.cat([blk_lat, blk_lat], dim=0)
+                    ts_in = torch.cat([ts, ts], dim=0)
+                else:
+                    blk_lat_in = blk_lat
+                    ts_in = ts
+
+                flow_both = self._forward_block(
+                    blk_lat_in, ts_in, viewmats_blk,
+                    full_action_emb, st, kv_batch, use_cfg,
+                    write_cache=False)
+
+                if use_cfg:
+                    flow_cond, flow_uncond = flow_both.chunk(2, dim=0)
+                    flow_pred = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
+                else:
+                    flow_pred = flow_both
+
+                # Scheduler step on the noisy portion only.
+                noisy_flow = flow_pred[:, local_partial_clean:]
+                noisy_lat = blk_lat[:, local_partial_clean:]
+                new_noisy = sched.step(noisy_flow, t, noisy_lat, return_dict=False)[0]
+                blk_lat = blk_lat.clone()
+                blk_lat[:, local_partial_clean:] = new_noisy
+
+            output[:, st:ed] = blk_lat
+
+            # Commit clean K/V for this block to cache (overwrites the last
+            # denoise step's writes). Reuses the PRoPE precompute from above.
+            blk_lat_clean = blk_lat
+            if use_cfg:
+                blk_lat_clean = torch.cat([blk_lat_clean, blk_lat_clean], dim=0)
+            t0 = torch.zeros(kv_batch, num_frame_per_block, dtype=torch.float32, device=device)
+            self._forward_block(blk_lat_clean, t0, viewmats_blk,
+                                full_action_emb, st, kv_batch, use_cfg)
+
+        # ─── Decode ──────────────────────────────────────────────────────────
+        # When `skip_vae_decode=True`, return only the latent so the caller can
+        # control where/when the VAE runs (used by the GPU/CPU offload path
+        # that swaps generator and VAE in/out of GPU memory).
+        if skip_vae_decode:
+            if return_latents:
+                return None, output
+            return None
+
+        video = self.vae.decode_to_pixel(output.to(dtype=dtype))
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        if return_latents:
+            return video, output
+        return video
+
+
+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
+        required_cache_size = (
+            num_output_frames * self.frame_seq_length
+            if self.local_attn_size == -1
+            else self.local_attn_size * self.frame_seq_length
+        )
+        need_reinit = (
+            self.kv_cache_pos is None
+            or self.kv_cache_pos[0]["k"].shape[1] < required_cache_size
+        )
+        if need_reinit:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+                num_total_frames=num_output_frames
+            )
+            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 step_idx, 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, num_total_frames=None):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache_pos = []
+        kv_cache_neg = []
+        if self.local_attn_size != -1:
+            # Sliding window: cache size equals the window size
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        elif num_total_frames is not None:
+            # Full history: allocate enough for every frame in this inference call
+            kv_cache_size = num_total_frames * self.frame_seq_length
+        else:
+            kv_cache_size = 32760
+
+        # For full-history mode, extend max_attention_size on every self-attention
+        # module so the attention op actually reads the full cache (default is 32760).
+        if self.local_attn_size == -1:
+            for module in self.generator.model.modules():
+                if hasattr(module, 'max_attention_size'):
+                    module.max_attention_size = kv_cache_size
+
+        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

DecMem/pipeline/causal_inference.py

@@ -0,0 +1,331 @@
+from typing import List, Optional
+import torch
+
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller, move_model_to_device_with_memory_preservation
+
+
+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_cache1 = 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,
+        profile: bool = False,
+        low_memory: 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
+        )
+
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(self.text_encoder, target_device=gpu, preserved_memory_gb=gpu_memory_preservation)
+
+        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
+        required_cache_size = (
+            num_output_frames * self.frame_seq_length
+            if self.local_attn_size == -1
+            else self.local_attn_size * self.frame_seq_length
+        )
+        need_reinit = (
+            self.kv_cache1 is None
+            or self.kv_cache1[0]["k"].shape[1] < required_cache_size
+        )
+        if need_reinit:
+            self._initialize_kv_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device,
+                num_total_frames=num_output_frames
+            )
+            self._initialize_crossattn_cache(
+                batch_size=batch_size,
+                dtype=noise.dtype,
+                device=noise.device
+            )
+        else:
+            # reset cross attn cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index]["is_init"] = False
+            # reset kv cache
+            for block_index in range(len(self.kv_cache1)):
+                self.kv_cache1[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+                self.kv_cache1[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=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_cache1,
+                    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()
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+
+        for current_num_frames in all_num_frames:
+            if profile:
+                block_start.record()
+
+            noisy_input = noise[
+                :, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                print(f"current_timestep: {current_timestep}")
+                # set current timestep
+                timestep = torch.ones(
+                    [batch_size, current_num_frames],
+                    device=noise.device,
+                    dtype=torch.int64) * current_timestep
+
+                if index < len(self.denoising_step_list) - 1:
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=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
+                    _, denoised_pred = self.generator(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length
+                    )
+
+                # 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 KV cache using clean context
+            context_timestep = torch.ones_like(timestep) * self.args.context_noise
+            self.generator(
+                noisy_image_or_video=denoised_pred,
+                conditional_dict=conditional_dict,
+                timestep=context_timestep,
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=current_start_frame * self.frame_seq_length,
+            )
+
+            if profile:
+                block_end.record()
+                torch.cuda.synchronize()
+                block_time = block_start.elapsed_time(block_end)
+                block_times.append(block_time)
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        if profile:
+            # End diffusion timing and synchronize CUDA
+            diffusion_end.record()
+            torch.cuda.synchronize()
+            diffusion_time = diffusion_start.elapsed_time(diffusion_end)
+            init_time = init_start.elapsed_time(init_end)
+            vae_start.record()
+
+        # Step 4: Decode the output
+        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, num_total_frames=None):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache1 = []
+        if self.local_attn_size != -1:
+            # Sliding window: cache size equals the window size
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        elif num_total_frames is not None:
+            # Full history: allocate enough for every frame in this inference call
+            kv_cache_size = num_total_frames * self.frame_seq_length
+        else:
+            kv_cache_size = 32760
+
+        # For full-history mode, extend max_attention_size on every self-attention
+        # module so the attention op actually reads the full cache (default is 32760).
+        if self.local_attn_size == -1:
+            for module in self.generator.model.modules():
+                if hasattr(module, 'max_attention_size'):
+                    module.max_attention_size = kv_cache_size
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append({
+                "k": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "v": torch.zeros([batch_size, kv_cache_size, 12, 128], dtype=dtype, device=device),
+                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index": torch.tensor([0], dtype=torch.long, device=device)
+            })
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(self, batch_size, dtype, device):
+        """
+        Initialize 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

DecMem/pipeline/self_forcing_training.py

@@ -0,0 +1,267 @@
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+class SelfForcingTrainingPipeline:
+    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_cache1 = 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 inference_with_trajectory(
+            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_cache1 is None:
+        #     self._initialize_kv_cache(
+        #         batch_size=batch_size,
+        #         dtype=noise.dtype,
+        #         device=noise.device,
+        #     )
+        #     self._initialize_crossattn_cache(
+        #         batch_size=batch_size,
+        #         dtype=noise.dtype,
+        #         device=noise.device
+        #     )
+        # else:
+        #     # reset cross attn cache
+        #     for block_index in range(self.num_transformer_blocks):
+        #         self.crossattn_cache[block_index]["is_init"] = False
+        #     # reset kv cache
+        #     for block_index in range(len(self.kv_cache1)):
+        #         self.kv_cache1[block_index]["global_end_index"] = torch.tensor(
+        #             [0], dtype=torch.long, device=noise.device)
+        #         self.kv_cache1[block_index]["local_end_index"] = torch.tensor(
+        #             [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            # 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_cache1,
+                    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_cache1,
+                            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_cache1,
+                                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_cache1,
+                            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_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length
+                )
+
+            # 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_cache1 = []
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.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_cache1 = kv_cache1  # 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

DecMem/requirements.txt

@@ -0,0 +1,26 @@
+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
+imageio-ffmpeg
+numpy
+wandb
+omegaconf
+einops
+av==13.1.0
+opencv-python
+pycocotools
+lmdb
+matplotlib
+pydantic==2.10.6
+scikit-image
+huggingface_hub[cli]
+dominate
+ninja
+pytest
+decord

DecMem/scripts/build_worldmem_csv.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""Build WorldMem train/test CSV index files from mp4/npz pairs.
+
+Each row: video_path, action_path (paths relative to the project cwd when possible).
+
+When --root is a symlink under the project (e.g. ./data), CSV rows use paths like
+data/training/.... Avoid resolving symlinks so relative paths stay valid.
+
+Example:
+    python scripts/build_worldmem_csv.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import random
+from pathlib import Path
+
+
+def _path_for_csv(path: Path) -> str:
+    """Path for CSV rows, relative to cwd when possible (preserves symlinks)."""
+    if not path.is_absolute():
+        return path.as_posix()
+    try:
+        return path.relative_to(Path.cwd()).as_posix()
+    except ValueError:
+        return path.as_posix()
+
+
+def collect_pairs(split_dir: Path, root: Path) -> list[tuple[str, str]]:
+    """Collect (video_path, action_path) for every mp4 with a matching npz."""
+    if not split_dir.is_dir():
+        raise FileNotFoundError(f"Split directory not found: {split_dir}")
+
+    pairs: list[tuple[str, str]] = []
+    missing_npz: list[str] = []
+
+    for mp4 in sorted(split_dir.rglob("*.mp4")):
+        npz = mp4.with_suffix(".npz")
+        if not npz.is_file():
+            missing_npz.append(mp4.relative_to(root).as_posix())
+            continue
+        pairs.append((
+            _path_for_csv(mp4),
+            _path_for_csv(npz),
+        ))
+
+    if missing_npz:
+        print(f"  warning: {len(missing_npz)} mp4 files without matching npz "
+              f"(first: {missing_npz[0]})")
+
+    return pairs
+
+
+def write_csv(rows: list[tuple[str, str]], output_path: Path) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        writer.writerow(["video_path", "action_path"])
+        writer.writerows(rows)
+    print(f"  wrote {len(rows)} rows -> {output_path}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Build WorldMem train/test CSV metadata.")
+    parser.add_argument(
+        "--root",
+        type=str,
+        default='./data',
+        help=f"Dataset root directory",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default='./data',
+        help="Directory for train.csv / test.csv (default: same as --root)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for shuffling rows",
+    )
+    args = parser.parse_args()
+
+    root = Path(args.root)
+    output_dir = Path(args.output_dir) if args.output_dir else root
+
+    splits = {
+        "train.csv": root / "training",
+        "test.csv": root / "test",
+    }
+
+    print(f"Dataset root : {root} (-> {root.resolve()})")
+    print(f"Output dir   : {output_dir}")
+    print(f"Shuffle seed : {args.seed}")
+
+    rng = random.Random(args.seed)
+
+    for csv_name, split_dir in splits.items():
+        print(f"\n[{csv_name}] scanning {split_dir.relative_to(root)}/")
+        pairs = collect_pairs(split_dir, root)
+        rng.shuffle(pairs)
+        write_csv(pairs, output_dir / csv_name)
+
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

DecMem/scripts/infer.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+set -o pipefail
+
+# CONFIG SETTING
+NUM_GPUS=${NUM_GPUS:-1}
+MASTER_PORT=${MASTER_PORT:-29110}
+OFFLOAD=${OFFLOAD:-0}   # 1 or true = mutually-exclusive GPU residency for generator/VAE
+EXP_NAME=${EXP_NAME:-decmem}
+STEP=${STEP:-7000}
+printf -v STEP "%06d" "$STEP"
+
+DATA_PATH=data/test.csv
+CONFIG_PATH=${CONFIG_PATH:-configs/${EXP_NAME}.yaml}
+CHECKPOINT_PATH=${CHECKPOINT_PATH:-exps/${EXP_NAME}/checkpoint_model_${STEP}/model.pt}
+NUM_CONDITION_FRAMES=${NUM_CONDITION_FRAMES:-56}
+NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-50}
+GUIDANCE_SCALE=${GUIDANCE_SCALE:-5}
+SEED=${SEED:-42}
+START_INDEX=${START_INDEX:-0}
+N=${N:-50}
+
+T_LAT=${T_LAT:-200}   # optional; empty = use config value
+T_TAG=${T_LAT:+_T${T_LAT}}
+
+OUTPUT_DIR=${OUTPUT_DIR:-test/res/${EXP_NAME}_step_${STEP}${T_TAG}_con_${NUM_CONDITION_FRAMES}_inf_${NUM_INFERENCE_STEPS}_gs_${GUIDANCE_SCALE}}
+mkdir -p "$OUTPUT_DIR"
+RANK0_LOG="${OUTPUT_DIR}/rank0.log"
+
+SCRIPT_ARGS=(
+    --config_path "$CONFIG_PATH"
+    --checkpoint_path "$CHECKPOINT_PATH"
+    --data_path "$DATA_PATH"
+    --output_dir "$OUTPUT_DIR"
+    --num_condition_frames "$NUM_CONDITION_FRAMES"
+    --guidance_scale "$GUIDANCE_SCALE"
+    --num_inference_steps "$NUM_INFERENCE_STEPS"
+    --seed "$SEED")
+
+SCRIPT_ARGS+=(--start_index "$START_INDEX")
+if [[ -n "$N" ]]; then
+    SCRIPT_ARGS+=(--num_entries "$N")
+fi
+
+if [[ -n "$T_LAT" ]]; then
+    SCRIPT_ARGS+=(--num_latent_frames "$T_LAT")
+fi
+
+if [[ "$NUM_GPUS" -gt 1 ]]; then
+    echo "Multi-GPU inference: NUM_GPUS=${NUM_GPUS}  MASTER_PORT=${MASTER_PORT}"
+    torchrun \
+        --standalone \
+        --nproc_per_node="$NUM_GPUS" \
+        --master_port="$MASTER_PORT" \
+        --local-ranks-filter=0 \
+        infer.py \
+        "${SCRIPT_ARGS[@]}" \
+        2>&1 | tee "$RANK0_LOG"
+else
+    torchrun infer.py "${SCRIPT_ARGS[@]}" 2>&1 | tee "$RANK0_LOG"
+fi

DecMem/scripts/infer_example.sh

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+set -o pipefail
+
+# CONFIG SETTING
+NUM_GPUS=${NUM_GPUS:-1}
+MASTER_PORT=${MASTER_PORT:-29110}
+OFFLOAD=${OFFLOAD:-0}   # 1 or true = mutually-exclusive GPU residency for generator/VAE
+
+EXP_NAME=decmem
+CONFIG_PATH=${CONFIG_PATH:-configs/${EXP_NAME}.yaml}
+CHECKPOINT_PATH=checkpoints/decmem.pt
+DATA_PATH=example/example.csv
+
+NUM_CONDITION_FRAMES=${NUM_CONDITION_FRAMES:-56}
+NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
+GUIDANCE_SCALE=${GUIDANCE_SCALE:-5}
+SEED=${SEED:-42}
+START_INDEX=${START_INDEX:-0}
+N=${N:-2}
+
+T_LAT=${T_LAT:-160}   # optional; empty = use config value
+T_TAG=${T_LAT:+_T${T_LAT}}
+
+OUTPUT_DIR=${OUTPUT_DIR:-test/res/${EXP_NAME}_${T_TAG}_con_${NUM_CONDITION_FRAMES}_inf_${NUM_INFERENCE_STEPS}_gs_${GUIDANCE_SCALE}}
+mkdir -p "$OUTPUT_DIR"
+RANK0_LOG="${OUTPUT_DIR}/rank0.log"
+
+SCRIPT_ARGS=(
+    --config_path "$CONFIG_PATH"
+    --checkpoint_path "$CHECKPOINT_PATH"
+    --data_path "$DATA_PATH"
+    --output_dir "$OUTPUT_DIR"
+    --num_condition_frames "$NUM_CONDITION_FRAMES"
+    --guidance_scale "$GUIDANCE_SCALE"
+    --num_inference_steps "$NUM_INFERENCE_STEPS"
+    --seed "$SEED")
+
+SCRIPT_ARGS+=(--start_index "$START_INDEX")
+if [[ -n "$N" ]]; then
+    SCRIPT_ARGS+=(--num_entries "$N")
+fi
+
+if [[ -n "$T_LAT" ]]; then
+    SCRIPT_ARGS+=(--num_latent_frames "$T_LAT")
+fi
+
+if [[ "$NUM_GPUS" -gt 1 ]]; then
+    echo "Multi-GPU inference: NUM_GPUS=${NUM_GPUS}  MASTER_PORT=${MASTER_PORT}"
+    torchrun \
+        --standalone \
+        --nproc_per_node="$NUM_GPUS" \
+        --master_port="$MASTER_PORT" \
+        --local-ranks-filter=0 \
+        infer.py \
+        "${SCRIPT_ARGS[@]}" \
+        2>&1 | tee "$RANK0_LOG"
+else
+    torchrun infer.py "${SCRIPT_ARGS[@]}" 2>&1 | tee "$RANK0_LOG"
+fi

DecMem/scripts/train_multinode.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+# WorldMem diffusion multi-node training, launched MANUALLY on each node.
+#
+# Run this on every node, passing that node's rank as the only argument.
+# Node 0 must be the machine whose IP matches MASTER_ADDR.
+#
+# Usage (run on each node):
+#   bash scripts/train_multinode.sh <node_rank> [exp_name]
+#
+# Example with 2 nodes:
+#   # on master node:
+#   bash scripts/train_multinode.sh 0 decmem
+#   # on worker node:
+#   bash scripts/train_multinode.sh 1 decmem
+#
+# Env overrides:
+#   EXP_NAME        experiment name / config stem (default: decmem)
+#   NPROC_PER_NODE  GPUs per node (default 8)
+#   MASTER_PORT     torchrun rendezvous port (default 29500)
+#   RESUME_CKPT     path to a checkpoint to fully resume from (optional)
+
+# ─── Topology (EDIT THESE) ───────────────────────────────────────────────────
+NNODES=2                      # TODO: total number of nodes
+MASTER_ADDR=""    # TODO: IP of node_rank=0
+MASTER_PORT=${MASTER_PORT:-29500}
+NPROC_PER_NODE=${NPROC_PER_NODE:-8}
+
+# ─── Node rank / experiment from CLI ─────────────────────────────────────────
+NODE_RANK=$1
+EXP_NAME=${2:-${EXP_NAME:-decmem}}
+if [ -z "$NODE_RANK" ]; then
+    echo "Error: NODE_RANK is required. Usage: bash $0 <node_rank> [exp_name]"
+    exit 1
+fi
+
+# ─── Paths / config ──────────────────────────────────────────────────────────
+CONFIG=configs/${EXP_NAME}.yaml
+LOGDIR=exps/${EXP_NAME}
+WANDB_SAVE_DIR=${LOGDIR}/wandb
+TORCHRUN=torchrun
+RESUME_CKPT=${RESUME_CKPT:-}
+
+mkdir -p "${LOGDIR}"
+
+# Optional resume: `RESUME_CKPT=/path/to/model.pt bash <this_script> <node_rank>`
+# Must be set on EVERY node (this script doesn't propagate env across nodes).
+RESUME_ARGS=""
+if [ -n "${RESUME_CKPT}" ]; then
+    RESUME_ARGS="--resume-ckpt ${RESUME_CKPT}"
+    echo "[resume] RESUME_CKPT = ${RESUME_CKPT}"
+fi
+
+# ─── NCCL (match train_worldmem_multinode_torchrun.sh) ───────────────────────
+export NCCL_IB_DISABLE=0
+export NCCL_IB_GID_INDEX=3
+export NCCL_MIN_NCHANNELS=16
+export NCCL_IB_HCA=mlx5
+export NCCL_IB_QPS_PER_CONNECTION=4
+export NCCL_IB_TIMEOUT=22
+export NCCL_DEBUG=WARN
+
+echo "EXP_NAME        = ${EXP_NAME}"
+echo "CONFIG          = ${CONFIG}"
+echo "LOGDIR          = ${LOGDIR}"
+echo "NNODES          = ${NNODES}"
+echo "NPROC_PER_NODE  = ${NPROC_PER_NODE}"
+echo "WORLD_SIZE      = $((NNODES * NPROC_PER_NODE))"
+echo "MASTER_ADDR     = ${MASTER_ADDR}"
+echo "MASTER_PORT     = ${MASTER_PORT}"
+echo "RESUME_CKPT     = ${RESUME_CKPT:-<none>}"
+echo "Starting node with NODE_RANK=${NODE_RANK}"
+
+# ─── Launch torchrun ─────────────────────────────────────────────────────────
+${TORCHRUN} \
+    --nnodes=${NNODES} \
+    --nproc_per_node=${NPROC_PER_NODE} \
+    --node_rank=${NODE_RANK} \
+    --master_addr=${MASTER_ADDR} \
+    --master_port=${MASTER_PORT} \
+    train.py \
+    --config_path ${CONFIG} \
+    --logdir ${LOGDIR} \
+    --wandb-save-dir ${WANDB_SAVE_DIR} \
+    ${RESUME_ARGS} \
+    > "${LOGDIR}/node_${NODE_RANK}_$(date +%Y%m%d_%H%M%S).log" 2>&1

DecMem/train.py

@@ -0,0 +1,45 @@
+import argparse
+import os
+from omegaconf import OmegaConf
+import wandb
+
+from trainer import DiffusionTrainer
+
+
+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", action="store_true")
+    parser.add_argument("--resume-ckpt", type=str, default=None,
+                        help="Path to a checkpoint to fully resume from (weights + optimizer + step + wandb_run_id). "
+                            "Overrides config.resume_ckpt if both are set.")
+
+    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 args.resume_ckpt is not None:
+        config.resume_ckpt = args.resume_ckpt
+
+    trainer = DiffusionTrainer(config)
+    trainer.train()
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main()

DecMem/trainer/__init__.py

@@ -0,0 +1,7 @@
+from .diffusion import Trainer as DiffusionTrainer
+from .gan import Trainer as GANTrainer
+
+__all__ = [
+    "DiffusionTrainer",
+    "GANTrainer",
+]

DecMem/trainer/diffusion.py

@@ -0,0 +1,432 @@
+import gc
+import logging
+
+from model import CausalDiffusion
+from utils.dataset import ShardingLMDBDataset, WorldMemCSVDataset, cycle
+from utils.pose_utils import poses_to_viewmats
+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 barrier, fsdp_wrap, fsdp_state_dict, fsdp_optim_state_dict, load_fsdp_optim_state_dict, launch_distributed_job
+# from utils.distributed import EMA_FSDP
+
+
+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
+
+        # If resume_ckpt is set, load it now so we can restore step / seed / wandb_run_id
+        # BEFORE seeding RNG and initializing wandb. The full state stays in CPU memory
+        # and is consumed (then freed) after the optimizer / EMA are constructed.
+        self._resume_state = None
+        resume_path = getattr(config, "resume_ckpt", None)
+        if resume_path:
+            if global_rank == 0:
+                print(f"Loading resume checkpoint: {resume_path}")
+            self._resume_state = torch.load(resume_path, map_location="cpu", weights_only=False)
+            self.step = int(self._resume_state.get("step", 0))
+            resumed_seed = self._resume_state.get("seed", None)
+            if resumed_seed is not None:
+                config.seed = int(resumed_seed)
+            if global_rank == 0:
+                print(f"Resumed: step={self.step}, seed={config.seed}, "
+                    f"wandb_run_id={self._resume_state.get('wandb_run_id')}")
+
+        # 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_kwargs = dict(
+                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,
+            )
+            wandb_run_id = self._resume_state.get("wandb_run_id") if self._resume_state else None
+            if wandb_run_id:
+                # resume="must" → wandb errors out if the id is missing, instead of
+                # silently creating a new run that diverges from the original.
+                wandb_init_kwargs["id"] = wandb_run_id
+                wandb_init_kwargs["resume"] = "must"
+            wandb.init(**wandb_init_kwargs)
+
+        self.output_path = config.logdir
+
+        # Step 2: Initialize the model and optimizer
+        # For WorldMem training (CSV data path), skip text encoder to save GPU memory
+        data_path = config.data_path
+        self.use_worldmem_data = isinstance(data_path, str) and data_path.endswith(".csv")
+        if self.use_worldmem_data:
+            config.use_text_encoder = False
+
+        self.model = CausalDiffusion(config, device=self.device)
+
+        if global_rank == 0:
+            def _fmt(n: int) -> str:
+                for unit in ("", "K", "M", "B"):
+                    if abs(n) < 1000:
+                        return f"{n:.2f}{unit}"
+                    n /= 1000
+                return f"{n:.2f}T"
+            total = sum(p.numel() for p in self.model.generator.parameters())
+            trainable = sum(p.numel() for p in self.model.generator.parameters() if p.requires_grad)
+            print(f"[generator] params total={_fmt(total)} ({total:,})  "
+                f"trainable={_fmt(trainable)} ({trainable:,})", flush=True)
+            if self.model.text_encoder is not None:
+                te = sum(p.numel() for p in self.model.text_encoder.parameters())
+                print(f"[text_encoder] params total={_fmt(te)} ({te:,})", flush=True)
+            vae = sum(p.numel() for p in self.model.vae.parameters())
+            print(f"[vae] params total={_fmt(vae)} ({vae:,})", flush=True)
+
+        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
+        )
+
+        if self.model.text_encoder is not None:
+            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
+        if self.use_worldmem_data:
+            num_latent_frames = getattr(config, "num_frames", 21)
+            sample_strategy = getattr(config, "sample_strategy", "random")
+            dataset = WorldMemCSVDataset(
+                data_path,
+                num_latent_frames=num_latent_frames,
+                sample_strategy=sample_strategy,
+            )
+        else:
+            dataset = ShardingLMDBDataset(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,
+            pin_memory=True,
+            persistent_workers=True)
+
+        if dist.get_rank() == 0:
+            world_size = dist.get_world_size()
+            per_rank_samples = len(sampler)  # honors drop_last + world_size split
+            global_batch_size = world_size * config.batch_size
+            print("=" * 80, flush=True)
+            print(f"[dataset]      total samples : {len(dataset)}", flush=True)
+            print(f"[dataset]      per-rank      : {per_rank_samples}  "
+                  f"(drop_last=True, world_size={world_size})", flush=True)
+            print(f"[batch]        per-rank      : {config.batch_size}", flush=True)
+            print(f"[batch]        global        : {global_batch_size}  "
+                  f"(= {world_size} ranks x {config.batch_size})", flush=True)
+            print(f"[steps/epoch]                : {per_rank_samples // config.batch_size}",
+                  flush=True)
+            print("-" * 80, flush=True)
+            print("[config]")
+            print(OmegaConf.to_yaml(config, resolve=True), flush=True)
+            print("=" * 80, flush=True)
+        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 statedicts
+        if self._resume_state is not None:
+            # Full resume: weights + optimizer + EMA (step already restored in __init__ top).
+            if global_rank == 0:
+                print("Restoring generator / optimizer from resume checkpoint.")
+            # Use FULL_STATE_DICT context so FSDP scatters the (CPU) full
+            # state-dict into local GPU shards instead of leaving params on CPU.
+            from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+            from torch.distributed.fsdp import StateDictType
+            with FSDP.state_dict_type(self.model.generator, StateDictType.FULL_STATE_DICT):
+                self.model.generator.load_state_dict(self._resume_state["generator"], strict=False)
+            if "optimizer" in self._resume_state:
+                load_fsdp_optim_state_dict(
+                    self.model.generator,
+                    self.generator_optimizer,
+                    self._resume_state["optimizer"],
+                )
+            # if "generator_ema" in self._resume_state and self.generator_ema is not None:
+                # self.generator_ema.load_state_dict(self._resume_state["generator_ema"])
+            self._resume_state = None
+            torch.cuda.empty_cache()
+        elif getattr(config, "generator_ckpt", False):
+            # Warm-start: weights only, step counter and optimizer stay fresh.
+            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"]
+            # strict=False: dense ckpts lack LTM-only params (e.g. mem_proj); keep init values.
+            missing, unexpected = self.model.generator.load_state_dict(
+                state_dict, strict=False
+            )
+            if missing:
+                print(
+                    f"  warm-start missing keys ({len(missing)}): "
+                    f"{missing[:5]}{'...' if len(missing) > 5 else ''}"
+                )
+            if unexpected:
+                print(
+                    f"  warm-start unexpected keys ({len(unexpected)}): "
+                    f"{unexpected[:5]}{'...' if len(unexpected) > 5 else ''}"
+                )
+
+        ##############################################################################################################
+
+        # # 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...")
+        # Both calls are collectives: every rank must participate, only rank 0 ends up
+        # holding the full CPU tensors (rank0_only=True inside the helpers).
+        generator_state_dict = fsdp_state_dict(self.model.generator)
+        optimizer_state_dict = fsdp_optim_state_dict(
+            self.model.generator, self.generator_optimizer)
+
+        wandb_run_id = (
+            wandb.run.id
+            if (self.is_main_process and not self.disable_wandb and wandb.run is not None)
+            else None
+        )
+        state_dict = {
+            "generator": generator_state_dict,
+            "optimizer": optimizer_state_dict,
+            "step": self.step,
+            "seed": int(self.config.seed),
+            "wandb_run_id": wandb_run_id,
+        }
+        # if self.step > self.config.ema_start_step and self.generator_ema is not None:
+        #     state_dict["generator_ema"] = self.generator_ema.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):
+        torch.cuda.synchronize()
+        t1 = time.time()
+        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: Build conditional_dict
+        proportion_empty_prompts = getattr(self.config, "proportion_empty_prompts", 0.0)
+
+        if self.use_worldmem_data:
+            # WorldMem: no text conditioning; only actions and viewmats
+            conditional_dict = {}
+            unconditional_dict = {}
+
+            actions = batch["actions"].to(device=self.device, dtype=self.dtype)  # (B, T, 25)
+            poses = batch["poses"].to(device=self.device, dtype=self.dtype)      # (B, T, 5)
+            viewmats = poses_to_viewmats(poses)                                  # (B, T, 4, 4)
+
+            # CFG: randomly drop action conditioning
+            if proportion_empty_prompts > 0.0:
+                drop_mask = torch.rand(batch_size, device=self.device) < proportion_empty_prompts
+                actions = actions.clone()
+                # viewmats = viewmats.clone()
+                actions[drop_mask] = 0.0
+                # eye = torch.eye(4, dtype=viewmats.dtype, device=viewmats.device)
+                # viewmats[drop_mask] = eye.unsqueeze(0).unsqueeze(0).expand(
+                #     drop_mask.sum(), viewmats.shape[1], -1, -1)
+
+            conditional_dict["actions"] = actions
+            conditional_dict["viewmats"] = viewmats
+        else:
+            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
+                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
+
+        loss_unweighted_val = log_dict["loss_unweighted"].item()
+        wandb_loss_dict = {
+            "generator_loss": generator_loss.item(),               # weighted, used for backprop
+            "generator_loss_unweighted": loss_unweighted_val,      # raw MSE, comparable across runs
+            "generator_grad_norm": generator_grad_norm.item(),
+        }
+
+        # Per-frame unweighted loss bucketed by frame position (10 frames per bucket)
+        bucket_str = ""
+        loss_per_frame = log_dict.get("loss_per_frame", None)
+        if loss_per_frame is not None:
+            per_frame = loss_per_frame.float().mean(dim=0)  # (F,)
+            num_frame = per_frame.shape[0]
+            bucket_size = 10
+            bucket_means = []
+            for start in range(0, num_frame, bucket_size):
+                end = min(start + bucket_size, num_frame)
+                v = per_frame[start:end].mean().item()
+                bucket_means.append((start, end, v))
+                wandb_loss_dict[f"loss_frames_{start:02d}-{end - 1:02d}"] = v
+            bucket_str = "  " + "  ".join(
+                f"f{start:02d}-{end - 1:02d}={v:.4f}" for start, end, v in bucket_means
+            )
+
+        # Step 4: Logging
+        torch.cuda.synchronize()
+        t2 = time.time()
+        if self.is_main_process:
+            print(
+                f"step={self.step}  loss={generator_loss.item():.4f}  "
+                f"loss_unweighted={loss_unweighted_val:.4f}  "
+                f"grad_norm={generator_grad_norm.item():.4f}",
+                f"time={t2 - t1:.2f} seconds",
+                flush=True,
+            )
+            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()
+
+    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

DecMem/trainer/distillation.py

@@ -0,0 +1,380 @@
+import gc
+import logging
+
+from utils.dataset import ShardingLMDBDataset, cycle
+from utils.dataset import TextDataset
+from utils.distributed import EMA_FSDP, fsdp_wrap, fsdp_state_dict, launch_distributed_job
+from utils.misc import (
+    set_seed,
+    merge_dict_list
+)
+import torch.distributed as dist
+from omegaconf import OmegaConf
+from model import CausVid, DMD, SiD
+import torch
+import wandb
+import time
+import os
+
+
+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.world_size = dist.get_world_size()
+
+        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
+        if config.distribution_loss == "causvid":
+            self.model = CausVid(config, device=self.device)
+        elif config.distribution_loss == "dmd":
+            self.model = DMD(config, device=self.device)
+        elif config.distribution_loss == "sid":
+            self.model = SiD(config, device=self.device)
+        else:
+            raise ValueError("Invalid distribution matching loss")
+
+        # Save pretrained model state_dicts to CPU
+        self.fake_score_state_dict_cpu = self.model.fake_score.state_dict()
+
+        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.real_score = fsdp_wrap(
+            self.model.real_score,
+            sharding_strategy=config.sharding_strategy,
+            mixed_precision=config.mixed_precision,
+            wrap_strategy=config.real_score_fsdp_wrap_strategy
+        )
+
+        self.model.fake_score = fsdp_wrap(
+            self.model.fake_score,
+            sharding_strategy=config.sharding_strategy,
+            mixed_precision=config.mixed_precision,
+            wrap_strategy=config.fake_score_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,
+            cpu_offload=getattr(config, "text_encoder_cpu_offload", False)
+        )
+
+        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
+        )
+
+        self.critic_optimizer = torch.optim.AdamW(
+            [param for param in self.model.fake_score.parameters()
+             if param.requires_grad],
+            lr=config.lr_critic if hasattr(config, "lr_critic") else config.lr,
+            betas=(config.beta1_critic, config.beta2_critic),
+            weight_decay=config.weight_decay
+        )
+
+        # Step 3: Initialize the dataloader
+        if self.config.i2v:
+            dataset = ShardingLMDBDataset(config.data_path, max_pair=int(1e8))
+        else:
+            dataset = TextDataset(config.data_path)
+        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_generator = getattr(config, "max_grad_norm_generator", 10.0)
+        self.max_grad_norm_critic = getattr(config, "max_grad_norm_critic", 10.0)
+        self.previous_time = None
+
+    def save(self):
+        print("Start gathering distributed model states...")
+        generator_state_dict = fsdp_state_dict(
+            self.model.generator)
+        critic_state_dict = fsdp_state_dict(
+            self.model.fake_score)
+        state_dict = {
+            "generator": generator_state_dict,
+            "critic": critic_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 fwdbwd_one_step(self, batch, train_generator):
+        self.model.eval()  # prevent any randomness (e.g. dropout)
+
+        if self.step % 20 == 0:
+            torch.cuda.empty_cache()
+
+        # Step 1: Get the next batch of text prompts
+        text_prompts = batch["prompts"]
+        if self.config.i2v:
+            clean_latent = None
+            image_latent = batch["ode_latent"][:, -1][:, 0:1, ].to(
+                device=self.device, dtype=self.dtype)
+        else:
+            clean_latent = None
+            image_latent = None
+
+        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: Store gradients for the generator (if training the generator)
+        if train_generator:
+            generator_loss, generator_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 if self.config.i2v else None
+            )
+
+            generator_loss.backward()
+            generator_grad_norm = self.model.generator.clip_grad_norm_(
+                self.max_grad_norm_generator)
+
+            generator_log_dict.update({"generator_loss": generator_loss,
+                                       "generator_grad_norm": generator_grad_norm})
+
+            return generator_log_dict
+        else:
+            generator_log_dict = {}
+
+        # Step 4: Store gradients for the critic (if training the critic)
+        critic_loss, critic_log_dict = self.model.critic_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 if self.config.i2v else None
+        )
+
+        critic_loss.backward()
+        critic_grad_norm = self.model.fake_score.clip_grad_norm_(
+            self.max_grad_norm_critic)
+
+        critic_log_dict.update({"critic_loss": critic_loss,
+                                "critic_grad_norm": critic_grad_norm})
+
+        return critic_log_dict
+
+    def generate_video(self, pipeline, prompts, image=None):
+        batch_size = len(prompts)
+        if image is not None:
+            image = image.squeeze(0).unsqueeze(0).unsqueeze(2).to(device="cuda", dtype=torch.bfloat16)
+
+            # Encode the input image as the first latent
+            initial_latent = pipeline.vae.encode_to_latent(image).to(device="cuda", dtype=torch.bfloat16)
+            initial_latent = initial_latent.repeat(batch_size, 1, 1, 1, 1)
+            sampled_noise = torch.randn(
+                [batch_size, self.model.num_training_frames - 1, 16, 60, 104],
+                device="cuda",
+                dtype=self.dtype
+            )
+        else:
+            initial_latent = None
+            sampled_noise = torch.randn(
+                [batch_size, self.model.num_training_frames, 16, 60, 104],
+                device="cuda",
+                dtype=self.dtype
+            )
+
+        video, _ = pipeline.inference(
+            noise=sampled_noise,
+            text_prompts=prompts,
+            return_latents=True,
+            initial_latent=initial_latent
+        )
+        current_video = video.permute(0, 1, 3, 4, 2).cpu().numpy() * 255.0
+        return current_video
+
+    def train(self):
+        start_step = self.step
+
+        while True:
+            TRAIN_GENERATOR = self.step % self.config.dfake_gen_update_ratio == 0
+
+            # Train the generator
+            if TRAIN_GENERATOR:
+                self.generator_optimizer.zero_grad(set_to_none=True)
+                extras_list = []
+                batch = next(self.dataloader)
+                extra = self.fwdbwd_one_step(batch, True)
+                extras_list.append(extra)
+                generator_log_dict = merge_dict_list(extras_list)
+                self.generator_optimizer.step()
+                # if self.generator_ema is not None:
+                #     self.generator_ema.update(self.model.generator)
+
+            # Train the critic
+            self.critic_optimizer.zero_grad(set_to_none=True)
+            extras_list = []
+            batch = next(self.dataloader)
+            extra = self.fwdbwd_one_step(batch, False)
+            extras_list.append(extra)
+            critic_log_dict = merge_dict_list(extras_list)
+            self.critic_optimizer.step()
+
+            # Increment the step since we finished gradient update
+            self.step += 1
+
+            # # Create EMA params (if not already created)
+            # if (self.step >= self.config.ema_start_step) and \
+            #         (self.generator_ema is None) and (self.config.ema_weight > 0):
+            #     self.generator_ema = EMA_FSDP(self.model.generator, decay=self.config.ema_weight)
+
+            # Save the model
+            if (not self.config.no_save) and (self.step - start_step) > 0 and self.step % self.config.log_iters == 0:
+                torch.cuda.empty_cache()
+                self.save()
+                torch.cuda.empty_cache()
+
+            # Logging
+            if self.is_main_process:
+                wandb_loss_dict = {}
+                if TRAIN_GENERATOR:
+                    wandb_loss_dict.update(
+                        {
+                            "generator_loss": generator_log_dict["generator_loss"].mean().item(),
+                            "generator_grad_norm": generator_log_dict["generator_grad_norm"].mean().item(),
+                            "dmdtrain_gradient_norm": generator_log_dict["dmdtrain_gradient_norm"].mean().item()
+                        }
+                    )
+
+                wandb_loss_dict.update(
+                    {
+                        "critic_loss": critic_log_dict["critic_loss"].mean().item(),
+                        "critic_grad_norm": critic_log_dict["critic_grad_norm"].mean().item()
+                    }
+                )
+
+                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()
+                torch.cuda.empty_cache()
+
+            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

DecMem/trainer/gan.py

@@ -0,0 +1,452 @@
+import gc
+import logging
+
+from utils.dataset import ShardingLMDBDataset, cycle
+from utils.distributed import EMA_FSDP, fsdp_wrap, fsdp_state_dict, launch_distributed_job
+from utils.misc import (
+    set_seed,
+    merge_dict_list
+)
+import torch.distributed as dist
+from omegaconf import OmegaConf
+from model import GAN
+import torch
+import wandb
+import time
+import os
+
+
+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.world_size = dist.get_world_size()
+
+        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
+
+        # Configuration for discriminator warmup
+        self.discriminator_warmup_steps = getattr(config, "discriminator_warmup_steps", 0)
+        self.in_discriminator_warmup = self.step < self.discriminator_warmup_steps
+        if self.in_discriminator_warmup and self.is_main_process:
+            print(f"Starting with discriminator warmup for {self.discriminator_warmup_steps} steps")
+        self.loss_scale = getattr(config, "loss_scale", 1.0)
+
+        # 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 = GAN(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.fake_score = fsdp_wrap(
+            self.model.fake_score,
+            sharding_strategy=config.sharding_strategy,
+            mixed_precision=config.mixed_precision,
+            wrap_strategy=config.fake_score_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,
+            cpu_offload=getattr(config, "text_encoder_cpu_offload", False)
+        )
+
+        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.gen_lr,
+            betas=(config.beta1, config.beta2)
+        )
+
+        # Create separate parameter groups for the fake_score network
+        # One group for parameters with "_cls_pred_branch" or "_gan_ca_blocks" in the name
+        # and another group for all other parameters
+        fake_score_params = []
+        discriminator_params = []
+
+        for name, param in self.model.fake_score.named_parameters():
+            if param.requires_grad:
+                if "_cls_pred_branch" in name or "_gan_ca_blocks" in name:
+                    discriminator_params.append(param)
+                else:
+                    fake_score_params.append(param)
+
+        # Use the special learning rate for the special parameter group
+        # and the default critic learning rate for other parameters
+        self.critic_param_groups = [
+            {'params': fake_score_params, 'lr': config.critic_lr},
+            {'params': discriminator_params, 'lr': config.critic_lr * config.discriminator_lr_multiplier}
+        ]
+        if self.in_discriminator_warmup:
+            self.critic_optimizer = torch.optim.AdamW(
+                self.critic_param_groups,
+                betas=(0.9, config.beta2_critic)
+            )
+        else:
+            self.critic_optimizer = torch.optim.AdamW(
+                self.critic_param_groups,
+                betas=(config.beta1_critic, config.beta2_critic)
+            )
+
+        # Step 3: Initialize the dataloader
+        self.data_path = config.data_path
+        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
+            )
+        if hasattr(config, "load"):
+            resume_ckpt_path_critic = os.path.join(config.load, "critic")
+            resume_ckpt_path_generator = os.path.join(config.load, "generator")
+        else:
+            resume_ckpt_path_critic = "none"
+            resume_ckpt_path_generator = "none"
+
+        _, _ = self.checkpointer_critic.try_best_load(
+            resume_ckpt_path=resume_ckpt_path_critic,
+        )
+        self.step, _ = self.checkpointer_generator.try_best_load(
+            resume_ckpt_path=resume_ckpt_path_generator,
+            force_start_w_ema=config.force_start_w_ema,
+            force_reset_zero_step=config.force_reset_zero_step,
+            force_reinit_ema=config.force_reinit_ema,
+            skip_optimizer_scheduler=config.skip_optimizer_scheduler,
+        )
+
+        ##############################################################################################################
+
+        # # 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_generator = getattr(config, "max_grad_norm_generator", 10.0)
+        self.max_grad_norm_critic = getattr(config, "max_grad_norm_critic", 10.0)
+        self.previous_time = None
+
+    def save(self):
+        print("Start gathering distributed model states...")
+        generator_state_dict = fsdp_state_dict(
+            self.model.generator)
+        critic_state_dict = fsdp_state_dict(
+            self.model.fake_score)
+
+        state_dict = {
+            "generator": generator_state_dict,
+            "critic": critic_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 fwdbwd_one_step(self, batch, train_generator):
+        self.model.eval()  # prevent any randomness (e.g. dropout)
+
+        if self.step % 20 == 0:
+            torch.cuda.empty_cache()
+
+        # Step 1: Get the next batch of text prompts
+        text_prompts = batch["prompts"]  # next(self.dataloader)
+        if "ode_latent" in batch:
+            clean_latent = batch["ode_latent"][:, -1].to(device=self.device, dtype=self.dtype)
+        else:
+            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
+
+        mini_bs, full_bs = (
+            batch["mini_bs"],
+            batch["full_bs"],
+        )
+
+        # Step 3: Store gradients for the generator (if training the generator)
+        if train_generator:
+            gan_G_loss = 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 if self.config.i2v else None
+            )
+
+            loss_ratio = mini_bs * self.world_size / full_bs
+            total_loss = gan_G_loss * loss_ratio * self.loss_scale
+
+            total_loss.backward()
+            generator_grad_norm = self.model.generator.clip_grad_norm_(
+                self.max_grad_norm_generator)
+
+            generator_log_dict = {"generator_grad_norm": generator_grad_norm,
+                                  "gan_G_loss": gan_G_loss}
+
+            return generator_log_dict
+        else:
+            generator_log_dict = {}
+
+        # Step 4: Store gradients for the critic (if training the critic)
+        (gan_D_loss, r1_loss, r2_loss), critic_log_dict = self.model.critic_loss(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            clean_latent=clean_latent,
+            real_image_or_video=clean_latent,
+            initial_latent=image_latent if self.config.i2v else None
+        )
+
+        loss_ratio = mini_bs * dist.get_world_size() / full_bs
+        total_loss = (gan_D_loss + 0.5 * (r1_loss + r2_loss)) * loss_ratio * self.loss_scale
+
+        total_loss.backward()
+        critic_grad_norm = self.model.fake_score.clip_grad_norm_(
+            self.max_grad_norm_critic)
+
+        critic_log_dict.update({"critic_grad_norm": critic_grad_norm,
+                                "gan_D_loss": gan_D_loss,
+                                "r1_loss": r1_loss,
+                                "r2_loss": r2_loss})
+
+        return critic_log_dict
+
+    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):
+        start_step = self.step
+
+        while True:
+            if self.step == self.discriminator_warmup_steps and self.discriminator_warmup_steps != 0:
+                print("Resetting critic optimizer")
+                del self.critic_optimizer
+                torch.cuda.empty_cache()
+                # Create new optimizers
+                self.critic_optimizer = torch.optim.AdamW(
+                    self.critic_param_groups,
+                    betas=(self.config.beta1_critic, self.config.beta2_critic)
+                )
+                # Update checkpointer references
+                self.checkpointer_critic.optimizer = self.critic_optimizer
+            # Check if we're in the discriminator warmup phase
+            self.in_discriminator_warmup = self.step < self.discriminator_warmup_steps
+
+            # Only update generator and critic outside the warmup phase
+            TRAIN_GENERATOR = not self.in_discriminator_warmup and self.step % self.config.dfake_gen_update_ratio == 0
+
+            # Train the generator (only outside warmup phase)
+            if TRAIN_GENERATOR:
+                self.model.fake_score.requires_grad_(False)
+                self.model.generator.requires_grad_(True)
+                self.generator_optimizer.zero_grad(set_to_none=True)
+                extras_list = []
+                for ii, mini_batch in enumerate(self.dataloader.next()):
+                    extra = self.fwdbwd_one_step(mini_batch, True)
+                    extras_list.append(extra)
+                generator_log_dict = merge_dict_list(extras_list)
+                self.generator_optimizer.step()
+                # if self.generator_ema is not None:
+                #     self.generator_ema.update(self.model.generator)
+            else:
+                generator_log_dict = {}
+
+            # Train the critic/discriminator
+            if self.in_discriminator_warmup:
+                # During warmup, only allow gradient for discriminator params
+                self.model.generator.requires_grad_(False)
+                self.model.fake_score.requires_grad_(False)
+
+                # Enable gradient only for discriminator params
+                for name, param in self.model.fake_score.named_parameters():
+                    if "_cls_pred_branch" in name or "_gan_ca_blocks" in name:
+                        param.requires_grad_(True)
+            else:
+                # Normal training mode
+                self.model.generator.requires_grad_(False)
+                self.model.fake_score.requires_grad_(True)
+
+            self.critic_optimizer.zero_grad(set_to_none=True)
+            extras_list = []
+            batch = next(self.dataloader)
+            extra = self.fwdbwd_one_step(batch, False)
+            extras_list.append(extra)
+            critic_log_dict = merge_dict_list(extras_list)
+            self.critic_optimizer.step()
+
+            # Increment the step since we finished gradient update
+            self.step += 1
+
+            # If we just finished warmup, print a message
+            if self.is_main_process and self.step == self.discriminator_warmup_steps:
+                print(f"Finished discriminator warmup after {self.discriminator_warmup_steps} steps")
+
+            # Save the model
+            if (not self.config.no_save) and (self.step - start_step) > 0 and self.step % self.config.log_iters == 0:
+                torch.cuda.empty_cache()
+                self.save()
+                torch.cuda.empty_cache()
+
+            # Logging
+            wandb_loss_dict = {
+                "generator_grad_norm": generator_log_dict["generator_grad_norm"],
+                "critic_grad_norm": critic_log_dict["critic_grad_norm"],
+                "real_logit": critic_log_dict["noisy_real_logit"],
+                "fake_logit": critic_log_dict["noisy_fake_logit"],
+                "r1_loss": critic_log_dict["r1_loss"],
+                "r2_loss": critic_log_dict["r2_loss"],
+            }
+            if TRAIN_GENERATOR:
+                wandb_loss_dict.update({
+                    "generator_grad_norm": generator_log_dict["generator_grad_norm"],
+                })
+            self.all_gather_dict(wandb_loss_dict)
+            wandb_loss_dict["diff_logit"] = wandb_loss_dict["real_logit"] - wandb_loss_dict["fake_logit"]
+            wandb_loss_dict["reg_loss"] = 0.5 * (wandb_loss_dict["r1_loss"] + wandb_loss_dict["r2_loss"])
+
+            if self.is_main_process:
+                if self.in_discriminator_warmup:
+                    warmup_status = f"[WARMUP {self.step}/{self.discriminator_warmup_steps}] Training only discriminator params"
+                    print(warmup_status)
+                    if not self.disable_wandb:
+                        wandb_loss_dict.update({"warmup_status": 1.0})
+
+                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()
+                torch.cuda.empty_cache()
+
+            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
+
+    def all_gather_dict(self, target_dict):
+        for key, value in target_dict.items():
+            gathered_value = torch.zeros(
+                [self.world_size, *value.shape],
+                dtype=value.dtype, device=self.device)
+            dist.all_gather_into_tensor(gathered_value, value)
+            avg_value = gathered_value.mean().item()
+            target_dict[key] = avg_value

DecMem/trainer/ode.py

@@ -0,0 +1,242 @@
+import gc
+import logging
+from utils.dataset import ODERegressionLMDBDataset, cycle
+from model import ODERegression
+from collections import defaultdict
+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 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.world_size = dist.get_world_size()
+
+        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.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
+
+        assert config.distribution_loss == "ode", "Only ODE loss is supported for ODE training"
+        self.model = ODERegression(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,
+            cpu_offload=getattr(config, "text_encoder_cpu_offload", False)
+        )
+
+        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 = ODERegressionLMDBDataset(
+            config.data_path, max_pair=getattr(config, "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)
+        total_batch_size = getattr(config, "total_batch_size", None)
+        if total_batch_size is not None:
+            assert total_batch_size == config.batch_size * self.world_size, "Gradient accumulation is not supported for ODE training"
+        self.dataloader = cycle(dataloader)
+
+        self.step = 0
+
+        ##############################################################################################################
+        # 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")[
+                'generator']
+            self.model.generator.load_state_dict(
+                state_dict, strict=True
+            )
+
+        ##############################################################################################################
+
+        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)
+        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):
+        VISUALIZE = self.step % 100 == 0
+        self.model.eval()  # prevent any randomness (e.g. dropout)
+
+        # Step 1: Get the next batch of text prompts
+        batch = next(self.dataloader)
+        text_prompts = batch["prompts"]
+        ode_latent = batch["ode_latent"].to(
+            device=self.device, dtype=self.dtype)
+
+        # Step 2: Extract the conditional infos
+        with torch.no_grad():
+            conditional_dict = self.model.text_encoder(
+                text_prompts=text_prompts)
+
+        # Step 3: Train the generator
+        generator_loss, log_dict = self.model.generator_loss(
+            ode_latent=ode_latent,
+            conditional_dict=conditional_dict
+        )
+
+        unnormalized_loss = log_dict["unnormalized_loss"]
+        timestep = log_dict["timestep"]
+
+        if self.world_size > 1:
+            gathered_unnormalized_loss = torch.zeros(
+                [self.world_size, *unnormalized_loss.shape],
+                dtype=unnormalized_loss.dtype, device=self.device)
+            gathered_timestep = torch.zeros(
+                [self.world_size, *timestep.shape],
+                dtype=timestep.dtype, device=self.device)
+
+            dist.all_gather_into_tensor(
+                gathered_unnormalized_loss, unnormalized_loss)
+            dist.all_gather_into_tensor(gathered_timestep, timestep)
+        else:
+            gathered_unnormalized_loss = unnormalized_loss
+            gathered_timestep = timestep
+
+        loss_breakdown = defaultdict(list)
+        stats = {}
+
+        for index, t in enumerate(timestep):
+            loss_breakdown[str(int(t.item()) // 250 * 250)].append(
+                unnormalized_loss[index].item())
+
+        for key_t in loss_breakdown.keys():
+            stats["loss_at_time_" + key_t] = sum(loss_breakdown[key_t]) / \
+                len(loss_breakdown[key_t])
+
+        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()
+
+        # Step 4: Visualization
+        if VISUALIZE and not self.config.no_visualize and not self.config.disable_wandb and self.is_main_process:
+            # Visualize the input, output, and ground truth
+            input = log_dict["input"]
+            output = log_dict["output"]
+            ground_truth = ode_latent[:, -1]
+
+            input_video = self.model.vae.decode_to_pixel(input)
+            output_video = self.model.vae.decode_to_pixel(output)
+            ground_truth_video = self.model.vae.decode_to_pixel(ground_truth)
+            input_video = 255.0 * (input_video.cpu().numpy() * 0.5 + 0.5)
+            output_video = 255.0 * (output_video.cpu().numpy() * 0.5 + 0.5)
+            ground_truth_video = 255.0 * (ground_truth_video.cpu().numpy() * 0.5 + 0.5)
+
+            # Visualize the input, output, and ground truth
+            wandb.log({
+                "input": wandb.Video(input_video, caption="Input", fps=16, format="mp4"),
+                "output": wandb.Video(output_video, caption="Output", fps=16, format="mp4"),
+                "ground_truth": wandb.Video(ground_truth_video, caption="Ground Truth", fps=16, format="mp4"),
+            }, step=self.step)
+
+        # Step 5: Logging
+        if self.is_main_process and not self.disable_wandb:
+            wandb_loss_dict = {
+                "generator_loss": generator_loss.item(),
+                "generator_grad_norm": generator_grad_norm.item(),
+                **stats
+            }
+            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()
+
+    def train(self):
+        while True:
+            self.train_one_step()
+            if (not self.config.no_save) and self.step % self.config.log_iters == 0:
+                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
+
+            self.step += 1

DecMem/utils/action_encoder.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class ActionEncoder(nn.Module):
+    """Encode raw per-video-frame action vectors and project to latent spatial dim.
+
+    Input:  (B, T_vid, action_dim) where T_vid = (T_lat - 1) * 4 + 1
+    Output: (B, T_lat, C, 1, 1) broadcastable latent-space embeddings.
+
+    Packing strategy:
+      - Frame 0 (first latent frame): the single raw action is repeated 4 times
+        so it has the same packed dimension as grouped frames.
+      - Frames 1..(T_vid-1): every 4 consecutive raw actions are flattened into
+        one latent frame, giving T_lat - 1 grouped frames.
+    The linear layer therefore receives 4 * action_dim features per latent frame.
+    """
+
+    def __init__(self, action_dim: int = 25, hidden_dim: int = 64, latent_channels: int = 16):
+        super().__init__()
+        self.latent_channels = latent_channels
+        # self.net = nn.Sequential(
+        #     nn.Linear(action_dim * 4, hidden_dim),
+        #     nn.SiLU(),
+        #     nn.Linear(hidden_dim, latent_channels),
+        # )
+        self.net = nn.Linear(action_dim * 4, latent_channels, bias=False)
+        
+    def forward(self, actions: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            actions: (B, T_vid, action_dim)  T_vid = (T_lat - 1) * 4 + 1
+        Returns:
+            (B, T_lat, C, 1, 1)
+        """
+        B, T_vid, D = actions.shape
+        # First frame: repeat raw action 4 times to match packed dim of grouped frames
+        action_first = actions[:, :1].unsqueeze(2).repeat(1, 1, 4, 1)  # (B, 1, 4, D)
+        # Remaining frames: group every 4 raw frames into one latent frame
+        actions_rest = actions[:, 1:].reshape(B, (T_vid - 1) // 4, 4, D)  # (B, T_lat-1, 4, D)
+        # Concatenate along the latent-frame dimension
+        actions = torch.cat([action_first, actions_rest], dim=1)  # (B, T_lat, 4, D)
+        # Flatten 4 raw actions into the feature dim
+        actions = rearrange(actions, "b f p d -> b f (p d)")  # (B, T_lat, 4*D)
+        emb = self.net(actions)                                # (B, T_lat, C)
+        return emb.unsqueeze(-1).unsqueeze(-1)                 # (B, T_lat, C, 1, 1)

DecMem/utils/dataset.py

@@ -0,0 +1,306 @@
+from utils.lmdb import get_array_shape_from_lmdb, retrieve_row_from_lmdb
+from utils.pose_utils import extract_pose_action_from_npz
+from torch.utils.data import Dataset
+import numpy as np
+import torch
+import lmdb
+import json
+import csv
+from pathlib import Path
+from PIL import Image
+import os
+
+
+class TextDataset(Dataset):
+    def __init__(self, prompt_path, extended_prompt_path=None):
+        with open(prompt_path, encoding="utf-8") as f:
+            self.prompt_list = [line.rstrip() for line in f]
+
+        if extended_prompt_path is not None:
+            with open(extended_prompt_path, encoding="utf-8") as f:
+                self.extended_prompt_list = [line.rstrip() for line in f]
+            assert len(self.extended_prompt_list) == len(self.prompt_list)
+        else:
+            self.extended_prompt_list = None
+
+    def __len__(self):
+        return len(self.prompt_list)
+
+    def __getitem__(self, idx):
+        batch = {
+            "prompts": self.prompt_list[idx],
+            "idx": idx,
+        }
+        if self.extended_prompt_list is not None:
+            batch["extended_prompts"] = self.extended_prompt_list[idx]
+        return batch
+
+
+class ODERegressionLMDBDataset(Dataset):
+    def __init__(self, data_path: str, max_pair: int = int(1e8)):
+        self.env = lmdb.open(data_path, readonly=True,
+                             lock=False, readahead=False, meminit=False)
+
+        self.latents_shape = get_array_shape_from_lmdb(self.env, 'latents')
+        self.max_pair = max_pair
+
+    def __len__(self):
+        return min(self.latents_shape[0], self.max_pair)
+
+    def __getitem__(self, idx):
+        """
+        Outputs:
+            - prompts: List of Strings
+            - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
+        """
+        latents = retrieve_row_from_lmdb(
+            self.env,
+            "latents", np.float16, idx, shape=self.latents_shape[1:]
+        )
+
+        if len(latents.shape) == 4:
+            latents = latents[None, ...]
+
+        prompts = retrieve_row_from_lmdb(
+            self.env,
+            "prompts", str, idx
+        )
+        return {
+            "prompts": prompts,
+            "ode_latent": torch.tensor(latents, dtype=torch.float32)
+        }
+
+
+class ShardingLMDBDataset(Dataset):
+    def __init__(self, data_path: str, max_pair: int = int(1e8)):
+        self.envs = []
+        self.index = []
+
+        for fname in sorted(os.listdir(data_path)):
+            path = os.path.join(data_path, fname)
+            env = lmdb.open(path,
+                            readonly=True,
+                            lock=False,
+                            readahead=False,
+                            meminit=False)
+            self.envs.append(env)
+
+        self.latents_shape = [None] * len(self.envs)
+        for shard_id, env in enumerate(self.envs):
+            self.latents_shape[shard_id] = get_array_shape_from_lmdb(env, 'latents')
+            for local_i in range(self.latents_shape[shard_id][0]):
+                self.index.append((shard_id, local_i))
+
+            # print("shard_id ", shard_id, " local_i ", local_i)
+
+        self.max_pair = max_pair
+
+    def __len__(self):
+        return len(self.index)
+
+    def __getitem__(self, idx):
+        """
+            Outputs:
+                - prompts: List of Strings
+                - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
+        """
+        shard_id, local_idx = self.index[idx]
+
+        latents = retrieve_row_from_lmdb(
+            self.envs[shard_id],
+            "latents", np.float16, local_idx,
+            shape=self.latents_shape[shard_id][1:]
+        )
+
+        if len(latents.shape) == 4:
+            latents = latents[None, ...]
+
+        prompts = retrieve_row_from_lmdb(
+            self.envs[shard_id],
+            "prompts", str, local_idx
+        )
+
+        return {
+            "prompts": prompts,
+            "ode_latent": torch.tensor(latents, dtype=torch.float32)
+        }
+
+
+class TextImagePairDataset(Dataset):
+    def __init__(
+        self,
+        data_dir,
+        transform=None,
+        eval_first_n=-1,
+        pad_to_multiple_of=None
+    ):
+        """
+        Args:
+            data_dir (str): Path to the directory containing:
+                - target_crop_info_*.json (metadata file)
+                - */ (subdirectory containing images with matching aspect ratio)
+            transform (callable, optional): Optional transform to be applied on the image
+        """
+        self.transform = transform
+        data_dir = Path(data_dir)
+
+        # Find the metadata JSON file
+        metadata_files = list(data_dir.glob('target_crop_info_*.json'))
+        if not metadata_files:
+            raise FileNotFoundError(f"No metadata file found in {data_dir}")
+        if len(metadata_files) > 1:
+            raise ValueError(f"Multiple metadata files found in {data_dir}")
+
+        metadata_path = metadata_files[0]
+        # Extract aspect ratio from metadata filename (e.g. target_crop_info_26-15.json -> 26-15)
+        aspect_ratio = metadata_path.stem.split('_')[-1]
+
+        # Use aspect ratio subfolder for images
+        self.image_dir = data_dir / aspect_ratio
+        if not self.image_dir.exists():
+            raise FileNotFoundError(f"Image directory not found: {self.image_dir}")
+
+        # Load metadata
+        with open(metadata_path, 'r') as f:
+            self.metadata = json.load(f)
+
+        eval_first_n = eval_first_n if eval_first_n != -1 else len(self.metadata)
+        self.metadata = self.metadata[:eval_first_n]
+
+        # Verify all images exist
+        for item in self.metadata:
+            image_path = self.image_dir / item['file_name']
+            if not image_path.exists():
+                raise FileNotFoundError(f"Image not found: {image_path}")
+
+        self.dummy_prompt = "DUMMY PROMPT"
+        self.pre_pad_len = len(self.metadata)
+        if pad_to_multiple_of is not None and len(self.metadata) % pad_to_multiple_of != 0:
+            # Duplicate the last entry
+            self.metadata += [self.metadata[-1]] * (
+                pad_to_multiple_of - len(self.metadata) % pad_to_multiple_of
+            )
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        """
+        Returns:
+            dict: A dictionary containing:
+                - image: PIL Image
+                - caption: str
+                - target_bbox: list of int [x1, y1, x2, y2]
+                - target_ratio: str
+                - type: str
+                - origin_size: tuple of int (width, height)
+        """
+        item = self.metadata[idx]
+
+        # Load image
+        image_path = self.image_dir / item['file_name']
+        image = Image.open(image_path).convert('RGB')
+
+        # Apply transform if specified
+        if self.transform:
+            image = self.transform(image)
+
+        return {
+            'image': image,
+            'prompts': item['caption'],
+            'target_bbox': item['target_crop']['target_bbox'],
+            'target_ratio': item['target_crop']['target_ratio'],
+            'type': item['type'],
+            'origin_size': (item['origin_width'], item['origin_height']),
+            'idx': idx
+        }
+
+
+class WorldMemCSVDataset(Dataset):
+    """Dataset for WorldMem-format data: CSV index + NPZ action/pose files + raw MP4 videos.
+
+    CSV columns (as produced by WorldMem):
+        video_path  — path to MP4 video (640×360 @10fps Minecraft footage)
+        action_path — path to NPZ file with 'actions' (N, 8) and 'poses' (N, 5)
+
+    Loads (num_latent_frames - 1) * 4 + 1 = 81 raw video frames per sample and returns
+    them as (C, T_vid, H, W). The trainer encodes with the WAN VAE to obtain
+    num_latent_frames=21 latent frames (first frame kept, rest 4× compressed).
+    Actions/poses are subsampled to match latent frame count (one per latent frame).
+    Resolution 352×640 avoids the odd latent-height issue (360/8=45 odd).
+    """
+
+    VIDEO_H = 352
+    VIDEO_W = 640
+    SKIP_FRAMES = 0      # skip unstable early frames
+
+    def __init__(self, csv_path: str, num_latent_frames: int = 21, sample_strategy: str = "random"):
+        assert sample_strategy in ("random", "first"), \
+            f"sample_strategy must be 'random' or 'first', got {sample_strategy!r}"
+        self.sample_strategy = sample_strategy
+        self.num_latent_frames = num_latent_frames
+        # WAN VAE: T_vid = (T_lat - 1) * 4 + 1
+        self.num_video_frames = (num_latent_frames - 1) * 4 + 1
+        # Pose subsample indices: frame 0, then first frame of each 4-frame group
+        self.pose_indices = [k * 4 for k in range(num_latent_frames)]
+        self.rows = []
+        with open(csv_path, newline="", encoding="utf-8") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                self.rows.append(row)
+
+    def __len__(self):
+        return len(self.rows)
+
+    def _load_video_frames(self, video_path: str, start_frame: int) -> torch.Tensor:
+        """Load num_video_frames frames from an MP4, returns (C, T, H, W) float32 in [-1,1]."""
+        import decord
+        decord.bridge.set_bridge("torch")
+        reader = decord.VideoReader(
+            video_path,
+            ctx=decord.cpu(0),
+            height=self.VIDEO_H,
+            width=self.VIDEO_W,
+        )
+        total = len(reader)
+        end_frame = min(start_frame + self.num_video_frames, total)
+        frame_indices = list(range(start_frame, end_frame))
+        frames = reader.get_batch(frame_indices)  # (T, H, W, C) uint8, torch.Tensor
+        # (T, H, W, C) uint8 -> (C, T, H, W) float32 in [-1, 1]
+        return frames.permute(3, 0, 1, 2).to(torch.float32).div(127.5).sub(1.0).contiguous()
+
+    def __getitem__(self, idx):
+        import decord
+        row = self.rows[idx]
+        video_path = row["video_path"]
+        action_path = row["action_path"]
+
+        reader = decord.VideoReader(video_path, ctx=decord.cpu(0))
+        total = len(reader)
+        if self.sample_strategy == "first":
+            start_frame = self.SKIP_FRAMES
+        else:
+            max_start = max(self.SKIP_FRAMES, total - self.num_video_frames)
+            start_frame = int(np.random.randint(self.SKIP_FRAMES, max_start + 1))
+
+        frames = self._load_video_frames(video_path, start_frame)  # (C, T_vid, H, W)
+
+        actions_full, poses_full = extract_pose_action_from_npz(
+            action_path, start_frame, self.num_video_frames
+        )  # (T_vid, 25), (T_vid, 5)
+
+        actions = actions_full                        # (T_vid, 25) — all raw frames
+        poses   = poses_full[self.pose_indices]       # (T_lat, 5)  — one per latent frame
+
+        return {
+            "prompts": "",        # unused for WorldMem
+            "frames": frames,     # (C, T_vid, H, W) float32 in [-1, 1]
+            "actions": actions,   # (T_vid, 25)
+            "poses": poses,       # (T_lat, 5)
+        }
+
+
+def cycle(dl):
+    while True:
+        for data in dl:
+            yield data

DecMem/utils/distributed.py

@@ -0,0 +1,158 @@
+from datetime import timedelta
+from functools import partial
+import os
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
+from torch.distributed.fsdp.api import CPUOffload
+from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
+
+
+def fsdp_state_dict(model):
+    fsdp_fullstate_save_policy = FullStateDictConfig(
+        offload_to_cpu=True, rank0_only=True
+    )
+    with FSDP.state_dict_type(
+        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy
+    ):
+        checkpoint = model.state_dict()
+
+    return checkpoint
+
+
+def fsdp_optim_state_dict(model, optimizer):
+    """Gather the full optimizer state on rank 0 (CPU). Other ranks return an empty dict."""
+    state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    optim_state_dict_config = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    with FSDP.state_dict_type(
+        model,
+        StateDictType.FULL_STATE_DICT,
+        state_dict_config=state_dict_config,
+        optim_state_dict_config=optim_state_dict_config,
+    ):
+        return FSDP.optim_state_dict(model, optimizer)
+
+
+def load_fsdp_optim_state_dict(model, optimizer, full_optim_state_dict):
+    """Scatter a full optimizer state (valid on rank 0) into the local FSDP shards."""
+    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
+        sharded = FSDP.optim_state_dict_to_load(
+            model=model,
+            optim=optimizer,
+            optim_state_dict=full_optim_state_dict,
+        )
+        optimizer.load_state_dict(sharded)
+
+
+def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
+    if mixed_precision:
+        mixed_precision_policy = MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            buffer_dtype=torch.float32,
+            cast_forward_inputs=False
+        )
+    else:
+        mixed_precision_policy = None
+
+    if wrap_strategy == "transformer":
+        auto_wrap_policy = partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls=transformer_module
+        )
+    elif wrap_strategy == "size":
+        auto_wrap_policy = partial(
+            size_based_auto_wrap_policy,
+            min_num_params=min_num_params
+        )
+    else:
+        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")
+
+    os.environ["NCCL_CROSS_NIC"] = "1"
+
+    sharding_strategy = {
+        "full": ShardingStrategy.FULL_SHARD,
+        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
+        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
+        "no_shard": ShardingStrategy.NO_SHARD,
+    }[sharding_strategy]
+
+    module = FSDP(
+        module,
+        auto_wrap_policy=auto_wrap_policy,
+        sharding_strategy=sharding_strategy,
+        mixed_precision=mixed_precision_policy,
+        device_id=torch.cuda.current_device(),
+        limit_all_gathers=True,
+        use_orig_params=True,
+        cpu_offload=CPUOffload(offload_params=cpu_offload),
+        sync_module_states=False  # Load ckpt on rank 0 and sync to other ranks
+    )
+    return module
+
+
+def barrier():
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def launch_distributed_job(backend: str = "nccl"):
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    host = os.environ["MASTER_ADDR"]
+    port = int(os.environ["MASTER_PORT"])
+
+    if ":" in host:  # IPv6
+        init_method = f"tcp://[{host}]:{port}"
+    else:  # IPv4
+        init_method = f"tcp://{host}:{port}"
+    dist.init_process_group(rank=rank, world_size=world_size, backend=backend,
+                            init_method=init_method, timeout=timedelta(minutes=30))
+    torch.cuda.set_device(local_rank)
+
+
+class EMA_FSDP:
+    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
+        self.decay = decay
+        self.shadow = {}
+        self._init_shadow(fsdp_module)
+
+    @torch.no_grad()
+    def _init_shadow(self, fsdp_module):
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        # Anchor shadow to the local CUDA device. summon_full_params can return
+        # CPU views of params under some FSDP/resume code paths, so don't
+        # implicitly inherit p's device.
+        target_device = torch.device("cuda", torch.cuda.current_device())
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                self.shadow[n] = p.detach().to(target_device).float().clone()
+
+    @torch.no_grad()
+    def update(self, fsdp_module):
+        d = self.decay
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                shadow = self.shadow[n]
+                # Always align the source tensor to shadow's device to avoid
+                # cuda/cpu mismatches when FSDP exposes params on CPU.
+                shadow.mul_(d).add_(
+                    p.detach().to(shadow.device).float(), alpha=1. - d
+                )
+
+    # Optional helpers ---------------------------------------------------
+    def state_dict(self):
+        return self.shadow            # picklable
+
+    def load_state_dict(self, sd):
+        self.shadow = {k: v.clone() for k, v in sd.items()}
+
+    def copy_to(self, fsdp_module):
+        # load EMA weights into an (unwrapped) copy of the generator
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=True):
+            for n, p in fsdp_module.module.named_parameters():
+                if n in self.shadow:
+                    p.data.copy_(self.shadow[n].to(p.dtype, device=p.device))

DecMem/utils/lmdb.py

@@ -0,0 +1,72 @@
+import numpy as np
+
+
+def get_array_shape_from_lmdb(env, array_name):
+    with env.begin() as txn:
+        image_shape = txn.get(f"{array_name}_shape".encode()).decode()
+        image_shape = tuple(map(int, image_shape.split()))
+    return image_shape
+
+
+def store_arrays_to_lmdb(env, arrays_dict, start_index=0):
+    """
+    Store rows of multiple numpy arrays in a single LMDB.
+    Each row is stored separately with a naming convention.
+    """
+    with env.begin(write=True) as txn:
+        for array_name, array in arrays_dict.items():
+            for i, row in enumerate(array):
+                # Convert row to bytes
+                if isinstance(row, str):
+                    row_bytes = row.encode()
+                else:
+                    row_bytes = row.tobytes()
+
+                data_key = f'{array_name}_{start_index + i}_data'.encode()
+
+                txn.put(data_key, row_bytes)
+
+
+def process_data_dict(data_dict, seen_prompts):
+    output_dict = {}
+
+    all_videos = []
+    all_prompts = []
+    for prompt, video in data_dict.items():
+        if prompt in seen_prompts:
+            continue
+        else:
+            seen_prompts.add(prompt)
+
+        video = video.half().numpy()
+        all_videos.append(video)
+        all_prompts.append(prompt)
+
+    if len(all_videos) == 0:
+        return {"latents": np.array([]), "prompts": np.array([])}
+
+    all_videos = np.concatenate(all_videos, axis=0)
+
+    output_dict['latents'] = all_videos
+    output_dict['prompts'] = np.array(all_prompts)
+
+    return output_dict
+
+
+def retrieve_row_from_lmdb(lmdb_env, array_name, dtype, row_index, shape=None):
+    """
+    Retrieve a specific row from a specific array in the LMDB.
+    """
+    data_key = f'{array_name}_{row_index}_data'.encode()
+
+    with lmdb_env.begin() as txn:
+        row_bytes = txn.get(data_key)
+
+    if dtype == str:
+        array = row_bytes.decode()
+    else:
+        array = np.frombuffer(row_bytes, dtype=dtype)
+
+    if shape is not None and len(shape) > 0:
+        array = array.reshape(shape)
+    return array

DecMem/utils/loss.py

@@ -0,0 +1,81 @@
+from abc import ABC, abstractmethod
+import torch
+
+
+class DenoisingLoss(ABC):
+    @abstractmethod
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Base class for denoising loss.
+        Input:
+            - x: the clean data with shape [B, F, C, H, W]
+            - x_pred: the predicted clean data with shape [B, F, C, H, W]
+            - noise: the noise with shape [B, F, C, H, W]
+            - noise_pred: the predicted noise with shape [B, F, C, H, W]
+            - alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T]
+            - timestep: the current timestep with shape [B, F]
+        """
+        pass
+
+
+class X0PredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return torch.mean((x - x_pred) ** 2)
+
+
+class VPredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1))
+        return torch.mean(weights * (x - x_pred) ** 2)
+
+
+class NoisePredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return torch.mean((noise - noise_pred) ** 2)
+
+
+class FlowPredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        **kwargs
+    ) -> torch.Tensor:
+        return torch.mean((kwargs["flow_pred"] - (noise - x)) ** 2)
+
+
+NAME_TO_CLASS = {
+    "x0": X0PredLoss,
+    "v": VPredLoss,
+    "noise": NoisePredLoss,
+    "flow": FlowPredLoss
+}
+
+
+def get_denoising_loss(loss_type: str) -> DenoisingLoss:
+    return NAME_TO_CLASS[loss_type]

DecMem/utils/misc.py

@@ -0,0 +1,39 @@
+import numpy as np
+import random
+import torch
+
+
+def set_seed(seed: int, deterministic: bool = False):
+    """
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+
+    Args:
+        seed (`int`):
+            The seed to set.
+        deterministic (`bool`, *optional*, defaults to `False`):
+            Whether to use deterministic algorithms where available. Can slow down training.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.use_deterministic_algorithms(True)
+
+
+def merge_dict_list(dict_list):
+    if len(dict_list) == 1:
+        return dict_list[0]
+
+    merged_dict = {}
+    for k, v in dict_list[0].items():
+        if isinstance(v, torch.Tensor):
+            if v.ndim == 0:
+                merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0)
+            else:
+                merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0)
+        else:
+            # for non-tensor values, we just copy the value from the first item
+            merged_dict[k] = v
+    return merged_dict

DecMem/utils/pose_utils.py

@@ -0,0 +1,203 @@
+"""
+Pose and action utility functions adapted from WorldMem/framear.
+"""
+import numpy as np
+import torch
+from einops import rearrange
+
+
+ACTION_KEYS = [
+    "inventory",    # 0
+    "ESC",          # 1
+    "hotbar.1",     # 2
+    "hotbar.2",     # 3
+    "hotbar.3",     # 4
+    "hotbar.4",     # 5
+    "hotbar.5",     # 6
+    "hotbar.6",     # 7
+    "hotbar.7",     # 8
+    "hotbar.8",     # 9
+    "hotbar.9",     # 10
+    "forward",      # 11
+    "back",         # 12
+    "left",         # 13
+    "right",        # 14
+    "cameraY",      # 15
+    "cameraX",      # 16
+    "jump",         # 17
+    "sneak",        # 18
+    "sprint",       # 19
+    "swapHands",    # 20
+    "attack",       # 21
+    "use",          # 22
+    "pickItem",     # 23
+    "drop",         # 24
+]
+
+
+def convert_action_space(actions: np.ndarray) -> torch.Tensor:
+    """Convert raw WorldMem action array to 25-dim one-hot vector.
+
+    Args:
+        actions: np.ndarray of shape (T, 8+), raw action indices from NPZ.
+
+    Returns:
+        Tensor of shape (T, 25) with float values.
+    """
+    vec_25 = torch.zeros(len(actions), len(ACTION_KEYS))
+    vec_25[actions[:, 0] == 1, 11] = 1   # forward
+    vec_25[actions[:, 0] == 2, 12] = 1   # backward
+    vec_25[actions[:, 4] == 11, 16] = -1  # camera X left
+    vec_25[actions[:, 4] == 13, 16] = 1   # camera X right
+    vec_25[actions[:, 3] == 11, 15] = -1  # camera Y up
+    vec_25[actions[:, 3] == 13, 15] = 1   # camera Y down
+    vec_25[actions[:, 5] == 6, 24] = 1   # place box
+    vec_25[actions[:, 5] == 1, 24] = 1   # place box
+    vec_25[actions[:, 1] == 1, 13] = 1   # left
+    vec_25[actions[:, 1] == 2, 14] = 1   # right
+    vec_25[actions[:, 7] == 1, 2] = 1    # hotbar 1
+    return vec_25
+
+
+def euler_to_rotation_matrix(pitch: torch.Tensor, yaw: torch.Tensor) -> torch.Tensor:
+    """Convert pitch and yaw angles (radians) to 3x3 rotation matrix.
+
+    Args:
+        pitch: Tensor of pitch angles, shape (...,).
+        yaw:   Tensor of yaw angles, shape (...,).
+
+    Returns:
+        Rotation matrix of shape (..., 3, 3).
+    """
+    cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
+    cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)
+
+    R_pitch = torch.stack([
+        torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
+        torch.zeros_like(pitch), cos_pitch, -sin_pitch,
+        torch.zeros_like(pitch), sin_pitch, cos_pitch
+    ], dim=-1).reshape(-1, 3, 3)
+
+    R_yaw = torch.stack([
+        cos_yaw, torch.zeros_like(yaw), sin_yaw,
+        torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
+        -sin_yaw, torch.zeros_like(yaw), cos_yaw
+    ], dim=-1).reshape(-1, 3, 3)
+
+    return torch.matmul(R_yaw, R_pitch)
+
+
+def euler_to_camera_to_world_matrix(pose: torch.Tensor) -> torch.Tensor:
+    """Convert (x, y, z, pitch, yaw) pose vector to 4x4 camera-to-world matrix.
+
+    Args:
+        pose: Tensor of shape (5,), (f, 5), or (f, b, 5).
+
+    Returns:
+        Tensor of shape (4,4), (f, 4, 4), or (f, b, 4, 4).
+    """
+    origin_dim = pose.ndim
+    if origin_dim == 1:
+        pose = pose.unsqueeze(0).unsqueeze(0)
+    elif origin_dim == 2:
+        pose = pose.unsqueeze(0)
+
+    x, y, z, pitch, yaw = (
+        pose[..., 0], pose[..., 1], pose[..., 2],
+        pose[..., 3], pose[..., 4]
+    )
+    pitch = torch.deg2rad(pitch)
+    yaw = torch.deg2rad(yaw)
+
+    R = euler_to_rotation_matrix(pitch, yaw)
+
+    eye = torch.eye(4, dtype=pose.dtype, device=pose.device)
+    c2w = eye.repeat(R.shape[0], 1, 1)
+    c2w[:, :3, :3] = R
+    c2w[:, :3, 3] = torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=-1)
+
+    if origin_dim == 3:
+        return c2w.view(pose.shape[0], pose.shape[1], 4, 4)
+    elif origin_dim == 2:
+        return c2w.view(pose.shape[0], 4, 4)
+    else:
+        return c2w.squeeze(0).squeeze(0)
+
+
+def c2w_to_w2c(c2w: torch.Tensor) -> torch.Tensor:
+    """Invert camera-to-world to world-to-camera matrix.
+
+    Args:
+        c2w: Tensor of shape (..., 4, 4).
+
+    Returns:
+        w2c: Tensor of shape (..., 4, 4).
+    """
+    R = c2w[..., :3, :3]
+    t = c2w[..., :3, 3]
+    Rt = R.transpose(-1, -2)
+    t_new = -Rt @ t.unsqueeze(-1)
+    w2c = torch.zeros_like(c2w)
+    w2c[..., :3, :3] = Rt
+    w2c[..., :3, 3] = t_new.squeeze(-1)
+    w2c[..., 3, 3] = 1.0
+    return w2c
+
+
+def extract_pose_action_from_npz(
+    data_path: str,
+    start_frame: int,
+    num_frames: int,
+    dtype=torch.float32
+) -> tuple:
+    """Load action and pose from WorldMem NPZ file.
+
+    Args:
+        data_path:   Path to .npz file containing 'actions' and 'poses'.
+        start_frame: Starting frame index.
+        num_frames:  Number of frames to load.
+        dtype:       Output tensor dtype.
+
+    Returns:
+        actions: Tensor (num_frames, 25) - encoded action vectors.
+        poses:   Tensor (num_frames, 5)  - raw pose vectors (x,y,z,pitch,yaw).
+    """
+    data = np.load(data_path)
+    actions_raw = data["actions"]          # (N, 8+)
+    poses_pool = data["poses"].copy()      # (N, 5)
+
+    # Fix potentially corrupted height in first frame
+    poses_pool[0, 1] = poses_pool[1, 1]
+
+    # Pad poses if shorter than actions
+    if len(poses_pool) < len(actions_raw):
+        poses_pool = np.pad(poses_pool, ((1, 0), (0, 0)))
+
+    actions_raw = actions_raw[start_frame: start_frame + num_frames]
+    poses = poses_pool[start_frame: start_frame + num_frames]
+
+    # Normalize pose: translate relative to first frame, negate yaw
+    poses[:, :3] -= poses[:1, :3]
+    poses[:, -1] = -poses[:, -1]
+    poses[:, 3:] %= 360
+
+    actions = convert_action_space(actions_raw)  # (T, 25)
+    return actions.to(dtype), torch.tensor(poses, dtype=dtype)
+
+
+def poses_to_viewmats(poses: torch.Tensor) -> torch.Tensor:
+    """Convert (B, T, 5) pose vectors to (B, T, 4, 4) world-to-camera matrices.
+
+    Args:
+        poses: Tensor of shape (B, T, 5).
+
+    Returns:
+        viewmats: Tensor of shape (B, T, 4, 4).
+    """
+    B, T, _ = poses.shape
+    # (B, T, 5) -> (T, B, 5) for euler_to_camera_to_world_matrix
+    pose_tbd = rearrange(poses, "b t d -> t b d")
+    c2w = euler_to_camera_to_world_matrix(pose_tbd)  # (T, B, 4, 4)
+    w2c = c2w_to_w2c(c2w)                            # (T, B, 4, 4)
+    viewmats = rearrange(w2c, "t b p q -> b t p q")  # (B, T, 4, 4)
+    return viewmats