fix(fp8): resolve compute dtype instead of reading model.dtype

[Pfannkuchensack]

Summary

SDXL with fp8_storage crashed before the UNet was ever called:

NotImplementedError: "pow_cuda" not implemented for 'Float8_e4m3fn'

After layerwise casting the UNet's weights are float8_e4m3fn, and diffusers derives model.dtype from the first parameter — so unet.dtype reports a storage-only dtype. The legacy SD/SDXL denoise path used it for every tensor it built, so the latents were created in float8 and the first bit of scheduler math (sigma ** 2 in add_noise) blew up. torch has no arithmetic kernels for float8; it is only valid for weights that the forward hooks cast up per layer.

Add get_model_compute_dtype(): returns model.dtype for normal models and the compute dtype for fp8 ones. The loader records the compute dtype on the model when it applies the cast; if the marker is missing (older cache entry, Krea2 encoder path) the resolver scans for the first non-fp8 float param, which works because the cast skips norm layers.

Converted every site that derived a tensor dtype from a possibly-fp8 model: latents, noise, mask, masked_latents, conditioning, IP-Adapter and LoRA patch weights in denoise_latents and tiled_multi_diffusion_denoise_latents, plus the LoRA and T2I-Adapter extensions on the modular path. ControlNet and T2I-Adapter control images had the same latent bug — those configs expose an fp8_storage toggle too, so their control image would have been built in float8.

Also point LayerPatcher at the shared FP8_STORAGE_DTYPES constant.

Regression test covers the real loader path: model.dtype is float8 while the resolver returns fp16, and the resolved dtype survives the scheduler arithmetic that crashed.

Related Issues / Discussions

n/a

QA Instructions

Needs a CUDA GPU — _should_use_fp8 returns False on any other device, so the whole path is inert elsewhere.

Setup: Model Manager → an SDXL model → Default Settings → enable FP8 Storage → Save. The model cache invalidates the entry on save, so the next generation reloads with fp8 storage (watch for FP8 layerwise casting enabled for <name> (storage=float8_e4m3fn, compute=torch.float16, ...) in the log).

Before this PR, step 1 alone fails with NotImplementedError: "pow_cuda" not implemented for 'Float8_e4m3fn'.

1. txt2img on Canvas — generates, VRAM for the UNet is roughly halved vs. the same model without the toggle.
2. img2img — non-zero denoising start, i.e. the add_noise path from the traceback.
3. Inpaint/outpaint with a mask — exercises mask / masked_latents.
4. SDXL + LoRA — fp8 weights force sidecar patching; output should match the non-fp8 run closely (fp8 rounding aside), not be noise.
5. ControlNet, with FP8 Storage also enabled on the ControlNet model itself — this is the second, previously untriggered instance of the bug.
6. T2I-Adapter — same, control image path.
7. Tiled Multi-Diffusion node (e.g. a tiled upscale workflow) — the second denoise implementation that read unet.dtype.
8. Toggle FP8 Storage back off and repeat 1 — unchanged behaviour, get_model_compute_dtype() just returns model.dtype.

Non-SDXL sanity (these never read unet.dtype, but the shared LayerPatcher constant moved): FLUX txt2img with and without a LoRA, plus a Krea2 generation if you have the fp8 encoder installed.

Unit tests:

uv run --extra cuda --extra test pytest tests/backend/util/test_fp8.py tests/backend/model_manager/load -q --no-cov

429 passed, 129 skipped locally. The fp8-specific tests run on CPU (float8 conversion works there; only arithmetic doesn't).

Merge Plan

Plain merge — no DB schema, no redux slice, no API schema change. No migration needed: the compute dtype is recorded at load time, so it is always in sync with the code that reads it, and the resolver falls back to a param scan if the marker is ever absent.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, no redux changes
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR) — n/a

[lstein]

Adversarial review

I attacked every claim in the PR description against the actual code — most survived; a few hardening gaps below.

Claims I attacked and could not break

