Ejecutando IF con difusores 🧨 en un Google Colab de nivel gratuito

Ejecutando IF con difusores en Google Colab gratis.

TL;DR: Mostramos cómo ejecutar uno de los modelos de texto a imagen de código abierto más potentes, IF, en un Google Colab de nivel gratuito con 🧨 difusores.

También puedes explorar las capacidades del modelo directamente en el Espacio de Hugging Face.

Imagen comprimida del repositorio oficial de IF en GitHub.

• Cómo instalar y usar la API de Unity de Hugging Face
• StarCoder Un LLM de última generación para el código
• Generación Asistida una nueva dirección hacia la generación de texto de baja latencia

Introducción

IF es un modelo de generación de imágenes a partir de texto basado en píxeles y fue lanzado a fines de abril de 2023 por DeepFloyd. La arquitectura del modelo está fuertemente inspirada en Imagen, propiedad de Google y no accesible al público.

IF tiene dos ventajas distintas en comparación con los modelos existentes de texto a imagen como Stable Diffusion:

• El modelo opera directamente en “espacio de píxeles” (es decir, en imágenes sin comprimir) en lugar de ejecutar el proceso de eliminación de ruido en el espacio latente, como hace Stable Diffusion.
• El modelo se entrena en salidas de T5-XXL, un codificador de texto más potente que CLIP, utilizado por Stable Diffusion como codificador de texto.

Como resultado, IF es mejor para generar imágenes con detalles de alta frecuencia (por ejemplo, rostros y manos humanas) y es el primer modelo de generación de imágenes de código abierto que puede generar imágenes de manera confiable con texto.

La desventaja de operar en el espacio de píxeles y utilizar un codificador de texto más potente es que IF tiene una cantidad significativamente mayor de parámetros. T5, el UNet de texto a imagen de IF y el UNet de escalado de IF tienen 4.5B, 4.3B y 1.2B de parámetros, respectivamente. En comparación, el codificador de texto y el UNet de Stable Diffusion 2.1 tienen solo 400M y 900M de parámetros, respectivamente.

No obstante, es posible ejecutar IF en hardware para consumidores si se optimiza el modelo para un uso de memoria reducido. Mostraremos cómo hacer esto con 🧨 difusores en esta publicación del blog.

En 1.), explicamos cómo utilizar IF para la generación de imágenes a partir de texto, y en 2.) y 3.), repasamos las capacidades de variación de imagen y relleno de imagen de IF.

💡 Nota: Estamos sacrificando ganancias en memoria por ganancias en velocidad aquí para poder ejecutar IF en un Google Colab de nivel gratuito. Si tienes acceso a GPUs de alta gama como una A100, recomendamos dejar todos los componentes del modelo en la GPU para obtener la máxima velocidad, como se hace en la demostración oficial de IF.

💡 Nota: Algunas de las imágenes más grandes se han comprimido para cargar más rápido en el formato del blog. ¡Cuando uses el modelo oficial, deberían tener una calidad aún mejor!

¡Vamos a sumergirnos 🚀!

Capacidades de generación de texto de IF

Tabla de contenidos

• Aceptación de la licencia
• Optimización de IF para ejecutarse en hardware con restricciones de memoria
• Recursos disponibles
• Instalación de dependencias
• Generación de imágenes a partir de texto
• Variación de imagen
• Relleno de imagen

Aceptación de la licencia

Antes de poder utilizar IF, debes aceptar sus condiciones de uso. Para hacerlo:

• 1. Asegúrate de tener una cuenta de Hugging Face y haber iniciado sesión
• 1. Acepta la licencia en la tarjeta de modelo de DeepFloyd/IF-I-XL-v1.0. Aceptar la licencia en la tarjeta del modelo de la etapa I aceptará automáticamente para los otros modelos de IF.
• 1. Asegúrate de haber iniciado sesión localmente. Instala huggingface_hub

pip install huggingface_hub --upgrade

ejecuta la función de inicio de sesión en una consola de Python

from huggingface_hub import login

login()

y ingresa tu token de acceso de Hugging Face Hub.

