Inversión textual

Textual Inversion es una técnica de entrenamiento para personalizar modelos de generación de imágenes con solo unas pocas imágenes de ejemplo de lo que se desea que aprendan. Esta técnica funciona aprendiendo y actualizando los embeddings de texto (los nuevos embeddings se asocian a una palabra especial que debe usarse en el prompt) para que coincidan con las imágenes de ejemplo que se proporcionan.

Para la inferencia con embeddings entrenados, consulte Inferencia de inversión textual.

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 textual_inversion.py para ayudarte a familiarizarte con él y cómo puedes adaptarlo a tu 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/textual_inversion
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

El script de entrenamiento tiene muchos parámetros para ayudarte a adaptar la ejecución de entrenamiento a tus necesidades. Todos los parámetros y sus descripciones están listados en la función parse_args(). Cuando sea aplicable, Diffusers proporciona valores predeterminados para cada parámetro, como el tamaño del lote de entrenamiento y la tasa de aprendizaje, pero puedes cambiar estos valores en el comando de entrenamiento si lo deseas.

Por ejemplo, para aumentar el número de pasos de acumulación de gradientes por encima del valor predeterminado de 1:

accelerate launch textual_inversion.py \
  --gradient_accumulation_steps=4

Entre los demás parámetros básicos e importantes que se deben especificar se incluyen:

--pretrained_model_name_or_path: el nombre del modelo en el Hub o una ruta local al modelo preentrenado
--train_data_dir: ruta a una carpeta que contiene el conjunto de datos de entrenamiento (imágenes de ejemplo)
--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
--num_vectors: el número de vectores para aprender los embeddings; aumentar este parámetro ayuda al modelo a aprender mejor, pero conlleva un aumento en los costos de entrenamiento
--placeholder_token: la palabra especial a la que se vinculan las incorporaciones aprendidas (debe usar la palabra en su indicación para la inferencia)
--initializer_token: una sola palabra que describe de forma aproximada el objeto o estilo sobre el que estás intentando entrenar
--learnable_property: si estás entrenando el modelo para aprender un nuevo “estilo” (por ejemplo, el estilo de pintura de Van Gogh) o “objeto” (por ejemplo, tu perro)

Script de entrenamiento

A diferencia de otros scripts de entrenamiento, TextualInversionDataset tiene una clase de dataset personalizada, TextualInversionDataset, para crear un dataset. Puedes personalizar el tamaño de la imagen, el token de marcador de posición, el método de interpolación, si recortar la imagen, y más. Si necesitas cambiar cómo se crea el dataset, puedes modificar TextualInversionDataset.

A continuación, encontrarás el código de preprocesamiento del conjunto de datos y el bucle de entrenamiento en la función main().

El script comienza cargando el tokenizer, el programador y el modelo:

# Load tokenizer
if args.tokenizer_name:
    tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
elif args.pretrained_model_name_or_path:
    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")

# Load scheduler and models
noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
text_encoder = CLIPTextModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision
)

El token de marcador de posición especial se añade junto al tokenizador y el embedding se reajusta para tener en cuenta el nuevo token.

Luego, el script crea un conjunto de datos a partir de TextualInversionDataset:

train_dataset = TextualInversionDataset(
    data_root=args.train_data_dir,
    tokenizer=tokenizer,
    size=args.resolution,
    placeholder_token=(" ".join(tokenizer.convert_ids_to_tokens(placeholder_token_ids))),
    repeats=args.repeats,
    learnable_property=args.learnable_property,
    center_crop=args.center_crop,
    set="train",
)
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
)

Por último, el bucle de entrenamiento se encarga de todo lo demás, desde predecir el residuo ruidoso hasta actualizar los pesos de incrustación del token especial de marcador de posición.

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

Para esta guía, descargará algunas imágenes de un juguete para gatos 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 = "./cat"
snapshot_download(
    "diffusers/cat_toy_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, y DATA_DIR en la ruta donde acaba de descargar las imágenes de gatos. El script crea y guarda los siguientes archivos en su repositorio:

learned_embeds.bin: los vectores de embedding aprendidos correspondientes a tus imágenes de ejemplo
token_identifier.txt: el token de marcador de posición especial
type_of_concept.txt: el tipo de concepto en el que estás entrenando (ya sea “object” o “style”)

Una ejecución completa de entrenamiento tarda aproximadamente 1 hora en una sola GPU V100.

Una cosa más antes de lanzar el script. 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 <cat-toy> train"
--num_validation_images=4
--validation_steps=100

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export DATA_DIR="./cat"

accelerate launch textual_inversion.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --learnable_property="object" \
  --placeholder_token="<cat-toy>" \
  --initializer_token="toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 \
  --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir="textual_inversion_cat" \
  --push_to_hub

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

from diffusers import StableDiffusionPipeline
import torch

pipeline = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_textual_inversion("sd-concepts-library/cat-toy")
image = pipeline("A <cat-toy> train", num_inference_steps=50).images[0]
image.save("cat-train.png")

Próximos pasos

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

Aprende a cargar embeddings de Textual Inversion y a utilizarlos también como embeddings negativos.

Update on GitHub