• Call-site completeness. Swept every model.dtype/unet.dtype/transformer.dtype read across invokeai/app/invocations/ and invokeai/backend/stable_diffusion/. All remaining hits are on models that can never be fp8-cast: Z-Image is explicitly excluded in _should_use_fp8, Anima's loader (model_loaders/anima.py) never calls _apply_fp8_layerwise_casting, and FLUX's image_encoder_model is a CLIP-vision model with no fp8_storage setting. The extension-based ControlNet/inpaint paths derive from ctx.latents.dtype, which inherits correctness from the fixed latent creation. Coverage is genuinely complete.
• The deepcopy-survival claim. Verified: _build_meta_shell deepcopies the module — __dict__ and the marker come along — and the shell is built after _load_model applies the cast, so adopted models on a second GPU carry both hooks and marker.
• No stale-marker risk. The marker lives on the model instance and is set exactly when the cast is applied, never persisted. Toggling fp8 off yields a fresh un-cast, un-marked model.
• Fallback-scan soundness. _apply_fp8_to_nn_module skips norm/embedding layers, so the first non-fp8 float param does reveal the compute dtype, as claimed.
• LayerPatcher change is a pure constant dedup; behavior identical.

Findings

1. Nothing prevents a poisoned marker (hardening — requesting this one before merge). _apply_fp8_layerwise_casting derives the compute dtype from the first parameter's dtype and is not idempotent. If it is ever called on an already-cast model (a future loader refactor, a cache re-put), the first param is already fp8, so it would record float8 as the compute dtype — and get_model_compute_dtype's isinstance(marked, torch.dtype) check would happily return it, silently reintroducing the exact crash this PR fixes. Two one-liners close the class: assert compute_dtype not in FP8_STORAGE_DTYPES in set_fp8_compute_dtype, and/or early-return in _apply_fp8_layerwise_casting when the marker is already present.

2. Marker-setting is duplicated at both _apply_fp8_to_nn_module call sites (design). load_default.py and krea2's _load_text_encoder must each remember to call set_fp8_compute_dtype after casting; a third caller will forget — the precise failure mode the fallback scan exists to paper over. Setting the marker inside _apply_fp8_to_nn_module itself removes the footgun. test_falls_back_to_scan_when_marker_missing currently locks in the split behavior; it could simulate a legacy model with delattr instead.

3. Last-resort fallback can silently pick the wrong dtype (minor). If every float param is fp8, get_model_compute_dtype returns TorchDevice.choose_torch_dtype() — fp16 on most CUDA devices. For a bf16-compute model that would produce a confusing dtype-mismatch deep in the forward pass. A warning log at that branch would make the failure diagnosable. Low likelihood given the skip list.

4. Vendored HiDiffusion pipeline retains the same bug (note only). invokeai/backend/hidiffusion/hidiffusion.py reads controlnet.dtype at four sites (525, 541, 561, 577). It's dead code today — InvokeAI imports only apply_hidiffusion/remove_hidiffusion and the threshold dicts, never the pipeline __call__ — but if that pipeline is ever wired up with an fp8 ControlNet, it reproduces this crash. Worth a follow-up or a comment.

Test coverage

• The regression test is well-aimed: it exercises the real loader cast path, asserts the naive model.dtype read is fp8 (the precondition that crashed), and confirms the resolved dtype survives the actual scheduler arithmetic. The marker-missing fallback is also covered.
• _fp8_supported() (hasattr(torch, "float8_e4m3fn")) is vacuously true on any supported torch — the skips will never fire; harmless.
• ModelLoader.__new__ + hand-set private attrs is brittle against loader-constructor refactors, but pragmatic for a unit test.
• No invocation-level test for the ControlNet/T2I control-image sites — acceptable given they'd need real models; the QA plan covers them manually.

Verdict

The core fix is correct and the call-site sweep is complete — I could not find a missed live path. Finding 1 (guard against recording an fp8 dtype as the compute dtype) is cheap insurance against silently reintroducing the bug and is the one change I'd ask for before merge; 2–4 are follow-up material.