Stable Diffusion 3 Medium TensorRT

Introduction

This repository hosts the TensorRT version of Stable Diffusion 3 Medium created in collaboration with NVIDIA. The optimized versions give substantial improvements in speed and efficiency.

Stable Diffusion 3 Medium is a fast generative text-to-image model with greatly improved performance in multi-subject prompts, image quality, and spelling abilities.

Model Details

Model Description

Stable Diffusion 3 Medium combines a diffusion transformer architecture and flow matching.

Developed by: Stability AI
Model type: MMDiT text-to-image model
Model Description: This is a conversion of the Stable Diffusion 3 Medium model

Performance using TensorRT 10.1

Timings for 50 steps at 1024x1024

Accelerator	CLIP-G	CLIP-L	T5XXL	MMDiT	VAE Decoder	Total
A100	11.95 ms	5.04 ms	21.39 ms	5468.17 ms	72.25 ms	5622.47 ms

Timings for 30 steps at 1024x1024 with input image conditioning

Accelerator	VAE Encoder	CLIP-G	CLIP-L	T5XXL	MMDiT	VAE Decoder	Total
A100	37.04 ms	12.07 ms	5.07 ms	21.49 ms	3340.69 ms	72.02 ms	3531.49 ms

Int8 quantization with TensorRT Model Optimizer

The MMDiT in Stable Diffusion 3 Medium can be further optimized with INT8 quantization using TensorRT Model Optimizer. The estimated end-to-end speedup comparing TensorRT fp16 and TensorRT int8 is 1.2x~1.4x on various NVidia GPUs. The memory saving is about 2x for the int8 MMDiT engine compared with the fp16 counterpart. The image quality can be maintained with minimal to negligible degradation.

Usage Example

Follow the setup instructions on launching a TensorRT NGC container.

git clone https://github.com/NVIDIA/TensorRT.git
cd TensorRT
git checkout release/sd3
docker run --rm -it --gpus all -v $PWD:/workspace nvcr.io/nvidia/pytorch:24.05-py3 /bin/bash

Download the Stable Diffusion 3 Medium TensorRT files from this repo

git lfs install 
git clone https://huggingface.co./stabilityai/stable-diffusion-3-medium-tensorrt
cd stable-diffusion-3-medium-tensorrt
git lfs pull
cd ..

Install libraries and requirements

cd demo/Diffusion
python3 -m pip install --upgrade pip
pip3 install -r requirements.txt
python3 -m pip install --pre --upgrade --extra-index-url https://pypi.nvidia.com tensorrt-cu12

Perform TensorRT optimized inference:

Stable Diffusion 3 Medium

Works best for 1024x1024 images. The first invocation produces plan files in --engine-dir specific to the accelerator being run on and are reused for later invocations.

python3 demo_txt2img_sd3.py \
  ""Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"" \
  --version=sd3 \
  --onnx-dir /workspace/sd3-tensorrt/ \
  --engine-dir /workspace/sd3-tensorrt/engine \
  --seed 42 \
  --width 1024 \
  --height 1024 \
  --build-static-batch \
  --use-cuda-graph

Stable Diffusion 3 Medium with input image conditioning

Provide an input image conditioning using below. Works best for 1024x1024 but may also work at 512x512.

wget https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png -O dog-on-bench.png

python3 demo_txt2img_sd3.py \
  """dog wearing a sweater and a blue collar""" \
  --version=sd3 \
  --onnx-dir /workspace/sd3-tensorrt/ \
  --engine-dir /workspace/sd3-tensorrt/engine \
  --seed 42 \
  --width 1024 \
  --height 1024 \
  --input-image dog-on-bench.png \
  --build-static-batch \
  --use-cuda-graph