(LoRA) Fine-Tuning FLUX.1-dev on Consumer Hardware

In our previous post, Exploring Quantization Backends in Diffusers, we dived into how various quantization techniques can shrink diffusion models like FLUX.1-dev, making them significantly more accessible for inference without drastically compromising performance. We saw how bitsandbytes, torchao, and others reduce memory footprints for generating images.

Performing inference is cool, but to make these models truly our own, we also need to be able to fine-tune them. Therefore, in this post, we tackle efficient fine-tuning of these models with peak memory use under ~10 GB of VRAM on a single GPU. This post will guide you through fine-tuning FLUX.1-dev using QLoRA with the diffusers library. We'll showcase results from an NVIDIA RTX 4090. We'll also highlight how FP8 training with torchao can further optimize speed on compatible hardware.

Table of Contents

• Dataset
• FLUX Architecture
• QLoRA Fine-tuning FLUX.1-dev with diffusers
  • Key Optimization Techniques
  • Setup & Results
• FP8 Fine-tuning with torchao
• Inference with Trained LoRA Adapters
  • Option 1: Loading LoRA Adapters
  • Option 2: Merging LoRA into Base Model
• Running on Google Colab
• Conclusion

Dataset

We aim to fine-tune black-forest-labs/FLUX.1-dev to adopt the artistic style of Alphonse Mucha, using a small dataset.

FLUX Architecture

The model consists of three main components:

• Text Encoders (CLIP and T5)
• Transformer (Main Model - Flux Transformer)
• Variational Auto-Encoder (VAE)

In our QLoRA approach, we focus exclusively on fine-tuning the transformer component. The text encoders and VAE remain frozen throughout training.

QLoRA Fine-tuning FLUX.1-dev with Diffusers

We used a diffusers training script (slightly modified from here designed for DreamBooth-style LoRA fine-tuning of FLUX models. Also, a shortened version to reproduce the results in this blogpost (and used in the Google Colab) is available here. Let's examine the crucial parts for QLoRA and memory efficiency:

Code Code Code