Fine-Tuning Diffusion Models - DreamBooth, LoRA, Textual Inversion, ControlNet

:::note Reading time: ~55 minutes | Interview relevance: High | Target roles: ML Engineer, AI Engineer, Applied Scientist, MLOps Engineer :::

---

The Real Interview Moment

A design lead at a product company has just asked your team to make Stable Diffusion reliably generate their brand mascot - a specific cartoon bear character with a distinctive style. The model has never seen this bear. If you generate "a cartoon bear in a spacesuit," you get a generic bear, not theirs.

The interviewer asks: "How would you approach this? Walk me through at least two methods, their tradeoffs, the hyperparameter choices you'd make, and how you'd evaluate whether it worked." This question tests whether you understand the personalization literature or just know how to run scripts. The right answer covers Textual Inversion and DreamBooth, explains the prior preservation loss, understands LoRA as a parameter efficiency technique, and has concrete evaluation metrics.

You have 20 reference photos of the mascot. DreamBooth with LoRA, 30 minutes of fine-tuning on one GPU, and a 100MB checkpoint later - the model generates the mascot reliably in any style, any setting, any action. This is how modern AI product teams move.

Then they push further: "Now the client wants to lock the pose of the character. They have a reference sketch showing the bear's exact stance. How do you add spatial control without retraining the whole model?" This is where ControlNet enters. And if they ask "What if they want to provide a reference image instead of a text description for the style?" - that is IP-Adapter territory.

Each of these techniques solves a distinct problem in the personalization space. By the end of this lesson you will understand all of them - not just what they do, but why they were designed the way they were, what they trade off, and when to reach for each one in a production context.

---

Why This Exists - The Personalization Problem

A pretrained Stable Diffusion model knows thousands of concepts: dogs, chairs, medieval castles, Van Gogh's style, anime faces. But it has never seen your specific dog, your product design, your brand's visual identity. The model's knowledge is statistical - it knows the distribution of "dog" images across the internet, not the specific visual identity of your pet.

Naively, you might think: just write a better prompt. "A golden retriever with floppy ears and a brown patch on the left eye, sitting on a red couch." The model tries, but generates a different dog every time - the description activates the general distribution, not a specific identity. Text embeddings in CLIP are learned on millions of images; no combination of existing tokens can uniquely identify a specific instance.

The solution is to give the model a few reference images and teach it to associate a new, rare token (like [V] or sks) with that specific concept. This is personalization: extending the model's vocabulary with a new concept defined by examples.

The fundamental challenge: you cannot train a large neural network on 5-20 images without immediately overfitting. Every fine-tuning approach in this lesson is fundamentally about finding ways to update the minimum set of parameters needed to encode the new concept while preserving the model's existing generative capabilities. Too many parameters updated, and the model forgets everything else. Too few, and the concept is not captured faithfully enough.

The secondary challenge for production deployment: checkpoints must be portable, shareable, and composable. A 3.5GB full-model checkpoint per concept is impractical when you serve thousands of user-defined concepts. LoRA's 50-100MB checkpoints that can be merged or stacked at inference are what make personalization-as-a-service commercially viable.

---

Historical Context

The personalization problem for generative models pre-dates diffusion. GAN inversion research in 2021 tried to find the latent code that reproduced a given face. The code often existed, but moving in that latent space was hard to control - you could interpolate between two faces but controlling attributes precisely was brittle.

Rinon Gal et al. at Tel Aviv University (2022) introduced Textual Inversion in the paper "An Image is Worth One Word." The key insight: instead of modifying model weights, learn a new text embedding. If CLIP's token vocabulary is fixed, learn a new pseudo-word embedding $v^*$ that, when inserted into a prompt, causes the model to generate the target concept.

Nataniel Ruiz et al. at Google Research (2022) introduced DreamBooth: "DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation." The insight: use a rare identifier token and fine-tune the full model, but simultaneously train on class-prior images to prevent catastrophic forgetting.

Edward Hu et al. had introduced LoRA for LLMs in 2021 ("LoRA: Low-Rank Adaptation of Large Language Models"). The diffusion community rapidly adapted it - community trainers like kohya-ss built LoRA training scripts that ran on consumer hardware, making DreamBooth accessible to anyone with an RTX 3090.

Lvmin Zhang and Maneesh Agrawala (2023) introduced ControlNet: "Adding Conditional Control to Text-to-Image Diffusion Models." The key architectural insight: freeze the base model completely, clone the encoder branch, and add conditioning through zero-initialized convolution projections. This enabled precise spatial control without touching the pretrained model.

Ye et al. (2023) introduced IP-Adapter: "IP-Adapter: Text Compatible Image Prompt Adapter for Text-to-Image Diffusion Models." The insight: add a separate decoupled cross-attention for image features, allowing an image to serve as a conditioning signal alongside or instead of text.

---

1. Textual Inversion - Learn a New Token

The Core Idea

Textual Inversion (Gal et al. 2022) does not modify any model weights. Instead, it learns a new token embedding $v^* \in \mathbb{R}^{d_{text}}$ that, when inserted into the text conditioning pipeline, causes the model to generate the target concept.

The token is a new entry in the CLIP vocabulary - a placeholder like <my-bear>. The embedding is randomly initialized and then optimized via gradient descent on the reconstruction loss, using only the target concept images.

The intuition: CLIP's text encoder maps a prompt like "a photo of a dog" to a sequence of embedding vectors, one per token. The model has learned to associate these embedding vectors with visual features in training images. Textual Inversion asks: what embedding vector, if placed in that sequence, would cause the model to reconstruct these specific reference images? The answer is the learned pseudoword $v^*$.

Training Objective

$v^* = \arg\min_v \mathbb{E}_{z \sim \mathcal{E}(x), y, \varepsilon, t}\!\left[\|\varepsilon - \varepsilon_\theta(z_t, t, \tau_\theta(y \oplus v))\|^2\right]$

where $y$ is a text template like "a photo of <my-bear>" and $v$ is the learned embedding for <my-bear>.

