Inversão textual

Textual Inversion é uma técnica de treinamento para personalizar modelos de geração de imagens com apenas algumas imagens de exemplo do que você deseja que ele aprenda. Essa técnica funciona aprendendo e atualizando os embeddings de texto (os novos embeddings estão vinculados a uma palavra especial que você deve usar no prompt) para corresponder às imagens de exemplo que você fornece.

Para inferência com embeddings treinados, consulte Inferência de inversão textual.

Se você estiver treinando em uma GPU com vRAM limitado, deve tentar habilitar os parâmetros gradient_checkpointing e mixed_precision no comando de treinamento. Você também pode reduzir o uso de memória usando atenção eficiente em memória com xFormers.

Este guia explorará o script textual_inversion.py para ajudá-lo a se familiarizar melhor com ele e como você pode adaptá-lo para o seu próprio caso de uso.

Antes de executar o script, certifique-se de instalar a biblioteca a partir do código-fonte:

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

Navegue até a pasta de exemplo com o script de treinamento e instale as dependências necessárias para o script que você está usando:

cd examples/textual_inversion
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 do script de treinamento que são importantes para entender como modificá-lo, mas não cobrem todos os aspectos do script em detalhes. Se você tiver interesse em aprender mais, sinta-se à vontade para ler o script e nos informar caso tenha alguma dúvida ou preocupação.

Parâmetros do script

O script de treinamento possui muitos parâmetros para ajudar você a ajustar a execução do treinamento às suas necessidades. Todos os parâmetros e suas descrições estão listados na função parse_args(). Onde aplicável, o Diffusers fornece valores padrão para cada parâmetro, como o tamanho do lote de treinamento e a taxa de aprendizado, mas sinta-se à vontade para alterar esses valores no comando de treinamento, se desejar.

Por exemplo, para aumentar o número de passos de acúmulo de gradiente acima do valor padrão de 1:

accelerate launch textual_inversion.py \
  --gradient_accumulation_steps=4

Outros parâmetros básicos e importantes a serem especificados incluem:

• --pretrained_model_name_or_path: o nome do modelo no Hub ou um caminho local para o modelo pré-treinado
• --train_data_dir: caminho para uma pasta contendo o conjunto de dados de treinamento (imagens de exemplo)
• --output_dir: onde salvar o modelo treinado
• --push_to_hub: se o modelo treinado deve ser enviado para o Hub
• --checkpointing_steps: frequência de salvamento de um checkpoint durante o treinamento do modelo; isso é útil caso, por algum motivo, o treinamento seja interrompido, pois você pode continuar o treinamento a partir desse checkpoint adicionando --resume_from_checkpoint ao seu comando de treinamento
• --num_vectors: o número de vetores para aprender os embeddings; aumentar esse parâmetro ajuda o modelo a aprender melhor, mas isso vem com custos de treinamento aumentados
• --placeholder_token: a palavra especial para vincular os embeddings aprendidos (você deve usar a palavra no seu prompt para inferência)
• --initializer_token: uma única palavra que descreve de forma aproximada o objeto ou estilo em que você está tentando treinar
• --learnable_property: se você está treinando o modelo para aprender um novo “estilo” (por exemplo, o estilo de pintura de Van Gogh) ou “objeto” (por exemplo, seu cachorro)

Script de treinamento

Diferentemente de alguns dos outros scripts de treinamento, textual_inversion.py possui uma classe de dataset personalizada, TextualInversionDataset, para a criação de um dataset. Você pode personalizar o tamanho da imagem, o token placeholder, o método de interpolação, se deve recortar a imagem, entre outras opções. Se precisar alterar a forma como o dataset é criado, você pode modificar TextualInversionDataset.

Em seguida, você encontrará o código de pré-processamento do conjunto de dados e o loop de treinamento na função main().

O script começa carregando o tokenizer, agendador e modelo:

# Load tokenizer
if args.tokenizer_name:
    tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
elif args.pretrained_model_name_or_path:
    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

O token de marcador de posição especial é adicionado ao lado do tokenizer, e o embedding é reajustado para levar em conta o novo token.

Em seguida, o script cria um dataset a partir do TextualInversionDataset:

train_dataset = TextualInversionDataset(
    data_root=args.train_data_dir,
    tokenizer=tokenizer,
    size=args.resolution,
    placeholder_token=(" ".join(tokenizer.convert_ids_to_tokens(placeholder_token_ids))),
    repeats=args.repeats,
    learnable_property=args.learnable_property,
    center_crop=args.center_crop,
    set="train",
)
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)

Por fim, o training loop cuida de tudo o mais, desde a previsão do residual ruidoso até a atualização dos pesos de incorporação do token especial de espaço reservado.

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

Neste guia, você baixará algumas imagens de um brinquedo de gato e as armazenará em um diretório. Mas lembre-se: você pode criar e usar seu próprio conjunto de dados, se desejar (veja o guia Criar um conjunto de dados para treinamento).

from huggingface_hub import snapshot_download

local_dir = "./cat"
snapshot_download(
    "diffusers/cat_toy_example", local_dir=local_dir, repo_type="dataset", ignore_patterns=".gitattributes"
)

Defina a variável de ambiente MODEL_NAME para um id de modelo no Hub ou um caminho para um modelo local, e DATA_DIR para o caminho onde você acabou de baixar as imagens de gatos. O script cria e salva os seguintes arquivos no seu repositório:

• learned_embeds.bin: os vetores de incorporação aprendidos correspondentes às suas imagens de exemplo
• token_identifier.txt: o token especial de marcador de posição
• type_of_concept.txt: o tipo de conceito em que você está treinando (seja “object” ou “style”)

Um treinamento completo leva cerca de 1 hora em uma única GPU V100.

Mais uma coisa antes de você executar o script. Se você estiver interessado em acompanhar o processo de treinamento, pode salvar periodicamente as imagens geradas à medida que o treinamento avança. Adicione os seguintes parâmetros ao comando de treinamento:

--validation_prompt="A <cat-toy> train"
--num_validation_images=4
--validation_steps=100

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export DATA_DIR="./cat"

accelerate launch textual_inversion.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --learnable_property="object" \
  --placeholder_token="<cat-toy>" \
  --initializer_token="toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_cat" \
  --push_to_hub

Após o treinamento ser concluído, você pode usar seu modelo recém-treinado para inferência da seguinte forma:

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_textual_inversion("sd-concepts-library/cat-toy")
image = pipeline("A <cat-toy> train", num_inference_steps=50).images[0]
image.save("cat-train.png")

Próximas etapas

Parabéns por treinar seu próprio modelo de Textual Inversion! 🎉 Para saber mais sobre como usar seu novo modelo, os seguintes guias podem ser úteis:

• Aprenda como carregar embeddings de Textual Inversion e também usá-los como embeddings negativos.