Дообучение Stable Diffusion DreamBooth: LoRA, SDXL, без переобучения

Клиент приносит 15 фотографий своего продукта — кроссовки новой модели. Нужно разместить их на рекламных макетах: на пляже, в горах, в студии. Базовая Stable Diffusion не знает этот объект — результат зависит от случайного seed. Если seed не зафиксировать, каждый промпт будет выдавать разный ракурс, цвет, текстуру. DreamBooth решает задачу: дообучает модель на 5–20 снимках, запоминая уникальный идентификатор субъекта (например, sks sneaker). Мы используем этот подход для брендовых аватаров, персонажей и художественных стилей. Опыт с SDXL, LoRA, ControlNet позволяет гарантировать качество генераций без переобучения. Команда имеет 5+ лет опыта в CV и NLP, выполнила более 100 проектов по дообучению моделей.

DreamBooth — метод, предложенный Google Research, для точной настройки text-to-image моделей под конкретный субъект.

Как DreamBooth сохраняет уникальность субъекта?

DreamBooth привязывает редкий токен sks к визуальным признакам объекта через prior preservation loss. Это предотвращает «языковой дрейф» — модель не забывает общие концепции класса (например, «кроссовки» в целом). Prior preservation loss использует изображения класса (например, 'sneaker' без субъекта) чтобы модель не забывала, как выглядят обычные кроссовки. Это реализуется через случайный сэмплинг из предобученной модели. Результат: субъект узнаваем в любом контексте.

Технически процесс состоит из двух этапов: подготовка датасета и обучение LoRA-весов. LoRA (Low-Rank Adaptation) замораживает исходные веса SD и добавляет адаптеры — это требует в 10–20 раз меньше VRAM, чем полный fine-tuning (8 ГБ против 24+ ГБ).

Почему LoRA эффективнее полного fine-tuning?

LoRA — стандарт для продакшена: быстрый деплой, малый размер, легко комбинируется с другими LoRA (например, стиль + субъект).

Как подготовить датасет для DreamBooth?

Первое, с чем сталкивается инженер — качество исходных изображений. Модель копирует ракурсы, освещение, фон. Если все снимки сделаны в одной студии — DreamBooth выучит студию как часть субъекта.

1. Сбор изображений. Нужно 10–20 снимков разных ракурсов (спереди, сбоку, сверху), разное освещение (естественное, искусственное). Объект должен занимать 50–80% кадра. Избегайте сильного перекрытия (рука, тень).
2. Обрезка и центрирование. Приводим все изображения к квадрату 1024x1024. Используем функцию из листинга ниже.
3. Аугментация. Для улучшения обобщения применяем случайное горизонтальное отражение, небольшой поворот (до 10°), изменение яркости/контраста. Сильные искажения ломают геометрию.
4. Сегментация (опционально). Если объект — человек, используйте RMBG 2.0 для изоляции.
5. Prior preservation. Генерируем 100–200 изображений класса (например, 'sneaker' без субъекта) с помощью базовой модели. Эти снимки используются в prior preservation loss.

from PIL import Image
import os

def prepare_dreambooth_dataset(
    source_images: list[str],
    output_dir: str,
    target_size: int = 1024
) -> None:
    os.makedirs(output_dir, exist_ok=True)
    for i, img_path in enumerate(source_images):
        img = Image.open(img_path).convert("RGB")
        # Центрируем и обрезаем до квадрата
        width, height = img.size
        min_dim = min(width, height)
        left = (width - min_dim) // 2
        top = (height - min_dim) // 2
        img_cropped = img.crop((left, top, left + min_dim, top + min_dim))
        img_resized = img_cropped.resize((target_size, target_size), Image.LANCZOS)
        img_resized.save(f"{output_dir}/{i:03d}.jpg", quality=95)
    print(f"Подготовлено {len(source_images)} изображений в {output_dir}")

Обучение: выбор гиперпараметров

Скрипт Diffusers для SDXL запускается через accelerate. Рекомендуем --mixed_precision="fp16" и --use_8bit_adam для экономии памяти.

accelerate launch train_dreambooth_lora_sdxl.py \
    --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
    --instance_data_dir="./training_images" \
    --output_dir="./dreambooth_output" \
    --instance_prompt="a photo of sks person" \
    --resolution=1024 \
    --train_batch_size=1 \
    --gradient_accumulation_steps=4 \
    --learning_rate=1e-4 \
    --lr_scheduler="constant" \
    --lr_warmup_steps=0 \
    --max_train_steps=500 \
    --seed=42 \
    --mixed_precision="fp16"

