ONNX Runtime

🤗 Optimum ofrece un pipeline de Stable Diffusion compatible con ONNX Runtime. Deberás instalar 🤗 Optimum con el siguiente comando para obtener soporte de ONNX Runtime:

pip install -q optimum["onnxruntime"]

Esta guía te mostrará cómo usar los pipelines de Stable Diffusion y Stable Diffusion XL (SDXL) con ONNX Runtime.

Stable Diffusion

Para cargar y ejecutar la inferencia, use ORTStableDiffusionPipeline. Si desea cargar un modelo de PyTorch y convertirlo al formato ONNX sobre la marcha, establezca export=True:

from optimum.onnxruntime import ORTStableDiffusionPipeline

model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
pipeline = ORTStableDiffusionPipeline.from_pretrained(model_id, export=True)
prompt = "sailing ship in storm by Leonardo da Vinci"
image = pipeline(prompt).images[0]
pipeline.save_pretrained("./onnx-stable-diffusion-v1-5")

Generar varios prompts en un lote parece consumir demasiada memoria. Mientras investigamos el problema, es posible que necesites iterar en lugar de usar lotes.

Para exportar el pipeline en formato ONNX de forma offline y usarlo más adelante para inferencia, utilice el comando optimum-cli export:

optimum-cli export onnx --model stable-diffusion-v1-5/stable-diffusion-v1-5 sd_v15_onnx/

Luego, para realizar la inferencia (no es necesario especificar export=True de nuevo):

from optimum.onnxruntime import ORTStableDiffusionPipeline

model_id = "sd_v15_onnx"
pipeline = ORTStableDiffusionPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Leonardo da Vinci"
image = pipeline(prompt).images[0]

Puedes encontrar más ejemplos en la documentación de 🤗 Optimum, y Stable Diffusion es compatible con texto a imagen, imagen a imagen e inpainting.

Stable Diffusion XL

Para cargar y ejecutar inferencia con SDXL, use ORTStableDiffusionXLPipeline:

from optimum.onnxruntime import ORTStableDiffusionXLPipeline

model_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipeline = ORTStableDiffusionXLPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Leonardo da Vinci"
image = pipeline(prompt).images[0]

Para exportar el pipeline en el formato ONNX y usarlo más adelante para inferencia, use el comando optimum-cli export:

optimum-cli export onnx --model stabilityai/stable-diffusion-xl-base-1.0 --task stable-diffusion-xl sd_xl_onnx/

SDXL en el formato ONNX es compatible con texto a imagen e imagen a imagen.