Accelerating SD Turbo and SDXL Turbo Inference with ONNX Runtime and Olive

TL;DR

Hugging Face and Microsoft have optimized SD Turbo and SDXL Turbo inference using ONNX Runtime and the Olive optimization tool, resulting in throughput gains of up to 229% for SDXL Turbo and 120% for SD Turbo over PyTorch. These optimizations enable high-performance image generation in as little as one step and expand model accessibility to non-Python languages like C# and Java.

Performance Benchmarks: ONNX Runtime vs. PyTorch

ONNX Runtime significantly outperforms PyTorch across all tested batch sizes and step counts for both SD Turbo and SDXL Turbo models on NVIDIA GPUs.

Key Throughput Gains

  • SDXL Turbo: Throughput gains reached as high as 229% compared to PyTorch.
  • SD Turbo: Throughput gains reached as high as 120% compared to PyTorch.
  • Execution Providers: Both CUDA and TensorRT execution providers showed marked improvements over PyTorch for both static and dynamic shapes.

Hardware and Configuration Results

Benchmarks were conducted using A100-SXM4-80GB and RTX-4090 GPUs with the LCM Scheduler and fp16 models.

  • Static vs. Dynamic Shapes: Static shapes generally offer faster performance when batch and image sizes are known at graph definition time. Dynamic shapes are more flexible, allowing users to change batch and image sizes during execution without rebuilding the engine.
  • GPU Specifics: On A100 GPUs, ONNX Runtime with the CUDA execution provider is often the better choice for dynamic shapes. On RTX-4090 GPUs, the TensorRT execution provider typically performs slightly better for dynamic shapes.

Technical Optimizations and Tooling

The performance increases are driven by the Olive model optimization tool and specific GPU-level enhancements within ONNX Runtime.

The Olive Optimization Tool

Models are generated using Olive, a hardware-aware model optimization tool. For maximum performance, fp16 VAE must be enabled via the command line.

GPU-Specific Enhancements

Beyond standard Stable Diffusion optimizations, the following specific technical improvements were implemented:

  • CUDA Graph: Enabled for static shape inputs to reduce overhead.
  • Flash Attention V2: Integrated to accelerate attention mechanisms.
  • Text Encoder Optimization: Extra outputs in the text encoder were removed, retaining only the hidden state output specified by the clip_skip parameter.
  • SkipGroupNorm Fusion: Group normalization was fused with preceding Add nodes.
  • LoRA Support: Added support for LoRA weights for latent consistency models (LCMs).

Cross-Platform Accessibility and Integration

By utilizing ONNX Runtime, SD Turbo and SDXL Turbo are no longer restricted to Python environments, enabling integration into a wider variety of software stacks.

Language Support

  • C#: Community projects like OnnxStack provide a .NET library for Stable Diffusion inference.
  • Java: Oracle has released a Stable Diffusion sample (sd4j) that runs inference on top of ONNX Runtime.

Web UI Integration

An ONNX Runtime Extension for Automatic1111’s SD WebUI allows for optimized execution of the Stable Diffusion UNet model on NVIDIA GPUs using the CUDA execution provider and Olive-optimized models.

Future Roadmap

Hugging Face and Microsoft plan to extend these optimizations to additional features and models, including:

  • Feature Support: Integration of IP Adapter and ControlNet.
  • Model Support: Expansion to Stable Video Diffusion.
  • UI Improvements: Further optimization of the SD Turbo and SDXL Turbo performance within the existing Stable Diffusion web UI extension and support for a community-developed Windows UI.

Sources