Text to image with classifier-free guidance

Open in Colab View on GitHub DownloadGPU

The model in notebook 02 learns what flowers look like, but we cannot tell it which flower to draw. In this notebook we give the model a caption with every image, so at the end we can type “a red rose” and get something red and rose-like. We also look at classifier-free guidance, the trick that makes a model follow its caption more closely.

The notebook expects one NVIDIA GPU. Training takes most of the time: about twenty minutes on an RTX 4080. The first run downloads the CLIP text encoder (about 1.7 GB) and the flower images (about 270 MB).

Install

%pip install -q "dew-ml[cuda13,streaming] @ git+https://github.com/AshishKumar4/dew"

Settings

UNCONDITIONAL_PROB is the fraction of training captions we replace with an empty one; the guidance section explains why. PROMPTS are the captions we sample at the end, and GUIDANCE_SCALES the guidance strengths we compare.

import os

IMAGE_SIZE = 64
BATCH_SIZE = 64
STEPS = 8000
LEARNING_RATE = 3e-4
UNCONDITIONAL_PROB = 0.12
SAMPLE_STEPS = 40
PROMPTS = ["a red rose", "a yellow sunflower", "a white daisy", "a purple iris"]
GUIDANCE_SCALES = [1.0, 3.0, 6.0]
RUN_DIR = "runs/03-text-to-image"
SEED = 0

# The tutorial test sets DEW_TUTORIAL_SMOKE=1 to run every cell in minutes on a CPU.
if os.environ.get("DEW_TUTORIAL_SMOKE") == "1":
    BATCH_SIZE, STEPS, SAMPLE_STEPS = 8, 4, 4

import jax
import matplotlib.pyplot as plt
import numpy as np

print(jax.devices())

import textwrap

def show_images(images, titles, columns=4):
    rows = (len(images) + columns - 1) // columns
    figure, axes = plt.subplots(rows, columns, figsize=(columns * 1.8, rows * 2.1))
    for axis, image, title in zip(np.ravel(axes), images, titles):
        axis.imshow(image)
        axis.set_title(textwrap.fill(title, 22), fontsize=7)
        axis.axis("off")
    plt.tight_layout()
    plt.show()

Turning text into numbers

A network cannot read words, so a text encoder turns each caption into a sequence of vectors first. We use the text half of CLIP (ViT-L/14), which was trained to match captions with images, so its vectors already say something about what a caption looks like. We keep CLIP frozen; only the diffusion model trains.

The InputSpec now has a condition as well as the image. Condition(encoder, field="text", unconditional="") says: read tokens from the batch field text, encode them with CLIP, and use the empty caption as “no caption”. The key textcontext is the keyword the model receives the encoded caption under.

from dew import Condition, Field, InputSpec
from dew.inputs import CLIPText

encoder = CLIPText.from_pretrained("openai/clip-vit-large-patch14")
inputs = InputSpec(
    sample=Field("image", (IMAGE_SIZE, IMAGE_SIZE, 3)),
    conditions={"textcontext": Condition(encoder, field="text", unconditional="")},
)
print("context length:", encoder.context)
print(inputs.tokenize(["a red rose"])["text"].keys())

Captioned data

The data is the same Hugging Face copy of Oxford Flowers as notebook 02, and every photo carries a caption written by an image captioning model (BLIP). Passing tokenize=inputs.tokenize to load turns each caption into CLIP tokens while the batch is built, so the batches carry token arrays under text instead of strings.

from dew.data import DataPartition, HFImages, Loading

data = HFImages(
    name="pranked03/flowers-blip-captions",
    image_size=IMAGE_SIZE,
    val_batches=0,
    loading=Loading(workers=0, threads=16, read_buffer=64),
).load(batch=BATCH_SIZE, tokenize=inputs.tokenize)

batch = next(iter(data.train(DataPartition())))
captions = encoder.captions(batch["text"])
show_images(batch["image"][:8], captions[:8])

The model and the objective

The model is the DiT from notebook 02. It takes the caption vectors through the textcontext keyword and mixes them into every layer.

DiffusionObjective gets two new arguments. unconditional_prob makes it swap the caption for the empty one on 12% of training rows, and guidance=CFG(3.0) is how its previews sample. The CLIP weights ride along in the training state as frozen parameters, so they are saved in the checkpoint but never updated.

