Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image by Pfannkuchensack · Pull Request #9281 · invoke-ai/InvokeAI

Pfannkuchensack · 2026-06-08T21:06:18Z

Summary

Adds PiD (Pixel Diffusion Decoder) support to InvokeAI — NVIDIA's few-step pixel-diffusion decoder that replaces the regular VAE decode with a caption-conditioned, 4× super-resolution decode (512→2048 in a single 4-step distill pass).

This PR vendors a minimal, inference-only subset of PiD at invokeai/backend/pid/ (upstream https://github.com/nv-tlabs/PiD, code Apache-2.0) and wires it end-to-end into the model manager, invocation nodes, starter models, and the generation UI.

What you get

A generic PiD backend (backend/pid/decode.py): per-backbone net config (_PER_BACKBONE), build_pid_net / load_pid_decoder / PiDDecoder, Gemma-2 caption encoding, and a working-memory estimator for the cache.
Generic loader nodes: pid_decoder_loader (→ PiDDecoderField) and gemma2_encoder_loader (→ Gemma2EncoderField), plus model-manager configs/loaders for PiD checkpoints and the shared Gemma-2 caption encoder.

PiD decode nodes for six base models, each replacing that base's VAE decode:

Base	Node	Notes
FLUX.1	`flux_pid_decode`	16ch / down-8
FLUX.2 Klein (4B/9B)	`flux2_pid_decode`	packs the stored 32ch latent → 128ch / down-16; BN-normalized (no scalar denorm)
SD3	`sd3_pid_decode`	16ch / down-8; fixed VAE constants
SDXL	`sdxl_pid_decode`	4ch / down-8; VAE `scaling_factor` read at runtime
Z-Image (+ Turbo)	`z_image_pid_decode`	reuses the FLUX decoder (shared 16ch VAE)
Qwen-Image	`qwen_image_pid_decode`	16ch / down-8; per-channel `latents_mean`/`latents_std` denorm + 5D→4D temporal squeeze

Generation UI: a new PiD mode selector (Off / Fit / Native) with PiD-decoder + Gemma-2-encoder pickers and a PiD-steps control, shown for any base that supports PiD. Decoder pickers are base-filtered (Z-Image shows FLUX decoders). Both txt2img and img2img (Canvas) are supported in both modes:
- Fit: generate at the target size, PiD decodes 4×, downscale back (compositing-safe).
- Native: the requested size is the 4× target — generate at target/4 and use PiD's full 4× output directly.
Starter models: NVIDIA PiD decoders (nvidia/PiD, per backbone; FLUX/FLUX.2/SD3 ship 2K + 2K-to-4K, SDXL/Qwen-Image ship 2K-to-4K only) plus the shared gemma-2-2b-it caption encoder.

Robustness details

Backbone identification is driven primarily by the checkpoint's latent channel count (4/16/128), with the filename/directory name as a tie-breaker. Because FLUX.1 / SD3 / Qwen-Image all share 16 channels, the config probe additionally trusts an explicit base override (which the starter installer sends) when the directory name is ambiguous — so single-file HF downloads are still identified correctly.
Readiness checks gate each supported base (decoder + Gemma-2 encoder present, "Scale Before Processing" off, SDXL refiner disabled with PiD).
The standard (non-PiD) FLUX/FLUX.2/SD3/SDXL/Z-Image/Qwen-Image paths are unchanged.

License note: the vendored PiD code is Apache-2.0. The pretrained PiD weights are released by NVIDIA under NSCLv1 (non-commercial / research) — relevant for anyone redistributing the checkpoints.

Related Issues / Discussions

Closes #9240

QA Instructions

Install models (Model Manager → Starter Models): a PiD Decoder for your base (e.g. "PiD Decoder FLUX (2K)") — its dependency, "Gemma 2 2B (PiD caption encoder)", installs automatically. Confirm the decoder is registered with the correct base (e.g. a Qwen-Image decoder shows as qwen-image, not FLUX).
Select a supported main model (FLUX, FLUX.2 Klein, SD3, SDXL, Z-Image/Turbo, or Qwen-Image). In the Generation settings expander, a PiD control appears; pick the PiD decoder + Gemma-2 encoder.
txt2img – Fit: set PiD = Fit, generate. Output is the requested size, refined via the 4× decode.
txt2img – Native: set PiD = Native, generate. The requested dimensions become the 4× target (image is generated at target/4).
Canvas img2img – Fit / Native: with "Scale Before Processing" = None, run a raster-layer generation in both modes.
Guards: verify inpaint/outpaint and (SDXL) an active refiner surface a clear "unsupported" toast / disabled Invoke with a reason.

Automated gates already green on the branch: backend imports (starter_models, config factory, every *_pid_decode node), and the frontend pnpm lint:tsc / lint:eslint / lint:knip / lint:dpdm.

Not yet hardware-verified (needs a GPU + downloaded checkpoints): full end-to-end image output per base, measured VRAM peak (the working-memory constant is calibrated to a 2048px output ≈ 4.3 GB), and Qwen-Image Edit-mode (reference image) + PiD.

Merge Plan

Self-contained feature; no redux migration beyond the already-included params slice bump.

Checklist

The PR has a short but descriptive title, suitable for a changelog
Tests added / updated (if applicable)
❗Changes to a redux slice have a corresponding migration (params slice _version bump)
Documentation added / updated (if applicable)
Updated What's New copy (if doing a release after this PR)

Adds a vendored subset of NVIDIA's PiD (Pixel Diffusion Decoder) at invokeai/backend/pid/ as the foundation for upcoming FLUX / FLUX.2 / SD3 / Z-Image PiD decode nodes plus a future PiD-based 4x upscale node. Upstream: https://github.com/nv-tlabs/PiD (Apache 2.0). Vendor scope: * _src/{networks,models,modules}: PidNet, PixDiT_T2I, LQProjection2D, PidModel, PidDistillModel, PixelDiTModel, GeneralConditioner. * _ext/imaginaire: minimal Imaginaire framework subset (lazy_config, model, utils/{log,misc,distributed,device,count_params}). * configs/, tokenizers/, checkpointer/, trainer.py, visualize/, _demo_*, from_*, easy_io/, S3/wandb training helpers were intentionally excluded. Dependency stripping (no new hard deps introduced): * loguru, termcolor -> stdlib logging shim * iopath PathManager -> stdlib pathlib stub * fvcore Registry -> minimal stdlib Registry * lazy_config/lazy.py: yaml/dill/cloudpickle/detectron2 save/load paths replaced with a minimal LazyCall stub * lazy_config/instantiate.py: omegaconf DictConfig/ListConfig branches removed; configs are plain dict / LazyCall mappings * megatron, pynvml, boto3/wandb imports are try/except-guarded or local to functions and stay inert in our inference path All pid.* imports rewritten to invokeai.backend.pid.*; SPDX-Apache-2.0 headers retained on vendored files; attribution and detailed list of local modifications added in LICENSE-PiD.txt. The pre-trained PiD checkpoints distributed by NVIDIA remain under NSCLv1 (non-commercial); this commit only vendors code. Smoke test: PidNet, PidModel, PidDistillModel, GeneralConditioner import cleanly; LazyCall -> instantiate round-trip resolves to the expected nn.Module. ruff check passes.

Adds the model-manager plumbing and workflow nodes needed to use the vendored PiD decoder (phase A) end-to-end with FLUX, SD3 and Z-Image. Model manager (Phase B + B.5): * taxonomy: ModelType.PiDDecoder, PiDDecoderVariantType (Res2k_Sr4x / Res2kTo4k_Sr4x), ModelType.Gemma2Encoder + ModelFormat.Gemma2Encoder, both added to AnyVariant + variant_type_adapter. * configs/pid_decoder.py: per-backbone PiD configs (FLUX / FLUX.2 / SD3) with state-dict probing on 'lq_proj' substring and backbone/variant detection from the official NVIDIA filenames. * configs/gemma2_encoder.py: Gemma-2 directory probing on Gemma2ForCausalLM architecture + tokenizer files. * AnyModelConfig union updated. * model_loaders/pid_decoder.py: loads .pth / .safetensors, strips the upstream 'net.' prefix, supports torch.load(weights_only=True). * model_loaders/gemma2_encoder.py: SubModelType.{Tokenizer, TextEncoder} dispatch; returns the causal LM's inner Gemma2Model (transformers 4.56's get_decoder() returns None for Gemma2). Decode pipeline (Phase C): * backend/pid/decode.py: build_pid_net + load_pid_decoder (per-backbone PixDiT_T2I hyperparams derived from PiD's pid_sr4x base + per-experiment overrides), encode_caption_for_pid (chi-prompt + Gemma encoding, mirrors PixelDiTModel._encode_text_raw), and a PiDDecoder wrapper with a reimplemented few-step distill sampler (no autocast / no distributed / no PixelDiTModel init paths from upstream). Invocations (Phase 6.x): * Gemma2EncoderField + PiDDecoderField in invocations/model.py. * gemma2_encoder_loader / pid_decoder_loader: thin ModelIdentifierField pickers that emit the corresponding fields. * z_image_pid_decode (pilot), flux_pid_decode, sd3_pid_decode: caption encode -> Gemma offload -> PiD state dict load -> PidNet construct -> decode. Per-backbone latent denormalisation (FLUX1 ae_params, SD3 hardcoded 1.5305/0.0609, Z-Image piggybacks on FLUX VAE). End-to-end validated with the released PiD_res2k_sr4x_official_flux_distill_4step.pth checkpoint and gemma-2-2b-it: PidNet rebuilds at exactly 456 keys / 1.36B params, sampler runs at ~5 GB VRAM peak (Gemma dominates), output shape and range match. FLUX.2 PiD decode is deliberately deferred: it needs BN-based latent denormalisation and 32->128 channel packing, and we have no FLUX.2 checkpoint to validate against yet.

Adds the NVIDIA PiD decoder as a 4x super-resolution alternative to the regular VAE/RAE decode path. Includes model-manager configs and loaders for both the PiD checkpoints and the Gemma-2 caption encoder they require, plus four invocations: latent-in decode for FLUX / SD3 / Z-Image and an image-in pid_upscale node. - Decode pipeline keeps PidNet params in fp32 and uses bf16 autocast only for matmuls; caption embeddings have outliers that overflow bf16 RMSNorm. - encode_caption_for_pid forces tokenizer padding_side="right" (Gemma defaults to left, PiD trained with right) and returns the attention mask as bool so it stays compatible with SDPA. - Z-Image reuses the FLUX-trained checkpoint and reads scale/shift from the VAE config at runtime (PiD upstream notes they are checkpoint-specific). - TextLLM config now excludes Gemma2ForCausalLM so it falls through to the dedicated Gemma2 encoder config instead of being misclassified. - Frontend: new model_type / model_format / variant enums, type guards and category metadata; schema.ts regenerated via pnpm typegen.

Read latent channel count from lq_proj.latent_proj.0.weight (FLUX.2=128, FLUX.1/SD3=16) as the primary discriminator; fall back to filename/dir name only to disambiguate the architecturally identical FLUX.1/SD3 pair. Fixes FLUX.2 checkpoints (model_ema_bf16.pth) not being recognised, and correctly rejects unsupported backbones (RAE/dinov2, 768ch). Fix Flux2 docstring 32->128.

Add a "PiD Decode" mode select (Off / Fit / Native) to the FLUX advanced settings with PiD decoder + Gemma-2 encoder pickers. In Fit mode the FLUX graph swaps the VAE decode for a PiD 4x super-resolution decode and downscales back to the requested size. Adds params state (pidMode, decoder, encoder, steps) with a v3->v4 migration, model hooks, readiness checks, and graph guards for the not-yet-wired Native and non-txt2img paths.

Make the generation dimension helpers PiD-aware via an optional pidScale: in Native mode the user-facing dimensions are the 4x target (grid 64, optimal 2048), generation runs at target/4, and PiD's 4x output is used directly with no downscale. Thread pidScale through the params dimension reducers and the optimal-dimension/grid-size selectors, resync dimensions when toggling Native, and wire the Native path in the FLUX graph builder. Add working_mem_bytes for PiD Decode

Extract the PiD decode chain into buildPidDecodeChain (loaders + decode + fit-downscale, no denoise setup) so it can substitute for the VAE decode across generation modes. Widen addImageToImage's l2i param to ImageOutputNodes (it only consumes .image) and wire the PiD chain into the img2img branch in Fit mode. Native stays txt2img-only (a 4x result can't composite onto the bbox); inpaint/outpaint remain gated off for now.

