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

invokeai/backend/model_manager/load/model_loaders/krea2.py

@@ -2,7 +2,7 @@
 """Class for Krea-2 model loading in InvokeAI."""
 
 from pathlib import Path
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import accelerate
 from transformers import AutoConfig, AutoTokenizer
@@ -27,6 +27,10 @@
 from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
 from invokeai.backend.util.devices import TorchDevice
 
+if TYPE_CHECKING:
+    # torch is imported lazily inside the helpers below; this is annotations-only.
+    import torch
+
 
 def _normalize_qwen3vl_rope_config(config: Any) -> Any:
     """Mirror Qwen3-VL rope_parameters into rope_scaling for Transformers compatibility."""
@@ -89,25 +93,33 @@ def _is_native_krea2_format(sd: dict[str, Any]) -> bool:
     )
 
 
-def _dequantize_scaled_fp8(sd: dict[str, Any]) -> dict[str, Any]:
+def _dequantize_scaled_fp8(sd: dict[str, Any], dtype: "torch.dtype | None" = None) -> dict[str, Any]:
     """Dequantize ComfyUI 'scaled fp8' weights: ``dequant = weight.float() * weight_scale``.
 
     Each quantized layer stores an fp8 ``<name>.weight`` plus a (usually scalar) ``<name>.weight_scale``.
-    Returns a new dict with the weights dequantized to float and the ``.weight_scale`` keys removed.
-    No-op if there are no scale keys.
+    Returns a new dict with the weights dequantized and the ``.weight_scale`` keys removed. No-op if
+    there are no scale keys.
+
+    The multiply runs in float32 for precision, but each result is stored as ``dtype`` immediately so
+    the *whole model* is never materialized in float32. Krea-2's ~12 GB fp8 checkpoint would otherwise
+    peak at ~50 GB of RAM (4 bytes/param) before the caller's later bf16 cast brings it down to ~25 GB,
+    which puts a 32 GB machine into swap during a cold load. This mirrors the same fix already applied
+    to the FLUX.2 loader.
     """
     import torch
 
     scale_keys = [k for k in sd if isinstance(k, str) and k.endswith(".weight_scale")]
     if not scale_keys:
         return sd
+    dtype = dtype or torch.bfloat16
     out = dict(sd)
     for scale_key in scale_keys:
         weight_key = scale_key.replace(".weight_scale", ".weight")
         if weight_key in out:
             weight = torch.as_tensor(_to_plain_tensor(out[weight_key])).float()
             scale = torch.as_tensor(_to_plain_tensor(out[scale_key])).float()
-            out[weight_key] = weight * scale
+            out[weight_key] = (weight * scale).to(dtype)
+            del weight
         del out[scale_key]
     return out
 
@@ -331,17 +343,19 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel:
             raise TypeError(f"Expected Main_Checkpoint_Krea2_Config, got {type(config).__name__}.")
         model_path = Path(config.path)
 
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+
         sd = load_file(model_path)
         sd = _strip_comfyui_prefix(sd)
-        # ComfyUI 'scaled fp8' checkpoints: fold the per-tensor weight_scale into the weights (→ float).
-        sd = _dequantize_scaled_fp8(sd)
+        # ComfyUI 'scaled fp8' checkpoints: fold the per-tensor weight_scale into the weights. The
+        # compute dtype is resolved first so the dequantized weights land there directly instead of
+        # transiently materializing the whole model in float32.
+        sd = _dequantize_scaled_fp8(sd, model_dtype)
         # Native/ComfyUI key naming → diffusers Krea2Transformer2DModel keys.
         if _is_native_krea2_format(sd):
             sd = _convert_krea2_native_to_diffusers(sd)
 
-        target_device = TorchDevice.choose_torch_device()
-        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
-
         with accelerate.init_empty_weights():
             model = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG)
 
@@ -568,7 +582,7 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod
             getattr(t, "dtype", None) in (torch.float8_e4m3fn, torch.float8_e5m2) for t in sd.values()
         )
         # ComfyUI 'scaled fp8': fold weight_scale into the weights, then drop quantization metadata.
-        sd = _dequantize_scaled_fp8(sd)
+        sd = _dequantize_scaled_fp8(sd, model_dtype)
         for k in list(sd.keys()):
             if isinstance(k, str) and (k.endswith(".comfy_quant") or "scale_input" in k):
                 del sd[k]

invokeai/backend/model_manager/load/model_loaders/z_image.py

@@ -620,7 +620,13 @@ def _load_from_singlefile(
                             if block_size > 1:
                                 # Repeat scale along this dimension to match weight shape
                                 scale = scale.repeat_interleave(block_size, dim=dim)
-                sd[weight_key] = weight_float * scale
+                # Multiply in float32 for precision, but store the compute dtype immediately so the
+                # *whole model* is never materialized in float32. Keeping every dequantized weight as
+                # float32 until the caller's later cast quadruples the per-parameter cost (4 bytes vs
+                # 1 on disk) and dominates the cold-load RAM peak — enough to swap a 32 GB machine.
+                # Same fix as in the FLUX.2 and Krea-2 loaders.
+                sd[weight_key] = (weight_float * scale).to(model_dtype)
+                del weight_float
                 dequantized_count += 1
 
         if dequantized_count > 0:

tests/backend/model_manager/load/test_krea2_state_dict_utils.py

@@ -92,7 +92,21 @@ def test_folds_scale_into_weight_and_drops_scale_key(self) -> None:
         }
         out = _dequantize_scaled_fp8(sd)
         assert "layer.weight_scale" not in out
-        assert torch.allclose(out["layer.weight"], torch.tensor([1.0, 2.0]))
+        assert torch.allclose(out["layer.weight"].float(), torch.tensor([1.0, 2.0]))
+
+    def test_result_is_stored_in_the_compute_dtype_not_float32(self) -> None:
+        """The whole model must never be materialized in float32.
+
+        The multiply runs in float32 for precision, but holding every dequantized weight there
+        costs 4 bytes per parameter: Krea-2's ~12 GB fp8 checkpoint peaked at ~50 GB of RAM before
+        the caller's later bf16 cast, which swaps a 32 GB machine during a cold load.
+        """
+        sd = {
+            "layer.weight": torch.tensor([2.0, 4.0]),
+            "layer.weight_scale": torch.tensor(0.5),
+        }
+        assert _dequantize_scaled_fp8(dict(sd))["layer.weight"].dtype is torch.bfloat16
+        assert _dequantize_scaled_fp8(dict(sd), torch.float16)["layer.weight"].dtype is torch.float16
 
     def test_noop_without_scale_keys(self) -> None:
         sd = {"layer.weight": torch.tensor([2.0, 4.0])}