Training ControlNet with Hugging Face Diffusers

Hugging Face has released a detailed guide and training script within the diffusers library that allows developers to train their own ControlNet models. ControlNet is a neural network structure that enables fine-grained control of diffusion models by adding extra conditions, such as pose estimations, depth maps, or sketches, to the generation process.

The Three-Step ControlNet Training Workflow

Training a custom ControlNet for Stable Diffusion involves three primary phases: planning the condition, building the dataset, and executing the training.

1. Planning the Condition

The first step is defining the specific conditioning required for the task. This involves determining the desired control mechanism and identifying if an existing model can convert standard images into that specific condition. For example, the Hugging Face team aimed to create a model for facial landmarks to allow Stable Diffusion to follow specific facial expressions or poses.

2. Building the Dataset

A ControlNet dataset requires three specific columns:

• Ground Truth Image: The target image (e.g., a face).
• Conditioning Image: The image representing the condition (e.g., a visualized facial landmark mask).
• Prompt: A text caption describing the image.

In the "Uncanny Faces" example, the team used the Microsoft FaceSynthetics dataset containing 100K synthetic faces. Because no existing model could convert faces directly into the dataset's specific landmark format, the team used the SPIGA state-of-the-art model to extract 68-facial landmarks in iBUG format, converted those landmarks into illustrated masks, and used BLIP captioning to generate descriptions for each image.

3. Training the Model

Training is facilitated by the train_controlnet.py script provided in the diffusers examples. The team used a single A100 GPU, training for 3 epochs with a batch size of 4.

Training Observations and Overfitting

The team discovered that training for 3 epochs led to overfitting, where the model began to ignore styles and forgot concepts that diverged from real faces (e.g., failing to generate a "cat" or "shrek" when prompted). Convergence was achieved after approximately 1 epoch (around 25K steps), at which point the model successfully followed poses without overfitting. Because the FaceSynthetics dataset consisted of synthetic images, the resulting model produced "uncanny" 3D-looking faces rather than photorealistic ones.

Technical Implementation and Hardware Optimization

Training Configuration

The train_controlnet.py script utilizes several key parameters to control the output:

• pretrained_model_name_or_path: The base Stable Diffusion model (v2-1-base was used for better face rendering).
• learning_rate: Set to 1e-5 for the example, though values between 1e-4 and 2e-6 are suggested.
• resolution: Set to 512x512 for both conditioning and ground truth images.
• validation_steps: Determines how often the model runs a validation prompt and image to track progress.

VRAM Optimization for Lower-End GPUs

While an A100 was used for the primary example, the diffusers script supports optimizations to fit training on GPUs with less VRAM:

GPU VRAM	Required Optimizations / Parameters
16GB	train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing, and use_8bit_adam (via bitsandbytes).
12GB	All 16GB optimizations plus set_grads_to_none.
8GB	Specific configurations detailed in the diffusers GitHub training guide.

By using a batch size of 1 with 4 gradient accumulation steps, users can simulate the effective batch size of 4 used in the A100 training run while significantly reducing memory overhead.