Optimización de IF para ejecutarse en hardware con restricciones de memoria

El aprendizaje automático de última generación no debería estar solo en manos de unos pocos élites. Democratizar el aprendizaje automático significa poner modelos disponibles para ejecutarse en algo más que el hardware más avanzado y mejor.

La comunidad de aprendizaje profundo ha creado herramientas de clase mundial para ejecutar modelos intensivos en recursos en hardware de consumo:

• 🤗 accelerate proporciona utilidades para trabajar con modelos grandes.
• bitsandbytes pone a disposición la cuantificación de 8 bits para todos los modelos de PyTorch.
• 🤗 safetensors no solo asegura que se ejecute el código guardado, sino que también acelera significativamente el tiempo de carga de modelos grandes.

Diffusers integra de forma transparente las bibliotecas anteriores para permitir una API sencilla al optimizar modelos grandes.

La versión gratuita de Google Colab tiene restricciones tanto en la CPU RAM (13 GB de RAM) como en la GPU VRAM (15 GB de RAM para T4), lo que dificulta la ejecución del modelo completo de más de 10 mil millones.

Veamos el tamaño de los componentes del modelo IF en precisión float32 completa:

• T5-XXL Text Encoder: 20 GB
• Stage 1 UNet: 17.2 GB
• Stage 2 Super Resolution UNet: 2.5 GB
• Stage 3 Super Resolution Model: 3.4 GB

No hay forma de ejecutar el modelo en float32 ya que los pesos de T5 y Stage 1 UNet son cada uno más grandes que la RAM de la CPU disponible.

En float16, los tamaños de los componentes son 11 GB, 8.6 GB y 1.25 GB para T5, Stage1 y Stage2 UNets, respectivamente, lo cual es factible para la GPU, pero aún nos encontramos con errores de desbordamiento de memoria de la CPU al cargar el T5 (alguna CPU está ocupada por otros procesos).

Por lo tanto, reducimos aún más la precisión de T5 utilizando la cuantificación de 8 bits de bitsandbytes, lo que permite guardar el punto de control de T5 con tan solo 8 GB.

Ahora que cada componente encaja individualmente tanto en la memoria de la CPU como en la de la GPU, debemos asegurarnos de que los componentes tengan toda la memoria de la CPU y la GPU para sí mismos cuando sea necesario.

Diffusers admite la carga modular de componentes individuales, es decir, podemos cargar el codificador de texto sin cargar la UNet. Esta carga modular garantizará que solo carguemos el componente que necesitamos en un paso dado en el flujo de trabajo para evitar agotar la RAM de la CPU y la VRAM de la GPU disponibles.

Probémoslo 🚀

Recursos disponibles

La versión gratuita de Google Colab viene con aproximadamente 13 GB de RAM de la CPU:

!grep MemTotal /proc/meminfo

MemTotal:       13297192 kB

Y una NVIDIA T4 con 15 GB de VRAM:

!nvidia-smi

Sun Apr 23 23:14:19 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   72C    P0    32W /  70W |   1335MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                                
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

Instalar dependencias

Algunas optimizaciones pueden requerir versiones actualizadas de las dependencias. Si tienes problemas, por favor verifica y actualiza las versiones.

! pip install --upgrade \
  diffusers~=0.16 \
  transformers~=4.28 \
  safetensors~=0.3 \
  sentencepiece~=0.1 \
  accelerate~=0.18 \
  bitsandbytes~=0.38 \
  torch~=2.0 -q

1. Generación de texto a imagen

Recorreremos paso a paso la generación de texto a imagen con IF utilizando Diffusers. Explicaremos brevemente las APIs y las optimizaciones, pero se pueden encontrar explicaciones más detalladas en la documentación oficial de Diffusers, Transformers, Accelerate y bitsandbytes.

1.1 Cargar el codificador de texto

Cargaremos T5 utilizando una cuantización de 8 bits. Transformers admite directamente bitsandbytes a través de la bandera load_in_8bit.

La bandera variant="8bit" descargará pesos pre-cuantizados.

