XB-BOX - UNet Block Swap

Big transformer-based diffusion models have a rude habit of being a few hundred MB too big for your VRAM. XB_UNetBlockSwap is the XB_ToolBox answer: it physically moves the first N blocks of a UNet/DiT model into system RAM at load time, then lets ComfyUI drag each block back onto the GPU only when a sampling step actually needs it. The result is a model that fits where it didn't before - at the cost of speed, because offloaded blocks take a trip across the PCIe bus on every pass.

This is the pack's "static physical block swap," and it's aimed at the same audience as everything in VRAM_Hacks: consumer cards running 14B+ video models (Wan, LTX, and friends) that would otherwise OOM the instant sampling starts. If your model loads but dies on the first step, this is a candidate fix.

Inputs

• unet_model - the model, and the tooltip is specific: "Accepts pure UNet/DiT model input only." Wire it right after a UNet/DiT loader, before any checkpoint-level wrapper. It's typed * (anything), but if you feed it a full checkpoint patcher, the node simply passes it through unmodified.
• blocks_to_swap - 0 to 200, default 10. Zero means "node off, pure passthrough." Higher = more VRAM saved, slower sampling. The source even clamps it if you ask for more blocks than the model actually has, so it won't error on an oversized number.

Output is the unet_model back out, ready for your sampler.

The mechanism, briefly

The node clones the model and registers an ON_LOAD callback. When ComfyUI loads the model, the callback walks the architecture looking for block lists (transformer_blocks, double_blocks, single_blocks, down_blocks, up_blocks, etc.), then moves blocks 0..N to the offload device (system RAM) and the rest to VRAM. It grabs the actual GPU device from the model patcher rather than hardcoding cuda:0 - a nice touch for multi-GPU setups.

The two guardrails

• Blacklist. Lumina, ZImage, HunyuanDiT, and ErnieImage models are refused - the pack determined their architecture breaks under physical block movement and logs a "拦截盾" (intercept shield) skip message instead of corrupting your run.
• --enable-dynamic-vram. If you launch ComfyUI with that flag, this node detects it and goes inert rather than double-offloading.

Where people get burned

The "is it a checkpoint or a UNet?" question is the number one confusion. Feed it a full checkpoint and it silently does nothing - no error, just no help. And don't expect it to be fast: the honest trade is explicitly time for space, and the pack's own framing is "trading time for space." It shines on the video models it's tuned for; for ordinary SDXL-style work, ComfyUI's built-in lowvram handling is usually the better (and smarter) tool.

Install: Manager → XB_ToolBox, or git clone https://github.com/wjluoxiao/XB_ToolBox.git into custom_nodes, restart. It's pure ComfyUI-patching - no extra pip dependencies.