Only $v$ is updated - all model weights ($\varepsilon_\theta$, $\tau_\theta$, $\mathcal{E}$, $\mathcal{D}$) are frozen. This is a gradient flowing backward through the frozen text encoder and U-Net to update a single embedding vector. The gradient path: loss → U-Net cross-attention → text encoder output → embedding lookup → $v$.

The Rare Token Trick

Initialization matters significantly. Initializing $v^*$ with the embedding of a semantically related token (e.g., "bear" for a bear mascot) converges faster than random initialization. The text prompt template matters too - using multiple templates during training ("a photo of {token}", "a {token} in a park", "a drawing of {token}") improves generalization.

Properties and Limitations

Property	Value
Parameters trained	1 embedding vector (~768 floats for CLIP-L)
Training images needed	5-10
Training time	30-60 minutes on a single GPU
Checkpoint size	~6KB
Quality ceiling	Low - limited to what a single embedding can express
Language drift	None - model weights unchanged
Style transfer	Good for style; weaker for fine-grained identity
Multiple tokens	Possible - 2-4 tokens improve quality for complex concepts

The limitation of Textual Inversion is expressive capacity. A single embedding vector must encode everything about the concept - its appearance, texture, style, geometry. Complex concepts (specific faces, detailed characters) exceed what one token can represent. Multiple tokens can help, but quality remains lower than DreamBooth.

When Textual Inversion is the right choice: style concepts (a distinctive illustration style, a specific artistic medium), simple textures, concepts you want embedded in a tiny shareable file. Not appropriate for: specific face identity preservation, detailed character appearances, complex multi-attribute concepts.

---

2. DreamBooth - Fine-Tune the Whole Model

Concept and Motivation

DreamBooth (Ruiz et al. 2022) takes the opposite approach: fine-tune the entire U-Net (and optionally the text encoder) to encode the concept directly into the model's weights. This gives the model much greater expressive capacity than a single embedding.

The key design choice: use a rare token identifier - a string like sks that appears extremely rarely in CLIP's training data, so the model has no prior meaning for it. Then associate this token with the concept: "a photo of a sks bear."

Why rare? If you used a common token like "bear," the gradient updates would need to fight against the model's strong existing prior for "bear." A rare token has no such prior - the model is free to associate it with whatever the training images show. After fine-tuning, the model has learned that sks bear = the specific mascot. You can then generate "a sks bear on Mars," "a sks bear in Van Gogh style," "a sks bear as a LEGO figure" - the concept generalizes across styles and settings.

The Prior Preservation Loss

The fundamental problem with fine-tuning a large model on 5-20 images: catastrophic forgetting. The model overfits to the few training images and forgets what a "bear" looks like in general. After fine-tuning, generating "a polar bear" might produce a version of your mascot instead of a polar bear. This is called language drift: the model's learned meaning of "bear" has drifted toward the fine-tuned concept.

DreamBooth solves this with the prior preservation loss: simultaneously fine-tune on the target concept images and on images generated by the model itself using the class prompt (e.g., "a bear"). These class prior images serve as anchors that prevent the model from forgetting what a generic bear looks like.

$L = \underbrace{L_{DM}(x, c_{text})}_{\text{target concept loss}} + \lambda \underbrace{L_{DM}(x_{pr}, c_{pr})}_{\text{class prior preservation loss}}$

where:

• $x$ are the target concept images (your 20 bear photos)
• $c_{text}$ is the identifier prompt: "a photo of a sks bear"
• $x_{pr}$ are images of class instances generated by the original frozen model: "a photo of a bear"
• $c_{pr}$ is the class prompt: "a photo of a bear"
• $\lambda$ is the prior preservation weight (typically 1.0)

The class prior images ($x_{pr}$) are generated before fine-tuning using the original model. During training, the model is simultaneously trained to generate your specific bear AND to preserve its knowledge of generic bears.

Training Recipe

1. Generate 200-300 class prior images using original SD:
   prompt: "a cartoon bear"
   These represent the class distribution to preserve.

2. Fine-tune U-Net (and optionally text encoder) on:
   - Your 20 concept images with prompt "a photo of sks bear"
   - 200-300 class images with prompt "a photo of a bear"
   Balanced in each batch.

3. Hyperparameters (SD 1.x, A100):
   - learning_rate: 1e-6 (U-Net full fine-tune)
   - text_encoder_learning_rate: 5e-7 (if training text encoder)
   - batch_size: 4 (2 concept + 2 class per batch)
   - num_train_steps: 800-1200
   - lr_scheduler: constant with linear warmup (50 steps)
   - prior_preservation_weight: 1.0
   - mixed_precision: fp16 (saves memory, marginal quality loss)

Text Encoder Fine-Tuning

Whether to fine-tune the text encoder alongside the U-Net is a significant design decision:

Frozen text encoder: the identifier token sks gets its visual meaning purely from U-Net weight updates. Faster, lower memory, less risk of corrupting text understanding. Good for concepts that can be described with existing language constructs.

Fine-tuned text encoder: the identifier embedding itself is updated. Can better capture concepts that are hard to describe in language. Risk: the text encoder can overfit and produce degraded performance on other prompts.

Rule of thumb: for DreamBooth without LoRA, fine-tuning the text encoder at a lower learning rate (5e-7 vs 1e-6 for U-Net) usually improves concept fidelity. For DreamBooth+LoRA on small datasets (fewer than 15 images), keep the text encoder frozen.

---

3. LoRA for Diffusion - Parameter Efficiency

The LoRA Decomposition

DreamBooth fine-tunes the full U-Net (860M parameters for SD 1.5). This requires significant GPU memory for optimizer states and results in a large checkpoint. LoRA (Low-Rank Adaptation, Hu et al. 2021, adapted to diffusion) solves this by approximating weight updates with low-rank matrices.

For each weight matrix $W_0 \in \mathbb{R}^{m \times n}$ in the U-Net (attention projections and FFN layers), LoRA adds a low-rank update:

$W = W_0 + \Delta W = W_0 + B A$

where $B \in \mathbb{R}^{m \times r}$ and $A \in \mathbb{R}^{r \times n}$ with rank $r \ll \min(m, n)$.

