Guía definitiva para optimizar Stable Diffusion XL

Versión de CUDA y PyTorch

Empecé esta prueba con la duda de si habría diferencia entre utilizar CUDA 11.8 o CUDA 12.1, así como posibles diferencias entre las distintas versiones de PyTorch, siempre por encima de la versión 2.0.

Pues... qué desilusión, todas tienen el mismo rendimiento. Las diferencias son tan pequeñas que tal vez desaparecen si hago un mayor número de pruebas.

Cuál utilizar: Aún así tengo una teoría. La versión 11.8 de CUDA ha estado entre nosotros más tiempo, por lo que tiene sentido que las librerías y aplicaciones desempeñen mejor en esta versión que en una más moderna. En cambio, en cuanto a PyTorch, cuanto más moderna sea la versión más funcionaliades debería ofrecer y menos fallos debería incluir. Por lo tanto, y aunque sea placebo, yo me quedo con CUDA 11.8 + PyTorch 2.2.0.

Mecanismos de atención

Antes había que optimizar los mecanismos de atención instalando librerías como xFormers o FlashAttention.

Si te preguntas por qué en este artículo no aparece mención a estas optimizaciones, es porque ya no hacen falta. Desde la llegada de PyTorch 2.0, la optimización de estos algoritmos está integrada en la propia librería a través de varias implementaciones (como estas dos mencionadas). PyTorch utilizará la implementación adecuada según los inputs y el hardware en uso.

FP16

Por defecto Stable Diffusion XL utiliza el formato de coma flotante de 32 bits (FP32) para representar los números con los que trabaja y realiza cálculos.

La pregunta obvia es... ¿se puede bajar la precisión? La respuesta es sí. Al utilizar el parámetro torch_dtype=torch.float16, el modelo se carga en memoria en formato de coma flotante de media precisión (FP16). Para evitar realizar esta conversión constantemente podemos descargar el modelo en formato FP16, ya que se distribuye esa variante. Basta con incluir el parámetro variant='fp16'.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

• Retrato
• Interior
• Macro
• 3D

• FP32
• FP16
• FP32
• FP16
• FP32
• FP16
• FP32
• FP16

Al trabajar con números que ocupan la mitad, el uso de memoria cae drásticamente y la velocidad con la que se realizan los cálculos aumenta considerablemente.

El único punto "negativo" es una pérdida de calidad en la imagen generada, pero es prácticamente imposible ver alguna diferencia porque FP16 sigue siendo suficiente.

Además, gracias al parámetro variant='fp16' ahorramos espacio en disco ya que esta variante ocupa la mitad de tamaño (5 GB en vez de 10 GB).

Cuándo utilizar: Siempre.

TF32

TensorFloat-32 es un formato a medio camino entre FP32 y FP16, que se utiliza en algunas tarjetas gráficas NVIDIA (como los modelos A100 o H100) para realizar cálculos utilizando los tensor cores. Utiliza los mismos bits que FP32 para representar el exponente y los mismos bits que FP16 para representar la parte decimal.

A pesar de que en nuestro banco de pruebas (RTX 3090) no se pueden realizar cálculos con este formato, sucede algo curioso que seguro no te vas a esperar.

Para activar este formato numérico se utilizan dos propiedades: torch.backends.cudnn.allow_tf32 (que viene activada por defecto) y torch.backends.cuda.matmul.allow_tf32 (que habría que activarla manualmente). La primera activa TF32 en operaciones de convolución realizadas por cuDNN y la segunda activa TF32 en operaciones de multiplicación de matrices.

Que la propiedad torch.backends.cudnn.allow_tf32 esté activada por defecto sea cual sea tu tarjeta gráfica es un poco extraño, ¿verdad? Veamos qué ocurre si desactivamos esta propiedad asignándole el valor False.

torch.backends.cudnn.allow_tf32 = False
# Ya está desactivada por defecto
# torch.backends.cuda.matmul.allow_tf32 = False

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Además, y por curiosidad, he realizado pruebas utilizando una tarjeta gráfica NVIDIA A100 activando TF32.

# Ya está activada por defecto
# torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Utilizando una RTX 3090, si desactivamos la propiedad torch.backends.cudnn.allow_tf32 disminuye el uso de memoria en un 7%. ¿Por qué? No lo sé, en principio diría que es un bug ya que no tiene sentido activar TF32 en una tarjeta gráfica que no lo soporta.

En el caso de utilizar una tarjeta gráfica A100, utilizando FP16 conseguimos reducir el tiempo de inferencia y el uso de memoria de manera muy sustancial. El uso de memoria se puede reducir aún más desactivando torch.backends.cudnn.allow_tf32 al igual que en la RTX 3090. En cuanto a utilizar TF32, al estar a medio camino entre FP32 y FP16... no consigue batir a FP16.

Cuándo utilizar: En el caso de tarjetas gráficas no compatibles con TF32, claramente es buena idea desactivar la propiedad que viene activada por defecto. Utilizando una A100 tampoco merece la pena utilizar TF32 si podemos utilizar FP16.

Model CPU Offload

Esta optimización proviene de la librería accelerate. Cuando se ejecuta un pipeline se cargan todos los modelos en memoria. Mediante esta optimización le decimos al pipeline que mueva a la memoria únicamente el modelo que se necesite en cada momento. Este orden podrás encontrarlo en el código fuente del pipeline, en el caso de Stable Diffusion XL encontraremos la siguiente línea:

model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"

El código para implementar Model CPU Offload es bastante sencillo:

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
)

pipe.enable_model_cpu_offload()

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

pipe = AutoPipelineForText2Image.from_pretrained(
  # ...
).to('cuda')

Utilizar esta técnica va a depender de la tarjeta gráfica con la que contemos. Nos será útil si nuestra tarjeta gráfica tiene 6-8 GB de memoria ya que el consumo de memoria se reduce exáctamente a la mitad.

En cuanto al tiempo de inferencia no se ve tan perjudicado como para que sea un problema.

Cuándo utilizar: Cuando necesitemos reducir el consumo de memoria. Como el componente que más memoria utiliza es el predictor de ruido (U-Net), no conseguiremos reducir más aún el consumo de memoria aplicando optimizaciones al VAE.

Sequential CPU Offload

