Implementación de Stable Diffusion en Vertex AI

Stable Diffusion es un modelo de generación de imágenes. Se convirtió en código abierto en 2022 y ha ganado popularidad gracias a su capacidad para generar imágenes de alta calidad a partir de descripciones de texto. Al igual que otros modelos de generación de imágenes, como Dall-E, Stable Diffusion utiliza técnicas de aprendizaje automático para generar imágenes a partir de una entrada determinada.

Moët Hennessy, la división de vinos y licores del conglomerado de lujo LVMH, gestiona una cartera de más de 26 marcas emblemáticas, como Moët & Chandon, Hennessy y Veuve Clicquot. Moët Hennessy ha colaborado con Artefact para investigar los posibles usos de la tecnología de vanguardia en la generación de contenido de marketing. Con un enfoque en la privacidad y la seguridad, el equipo decidió explorar la implementación de Stable Diffusion en Google Cloud (GCP) para permitir a Moët Hennessy ajustar y ejecutar el modelo dentro de su propia infraestructura, proporcionando una experiencia fluida desde el ajuste del modelo hasta la exposición de la API.

Antes de empezar, es importante señalar que este artículo da por hecho que tienes conocimientos previos sobre Google Cloud (GCP) y, en concreto, sobre Vertex AI. Esto incluye conceptos como el registro de modelos y los puntos de conexión de Vertex. Además, necesitas tener experiencia previa con Docker para seguir algunos de los pasos. Si no estás familiarizado con estos conceptos, te recomendamos que te familiarices con ellos antes de continuar.
Además, para descargar los pesos de Stable Diffusion necesitas tener una cuenta de Hugging Face; si aún no tienes una, puedes crearla fácilmente enel sitio web de Hugging Face.
Dicho esto, ¡comencemos!

Te recomiendo que clones elrepositorio de GitHubque he preparado para poder seguir los pasos del artículo.

Siempre es importante crear un entorno virtual para instalar los paquetes. Por mi parte, voy a utilizar Anaconda e instalar el archivo requirements.txt con todas las dependencias:

Ya estás listo para descargar los pesos de Stable Diffusion; en este artículo utilizaremosStable Diffusion 1.5. Debes aceptar la licencia en la página del modelo; de lo contrario, se producirán errores al descargar los pesos del modelo.
Ve a tu cuenta → configuración → token de acceso → nuevo token (acceso de lectura)

Puedes añadir el token como variable de entorno. Personalmente, recomiendo utilizar un archivo .env y cargar la variable de entorno mediante labiblioteca python-dotenv.
Debes ir a `src/stable_diffusion` y ejecutar:

Esto descargará los pesos que hay dentro

Torchservees un marco de trabajo para el servicio de modelos de PyTorch. Permite implementar modelos de PyTorch en un entorno de producción y ofrece funciones como el control de versiones de los modelos y el servicio de múltiples modelos. Está diseñado para ser fácil de usar y te permite centrarte en crear e implementar tus modelos, en lugar de preocuparte por la infraestructura.

Un controlador personalizado es una clase de Python que define cómo preprocesar data de entrada, cómo ejecutar el modelo y cómo posprocesar los resultados. Para crear un controlador personalizado para tu modelo, tendrás que crear una clase de Python que siga el formato de TorchServe.
En el repositorio de TorchServe ya se incluye un controlador personalizado para Stable Diffusion. Sin embargo, el punto final de Vertex espera un formato específico para las solicitudes, por lo que debemos adaptar el método preprocess() del controlador para que se ajuste al formato de Vertex. Puedes utilizar la versión modificada del controlador denominada `stable_diffusion_handler.py` que se incluye en el repositorio de GitHub de este artículo.

Una vez que hayas creado tu controlador personalizado, deberás empaquetarlo junto con las dependencias y el propio modelo en un archivo .mar utilizando la herramienta model-archiver. La herramienta model-archiver es una herramienta de línea de comandos que te permite empaquetar el modelo, el controlador y las dependencias en un único archivo.

Esto creará un archivo .mar llamado output.mar que contiene tu modelo, el controlador y las dependencias.

Es posible que tengas que modificar la ruta en función de la ubicación de los archivos en tu ordenador:

Una vez que hayas creado tu archivo .mar, puedes iniciar el servidor TorchServe mediante el comando torchserve. Para ello, tendrás que ejecutar el siguiente comando:

El archivo config.properties te permite especificar las configuraciones de tu servidor TorchServe, como el puerto para la inferencia, las comprobaciones de estado, el número de trabajadores, etc.
ADVERTENCIA: Todos los scripts deben ejecutarse en el directorio donde se encuentra el archivo para evitar errores de ruta.

