🧨 ¡Difusión estable en JAX / Flax!

'🧨 ¡Difusión estable en JAX / Flax!' can be condensed to '🧨¡Difusión estable en JAX/Flax!'

🤗 Hugging Face Diffusers soporta Flax desde la versión 0.5.1! Esto permite una inferencia súper rápida en las TPUs de Google, como las disponibles en Colab, Kaggle o Google Cloud Platform.

Esta publicación muestra cómo ejecutar inferencia usando JAX / Flax. Si desea más detalles sobre cómo funciona Stable Diffusion o desea ejecutarlo en GPU, consulte este notebook de Colab.

Si desea seguir, haga clic en el botón de arriba para abrir esta publicación como un notebook de Colab.

• MTEB Referente de Evaluación de Incrustación de Texto Masivo
• De PyTorch DDP a Accelerate Trainer, dominio del entrenamiento distribuido con facilidad.
• Evaluando el sesgo del modelo de lenguaje con 🤗 Evaluate

Primero, asegúrese de estar utilizando un backend de TPU. Si está ejecutando este notebook en Colab, seleccione Entorno de ejecución en el menú de arriba, luego seleccione la opción “Cambiar tipo de entorno de ejecución” y luego seleccione TPU en la configuración de Acelerador de hardware.

Tenga en cuenta que JAX no es exclusivo de las TPUs, pero brilla en ese hardware porque cada servidor TPU tiene 8 aceleradores TPU trabajando en paralelo.

Configuración

import jax
num_devices = jax.device_count()
device_type = jax.devices()[0].device_kind

print(f"Se encontraron {num_devices} dispositivos JAX de tipo {device_type}.")
assert "TPU" in device_type, "El dispositivo disponible no es una TPU, por favor seleccione TPU desde Editar > Configuración del cuaderno > Acelerador de hardware"

Salida:

    Se encontraron 8 dispositivos JAX de tipo TPU v2.

Asegúrese de tener instalado diffusers.

!pip install diffusers==0.5.1

Luego importamos todas las dependencias.

import numpy as np
import jax
import jax.numpy as jnp

from pathlib import Path
from jax import pmap
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from PIL import Image

from huggingface_hub import notebook_login
from diffusers import FlaxStableDiffusionPipeline

Carga del modelo

Antes de usar el modelo, debes aceptar la licencia del modelo para descargar y usar los pesos.

La licencia está diseñada para mitigar los posibles efectos perjudiciales de un sistema de aprendizaje automático tan potente. Solicitamos a los usuarios que lean la licencia completa y cuidadosamente. Aquí ofrecemos un resumen:

1. No puedes utilizar el modelo para producir ni compartir deliberadamente salidas o contenido ilegal o perjudicial,
2. No reclamamos derechos sobre las salidas que generes, eres libre de usarlas y eres responsable de su uso, el cual no debe ir en contra de las disposiciones establecidas en la licencia, y
3. Puedes redistribuir los pesos y utilizar el modelo comercialmente y/o como servicio. Si lo haces, ten en cuenta que debes incluir las mismas restricciones de uso que las establecidas en la licencia y compartir una copia de CreativeML OpenRAIL-M con todos tus usuarios.

Los pesos de Flax están disponibles en Hugging Face Hub como parte del repositorio Stable Diffusion. El modelo Stable Diffusion se distribuye bajo la licencia CreateML OpenRail-M. Es una licencia abierta que no reclama derechos sobre las salidas que generas y te prohíbe producir deliberadamente contenido ilegal o perjudicial. La tarjeta del modelo proporciona más detalles, así que tómate un momento para leerlos y considera cuidadosamente si aceptas la licencia. Si lo haces, debes ser un usuario registrado en el Hub y usar un token de acceso para que funcione el código. Tienes dos opciones para proporcionar tu token de acceso:

• Usa la herramienta de línea de comandos huggingface-cli login en tu terminal y pega tu token cuando se te solicite. Se guardará en un archivo en tu computadora.
• O usa notebook_login() en un notebook, que hace lo mismo.

La siguiente celda presentará una interfaz de inicio de sesión a menos que ya hayas autenticado antes en esta computadora. Deberás pegar tu token de acceso.

if not (Path.home()/'.huggingface'/'token').exists(): notebook_login()

Los dispositivos TPU admiten bfloat16, un tipo de semi-float eficiente. Lo usaremos para nuestras pruebas, pero también puedes usar float32 para usar precisión completa en su lugar.

dtype = jnp.bfloat16

