Hugging Face's logo

Diffusers documentation

Kandinsky 2.2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Kandinsky 2.2

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

Kandinsky 2.2 — это мультиязычная модель генерации изображений по текстовому описанию, способная создавать более фотореалистичные изображения. Модель включает в себя модель image prior для создания эмбеддингов изображений из текстовых промптов, а также модель декодера, которая генерирует изображения на основе эмбеддингов модели prior. Именно поэтому в Diffusers для Kandinsky 2.2 представлены два отдельных скрипта: один для обучения модели prior, а другой — для обучения модели декодера. Вы можете обучать обе модели независимо, однако для достижения наилучших результатов рекомендуется обучать как модель prior, так и модель декодера.

В зависимости от вашей GPU может потребоваться включить gradient_checkpointing (⚠️ не поддерживается для предыдущей модели!), mixed_precision и gradient_accumulation_steps, чтобы уместить модель в память и ускорить обучение. Вы можете ещё больше снизить потребление памяти, включив энергоэффективное внимание с помощью xFormers (версия v0.0.16 не работает для обучения на некоторых GPU, поэтому вам может потребоваться установить версию для разработки).

This guide explores the train_text_to_image_prior.py and the train_text_to_image_decoder.py scripts to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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

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

cd examples/kandinsky2_2/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_prior.py \
  --mixed_precision="fp16"

Большинство параметров совпадают с параметрами из руководства по обучению Text-to-image, поэтому перейдём сразу к разбору скриптов обучения Kandinsky!

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

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

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

accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0

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

The training script is also similar to the Text-to-image training guide, but it’s been modified to support training the prior and decoder models. This guide focuses on the code that is unique to the Kandinsky 2.2 training scripts.

предыдущая модель
модель декодера

Функция main() содержит код для подготовки набора данных и обучения модели.

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

noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", dtype=weight_dtype
    ).eval()

Kandinsky использует PriorTransformer для генерации эмбеддингов изображений, поэтому вам нужно настроить оптимизатор для обучения параметров режима prior.

prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

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

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples

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

model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding

Если вы хотите узнать больше о том, как работает цикл обучения, обратитесь к учебнику Understanding pipelines, models and schedulers, в котором подробно разбирается базовый паттерн процесса денойзинга.

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

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

Вы будете обучать модель на датасете Naruto BLIP captions, чтобы генерировать собственных персонажей Наруто, но вы также можете создать и обучить модель на собственном датасете, следуя руководству Create a dataset for training. Установите переменную окружения DATASET_NAME в имя датасета на Hub или, если вы обучаете на собственных файлах, установите переменную окружения TRAIN_DIR в путь к вашему датасету.

Если вы обучаете модель на нескольких GPU, добавьте параметр --multi_gpu к команде accelerate launch.

Для отслеживания хода обучения с помощью Weights & Biases добавьте параметр --report_to=wandb к команде обучения. Также необходимо добавить --validation_prompt в команду обучения, чтобы фиксировать результаты. Это может быть очень полезно для отладки модели и просмотра промежуточных результатов.

предыдущая модель
модель декодера
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_prior.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-prior-naruto-model"

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

предыдущая модель
модель декодера
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
import torch

prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, dtype=torch.float16)
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, dtype=torch.float16)

pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]

Вы можете заменить kandinsky-community/kandinsky-2-2-decoder на собственный обученный чекпоинт декодера!

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

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

  • Ознакомьтесь с руководством по Kandinsky, чтобы узнать, как использовать его для различных задач (генерация изображений по тексту, преобразование изображения в изображение, инпейнтинг, интерполяция) и как его можно комбинировать с ControlNet.
  • Ознакомьтесь с руководствами по обучению DreamBooth и LoRA, чтобы узнать, как обучить персонализированную модель Kandinsky всего на нескольких примерах изображений. Эти два метода обучения можно даже комбинировать!
Update on GitHub

←Stable Diffusion XL ControlNet→