stable diffusion keras ft
1.0.0
Este repositorio proporciona código para ajustar la difusión estable en Keras. Está adaptado de este guión abrazando la cara. El modelo previamente capacitado utilizado para el ajuste fino proviene de Kerascv. Para saber sobre el modelo original, consulte esta documentación.
El código proporcionado en este repositorio es solo para fines de investigación . Consulte esta sección para saber más sobre los posibles casos de uso y 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 adjuntos de este repositorio, aquí están los enlaces:
Tabla de contenido :
Este repositorio tiene un repositorio hermano (servidor de keras-SD) que cubre varios patrones de implementación para la difusión estable.
ACTUALIZACIÓN 13 de enero de 2023 : Este proyecto aseguró el segundo lugar en la primera competencia de premios comunitarios de Keras organizada por Google.
Siguiendo el script original de Hugging Face, este repositorio también usa el conjunto de datos Pokemon. Pero se regeneró para funcionar mejor con tf.data . La versión regenerada del conjunto de datos se aloja aquí. Consulte ese enlace para obtener más detalles.
El código de ajuste fino se proporciona en finetune.py . Antes de ejecutar la capacitación, asegúrese de tener las dependencias (consulte requirements.txt ) instalados.
Puede iniciar capacitación con los argumentos predeterminados ejecutando python finetune.py . Ejecute python finetune.py -h para conocer los argumentos de línea de comandos compatibles. Puede habilitar el entrenamiento de precisión mixta pasando la bandera --mp .
Cuando lance el entrenamiento, se generará un punto de control del modelo de difusión en cuanto a la época solo si la pérdida de corriente es menor que la anterior.
Para evitar OOM y un entrenamiento más rápido, se recomienda usar una GPU V100 al menos. Usamos un A100.
Algunos detalles importantes a tener en cuenta :
Detalles de capacitación :
Afinamos el modelo en dos resoluciones diferentes: 256x256 y 512x512. Solo variamos el tamaño del lote y el número de épocas para ajustar con estas dos resoluciones diferentes. Como no usamos la acumulación de gradiente, usamos este fragmento de código para derivar el número de épocas.
python finetune.py --batch_size 4 --num_epochs 577python finetune.py --img_height 512 --img_width 512 --batch_size 1 --num_epochs 72 --mpPara la resolución de 256x256, reducimos intencionalmente el número de épocas para ahorrar tiempo de cómputo.
Pesos ajustados :
Puede encontrar los pesos del modelo de difusión ajustado aquí.
El conjunto de datos Pokemon predeterminado utilizado en este repositorio viene con la siguiente estructura:
pokemon_dataset/
data.csv
image_24.png
image_3.png
image_550.png
image_700.png
... data.csv parece así:
Mientras su conjunto de datos personalizado siga esta estructura, no necesita cambiar nada en la base de código actual, excepto el dataset_archive .
En caso de que su conjunto de datos tenga múltiples subtítulos por imagen, puede seleccionar aleatoriamente uno del grupo de subtítulos por imagen durante el entrenamiento.
Según el conjunto de datos, es posible que deba ajustar los hiperparámetros.
import keras_cv
import matplotlib . pyplot as plt
from tensorflow import keras
IMG_HEIGHT = IMG_WIDTH = 512
def plot_images ( images , title ):
plt . figure ( figsize = ( 20 , 20 ))
for i in range ( len ( images )):
ax = plt . subplot ( 1 , len ( images ), i + 1 )
plt . title ( title )
plt . imshow ( images [ i ])
plt . axis ( "off" )
# We just have to load the fine-tuned weights into the diffusion model.
weights_path = keras . utils . get_file (
origin = "https://huggingface.co/sayakpaul/kerascv_sd_pokemon_finetuned/resolve/main/ckpt_epochs_72_res_512_mp_True.h5"
)
pokemon_model = keras_cv . models . StableDiffusion (
img_height = IMG_HEIGHT , img_width = IMG_WIDTH
)
pokemon_model . diffusion_model . load_weights ( weights_path )
# Generate images.
generated_images = pokemon_model . text_to_image ( "Yoda" , batch_size = 3 )
plot_images ( generated_images , "Fine-tuned on the Pokemon dataset" ) Puede traer sus weights_path (debe ser compatible con la diffusion_model ) y reutilizar el fragmento de código.
Echa un vistazo a este cuaderno de Colab para jugar con el código de inferencia.
Inicialmente, afinamos el modelo en una resolución de 256x256. Aquí hay algunos resultados junto con las comparaciones con los resultados del modelo original.
| Imágenes | Indicaciones |
|---|---|
| Yoda | |
| gato robótico con alas | |
| Hello Kitty |
Podemos ver que el modelo ajustado tiene salidas más estables que el modelo original. Aunque los resultados pueden mejorarse estéticamente mucho más, los efectos de ajuste fino son visibles. Además, seguimos los mismos hiperparámetros del guión de Hugging Face para la resolución 256x256 (aparte de la cantidad de épocas y el tamaño de lotes). Con mejores hiperparámetros, los resultados probablemente mejorarán.
Para la resolución 512x512, observamos algo similar. Entonces, experimentamos con el parámetro unconditional_guidance_scale y notamos que cuando está configurado en 40 (mientras mantiene los otros argumentos fijos), los resultados salieron mejor.
| Imágenes | Indicaciones |
|---|---|
| Yoda | |
| gato robótico con alas | |
| Hello Kitty |
Nota : El ajuste fino en el 512x512 todavía está en progreso a partir de este escrito. Pero lleva mucho tiempo completar una sola época sin la presencia de entrenamiento distribuido y acumulación de gradiente. Los resultados anteriores son del punto de control derivado después de la 60ª época.
Con una receta similar (pero entrenada para más pasos de optimización), Lambda Labs demuestra resultados sorprendentes.
