PiD

PiD — Pixel Diffusion Decoder

TL;DR — PiD is a plug-and-play diffusion decoder that replaces VAE/RAE decoders, turning latent representations directly into super-resolved pixels in a single pass.

https://github.com/user-attachments/assets/a556e2d4-5de5-4bcf-9daa-80f7ea6b2124

PiD reformulates the latent-to-pixel decoder as a conditional pixel-space diffusion model, unifying decoding and upsampling into a single generative module. It directly denoises in high-resolution pixel space and produces a super-resolved image in one pass.

Paper, Project Page, Model Weights

Yifan Lu, Qi Wu, Jay Zhangjie Wu, Zian Wang, Huan Ling, Sanja Fidler, Xuanchi Ren

News

• 🚀 [July 14, 2026] PiD adds optional Boogu-Image text-to-image support, including Boogu native generation and PiD decode from Boogu’s Flux-style VAE latents.
• 🚀 [July 9, 2026] PiD Training code released, with PixelDiT and PiD v1.5 (2kto4k) distilled and undistilled checkpoints!
• 🚀 [July 9, 2026] PiD v1.5 checkpoints for FLUX (Z-Image, Z-Image-Turbo), FLUX.2, and Qwen-Image are released. Check release page to see improvements!
• 🔥 [June 2, 2026] PiD checkpoints for SDXL, Qwen-Image and Qwen-Image-2512 are released. Check HuggingFace.
• 🔥 [June 2, 2026] We clean up the codebase and remove useless code. Torch.compile mode is also available now.
• 🚀 [May 27, 2026] PiD is now in ComfyUI!
• 🚀 [May 25, 2026] Paper, code, and model weights released, with PiD options for FLUX, FLUX.2, Z-Image, Z-Image-Turbo, SD3, DINOv2, and SigLIP.

Table of Contents

• Installation
• Download Checkpoints
• Running inference with released checkpoints
  • LDM → PiD decode
  • image → PiD decode
  • Optional Boogu-Image
• Training
• Repository layout

Installation

pip install hydra-core omegaconf pyyaml \
    attrs einops loguru termcolor fvcore iopath wandb \
    imageio opencv-python-headless pandas \
    safetensors sentencepiece boto3 botocore

If you want to create a new inference environment, uv is fast and easy to use. It create a project-local .venv with the locked base dependencies:

# install uv: https://docs.astral.sh/uv/getting-started/installation/
# You can simply run `pip install uv`
uv python install 3.12
uv sync --frozen
source .venv/bin/activate
PYTHONPATH=. python verify_env.py

Download Checkpoints

Checkpoints are hosted at nvidia/PiD on the HuggingFace. Pull the checkpoints/ folder into this repo:

hf download nvidia/PiD --local-dir . --include "checkpoints/*"

Running inference with released checkpoints

PiD ships two main entry points, each selecting a backbone with --backbone:

• from_ldm.py — text/class → latent diffusion → PiD decode
• from_clean.py — image → VAE encode → PiD decode

[!IMPORTANT] Picking the checkpoint variant — --pid_ckpt_type Every entry point accepts --pid_ckpt_type {2k,2kto4k,2kto4k_v1pt5} (default 2k):

• 2k — the original 2048px-trained decoder, trained with 2K resolution only. Multiple aspect ratios are supported, typically 2048 × 2048 (1:1), 2304 × 1728 (4:3), 1728 × 2304 (3:4), 2688 × 1536 (16:9), and 1536 × 2688 (9:16).
• 2kto4k — the v1 up-to-4K-resolution decoder, trained with varying resolution (range from 2K to 4K). Multiple aspect ratios are supported. Less sharp than 2k at 2048px resolution.
• 2kto4k_v1pt5 — the v1.5 up-to-4K-resolution decoder for FLUX, FLUX.2, Qwen-Image (WAN2.1) VAE (range from 2K to 4K). Better color accuracy, no grid artifacts in the corners, trained with more anime data and small-face data. Better than 2kto4k overall but less sharp than 2k at 2048px resolution.

For the exact checkpoint path for each backbone, see docs/checkpoints.md and checkpoint registry.

For the exact checkpoint path behind each (backbone, --pid_ckpt_type), see docs/checkpoints.md.

📕 from_ldm: text / class → latent diffusion → PiD decode

Runs the chosen --backbone on a prompt, captures the intermediate x_t at user-specified denoising steps (early LDM termination) and the final clean x_0, then decodes each captured latent with both the native VAE / RAE decoder (baseline) and PiD.

Example 1 — Single-GPU, single prompt (Flux, default 2k decoder)

Generating a 2048px image with Flux + PiD decode. Decoding latent from 24 and 28 (full) LDM steps.

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

Example 2 — Single-GPU, 4K decode with 4:3 aspect ratio (Flux, 2kto4k_v1pt5 decoder)

Same backbone as Example 1 but with --resolution 4096,3072 --pid_ckpt_type 2kto4k_v1pt5. --resolution is the final output size, so the LDM runs at 1024,768 and PiD decodes it to 4K.

