YOLOE-26 Class Masks

Prompt Segment merges every detection into one mask. Instance Masks gives you one mask per object. This node splits the difference in the way you'll actually want most often: one mask per class you typed. Prompt person, car, and you get a mask covering every person and a separate mask covering every car. Grouped, not merged, not per-instance.

That's the shape that makes per-category routing practical. Give the person mask to one inpaint pass and the car mask to another. Mask out one category for a background replacement while leaving the other alone. Count how many classes actually showed up. When your workflow treats "people" and "cars" as different editing targets, this node hands you both masks at once.

How it works

The segmentation runs the same way as everywhere else in the pack - prompt split into classes, embeddings computed once, one predict pass. The difference is the grouping step. Every detected instance is assigned to the class it matched, and the masks of all instances in a class get merged into a single mask, in the order you listed the classes in the prompt. Class one from your prompt is mask one in the output. Simple and predictable.

One behavior worth knowing: if a class gets zero detections, it still gets a slot in the output - an all-zero mask - so the class-to-mask mapping never shifts. output_mask_count tells you how many class masks came back, which equals how many classes you prompted. That consistency is what makes downstream routing safe; you never have to guess which index is which.

Inputs

Standard set: model, image, prompt (comma-separated classes - the order defines the output order), conf, iou, max_det, mask_threshold, imgsz. Keep the defaults (conf 0.1, iou 0.7, imgsz 640) unless you have a reason not to. The mask_threshold quirk applies as everywhere: recent Ultralytics already returns binary masks, so values below 1.0 do nothing and exactly 1.0 empties everything.

Outputs

• class_masks - a MASK batch, one merged mask per prompt class, in prompt order.
• class_metadata_json - class-to-mask mapping, including which instance indices contributed to each class mask. Useful for auditing what the grouping actually did.
• output_mask_count - the number of class masks.

Installing it

It's one of seven nodes in Rinne414/ComfyUI-YOLOE26. Install through ComfyUI Manager (search YOLOE-26) or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/Rinne414/ComfyUI-YOLOE26.git
pip install -r ComfyUI-YOLOE26/requirements.txt

Restart, and remember the first-run ritual: model downloads once, then the ~250 MB MobileCLIP text encoder downloads on the first prompt inference, so your first execution needs network.

When to reach for it

Reach for Class Masks when your prompt names multiple categories and you want to edit them differently - replace the background but keep the people, blur the cars but not the pedestrians. It's also the cheapest way to get a per-class count for dataset work. If you only care about one category, Prompt Segment's single merged mask is simpler. If you need to target one specific object out of a group, you want Instance Masks instead. This one owns the middle: everything of a kind, grouped and indexed.