Only $B$ and $A$ are trained - $W_0$ is frozen. The number of parameters is $(m + n) \times r$ instead of $m \times n$.

At rank $r=4$ for a $768 \times 768$ attention layer:

• Full weight matrix: $768 \times 768 = 589,824$ parameters
• LoRA matrices: $(768 + 768) \times 4 = 6,144$ parameters - 96% reduction

The $A$ matrix is typically initialized from a Gaussian distribution; $B$ is initialized to zero. This ensures that at initialization, $\Delta W = BA = 0$, so the model starts as the original pretrained model and gradually deviates.

Which Layers to Target

This is a commonly asked interview question. Not all layers benefit equally from LoRA adaptation:

Cross-attention layers (query, key, value, output projections): these are where the text conditioning enters the U-Net. Fine-tuning these teaches the model how the identifier token maps to visual features. Most critical for concept learning.

Self-attention layers: these control spatial coherence and structure within generated images. Optional but helpful for style concepts.

Feed-forward layers: contain style and texture knowledge. Useful for style fine-tuning but less critical for concept/identity learning. Adding FFN LoRA increases checkpoint size.

For DreamBooth-style personalization (concept identity), target cross-attention to_q, to_k, to_v, to_out.0. For style fine-tuning, also include FFN layers (ff.net.0.proj, ff.net.2).

LoRA Alpha and Rank Selection

The LoRA contribution is scaled by $\alpha$:

$W' = W_0 + \frac{\alpha}{r} \cdot BA$

When $\alpha = r$, the scaling is 1.0. When $\alpha < r$, the contribution is dampened. In practice, setting lora_alpha = rank (so the scaling ratio is 1.0) is the default, and you control strength via the weight multiplier at inference.

Rank selection guide:

Dataset Size	Concept Type	Recommended Rank
5-15 images	Specific identity	4-8
15-50 images	Concept + style	8-16
50-500 images	Style fine-tuning	16-32
500+ images	Full style transfer	32-64
1000+ images	Domain adaptation	64-128

Higher rank gives more expressiveness but higher overfitting risk on small datasets. Always validate with the DINO-I / CLIP-T tradeoff - if CLIP-T drops significantly as you increase rank, you are overfitting.

Merging at Inference

After training, LoRA weights can be merged into the base model for zero-overhead inference:

$W' = W_0 + \alpha \cdot BA$

The merged $W'$ has the same shape as $W_0$ - no architectural change, no inference overhead. This is ideal for serving a specific concept at scale.

For multi-concept serving (many users, many LoRAs), keep base weights fixed and load/merge LoRA per request - the merge is fast (simple matrix multiplication).

DreamBooth + LoRA Comparison

The most practical personalization approach combines DreamBooth's semantic specificity with LoRA's parameter efficiency:

Approach	Parameters Trained	VRAM Needed	Checkpoint Size	Time (A100)
Textual Inversion	~768 floats	8GB	6KB	1 hour
DreamBooth (full)	860M	24GB+	3.5GB	3-5 hours
DreamBooth + LoRA	~3M (rank 4)	8-12GB	50-100MB	30-60 min
LoRA-only (no DB loss)	~3M	8GB	50-100MB	30-60 min

DreamBooth + LoRA at rank 4 with prior preservation loss is the current standard for personalization tasks: competitive quality, runs on consumer hardware, portable checkpoints.

---

4. Architecture Overview - Fine-Tuning Methods

---

5. ControlNet - Spatial Conditioning

The Problem ControlNet Solves

Fine-tuning with DreamBooth/LoRA teaches the model what to generate (a specific concept). But the text prompt still has only loose control over how to arrange the scene spatially. A prompt like "a sks bear standing with arms raised" can produce many different interpretations of "arms raised."

ControlNet (Zhang et al. 2023) solves this: given a spatial conditioning signal - an edge map, depth map, human pose skeleton, or segmentation mask - the model generates an image that precisely follows that spatial structure while the text controls style and appearance.

A real production example: you want to generate your mascot bear in an exact pose matching a reference sketch from the design team. Without ControlNet, you would write detailed pose descriptions and iterate with seeds. With ControlNet, you extract the Canny edges from the sketch and use them as the conditioning signal. The generated image follows the sketch's structure precisely.

ControlNet Architecture

The key architectural insight: freeze the base SD model completely and add a trainable copy of the U-Net encoder that processes the conditioning input.

Conditioning input (Canny edge map, 512x512):
  → ControlNet trainable encoder copy (same architecture as SD encoder)
  → Produces feature maps at each resolution level
  → These feature maps are injected into the frozen SD decoder
    via zero-initialized convolution projections

Specifically, ControlNet output at each resolution is added
to the frozen SD encoder/decoder skip connections.

The zero-initialization of the output convolutions is the critical design choice. At training start, the ControlNet output is zero - the base model generates exactly as it did before ControlNet was added. This means:

1. Training starts from a stable, well-defined initialization (the pretrained SD model)
2. The ControlNet must "earn its contribution" by learning to produce useful features
3. The gradient signal is clean from the start because the base model behaves normally

Without zero-initialization, random ControlNet outputs would immediately corrupt the pretrained model's activations at training step 1, creating a noisy gradient signal that makes convergence slow and unstable.

How ControlNet Processes Conditioning Inputs

The conditioning input is a spatial image (same resolution as the target) but with a different semantic meaning:

Canny edges: binary edge map extracted from a reference image or sketch. Black background, white edges. Preserves the structural outline of objects.

Depth map (MiDaS): grayscale depth estimation where brighter = closer to camera. Captures 3D spatial arrangement without color or texture.

Human pose (OpenPose): skeleton image showing body keypoints (shoulders, elbows, wrists, hips, knees, ankles) as colored circles connected by lines. Does not capture person appearance - only pose.

HED edges: soft edge detection preserving more detail than Canny, useful for architectural and texture-rich scenes.

Segmentation mask: semantic regions colored by category. Provides region-level control over what appears where.

Normal map: encodes surface orientation. Each pixel's color encodes the 3D normal vector direction. Used for relighting and material editing.