Esta optimización funciona de manera similar a Model CPU Offload, solo que es mucho más agresiva. En vez de mover componentes enteros a la memoria, se mueven submódulos de cada componente. Por ejemplo, en vez de mover el modelo U-Net entero a la memoria, se mueven ciertas partes mientras se trabaja con ellas, ocupando la mínima cantidad de memoria posible. Esto significa que si el predictor de ruido tiene que limpiar un tensor a lo largo de 50 pasos, los submódulos tienen que entrar y salir de la memoria 50 veces.

Se añade también con una sola línea:

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
)

pipe.enable_sequential_cpu_offload()

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Esta optimización pondrá a prueba nuestra paciencia. El tiempo de inferencia aumenta una barbaridad a cambio de reducir el uso de memoria lo máximo posible.

Cuándo utilizar: Si necesitas utilizar menos de 4 GB de memoria, utilizar esta optimización junto con VAE FP16 fix o Tiny VAE es tu única opción, pero mejor si no lo necesitas.

Proceso por lotes

Esta técnica es el resultado del aprendizaje obtenido en 2 artículos de este blog: Cómo implementar Stable Diffusion y PixArt-α con menos de 8GB de VRAM. En estos artículos encontrarás información sobre algunas líneas de código que utilizaré aquí pero no explicaré nuevamente.

Se trata de ejecutar los componentes por lotes, el problema es que el pipeline oficial no está implementado de manera óptima para reducir el uso de memoria al máximo. Cuando inicias el pipeline y solo deseas obtener los text encoders... no puedes.

Es decir, deberíamos poder hacer esto:

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
  unet=None,
  vae=None,
).to('cuda')

Pero no se puede. Cuando inicias el pipeline, éste necesita acceder a la configuración del modelo U-Net (self.unet.config.*), así como a la del VAE (self.vae.config.*).

Por lo tanto (y sin necesidad de crear un fork), vamos a utilizar los text encoders a mano sin depender del pipeline.

El primer paso es copiar la función encode_prompt del pipeline y adaptarla/simplificarla.

Esta función se encarga de tokenizar un prompt y procesarlo para obtener los tensores de embedding ya transformados. Encontrarás una explicación de este proceso en este artículo.

def encode_prompt(prompts, tokenizers, text_encoders):
  embeddings_list = []

  for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
    cond_input = tokenizer(
      prompt,
      max_length=tokenizer.model_max_length,
      padding='max_length',
      truncation=True,
      return_tensors='pt',
    )

    prompt_embeds = text_encoder(cond_input.input_ids.to('cuda'), output_hidden_states=True)

    pooled_prompt_embeds = prompt_embeds[0]
    embeddings_list.append(prompt_embeds.hidden_states[-2])

  prompt_embeds = torch.concat(embeddings_list, dim=-1)

  negative_prompt_embeds = torch.zeros_like(prompt_embeds)
  negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)

  bs_embed, seq_len, _ = prompt_embeds.shape
  prompt_embeds = prompt_embeds.repeat(1, 1, 1)
  prompt_embeds = prompt_embeds.view(bs_embed * 1, seq_len, -1)

  seq_len = negative_prompt_embeds.shape[1]
  negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1)
  negative_prompt_embeds = negative_prompt_embeds.view(1 * 1, seq_len, -1)

  pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)
  negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)

  return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds

El siguiente paso es instanciar todos componentes y modelos que necesitamos. También necesitaremos más adelante el garbage collector (gc).

import gc
from transformers import CLIPTokenizer, CLIPTextModel, CLIPTextModelWithProjection

# ...

tokenizer = CLIPTokenizer.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  subfolder='tokenizer',
)

text_encoder = CLIPTextModel.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  subfolder='text_encoder',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

tokenizer_2 = CLIPTokenizer.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  subfolder='tokenizer_2',
)

text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  subfolder='text_encoder_2',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

Y ahora queda unir estas dos partes. Llamamos a la función encode_prompt y le pasamos el mismo prompt tanto al primer text encoder como al segundo. También le damos los componentes para que pueda utilizarlos.

with torch.no_grad():
  for generation in queue:
    generation['embeddings'] = encode_prompt(
      [generation['prompt'], generation['prompt']],
      [tokenizer, tokenizer_2],
      [text_encoder, text_encoder_2],
    )

Los tensores que obtenemos como resultado los guardamos en una variable para utilizarlos después.

Como ya tenemos todos los prompts procesados podemos eliminar de la memoria estos componentes:

del tokenizer, text_encoder, tokenizer_2, text_encoder_2
gc.collect()
torch.cuda.empty_cache()

Ahora vamos a crear un pipeline que tendrá acceso solamente al U-Net y al VAE, ahorrándo memoria al no tener que instanciar los text encoders.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
  tokenizer=None,
  text_encoder=None,
  tokenizer_2=None,
  text_encoder_2=None,
).to('cuda')

for generation in queue:
  pipe(
    prompt_embeds=generation['embeddings'][0],
    negative_prompt_embeds =generation['embeddings'][1],
    pooled_prompt_embeds=generation['embeddings'][2],
    negative_pooled_prompt_embeds=generation['embeddings'][3],
    output_type='latent',
  )

Utilizamos el pipeline para procesar los tensores de embedding que tenemos guardados del paso anterior. Recuerda que en esta parte el pipeline crea un tensor lleno de ruido y lo limpia a lo largo de 50 pasos, dejándose guiar por nuestros embeddings.

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  generation['latents'] = pipe(
    prompt_embeds=generation['embeddings'][0],
    negative_prompt_embeds =generation['embeddings'][1],
    pooled_prompt_embeds=generation['embeddings'][2],
    negative_pooled_prompt_embeds=generation['embeddings'][3],
    generator=generator,
    output_type='latent',
  ).images # No accedemos a images[0], sino el tensor completo

Como puedes ver, le hemos indicado al pipeline que nos devuelva el tensor que está en el espacio latente (output_type='latent'). Esto lo hacemos porque de no ser así, se cargaría el VAE en memoria para devolvernos una imagen y esto ocasionaría que ambos modelos estarían ocupando recursos a la vez. Así que vamos a eliminar primero el modelo U-Net como hicimos con los text encoders:

del pipe.unet
gc.collect()
torch.cuda.empty_cache()

