Diffusion Models Course, юнит 2: дообучение, guidance и условная генерация

Перевод юнита 2 курса Hugging Face о диффузионных моделях: быстрое сэмплирование DDIM, дообучение готовой модели на новых данных с накоплением градиентов, guidance по цвету и по CLIP, демо на Gradio и модель MNIST с условием на класс.

Полный переводс языка: английский

Оригинал: Diffusion Models Course — Unit 2: Fine-Tuning, Guidance and Conditioning (README + 2 notebooks) · курс, юнит 2: теория + 2 ноутбука, ≈3–4 ч с обучением · репозиторий huggingface/diffusion-models-class на 2026-09-17 (материал 2022 года)
Автор: Джонатан Уитакер (Jonathan Whitaker) и Льюис Танстолл (Lewis Tunstall), Hugging Face
Лицензия оригинала: Apache-2.0

Текст оригинала переведён целиком, его лицензия это разрешает.

К урокам курса

Это перевод юнита 2 курса Diffusion Models Course: вводный текст юнита, ноутбук «Fine-Tuning and Guidance» и ноутбук «Making a Class-Conditioned Diffusion Model». Предыдущий юнит — введение и модель с нуля, карта курса — в README оригинала.

Юнит отвечает на вопрос «как управлять тем, что рисует модель». Три ответа: дообучить на своих данных, подталкивать сэмплирование градиентом произвольной функции потерь (guidance) и обучить модель с условием (метка класса, текст). Guidance по CLIP — предок того, как сегодня работают текстовые генераторы; сам CLIP разобран в уроке 9.1.1, а classifier-free guidance, которым пользуются современные модели, — в уроке 9.1.2. Модели 256×256 из этого юнита на одной RTX 4080 сэмплируются за секунды, дообучение занимает минуты. В коде есть пара ошибок — мы их отметили.

Юнит 2: дообучение, guidance и условная генерация

Добро пожаловать в юнит 2 курса Hugging Face о диффузионных моделях! В этом юните вы научитесь использовать и адаптировать предобученные диффузионные модели по-новому. Вы также увидите, как создавать диффузионные модели, которые принимают дополнительные входы в качестве условия (conditioning), чтобы управлять процессом генерации.

Начинаем юнит 🚀

Шаги этого юнита:

Убедитесь, что записались на курс, чтобы получать уведомления о новых материалах.
Прочитайте материал ниже — обзор ключевых идей юнита.
Откройте ноутбук Fine-tuning and Guidance, чтобы дообучить существующую диффузионную модель на новом датасете с помощью библиотеки 🤗 Diffusers и изменить процедуру сэмплирования с помощью guidance.
По примеру из ноутбука опубликуйте демо на Gradio для своей модели.
(Необязательно) Откройте ноутбук Class-conditioned Diffusion Model Example, чтобы увидеть, как добавить дополнительный контроль над процессом генерации.
(Необязательно) Посмотрите это видео — неформальный разбор материала юнита.

📢 Не забудьте присоединиться к Discord, где можно обсуждать материал и делиться тем, что получилось, в канале #diffusion-models-class.

Дообучение

Как вы, возможно, видели в юните 1, обучать диффузионные модели с нуля бывает долго! Особенно когда мы переходим к высоким разрешениям: время и объём данных для обучения модели с нуля могут стать неподъёмными. К счастью, есть решение: начать с уже обученной модели! Тогда мы стартуем с модели, которая уже научилась очищать от шума какие-то изображения, и надеемся, что это лучшая отправная точка, чем модель со случайной инициализацией.

(В оригинале — примеры изображений модели, обученной на спальнях из LSUN и дообученной 500 шагов на WikiArt: размытые «картины», в которых ещё угадываются очертания комнат.)

Дообучение обычно работает лучше всего, если новые данные чем-то похожи на исходные обучающие данные базовой модели (например, если вы хотите генерировать мультяшные лица, начать с модели, обученной на лицах, — наверное, хорошая идея), но, как ни удивительно, польза сохраняется, даже если область меняется очень сильно. Изображение выше сгенерировано моделью, обученной на датасете LSUN Bedrooms и дообученной 500 шагов на датасете WikiArt. Скрипт обучения для справки лежит рядом с ноутбуками этого юнита.

Guidance («направление»)

Безусловные модели почти не дают контроля над тем, что генерируется. Можно обучить условную модель (о ней — в следующем разделе), которая принимает дополнительные входы, помогающие направлять генерацию. Но что если у нас уже есть обученная безусловная модель, которой мы хотим воспользоваться? Здесь на сцену выходит guidance — процесс, при котором предсказания модели на каждом шаге генерации оцениваются некоторой функцией направления (guidance function) и изменяются так, чтобы итоговое изображение больше нам нравилось.

(В оригинале — пример: одни и те же сэмплы без guidance и с guidance к бирюзовому цвету.)

Функцией направления может быть почти что угодно, и это делает технику мощной! В ноутбуке мы идём от простого примера (управление цветом, как на картинке выше) к примеру с мощной предобученной моделью CLIP, которая позволяет направлять генерацию по текстовому описанию.

Условие (conditioning)

Guidance — отличный способ выжать дополнительную пользу из безусловной диффузионной модели, но если при обучении у нас есть дополнительная информация (например, метка класса или подпись к изображению), её тоже можно подавать в модель, чтобы она использовала её в своих предсказаниях. Так мы получаем условную модель, которой можно управлять при инференсе, выбирая, что подать в качестве условия. В ноутбуке показан пример модели с условием на класс, которая учится генерировать изображения по метке класса.

(В оригинале — сетка цифр MNIST, сгенерированных условной моделью: в каждой строке — одна и та же цифра от 0 до 9.)

Передать информацию об условии можно разными способами, например:

Подать её дополнительными каналами на вход U-Net. Так часто делают, когда условие имеет ту же форму, что и изображение: маска сегментации, карта глубины или размытая версия изображения (в модели восстановления или сверхразрешения). Это работает и для других видов условий. Например, в ноутбуке метка класса отображается в эмбеддинг, а затем растягивается до ширины и высоты входного изображения, чтобы её можно было подать дополнительными каналами.
Создать эмбеддинг, спроецировать его до размера, равного числу каналов на выходе одного или нескольких внутренних слоёв U-Net, и прибавить к этим выходам. Так, например, устроено условие на временной шаг: к выходу каждого блока ResNet прибавляется спроецированный эмбеддинг шага. Это удобно, когда условие — вектор, например эмбеддинг изображения CLIP. Заметный пример — версия Stable Diffusion «Image Variations», которая делает именно это.
Добавить слои перекрёстного внимания (cross-attention), которые «смотрят» на последовательность, поданную как условие. Это полезнее всего, когда условие — текст: текст превращается трансформером в последовательность эмбеддингов, а слои перекрёстного внимания в U-Net встраивают эту информацию в путь очистки от шума. Мы увидим это в действии в юните 3, когда будем разбирать, как Stable Diffusion работает с текстовым условием.

flowchart TB
subgraph A["1. Доп. каналы на входе"]
  a1["картинка 1 канал"] --> a3["concat → U-Net"]
  a2["условие, растянутое до H×W"] --> a3
end
subgraph B["2. Сложение с признаками"]
  b1["эмбеддинг условия"] --> b2["линейная проекция"] --> b3["+ к выходу блока ResNet"]
end
subgraph C["3. Перекрёстное внимание"]
  c1["текст → токены → эмбеддинги"] --> c2["K, V для cross-attention"]
  c3["признаки картинки"] --> c4["Q"] --> c5["внимание: картинка смотрит на текст"]
  c2 --> c5
end

