Custom Diffusion

Custom Diffusion es una técnica de entrenamiento para personalizar modelos de generación de imágenes. Al igual que Textual Inversion, DreamBooth y LoRA, Custom Diffusion solo requiere unas pocas (~4-5) imágenes de ejemplo. Esta técnica funciona entrenando únicamente los pesos en las capas de atención cruzada y utiliza una palabra especial para representar el concepto recién aprendido. Custom Diffusion es única porque también puede aprender múltiples conceptos al mismo tiempo.

Si estás entrenando en una GPU con vRAM limitada, deberías intentar habilitar xFormers con --enable_xformers_memory_efficient_attention para un entrenamiento más rápido con requisitos de vRAM más bajos (16GB). Para ahorrar aún más memoria, añade --set_grads_to_none en el argumento de entrenamiento para establecer los gradientes en None en lugar de cero (esta opción puede causar algunos problemas, así que si experimentas alguno, intenta eliminar este parámetro).

Esta guía explorará el script train_custom_diffusion.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:

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 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 contiene todos los parámetros para ayudarte a personalizar tu ejecución de entrenamiento. Estos se encuentran en la función parse_args(). La función viene con valores predeterminados, pero también puedes establecer tus propios valores en el comando de entrenamiento si lo deseas.

Por ejemplo, para cambiar la resolución de la imagen de entrada:

accelerate launch train_custom_diffusion.py \
  --resolution=256

Muchos de los parámetros básicos se describen en la guía de entrenamiento de DreamBooth, por lo que esta guía se centra en los parámetros exclusivos de Custom Diffusion:

--freeze_model: congela los parámetros de clave y valor en la capa de atención cruzada; el valor predeterminado es crossattn_kv, pero puedes configurarlo como crossattn para entrenar todos los parámetros de la capa de atención cruzada
--concepts_list: para aprender varios conceptos, proporciona una ruta a un archivo JSON que contenga los conceptos
--modifier_token: una palabra especial utilizada para representar el concepto aprendido
--initializer_token: una palabra especial utilizada para inicializar los embeddings del modifier_token

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.

Muchos de los parámetros para la pérdida de preservación previa se describen en la guía de entrenamiento de DreamBooth.

Regularización

Custom Diffusion incluye el entrenamiento de las imágenes objetivo con un pequeño conjunto de imágenes reales para evitar el sobreajuste. Como puedes imaginar, esto puede ser fácil de hacer cuando solo entrenas con unas pocas imágenes. Descarga 200 imágenes reales con clip_retrieval. El class_prompt debe ser la misma categoría que las imágenes objetivo. Estas imágenes se almacenan en class_data_dir.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

Para habilitar la regularización, añade los siguientes parámetros:

--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
--real_prior: si se debe usar un pequeño conjunto de imágenes reales para evitar el sobreajuste

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

Script de entrenamiento

Gran parte del código del script de entrenamiento de Custom Diffusion es similar al script de DreamBooth. Esta guía se centra, en cambio, en el código relevante para Custom Diffusion.

El script de entrenamiento de Custom Diffusion tiene dos clases de conjunto de datos:

CustomDiffusionDataset: preprocesa las imágenes, las imágenes de clase y los prompts para el entrenamiento
PromptDataset: prepara los prompts para generar imágenes de clases

A continuación, el modifier_token se añade al tokenizer, se convierte en ids de tokens y se redimensionan las incorporaciones de tokens para tener en cuenta el nuevo modifier_token. Luego, las incorporaciones del modifier_token se inicializan con las incorporaciones del initializer_token. Todos los parámetros del codificador de texto se congelan, excepto las incorporaciones de tokens, ya que es lo que el modelo intenta aprender a asociar con los conceptos.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

Ahora necesitarás añadir los pesos de Custom Diffusion a las capas de atención. Este es un paso muy importante para obtener la forma y el tamaño correctos de los pesos de atención y para establecer el número adecuado de procesadores de atención en cada bloque de UNet.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

El optimizer se inicializa para actualizar los parámetros de la capa de atención cruzada:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

En el bucle de entrenamiento, es importante actualizar solo los embeddings del concepto que se está intentando aprender. Esto implica establecer en cero los gradientes de todos los demás embeddings de tokens:

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

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

En esta guía, descargará y usará estas imágenes de gatos de ejemplo. También puede crear y usar su propio conjunto de datos si lo desea (consulte la guía Crear un conjunto de datos para el entrenamiento).

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 gatos y OUTPUT_DIR en el lugar donde desea guardar el modelo. Utilizará <new1> como palabra especial para vincular los embeddings recién aprendidos. El script crea y guarda puntos de control del modelo y un archivo pytorch_custom_diffusion_weights.bin en su repositorio.

Para monitorear el progreso del entrenamiento con Weights and Biases, agregue el parámetro --report_to=wandb al comando de entrenamiento y especifique un prompt de validación con --validation_prompt. Esto es útil para depurar y guardar resultados intermedios.

Si estás entrenando con rostros humanos, el equipo de Custom Diffusion ha encontrado que los siguientes parámetros funcionan bien:

--learning_rate=5e-6

--max_train_steps puede estar entre 1000 y 2000

--freeze_model=crossattn

use al menos 15-20 imágenes para entrenar

un solo concepto

varios conceptos

Una vez finalizado el entrenamiento, puedes usar tu nuevo modelo de Custom Diffusion para la inferencia.

un solo concepto

varios conceptos

Próximos pasos

¡Felicidades por entrenar un modelo con Custom Diffusion! 🎉 Para saber más:

Lee la publicación de blog Multi-Concept Customization of Text-to-Image Diffusion para conocer más detalles sobre los resultados experimentales del equipo de Custom Diffusion.

Update on GitHub

Diffusers

Custom Diffusion

Parámetros del script

Pérdida de preservación previa

Regularización

Script de entrenamiento

Ejecutar el script

Próximos pasos