ControlNet

Los modelos ControlNet son adaptadores entrenados sobre otro modelo preentrenado. Permiten un mayor grado de control sobre la generación de imágenes al condicionar el modelo con una imagen de entrada adicional. La imagen de entrada puede ser un borde canny, un mapa de profundidad, una pose humana y muchas más.

Si estás entrenando en una GPU con vRAM limitada, debes intentar habilitar los parámetros gradient_checkpointing, gradient_accumulation_steps 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 de entrenamiento train_controlnet.py para que se familiarice 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 .

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/controlnet
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 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_controlnet.py \
  --mixed_precision="fp16"

Muchos de los parámetros básicos e importantes se describen en la guía de entrenamiento de Texto a imagen, por lo que esta guía se centra solo en los parámetros relevantes para ControlNet:

--max_train_samples: el número de muestras de entrenamiento; este valor puede reducirse para un entrenamiento más rápido, pero si deseas transmitir conjuntos de datos realmente grandes, deberás incluir este parámetro y el parámetro --streaming en tu comando de entrenamiento
--gradient_accumulation_steps: número de pasos de actualización que se acumulan antes de la propagación hacia atrás; esto permite entrenar con un tamaño de lote mayor del que la memoria de la GPU puede manejar normalmente

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_controlnet.py \
  --snr_gamma=5.0

Script de entrenamiento

Al igual que con los parámetros del script, un recorrido general por el script de entrenamiento se proporciona en la guía de entrenamiento Text-to-image. En su lugar, esta guía examina las partes relevantes del script de ControlNet.

El script de entrenamiento tiene una función make_train_dataset para preprocesar el conjunto de datos con transformaciones de imagen y tokenización de las descripciones. Verás que, además de la tokenización de descripciones y las transformaciones de imagen habituales, el script también incluye transformaciones para la imagen de condicionamiento.

Si estás transmitiendo un conjunto de datos en un TPU, el rendimiento puede verse limitado por la biblioteca 🤗 Datasets, que no está optimizada para imágenes. Para garantizar el máximo rendimiento, se recomienda explorar otros formatos de conjuntos de datos como WebDataset, TorchData y TensorFlow Datasets.

conditioning_image_transforms = transforms.Compose(
    [
        transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
        transforms.CenterCrop(args.resolution),
        transforms.ToTensor(),
    ]
)

Dentro de la función main(), encontrará el código para cargar el tokenizer, el text encoder, el scheduler y los modelos. Aquí también se carga el modelo ControlNet, ya sea desde pesos existentes o inicializándolo aleatoriamente desde un UNet:

if args.controlnet_model_name_or_path:
    logger.info("Loading existing controlnet weights")
    controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path)
else:
    logger.info("Initializing controlnet weights from unet")
    controlnet = ControlNetModel.from_unet(unet)

El optimizer se configura para actualizar los parámetros de ControlNet:

params_to_optimize = controlnet.parameters()
optimizer = optimizer_class(
    params_to_optimize,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Por último, en el bucle de entrenamiento, los embeddings del texto de condicionamiento y la imagen se pasan a los bloques inferiores y medios del modelo ControlNet:

encoder_hidden_states = text_encoder(batch["input_ids"])[0]
controlnet_image = batch["conditioning_pixel_values"].to(dtype=weight_dtype)

down_block_res_samples, mid_block_res_sample = controlnet(
    noisy_latents,
    timesteps,
    encoder_hidden_states=encoder_hidden_states,
    controlnet_cond=controlnet_image,
    return_dict=False,
)

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

¡Ahora estás listo para lanzar el script de entrenamiento! 🚀

Esta guía utiliza el conjunto de datos fusing/fill50k, pero recuerda que puedes crear y utilizar tu propio conjunto de datos si lo deseas (consulta la guía Create a dataset for training).

Establezca la variable de entorno MODEL_NAME en un id de modelo en el Hub o en una ruta a un modelo local y OUTPUT_DIR en el lugar donde desea guardar el modelo.

Descarga las siguientes imágenes para condicionar tu entrenamiento con:

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

¡Una cosa más antes de lanzar el script! Dependiendo de la GPU que tengas, es posible que necesites habilitar ciertas optimizaciones para entrenar un ControlNet. La configuración predeterminada de este script requiere ~38GB de vRAM. Si estás entrenando en más de una GPU, añade el parámetro --multi_gpu al comando accelerate launch.

16GB

12GB

8GB

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/save/model"

accelerate launch train_controlnet.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --resolution=512 \
 --learning_rate=1e-5 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --push_to_hub

Una vez que el entrenamiento esté completo, ¡puedes usar tu modelo recién entrenado para la inferencia!

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch

controlnet = ControlNetModel.from_pretrained("path/to/controlnet", dtype=torch.float16)
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
    "path/to/base/model", controlnet=controlnet, dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

control_image = load_image("./conditioning_image_1.png")
prompt = "pale golden rod circle with old lace background"

generator = torch.manual_seed(0)
image = pipeline(prompt, num_inference_steps=20, generator=generator, image=control_image).images[0]
image.save("./output.png")

Stable Diffusion XL

Stable Diffusion XL (SDXL) es un potente modelo de texto a imagen que genera imágenes de alta resolución y añade un segundo codificador de texto a su arquitectura. Utilice el script train_controlnet_sdxl.py para entrenar un adaptador ControlNet para el modelo SDXL.

El script de entrenamiento de SDXL se describe con más detalle en la guía de entrenamiento de SDXL.

Próximos pasos

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

Aprende a usar un ControlNet para la inferencia en una variedad de tareas.

Update on GitHub

Diffusers

ControlNet

Parámetros del script

Ponderación Min-SNR

Script de entrenamiento

Ejecutar el script

Stable Diffusion XL

Próximos pasos