Add stage3_new entry to README

README.md

@@ -19,8 +19,15 @@ Training checkpoints for [ForgeWM](https://github.com/asdfo123/ForgeWM).
 | `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
 
@@ -28,12 +35,13 @@ Training checkpoints for [ForgeWM](https://github.com/asdfo123/ForgeWM).
 # Download
 huggingface-cli download asdfo123/ForgeWM --local-dir ./ckpts --repo-type model
 
-# Inference
-python inference.py --checkpoint_path ckpts/stage3/model.pt --image_path demo_images/forest.png --action_type forward
+# 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
 
@@ -47,4 +55,4 @@ python inference.py --checkpoint_path ckpts/stage3/model.pt --image_path demo_im
 ```
 
 ## Contact
-[Xinye Li](https://asdfo123.github.io) leeasdfo123@gmail.com
+[Xinye Li](https://asdfo123.github.io) leeasdfo123@gmail.com