Junhao Zhuang, Shiyi Zhang, Yuxuan Bian, Yaowei Li, Yawen Luo, Yijun Liu, Weiyang Jin, Songchun Zhang, Xianglong He, Xuying Zhang, Haoran Li, Haoyang Huang, Zeyue Xue, Nan Duan
Joy Future Academy, JD
⭐ If Self Gradient Forcing is useful for your research, please consider starring this repository.
- 2026-07-23: Paper, model checkpoints, inference scripts, and training code are publicly released.
Self Gradient Forcing (SGF) recovers the missing context-gradient path for self-generated causal memory through a bounded two-pass replay, enabling models trained with only a 5-second window to extrapolate to minute-scale videos with stronger identity, layout, and temporal stability.
The environment follows the Causal-Forcing setup.
conda create -n self_gradient_forcing python=3.10 -y
conda activate self_gradient_forcing
pip install -r requirements.txt
pip install flash-attn --no-build-isolation
python setup.py developbash scripts/download_weights.shThe script uses the Hugging Face CLI command hf by default. Set HF_CLI=huggingface-cli if your environment still uses the older command name.
It downloads:
- Wan base models to
wan_models/Wan2.1-T2V-1.3Bandwan_models/Wan2.1-T2V-14B. - All Causal-Forcing initialization checkpoints under
checkpoints/init/framewise/andcheckpoints/init/chunkwise/:ar_diffusion.pt,causal_cd.pt, andcausal_ode.pt. - Released SGF inference checkpoints to
checkpoints/framewise/ar/model.ptandcheckpoints/chunkwise/ar/model.pt. - The training prompt list to
prompts/vidprom_filtered_extended.txt.
The default prompt file is prompts/test_prompt.txt with 8 prompts. The launcher uses 8 GPUs when at least 8 GPUs are visible; otherwise it falls back to single-GPU serial inference. By default it generates 963 latent frames, which decode to about 240 seconds of video at 16 fps.
The inference script takes the release setting name (framewise or chunkwise) and selects the matching config and checkpoint automatically:
- framewise config:
configs/self_gradient_forcing_framewise.yaml - chunkwise config:
configs/self_gradient_forcing_chunkwise.yaml
The long-video KV-cache geometry is set in scripts/infer_self_gradient_forcing.sh. Framewise defaults to KV_CACHE_SINK=4, KV_CACHE_FIFO_FRAMES=16, and KV_CACHE_CURRENT_FRAMES=1, so the actual --kv_cache_max_frames passed to inference.py is 4 + 16 + 1 = 21. Chunkwise defaults to KV_CACHE_SINK=3, KV_CACHE_FIFO_FRAMES=6, and KV_CACHE_CURRENT_FRAMES=3, so --kv_cache_max_frames is 12.
bash scripts/infer_self_gradient_forcing.sh framewiseThis uses:
configs/self_gradient_forcing_framewise.yaml
checkpoints/framewise/ar/model.pt
bash scripts/infer_self_gradient_forcing.sh chunkwiseThis uses:
configs/self_gradient_forcing_chunkwise.yaml
checkpoints/chunkwise/ar/model.pt
bash scripts/infer_self_gradient_forcing.sh \
framewise \
checkpoints/framewise/ar/model.pt \
prompts/test_prompt.txtUseful overrides:
NUM_OUTPUT_FRAMES=963 SEED=42 OUTPUT_ROOT=outputs/demo \
bash scripts/infer_self_gradient_forcing.sh framewiseFor trained checkpoints, pass the release setting first and the produced logs/.../checkpoint_model_*/model.pt path as the second argument. The script uses EMA weights by default; set USE_EMA=0 if you explicitly want the non-EMA generator weights.
bash scripts/train_self_gradient_forcing_framewise.shEquivalent explicit form:
bash scripts/train_self_gradient_forcing_framewise.sh \
configs/self_gradient_forcing_framewise.yaml \
logs/sgf_framewisebash scripts/train_self_gradient_forcing_chunkwise.shEquivalent explicit form:
bash scripts/train_self_gradient_forcing_chunkwise.sh \
configs/self_gradient_forcing_chunkwise.yaml \
logs/sgf_chunkwiseThe launchers accept [config.yaml] [logdir] [extra train.py args...], matching the multi-node launcher convention used by the reference training scripts. They support single-node and multi-node training. For multi-node jobs, run the same command on every node within the gather window. The scripts auto-register nodes through .rendezvous/ on the shared filesystem and launch static torchrun with an IP master address.
Useful overrides:
GATHER_WINDOW=90 NUM_GPUS=8 MASTER_PORT=29501 ENABLE_WANDB=1 \
bash scripts/train_self_gradient_forcing_framewise.sh logs/sgf_framewise
NNODES=2 NODE_RANK=0 MASTER_ADDR=10.0.0.1 NUM_GPUS=8 \
bash scripts/train_self_gradient_forcing_chunkwise.sh logs/sgf_chunkwiseThis implementation builds on the Wan video model ecosystem and follows the installation conventions of thu-ml/Causal-Forcing. We thank the open-source community for the infrastructure that made this release possible.
For questions, please contact Junhao Zhuang at zhuangjh23@mails.tsinghua.edu.cn.
This project is released under the Apache-2.0 license.
@article{zhuang2026sgf,
title = {Self Gradient Forcing: Native Long Video Extrapolation},
author = {Zhuang, Junhao and Zhang, Shiyi and Bian, Yuxuan and Li, Yaowei
and Luo, Yawen and Jin, Weiyang and Zhang, Songchun and He, Xianglong
and Zhang, Xuying and Li, Haoran and Huang, Haoyang and Xue, Zeyue
and Duan, Nan},
journal = {arXiv preprint},
year = {2026}
}
