Este repositorio proporciona una implementación de DreamBooth usando KerasCV y TensorFlow. La implementación se basa en gran medida en el ejemplo de diffusers
de Hugging Face.
DreamBooth es una forma de enseñar (afinar) rápidamente Difusión estable sobre nuevos conceptos visuales. Para obtener más detalles, consulte este documento.
El código proporcionado en este repositorio es sólo para fines de investigación . Consulte esta sección para obtener más información sobre los posibles casos de uso y las limitaciones.
Al cargar este modelo, acepta la licencia CreativeML Open RAIL-M en https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE.
Si solo está buscando los recursos que acompañan a este repositorio, aquí están los enlaces:
Actualización 15/02/2023 : Gracias a Soumik Rakshit; ahora tenemos mejores utilidades para soportar pesos y sesgos (ver n.° 22).
Instale los requisitos previos: pip install -r requirements.txt
.
Primero debe elegir una clase a la que se le agrega un identificador único. El código base de este repositorio se probó utilizando sks
como identificador único y dog
como clase.
Luego se generan dos tipos de mensajes:
(a) mensaje de instancia : f"una foto de {self.unique_id} {self.class_category}" (b) mensaje de clase : f"una foto de {self.class_category}"
Imágenes de instancia
Obtenga algunas imágenes (3 - 10) que sean representativas del concepto con el que se va a afinar el modelo. Estas imágenes estarían asociadas con instance_prompt
. Estas imágenes se denominan instance_images
del código base. Archive estas imágenes y guárdelas en algún lugar en línea de modo que el archivo se pueda descargar usando la función tf.keras.utils.get_file()
internamente.
Imágenes de clase
DreamBooth utiliza la pérdida de preservación previa para regularizar el entrenamiento. En pocas palabras, la pérdida por preservación previa ayuda al modelo a adaptarse lentamente al nuevo concepto bajo consideración a partir de cualquier conocimiento previo que pueda haber tenido sobre el concepto. Para utilizar la pérdida de preservación previa, necesitamos el mensaje de clase como se muestra arriba. El mensaje de clase se utiliza para generar una cantidad predefinida de imágenes que se utilizan para calcular la pérdida final utilizada para el entrenamiento de DreamBooth.
Según este recurso, entre 200 y 300 imágenes generadas mediante el mensaje de clase funcionan bien en la mayoría de los casos.
Entonces, después de haber decidido instance_prompt
y class_prompt
, use este Colab Notebook para generar algunas imágenes que se usarían para el entrenamiento con la pérdida de preservación previa. Luego archive las imágenes generadas como un archivo único y alójelo en línea de modo que pueda descargarse utilizando la función tf.keras.utils.get_file()
internamente. En el código base, simplemente nos referimos a estas imágenes como class_images
.
Es posible realizar capacitación DreamBooth SIN utilizar una pérdida de preservación previa. Este repositorio siempre lo usa. Para que las personas puedan probar fácilmente este código base, hospedamos las imágenes de instancia y clase aquí.
¡Lanza el entrenamiento! Hay una serie de hiperparámetros con los que puedes jugar. Consulte el script train_dreambooth.py
para saber más sobre ellos. Aquí hay un comando que inicia el entrenamiento con precisión mixta y otros valores predeterminados:
python train_dreambooth.py --mp
También puede ajustar el codificador de texto especificando la opción --train_text_encoder
.
Además, el script admite la integración con pesos y sesgos ( wandb
). Si especifica --log_wandb
,
wandb
utilizando la devolución de llamada WandbMetricsLogger
.wandb
como artefactos para el control de versiones del modelo. Esto se hace usando DreamBoothCheckpointCallback
que se creó usando la devolución de llamada WandbModelCheckpoint
.wandb.Table
en su panel wandb
. Esto se hace utilizando QualitativeValidationCallback
, que también registra las imágenes generadas en un panel multimedia en su panel wandb
al final de la capacitación.Aquí hay un comando que inicia el entrenamiento y registra las métricas de entrenamiento y las imágenes generadas en su espacio de trabajo de Weights & Biases:
python train_dreambooth.py
--log_wandb
--validation_prompts
" a photo of sks dog with a cat "
" a photo of sks dog riding a bicycle "
" a photo of sks dog peeing "
" a photo of sks dog playing cricket "
" a photo of sks dog as an astronaut "
Aquí hay un ejemplo de ejecución wandb
donde puede encontrar las imágenes generadas, así como los puntos de control del modelo.
Hemos probado nuestra implementación en dos métodos diferentes: (a) ajustar el modelo de difusión (UNet) únicamente, (b) ajustar el modelo de difusión junto con el codificador de texto. Los experimentos se realizaron sobre una amplia gama de hiperparámetros para learning rate
y training steps
durante el entrenamiento y para number of steps
y unconditional guidance scale
(ugs) durante la inferencia. Pero aquí sólo se incluyen los resultados más destacados (desde nuestra perspectiva). Si tiene curiosidad acerca de cómo los diferentes hiperparámetros afectan la calidad de la imagen generada, busque el enlace a los informes completos en cada sección.
Tenga en cuenta que nuestros experimentos se basaron en esta publicación de blog de Hugging Face.
A continuación se muestran algunos resultados seleccionados de varios experimentos que realizamos. Nuestros registros experimentales para esta configuración están disponibles aquí. Más imágenes de visualización (generadas con los puntos de control de estos experimentos) están disponibles aquí.
Imágenes | Pasos | UGS | Configuración |
---|---|---|---|
50 | 30 | LR: 1e-6 Pasos de entrenamiento: 800 (Pesos) | |
25 | 15 | LR: 1e-6 Pasos de entrenamiento: 1000 (Pesos) | |
75 | 15 | LR: 3e-6 Pasos de entrenamiento: 1200 (Pesos) |
Imágenes | Pasos | ugs |
---|---|---|
75 | 15 | |
75 | 30 |
con tasa de aprendizaje = 9e-06, pasos máximos del tren = 200 (pesos | informes)
Imágenes | Pasos | ugs |
---|---|---|
150 | 15 | |
75 | 30 |
con tasa de aprendizaje = 9e-06, pasos máximos del tren = 200 (conjuntos de datos | informes)
La biblioteca diffusers
proporciona herramientas de última generación para experimentar con diferentes modelos de difusión, incluida la difusión estable. Incluye diferentes técnicas de optimización que se pueden aprovechar para realizar inferencias eficientes con diffusers
cuando se utilizan grandes puntos de control de difusión estable. Una característica particularmente ventajosa que tienen diffusers
es su soporte para diferentes programadores que se pueden configurar durante el tiempo de ejecución y se pueden integrar en cualquier modelo de Diffusion compatible.
Una vez que haya obtenido los puntos de control ajustados de DreamBooth utilizando este código base, puede exportarlos a un práctico StableDiffusionPipeline
y usarlo directamente desde la biblioteca diffusers
.
Considere este repositorio: chansung/dreambooth-dog. Puede utilizar los puntos de control de este repositorio en StableDiffusionPipeline
después de ejecutar algunos pequeños pasos:
from diffusers import StableDiffusionPipeline
# checkpoint of the converted Stable Diffusion from KerasCV
model_ckpt = "sayakpaul/text-unet-dogs-kerascv_sd_diffusers_pipeline"
pipeline = StableDiffusionPipeline . from_pretrained ( model_ckpt )
pipeline . to ( "cuda" )
unique_id = "sks"
class_label = "dog"
prompt = f"A photo of { unique_id } { class_label } in a bucket"
image = pipeline ( prompt , num_inference_steps = 50 ). images [ 0 ]
Siga esta guía para saber más.
Hemos convertido el punto de control ajustado para las imágenes de perros en StableDiffusionPipeline compatible con difusores y hemos realizado varios experimentos con diferentes configuraciones del programador. Por ejemplo, los siguientes parámetros de DDIMScheduler
se prueban en un conjunto diferente de guidance_scale
y num_inference_steps
.
num_inference_steps_list = [ 25 , 50 , 75 , 100 ]
guidance_scale_list = [ 7.5 , 15 , 30 ]
scheduler_configs = {
"DDIMScheduler" : {
"beta_value" : [
[ 0.000001 , 0.02 ],
[ 0.000005 , 0.02 ],
[ 0.00001 , 0.02 ],
[ 0.00005 , 0.02 ],
[ 0.0001 , 0.02 ],
[ 0.0005 , 0.02 ]
],
"beta_schedule" : [
"linear" ,
"scaled_linear" ,
"squaredcos_cap_v2"
],
"clip_sample" : [ True , False ],
"set_alpha_to_one" : [ True , False ],
"prediction_type" : [
"epsilon" ,
"sample" ,
"v_prediction"
]
}
}
A continuación se muestra la comparación entre diferentes valores de los parámetros beta_schedule
mientras que otros están fijos en sus valores predeterminados. Eche un vistazo al informe original que incluye los resultados de otros programadores como PNDMScheduler
y LMSDiscreteScheduler
.
A menudo se observa que la configuración predeterminada garantiza generar imágenes de mejor calidad. Por ejemplo, los valores predeterminados de guidance_scale
y beta_schedule
están establecidos en 7,5 y linear
. Sin embargo, cuando guidance_scale
se establece en 7.5, scaled_linear
de beta_schedule
parece funcionar mejor. O, cuando beta_schedule
está configurado en linear
, una mayor guidance_scale
parece funcionar mejor.
Realizamos 4.800 experimentos que generaron 38.400 imágenes en total. Esos experimentos se registran en Pesos y sesgos. Si tiene curiosidad, consúltelos aquí, así como el script que se utilizó para ejecutar los experimentos.
Además de los consejos y trucos compartidos en esta publicación de blog, seguimos estas cosas mientras preparábamos las instancias para realizar la capacitación de DreamBooth en rostros humanos:
Gracias a Abhishek Thakur por compartir estos consejos.