+
+ Zhenhao Yang1
+ Xiaoshi Wu2
+ Zhengyao Lv1
+ Xiaoyu Shi2,†
+ Xintao Wang2
+ Pengfei Wan2
+ Kun Gai2
+ Kwan-Yee K. Wong1,†
+
+
+ 1The University of Hong Kong
+ 2Kling Team, Kuaishou Technology
+ †Corresponding author
+
+
+
+[](https://www.youtube.com/watch?v=f6jBmFXxgzg)
+
+
+
+## ⚙️ 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},
+}
+```
\ No newline at end of file
diff --git a/DecMem/configs/decmem.yaml b/DecMem/configs/decmem.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..6f9b9ac210343653e3c352603fad1cd5c113c019
--- /dev/null
+++ b/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
diff --git a/DecMem/configs/default_config.yaml b/DecMem/configs/default_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..3ca7ae3f11d8250f2ad8134767acc7e5ba37fded
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/configs/dense.yaml b/DecMem/configs/dense.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..7d58eb178b62dfabfae5fa3138714eb2f3f96cdd
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/configs/dense_short.yaml b/DecMem/configs/dense_short.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..061f59ddad41c288676a0c9b80daa60e71a31d02
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/configs/worldmem_diffusion.yaml b/DecMem/configs/worldmem_diffusion.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..114a8f5548e4798610e314a8caf070b6840280b9
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/demo_utils/constant.py b/DecMem/demo_utils/constant.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ef237829ca620a5b0b4b00ef7a795fb36a0ca94
--- /dev/null
+++ b/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
diff --git a/DecMem/demo_utils/memory.py b/DecMem/demo_utils/memory.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9758df6dffaed70887f2ec3bf14d9fa49181c32
--- /dev/null
+++ b/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
diff --git a/DecMem/demo_utils/taehv.py b/DecMem/demo_utils/taehv.py
new file mode 100644
index 0000000000000000000000000000000000000000..8531563e7d8da9cf5b1f93f46fda8215f50ed769
--- /dev/null
+++ b/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()
diff --git a/DecMem/demo_utils/utils.py b/DecMem/demo_utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a24aefb4d8074acdce8fe3d32f6e37c07c4d6baa
--- /dev/null
+++ b/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)
diff --git a/DecMem/demo_utils/vae.py b/DecMem/demo_utils/vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c80925194127bc6f96759b1893a08c3d440e0bc
--- /dev/null
+++ b/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
diff --git a/DecMem/demo_utils/vae_block3.py b/DecMem/demo_utils/vae_block3.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5b80a025f7b4b81356b1be93cc1260e10a841fc
--- /dev/null
+++ b/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
diff --git a/DecMem/demo_utils/vae_torch2trt.py b/DecMem/demo_utils/vae_torch2trt.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0e08f258a199cbfb6eaf0b3a0782f29f2d8faf3
--- /dev/null
+++ b/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)
diff --git a/DecMem/evaluation/method_manifest.json b/DecMem/evaluation/method_manifest.json
new file mode 100644
index 0000000000000000000000000000000000000000..8868e1e3fbfd09754d2d38e3607b1c37ba062fde
--- /dev/null
+++ b/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"
+}
diff --git a/DecMem/evaluation/worldmem_adapter.py b/DecMem/evaluation/worldmem_adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..7177c016fe0e3baed7eaf55370a25a47caf4e701
--- /dev/null
+++ b/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()
diff --git a/DecMem/example/example.csv b/DecMem/example/example.csv
new file mode 100644
index 0000000000000000000000000000000000000000..45a3160800ef8bc83e37b4303c9bf60ca7a4b7e1
--- /dev/null
+++ b/DecMem/example/example.csv
@@ -0,0 +1,3 @@
+video_path,action_path
+example/000045.mp4,example/000045.npz
+example/000046.mp4,example/000046.npz
\ No newline at end of file
diff --git a/DecMem/infer.py b/DecMem/infer.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f3eab7742b04e466c92cfb96d2989d1af422f5a
--- /dev/null
+++ b/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.")
\ No newline at end of file
diff --git a/DecMem/model/__init__.py b/DecMem/model/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f858161d1f9118fb2ddb1b33205fe5f09bfcec7
--- /dev/null
+++ b/DecMem/model/__init__.py
@@ -0,0 +1,9 @@
+from .diffusion import CausalDiffusion
+from .causvid import CausVid
+from .gan import GAN
+
+__all__ = [
+ "CausalDiffusion",
+ "CausVid",
+ "GAN",
+]
diff --git a/DecMem/model/base.py b/DecMem/model/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..2714042bea30c10a50640b447ffafb2962ea4656
--- /dev/null
+++ b/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
+ )
diff --git a/DecMem/model/causvid.py b/DecMem/model/causvid.py
new file mode 100644
index 0000000000000000000000000000000000000000..3abdac190f3fb596cf5e8391687d185f47661e48
--- /dev/null
+++ b/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
diff --git a/DecMem/model/diffusion.py b/DecMem/model/diffusion.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1a00fba235e9b2f6e22bcce2fdfe57f0a7f15ff
--- /dev/null
+++ b/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
diff --git a/DecMem/model/gan.py b/DecMem/model/gan.py
new file mode 100644
index 0000000000000000000000000000000000000000..92367d1835728133a66e20999786f96d4b640397
--- /dev/null
+++ b/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
diff --git a/DecMem/pipeline/__init__.py b/DecMem/pipeline/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ad80477f22b133c1bac1d08bf1ec6fa17f36ff62
--- /dev/null
+++ b/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",
+]
diff --git a/DecMem/pipeline/bidirectional_diffusion_inference.py b/DecMem/pipeline/bidirectional_diffusion_inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..31cce4d1553d88261442936171df1ac8a0cf4f2c
--- /dev/null
+++ b/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
diff --git a/DecMem/pipeline/bidirectional_inference.py b/DecMem/pipeline/bidirectional_inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..b523ec0d0ac0082daf310b159aed6f286e0b0efb
--- /dev/null
+++ b/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
diff --git a/DecMem/pipeline/causal_diffusion_inference.py b/DecMem/pipeline/causal_diffusion_inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..15d56ae2d8beabf75dd1605da6eb40d02e3b88b6
--- /dev/null
+++ b/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
diff --git a/DecMem/pipeline/causal_inference.py b/DecMem/pipeline/causal_inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ffbf6214d78c2377204d4f6871eca650a659e51
--- /dev/null
+++ b/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
diff --git a/DecMem/pipeline/self_forcing_training.py b/DecMem/pipeline/self_forcing_training.py
new file mode 100644
index 0000000000000000000000000000000000000000..ab989fe620b2441f2b01c5bd656628e34687db85
--- /dev/null
+++ b/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
diff --git a/DecMem/requirements.txt b/DecMem/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..ac932081f345511ad4fd53b1d417922ebcab3f2f
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/scripts/build_worldmem_csv.py b/DecMem/scripts/build_worldmem_csv.py
new file mode 100644
index 0000000000000000000000000000000000000000..9bb2eab8835c450583f6f1b6042399f61fb11443
--- /dev/null
+++ b/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()
diff --git a/DecMem/scripts/infer.sh b/DecMem/scripts/infer.sh
new file mode 100644
index 0000000000000000000000000000000000000000..aa7c2537fbe0155ecd5b5ddb3e76ec4590dbb11f
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/scripts/infer_example.sh b/DecMem/scripts/infer_example.sh
new file mode 100644
index 0000000000000000000000000000000000000000..491724664f9234a76e57b5ff5fa5b8f202d6f8ec
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/DecMem/scripts/train_multinode.sh b/DecMem/scripts/train_multinode.sh
new file mode 100644
index 0000000000000000000000000000000000000000..4b6d035fc15ff602ed5fc3d347e2a4d1df2ae0a9
--- /dev/null
+++ b/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 
+
+ 
+
+ 
+
+ 
+