Stable Diffusion 3 ONNX support

[gmarcosf]

Feature request

Support Stable Diffusion 3's pipeline ONNX exportation.

• Paper: https://arxiv.org/pdf/2403.03206
• Repository: https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers

Motivation

Stable Diffusion 3 is a new SOTA SD model which greatly improves its previous versions. It uses a new backbone for the image denoising via a Transformer architecture (namely MM-DiT) rather than the U-Net, used in previous SD pipelines. It also uses T5 for text encoding. Thus, we currently cannot convert SD3 to ONNX because we need a new pipeline called ORTStableDiffusion3Pipeline.

We cannot use the current ORTStableDiffusionPipeline with SD3 because of the new architecture. If we do either:

from optimum.onnxruntime import ORTStableDiffusionPipeline

# The following is not yet implemented:
# from optimum.onnxruntime import ORTStableDiffusion3Pipeline
model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
pipeline = ORTStableDiffusionPipeline.from_pretrained(model_id, export=True)
# The following is not yet implemented:
# pipeline = ORTStableDiffusion3Pipeline.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-3-medium-diffusers")

or via CLI:

optimum-cli export onnx --model stabilityai/stable-diffusion-3-medium-diffusers sd_v3_onnx/

We get the following error (expected as SD3 doesn't use U-Net anymore):

AttributeError: 'StableDiffusion3Pipeline' object has no attribute 'unet'

Your contribution

I can get started implementing such pipeline ASAP.

[IlyasMoutawwakil]

Hi ! this will be available once #2073 is merged 🤗

[IlyasMoutawwakil]

One small detail tho, ORTStableDiffusionPipeline will not work with this architecture the same way diffusers' StableDiffusionPipeline won't. We implement a pipeline specific for each architecture the same way in diffusers, so you'll have to use ORTStableDiffusion3Pipeline. But the best way would be to use task-specific auto-pipelines ; ORTPipelineForText2Image (zquivalent to diffusers' AutoPipelineForText2Image).

[gmarcosf]

Hi Ilyas, thanks for your clear response! I'm closing the issue - look forward the merge 🤗!

[IlyasMoutawwakil]

merged ! 🤗
please try it out extensively 🔥