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 == []