Три способа подать условие в U-Net (рисунок переводчика)

Практические ноутбуки

Глава	Где запустить
Fine-tuning and Guidance	Colab, Kaggle
Class-conditioned Diffusion Model Example	Colab, Kaggle

Теперь вы знаете достаточно, чтобы взяться за ноутбуки! Откройте их на удобной платформе по ссылкам выше. Дообучение довольно требовательно к вычислениям, поэтому в Kaggle или Google Colab обязательно выберите тип среды «GPU».

Основная часть материала — в Fine-tuning and Guidance: там обе темы разобраны на примерах. Ноутбук показывает, как дообучить существующую модель на новых данных, добавить guidance и опубликовать результат как демо на Gradio. К нему прилагается скрипт (finetune_model.py), с которым легко экспериментировать с разными настройками дообучения, и пример Space, который можно взять за шаблон своего демо на 🤗 Spaces.

В Class-conditioned Diffusion Model Example мы кратко показываем на датасете MNIST, как создать диффузионную модель с условием на метки классов. Цель — показать главную идею как можно проще: давая модели дополнительную информацию о том, что она должна очищать от шума, мы потом можем управлять тем, какие изображения генерируются при инференсе.

Время проекта

По примерам из ноутбука Fine-tuning and Guidance дообучите свою модель или возьмите готовую и сделайте демо на Gradio, которое показывает ваши новые навыки guidance. Не забудьте поделиться демо в Discord, Twitter и т. д., чтобы мы могли полюбоваться вашей работой!

Дополнительные материалы

Denoising Diffusion Implicit Models — статья, где предложен метод сэмплирования DDIM (его использует DDIMScheduler).

GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models — предложены способы подавать в диффузионные модели текстовое условие.

eDiffi: Text-to-Image Diffusion Models with an Ensemble of Expert Denoisers — показано, как сочетать много разных видов условий, чтобы ещё сильнее управлять генерацией.

Нашли ещё отличные материалы? Сообщите нам, и мы добавим их в список.

Ноутбук: дообучение и guidance

В этом ноутбуке мы разберём два основных подхода к адаптации существующих диффузионных моделей:

с помощью дообучения (fine-tuning) мы заново обучим существующие модели на новых данных, чтобы изменить тип результатов, которые они выдают;
с помощью guidance мы возьмём существующую модель и будем направлять процесс генерации во время инференса, чтобы получить дополнительный контроль.

Чему вы научитесь

К концу ноутбука вы будете уметь:

создавать цикл сэмплирования и быстрее генерировать примеры с новым планировщиком;
дообучать существующую диффузионную модель на новых данных, в том числе:
- с накоплением градиентов, чтобы обойти часть проблем маленьких батчей;
- с логированием примеров в Weights and Biases во время обучения, чтобы следить за прогрессом (через прилагаемый скрипт);
- сохранять получившийся конвейер и выгружать его на Hub;
направлять процесс сэмплирования дополнительными функциями потерь, чтобы управлять существующими моделями, в том числе:
- исследовать разные подходы к guidance на простой функции потерь по цвету;
- направлять генерацию текстовым промптом с помощью CLIP;
- делиться своим циклом сэмплирования через Gradio и 🤗 Spaces.

❓ Если появятся вопросы, пишите в канал #diffusion-models-class на Discord-сервере Hugging Face. Если ещё не зарегистрировались, это можно сделать здесь: https://huggingface.co/join/discord

Подготовка и импорты

Чтобы сохранять дообученные модели на Hugging Face Hub, нужно войти с токеном, у которого есть право записи. Код ниже попросит его и даст ссылку на страницу токенов вашего аккаунта. Аккаунт Weights and Biases понадобится, если вы захотите логировать примеры по ходу обучения через скрипт, — код тоже предложит войти, где нужно.

Кроме этого, нужно только установить несколько зависимостей, импортировать всё необходимое и указать устройство:

%pip install -qq diffusers datasets accelerate wandb open-clip-torch

# Код для входа в Hugging Face Hub — нужен, чтобы делиться моделями
# Используйте токен с правом записи (WRITE)
from huggingface_hub import notebook_login

notebook_login()

import numpy as np
import torch
import torch.nn.functional as F
import torchvision
from datasets import load_dataset
from diffusers import DDIMScheduler, DDPMPipeline
from matplotlib import pyplot as plt
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm

device = (
    "mps"
    if torch.backends.mps.is_available()
    else "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

Загружаем предобученный конвейер

Для начала загрузим существующий конвейер и посмотрим, что с ним можно сделать:

image_pipe = DDPMPipeline.from_pretrained("google/ddpm-celebahq-256")
image_pipe.to(device);

Сгенерировать изображения так же просто, как вызвать метод __call__ конвейера — вызвать сам конвейер как функцию:

images = image_pipe().images
images[0]

(Модель обучена на CelebA-HQ — фотографиях лиц знаменитостей, — поэтому генерирует лица 256×256.)

Красиво, но МЕДЛЕННО! Поэтому, прежде чем переходить к главным темам, заглянем в сам цикл сэмплирования и посмотрим, как ускорить его более продвинутым сэмплером.

Быстрее сэмплирование с DDIM

На каждом шаге модели подают зашумлённый вход и просят предсказать шум (а значит, и оценку того, как может выглядеть полностью очищенное изображение). Поначалу эти предсказания не очень хороши, поэтому процесс и разбивают на много шагов. Однако оказалось, что 1000+ шагов не нужны, и волна недавних исследований изучает, как получать хорошие примеры за как можно меньшее число шагов.

В библиотеке 🤗 Diffusers методы сэмплирования реализует планировщик, который делает каждое обновление функцией step(). Чтобы сгенерировать изображение, мы начинаем со случайного шума $x$ . Затем на каждом шаге расписания шума планировщика подаём зашумлённый вход $x$ в модель и передаём полученное предсказание в функцию step(). Она возвращает результат с атрибутом prev_sample — «предыдущий», потому что мы идём «назад» во времени, от сильного шума к слабому (противоположно прямому процессу диффузии).

Посмотрим на это в действии! Сначала загрузим планировщик — здесь DDIMScheduler по статье Denoising Diffusion Implicit Models, который даёт приличные примеры за гораздо меньшее число шагов, чем исходная реализация DDPM:

# Создаём новый планировщик и задаём число шагов инференса
scheduler = DDIMScheduler.from_pretrained("google/ddpm-celebahq-256")
scheduler.set_timesteps(num_inference_steps=40)

Видно, что модель делает всего 40 шагов, каждый из которых перепрыгивает 25 шагов исходного расписания из 1000:

scheduler.timesteps

Вывод:

tensor([975, 950, 925, 900, 875, 850, 825, 800, 775, 750, 725, 700, 675, 650,
        625, 600, 575, 550, 525, 500, 475, 450, 425, 400, 375, 350, 325, 300,
        275, 250, 225, 200, 175, 150, 125, 100,  75,  50,  25,   0])

Создадим 4 случайных изображения и пройдём цикл сэмплирования, глядя и на текущий $x$ , и на предсказанную очищенную версию по мере продвижения:

# Случайная стартовая точка
x = torch.randn(4, 3, 256, 256).to(device)  # батч из 4 трёхканальных изображений 256×256

# Проходим по шагам сэмплирования
for i, t in tqdm(enumerate(scheduler.timesteps)):

    # Готовим вход модели
    model_input = scheduler.scale_model_input(x, t)

    # Получаем предсказание
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]

    # Считаем планировщиком, как должен выглядеть обновлённый сэмпл
    scheduler_output = scheduler.step(noise_pred, t, x)

    # Обновляем x
    x = scheduler_output.prev_sample

    # Время от времени показываем и x, и предсказанные очищенные изображения
    if i % 10 == 0 or i == len(scheduler.timesteps) - 1:
        fig, axs = plt.subplots(1, 2, figsize=(12, 5))

        grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)
        axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[0].set_title(f"Current x (step {i})")

        pred_x0 = (
            scheduler_output.pred_original_sample
        )  # есть не у всех планировщиков
        grid = torchvision.utils.make_grid(pred_x0, nrow=4).permute(1, 2, 0)
        axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
        axs[1].set_title(f"Predicted denoised images (step {i})")
        plt.show()