Y ahora sí, convertimos en una imagen los tensores libres de ruido que tenemos almacenados:

pipe.upcast_vae()

with torch.no_grad():
  for i, generation in enumerate(queue, start=1):
    generation['latents'] = generation['latents'].to(next(iter(pipe.vae.post_quant_conv.parameters())).dtype)

    image = pipe.vae.decode(
      generation['latents'] / pipe.vae.config.scaling_factor,
      return_dict=False,
    )[0]

    image = pipe.image_processor.postprocess(image, output_type='pil')[0]
    image.save(f'image_{i}.png')

Este bucle se encarga de decodificar los tensores del espacio latente para convertirlos al espacio de imagen. Después, mediante el método pipe.image_processor.postprocess, se convierten en una imagen y se guarda.

Este es uno de los motivos por el que me animé a escribir este artículo. Sin penalización en tiempo de inferencia hemos conseguido reducir a la mitad el uso de memoria. Ahora podríamos incluso generar imágenes con una tarjeta gráfica con 6 GB de memoria. También se puede añadir el modelo refiner mediante el método Ensemble of Expert Denoisers y el consumo de memoria sería el mismo.

Cuándo utilizar: Es cierto que Model CPU Offload es solo una línea de código extra, pero hay un pequeño incremento en tiempo de inferencia. Por lo que, si no te importa escribir algo más de código, con esta técnica tendrás un control absoluto y mejor rendimiento.

Stable Fast

Stable Fast es un proyecto que acelera cualquier modelo de difusión mediante una serie de técnicas, como por ejemplo: trazar los modelos mediante una versión mejorada de torch.jit.trace, xFormers, implementación avanzada de Channels-last memory format, entre otras. La verdad es que han realizado un trabajo impresionante.

El resultado que prometen es un tiempo de inferencia de récord, batiendo con bastante diferencia al API torch.compile y poniéndose a la par de TensorRT. Lo más gracioso de todo es que, al tratarse de optimizaciones en tiempo de ejecución, no hace falta esperar decenas de minutos para realizar la compilación inicial.

Para integrarlo, primero instalamos la librería del proyecto además de Triton y una versión de xFormers compatible con la versión de PyTorch que estemos utilizando.

pip install stable-fast
pip install torch torchvision triton xformers --index-url https://download.pytorch.org/whl/cu118

Y ahora sí, modificamos el script para importar y activar estas librerías y hacer uso de Stable Fast:

import xformers
import triton
from sfast.compilers.diffusion_pipeline_compiler import (compile, CompilationConfig)

# ...

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

config = CompilationConfig.Default()

config.enable_xformers = True
config.enable_triton = True
config.enable_cuda_graph = True

pipe = compile(pipe, config)

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Además, este proyecto también destaca por su sencillez, unas cuantas líneas y todo en marcha. Veámos ahora si cumple con las expectativas.

• Retrato
• Interior
• Macro
• 3D

• Base
• Stable Fast
• Base
• Stable Fast
• Base
• Stable Fast
• Base
• Stable Fast

Cumple sobradamente, se nota el gran trabajo que hay detrás de este proyecto.

El incremento de velocidad es de los más notables de este artículo. La primera imagen que generamos tarda un poco más (19s), pero no tiene mayor importancia si hacemos warm up como en estas pruebas.

El uso de memoria aumenta un poco pero es bastante asumible.

En cuanto al aspecto visual, la composición cambia ligeramente. En algunas imágenes incluso diría que se ha incrementado la calidad de ciertos elementos, así que... ver para creer.

Cuándo utilizar: Yo diría que siempre.

DeepCache

DeepCache promete ser una de las mejores optimizaciones que podemos implementar, sin apenas desventajas y además es sencilla de añadir. Se hace uso de un sistema de almacenamiento en caché para reutilizar las funciones de alto nivel además de actualizar las funciones de bajo nivel de una manera más eficiente.

Primero instalamos la librería necesaria:

pip install deepcache

Y después integramos el siguiente código en nuestro pipeline:

from DeepCache import DeepCacheSDHelper

# ...

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

helper = DeepCacheSDHelper(pipe=pipe)
helper.set_params(cache_interval=3, cache_branch_id=0)
helper.enable()

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Hay dos parámetros que se pueden modificar para conseguir mayor velocidad, aunque introduciendo mayor pérdida de calidad en el resultado.

• cache_interval=3: Especifica cada cuántos pasos se actualiza la caché.
• cache_branch_id=0: Especifica qué rama de la red neuronal es responsable de ejecutar los procesos de almacenamiento en caché (en orden descendente, 0 es la primera capa).

Veamos el resultado con los parámetros recomendados por defecto.

• Retrato
• Interior
• Macro
• 3D

• Base
• DeepCache
• Base
• DeepCache
• Base
• DeepCache
• Base
• DeepCache

Wow. Con una pequeña penalización en el uso de memoria, se consigue reducir a más de la mitad el tiempo de inferencia.

En cuanto a la calidad de imagen, habrás podido comprobar que cambia bastante y, lamentablemente, a peor. Dependiendo del estilo de la imagen importará más o menos, pero la desventaja ahí está (en imágenes de objetos parece no reducir mucho la calidad).

Aumentando el valor de cache_branch_id parece que se obtiene un poco más de calidad visual, aunque puede que no la suficiente.

Cuándo utilizar: Ya que reduce tanto el tiempo de inferencia, es comprensible que reduzca un poco la calidad del resultado. Sin duda, para probar prompts o parámetros es una optimización muy útil. Para utilizarla en un proceso en el que buscamos un buen resultado... diría que no.

TensorRT

TensorRT es un entorno de ejecución y optimizador de inferencia de alto rendimiento creado por NVIDIA. Promete acelerar el proceso de inferencia de redes neuronales batiendo récords.

Pero ya de inicio tenemos un problema. Para las pruebas estamos utilizando pipelines de la librería diffusers, y por el momento no existe un pipeline compatible con TensorRT para Stable Diffusion XL. Sí que existen pipelines comunitarios para Stable Diffusion 2.x (txt2img, img2img o inpainting). También he visto algunas cosas para Stable Diffusion 1.x, pero como decía, no para SDXL.

