Yogurt VITMatte Model Loader

Every prompt-based segmentation model has the same weakness: it finds the object, but the mask edge looks like it was cut with safety scissors. Hair, fur, whiskers, glass - anything with fine or semi-transparent edges comes back as a lumpy blob. Sa2VA is no exception. This node loads the model that fixes that: VITMatte, a trimap-based matting model from HUST that turns a rough mask into a real per-pixel alpha matte.

The name "Model Loader" undersells the point. This isn't just a load node - it's the reuse node. In the upstream ComfyUI-Sa2VA-XJ pack, VITMatte was loaded inside the V2 segmentation node, so every run paid the load cost. This fork pulled it out into its own loader that caches the model and hands you a reusable object. Run it once, wire the output to as many Yogurt Sa2VA Image Segmentation V2 nodes as you like, and nobody reloads anything.

What it loads and how

VITMatte (VitMatteForImageMatting from transformers) is a trimap-based matting model. It doesn't guess where the subject is; you give it a coarse segmentation plus a trimap - the ring of "unknown" pixels around the edge - and it predicts fractional alpha for just those pixels. That's the trick: by only solving the hard boundary, it spends its capacity on hair and fur instead of re-segmenting the whole image. The loader runs it through a VitMatteImageProcessor and registers the model with ComfyUI's model_management so it gets offloaded from VRAM when other nodes need room and pulled back when the V2 node runs.

Inputs that matter

• model_name - the matting model. Two real choices: hustvl/vitmatte-base-composition-1k and hustvl/vitmatte-small-composition-1k (the default). The small one is noticeably lighter and genuinely fine for hair edges; base buys a little more quality on gnarly semi-transparency. The dropdown scans ComfyUI/models/vitmatte for folders containing a config.json.
• device - auto (default), cuda, or cpu. Auto is right; cpu exists if you want to keep the 80MB of VRAM this thing uses off the GPU, and it's fast enough on stills.
• force_reload - unloads and rebuilds the cached model for this model/device combo. Only reach for it if you swap model files underneath a running ComfyUI.

Output is a single vitmatte_model (YOGURT_VITMATTE_MODEL), which plugs into the vitmatte_model input of Yogurt Sa2VA Image Segmentation V2. It has no other job.

Install

This is a shared-pack install - you get all five Yogurt Sa2VA nodes at once:

cd ComfyUI\custom_nodes
git clone https://github.com/yogurt7771/ComfyUI-YogurtSa2VA
cd ComfyUI-YogurtSa2VA
python -m pip install -r requirements.txt
python -m pip install opencv-python

opencv-python isn't in requirements.txt but the README calls it out separately: VITMatte and the morphology helpers depend on it. Grab the model with huggingface-cli download --local-dir, or just let transformers pull it from Hugging Face on first run - the loader falls back to that automatically if the folder isn't there.

Gotchas

The two that bite, both real and both in the README:

1. The dropdown doesn't show your model. You dropped an HF cache structure (models--org--name/snapshots/...) into models/vitmatte. It needs a plain directory with a config.json inside. Use huggingface-cli download --local-dir and restart ComfyUI.
2. "VITMatte node not available" - almost always missing opencv-python or a model dir without config.json.

This loader only earns its keep if you're doing detail refinement, so if you're on the plain Yogurt Sa2VA Image Segmentation node, skip it. If you want hair that survives a cutout, it's the node that makes the V2 chain work.