diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py
index d646d4466d3..5d547193d6f 100644
--- a/invokeai/backend/model_manager/load/load_default.py
+++ b/invokeai/backend/model_manager/load/load_default.py
@@ -344,12 +344,7 @@ def get_size_fs(
def _should_use_fp8(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> bool:
"""Check if FP8 layerwise casting should be applied to a model."""
- # 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:
- return False
+ from invokeai.backend.model_manager.taxonomy import ModelType
# VAEs are excluded — fp8 storage causes noticeable quality degradation in decode.
if hasattr(config, "type") and config.type == ModelType.VAE:
@@ -430,7 +425,19 @@ def _apply_fp8_layerwise_casting(
# `register_forward_hook` path fires around `nn.Module._call_impl` without replacing
# `forward`, so `CustomLinear.forward` is still reached.
if isinstance(model, torch.nn.Module):
- self._apply_fp8_to_nn_module(model, storage_dtype=storage_dtype, compute_dtype=compute_dtype)
+ # Diffusers models declare their own precision-sensitive modules in
+ # `_skip_layerwise_casting_patterns`, and `enable_layerwise_casting()` honors them. Since
+ # we no longer call it, we have to apply that list ourselves — it is not cosmetic. Z-Image's
+ # `TimestepEmbedder.forward` reads `self.mlp[0].weight.dtype` and casts its *input* to it;
+ # with an fp8 weight the input becomes float8 before our pre-hook can restore the weight,
+ # and `F.linear` dies with `"addmm_cuda" not implemented for 'Float8_e4m3fn'`. Hence
+ # `['t_embedder', 'cap_embedder']` for that model.
+ self._apply_fp8_to_nn_module(
+ model,
+ storage_dtype=storage_dtype,
+ compute_dtype=compute_dtype,
+ extra_skip_patterns=tuple(getattr(model, "_skip_layerwise_casting_patterns", None) or ()),
+ )
else:
return model
@@ -443,7 +450,12 @@ def _apply_fp8_layerwise_casting(
return model
@staticmethod
- def _apply_fp8_to_nn_module(model: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype) -> None:
+ def _apply_fp8_to_nn_module(
+ model: torch.nn.Module,
+ storage_dtype: torch.dtype,
+ compute_dtype: torch.dtype,
+ extra_skip_patterns: tuple[str, ...] = (),
+ ) -> None:
"""Apply FP8 layerwise casting to a plain nn.Module.
Mirrors diffusers' `apply_layerwise_casting` semantics: only the layer classes in
@@ -452,6 +464,10 @@ def _apply_fp8_to_nn_module(model: torch.nn.Module, storage_dtype: torch.dtype,
Without the skip list, precision-sensitive tiny learned scalars (e.g. FLUX RMSNorm.scale)
get crushed to FP8 and quality degrades noticeably.
+ `extra_skip_patterns` carries the model's own declared exclusions (diffusers'
+ `_skip_layerwise_casting_patterns`), which are model-specific and cannot be inferred from
+ layer types or generic name patterns.
+
Records the compute dtype on the model. After the cast, `model.dtype` reports the float8
storage dtype, which must never be used to create or cast tensors — torch has no arithmetic
kernels for it (see `get_model_compute_dtype`). The marker is set here rather than at the
@@ -459,10 +475,11 @@ def _apply_fp8_to_nn_module(model: torch.nn.Module, storage_dtype: torch.dtype,
"""
set_fp8_compute_dtype(model, compute_dtype)
+ skip_patterns = _FP8_DEFAULT_SKIP_PATTERNS + tuple(extra_skip_patterns)
for module_name, module in model.named_modules():
if not isinstance(module, _FP8_SUPPORTED_PYTORCH_LAYERS):
continue
- if any(re.search(pattern, module_name) for pattern in _FP8_DEFAULT_SKIP_PATTERNS):
+ if any(re.search(pattern, module_name) for pattern in skip_patterns):
continue
params = list(module.parameters(recurse=False))
if not params:
diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py
index 534db4e8bd2..0ca5dfb010a 100644
--- a/invokeai/backend/model_manager/load/model_loaders/z_image.py
+++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py
@@ -482,6 +482,11 @@ def _load_from_singlefile(
sd[k] = sd[k].to(model_dtype)
model.load_state_dict(sd, assign=True)
+
+ # Every param is uniform `model_dtype` at this point (the loop above casts the whole state
+ # dict, including ComfyUI fp8 checkpoints, whose scale metadata was filtered out above), so
+ # the layerwise cast has a single unambiguous compute dtype to restore to.
+ model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer)
return model
diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx
index b69bf117bc6..91d255f5a49 100644
--- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx
+++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx
@@ -56,10 +56,6 @@ export const MainModelDefaultSettings = memo(({ modelConfig }: Props) => {
return ['flux', 'flux2'].includes(modelConfig.base);
}, [modelConfig]);
- const isZImage = useMemo(() => {
- return modelConfig.base === 'z-image';
- }, [modelConfig]);
-
const defaultSettingsDefaults = useMainModelDefaultSettings(modelConfig);
const optimalDimension = useMemo(() => {
const modelBase = modelConfig?.base;
@@ -148,7 +144,7 @@ export const MainModelDefaultSettings = memo(({ modelConfig }: Props) => {
{!isFluxFamily && }
- {!isZImage && }
+
);
diff --git a/tests/backend/model_manager/load/test_load_default_fp8.py b/tests/backend/model_manager/load/test_load_default_fp8.py
index 7909b5620e5..17b96c3ad93 100644
--- a/tests/backend/model_manager/load/test_load_default_fp8.py
+++ b/tests/backend/model_manager/load/test_load_default_fp8.py
@@ -62,7 +62,8 @@ def _make_config(model_type: ModelType, fp8: bool, base: BaseModelType = BaseMod
[
(_make_config(ModelType.VAE, fp8=True), None),
(_make_config(ModelType.LoRA, fp8=True), None),
- (_make_config(ModelType.Main, fp8=True, base=BaseModelType.ZImage), None),
+ # Z-Image used to be listed here. It is no longer excluded — see
+ # `test_should_use_fp8_allows_z_image` for why the exclusion became obsolete.
(_make_config(ModelType.Main, fp8=True), SubModelType.Tokenizer),
(_make_config(ModelType.Main, fp8=False), None),
],
@@ -313,6 +314,72 @@ def __init__(self):
assert model.rms.scale.dtype == compute_dtype
+def test_apply_fp8_to_nn_module_honors_extra_skip_patterns():
+ """A model's own `_skip_layerwise_casting_patterns` must be applied on top of our defaults."""
+
+ class _Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.t_embedder = torch.nn.Linear(4, 4)
+ self.attn = torch.nn.Linear(4, 4)
+
+ storage_dtype = torch.float16
+ compute_dtype = torch.float32
+ model = _Model()
+ for p in model.parameters():
+ p.data = p.data.to(compute_dtype)
+
+ ModelLoader._apply_fp8_to_nn_module(
+ model, storage_dtype, compute_dtype, extra_skip_patterns=("t_embedder", "cap_embedder")
+ )
+
+ assert model.attn.weight.dtype == storage_dtype
+ assert model.t_embedder.weight.dtype == compute_dtype
+
+
+def test_apply_fp8_layerwise_casting_passes_model_declared_skip_patterns():
+ """Regression test for Z-Image + fp8 crashing with
+ `RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn'`.
+
+ Diffusers models declare precision-sensitive modules in `_skip_layerwise_casting_patterns`, and
+ `enable_layerwise_casting()` honors them. Our hook-based replacement must read that list too —
+ it is not redundant with `_FP8_DEFAULT_SKIP_PATTERNS`. `ZImageTransformer2DModel` declares
+ `['t_embedder', 'cap_embedder']` because `TimestepEmbedder.forward` reads
+ `self.mlp[0].weight.dtype` and casts its *input* to it: with an fp8 weight the input becomes
+ float8 before the pre-hook restores the weight, and `F.linear` has no float8 kernel.
+ """
+
+ class _FakeZImage(torch.nn.Module):
+ _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]
+
+ def __init__(self):
+ super().__init__()
+ self.t_embedder = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.Linear(4, 4))
+ self.cap_embedder = torch.nn.Linear(4, 4)
+ self.layers = torch.nn.Linear(4, 4)
+
+ loader = _make_loader(device="cuda")
+ model = _FakeZImage().to(torch.bfloat16)
+
+ with patch.object(ModelLoader, "_should_use_fp8", return_value=True):
+ loader._apply_fp8_layerwise_casting(model, _make_config(ModelType.Main, fp8=True, base=BaseModelType.ZImage))
+
+ # The declared modules keep their compute dtype...
+ assert model.t_embedder[0].weight.dtype == torch.bfloat16
+ assert model.cap_embedder.weight.dtype == torch.bfloat16
+ # ...while everything else is stored in fp8, so the toggle still saves VRAM.
+ assert model.layers.weight.dtype == torch.float8_e4m3fn
+
+
+def test_should_use_fp8_allows_z_image():
+ """Z-Image was excluded while we used diffusers' `enable_layerwise_casting()` with the global
+ torch dtype (fp16) as compute dtype, which clashed with the model's bf16 weights. The compute
+ dtype now comes from the model itself, so the exclusion is obsolete.
+ """
+ loader = _make_loader(device="cuda")
+ assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=True, base=BaseModelType.ZImage)) is True
+
+
def test_wrap_forward_reaches_custom_linear_after_apply_custom_layers():
"""Production order: `_load_model` applies FP8 wrapping, THEN `ModelCache.put()` calls
`apply_custom_layers_to_model` which constructs a NEW `CustomLinear` object via