Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image - #9281
Merged
Conversation
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
requested changes
Jul 26, 2026
JPPhoto
left a comment
Collaborator
There was a problem hiding this comment.
Latest review:
Merge blockers
invokeai/backend/model_manager/configs/text_llm.py:44-53andinvokeai/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, whileTextLLM_Diffusers_Configdefers everyGemma2ForCausalLMregardless of size.ModelConfigFactoryconsequently classifies these valid causal LMs asUnknowninstead of preserving their previousTextLLMbehavior. Test: Classify Gemma 2 configurations with hidden sizes 2304, 3584, and 4608 without overrides; expect 2304 to becomeGemma2Encoderand the larger variants to remainTextLLM.
Follow-up PR candidates
-
invokeai/backend/pid/decode.py:202-221andtests/backend/pid/test_pid_decode.py:31-34: The new invalid-schedule safety net usesassert, which is removed underPYTHONOPTIMIZE=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 underpython -Oand require an explicitValueErroror validated configuration rather than an assertion. -
invokeai/app/invocations/flux_pid_decode.py:97-110andinvokeai/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 providingLoadedModel.compute_devicefor 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 usesgemma_text_encoder_info.compute_deviceand completes without device mismatches. -
invokeai/app/invocations/pid_upscale.py:72-75andinvokeai/app/invocations/flux_vae_encode.py:39-55: The new upscale node advertises Z-Image and other 16-channel-compatible VAEs, but delegates encoding toFluxVaeEncodeInvocation.vae_encode(), which accepts only InvokeAI's FLUXAutoEncoder. A Z-Image DiffusersAutoencoderKL, which the existing Z-Image encode path explicitly supports, cannot satisfy this implementation and fails instead of upscaling. Test: Connect both a FLUXAutoEncoderand a Z-Image DiffusersAutoencoderKLtopid_upscale; either support both advertised cases with their correct scaling rules or narrow the field description and validation to FLUXAutoEncoderonly.
…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
self-requested a review
July 26, 2026 20:31
JPPhoto
approved these changes
Jul 26, 2026
JPPhoto
left a comment
Collaborator
There was a problem hiding this comment.
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, andinvokeai/backend/model_manager/taxonomy.py:186still 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 atpid_decoder.py:169andpid_decoder.py:191repeat 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
requested changes
Jul 27, 2026
JPPhoto
left a comment
Collaborator
There was a problem hiding this comment.
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, andinvokeai/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, butbuild_pid_net()still constructs only the legacy architecture. V1.5 useslq_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=Falsedoes not ignore tensor shape mismatches. The current variant enum also conflates legacy 2K-to-4K and v1.5, so selecting architecture solely fromRes2kTo4k_Sr4xwould 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 indocs/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 andmodel_ema_bf16.pth;_variant_from_filename()therefore labels a v1.5 2K-to-4K checkpoint asres2k_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 ismodel_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, andinvokeai/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.
Collaborator
|
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
AutoTokenizer.from_pretrained(model_dir, gguf_file=gguf_file, local_files_only=True)This parses tokenizer metadata without loading model tensors.
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.
compute_dtype = TorchDevice.choose_bfloat16_safe_dtype(
TorchDevice.choose_torch_device()
)
sd = gguf_sd_loader(gguf_path, compute_dtype=compute_dtype)
The essential mapping is: Use names without the
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 This follows the working patterns in:
A safe initial rule is to dequantize the embedding and every one-dimensional 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
Finally, add tests that:
The main implementation belongs in |
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
approved these changes
Jul 30, 2026
JPPhoto
left a comment
Collaborator
There was a problem hiding this comment.
@Pfannkuchensack Approved! It needs a follow-up PR with (at least) the following:
-
invokeai/backend/model_manager/configs/pid_decoder.py:_looks_like_pid_decoderandinvokeai/backend/pid/decode.py:load_pid_decoder: onelq_projkey is enough for identification, while all missinglq_proj.*parameters are tolerated. Model creation runs underskip_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.pthdefault toRes2k_Sr4x. SDXL and Qwen only provideRes2kTo4k_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 remainGGMLTensor, embedding/norms materialize, no meta parameters remain, and forward output is finite. -
docs/src/content/docs/features/pid-decode.mdx:15,47andinvokeai/backend/model_manager/configs/gemma2_encoder.py:Gemma2Encoder_GGUF_Config: documentation says PiD cannot upscale an existing image despite the newpid_upscalenode, says all legacy starters live undercheckpoints_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.
Pfannkuchensack
added a commit
to Pfannkuchensack/InvokeAI
that referenced
this pull request
Aug 6, 2026
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.
lstein
pushed a commit
that referenced
this pull request
Aug 13, 2026
* fix(pid): harden PiD decoder identification, GGUF loader tests and docs Follow-up to #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. * fix(pid): check LQ completeness before the backbone, stop probing the RNG Review follow-ups for #9474. 1. `_raise_if_lq_projection_incomplete` ran after `_validate_base`, but the backbone is read from `lq_proj.latent_proj.0.weight` — one of the weights a truncated file may be missing. Such a file therefore failed with "cannot determine PiD decoder backbone" instead of the "missing … LQ projection weights" message the install flow promises. Completeness is now checked first, against `common_required_lq_proj_keys()`, the key set every backbone requires. The per-backbone check stays after the backbone is known: the two sets are the same 71 keys today, so it is a no-op that keeps the check from silently weakening to the intersection if a backbone ever adds LQ parameters of its own. Note this does not change the `Unknown_Config` fallback: `ModelConfigFactory` applies that to any file no config matches, and `allow_unknown_models` defaults to true. With it disabled the truncated file is rejected outright. The PR's QA step is worded as if rejection were unconditional; it is not, and that is a model-manager-wide behaviour rather than anything PiD-specific. 2. `required_lq_proj_keys()` built a real `LQProjection2D` just to read parameter names, running every `reset_parameters()` and so drawing from the global CPU RNG during model identification — leaving later unseeded randomness dependent on how many candidate files were probed. It is now built on the meta device inside `torch.random.fork_rng`. 3. The `net.` normalisation existed twice, and the copies had already diverged: only the loader dropped the distill-only submodules (`net_ema.`, `fake_score.`, `discriminator.`). Since `net_ema.*` shadows PidNet's own parameter names, the drift ran in the direction where identification accepts what the loader then refuses. Both now share `backend/pid/state_dict_utils.py`. * fix(pid): reject a recognised-but-broken PiD checkpoint instead of registering it Identification rejected a truncated PiD checkpoint with NotAMatchError, which only means "not my kind of model": ModelConfigFactory collects those, finds no match, and with allow_unknown_models (default: true) falls back to Unknown_Config. A file that had already identified itself as a PiD decoder and was then found to be missing LQ projection weights was therefore still installed, as an unknown model with a database record, and only failed once something tried to load it. Add InvalidMatchError for "recognised, and unusable". It is deliberately not a NotAMatchError subclass, since the factory catches that one per candidate class and would swallow it. When no config class matched and at least one raised it, classification returns no config regardless of allow_unknown, and ModelInstallService._probe reports the specific reason rather than the misleading "could not identify model". Order the architecture check ahead of the completeness check. A v1.5 checkpoint is intact, just built to a shape InvokeAI cannot construct; judged against the legacy key set it would be misreported as truncated and now hard-rejected on top of that. It stays a plain no-match, so it remains registrable as an unknown model - only a broken file is fatal. This costs the completeness check nothing: the hidden dim is read from lq_proj.latent_proj.0.weight, so a file truncated past that weight falls straight through to it. Collapse the LQ key contract to one entry point. common_required_lq_proj_keys() and the per-backbone re-check are gone; required_lq_proj_keys() takes no backbone, the probe lives in the private _probe_lq_proj_keys(), and test_pid_decode.py pins that every backbone agrees, so key drift fails in CI instead of silently weakening the install-time check. * fix(pid): make every backbone-independent PiD rejection final The previous commit made a truncated checkpoint fatal but left the architecture check a plain no-match, on the reasoning that an intact v1.5 file is merely unsupported and should stay registrable. That opened a hole: a file that is both 1024-dim and truncated is rejected by the architecture check first, never reaches the completeness check, and lands back in Unknown_Config - the exact outcome the previous commit set out to prevent. The distinction does not survive contact with the failure mode. Once a file has identified itself as a PiD decoder, any rejection that does not depend on which backbone it is will be raised identically by all five config classes, so the file ends up with no match and is registered through the Unknown_Config fallback. Those rejections are now all InvalidMatchError: unsupported lq_hidden_dim, an incomplete LQ projection, a latent channel count no backbone uses, and a checkpoint whose backbone cannot be determined at all. Splitting them out of _validate_base is what makes that legible. _validate_base now only ever answers "not *this* backbone", which four of the five classes are supposed to say about every valid checkpoint, and every rejection in it stays a NotAMatchError. The backbone-independent checks run ahead of it in from_model_on_disk, architecture first so an intact v1.5 file is diagnosed as unsupported rather than judged against the legacy key set and misreported as truncated. Also handle InvalidModelConfigException in the startup orphan scan. ModelSearch._walk_directory already contains anything the on_model_found callback raises, so startup was never actually at risk; catching it in the callback makes skipping a bad file a property of the scan rather than of its caller, and names the file and the reason in the log. * fix(pid): hold a checkpoint to PidNet's whole contract, and stop guessing from paths Identification checked less than the loader demands and inferred the rest from file paths. Three consequences, all reported in review: A checkpoint with every lq_proj weight and none of the 385 backbone weights was registered, then refused by load_pid_decoder. Only the 71-key LQ projection was ever checked. A subset check is not a milder version of the same guarantee: loaders run under skip_torch_weight_init(), so a weight the checkpoint does not supply is uninitialised memory rather than a default. required_pid_net_shapes() now derives the whole contract - 456 keys and their shapes - from a meta-device PidNet, the same trick the LQ probe already used but applied to the real network instead of one submodule. Missing keys, unexpected keys and wrong shapes are all fatal, because all three are fatal in load_pid_decoder; a stricter installer cannot reject a file that would have loaded. Probing the real net also removes the reason _LQ_PROBE_DIM, _LQ_NUM_RES_BLOCKS_DEFAULT and the hand-copied num_outputs derivation existed, along with the test that kept them in sync. Wrong-shaped tensors were accepted when a filename supplied a backbone. The architecture, the backbone and the kernel are all read off lq_proj.latent_proj.0 .weight, and each read answered None when it was not a 4D conv - so one malformed tensor made all three abstain at once and the file fell through to name-only matching. That weight is now validated first, and the three reads only run when it is there; its absence is a truncation, which the contract check diagnoses better than a guess about the architecture. Backbone detection concatenated the install source, the parent directory and the filename into one string and substring-matched it with a fixed precedence, so /flux/model_sd3.pth matched flux first and was registered as FLUX although the file says sd3. Name components are now matched most-specific-first, a component naming two different backbones decides nothing rather than being resolved by precedence, and a local-path source is not evidence at all - the model manager sets source to the file's own path when there is no remote one, so trusting it means matching arbitrary ancestor directories of the user's model library. Nothing is lost: install_path identifies a local file before it moves it. Requiring the full contract also makes the latent channel count always readable, which retires the name-only backbone path entirely. The name can now only break the FLUX.1 / SD3 / Qwen-Image tie, never pick a backbone outright, and an explicit base override - already validated against one class's Literal - beats it. The checks that would rule out all five configs move out of _validate_base, leaving it to answer only "not this backbone". Also fixes, orthogonally: from_model_on_disk popped `variant` out of the override dict the factory builds once and shares across every candidate class, so the first PiD class to run consumed it and a later one that actually matched fell back to name inference. Verified against all 11 NVIDIA checkpoints: every one matches the contract exactly (missing=0, unexpected=0, no shape mismatch), the 9 supported decoders identify with the right base and variant both in place and as a direct single-file install, and the dinov2 / siglip decoders are rejected by latent channel count rather than registered as unknown models. * fix(pid): tolerate a checkpoint whose keys are not all strings A bare (un-prefixed) PidNet checkpoint is passed through strip_net_prefix untouched, on purpose: without the net. prefix there is no evidence the file is a distill serialisation, so a stray key must reach the unexpected-key checks rather than be dropped. A .pth unpickles to whatever it contains, so those keys need not all be strings - and reporting the unexpected ones sorts them. Sorting {1, "not_a_pid_key"} raises TypeError. That failure does not surface as a failure. ModelConfigFactory catches an unexpected exception from a candidate class as a generic no-match, so all five PiD configs drop out and allow_unknown_models registers the file as Unknown_Config - the exact fallback these checks exist to close. A complete bare contract plus one non-string key and one unexpected string key was therefore installed as an unknown model. Sort both key sets with key=str, and stop the type annotations claiming otherwise: strip_net_prefix and pid_net_shapes return dict[Any, ...], not dict[str, ...], and _Shapes follows. The old signature was not merely imprecise - it carried a type: ignore for the pass-through return, which is what let a str-only assumption look checked. Verified against the eleven real NVIDIA checkpoints: unchanged, all five supported backbones identify in place and as a direct single-file install, and the dinov2 / siglip decoders are still rejected by latent channel count. * fix(pid): reject non-string keys before torch trips over them The previous commit taught identification to tolerate a bare checkpoint whose keys are not all strings, and justified keeping those keys with a claim about the loader that is simply wrong: nn.Module.load_state_dict calls .startswith() on every key, so a non-string one raises AttributeError from inside torch before any unexpected key is reported. Passing a complete state dict plus {1: tensor} to load_pid_decoder raised that AttributeError rather than the RuntimeError the function reports every other unusable checkpoint with. load_pid_decoder now checks for non-string keys before it hands anything to torch, and says what is actually wrong with the file. Identification already rejects such a checkpoint, so this is the second line rather than the first - but load_pid_decoder is public, the model cache reaches it for records written before this PR, and a file can be swapped on disk after install. The reasoning in strip_net_prefix and its test is corrected to match what torch does. Keeping non-string keys is still right - dropping them would hide a malformed file from the checks meant to catch it - but the burden it puts on consumers is the opposite of what was written there: neither may assume the key type, so identification sorts its key reports with key=str and the loader rejects non-strings up front. Verified: the reviewer's repro now raises "PiD checkpoint has 1 keys that are not strings and so cannot name a PidNet parameter: [1]". The eleven real NVIDIA checkpoints are unaffected, 22/22 as before. --------- Co-authored-by: Jonathan <34005131+JPPhoto@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
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.pid_decoder_loader(→PiDDecoderField) andgemma2_encoder_loader(→Gemma2EncoderField), plus model-manager configs/loaders for PiD checkpoints and the shared Gemma-2 caption encoder.flux_pid_decodeflux2_pid_decodesd3_pid_decodesdxl_pid_decodescaling_factorread at runtimez_image_pid_decodeqwen_image_pid_decodelatents_mean/latents_stddenorm + 5D→4D temporal squeezenvidia/PiD, per backbone; FLUX/FLUX.2/SD3 ship 2K + 2K-to-4K, SDXL/Qwen-Image ship 2K-to-4K only) plus the sharedgemma-2-2b-itcaption encoder.Robustness details
baseoverride (which the starter installer sends) when the directory name is ambiguous — so single-file HF downloads are still identified correctly.Related Issues / Discussions
Closes #9240
QA Instructions
base(e.g. a Qwen-Image decoder shows asqwen-image, not FLUX).Automated gates already green on the branch: backend imports (
starter_models, config factory, every*_pid_decodenode), and the frontendpnpm 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
paramsslice bump.Checklist
_versionbump)What's Newcopy (if doing a release after this PR)