Crea un avatar personalizado con IA generativa utilizando Amazon SageMaker

Crea un avatar personalizado con IA en Amazon SageMaker

La IA generativa se ha convertido en una herramienta común para mejorar y acelerar el proceso creativo en diversas industrias, incluyendo el entretenimiento, la publicidad y el diseño gráfico. Permite experiencias más personalizadas para las audiencias y mejora la calidad general de los productos finales.

Un beneficio significativo de la IA generativa es la creación de experiencias únicas y personalizadas para los usuarios. Por ejemplo, los servicios de streaming utilizan la IA generativa para generar títulos de películas y visuales personalizados con el fin de aumentar el compromiso de los espectadores y crear visuales para títulos basados en el historial de visualización y las preferencias de un usuario. El sistema genera entonces miles de variaciones de la ilustración de un título y las prueba para determinar qué versión atrae más la atención del usuario. En algunos casos, las ilustraciones personalizadas para series de televisión aumentaron significativamente las tasas de clics y visualizaciones en comparación con programas sin ilustraciones personalizadas.

En esta publicación, demostramos cómo puedes utilizar modelos de IA generativa como Stable Diffusion para construir una solución de avatar personalizado en Amazon SageMaker y ahorrar costes de inferencia utilizando puntos finales de múltiples modelos (MMEs) al mismo tiempo. La solución demuestra cómo, al subir 10-12 imágenes de ti mismo, puedes ajustar finamente un modelo personalizado que luego puede generar avatares basados en cualquier indicación de texto, como se muestra en las siguientes capturas de pantalla. Aunque este ejemplo genera avatares personalizados, puedes aplicar la técnica a cualquier generación de arte creativo al ajustar finamente objetos o estilos específicos.

Descripción general de la solución

El siguiente diagrama de arquitectura describe la solución de extremo a extremo para nuestro generador de avatares.

El alcance de esta publicación y el código de ejemplo que proporcionamos en GitHub se centran únicamente en el entrenamiento del modelo y la orquestación de la inferencia (la sección verde en el diagrama anterior). Puedes consultar la arquitectura completa de la solución y construir sobre el ejemplo que proporcionamos.

El entrenamiento del modelo y la inferencia se pueden dividir en cuatro pasos:

  1. Subir imágenes a Amazon Simple Storage Service (Amazon S3). En este paso, te pedimos que proporciones un mínimo de 10 imágenes de alta resolución de ti mismo. Cuantas más imágenes, mejor será el resultado, pero tomará más tiempo entrenar.
  2. Ajustar finamente un modelo base de Stable Diffusion 2.1 utilizando la inferencia asíncrona de SageMaker. Explicamos la justificación de utilizar un punto final de inferencia para el entrenamiento más adelante en esta publicación. El proceso de ajuste fino comienza con la preparación de las imágenes, incluyendo el recorte de rostros, la variación de fondo y el redimensionamiento para el modelo. Luego utilizamos Low-Rank Adaptation (LoRA), una técnica de ajuste fino eficiente en términos de parámetros para modelos de lenguaje grandes (LLMs), para ajustar finamente el modelo. Finalmente, en el postprocesamiento, empaquetamos los pesos de LoRA ajustados finamente con el script de inferencia y los archivos de configuración (tar.gz) y los subimos a una ubicación del bucket de S3 para los MMEs de SageMaker.
  3. Alojar los modelos ajustados finamente utilizando los MMEs de SageMaker con GPU. SageMaker cargará y almacenará en caché dinámicamente el modelo desde la ubicación de Amazon S3 en función del tráfico de inferencia de cada modelo.
  4. Utilizar el modelo ajustado finamente para la inferencia. Después de recibir la notificación del servicio de notificación simple de Amazon (Amazon SNS) que indica que el ajuste fino se ha completado, puedes utilizar inmediatamente ese modelo suministrando un parámetro target_model al invocar el MME para crear tu avatar.

Explicamos cada paso con más detalle en las siguientes secciones y presentamos algunos fragmentos de código de ejemplo.

Preparar las imágenes

