Дообучение Stable Diffusion методом Dreambooth

Fine-tuning Stable Diffusion через DreamBooth

DreamBooth позволяет обучить SD на конкретном субъекте (человек, продукт, стиль, персонаж) по 5–20 фотографиям. После обучения модель генерирует субъект в произвольных сценариях, сохраняя узнаваемость.

Применения

• Брендовый продукт: кроссовки в разных сценах — на природе, в городе, в студии
• Аватары: лицо человека в разных стилях — anime, oil painting, cartoon
• Персонаж: игровой герой в новых ситуациях
• Стиль художника: перенос художественного стиля на новые сцены

Diffusers DreamBooth обучение

pip install accelerate diffusers transformers bitsandbytes

# Скрипт обучения для SDXL DreamBooth
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"

Подготовка датасета

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}")

Python API обучения

from diffusers import DiffusionPipeline
import torch

def train_dreambooth_sdxl(
    instance_images_dir: str,
    instance_prompt: str,        # "sks dog" — уникальный токен + класс
    class_prompt: str,           # "dog" — только класс для prior preservation
    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()

Гиперпараметры и советы

Признак переобучения: модель генерирует только один ракурс субъекта, отказывается от фоновых описаний.

Сроки: обучение DreamBooth LoRA (~500 шагов на RTX 3090) — 15–30 минут. Сервис с загрузкой изображений пользователем и генерацией аватаров — 2–4 недели.