Por otro lado, en HuggingFace podemos encontrar el repositorio oficial stabilityai/stable-diffusion-xl-1.0-tensorrt. Contiene instrucciones para realizar el proceso de inferencia con TensorRT, pero lamentablemente se utilizan scripts bastante complejos y prácticamente imposible de adaptar para estas pruebas.

Los resultados se van a ver bastante distintos porque, para empezar, ni siquiera está disponible el mismo scheduler (Euler) en los scripts que he utilizado. Aún así, he reutilizado todos los valores que he podido, incluyendo prompt positivo, la ausencia del prompt negativo, la misma semilla, el mismo valor de CFG y el mismo tamaño de imagen.

Te dejo las instrucciones para utilizar este script por si quieres indagar:

# Clonar el repositorio entero o descargar los ficheros de esta carpeta
# https://github.com/rajeevsrao/TensorRT/tree/release/8.6/demo/Diffusion

# Crear y activar un entorno virtual, como de costumbre
python -m venv .venv

## Unix
source .venv/bin/activate
## Windows
.venv\Scripts\activate

# Instalar librerías necesarias
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install transformers accelerate diffusers cuda-python nvtx onnx colored scipy polygraphy
pip install --pre --extra-index-url https://pypi.nvidia.com tensorrt
pip install --pre --extra-index-url https://pypi.ngc.nvidia.com onnx_graphsurgeon

# Podemos comprobar que TensorRT está instalado correctamente con la siguiente línea
python -c "import tensorrt; print(tensorrt.__version__)"
# 9.3.0.post12.dev1

# Realizar inferencia
python demo_txt2img_xl.py "macro shot of a bee collecting nectar from lavender flowers"

• Retrato
• Interior
• Macro
• 3D

• Base
• TensorRT
• Base
• TensorRT
• Base
• TensorRT
• Base
• TensorRT

Tras preparar los modelos (algo que lleva una media hora y solo ocurre la primera vez), el proceso de inferencia parece acelerarse bastante, consiguiendo generar cada imagen en apenas 8 segundos, a diferencia de los 14 segundos que tarda el código no optimizado. No puedo hablar del consumo de memoria porque diría que TensorRT utiliza otros APIs distintos.

Sobre la calidad de las imágenes... se ven asombrosas de serie.

Cuándo utilizar: Si en tu proceso puedes integrar TensorRT, adelante. Parece una buena optimización y al menos deberías probarla.

torch.compile

Utilizando PyTorch 2 o superior podemos compilar modelos para obtener un mejor rendimiento gracias al API torch.compile. Si bien es cierto que la compilación lleva su tiempo, las sucesivas llamadas se beneficiarán de un extra de velocidad.

En versiones anteriores de PyTorch también se podían compilar los modelos con la técnica de tracing, a través del API torch.jit.trace. Esta compilación en tiempo de ejecución (just-in-time / JIT) es menos eficiente que el nuevo método, así que ya nos podemos ir olvidando de este API.

En el método torch.compile, el parámetro mode acepta los siguientes valores: default, reduce-overhead, max-autotune y max-autotune-no-cudagraphs. En teoría son distintos pero no he visto diferencia, así que vamos a utilizar reduce-overhead.

RuntimeError: Windows not yet supported for torch.compile

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

pipe.unet = torch.compile(pipe.unet, mode='reduce-overhead', fullgraph=True)

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Vamos a evaluar tanto el tiempo que tarda en compilar el modelo como el tiempo que tarda cada sucesiva generación.

Una simple optimización que enseguida se vuelve rentable.

Cuándo utilizar: Siempre que vayas a generar las suficientes imágenes como para que haya merecido la pena el tiempo de compilación.

OneDiff

OneDiff es una librería de optimización compatible con diffusers, ComfyUI y Stable Diffusion web UI de Automatic1111. El nombre significa literalmente: una línea de código para acelerar modelos de diffusión.

Utiliza técnicas como la cuantización, mejoras en los mecanismos de atención y compilación de modelos.

La instalación se realiza añadiendo un par de librerías, aunque si utilizas otra versión de CUDA o quieres utilizar otro método de instalación, consulta la documentación.

pip install --pre oneflow -f https://github.com/siliconflow/oneflow_releases/releases/expanded_assets/community_cu118

pip install --pre onediff

RuntimeError: This package is a placeholder. Please install oneflow following the instructions in https://github.com/Oneflow-Inc/oneflow#install-oneflow

Los creadores también ofrecen una versión Enterprise que promete un 20% extra de velocidad (o incluso más), aunque no lo puedo comprobar y tampoco dan muchos detalles.

Al igual que torch.compile, el código necesario es una única línea que altera el comportamiento de pipe.unet.

import oneflow as flow
from onediff.infer_compiler import oneflow_compile

# ...

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

pipe.unet = oneflow_compile(pipe.unet)

generator = torch.Generator(device='cuda')

with flow.autocast('cuda'):
  for i, generation in enumerate(queue, start=1):
    generator.manual_seed(generation['seed'])

    image = pipe(
      prompt=generation['prompt'],
      generator=generator,
    ).images[0]

    image.save(f'image_{i}.png')

Vamos a ver si cumple con las expectativas.

• Retrato
• Interior
• Macro
• 3D

• Base
• OneDiff
• Base
• OneDiff
• Base
• OneDiff
• Base
• OneDiff

OneDiff introduce un ligero cambio en la estructura de la imagen, pero es un cambio favorable. En la imagen de interior se puede apreciar cómo se arregla un fallo convirtiéndolo en una sombra.

El tiempo de compilación es muy bajo, bastante más rápido que torch.compile.

En cuanto al tiempo de inferencia se consigue una impresionante reducción del 45%, batiendo a todas las optimizaciones con las que compite (Stable Fast, TensorRT y torch.compile).

Sorprendentemente (y a diferencia de Stable Fast), no hay aumento en el uso de memoria.

Cuándo utilizar: ¿Siempre? Mejora la calidad visual del resultado, reduce casi a la mitad el tiempo de inferencia y la única penalización es una pequeña espera en el tiempo de compilación. ¡Qué gran trabajo!

Channels-last memory format

El formato de memoria channels-last organiza los datos de manera que los canales de color se almacenen en la última dimensión del tensor.

Por defecto el tensor tiene el formato NCHW, que corresponden a las siguientes cuatro dimensiones:

