Implementation:Microsoft DeepSpeedExamples Stable Diffusion LoRA Distillation

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Generative AI, Knowledge Distillation, Parameter Efficient Fine-Tuning
Last Updated	2026-02-07 12:00 GMT

Overview

A training script for Stable Diffusion LoRA distillation using DreamBooth-style training with Accelerate, supporting knowledge distillation from a teacher model to a student UNet with low-rank adaptation.

Description

This script implements a comprehensive LoRA distillation training pipeline for Stable Diffusion models. It combines the DreamBooth training paradigm (learning from instance-specific images with optional prior preservation) with LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning and knowledge distillation where a teacher model guides the training of a student model's UNet.

The training workflow involves several key components. The DreamBoothDataset class handles loading instance images with their prompts, applying image transforms (resize, crop, normalize), and tokenizing both conditional and unconditional prompts for classifier-free guidance. The main function orchestrates the entire pipeline: setting up the Accelerator for distributed training, loading pretrained VAE/text encoder/UNet/tokenizer components, optionally generating class images for prior preservation, creating LoRA layers on the UNet, and running the training loop with noise prediction loss.

The script supports advanced training features including pre-computed text embeddings to avoid repeated text encoder forward passes, optional text encoder training, mixed precision (fp16/bf16/fp32), gradient checkpointing for memory efficiency, and integration with Weights & Biases for experiment tracking. The collate_fn handles batching of instance and class images with their corresponding prompt IDs and attention masks. Validation is performed periodically by generating sample images with the current model state via the log_validation function.

Usage

Use this script to perform knowledge distillation from a full Stable Diffusion model into a LoRA-adapted variant, combining DreamBooth-style subject-driven generation with parameter-efficient training. Launch via accelerate launch train_sd_distil_lora.py --pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 --default_prompt "a photo of sks dog".

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/stable_diffusion/train_sd_distil_lora.py
• Lines: 1-1288

Signature

def parse_args(input_args=None):
    """Parse training arguments for LoRA distillation."""
    ...

class DreamBoothDataset(Dataset):
    """Dataset for DreamBooth-style training with instance and class images."""
    def __init__(self, instance_prompts, instance_images, tokenizer,
                 class_data_root=None, class_prompt=None, class_num=None,
                 size=512, center_crop=False, encoder_hidden_states=None,
                 instance_prompt_encoder_hidden_states=None,
                 tokenizer_max_length=None):
        ...

class PromptDataset(Dataset):
    """Simple dataset for generating class images from prompts."""
    ...

def collate_fn(examples, with_prior_preservation=False):
    """Collate function for batching training examples."""
    ...

def main(args):
    """Main training function: setup Accelerator, load models, train LoRA distillation."""
    ...

def log_validation(text_encoder, tokenizer, unet, vae, args, accelerator,
                   weight_dtype, epoch, prompt_embeds, negative_prompt_embeds):
    """Generate validation images during training."""
    ...

def save_model_card(repo_id, images=None, base_model=str,
                    train_text_encoder=False, prompt=str, repo_folder=None):
    """Save a model card with training details."""
    ...

Import

# This is a standalone training script, not typically imported.
# Run via: accelerate launch train_sd_distil_lora.py [args]
from accelerate import Accelerator
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel

I/O Contract

Inputs

Name	Type	Required	Description
--pretrained_model_name_or_path	str	Yes	Path to pretrained Stable Diffusion model or HuggingFace model ID
--default_prompt	str	Yes	Default prompt for training instance images
--resolution	int	No	Resolution for input images (default: 512)
--train_batch_size	int	No	Per-device training batch size (default: 4)
--num_train_epochs	int	No	Number of training epochs (default: 1)
--learning_rate	float	No	Initial learning rate (default: 5e-4)
--train_text_encoder	flag	No	Whether to also train the text encoder
--with_prior_preservation	flag	No	Enable prior preservation loss with class images
--class_data_dir	str	No	Directory containing class images for prior preservation
--output_dir	str	No	Directory for saving model checkpoints (default: 'text-inversion-model')
--mixed_precision	str	No	Mixed precision mode: 'no', 'fp16', or 'bf16'

Outputs

Name	Type	Description
LoRA weights	safetensors	Trained LoRA adapter weights saved to output directory
Validation images	PIL.Image	Sample images generated during training for quality monitoring
Model card	README.md	Auto-generated model card with training details
Training logs	stdout/wandb	Loss metrics and training progress

Usage Examples

# Train LoRA distillation on Stable Diffusion v1.5
# accelerate launch training/stable_diffusion/train_sd_distil_lora.py \
#     --pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \
#     --default_prompt "a photo of sks dog" \
#     --resolution 512 \
#     --train_batch_size 4 \
#     --learning_rate 5e-4 \
#     --num_train_epochs 100 \
#     --output_dir ./lora_output \
#     --mixed_precision fp16

# With prior preservation
# accelerate launch train_sd_distil_lora.py \
#     --pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \
#     --default_prompt "a photo of sks dog" \
#     --with_prior_preservation \
#     --class_data_dir ./class_images \
#     --class_prompt "a photo of a dog" \
#     --num_class_images 200 \
#     --output_dir ./lora_output