kongo_junandClaude Fable 5
ea28c6eb69
Avoid collision with existing *Blockswap custom-node repos (wanBlockswap, HidreamBlockswap, hvBlockswap, NativeBlockSwap). Node class name (BlockSwap) and display name are unchanged so existing workflows keep working. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ComfyUI-JITBlockSwap
Block Swap (RAM Offload) node for native ComfyUI MODEL — run DiT models
larger than VRAM (e.g. Wan 2.2 / Bernini-R 14B fp16, 28.6 GB on a 24 GB GPU)
by keeping the first N transformer blocks in system RAM and streaming each
block to the GPU only while its forward runs.
Node
advanced/model → Block Swap (RAM Offload) (BlockSwap)
| input | default | meaning |
|---|---|---|
model |
— | MODEL from UNETLoader (place AFTER LoRA loaders) |
blocks_to_swap |
20 | blocks kept in RAM (Wan 14B has 40). Raised automatically if the resident part still exceeds the VRAM weight budget |
pin_memory |
true | page-lock CPU copies for fast PCIe transfer |
How it works
- Hooks
ModelPatcherON_LOAD: after ComfyUI's normal load pass it re-organizes the diffusion model — swap blocks get their weight patches (LoRA) baked in once, moved to (pinned) RAM; everything else is made fully GPU-resident, removing the slow per-layer LowVramPatch cast path. - Each swapped block's
forwardis wrapped: params are repointed to a GPU copy just-in-time and repointed back to the CPU master afterwards (no device-to-host copy needed — weights never change during inference). Transfers are synchronized before compute: async submission of large pinned H2D copies interleaved with DiT kernels busy-loop hangs the GPU on Windows/WDDM (observed: 100% util at ~100 W). Cost is ~0 when compute-bound (measured 28-29 s/it either way on Wan 14B). ON_DETACHrestores state; placement is re-applied on every load, so model switching (high/low noise) and LoRA changes are safe.
Notes / limitations
- Chain:
UNETLoader → LoraLoaderModelOnly → BlockSwap → ModelSamplingSD3 → KSampler. - Requires
--disable-dynamic-vram(legacy ModelPatcher). With dynamic VRAM enabled the node logs a warning and does nothing. - Works on any model whose diffusion model exposes
.blocks(Wan, LTX-style DiTs, Flux double blocks are NOT covered — Wan-family tested). - Do not chain two BlockSwap nodes on the same model.
- Approx. cost: one PCIe H2D transfer per swapped block per forward call (~25-30 ms per 660 MB fp16 block on PCIe 4.0 x16).