DreamBooth

DreamBooth es una técnica de entrenamiento que actualiza todo el modelo de difusión entrenando con solo unas pocas imágenes de un sujeto o estilo. Funciona asociando una palabra especial en el prompt con las imágenes de ejemplo.

Para cargar un punto de control entrenado para inferencia, consulte Cargar un modelo DreamBooth para inferencia.

Si estás entrenando en una GPU con vRAM limitada, debes intentar habilitar los parámetros gradient_checkpointing y mixed_precision en el comando de entrenamiento. También puedes reducir tu huella de memoria usando atención eficiente en memoria con xFormers.

Esta guía explorará el script train_dreambooth.py para que se familiarice más con él y para que pueda adaptarlo a su propio caso de uso.

Antes de ejecutar el script, asegúrate de instalar la biblioteca desde el código fuente:

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

Navegue hasta la carpeta de ejemplo con el script de entrenamiento e instale las dependencias requeridas para el script que esté utilizando:

cd examples/dreambooth
pip install -r requirements.txt

🤗 Accelerate es una biblioteca para ayudarte a entrenar en varias GPUs/TPUs o con precisión mixta. Configura automáticamente tu configuración de entrenamiento en función de tu hardware y entorno. Echa un vistazo al Quick tour de 🤗 Accelerate para saber más.

Inicializa un entorno de 🤗 Accelerate:

accelerate config

Para configurar un entorno predeterminado de 🤗 Accelerate sin elegir ninguna configuración:

accelerate config default

O si su entorno no admite una shell interactiva, como un notebook, puede usar:

from accelerate.utils import write_basic_config

write_basic_config()

Por último, si deseas entrenar un modelo con tu propio conjunto de datos, consulta la guía Crear un conjunto de datos para el entrenamiento para aprender a crear un conjunto de datos que funcione con el script de entrenamiento.

Las siguientes secciones resaltan las partes del script de entrenamiento que son importantes para entender cómo modificarlo, pero no abordan todos los aspectos del script en detalle. Si te interesa aprender más, no dudes en leer el script e informarnos si tienes alguna pregunta o inquietud.

Parámetros del script

DreamBooth es muy sensible a los hiperparámetros de entrenamiento y es fácil sobreajustar. Lee el artículo de blog Training Stable Diffusion with Dreambooth using 🧨 Diffusers para conocer las configuraciones recomendadas para distintos sujetos y así elegir los hiperparámetros adecuados.

El script de entrenamiento ofrece muchos parámetros para personalizar tu ejecución de entrenamiento. Todos los parámetros y sus descripciones se encuentran en la función parse_args(). Los parámetros se establecen con valores predeterminados que deberían funcionar bastante bien sin configuración adicional, pero también puedes establecer tus propios valores en el comando de entrenamiento si lo deseas.

Por ejemplo, para entrenar en el formato bf16:

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

Algunos parámetros básicos e importantes que conviene conocer y especificar son:

--pretrained_model_name_or_path: el nombre del modelo en el Hub o una ruta local al modelo preentrenado
--instance_data_dir: ruta a una carpeta que contiene el conjunto de datos de entrenamiento (imágenes de ejemplo)
--instance_prompt: el texto de la instrucción que contiene la palabra especial para las imágenes de ejemplo
--train_text_encoder: si se debe entrenar también el codificador de texto
--output_dir: dónde guardar el modelo entrenado
--push_to_hub: si se debe subir el modelo entrenado al Hub
--checkpointing_steps: frecuencia de guardado de un punto de control durante el entrenamiento del modelo; esto es útil si, por alguna razón, el entrenamiento se interrumpe, ya que puedes continuar el entrenamiento desde ese punto de control añadiendo --resume_from_checkpoint a tu comando de entrenamiento

Ponderación Min-SNR

La estrategia de ponderación Min-SNR puede ayudar en el entrenamiento al reequilibrar la pérdida para lograr una convergencia más rápida. El script de entrenamiento admite la predicción de epsilon (ruido) o v_prediction, pero Min-SNR es compatible con ambos tipos de predicción. Esta estrategia de ponderación solo está soportada por PyTorch.

