Texto para imagem

O script de texto para imagem é experimental, e é fácil sofrer overfitting e encontrar problemas como o esquecimento catastrófico. Tente explorar diferentes hiperparâmetros para obter os melhores resultados no seu conjunto de dados.

Modelos de texto para imagem, como o Stable Diffusion, são condicionados para gerar imagens a partir de um prompt de texto.

Treinar um modelo pode ser exigente para o seu hardware, mas, se você ativar gradient_checkpointing e mixed_precision, é possível treinar um modelo em uma única GPU de 24GB. Se você estiver treinando com tamanhos de lote maiores ou quiser treinar mais rápido, é melhor usar GPUs com mais de 30GB de memória. Você pode reduzir a sua pegada de memória ativando a atenção eficiente em memória com xFormers.

Este guia explorará o script de treinamento train_text_to_image.py para ajudá-lo a se familiarizar 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 .

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

Parâmetros do script

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.

O script de treinamento fornece muitos parâmetros para ajudar você a personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão encontrados na função parse_args(). Esta função fornece 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.py \
  --mixed_precision="fp16"

Alguns parâmetros básicos e importantes incluem:

• --pretrained_model_name_or_path: o nome do modelo no Hub ou um caminho local para o modelo pré-treinado
• --dataset_name: o nome do conjunto de dados no Hub ou um caminho local para o conjunto de dados em que o treinamento será realizado
• --image_column: o nome da coluna de imagens no conjunto de dados para treinamento
• --caption_column: o nome da coluna de texto no conjunto de dados para treinamento
• --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

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.py \
  --snr_gamma=5.0

Você pode comparar as superfícies de perda para diferentes valores de snr_gamma neste relatório do Weights and Biases. Para conjuntos de dados menores, os efeitos do Min-SNR podem não ser tão evidentes em comparação com conjuntos de dados maiores.

Script de treinamento

O código de pré-processamento do conjunto de dados e o loop de treinamento estão na função main(). Se você precisar adaptar o script de treinamento, é aqui que você precisará fazer suas alterações.

O script train_text_to_image começa carregando um scheduler e um tokenizer. Você pode escolher usar um scheduler diferente aqui, se desejar:

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

Em seguida, o script carrega o modelo UNet:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

Em seguida, as colunas de texto e imagem do conjunto de dados precisam ser pré-processadas. A função tokenize_captions trata a tokenização das entradas, e a função train_transforms especifica o tipo de transformações a serem aplicadas à imagem. Ambas as funções são agrupadas em preprocess_train:

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["input_ids"] = tokenize_captions(examples)
    return examples

Por fim, o training loop cuida de tudo o mais. Ele codifica as imagens no espaço latente, adiciona ruído aos latentes, calcula os embeddings de texto para condicionar, atualiza os parâmetros do modelo e salva e envia o modelo para o Hub. Se você quiser aprender mais sobre como o training loop funciona, confira o tutorial Understanding pipelines, models and schedulers, 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! 🚀

Vamos treinar no conjunto de dados Naruto BLIP captions para gerar seus próprios personagens de Naruto. Defina as variáveis de ambiente MODEL_NAME e dataset_name para o modelo e o conjunto de dados (seja do Hub ou um caminho local). Se você estiver treinando em mais de uma GPU, adicione o parâmetro --multi_gpu ao comando accelerate launch.

Para treinar em um conjunto de dados local, defina as variáveis de ambiente TRAIN_DIR e OUTPUT_DIR com o caminho do conjunto de dados e o local onde o modelo será salvo.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

Quando o treinamento estiver concluído, você pode usar o seu modelo recém-treinado para inferência:

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

Próximas etapas

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

• Aprenda como carregar pesos LoRA para inferência se você treinou seu modelo com LoRA.
• Saiba mais sobre como certos parâmetros, como a escala de orientação, ou técnicas, como o ponderamento de prompts, podem ajudar você a controlar a inferência no guia de tarefas de Texto para imagem.