XB-BOX - Checkpoint Block Swap

The oldest trick in the VRAM book: if the model doesn't fit, don't load all of it. XB_CheckpointBlockSwap does exactly that - it takes a loaded model and, right as ComfyUI moves it onto the GPU, shuffles the first N transformer blocks plus the text/image embedding layers over to system RAM. The model stays addressable, but the GPU only holds part of it, trading a big chunk of VRAM for a bit of swap latency. That's the difference between "OOM at startup" and "runs slowly, but runs."

It's part of XB_ToolBox's VRAM_Hacks category, and it's squarely aimed at the pack's headline use case: running 14B–22B video models (Wan-class) on consumer cards. Think of it as a manual, coarse-grained version of ComfyUI's own low-VRAM offloading, one you can point at any checkpoint-shaped model.

How it works

The node attaches an ON_LOAD callback to the model patcher. When ComfyUI loads the model onto the GPU, the callback:

1. Scans the diffusion model for its block containers - it knows the common names (transformer_blocks, blocks, down_blocks, up_blocks, mid_block, layers, double_blocks, single_blocks, joint_blocks - it's written to handle both DiT and U-Net layouts).
2. Moves the first blocks_to_swap of them to the offload device (system RAM).
3. If offload_txt_emb / offload_img_emb are on, also ships the text/image embedding encoders to RAM.
4. Runs soft_empty_cache() + gc.collect() to actually free the VRAM.

A few built-in safeguards matter: if blocks_to_swap = 0 it passes the model through untouched; if ComfyUI's dynamic VRAM is already active it steps aside (no double-offloading); and if the model architecture is one it doesn't recognize as safely splittable, it logs an "拦截盾" (interception shield) message and skips rather than corrupting the model. use_non_blocking toggles non-blocking H2D copies on the embedding offload (can help on some AMD drivers).

The inputs that matter

• checkpoint_model - accepts any model (the tooltip literally says "including CLIP/VAE"). Wire your checkpoint, UNet, or model output here.
• blocks_to_swap - how many blocks to park in RAM (default 15). This is your main VRAM/trade knob: more swapped = less VRAM, slower.
• offload_img_emb / offload_txt_emb - also offload the image/text embedding layers (default on).
• use_non_blocking - non-blocking copies (default off).

Output: checkpoint_model - the patched model, same type, ready to feed your sampler.

Installing

Part of XB_ToolBox - ComfyUI Manager → XB_ToolBox, or

cd ComfyUI/custom_nodes
git clone https://github.com/WJLUOXIAO/XB_ToolBox.git

restart, zero extra deps.

Common issues

• No visible effect - if ComfyUI's --lowvram/dynamic mode is active, this node deliberately no-ops. Check the console for the skip message.
• Way too slow - too many blocks swapped. Block-swap is "time for space," and video models swap a lot of weights every step; start low (5–10) and only raise it until you stop OOMing.
• Model looks fine but output is garbage - this shouldn't happen (it's a pure device placement, no quantization), but if it does, you've likely hit an architecture it partially mis-detected; the interception shield logs which models it refuses.

Honest take: this is a blunt instrument next to proper offload and GGUF quantization, and for most users a Q4 GGUF of the same model is the cleaner answer. But it works on models that have no GGUF, and it's the one knob you can turn per-workflow without re-downloading anything.