Musubi UNET JustInTime Loader

MusubiUNETJitLoader is the interesting half of the polym/ComfyUI-musubi pack, because it doesn't need a musubi model at all. Feed it any diffusion model file and it converts the weights to musubi block-scaled fp8 - right there, on your machine, at load time. The "Jit" is just-in-time: the quantization happens while the node runs, instead of being a file you had to download pre-quantized. You can even bake a LoRA in at the same time. It's the least-fiddly path I know to get a regular fp16 model down to fp8-class VRAM usage with per-block scaling, and it's a nice trick: one node replacing what would otherwise be a loader plus a LoraLoader plus a separate quantization step.

How it works

The loader reads the state dict from your model file, optionally merges a LoRA's weights into it, then runs optimize_state_dict_with_fp8 over every tensor. Each weight gets split into blocks of block_size (64 by default) and replaced by an fp8 version plus a per-block scale_weight - the same scheme musubi-tuner produces when you train with it. The node then builds a ModelPatcher from that, so the rest of your graph sees a completely ordinary MODEL.

Two details worth knowing. First, the LoRA merge happens before quantization, so the LoRA's effect is baked into the fp8 weights - no separate LoRA node needed downstream, but also no LoRA strength dial after the fact. Set it and commit. Second, this does the conversion in RAM, watching a progress bar crawl through every layer, and it doesn't write the result back to disk - every cold start re-does the whole pass.

The inputs that matter

• unet_name - unlike its sibling MusubiUNETLoader, this lists every model in diffusion_models, no "musubi" filename filter. Any standard safetensors works.
• weight_dtype - fp8_e4m3fn is the default choice for diffusion (4 exponent, 3 mantissa, most precision); fp8_e5m2 trades precision for range; fp8_e4m3fnuz is the niche zero-biased variant. When in doubt, e4m3fn.
• block_size - default 64. Per-block scale factors; smaller blocks quantize more tightly at the cost of more scale overhead. Leave it alone unless you have a reason.
• lora_name - set to "none" to skip, or pick a LoRA from your loras folder to merge it into the weights on load.

Output is a single MODEL, straight into the KSampler - exactly like the built-in UNETLoader's, so swapping is a one-wire change.

Installing it

ComfyUI Manager: search "ComfyUI-musubi" and install. Or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/polym/ComfyUI-musubi

Restart ComfyUI and you're done - there are no pip dependencies in the pack's pyproject.toml, so nothing else to install. You'll want a recent PyTorch build, though (see below), and your model file stays in ComfyUI/models/diffusion_models.

Where people get burned

• Slow first load, and every load after it. Quantizing the whole state dict takes a while and nothing is cached. A 12B model is a legit several-minutes wait with a progress bar. This is the honest price of "just in time": you pay the conversion cost in exchange for never downloading a special model file.
• torch has no attribute float8_e4m3fn. The node calls the fp8 dtypes directly, so your PyTorch needs to be recent enough (2.2+) to have them. Old installs fail instantly with an AttributeError that looks like a broken model but is really a broken environment.
• block_size must divide a layer's input width. When it doesn't, the code logs a warning and silently falls back to per-channel quantization for that layer. Harmless, but if you're chasing quality and changed block_size, the first thing to check is those console warnings.
• It's not for pre-quantized musubi files. If the file already carries scale_weight tensors, use MusubiUNETLoader instead - that one expects them, this one creates them. The two nodes are two halves of the same pack and they don't swap cleanly.