Texto a imagen

El script de texto a imagen 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.

Los modelos de texto a imagen, como Stable Diffusion, están condicionados para generar imágenes a partir de un prompt de texto.

Entrenar un modelo puede ser exigente para tu hardware, pero si activas gradient_checkpointing y mixed_precision, es posible entrenar un modelo en una sola GPU de 24 GB. Si entrenas con tamaños de lote más grandes o quieres entrenar más rápido, es mejor usar GPUs con más de 30 GB de memoria. Puedes reducir tu huella de memoria activando la atención eficiente en memoria con xFormers.

Esta guía explorará el script de entrenamiento train_text_to_image.py para que se familiarice 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.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 fp16, añade el parámetro --mixed_precision al comando de entrenamiento:

accelerate launch train_text_to_image.py \
  --mixed_precision="fp16"

Algunos parámetros básicos e importantes incluyen:

• --pretrained_model_name_or_path: el nombre del modelo en el Hub o una ruta local al modelo preentrenado
• --dataset_name: el nombre del conjunto de datos en el Hub o una ruta local al conjunto de datos en el que se va a entrenar
• --image_column: el nombre de la columna de imágenes en el conjunto de datos sobre el que se entrena
• --caption_column: el nombre de la columna de texto en el conjunto de datos sobre el que se va a entrenar
• --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_text_to_image.py \
  --snr_gamma=5.0

Puede comparar las superficies de pérdida para diferentes valores de snr_gamma en este informe de Weights and Biases. Para conjuntos de datos más pequeños, los efectos de Min-SNR pueden no ser tan evidentes en comparación con conjuntos de datos más grandes.

Script de entrenamiento

El código de preprocesamiento del conjunto de datos y el bucle de entrenamiento se encuentran en la función main(). Si necesitas adaptar el script de entrenamiento, este es el lugar donde debes realizar los cambios.

El script train_text_to_image comienza cargando un scheduler y un tokenizer. Puedes optar por usar un scheduler diferente aquí si lo deseas:

noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
tokenizer = CLIPTokenizer.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
)

Luego, el script carga el modelo UNet:

load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet")
model.register_to_config(**load_model.config)

model.load_state_dict(load_model.state_dict())

A continuación, las columnas de texto e imagen del conjunto de datos deben preprocesarse. La función tokenize_captions se encarga de tokenizar las entradas, y la función train_transforms especifica el tipo de transformaciones que se aplican a la imagen. Ambas funciones se agrupan en preprocess_train:

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["input_ids"] = tokenize_captions(examples)
    return examples

Por último, el bucle de entrenamiento se encarga de todo lo demás. Codifica las imágenes en el espacio latente, añade ruido a los latentes, calcula los embeddings de texto sobre los que condicionar, actualiza los parámetros del modelo y guarda y publica el modelo en el Hub. Si quieres aprender 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 denoising.

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). Si entrenas en más de una GPU, añade el parámetro --multi_gpu al comando accelerate launch.

Para entrenar con un conjunto de datos local, establezca las variables de entorno TRAIN_DIR y OUTPUT_DIR en la ruta del conjunto de datos y en el lugar donde se guardará el modelo.

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export dataset_name="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$dataset_name \
  --use_ema \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --enable_xformers_memory_efficient_attention \
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-naruto-model" \
  --push_to_hub

Una vez completado el entrenamiento, puedes usar tu modelo recién entrenado para inferencia:

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("path/to/saved_model", dtype=torch.float16, use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

image = pipeline(prompt="yoda").images[0]
image.save("yoda-naruto.png")

Próximos pasos

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

• Aprende a cargar los pesos de LoRA para la inferencia si entrenaste tu modelo con LoRA.
• Obtenga más información sobre cómo ciertos parámetros, como la escala de orientación, o técnicas como el ponderado de prompts pueden ayudarle a controlar la inferencia en la guía de la tarea Texto a imagen.