The ControlNet encoder is trained to map each conditioning type into feature maps that guide the frozen SD decoder to follow the spatial structure.

Supported Conditioning Types

Conditioning	What It Controls	Extraction Tool	Example Use
Canny Edges	Object boundaries and structure	OpenCV Canny	Generate variations of a sketch
Depth Map (MiDaS)	3D spatial arrangement	MiDaS depth estimator	Match depth structure of reference photo
Human Pose (OpenPose)	Body keypoints and skeleton	OpenPose / DWPose	Generate people in specific poses
Segmentation Mask	Region-level semantic control	Segment Anything	Change sky while keeping foreground
Scribble	Rough user sketches	User draw or XDoG	Turn a doodle into a realistic image
Normal Map	Surface orientation	BAE, ZoeDepth	Architectural and product design
MLSD Lines	Straight line structure	MLSD	Architectural layout control

Multiple ControlNets can be combined at inference with additive contributions weighted by individual strength parameters: $\text{SD\_output} = f_{SD}(z_t) + w_1 \cdot \text{CN}_1(c_1) + w_2 \cdot \text{CN}_2(c_2)$

ControlNet vs Fine-Tuning - Complementary, Not Competing

ControlNet is not mutually exclusive with DreamBooth/LoRA:

• DreamBooth+LoRA answers: "what concept is this?"
• ControlNet answers: "what pose/structure/layout?"
• Combined: generate your specific mascot character (DreamBooth identity) in an exact pose from a reference sketch (ControlNet spatial control)

To combine: load your DreamBooth+LoRA checkpoint, then load a ControlNet checkpoint for the desired conditioning type. At inference, pass both the identifier prompt and the conditioning image.

---

6. IP-Adapter - Image as Conditioning Signal

The Concept

IP-Adapter (Ye et al. 2023) conditions generation on a reference image (not just text). Instead of describing "a painting in the style of Klimt," you provide an image of a Klimt painting. The model generates images that match the visual style of the reference.

The key difference from DreamBooth: IP-Adapter does not fine-tune for a specific concept. It adapts at inference time using a reference image as a conditioning signal. No training on your specific images is required.

Architecture

IP-Adapter adds a decoupled cross-attention mechanism:

Standard SD: text cross-attention → spatial features
IP-Adapter:  text cross-attention → spatial features
             +
             image cross-attention → spatial features (separate)

Image features come from CLIP image encoder (not text encoder).
The two cross-attention outputs are added.

The image cross-attention uses CLIP image embeddings (not text embeddings) as the key-value pairs. This allows the model to attend to image features using the same attention mechanism that drives text conditioning.

IP-Adapter weights consist only of the image cross-attention layers - approximately 100MB. They can be loaded on top of any compatible base model checkpoint without modifying the base weights.

IP-Adapter Variants

IP-Adapter (basic): style and coarse content transfer from reference image. Good for artistic style, general composition.

IP-Adapter-Plus: uses image features from a larger CLIP model and more fine-grained attention. Better at preserving fine details from the reference.

IP-Adapter-FaceID: specialized for face identity transfer. Uses ArcFace facial recognition embeddings instead of CLIP. Preserves face identity more precisely than CLIP-based approaches.

When to Use IP-Adapter vs DreamBooth

Scenario	Recommended Approach
Replicate a specific art style from one image	IP-Adapter
Generate consistent character across many images	DreamBooth + LoRA
Transfer face identity to new scenes	IP-Adapter FaceID
Apply product appearance to new contexts	DreamBooth + LoRA
Style reference at inference without training	IP-Adapter
Need highest identity fidelity	DreamBooth + LoRA

IP-Adapter is inference-time only - no training required. This makes it operationally simpler but less precise for identity preservation than DreamBooth. Think of IP-Adapter as "style/content transfer at inference" and DreamBooth as "concept embedding through training."

---

7. Full DreamBooth + LoRA Training Script

"""
DreamBooth + LoRA training for Stable Diffusion.
Adapted from HuggingFace diffusers training example.
Demonstrates prior preservation loss, LoRA on cross-attention,
and the full training loop with evaluation.
"""

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    DDPMScheduler,
    StableDiffusionPipeline,
)
from transformers import CLIPTextModel, CLIPTokenizer
from peft import LoraConfig, get_peft_model
from PIL import Image
from pathlib import Path
import numpy as np
from typing import Optional, List


# ============================================================
# Dataset: concept images + class prior images
# ============================================================

class DreamBoothDataset(Dataset):
    """
    Dataset for DreamBooth training with prior preservation.

    Returns batches containing both concept images (identifier prompt)
    and class prior images (class prompt) for the prior preservation loss.
    The dataset cycles concept images to match the class image count.
    """

    def __init__(
        self,
        instance_data_root: str,      # Your 20 bear photos
        instance_prompt: str,          # "a photo of sks bear"
        class_data_root: Optional[str],  # 200 generated "a bear" images
        class_prompt: Optional[str],   # "a photo of a bear"
        tokenizer: CLIPTokenizer,
        size: int = 512,
    ):
        self.instance_data_root = Path(instance_data_root)
        self.instance_prompt = instance_prompt
        self.class_data_root = Path(class_data_root) if class_data_root else None
        self.class_prompt = class_prompt
        self.tokenizer = tokenizer
        self.size = size

        self.instance_images = sorted([
            f for f in self.instance_data_root.iterdir()
            if f.suffix.lower() in {'.jpg', '.jpeg', '.png', '.webp'}
        ])
        self.class_images = []
        if self.class_data_root and self.class_data_root.exists():
            self.class_images = sorted([
                f for f in self.class_data_root.iterdir()
                if f.suffix.lower() in {'.jpg', '.jpeg', '.png', '.webp'}
            ])

        self.num_instance = len(self.instance_images)
        self.num_class = len(self.class_images)
        print(f"Dataset: {self.num_instance} concept images, {self.num_class} class images")

    def __len__(self):
        return max(self.num_instance, self.num_class) if self.class_images else self.num_instance

    def load_and_preprocess(self, path: Path) -> torch.Tensor:
        """Load image, center-crop to square, resize, normalize to [-1, 1]."""
        image = Image.open(path).convert("RGB")
        w, h = image.size
        min_side = min(w, h)
        left = (w - min_side) // 2
        top = (h - min_side) // 2
        image = image.crop((left, top, left + min_side, top + min_side))
        image = image.resize((self.size, self.size), Image.LANCZOS)
        image = np.array(image).astype(np.float32) / 255.0
        image = (image - 0.5) * 2.0  # [0, 1] -> [-1, 1]
        return torch.from_numpy(image).permute(2, 0, 1)  # (3, H, W)

    def tokenize(self, prompt: str) -> torch.Tensor:
        return self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt"
        ).input_ids[0]

    def __getitem__(self, idx):
        instance_path = self.instance_images[idx % self.num_instance]
        result = {
            "instance_pixel_values": self.load_and_preprocess(instance_path),
            "instance_input_ids": self.tokenize(self.instance_prompt),
        }
        if self.class_images:
            class_path = self.class_images[idx % self.num_class]
            result["class_pixel_values"] = self.load_and_preprocess(class_path)
            result["class_input_ids"] = self.tokenize(self.class_prompt)
        return result