1. N (Number / Número): Cuántas imágenes generar al mismo tiempo (batch size).
2. C (Channels / Canales): Cuántos canales tiene la imagen.
3. H (Height / Altura): La altura de la imagen en pixeles.
4. W (Width / Anchura): La anchura de la imagen en pixeles.

En cambio con esta técnica se reordenan los datos del tensor para que estén en formato NHWC, poniendo el número de canales al final.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

pipe.unet.to(memory_format=torch.channels_last)

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Se puede comprobar si el tensor se ha reordenado mediante la siguiente línea (poniéndola antes y después):

print(pipe.unet.conv_out.state_dict()['weight'].stride())

Aunque puede mejorar la eficiencia en algunos casos y reducir el uso de memoria, no es compatible con algunas redes neuronales e incluso podría empeorar el rendimiento. ¡Vamos a salir de dudas!

El modelo U-Net de Stable Diffusion XL parece que no se beneficia de esta optimización, pero el conocimiento no ocupa lugar, ¿verdad?

Cuándo utilizar: Nunca, supongo.

FreeU

FreeU es la primera y única optimización que no mejora el tiempo de inferencia o el uso de memoria, si no la calidad del resultado.

Esta técnica equilibra la contribución de dos elementos clave en la arquitectura U-Net: las conexiones de omisión o skip connections (que introducen detalles de alta frecuencia) y los mapas de características del backbone (que aportan semántica).

En otras palabras, FreeU contrarresta la introducción de detalles poco naturales en las imágenes, ofreciendo resultados visuales más realistas.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.3, b2=1.4)

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Puedes jugar con estos valores aunque estos son los recomendados para Stable Diffusion XL. Más información en el readme del proyecto.

• Retrato
• Interior
• Macro
• 3D

• Base
• FreeU
• Base
• FreeU
• Base
• FreeU
• Base
• FreeU

Nunca había probado FreeU y la verdad es que no esperaba este resultado. Me parecen imágenes bastante impactantes aunque la estructura sea algo distinta a la original. Parece como si la imagen fuera más fiel al prompt y se centrase en ofrecer la máxima calidad visual en vez de perderse en los pequeños detalles.

El único punto negativo que veo es que pierde algo de coherencia la imagen. Por ejemplo, el sofá tiene una planta encima y la abeja tiene 3 alas. No lo sé, Rick...

Cuándo utilizar: Cuando queramos obtener un resultado más creativo y con mayor calidad visual (aunque también depende del estilo que busquemos).

VAE FP16 fix

Como vimos en la optimización de proceso por lotes, el VAE que incluye Stable Diffusion XL de serie no funciona en formato FP16. Antes de decodificar imágenes, el pipeline ejecuta un método que fuerza al modelo a trabajar en formato FP32 (pipe.upcast_vae()). Y tal como vimos en la optimización FP16, ejecutar un modelo en formato FP32 es un gasto innecesario de recursos.

El usuario madebyollin (creador también de TAESD, algo que veremos más abajo) ha creado una versión parcheada de este modelo para que funcione en formato FP16.

Solo tenemos que importar este VAE y reemplazar el original:

from diffusers import AutoPipelineForText2Image, AutoencoderKL

# ...

vae = AutoencoderKL.from_pretrained(
  'madebyollin/sdxl-vae-fp16-fix',
  use_safetensors=True,
  torch_dtype=torch.float16,
).to('cuda')

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
  vae=vae,
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

• Retrato
• Interior
• Macro
• 3D

• Base
• VAE FP16
• Base
• VAE FP16
• Base
• VAE FP16
• Base
• VAE FP16

No hay pérdida de calidad, las imágenes son prácticamente iguales.

En cuanto al uso de memoria se ha conseguido rascar casi un 15%, no está nada mal para este simple cambio.

Cuándo utilizar: Siempre, salvo que prefieras utilizar la optimización Tiny VAE.

VAE slicing

Cuando generamos varias imágenes al mismo tiempo (aumentando el batch size), el VAE decodifica todos los tensores al mismo tiempo (en paralelo). Esto incrementa considerablemente el uso de memoria. Para evitarlo, se puede utilizar la técnica de VAE slicing para decodificar los tensores uno por uno (en serie). Más o menos lo mismo que hicimos manualmente en la optimización Proceso por lotes.

Aunque se utilice, por ejemplo, un batch size de 1, 2, 8 o 32, el consumo de memoria por parte del VAE será el mismo, a cambio de una pequeña penalización de tiempo que apenas se notará.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

pipe.enable_vae_slicing()

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Cuando el batch size es 1 esta optimización no hace nada. Y ya que en estas pruebas estamos utilizando un batch size de 1, nos saltaremos los resultados de la prueba para ver directamente el veredicto.

Esta optimización intenta reducir el uso de memoria cuando aumentas el batch size, que es precisamente lo que más aumenta el uso de memoria. Es una contradicción en sí misma.

Cuándo utilizar: Solamente cuando tengas un proceso bien establecido, de manera que puedas generar varias imágenes al mismo tiempo y la ejecución del VAE sea el cuello de botella. O sea, rara vez.

VAE tiling

Cuando generamos una imagen de gran resolución (4K / 8K), el VAE se convierte claramente en un cuello de botella. Decodificar una imagen de este tamaño no solo lleva varios minutos, si no que además utiliza una cantidad de memoria desorbitada. No es raro acabar recibiendo el infame error: torch.cuda.OutOfMemoryError: CUDA out of memory.

Mediante esta optimización se divide el tensor en varias partes (como si fuesen baldosas), después se decodifican una a una y por último se vuelven a unir para formar la imagen. De esta manera el VAE no tiene que decodificar todo de golpe.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

pipe.enable_vae_tiling()

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
    height=4096,
    width=4096,
  ).images[0]

  image.save(f'image_{i}.png')

Pueden llegar a notarse las zonas donde se han unido estas partes debido a algunas diferencias en el color, pero no es común ni se percibe fácilmente.

Esta optimización es bastante sencilla de entender: si necesitas generar imágenes de muy alta resolución y tu tarjeta gráfica no tiene memoria suficiente, ésta será la única opción para lograrlo.

Cuándo utilizar: Nunca. Las imágenes de muy alta resolución contienen fallos porque Stable Diffusion no ha sido entrenado para esta tarea. Si necesitas aumentar la resolución utiliza un upscaler.