Para obtener los mejores resultados al ajustar finamente Stable Diffusion para generar imágenes de ti mismo, generalmente necesitas proporcionar una gran cantidad y variedad de fotos de ti mismo desde diferentes ángulos, con diferentes expresiones y en diferentes fondos. Sin embargo, con nuestra implementación, ahora puedes obtener un resultado de alta calidad con tan solo 10 imágenes de entrada. También hemos añadido un preprocesamiento automatizado para extraer tu rostro de cada foto. Todo lo que necesitas es capturar la esencia de cómo te ves claramente desde múltiples perspectivas. Incluye una foto de frente, una toma de perfil desde cada lado y fotos desde ángulos intermedios. También debes incluir fotos con diferentes expresiones faciales como sonreír, fruncir el ceño y una expresión neutral. Tener una mezcla de expresiones permitirá que el modelo reproduzca mejor tus rasgos faciales únicos. Las imágenes de entrada determinan la calidad del avatar que puedes generar. Para asegurarte de que esto se haga correctamente, recomendamos una experiencia de interfaz de usuario frontal intuitiva para guiar al usuario a través del proceso de captura y carga de imágenes.

Estas son imágenes de ejemplo de selfies tomadas desde diferentes ángulos y con distintas expresiones faciales.

Ajustar un modelo de difusión estable

Una vez que las imágenes se cargan en Amazon S3, podemos invocar el endpoint de inferencia asíncrona de SageMaker para iniciar nuestro proceso de entrenamiento. Los endpoints asíncronos están diseñados para casos de uso de inferencia con cargas útiles grandes (hasta 1 GB) y tiempos de procesamiento prolongados (hasta 1 hora). También proporciona un mecanismo de encolado incorporado para encolar solicitudes y un mecanismo de notificación de finalización de tareas a través de Amazon SNS, además de otras características nativas de alojamiento de SageMaker, como el escalado automático.

Aunque el ajuste fino no es un caso de uso de inferencia, elegimos utilizarlo aquí en lugar de los trabajos de entrenamiento de SageMaker debido a sus mecanismos de encolado y notificación incorporados y su escalado automático administrado, incluida la capacidad de reducir el número de instancias a 0 cuando el servicio no está en uso. Esto nos permite escalar fácilmente el servicio de ajuste fino a un gran número de usuarios concurrentes y elimina la necesidad de implementar y gestionar los componentes adicionales. Sin embargo, tiene la limitación de una carga útil de 1 GB y un tiempo máximo de procesamiento de 1 hora. En nuestras pruebas, encontramos que 20 minutos es un tiempo suficiente para obtener resultados razonablemente buenos con aproximadamente 10 imágenes de entrada en una instancia ml.g5.2xlarge. Sin embargo, para trabajos de ajuste fino a mayor escala, se recomendaría utilizar el entrenamiento de SageMaker.

Para alojar el endpoint asíncrono, debemos completar varios pasos. El primero es definir nuestro servidor de modelos. Para esta publicación, utilizamos el contenedor de inferencia de modelos grandes (LMI). LMI está impulsado por DJL Serving, que es una solución de servicio de modelos de alto rendimiento y agnóstica del lenguaje de programación. Elegimos esta opción porque el contenedor de inferencia administrado de SageMaker ya tiene muchas de las bibliotecas de entrenamiento que necesitamos, como Hugging Face Diffusers y Accelerate. Esto reduce en gran medida la cantidad de trabajo necesario para personalizar el contenedor para nuestro trabajo de ajuste fino.

El siguiente fragmento de código muestra la versión del contenedor LMI que utilizamos en nuestro ejemplo:

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0.21.0-deepspeed0.8.3-cu117"
)
print(f"La imagen que se utilizará es ----> {inference_image_uri}")

Además de eso, necesitamos tener un archivo serving.properties que configure las propiedades de servicio, incluido el motor de inferencia a utilizar, la ubicación del artefacto del modelo y el agrupamiento dinámico. Por último, debemos tener un archivo model.py que cargue el modelo en el motor de inferencia y prepare la entrada y salida de datos del modelo. En nuestro ejemplo, utilizamos el archivo model.py para iniciar el trabajo de ajuste fino, que explicamos con más detalle en una sección posterior. Tanto los archivos serving.properties como model.py se proporcionan en la carpeta training_service.

