ComfyUI Node

Xiaomi-Lens-0 Loader

The 9 GB handshake you do before anything gets masked

By hiroki-abe-58·Created 23 days ago·Updated 23 days ago· 1
Xiaomi-Lens-0 Loader
    • model
    ◄model▾►
    ◄deviceauto►
    ◄attentionauto►
    ◄decoder_attentionauto►
    ◄after_runkeep on device►

    Most loaders in ComfyUI are boring on purpose. This one is too, but the thing it loads isn't: a 4B vision-language model that reads your sentence and returns masks, about 9 GB of bfloat16 weights in models/xiaomi_lens/. The Loader gets those weights into a usable state and hands you a XIAOMI_LENS_MODEL for one destination - the Segment node.

    If you've been getting text-described masks out of GroundingDINO → SAM, or SAM 3's text prompt, this is the same job by a different route: one model does the grounding and the mask in a single pass, and it can answer a request instead of just finding a noun. The cost is memory - a 4B VLM is a chunkier dependency than a detector plus a segmenter.

    What actually happens when you queue

    The weights are not downloaded by this node - the checkpoint opens with local_files_only=True, so an empty models/xiaomi_lens/Xiaomi-Lens-0/ gives you an error, not a quiet 9 GB download mid-queue.

    On execute it loads the checkpoint in bfloat16 onto the CPU and validates it. That's deliberate - the source carries the comment "fail here, not later, when a file is missing or broken" - so a busted download shows up as a red node at the top of the graph instead of halfway through a long run.

    The GPU part happens downstream: the Segment node moves the model onto the device, and first asks ComfyUI to free roughly the model size plus 2 GB so it can evict whatever image model was resident. On a 24 GB card that dance is the difference between running and not.

    The inputs you'll actually touch

    model is a dropdown of sub-folders in models/xiaomi_lens/, stricter than it looks: a folder only appears if it holds config.json, tokenizer.json, processor_config.json and at least one .safetensors shard. Half-finished download, no dropdown entry - that's your first diagnostic.

    device defaults to auto, meaning ComfyUI's GPU. cpu genuinely works - the author measured a whole run on 24 threads with about 10 GB of RAM - and it's the fallback when you don't have the VRAM. Slower, and slower still as the image grows.

    attention and decoder_attention are the two settings that will bite you. Both default to auto: FlashAttention-2 if flash-attn is installed and you're on CUDA, otherwise PyTorch SDPA. Leave them alone unless you're debugging. The one to understand is decoder_attention = mha, which is what upstream falls back to without flash-attn: it holds the whole patch-by-patch attention matrix, about 2 GB at a 2 MP image and over 30 GB at the 8 MP ceiling, where it ran out of memory on the author's 5090. SDPA and MHA produce identical boxes and masks agreeing at an IoU of 0.997 or better, so there is no reason to fight for MHA.

    after_run is the practical one: keep on device for fast repeat runs, move to CPU to hand roughly 9 GB of VRAM back after every run. If your workflow segments and then loads a big image model, set it to move to CPU and stop thinking about it.

    One output, model, typed XIAOMI_LENS_MODEL - it goes into the Segment node and nowhere else.

    Install

    ComfyUI v0.38.0 or newer (this pack is written against the V3 node API - no NODE_CLASS_MAPPINGS in the source, which is normal for a new pack, not a symptom), Python 3.10+, and Transformers 5.4.0 or newer.

    Search "Xiaomi-Lens-0 Segmentation for ComfyUI" in ComfyUI Manager, or do it by hand:

    cd ComfyUI
    git clone https://github.com/hiroki-abe-58/ComfyUI-XiaomiLens custom_nodes/ComfyUI-XiaomiLens
    python -m pip install -r custom_nodes/ComfyUI-XiaomiLens/requirements.txt
    

    That requirements file is one line - transformers>=5.4.0 - and it's the sharp edge. If you're still on Transformers 4.x, this upgrade is mandatory here and potentially fatal for other nodes pinned to 4.x. The README says so plainly, which is more honesty than most packs manage.

    Then fetch the weights, SHA-256 verified per file:

    python custom_nodes/ComfyUI-XiaomiLens/tools/download_model.py --models-dir models/xiaomi_lens
    

    You get models/xiaomi_lens/Xiaomi-Lens-0/ - 9.2 GB across three shards. Storing models elsewhere works through extra_model_paths.yaml with the key xiaomi_lens.

    Where it goes wrong

    Your VRAM budget is the headline number, because it's bigger than people assume for a "segmentation tool": peak usage measured 9.5 GB at 1 MP, 11 GB at 3 MP and 15 GB at 8 MP. Cards in the 8 GB class won't run this on the GPU; use device = cpu and accept the wait.

    If the dropdown is empty, run the download script with --verify-only - it re-hashes what's on disk and tells you which file is wrong.

    A RuntimeError naming Transformers 5.4.0 means your install didn't take; check you piped into ComfyUI's Python. A "no checkpoint named …" FileNotFoundError is the dropdown's strictness again: right name, failed file check.

    The author also ships ComfyUI-Win-Blackwell, a one-click Windows setup for 50-series cards - which is why this README's "verified on a 5090, here are the numbers" framing reads in character.

    Categoryloaders/xiaomi_lens

    Inputs (5)

    NameTypeDefaultDescription
    modelCOMBOa folder in models/xiaomi_lens/
    deviceCOMBOautoauto: ComfyUI's GPU. cpu: works, slower (bfloat16, about 10 GB of RAM)
    attentionCOMBOautoauto: FlashAttention-2 (what upstream uses) when flash-attn is installed and the model runs on CUDA, otherwise PyTorch SDPA
    decoder_attentionCOMBOautothe mask decoder's attention. auto: flash-attn (upstream's default) when installed and on CUDA, otherwise PyTorch SDPA. mha: upstream's fallback without flash-attn (identical results to upstream there), but its memory grows with the square of the image size: about 32 GB for an 8 MP image
    after_runCOMBOkeep on devicekeep on device: fast repeated runs. move to CPU: frees about 9 GB of VRAM after each run

    Outputs (1)

    NameTypeDescription
    modelXIAOMI_LENS_MODEL—