Agrega el parámetro --snr_gamma y establécelo en el valor recomendado de 5.0:

accelerate launch train_dreambooth.py \
  --snr_gamma=5.0

Pérdida de preservación previa

La pérdida de preservación previa es un método que utiliza las propias muestras generadas por un modelo para ayudarlo a aprender a generar imágenes más diversas. Dado que estas imágenes de muestra generadas pertenecen a la misma clase que las imágenes que proporcionaste, ayudan al modelo a retener lo que ha aprendido sobre la clase y a utilizar lo que ya sabe sobre la clase para crear nuevas composiciones.

--with_prior_preservation: si se debe usar la pérdida de preservación de la distribución previa
--prior_loss_weight: controla la influencia de la pérdida de preservación del prior en el modelo
--class_data_dir: ruta a una carpeta que contiene las imágenes de muestra de clase generadas
--class_prompt: el texto de la instrucción que describe la clase de las imágenes de muestra generadas

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"

Entrenar el codificador de texto

Para mejorar la calidad de las salidas generadas, también puedes entrenar el codificador de texto además de la UNet. Esto requiere memoria adicional y necesitarás una GPU con al menos 24 GB de vRAM. Si tienes el hardware necesario, entrenar el codificador de texto produce mejores resultados, especialmente al generar imágenes de rostros. Habilita esta opción mediante:

accelerate launch train_dreambooth.py \
  --train_text_encoder

Script de entrenamiento

DreamBooth incluye sus propias clases de conjunto de datos:

DreamBoothDataset: preprocesa las imágenes y las imágenes de clase, y tokeniza los prompts para el entrenamiento
PromptDataset: genera los embeddings de los prompts para generar las imágenes de clase

Si habilitaste la pérdida de preservación previa, las imágenes de clase se generan aquí:

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 continuación se presenta la función main(), que se encarga de configurar el conjunto de datos para el entrenamiento y del propio bucle de entrenamiento. El script carga el tokenizer, el programador y los 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
)

Luego, es momento de crear el conjunto de datos de entrenamiento y el 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 último, el bucle de entrenamiento se encarga de los pasos restantes, como convertir las imágenes al espacio latente, añadir ruido a la entrada, predecir el residuo de ruido y calcular la pérdida.

Si quieres saber más sobre cómo funciona el bucle de entrenamiento, consulta el tutorial Understanding pipelines, models and schedulers, que desglosa el patrón básico del proceso de desruido.

Ejecutar el script

¡Ahora estás listo para lanzar el script de entrenamiento! 🚀

Para esta guía, descargará algunas imágenes de un perro y las guardará en un directorio. Pero recuerde que puede crear y usar su propio conjunto de datos si lo desea (consulte la guía Crear un conjunto de datos para el entrenamiento).

from huggingface_hub import snapshot_download

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

Establezca la variable de entorno MODEL_NAME en un id de modelo del Hub o en una ruta a un modelo local, INSTANCE_DIR en la ruta donde acaba de descargar las imágenes de perros y OUTPUT_DIR en el lugar donde desea guardar el modelo. Usará sks como palabra especial para vincular el entrenamiento.

Si te interesa seguir el proceso de entrenamiento, puedes guardar periódicamente las imágenes generadas a medida que avanza el entrenamiento. Añade los siguientes parámetros al comando de entrenamiento:

--validation_prompt="a photo of a sks dog"
--num_validation_images=4
--validation_steps=100

¡Una cosa más antes de lanzar el script! Dependiendo de la GPU que tengas, es posible que necesites habilitar ciertas optimizaciones para entrenar 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

Una vez que el entrenamiento esté completo, ¡puedes usar tu modelo recién entrenado para la inferencia!

¿No puedes esperar a probar tu modelo para inferencia antes de que termine el entrenamiento? 🤭 Asegúrate de tener instalada la última versión de 🤗 Accelerate.

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 es una técnica de entrenamiento para reducir significativamente el número de parámetros entrenables. Como resultado, el entrenamiento es más rápido y es más fácil almacenar los pesos resultantes porque son mucho más pequeños (~100MB). Utilice el script train_dreambooth_lora.py para entrenar con LoRA.