El siguiente paso después de definir nuestro servidor de modelos es crear una configuración de endpoint que defina cómo se servirá nuestra inferencia asíncrona. Para nuestro ejemplo, solo estamos definiendo el límite máximo de invocaciones concurrentes y la ubicación de salida en S3. Con la instancia ml.g5.2xlarge, hemos descubierto que podemos ajustar fino hasta dos modelos simultáneamente sin encontrar una excepción de falta de memoria (OOM), por lo que establecemos max_concurrent_invocations_per_instance en 2. Es posible que este número deba ajustarse si usamos un conjunto diferente de parámetros de ajuste o un tipo de instancia más pequeño. Recomendamos establecer esto en 1 inicialmente y monitorear la utilización de memoria de la GPU en Amazon CloudWatch.

# crear configuración de endpoint asíncrono
async_config = AsyncInferenceConfig(
    output_path=f"s3://{bucket}/{s3_prefix}/async_inference/output" , # Donde se almacenarán nuestros resultados
    max_concurrent_invocations_per_instance=2,
    notification_config={
      "SuccessTopic": "...",
      "ErrorTopic": "...",
    }, # Configuración de notificación
)

Finalmente, creamos un modelo de SageMaker que empaca la información del contenedor, los archivos del modelo y el rol de IAM de AWS en un solo objeto. El modelo se implementa utilizando la configuración de endpoint que definimos anteriormente:

modelo = Modelo(
    imagen_uri=imagen_uri,
    datos_modelo=datos_modelo,
    rol=papel,
    entorno=entorno
)

modelo.desplegar(
    cantidad_instancias_iniciales=1,
    tipo_instancia=tipo_instancia,
    nombre_endpoint=nombre_endpoint,
    configuracion_inferencia_asincrona=configuracion_inferencia_asincrona
)

predictor = sagemaker.Predictor(
    nombre_endpoint=nombre_endpoint,
    sesion_sagemaker=sesion_sagemaker
)

Cuando el endpoint esté listo, utilizamos el siguiente código de muestra para invocar el endpoint asíncrono y comenzar el proceso de ajuste fino:

sm_runtime = boto3.client("sagemaker-runtime")

ubicacion_entrada_s3 = sess.upload_data("data/jw.tar.gz", bucket, s3_prefix)

respuesta = sm_runtime.invoke_endpoint_async(
    NombreEndpoint=sd_tuning.nombre_endpoint,
    UbicacionEntrada=ubicacion_entrada_s3)

Para obtener más detalles sobre LMI en SageMaker, consulte Implementación de modelos grandes en Amazon SageMaker utilizando DJLServing y la inferencia paralela de modelos DeepSpeed.

Después de la invocación, el endpoint asíncrono comienza a encolar nuestro trabajo de ajuste fino. Cada trabajo pasa por los siguientes pasos: preparar las imágenes, realizar el ajuste fino de Dreambooth y LoRA, y preparar los artefactos del modelo. Profundicemos más en el proceso de ajuste fino.

Preparar las imágenes

Como mencionamos anteriormente, la calidad de las imágenes de entrada afecta directamente la calidad del modelo ajustado. Para el caso de uso de avatar, queremos que el modelo se centre en las características faciales. En lugar de requerir que los usuarios proporcionen imágenes cuidadosamente seleccionadas de tamaño y contenido exactos, implementamos un paso de preprocesamiento utilizando técnicas de visión por computadora para aliviar esta carga. En el paso de preprocesamiento, primero utilizamos un modelo de detección de rostros para aislar el rostro más grande de cada imagen. Luego recortamos y ajustamos el tamaño de la imagen al tamaño requerido de 512 x 512 píxeles para nuestro modelo. Finalmente, segmentamos el rostro del fondo y agregamos variaciones aleatorias de fondo. Esto ayuda a resaltar las características faciales, permitiendo que nuestro modelo aprenda del propio rostro en lugar del fondo. Las siguientes imágenes ilustran los tres pasos de este proceso.

Paso 1: Detección de rostros utilizando visión por computadora Paso 2: Recortar y ajustar el tamaño de la imagen a 512 x 512 píxeles Paso 3 (Opcional): Segmentar y agregar variaciones de fondo