Add addPidImageToImageNative: the canvas bbox is the 4x target, so the init image is downscaled to bbox/4, denoised at that resolution, and PiD decodes straight back up to the full bbox with no post-decode downscale - preserving all PiD detail while still compositing cleanly onto the region. Wire it into the img2img branch of buildFLUXGraph (native vs fit vs off) and drop the native-txt2img-only guard. Make the canvas FLUX grid check PiD-aware so a native bbox must be a multiple of 64 (16 * 4) for bbox/4 to land on the grid.

Explain PiD usage on hover, mirroring the DyPE popover: what the decoder is (NVIDIA Pixel Diffusion Decoder, 4x SR, needs a PiD decoder + Gemma-2 encoder), Fit vs Native modes, the 2K / 2K-to-4K target resolutions, that Steps can be lowered, and that Scale Before Processing must be off. Links to nv-tlabs/PiD.

Register NVIDIA's PiD FLUX decoders (2K and 2K-to-4K presets, from nvidia/PiD) and the Efficient-Large-Model/gemma-2-2b-it caption encoder as starter models so they can be installed from the Model Manager. The Gemma-2 encoder is wired as a dependency of each decoder (and offered standalone).

Add a flux2_pid_decode node that packs the stored FLUX.2 latent (32ch @ H/8) into PiD's 128ch @ H/16 layout before decoding; FLUX.2's BatchNorm denormalization is already applied in flux2_denoise, so no scalar denorm is needed (optional vae input reads identity constants). Generalize the frontend PiD decode chain (decodeNodeType, optional vaeSource) and wire the isFlux2 graph path for txt2img/img2img (Fit & Native). Base-aware PiD gating/decoder-filter, FLUX.2 readiness checks, and two nvidia/PiD FLUX.2 starter decoders (2K, 2Kto4K). Standard FLUX PiD path unchanged.

