feat(architectures): move the fp8-storage exclusion into a loader-flags facet

invokeai/backend/architectures/__init__.py

@@ -43,6 +43,11 @@
     get_latent_space,
     resolve_latent_space,
 )
+from invokeai.backend.architectures.facets.loader import (
+    DEFAULT_LOADER_FLAGS,
+    LoaderFlagsFacet,
+    get_loader_flags,
+)
 from invokeai.backend.architectures.facets.unet import UNetDownscaleFacet, get_max_unet_downscale
 from invokeai.backend.architectures.facets.variant import (
     VariantFacet,
@@ -61,11 +66,13 @@
 )
 
 __all__ = [
+    "DEFAULT_LOADER_FLAGS",
     "ArchitectureError",
     "ConditioningFacet",
     "Facet",
     "LatentSpace",
     "LatentSpaceFacet",
+    "LoaderFlagsFacet",
     "UNetDownscaleFacet",
     "VariantFacet",
     "conditioning_infos",
@@ -76,6 +83,7 @@
     "get",
     "get_conditioning_info",
     "get_latent_space",
+    "get_loader_flags",
     "get_max_unet_downscale",
     "get_variant_enum",
     "register",

invokeai/backend/architectures/defs/z_image.py

@@ -1,5 +1,6 @@
 from invokeai.backend.architectures.facets.conditioning import ConditioningFacet
 from invokeai.backend.architectures.facets.latent_space import FLUX_16, LatentSpaceFacet
+from invokeai.backend.architectures.facets.loader import LoaderFlagsFacet
 from invokeai.backend.architectures.facets.variant import VariantFacet
 from invokeai.backend.architectures.registry import register
 from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, ZImageVariantType
@@ -10,6 +11,9 @@
     ConditioningFacet(ZImageConditioningInfo),
     # Z-Image uses a FLUX-compatible VAE with 16 latent channels.
     LatentSpaceFacet(FLUX_16),
+    # Diffusers' layerwise casting hits a dtype mismatch here: skipped modules produce bf16 while
+    # hooked modules expect fp16.
+    LoaderFlagsFacet(supports_fp8_storage=False),
     VariantFacet(
         {
             ModelType.Main: ZImageVariantType,

invokeai/backend/architectures/facets/loader.py

@@ -0,0 +1,49 @@
+"""Per-architecture exceptions to the generic model-loading policy.
+
+Generic loading code should not know architecture names. `ModelLoader._should_use_fp8` did: it
+carried a `config.base == BaseModelType.ZImage` branch in the middle of otherwise architecture-blind
+logic, behind a function-local import.
+
+The flags here are scalars an architecture *declares*, not implementations it *provides*. Choosing a
+different loader or a different conversion function per architecture is dispatch, and belongs in the
+`ModelLoaderRegistry` where it already lives -- not here. `model_loaders/lora.py` is the clearest
+example of that distinction: its base chain picks a different state-dict conversion per
+architecture, so it stays where it is.
+
+Optional by design, with defaults describing the ordinary case. An architecture that needs no
+exception declares nothing, and `get_loader_flags()` hands back the defaults -- the Null Object of
+`.ideas/ArchitectureSpec.md` §3, which is what keeps callers free of `is None` checks.
+"""
+
+from dataclasses import dataclass
+
+from invokeai.backend.architectures.facet import Facet
+from invokeai.backend.architectures.registry import get
+from invokeai.backend.model_manager.taxonomy import BaseModelType
+
+
+@dataclass(frozen=True)
+class LoaderFlagsFacet(Facet):
+    """Architecture-level answers the generic loader needs, where the general rule does not hold."""
+
+    supports_fp8_storage: bool = True
+    """Whether fp8 layerwise casting may be applied to this architecture's weights.
+
+    The other exclusions in `_should_use_fp8` are keyed on model *type* (VAEs degrade in decode,
+    LoRAs are patched in rather than run) or on submodel type, and stay generic -- they hold for
+    every architecture.
+    """
+
+
+DEFAULT_LOADER_FLAGS = LoaderFlagsFacet()
+"""What an architecture gets when it declares nothing: the general rule, unqualified."""
+
+
+def get_loader_flags(base: BaseModelType) -> LoaderFlagsFacet:
+    """The loader flags for `base`, or the defaults.
+
+    Also returns the defaults for `Any`, `External` and `Unknown`, which are never registered --
+    matching the previous behaviour, where only an exact match on Z-Image changed anything.
+    """
+    facet = get(base, LoaderFlagsFacet)
+    return facet if facet is not None else DEFAULT_LOADER_FLAGS

invokeai/backend/model_manager/load/load_default.py

@@ -11,6 +11,7 @@
 import torch
 
 from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.architectures import get_loader_flags
 from invokeai.backend.model_manager.configs.base import Diffusers_Config_Base
 from invokeai.backend.model_manager.configs.factory import AnyModelConfig
 from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase
@@ -25,6 +26,7 @@
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.model_manager.taxonomy import (
     AnyModel,
+    ModelType,
     SubModelType,
 )
 from invokeai.backend.util.devices import TorchDevice
@@ -234,21 +236,18 @@ def _should_use_fp8(self, config: AnyModelConfig, submodel_type: Optional[SubMod
         if self._torch_device.type != "cuda":
             return False
 
-        # Z-Image has dtype mismatch issues with diffusers' layerwise casting
-        # (skipped modules produce bf16, hooked modules expect fp16).
-        from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
-
-        if hasattr(config, "base") and config.base == BaseModelType.ZImage:
+        # Some architectures cannot take fp8 storage at all; they say so themselves.
+        if not get_loader_flags(config.base).supports_fp8_storage:
             return False
 
         # VAEs are excluded — fp8 storage causes noticeable quality degradation in decode.
-        if hasattr(config, "type") and config.type == ModelType.VAE:
+        if config.type == ModelType.VAE:
             return False
 
         # LoRAs (including ControlLoRA) are excluded — they are not run as a standalone forward pass,
         # they are patched into a base model, so the layerwise-casting hooks would never fire. The
         # toggle is also hidden in the UI for ControlLoRA; this guard handles legacy persisted values.
-        if hasattr(config, "type") and config.type in (ModelType.LoRA, ModelType.ControlLoRa):
+        if config.type in (ModelType.LoRA, ModelType.ControlLoRa):
             return False
 
         # Don't apply FP8 to text encoders, tokenizers, schedulers, VAEs, etc.

tests/backend/architectures/test_loader_flags.py

@@ -0,0 +1,81 @@
+"""`_should_use_fp8` no longer knows any architecture by name.
+
+It carried `config.base == BaseModelType.ZImage` in the middle of otherwise architecture-blind
+logic, behind a function-local import of a module the file already imports at module level.
+
+An audit of the whole loading path found that branch to be the only per-architecture *policy* in
+generic loader code. The other base-keyed code there is dispatch -- `model_loaders/lora.py` picks a
+different state-dict conversion per architecture, `model_loader_registry.py` builds a lookup key --
+and dispatch is not a flag, so it stays where it is. These tests pin that boundary.
+"""
+
+import pytest
+
+from invokeai.backend.architectures import (
+    DEFAULT_LOADER_FLAGS,
+    LoaderFlagsFacet,
+    generative_bases,
+    get,
+    get_loader_flags,
+)
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig  # noqa: F401  (registers every config class)
+from invokeai.backend.model_manager.taxonomy import BaseModelType
+
+NO_FP8_STORAGE = {BaseModelType.ZImage}
+"""The architectures that opt out. Kept as a set so adding one is a visible change here too."""
+
+
+def test_z_image_declares_no_fp8_storage() -> None:
+    assert get_loader_flags(BaseModelType.ZImage).supports_fp8_storage is False
+
+
+@pytest.mark.parametrize(
+    "base", sorted(set(generative_bases()) - NO_FP8_STORAGE, key=lambda b: b.value), ids=lambda b: b.value
+)
+def test_every_other_architecture_gets_the_general_rule(base: BaseModelType) -> None:
+    assert get_loader_flags(base).supports_fp8_storage is True
+
+
+@pytest.mark.parametrize("base", [BaseModelType.Any, BaseModelType.External, BaseModelType.Unknown])
+def test_sentinel_bases_get_the_defaults_rather_than_raising(base: BaseModelType) -> None:
+    """Models with `base=Any` -- CLIP embedders, T5 encoders -- reach `_should_use_fp8` too.
+
+    They are never registered, so the accessor must fall through to the defaults instead of raising,
+    exactly as the old `config.base == ZImage` comparison simply evaluated to False for them.
+    """
+    assert get_loader_flags(base) is DEFAULT_LOADER_FLAGS
+
+
+def test_only_the_opting_out_architectures_declare_the_facet() -> None:
+    """The facet is optional: declaring nothing means the general rule applies.
+
+    An architecture declaring `LoaderFlagsFacet()` with all defaults would be indistinguishable in
+    behaviour but misleading to read, so it should not happen.
+    """
+    declaring = {base for base in generative_bases() if get(base, LoaderFlagsFacet) is not None}
+
+    assert declaring == NO_FP8_STORAGE
+
+
+def test_the_defaults_describe_the_ordinary_case() -> None:
+    assert DEFAULT_LOADER_FLAGS.supports_fp8_storage is True
+    assert LoaderFlagsFacet() == DEFAULT_LOADER_FLAGS
+
+
+def test_every_config_declares_base_and_type() -> None:
+    """What made the `hasattr(config, "base")` / `hasattr(config, "type")` guards dead code.
+
+    `Config_Base.__pydantic_init_subclass__` refuses to create a concrete config class that does not
+    declare `type`, `base` and `format` with defaults, so no instance reaching `_should_use_fp8` can
+    lack them. Removing the guards is safe only while this holds, so it is asserted rather than
+    assumed.
+    """
+    missing = sorted(
+        f"{cls.__name__}.{field}"
+        for cls in Config_Base.CONFIG_CLASSES
+        for field in ("base", "type", "format")
+        if field not in cls.model_fields
+    )
+
+    assert missing == []