fix(model-loaders): stop materializing scaled-fp8 checkpoints in float32

[Pfannkuchensack]

Summary

The Krea-2 and Z-Image loaders dequantize ComfyUI "scaled fp8" checkpoints with weight.float() * scale and leave the result in float32 until a much later cast to the compute dtype. That holds the entire model at 4 bytes per parameter during the cold load.

For a ~12.2 GB Krea-2 fp8 checkpoint (≈12.2B params at 1 byte each) that means:

Load stage	before	after
fp8 file read into RAM	~12 GB	~12 GB
after dequantization	~49 GB (float32)	~25 GB (bf16)
after the caller's later cast	~25 GB	— (already there)

The ~25 GB figure is corroborated by the official Krea-2-Turbo diffusers transformer, which is 25 GB on disk in bf16.

Both loaders now multiply in float32 for precision but store the compute dtype immediately, so the whole model is never materialized in float32. This is the same fix the FLUX.2 loader already carries — its comment documents the identical symptom ("~36GB vs ~17GB for a 9B model … was the dominant cold-load spike"). The Qwen-Image loader was already correct (it multiplies in compute_dtype).

_dequantize_scaled_fp8 gains a target-dtype parameter (defaulting to bfloat16), and the Krea-2 single-file loader resolves the compute dtype before calling it, so weights land in their final type directly instead of being cast twice.

Related Issues / Discussions

Reported by a user with 32 GB RAM (~22 GB free) loading a Krea-2 scaled-fp8 checkpoint: RAM filled up during load, before anything moved toward VRAM.

QA Instructions

Requires a ComfyUI "scaled fp8" checkpoint — one with <layer>.weight in float8_e4m3fn plus a <layer>.weight_scale sibling. Krea-2 and Z-Image single-file checkpoints from the usual sources qualify.

1. Watch process RSS (Task Manager / htop) while loading such a model on main. Peak RAM is ~4× the file size.
2. Repeat with this PR. The peak should be ~2× the file size, and the model should still load and generate identically.
3. Generate an image and compare against main — output is unchanged. The multiply still happens in float32; only the storage dtype of the result changed, and it was going to be cast to that dtype a few steps later anyway.
4. Regression check for the Z-Image path specifically, since its loader has its own copy of this logic.

Automated: pytest tests/backend/model_manager/load/ — 426 passed. Includes a new case asserting the result lands in the compute dtype rather than float32, which is the property that bounds the peak.

Not covered: the RAM peak itself is not asserted in a test; the guard is the dtype assertion. Measuring peak RSS in CI would be flaky.

Merge Plan

Nothing special. No schema change, no node versions affected — this only touches loader internals.

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, backend only
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR)

[joshistoast]

This would fix my issues with krea thus far

[lstein]

Adversarial review — I set out to prove this broken and couldn't. Approving.

What I attacked, and why each attack failed:

• Dtype resolved too late / wrong value. All three call sites compute model_dtype before the dequant call; moving the resolution above load_file in the Krea-2 single-file loader has no ordering side effects (choose_torch_device is config-based).
• Downstream code assuming float32. Everything between the dequant and the old cast point is dtype-agnostic: _is_native_krea2_format / _convert_krea2_native_to_diffusers only rename keys and reshape (dtype-preserving, no arithmetic); the Z-Image variant detection reads only shapes; the encoder key remap and metadata stripping are key-only.
• Numerical divergence. The old per-key .to(model_dtype) loops still exist and are now no-ops for the dequantized weights, so the same float32 product gets the same deterministic cast, just earlier — outputs are bit-identical to main.
• fp16 overflow on bf16-unsafe devices (where choose_bfloat16_safe_dtype returns float16): any value that overflows now overflowed identically at the later cast on main.
• fp8 detection ordering. The Qwen3-VL encoder's source_is_fp8 check still runs before dequantization, and the requantize-to-fp8 path receives the same values as before.
• Missed sibling sites. Swept the repo for weight_scale handling: flux.py already carries the fix, qwen_image.py is confirmed already correct (multiplies in compute_dtype directly), the Z-Image encoder is fixed here, ideogram4 keeps fp8 resident by design, and the Z-Image transformer loader has no dequant logic at all.
• Edge cases — scale-without-weight keys, non-string keys, GGUF paths, make_room accounting (now more accurate), and the TYPE_CHECKING torch import (string annotation, never evaluated; no top-level torch existed) — all fine.

pytest tests/backend/model_manager/load/ passes locally (240 passed, 86 GPU-gated skips). The RAM math in the description also checks out: peak drops from fp8-file + full-float32 copy to fp8-file + bf16 copy + one transient float32 tensor.

One non-blocking nit, fine to leave as-is: _dequantize_scaled_fp8's dtype=None default silently means bfloat16. Every current caller passes the dtype explicitly, but a future caller on an fp16-only device that omits it would land weights in bf16 and pick up an extra rounding step on the later fp16 cast. If you want to harden it, making the parameter required (or asserting it) would close that door; not worth a new round on its own.