ComfyUI Flux Trainer
Currently supports LoRA training, and untested full finetune with code from kohya's scripts:…
Nodes (32)
Turn a full fine-tune into a shippable LoRA
Test your trained LoRA with Kohya's own sampler
Run training and preview it as it goes
Train only the transformer blocks you name
Finalize the run and write out your finished LoRA
The huber loss knobs (and why you can mostly ignore them)
The node that actually grinds training steps
Point the trainer at your Flux model, VAE, and text encoders
Pick a training run back up where it died
Write the LoRA to disk (and straight into ComfyUI)
Save intermediate checkpoints mid-train (so you can pick the best one)
Generate preview images mid-training so you can see if it's working
Control the preview images during training
The node that actually runs a Flux LoRA train inside ComfyUI
Full Flux fine-tuning (DreamBooth) in ComfyUI, not a LoRA
LoRA training for SD3.5 inside ComfyUI
Train an SDXL or Illustrious LoRA without leaving ComfyUI
The default optimizer picker for Flux Trainer (AdamW8bit and friends)
The low-VRAM optimizer for when memory is tight
The optimizer that finds its own learning rate
The kitchen-sink adaptive optimizer
Load the model files for an SD3 training run
Preview settings for an SD3 training run
Pick the checkpoint your SDXL LoRA trains on
Preview images while you train an SDXL LoRA
Control the preview images during SDXL training
Point FluxTrainer at your images
The dataset-wide settings for a FluxTrainer run
Add reg images to fight overfitting
Switch from plain LoRA to LyCORIS/LoKr and friends
Push your finished LoRA straight to the Hub from ComfyUI
Plot the training loss curve right in your graph
ComfyUI Flux Trainer
Wrapper for slightly modified kohya's training scripts: https://github.com/kohya-ss/sd-scripts
Including code from: https://github.com/KohakuBlueleaf/Lycoris
And https://github.com/LoganBooker/prodigy-plus-schedule-free
DISCLAIMER:
I have very little previous experience in training anything, Flux is basically first model I've been inspired to learn. Previously I've only trained AnimateDiff Motion Loras, and built similar training nodes for it.
DO NOT ASK ME FOR TRAINING ADVICE
I can not emphasize this enough, this repository is not for raising questions related to the training itself, that would be better done to kohya's repo. Even so keep in mind my implementation may have mistakes.
The default settings aren't necessarily any good, they are just the last (out of many) I've tried and worked for my dataset.
THIS IS EXPERIMENTAL
Both these nodes and the underlaying implementation by kohya is work in progress and expected to change.
Installation
- Clone this repo into
custom_nodesfolder. - Install dependencies:
pip install -r requirements.txtor if you use the portable install, run this in ComfyUI_windows_portable -folder:
python_embeded\python.exe -m pip install -r ComfyUI\custom_nodes\ComfyUI-FluxTrainer\requirements.txt
In addition torch version 2.4.0 or higher is highly recommended.
Example workflow for LoRA training can be found in the examples folder, it utilizes additional nodes from:
https://github.com/kijai/ComfyUI-KJNodes
And some (optional) debugging nodes from:
https://github.com/rgthree/rgthree-comfy
For LoRA training the models need to be the normal fp8 or fp16 versions, also make sure the VAE is the non-diffusers version:
https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/ae.safetensors
For full model training the fp16 version of the main model needs to be used.
Why train in ComfyUI?
- Familiar UI (obviously only if you are a Comfy user already)
- You can use same models you use for inference
- You can use same python environment, I faced no incompabilities
- You can build workflows to compare settings etc.
Currently supports LoRA training, and untested full finetune with code from kohya's scripts: https://github.com/kohya-ss/sd-scripts
Experimental support for LyCORIS training has been added as well, using code from: https://github.com/KohakuBlueleaf/Lycoris