Nodes/ComfyUI-Imatrix/Imatrix UNet Loader
ComfyUI Node

Imatrix UNet Loader

The model loader that's secretly measuring your Flux

By laksjdjf·Created about a year ago·Updated about a year ago· 6
Imatrix UNet Loader
    • MODEL
    unet_name
    weight_dtype

    On the surface, Imatrix UNet Loader looks like any other model loader: pick a UNet, get a MODEL out, move on. The name is the tell. This one is also listening. Every time the loaded model runs a Linear or Conv2d layer, the node records the mean-squared magnitude of the activations per input channel and folds it into a running average. That quiet tally is the raw material for the rest of this pack: the imatrix file that GGUF quantizers use to decide where to spend their scarce bits.

    Why would you want that? Because this pack exists for the bottom of the GGUF ladder. Q8 and Q6 are nearly free - the KB's consensus is that Q8 is "99% identical to fp16" - but the moment you drop to Q2_K or Q3_K to squeeze Flux-dev onto a midrange card, quality falls apart. A good imatrix is one of the few things that claws some of it back, because it tells the quantizer which channels genuinely matter and deserve precision. LLMs get their imatrix from calibration text corpora. Image models don't ship one, and nobody's about to calibrate your workflow for you. So this node turns your own generations into the calibration set.

    How it works

    The node loads the UNet through a custom operation class (ImatrixOps) that wraps every Linear and Conv2d with a recorder. On each forward pass it squares the activations, averages them over the batch and spatial dims, and updates a running average that gets more stable the more you run. It's deliberately dumb and it's the right kind of dumb: you just run your normal workflow and the model quietly builds the statistics.

    Inputs are refreshingly few:

    • unet_name - any model in your diffusion_models folder.
    • weight_dtype - default, or one of the fp8 options (fp8_e4m3fn, fp8_e4m3fn_fast, fp8_e5m2). Loading fp8 mimics how you'll actually run the quant and halves VRAM while you collect data.

    Output is a single MODEL - wire it exactly where you'd wire any UNet, into your KSampler, LoRA stack, the works.

    How to install

    ComfyUI Manager, search ComfyUI-Imatrix, or:

    cd ComfyUI/custom_nodes
    git clone https://github.com/laksjdjf/ComfyUI-Imatrix
    # restart ComfyUI
    

    One wrinkle: the pack has no requirements.txt, yet its code imports the gguf Python package at load time. If you already run ComfyUI-GGUF you're fine; otherwise pip install gguf first or the node errors on import.

    Using it in the pipeline

    This is step one of a two-step job. Load through ImatrixUNETLoader, generate a batch of images, then export the stats with Save Imatrix and feed the .dat into llama-quantize --imatrix when you build your low-bit GGUF. Keep the model loaded for the whole session - the stats live in the model object, so closing the workflow throws them away.

    Where people get burned: the running average mixes everything you run, so don't collect data at one resolution then save after a different-resolution run. And remember it's a recorder, not a magic quality button - you still have to actually do the quantization step with the file it produces.

    Categoryimatrix

    Inputs (2)

    NameTypeDefaultDescription
    unet_nameCOMBO0 options:
    weight_dtypeCOMBO4 options: default, fp8_e4m3fn, fp8_e4m3fn_fast, fp8_e5m2

    Outputs (1)

    NameTypeDescription
    MODELMODEL