También usaremos la bandera device_map para permitir que transformers transfiera capas del modelo a la CPU o al disco. Transformers big modeling admite mapas de dispositivos arbitrarios, que se pueden utilizar para cargar parámetros del modelo por separado directamente en los dispositivos disponibles. Pasar "auto" creará automáticamente un mapa de dispositivos. Consulta la documentación de transformers para obtener más información.

from transformers import T5EncoderModel

text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder", 
    device_map="auto", 
    load_in_8bit=True, 
    variant="8bit"
)

1.2 Crear embeddings de texto

La API de Diffusers para acceder a modelos de difusión es la clase DiffusionPipeline y sus subclases. Cada instancia de DiffusionPipeline es un conjunto completamente autocontenido de métodos y modelos para ejecutar redes de difusión. Podemos anular los modelos que utiliza pasando instancias alternativas como argumentos de palabra clave a from_pretrained.

En este caso, pasamos None para el argumento unet, por lo que no se cargará ninguna UNet. Esto nos permite ejecutar la parte de incrustación de texto del proceso de difusión sin cargar la UNet en la memoria.

from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=text_encoder, # pasamos el codificador de texto de 8 bits previamente instanciado
    unet=None, 
    device_map="auto"
)

IF también viene con un pipeline de super resolución. Guardaremos los embeddings de la consulta para poder pasarlos directamente al pipeline de super resolución más adelante. Esto permitirá que el pipeline de super resolución se cargue sin un codificador de texto.

En lugar de que un astronauta simplemente monte un caballo, ¡también le daremos un cartel!

Definamos una consulta adecuada:

prompt = "una fotografía de un astronauta montando un caballo sosteniendo un cartel que dice Pixel en el espacio"

y ejecutémoslo a través del modelo T5 cuantizado de 8 bits:

prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

1.3 Liberar memoria

Una vez que se hayan creado los embeddings de la consulta, ya no necesitaremos el codificador de texto. Sin embargo, todavía está en la memoria de la GPU. Necesitamos eliminarlo para poder cargar la UNet.

No es trivial liberar la memoria de PyTorch. Debemos recolectar basura de los objetos de Python que apuntan a la memoria asignada real en la GPU.

Primero, utiliza la palabra clave de Python del para eliminar todos los objetos de Python que hacen referencia a la memoria asignada en la GPU.

del text_encoder
del pipe

Eliminar el objeto de Python no es suficiente para liberar la memoria de la GPU. La recolección de basura es cuando se libera la memoria de la GPU.

Además, llamaremos a torch.cuda.empty_cache(). Este método no es estrictamente necesario, ya que la memoria caché de CUDA se liberará inmediatamente para asignaciones posteriores. Vaciar la caché nos permite verificar en la interfaz de Colab que la memoria está disponible.

Usaremos una función auxiliar flush() para vaciar la memoria.

import gc
import torch

def flush():
    gc.collect()
    torch.cuda.empty_cache()

y ejecutémoslo

flush()

1.4 Etapa 1: El proceso principal de difusión

Con nuestra memoria de GPU ahora disponible, podemos volver a cargar el DiffusionPipeline solo con el UNet para ejecutar el proceso de difusión principal.

Las banderas variant y torch_dtype se utilizan por los difusores para descargar y cargar los pesos en formato de punto flotante de 16 bits.

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

A menudo, pasamos directamente la indicación de texto a DiffusionPipeline.__call__. Sin embargo, previamente calculamos nuestras incrustaciones de texto que podemos pasar en su lugar.

IF también viene con un proceso de difusión de súper resolución. Configurando output_type="pt" se devolverán tensores de PyTorch en bruto en lugar de una imagen PIL. De esta manera, podemos mantener los tensores de PyTorch en la GPU y pasarlos directamente al pipeline de súper resolución de la etapa 2.

Definamos un generador aleatorio y ejecutemos el proceso de difusión de la etapa 1.

generator = torch.Generator().manual_seed(1)
image = pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Conviertamos manualmente los tensores en bruto a PIL y echemos un vistazo al resultado final. La salida de la etapa 1 es una imagen de 64×64.

