Segment for Mask (ZenTrainer)

Masked training is the feature that used to force people off ComfyUI and into OneTrainer: paint over the thing you want to teach, and the loss ignores everything else. ZenTrainer gives you that inside the dataset editor, and this node is what turns your click - or the word "face" - into the mask.

What it's for

Masks solve a specific training problem. You have 30 photos of a person in a busy street, and you want the LoRA to learn the face and not the pedestrians, the signage, the wall. Full-image training learns all of it. A mask says "this region counts, the rest is ignored", which cleans up the identity and quiets the background bleed.

Concretely: pick your subject, get a greyscale PNG, and ZenTrainer saves it as name-masklabel.png next to the image - the OneTrainer convention, also accepted as <stem>.mask.<ext> or under a _masks/ folder. Existing masks from another trainer are picked up as-is. You don't have to redo them.

How it works

ZenSegmentSession hosts one SAM3 request. job_id is its only input, and like the other session nodes in this pack, you don't fill it in: the dataset mask editor creates a segment job (the image path, an optional text prompt, keep/exclude click points, a threshold, a refine count) and submits this node as a one-node prompt so the work takes a turn on the GPU instead of racing a training run for it.

Inside, the node loads a SAM3 checkpoint from models/checkpoints - whatever file has "sam3" in the name - through ComfyUI's own comfy_extras/nodes_sam3.py. Your text goes through CLIPTextEncode as conditioning, the click points go in as positive and negative coordinate sets, and SAM3_Detect returns the mask. That mask is flattened to greyscale, resized to the original image's dimensions, and written to user/zentrainer/segment/<job_id>.png, which is what the editor shows you to feather, grow and brush before saving.

Two behaviours that explain what you see in the UI. The model stays loaded between requests, so the second click on the same photo is quick instead of a reload. And the route waits about twenty seconds for the result; if it's still running (or a training run currently holds the queue) you get a job back to poll rather than a hang, with a message saying it's waiting for training to finish.

How you use it

You set nothing on the node. In the dataset editor you click the subject - or type something like person or face:1 - and the pack does the rest. You can combine both: a word to narrow the search, clicks to pick which candidate you meant. Threshold and refine iterations exist and default to 0.5 and 2, which is sane for most photos; raise refine if edges are ragged.

It's an output node with no outputs on purpose. The result is a file on disk that training reads, not a MASK socket you wire into something. If you want the mask elsewhere in a graph, use SAM3 directly.

Installing it

Same as everything else in ZenTrainer - Manager, search ZenTrainer, install; or:

cd ComfyUI/custom_nodes
git clone https://github.com/Nynxz/ZenTrainer

Restart and open /zentrainer/. ComfyUI 0.41 or newer, and if you cloned from source, build the web app once (cd web && pnpm install && pnpm build) or the dataset editor won't load.

The one real download is the SAM3 checkpoint: put sam3.1_multiplex_fp16.safetensors in ComfyUI/models/checkpoints. The pack's Python dependencies are nil - SAM3 support comes from ComfyUI itself, which is also where the friction lives: SAM3's ComfyUI nodes want Triton, and Triton on Windows is a genuine obstacle rather than a minor annoyance. If SAM3 works in a normal ComfyUI graph on your machine, it works here.

When it doesn't work

"no SAM3 model: put sam3.1_multiplex_fp16.safetensors in models/checkpoints (Models page)." Exactly that. The filename must contain "sam3", and the app's Models page can fetch it for you.

"give it something to find: a word, or click on the subject." You sent neither - an empty text prompt and zero points. SAM3 has nothing to segment on.

"segment job seg-xxxx is gone (ComfyUI restarted?); try again." Job records live in memory, last 50 kept. Restart clears them; the PNGs already saved as dataset masks are untouched.

Clicks feel slow, or the mask is empty. Kept models get evicted when VRAM is tight, so a mask request right after a big generation may pay a load. And a text prompt that's too abstract ("the mood") segments nothing - describe an object.

If masked training is the whole reason you're here, know what you're buying: it changes what the loss sees, not how the model converges. Dataset quality still dominates.