ONNX Runtime

🤗 Optimum fornece um pipeline de Stable Diffusion compatível com ONNX Runtime. Você precisará instalar o 🤗 Optimum com o seguinte comando para suporte ao ONNX Runtime:

pip install -q optimum["onnxruntime"]

Este guia mostrará como usar os pipelines Stable Diffusion e Stable Diffusion XL (SDXL) com o ONNX Runtime.

Stable Diffusion

Para carregar e executar inferência, use o ORTStableDiffusionPipeline. Se você quiser carregar um modelo do PyTorch e convertê-lo para o formato ONNX em tempo de execução, defina 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")

Gerar múltiplos prompts em um lote parece consumir muita memória. Enquanto investigamos o problema, você pode precisar iterar em vez de usar lotes.

Para exportar o pipeline no formato ONNX offline e usá-lo posteriormente para inferência, use o comando optimum-cli export:

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

Em seguida, para executar a inferência (não é necessário especificar export=True novamente):

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]

Você pode encontrar mais exemplos na documentação do 🤗 Optimum, e o Stable Diffusion é suportado para texto-para-imagem, imagem-para-imagem e inpainting.

Stable Diffusion XL

Para carregar e executar inferência com SDXL, use o 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 o pipeline no formato ONNX e usá-lo posteriormente para inferência, use o comando optimum-cli export:

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

O SDXL no formato ONNX é suportado para texto-para-imagem e imagem-para-imagem.