Nodes/ComfyUI-musubi/Musubi UNET Loader
ComfyUI Node

Musubi UNET Loader

The loader that finally runs musubi block-scaled fp8 models in ComfyUI

By polym·Created 11 months ago·Updated 10 months ago· 3
Musubi UNET Loader
    • MODEL
    unet_name
    weight_dtype
    block_size64

    If you've trained a LoRA with kohya's musubi-tuner, you already know the end of this story: the output is a model quantized with block-wise scaling, and the stock UNETLoader will not touch it. That's the whole reason MusubiUNETLoader exists. It's a drop-in replacement for ComfyUI's built-in loader, minus the "can't read this file" error, for models that ship with per-block fp8 scale factors baked in.

    What it's for

    Musubi block-wise scaling is the quantization scheme behind the musubi-tuner ecosystem: weights are stored as fp8 but split into blocks of (by default) 64 channels, each block carrying its own float32 scale factor instead of one scale for the whole tensor. That buys most of fp8's memory savings with noticeably less of its quality hit - per-block scaling is a coarser, cheaper cousin of per-channel quantization. This node is the ComfyUI inference side of that scheme. The unet_name dropdown doesn't even list every model in your diffusion_models folder - it filters to files whose name contains "musubi", because this loader expects weights that already have those scale_weight tensors inside them.

    The pack is a small, mostly-undiscussed one from polym (the same account hosting polym/COMFY_MODELS on Hugging Face). No drama there - it just does one job, and you reach for it only when you're holding a musubi-quantized file.

    How it works

    Under the hood it's ComfyUI's own load_diffusion_model called with a custom_operations override (block_scaled_fp8_ops). That override swaps every Linear in the model for one that carries an extra scale_weight parameter - one scale per block - and dequantizes weight-by-scale right before each matmul during sampling. So the fp8 tensors stay compact in VRAM, and the scaling math happens on the fly. The model file itself never gets rewritten.

    The inputs that matter

    There are only three, and you'll touch two of them:

    • unet_name - the musubi model file from ComfyUI/models/diffusion_models. Only filenames containing "musubi" show up.
    • weight_dtype - fp8_e4m3fn (the standard for diffusion, most precision), fp8_e5m2 (more range, less precision), or fp8_e4m3fnuz (a niche zero-biased variant). Don't overthink this: if the model came with a stated dtype, match it; otherwise e4m3fn is the one to reach for.
    • block_size - default 64. Bigger blocks mean fewer scale factors, less overhead and a touch less accuracy; smaller means more scales, tighter quantization. 64 is fine unless a file's README says otherwise.

    It outputs a single MODEL that wires straight into the KSampler like any loader output. If your workflow already uses a normal UNETLoader, swap this node in and nothing downstream changes.

    Installing it

    Via ComfyUI Manager, search "ComfyUI-musubi" and install. Or the manual way:

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

    Then restart ComfyUI. No pip install, no extra Python dependencies - the pack's pyproject.toml lists none, it only needs what ComfyUI already ships. What you do need is the model: grab the musubi-quantized checkpoint from https://huggingface.co/polym/COMFY_MODELS and drop it in ComfyUI/models/diffusion_models.

    Where people get burned

    • Empty dropdown. If nothing appears under unet_name, the "musubi" filename filter is eating your file. Either rename the file to include "musubi" or you've placed it in the wrong folder - diffusion_models, not checkpoints.
    • AttributeError: 'torch' has no attribute 'float8_e5m2' (or similar). The nodes call the float8 dtypes directly, so you need a PyTorch build with fp8 support - roughly 2.2 or newer. On an old environment nothing in this pack will run, and the error is easy to mistake for a broken model.
    • Don't use this loader for normal models. It's built for pre-quantized musubi files. If you feed it a plain fp16 checkpoint, the scale tensors won't be there and the result is undefined - for converting a regular model to block-scaled fp8 yourself, use the pack's other node, MusubiUNETJitLoader.
    Categorymusubi

    Inputs (3)

    NameTypeDefaultDescription
    unet_nameCOMBO0 options:
    weight_dtypeCOMBO3 options: fp8_e5m2, fp8_e4m3fn, fp8_e4m3fnuz
    block_sizeoptINT64

    Outputs (1)

    NameTypeDescription
    MODELMODEL