Train Dataset Regularization

This node builds a regularization image set - a second pool of pictures whose whole job is to stop your LoRA from forgetting everything else the model knows while it learns your subject. It's an optional, opinionated technique borrowed from the DreamBooth lineage, and it plugs into TrainDatasetAdd as an extra input. You don't need it for every run, but when overfitting or "concept bleed" is the problem, this is one of the classic tools for it.

What regularization images actually do

Here's the intuition. When you train a character LoRA on 20 photos of one person and nothing else, the model can start associating the class - "person," "man," "woman" - with your specific subject, so suddenly every person it draws drifts toward your character. Regularization images push back: you feed it a batch of generic examples of the same class (lots of different people, if you're training a person) tagged with just the class word. Training on those alongside your subject preserves the general concept, so your trigger word learns the individual without hijacking the whole category.

The LoRA-training playbook lists regularization images as a standard fix for overfitting, right next to lowering the learning rate and cutting steps. It's not magic and it roughly doubles your effective dataset (more images, more steps, more time), which is why it's optional rather than default. Reach for it when your LoRA is bleeding into everything, not as a reflex.

The inputs that matter

Three inputs, all straightforward:

• dataset_path - the folder of regularization images. These are your generic class examples, separate from your subject folder.
• class_tokens - the class word for these images, e.g. person, man, dog, style. This is what you want preserved. Keep it generic; the reg set is teaching the model to hold onto the broad concept.
• num_repeats - how many times each reg image is seen per epoch, default 1. Usually you keep reg repeats low relative to your subject so the training still leans toward learning your concept, not memorizing the generic pool.

Output is subset (JSON), which you connect to the optional regularization input on TrainDatasetAdd.

Installing ComfyUI-FluxTrainer

ComfyUI Manager: search "ComfyUI Flux Trainer", install, restart. Manual:

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

Restart ComfyUI. Torch 2.4.0 or newer recommended. The example workflow uses ComfyUI-KJNodes.

Common issues

It's a cost, not a free win. Reg images add to your total training - more steps, longer runs, more VRAM churn. If your LoRA isn't actually overfitting, adding regularization just makes the run slower for no benefit, and can even dilute a subject that was training fine. Diagnose first: are outputs looking like copies of your training images, or is the class bleeding? If neither, you probably don't need this.

Balance matters. If your reg set is huge or its num_repeats is high relative to your subject, training tips too far toward the generic images and your concept comes out weak. The subject should still dominate. Keep reg repeats modest.

Class token, not trigger word. A common mix-up: class_tokens here should be the generic class you want to preserve (person), not your rare trigger token. Putting your trigger word on the reg images defeats the purpose - you'd be reinforcing the wrong thing.

Honestly, regularization is a somewhat old-school, DreamBooth-era move, and plenty of modern LoRAs skip it in favor of just curating a tighter dataset and watching the step count. But it's a legitimate tool, and if a run keeps overfitting no matter how you tune it, a well-built reg set is worth trying.