yilwang/PiD-QDM-Halo2-4K

PiD-QDM Halo2

Research checkpoints that give NVIDIA's PiD decoder (PiD_v1pt5_res2kto4k_sr4x_official_flux_undistilled) linear-cost attention for 4K decoding. In all 14 MMDiT blocks, global softmax attention is replaced by two branches: exact softmax inside 8×8 query windows that read a 12×12 neighbourhood plus all text tokens, and a quadrant delta memory (QDM) with one closed-form state per head and image quadrant. The PiT pixel blocks are unchanged. The original attention stays inside each block and served as the distillation teacher.

This is not an official NVIDIA release.

Files

Folder	Stage	Updates	Effective batch	Parameters
a1/	Per-layer alignment to the original attention	1,000 of 1,000	512	41.2M
a2/	Output distillation	2,000 of 2,000	512	41.2M
b/	Output distillation with rank-64 LoRA on the original attention	3,000 of 6,000	512	63.2M

Each folder holds model.safetensors (fp32; only the new and changed parameters) and config.json (architecture, training settings and the SHA-256 of the PiD checkpoint the weights apply to). Use b/ for decoding. Earlier stages are intermediate and can start later stages through --init-from.

Usage

The weights apply on top of the original PiD checkpoint and need the Linear-PiD code:

python download/download_models.py --models pid vae gemma --weights-root weights
hf download yilwang/PiD-QDM-Halo2-4K --include "b/*" --local-dir checkpoints/pid-qdm

Pass checkpoints/pid-qdm/b as --checkpoint to scripts/vdn_pid/sample_pid_vdn.sh; the PiD-QDM section of the code README has the full command. The PiD checkpoint's SHA-256 is checked when the weights load. Gemma-2-2b-it is gated: accept its license on Hugging Face first.

Training

Each stage distills from the model's own original attention on MultiAspect-4K-1M images at native size, with captions encoded by Gemma-2-2b-it. The stages run in order, A1 → A2 → B, each starting from the previous one.

License

The weights in this repository are derivative works of nvidia/PiD and are distributed under the NVIDIA License (NSCLv1), included in full in LICENSE. They may be used only non-commercially, which the license defines as research or evaluation purposes. The PiD checkpoint, the FLUX VAE and Gemma-2-2b-it are not included; obtain them under their own terms. The Linear-PiD code is licensed separately (Apache-2.0).

Citation

Please cite PiD:

@article{lu2026pid,
    title={PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion},
    author={Lu, Yifan and Wu, Qi and Wu, Jay Zhangjie and Wang, Zian and Ling, Huan and Fidler, Sanja and Ren, Xuanchi},
    journal={arXiv preprint arXiv:2605.23902},
    year={2026}
}