Hugging Face SDXL Dreambooth LoRA Advanced Training Guide

Hugging Face has released an advanced training script for Stable Diffusion XL (SDXL) Dreambooth LoRAs, combining the Pivotal Tuning technique from Replicate's Cog trainer and the Prodigy optimizer from the Kohya trainer. This integration aims to achieve high-quality concept capture using minimal images while maintaining the aesthetic quality of the base SDXL model.

Pivotal Tuning for Concept Representation

Pivotal Tuning combines Textual Inversion with standard diffusion fine-tuning to prevent semantic interference from existing tokens. Instead of reusing a rare token (e.g., "sks"), which may have pre-existing associations in the model's embedding space, Pivotal Tuning inserts new tokens into the text encoders.

These new tokens are optimized to represent the new concept. The training process typically involves performing textual inversion for the first half of the training epochs (controlled by --train_text_encoder_ti_frac) before continuing with UNet optimization. This ensures the model learns a clean representation of the concept before fine-tuning the weights.

Adaptive Optimizers and Prodigy

To reduce the need for manual hyperparameter tuning of learning rates and weight decay, Hugging Face recommends adaptive optimizers. While Adafactor is an option, the guide highlights Prodigy as particularly beneficial for Dreambooth LoRA training.

Prodigy dynamically adjusts the learning rate for each parameter based on past gradients. When using Prodigy, the following settings are recommended:

  • Learning Rate: Set to 1.0.
  • Additional Settings: Enable --prodigy_safeguard_warmup and --prodigy_use_bias_correction, with adam_beta2 set to 0.99 and adam_weight_decay set to 0.01.

Advanced Training Practices

Several additional techniques are incorporated into the diffusers training script to enhance LoRA quality:

Independent Learning Rates

Setting a lower learning rate for the text encoder compared to the UNet can prevent the text encoder from overfitting too quickly. However, when using adaptive optimizers like Prodigy, the optimizer manages these adjustments automatically from an identical initial learning rate.

Custom Captioning

Using a single instance prompt for all images can be suboptimal. The script supports custom captioning via the datasets library, allowing users to provide unique prompts for each image. This can be achieved by using a dataset from the Hugging Face Hub or by creating a local ImageFolder with metadata.

Min-SNR Gamma Weighting

Min-SNR gamma weighting balances conflicts between timesteps during training by adapting loss weights based on clamped signal-to-noise ratios. This is particularly effective for larger datasets; the recommended value is --snr_gamma=5.0.

Training Set Curation

High-quality, diverse data is critical for LoRA performance. Key recommendations include:

  • Faces: Use high-resolution images, avoid other faces in the training set, and include a mix of close-ups and full-body shots while avoiding distant shots.
  • Variety: Ensure diversity in lighting, poses, backgrounds, and facial expressions to improve generalization.
  • Prior Preservation Loss: Using real portrait images for regularization (rather than model-generated ones) was found to reduce language drift and maintain realism.

Experimental Results and Benchmarks

Hugging Face conducted experiments across three categories to validate these techniques:

  1. Style and Character (Huggy LoRA): Pivotal Tuning was found to be competitive with or better than full text encoder training. The use of snr_gamma=5.0 and the Prodigy optimizer improved results over AdamW.
  2. Style (Y2K Webpage LoRA): This experiment showed that style LoRAs can overfit more easily than character LoRAs. Adjusting max_train_steps, repeats, and train_batch_size was necessary to find the balance between concept capture and flexibility.
  3. Faces (Face LoRA): Experiments indicated that a rank of 32 is optimal; higher ranks (e.g., 64) often resulted in an "air-brushed" appearance with less realistic skin texture. A training multiplier of 120x the number of images was found to be effective for diverse datasets.

Inference and Compatibility

Models trained with Pivotal Tuning require both the LoRA weights (*.safetensors) and the trained text embeddings (*.safetensors).

Diffusers Inference

In diffusers, users must first load the embeddings into both text encoders (CLIP ViT-L/14 and CLIP ViT-G/14) using pipe.load_textual_inversion before loading the LoRA weights with pipe.load_lora_weights.

ComfyUI and AUTOMATIC1111

The training script generates WebUI-compatible LoRAs and embeddings. In AUTOMATIC1111, users can prompt using the embedding token and the LoRA tag (e.g., a y2k_emb webpage <lora:y2k:0.9>). In ComfyUI, the LoRA is loaded via the LoRALoader node, and the embedding is placed in the models/embeddings directory.

Sources