Дообучення 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) за запитом
• Місяць технічної підтримки

Зв'яжіться з нами для консультації по вашому проекту. Отримайте оцінку термінів та вартості — напишіть, і ми підготуємо пропозицію протягом дня.

Замовте донавчання моделі, щоб отримати стабільно впізнаваний суб'єкт у будь-якому контексті.