Flax es un marco de trabajo funcional, por lo que los modelos son sin estado y los parámetros se almacenan fuera de ellos. Al cargar el pipeline pre-entrenado de Flax, se devolverá tanto el pipeline en sí como los pesos del modelo (o parámetros). Estamos utilizando una versión bf16 de los pesos, lo cual genera advertencias de tipo que se pueden ignorar de forma segura.

pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="bf16",
    dtype=dtype,
)

Inferencia

Dado que las TPUs generalmente tienen 8 dispositivos trabajando en paralelo, replicaremos nuestra entrada tantas veces como dispositivos tengamos. Luego realizaremos la inferencia en los 8 dispositivos al mismo tiempo, cada uno responsable de generar una imagen. Por lo tanto, obtendremos 8 imágenes en el mismo tiempo que tarda un chip en generar una sola.

Después de replicar la entrada, obtenemos los identificadores de texto tokenizados invocando la función prepare_inputs del pipeline. La longitud del texto tokenizado se establece en 77 tokens, como lo requiere la configuración del modelo de texto CLIP subyacente.

prompt = "Una imagen fija de una película de Morgan Freeman interpretando a Jimi Hendrix, retrato, lente de 40mm, poca profundidad de campo, primer plano, iluminación dividida, cinematográfica"
prompt = [prompt] * jax.device_count()
prompt_ids = pipeline.prepare_inputs(prompt)
prompt_ids.shape

Salida :

(8, 77)

Replicación y paralelización

Los parámetros del modelo y las entradas deben replicarse en los 8 dispositivos paralelos que tenemos. El diccionario de parámetros se replica utilizando flax.jax_utils.replicate, que recorre el diccionario y cambia la forma de los pesos para que se repitan 8 veces. Los arrays se replican utilizando shard.

p_params = replicate(params)

prompt_ids = shard(prompt_ids)
prompt_ids.shape

Salida :

(8, 1, 77)

Esa forma significa que cada uno de los 8 dispositivos recibirá como entrada un array jnp con una forma de (1, 77). Por lo tanto, 1 es el tamaño del lote por dispositivo. En TPUs con suficiente memoria, podría ser mayor que 1 si quisiéramos generar varias imágenes (por chip) a la vez.

¡Ya casi estamos listos para generar imágenes! Solo necesitamos crear un generador de números aleatorios para pasar a la función de generación. Este es el procedimiento estándar en Flax, que es muy serio y tiene opiniones sobre los números aleatorios: se espera que todas las funciones que tratan con números aleatorios reciban un generador. Esto garantiza la reproducibilidad, incluso cuando estamos entrenando en múltiples dispositivos distribuidos.

La función auxiliar a continuación utiliza una semilla para inicializar un generador de números aleatorios. Siempre que usemos la misma semilla, obtendremos los mismos resultados exactos. Siéntete libre de usar diferentes semillas al explorar los resultados más adelante en el cuaderno.

def create_key(seed=0):
    return jax.random.PRNGKey(seed)

Obtenemos un rng y luego lo “dividimos” 8 veces para que cada dispositivo reciba un generador diferente. Por lo tanto, cada dispositivo creará una imagen diferente y todo el proceso es reproducible.

rng = create_key(0)
rng = jax.random.split(rng, jax.device_count())

El código de JAX se puede compilar en una representación eficiente que se ejecuta muy rápido. Sin embargo, debemos asegurarnos de que todas las entradas tengan la misma forma en llamadas posteriores; de lo contrario, JAX tendrá que volver a compilar el código y no podremos aprovechar la velocidad optimizada.

El pipeline de Flax puede compilar el código por nosotros si pasamos jit = True como argumento. También se asegurará de que el modelo se ejecute en paralelo en los 8 dispositivos disponibles.

La primera vez que ejecutemos la siguiente celda, llevará mucho tiempo compilar, pero las llamadas posteriores (incluso con entradas diferentes) serán mucho más rápidas. Por ejemplo, tomó más de un minuto compilar en una TPU v2-8 cuando probé, pero luego tarda aproximadamente 7s para futuras ejecuciones de inferencia.

images = pipeline(prompt_ids, p_params, rng, jit=True)[0]

Salida:

    Tiempos de CPU: user 464 ms, sys: 105 ms, total: 569 ms
    Tiempo de ejecución: 7.07 s

El arreglo devuelto tiene forma (8, 1, 512, 512, 3). Lo reestructuramos para eliminar la segunda dimensión y obtener 8 imágenes de 512 × 512 × 3 y luego las convertimos a formato PIL.

images = images.reshape((images.shape[0],) + images.shape[-3:])
images = pipeline.numpy_to_pil(images)