Подробнее о скрипте см. в официальной документации Diffusers. Основные гиперпараметры:

Оптимальное количество шагов зависит от сложности субъекта. Для простых объектов (продукт на белом фоне) достаточно 300-500 шагов. Для сложных (человек с деталями одежды) — до 800-1000. Learning rate лучше начинать с 1e-4 и уменьшать по cosine schedule.

Если после обучения модель генерирует только один ракурс или игнорирует фон — это признак переобучения. Решение: увеличить prior preservation weight, уменьшить шаги, добавить аугментацию.

Интеграция и запуск в продакшен

После обучения получаем LoRA-весы (обычно ~150 МБ). Загружаем в кастомный пайплайн StableDiffusionXLPipeline:

from diffusers import DiffusionPipeline
import torch


def train_dreambooth_sdxl(
    instance_images_dir: str,
    instance_prompt: str,
    class_prompt: str,
    output_dir: str,
    num_steps: int = 800,
    learning_rate: float = 1e-4
) -> str:
    import subprocess
    result = subprocess.run([
        "accelerate",
        "launch",
        "train_dreambooth_lora_sdxl.py",
        "--pretrained_model_name_or_path",
        "stabilityai/stable-diffusion-xl-base-1.0",
        "--instance_data_dir",
        instance_images_dir,
        "--instance_prompt",
        instance_prompt,
        "--class_prompt",
        class_prompt,
        "--output_dir",
        output_dir,
        "--max_train_steps",
        str(num_steps),
        "--learning_rate",
        str(learning_rate),
        "--resolution",
        "1024",
        "--train_batch_size",
        "1",
        "--gradient_checkpointing",
        "--mixed_precision",
        "fp16",
        "--use_8bit_adam",
    ], capture_output=True)
    return output_dir


def generate_with_dreambooth(
    lora_path: str,
    prompt_template: str,
    subject_token: str = "sks"
) -> bytes:
    pipe = DiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16
    ).to("cuda")
    pipe.load_lora_weights(lora_path)
    prompt = prompt_template.replace("{subject}", subject_token)
    image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
    import io
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    return buf.getvalue()

После обучения LoRA можно комбинировать с ControlNet для точного управления позой, глубиной или краями. Например, задать позу персонажа через OpenPose, оставив внешность обученной DreamBooth.

Процесс работы с нами

При заказе дообучения мы проводим следующие этапы:

• Анализ задачи — изучаем ваши референсы, определяем класс субъекта, выбираем базовую модель (SD 2.1, SDXL, или SD 3).
• Подготовка датасета — помогаем с очисткой и аугментацией изображений.
• Обучение LoRA — подбираем гиперпараметры, проводим обучение, проверяем на переобучение.
• Тестирование — генерируем 50+ вариантов в разных контекстах, отбираем лучший чекпоинт.
• Деплой — выкладываем модель в облачную инфраструктуру (SageMaker, Vertex AI) или передаём файлы для локального запуска.
• Документация и поддержка — передаём API-документацию, примеры инференса, и месяц сопровождения.

Сроки: от 2 дней для простых объектов до 3 недель для персонажа с анимацией (последовательные LoRA). Стоимость рассчитывается индивидуально после оценки.

Типичные ошибки и их предотвращение

• Переобучение — модель генерирует только один ракурс. Решение: уменьшить число шагов, увеличить prior preservation, добавить аугментацию.
• Неправильный токен — использование частотного слова (например, person) приводит к смешиванию с другими субъектами. Выбирайте редкий токен вроде sks.
• Малый датасет — меньше 5 изображений не дают модели выучить объект. Минимум 10.
• Плохой фон — если фон не разнообразен, модель "привязывает" субъект к одному окружению. Используйте снимки с разными фонами.

Что входит в работу

• Подготовленный и аугментированный датасет (до 20 изображений)
• Обученная LoRA-модель (файл ~150 МБ)
• Чекпоинт с наилучшим качеством (отбор по 50+ генерациям)
• API-документация и пример инференса на Python
• Деплой в облако (SageMaker/Vertex AI) по запросу
• Месяц технической поддержки

Свяжитесь с нами для консультации по вашему проекту. Получите оценку сроков и стоимости — напишите, и мы подготовим предложение в течение дня.

Закажите дообучение модели, чтобы получить стабильно узнаваемый субъект в любом контексте.