Tiny VAE

En el caso de Stable Diffusion XL se utiliza un VAE de 32 bits con 50M de parámetros. Ya que este componente se puede intercambiar vamos a utilizar un VAE llamado TAESD. Este pequeño modelo de apenas 1M de parámetros es una versión destilada del VAE original que, además, es capaz de funcionar en formato de 16 bits.

from diffusers import AutoPipelineForText2Image, AutoencoderTiny

# ...

vae = AutoencoderTiny.from_pretrained(
  'madebyollin/taesdxl',
  use_safetensors=True,
  torch_dtype=torch.float16,
).to('cuda')

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
  vae=vae,
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

¿Merece la pena sacrificar calidad de imagen para obtener más velocidad y menor uso de memoria? Veámoslo.

• Retrato
• Interior
• Macro
• 3D

• Base
• TAESD
• Base
• TAESD
• Base
• TAESD
• Base
• TAESD

La reducción en el uso de memoria es bastante espectacular gracias a ser un modelo más pequeño y capaz de funcionar en 16 bits.

La reducción del tiempo de inferencia es insignificante.

¿Y qué hay de la pérdida de calidad? Pues como puedes comprobar, no se aprecia. Es cierto que cambia un poco la imagen, sobre todo parece que añade algo más de contraste y cambian un poco las texturas, pero sinceramente yo no sabría distinguirlos.

Cuándo utilizar: Si necesitas reducir el uso de memoria, siempre. Con este modelo y sin utilizar ninguna otra optimización, ya se consigue ejecutar el proceso de inferencia en una tarjeta gráfica de 8 GB. Si no estás forzado a reducir el uso de memoria, pues tampoco sería mala idea utilizarlo porque no parece afectar negativamente.

Steps

Por defecto SDXL utiliza 50 pasos para limpiar un tensor lleno de ruido. A más pasos... más tiempo de inferencia, así que aquí hay margen de mejora. Mediante el parámetro num_inference_steps podemos especificar cuántos pasos queremos utilizar.

Vamos a generar imágenes con los siguientes pasos: 30, 25, 20 y 15. Utilizaremos el valor por defecto (50) como base para la comparativa.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    num_inference_steps=30,
    generator=generator,
  ).images[0]

  image.save(f'image_{i}.png')

Por supuesto... a menos pasos, menor tiempo de inferencia. Lo que nos interesa es quedarnos en el rango de pasos donde la calidad y la estructura se mantenga lo máximo posible. De nada sirve ahorrar mucho tiempo si vamos a obtener imágenes dignas del museo de los horrores. Vamos a ver dónde está el límite.

• Retrato
• Interior
• Macro
• 3D

• Base (50 pasos)
• 40 pasos
• 30 pasos
• 25 pasos
• 20 pasos
• 15 pasos
• Base (50 pasos)
• 40 pasos
• 30 pasos
• 25 pasos
• 20 pasos
• 15 pasos
• Base (50 pasos)
• 40 pasos
• 30 pasos
• 25 pasos
• 20 pasos
• 15 pasos
• Base (50 pasos)
• 40 pasos
• 30 pasos
• 25 pasos
• 20 pasos
• 15 pasos

Retrato: En 15 y 20 pasos la calidad no es tan mala, pero la estructura es distinta. A partir de 25 pasos me parece una imagen bastante correcta.

Interior: En 15 pasos aún no se obtiene la estructura deseada. En 20 pasos el resultado es bastante decente al tratarse de una imagen creativa, pero faltan algunos elementos. Así que aquí también considero que mínimo hay que hacer 25.

Macro: En la fotografía macro es sorprendente el nivel de detalle con tan solo 15 pasos. No sabría con cuál quedarme, todas son válidas y correctas.

3D: En una imagen estilo renderizado 3D hay demasiados defectos con pocos pasos, incluso se ve borrosa en ciertas zonas. Aunque la imagen en 30 pasos es decente, aquí me quedaría con el resultado tras 50 pasos (o puede que 40).

Así pues, pepende del estilo de imagen que estés generando puedes utilizar más o menos pasos, pero en general se obtiene bastante calidad con 25-30 pasos, lo que supone una reducción en tiempo de inferencia en torno al 40%, ¡lo cual es mucho!

Cuándo utilizar: Es una magnífica optimización cuando estás probando prompts o ajustando parámetros y quieres generar imágenes rápidamente. Cuando ya tienes todo ajustado, puedes ir en busca de la máxima calidad aumentando el número de pasos. Depende del caso de uso incluso también puede ser permanente.

Desactivar CFG

Como vimos en el artículo "Cómo funciona Stable Diffusion", la técnica de classifier-free guidance (CFG) se encarga de acercar o alejar al predictor de ruido de ciertas etiquetas.

Por ejemplo, si en el prompt positivo añadimos la palabra coche y en el prompt negativo añadimos la palabra juguete, el valor de CFG controla cuánto debe acercarse el predictor de ruido a las imágenes asociadas con el términocoche y cuánto debe alejarse del territorio donde se encuentran las imágenes asociadas con la palabra juguete. Es una manera muy efectiva de controlar el condicionamiento a la hora de generar imágenes.

Después, en el artículo "Cómo implementar Stable Diffusion", vimos como implementar la técnica de CFG y cómo introduce la necesidad de duplicar los tensores:

# Como estamos utilizando classifier-free guidance, duplicamos el tensor para evitar hacer dos pasadas
# Una pasada será por los valores condicionados y otra por los incondicionados
latent_model_input = torch.cat([latents] * 2)

Esto quiere decir que el predictor de ruido tarda el doble de tiempo en realizar cada paso.

Durante los primeros pasos es indispensable tener CFG activo para obtener una imagen con buena calidad y fiel a nuestro prompt. Una vez que el predictor de ruido va por buen camino, ¿necesitamos seguir utilizando CFG? En esta optimización vamos a explorar qué ocurre cuando dejamos de utilizar CFG durante el proceso.

El código es sencillo, creamos una función que se encarga de desactivar el CFG (pipe._guidance_scale = 0.0) cuando el número de pasos haya alcanzado el valor. Además, los tensores dejarán de estar duplicados a partir de este punto.

