PiD Model Loader

PiD Model Loader is the boring half of an actually exciting pair. It does one job: find an NVIDIA PiD checkpoint on disk and hand you a PID_MODEL you can feed into PiD Decode. The Decode node does all the magic; this node is the gatekeeper that makes sure you loaded the right weights, in the right precision, for the right model family.

If you haven't met PiD yet: it's NVIDIA's diffusion-based decoder, the thing that quietly became the answer to "VAE decode then upscale" - instead of one pass through the autoencoder you get a distilled 4-step diffusion decode that renders a 2048px (or 3840px) image straight from the latent, generating detail along the way. It's big news in the Flux world, and it's heavy, so the loading side matters.

What actually loads

The loader reads from a folder it registers on startup: ComfyUI/models/pid/. Drop the official files there and the node's pid_checkpoint dropdown fills itself in. You need three things:

ComfyUI/models/pid/checkpoints/ae.safetensors
ComfyUI/models/pid/checkpoints/PiD_res2k_sr4x_official_flux_distill_4step/model_ema_bf16.pth
ComfyUI/models/pid/checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step/model_ema_bf16.pth

All from the nvidia/PiD Hugging Face repo. The ae.safetensors is the Flux autoencoder that PiD's own internals need, and it's a hard requirement - the loader will refuse to build the model without it, even though your workflow already has a VAE plugged in elsewhere.

Mechanically, the node is doing more than a load_state_dict. It vendors NVIDIA's whole pid package (under third_party/PiD) and slaps it onto sys.path, instantiates the official config, and then - the bit people find surprising - it doesn't actually load the weights yet. PiD Model Loader returns a handle, and the model materializes lazily on the first decode, then stays cached keyed by checkpoint + backbone + precision + device. So your first run pays a real one-time cost (including a possible download of text-encoder assets through transformers), and re-runs are fast.

The inputs that matter

• pid_checkpoint - the dropdown of .pth files it found under models/pid. Empty folder? You get the placeholder string <put model_ema_bf16.pth under ComfyUI/models/pid> and an explicit error telling you where to put it.
• backbone_type - flux1 (default) or zimage. This must match the latent you're decoding; PiD Decode validates it, but get it right here.
• precision - bf16 / fp16 / fp32, default bf16. Leave it. bf16 is what these checkpoints are distilled for.
• device - auto or cuda. Both end up on CUDA; this pack is CUDA-only, no CPU, no Apple Silicon.

Output is the pid_model (PID_MODEL) handle that wires straight into PiD Decode's pid_model input.

Install

ComfyUI Manager: search PiD and install "ComfyUI-PiD". Or the manual route:

cd ComfyUI/custom_nodes
git clone https://github.com/npiriou/ComfyUI-PiD
pip install -r requirements.txt

Then drop the three model files above in place and restart ComfyUI. Note the dependency list is heavier than most nodes - pinned hydra-core/omegaconf, plus transformers>=4.57, diffusers>=0.37, even boto3. On a clean ComfyUI Python env that's fine, but if you're sharing a venv with other tools, watch for version clashes.

Where people get burned

• Checkpoint family mismatches. The 2k checkpoint expects a latent that maps to 512px input; 2kto4k expects 1024px. The loader auto-detects the variant from the filename, and Decode will error loudly if your latent doesn't match - that's by design, not a bug.
• It rejects non-Flux1/Z-Image weights on purpose. NVIDIA also ships PiD checkpoints for Flux2, SD3, DINOv2 and SIGLIP backbones; this loader scans the filename and refuses those. Don't fight it - you're on the wrong pack for that.
• First load hits the network and can look like a hang. It's downloading text-encoder assets. Give it a minute.

If you're hitting a wrong-backbone error while the checkpoint path is fine, you've picked zimage for a flux checkpoint (or vice versa) - the loader checks that too.