From 8912f070d09b821d594a0b0a85b043a95c4b1046 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 2 Aug 2026 20:21:14 +0200 Subject: [PATCH] feat(architectures): move the fp8-storage exclusion into a loader-flags facet `ModelLoader._should_use_fp8` knew one architecture by name. In the middle of otherwise architecture-blind logic sat from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType if hasattr(config, "base") and config.base == BaseModelType.ZImage: return False Z-Image now declares that itself, as LoaderFlagsFacet(supports_fp8_storage=False), and the loader asks the registry. The facet is optional and its defaults describe the ordinary case, so an architecture that needs no exception declares nothing and the accessor hands back the defaults. That also covers models with base=Any -- CLIP embedders, T5 encoders -- which reach this function too and are never registered: they get the defaults rather than an error, exactly as the old equality comparison simply evaluated to False for them. Two pieces of dead code go with it. The function-local import was importing from a module load_default.py already imports at module level (taxonomy, line 26), so it was incidental rather than cycle-avoidance. And the three hasattr guards cannot fire: Config_Base.__pydantic_init_subclass__ refuses to create a concrete config class that does not declare type, base and format with defaults. That invariant is now asserted in a test rather than assumed, since removing the guards depends on it. An audit of the loading path answers the question the plan left open in section 7 -- whether other base-specific loader special cases exist. They do not. The only other base-keyed code in the load package is dispatch or key construction: model_loaders/lora.py picks a different state-dict conversion per architecture, model_loader_registry.py builds a lookup key, and krea2.py's fp8 call keys off the checkpoint's own dtype rather than the base. Dispatch is not a flag, so none of it belongs here. LoaderFlagsFacet therefore has exactly one field and one declaration, and the tests pin that boundary. Known interaction with two in-flight upstream PRs, flagged deliberately. invoke-ai/InvokeAI#9414 removes this exclusion rather than relocating it: the root cause was fixed inside #8945 and the branch is obsolete, replaced by an extra_skip_patterns mechanism reading the model's own _skip_layerwise_casting_patterns. #9415 gives Anima the same declaration. Both touch _should_use_fp8, so this conflicts textually, and once they land supports_fp8_storage has no remaining declaration. Landing this first is a deliberate choice; see the PR description. openapi.json is unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- invokeai/backend/architectures/__init__.py | 8 ++ .../backend/architectures/defs/z_image.py | 4 + .../backend/architectures/facets/loader.py | 49 +++++++++++ .../model_manager/load/load_default.py | 13 ++- .../architectures/test_loader_flags.py | 81 +++++++++++++++++++ 5 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 invokeai/backend/architectures/facets/loader.py create mode 100644 tests/backend/architectures/test_loader_flags.py diff --git a/invokeai/backend/architectures/__init__.py b/invokeai/backend/architectures/__init__.py index 6547ccdaaed..ff67c60b845 100644 --- a/invokeai/backend/architectures/__init__.py +++ b/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", diff --git a/invokeai/backend/architectures/defs/z_image.py b/invokeai/backend/architectures/defs/z_image.py index 4e14f59f8bc..018842c67e1 100644 --- a/invokeai/backend/architectures/defs/z_image.py +++ b/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, diff --git a/invokeai/backend/architectures/facets/loader.py b/invokeai/backend/architectures/facets/loader.py new file mode 100644 index 00000000000..5bd5d31c378 --- /dev/null +++ b/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 diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index ba617b2b55a..e1ff07df1ba 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/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. diff --git a/tests/backend/architectures/test_loader_flags.py b/tests/backend/architectures/test_loader_flags.py new file mode 100644 index 00000000000..7c7cc358d8e --- /dev/null +++ b/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 == []