FLUX.1 LoRA training (20240824)

• An Nvidia GPU with at least 24GB of Video RAM

• Using Flux.1-Dev, requires acceptance of the terms of use.

• A HuggingFace account to download the weights.

Preliminary steps

Code retrieval and virtualenv creation

mkdir ~/Flux_Lora
cd ~/Flux_Lora

# Obtain the source and the code's submodules
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive

# Create and enable a python virtual environment to place all the needed packages
python3 -m venv venv
source venv/bin/activate
# for future use, re-enabling the venv can be done by running the source command again

# Install the required packages
pip3 install torch
pip3 install -r requirements.txt
pip3 install mediapipe peft 
# mediapipe appears to be needed for Dev, while peft will be needed for Schnell

Hugging Face token

# Install the CLI using brew
brew install huggingface-cli

# Confirm the token is valid and add it to the default HF_HOME at ~/.cache/huggingface
huggingface-cli login
# Answer no to "Add token as git credential" as this is a "download-only" token
# This will store the token in ~/.cache/huggingface/token
# All models retrieved from HF's hub will end up in ~/.cache/huggingface/hub

Dataset preparation

mkdir ~/Flux_Lora/training_images-tok

• have at least 12 face photographs

• each photograph should be with altering conditions so the model training is able to differentiate the person from the background or its clothings. As such, use shots with altering angles of the face, lighting conditions, background, clothes, ornaments (limit the ones hiding your hair, but if wearing glass, take images with and without those).

• the script requests that you name each file image[number].[extension] increasing number as needed: image1.jpg, image2.png … image25.jpg

• for each image file there must be a corresponding .txt file with a description of the image. Because the developer allowed for text replacement using the [trigger] keyword, each file can contain the same text, for example, ours contains a photo of a [trigger] man. The content can be extended to be more descriptive but this is sufficient to train the LoRA.
• Each image must have a matching .txt file: image1.jpg with image1.txt … image25.jpg and image25.txt

LoRA Training

Flux.1-Dev

cd  ~/Flux_Lora/ai-toolkit
# create an output directory for the trained LoRA(s) 
mkdir ../trained_LoRA
# copy the example training yaml and rename it
cp config/examples/train_lora_flux_24gb.yaml config/flux_dev-tok.yaml

• In config: we modified name: "flux_dev-tok" to reflect that we are creating a flux.1-dev LoRA for a tok prompt trigger word.

• We are modifying training_folder: "../trained_LoRA" to use our created directory. The script will create a flux_dev-tok folder within (it is using the value of name), to store the training and samples.

• Uncomment and set the trigger_word to a value that you will use in your prompts. p3r5on or tok are commonly used values. This value will be recognized by the image generation text prompts. Using different values for different persons (it is recommended to always use a single non-existing word). Here we will use trigger_word: "tok” . This trigger word will be replaced inside the imageNUM.txt automatically as well as in the sample prompts section.

• The save section specifies how often the training saves .safetensors checkpoints. Each of those are usable as is, so if you decide a certain checkpoint gives better results (based on the samples generated), you can use that file instead of the final file. It is recommended to keep the sample_every generation matching the save_every value.

• In the datasets section we adapt the folder_path to where our training images are located: folder_path: "../training_images-tok"

• In the train section, we increase the number of steps to steps: 4000.
• The tool will generate checkpoints at save_every steps. If the samples seem to deteriorate as training steps increase, we can select an earlier checkpoint file instead of the file one.

• We are not altering the model section.
• The file if already present in the HF cache will not be redownloaded.
• You must have agreed to the terms for the FLUX.1-dev model for it to be accessible.

• in the sample section, we keep the sample_every to match the save_every value. We can comment out (#) some prompt while adding [trigger] to some. The idea is to see how good the model is at training the LoRA on the images we provided, not to demonstrate the capabilities of FLUX.1-dev. For example:

        prompts:
          - "a woman holding a coffee cup, in a beanie, sitting at a cafe"
          - "a [trigger] man holding a sign that says 'AI is fun'"
          - "a [trigger] man as a Jedi warrior, deep blue lightsaber, dark purple background with some multicolor neon light reflecting"
          - "a [trigger] man as a green lantern, cosmic background, creating a dragon from his ring, majestic, rule of third"
          - "a [trigger] man as Conan the barbarian, snow-covered fields, sword in hand, defiant, dynamic pose, golden ratio, asymmetric composition, photorealism, cinematic realism"
          - "a portrait of a [trigger] man as Superman flying in space, yellow sun in the background, majestic, comic style"
          - "a [trigger] man as a cyberpunk warrior, anime style"
          - "professional photograph portrait of a [trigger] man, black and white"
          - "a [trigger] man wearing sunglasses riding a high-speed motorbike, hyper-maximalist, octane render"
          - "a [trigger] man, lucid dream-like 3d model, game asset, blender, unreal engine, rule of thirds, wide angle shot, looking off in distance style, glowing background, vivid neon wonderland, particles, blue, green, orange"

Here, only our first prompt will be consistent from sample generation to sample generation (the seed being set and not random). All the other generations will show us the improvements of the model at generating a [trigger] man. We have a total of 10 prompts.

# Make sure you are running the create virtual env
# if needed: cd ~/Flux_Lora/ai-toolkit; source venv/bin/activate
time python3 ./run.py config/flux_dev-tok.yaml

Flux.1-Schnell

cd  ~/Flux_Lora/ai-toolkit
# create an output directory for the trained LoRA(s) if it does not already exist
mkdir ../trained_LoRA
cp config/examples/train_lora_flux_schnell_24gb.yaml config/flux_schnell-tok.yaml

• training_folder: "../trained_LoRA"

• trigger_word: "tok”

• folder_path: "../training_images-tok"

• steps: 4000

# Make sure you are using the created virtualenv
# if needed: ~/Flux_Lora/ai-toolkit; source venv/bin/activate

# Run the training
time python3 ./run.py config/flux_schnell-tok.yaml

Using the trained models

With ComfyUI

Flux.1-Dev LoRA with ComfyUI

Flux.1-Schnell LoRA with ComfyUI

Extended workflow

Revision History

• 20240824-0: Added more complex workflow, integrating the use of ComfyUI Manager.

• 20240818-0: Initial release