Fine-Tuning Text-to-Image Models

"I already know ten million faces. You want me to learn an eleventh, your dog, from five photos, without forgetting the other ten million. People call this 'fine-tuning'. I call it learning a single new word and promising not to bulldoze the dictionary."

A Pretrained Generator Negotiating Its Next Lesson

Prompting (Section 34.5) steers a model toward concepts it already knows. Fine-tuning teaches it concepts it does not. When you need your dog, this product, or a consistent character across a children's book (the studio from Section 34.3), no prompt suffices, because the concept is not in the training distribution. This section covers how to inject it. The transfer-learning principle is the one from Chapter 21: start from strong pretrained weights and adapt minimally, because adapting too much destroys the general capability you started from.

1. The Forgetting Problem and the Method Ladder Beginner

Naive fine-tuning has a signature failure: catastrophic forgetting. Train a billion-parameter generator on five photos of your dog and it will learn your dog and forget how to draw anything else, because the gradient from five images, applied to all parameters, overwrites the broad knowledge with the narrow. Every method on the ladder is a different answer to "how do I add the new concept without overwriting the old?" Table 34.6.1 lays out the ladder; the rest of the section explains each rung. The illustration below pictures the bargain: the heavier the change, the more of the existing library you risk smudging.

2. Textual Inversion: Learn One New Word Intermediate

The lightest method changes nothing in the model. Textual inversion freezes the entire generator and learns a single new embedding vector, a new "word" in the text encoder's vocabulary, that, when placed in a prompt, evokes the target concept. You introduce a pseudo-token like <my-dog>, initialize its embedding, and optimize only that vector so that "a photo of <my-dog>" reconstructs your training images under the frozen diffusion objective. Because the model is untouched, there is zero forgetting and the artifact is a few-kilobyte vector you can email. The cost is capacity: a single vector cannot capture fine appearance detail, so textual inversion is best for concepts that genuinely behave like a word (a recognizable style, a simple object) rather than a specific complex subject.

Mechanically it is the cleanest illustration of the whole chapter: it manipulates exactly the conditioning sequence of Section 34.1 and nothing else. The new word slots into the per-token embeddings, cross-attention (Section 34.2) routes it to the right spatial regions, and the frozen denoiser does the rest.

3. DreamBooth: Bind a Subject to a Rare Token Advanced

DreamBooth aims higher: faithful reproduction of a specific subject. It fine-tunes the whole model (or, in practice, combines with LoRA) to bind the subject to a rare token, "a photo of sks dog", where sks is chosen to have little prior meaning. The key idea that prevents forgetting is the prior-preservation loss. Alongside the few subject images, DreamBooth generates a batch of generic images from the model itself ("a photo of a dog") and includes them in training with their original captions. The model is thus simultaneously taught the new subject and reminded of the broad class, so it learns "sks dog is this particular dog" without collapsing the entire concept of "dog" onto your pet. The training objective is the standard diffusion loss plus the prior-preservation term:

$$ \mathcal{L} = \mathbb{E}\big[\|\epsilon - \epsilon_\theta(z_t, t, c_{\text{sks}})\|^2\big] + \lambda\,\mathbb{E}\big[\|\epsilon' - \epsilon_\theta(z'_{t'}, t', c_{\text{class}})\|^2\big], $$

where the first term learns the subject from your photos and the second, weighted by $\lambda$, preserves the class prior using the model's own generated class images. The prior-preservation term is the regularizer that guards the dictionary the epigraph worries about.

4. LoRA: The Practical Default Intermediate

LoRA (Low-Rank Adaptation) is where most practitioners live, because it is cheap, shareable, and resistant to forgetting. The insight, borrowed from language models, is that the weight update needed to adapt a model to a new concept is low-rank: instead of learning a full update $\Delta W$ to a weight matrix $W \in \mathbb{R}^{d \times k}$, learn two small matrices $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$ with rank $r \ll \min(d, k)$, and use $W + BA$ at inference. For a $1024 \times 1024$ matrix with $r = 8$, the two adapters hold $1024 \times 8$ entries in $B$ plus $8 \times 1024$ in $A$, so $2 \times 8192 = 16384$ trainable parameters replace the roughly one million in the full matrix, a 64-fold reduction. The base weights are frozen, so forgetting is minimal, and the adapter is a few megabytes you can publish and combine. LoRA is applied specifically to the cross-attention layers of Section 34.2, since that is where the text-to-image binding lives.

