JunhaoZhuang/Self_Gradient_Forcing

⚡ Self Gradient Forcing

Native Long Video Extrapolation

       

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.

🔥 News

• 2026-07-23: Paper, model checkpoints, inference scripts, and training code are publicly released.

🧠 Method Overview

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.

🛠️ Installation

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 develop

⬇️ Download Weights

bash scripts/download_weights.sh

The 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.3B and wan_models/Wan2.1-T2V-14B.
• All Causal-Forcing initialization checkpoints under checkpoints/init/framewise/ and checkpoints/init/chunkwise/: ar_diffusion.pt, causal_cd.pt, and causal_ode.pt.
• Released SGF inference checkpoints to checkpoints/framewise/ar/model.pt and checkpoints/chunkwise/ar/model.pt.
• The training prompt list to prompts/vidprom_filtered_extended.txt.

🚀 Inference

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.

Framewise

bash scripts/infer_self_gradient_forcing.sh framewise

This uses:

configs/self_gradient_forcing_framewise.yaml
checkpoints/framewise/ar/model.pt

Chunkwise

bash scripts/infer_self_gradient_forcing.sh chunkwise

This uses:

configs/self_gradient_forcing_chunkwise.yaml
checkpoints/chunkwise/ar/model.pt

Custom checkpoint or prompt file

bash scripts/infer_self_gradient_forcing.sh \
  framewise \
  checkpoints/framewise/ar/model.pt \
  prompts/test_prompt.txt

Useful overrides:

NUM_OUTPUT_FRAMES=963 SEED=42 OUTPUT_ROOT=outputs/demo \
  bash scripts/infer_self_gradient_forcing.sh framewise

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

🏋️ Training

Framewise SGF

bash scripts/train_self_gradient_forcing_framewise.sh

Equivalent explicit form:

bash scripts/train_self_gradient_forcing_framewise.sh \
  configs/self_gradient_forcing_framewise.yaml \
  logs/sgf_framewise

Chunkwise SGF

bash scripts/train_self_gradient_forcing_chunkwise.sh

Equivalent explicit form:

bash scripts/train_self_gradient_forcing_chunkwise.sh \
  configs/self_gradient_forcing_chunkwise.yaml \
  logs/sgf_chunkwise

The 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_chunkwise

🙏 Acknowledgements

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

📮 Contact

For questions, please contact Junhao Zhuang at zhuangjh23@mails.tsinghua.edu.cn.

📜 License

This project is released under the Apache-2.0 license.

📜 Citation

@misc{zhuang2026selfgradientforcingnative,
      title={Self Gradient Forcing: Native Long Video Extrapolation}, 
      author={Junhao Zhuang and Shiyi Zhang and Yuxuan Bian and Yaowei Li and Yawen Luo and Yijun Liu and Weiyang Jin and Songchun Zhang and Xianglong He and Xuying Zhang and Haoran Li and Haoyang Huang and Zeyue Xue and Nan Duan},
      year={2026},
      eprint={2607.20368},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2607.20368}, 
}