Stable Diffusion XL

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.

Stable Diffusion XL (SDXL) é uma iteração maior e mais poderosa do modelo Stable Diffusion, capaz de produzir imagens com resolução mais alta.

O UNet do SDXL é 3x maior e o modelo adiciona um segundo codificador de texto à arquitetura. Dependendo do hardware disponível para você, isso pode ser muito intensivo computacionalmente e pode não rodar em uma GPU de consumidor como uma Tesla T4. Para ajudar a caber esse modelo maior na memória e acelerar o treinamento, tente habilitar gradient_checkpointing, mixed_precision e gradient_accumulation_steps. Você pode reduzir ainda mais o uso de memória habilitando a atenção eficiente em memória com xFormers e usando o otimizador de 8 bits do bitsandbytes.

Este guia explorará o script de treinamento train_text_to_image_sdxl.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 .

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_sdxl.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 bf16, adicione o parâmetro --mixed_precision ao comando de treinamento:

accelerate launch train_text_to_image_sdxl.py \
  --mixed_precision="bf16"

A maioria dos parâmetros é idêntica aos parâmetros do guia de treinamento Texto-para-imagem, então você se concentrará nos parâmetros relevantes para o treinamento do SDXL neste guia.

• --pretrained_vae_model_name_or_path: caminho para um VAE pré-treinado; o VAE do SDXL é conhecido por sofrer de instabilidade numérica, portanto, este parâmetro permite especificar um VAE melhor
• --proportion_empty_prompts: a proporção de prompts de imagem a serem substituídos por strings vazias
• --timestep_bias_strategy: onde (mais cedo ou mais tarde) no timestep aplicar um viés, o que pode incentivar o modelo a aprender detalhes de baixa ou alta frequência
• --timestep_bias_multiplier: o peso do viés a ser aplicado ao timestep
• --timestep_bias_begin: o timestep em que se começa a aplicar o viés
• --timestep_bias_end: o timestep em que a aplicação do viés termina
• --timestep_bias_portion: a proporção de timesteps para aplicar o viés

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_sdxl.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 do SDXL. Este guia se concentrará no código que é exclusivo do script de treinamento do SDXL.

Começa criando funções para tokenizar os prompts para calcular os embeddings dos prompts e para calcular os embeddings das imagens com o VAE. Em seguida, você criará uma função para gerar os pesos dos timesteps dependendo do número de timesteps e da estratégia de viés de timestep a ser aplicada.

Dentro da função main(), além de carregar um tokenizer, o script carrega um segundo tokenizer e um codificador de texto, pois a arquitetura SDXL usa dois de cada:

tokenizer_one = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False
)
tokenizer_two = AutoTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer_2", revision=args.revision, use_fast=False
)

text_encoder_cls_one = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision
)
text_encoder_cls_two = import_model_class_from_model_name_or_path(
    args.pretrained_model_name_or_path, args.revision, subfolder="text_encoder_2"
)

Os embeddings do prompt e da imagem são calculados primeiro e mantidos na memória, o que normalmente não é um problema para um conjunto de dados menor, mas para conjuntos de dados maiores pode levar a problemas de memória. Se esse for o caso, você deve salvar os embeddings pré-calculados no disco separadamente e carregá-los na memória durante o processo de treinamento (veja este PR para mais discussões sobre esse tópico).

text_encoders = [text_encoder_one, text_encoder_two]
tokenizers = [tokenizer_one, tokenizer_two]
compute_embeddings_fn = functools.partial(
    encode_prompt,
    text_encoders=text_encoders,
    tokenizers=tokenizers,
    proportion_empty_prompts=args.proportion_empty_prompts,
    caption_column=args.caption_column,
)

train_dataset = train_dataset.map(compute_embeddings_fn, batched=True, new_fingerprint=new_fingerprint)
train_dataset = train_dataset.map(
    compute_vae_encodings_fn,
    batched=True,
    batch_size=args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps,
    new_fingerprint=new_fingerprint_for_vae,
)

Após calcular os embeddings, o codificador de texto, o VAE e o tokenizer são removidos para liberar um pouco de memória:

del text_encoders, tokenizers, vae
gc.collect()
torch.cuda.empty_cache()

Por fim, o training loop cuida do restante. Se você optou por aplicar uma estratégia de viés de timestep, verá que os pesos dos timesteps são calculados e adicionados como ruído:

weights = generate_timestep_weights(args, noise_scheduler.config.num_train_timesteps).to(
        model_input.device
    )
    timesteps = torch.multinomial(weights, bsz, replacement=True).long()

noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)

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

Vamos treinar no dataset 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 dataset (seja do Hub ou um caminho local). Você também deve especificar um VAE diferente do VAE do SDXL (seja do Hub ou um caminho local) com VAE_NAME para evitar instabilidades numéricas.

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 --validation_prompt e --validation_epochs ao comando de treinamento para acompanhar os resultados. Isso pode ser muito útil para depurar o modelo e visualizar resultados intermediários.

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" \
  --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="sdxl-naruto-model" \
  --push_to_hub

Depois de terminar o treinamento, você pode usar seu novo modelo SDXL treinado para inferência!

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path/to/your/model", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A naruto with green eyes and red legs."
image = pipeline(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("naruto.png")

from diffusers import DiffusionPipeline
import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()
pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to(device)

prompt = "A naruto with green eyes and red legs."
start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Compilation time is {time()-start} sec')
image.save("naruto.png")

start = time()
image = pipeline(prompt, num_inference_steps=inference_steps).images[0]
print(f'Inference time is {time()-start} sec after compilation')

Próximas etapas

Parabéns por treinar um modelo SDXL! Para saber mais sobre como usar seu novo modelo, os seguintes guias podem ser úteis:

• Leia o guia do Stable Diffusion XL para aprender a usá-lo em uma variedade de tarefas diferentes (texto-para-imagem, imagem-para-imagem, inpainting), como usar seu modelo refinador e os diferentes tipos de micro-condicionamentos.
• Confira os guias de treinamento de DreamBooth e LoRA para aprender a treinar um modelo SDXL personalizado com apenas algumas imagens de exemplo. Essas duas técnicas de treinamento podem até ser combinadas!