Feat: pid followup

[Pfannkuchensack]

Summary

Follow-up to #9281 (PiD decode). Backend, tests and docs only — no frontend behaviour changes (the schema.ts diff is a regenerated docstring).

PiD decoder identification checked less than the loader demands and inferred the rest from file paths, so a checkpoint InvokeAI cannot load could be installed and only fail later. This PR closes that gap and the follow-on problem it exposed: a file that identification does reject was still being registered, because a rejection had no way to say "this is my kind of model, and it is broken".

1. Hold a checkpoint to the whole PidNet contract

Only the 71-key LQ projection was ever checked, so a checkpoint carrying every lq_proj weight and none of the 385 backbone weights was registered and then refused by load_pid_decoder. 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 — such a file decodes to garbage or NaNs instead of failing.

required_pid_net_shapes() now derives the full contract — 456 keys and their shapes — from a meta-device PidNet, the same probe trick the LQ check already used, 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 network also retires _LQ_PROBE_DIM, _LQ_NUM_RES_BLOCKS_DEFAULT and the hand-copied num_outputs derivation, along with the test that kept them in sync with the vendored module.

2. Stop guessing the backbone from paths

• A malformed discriminator was accepted when a filename supplied a backbone. The architecture version, 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 far better than a guess about the architecture.
• Component-wise name matching. Backbone detection concatenated the install source, the parent directory and the filename into one string and substring-matched it with a fixed backbone precedence, so /flux/model_sd3.pth matched flux first and was registered as FLUX although the file says sd3. Components are now matched most-specific-first, and a component naming two different backbones decides nothing rather than being resolved by precedence.
• A local-path source is not evidence. 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 wherever the user keeps their models. Nothing is lost: install_path identifies a local file before it moves it.
• Requiring the full contract 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 beats it.
• Orthogonal fix found on the way: 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.

3. Make a rejection stick (InvalidMatchError)

Every config class signals "not mine" with NotAMatchError. The factory collects those and, with allow_unknown_models (default: true), answers no-match-at-all by registering Unknown_Config. That is right for a file nobody recognises and wrong for one that was recognised and found broken: a truncated PiD checkpoint was stored as a normal record and failed only when something tried to load it.

InvalidMatchError is the new signal 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 class matched and at least one raised it, ModelConfigFactory.from_model_on_disk returns no config regardless of allow_unknown, and ModelInstallService._probe reports the specific reason instead of a misleading "could not identify model".

Every rejection that would rule out all five backbone configs raises it — unsupported lq_hidden_dim, an unmet contract, a latent channel count no backbone uses, a malformed discriminator. Those checks moved out of _validate_base, which is left answering only "not this backbone" and stays entirely NotAMatchError.

Two hazards this class of change surfaced, both fixed:

• The architecture check must run before the contract check. A v1.5 checkpoint is intact, just built to a shape build_pid_net cannot construct; judged against the legacy contract it would be reported as a pile of missing and unexpected keys rather than as the newer architecture it is. Ordering it first costs the contract check nothing — the hidden dim comes from the same weight, so a file truncated past it falls straight through.
• A crash in an unusability check fails as a silent accept. A bare (un-prefixed) checkpoint is passed through strip_net_prefix untouched by design, so a .pth can hand identification keys that are not strings — and sorting {1, "not_a_pid_key"} for the unexpected-key report raises TypeError. The factory catches that as a generic candidate failure, so all five configs drop out and the file is registered as Unknown_Config. Both key sets now sort with key=str, and strip_net_prefix / pid_net_shapes / _Shapes no longer annotate a key type they do not guarantee (the old signature carried a type: ignore for exactly that pass-through, which is what made a str-only assumption look checked).

4. Startup orphan scan

_register_orphaned_models runs inside ModelInstallService.start(), which re-raises anything that escapes it. ModelSearch._walk_directory already contains every on_model_found exception, so startup was never actually at risk — verified. Handling InvalidModelConfigException 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.

5. Gemma-2 GGUF loader tests

test_gemma2_encoder_gguf_loader.py's only load test pointed at a hardcoded local Windows path and therefore always skipped in CI. It is replaced by a synthetic 2-layer Gemma-2 built from mocked GGUF tensors (Q8_0 projections, F32 norms), exercising the real loader path: projections stay GGMLTensor with the correct dequantized shape, embedding and RMSNorms are materialized (including llama.cpp's folded +1 norm convention), the rotary inv_freq meta-buffer is rebuilt, no meta parameters or buffers remain, a forward pass with quantized weights in place returns finite values, and an unexpected tensor fails loudly. The real-file comparison against transformers' own GGUF loader is kept as an opt-in test driven by INVOKEAI_TEST_GEMMA2_GGUF.

6. Docs

