ForgeWM Checkpoints

Training checkpoints for ForgeWM.

Available Checkpoints

File	Stage	Description
stage0/model.pt	Stage 0	Bidirectional SFT (domain adaptation, 4000 steps)
stage1/model.pt	Stage 1	Teacher-Forcing Causal AR (10000 steps)
stage2/model.pt	Stage 2	Consistency Distillation (6000 steps)
stage3/model.pt	Stage 3	DMD final model (4-step real-time inference)
stage3_new/model.pt	Stage 3 (improved)	DMD with cache-refresh fix in training pipeline + sliding-window (las=6, sink=0) inference alignment (4000 steps). Better long-video stability and reduced HUD shrinkage / OOD drift on 22s rollouts.

Notes on stage3_new

stage3_new was retrained from the same Stage 2 starting point as the original stage3, but with two pipeline-level fixes:

1. Train-time KV cache refresh — the per-chunk t=0 cache rerun is now applied unconditionally. Previously it was skipped for grad-enabled chunks, creating a train/inference distribution shift on the KV cache (training cache was noisy at the exit_flags step, while inference cache reaches near-clean t≈0). See https://github.com/asdfo123/ForgeWM/blob/main/pipeline/self_forcing_training.py for the fix.
2. Inference-time sliding-window alignment — the inference pipeline now correctly forwards the top-level local_attn_size / sink_size from the training config into WanDiffusionWrapper. Previously these defaulted to -1 / 0 and inference silently ran with full causal attention while the model was trained with a 6-frame window.

These changes are most visible on long-horizon (22s+) rollouts.

Usage

# Download
huggingface-cli download asdfo123/ForgeWM --local-dir ./ckpts --repo-type model

# Inference (default 4-step recipe)
python inference.py --checkpoint_path ckpts/stage3_new/model.pt --image_path demo_images/forest.png --action_type forward

🚧 Coming Soon

• More intermediate checkpoints
• Few-step variants (2-step / 1-step with first-chunk FFE schedule)

Citation

@misc{forgewm2026,
  title={ForgeWM: A Reproducible Training Recipe for Action-Controllable World Models},
  author=ForgeWM Team,
  year={2026},
  url={https://github.com/asdfo123/ForgeWM}
}

Contact

Xinye Li leeasdfo123@gmail.com