Buckets:

mishig's picture
|
download
raw
20.9 kB
# 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 всего на нескольких примерах изображений. Эти два метода обучения можно даже комбинировать!

Xet Storage Details

Size:
20.9 kB
·
Xet hash:
a139b4899ba08aaa7596d17e12a48a4c8cc733db1103cd8142176a0ffa965f33

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.