Servir Stable Diffusion XL (SDXL) mediante TPUs en GKE con MaxDiffusion

En este tutorial se muestra cómo servir un modelo de generación de imágenes SDXL mediante unidades de procesamiento de tensor (TPUs) en Google Kubernetes Engine (GKE) con MaxDiffusion. En este tutorial, descargarás el modelo de Hugging Face y lo desplegarás en un clúster Autopilot o Standard con un contenedor que ejecute MaxDiffusion.

Esta guía es un buen punto de partida si necesitas el control granular, la personalización, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes gestionado al desplegar y servir tus cargas de trabajo de IA o aprendizaje automático. Si necesitas una plataforma de IA gestionada unificada para crear y ofrecer modelos de aprendizaje automático rápidamente y de forma rentable, te recomendamos que pruebes nuestra solución de despliegue Vertex AI.

Fondo

Si sirves SDXL mediante TPUs en GKE con MaxDiffusion, puedes crear una solución de servicio estable y lista para producción con todas las ventajas de Kubernetes gestionado, como la rentabilidad, la escalabilidad y una mayor disponibilidad. En esta sección se describen las tecnologías clave que se usan en este tutorial.

Stable Diffusion XL (SDXL)

Stable Diffusion XL (SDXL) es un tipo de modelo de difusión latente (LDM) compatible con MaxDiffusion para la inferencia. En el caso de la IA generativa, puedes usar los modelos de difusión latente para generar imágenes de alta calidad a partir de descripciones de texto. Los modelos de difusión latente son útiles para aplicaciones como la búsqueda de imágenes y la generación de títulos de imágenes.

SDXL admite la inferencia de un solo host o de varios hosts con anotaciones de fragmentación. Esto permite que SDXL se entrene y se ejecute en varias máquinas, lo que puede mejorar la eficiencia.

Para obtener más información, consulta el repositorio de modelos generativos de Stability AI y el artículo sobre SDXL.

TPUs

Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google que se utilizan para acelerar los modelos de aprendizaje automático y de IA creados con frameworks como TensorFlow, PyTorch y JAX.

Antes de usar las TPUs en GKE, te recomendamos que completes el siguiente plan de formación:

  1. Consulta la arquitectura del sistema de las TPU de Cloud para obtener información sobre la disponibilidad de las versiones actuales de las TPU.
  2. Consulta información sobre las TPUs en GKE.

En este tutorial se explica cómo servir el modelo SDXL. GKE implementa el modelo en nodos de TPU v5e de un solo host con topologías de TPU configuradas en función de los requisitos del modelo para servir peticiones con baja latencia. En esta guía, el modelo usa un chip de TPU v5e con una topología 1x1.

MaxDiffusion

MaxDiffusion es una colección de implementaciones de referencia, escritas en Python y Jax, de varios modelos de difusión latente que se ejecutan en dispositivos XLA, incluidas las TPUs y las GPUs. MaxDiffusion es un punto de partida para proyectos de difusión, tanto de investigación como de producción.

Para obtener más información, consulta el repositorio MaxDiffusion.