Wire the existing sd3_pid_decode node into the SD3 graph builder (txt2img and img2img, Fit & Native) with a PiD guard, base-aware gating/decoder-filter (sd-3), and SD3 readiness checks. Add two nvidia/PiD SD3 starter decoders (2K, 2Kto4K). Harden the PiD config probe against the 16-channel FLUX.1/SD3 ambiguity: when the checkpoint's directory name is silent (the HF single-file download renames it), trust an explicit base override so SD3 checkpoints are not misidentified as FLUX.1. Also benefits Qwen. FLUX / FLUX.2 identification is unchanged.

Build the full SDXL PiD backend stack: _PER_BACKBONE[SDXL] (4ch/down8), PiDDecoder_Checkpoint_SDXL_Config with a 4-channel latent-map entry, factory union + loader registration, and a new sdxl_pid_decode node (reads the VAE's scaling_factor/shift at runtime; SDXL fallbacks 0.13025/0.0). 4-channel latents are unambiguous, so no directory-name disambiguation is needed. Generalize the shared PiD decode chain to support SD-family denoise: denoise_latents has no width/height, so thread an optional noise node for sizing and round to the model's native grid (8 for SDXL, 16 for FLUX). Wire buildSDXLGraph (txt2img + img2img, Fit & Native) with the VAE as the decode's scaling source, base-aware gating/readiness, and a starter decoder (SDXL 2Kto4K only). PiD + SDXL refiner is blocked for now via a graph guard and a readiness reason. FLUX/FLUX.2/SD3 paths are unchanged.

JPPhoto

Latest review:

Merge blockers

invokeai/backend/model_manager/configs/text_llm.py:44-53 and invokeai/backend/model_manager/configs/gemma2_encoder.py:70-78: Automatic classification now rejects Gemma 2 9B and 27B models from both candidate configurations. The PiD config rejects their non-2304 hidden size, while TextLLM_Diffusers_Config defers every Gemma2ForCausalLM regardless of size. ModelConfigFactory consequently classifies these valid causal LMs as Unknown instead of preserving their previous TextLLM behavior. Test: Classify Gemma 2 configurations with hidden sizes 2304, 3584, and 4608 without overrides; expect 2304 to become Gemma2Encoder and the larger variants to remain TextLLM.

Follow-up PR candidates

invokeai/backend/pid/decode.py:202-221 and tests/backend/pid/test_pid_decode.py:31-34: The new invalid-schedule safety net uses assert, which is removed under PYTHONOPTIMIZE=1. In that supported runtime mode, _get_t_list(num_steps=5) returns the duplicate schedule instead of raising, and the new regression test fails. The invocation and UI bounds protect normal graphs, so this is no longer the original user-facing blocker, but the claimed backend guard is ineffective. Test: Run the schedule test under python -O and require an explicit ValueError or validated configuration rather than an assertion.
invokeai/app/invocations/flux_pid_decode.py:97-110 and invokeai/backend/model_manager/load/load_base.py:92-101: All seven PiD paths determine the Gemma input device from the first parameter, despite the cache contract explicitly providing LoadedModel.compute_device for partially loaded models. If the first large parameter remains on CPU while later modules are loaded on CUDA, caption inputs are incorrectly placed on CPU instead of the intended execution device, causing avoidable transfers or a device failure depending on the patched module boundary. Test: Partially load Gemma with its first parameter on CPU and later parameters on CUDA, then verify every PiD caption path uses gemma_text_encoder_info.compute_device and completes without device mismatches.
invokeai/app/invocations/pid_upscale.py:72-75 and invokeai/app/invocations/flux_vae_encode.py:39-55: The new upscale node advertises Z-Image and other 16-channel-compatible VAEs, but delegates encoding to FluxVaeEncodeInvocation.vae_encode(), which accepts only InvokeAI's FLUX AutoEncoder. A Z-Image Diffusers AutoencoderKL, which the existing Z-Image encode path explicitly supports, cannot satisfy this implementation and fails instead of upscaling. Test: Connect both a FLUX AutoEncoder and a Z-Image Diffusers AutoencoderKL to pid_upscale; either support both advertised cases with their correct scaling rules or narrow the field description and validation to FLUX AutoEncoder only.

…uard, compute_device, narrow pid_upscale VAE Address the latest review on the PiD PR: - Merge blocker: automatic classification sent Gemma 2 9B/27B to Unknown. The PiD Gemma2 encoder config rejects their non-2304 hidden size, and TextLLM deferred *every* Gemma2ForCausalLM, so neither matched. TextLLM now defers only the size the encoder config accepts (2304 = Gemma-2-2b); larger variants stay TextLLM. - Schedule safety net used assert, which is stripped under `python -O`, leaving _get_t_list(num_steps=5) returning a duplicate schedule. Raise ValueError instead so the guard holds in optimized runtimes; the regression test now asserts ValueError and passes under `python -O`. - All seven PiD caption paths derived the Gemma device from the first parameter, which is wrong under partial loading (first param on CPU, later modules on CUDA). Use the cache contract's LoadedModel.compute_device instead. - pid_upscale advertised Z-Image / 16-channel VAEs but delegates to the FLUX-only vae_encode. Narrow the field description and validate the VAE is a FLUX AutoEncoder up front (a diffusers AutoencoderKL now fails with a clear error instead of a stripped-assert failure inside vae_encode). Update the TextLLM/Gemma2 tests (per-size config-level + a factory-level check that 2304 -> Gemma2Encoder and 3584/4608 -> TextLLM) and the schedule test (ValueError, green under python -O).

JPPhoto

One minor thing left, not a merge blocker but maybe you can squeeze it in:

invokeai/app/invocations/pid_decoder_loader.py:24, invokeai/backend/model_manager/configs/pid_decoder.py:4, and invokeai/backend/model_manager/taxonomy.py:186 still describe support as FLUX.1, FLUX.2, and SD3 only. The PR also supports SDXL and Qwen-Image, and those latter checkpoints do not both ship in two presets. The classifier errors at pid_decoder.py:169 and pid_decoder.py:191 repeat the incomplete list, producing misleading installation diagnostics. Update these strings and regenerate the schema. Test: assert the invocation title and classifier diagnostics enumerate every supported backbone and accurately describe available variants.

All previously reported merge blockers and follow-up candidates have been addressed in the current PR head!

The PiD Gemma encoder was directory + HuggingFace only, so a llama.cpp GGUF (e.g. gemma-2-2b-it-Q4_K_M.gguf) could not be used. Add GGUF support: - Gemma2Encoder_GGUF_Config: identifies a single .gguf file, reads the GGUF metadata and requires general.architecture == "gemma2" and <arch>.embedding_length == 2304 (Gemma-2-2b), rejecting 9B/27B as the directory config does. - Gemma2EncoderGGUFLoader (format gguf_quantized): loads via transformers from_pretrained(<dir>, gguf_file=<name>) — transformers dequantizes gemma2 GGUFs and reads the tokenizer from the GGUF metadata — then exposes the Gemma2Model decoder, matching the directory loader. PiD encodes the caption once and offloads the encoder, so dequantizing at load is acceptable. - Register the config in the AnyModelConfig union. No frontend change: the PiD encoder picker filters by type=gemma2_encoder, so the GGUF variant appears automatically. Verified end-to-end against a real q4_k_m file: it classifies as Gemma2Encoder_GGUF_Config and loads to a Gemma2Model producing 2304-dim hidden states. Adds config identification tests (match, 9B/27B rejected, non-gemma2 rejected, non-.gguf rejected).

…/InvokeAI into feat/pid-decoder

A Gemma-2 GGUF satisfies the generic Qwen3 GGUF key heuristic (token_embd.weight + blk.* keys), so it matched both Qwen3Encoder_GGUF_Config and the intended Gemma2Encoder_GGUF_Config. On a fresh install the Gemma config happened to win, but re-identification could pick Qwen3, mis-classifying the model. Add _has_gemma2_keys (Gemma uses blk.*.post_attention_norm / post_ffw_norm, which a Qwen3 encoder never has — Qwen3 has attn_q_norm/attn_k_norm instead) and reject such state dicts in both Qwen3 encoder configs' _validate_looks_like_qwen3_model (GGUF and checkpoint), mirroring the existing T5 / Qwen-VL exclusions. The Gemma config already rejects Qwen3 GGUFs via the general.architecture metadata, so the two are now mutually exclusive and identification is deterministic. Add regression tests: _has_gemma2_keys detection and that the Qwen3 GGUF config rejects a Gemma-keyed state dict.

…ariant The Gemma2 GGUF encoder config has no `variant` field, so re-identifying a model previously mis-detected as a Qwen3 GGUF (which carries a variant) drops it — the serialized record has no variant key and replace_model overwrites it away. Assert this explicitly in the Gemma GGUF identification test.

NVIDIA deprecated the FLUX / FLUX.2 / Qwen-Image `res2kto4k_sr4x` PiD decoders and moved them to `checkpoints_deprecated/`, replacing them with the recommended `v1pt5_res2kto4k_sr4x` checkpoints. Our starter models still pointed at the old `checkpoints/` paths, which now 404 on install. Repoint the three affected 2K-to-4K starters (FLUX, FLUX.2, Qwen-Image) to the v1.5 successors and note the upgrade in their descriptions. The 2K (`res2k_sr4x`) decoders and the SD3 / SDXL 2K-to-4K decoders are not deprecated and are unchanged. Base and variant are still sent as explicit overrides, so config identification is unaffected by the new directory name (res2kto4k -> Res2kTo4k_Sr4x).

JPPhoto

I'm hitting this error (scroll on to the bottom for my diagnosis):

[2026-07-26 20:32:10,907]::[InvokeAI]::ERROR --> Error while invoking session 84dbd9f1-c3e8-4c26-a819-29599d07150e, invocation f83e0105-6888-40da-ade0-a49ec7912328 (flux_pid_decode): Error(s) in loading state_dict for PidNet:
        size mismatch for lq_proj.latent_proj.0.weight: copying a param with shape torch.Size([1024, 16, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 16, 3, 3]).
        size mismatch for lq_proj.latent_proj.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.3.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.3.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.4.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.4.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.5.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.5.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.6.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.6.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.output_heads.0.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.1.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.2.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.3.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.4.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.5.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.6.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.gate_modules.0.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.0.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.1.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.1.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.2.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.2.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.3.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.3.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.4.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.4.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.5.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.5.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.6.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.6.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
[2026-07-26 20:32:10,907]::[InvokeAI]::ERROR --> Traceback (most recent call last):
  File "/mnt/AI/InvokeAI3/src/invokeai/app/services/session_processor/session_processor_default.py", line 143, in run_node
    output = invocation.invoke_internal(context=context, services=self._services)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/app/invocations/baseinvocation.py", line 244, in invoke_internal
    output = self.invoke(context)
             ^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/app/invocations/flux_pid_decode.py", line 131, in invoke
    pid_info = context.models.load(self.pid_decoder.decoder)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/app/services/shared/invocation_context.py", line 397, in load
    return self._services.model_manager.load.load_model(model, submodel_type, user_id=self._data.queue_item.user_id)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/app/services/model_load/model_load_default.py", line 78, in load_model
    ).load_model(model_config, submodel_type)
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/backend/model_manager/load/load_default.py", line 89, in load_model
    cache_record = self._load_and_cache(model_config, submodel_type)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/backend/model_manager/load/load_default.py", line 134, in _load_and_cache
    loaded_model = self._load_model(config, submodel_type)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/backend/model_manager/load/model_loaders/pid_decoder.py", line 88, in _load_model
    pid_net = load_pid_decoder(raw_sd, backbone)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/src/invokeai/backend/pid/decode.py", line 178, in load_pid_decoder
    missing, unexpected = net.load_state_dict(state_dict, strict=False)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/AI/InvokeAI3/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2593, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for PidNet:
        size mismatch for lq_proj.latent_proj.0.weight: copying a param with shape torch.Size([1024, 16, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 16, 3, 3]).
        size mismatch for lq_proj.latent_proj.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.3.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.3.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.3.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.4.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.4.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.4.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.5.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.5.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.5.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.0.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.0.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.2.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.6.block.2.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.3.weight: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.3.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.latent_proj.6.block.5.weight: copying a param with shape torch.Size([1024, 1024, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
        size mismatch for lq_proj.latent_proj.6.block.5.bias: copying a param with shape torch.Size([1024]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for lq_proj.output_heads.0.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.1.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.2.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.3.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.4.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.5.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.output_heads.6.weight: copying a param with shape torch.Size([1536, 1024]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for lq_proj.gate_modules.0.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.0.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.1.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.1.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.2.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.2.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.3.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.3.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.4.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.4.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.5.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.5.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for lq_proj.gate_modules.6.content_proj.weight: copying a param with shape torch.Size([1, 3072]) from checkpoint, the shape in current model is torch.Size([1536, 3072]).
        size mismatch for lq_proj.gate_modules.6.content_proj.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([1536]).

The reported error is not caused by the quantized Gemma model I installed. The failure occurs afterward when the PiD decoder is loaded at line 131.

Thus:

invokeai/backend/model_manager/starter_models.py:175, invokeai/backend/pid/decode.py:37, and invokeai/backend/model_manager/taxonomy.py:183: The latest commit points the FLUX, FLUX.2, and Qwen-Image 2K-to-4K starters at NVIDIA's v1.5 checkpoints, but build_pid_net() still constructs only the legacy architecture. V1.5 uses lq_hidden_dim=1024, scalar per-token gates, PiT LQ injection, replicate padding, and additional heads; InvokeAI constructs a 512-channel, per-dimension-gate network without PiT injection. This exactly explains the reported 1024-vs-512 and 1-vs-1536 mismatches. strict=False does not ignore tensor shape mismatches. The current variant enum also conflates legacy 2K-to-4K and v1.5, so selecting architecture solely from Res2kTo4k_Sr4x would break the legacy SD3/SDXL checkpoints. The official differences are visible in the NVIDIA v1.5 network configuration. Add an explicit architecture/version discriminator, implement the complete v1.5 network configuration, and document the distinction in docs/src/content/docs/features/pid-decode.mdx. Test: load representative legacy and v1.5 checkpoints for FLUX, FLUX.2, Qwen-Image, SD3, and SDXL and require every checkpoint tensor to match the selected network before running a minimal decode.
invokeai/backend/model_manager/configs/pid_decoder.py:72-111: Direct single-file installs lose NVIDIA's directory name because _name_for_matching() sees only the UUID directory and model_ema_bf16.pth; _variant_from_filename() therefore labels a v1.5 2K-to-4K checkpoint as res2k_sr4x. The reported local decoder record demonstrates this exact state. The classifier also validates only latent input channels, so it accepts the incompatible v1.5 architecture and defers the failure until execution. Test: install each checkpoint both through Starter Models and through a direct URL whose local filename is model_ema_bf16.pth; require identical base, resolution variant, architecture version, and successful strict loading.
invokeai/app/invocations/pid_decoder_loader.py:24, invokeai/backend/model_manager/configs/pid_decoder.py:169, and invokeai/backend/model_manager/taxonomy.py:186: User-facing titles, classifier errors, and generated schema text still describe only FLUX.1, FLUX.2, and SD3, despite support for SDXL and Qwen-Image. They also incorrectly imply every backbone has both presets. Test: assert the invocation title and classifier diagnostics enumerate all supported backbones and accurately describe their available legacy/v1.5 presets.

JPPhoto · 2026-07-27T02:30:25Z

Also, this combination of models takes up more VRAM than I have so I can't really test it. It seems that the quantized Gemma 2 model is dequantized before running so it takes up the same amount of VRAM. A path forward, either in this PR or a follow-up (with a note here that quantized Gemma 2 is not yet fully supported):

Implement it by replacing the Transformers GGUF model load with InvokeAI's native GGMLTensor path.

Keep the existing tokenizer load unchanged:

AutoTokenizer.from_pretrained(model_dir, gguf_file=gguf_file, local_files_only=True)

This parses tokenizer metadata without loading model tensors.

Read the model configuration without tensor materialization:

from transformers import Gemma2Config
from transformers.modeling_gguf_pytorch_utils import load_gguf_checkpoint

metadata = load_gguf_checkpoint(gguf_path, return_tensors=False)
gemma_config = Gemma2Config(**metadata["config"])

Alternatively, infer the small set of dimensions from tensor shapes, as the T5 and Qwen loaders do. Using Transformers' metadata parser avoids duplicating Gemma defaults.

Load quantized storage through InvokeAI:

compute_dtype = TorchDevice.choose_bfloat16_safe_dtype(
    TorchDevice.choose_torch_device()
)
sd = gguf_sd_loader(gguf_path, compute_dtype=compute_dtype)

gguf_sd_loader() returns GGMLTensor values rather than materialized BF16 tensors. See invokeai/backend/quantization/gguf/loaders.py.

Convert llama.cpp names to Gemma2Model names. Follow _convert_llamacpp_to_pytorch() in invokeai/backend/model_manager/load/model_loaders/z_image.py.

The essential mapping is:

token_embd.weight                   -> embed_tokens.weight
output_norm.weight                  -> norm.weight

blk.N.attn_q.weight                 -> layers.N.self_attn.q_proj.weight
blk.N.attn_k.weight                 -> layers.N.self_attn.k_proj.weight
blk.N.attn_v.weight                 -> layers.N.self_attn.v_proj.weight
blk.N.attn_output.weight            -> layers.N.self_attn.o_proj.weight

blk.N.ffn_gate.weight               -> layers.N.mlp.gate_proj.weight
blk.N.ffn_up.weight                 -> layers.N.mlp.up_proj.weight
blk.N.ffn_down.weight               -> layers.N.mlp.down_proj.weight

blk.N.attn_norm.weight              -> layers.N.input_layernorm.weight
blk.N.post_attention_norm.weight    -> layers.N.post_attention_layernorm.weight
blk.N.ffn_norm.weight               -> layers.N.pre_feedforward_layernorm.weight
blk.N.post_ffw_norm.weight          -> layers.N.post_feedforward_layernorm.weight

Use names without the model. prefix because PiD needs Gemma2Model, not Gemma2ForCausalLM.

Construct and populate an empty decoder:

from transformers import Gemma2Model

with accelerate.init_empty_weights():
    model = Gemma2Model(gemma_config)

incompatible = model.load_state_dict(sd, strict=False, assign=True)

Reject unexpected keys and any missing parameter other than explicitly understood nonpersistent buffers. Then assert that no parameter remains on meta.

This follows the working patterns in:

T5: invokeai/backend/model_manager/load/model_loaders/flux.py, T5EncoderGGUFModel._load_from_gguf()
Qwen3: invokeai/backend/model_manager/load/model_loaders/z_image.py, Qwen3EncoderGGUFLoader._load_from_gguf()

Eagerly materialize weights that cannot remain GGMLTensor:

embed_tokens.weight, because nn.Embedding requires indexed access.
Every RMSNorm weight. Gemma 2 calls self.weight.float() inside Gemma2RMSNorm; leaving those wrapped would trigger an unsupported dtype conversion.
Any other non-linear parameter discovered by a forward test.

A safe initial rule is to dequantize the embedding and every one-dimensional GGMLTensor, leaving the large two-dimensional projection weights quantized:

for module in model.modules():
    for name, param in list(module.named_parameters(recurse=False)):
        if isinstance(param, GGMLTensor) and (
            isinstance(module, torch.nn.Embedding) or param.ndim == 1
        ):
            setattr(
                module,
                name,
                torch.nn.Parameter(
                    param.get_dequantized_tensor(),
                    requires_grad=False,
                ),
            )

The remaining linear weights will be dequantized on demand by GGMLTensor and the model cache's custom linear handling. See:

invokeai/backend/quantization/gguf/ggml_tensor.py
invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py
invokeai/backend/model_manager/load/model_cache/model_cache.py

Finally, add tests that:

Verify every GGUF key mapping.
Assert a representative q_proj.weight remains a GGMLTensor.
Assert embeddings and norm weights are ordinary tensors.
Assert no parameters remain on meta.
Run a short forward pass and compare its hidden states with the current fully dequantized loader within a quantization-appropriate tolerance.

The main implementation belongs in invokeai/backend/model_manager/load/model_loaders/gemma2_encoder.py. No configuration, schema, or frontend changes should be necessary.

The GGUF Gemma encoder used transformers' from_pretrained(gguf_file=...), which dequantizes every weight at load — so a quantized Gemma cost the same VRAM as the unquantized model. Load it via InvokeAI's GGMLTensor path instead: read the config from GGUF metadata, map llama.cpp tensor names to Gemma2Model, and keep the 2D projection weights as GGMLTensor (dequantized on demand by the model cache). Materialize only the embedding and the RMSNorm weights, subtracting 1 from the norms (llama.cpp folds +1 in; Gemma2RMSNorm re-adds it), and assert nothing is left on meta. Verified: hidden states match the fully-dequantized loader within quantization tolerance. Adds key-mapping tests and a local load/compare test. NVIDIA's v1.5 decoders use a different network (lq_hidden_dim=1024, PiT injection) that build_pid_net (512-dim legacy) cannot load, causing a size-mismatch crash. - Point the FLUX/FLUX.2/Qwen 2K-to-4K starters back at the legacy checkpoints (moved to checkpoints_deprecated/) that the current network loads. - Reject a checkpoint whose lq_proj hidden dim is not the supported 512 at identification time, instead of accepting it and failing inside the decode. - Enumerate all supported backbones (add SDXL, Qwen-Image) in the loader title and correct the variant enum docs (not every backbone ships both presets). Full v1.5 architecture support is planned as a follow-up. Adds PiD decoder identification tests (legacy accepted, 1024-dim v1.5 rejected).

JPPhoto

@Pfannkuchensack Approved! It needs a follow-up PR with (at least) the following:

invokeai/backend/model_manager/configs/pid_decoder.py:_looks_like_pid_decoder and invokeai/backend/pid/decode.py:load_pid_decoder: one lq_proj key is enough for identification, while all missing lq_proj.* parameters are tolerated. Model creation runs under skip_torch_weight_init, so a partial checkpoint can leave uninitialized Conv/Linear weights and produce garbage or NaNs. Test: remove required LQ keys from a valid checkpoint and require identification/load failure.
invokeai/backend/model_manager/configs/pid_decoder.py:_variant_from_filename: direct single-file installs stored as <uuid>/model_ema_bf16.pth default to Res2k_Sr4x. SDXL and Qwen only provide Res2kTo4k_Sr4x, so records can be mislabeled. Runtime currently ignores this field, hence follow-up. Test: directly install each decoder and compare base/variant with its Starter Model record.
tests/backend/model_manager/load/test_gemma2_encoder_gguf_loader.py:17: the only full native-GGUF loader test uses an author's hardcoded Windows path and always skips in CI. Mapping tests do not exercise quantized retention, norm conversion, meta-buffer repair, or forward execution. Test: use a tiny synthetic Gemma configuration plus mocked GGUF tensors; assert projection weights remain GGMLTensor, embedding/norms materialize, no meta parameters remain, and forward output is finite.
docs/src/content/docs/features/pid-decode.mdx:15,47 and invokeai/backend/model_manager/configs/gemma2_encoder.py:Gemma2Encoder_GGUF_Config: documentation says PiD cannot upscale an existing image despite the new pid_upscale node, says all legacy starters live under checkpoints_deprecated/ despite mixed starter paths, and still says Transformers dequantizes GGUF despite the new native loader. Test: remove these stale claims, build docs/schema, and document Generation PiD separately from the prototype upscale node.

Follow-up to invoke-ai#9281, addressing the review items: - Require the complete LQ projection. Identification accepted a checkpoint with a single `lq_proj.*` key and `load_pid_decoder` tolerated every missing `lq_proj.*`. Models are built under `skip_torch_weight_init()`, so those weights stayed uninitialised and would decode to garbage/NaNs. `required_lq_proj_keys()` derives the expected key set from the vendored network, and both identification and load now reject any missing key. - Match the install source when identifying backbone and variant. A direct single-file install is stored as `<uuid>/model_ema_bf16.pth`, so the name carried no `res2k…` marker and no backbone hint: SDXL/Qwen-Image decoders were labelled `res2k_sr4x` although only the 2K-to-4K preset exists, and SD3/Qwen-Image decoders were registered as `flux`, which their decode node then rejects. The source (HF path/URL) survives the download and is now matched alongside the on-disk name, with a per-backbone variant fallback. - Replace the hardcoded-path Gemma-2 GGUF loader test (always skipped in CI) with a synthetic tiny Gemma-2 built from mocked GGUF tensors: asserts quantized retention, norm materialisation, meta-buffer repair, absence of meta parameters and a finite forward. The real-file comparison is now opt-in via INVOKEAI_TEST_GEMMA2_GGUF. - Drop stale doc claims: PiD as a decode is now documented separately from the prototype `pid_upscale` node, the starter checkpoints are spread over `checkpoints/` and `checkpoints_deprecated/`, and the GGUF encoder is loaded natively instead of being dequantized by transformers.

Pfannkuchensack added 3 commits May 29, 2026 05:03

github-actions Bot added python PRs that change python files Root invocations PRs that change invocations backend PRs that change backend files services PRs that change app services frontend PRs that change frontend files labels Jun 8, 2026

Pfannkuchensack added the DO NOT MERGE label Jun 8, 2026

Merge branch 'main' into feat/pid-decoder

6b72ef0

lstein assigned JPPhoto Jun 17, 2026

lstein added the 6.14.0 label Jun 17, 2026

lstein added this to Invoke - Community Roadmap Jun 17, 2026

lstein moved this to 6.14.x Theme: USER EXPERIENCE in Invoke - Community Roadmap Jun 17, 2026

Pfannkuchensack and others added 15 commits June 21, 2026 23:31

Merge branch 'invoke-ai:main' into feat/pid-decoder

bc79446

Chore Ruff

47c7513

Chore Typegen

51e6a93

Chore Knip

9029c12

fix(pid): remove unused vendored models/utils.py (broken easy_io import)

2aeaa97

Chore fix

01a0c0f

JPPhoto self-requested a review July 25, 2026 23:10

JPPhoto requested changes Jul 26, 2026

View reviewed changes

Pfannkuchensack and others added 2 commits July 26, 2026 21:09

chore: regenerate OpenAPI schema and frontend types

9ebeb90

JPPhoto self-requested a review July 26, 2026 20:31

JPPhoto approved these changes Jul 26, 2026

View reviewed changes

Pfannkuchensack added 6 commits July 26, 2026 23:34

Merge branch 'feat/pid-decoder' of https://github.com/Pfannkuchensack…

5183eb8

…/InvokeAI into feat/pid-decoder

Chore Ruff

f523fb9

JPPhoto requested changes Jul 27, 2026

View reviewed changes

Pfannkuchensack added 2 commits July 27, 2026 18:39

Chore openapi + typegen

916573b

Pfannkuchensack requested a review from JPPhoto July 27, 2026 17:24

Pfannkuchensack added 6 commits July 27, 2026 22:15

Merge branch 'main' into feat/pid-decoder

e327b8b

Docs update

e695d21

Merge remote-tracking branch 'upstream/main' into feat/pid-decoder

efa5525

Merge branch 'main' into feat/pid-decoder

dbf79b6

Merge remote-tracking branch 'upstream/main' into feat/pid-decoder

1ef5459

Merge branch 'main' into feat/pid-decoder

9f93e98

JPPhoto approved these changes Jul 30, 2026

View reviewed changes

Merge branch 'main' into feat/pid-decoder

5a96094

lstein merged commit 3f5588f into invoke-ai:main Jul 30, 2026
17 checks passed

Pfannkuchensack deleted the feat/pid-decoder branch July 31, 2026 00:54

Pfannkuchensack mentioned this pull request Aug 6, 2026

Feat: pid followup #9474

Open

5 tasks

Navigation Menu

Uh oh!

Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image - #9281

Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image#9281
lstein merged 52 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/pid-decoder

Pfannkuchensack commented Jun 8, 2026 •

edited

Loading

Uh oh!

JPPhoto left a comment

Uh oh!

JPPhoto left a comment

Uh oh!

JPPhoto left a comment

Uh oh!

JPPhoto commented Jul 27, 2026 •

edited

Loading

Uh oh!

JPPhoto left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Pfannkuchensack commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Related Issues / Discussions

QA Instructions

Merge Plan

Checklist

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Merge blockers

Follow-up PR candidates

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

JPPhoto commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JPPhoto left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Pfannkuchensack commented Jun 8, 2026 •

edited

Loading

JPPhoto commented Jul 27, 2026 •

edited

Loading