Krea2 Block Swap

Krea 2 is a 12B diffusion transformer, and the BF16 weights alone take 24.76GB. If you're on an 8GB or 12GB card, the stock experience is not great. People on r/StableDiffusion describe the FP8 version running fine and then, the moment they load a LoRA or try the full-precision weights, the run "gets stuck in infinite memory swapping" and a single image takes eleven minutes. ComfyUI's automatic offloading keeps things technically alive, but it hammers the PCIe bus shuffling the whole model around.

Krea2 Block Swap is the surgical fix. It sits between your Krea 2 loader and the KSampler and does exactly one job: park the first N transformer blocks in system RAM while they're idle, and only move them to the GPU for the instant each one actually runs.

The trick that makes it free

Most offload approaches park blocks on the CPU and leave them there, which either runs them on the CPU (very slow) or blows up with device-mismatch errors when a layer expects its weights on the GPU. This node does something simpler and smarter. Krea 2's SingleStreamDiT keeps its stack of transformer blocks in diffusion_model.blocks, and the node registers a forward pre-hook on each swappable block that moves it to the GPU right before its forward pass, plus a forward hook that moves it straight back to CPU the moment the pass finishes.

Every block still executes on the GPU, with its original weights, in their original order. Moving a tensor between devices doesn't change its values, so the output is bit-identical to a run with no swapping at all. The only thing you pay is PCIe transfer time - one CPU↔GPU round trip per swapped block per step. That's the entire tradeoff, and it's why the README can promise "without any quality loss" without hedging.

The inputs, all three of them

This is a genuinely thin node, which is part of its charm:

• model - your loaded Krea 2 model, straight from the UNETLoader.
• blocks_to_swap (0–48, default 8) - how many blocks, counting from block 0, live in RAM while idle. This is the only dial you'll touch. Start at 8 and raise it until the run fits; each extra block frees memory but adds another transfer per step.
• use_non_blocking (default on) - asynchronous CPU↔GPU transfers. Leave it on. Turn it off only if you hit genuine instability.

The single output is a model, which wires straight into your KSampler. Drop the node between loader and sampler and you're done - it doesn't touch the Qwen3-VL text encoder or the VAE at all.

Installing it

No dependencies, no model downloads, nothing heavy - the repo's requirements file is empty, just torch and what ComfyUI already ships. Either search "ComfyUI-Krea2BlockSwap" in ComfyUI Manager, or:

cd ComfyUI/custom_nodes
git clone https://github.com/olob0/ComfyUI-Krea2BlockSwap

Then restart ComfyUI. It's GPL-3.0, so if you redistribute it the license tag follows.

Where people get burned

• Over-swapping won't hurt quality, but it will slow you down. There's no correctness cliff, just a PCIe toll bridge - swap too many blocks and every step crawls through round trips. The node logs how many MB it keeps out of VRAM, so watch the console and back off if steps get sluggish.
• If it's still slow after swapping, you're at the memory wall, not the node. On 8GB the FP8 Krea 2 already renders a 1024×1024 in around ten seconds on a fast card, and the community fix for the thrashing reports is the INT8 variant on a current ComfyUI. This node complements those - it stops the uncontrolled swapping, not the lack of VRAM.
• It works with any quantization - int8, fp8, nvfp4, mxfp8 - because it moves whatever modules are loaded regardless of precision.
• If a future ComfyUI release renames the .blocks attribute, the node errors out and the fix is one line in _find_blocks. That's basically the whole failure surface.

If you're on a 12GB card running FP8 comfortably, you may not need this at all. If you're under 12GB - or the moment you add a LoRA and the sampler starts gasping - this is the difference between "technically runs" and "actually usable."