# ============================================================
# LoRA configuration - target cross-attention layers
# ============================================================

def apply_lora_to_unet(unet: UNet2DConditionModel, rank: int = 4) -> UNet2DConditionModel:
    """
    Apply LoRA to U-Net attention layers.

    Key design decision - which modules to target:
    - to_q, to_k, to_v, to_out.0: cross-attention projections
      These mediate how text conditioning influences spatial features.
      Most important for concept identity learning.
    - Optionally add self-attention for style fine-tuning.

    Rank selection: 4-8 for small datasets (5-20 images),
                    16-32 for larger style datasets (100+ images).
    """
    lora_config = LoraConfig(
        r=rank,
        lora_alpha=rank,  # alpha/r = 1.0 - standard scaling
        target_modules=[
            "to_q",      # Query projection (cross-attention)
            "to_k",      # Key projection (cross-attention)
            "to_v",      # Value projection (cross-attention)
            "to_out.0",  # Output projection (cross-attention)
        ],
        lora_dropout=0.0,
        bias="none",
    )
    unet = get_peft_model(unet, lora_config)
    unet.print_trainable_parameters()
    # Expected: ~3M trainable out of 860M total (SD 1.5, rank=4, cross-attn only)
    return unet


# ============================================================
# Training loop with prior preservation loss
# ============================================================

