Stable Diffusion XL

Este script es experimental y es fácil sobreajustar y encontrarse con problemas como el olvido catastrófico. Pruebe explorar diferentes hiperparámetros para obtener los mejores resultados en su conjunto de datos.

Stable Diffusion XL (SDXL) es una iteración mayor y más potente del modelo Stable Diffusion, capaz de generar imágenes con una resolución más alta.

El UNet de SDXL es 3 veces más grande y el modelo añade un segundo codificador de texto a la arquitectura. Dependiendo del hardware disponible, esto puede ser muy intensivo en términos computacionales y es posible que no funcione en una GPU de consumo como una Tesla T4. Para ayudar a que este modelo más grande quepa en la memoria y acelerar el entrenamiento, intente habilitar gradient_checkpointing, mixed_precision y gradient_accumulation_steps. Puede reducir aún más el uso de memoria habilitando la atención eficiente en memoria con xFormers y utilizando el optimizador de 8 bits de bitsandbytes.

Esta guía explorará el script de entrenamiento train_text_to_image_sdxl.py para que se familiarice más con él y para que pueda adaptarlo a su 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 .

Luego, navegue a la carpeta de ejemplo que contiene el script de entrenamiento e instale las dependencias requeridas para el script que esté utilizando:

cd examples/text_to_image
pip install -r requirements_sdxl.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.

Parámetros del script

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.

El script de entrenamiento ofrece muchos parámetros para ayudarte a personalizar tu ejecución de entrenamiento. Todos los parámetros y sus descripciones se encuentran en la función parse_args(). Esta función proporciona valores predeterminados para cada parámetro, como el tamaño del lote de entrenamiento y la tasa de aprendizaje, pero también puedes establecer tus propios valores en el comando de entrenamiento si lo deseas.

Por ejemplo, para acelerar el entrenamiento con precisión mixta usando el formato bf16, añade el parámetro --mixed_precision al comando de entrenamiento:

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

La mayoría de los parámetros son idénticos a los de la guía de entrenamiento Texto a imagen, por lo que en esta guía te centrarás en los parámetros relevantes para el entrenamiento de SDXL.

• --pretrained_vae_model_name_or_path: ruta a un VAE preentrenado; se sabe que el VAE de SDXL sufre de inestabilidad numérica, por lo que este parámetro le permite especificar un VAE mejor
• --proportion_empty_prompts: la proporción de prompts de imagen que se reemplazarán con cadenas vacías
• --timestep_bias_strategy: dónde (en pasos anteriores o posteriores) aplicar un sesgo, lo que puede animar al modelo a aprender detalles de baja o alta frecuencia
• --timestep_bias_multiplier: el peso del sesgo a aplicar al timestep
• --timestep_bias_begin: el paso de tiempo en el que se comienza a aplicar el sesgo
• --timestep_bias_end: el paso de tiempo en el que se deja de aplicar el sesgo
• --timestep_bias_portion: la proporción de pasos de tiempo a los que se aplica el sesgo

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

Script de entrenamiento

El script de entrenamiento también es similar a la guía de entrenamiento de Texto a imagen, pero se ha modificado para admitir el entrenamiento de SDXL. Esta guía se centrará en el código que es exclusivo del script de entrenamiento de SDXL.

Comienza creando funciones para tokenizar los prompts y calcular los embeddings de los prompts, y para calcular los embeddings de las imágenes con el VAE. A continuación, creará una función para generar los pesos de los timesteps en función del número de timesteps y la estrategia de sesgo de timestep a aplicar.

Dentro de la función main(), además de cargar un tokenizador, el script carga un segundo tokenizador y un codificador de texto porque la arquitectura de SDXL utiliza dos de cada uno:

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

Los embeddings del prompt y de la imagen se calculan primero y se mantienen en memoria, lo cual normalmente no es un problema para un conjunto de datos más pequeño, pero para conjuntos de datos más grandes puede provocar problemas de memoria. Si este es el caso, debe guardar los embeddings precalculados en el disco de forma separada y cargarlos en memoria durante el proceso de entrenamiento (vea esta PR para más información sobre este tema).

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

Después de calcular los embeddings, se eliminan el codificador de texto, el VAE y el tokenizador para liberar algo de memoria:

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

Por último, el bucle de entrenamiento se encarga del resto. Si elegiste aplicar una estrategia de sesgo de paso de tiempo, verás que los pesos de los pasos de tiempo se calculan y se añaden como ruido:

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)

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

Una vez que hayas realizado todos tus cambios o estés conforme con la configuración predeterminada, ¡estás listo para lanzar el script de entrenamiento! 🚀

Entrenemos con el conjunto de datos Naruto BLIP captions para generar tus propios personajes de Naruto. Establece las variables de entorno MODEL_NAME y DATASET_NAME con el modelo y el conjunto de datos (ya sea desde el Hub o una ruta local). También debes especificar un VAE distinto del VAE de SDXL (ya sea desde el Hub o una ruta local) con VAE_NAME para evitar inestabilidades numéricas.

Para monitorear el progreso del entrenamiento con Weights & Biases, añade el parámetro --report_to=wandb al comando de entrenamiento. También debes añadir --validation_prompt y --validation_epochs al comando de entrenamiento para llevar un registro de los resultados. Esto puede ser muy útil para depurar el modelo y ver resultados intermedios.

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

Después de que hayas terminado el entrenamiento, ¡puedes usar tu nuevo modelo SDXL entrenado para la inferencia!

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óximos pasos

¡Felicidades por entrenar un modelo de SDXL! Para aprender más sobre cómo usar tu nuevo modelo, las siguientes guías pueden ser útiles:

• Lee la guía de Stable Diffusion XL para aprender a usarla para una variedad de tareas diferentes (texto a imagen, imagen a imagen, inpainting), cómo usar su modelo refiner y los diferentes tipos de micro-condicionamientos.
• Consulta las guías de entrenamiento de DreamBooth y LoRA para aprender a entrenar un modelo de SDXL personalizado con solo unas pocas imágenes de ejemplo. ¡Incluso se pueden combinar estas dos técnicas de entrenamiento!