pipe = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
  if step_index == int(pipe.num_timesteps * 0.5):
    callback_kwargs['prompt_embeds'] = callback_kwargs['prompt_embeds'].chunk(2)[-1]
    callback_kwargs['add_text_embeds'] = callback_kwargs['add_text_embeds'].chunk(2)[-1]
    callback_kwargs['add_time_ids'] = callback_kwargs['add_time_ids'].chunk(2)[-1]
    pipe._guidance_scale = 0.0

  return callback_kwargs

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = pipe(
    prompt=generation['prompt'],
    generator=generator,
    callback_on_step_end=callback_dynamic_cfg,
    callback_on_step_end_tensor_inputs=['prompt_embeds', 'add_text_embeds', 'add_time_ids'],
  ).images[0]

  image.save(f'image_{i}.png')

Esta función se ejecuta al final de cada paso en forma de callback, gracias al parámetro callback_on_step_end. También hay que especificar los tensores que vamos a modificar dentro del callback mediante el parámetro callback_on_step_end_tensor_inputs.

Vamos a explorar qué ocurre cuando dejamos de utilizar CFG en el último cuarto (75%) y en la segunda mitad del proceso (50%).

• Retrato
• Interior
• Macro
• 3D

• Base
• Disable CFG 75%
• Disable CFG 50%
• Base
• Disable CFG 75%
• Disable CFG 50%
• Base
• Disable CFG 75%
• Disable CFG 50%
• Base
• Disable CFG 75%
• Disable CFG 50%

Como era de esperar, desactivar CFG al 50% produce una reducción del 25% en el tiempo de inferencia de cada imagen (no del total, porque la carga de modelos también cuenta). Esto se debe a que si realizamos 50 pasos con CFG activo, en realidad el modelo limpia 100 tensores. En cambio, utilizando esta optimización, el modelo limpia 50 tensores en los primeros 25 pasos y la mitad (25 tensores) en los últimos 25 pasos. Así que... 75/100 equivale a saltarse un 25% del trabajo. En el caso de desactivar CFG al 75%, la reducción es del 12.5% en cada imagen.

En cuanto a la calidad del resultado, parece que baja un poco pero no demasiado. Esto también puede deberse a no haber utilizado un prompt negativo, que es la principal ventaja de utilizar CFG. Al utilizar mejores prompts seguro que aumenta la calidad. Al 75% es prácticamente imperceptible.

Cuándo utilizar: Agresivamente, cuando quieras generar imágenes de manera más rápida y no te importe perder una pizca de calidad (por ejemplo, para probar prompts o parámetros). Al desactivar CFG un poco más tarde, la velocidad aumenta sin sacrificar calidad.

Refiner

¿Y qué hay del modelo refiner? En todo momento hemos optimizado el modelo base, pero una de las principales ventajas de Stable Diffusion XL es que cuenta con un modelo especializado que refina los pequeños detalles finales. Este modelo incrementa notablemente la calidad del resultado.

Por defecto, el modelo base utiliza 11.24 GB de memoria. Al utilizar también el modelo refiner, los requerimientos de memoria suben hasta 17.38 GB. Pero recuerda, la mayoría de optimizaciones también se pueden aplicar a este modelo ya que cuenta con los mismos componentes (excepto el primer text encoder).

for generation in queue:
  image = base(generation['prompt'], output_type='latent').images
  refiner(generation['prompt'], image=image)

El modelo refiner se puede utilizar de dos maneras distintas, así que vamos a verlas individualmente.

Ensemble of Expert Denoisers

El método Ensemble of Expert Denoisers se refiere al enfoque en el cual comienza la generación de una imagen con el modelo base y termina mediante el modelo refiner. Durante todo el proceso no se genera ninguna imagen, si no que el modelo base limpia el tensor durante una cantidad específica de pasos (un porcentaje del total), y después le entrega el tensor al modelo refiner para que concluya el trabajo.

Se podría decir que trabajan de manera conjunta para generar el resultado (base+refiner).

En cuanto al código, el modelo base detiene su trabajo al 80% del proceso mediante el parámetro denoising_end=0.8 y devuelve el tensor gracias a output_type='latent'.

El modelo refiner recibe este tensor mediante el parámetro image (es irónico, pero no es una imagen). Después empieza a limpiarlo asumiendo que ya se ha realizado un 80% del trabajo, indicado por el parámetro denoising_start=0.8. También especificamos nuevamente cuántos pasos tiene el proceso en total (num_inference_steps) para que calcule cuántos pasos le quedan por limpiar. Es decir, si utilizamos 50 pasos con un cambio al 80%, el modelo base limpiará el tensor durante 40 pasos y el modelo refiner durante los últimos 10, refinando los detalles que falten por pulir.

• Python

from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image

# ...

base = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

refiner = AutoPipelineForImage2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-refiner-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = base(
    prompt=generation['prompt'],
    generator=generator,
    num_inference_steps=50,
    denoising_end=0.8,
    output_type='latent',
  ).images # Recuerda que aquí no accedemos a images[0], sino el tensor completo

  image = refiner(
    prompt=generation['prompt'],
    generator=generator,
    num_inference_steps=50,
    denoising_start=0.8,
    image=image,
  ).images[0]

  image.save(f'image_{i}.png')

Vamos a generar imágenes en 50, 40, 30 y 20 pasos, cambiando al modelo refiner en 0.9 y 0.8.

Para tener una referencia, también se incluirá la imagen que se ha utilizado como base para todas las comparativas (solo modelo base, 50 pasos).

🏆 Resultados 🏆

• Retrato
• Interior
• Macro
• 3D

• Base
• 50 pasos, 0.9
• 50 pasos, 0.8
• 40 pasos, 0.9
• 40 pasos, 0.8
• 30 pasos, 0.9
• 30 pasos, 0.8
• 20 pasos, 0.9
• 20 pasos, 0.8
• Base
• 50 pasos, 0.9
• 50 pasos, 0.8
• 40 pasos, 0.9
• 40 pasos, 0.8
• 30 pasos, 0.9
• 30 pasos, 0.8
• 20 pasos, 0.9
• 20 pasos, 0.8
• Base
• 50 pasos, 0.9
• 50 pasos, 0.8
• 40 pasos, 0.9
• 40 pasos, 0.8
• 30 pasos, 0.9
• 30 pasos, 0.8
• 20 pasos, 0.9
• 20 pasos, 0.8
• Base
• 50 pasos, 0.9
• 50 pasos, 0.8
• 40 pasos, 0.9
• 40 pasos, 0.8
• 30 pasos, 0.9
• 30 pasos, 0.8
• 20 pasos, 0.9
• 20 pasos, 0.8