Figure 34.6.1 shows why LoRA is forgiving: the bulk of the model is frozen by construction, so the worst a bad LoRA can do is be ignored. The following script trains a LoRA on the cross-attention projections of an SD model, runnable on an 8 GB GPU.

import torch
from diffusers import StableDiffusionPipeline
from peft import LoraConfig, get_peft_model

pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
unet = pipe.unet

# Inject rank-8 LoRA adapters into the cross-attention projection matrices only.
lora = LoraConfig(r=8, lora_alpha=16,
                  target_modules=["to_q", "to_k", "to_v", "to_out.0"])
unet = get_peft_model(unet, lora)
unet.print_trainable_parameters()
# e.g. trainable params: 797,184 || all params: 860,318,148 || trainable%: 0.09

opt = torch.optim.AdamW(
    [p for p in unet.parameters() if p.requires_grad], lr=1e-4)

for step, (latents, text_embeds) in enumerate(my_dataloader):  # your subject images
    noise = torch.randn_like(latents)
    t = torch.randint(0, 1000, (latents.size(0),), device="cuda")
    noisy = pipe.scheduler.add_noise(latents, noise, t)         # forward process
    pred = unet(noisy, t, encoder_hidden_states=text_embeds).sample
    loss = torch.nn.functional.mse_loss(pred, noise)            # noise-prediction loss
    loss.backward(); opt.step(); opt.zero_grad()

unet.save_pretrained("my-subject-lora")   # a few MB, not the whole model

The print_trainable_parameters line is the whole pitch: under one percent of the model trains, which is why this fits on a consumer GPU in an afternoon and why the result is a shareable file rather than a model fork. Loading it back is one call, and adapters compose: a subject LoRA and a style LoRA can be applied together.

# Train via the official script (one shell command, abbreviated):
# accelerate launch train_dreambooth_lora.py \
#   --pretrained_model_name_or_path=stable-diffusion-v1-5/stable-diffusion-v1-5 \
#   --instance_data_dir=./my_dog --instance_prompt="a photo of sks dog" \
#   --rank=8 --output_dir=./my-dog-lora

from diffusers import StableDiffusionPipeline
import torch
pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("./my-dog-lora")            # one line to attach the adapter
image = pipe("a photo of sks dog wearing sunglasses on the beach").images[0]

The chapter ends where it began, at the assembly line. The Hands-On Lab below turns the three stations and three knobs of the opening key insight into a single program you build and run: it loads a generator, sweeps the prompt-engineering knobs of Section 34.5, swaps the model of Section 34.3, and finally bolts on a LoRA from this section so the same pipeline renders a concept the base model never saw.

pip install diffusers transformers accelerate torch pillow

import torch
from diffusers import StableDiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=dtype
).to(device)
pipe.set_progress_bar_config(disable=True)

# TODO: print the class name of the three stations so the pipeline is concrete:
#       the text encoder (pipe.text_encoder), the generator (pipe.unet),
#       and the decoder (pipe.vae). Confirm they match Section 34.2's diagram.

def generate(prompt, guidance=7.5, negative=None, seed=0, steps=30):
    gen = torch.Generator(device=device).manual_seed(seed)
    # TODO: call pipe(...) with prompt, negative_prompt=negative,
    #       guidance_scale=guidance, num_inference_steps=steps, generator=gen.
    #       Return result.images[0].
    ...

img = generate("a watercolor painting of a lighthouse at dawn", seed=0)
img.save("step2_base.png")
print("saved step2_base.png")

prompt = "a watercolor painting of a lighthouse at dawn"
scales = [1.5, 4.0, 7.5, 12.0, 20.0]

# TODO: generate one image per scale with the SAME seed (e.g. 0), collect
#       them in a list `row`, and label each with its guidance value.
row = []
for g in scales:
    ...
print(f"collected {len(row)} images across guidance {scales}")

neg = "blurry, low quality, extra fingers, watermark, text"

# TODO: generate two images on the same seed and guidance: one with
#       negative=None and one with negative=neg. Store them as a labeled pair.
pair = ...
print("generated the with/without negative-prompt pair")

from diffusers import DiffusionPipeline

pipe2 = DiffusionPipeline.from_pretrained(
    "stabilityai/sd-turbo", torch_dtype=dtype
).to(device)
pipe2.set_progress_bar_config(disable=True)

