Model Block Restore

MOSS-Image SFT draws faint vertical banding across skies and smooth backgrounds. The fix is not a prompt tweak or a denoise setting - it's swapping the last ten transformer blocks back to the Base weights while keeping the fine-tune's character everywhere else. That's what Model Block Restore does, and it does it without touching your LoRAs.

What it's for

Block-level surgery has been a thing since the early days: block merge, block swap, "these layers are where the model went wrong". It was always a tool for people editing checkpoint files. This is the live-graph version - you point a node at a second checkpoint, name a range of blocks, and sample a hybrid, no merging on disk, no 12 GB file to clean up afterwards.

The concrete, documented use is the MOSS-Image one. Restore blocks 32 to 41 from moss_image_base_int8_convrot under the SFT model and the banding goes away while the SFT's look stays. The general use is broader: a fine-tune that broke one behaviour (hands, text, skies, colour cast) can often lend you the good blocks from its own base, because the base and the fine-tune share a parameter layout.

How it works

It walks the model's transformer, finds its block list (one of layers, blocks, transformer_blocks, joint_blocks, and it insists on exactly one), and for each block in your range reads the matching tensors out of the source checkpoint with safe_open - so the file is read, not loaded into VRAM. Then each weight becomes a straight interpolation:

(1 - strength) * own + strength * source

Those are applied ahead of the model's other patches, which is why LoRAs applied upstream stay applied on top of the restored blocks. If the source weights are quantized, they get dequantized to float32 first using ComfyUI's own ops, so you can restore from an int8 file into a bf16 model.

Before it touches anything it checks shapes and key names. A mismatch stops with a message naming the exact tensor.

The inputs

model takes the MODEL you want repaired - MOSS-Image SFT with its LoRAs already on, in the intended workflow.

checkpoint is a menu of safetensors in ComfyUI/models/diffusion_models. Both files have to be there; the node reads only the blocks you name, but it does need the file present on disk.

first_block and last_block are counted from 0. The defaults, 32 and 41, are the banding range on MOSS-Image SFT. last_block 41 is that model's final block.

strength dials the whole thing: 1.0 is the source checkpoint's blocks exactly, 0.5 is halfway, and 0 returns your model untouched (it short-circuits, so it costs nothing).

Output is a single patched model. Feed it to the sampler, or to whatever patch comes next.

Installing it

Manager, search WAS Node Suite v3, restart. Manually:

cd ComfyUI/custom_nodes
git clone https://github.com/WASasquatch/was-node-suite-comfyui.git

ComfyUI 0.14.0+, Python 3.10+, no pip step and no downloads - this node just reads files you already have. Then put both MOSS-Image files in models/diffusion_models: the SFT you're sampling and the Base you're restoring from, 5.4 GB each.

Where it goes wrong

holds no layers.32 weights. The checkpoint in the menu isn't the same model. You've picked another architecture or a different family - the message names the block it looked for and the fix is to point it at the Base of the same model.

Still banding. Either the range is wrong or the layer isn't there at all. first_block at 37 leaves part of the banding, because the problem lives across 32–41, not the top few. And if the node isn't in the graph at all, no amount of parameter tuning will help.

This model has blocks 0 to N. Self-explanatory once you see it: MOSS-Image ranges 0–41, and a request for block 60 can't be honoured. Check last_block isn't past the end.

Expecting a bigger visual change than you get. Restoring ten mid-to-late blocks is a targeted repair, not a re-style. If you want a different look, you want a LoRA or a different checkpoint. If you want the banding gone, this is the exact tool.