DreamBooth

DreamBooth é uma técnica de treinamento que atualiza o modelo de difusão inteiro treinando em apenas algumas imagens de um sujeito ou estilo. Ele funciona associando uma palavra especial no prompt às imagens de exemplo.

Para carregar um checkpoint treinado para inferência, consulte Carregar um modelo DreamBooth para inferência.

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

DreamBooth é muito sensível aos hiperparâmetros de treinamento e é fácil sofrer overfitting. Leia o post no blog Training Stable Diffusion with Dreambooth using 🧨 Diffusers para obter configurações recomendadas para diferentes sujeitos e ajudá-lo a escolher os hiperparâmetros apropriados.

O script de treinamento oferece muitos parâmetros para personalizar sua execução de treinamento. Todos os parâmetros e suas descrições estão na função parse_args(). Os parâmetros são definidos com valores padrão que devem funcionar bem sem ajustes, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.

Por exemplo, para treinar no formato bf16:

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

Alguns parâmetros básicos e importantes que você deve conhecer e especificar são:

--pretrained_model_name_or_path: o nome do modelo no Hub ou um caminho local para o modelo pré-treinado
--instance_data_dir: caminho para uma pasta contendo o conjunto de dados de treinamento (imagens de exemplo)
--instance_prompt: o prompt de texto que contém a palavra especial para as imagens de exemplo
--train_text_encoder: se o text encoder também deve ser treinado
--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_dreambooth.py \
  --snr_gamma=5.0

Perda de preservação de prioridade

A perda de preservação de prioridade é um método que utiliza os próprios amostras geradas por um modelo para ajudá-lo a aprender a gerar imagens mais diversas. Como essas imagens de amostra geradas pertencem à mesma classe das imagens fornecidas, elas ajudam o modelo a reter o que ele aprendeu sobre a classe e como pode usar o que já sabe sobre a classe para criar novas composições.

--with_prior_preservation: se deve usar a perda de preservação de prior
--prior_loss_weight: controla a influência da perda de preservação do prior no modelo
--class_data_dir: caminho para uma pasta contendo as imagens de exemplo das classes geradas
--class_prompt: o prompt de texto descrevendo a classe das imagens de amostra geradas

accelerate launch train_dreambooth.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="path/to/class/images" \
  --class_prompt="text prompt describing class"

Treinar o codificador de texto

Para melhorar a qualidade das saídas geradas, você também pode treinar o text encoder além do UNet. Isso requer memória adicional e você precisará de uma GPU com pelo menos 24GB de vRAM. Se você tiver o hardware necessário, treinar o text encoder produz melhores resultados, especialmente ao gerar imagens de rostos. Ative essa opção por:

accelerate launch train_dreambooth.py \
  --train_text_encoder

Script de treinamento

O DreamBooth vem com suas próprias classes de dataset:

DreamBoothDataset: pré-processa as imagens e as imagens de classe, e tokeniza os prompts para o treinamento
PromptDataset: gera os embeddings de prompt para gerar as imagens de classe

Se você ativou a prior preservation loss, as imagens de classe são geradas aqui:

sample_dataset = PromptDataset(args.class_prompt, num_new_images)
sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)

sample_dataloader = accelerator.prepare(sample_dataloader)
pipeline.to(accelerator.device)

for example in tqdm(
    sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process
):
    images = pipeline(example["prompt"]).images

A seguir está a função main(), que trata da configuração do conjunto de dados para o treinamento e do próprio loop de treinamento. O script carrega o tokenizer, o agendador e os modelos:

# Load the tokenizer
if args.tokenizer_name:
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, revision=args.revision, use_fast=False)
elif args.pretrained_model_name_or_path:
    tokenizer = AutoTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
        use_fast=False,
    )

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = text_encoder_cls.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)

if model_has_vae(args):
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision
    )
else:
    vae = None

unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

Em seguida, é hora de criar o conjunto de dados de treinamento e o DataLoader a partir de DreamBoothDataset:

train_dataset = DreamBoothDataset(
    instance_data_root=args.instance_data_dir,
    instance_prompt=args.instance_prompt,
    class_data_root=args.class_data_dir if args.with_prior_preservation else None,
    class_prompt=args.class_prompt,
    class_num=args.num_class_images,
    tokenizer=tokenizer,
    size=args.resolution,
    center_crop=args.center_crop,
    encoder_hidden_states=pre_computed_encoder_hidden_states,
    class_prompt_encoder_hidden_states=pre_computed_class_prompt_encoder_hidden_states,
    tokenizer_max_length=args.tokenizer_max_length,
)

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=args.train_batch_size,
    shuffle=True,
    collate_fn=lambda examples: collate_fn(examples, args.with_prior_preservation),
    num_workers=args.dataloader_num_workers,
)

Por fim, o training loop cuida das etapas restantes, como converter as imagens para o espaço latente, adicionar ruído à entrada, prever o resíduo de ruído e calcular a perda.

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

Agora você está pronto para executar o script de treinamento! 🚀

Neste guia, você baixará algumas imagens de um cachorro 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 = "./dog"
snapshot_download(
    "diffusers/dog-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, INSTANCE_DIR para o caminho onde você acabou de baixar as imagens de cachorro e OUTPUT_DIR para onde você deseja salvar o modelo. Você usará sks como a palavra especial para vincular o treinamento.

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 photo of a sks dog"
--num_validation_images=4
--validation_steps=100

Mais uma coisa antes de você executar o script! Dependendo da GPU que você possui, pode ser necessário habilitar certas otimizações para treinar o DreamBooth.

16GB

12GB

8GB

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export INSTANCE_DIR="./dog"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub

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

Não consegue esperar para testar seu modelo para inferência antes de o treinamento ser concluído? 🤭 Certifique-se de ter a versão mais recente do 🤗 Accelerate instalada.

from diffusers import DiffusionPipeline, UNet2DConditionModel
from transformers import CLIPTextModel
import torch

unet = UNet2DConditionModel.from_pretrained("path/to/model/checkpoint-100/unet")

# if you have trained with `--args.train_text_encoder` make sure to also load the text encoder
text_encoder = CLIPTextModel.from_pretrained("path/to/model/checkpoint-100/checkpoint-100/text_encoder")

pipeline = DiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, text_encoder=text_encoder, dtype=torch.float16,
).to("cuda")

