Duplicate from nvidia/PiD

Co-authored-by: Xuanchi Ren <xrenaa@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+figures/teaser.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,142 @@
+---
+library_name: pytorch
+tags:
+- super-resolution
+- diffusion
+- pixel-diffusion-decoder
+- vae-decoder
+pipeline_tag: image-to-image
+base_model:
+- nvidia/PixelDiT-1300M-1024px
+- Tongyi-MAI/Z-Image
+- black-forest-labs/FLUX.1-dev
+- black-forest-labs/FLUX.2-dev
+- nyu-visionx/Scale-RAE-Qwen7B_DiT9.8B
+---
+
+# PiD — Pixel Diffusion Decoder
+
+<p align="center">
+  <img src="figures/teaser.jpg" alt="PiD teaser" width="100%">
+</p>
+
+
+**[Paper](https://arxiv.org/abs/2605.23902), [Project Page](https://research.nvidia.com/labs/sil/projects/pid/)**
+
+[Yifan Lu](https://yifanlu0227.github.io/),
+[Qi Wu](https://wilsoncernwq.github.io/),
+[Jay Zhangjie Wu](https://zhangjiewu.github.io/),
+[Zian Wang](https://www.cs.toronto.edu/~zianwang/),
+[Huan Ling](https://www.cs.toronto.edu/~linghuan/),
+[Sanja Fidler](https://www.cs.utoronto.ca/~fidler/),
+[Xuanchi Ren](https://xuanchiren.com/) <br>
+
+
+PiD reformulates the latent-to-pixel decoder as a conditional pixel-space
+diffusion model, unifying decoding and upsampling into a single generative
+module. It denoises directly in high-resolution pixel space and produces a
+super-resolved image in one pass. This repository hosts the released decoder
+checkpoints, plus the encoder/decoder ("VAE") weights they depend on.
+
+All `PiD_*` checkpoints in this repo are **4-step distilled**. The non-`PiD_*`
+entries (`ae.safetensors`, `flux2_ae.safetensors`, `sdxl_vae.safetensors`, `QwenImage_VAE_2d.pth`, `sd3_vae/`, `rae/`,
+`scale_rae/`) are **the corresponding encoder/decoder VAE weights** that PiD
+plugs into — they're not PiD checkpoints themselves.
+
+### License/Terms of Use
+
+This model is released under the [NSCLv1](https://huggingface.co/nvidia/PixelDiT-1300M-1024px/blob/main/LICENSE) License. The work and any derivative works may only be used for non-commercial (research or evaluation) purposes.
+
+### Deployment Geography:
+Global
+
+## PiD checkpoints
+
+Two variants are released for each diffusers-style backbone:
+
+- **`2k`** — trained at 2048px, used as a 4× decoder (512 LDM → 2048 px), or as
+  an 8× decoder for the Scale-RAE backbone (256 → 2048).
+- **`2kto4k`** — trained with multi-resolution data bucketing 2048→4096 and an
+  SD3-style dynamic shift; designed for 1024 LDM → 4K (4096 px) decoding.
+
+Both checkpoint variants support multiple aspect ratios.
+
+| Path                                                            | Latent space | SR factor | Variant |
+|-----------------------------------------------------------------|--------------|-----------|---------|
+| `checkpoints/PiD_res2k_sr4x_official_flux_distill_4step`        | Flux1-dev    | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr4x_official_flux2_distill_4step`       | Flux2-dev    | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step`         | SD3 medium   | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr4x_official_dinov2_distill_4step`      | DINOv2-B     | 4×        | 2k      |
+| `checkpoints/PiD_res2k_sr8x_official_siglip_distill_4step`      | SigLIP-2     | 8×        | 2k      |
+| `checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step`    | Flux1-dev    | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step_2606` | Flux2-dev  | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_sd3_distill_4step`     | SD3 medium   | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_sdxl_distill_4step`    | SDXL         | 4×        | 2kto4k  |
+| `checkpoints/PiD_res2kto4k_sr4x_official_qwenimage_distill_4step` | Qwen-Image  | 4×        | 2kto4k  |
+
+Each directory contains a single file, `model_ema_bf16.pth`, which is the EMA
+weights cast to bfloat16 — the format the inference scripts load by default.
+
+> **⚠️ Flux2-dev `2kto4k` — use the new `_2606` checkpoint.** The previous
+> `PiD_res2kto4k_sr4x_official_flux2_distill_4step` (without the `_2606` suffix)
+> suffered from a color-drifting issue. The new
+> `PiD_res2kto4k_sr4x_official_flux2_distill_4step_2606` fixes it — please use it
+> and do **not** use the old one. See the
+> [comparison](https://github.com/nv-tlabs/pid/blob/main/docs/FLUX2_2kto4k_new_ckpt_compare.md)
+> for details.
+
+### Latent space → compatible LDMs
+
+A PiD decoder is tied to a *latent space*, not to a single generative model. Any
+LDM that produces latents in that space can reuse the same checkpoint. The
+`--backbone` aliases below pick the right LDM pipeline; they all decode through
+the latent space's checkpoint above.
+
+| Latent space | VAE / vision encoder weights       | compatible `--backbone`                   | Corresponding LDM Links |
+|--------------|------------------------------------|-------------------------------------------|-----------------|
+| Flux1-dev    | `checkpoints/ae.safetensors`       | `flux`, `zimage`, `zimage-turbo`          | [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev), [Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image), [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) |
+| Flux2-dev    | `checkpoints/flux2_ae.safetensors` | `flux2`, `flux2-klein-4b`, `flux2-klein-9b` | [FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev), [FLUX.2-klein-4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B), [FLUX.2-klein-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B) |
+| SD3 medium   | `checkpoints/sd3_vae/`             | `sd3`                                     | [SD3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) |
+| SDXL         | `checkpoints/sdxl_vae.safetensors` | `sdxl`                                    | [SDXL-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) |
+| Qwen-Image   | `checkpoints/QwenImage_VAE_2d.pth` | `qwenimage`, `qwenimage-2512`             | [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image), [Qwen-Image-2512](https://huggingface.co/Qwen/Qwen-Image-2512) |
+| DINOv2-B     | `checkpoints/rae/`                 | `dinov2`                                  | [RAE](https://github.com/bytetriper/RAE) (class-conditional; DINOv2-B) |
+| SigLIP-2     | `checkpoints/scale_rae/`           | `siglip`                                  | [Scale-RAE](https://github.com/ZitengWangNYU/Scale-RAE) (text-conditional; nyu-visionx/Scale-RAE-Qwen1.5B_DiT2.4B) |
+
+For example, Z-Image and Z-Image-Turbo share Flux1-dev's VAE, so they reuse the
+`flux` checkpoints (both `2k` and `2kto4k`) — no separate `zimage` checkpoint is
+shipped. Likewise `qwenimage-2512` reuses the `qwenimage` decoder (same VAE,
+different transformer). 
+
+## Usage
+
+The decoder checkpoints are loaded by the inference scripts in the [PiD
+codebase](https://github.com/nv-tlabs/pid). The exact `(backbone, ckpt_type) → path` mapping is the single source
+of truth in
+[`pid/_src/inference/checkpoint_registry.py`](https://github.com/nv-tlabs/PiD/blob/main/pid/_src/inference/checkpoint_registry.py) — clone the
+repo, point it at this snapshot, and the demos pick the right file
+automatically:
+
+```bash
+# Pull just the checkpoints/ tree into the repo root (skips this README and
+# the teaser figure so they don't clobber the files in the source repo).
+hf download nvidia/PiD --local-dir . --include "checkpoints/*"
+
+# Then run any of the demos, e.g.:
+PYTHONPATH=. python -m pid._src.inference.from_ldm --backbone flux \
+    --prompt "A photorealistic half-body portrait of a brown tabby cat with bold stripes sitting attentively on a rustic wooden kitchen table, soft morning light streaming sideways through a large window, fine fur detail and stripe patterns sharply visible, intense amber-green eyes in razor-sharp focus, warm farmhouse kitchen softly out of focus, cinematic shallow depth of field, ultra-detailed fur texture, photorealistic" \
+    --ldm_inference_steps 28 --save_xt_steps 24 \
+    --output_dir ./results/official_demo/flux \
+    --pid_inference_steps 4
+```
+
+Pick the `2kto4k` variant via `--pid_ckpt_type 2kto4k` when decoding at 4K.
+
+## Citation
+```
+@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}
+}
+```

checkpoints/PiD_res2k_sr4x_official_dinov2_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4ce25b6fd2c953720468cd88a7c6a192b4ce908bf085659a1324c186949eab0
+size 2731773393

checkpoints/PiD_res2k_sr4x_official_flux2_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0e3a7a19d4738e0b53c1267815a87e77295fddbee117b5aed802b9b62030cac
+size 2725875153

checkpoints/PiD_res2k_sr4x_official_flux_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47dd165bf3ea85df08deb152e6fcded19ddf5a35b83832abcd99d403ffca6ac3
+size 2724842961

checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:835a069763903bfee317e524b83d0974639a6c1d9f79e00728d338fcc249fa27
+size 2724842961

checkpoints/PiD_res2k_sr8x_official_siglip_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9b735ca95044c9d4b5777f1a398fb30efcf4c28021622f13964ba785c33495c
+size 2735312130

checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca75b0b2712f0872d8225fb0c2d520f817c7185eb4a5dbedfe2d164239df044a
+size 2725875153

checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step_2606/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da81ef12a56e18b1e978afc0413fb9a0daeafa67da78f9c7c47517e2b57feea1
+size 2725875153

checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f3eabf4f2f83320472e6146f6545d8237e1423849e62148d5d656bfb571d00e
+size 2724842961

checkpoints/PiD_res2kto4k_sr4x_official_qwenimage_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d95ca9c5f6c7c055954cac836de2a6b4659610894fbc5bb55b807f4dbb4b8c2
+size 2724842961

checkpoints/PiD_res2kto4k_sr4x_official_sd3_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22d1224cae36041b61517c2145c954fca2a9624a30a01d57b7bb044304b9dc31
+size 2724842961

checkpoints/PiD_res2kto4k_sr4x_official_sdxl_distill_4step/model_ema_bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26933047e73753842d73f3d45e1fdb61df6448f3ed82b9ee8c26037431cf4db0
+size 2724732369

checkpoints/QwenImage_VAE_2d.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1658b24b9248dc2782c9ce1ce6ca226f32aceef5e52788101720d55dd9a70bc5
+size 498332666

checkpoints/ae.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:afc8e28272cd15db3919bacdb6918ce9c1ed22e96cb12c4d5ed0fba823529e38
+size 335304388

checkpoints/flux2_ae.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:868fe7b343cc8f3a19dbcfcafbc3d5f888802be3f89bd81b65b3621a066ce8f3
+size 336211292

checkpoints/rae/decoders/dinov2/wReg_base/ViTXL_n08_i512/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8b83b058cd84567ce1f671c9ea32d6ec0e48532c650bbf1ba04bb4739292630
+size 1665128758

checkpoints/rae/stats/dinov2/wReg_base/imagenet1k_512/stat.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe0ac4e914a28708ec85f0dbd3be436dcec48e556d5d745f78a13625dc8fe2c7
+size 6292742

checkpoints/scale_rae/decoder/XL_decoder_config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "facebook/vit-mae-base",
+  "architectures": [
+    "ViTMAEForPreTraining"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "decoder_hidden_size": 1152,
+  "decoder_intermediate_size": 4096,
+  "decoder_num_attention_heads": 16,
+  "decoder_num_hidden_layers": 28,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1152,
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "mask_ratio": 0.75,
+  "model_type": "vit_mae",
+  "norm_pix_loss": false,
+  "num_attention_heads": 12,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "patch_size": 14,
+  "qkv_bias": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.3"
+}

checkpoints/scale_rae/decoder/siglip2_sop14_i224_web73M_ganw3_decXL.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca7e6b907bb51455a12eea39b6acb1999c2133f325c123cda20ceb206d1ef3cb
+size 1662529538

checkpoints/sd3_vae/vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9b67a279283625caee39d61eacb5324243848477b4eb535355eaaa8423d4e09
+size 167666654

checkpoints/sdxl_vae.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63aeecb90ff7bc1c115395962d3e803571385b61938377bc7089b36e81e92e2e
+size 334641164

config.json

@@ -0,0 +1,6 @@
+{
+    "input_types": [
+        "PiD"
+    ],
+    "model_size": "1.3B"
+}