The Flux Dev AI model is a great leap forward in local Diffusion models. It delivers quality surpassing Stable Diffusion 1.5 and XL models. Like Stable Diffusion models, you can train LoRA models on top of Flux to inject custom characters and styles.
In this post, I will provide a Google Colab notebook for training Flux LoRA models. With some setup, you can also run the training workflow locally if you have a good GPU card.
Table of Contents
Software
This is an advanced tutorial for training your own AI model. If you are new to Stable Diffusion or Flux AI models, get the Quick Start Guide for a newbie setup guide.
This training workflow uses ComfyUI as GUI. It uses the ComfyUI Flux Trainer custom node with modifications for ease of use, which calls the try-and-true Kohya LoRA Trainer under the hood.
The training workflow consists of two parts:
- Generate captions of the training images automatically.
- Apply those images to Flux.1 dev model to train a custom LoRA model.
Training on Google Colab
What you need before you start
Google Colab Paid Plan
Due to the demanding computational resources required for training Flux LoRAs, this notebook requires a paid Google Colab plan. I use the Colab Pro Plan but any paid plan should work.
It typically takes ~4.5 hours to train a LoRA on an L4 instance. At the time of writing, the Colab Plus plan costs $10 a month, and you can use an L4 for about 20 hours. So, training a LoRA set you back $2.5 on Colab.
Training images
To train a Flux LoRA model, you need a set of training images. 10-20 images should do the trick for training a face.
If you are a member of this site, you can download the example training images below. They are all cropped to 1024×1024, but the training notebook supports different sizes.
Tips for good training images
- The ideal size is 1024×1024. It is OK to use images with different sizes but be sure to include some close to 1024×1024.
- Diversity is the key. You want your subject to be in different scenes, settings, and clothing. Otherwise, the model will be confused about what you are trying to train.
- If you are training a face, include a few high-resolution headshots.
- The default parameters work for 10 – 20 images.
Step 1: Upload images to Google Drive
Put your training images in the folder AI_PICS > Flux_trainer_input in your Google Drive. (The folder names are case-sensitive.)
Step 2: Run the training notebook
Open the Google Colab Notebook below.
Click the Run button to start running the notebook.
You will be asked to grant permission to access your Google Drive. Grant the permission, as it is necessary for loading the training images and saving the LoRA model.
It will take a few minutes to load. When it is done, you should see an URL and a Tunnel Password like the ones below.
Visit the URL and enter the Tunnel Password to access ComfyUI.
Step 3: Load the workflow
Download the Easy Flux Trainer workflow below and drop it on the ComfyUI browser page to load it.
You should see the workflow loaded like the screenshot below.
Step 4: Review input parameters
Review the input parameters in the Inputs group.
- LoRA name: The name of your LoRA. Pick a name that matches what you are training for.
- Token: The trigger keyword of your LoRA. Put this in the prompt when using your LoRA.
- Image Input Path: The folder path of your training images in Google Drive. You don’t need to change this unless you have put the images in another folder.
- LoRA Output Path: The folder where the LoRA model will be stored.
- Test prompts: The prompts that will be used to test the LoRA model during training. The prompts are separated by a vertical line (|). Include the token in the prompt if you want to test the training. Include a prompt to test not using the token to monitor over-training.
Step 5: Generate captions
The workflow can generate captions automatically for your training images using the BLIP Vision-language model.
In the Step Selector Input node, set
- Enable Captioning: Yes
- Enable Training: No
Click Queue Prompt to start captioning.
After it is done, you will see text files with the same name of the training images created in the image input folder.
You can optionally revise them to better match the images.
Step 6: Start training
Now, we get to the fun part of training.
In the Step Selector input node, set:
- Enable Captioning: No
- Enable Training: Yes
The workflow is configured to test and save the LoRA with the prompts you specified every 400 steps. You should see 4 blocks of Flux Train Loop nodes like the one shown below.
You can change the steps values to change intervals between saving/testing the LoRA.
Click Queue Prompt to run the workflow.
Running on an L4 GPU instance (default) takes ~1.5 hours per 400 steps. If you can visualize the training results in ComfyUI like below.
If you don’t see the images, you can find them in the samples folder inside the LoRA output folder. The default output folder is AI_PICS > Flux_trainer_output in your Google Drive.
It typically takes 1,000-1,500 steps to train a LoRA for Flux.1 Dev.
Feel free to stop the workflow early if you have achieved the results already.
Step 7: Test the LoRA
Follow the tutorial How to use LoRA with Flux AI model to test your LoRA.
Training locally on Windows/Linux
You can use this workflow locally, although I won’t support it.
Software setup
It is better to install a fresh ComfyUI just for this training.
Install ComfyUI Manager.
Drop the workflow to the ComfyUI browser page and install the missing custom nodes using the ComfyUI Manager.
I am not able to keep up with the breaking changes in the custom nodes. So, the workflow only works for particular versions of ComfyUI and custom nodes.
Use git checkout to check out the following commits accordingly.
- ComfyUI:
9c5fca75f46f7b9f18c07385925f151a7629a94f - ComfyUI-Manager:
ce874d5c624d5713e7db334d1e0c50aeddb90d82 - was-node-suite-comfyui:
bb34bd429ab74a22a7f58551429a0f3046e1464e - rgthree-comfy:
cae8e2ad28ddb933a916b852d26b95726f60529f - ComfyUI-KJNodes:
7aa591b3a64a3f83ec2c3e92758d0bb0926a6fe0 - ComfyUI-FluxTrainer:
c3aa4ea889153519f7be40636b44d5a03a060816 - Image-Captioning-in-ComfyUI:
9b24deea8eef830da059aa91cac9690ecde19fda
For example, run the following under the ComfyUI folder.
git checkout 9c5fca75f46f7b9f18c07385925f151a7629a94fAnd the following under the custom_nodes\ComfyUI-FluxTrainer folder.
git checkout c3aa4ea889153519f7be40636b44d5a03a060816