def train_dreambooth_lora(
    model_id: str = "runwayml/stable-diffusion-v1-5",
    instance_data_dir: str = "./concept_images",
    class_data_dir: str = "./class_images",
    instance_prompt: str = "a photo of sks bear",
    class_prompt: str = "a photo of a bear",
    output_dir: str = "./dreambooth-lora-output",
    num_train_steps: int = 1000,
    learning_rate: float = 1e-4,     # 1e-4 for LoRA; 1e-6 for full fine-tune
    lora_rank: int = 4,
    prior_loss_weight: float = 1.0,
    train_batch_size: int = 1,
    seed: int = 42,
):
    """
    Hyperparameter notes:
    - learning_rate 1e-4: appropriate for LoRA (small parameter count)
    - learning_rate 1e-6: appropriate for full U-Net fine-tune (large parameter count)
    - learning_rate 5e-7: appropriate for text encoder fine-tuning
    - Higher LR with LoRA is fine because LoRA has fewer parameters to diverge.
    - num_train_steps 800-1200: sweet spot for 15-25 concept images with prior preservation.
      Below 500: underfitting. Above 2000: overfitting.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(seed)

    # ---- Load models ----
    tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder")
    vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae")
    unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet")
    noise_scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")

    # ---- Apply LoRA ----
    unet = apply_lora_to_unet(unet, rank=lora_rank)

    # ---- Freeze everything except LoRA params ----
    vae.requires_grad_(False)
    text_encoder.requires_grad_(False)
    # U-Net: only LoRA params are trainable (PEFT manages this)

    vae = vae.to(device, dtype=torch.float32)
    text_encoder = text_encoder.to(device, dtype=torch.float32)
    unet = unet.to(device, dtype=torch.float32)

    # ---- Dataset ----
    dataset = DreamBoothDataset(
        instance_data_root=instance_data_dir,
        instance_prompt=instance_prompt,
        class_data_root=class_data_dir,
        class_prompt=class_prompt,
        tokenizer=tokenizer,
        size=512,
    )
    dataloader = DataLoader(dataset, batch_size=train_batch_size, shuffle=True)

    # ---- Optimizer: only LoRA params ----
    lora_params = [p for p in unet.parameters() if p.requires_grad]
    print(f"LoRA trainable parameters: {sum(p.numel() for p in lora_params):,}")

    optimizer = torch.optim.AdamW(
        lora_params,
        lr=learning_rate,
        betas=(0.9, 0.999),
        weight_decay=1e-2,
        eps=1e-8,
    )

    # ---- Training loop ----
    unet.train()
    global_step = 0

    while global_step < num_train_steps:
        for batch in dataloader:
            if global_step >= num_train_steps:
                break

            # ---- Encode concept images to latents ----
            pixel_values = batch["instance_pixel_values"].to(device)
            with torch.no_grad():
                latents = vae.encode(pixel_values).latent_dist.sample()
                latents = latents * vae.config.scaling_factor  # (B, 4, 64, 64)

            # ---- Sample noise and timestep ----
            noise = torch.randn_like(latents)
            timesteps = torch.randint(
                0, noise_scheduler.config.num_train_timesteps,
                (latents.shape[0],), device=device
            ).long()

            # ---- Forward diffusion: add noise ----
            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

            # ---- Text conditioning ----
            with torch.no_grad():
                encoder_hidden_states = text_encoder(
                    batch["instance_input_ids"].to(device)
                )[0]  # (B, 77, 768)

            # ---- Denoise prediction ----
            noise_pred = unet(
                noisy_latents, timesteps,
                encoder_hidden_states=encoder_hidden_states
            ).sample

            # ---- Concept loss (MSE on noise prediction) ----
            loss = F.mse_loss(noise_pred, noise, reduction="mean")

            # ---- Prior preservation loss ----
            if "class_pixel_values" in batch:
                class_pixel_values = batch["class_pixel_values"].to(device)
                with torch.no_grad():
                    class_latents = vae.encode(class_pixel_values).latent_dist.sample()
                    class_latents = class_latents * vae.config.scaling_factor

                class_noise = torch.randn_like(class_latents)
                class_timesteps = torch.randint(
                    0, noise_scheduler.config.num_train_timesteps,
                    (class_latents.shape[0],), device=device
                ).long()
                class_noisy = noise_scheduler.add_noise(class_latents, class_noise, class_timesteps)

                with torch.no_grad():
                    class_hidden_states = text_encoder(
                        batch["class_input_ids"].to(device)
                    )[0]

                class_noise_pred = unet(
                    class_noisy, class_timesteps,
                    encoder_hidden_states=class_hidden_states
                ).sample

                prior_loss = F.mse_loss(class_noise_pred, class_noise, reduction="mean")
                loss = loss + prior_loss_weight * prior_loss

            # ---- Optimization step ----
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(lora_params, 1.0)
            optimizer.step()

            global_step += 1
            if global_step % 100 == 0:
                print(f"Step {global_step}/{num_train_steps}, Loss: {loss.item():.4f}")

    # ---- Save LoRA weights ----
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    unet.save_pretrained(output_dir)
    print(f"LoRA weights saved to {output_dir}")

    return unet


# ============================================================
# Evaluation: DINO similarity for concept fidelity
# ============================================================

def compute_dino_similarity(
    reference_images: List[Image.Image],
    generated_images: List[Image.Image],
    device: str = "cuda"
) -> float:
    """
    Compute DINO ViT-B/16 cosine similarity between generated and reference images.
    Higher = better concept preservation. Target > 0.6 for good personalization.

    DINO features capture visual identity better than CLIP for fine-grained
    identity comparison (specific character, face, object instance).
    """
    import torchvision.transforms as T

    # Load DINO ViT-B/16 (self-supervised, better at identity than CLIP)
    dino = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')
    dino.eval().to(device)

    preprocess = T.Compose([
        T.Resize(256),
        T.CenterCrop(224),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    def extract_features(images):
        features = []
        for img in images:
            tensor = preprocess(img).unsqueeze(0).to(device)
            with torch.no_grad():
                feat = dino(tensor)  # (1, 768)
            feat = feat / feat.norm(dim=-1, keepdim=True)
            features.append(feat.cpu())
        return torch.cat(features, dim=0)  # (N, 768)

    ref_feats = extract_features(reference_images)     # (N_ref, 768)
    gen_feats = extract_features(generated_images)     # (N_gen, 768)

    # Average cosine similarity between all generated and all reference pairs
    sim_matrix = gen_feats @ ref_feats.T  # (N_gen, N_ref)
    mean_sim = sim_matrix.mean().item()
    return mean_sim

---

8. Evaluation - Measuring Personalization Quality

Evaluating personalization quality requires measuring two competing objectives that can trade off against each other:

Concept Preservation (did the model learn the specific identity?)

DINO-I similarity: extract ViT-B/16 (DINO) features from generated images and reference images, measure cosine similarity. Higher is better. DINO features capture semantic and visual identity better than CLIP for this task because DINO was trained with self-supervised objectives that preserve fine-grained visual structure.

Target: DINO-I > 0.6 indicates good identity preservation. Below 0.4 suggests the model has not captured the concept. Above 0.8 suggests possible overfitting (the generated images look too similar to training images).

CLIP-I (CLIP image-image similarity): cosine similarity between CLIP image embeddings of generated vs reference images. Less sensitive than DINO for fine-grained identity, but correlates with human judgments for coarser concepts.

Text Alignment (can the model still follow prompts?)

CLIP-T (text-image CLIP score): cosine similarity between the CLIP text embedding of the prompt and the CLIP image embedding of the generated image. Should be similar to or higher than the base model on the same prompts.

The key warning sign: CLIP-T significantly lower after fine-tuning than before. This means the model has overfit to the concept and can no longer interpret prompts correctly - "a sks bear on Mars" generates something that does not look like Mars.

The Pareto Tradeoff

Text Alignment (CLIP-T)
    ^
    |    Underfitting zone
    |    (model ignores concept,
    |     high text alignment,
    |     low DINO-I)
    |
    |         ★ Target zone
    |         (both DINO-I > 0.6
    |          and CLIP-T high)
    |
    |                    Overfitting zone
    |                    (model copies training images,
    |                     high DINO-I,
    |                     low CLIP-T)
    +---------------------------------> Concept Preservation (DINO-I)

The prior preservation loss shifts the operating point toward the upper-right by preventing overfitting. Without it, training always drifts toward the lower-right.

Practical Evaluation Protocol

1. Generate 50 images per test category:

   • Identity test: "a sks bear sitting outdoors" (5 seeds × 10 prompts)
   • Style generalization: "a sks bear in Van Gogh style", "a sks bear as LEGO"
   • Language drift test: "a polar bear" (no identifier - should be a generic polar bear)

2. Compute metrics:

   • DINO-I on identity test images vs reference images (target > 0.6)
   • CLIP-T on all generated images vs their prompts (compare to base model baseline)
   • Visual inspection of language drift test (should not show your mascot)

3. Declare success when DINO-I > 0.6 AND CLIP-T is within 0.03 of the base model AND the language drift test shows generic bears, not the mascot.

---

9. YouTube Resources

Video	Channel	What You Learn
DreamBooth Explained	Yannic Kilcher	Full paper walkthrough with prior preservation intuition and math
LoRA for Stable Diffusion	Fast.ai	LoRA mechanics, rank selection, practical tips, weight merging
ControlNet Deep Dive	AI Coffee Break	ControlNet architecture, zero-initialization, conditioning types
Textual Inversion Paper	Yannic Kilcher	Learning concept embeddings, CLIP gradient flow, limitations
DreamBooth + LoRA Training	Sebastian Raschka	Practical training tutorial with kohya-ss and diffusers

---

10. Production Engineering Notes

Checkpoint Formats and Portability

LoRA checkpoint (.safetensors): only the LoRA matrices, ~50-200MB. Can be combined with any compatible base model. Standard format for community sharing on CivitAI and HuggingFace. Load with load_attn_procs() in diffusers.

Full model checkpoint (.safetensors): complete fine-tuned weights (~1.7GB for SD 1.5). Required when full DreamBooth was used. Slowest to load, largest file, but most compatible with all inference tools.

Diffusers format (directory): model saved as a directory with separate files for U-Net, VAE, text encoder. Human-readable structure. Standard output of HuggingFace training scripts.

Compatibility: LoRAs trained for SD 1.5 are not compatible with SDXL. The U-Net architectures have different channel dimensions, attention layer structure (2 cross-attention in SD 1.5 vs 3 in SDXL), and layer count. Always check base model version before downloading community LoRAs.

Multi-Tenant LoRA Serving

For systems serving many user-defined LoRAs (e.g., an AI avatar platform with 50,000 user concepts):

1. Load base SD 1.5 U-Net weights once into GPU memory (~3.4GB fp16)
2. Per request: load user's LoRA adapter (~100MB) and merge:
   W' = W_0 + alpha * B @ A
3. Run inference with merged weights
4. LRU cache: keep top-N recently used merged models in GPU memory
5. On cache miss: unmerge (subtract LoRA) or reload base, then merge new LoRA

Unmerging is possible: $W_0 = W' - \alpha BA$ - you can revert to base model without reloading from disk.

LoRA Stacking at Inference

Multiple LoRA checkpoints can be combined at inference:

$W' = W_0 + \alpha_1 B_1 A_1 + \alpha_2 B_2 A_2 + \alpha_3 B_3 A_3$

This enables mixing a style LoRA ($\alpha_1$), a character LoRA ($\alpha_2$), and a lighting LoRA ($\alpha_3$). The mixing weights $\alpha_i$ can be tuned interactively at inference time.

Caveats: LoRAs trained for different purposes may interfere when they modify the same attention layers in conflicting ways. A face LoRA and a style LoRA may produce inconsistent results. Always test combinations before deploying.

Dataset Preparation Best Practices

Image quality: reference images should be sharp, well-lit, and at target resolution (512px or 768px). Blurry or JPEG-compressed references produce blurry outputs - the model can only learn what is clearly visible.

Diversity: use images from different angles, lighting conditions, and contexts. All-frontal-pose references will produce a model that struggles with side views. Aim for 30-40% profile, 30-40% three-quarter, 20-30% frontal for character concepts.

Caption quality: the training prompt should describe only the concept, not background details. "A photo of sks bear outdoors" is better than "A photo of sks bear sitting on a wooden bench in a green park with a fountain in the background." Over-descriptive captions cause the model to associate incidental details with the identifier token.

Aspect ratio: center-crop all images to square before training, or use Kohya-ss's multi-aspect ratio training mode which trains on multiple resolutions simultaneously for better generalization.

---

11. Common Pitfalls

:::danger Prior preservation images must come from the SAME base model The class prior images used in DreamBooth must be generated by the same base model you are fine-tuning, using the class prompt ("a bear"). Do not use real images from the internet as class priors - the domain mismatch between internet photos and SD-generated images will create a confounding gradient signal. Generate 200-300 class images before starting training using the original frozen model, then use those as the prior preservation set. :::

:::warning LoRA rank is not "higher is always better" Higher rank increases expressiveness but also overfitting risk on small datasets. For personalization with 5-20 images, rank 4-8 is typically optimal. Rank 64+ is appropriate for style transfer on larger datasets (1000+ images). The right rank is the lowest rank that achieves your target DINO-I score. Measure DINO-I and CLIP-T on a validation set rather than guessing. Overfitting manifests as DINO-I above 0.8 with CLIP-T dropping below baseline - the model generates training-image-like outputs regardless of prompt. :::

:::warning Overfitting without prior preservation looks like success initially After 800 steps of DreamBooth without prior preservation, generated images of your concept look excellent. But generating "a bear" now produces your mascot. Generating "a dog" may also produce something that resembles your mascot. This is language drift - the model's representation of bear has shifted toward the fine-tuned concept. Prior preservation loss is not optional for DreamBooth; it is essential for maintaining model utility beyond the specific concept. :::

:::danger Learning rates differ by 100x between full fine-tune and LoRA Full DreamBooth U-Net fine-tuning uses 1e-6 learning rate. LoRA uses 1e-4 (or sometimes higher). These are NOT interchangeable. Using 1e-4 for full fine-tune will cause catastrophic model divergence in the first 100 steps. Using 1e-6 for LoRA will cause extremely slow convergence and underfitting. The reason for the difference: LoRA trains ~3M parameters vs 860M. The effective learning rate per parameter is very different at the same learning rate value. :::

---

12. Interview Q&A

Q1: What is the prior preservation loss in DreamBooth and why is it necessary?

DreamBooth fine-tunes the full U-Net on 5-20 images. Without any regularization, the model quickly overfits - after 500 steps, it may only generate variations of the training images regardless of the prompt. More subtly, it undergoes "language drift": the token bear starts to become associated with your specific bear, so generating "a polar bear" or "a teddy bear" produces corrupted results.

The prior preservation loss addresses this by jointly training on class prior images (generated by the original frozen model using the class prompt) alongside the concept images. The objective has two terms: a concept loss on identifier-prompt pairs that teaches the model what sks bear looks like, and a prior loss on class-prompt pairs that maintains the model's knowledge of what a bear looks like in general. The prior loss acts as a regularizer that constrains the model updates to not deviate too far from the original class distribution. Setting $\lambda=1.0$ gives equal weight to both; increasing $\lambda$ prioritizes prior preservation at the cost of concept fidelity.

---

Q2: How does LoRA reduce memory requirements during DreamBooth training?

Full DreamBooth fine-tuning stores gradients and optimizer states (AdamW: gradient + two moment tensors) for all 860M U-Net parameters. At float32, optimizer states alone require approximately 860M × 3 × 4 bytes = ~10GB just for optimizer states, plus activations for backpropagation through all layers.

LoRA freezes $W_0$ and only trains $A$ and $B$ matrices. At rank 4 for SD 1.5, this is approximately 3M trainable parameters. Optimizer memory drops proportionally: ~3M × 3 × 4 bytes = ~36MB. The frozen $W_0$ does not require gradient computation - activations do not need to be stored for frozen layers, further reducing memory. Additionally, because LoRA has far fewer parameters, it fits comfortably within a 8-10GB VRAM budget (RTX 3080/3090) whereas full DreamBooth requires 24GB+ (A100).

---

Q3: Why does ControlNet use zero-initialized output convolutions?

ControlNet adds a trainable encoder copy that processes the conditioning input (e.g., Canny edges) and produces feature maps that are injected into the frozen SD decoder skip connections via zero-initialized $1\times1$ convolution projections.

At initialization, these projections output exactly zero - ControlNet contributes nothing to the base model. The base model generates normally. This achieves two goals: (1) the training starts from a stable, well-understood state where the pretrained model's capabilities are fully intact, giving the gradient a clean signal to work from; (2) the ControlNet must learn to produce useful features before any contribution enters the base model - it cannot hurt the base model until it has learned to help it. Without zero-initialization, random initial ControlNet outputs would corrupt the pretrained activations from step 1, creating a chaotic gradient signal that makes convergence slow and unstable.

---

Q4: How would you decide between Textual Inversion, DreamBooth+LoRA, and ControlNet for a production use case?

These techniques address different problems and are often combined rather than competing.

Textual Inversion: use when the concept is a style or texture (not a specific identity), checkpoint size must be minimal (~6KB), and you have no GPU budget beyond inference. Quality ceiling is low for complex identities.

DreamBooth + LoRA: use when you need to preserve a specific visual identity (character, product, face) across diverse contexts and prompts. The production standard - ~100MB checkpoints, runs on consumer hardware, stackable. Best absolute quality for identity-critical applications.

ControlNet: use when you need precise spatial control - exact pose, specific layout from a reference, edges from a sketch. Does not require training for your specific concept. Combined with DreamBooth+LoRA: generate your specific concept (DB+LoRA) in an exact pose from a reference sketch (ControlNet).

For the mascot example: DreamBooth+LoRA for the mascot identity. If they also need pose control from design sketches, add ControlNet (Canny or Scribble conditioning). These are loaded together at inference - they are not mutually exclusive.

---

Q5: How do you evaluate whether a DreamBooth/LoRA fine-tune has learned the concept correctly?

Use a three-part evaluation combining automatic metrics and visual inspection:

Identity preservation: generate 50 images using the identifier prompt in diverse contexts ("sks bear on Mars," "sks bear in Van Gogh style," "sks bear as a LEGO figure"). Compute DINO ViT-B/16 cosine similarity between these images and the reference images. Target DINO-I > 0.6. DINO is better than CLIP for fine-grained identity comparison.

Text alignment: compute CLIP-T (CLIP text-image cosine similarity) on the same 50 generated images using their respective prompts. Compare to base model CLIP-T on the same prompts. CLIP-T should not drop more than 0.03 below baseline - a larger drop indicates the model has overfit and can no longer follow prompts correctly.

Language drift test: generate 30 images with the class prompt only ("a bear," "a cartoon bear") using no identifier token. These should show generic bears, not your mascot. If your mascot appears, prior preservation has failed - reduce training steps or increase $\lambda$ (prior preservation weight).

---

Q6: What happens when LoRA rank is set too low or too high?

Rank too low (e.g., rank 1 for a complex concept): the LoRA update has insufficient expressiveness to capture all visual features of the concept. Generated images may capture the general category (a bear) but miss distinctive features (specific color pattern, distinctive face shape). The model effectively underfits. DINO-I will be low (below 0.5) even after full training.

Rank too high (e.g., rank 64 for a 10-image dataset): the LoRA update has enough capacity to memorize the training images. The model overfits rapidly - at 300-400 training steps, generated images start to look like direct variations of training images regardless of the prompt. DINO-I climbs above 0.8 while CLIP-T drops. The model loses prompt-following ability because the high-capacity LoRA update has overwritten the attention weights' ability to respond to diverse text conditioning. The fix: reduce rank or reduce training steps, then validate on the DINO-I / CLIP-T tradeoff.

---

13. Advanced Techniques

Kohya-ss Multi-Concept Training

When you need multiple concepts in one model (e.g., "person A" and "dog B"), single-concept DreamBooth is insufficient. Kohya-ss implements multi-concept training with separate identifier tokens per concept and balanced sampling across concept datasets.

The key challenge: preventing one concept from dominating training and causing the other to be poorly learned. Solution: balance sampling by concept (equal steps per concept per epoch), use separate class prior images per concept, and monitor DINO-I independently per concept.

Hypernetworks

A hypernetwork is a small secondary network that generates weight offsets for the main generation network. Instead of modifying U-Net weights directly (DreamBooth) or learning low-rank matrices (LoRA), the hypernetwork outputs weight adjustments conditioned on the current activation. Hypernetworks were popular in the NovelAI ecosystem. They have largely been superseded by LoRA in practice due to LoRA's better quality/size tradeoff and cleaner theoretical motivation.

Concept Mixing and Negative Concepts

LoRA can be applied with negative strength at inference to remove a concept from generation. Using a negative LoRA weight ($\alpha < 0$) pushes generated images away from the LoRA's learned concept. This is used in community workflows to suppress specific styles (e.g., "remove photorealism bias") or features (e.g., "remove watermarks") from base models.

---

This lesson is part of the Diffusion Models module. Next: Diffusion Models Beyond Images - Audio, Video, 3D, Molecules, Text.

:::tip 🎮 Interactive Playground

Visualize this concept: Try the Training Dynamics demo on the EngineersOfAI Playground - no code required.

:::