---

	Tiempo inferencia
Base	14.1s
50 pasos, 0.9	14.3s +1.42%
50 pasos, 0.8	14.5s +2.84%
40 pasos, 0.9	11.4s -19.15%
40 pasos, 0.8	11.5s -18.44%
30 pasos, 0.9	9s -36.17%
30 pasos, 0.8	9.1s -35.46%
20 pasos, 0.9	6.2s -56.03%
20 pasos, 0.8	6.3s -55.32%

Veredicto ⚖️

Sin duda utilizar el modelo refiner mejora mucho el resultado.

¿Cuándo es mejor que tome el relevo? Claramente se puede ver como los resultados en 0.9 son mejores que en 0.8, lo que tiene sentido al tratarse de un modelo que refina los últimos detalles y no debería utilizarse para actuar sobre la estructura de la imagen.

En cuanto a la cantidad de pasos, mi percepción es que el modelo es capaz de ofrecer un resultado de altísima calidad visual sea cuál sea el número de pasos. Lo único que parece cambiar es la estructura/composición de la imagen, pero calidad visual es alta con tan solo 30 pasos.

Y sin que sea menos importante, también hay que tener en cuenta la considerable reducción de tiempo al bajar de 40 pasos.

Cuándo utilizar: Siempre que queramos utilizar el modelo refiner para aumentar la calidad visual de la imagen. En cuánto a los parámetros, podríamos utilizar 30 o 40 pasos siempre que no busquemos la mejor calidad posible. Por supuesto, siempre haciendo el cambio en 0.9.

Image-to-image

El clásico img2img no es novedad en Stable Diffusion XL. Se trata del método en el cual se produce una imagen completa con el modelo base, para después entregarle al modelo refiner tanto esa imagen como el prompt original, para que genere una nueva imagen con estos condicionamientos.

En otras palabras, en img2img los modelos trabajan independientemente (base->refiner).

Al ser dos procesos independientes es un poco más sencillo aplicar las optimizaciones de este artículo. Aún así, el código no es muy distinto, simplemente se genera una imagen y se utiliza como parámetro en el modelo refiner.

• Python

from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image

# ...

base = AutoPipelineForText2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-base-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

refiner = AutoPipelineForImage2Image.from_pretrained(
  'stabilityai/stable-diffusion-xl-refiner-1.0',
  use_safetensors=True,
  torch_dtype=torch.float16,
  variant='fp16',
).to('cuda')

generator = torch.Generator(device='cuda')

for i, generation in enumerate(queue, start=1):
  generator.manual_seed(generation['seed'])

  image = base(
    prompt=generation['prompt'],
    generator=generator,
    num_inference_steps=50,
  ).images[0]

  image = refiner(
    prompt=generation['prompt'],
    generator=generator,
    num_inference_steps=10,
    image=image,
  ).images[0]

  image.save(f'image_{i}.png')

Vamos a generar imágenes con el modelo base en 50, 40, 30 y 20 pasos, para añadir después una combinación de 20 y 10 pasos extra mediante el modelo refiner.

Para tener una referencia, también se incluirá la imagen que se ha utilizado como base para todas las comparativas (solo modelo base, 50 pasos).

🏆 Resultados 🏆

• Retrato
• Interior
• Macro
• 3D

• Base
• 50 base + 20 refiner
• 50 base + 10 refiner
• 40 base + 20 refiner
• 40 base + 10 refiner
• 30 base + 20 refiner
• 30 base + 10 refiner
• 20 base + 20 refiner
• 20 base + 10 refiner
• Base
• 50 base + 20 refiner
• 50 base + 10 refiner
• 40 base + 20 refiner
• 40 base + 10 refiner
• 30 base + 20 refiner
• 30 base + 10 refiner
• 20 base + 20 refiner
• 20 base + 10 refiner
• Base
• 50 base + 20 refiner
• 50 base + 10 refiner
• 40 base + 20 refiner
• 40 base + 10 refiner
• 30 base + 20 refiner
• 30 base + 10 refiner
• 20 base + 20 refiner
• 20 base + 10 refiner
• Base
• 50 base + 20 refiner
• 50 base + 10 refiner
• 40 base + 20 refiner
• 40 base + 10 refiner
• 30 base + 20 refiner
• 30 base + 10 refiner
• 20 base + 20 refiner
• 20 base + 10 refiner

---

	Tiempo inferencia
Base	14.1s
50 base + 20 refiner	16.8s +19.15%
50 base + 10 refiner	15.6s +10.64%
40 base + 20 refiner	14.1s
40 base + 10 refiner	12.9s -8.51%
30 base + 20 refiner	11.5s -18.44%
30 base + 10 refiner	10.5s -25.53%
20 base + 20 refiner	8.9s -36.88%
20 base + 10 refiner	7.9s -43.97%

Veredicto ⚖️

En modo img2img el modelo refiner no funciona tan bien.

Cuando utilizamos suficientes pasos con el modelo base, parece como el modelo refiner estuviera forzado a añadir detalles en algo que no lo necesita. Como decimos en nuestro mundo: si algo funciona, no lo toques.

En cambio, si utilizamos pocos pasos en el modelo base el resultado es algo mejor. Esto sucede porque con tan pocos pasos, el modelo base no es capaz de añadir los pequeños detalles y le da algo más de margen de trabajo al modelo refiner.

Aquí también hay que tener en cuenta la reducción de tiempo al disminuir los pasos. Si utilizamos muchos pasos la penalización es importante.

Cuándo utilizar: Lo primero es tener en cuenta que cuando utilizamos el modelo refiner es para exprimir al máximo la calidad visual. En este caso de uso conviene incrementar el número de pasos y por lo tanto el método "Ensemble of Expert Denoisers" es la mejor opción. Con pocos pasos tampoco creo que haya mejor calidad visual, ni tampoco aumenta la velocidad de generación con respecto al otro método. Por lo tanto, utilizar el modelo refiner en modo img2img queda en tierra de nadie, tiene sus ventajas pero no destaca en nada.