(В оригинале — четыре пары картинок: на шаге 0 текущий $x$ — чистый шум, а предсказание — размытое пятно в форме лица; к последнему шагу оба превращаются в чёткие лица.)

Видно, что начальные предсказания не очень хороши, но по ходу процесса предсказанные результаты становятся всё точнее. Если интересно, какая математика происходит внутри функции step(), посмотрите её (хорошо прокомментированный) код:

# ??scheduler.step

Внутри step() у DDIM две строки математики. Сначала по предсказанному шуму восстанавливается оценка чистой картинки: $\hat x_0 = (x_t - \sqrt{1-\bar\alpha_t}\,\hat\varepsilon)/\sqrt{\bar\alpha_t}$ . Потом эта оценка снова «зашумляется», но уже до более раннего шага $s$ , с тем же предсказанным шумом: $x_s = \sqrt{\bar\alpha_s}\,\hat x_0 + \sqrt{1-\bar\alpha_s}\,\hat\varepsilon$ . Случайного шума при этом не добавляется (при eta=0), поэтому DDIM детерминирован и может перепрыгивать через много шагов сразу. Именно это свойство используется в юните 4 для «инверсии» DDIM.

Новый планировщик можно подставить вместо исходного, который шёл в комплекте с конвейером, и сэмплировать так:

image_pipe.scheduler = scheduler
images = image_pipe(num_inference_steps=40).images
images[0]

Отлично — теперь мы получаем примеры за разумное время! Это ускорит работу в оставшейся части ноутбука :)

Дообучение

Самое интересное! Как, имея этот предобученный конвейер, переобучить модель генерировать изображения по новым обучающим данным?

Оказывается, это выглядит почти так же, как обучение модели с нуля (как мы видели в юните 1), только мы начинаем с существующей модели. Посмотрим на это в действии и по ходу обсудим несколько дополнительных соображений.

Сначала датасет: можно попробовать этот датасет винтажных лиц или эти лица в стиле аниме — что-то ближе к исходным обучающим данным этой модели лиц, — но ради интереса возьмём тот же маленький датасет бабочек, на котором мы обучали модель с нуля в юните 1. Выполните код ниже, чтобы скачать датасет бабочек и создать загрузчик, из которого можно брать батчи изображений:

# @markdown загружаем и готовим датасет:
# Не в Colab? Комментарии с #@ включают в интерфейсе Colab заголовки и поля ввода,
# на других платформах их можно спокойно игнорировать.

dataset_name = "huggan/smithsonian_butterflies_subset"  # @param
dataset = load_dataset(dataset_name, split="train")
image_size = 256  # @param
batch_size = 4  # @param
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ]
)


def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}


dataset.set_transform(transform)

train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True
)

print("Previewing batch:")
batch = next(iter(train_dataloader))
grid = torchvision.utils.make_grid(batch["images"], nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);

Соображение 1: размер батча (4) здесь довольно маленький, потому что мы обучаем на крупных изображениях (256 пикселей) довольно большую модель и упрёмся в память GPU, если сильно увеличить батч. Можно уменьшить размер изображений, чтобы ускорить обучение и взять батч побольше, но эти модели спроектированы и исходно обучены для генерации 256×256.

Теперь цикл обучения. Мы будем обновлять веса предобученной модели, указав целью оптимизации image_pipe.unet.parameters(). Остальное почти совпадает с примером цикла обучения из юнита 1. В Colab это занимает около 10 минут, так что самое время сходить за кофе или чаем:

num_epochs = 2  # @param
lr = 1e-5  # @param
grad_accumulation_steps = 2  # @param

optimizer = torch.optim.AdamW(image_pipe.unet.parameters(), lr=lr)

losses = []

for epoch in range(num_epochs):
    for step, batch in tqdm(enumerate(train_dataloader), total=len(train_dataloader)):
        clean_images = batch["images"].to(device)
        # Выбираем шум, который добавим к изображениям
        noise = torch.randn(clean_images.shape).to(clean_images.device)
        bs = clean_images.shape[0]

        # Выбираем случайный временной шаг для каждого изображения
        timesteps = torch.randint(
            0,
            image_pipe.scheduler.num_train_timesteps,
            (bs,),
            device=clean_images.device,
        ).long()

        # Добавляем к чистым изображениям шум с силой, соответствующей шагу
        # (это прямой процесс диффузии)
        noisy_images = image_pipe.scheduler.add_noise(clean_images, noise, timesteps)

        # Получаем предсказание шума моделью
        noise_pred = image_pipe.unet(noisy_images, timesteps, return_dict=False)[0]

        # Сравниваем предсказание с настоящим шумом:
        loss = F.mse_loss(
            noise_pred, noise
        )  # NB: предсказываем шум (eps), а не (noisy_ims - clean_ims) и не просто (clean_ims)

        # Сохраняем для графика
        losses.append(loss.item())

        # Обновляем параметры модели оптимизатором по этому loss
        loss.backward(loss)

        # Накопление градиентов:
        if (step + 1) % grad_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

    print(
        f"Epoch {epoch} average loss: {sum(losses[-len(train_dataloader):])/len(train_dataloader)}"
    )

# Рисуем кривую loss:
plt.plot(losses)

Вывод:

Epoch 0 average loss: 0.013324214214226231
Epoch 1 average loss: 0.014018508377484978

(Примечание переводчика: та же ошибка, что в юните 1, — loss.backward(loss) умножает градиенты на значение loss; пишите loss.backward(). И image_pipe.scheduler.num_train_timesteps в новых версиях diffusers лучше писать как image_pipe.scheduler.config.num_train_timesteps. Полное дообучение U-Net на 256×256 с батчем 4 помещается на одну RTX 4080; две эпохи по 1000 картинок — несколько минут.)

Соображение 2: сигнал loss у нас очень шумный, ведь на каждом шаге мы работаем всего с четырьмя примерами при случайных уровнях шума. Для обучения это не идеально. Одно из решений — очень низкая скорость обучения, чтобы ограничить размер обновления на каждом шаге. А ещё лучше найти способ получить ту же пользу, что дал бы больший батч, не взвинчивая требования к памяти…

