Buckets:

mishig's picture
|
download
raw
16.1 kB
# Kandinsky 2.2
> [!WARNING]
> Este script é experimental e é fácil sofrer overfitting e enfrentar problemas como o esquecimento catastrófico. Tente explorar diferentes hiperparâmetros para obter os melhores resultados no seu conjunto de dados.
Kandinsky 2.2 é um modelo multilíngue de texto para imagem capaz de produzir imagens mais fotorrealistas. O modelo inclui um modelo de prior de imagem para criar embeddings de imagem a partir de prompts de texto, e um modelo decodificador que gera imagens com base nos embeddings do modelo de prior. É por isso que você encontrará dois scripts separados no Diffusers para o Kandinsky 2.2, um para treinar o modelo de prior e outro para treinar o modelo decodificador. Você pode treinar ambos os modelos separadamente, mas para obter os melhores resultados, você deve treinar tanto o modelo de prior quanto o modelo decodificador.
Dependendo da sua GPU, você pode precisar habilitar `gradient_checkpointing` (⚠️ não suportado para o modelo anterior!), `mixed_precision` e `gradient_accumulation_steps` para ajudar a caber o modelo na memória e acelerar o treinamento. Você pode reduzir ainda mais o uso de memória habilitando a atenção eficiente em memória com [xFormers](../optimization/attention_backends) (a versão [v0.0.16](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212) falha no treinamento em algumas GPUs, então você pode precisar instalar uma versão de desenvolvimento em vez disso).
Este guia explora os scripts [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) e [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) para ajudá-lo a se familiarizar com eles e entender como adaptá-los ao seu próprio caso de uso.
Antes de executar os scripts, certifique-se de instalar a biblioteca a partir do código-fonte:
```bash
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
```
Em seguida, navegue até a pasta de exemplo que contém o script de treinamento e instale as dependências necessárias para o script que você está usando:
```bash
cd examples/kandinsky2_2/text_to_image
pip install -r requirements.txt
```
> [!TIP]
> 🤗 Accelerate é uma biblioteca para ajudar você a treinar em múltiplos GPUs/TPUs ou com precisão mista. Ela configura automaticamente seu ambiente de treinamento com base no seu hardware e ambiente. Dê uma olhada no [Quick tour](https://huggingface.co/docs/accelerate/quicktour) do 🤗 Accelerate para saber mais.
Inicialize um ambiente do 🤗 Accelerate:
```bash
accelerate config
```
Para configurar um ambiente padrão do 🤗 Accelerate sem escolher nenhuma configuração:
```bash
accelerate config default
```
Ou, se o seu ambiente não suporta um shell interativo, como um notebook, você pode usar:
```py
from accelerate.utils import write_basic_config
write_basic_config()
```
Por fim, se você deseja treinar um modelo em seu próprio conjunto de dados, consulte o guia [Criar um conjunto de dados para treinamento](create_dataset) para aprender a criar um conjunto de dados que funcione com o script de treinamento.
> [!TIP]
> As seções a seguir destacam partes dos scripts de treinamento que são importantes para entender como modificá-los, mas não cobrem todos os aspectos dos scripts em detalhes. Se você estiver interessado em aprender mais, sinta-se à vontade para ler os scripts e nos informar caso tenha alguma dúvida ou preocupação.
## Parâmetros do script
Os scripts de treinamento oferecem muitos parâmetros para ajudar você a personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão na função [`parse_args()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L190). Os scripts de treinamento fornecem valores padrão para cada parâmetro, como o tamanho do lote de treinamento e a taxa de aprendizado, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.
Por exemplo, para acelerar o treinamento com precisão mista usando o formato fp16, adicione o parâmetro `--mixed_precision` ao comando de treinamento:
```bash
accelerate launch train_text_to_image_prior.py \
--mixed_precision="fp16"
```
A maioria dos parâmetros é idêntica aos parâmetros do guia de treinamento [Texto-para-imagem](text2image#script-parameters), então vamos direto para uma explicação dos scripts de treinamento do Kandinsky!
### Ponderação Min-SNR
A estratégia de ponderação [Min-SNR](https://huggingface.co/papers/2303.09556) pode ajudar no treinamento ao reequilibrar a perda para alcançar uma convergência mais rápida. O script de treinamento suporta a previsão de `epsilon` (ruído) ou `v_prediction`, mas o Min-SNR é compatível com ambos os tipos de previsão. Essa estratégia de ponderação é suportada apenas pelo PyTorch.
Adicione o parâmetro `--snr_gamma` e defina-o com o valor recomendado de 5.0:
```bash
accelerate launch train_text_to_image_prior.py \
--snr_gamma=5.0
```
## Script de treinamento
O script de treinamento também é semelhante ao guia de treinamento [Texto-para-imagem](text2image#training-script), mas foi modificado para suportar o treinamento dos modelos prior e decoder. Este guia se concentra no código que é exclusivo dos scripts de treinamento do Kandinsky 2.2.
A função [`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L441) contém o código para preparar o conjunto de dados e treinar o modelo.
Uma das principais diferenças que você notará de imediato é que o script de treinamento também carrega um [CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor) - além de um scheduler e um tokenizer - para pré-processar as imagens e um modelo [CLIPVisionModelWithProjection](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPVisionModelWithProjection) para codificar as imagens:
```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()
```
O Kandinsky usa um [PriorTransformer](/docs/diffusers/main/pt/api/models/prior_transformer#diffusers.PriorTransformer) para gerar os embeddings da imagem, portanto, você deve configurar o otimizador para aprender os parâmetros do modo 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,
)
```
Em seguida, as legendas de entrada são tokenizadas e as imagens são [pré-processadas](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L632) pelo [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
```
Por fim, o [training loop](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L718) converte as imagens de entrada em latentes, adiciona ruído aos embeddings das imagens e faz uma previsão:
```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
```
Se você quiser saber mais sobre como funciona o loop de treinamento, confira o tutorial [Entendendo pipelines, modelos e agendadores](../using-diffusers/write_own_pipeline), que detalha o padrão básico do processo de remoção de ruído.
A função [`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L440) contém o código para preparar o conjunto de dados e treinar o modelo.
Diferentemente do modelo anterior, o decodificador inicializa um [VQModel](/docs/diffusers/main/pt/api/models/vq#diffusers.VQModel) para decodificar os latentes em imagens e utiliza um [UNet2DConditionModel](/docs/diffusers/main/pt/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")
```
Em seguida, o script inclui várias transformações de imagem e uma função de [pré-processamento](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L622) para aplicar as transformações às imagens e retornar os valores de pixels:
```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
```
Por fim, o [training loop](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L706) é responsável por converter as imagens em latentes, adicionar ruído e prever o resíduo de ruído.
Se você quiser saber mais sobre como funciona o loop de treinamento, confira o tutorial [Entendendo pipelines, modelos e agendadores](../using-diffusers/write_own_pipeline), que detalha o padrão básico do processo de remoção de ruído.
```py
model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]
```
## Inicie o script
Depois de fazer todas as alterações ou se você estiver satisfeito com a configuração padrão, está pronto para executar o script de treinamento! 🚀
Você treinará no conjunto de dados [Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) para gerar seus próprios personagens de Naruto, mas também pode criar e treinar em seu próprio conjunto de dados seguindo o guia [Create a dataset for training](create_dataset). Defina a variável de ambiente `DATASET_NAME` com o nome do conjunto de dados no Hub ou, se estiver treinando em seus próprios arquivos, defina a variável de ambiente `TRAIN_DIR` com um caminho para o seu conjunto de dados.
Se você estiver treinando em mais de uma GPU, adicione o parâmetro `--multi_gpu` ao comando `accelerate launch`.
> [!TIP]
> Para monitorar o progresso do treinamento com Weights & Biases, adicione o parâmetro `--report_to=wandb` ao comando de treinamento. Você também precisará adicionar o `--validation_prompt` ao comando de treinamento para acompanhar os resultados. Isso pode ser muito útil para depurar o modelo e visualizar resultados intermediários.
```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"
```
Quando o treinamento for concluído, você poderá usar o seu modelo recém-treinado para inferência!
```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]
> Sinta-se à vontade para substituir `kandinsky-community/kandinsky-2-2-decoder` pelo seu próprio checkpoint de decoder treinado!
```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]
```
Para o modelo decodificador, você também pode executar inferência a partir de um checkpoint salvo, o que pode ser útil para visualizar resultados intermediários. Nesse caso, carregue o checkpoint no 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]
```
## Próximas etapas
Parabéns por treinar um modelo Kandinsky 2.2! Para saber mais sobre como usar seu novo modelo, os seguintes guias podem ser úteis:
- Leia o guia do [Kandinsky](../api/pipelines/kandinsky) para aprender a usá-lo em uma variedade de tarefas diferentes (texto-para-imagem, imagem-para-imagem, inpainting, interpolação) e como ele pode ser combinado com um ControlNet.
- Confira os guias de treinamento de [DreamBooth](dreambooth) e [LoRA](lora) para aprender a treinar um modelo Kandinsky personalizado com apenas algumas imagens de exemplo. Essas duas técnicas de treinamento podem até ser combinadas!

Xet Storage Details

Size:
16.1 kB
·
Xet hash:
bd0ae2fdb80a7d669a5d23e814aeb8458b7da21540f9fc37dcdd58e25ee7ec6f

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