• pid-decode.mdx claimed PiD "is not a separate upscale pass" although Feat: Add PiD (Pixel Diffusion Decoder) 4× super-resolution decode for FLUX / FLUX.2 / SD3 / SDXL / Z-Image / Qwen-Image #9281 added the pid_upscale node. Generation PiD and the prototype node are now documented separately, with the node's FLUX-AutoEncoder-only constraint spelled out.
• The starter checkpoints are spread over checkpoints/ and checkpoints_deprecated/ on Hugging Face, not all under the latter.
• Gemma2Encoder_GGUF_Config still said transformers dequantizes the GGUF; it is loaded natively. Added a note that a Gemma-2-2b GGUF encoder is a supported manual install. openapi.json / schema.ts regenerated for the docstring.

Related Issues / Discussions

Follow-up to #9281 — #9281 (review)

Also addresses the reviews on this PR: 4889406183, 4889527773, 4891646254, 4893418632, 4896390812.

QA Instructions

Automated

• pytest tests/backend tests/model_identification tests/app/services/model_install — green. 83 tests across the four touched test files (tests/backend/pid/, test_pid_decoder_config.py, test_orphan_scan.py, test_gemma2_encoder_gguf_loader.py).
• The 4 failures in tests/app/services/model_install/test_model_install.py on my machine are network-bound (they resolve a local HF_ENDPOINT) and are identical on main.
• ruff check / ruff format --check clean; schema.ts regenerated via the standard typegen.

Every rejection path has a factory-level test asserting result.config is None, that no Unknown_Config was produced, and that the specific reason survives for _probe to report — partial file, missing backbone weight, unexpected key, wrong shape, intact v1.5, truncated v1.5, unsupported latent channels, malformed discriminator, and a bare checkpoint with a non-string key.

Verified against the real NVIDIA checkpoints

11 files locally — all five supported backbones, both presets, plus the two unsupported (dinov2, siglip) decoders. 22/22 assertions pass (each checkpoint both in place and as a direct single-file install):

1. No regression from the stricter check: every checkpoint matches the derived contract exactly — missing=0, unexpected=0, shape-mismatch=0.
2. In-place install: all 9 supported decoders identify with the correct base and variant.
3. Direct single-file install: hard-linked into a fresh <uuid>/model_ema_bf16.pth with no base / variant override, all 9 still come out correct (before this work: SD3 2K, SD3 2K-to-4K and Qwen-Image were identified as flux).
4. Unsupported backbones: the dinov2 (768ch) and siglip (1152ch) decoders are rejected by latent channel count rather than registered as unknown models.

Manual

• Install a PiD decoder from Starter Models — unchanged behaviour, base/variant as before.
• Install a decoder directly (paste e.g. nvidia/PiD::checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step/model_ema_bf16.pth) into Model Manager → registers as sd-3 / 2K and is accepted by the SD3 PiD decode node.
• Truncate a checkpoint (drop some tensors) and install it: identification must reject it with a "missing N of the weights required by PidNet" message and no model record, instead of registering it as unknown.

Merge Plan

Self-contained.

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 (n/a — no slice changes)
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[JPPhoto]

Merge blockers:

• invokeai/backend/model_manager/configs/pid_decoder.py:218: Completeness check runs after backbone inference. Anonymous truncated files missing lq_proj.latent_proj.0.weight fail with “cannot determine backbone,” then default allow_unknown_models=True registers Unknown_Config, contradicting PR QA requirement to reject them. Test: classify a temp file containing only net.lq_proj.latent_proj.1.weight with no base/source and allow_unknown=True; current result is Unknown_Config.

Other findings/issues:

• invokeai/backend/pid/decode.py:195: required_lq_proj_keys() initializes Conv/Linear weights only to inspect names, consuming global CPU RNG during model identification. Later unseeded randomness changes with install order. Test: compare the next torch.rand() after calling required_lq_proj_keys.cache_clear(); required_lq_proj_keys(...) against a seeded control; values differ.

Alternative implementation ideas

• Instead of constructing initialized LQProjection2D, build it under meta initialization or torch.random.fork_rng; this keeps key probing deterministic and side-effect free.

• Instead of validating LQ keys after backbone inference, validate normalized LQ keys first using a base-independent expected set; this catches missing diagnostic keys and prevents Unknown_Config fallback.

• Instead of duplicating net. normalization in invokeai/backend/model_manager/configs/pid_decoder.py and the loader, share one normalizer; this prevents future identification/runtime key-contract drift.

[JPPhoto]

Prior RNG and key-normalization findings are fixed!

This has surfaced:

• invokeai/backend/model_manager/configs/pid_decoder.py:197: NotAMatchError is caught by factory.py:665-688; default allow_unknown_models=true returns Unknown_Config, and _probe persists it. Test: partial .pth containing only net.lq_proj.latent_proj.1.weight with allow_unknown=True returns Unknown_Config. Violates PR QA contract.

Other ideas:

