How to Fine-Tune a FLUX Model in under an hour with AI Toolkit and a DigitalOcean H100 GPU

2024-08-21 · Source: Paperspace by DigitalOcean Blog · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Cloud Computing & IT Infrastructure, Data Science & Analytics · Depth: Intermediate, medium

Summary

This tutorial outlines a step-by-step process for fine-tuning a FLUX.1 image generation model using a Low Rank Adaptation (LoRA) on a DigitalOcean H100 GPU via Paperspace. It details setting up the H100 machine with the ML-In-A-Box template and sufficient storage (over 250 GB). The process involves preparing custom image data by cropping, renaming, and automatically captioning images using Florence-2, then setting up the AI Toolkit environment. Users configure the training loop by editing the `train_lora_flux_24gb.yaml` file, specifying parameters like model name (FLUX.1-dev or FLUX.1-schnell), image directory path, resolution, batch size, and training steps (e.g., 2500 steps). The article demonstrates running the training, which took approximately 45 minutes for 60 images on an H100, and concludes with an inference script to use the newly trained LoRA.

Key takeaway

For Machine Learning Engineers looking to customize FLUX.1 for specific subjects or styles, you should leverage the AI Toolkit on a DigitalOcean H100 GPU. Focus on preparing high-quality, diverse image datasets and carefully configure the `train_lora_flux_24gb.yaml` file, paying attention to image paths, resolutions, and training steps. This approach enables rapid LoRA fine-tuning, allowing you to quickly adapt FLUX.1 for unique image generation tasks.

Key insights

Fine-tuning FLUX.1 LoRAs on custom data is achievable in under an hour using cloud GPUs and AI Toolkit.

Principles

• FLUX.1-dev offers superior prompt understanding over FLUX.1-schnell.
• LoRA and ControlNet enhance FLUX model guidance.
• Diverse, high-quality training images improve LoRA accuracy.

Method

Set up a cloud H100 GPU, prepare captioned image data using Florence-2, configure AI Toolkit's `train_lora_flux_24gb.yaml` with custom parameters, and execute `run.py` for LoRA training.

In practice

• Use ML-In-A-Box template for pre-installed packages.
• Caption images automatically with Florence-2.
• Adjust `train_lora_flux_24gb.yaml` for custom data and prompts.

Topics

• FLUX.1 Model
• LoRA Fine-tuning
• Image Generation
• AI Toolkit
• GPU Training

Code references

• ostris/ai-toolkit
• ostris/ai-toolkit

Best for: Machine Learning Engineer, Deep Learning Engineer, AI Student

Related on AIssential

• Fine-Tuning Biological Foundation Models with LoRA Using NVIDIA BioNeMo Recipes — LoRA enables efficient fine-tuning of large biological foundation models on single GPUs by adapting a minimal parameter subset.
• RIP NANO BANANA! MOST VERSATILE NSFW COMPETITOR RIGHT NOW! — Flux 2 Klein is a versatile, fast, and low-VRAM AI model for comprehensive image generation and editing.
• Fine-Tuning NVIDIA Cosmos Predict 2.5 with LoRA/DoRA for Robot Video Generation — LoRA/DoRA fine-tuning effectively adapts large world models like Cosmos Predict 2.5 for domain-specific robot video generation.
• I Fine-Tuned One Model 3 Ways: The $50,000 Run Forgot More Than the $1,500 One — Expensive full fine-tuning can degrade a model's general ability more than cheaper parameter-efficient methods.
• LoRA-Muon: Spectral Steepest Descent on the Low-Rank Manifold — LoRA-Muon applies spectral steepest descent to LoRA, improving finetuning stability, learning rate transferability, and performance over dense baselines.
• LearniBridge: Learnable Calibration of Feature Caching for Diffusion Models Acceleration — LearniBridge accelerates DiTs via learnable LoRA calibration of feature caching, correcting errors with few samples for high acceleration.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Paperspace by DigitalOcean Blog.