feat(fp8): enable FP8 storage for Z-Image

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,17 +464,22 @@ 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
         call sites so a new caller cannot forget it.
         """
         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:

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
 
 

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 && <DefaultCfgRescaleMultiplier control={control} name="cfgRescaleMultiplier" />}
         <DefaultWidth control={control} optimalDimension={optimalDimension} />
         <DefaultHeight control={control} optimalDimension={optimalDimension} />
-        {!isZImage && <DefaultFp8Storage control={control} name="fp8Storage" />}
+        <DefaultFp8Storage control={control} name="fp8Storage" />
       </SimpleGrid>
     </>
   );

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