• Instead of routing invalid PiD files through generic Unknown_Config, make recognized-but-invalid PiD classification fatal; this prevents bad database records.

• Instead of intersecting per-backbone probes, expose one normalized expected-key contract; this removes duplicate checks and future key drift.

[JPPhoto]

To fix:

• invokeai/backend/model_manager/configs/pid_decoder.py:196-211 / invokeai/backend/model_manager/configs/factory.py:693-706: A truncated 1024-dim/v1.5 PiD file bypasses completeness, falls back to Unknown_Config, and is registered despite PR metadata promising partial checkpoints are rejected. Test: remove one LQ key from a 1024-dim checkpoint; factory currently returns Unknown_Config.

• invokeai/app/services/model_install/model_install_default.py:1122-1153: Startup orphan scanning catches only duplicate errors; the new invalid-checkpoint exception propagates and aborts start() when one corrupt PiD file exists. Test: enable scan_models_on_startup, add a partial PiD orphan, and verify startup continues.

Ideas:

• Instead of short-circuiting on non-512 dimensions, validate the relevant architecture contract first (or reject unsupported PiD files explicitly); this prevents truncated files entering the unknown-model fallback.

• Instead of propagating InvalidModelConfigException from orphan scanning, catch/log it and return False; this skips one bad file without aborting startup.

[JPPhoto]

Merge blockers:

• invokeai/backend/model_manager/configs/pid_decoder.py:241: Arbitrary local source paths affect backbone detection; /flux/model_sd3.pth is classified as FLUX because flux matches first. Test: classify that path with a complete 16-channel PiD fixture; current result is FLUX, not SD3.

• invokeai/backend/model_manager/configs/pid_decoder.py:45: Completeness checks only LQ keys. A file with all LQ keys but no 385 backbone keys is registered, then load_pid_decoder() rejects it. Test: factory-classify an LQ-only checkpoint; expect no config, but current result is a FLUX config.

Other findings/issues:

• invokeai/backend/model_manager/configs/pid_decoder.py:87: Wrong-shaped required tensors are accepted when the filename supplies a backbone. Loading later fails with a shape mismatch. Test: use complete key names with rank-1 lq_proj.latent_proj.0.weight; classification currently succeeds.

Alternative implementation ideas:

• Instead of concatenating full source paths with model names, parse only trusted HF metadata and prefer explicit overrides; this prevents incidental directory names from changing base or variant.

• Instead of checking only LQ keys, derive and compare the full normalized PidNet key contract on meta tensors; this keeps installation and loading acceptance identical.

• Instead of accepting malformed diagnostic tensors by filename, validate required tensor ranks and shapes during identification; this turns runtime load failures into clear install errors.

[JPPhoto]

Merge blockers:

• invokeai/backend/model_manager/configs/pid_decoder.py:155 and invokeai/backend/pid/state_dict_utils.py:76: A complete bare-PidNet state dict with both a non-string key and an unexpected string key reaches sorted(shapes.keys() - contract.keys()); bare checkpoints are intentionally preserved, so mixed int/str keys raise TypeError. ModelConfigFactory catches that as a generic candidate failure, then allow_unknown=True returns and registers Unknown_Config, bypassing this PR's unusable-checkpoint rejection. Test: build a bare fixture from required_pid_net_shapes() plus {1: torch.zeros(1), "not_a_pid_key": torch.zeros(1)}, call ModelConfigFactory.from_model_on_disk(..., allow_unknown=True), and assert result.config is None; current code returns Unknown_Config after the TypeError. Exact-head execution was unavailable because the immutable head is absent from the read-only local Git objects.

[JPPhoto]

• invokeai/backend/pid/decode.py:219 (fed by invokeai/backend/pid/state_dict_utils.py:48): Bare checkpoints preserve non-string keys, but the loader passes them directly to PyTorch, which raises AttributeError on .startswith() before reporting an unexpected key. Test: pass a complete state dict plus {1: torch.zeros(1)} to load_pid_decoder(); current head raises AttributeError instead of InvokeAI's RuntimeError.

  • Instead of passing non-string keys into load_state_dict(), reject them in load_pid_decoder() first; this keeps malformed-checkpoint handling explicit and avoids PyTorch's low-level error.

[JPPhoto]

You might want to look into this or make another followup:

• invokeai/backend/pid/state_dict_utils.py:34 and invokeai/backend/pid/state_dict_utils.py:67: Prefixed checkpoints silently drop non-string keys before contract validation and loader rejection. Test: add {1: torch.zeros(1)} to a complete net.-prefixed fixture and classify with allow_unknown=True; current head returns a PiD config instead of rejecting the malformed checkpoint.

Alternative implementation ideas:

• Instead of dropping non-string keys during prefixed normalization, preserve and reject them before loading; this gives bare and prefixed checkpoints one consistent contract and prevents silent acceptance.