PYTHONPATH=. python -m pid._src.inference.from_ldm --backbone flux \
    --prompt "A close photograph of a cat looking through frosted glass beside a small pine branch, winter light, soft condensation, simple cozy composition, expressive eyes." \
    --resolution 4096,3072 --pid_ckpt_type 2kto4k_v1pt5 \
    --ldm_inference_steps 28 --save_xt_steps 24 26 \
    --output_dir ./results/official_demo/flux_4k_ar4_3

Example 3 — Multi-GPU with a prompt file (Z-Image) with torch.compile

torchrun shards --prompt_file across ranks; each rank writes to --output_dir independently. We use --compile to enable torch.compile for faster inference, the first call will be slow due to the compilation. We use default compilation mode, to get further speedup, change to the max-autotune mode in _maybe_compile_net (pid/_src/models/pixeldit_model.py). Note that extra cudatoolkit like nvcc is required for the compilation.

PYTHONPATH=. torchrun --nproc_per_node=4 \
    -m pid._src.inference.from_ldm --backbone zimage \
    --prompt_file pid/_src/inference/prompts/prompt_creative.txt \
    --ldm_inference_steps 50 --save_xt_steps 46 \
    --compile \
    --output_dir ./results/official_demo/zimage

Example 4 — Multi-GPU, 4K decode (Z-Image-Turbo, 2kto4k_v1pt5 decoder)

Z-Image-Turbo defaults to 9 diffusers steps with guidance_scale=0.0. The final clean latent x0 is always saved and is the recommended Turbo output to inspect. --save_xt_steps 7 is optional; it saves an additional near-final x_t sample for comparison. --resolution 4096 means H=4096, W=4096 and the LDM runs at 1024,1024.

PYTHONPATH=. torchrun --nproc_per_node=4 \
    -m pid._src.inference.from_ldm --backbone zimage-turbo \
    --prompt_file pid/_src/inference/prompts/prompt_zimage_turbo.txt \
    --resolution 4096 --pid_ckpt_type 2kto4k_v1pt5 \
    --output_dir ./results/official_demo/zimage_turbo_4k

dinov2 / siglip backbones

The upstream RAE / Scale-RAE LDMs don’t live in diffusers — see docs/dinov2_siglip.md for setup and end-to-end examples.

Suggested step settings per diffusers backbone

(See each script’s docstring for the exact recipe.)

📗 from_clean: image → VAE encode → PiD decode

No latent diffusion model is run. The input image is fed at its native resolution (only center-cropped so each side is a multiple of 16), encoded by VAE, optionally corrupted with Gaussian noise at each sigma in --degrade_sigmas, then decoded by PiD at --scale * vae_native_resolution.

Single-GPU example (Flux):

PYTHONPATH=. python -m pid._src.inference.from_clean --backbone flux \
    --manifest assets/clean_image_manifest.jsonl \
    --degrade_sigmas 0.0 \
    --output_dir ./results/official_demo_from_clean/flux \
    --cfg_scale 1 --pid_inference_steps 4 --scale 4

You can pass a single image with --input_path and a prompt with --prompt instead of --manifest, and a sigma sweep such as --degrade_sigmas 0.0 0.2 0.4 0.8 to decode noise-corrupted latents. Swap --backbone to use a different VAE (flux2 / sd3 / sdxl / qwenimage); sdxl automatically uses its variance-preserving noising form.

The dinov2 / siglip from_clean flows take the same flags but with a different --scale (8 for siglip); their encoders resize internally to their fixed native interface (512 / 256) regardless of the input image size — see docs/dinov2_siglip.md.

from_boogu: Boogu-Image text → image

Boogu-Image uses its own custom pipeline and a Flux1-style VAE, so it is kept as an optional integration rather than a base dependency. After following docs/boogu_image.md, you can run Boogu native T2I or add --pid_decode to capture Boogu’s final normalized Flux1 latent and run the Flux1 PiD pixel decoder:

PYTHONPATH=.:third_party/Boogu-Image python -m pid._src.inference.from_boogu \
    --variant turbo \
    --prompt "A tiny ceramic robot watering a bonsai tree" \
    --resolution 512 \
    --enable_model_cpu_offload \
    --pid_decode \
    --output_dir ./results/boogu/turbo_pid

See docs/boogu_image.md for optional dependency setup, batch/torchrun examples, and environment notes.

Training

Create the training environment

We still use conda since it is easy to install CUDA toolkit.

conda env create -f environment.yml
conda activate pid
python -m pip install -e . --group full \
    --extra-index-url https://download.pytorch.org/whl/cu128

Check Docker to create a training docker image.

Training Tutorial

See the training guide for dataset preparation and commands to finetune PixelDiT, train a PiD teacher, and distill a PiD student.

To evaluate training experiment checkpoints, see inference with training checkpoints for PixelDiT, PiD teacher, and PiD student examples.

License

PiD codebase is licensed under the Apache License 2.0.

Contributing

See CONTRIBUTING.md for development setup, code style, and the DCO sign-off requirement.

Acknowledgments

The authors would like to acknowledge Yongsheng Yu and Wei Xiong for open-sourcing PixelDiT’s model and weights, and thank Product Managers Aditya Mahajan and Matt Cragun for their valuable support and guidance.

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}
}