Ajuste fino de Dreambooth y LoRA

Para el ajuste fino, combinamos las técnicas de Dreambooth y LoRA. Dreambooth te permite personalizar tu modelo de Difusión Estable, incrustando un sujeto en el dominio de salida del modelo utilizando un identificador único y ampliando el diccionario de visión del lenguaje del modelo. Utiliza un método llamado preservación previa para preservar el conocimiento semántico del modelo sobre la clase del sujeto, en este caso una persona, y utilizar otros objetos de la clase para mejorar la salida final de la imagen. De esta manera, Dreambooth puede lograr resultados de alta calidad con solo unas pocas imágenes de entrada del sujeto.

El siguiente fragmento de código muestra las entradas de nuestra clase trainer.py para nuestra solución de avatar. Observa que elegimos <<TOK>> como identificador único. Esto se hace a propósito para evitar elegir un nombre que pueda estar en el diccionario del modelo. Si el nombre ya existe, el modelo tiene que desaprender y luego volver a aprender el sujeto, lo que puede conducir a resultados deficientes en el ajuste fino. La clase del sujeto se establece en "una foto de una persona", lo que permite la preservación previa mediante la generación de fotos de personas para alimentarlas como entradas adicionales durante el proceso de ajuste fino. Esto ayudará a reducir el sobreajuste a medida que el modelo intenta preservar el conocimiento previo de una persona utilizando el método de preservación previa.

status = trn.run(base_model="stabilityai/stable-diffusion-2-1-base",
    resolution=512,
    n_steps=1000,
    concept_prompt="foto de <<TOK>>", # << identificador único del sujeto
    learning_rate=1e-4,
    gradient_accumulation=1,
    fp16=True,
    use_8bit_adam=True,
    gradient_checkpointing=True,
    train_text_encoder=True,
    with_prior_preservation=True,
    prior_loss_weight=1.0,
    class_prompt="una foto de persona", # << clase del sujeto
    num_class_images=50,
    class_data_dir=class_data_dir,
    lora_r=128,
    lora_alpha=1,
    lora_bias="ninguno",
    lora_dropout=0.05,
    lora_text_encoder_r=64,
    lora_text_encoder_alpha=1,
    lora_text_encoder_bias="ninguno",
    lora_text_encoder_dropout=0.05
)

Se han habilitado varias opciones para ahorrar memoria en la configuración, incluyendo fp16, use_8bit_adam y la acumulación de gradientes. Esto reduce la huella de memoria a menos de 12 GB, lo que permite el ajuste fino de hasta dos modelos simultáneamente en una instancia ml.g5.2xlarge.

LoRA es una técnica eficiente de ajuste fino para LLMs que congela la mayoría de los pesos y conecta una pequeña red adaptadora a capas específicas del LLM pre-entrenado, lo que permite un entrenamiento más rápido y un almacenamiento optimizado. Para Stable Diffusion, el adaptador se conecta al codificador de texto y a los componentes U-Net del pipeline de inferencia. El codificador de texto convierte la entrada de la consulta en un espacio latente que es entendido por el modelo U-Net, y el modelo U-Net utiliza el significado latente para generar la imagen en el proceso de difusión posterior. El resultado del ajuste fino son solo los pesos del codificador_de_texto y del adaptador U-Net. En el momento de la inferencia, estos pesos se pueden volver a conectar al modelo base Stable Diffusion para reproducir los resultados del ajuste fino.

A continuación se muestran diagramas detallados del ajuste fino de LoRA proporcionados por el autor original: Cheng-Han Chiang, Yung-Sung Chuang, Hung-yi Lee, “AACL_2022_tutorial_PLMs,” 2022

Al combinar ambos métodos, pudimos generar un modelo personalizado mientras ajustábamos una cantidad de parámetros diez veces menor. Esto resultó en un tiempo de entrenamiento mucho más rápido y una utilización reducida de la GPU. Además, el almacenamiento se optimizó con un peso de adaptador de solo 70 MB, en comparación con 6 GB para un modelo Stable Diffusion completo, lo que representa una reducción de tamaño del 99%.

Preparar los artefactos del modelo