Es importante probar el código localmente antes de empezar a implementar nuestra aplicación en Docker; ya he preparado un script de Bash que creará el archivo .mar e iniciará el servidor TorchServe.

Puedes ejecutarlo con:

Debes esperar unos minutos a que el servidor se inicialice y el worker cargue el modelo; podrás empezar a ejecutar la inferencia cuando veas el siguiente registro:

A continuación, puedes utilizar el siguiente código para realizar solicitudes de inferencia a tu modelo:

Puedes comprobar si el servidor ha recibido tu solicitud consultando los registros del servidor:

Stable Diffusion necesita una GPU para funcionar correctamente, por lo que, de momento, no obtendrás ningún resultado. Puedes detener el servidor torchserve con el comando `torchserve –stop`.

Tu servidor TorchServe funciona de forma local; para implementar Stable Diffusion en Vertex AI, tendrás que convertirlo a Docker. Esto implica crear una imagen de Docker que contenga el modelo, el controlador personalizado y todas las dependencias necesarias. Se trata, sencillamente, de todos los pasos que hemos seguido anteriormente dentro de un archivo Dockerfile.
Por suerte, ya está preparado y listo para usar aquí:Dockerfile.

Es importante ejecutar el contenedor localmente para comprobar que funciona correctamente. Yo lo voy a compilar localmente, pero tú puedes compilarlo con cloud y descargarlo en tu equipo.

Compila la imagen localmente (necesitas tener el demonio de Docker en ejecución):

La creación de la imagen tardará entre 20 y 30 minutos; la fase de creación es larga porque los pesos del modelo deben copiarse dentro de la imagen antes de que el archivador de modelos la empaquete.

Puedes ejecutar un contenedor de Docker y escuchar en el puerto 7080 con:

Para comprobar que todo funciona correctamente, puedes esperar a que el worker haya cargado el modelo y ejecutar el mismo código de inferencia que antes, ya que estamos utilizando el mismo puerto 7080.

Ahora que el Dockerfile está listo y funciona, tenemos que:

Esto es precisamente lo que va a hacer el script de Bash«deploy.sh»; puedes ejecutarlo con:

La instalación tarda aproximadamente una hora con una buena conexión a Internet:

Una vez que el modelo se haya vinculado correctamente al punto final, puedes consultar el punto final con el siguiente código:

PROJECT_NAME = ""
REGION = ""
ENDPOINT_ID = ""
aip.init(project=PROJECT_NAME, location=REGION)
endpoint = aip.Endpoint(endpoint_name=ENDPOINT_ID)
text_input = """Una botella de coñac añejo y exclusivo
descansa sobre una superficie reflectante, frente a una animada barra,
hiperdetallada, 4K, bokeh"""

def query_endpoint(endpoint, text_input):
payload =
response = endpoint.predict(instances=[payload])
return response

Una botella de coñac añejo y exclusivo descansa sobre una superficie reflectante, frente a una animada barra, con gran nivel de detalle, en 4K y con efecto bokeh — Stable Diffusion 1.5

La inferencia debería tardar entre 10 y 15 segundos en una GPU T4. Puedes aumentar la velocidad eligiendo una GPU más potente; para ello, modifica la variable ACCELERATOR_TYPE en el archivo deploy.sh.

Quizás quieras llevar un registro del historial de las imágenes, o bien estés encontrando errores debido al límite de tamaño de 1,5 MB para la respuesta del punto final. En ese caso, te recomiendo utilizar el método de posprocesamiento del controlador para guardar la imagen en GCS y devolver únicamente la ruta de GCS de la imagen.

Por suerte, ya he preparado un controlador que hace precisamente eso por ti; basta con cambiar de stable_diffusion_handler.py astable_diffusion_handler_gcs.pypara que funcione.

ADVERTENCIA: Antes de ejecutar la implementación con el nuevo controlador, debes:

En este artículo se ofrece una guía completa para implementar el modelo Stable Diffusion en Google Cloud mediante Vertex AI.

La guía abordaba pasos fundamentales como:

Es importante recordar que, aunque el punto de conexión de Vertex puede ser una solución eficaz, no permite reducir el número de instancias hasta 0, lo que podría suponer un aumento de los costes, ya que la GPU sigue en uso.

Además, una vez completada la implementación del modelo Stable Diffusion, actualmente estamos explorando otras posibilidades, como el ajuste fino del modelo para productos específicos de Moët Hennessy, con el fin de mejorar aún más sus capacidades.