image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("path_to_saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
image = pipeline("A photo of sks dog in a bucket", num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("dog-bucket.png")

LoRA

LoRA é uma técnica de treinamento para reduzir significativamente o número de parâmetros treináveis. Como resultado, o treinamento é mais rápido e fica mais fácil armazenar os pesos resultantes, pois eles são muito menores (~100MBs). Use o script train_dreambooth_lora.py para treinar com LoRA.

O script de treinamento LoRA é discutido em mais detalhes no guia treinamento LoRA.

Stable Diffusion XL

Stable Diffusion XL (SDXL) é um modelo poderoso de texto para imagem que gera imagens em alta resolução e adiciona um segundo text-encoder à sua arquitetura. Use o script train_dreambooth_lora_sdxl.py para treinar um modelo SDXL com LoRA.

O script de treinamento do SDXL é discutido em mais detalhes no guia de treinamento do SDXL.

DeepFloyd IF

DeepFloyd IF é um modelo de difusão de pixels em cascata com três estágios. O primeiro estágio gera uma imagem base e os segundo e terceiro estágios aumentam progressivamente a resolução da imagem base até uma imagem de alta resolução de 1024x1024. Use os scripts train_dreambooth_lora.py ou train_dreambooth.py para treinar um modelo DeepFloyd IF com LoRA ou o modelo completo.

O DeepFloyd IF usa variância prevista, mas os scripts de treinamento do Diffusers usam erro previsto, de modo que os modelos treinados do DeepFloyd IF são alterados para usar um agendamento de variância fixo. Os scripts de treinamento atualizam a configuração do agendador do modelo totalmente treinado para você. No entanto, ao carregar os pesos LoRA salvos, você também deve atualizar a configuração do agendador do pipeline.

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", use_safetensors=True)

pipe.load_lora_weights("<lora weights path>")

# Update scheduler config to fixed variance schedule
pipe.scheduler = pipe.scheduler.__class__.from_config(pipe.scheduler.config, variance_type="fixed_small")

O modelo da etapa 2 requer imagens de validação adicionais para aumentar a resolução. Você pode baixar e usar uma versão reduzida das imagens de treinamento para isso.

from huggingface_hub import snapshot_download

local_dir = "./dog_downsized"
snapshot_download(
    "diffusers/dog-example-downsized",
    local_dir=local_dir,
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

Os exemplos de código abaixo fornecem uma visão geral de como treinar um modelo DeepFloyd IF com uma combinação de DreamBooth e LoRA. Alguns parâmetros importantes a observar são:

--resolution=64, uma resolução muito menor é necessária porque o DeepFloyd IF é um modelo de difusão de pixels e, para trabalhar com pixels não comprimidos, as imagens de entrada devem ser menores
--pre_compute_text_embeddings, calcula os embeddings de texto antecipadamente para economizar memória, pois o T5Model pode consumir bastante memória
--tokenizer_max_length=77, você pode usar um comprimento de texto padrão maior com o T5 como codificador de texto, mas o procedimento de codificação do modelo padrão usa um comprimento de texto menor
--text_encoder_use_attention_mask, para passar a máscara de atenção para o codificador de texto

Etapa 1 LoRA DreamBooth

Etapa 2 LoRA DreamBooth

Etapa 1 DreamBooth

Etapa 2 DreamBooth

O estágio 1 de treinamento do DeepFloyd IF com LoRA e DreamBooth requer ~28GB de memória.

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="dreambooth_dog_lora"

accelerate launch train_dreambooth_lora.py \
  --report_to wandb \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks dog" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --validation_prompt="a sks dog" \
  --validation_epochs=25 \
  --checkpointing_steps=100 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

Dicas de treinamento

Treinar o modelo DeepFloyd IF pode ser desafiador, mas aqui estão algumas dicas que nos foram úteis:

LoRA é suficiente para treinar o modelo da etapa 1, pois a baixa resolução do modelo dificulta a representação de detalhes mais finos, independentemente do método utilizado.
Para objetos comuns ou simples, você não necessariamente precisa fazer o ajuste fino do upscaler. Certifique-se de que o prompt passado para o upscaler seja ajustado para remover o novo token do prompt da instância. Por exemplo, se o seu prompt da etapa 1 for “a sks dog”, então o seu prompt da etapa 2 deve ser “a dog”.
Para detalhes mais finos, como rostos, treinar completamente o upscaler da etapa 2 é melhor do que treinar o modelo da etapa 2 com LoRA. Também ajuda usar taxas de aprendizado mais baixas com tamanhos de lote maiores.
Taxas de aprendizado mais baixas devem ser usadas para treinar o modelo da etapa 2.
O DDPMScheduler funciona melhor do que o DPMSolver usado nos scripts de treinamento.

Próximas etapas

Parabéns por treinar seu modelo DreamBooth! Para saber mais sobre como usar seu novo modelo, o guia a seguir pode ser útil:

Aprenda como carregar um checkpoint do DreamBooth para inferência.
Se você treinou com LoRA, consulte LoRA para inferência.

Update on GitHub