from diffusers.utils import pt_to_pil

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

Y de nuevo, eliminamos el puntero de Python y liberamos la memoria de la CPU y la GPU:

del pipe
flush()

1.5 Etapa 2: Super Resolución de 64×64 a 256×256

IF viene con un proceso de difusión separado para el aumento de resolución.

Ejecutamos cada proceso de difusión con un pipeline separado.

El pipeline de súper resolución puede cargarse con un codificador de texto si es necesario. Sin embargo, normalmente tendremos incrustaciones de texto precalculadas del primer pipeline de IF. Si es así, cargue el pipeline sin el codificador de texto.

Creamos el pipeline

pipe = DiffusionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", 
    text_encoder=None, # sin uso de codificador de texto => ¡ahorro de memoria!
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

y lo ejecutamos, reutilizando las incrustaciones de texto precalculadas

image = pipe(
    image=image, 
    prompt_embeds=prompt_embeds, 
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Nuevamente, podemos inspeccionar los resultados intermedios.

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

Y de nuevo, eliminamos el puntero de Python y liberamos la memoria

del pipe
flush()

1.6 Etapa 3: Super Resolución de 256×256 a 1024×1024

El segundo modelo de súper resolución para IF es el Upscaler x4 de Stability AI previamente lanzado.

Creemos el pipeline y cárguelo directamente en la GPU con device_map="auto".

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-x4-upscaler", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

🧨 diffusers permite componer fácilmente modelos de difusión desarrollados independientemente, ya que los pipelines se pueden encadenar. Aquí simplemente podemos tomar la salida anterior de tensor de PyTorch y pasarla al pipeline de la etapa 3 como image=image.

💡 Nota: El x4 Upscaler no utiliza T5 y tiene su propio codificador de texto. Por lo tanto, no podemos usar los embeddings de la solicitud previamente creados y en su lugar debemos pasar la solicitud original.

pil_image = pipe(prompt, generator=generator, image=image).images

A diferencia de los pipelines IF, la marca de agua IF no se agregará automáticamente a las salidas del pipeline de escalado x4 de difusión estable.

En su lugar, podemos aplicar manualmente la marca de agua.

from diffusers.pipelines.deepfloyd_if import IFWatermarker

watermarker = IFWatermarker.from_pretrained("DeepFloyd/IF-I-XL-v1.0", subfolder="watermarker")
watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

Ver imagen de salida

pil_image[0]

¡Et voilà! Una hermosa imagen de 1024×1024 en un Google Colab de nivel gratuito.

Hemos mostrado cómo 🧨 diffusers facilita la descomposición y carga modular de modelos de difusión intensivos en recursos.

💡 Nota: No recomendamos utilizar la configuración anterior en producción. La cuantización de 8 bits, la desasignación manual de los pesos del modelo y la descarga en disco sacrifican memoria por tiempo (es decir, velocidad de inferencia). Esto puede ser especialmente notable si se reutiliza el pipeline de difusión. En producción, recomendamos utilizar una A100 de 40 GB con todos los componentes del modelo dejados en la GPU. Consulte la demo oficial de IF.

2. Variación de imagen

Los mismos checkpoints de IF también se pueden utilizar para la variación de imagen guiada por texto y el rellenado. El proceso de difusión central es el mismo que la generación de texto a imagen, excepto que la imagen inicial con ruido se crea a partir de la imagen que se va a variar o rellenar.

Para ejecutar la variación de imagen, cargue los mismos checkpoints con IFImg2ImgPipeline.from_pretrained() y IFImg2ImgSuperResolution.from_pretrained().

¡Las APIs de optimización de memoria son todas las mismas!

Liberemos la memoria de la sección anterior.

del pipe
flush()

Para la variación de imagen, comenzamos con una imagen inicial que queremos adaptar.

Para esta sección, adaptaremos el famoso meme “Slaps Roof of Car”. Descarguémoslo de internet.

import requests

url = "https://i.kym-cdn.com/entries/icons/original/000/026/561/car.jpg"
response = requests.get(url)

y carguémoslo en una imagen PIL

from PIL import Image
from io import BytesIO

original_image = Image.open(BytesIO(response.content)).convert("RGB")
original_image = original_image.resize((768, 512))
original_image

El pipeline de variación de imagen acepta tanto imágenes PIL como tensores en bruto. Consulte las cadenas de documentación para obtener más información detallada sobre las entradas esperadas, aquí .

2.1 Codificador de texto

La variación de imagen está guiada por texto, por lo que podemos definir una solicitud y codificarla con el codificador de texto de T5.

Nuevamente, cargamos el codificador de texto en una precisión de 8 bits.

from transformers import T5EncoderModel

text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder",
    device_map="auto",
    load_in_8bit=True,
    variant="8bit"
)