Después de completar el ajuste fino, el paso de postprocesamiento comprimirá los pesos de LoRA junto con el resto de los archivos de servicio del modelo para NVIDIA Triton. Utilizamos un backend de Python, lo que significa que se requieren el archivo de configuración de Triton y el script de Python utilizado para la inferencia. Tenga en cuenta que el script de Python debe tener el nombre model.py. El archivo TAR final del modelo debe tener la siguiente estructura de archivos:

|--sd_lora
   |--config.pbtxt
   |--1\
      |--model.py
      |--output #Pesos de LoRA
         |--codificador_de_texto\
         |--unet\
         |--train.sh

Hospedar los modelos ajustados finamente usando SageMaker MMEs con GPU

Después de ajustar finamente los modelos, hospedamos los modelos personalizados de Stable Diffusion usando un SageMaker MME. Un SageMaker MME es una poderosa función de implementación que permite hospedar múltiples modelos en un solo contenedor detrás de un único punto de conexión. Gestiona automáticamente el tráfico y el enrutamiento hacia sus modelos para optimizar la utilización de recursos, ahorrar costos y minimizar la carga operativa de gestionar miles de puntos de conexión. En nuestro ejemplo, ejecutamos en instancias de GPU, y los SageMaker MME admiten GPU mediante Triton Server. Esto le permite ejecutar varios modelos en un solo dispositivo GPU y aprovechar la computación acelerada. Para obtener más detalles sobre cómo hospedar Stable Diffusion en SageMaker MMEs, consulte Crear imágenes de alta calidad con modelos Stable Diffusion y implementarlos de manera rentable con Amazon SageMaker.

Para nuestro ejemplo, hemos realizado una optimización adicional para cargar los modelos ajustados más rápido durante situaciones de inicio en frío. Esto es posible debido al diseño del adaptador de LoRA. Debido a que los pesos del modelo base y los entornos de Conda son los mismos para todos los modelos ajustados, podemos compartir estos recursos comunes al precargarlos en el contenedor de alojamiento. Esto deja solo el archivo de configuración de Triton, el backend de Python (model.py) y los pesos del adaptador de LoRA para cargar dinámicamente desde Amazon S3 después de la primera invocación. El siguiente diagrama proporciona una comparación lado a lado.

Esto reduce significativamente el tamaño del archivo TAR del modelo de aproximadamente 6 GB a 70 MB, y, por lo tanto, es mucho más rápido de cargar y descomprimir. Para realizar la precarga en nuestro ejemplo, creamos un modelo de backend de Python de utilidad en models/model_setup. El script simplemente copia el modelo base de Difusión Estable y el entorno de Conda desde Amazon S3 a una ubicación común para compartir entre todos los modelos ajustados. El siguiente es el fragmento de código que realiza la tarea:

def initialize(self, args):
          
        # configuración del entorno de Conda
        self.conda_pack_path = Path(args['model_repository']) / "sd_env.tar.gz"
        self.conda_target_path = Path("/tmp/conda")
        
        self.conda_env_path = self.conda_target_path / "sd_env.tar.gz"
             
        if not self.conda_env_path.exists():
            self.conda_env_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(self.conda_pack_path, self.conda_env_path)
        
        # configuración del modelo base de difusión
        self.base_model_path = Path(args['model_repository']) / "stable_diff.tar.gz"
        
        try:
            with tarfile.open(self.base_model_path) as tar:
                tar.extractall('/tmp')
                
            self.response_message = "Configuración del entorno del modelo exitosa."
        
        except Exception as e:
            # imprimir el mensaje de excepción
            print(f"Se produjo una excepción: {e}")
            self.response_message = f"Se produjo una excepción: {e}"

Luego, cada modelo ajustado apuntará a la ubicación compartida en el contenedor. El entorno de Conda se referencia en el archivo config.pbtxt.

name: "pipeline_0"
backend: "python"
max_batch_size: 1

...

parameters: {
  key: "EXECUTION_ENV_PATH",
  value: {string_value: "/tmp/conda/sd_env.tar.gz"}
}

El modelo base de Difusión Estable se carga desde la función initialize() de cada archivo model.py. Luego, aplicamos los pesos personalizados de LoRA al modelo unet y al modelo text_encoder para reproducir cada modelo ajustado:

...

class TritonPythonModel:

    def initialize(self, args):
        self.output_dtype = pb_utils.triton_string_to_numpy(
            pb_utils.get_output_config_by_name(json.loads(args["model_config"]),
                                               "generated_image")["data_type"])
        
        self.model_dir = args['model_repository']
    
        device='cuda'
        self.pipe = StableDiffusionPipeline.from_pretrained('/tmp/stable_diff',
                                                            torch_dtype=torch.float16,
                                                            revision="fp16").to(device)
                                                            
        # Cargar los pesos de LoRA
        self.pipe.unet = PeftModel.from_pretrained(self.pipe.unet, unet_sub_dir)

        if os.path.exists(text_encoder_sub_dir):
            self.pipe.text_encoder = PeftModel.from_pretrained(self.pipe.text_encoder, text_encoder_sub_dir)

Usar el modelo ajustado para inferencia

Ahora podemos probar nuestro modelo ajustado invocando el punto final de MME. Los parámetros de entrada que expusimos en nuestro ejemplo incluyen prompt, negative_prompt y gen_args, como se muestra en el siguiente fragmento de código. Establecemos el tipo de datos y la forma de cada elemento de entrada en el diccionario y los convertimos en una cadena JSON. Finalmente, la carga útil de cadena y el TargetModel se pasan a la solicitud para generar tu imagen de avatar.

import random

prompt = """<<TOK>> epic portrait, zoomed out, blurred background cityscape, bokeh,
 perfect symmetry, by artgem, artstation ,concept art,cinematic lighting, highly 
 detailed, octane, concept art, sharp focus, rockstar games, post processing, 
 picture of the day, ambient lighting, epic composition"""

negative_prompt = """
beard, goatee, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, 
watermark, grainy, signature, cut off, draft, amateur, multiple, gross, weird, uneven, furnishing, decorating, decoration, furniture, text, poor, low, basic, worst, juvenile, 
unprofessional, failure, crayon, oil, label, thousand hands
"""

seed = random.randint(1, 1000000000)

gen_args = json.dumps(dict(num_inference_steps=50, guidance_scale=7, seed=seed))

inputs = dict(prompt = prompt, 
              negative_prompt = negative_prompt, 
              gen_args = gen_args)

payload = {
    "inputs":
        [{"name": name, "shape": [1,1], "datatype": "BYTES", "data": [data]} for name, data in inputs.items()]
}

response = sm_runtime.invoke_endpoint(
    EndpointName=endpoint_name,
    ContentType="application/octet-stream",
    Body=json.dumps(payload),
    TargetModel="sd_lora.tar.gz",
)
output = json.loads(response["Body"].read().decode("utf8"))["outputs"]
original_image = decode_image(output[0]["data"][0])
original_image

Limpiar

Siga las instrucciones de la sección de limpieza del cuaderno para eliminar los recursos provisionados como parte de esta publicación y evitar cargos innecesarios. Consulte la información sobre precios de Amazon SageMaker para obtener detalles sobre el costo de las instancias de inferencia.

Conclusión

En esta publicación, demostramos cómo crear una solución de avatar personalizado utilizando Stable Diffusion en SageMaker. Al ajustar un modelo preentrenado con solo algunas imágenes, podemos generar avatares que reflejen la individualidad y personalidad de cada usuario. Este es solo uno de muchos ejemplos de cómo podemos utilizar la inteligencia artificial generativa para crear experiencias personalizadas y únicas para los usuarios. Las posibilidades son infinitas y te animamos a experimentar con esta tecnología y explorar su potencial para mejorar el proceso creativo. Esperamos que esta publicación haya sido informativa e inspiradora. Te animamos a probar el ejemplo y compartir tus creaciones con nosotros usando los hashtags #sagemaker #mme #genai en las plataformas sociales. Nos encantaría ver lo que creas.

Además de Stable Diffusion, hay muchos otros modelos de inteligencia artificial generativa disponibles en Amazon SageMaker JumpStart. Consulta la guía “Getting started with Amazon SageMaker JumpStart” para explorar sus capacidades.