Kandinsky 2.2

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.

Kandinsky 2.2 es un modelo de texto a imagen multilingüe capaz de generar imágenes más fotorrealistas. El modelo incluye un modelo de prior de imagen para crear embeddings de imagen a partir de prompts de texto, y un modelo decodificador que genera imágenes basándose en los embeddings del modelo de prior. Por eso, encontrarás dos scripts separados en Diffusers para Kandinsky 2.2, uno para entrenar el modelo de prior y otro para entrenar el modelo decodificador. Puedes entrenar ambos modelos por separado, pero para obtener los mejores resultados, debes entrenar tanto el modelo de prior como el modelo decodificador.

Dependiendo de tu GPU, es posible que necesites habilitar gradient_checkpointing (⚠️ no compatible con el modelo anterior), mixed_precision y gradient_accumulation_steps para ayudar a que el modelo quepa en la memoria y acelerar el entrenamiento. Puedes reducir aún más el uso de memoria habilitando la atención eficiente en memoria con xFormers (la versión v0.0.16 falla en el entrenamiento en algunas GPUs, por lo que es posible que necesites instalar una versión de desarrollo en su lugar).

Esta guía explora los scripts train_text_to_image_prior.py y train_text_to_image_decoder.py para ayudarte a familiarizarte con ellos y con la forma de adaptarlos a tu caso de uso.

Antes de ejecutar los scripts, 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/kandinsky2_2/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.

Las siguientes secciones resaltan las partes de los scripts de entrenamiento que son importantes para entender cómo modificarlos, pero no abordan todos los aspectos de los scripts en detalle. Si te interesa aprender más, no dudes en leer los scripts y hacernos saber si tienes alguna pregunta o inquietud.

Parámetros del script

Los scripts de entrenamiento ofrecen 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(). Los scripts de entrenamiento proporcionan 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_prior.py \
  --mixed_precision="fp16"

La mayoría de los parámetros son idénticos a los de la guía de entrenamiento Text-to-image, así que vamos directamente al recorrido por los scripts de entrenamiento de Kandinsky.

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_prior.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 los modelos prior y decoder. Esta guía se centra en el código que es único para los scripts de entrenamiento de Kandinsky 2.2.

modelo anterior

modelo decodificador

La función main() contiene el código para preparar el conjunto de datos y entrenar el modelo.

Una de las principales diferencias que notará de inmediato es que el script de entrenamiento también carga un CLIPImageProcessor - además de un scheduler y un tokenizer - para preprocesar las imágenes y un modelo CLIPVisionModelWithProjection para codificar las imágenes:

noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", dtype=weight_dtype
    ).eval()

Kandinsky utiliza un PriorTransformer para generar los embeddings de la imagen, por lo que querrás configurar el optimizador para aprender los parámetros del modo prior.

prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

A continuación, las descripciones de entrada se tokenizan y las imágenes se preprocesan mediante el CLIPImageProcessor:

def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples

Por último, el bucle de entrenamiento convierte las imágenes de entrada en latentes, añade ruido a los embeddings de las imágenes y realiza una predicción:

model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding

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! 🚀

Entrenarás con el conjunto de datos Naruto BLIP captions para generar tus propios personajes de Naruto, pero también puedes crear y entrenar con tu propio conjunto de datos siguiendo la guía Create a dataset for training. Establece la variable de entorno DATASET_NAME con el nombre del conjunto de datos en el Hub o, si estás entrenando con tus propios archivos, establece la variable de entorno TRAIN_DIR con la ruta a tu conjunto de datos.

Si estás entrenando en más de una GPU, añade el parámetro --multi_gpu al comando accelerate launch.

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 al comando de entrenamiento para llevar un registro de los resultados. Esto puede ser muy útil para depurar el modelo y ver resultados intermedios.

modelo anterior

modelo decodificador

Una vez finalizado el entrenamiento, ¡puedes usar tu modelo recién entrenado para inferencia!

modelo anterior

modelo decodificador

Próximos pasos

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

Lee la guía de Kandinsky para aprender a usarla para una variedad de tareas diferentes (texto a imagen, imagen a imagen, inpainting, interpolación) y cómo puede combinarse con un ControlNet.
Consulta las guías de entrenamiento de DreamBooth y LoRA para aprender a entrenar un modelo de Kandinsky personalizado con solo unas pocas imágenes de ejemplo. ¡Incluso se pueden combinar estas dos técnicas de entrenamiento!

Update on GitHub