UniBlockSwap
Squeezing 9B diffusion models onto a 6GB card
- model
- model
You downloaded Klein 9B, or Bernini, or Minimax, watched the VRAM bar slam into the ceiling, and ComfyUI died mid-sample. That's the exact problem UniBlockSwap exists for. It's a passthrough node from smthemex that sits between your model loader and the KSampler and swaps the model's transformer blocks between VRAM and system RAM one at a time, so a DiT that "needs" 20GB+ runs on 4-6GB. The README's numbers are the whole pitch: Klein 9B at Q8 in 4.8GB, Bernini int4 with LoRAs and 120 frames in 9-10GB, Minimax in 4.5GB.
What it is and where it fits
It's a model/loaders node for ComfyUI's native workflow - no custom sampler, no API, no key. ComfyUI already has its own low-vram offloading and GGUF lazy dequantization, and for most models that's plenty. The gap it fills is the newest giant DiTs (Flux 2 Klein, the Wan-derived stuff like Bernini, Krea, Minimax), whose transformer stacks are simply bigger than your whole card. Block swap is the sharper lever: instead of shuffling whole layers around, keep a handful of blocks pinned in VRAM and let the rest load on demand. Think of it as ComfyUI's built-in offloading turned up a notch, aimed at models that arrived after that built-in path stopped being enough.
How it works
Wire it in and the node clones your model, then finds the block containers (names like double_blocks, single_blocks, transformer_blocks - whatever the model actually ships) and wraps them in a swap-aware module list. The num_blocks value sets a resident prefix: blocks 0..N-1 are pushed into VRAM once and stay there for the whole inference, while the tail blocks stay on the model's normal lazy path - quantized GGUF blocks get dequantized on the fly, safetensor blocks are pulled in when the loop reaches them. When inference ends, everything is released.
The clever bit is the release: blocks are pointed back at their original CPU/mmap data rather than copied around, so the swap is mostly RAM-backed, not a reallocation festival. It works with both safetensor and GGUF weights, and it keeps LoRA patches applied at cast time, which matters because this is a pack that literally fixed a "changing num_blocks silently killed my LoRA" bug.
The input that actually matters
The node takes model (MODEL) in and gives model out. The only knob is num_blocks (INT, default -1):
- -1 - most aggressive: a 1-block resident prefix. Least VRAM, slowest.
- 0 - swap off, pure passthrough.
- N - keep N leading blocks resident; more VRAM, less swapping.
- ≥ total blocks - swap off again (same as 0).
The output model feeds the KSampler. The pack's example workflows chain it as GGUF loader → LoRA → UniBlockSwap → UniBlockSwapCacheControl → KSampler; the CacheControl sibling clears the text encoder and VAE off VRAM every inference so they don't pile up between runs.
Installing it
ComfyUI Manager (search "ComfyUI_UniBlockSwap") or the README's way:
cd ComfyUI/custom_nodes
git clone https://github.com/smthemex/ComfyUI_UniBlockSwap
Then restart ComfyUI. Good news: there's no requirements.txt - it's pure Python on top of ComfyUI's own APIs, so no heavy pip dependencies to fight. Do update ComfyUI itself first, because the node leans on recent patcher callback APIs.
Common issues
- The num_blocks trap. Crank it high thinking "more resident = better" and you've silently disabled the node once it hits the total block count. Start at -1 and dial up only if you've got headroom.
- System RAM is the real cost. The swap parks blocks in RAM. On 16GB of system memory, aggressive model swap plus TE swap gets tight fast.
- Official template batch size. The README warns that unpacking official templates leaves the batch size pointing at the wrong value - set it to 1.
- Spilling into shared memory. With GGUF + very large resolutions or long videos, the author warns you can hit shared GPU memory and slow to a crawl. It'll run; it just won't be worth it.
- Speed. Swapping is slower than owning a bigger card. This node's job is making a workflow run, not making it fast - the first generation also warms up.
If you have an old clone, update it: the current resident-prefix design replaced an older sliding-window approach, and the LoRA bug fixes came with it.
Inputs (2)
| Name | Type | Default | Description |
|---|---|---|---|
| model | MODEL | — | |
| num_blocksopt | INT | -1-1–10000 | 前缀常驻块数: 一次性把 block 0..N-1 推送进 CUDA, 常驻到本次推理结束才释放; 其余块按需逐块懒加载。-1 = 单块前缀常驻(最省显存); N = N 块前缀常驻; >= 总块数 = 不 swap 全部驻留 |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| model | MODEL | — |