Duplicate from 9r4n4y/PiD_upscaler_image_restoration_tutorial_video

Co-authored-by: official <9r4n4y@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.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
+tutorial.webm filter=lfs diff=lfs merge=lfs -text
+tutorial.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,136 @@
+---
+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
+
+# tutorial video 
+
+<video controls autoplay muted loop width="100%">
+  <source src="https://huggingface.co/9r4n4y/PiD_upscaler_image_restoration_tutorial_video/resolve/main/tutorial.mp4" type="video/webm">
+  Your browser does not support the video tag.
+</video>
+
+---
+
+<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`, `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→3840 and an
+  SD3-style dynamic shift; designed for 1024 LDM → 4K (4096 px) decoding.
+
+Both checkpoint variants support multiple aspect ratios.
+
+| Path                                                          | Backbone (encoder side)                    | SR factor | Variant   |
+|---------------------------------------------------------------|--------------------------------------------|-----------|-----------|
+| `checkpoints/PiD_res2k_sr4x_official_flux_distill_4step`      | Flux1-dev (16-ch VAE)                      | 4×        | 2k        |
+| `checkpoints/PiD_res2k_sr4x_official_flux2_distill_4step`     | Flux2-dev (128-ch BN VAE)                  | 4×        | 2k        |
+| `checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step`       | SD3 medium (16-ch VAE)                     | 4×        | 2k        |
+| `checkpoints/PiD_res2k_sr4x_official_dinov2_distill_4step`    | DINOv2-B + RAE ViT-XL (768-ch)             | 4×        | 2k        |
+| `checkpoints/PiD_res2k_sr8x_official_siglip_distill_4step`    | SigLIP-2 So400M + Scale-RAE ViT-XL (1152)  | 8×        | 2k        |
+| `checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step`  | Flux1-dev (16-ch VAE)                      | 4×        | 2kto4k    |
+| `checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step` | Flux2-dev (128-ch BN VAE)                  | 4×        | 2kto4k    |
+| `checkpoints/PiD_res2kto4k_sr4x_official_sd3_distill_4step`   | SD3 medium (16-ch VAE)                     | 4×        | 2kto4k    |
+
+Z-Image shares Flux1's VAE, so its inference path reuses the `flux` checkpoints
+(both `2k` and `2kto4k`) — no separate `zimage` checkpoint is shipped.
+
+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.
+
+## VAE / encoder weights
+
+These are the per-backbone encoder (and, where applicable, original decoder)
+weights that PiD pairs with. They're hosted here so a single download brings
+everything needed end-to-end.
+
+| Path                            | Description                                                                          |
+|---------------------------------|--------------------------------------------------------------------------------------|
+| `checkpoints/ae.safetensors`    | Flux1-dev / Z-Image 16-ch VAE (encoder + original Flux decoder).                     |
+| `checkpoints/flux2_ae.safetensors` | Flux2-dev 128-ch BN VAE.                                                          |
+| `checkpoints/sd3_vae/`          | SD3 medium 16-ch VAE in diffusers format.                                            |
+| `checkpoints/rae/`              | DINOv2-B image encoder + RAE ViT-XL decoder + ImageNet-512 normalization statistics. |
+| `checkpoints/scale_rae/`        | SigLIP-2 So400M encoder + Scale-RAE ViT-XL decoder + decoder config.                 |
+
+## 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_flux \
+    --prompt "A photorealistic cat" \
+    --ldm_inference_steps 28 --save_xt_steps 22 24 26 \
+    --output_dir ./results/demo \
+    --cfg_scale 1 --pid_inference_steps 4 --scale 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_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_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/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

config.json

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

tutorial.mp4

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