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

Timings for 30 steps at 1024x1024 with input image conditioning

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

1. 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

2. 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 ..

3. 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

4. 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/stable-diffusion-3-medium-tensorrt/ \
    --engine-dir /workspace/stable-diffusion-3-medium-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/stable-diffusion-3-medium-tensorrt/ \
    --engine-dir /workspace/stable-diffusion-3-medium-tensorrt/engine \
    --seed 42 \
    --width 1024 \
    --height 1024 \
    --input-image dog-on-bench.png \
    --build-static-batch \
    --use-cuda-graph