Para la variación de imagen, cargamos el checkpoint con IFImg2ImgPipeline. Al usar DiffusionPipeline.from_pretrained(...), los checkpoints se cargan en su pipeline predeterminado. El pipeline predeterminado para el IF es el pipeline de texto a imagen IFPipeline. Cuando se cargan checkpoints con un pipeline que no es el predeterminado, el pipeline debe especificarse explícitamente.

from diffusers import IFImg2ImgPipeline

pipe = IFImg2ImgPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    text_encoder=text_encoder,
    unet=None,
    device_map="auto"
)

Vamos a convertir a nuestro vendedor en un personaje de anime.

prompt = "estilo de anime"

Como antes, creamos los embeddings de texto con T5

prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

y liberamos la memoria de la GPU y la CPU.

Primero, eliminamos los punteros de Python

del text_encoder
del pipe

y luego liberamos la memoria

flush()

2.2 Etapa 1: El proceso de difusión principal

A continuación, cargamos solo los pesos de la etapa 1 de UNet en el objeto de la canalización, como lo hicimos en la sección anterior.

pipe = IFImg2ImgPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

La canalización de variación de imagen requiere tanto la imagen original como los embeddings de texto.

Opcionalmente, podemos usar el argumento strength para configurar la cantidad de variación. strength controla directamente la cantidad de ruido agregado. Mayor strength significa más ruido, lo que significa más variación.

generator = torch.Generator().manual_seed(0)
image = pipe(
    image=original_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Veamos de nuevo el intermedio 64×64.

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

¡Se ve bien! Podemos liberar la memoria y volver a escalar la imagen.

del pipe
flush()

2.3 Etapa 2: Super Resolución

Para la super resolución, cargamos el punto de control con IFImg2ImgSuperResolutionPipeline y el mismo punto de control que antes.

from diffusers import IFImg2ImgSuperResolutionPipeline

pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

💡 Nota: La canalización de super resolución de variación de imagen requiere la imagen generada y la imagen original.

También puedes usar el escalador de difusión estable x4 en esta imagen. Siéntete libre de probarlo usando los fragmentos de código en la sección 1.6.

image = pipe(
    image=image,
    original_image=original_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    generator=generator,
).images[0]
image

¡Genial! Liberemos la memoria y veamos las canalizaciones de rellenado finales.

del pipe
flush()

3. Rellenado

La canalización de rellenado de IF es similar a la variación de imagen, excepto que solo se desenruida una área selecta de la imagen.

Especificamos el área a rellenar con una máscara de imagen.

Vamos a mostrar las increíbles capacidades de “generación de letras” de IF. Podemos reemplazar este texto de letrero con un eslogan diferente.

Primero, descarguemos la imagen

import requests

url = "https://i.imgflip.com/5j6x75.jpg"
response = requests.get(url)

y convirtámosla en una imagen de PIL

from PIL import Image
from io import BytesIO

original_image = Image.open(BytesIO(response.content)).convert("RGB")
original_image = original_image.resize((512, 768))
original_image

Vamos a enmascarar la señal para poder reemplazar su texto.

Para mayor comodidad, hemos generado previamente la máscara y la hemos cargado en un conjunto de datos de HF.

Vamos a descargarla.

from huggingface_hub import hf_hub_download

mask_image = hf_hub_download("diffusers/docs-images", repo_type="dataset", filename="if/sign_man_mask.png")
mask_image = Image.open(mask_image)

mask_image

💡 Nota : También puedes crear máscaras manualmente creando una imagen en escala de grises.

from PIL import Image
import numpy as np

height = 64
width = 64

example_mask = np.zeros((height, width), dtype=np.int8)

# Establece los píxeles enmascarados en 255
example_mask[20:30, 30:40] = 255

# Asegúrate de crear la imagen en modo 'L'
# que significa un solo canal en escala de grises
example_mask = Image.fromarray(example_mask, mode='L')

example_mask

Ahora podemos comenzar el proceso de relleno de espacios en blanco 🎨🖌

3.1. Codificador de Texto

Nuevamente, primero cargamos el codificador de texto

from transformers import T5EncoderModel

text_encoder = T5EncoderModel.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0",
    subfolder="text_encoder", 
    device_map="auto", 
    load_in_8bit=True, 
    variant="8bit"
)

