Destilação de Consistência Latente

Latent Consistency Models (LCMs) são capazes de gerar imagens de alta qualidade em apenas algumas etapas, representando um grande avanço, pois muitos pipelines exigem pelo menos 25+ etapas. Os LCMs são produzidos aplicando o método de destilação de consistência latente a qualquer modelo Stable Diffusion. Esse método funciona aplicando one-stage guided distillation ao espaço latente e incorporando um método de skipping-step para pular consistentemente os timesteps e acelerar o processo de destilação (consulte as seções 4.1, 4.2 e 4.3 do paper para mais detalhes).

Se você está treinando em uma GPU com vRAM limitada, tente habilitar gradient_checkpointing, gradient_accumulation_steps e mixed_precision para reduzir o uso de memória e acelerar o treinamento. Você pode reduzir ainda mais o uso de memória habilitando a atenção de eficiência de memória com xFormers e o otimizador de 8 bits do bitsandbytes.

Este guia explorará o script train_lcm_distill_sd_wds.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/consistency_distillation
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 🤗 Accelerate (tente habilitar torch.compile para acelerar significativamente o treinamento):

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_lcm_distill_sd_wds.py \
  --mixed_precision="fp16"

A maioria dos parâmetros é idêntica aos parâmetros do guia de treinamento Texto-para-imagem, portanto, você se concentrará nos parâmetros relevantes para a destilação de consistência latente neste guia.

--pretrained_teacher_model: o caminho para um modelo de difusão latente pré-treinado para usar como modelo professor
--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 alternativo (como este VAE) criado por madebyollin que funciona em fp16)
--w_min e --w_max: os valores mínimo e máximo da escala de orientação para a amostragem da escala de orientação
--num_ddim_timesteps: o número de timesteps para a amostragem DDIM
--loss_type: o tipo de perda (L2 ou Huber) a ser calculada para a destilação de consistência latente; a perda Huber é geralmente preferida porque é mais robusta a valores atípicos
--huber_c: o parâmetro da perda de Huber

Script de treinamento

O script de treinamento começa criando uma classe de conjunto de dados - Text2ImageDataset - para pré-processar as imagens e criar um conjunto de dados de treinamento.

def transform(example):
    image = example["image"]
    image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)

    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
    image = TF.crop(image, c_top, c_left, resolution, resolution)
    image = TF.to_tensor(image)
    image = TF.normalize(image, [0.5], [0.5])

    example["image"] = image
    return example

Para melhorar o desempenho na leitura e gravação de grandes conjuntos de dados armazenados na nuvem, este script usa o formato WebDataset para criar um pipeline de pré-processamento que aplica transformações e cria um dataset e um dataloader para o treinamento. As imagens são processadas e enviadas para o loop de treinamento sem a necessidade de baixar o conjunto de dados completo primeiro.

processing_pipeline = [
    wds.decode("pil", handler=wds.ignore_and_continue),
    wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
    wds.map(filter_keys({"image", "text"})),
    wds.map(transform),
    wds.to_tuple("image", "text"),
]

Na função main(), todos os componentes necessários, como o noise scheduler, os tokenizers, os text encoders e o VAE, são carregados. O UNet teacher também é carregado aqui e, em seguida, é possível criar um UNet student a partir do UNet teacher. O UNet student é atualizado pelo optimizer durante o treinamento.

teacher_unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)

unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()

Agora você pode criar o optimizer para atualizar os parâmetros do UNet:

optimizer = optimizer_class(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Crie o dataset:

dataset = Text2ImageDataset(
    train_shards_path_or_url=args.train_shards_path_or_url,
    num_train_examples=args.max_train_samples,
    per_gpu_batch_size=args.train_batch_size,
    global_batch_size=args.train_batch_size * accelerator.num_processes,
    num_workers=args.dataloader_num_workers,
    resolution=args.resolution,
    shuffle_buffer_size=1000,
    pin_memory=True,
    persistent_workers=True,
)
train_dataloader = dataset.train_dataloader

A seguir, você está pronto para configurar o training loop e implementar o método de destilação de consistência latente (veja o Algoritmo 1 no artigo para mais detalhes). Esta seção do script cuida de adicionar ruído aos latentes, amostrar e criar um embedding de escala de orientação, e prever a imagem original a partir do ruído.

pred_x_0 = predicted_origin(
    noise_pred,
    start_timesteps,
    noisy_model_input,
    noise_scheduler.config.prediction_type,
    alpha_schedule,
    sigma_schedule,
)

model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0

Em seguida, ele obtém as previsões do modelo professor e as previsões do LCM, calcula a perda e, então, a retropropaga para o LCM.

if args.loss_type == "l2":
    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
    loss = torch.mean(
        torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
    )

Se você quiser saber mais sobre como o loop de treinamento funciona, confira o tutorial sobre o entendimento de 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 e começar a destilar!

Neste guia, você usará o --train_shards_path_or_url para especificar o caminho para o conjunto de dados Conceptual Captions 12M armazenado no Hub aqui. Defina a variável de ambiente MODEL_DIR com o nome do modelo teacher e OUTPUT_DIR com o local em que deseja salvar o modelo.

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"

accelerate launch train_lcm_distill_sd_wds.py \
    --pretrained_teacher_model=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --mixed_precision=fp16 \
    --resolution=512 \
    --learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
    --max_train_steps=1000 \
    --max_train_samples=4000000 \
    --dataloader_num_workers=8 \
    --train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
    --validation_steps=200 \
    --checkpointing_steps=200 --checkpoints_total_limit=10 \
    --train_batch_size=12 \
    --gradient_checkpointing --enable_xformers_memory_efficient_attention \
    --gradient_accumulation_steps=1 \
    --use_8bit_adam \
    --resume_from_checkpoint=latest \
    --report_to=wandb \
    --seed=453645634 \
    --push_to_hub

Quando o treinamento estiver concluído, você pode usar seu novo LCM para inferência.

from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch

unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")

pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "sushi rolls in the form of panda heads, sushi platter"

image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]

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_lcm_distill_lora_sd_wds.py ou train_lcm_distill_lora_sdxl.wds.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_lcm_distill_sdxl_wds.py para treinar um modelo SDXL com LoRA.

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

Próximas etapas

Parabéns por destilar um modelo LCM! Para saber mais sobre o LCM, o seguinte pode ser útil:

Aprenda a usar LCMs para inferência em texto-para-imagem, imagem-para-imagem e com checkpoints LoRA.
Leia o post no blog SDXL em 4 etapas com Latent Consistency LoRAs para saber mais sobre os SDXL LCM-LoRA para inferência ultrarrápida, comparações de qualidade, benchmarks e muito mais.

Update on GitHub

Diffusers

Destilação de Consistência Latente

Parâmetros do script

Script de treinamento

Inicie o script

LoRA

Stable Diffusion XL

Próximas etapas