Здесь на сцену выходит накопление градиентов (gradient accumulation). Если вызвать loss.backward() несколько раз, прежде чем вызвать optimizer.step() и optimizer.zero_grad(), PyTorch накапливает (суммирует) градиенты, по сути объединяя сигнал нескольких батчей в одну (лучшую) оценку, по которой и обновляются параметры. Обновлений получается меньше — так же, как при большем батче. Многие фреймворки делают это за вас (например, 🤗 Accelerate позволяет это легко), но полезно увидеть реализацию с нуля — это полезная техника для обучения при ограниченной памяти GPU! Как видно из кода выше (после комментария # Накопление градиентов), кода для этого почти не нужно.

# Упражнение: попробуйте добавить накопление градиентов в цикл обучения из юнита 1.
# Как оно работает? Подумайте, как стоит менять скорость обучения в зависимости
# от числа шагов накопления, — должна ли она остаться прежней?

(Подсказка от переводчика: градиенты суммируются, поэтому при $k$ шагах накопления и loss, усреднённом внутри батча, градиент в $k$ раз больше, чем у одного батча. Чтобы получить ровно «большой батч», делят loss на $k$ : (loss / k).backward(). С AdamW это почти не важно — он нормирует шаг, — а с SGD без деления пришлось бы уменьшать скорость обучения в $k$ раз. Подробнее — в уроке 4.6.3.)

Соображение 3: это всё ещё долго, и строки раз в эпоху мало, чтобы понимать, что происходит. Стоит:

время от времени генерировать примеры, чтобы глазами оценивать качество по ходу обучения;
логировать loss и сгенерированные примеры во время обучения, например в Weights and Biases или tensorboard.

Я сделал небольшой скрипт (finetune_model.py), который берёт код обучения выше и добавляет минимальное логирование. Ниже — логи одного запуска обучения:

%wandb johnowhitaker/dm_finetune/2upaa341 # нужен аккаунт W&B — пропустите, если не хотите входить

Интересно смотреть, как меняются сгенерированные примеры по ходу обучения: хотя loss почти не улучшается, видно движение от исходной области (изображения спален) к новым обучающим данным (wikiart). В конце этого ноутбука есть закомментированный код для дообучения модели этим скриптом — как альтернатива ячейке выше.

# Упражнение: попробуйте изменить официальный учебный скрипт, который мы видели
# в юните 1, чтобы он начинал с предобученной модели, а не обучал с нуля.
# Сравните с минимальным скриптом по ссылке выше — каких возможностей в минимальном скрипте не хватает?

Сгенерировав несколько изображений этой моделью, видим, что лица уже выглядят очень странно!

# @markdown Генерируем и рисуем несколько изображений:
x = torch.randn(8, 3, 256, 256).to(device)  # батч из 8
for i, t in tqdm(enumerate(scheduler.timesteps)):
    model_input = scheduler.scale_model_input(x, t)
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]
    x = scheduler.step(noise_pred, t, x).prev_sample
grid = torchvision.utils.make_grid(x, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);

(В оригинале — восемь картинок: смесь лиц и крыльев бабочек, яркие пятна на месте глаз.)

Соображение 4: дообучение бывает весьма непредсказуемым! Если обучать гораздо дольше, возможно, получатся идеальные бабочки. Но промежуточные этапы сами по себе бывают очень интересны, особенно если вас больше привлекает художественная сторона! Попробуйте обучать совсем недолго или очень долго, меняйте скорость обучения и смотрите, как это влияет на то, что выдаёт итоговая модель.

Код дообучения минимальным скриптом, которым сделана демо-модель на WikiArt

Если хотите обучить модель, похожую на мою на WikiArt, раскомментируйте и выполните ячейки ниже. Это занимает время и может исчерпать память GPU, поэтому рекомендую делать это после того, как пройдёте остальной ноутбук.

## Скачать скрипт дообучения:
# !wget https://github.com/huggingface/diffusion-models-class/raw/main/unit2/finetune_model.py

## Запуск скрипта: дообучаем модель лиц на винтажных лицах
## (лучше запускать в терминале):
# !python finetune_model.py --image_size 128 --batch_size 8 --num_epochs 16\
#     --grad_accumulation_steps 2 --start_model "google/ddpm-celebahq-256"\
#     --dataset_name "Norod78/Vintage-Faces-FFHQAligned" --wandb_project 'dm-finetune'\
#     --log_samples_every 100 --save_model_every 1000 --model_save_name 'vintageface'

Сохранение и загрузка дообученных конвейеров

Теперь, когда мы дообучили U-Net в нашей диффузионной модели, сохраним её в локальную папку:

image_pipe.save_pretrained("my-finetuned-model")

Как мы видели в юните 1, сохранятся конфигурация, модель и планировщик:

!ls {"my-finetuned-model"}

Вывод:

model_index.json  scheduler  unet

Дальше можно повторить шаги из ноутбука Введение в Diffusers юнита 1, чтобы выгрузить модель на Hub для дальнейшего использования:

# @title Выгружаем на Hub конвейер, сохранённый локально

# Код выгрузки на Hub конвейера, сохранённого локально
from huggingface_hub import HfApi, ModelCard, create_repo, get_full_repo_name

