# Kandinsky 2.2

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

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

В зависимости от вашей GPU может потребоваться включить `gradient_checkpointing` (⚠️ не поддерживается для предыдущей модели!), `mixed_precision` и `gradient_accumulation_steps`, чтобы уместить модель в память и ускорить обучение. Вы можете ещё больше снизить потребление памяти, включив энергоэффективное внимание с помощью [xFormers](../optimization/attention_backends) (версия [v0.0.16](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212) не работает для обучения на некоторых GPU, поэтому вам может потребоваться установить версию для разработки).

This guide explores the [train_text_to_image_prior.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py) and the [train_text_to_image_decoder.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/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.

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

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

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

```bash
cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt
```

> [!TIP]
> 🤗 Accelerate — это библиотека, которая помогает обучать модели на нескольких GPU/TPU или с использованием смешанной точности. Она автоматически настраивает среду обучения в зависимости от вашего оборудования и окружения. Подробнее можно узнать в [Кратком обзоре](https://huggingface.co/docs/accelerate/quicktour) 🤗 Accelerate.

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

```bash
accelerate config
```

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

```bash
accelerate config default
```

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

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

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

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

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

Скрипты обучения предоставляют множество параметров, которые помогают настроить процесс обучения. Все параметры и их описания находятся в функции [`parse_args()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L190). Скрипты обучения задают значения по умолчанию для каждого параметра, например, размер обучающей выборки и скорость обучения, но вы также можете указать свои значения в команде обучения.

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

```bash
accelerate launch train_text_to_image_prior.py \
  --mixed_precision="fp16"
```

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

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

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

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

```bash
accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0
```

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

The training script is also similar to the [Text-to-image](text2image#training-script) 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()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L441) содержит код для подготовки набора данных и обучения модели.

Одно из основных отличий, которое вы заметите сразу, состоит в том, что скрипт обучения также загружает [CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor) — помимо планировщика и токенизатора — для предобработки изображений, а также модель [CLIPVisionModelWithProjection](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPVisionModelWithProjection) для кодирования изображений:

```py
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](/docs/diffusers/main/ru/api/models/prior_transformer#diffusers.PriorTransformer) для генерации эмбеддингов изображений, поэтому вам нужно настроить оптимизатор для обучения параметров режима prior.

```py
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,
)
```

Далее входные подписи токенизируются, а изображения [предобрабатываются](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L632) с помощью [CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor):

```py
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
```

Наконец, [цикл обучения](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L718) преобразует входные изображения в латенты, добавляет шум к эмбеддингам изображений и делает предсказание:

```py
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](../using-diffusers/write_own_pipeline), в котором подробно разбирается базовый паттерн процесса денойзинга.

Функция [`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L440) содержит код для подготовки набора данных и обучения модели.

В отличие от предыдущей модели, декодер инициализирует [VQModel](/docs/diffusers/main/ru/api/models/vq#diffusers.VQModel) для декодирования латентов в изображения и использует [UNet2DConditionModel](/docs/diffusers/main/ru/api/models/unet2d-cond#diffusers.UNet2DConditionModel):

```py
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    vae = VQModel.from_pretrained(
        args.pretrained_decoder_model_name_or_path, subfolder="movq", dtype=weight_dtype
    ).eval()
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")
```

Далее в скрипт включены несколько преобразований изображений и функция [preprocessing](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L622) для применения этих преобразований к изображениям и возврата значений пикселей:

```py
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["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    return examples
```

Наконец, [цикл обучения](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L706) отвечает за преобразование изображений в латенты, добавление шума и предсказание шумового остатка.

Если вы хотите узнать больше о том, как работает цикл обучения, обратитесь к учебнику [Understanding pipelines, models and schedulers](../using-diffusers/write_own_pipeline), в котором подробно разбирается базовый паттерн процесса денойзинга.

```py
model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]
```

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

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

Вы будете обучать модель на датасете [Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions), чтобы генерировать собственных персонажей Наруто, но вы также можете создать и обучить модель на собственном датасете, следуя руководству [Create a dataset for training](create_dataset). Установите переменную окружения `DATASET_NAME` в имя датасета на Hub или, если вы обучаете на собственных файлах, установите переменную окружения `TRAIN_DIR` в путь к вашему датасету.

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

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

```bash
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"
```

```bash
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_decoder.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --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-decoder-naruto-model"
```

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

```py
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]
```

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

```py
from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", dtype=torch.float16)
pipeline.enable_model_cpu_offload()

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

Для модели декодера также можно выполнять инференс из сохранённой контрольной точки, что может быть полезно для просмотра промежуточных результатов. В этом случае загрузите контрольную точку в UNet:

```py
from diffusers import AutoPipelineForText2Image, UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")

pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, dtype=torch.float16)
pipeline.enable_model_cpu_offload()

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

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

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

- Ознакомьтесь с руководством по [Kandinsky](../api/pipelines/kandinsky), чтобы узнать, как использовать его для различных задач (генерация изображений по тексту, преобразование изображения в изображение, инпейнтинг, интерполяция) и как его можно комбинировать с ControlNet.
- Ознакомьтесь с руководствами по обучению [DreamBooth](dreambooth) и [LoRA](lora), чтобы узнать, как обучить персонализированную модель Kandinsky всего на нескольких примерах изображений. Эти два метода обучения можно даже комбинировать!