Текст в изображение

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

Модели «текст в изображение», такие как Stable Diffusion, обучены генерировать изображения на основе текстового запроса.

Обучение модели может сильно нагружать ваше оборудование, но если включить gradient_checkpointing и mixed_precision, можно обучать модель на одном GPU с 24 ГБ памяти. Если вы обучаете модель с большими размерами батчей или хотите ускорить процесс, лучше использовать GPU с памятью более 30 ГБ. Вы можете уменьшить потребление памяти, включив энергоэффективное внимание с помощью xFormers.

В этом руководстве мы рассмотрим скрипт обучения train_text_to_image.py, чтобы вы могли ознакомиться с ним и понять, как адаптировать его под собственные задачи.

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

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

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

cd examples/text_to_image
pip install -r requirements.txt

🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в Кратком обзоре 🤗 Accelerate.

Инициализируйте среду 🤗 Accelerate:

accelerate config

Чтобы настроить среду по умолчанию 🤗 Accelerate без выбора каких-либо конфигураций:

accelerate config default

Или, если ваша среда не поддерживает интерактивную оболочку, например, ноутбук, вы можете использовать:

from accelerate.utils import write_basic_config

write_basic_config()

Наконец, если вы хотите обучить модель на собственном наборе данных, ознакомьтесь с руководством Create a dataset for training, чтобы узнать, как создать набор данных, совместимый со скриптом обучения.

Параметры скрипта

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

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

Например, чтобы ускорить обучение с использованием смешанной точности в формате fp16, добавьте параметр --mixed_precision к команде обучения:

accelerate launch train_text_to_image.py \
  --mixed_precision="fp16"

К числу основных и важных параметров относятся:

• --pretrained_model_name_or_path: название модели на Hub или локальный путь к предобученной модели
• --dataset_name: название набора данных на Hub или локальный путь к набору данных для обучения
• --image_column: имя столбца с изображениями в наборе данных для обучения
• --caption_column: имя текстовой колонки в наборе данных для обучения
• --output_dir: место для сохранения обученной модели
• --push_to_hub: нужно ли отправить обученную модель в Hub
• --checkpointing_steps: частота сохранения контрольных точек в процессе обучения модели; это полезно, если обучение по какой-либо причине прерывается, вы можете продолжить обучение с этой контрольной точки, добавив --resume_from_checkpoint к команде обучения

Взвешивание Min-SNR

Стратегия взвешивания Min-SNR может помочь в обучении, перераспределяя функцию потерь для достижения более быстрой сходимости. Скрипт обучения поддерживает предсказание epsilon (шума) или v_prediction, при этом Min-SNR совместима с обоими типами предсказаний. Данная стратегия взвешивания поддерживается только в PyTorch.

Добавьте параметр --snr_gamma и установите для него рекомендуемое значение 5.0:

accelerate launch train_text_to_image.py \
  --snr_gamma=5.0

Вы можете сравнить поверхности потерь для различных значений snr_gamma в этом отчёте Weights and Biases. Для небольших наборов данных эффекты Min-SNR могут быть менее очевидными по сравнению с большими наборами данных.

Скрипт обучения

Код предобработки набора данных и цикл обучения находятся в функции main(). Если вам нужно адаптировать скрипт обучения, именно здесь следует внести изменения.

Скрипт train_text_to_image начинает с загрузки планировщика и токенизатора. При желании вы можете использовать здесь другой планировщик:

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

Затем скрипт загружает модель UNet:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

Далее необходимо выполнить предобработку текстовых и изображений столбцов набора данных. Функция tokenize_captions отвечает за токенизацию входных данных, а функция train_transforms определяет тип преобразований, применяемых к изображению. Обе эти функции объединены в preprocess_train:

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

Lastly, the training loop handles everything else. It encodes images into latent space, adds noise to the latents, computes the text embeddings to condition on, updates the model parameters, and saves and pushes the model to the Hub. If you want to learn more about how the training loop works, check out the Understanding pipelines, models and schedulers tutorial which breaks down the basic pattern of the denoising process.

Запустите скрипт

Когда вы внесёте все изменения или решите оставить конфигурацию по умолчанию, можно запускать скрипт обучения! 🚀

Обучимся на датасете Naruto BLIP captions, чтобы генерировать собственных персонажей Наруто. Установите переменные окружения MODEL_NAME и dataset_name в значения модели и датасета (либо с Hub, либо локальный путь). Если вы обучаете на нескольких GPU, добавьте параметр --multi_gpu к команде accelerate launch.

Чтобы обучать модель на локальном наборе данных, установите переменные окружения TRAIN_DIR и OUTPUT_DIR в путь к набору данных и каталог для сохранения модели.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

Дальнейшие шаги

Поздравляем с обучением собственной модели text-to-image! Чтобы узнать, как использовать новую модель, вам могут помочь следующие руководства:

• Узнайте, как загрузить веса LoRA для инференса, если вы обучали модель с использованием LoRA.
• Подробнее о том, как определённые параметры, такие как guidance scale, или приёмы вроде prompt weighting помогают управлять процессом inference, читайте в руководстве по задаче Text-to-image.