fix(model loaders): ignore unexpected checkpoint keys, report them at debug only

invokeai/backend/ideogram4/quantized_loading.py

@@ -7,6 +7,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from invokeai.backend.util.state_dict_loading import log_unexpected_keys
+
 if TYPE_CHECKING:
     pass
 
@@ -104,8 +106,7 @@ def load_bnb4bit_state_dict(
     real_missing = [m for m in missing if m not in consumed]
     if real_missing:
         raise RuntimeError(f"missing keys after quantized load: {real_missing[:10]}")
-    if unexpected:
-        raise RuntimeError(f"unexpected keys after quantized load: {unexpected[:10]}")
+    log_unexpected_keys("Ideogram 4 quantized checkpoint", unexpected)
 
     for p in model.parameters():
         if isinstance(p, bnb.nn.Params4bit):
@@ -258,7 +259,7 @@ def load_fp8_state_dict(
     the caller must have already put the unquantized params in ``dtype``.
 
     ``strict=False`` downgrades missing keys to a warning (e.g. tied weights that a
-    ``transformers`` model resolves itself); unexpected keys always raise.
+    ``transformers`` model resolves itself); unexpected keys are logged at DEBUG and ignored.
     """
     prepared: dict[str, torch.Tensor] = {}
     for k, v in state_dict.items():
@@ -272,8 +273,7 @@ def load_fp8_state_dict(
             prepared[k] = v.to(device=device)
 
     missing, unexpected = model.load_state_dict(prepared, strict=False, assign=assign)
-    if unexpected:
-        raise RuntimeError(f"unexpected keys after fp8 load: {unexpected[:10]}")
+    log_unexpected_keys("Ideogram 4 fp8 checkpoint", unexpected)
     if missing:
         if strict:
             raise RuntimeError(f"missing keys after fp8 load: {missing[:10]}")

invokeai/backend/model_manager/configs/pid_decoder.py

@@ -137,9 +137,14 @@ def _raise_if_pid_net_contract_unmet(shapes: _Shapes, contract: Mapping[str, tup
     guarantee — loaders run under `skip_torch_weight_init()`, so a weight the checkpoint does not
     supply is uninitialised memory rather than a default.
 
-    Missing *and* unexpected keys are fatal here because both are fatal there, which is what makes
-    installation and loading accept the same set of files. A stricter installer cannot reject a file
-    that would have loaded: the loader already refuses everything rejected here.
+    Missing keys are fatal here because they are fatal there, which is what makes installation and
+    loading accept the same set of files. A stricter installer cannot reject a file that would have
+    loaded: the loader already refuses everything rejected here.
+
+    Extra keys are *not* fatal — `load_pid_decoder` ignores them (issue #9437), so rejecting them
+    here would refuse to install a file that loads fine. The one kind of extra key the loader still
+    cannot survive is a non-string one, which makes `nn.Module.load_state_dict` raise from inside
+    torch, so that is the extra this check keeps.
 
     `_LATENT_PROJ_KEY` is excluded from the shape comparison, and only from that: it is the one
     parameter whose shape legitimately varies by backbone, and its variable dimensions each have a
@@ -151,18 +156,18 @@ def _raise_if_pid_net_contract_unmet(shapes: _Shapes, contract: Mapping[str, tup
     # Both sorts take `key=str`: a bare checkpoint's keys need not all be strings (see
     # `strip_net_prefix`), and sorting a mixed set raises TypeError — which the factory answers with
     # the `Unknown_Config` registration these checks exist to prevent, so the crash fails as a silent
-    # accept rather than loudly. Only `unexpected` can hold one today; sorting both the same way keeps
-    # that from depending on which set is on which side of the subtraction.
+    # accept rather than loudly.
     if missing := sorted(contract.keys() - shapes.keys(), key=str):
         raise InvalidMatchError(
             f"PiD checkpoint is missing {len(missing)} of the weights required by PidNet; the file is "
             f"incomplete and cannot be used as a PiD decoder: {missing[:5]}{_and_more(missing)}"
         )
 
-    if unexpected := sorted(shapes.keys() - contract.keys(), key=str):
+    if not_strings := sorted((k for k in shapes if not isinstance(k, str)), key=str):
         raise InvalidMatchError(
-            f"PiD checkpoint has {len(unexpected)} keys PidNet does not expect, which `load_pid_decoder` "
-            f"rejects too: {unexpected[:5]}{_and_more(unexpected)}"
+            f"PiD checkpoint has {len(not_strings)} keys that are not strings and so cannot name a "
+            f"PidNet parameter, which `load_pid_decoder` rejects too: "
+            f"{not_strings[:5]}{_and_more(not_strings)}"
         )
 
     mismatched = [(k, shapes[k], want) for k, want in contract.items() if k != _LATENT_PROJ_KEY and shapes[k] != want]

invokeai/backend/model_manager/load/model_loaders/anima.py

@@ -21,6 +21,7 @@
 )
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.util.state_dict_loading import log_unexpected_keys, reject_incomplete_load
 
 logger = InvokeAILogger.get_logger(__name__)
 
@@ -175,17 +176,13 @@ def _load_from_singlefile(
                 sd[k] = sd[k].to(model_dtype)
 
         load_result = model.load_state_dict(sd, assign=True, strict=False)
-        if load_result.unexpected_keys:
-            raise RuntimeError(
-                f"Checkpoint contains {len(load_result.unexpected_keys)} unexpected keys. "
-                f"This may indicate a corrupted or incompatible checkpoint. "
-                f"First 5 unexpected keys: {load_result.unexpected_keys[:5]}"
-            )
-        if load_result.missing_keys:
-            logger.warning(
-                f"Checkpoint is missing {len(load_result.missing_keys)} keys "
-                f"(expected for inv_freq buffers). First 5: {load_result.missing_keys[:5]}"
-            )
+        log_unexpected_keys("Anima transformer checkpoint", load_result.unexpected_keys)
+        # `missing_keys` alone cannot police completeness here: AnimaTransformer's only three buffers
+        # are registered `persistent=False`, so they never appear in it (the old warning claiming
+        # otherwise was misleading). Sweep for tensors the checkpoint left on the meta device instead
+        # — that is the failure worth catching, and it is what the removed unexpected-key
+        # `RuntimeError` was really standing in for.
+        reject_incomplete_load(model, what="Anima transformer checkpoint")
 
         # Without this the `fp8_storage` toggle is shown for Anima models but does nothing. The
         # state dict was cast to a single `model_dtype` above, so the layerwise cast has one

invokeai/backend/model_manager/load/model_loaders/flux.py

@@ -86,6 +86,7 @@
 from invokeai.backend.quantization.sdnq.loaders import raise_on_incomplete_sdnq_load, sdnq_sd_loader
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.backend.util.silence_warnings import SilenceWarnings
+from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras
 
 logger = InvokeAILogger.get_logger(__name__)
 
@@ -117,7 +118,7 @@ def _load_model(
         with accelerate.init_empty_weights():
             model = AutoEncoder(get_flux_ae_params())
         sd = load_file(model_path)
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX VAE checkpoint", assign=True)
         # VAE is broken in float16, which mps defaults to
         if self._torch_dtype == torch.float16:
             try:
@@ -241,7 +242,7 @@ def _load_model(
         for k in sd.keys():
             sd[k] = sd[k].to(torch.bfloat16)
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX.2 VAE checkpoint", assign=True)
 
         # VAE is broken in float16, which mps defaults to
         if self._torch_dtype == torch.float16:
@@ -322,9 +323,13 @@ def _load_state_dict_into_t5(cls, model: T5EncoderModel, state_dict: dict[str, t
         # There is a shared reference to a single weight tensor in the model.
         # Both "encoder.embed_tokens.weight" and "shared.weight" refer to the same tensor, so only the latter should
         # be present in the state_dict.
-        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False, assign=True)
-        assert len(unexpected_keys) == 0
-        assert set(missing_keys) == {"encoder.embed_tokens.weight"}
+        load_state_dict_ignoring_extras(
+            model,
+            state_dict,
+            source="FLUX bnb-int8 T5 encoder",
+            assign=True,
+            allowed_missing={"encoder.embed_tokens.weight"},
+        )
         # Re-tie shared weights. In transformers 5.x, weight tying is implemented at the
         # parameter level (via _tie_weights / tie_weights) rather than as a Python object
         # alias.  load_state_dict(assign=True) replaces parameters in-place, which severs
@@ -467,7 +472,7 @@ def _load_from_gguf(self, config: T5Encoder_GGUF_Config) -> AnyModel:
             model = T5EncoderModel(t5_config)
 
         # Leave transformer Linear weights as GGMLTensors; the autocast cache handles them.
-        model.load_state_dict(sd, strict=False, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX GGUF T5 encoder", assign=True, allow_missing=True)
 
         # Embedding lookups can't run on quantized GGMLTensors, so dequantize the token embeddings and
         # re-tie the encoder's embed_tokens to the shared embedding.
@@ -720,7 +725,7 @@ def _load_from_singlefile(
         for k in sd.keys():
             # We need to cast to bfloat16 due to it being the only currently supported dtype for inference
             sd[k] = sd[k].to(torch.bfloat16)
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX transformer checkpoint", assign=True)
         return model
 
 
@@ -766,7 +771,7 @@ def _load_from_singlefile(
             img_in_weight.quantized_data = img_in_weight.quantized_data.view(expected_img_in_weight_shape)
             img_in_weight.tensor_shape = expected_img_in_weight_shape
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX GGUF transformer checkpoint", assign=True)
         return model
 
 
@@ -808,7 +813,7 @@ def _load_from_singlefile(
             sd = load_file(model_path)
             if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in sd:
                 sd = convert_bundle_to_flux_transformer_checkpoint(sd)
-            model.load_state_dict(sd, assign=True)
+            load_state_dict_ignoring_extras(model, sd, source="FLUX nf4 transformer checkpoint", assign=True)
         return model
 
 
@@ -1068,7 +1073,7 @@ def _load_from_singlefile(
             converted_sd[k] = converted_sd[k].to(torch.bfloat16)
 
         # Load the state dict - guidance weights were already initialized above if missing
-        model.load_state_dict(converted_sd, assign=True)
+        load_state_dict_ignoring_extras(model, converted_sd, source="FLUX.2 transformer checkpoint", assign=True)
 
         return model
 
@@ -1223,11 +1228,9 @@ def _load_text_encoder(self, config: Main_SDNQ_Diffusers_Flux2_Config) -> AnyMod
             model = Qwen3ForCausalLM(te_config)
 
         sd = sdnq_sd_loader(te_dir, compute_dtype=torch.bfloat16)
-        missing, unexpected = model.load_state_dict(sd, assign=True, strict=False)
-        if unexpected:
-            raise ValueError(f"Unexpected keys loading SDNQ Qwen3 text encoder: {unexpected}")
-        if missing and missing != ["lm_head.weight"]:
-            raise ValueError(f"Unexpected missing keys loading SDNQ Qwen3 text encoder: {missing}")
+        missing = load_state_dict_ignoring_extras(
+            model, sd, source="SDNQ Qwen3 text encoder", assign=True, allowed_missing={"lm_head.weight"}
+        )
         if missing == ["lm_head.weight"]:
             model.lm_head.weight = model.model.embed_tokens.weight
         return model
@@ -1316,7 +1319,7 @@ def _load_from_singlefile(self, config: Main_SDNQ_Flux2_Config) -> AnyModel:
                         out2, in2, dtype=torch.bfloat16
                     )
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="SDNQ FLUX.2 transformer checkpoint", assign=True)
         return model
 
 
@@ -1463,7 +1466,7 @@ def _load_from_singlefile(
                         out_features2, in_features2, dtype=torch.bfloat16
                     )
 
-        model.load_state_dict(converted_sd, assign=True)
+        load_state_dict_ignoring_extras(model, converted_sd, source="FLUX.2 GGUF transformer checkpoint", assign=True)
         return model
 
 
@@ -1500,7 +1503,7 @@ def _load_xlabs_controlnet(self, sd: dict[str, torch.Tensor]) -> AnyModel:
             # HACK(ryand): Is it safe to assume dev here?
             model = XLabsControlNetFlux(get_flux_transformers_params(FluxVariantType.Dev))
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX XLabs ControlNet checkpoint", assign=True)
         return model
 
     def _load_instantx_controlnet(self, sd: dict[str, torch.Tensor]) -> AnyModel:
@@ -1511,7 +1514,7 @@ def _load_instantx_controlnet(self, sd: dict[str, torch.Tensor]) -> AnyModel:
         with accelerate.init_empty_weights():
             model = InstantXControlNetFlux(flux_params, num_control_modes)
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX InstantX ControlNet checkpoint", assign=True)
         return model
 
 
@@ -1555,7 +1558,7 @@ def _load_model(
         with accelerate.init_empty_weights():
             model = FluxReduxModel()
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="FLUX Redux checkpoint", assign=True)
         model.to(dtype=torch.bfloat16)
         return model
 
@@ -1635,7 +1638,7 @@ def _load_sdnq_transformer_checkpoint(self, config: Main_SDNQ_FLUX_Config) -> An
         if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in sd:
             sd = convert_bundle_to_flux_transformer_checkpoint(sd)
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="SDNQ FLUX transformer checkpoint", assign=True)
         return model
 
     def _load_sdnq_transformer(self, transformer_path: Path, config: Main_SDNQ_Diffusers_FLUX_Config) -> AnyModel:
@@ -1649,7 +1652,7 @@ def _load_sdnq_transformer(self, transformer_path: Path, config: Main_SDNQ_Diffu
         # Convert from diffusers format to BFL format
         sd = self._convert_diffusers_sd_to_bfl(sd)
 
-        model.load_state_dict(sd, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="SDNQ FLUX transformer", assign=True)
         return model
 
     def _convert_diffusers_sd_to_bfl(self, sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:

invokeai/backend/model_manager/load/model_loaders/gemma2_encoder.py

@@ -23,6 +23,7 @@
 from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
 from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.state_dict_loading import log_unexpected_keys
 
 # llama.cpp GGUF tensor-name component -> Gemma2Model (decoder-only) component. PiD consumes only the
 # decoder stack, so we target Gemma2Model directly: no `model.` prefix and no lm_head. Gemma-2 has no
@@ -168,8 +169,7 @@ def load_gemma2_model_from_gguf(gguf_path: Path, compute_dtype: "torch.dtype") -
         model = Gemma2Model(gemma_config)
 
     _missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
-    if unexpected:
-        raise RuntimeError(f"Unexpected keys loading Gemma-2 GGUF encoder: {unexpected[:10]}")
+    log_unexpected_keys("Gemma-2 GGUF encoder", unexpected)
 
     # Materialize the weights that cannot remain quantized:
     #  - the token embedding, because nn.Embedding needs indexed access, and

invokeai/backend/model_manager/load/model_loaders/ideogram4.py

@@ -32,6 +32,7 @@
     SubModelType,
 )
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras, log_unexpected_keys
 
 
 def _load_local_state_dict(folder: Path, basename: str) -> dict[str, torch.Tensor]:
@@ -133,7 +134,7 @@ def _load_one_transformer(self, folder: Path) -> torch.nn.Module:
             with accelerate.init_empty_weights():
                 model: torch.nn.Module = Ideogram4Transformer(Ideogram4Config())
                 model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=compute_dtype)
-            model.load_state_dict(sd, strict=True, assign=True)
+            load_state_dict_ignoring_extras(model, sd, source="Ideogram 4 nf4 transformer", assign=True)
             return model
 
         if is_fp8_state_dict(sd):
@@ -148,7 +149,7 @@ def _load_one_transformer(self, folder: Path) -> torch.nn.Module:
         # Unquantized fallback.
         with accelerate.init_empty_weights():
             model = Ideogram4Transformer(Ideogram4Config())
-        model.load_state_dict(sd, strict=True, assign=True)
+        load_state_dict_ignoring_extras(model, sd, source="Ideogram 4 transformer", assign=True)
         return model.to(compute_dtype)
 
     def _load_text_encoder(self, model_path: Path) -> AnyModel:
@@ -201,11 +202,10 @@ def _load_text_encoder(self, model_path: Path) -> AnyModel:
                 model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=compute_dtype)
 
         _, unexpected = model.load_state_dict(sd, strict=False, assign=True)
-        # Unexpected keys signal a wrong or contaminated checkpoint and must hard-fail. Missing keys are
-        # acceptable only for tied weights (resolved by _verify_encoder_fully_materialized via
+        # Extra keys are exporter noise, not a correctness signal - log them and move on. Missing keys
+        # are acceptable only for tied weights (resolved by _verify_encoder_fully_materialized via
         # tie_weights); any genuinely missing non-tied weight is caught there as a leftover meta tensor.
-        if unexpected:
-            raise RuntimeError(f"unexpected keys loading Ideogram 4 text encoder: {unexpected[:10]}")
+        log_unexpected_keys("Ideogram 4 text encoder", unexpected)
         _verify_encoder_fully_materialized(model, context="Ideogram 4 text encoder")
         if not is_bnb_nf4:
             model = model.to(compute_dtype)
@@ -225,6 +225,6 @@ def _load_vae(self, model_path: Path) -> AnyModel:
         sd = load_file(model_path / "vae" / "diffusion_pytorch_model.safetensors")
         sd = convert_diffusers_state_dict(sd)
         ae = AutoEncoder(AutoEncoderParams())
-        ae.load_state_dict(sd)
+        load_state_dict_ignoring_extras(ae, sd, source="Ideogram 4 VAE")
         ae.eval()
         return ae.to(model_dtype)