El script de entrenamiento LoRA se describe con más detalle en la guía de entrenamiento LoRA.

Stable Diffusion XL

Stable Diffusion XL (SDXL) es un potente modelo de texto a imagen que genera imágenes de alta resolución y añade un segundo text-encoder a su arquitectura. Utilice el script train_dreambooth_lora_sdxl.py para entrenar un modelo SDXL con LoRA.

El script de entrenamiento de SDXL se describe con más detalle en la guía de entrenamiento de SDXL.

DeepFloyd IF

DeepFloyd IF es un modelo de difusión de píxeles en cascada con tres etapas. La primera etapa genera una imagen base y las segunda y tercera etapas escalan progresivamente la imagen base hasta obtener una imagen de alta resolución de 1024x1024. Utilice los scripts train_dreambooth_lora.py o train_dreambooth.py para entrenar un modelo DeepFloyd IF con LoRA o el modelo completo.

DeepFloyd IF utiliza varianza predicha, pero los scripts de entrenamiento de Diffusers utilizan error predicho, por lo que los modelos de DeepFloyd IF entrenados se cambian a un programa de varianza fija. Los scripts de entrenamiento actualizarán la configuración del programador del modelo completamente entrenado por ti. Sin embargo, al cargar los pesos de LoRA guardados, también debes actualizar la configuración del programador del 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")

El modelo de la etapa 2 requiere imágenes de validación adicionales para aumentar la resolución. Puedes descargar y usar una versión reducida de las imágenes de entrenamiento para esto.

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",
)

Las muestras de código a continuación ofrecen una breve descripción de cómo entrenar un modelo DeepFloyd IF con una combinación de DreamBooth y LoRA. Algunos parámetros importantes a tener en cuenta son:

--resolution=64, se requiere una resolución mucho menor porque DeepFloyd IF es un modelo de difusión de píxeles y, para trabajar con píxeles sin comprimir, las imágenes de entrada deben ser más pequeñas
--pre_compute_text_embeddings, calcula los embeddings de texto de antemano para ahorrar memoria, ya que T5Model puede consumir mucha memoria
--tokenizer_max_length=77, puede usar una longitud de texto predeterminada más larga con T5 como codificador de texto, pero el procedimiento de codificación del modelo predeterminado usa una longitud de texto más corta
--text_encoder_use_attention_mask, para pasar la máscara de atención al codificador de texto

Etapa 1 LoRA DreamBooth

Etapa 2 LoRA DreamBooth

Etapa 1 DreamBooth

Etapa 2 DreamBooth

La etapa 1 de entrenamiento de DeepFloyd IF con LoRA y DreamBooth requiere ~28GB de memoria.

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

Consejos de entrenamiento

Entrenar el modelo DeepFloyd IF puede ser un desafío, pero aquí hay algunos consejos que nos han resultado útiles:

LoRA es suficiente para entrenar el modelo de la etapa 1, ya que la baja resolución del modelo dificulta la representación de detalles más finos de todos modos.
Para objetos comunes o simples, no es necesario ajustar el upscaler. Asegúrate de que el prompt pasado al upscaler esté ajustado para eliminar el nuevo token del prompt de la instancia. Por ejemplo, si tu prompt de la etapa 1 es “a sks dog”, entonces tu prompt de la etapa 2 debe ser “a dog”.
Para detalles más finos, como los rostros, es mejor entrenar por completo el upscaler de la etapa 2 que entrenar el modelo de la etapa 2 con LoRA. También ayuda utilizar tasas de aprendizaje más bajas con tamaños de lote más grandes.
Se deben usar tasas de aprendizaje más bajas para entrenar el modelo de la etapa 2.
El DDPMScheduler funciona mejor que el DPMSolver utilizado en los scripts de entrenamiento.

Próximos pasos

¡Felicidades por entrenar tu modelo de DreamBooth! Para aprender más sobre cómo usar tu nuevo modelo, la siguiente guía puede ser útil:

Aprende a cargar un checkpoint de DreamBooth para la inferencia.
Si entrenaste con LoRA, consulta LoRA para inferencia.

Update on GitHub