Esta vez, inicializamos la tubería de relleno de espacios en blanco IFInpaintingPipeline con los pesos del codificador de texto.

from diffusers import IFInpaintingPipeline

pipe = IFInpaintingPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=text_encoder, 
    unet=None, 
    device_map="auto"
)

Listo, hagamos que el hombre promocione más capas en su lugar.

prompt = 'el texto, "solo agrega más capas"'

Una vez definido el texto, podemos crear los embeddings del texto de entrada

prompt_embeds, negative_embeds = pipe.encode_prompt(prompt)

Al igual que antes, liberamos la memoria

del text_encoder
del pipe
flush()

3.2 Etapa 1: El proceso principal de difusión

Al igual que antes, ahora cargamos la tubería de la etapa 1 solo con la UNet.

pipe = IFInpaintingPipeline.from_pretrained(
    "DeepFloyd/IF-I-XL-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

Ahora, necesitamos pasar la imagen de entrada, la imagen de la máscara y los embeddings del texto de entrada.

image = pipe(
    image=original_image,
    mask_image=mask_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    output_type="pt",
    generator=generator,
).images

Echemos un vistazo a la salida intermedia.

pil_image = pt_to_pil(image)
pipe.watermarker.apply_watermark(pil_image, pipe.unet.config.sample_size)

pil_image[0]

¡Se ve bien! El texto es bastante consistente.

Liberemos la memoria para poder ajustar la escala de la imagen

del pipe
flush()

3.3 Etapa 2: Super Resolución

Para la super resolución, carga el punto de control con IFInpaintingSuperResolutionPipeline.

from diffusers import IFInpaintingSuperResolutionPipeline

pipe = IFInpaintingSuperResolutionPipeline.from_pretrained(
    "DeepFloyd/IF-II-L-v1.0", 
    text_encoder=None, 
    variant="fp16", 
    torch_dtype=torch.float16, 
    device_map="auto"
)

El pipeline de super resolución de inpainting requiere la imagen generada, la imagen original, la imagen de la máscara y los embeddings de la consulta.

Vamos a realizar una última ejecución de denoising.

image = pipe(
    image=image,
    original_image=original_image,
    mask_image=mask_image,
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=negative_embeds, 
    generator=generator,
).images[0]
image

¡Genial, el modelo generó texto sin cometer ni un solo error de ortografía!

Conclusión

IF en precisión de punto flotante de 32 bits utiliza un total de 40 GB de pesos. Mostramos cómo, utilizando solo modelos y bibliotecas de código abierto, IF se puede ejecutar en una instancia gratuita de Google Colab.

El ecosistema de ML se beneficia enormemente del intercambio de herramientas y modelos abiertos. En este cuaderno, se utilizaron modelos de DeepFloyd, StabilityAI y Google. Las bibliotecas utilizadas, como Diffusers, Transformers, Accelerate y bitsandbytes, se benefician de innumerables contribuyentes de diferentes organizaciones.

Un agradecimiento enorme al equipo de DeepFloyd por la creación y la liberación de IF, y por contribuir a la democratización del buen aprendizaje automático 🤗.

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