Summary

Loading a custom Krea 2 safetensors from the custom models folder (the ComfyUI diffusion_models file, usually fp8 or fp8_scaled, 13 GB) could take the whole app down (#14474). Studio already supports these files: the family is detected from the header, the DiT comes from the file and the text encoder / VAE from krea/Krea-2-Turbo. What kills the process is host RAM:

• A plain fp8 file goes through load_original_layout_transformer, which widens every fp8 tensor to the compute dtype in host RAM (value.to(want)). That is about 23 to 26 GB of new memory for the 13B Krea 2 DiT.
• An fp8_scaled ComfyUI file goes through load_comfy_quant_transformer. On a GPU with no fp8 runtime (anything below sm_89, such as the reporter's RTX 3060 Ti), every layer is dequantized to bf16 the same way.
• Nothing checks host RAM on this branch (the small-host route only runs for full pipelines). The load runs in a thread of the server process, so running out of memory is not an exception it can catch: Linux's OOM killer ends the server, and on Windows the commit limit runs out inside native code.

This PR prices the new host memory from the safetensors header before any widening. Tensors already at the compute dtype stay memory-mapped and cost nothing. When the widened weights plus the small-host reserve do not fit in free RAM, the load is refused with a readable error. This applies to both loaders and to every family that uses them:

Not enough system RAM to load k4_fp8.safetensors: widening its weights to float32 needs about 11.9 GB of free host memory (8.9 GB of weights + 3.0 GB reserve) and 2.4 GB is free. Close other applications, use a GGUF or a smaller quant of this model, or set UNSLOTH_DIFFUSION_HOST_RAM_CHECK=0 to try anyway.

UNSLOTH_DIFFUSION_HOST_RAM_CHECK=0 (the existing small-host switch) skips the check, and loads under 1 GiB of widening are never checked.

Evidence

• Five tiny Krea 2 single files (plain bf16, fp8, fp8_scaled with scale_weight / scaled_fp8, model.diffusion_model. prefix, diffusers layout) all load through DiffusionBackend.load_pipeline from a custom folder on main. The load path itself is fine.
• I used a realistic file: real Krea 2 shapes, 4 of the 28 blocks, fp8 linears, 2.9 GB. Widening it to bf16 added 3318 MiB of private RAM, against 3312 MiB predicted from the header.
• I loaded that same file through the real backend on CPU inside a 3 GB memory cgroup:
  • main: the process is SIGKILLed (exit 137) with no Python error.
  • This branch: the readable refusal above.
  • With 20 GB: it loads normally.
• Tests, main vs this branch:
  • New tests/test_diffusion_single_file_host_ram.py. Four tests fail on main and pass here: the fp8 file is refused before any weight is read, the fp8_scaled file with no fp8 runtime is refused, a bf16 file is never refused, and the widened byte count matches. Two pass on both: the same files load when there is room, and the kill switch / unreadable RAM never refuse.
  • Related files: test_diffusion_comfy_quant.py, test_diffusion_krea2.py, test_diffusion_single_file_krea2_hunyuanimage.py, test_diffusion_single_file_registry.py, test_diffusion_small_host.py, test_hub_safetensors_single_file.py, test_video_comfy_quant.py. No change between base and head.

Review follow-ups

The header estimate alone kept missing per-layer choices the loader makes later (layers a runtime filter rejects, fp16 deferred dequantization, the packed NVFP4 width, FNUZ float8, nibble-swapped NVFP4 copies, key-map row reorders, tiled scale buffers). So the header estimate is now only the first gate, refusing before any read. During the load, every helper that allocates host memory (_dequant, dequant_block, swap_nibbles, untile_scales, tile_scales, the dtype casts in both loaders, the segment join, the key-map torch.cat) checks its allocation against the host's live available memory (cgroup-capped, about 0.5 ms per reading) minus the reserve, and the load is refused before the allocation that would not fit. Because it reads live memory rather than keeping a running total, freed buffers count as free again and a copy made where nothing checks still shows up in the next reading. Files under 1 GiB are not checked, and UNSLOTH_DIFFUSION_HOST_RAM_CHECK=0 installs no check at all.

Real-size check: a 13.2 GiB Krea 2 fp8_scaled ComfyUI file (rebuilt from the bf16 transformer) loading at bf16 on CPU, with a simulated 64 GiB host whose free memory falls as the process's anonymous memory grows (measured peak 23.26 GiB, reserve 6.4 GiB):

Each review item has a test that fails before its fix and passes after, including a 48-layer NVFP4 fallback load that fits and was refused by the earlier running-total version.

Not covered

• Not run on Windows. The check reads free physical RAM, not the available commit (RAM plus page file). A machine with a large page file that used to finish such a load slowly is now refused, and the message names the switch that allows it.
• Only the DiT is priced. The text encoder (4.8 to 8.9 GB) and the VAE load afterwards, so a load that just passes can still run short.
• ComfyUI loads these files without widening: the fp8 weights stay fp8 in memory and are upcast per layer at compute time. Doing the same in Studio (layerwise casting on the denoiser on GPUs without an fp8 runtime) is the real fix for low-RAM machines. It touches offload, compile and LoRA, so it is left as a follow-up.
• Krea 2 offers only Create (txt2img) in Studio by design today: diffusers ships no Krea 2 img2img / inpaint pipeline. ComfyUI does img2img / inpaint for Krea 2 through its generic latent denoise, so that is a follow-up, not a detection bug. The Hub models that offer more modes are FLUX.1 Krea dev, a different (FLUX.1) family.
• A kept NVFP4 load (FlashInfer runtime on Blackwell) was covered by unit tests only, not measured on a real file.