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