# Создаём репозиторий и выгружаем файлы
model_name = "ddpm-celebahq-finetuned-butterflies-2epochs"  # @param как модель будет называться на Hub
local_folder_name = "my-finetuned-model"  # @param создана скриптом или вами через image_pipe.save_pretrained('save_name')
description = "Describe your model here"  # @param
hub_model_id = get_full_repo_name(model_name)
create_repo(hub_model_id)
api = HfApi()
api.upload_folder(
    folder_path=f"{local_folder_name}/scheduler", path_in_repo="", repo_id=hub_model_id
)
api.upload_folder(
    folder_path=f"{local_folder_name}/unet", path_in_repo="", repo_id=hub_model_id
)
api.upload_file(
    path_or_fileobj=f"{local_folder_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

# Добавляем карточку модели (необязательно, но приятно!)
content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

# Example Fine-Tuned Model for Unit 2 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

{description}

## Usage

```python
from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained('{hub_model_id}')
image = pipeline().images[0]
image
```
"""

card = ModelCard(content)
card.push_to_hub(hub_model_id)

(Как и в юните 1: надёжнее выгружать всю папку целиком — api.upload_folder(folder_path=local_folder_name, repo_id=hub_model_id), — иначе файлы config.json планировщика и U-Net перезапишут друг друга в корне репозитория.)

Поздравляем, вы дообучили свою первую диффузионную модель!

Дальше в ноутбуке мы будем использовать модель, которую я дообучил примерно одну эпоху на датасете WikiArt, начав с этой модели, обученной на спальнях LSUN. Если хотите, пропустите эту ячейку и используйте конвейер лиц/бабочек, который мы дообучили в прошлом разделе, или загрузите любой с Hub:

# Загружаем предобученный конвейер
pipeline_name = "johnowhitaker/sd-class-wikiart-from-bedrooms"
image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)

# Сэмплируем несколько изображений планировщиком DDIM за 40 шагов
scheduler = DDIMScheduler.from_pretrained(pipeline_name)
scheduler.set_timesteps(num_inference_steps=40)

# Случайная стартовая точка (батч из 8 изображений)
x = torch.randn(8, 3, 256, 256).to(device)

# Минимальный цикл сэмплирования
for i, t in tqdm(enumerate(scheduler.timesteps)):
    model_input = scheduler.scale_model_input(x, t)
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]
    x = scheduler.step(noise_pred, t, x).prev_sample

# Смотрим результат
grid = torchvision.utils.make_grid(x, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);

Соображение 5: часто трудно понять, насколько хорошо работает дообучение, а что такое «хорошее качество», зависит от задачи. Например, дообучая модель с текстовым условием вроде Stable Diffusion на маленьком датасете, вы, скорее всего, хотите, чтобы она сохранила большую часть исходного обучения и понимала произвольные промпты, которых нет в новом датасете, и при этом адаптировалась к стилю новых данных. Это может означать низкую скорость обучения вместе с чем-то вроде экспоненциального усреднения модели, как показано в отличном блог-посте о создании «покемонной» версии Stable Diffusion. В другой ситуации, наоборот, нужно полностью переобучить модель на новых данных (как в нашем примере спальни → wikiart), и тогда разумны большая скорость обучения и долгое обучение. Хотя график loss почти не показывает улучшения, примеры явно уходят от исходных данных к более «художественным» результатам, хотя и остаются в основном бессвязными.

Это подводит нас к следующему разделу: посмотрим, как добавить к такой модели guidance, чтобы лучше управлять результатами…

Guidance

Что делать, если мы хотим управлять генерируемыми примерами? Скажем, хотим, чтобы изображения тяготели к определённому цвету. Как это сделать? Здесь на сцену выходит guidance — техника добавления дополнительного контроля в процесс сэмплирования.

Первый шаг — создать функцию условия: некоторую меру (loss), которую мы хотим минимизировать. Вот такая функция для примера с цветом: она сравнивает пиксели изображения с целевым цветом (по умолчанию — что-то вроде светлой бирюзы) и возвращает среднюю ошибку:

def color_loss(images, target_color=(0.1, 0.9, 0.5)):
    """По целевому цвету (R, G, B) возвращает loss: насколько в среднем
    пиксели изображений далеки от этого цвета. По умолчанию — светлая бирюза: (0.1, 0.9, 0.5)"""
    target = (
        torch.tensor(target_color).to(images.device) * 2 - 1
    )  # переводим целевой цвет в (-1, 1)
    target = target[
        None, :, None, None
    ]  # меняем форму, чтобы она подходила к изображениям (b, c, h, w)
    error = torch.abs(
        images - target
    ).mean()  # среднее абсолютное отличие пикселей изображения от целевого цвета
    return error

Затем сделаем изменённую версию цикла сэмплирования, где на каждом шаге:

создаём новую версию x с requires_grad = True;
вычисляем очищенную версию (x0);
пропускаем предсказанный x0 через функцию потерь;
находим градиент этой функции потерь по x;
с помощью этого градиента условия меняем x, прежде чем сделать шаг планировщиком, — в надежде сдвинуть x в направлении, которое по нашей функции направления даст меньший loss.

Есть два варианта, которые можно исследовать. В первом мы ставим requires_grad на x после того, как получили предсказание шума от U-Net, — это экономнее по памяти (не нужно протягивать градиенты обратно через диффузионную модель), но градиент получается менее точным. Во втором мы сначала ставим requires_grad на x, затем пропускаем его через U-Net и вычисляем предсказанный x0.

# Вариант 1: короткий путь

# Масштаб guidance определяет силу эффекта
guidance_loss_scale = 40  # попробуйте 5 или 100

x = torch.randn(8, 3, 256, 256).to(device)

for i, t in tqdm(enumerate(scheduler.timesteps)):

    # Готовим вход модели
    model_input = scheduler.scale_model_input(x, t)

    # Предсказываем шум
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]

    # Ставим x.requires_grad = True
    x = x.detach().requires_grad_()

    # Получаем предсказанный x0
    x0 = scheduler.step(noise_pred, t, x).pred_original_sample

    # Считаем loss
    loss = color_loss(x0) * guidance_loss_scale
    if i % 10 == 0:
        print(i, "loss:", loss.item())

    # Получаем градиент
    cond_grad = -torch.autograd.grad(loss, x)[0]

    # Меняем x по этому градиенту
    x = x.detach() + cond_grad

    # Теперь шаг планировщиком
    x = scheduler.step(noise_pred, t, x).prev_sample

# Смотрим результат
grid = torchvision.utils.make_grid(x, nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))

Вывод:

0 loss: 27.279136657714844
10 loss: 11.286816596984863
20 loss: 10.683112144470215
30 loss: 10.942476272583008

flowchart LR
X["x_t"] --> U["U-Net:<br/>шум ε̂"]
U --> X0["оценка x̂₀"]
X0 --> L["loss(x̂₀)<br/>цвет, CLIP…"]
L -- "−∇ₓ loss" --> XN["x_t + поправка"]
XN --> S["шаг планировщика"] --> XP["x_(t−1)"]

Один шаг сэмплирования с guidance (рисунок переводчика)

Второй вариант требует почти вдвое больше памяти GPU, хотя мы генерируем батч всего из четырёх изображений вместо восьми. Попробуйте найти отличие и подумайте, почему этот способ «точнее»:

# Вариант 2: ставим x.requires_grad до вычисления предсказаний модели

guidance_loss_scale = 40
x = torch.randn(4, 3, 256, 256).to(device)

for i, t in tqdm(enumerate(scheduler.timesteps)):

    # Ставим requires_grad до прямого прохода модели
    x = x.detach().requires_grad_()
    model_input = scheduler.scale_model_input(x, t)

    # Предсказываем (на этот раз с градиентом)
    noise_pred = image_pipe.unet(model_input, t)["sample"]

    # Получаем предсказанный x0:
    x0 = scheduler.step(noise_pred, t, x).pred_original_sample

    # Считаем loss
    loss = color_loss(x0) * guidance_loss_scale
    if i % 10 == 0:
        print(i, "loss:", loss.item())

    # Получаем градиент
    cond_grad = -torch.autograd.grad(loss, x)[0]

    # Меняем x по этому градиенту
    x = x.detach() + cond_grad

    # Теперь шаг планировщиком
    x = scheduler.step(noise_pred, t, x).prev_sample


grid = torchvision.utils.make_grid(x, nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))

Вывод:

0 loss: 30.750328063964844
10 loss: 18.550724029541016
20 loss: 17.515094757080078
30 loss: 17.55681037902832

Во втором варианте требования к памяти выше, а эффект слабее, так что может показаться, что он хуже. Однако результаты, пожалуй, ближе к тем изображениям, на которых обучалась модель, а для более сильного эффекта всегда можно увеличить масштаб guidance. Какой подход выбрать, в итоге решает эксперимент.

(Почему второй вариант «точнее»: $\hat x_0$ зависит от $x$ двумя путями — напрямую и через предсказание шума $\hat\varepsilon(x)$ . Первый вариант считает предсказание шума константой и учитывает только прямой путь, второй проводит градиент и через U-Net. Отсюда и удвоенная память: нужно хранить активации U-Net для обратного прохода.)

# Упражнение: выберите любимый цвет и найдите его значения в пространстве RGB.
# Передайте эти значения в color_loss() в ячейке выше и посмотрите на результат — совпадает ли он с ожиданием?

Guidance по CLIP

Направление к цвету даёт немного контроля, но что если бы можно было просто написать текстом, чего мы хотим?

CLIP — модель OpenAI, которая позволяет сравнивать изображения с текстовыми подписями. Это очень мощно: так можно численно оценить, насколько изображение соответствует промпту. А поскольку процесс дифференцируемый, его можно использовать как функцию потерь, чтобы направлять нашу диффузионную модель!

Не будем слишком вдаваться в детали. Базовый подход такой:

получаем 512-мерный эмбеддинг CLIP для текстового промпта;
на каждом шаге диффузионного процесса:
- делаем несколько вариантов предсказанного очищенного изображения (несколько вариантов дают более чистый сигнал loss);
- каждый из них превращаем в эмбеддинг CLIP и сравниваем с текстовым эмбеддингом промпта (мерой, которая называется «квадрат расстояния по большому кругу», Great Circle Distance Squared);
вычисляем градиент этого loss по текущему зашумлённому x и с его помощью меняем x перед обновлением планировщиком.

Подробнее о CLIP — в этом уроке или в этом отчёте о проекте OpenCLIP, через который мы загружаем модель CLIP. Выполните следующую ячейку, чтобы загрузить модель CLIP:

# @markdown загружаем модель CLIP и задаём функцию потерь
import open_clip

clip_model, _, preprocess = open_clip.create_model_and_transforms(
    "ViT-B-32", pretrained="openai"
)
clip_model.to(device)

# Преобразования, чтобы изменить размер и аугментировать изображение + нормировать под обучающие данные CLIP
tfms = torchvision.transforms.Compose(
    [
        torchvision.transforms.RandomResizedCrop(224),  # каждый раз случайный ВЫРЕЗ
        torchvision.transforms.RandomAffine(
            5
        ),  # одна из возможных случайных аугментаций: небольшой наклон изображения
        torchvision.transforms.RandomHorizontalFlip(),  # можно добавить и другие аугментации
        torchvision.transforms.Normalize(
            mean=(0.48145466, 0.4578275, 0.40821073),
            std=(0.26862954, 0.26130258, 0.27577711),
        ),
    ]
)

# Функция потерь: превращает изображение в эмбеддинг и сравнивает
# с текстовыми признаками промпта
def clip_loss(image, text_features):
    image_features = clip_model.encode_image(
        tfms(image)
    )  # NB: применяет преобразования выше
    input_normed = torch.nn.functional.normalize(image_features.unsqueeze(1), dim=2)
    embed_normed = torch.nn.functional.normalize(text_features.unsqueeze(0), dim=2)
    dists = (
        input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)
    )  # квадрат расстояния по большому кругу
    return dists.mean()

(Что считает эта цепочка: после нормировки оба эмбеддинга лежат на единичной сфере; $d = \lVert a - b \rVert$ — длина хорды между ними, $\arcsin(d/2)$ — половина угла между векторами, и итог $2\arcsin^2(d/2)$ пропорционален квадрату угла. Чем ближе картинка к тексту по смыслу, тем меньше угол и loss.)

С функцией потерь цикл сэмплирования с guidance похож на прежние примеры — заменяем color_loss() на новую функцию на основе CLIP:

# @markdown применяем guidance по CLIP

prompt = "Red Rose (still life), red flower painting"  # @param

# Попробуйте поменять
guidance_scale = 8  # @param
n_cuts = 4  # @param

# Больше шагов -> больше времени, чтобы guidance подействовал
scheduler.set_timesteps(50)

# Целевой текст — эмбеддинг промпта в CLIP
text = open_clip.tokenize([prompt]).to(device)
with torch.no_grad(), torch.cuda.amp.autocast():
    text_features = clip_model.encode_text(text)


x = torch.randn(4, 3, 256, 256).to(
    device
)  # расход памяти большой, можно генерировать по одному изображению

for i, t in tqdm(enumerate(scheduler.timesteps)):

    model_input = scheduler.scale_model_input(x, t)

    # Предсказываем шум
    with torch.no_grad():
        noise_pred = image_pipe.unet(model_input, t)["sample"]

    cond_grad = 0

    for cut in range(n_cuts):

        # Ставим requires_grad на x
        x = x.detach().requires_grad_()

        # Получаем предсказанный x0:
        x0 = scheduler.step(noise_pred, t, x).pred_original_sample

        # Считаем loss
        loss = clip_loss(x0, text_features) * guidance_scale

        # Получаем градиент (делим на n_cuts, потому что нужно среднее)
        cond_grad -= torch.autograd.grad(loss, x)[0] / n_cuts

    if i % 25 == 0:
        print("Step:", i, ", Guidance loss:", loss.item())

    # Меняем x по этому градиенту
    alpha_bar = scheduler.alphas_cumprod[i]
    x = (
        x.detach() + cond_grad * alpha_bar.sqrt()
    )  # обратите внимание на дополнительный масштабный множитель!

    # Теперь шаг планировщиком
    x = scheduler.step(noise_pred, t, x).prev_sample


grid = torchvision.utils.make_grid(x.detach(), nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))

Вывод:

Step: 0 , Guidance loss: 7.437869548797607
Step: 25 , Guidance loss: 7.174620628356934

(Промпт: «Красная роза (натюрморт), картина с красным цветком».)

Это немного похоже на розы! Не идеально, но если поиграть с настройками, можно получить симпатичные изображения.

scheduler.alphas_cumprod[i] берёт значение по номеру итерации $i$ (0…49), а не по временному шагу $t$ (980…0). Для $i$ меньше 50 $\bar\alpha_i$ почти равно 1, так что множитель фактически не действует. По смыслу, судя по тексту ниже, имелось в виду scheduler.alphas_cumprod[t] — тогда в начале генерации (большие $t$ , мало сигнала) поправка слабая, а к концу усиливается. Попробуйте оба варианта и сравните — это хорошее упражнение. Ещё: torch.cuda.amp.autocast() в новых версиях PyTorch заменён на torch.autocast("cuda").

Если присмотреться к коду выше, видно, что я масштабирую градиент условия множителем alpha_bar.sqrt(). Есть теория о том, как «правильно» масштабировать такие градиенты, но на практике с этим тоже можно экспериментировать. Для одних видов guidance хочется, чтобы основной эффект приходился на ранние шаги, для других (скажем, loss стиля, нацеленный на текстуры) — чтобы они включались только ближе к концу генерации. Несколько возможных расписаний показаны ниже:

# @markdown Рисуем несколько возможных расписаний:
plt.plot([1 for a in scheduler.alphas_cumprod], label="no scaling")
plt.plot([a for a in scheduler.alphas_cumprod], label="alpha_bar")
plt.plot([a.sqrt() for a in scheduler.alphas_cumprod], label="alpha_bar.sqrt()")
plt.plot(
    [(1 - a).sqrt() for a in scheduler.alphas_cumprod], label="(1-alpha_bar).sqrt()"
)
plt.legend()
plt.title("Possible guidance scaling schedules");

(В оригинале — четыре кривые по 1000 шагам: постоянная 1 («без масштабирования»), убывающие $\bar\alpha$ и $\sqrt{\bar\alpha}$ и растущая $\sqrt{1-\bar\alpha}$ .)

Экспериментируйте с разными расписаниями, масштабами guidance и любыми другими приёмами, которые придут в голову (популярная модификация — обрезать градиенты в некотором диапазоне), и посмотрите, насколько хорошего результата удастся добиться! Обязательно попробуйте и другие модели. Например, модель лиц, которую мы загрузили в начале: удастся ли надёжно направить её на генерацию мужского лица? А если сочетать guidance по CLIP с loss по цвету? И так далее.

Если посмотреть на код guidance по CLIP на практике, вы увидите более сложный подход — с лучшим классом для выбора случайных вырезов из изображений и множеством дополнительных доработок функции потерь для лучшего качества. До появления диффузионных моделей с текстовым условием это была лучшая система «текст → изображение»! Нашей маленькой игрушечной версии есть куда расти, но главную идею она передаёт: благодаря guidance и удивительным возможностям CLIP можно добавить управление текстом к безусловной диффузионной модели 🎨.

Делимся своим циклом сэмплирования через демо на Gradio

Может быть, вы придумали интересную функцию потерь для guidance и теперь хотите поделиться с миром и дообученной моделью, и этой стратегией сэмплирования…

Здесь на сцену выходит Gradio — бесплатный открытый инструмент, который позволяет легко создавать интерактивные модели машинного обучения с простым веб-интерфейсом и делиться ими. В Gradio можно построить свой интерфейс к модели и поделиться им по уникальной ссылке. Он также встроен в 🤗 Spaces, где легко размещать демо и делиться ими.

Основную логику поместим в функцию, которая принимает входы и выдаёт изображение. Её можно обернуть в простой интерфейс, где пользователь задаёт параметры (они передаются на вход основной функции генерации). Компонентов много — в этом примере возьмём ползунок для масштаба guidance и палитру для выбора целевого цвета.

%pip install -q gradio # устанавливаем библиотеку

import gradio as gr
from PIL import Image, ImageColor


# Функция, которая делает всю тяжёлую работу
def generate(color, guidance_loss_scale):
    target_color = ImageColor.getcolor(color, "RGB")  # целевой цвет в RGB
    target_color = [a / 255 for a in target_color]  # переводим из (0, 255) в (0, 1)
    x = torch.randn(1, 3, 256, 256).to(device)
    for i, t in tqdm(enumerate(scheduler.timesteps)):
        model_input = scheduler.scale_model_input(x, t)
        with torch.no_grad():
            noise_pred = image_pipe.unet(model_input, t)["sample"]
        x = x.detach().requires_grad_()
        x0 = scheduler.step(noise_pred, t, x).pred_original_sample
        loss = color_loss(x0, target_color) * guidance_loss_scale
        cond_grad = -torch.autograd.grad(loss, x)[0]
        x = x.detach() + cond_grad
        x = scheduler.step(noise_pred, t, x).prev_sample
    grid = torchvision.utils.make_grid(x, nrow=4)
    im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
    im = Image.fromarray(np.array(im * 255).astype(np.uint8))
    im.save("test.jpeg")
    return im


# Какие бывают входы и выходы — смотрите в документации gradio
inputs = [
    gr.ColorPicker(label="color", value="55FFAA"),  # добавьте нужные входы
    gr.Slider(label="guidance_scale", minimum=0, maximum=30, value=3),
]
outputs = gr.Image(label="result")

# И минимальный интерфейс
demo = gr.Interface(
    fn=generate,
    inputs=inputs,
    outputs=outputs,
    examples=[
        ["#BB2266", 3],
        ["#44CCAA", 5],  # можно дать примеры входов, чтобы было с чего начать
    ],
)
demo.launch(debug=True)  # debug=True показывает ошибки и вывод в Colab

(Примечание переводчика: значение по умолчанию у палитры лучше писать с решёткой — "#55FFAA", иначе ImageColor.getcolor не разберёт цвет. Подробнее о Gradio — в главе 9 LLM Course на русском.)

Можно строить гораздо более сложные интерфейсы — с красивым оформлением и множеством разных входов, — но для этого демо мы держимся максимально просто.

Демо на 🤗 Spaces по умолчанию работают на CPU, поэтому удобно отладить интерфейс в Colab (как выше), прежде чем переносить. Когда будете готовы поделиться демо, создайте Space, добавьте файл requirements.txt со списком библиотек, которые использует код, и поместите весь код в файл app.py, где заданы нужные функции и интерфейс.

(В оригинале — скриншот Space с демо: палитра, ползунок и сгенерированная картинка.)

К счастью, Space можно ещё и «Дублировать» (Duplicate). Откройте моё демо здесь (оно показано выше) и нажмите «Duplicate this Space» — получите шаблон, который можно изменить под свою модель и функцию направления.

В настройках Space можно выбрать более мощное железо (оно оплачивается почасово). Сделали что-то потрясающее и хотите показать это на хорошем железе, но нет денег? Напишите нам в Discord — посмотрим, сможем ли помочь!

Итоги и следующие шаги

Мы разобрали в этом ноутбуке много всего! Повторим главные идеи:

загружать существующие модели и сэмплировать из них с разными планировщиками сравнительно легко;
дообучение выглядит так же, как обучение с нуля, только, начиная с существующей модели, мы надеемся быстрее получить лучший результат;
чтобы дообучать большие модели на крупных изображениях, можно использовать приёмы вроде накопления градиентов и так обходить ограничения на размер батча;
при дообучении важно логировать примеры изображений: кривая loss может почти ничего не показывать;
guidance позволяет взять безусловную модель и направлять процесс генерации некоторой функцией направления/потерь: на каждом шаге мы находим градиент loss по зашумлённому изображению x и обновляем его по этому градиенту, прежде чем перейти к следующему шагу;
guidance по CLIP позволяет управлять безусловными моделями с помощью текста!

Чтобы закрепить это на практике, вот конкретные следующие шаги:

Дообучите свою модель и выгрузите её на Hub. Для этого выберите отправную точку (например, модель, обученную на лицах, спальнях, кошках или пример с wikiart выше) и датасет (например, эти мордочки животных или свои изображения), а затем запустите код этого ноутбука или учебный скрипт.
Исследуйте guidance на своей дообученной модели — с одной из готовых функций направления (color_loss или CLIP) или со своей собственной.
Поделитесь демо на Gradio: измените пример Space под свою модель или создайте свою версию с большими возможностями.

Ждём ваших результатов в Discord, Twitter и других местах 🤗!

Ноутбук: диффузионная модель с условием на класс

В этом ноутбуке мы покажем один из способов добавить в диффузионную модель информацию об условии. А именно — обучим на MNIST диффузионную модель с условием на класс, продолжая пример «с нуля» из юнита 1: при инференсе можно будет указать, какую цифру модель должна сгенерировать.

Как говорилось во введении к юниту, это лишь один из многих способов добавить в диффузионную модель дополнительную информацию об условии, и выбран он за сравнительную простоту. Как и ноутбук «с нуля» в юните 1, этот ноутбук в основном иллюстративный, и его можно спокойно пропустить.

Подготовка и данные

%pip install -q diffusers

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt
from tqdm.auto import tqdm

device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')

Вывод:

Using device: cuda

# Загружаем датасет
dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())

# Подаём его в загрузчик (батч из 8 — только для демонстрации)
train_dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

# Смотрим на несколько примеров
x, y = next(iter(train_dataloader))
print('Input shape:', x.shape)
print('Labels:', y)
plt.imshow(torchvision.utils.make_grid(x)[0], cmap='Greys');

Вывод:

Input shape: torch.Size([8, 1, 28, 28])
Labels: tensor([8, 1, 5, 9, 7, 6, 2, 2])

U-Net с условием на класс

Условие на класс мы подаём так:

создаём обычную UNet2DModel с дополнительными входными каналами;
отображаем метку класса в обучаемый вектор формы (class_emb_size) слоем эмбеддингов;
присоединяем эту информацию дополнительными каналами ко входу внутренней U-Net: net_input = torch.cat((x, class_cond), 1);
подаём этот net_input (в нём всего class_emb_size+1 каналов) в U-Net и получаем итоговое предсказание.

В этом примере я взял class_emb_size равным 4, но это совершенно произвольно: можно попробовать размер 1 (посмотреть, будет ли работать), 10 (по числу классов) или заменить обучаемый nn.Embedding простым one-hot кодированием метки класса.

Вот как выглядит реализация:

class ClassConditionedUnet(nn.Module):
  def __init__(self, num_classes=10, class_emb_size=4):
    super().__init__()

    # Слой эмбеддингов отображает метку класса в вектор размера class_emb_size
    self.class_emb = nn.Embedding(num_classes, class_emb_size)

    # self.model — безусловная U-Net с дополнительными входными каналами для информации об условии (эмбеддинга класса)
    self.model = UNet2DModel(
        sample_size=28,           # целевое разрешение изображения
        in_channels=1 + class_emb_size, # дополнительные входные каналы для условия на класс
        out_channels=1,           # число выходных каналов
        layers_per_block=2,       # сколько слоёв ResNet в каждом блоке U-Net
        block_out_channels=(32, 64, 64),
        down_block_types=(
            "DownBlock2D",        # обычный блок уменьшения на ResNet
            "AttnDownBlock2D",    # блок уменьшения на ResNet с пространственным self-attention
            "AttnDownBlock2D",
        ),
        up_block_types=(
            "AttnUpBlock2D",
            "AttnUpBlock2D",      # блок увеличения на ResNet с пространственным self-attention
            "UpBlock2D",          # обычный блок увеличения на ResNet
          ),
    )

  # Метод forward теперь принимает метки классов дополнительным аргументом
  def forward(self, x, t, class_labels):
    # Форма x:
    bs, ch, w, h = x.shape

    # условие на класс в нужной форме, чтобы добавить его дополнительными входными каналами
    class_cond = self.class_emb(class_labels) # отображаем в размерность эмбеддинга
    class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w, h)
    # x имеет форму (bs, 1, 28, 28), а class_cond теперь (bs, 4, 28, 28)

    # Вход сети — x и условие на класс, склеенные по измерению 1
    net_input = torch.cat((x, class_cond), 1) # (bs, 5, 28, 28)

    # Подаём это в U-Net вместе с временным шагом и возвращаем предсказание
    return self.model(net_input, t).sample # (bs, 1, 28, 28)

Если какие-то формы или преобразования непонятны, добавьте print, чтобы увидеть нужные формы и проверить, что они совпадают с ожиданиями. Я также подписал формы некоторых промежуточных переменных — надеюсь, так понятнее.

(Эмбеддинги — обучаемая таблица векторов, по строке на класс, как эмбеддинги токенов в уроке 4.3.)

Обучение и сэмплирование

Раньше мы делали что-то вроде prediction = unet(x, t), а теперь при обучении добавим третьим аргументом правильные метки (prediction = unet(x, t, y)), а при инференсе сможем подавать любые метки, и, если всё хорошо, модель должна генерировать подходящие изображения. y здесь — метки цифр MNIST, значения от 0 до 9.

Цикл обучения очень похож на пример из юнита 1. Теперь мы предсказываем шум (а не очищенное изображение, как в юните 1), чтобы цель соответствовала тому, чего ждёт стандартный DDPMScheduler, которым мы добавляем шум при обучении и генерируем примеры при инференсе. Обучение идёт довольно долго — ускорить его может быть интересным мини-проектом, но большинству из вас, наверное, достаточно просмотреть код (да и весь ноутбук), не запуская его: мы лишь иллюстрируем идею.

# Создаём планировщик
noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')

#@markdown Цикл обучения (10 эпох):

# Заново создаём загрузчик данных с батчем побольше, чем 8 в демонстрации
train_dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

# Сколько раз пройти по данным?
n_epochs = 10

# Наша сеть
net = ClassConditionedUnet().to(device)

# Функция потерь
loss_fn = nn.MSELoss()

# Оптимизатор
opt = torch.optim.Adam(net.parameters(), lr=1e-3)

# Сохраняем значения loss, чтобы потом посмотреть
losses = []

# Цикл обучения
for epoch in range(n_epochs):
    for x, y in tqdm(train_dataloader):

        # Берём данные и готовим испорченную версию
        x = x.to(device) * 2 - 1 # данные на GPU (переведены в (-1, 1))
        y = y.to(device)
        noise = torch.randn_like(x)
        timesteps = torch.randint(0, 999, (x.shape[0],)).long().to(device)
        noisy_x = noise_scheduler.add_noise(x, noise, timesteps)

        # Получаем предсказание модели
        pred = net(noisy_x, timesteps, y) # обратите внимание: передаём метки y

        # Считаем loss
        loss = loss_fn(pred, noise) # насколько выход близок к шуму

        # Обратное распространение и обновление параметров:
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Сохраняем loss на потом
        losses.append(loss.item())

    # Печатаем среднее последних 100 значений loss, чтобы видеть прогресс:
    avg_loss = sum(losses[-100:])/100
    print(f'Finished epoch {epoch}. Average of the last 100 loss values: {avg_loss:05f}')

# Смотрим на кривую loss
plt.plot(losses)

Вывод:

Finished epoch 0. Average of the last 100 loss values: 0.052451
Finished epoch 1. Average of the last 100 loss values: 0.045999
Finished epoch 2. Average of the last 100 loss values: 0.043344
Finished epoch 3. Average of the last 100 loss values: 0.042347
Finished epoch 4. Average of the last 100 loss values: 0.041174
Finished epoch 5. Average of the last 100 loss values: 0.040736
Finished epoch 6. Average of the last 100 loss values: 0.040386
Finished epoch 7. Average of the last 100 loss values: 0.039372
Finished epoch 8. Average of the last 100 loss values: 0.039056
Finished epoch 9. Average of the last 100 loss values: 0.039024

(Мелкая неточность: torch.randint(0, 999, …) выбирает шаги от 0 до 998 — верхняя граница не включается, и шаг 999 никогда не встречается. Правильнее torch.randint(0, 1000, …). На RTX 4080 десять эпох — порядка 5–10 минут.)

Когда обучение закончится, можно сэмплировать изображения, подавая разные метки в качестве условия:

#@markdown Сэмплируем разные цифры:

# Готовим случайный x, с которого начнём, и нужные метки y
x = torch.randn(80, 1, 28, 28).to(device)
y = torch.tensor([[i]*8 for i in range(10)]).flatten().to(device)

# Цикл сэмплирования
for i, t in tqdm(enumerate(noise_scheduler.timesteps)):

    # Получаем предсказание модели
    with torch.no_grad():
        residual = net(x, t, y)  # снова передаём метки y

    # Обновляем сэмпл шагом
    x = noise_scheduler.step(residual, t, x).prev_sample

# Показываем результат
fig, ax = plt.subplots(1, 1, figsize=(12, 12))
ax.imshow(torchvision.utils.make_grid(x.detach().cpu().clip(-1, 1), nrow=8)[0], cmap='Greys')

(y — это [0,0,…,0, 1,1,…,1, …, 9,…,9], по 8 раз каждая цифра; в сетке по 8 картинок в строке получается строка нулей, строка единиц и так далее. В оригинале — именно такая сетка с узнаваемыми цифрами.)

Вот и всё! Теперь мы можем в какой-то мере управлять тем, какие изображения получаются.

Надеюсь, вам понравился этот пример. Как всегда, задавайте вопросы в Discord.

# Упражнение (необязательно): попробуйте то же на FashionMNIST. Подберите скорость обучения, размер батча и число эпох.
# Удастся ли получить приличные изображения одежды за меньшее время обучения, чем в примере выше?

Эту модель легко превратить в модель с classifier-free guidance, которой пользуются все современные генераторы. Добавьте 11-й «пустой» класс и при обучении в 10–20% случаев подменяйте настоящую метку на пустую. При сэмплировании считайте два предсказания — с меткой и с пустой — и смешивайте: $\hat\varepsilon = \hat\varepsilon_{\varnothing} + w(\hat\varepsilon_{y} - \hat\varepsilon_{\varnothing})$ с $w$ около 3–7. Цифры станут чётче и «типичнее» за счёт разнообразия. Подробно — в уроке 9.1.2 и в юните 3 этого курса.

Словарик

English	Русский
fine-tuning	дообучение
guidance, guidance function, guidance scale	направление (guidance), функция направления, масштаб guidance
conditioning, conditional / unconditional model	условие, условная / безусловная модель
DDIM (Denoising Diffusion Implicit Models)	DDIM — неявные диффузионные модели с шумоподавлением
predicted denoised image ( $\hat x_0$ )	предсказанное очищенное изображение
gradient accumulation	накопление градиентов
exponential model averaging (EMA)	экспоненциальное усреднение весов модели
cross-attention	перекрёстное внимание
CLIP guidance	направление по CLIP
great circle distance	расстояние по большому кругу
class embedding, one-hot encoding	эмбеддинг класса, one-hot кодирование
classifier-free guidance (CFG)	guidance без классификатора

Перевод выполнен по лицензии оригинала (Apache-2.0); автор оригинала — Джонатан Уитакер (Jonathan Whitaker) и Льюис Танстолл (Lewis Tunstall), Hugging Face. Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04