Fine-Tune FLUX.2 Klein for Custom Image Generation

Configure Docker permissions

To manage containers without sudo, your user must be in the docker group. If you skip this step, run Docker commands with sudo.

Open a terminal and test Docker access:

docker ps

If you see a permission-denied error (for example, permission denied while trying to connect to the Docker daemon socket), add your user to the docker group:

sudo usermod -aG docker $USER
newgrp docker

Clone the repository

Clone the playbook assets and open the assets directory:

git clone https://github.com/NVIDIA/dgx-spark-playbooks
cd dgx-spark-playbooks/nvidia/playbook-flux-finetuning/assets

Model download

FLUX.2 Klein is gated. Open the Klein 4B model card for a 64 GB system or the Klein 9B model card for a 128 GB system, accept the terms, and gain access to the checkpoints. If you do not already have an HF_TOKEN, follow the Hugging Face token guide, then authenticate:

export HF_TOKEN=<YOUR_HF_TOKEN>
sh download.sh

The download can take about 30–45 minutes depending on network speed. The script detects physical memory and pulls:

• 64 GB system: FLUX.2 Klein 4B and Qwen3-4B
• 128 GB system: FLUX.2 Klein 9B and Qwen3-8B
• flux2-vae.safetensors

After download, models/ should look like:

models/
├── diffusion_models/
│   └── flux-2-klein-base.safetensors
├── loras/
├── text_encoders/
│   └── qwen_3.safetensors
└── vae/
    └── flux2-vae.safetensors

AI Toolkit downloads its training copy of the model into hf_cache/ when training starts.

If you already have fine-tuned LoRAs, place them in models/loras/. Otherwise continue to training in Step 6.

Base model inference

Generate an image with the base FLUX.2 Klein model for the sample concepts (Toy Jensen and a custom GPU) before training.

# Build the inference Docker image (run from assets/)
docker build -f Dockerfile.inference -t flux2-comfyui .

# Launch ComfyUI
sh launch_comfyui.sh

Open ComfyUI at http://localhost:8188 (or http://<HARDWARE_IP>:8188 from another device). Do not select a pre-existing template.

Open the workflow panel (left side, or press w) and load base_flux2.json. Enter a prompt in the Prompt node — for example, tjtoy toy holding sparkgpu gpu in a datacenter. High-resolution 1024px generation can take a few minutes.

Next steps:

• If you already placed LoRAs in models/loras/, skip to Step 7.
• If you will train, stop the ComfyUI container with Ctrl+C first.

sudo sh -c 'sync; echo 3 > /proc/sys/vm/drop_caches'

Dataset preparation

Prepare a dataset for concept LoRA fine-tuning on FLUX.2 Klein. This playbook ships a two-concept sample dataset of public-domain images. If you use those concepts as-is, you do not need to edit the files under config/.

TJToy concept

• Trigger phrase: tjtoy toy
• Training images: 6 images of Toy Jensen figures
• Use case: Generate scenes featuring that toy character

SparkGPU concept

• Trigger phrase: sparkgpu gpu
• Training images: 7 images of a custom GPU design
• Use case: Generate scenes featuring that GPU

For your own concepts, collect about 5–10 images per concept. Create one folder per concept under flux_data/ (this playbook uses tjtoy and sparkgpu). Update the matching file under config/ so each dataset entry has the correct folder_path, default_caption, and trigger_word. Appending a class token (for example toy or gpu) usually improves fine-tuning.

Training

Build the training image and start concept LoRA training with AI Toolkit:

docker build -f Dockerfile.train -t flux2-train .
sh launch_train.sh

launch_train.sh selects Klein 4B on 64 GB systems and Klein 9B on 128 GB systems. Both profiles run 1,000 optimizer steps at 1024×1024 and create the following stable LoRA path for the supplied workflow:

models/loras/flux2_klein_tjtoy_sparkgpu.safetensors

Useful knobs in the selected file under config/ include LoRA dimension and alpha (256), learning rate (0.0001 with AdamW 8-bit), batch size (2), mixed precision (bfloat16), and caching / torch compile options. If you reduce steps, also reduce save_every so a checkpoint is written.

Fine-tuned model inference

Generate images with your trained LoRAs:

# Launch ComfyUI (from assets/)
sh launch_comfyui.sh

Open http://localhost:8188, skip pre-existing templates, open the workflow panel (w), and load finetuned_flux2.json. Confirm that Load LoRA Model Only points to flux2_klein_tjtoy_sparkgpu.safetensors.

Prompt with your trigger phrases — for example, tjtoy toy holding sparkgpu gpu in a datacenter. High-resolution 1024px generation can take a few minutes. The fine-tuned path can combine multiple concepts in one image. Use ComfyUI nodes to adjust LoRA strength, resolution, seed, sampler, scheduler, and steps.

Cleanup (optional)

Stop running containers with Ctrl+C. Remove local images if you no longer need them:

docker rmi flux2-comfyui flux2-train