Kandinsky 2.2

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 (a versão v0.0.16 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 e 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:

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:

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

🤗 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 do 🤗 Accelerate para saber mais.

Inicialize um ambiente do 🤗 Accelerate:

accelerate config

Para configurar um ambiente padrão do 🤗 Accelerate sem escolher nenhuma configuração:

accelerate config default

Ou, se o seu ambiente não suporta um shell interativo, como um notebook, você pode usar:

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 para aprender a criar um conjunto de dados que funcione com o script de treinamento.

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(). 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:

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, 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 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:

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, 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.

modelo anterior

modelo decodificador

A função main() 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 - além de um scheduler e um tokenizer - para pré-processar as imagens e um modelo CLIPVisionModelWithProjection para codificar as imagens:

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 para gerar os embeddings da imagem, portanto, você deve configurar o otimizador para aprender os parâmetros do modo 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,
)

Em seguida, as legendas de entrada são tokenizadas e as imagens são pré-processadas pelo 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

Por fim, o training loop converte as imagens de entrada em latentes, adiciona ruído aos embeddings das imagens e faz uma previsão:

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, que detalha o padrão básico do processo de remoção de ruído.

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 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. 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.

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.

modelo anterior

modelo decodificador

Quando o treinamento for concluído, você poderá usar o seu modelo recém-treinado para inferência!

modelo anterior

modelo decodificador

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 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 e 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!

Update on GitHub