Update optimize_pipeline for SDXL

[tianleiwu]

Description

• Optimize SDXL models exported by optimum.
• Enable it to run locally instead of using module.
• Detect external data file in original model, and save with same format by default.
• Add tests

Example

pip install optimum transformers diffusers onnx onnxruntime-gpu>=1.16
optimum-cli export onnx --model stabilityai/stable-diffusion-xl-base-1.0 --task stable-diffusion-xl ./sd_xl_base_onnx
python -m  onnxruntime.transformers.models.stable_diffusion.optimize_pipeline -i ./sd_xl_base_onnx -o ./sd_xl_base_fp16 --float16

Known issues

(1) VAE decoder cannot be converted to float16. Otherwise, there will be black image in output.
(2) To use the float16 models, need a minor change in optimum to convert the inputs for VAE decoder from float16 to float32 since we keep VAE decoder as float32. The change is to append a line like the following after this line

latents = latents.astype(np.float32)

Motivation and Context

[SuperSecureHuman]

Just to maintain consistency, requirements of tenssort should match the requirement defined in other providers

Tensorrt --> diffusers>=0.16.0
Others --> diffusers>=0.15.1

Might as well keep it uniform and require >=0.16.0 for all the backends.

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements-tensorrt.txt#L1

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt#L1

[SuperSecureHuman]

Given the standalone capablity of this optimization, could even make a dockerised variant.

I can help with it, if its needed.

[tianleiwu]

Just to maintain consistency, requirements of tenssort should match the requirement defined in other providers

You are right. I will add support of SD-XL for TensorRT EP in another pull request. Current onnxruntime_tensorrt_txt2img.py does not support SD-XL. It need some change (similar to demo diffusion of TensorRT as in https://github.com/rajeevsrao/TensorRT/tree/release/8.6/demo/Diffusion)

[kousun12]

@tianleiwu I'm having trouble getting the basic version of this to run:

pip install optimum transformers diffusers onnx onnxruntime-gpu>=1.16
optimum-cli export onnx --model stabilityai/stable-diffusion-xl-base-1.0 --task stable-diffusion-xl ./sd_xl_base_onnx
python -m  onnxruntime.transformers.models.stable_diffusion.optimize_pipeline -i ./sd_xl_base_onnx -o ./sd_xl_base_fp16 --float16

works fine to create an optimized version, but when i go to load it in, it seems to look for a TensorRT op despite using a CudaExecutionProvider:

from optimum.onnxruntime import (
    ORTStableDiffusionXLPipeline,
    ORTStableDiffusionXLImg2ImgPipeline,
)

MODEL_DIR = '/sdxl_onnx'

pipe = ORTStableDiffusionXLPipeline.from_pretrained(
    MODEL_DIR,
    provider="CUDAExecutionProvider",
)
assert pipe.providers == ["CUDAExecutionProvider", "CPUExecutionProvider"]

it all loads fine, but running results in an error

im = pipe("vase of flowers")

2023-09-24 16:15:31.206037898 [E:onnxruntime:, sequential_executor.cc:514 ExecuteKernel] Non-zero status code returned while running Attention node. Name:'Attention_0' Status Message: /onnxruntime_src/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu:207 bool onnxruntime::contrib::cuda::FusedMHARunnerFP16v2::mhaImpl::is_flash_attention(int) const interface->mHasCausalMask == false was false. 

---------------------------------------------------------------------------
RuntimeException                          Traceback (most recent call last)
Cell In[4], line 1
----> 1 im = pipe("vase of flowers")

Any idea what is wrong here?

[SuperSecureHuman]

Can you check if it works for SD2.1/1.4

[tianleiwu]

@kousun12, the error is caused by a bug in Attention operator, see #17729.
You can build from source with that fix.

One walkaround is to set an environment variable ORT_ENABLE_FUSED_CAUSAL_ATTENTION=1. But it might bring slightly different result since the fused causal attention used fp16 in accumulation.

Another way to walkaround is to disable attention fusion in CLIP model. That might be slower but could avoid the runtime error.

[kousun12]

@tianleiwu - ok that brings me a little closer but now I am seeing an issue that the vae was not converted to fp16 with this pipeline. seems to be because the optimize function skips it (https://github.com/microsoft/onnxruntime/pull/17536/files#diff-59f8bc82830a0218e203a2e6b1f0bfe17bac0063bc0df70dcbf1342e47c6c67eR151)

im = pipe("vase of flowers")

100%
50/50 [00:06<00:00, 10.27it/s]
---------------------------------------------------------------------------
InvalidArgument                           Traceback (most recent call last)
Cell In[7], line 1
----> 1 im = pipe("vase of flowers")

File /usr/local/lib/python3.10/site-packages/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl.py:486, in StableDiffusionXLPipelineMixin.__call__(self, prompt, height, width, num_inference_steps, guidance_scale, negative_prompt, num_images_per_prompt, eta, generator, latents, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, output_type, return_dict, callback, callback_steps, cross_attention_kwargs, guidance_rescale, original_size, crops_coords_top_left, target_size)
    483 latents /= self.vae_decoder.config.get("scaling_factor", 0.18215)
    484 # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
    485 image = np.concatenate(
--> 486     [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]
    487 )
    488 # apply watermark if available
    489 if self.watermark is not None:

File /usr/local/lib/python3.10/site-packages/optimum/pipelines/diffusers/pipeline_stable_diffusion_xl.py:486, in <listcomp>(.0)
    483 latents /= self.vae_decoder.config.get("scaling_factor", 0.18215)
    484 # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
    485 image = np.concatenate(
--> 486     [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])]
    487 )
    488 # apply watermark if available
    489 if self.watermark is not None:

File /usr/local/lib/python3.10/site-packages/optimum/onnxruntime/modeling_diffusion.py:481, in _ORTDiffusionModelPart.__call__(self, *args, **kwargs)
    480 def __call__(self, *args, **kwargs):
--> 481     return self.forward(*args, **kwargs)

File /usr/local/lib/python3.10/site-packages/optimum/onnxruntime/modeling_diffusion.py:525, in ORTModelVaeDecoder.forward(self, latent_sample)
    521 def forward(self, latent_sample: np.ndarray):
    522     onnx_inputs = {
    523         "latent_sample": latent_sample,
    524     }
--> 525     outputs = self.session.run(None, onnx_inputs)
    526     return outputs

File /usr/local/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:220, in Session.run(self, output_names, input_feed, run_options)
    218     output_names = [output.name for output in self._outputs_meta]
    219 try:
--> 220     return self._sess.run(output_names, input_feed, run_options)
    221 except C.EPFail as err:
    222     if self._enable_fallback:

InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(float16)) , expected: (tensor(float))

[tianleiwu]

@kousun12, see the Known issues in the description of this PR. If VAE decoder is converted to fp16, there will be black image in output due to overflow. I noticed that TensorRT also does not convert VAE for SDXL, see https://github.com/NVIDIA/TensorRT/blob/109862c8fa19e14058de085d64cbb0b580ece8d4/demo/Diffusion/stable_diffusion_pipeline.py#L142-L144.

To walkaround, you can change optimum to append a line like the following after this line

latents = latents.astype(np.float32)

I am going to add another pull request to add an option to let VAE decoder ONNX model accept fp16 latent inputs, then there is no need to change optimum after that.

[SuperSecureHuman]

Exported using the SDXL export instructions from docs

I find that onnxruntime fp16 takes more time than diffusers version.

Resolution 1024 x 1024

Diffusers - 38sec for 50 steps
Optimum Onnxruntime - 44 sec for 50 steps

While for 512, it remained the same 9 sec

Experiment on T4 (Colab)

EDIT - Cuda runtime

[tianleiwu]

@SuperSecureHuman, thanks for sharing the results.

Optimum OnnxRuntime for SDXL is not fully optimized, like I/O Binding is not used. That could slow the pipeline since model output is copied to CPU, then copied back to GPU for next step.

Right now, I am working on a demo pipeline (similar to SDXL with TensorRT EP in pull request) that could use IO Binding for cuda EP. After that, I could revisit optimum performance issue.

[SuperSecureHuman]

@tianleiwu

Thanks for the update. Do let me know if I could help in testing.

Also, would love to see the work you are doing?

Are u working on a branch or a fork for this?

[tianleiwu]

@SuperSecureHuman, it is still in progress (Need some change in ArgMax, BiasAdd, BiasSplitGelu, GroupNorm cuda kernels to support SD XL Refiner; And make change in VAE part to work with FP16). Hopefully, I can send out a pull request next week.

For non XL version, you can try the demo_txt2img.py for CUDA EP.

[SuperSecureHuman]

SD 2.1 works perfectly. Got inference time from 9.2 sec to 5.9 Sec. I would want to see similar results with SDXL. This would make SDXL viable for deployment in low resource senario.