Block Swap Krea2 (RAM Offload)

This node is the creative option. BlockSwapKrea2 applies the same RAM-offload scheme as the rest of this pack - first N transformer blocks parked in system RAM, streamed to the GPU per forward pass, LoRAs baked into the CPU masters once - but rebuilt around module hooks so it survives the Ostris edit node that breaks the plain forward-wrapper approach.

Krea 2 is the single-stream DiT that took over the open image ecosystem in 2026: 12.9 B params, 28 blocks, and a bf16 checkpoint (krea2_turbo_bf16, ~25.8 GB) that plainly does not fit 24 GB. fp8 squeezes it to 12.9 GB and fits fine - but if you want the full-precision model, or you're running the fp8 plus a style LoRA, this is how you stay on the better weights instead of the smaller ones.

Why it's hook-based, and why that matters

The plain BlockSwap node wraps each block's forward. That's fine for Wan, and it completely breaks with ComfyUI-Krea2-Ostris-Edit: that node's Krea2OstrisEditModelPatch replaces diffusion_model.forward and calls each block's submodules directly (block.mod, block.attn, block.mlp), never block.forward. A forward-wrapper swap would silently leave the swapped weights on the CPU, and the reference-latent paths would crash with a device mismatch.

So BlockSwapKrea2 triggers the swap with PyTorch hooks instead. Every Krea 2 block execution path - native SingleStreamBlock.forward, the Ostris per-span reference forward, and the Ostris kv_cache forward - starts by calling block.mod and ends with block.mlp, so a forward_pre_hook on mod streams the block's weights to the GPU and a forward_hook on mlp repoints them back at the CPU masters. All three paths covered, including LoRA + multi-reference-image edit workflows.

The payoff is verified in the README: krea2_turbo_fp8 with 14 blocks swapped runs bit-exact against no-swap runs of the same seed - max pixel diff 0 across all three paths - while using 6.4 GB resident instead of 12.9 GB.

Inputs and output

Three inputs, all shared with the pack: model (MODEL, from UNETLoader, placed after the LoRA loader and the Ostris patch), blocks_to_swap (default 14 - note the lower default, Krea 2 only has 28 blocks total), and pin_memory (default true, page-locks the CPU copies). Output is a single MODEL. Chain:

UNETLoader → Krea2OstrisEditModelPatch → LoraLoaderModelOnly → BlockSwapKrea2 → KSampler

Installing

Same pack as the whole family - Manager → search "ComfyUI-JITBlockSwap", or git clone https://github.com/lovemachine100/ComfyUI-JITBlockSwap into custom_nodes/, then restart. No dependencies, no model downloads, no API.

Where people get burned

--disable-dynamic-vram is required, same as every node in this pack - without it the node silently passes the model through and only logs dynamic VRAM patcher detected, skipping. to the console.

bf16 + a runtime LoRA on Windows wants --disable-pinned-memory too. This was the pack's root-cause marathon, and it's worth knowing even if you're not on Windows. ComfyUI's mmap weight loader makes checkpoint tensors frombuffer views over file-backed pages; registering those as pinned host memory can queue asynchronous CUDA errors. Symptoms: silent H2D transfer corruption (psychedelic weight garbage in outputs) and a deferred CUDA error: invalid argument at the next load, VAE decode, or /free. Small pin footprints (the fp8 case, 5.8 GB) survive by luck; 11.6 GB of bf16 swap masters plus ComfyUI's own lowvram pins do not. The verified workaround is --disable-dynamic-vram --disable-pinned-memory together: clean styled output, stable across reloads, 12.3 s/it at 1024² - versus 8.5 s/it pinned right before it poisons itself, or ~4.5 min/it with ComfyUI's own lowvram path without this node. Unpinned costs about +5% wall clock on PCIe 4.0 x16.

fp8 + runtime LoRA is a separate trap: with --disable-dynamic-vram, ComfyUI's legacy loader merges the LoRA into fp8 weights per-key on a VRAM-packed GPU and can crawl for hours before this node's hook ever runs. Pre-merge the LoRA (bf16 merge, then requantize to fp8) and drop the loader, or run that combo under dynamic VRAM where this node is a no-op. fp16/bf16 + LoRA are unaffected.

Worth it? If you're on the Ostris edit workflows that a lot of the Krea 2 scene has moved to, this is currently the only block-swap implementation that doesn't fall over the moment reference latents enter the graph.