AnimaFastTrain - MaskTrain Context Tokens

The plain train node has a known flaw, and the author admits it in the release thread: the context tokens learn the character and the pose and the composition. Point it at a reference with a busy background and half your training budget goes into memorizing scenery you don't want. AnimaFastTrain - MaskTrain Context Tokens is the same training loop with one extra input - a MASK - so the loss only cares about the part of the image you actually care about. It's the version you reach for when the reference is a face on a cluttered canvas, which is most of the time.

How it works

Mechanically it's the sibling of the regular train node: reference image VAE-encoded into latents, per-block context tokens (28 blocks × 16 tokens × 1024 dim by default) optimized with AdamW inside a ComfyUI TrainGuider loop. The difference is the loss. Instead of plain MSE between predicted and target latents, it computes a mask-weighted MSE in latent space: white (1.0) mask areas contribute to the loss, black (0.0) areas are ignored entirely. The mask gets resized to your latent resolution with bilinear interpolation, so you can hand it a loose mask and it does the downsampling.

One thing the README is careful to state, because it's easy to misread: this is mask-weighted training, not a runtime cross-attention mask. The mask shapes what the tokens learn; it does nothing at sampling time. You still control runtime strength on the Patch Model node.

The inputs that matter

Everything from the regular train node applies - model, clip, vae, reference_image, training_prompt, seed, training_steps, learning_rate, training_image_size, num_tokens, init_std, training_dtype - plus the one that defines this node:

• training_mask - a MASK covering what you want the tokens to learn. Face-dominant character, mask the face and hair; keep the background black so it can't drag the tokens. The default prompt even swaps in multicolored hair, gray eyes, which tells you this node was built with face identity in mind.
• Defaults shift to match the job: training_steps is 160 (double the plain node's 80) and learning_rate drops to 0.01 - more careful, slower training for a smaller target region.

The optional inputs round it out: reference_image_2 / reference_image_3 with matching reference_mask_2 / reference_mask_3. Useful rule from the source: if you connect only one mask and give it multiple reference images, that one mask is reused for every image. Connect a mask per image only when the regions genuinely differ.

Outputs are the same pair as the regular node: context (ANIMA_REFERENCE_CONTEXT → Patch Model) and info (STRING, now including the mask's mean coverage in its training summary).

Install

Identical to the rest of the pack. ComfyUI Manager → search ComfyUI-AnimaFastTrain, or:

cd ComfyUI/custom_nodes
git clone https://github.com/quinteroac/ComfyUI-AnimaFastTrain.git

Restart ComfyUI. No extra pip packages, no model downloads - the pack ships no weights, just the training/injection logic. Nodes appear under Anima/AnimaFastTrain.

Where people get burned

All-black mask. The source raises if the resized mask contains zero non-zero pixels - so a mask that got shrunk out of existence, or one that never covered the subject, fails immediately with a clear error. Make sure your mask actually overlaps the reference.

Expecting the mask to constrain output. It won't. This node only decides what the training loss looks at. If you want the model to only generate inside a region, that's a different tool (inpainting/ControlNet territory) - and Anima's ControlNet story is thin, so don't come looking for it here.

Same VRAM ceiling as the sibling. The whole model is loaded for training; the author isn't sure anything under 12GB handles it. The mask doesn't change that.

Overfocused tokens. A tight face mask can make the tokens all-in on the face at the expense of the rest of the design - for full-body character work, a looser mask or the plain train node may honestly serve you better. The mask is a dial, not a mandate.