# TODO: render the same prompt and seed through pipe2. SD-Turbo is distilled
#       for very few steps, so use steps=1 and guidance=0.0. Note how the same
#       sentence yields a different look because the model knob changed.
gen2 = torch.Generator(device=device).manual_seed(0)
img_turbo = ...
img_turbo.save("step5_turbo.png")

# Use any SD-1.5-compatible LoRA on the Hub, or one you trained in Exercise 34.6.2.
# TODO: call pipe.load_lora_weights(REPO_OR_PATH) to attach the adapter, then
#       generate a prompt containing the LoRA's trigger token. Save the result.
...
img_lora = generate("a watercolor painting of a lighthouse, <trigger> style", seed=0)
img_lora.save("step6_lora.png")
pipe.unload_lora_weights()  # detach so later cells use the clean base model
print("saved step6_lora.png")

from PIL import Image, ImageDraw

def contact_sheet(labeled_images, cols=5, cell=256):
    rows = (len(labeled_images) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * cell, rows * (cell + 22)), "white")
    draw = ImageDraw.Draw(sheet)
    for i, (label, im) in enumerate(labeled_images):
        x, y = (i % cols) * cell, (i // cols) * (cell + 22)
        sheet.paste(im.resize((cell, cell)), (x, y + 22))
        draw.text((x + 4, y + 4), label, fill="black")
    return sheet

# TODO: build a list of (label, image) tuples from the guidance sweep (Step 3),
#       the negative-prompt pair (Step 4), and the model swap (Step 5), then
#       call contact_sheet(...) and save it as studio_contact_sheet.png.
...

import torch
from diffusers import StableDiffusionPipeline, DiffusionPipeline
from PIL import Image, ImageDraw

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# Step 1: load the pipeline and name the three stations.
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=dtype
).to(device)
pipe.set_progress_bar_config(disable=True)
print(type(pipe.text_encoder).__name__,
      type(pipe.unet).__name__,
      type(pipe.vae).__name__)

# Step 2: a reproducible generation helper.
def generate(prompt, guidance=7.5, negative=None, seed=0, steps=30):
    gen = torch.Generator(device=device).manual_seed(seed)
    return pipe(prompt, negative_prompt=negative, guidance_scale=guidance,
                num_inference_steps=steps, generator=gen).images[0]

prompt = "a watercolor painting of a lighthouse at dawn"

# Step 3: guidance sweep on one fixed seed.
scales = [1.5, 4.0, 7.5, 12.0, 20.0]
row = [generate(prompt, guidance=g, seed=0) for g in scales]

# Step 4: negative-prompt comparison on a fixed seed.
neg = "blurry, low quality, extra fingers, watermark, text"
pair = [("no negative", generate(prompt, seed=1, negative=None)),
        ("with negative", generate(prompt, seed=1, negative=neg))]

# Step 5: swap the generator behind the same idea.
pipe2 = DiffusionPipeline.from_pretrained(
    "stabilityai/sd-turbo", torch_dtype=dtype
).to(device)
pipe2.set_progress_bar_config(disable=True)
gen2 = torch.Generator(device=device).manual_seed(0)
img_turbo = pipe2(prompt, num_inference_steps=1, guidance_scale=0.0,
                  generator=gen2).images[0]

# Step 6: attach a LoRA and render a new concept.
# Replace the repo id and trigger token with a real SD-1.5 LoRA.
try:
    pipe.load_lora_weights("ostris/watercolor_style_lora_sd15")
    img_lora = generate(prompt + ", watercolor style", seed=0)
    img_lora.save("step6_lora.png")
    pipe.unload_lora_weights()
except Exception as e:
    print("LoRA step skipped (supply a valid adapter):", e)

# Step 7: assemble the labeled contact sheet.
def contact_sheet(labeled_images, cols=5, cell=256):
    rows = (len(labeled_images) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * cell, rows * (cell + 22)), "white")
    draw = ImageDraw.Draw(sheet)
    for i, (label, im) in enumerate(labeled_images):
        x, y = (i % cols) * cell, (i // cols) * (cell + 22)
        sheet.paste(im.resize((cell, cell)), (x, y + 22))
        draw.text((x + 4, y + 4), label, fill="black")
    return sheet

cells = [(f"g={g}", im) for g, im in zip(scales, row)]
cells += pair
cells += [("sd-turbo", img_turbo)]
contact_sheet(cells).save("studio_contact_sheet.png")
print("saved studio_contact_sheet.png")