import optax
from dew import Checkpoints, LocalTracker, Trainer, models, presets
from dew.objectives.diffusion import DiffusionObjective
from dew.sampling import CFG, EulerAncestral

process = presets.EDM()()
model = models.build(
    "simple_dit",
    patch_size=4,
    emb_features=256,
    num_layers=6,
    num_heads=4,
    output_channels=3,
    dtype="bfloat16",
    attention_impl="auto",
)
objective = DiffusionObjective(
    model, process, inputs,
    unconditional_prob=UNCONDITIONAL_PROB,
    ema_decay=0.999,
    sampler=EulerAncestral(), guidance=CFG(3.0), steps=SAMPLE_STEPS,
)
trainer = Trainer(
    objective, optax.adamw(LEARNING_RATE),
    key=jax.random.key(SEED),
    checkpoints=Checkpoints(RUN_DIR),
    tracker=LocalTracker(f"{RUN_DIR}/tracking"),
)

Training

state = trainer.fit(data, steps=STEPS, log_every=1000, checkpoint_every=STEPS)

import json

rows = [json.loads(line) for line in open(f"{RUN_DIR}/tracking/scalars.jsonl")]
rows = [row for row in rows if "train/loss" in row["scalars"]]
plt.figure(figsize=(6, 3))
plt.plot([row["step"] for row in rows], [row["scalars"]["train/loss"] for row in rows], marker=".")
plt.xlabel("step")
plt.ylabel("loss")
plt.show()

Classifier-free guidance

Because some captions were blanked during training, one model has learned two things: how to denoise a flower given its caption, and how to denoise a flower with no caption at all. At sampling time we ask it both questions at every step and push the answer away from the uncaptioned one:

$\hat{\epsilon} = \epsilon_{\text{uncond}} + s \, (\epsilon_{\text{cond}} - \epsilon_{\text{uncond}})$

With scale $s = 1$ this is the plain captioned prediction. Larger scales follow the caption harder and give up some variety. This is classifier-free guidance (Ho and Salimans, 2022), and CFG(scale) in Dew computes it.

objective.pipeline(state) packages the trained model, the process, the text encoder and the EMA weights into a TextToImage pipeline. Calling it with a list of prompts returns the images.

pipe = objective.pipeline(state)

rows = []
for scale in GUIDANCE_SCALES:
    out = pipe(PROMPTS * 2, steps=SAMPLE_STEPS, guidance=CFG(scale), sampler=EulerAncestral(), seed=1)
    rows.append(np.asarray(out.images))

Each row below is one guidance scale, from 1 at the top to 6 at the bottom. Each column is one prompt, and each prompt appears twice per row with different starting noise.

columns = len(PROMPTS) * 2
figure, axes = plt.subplots(len(GUIDANCE_SCALES), columns, figsize=(columns * 1.2, len(GUIDANCE_SCALES) * 1.3))
for row, (scale, images) in enumerate(zip(GUIDANCE_SCALES, rows)):
    for column in range(columns):
        axis = axes[row, column]
        axis.imshow(np.clip((images[column] + 1) / 2, 0, 1))
        axis.set_xticks([])
        axis.set_yticks([])
        if row == 0:
            axis.set_title(textwrap.fill((PROMPTS * 2)[column], 12), fontsize=7)
    axes[row, 0].set_ylabel(f"scale {scale:g}", fontsize=8)
plt.tight_layout()
plt.show()

When we ran this notebook, the colours at scale 1 only loosely followed the prompts: one “white daisy” came out dark red. At scale 3 every image had the colour its prompt asked for, and the sunflowers got a dark centre. At scale 6 the colours were the most saturated and the two samples of each prompt looked more alike, which is the variety that strong guidance gives up. Twenty minutes of training on 6,500 images is not enough for sharp flowers, but the caption already steered the colour and the rough shape.

Where to go next

recipes/diffusion/train.py runs this setup from the command line and writes run.json next to the checkpoints, so TextToImage.from_run(directory) can rebuild the pipeline later. For larger images, StableDiffusionVAE in dew.nn.autoencoders.sd_vae lets the model denoise small latents instead of pixels. Notebook 04 compares the samplers.