Visualización

Creemos una función auxiliar para mostrar las imágenes en una cuadrícula.

def image_grid(imgs, filas, columnas):
    w, h = imgs[0].size
    grid = Image.new('RGB', size=(columnas*w, filas*h))
    for i, img in enumerate(imgs): grid.paste(img, box=(i%columnas*w, i//columnas*h))
    return grid

image_grid(images, 2, 4)

Usando diferentes indicaciones

No es necesario replicar la misma indicación en todos los dispositivos. Podemos hacer lo que queramos: generar 2 indicaciones 4 veces cada una, o incluso generar 8 indicaciones diferentes al mismo tiempo. ¡Hagámoslo!

Primero, refactorizaremos el código de preparación de la entrada en una función práctica:

indicaciones = [
    "Labrador en el estilo de Hokusai",
    "Pintura de una ardilla patinando en Nueva York",
    "HAL-9000 en el estilo de Van Gogh",
    "Times Square bajo el agua, con peces y un delfín nadando",
    "Fresco romano antiguo que muestra a un hombre trabajando en su computadora portátil",
    "Fotografía en primer plano de una joven mujer negra contra un fondo urbano, alta calidad, bokeh",
    "Sillón en forma de aguacate",
    "Astronauta payaso en el espacio, con la Tierra en el fondo",
]

prompt_ids = pipeline.prepare_inputs(indicaciones)
prompt_ids = shard(prompt_ids)
images = pipeline(prompt_ids, p_params, rng, jit=True).images
images = images.reshape((images.shape[0], ) + images.shape[-3:])
images = pipeline.numpy_to_pil(images)
image_grid(images, 2, 4)

---

¿Cómo funciona la paralelización?

Dijimos antes que el flujo de trabajo diffusers de Flax compila automáticamente el modelo y lo ejecuta en paralelo en todos los dispositivos disponibles. Ahora veremos brevemente cómo funciona ese proceso.

La paralelización de JAX se puede hacer de varias formas. La más sencilla utiliza la función jax.pmap para lograr la paralelización de un solo programa y múltiples datos (SPMD, por sus siglas en inglés). Esto significa que ejecutaremos varias copias del mismo código, cada una con datos diferentes. También es posible utilizar enfoques más sofisticados, te invitamos a revisar la documentación de JAX y las páginas de pjit para explorar este tema si estás interesado.

jax.pmap hace dos cosas por nosotros:

• Compila (o realiza jit) el código, como si hubiéramos invocado jax.jit(). Esto no ocurre cuando llamamos a pmap, sino la primera vez que se invoca la función con pmap.
• Asegura que el código compilado se ejecute en paralelo en todos los dispositivos disponibles.

Para mostrar cómo funciona, aplicamos pmap al método _generate del flujo de trabajo, que es el método privado que genera las imágenes. Ten en cuenta que este método puede cambiar de nombre o eliminarse en futuras versiones de diffusers.

p_generate = pmap(pipeline._generate)

Después de utilizar pmap, la función preparada p_generate hará conceptualmente lo siguiente:

• Invocará una copia de la función subyacente pipeline._generate en cada dispositivo.
• Enviar a cada dispositivo una porción diferente de los argumentos de entrada. Para eso se utiliza el fragmentado. En nuestro caso, prompt_ids tiene forma (8, 1, 77, 768). Este arreglo se dividirá en 8 y cada copia de _generate recibirá una entrada con forma (1, 77, 768).

Podemos codificar _generate ignorando por completo el hecho de que se invocará en paralelo. Solo nos importa nuestro tamaño de lote (1 en este ejemplo) y las dimensiones que tienen sentido para nuestro código, y no tenemos que cambiar nada para hacer que funcione en paralelo.

De la misma manera que cuando usamos la llamada de canalización, la primera vez que ejecutemos la siguiente celda llevará un tiempo, pero luego será mucho más rápido.

images = p_generate(prompt_ids, p_params, rng)
images = images.block_until_ready()
images.shape

Salida :

    Tiempo de CPU: usuario 118 ms, sistema: 83.9 ms, total: 202 ms
    Tiempo de ejecución: 6.82 s

    (8, 1, 512, 512, 3)

Usamos block_until_ready() para medir correctamente el tiempo de inferencia, porque JAX utiliza un despacho asíncrono y devuelve el control al bucle de Python tan pronto como puede. No es necesario usar eso en tu código; el bloqueo ocurrirá automáticamente cuando quieras usar el resultado de un cálculo que aún no se ha materializado.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!