search

stochasticai / x-stable-diffusion

Real-time inference for Stable Diffusion - 0.88s latency. Covers AITemplate, nvFuser, TensorRT, FlashAttention. Join our Discord communty: https://discord.com/invite/TgHXuSJEk6
https://stochastic.ai
Apache License 2.0
553 stars 35 forks source link
aitemplate automl cuda docker inference notebook nvfuser onnx onnxruntime pytorch stable-diffusion tensorrt

readme

Stochastic.ai Stochastic.ai


Welcome to x-stable-diffusion by Stochastic!

This project is a compilation of acceleration techniques for the Stable Diffusion model to help you generate images faster and more efficiently, saving you both time and money.

With example images and a comprehensive benchmark, you can easily choose the best technique for your needs. When you're ready to deploy, our CLI called stochasticx makes it easy to get started on your local machine. Try x-stable-diffusion and see the difference it can make for your image generation performance and cost savings.

🚀 Installation

Quickstart

Make sure you have Python and Docker installed on your system

  1. Install the latest version of stochasticx library.

    pip install stochasticx

  2. Deploy the Stable Diffusion model

    stochasticx stable-diffusion deploy --type aitemplate

Alternatively, you can deploy stable diffusion without our CLI by checking the steps here.

  1. To perform inference with this deployed model:

    stochasticx stable-diffusion inference --prompt "Riding a horse"

    Check all the options of the inference command:

    stochasticx stable-diffusion inference --help

  2. You can get the logs of the deployment executing the following command:

    stochasticx stable-diffusion logs

  3. Stop and remove the deployment with this command:

    stochasticx stable-diffusion stop

How to get less than 1s latency?

Change the num_inference_steps to 30. With this, you can get an image generated in 0.88 seconds. 

{
  'max_seq_length': 64,
  'num_inference_steps': 30, 
  'image_size': (512, 512) 
}

You can also experiment with reducing the image_size.

How to run on Google Colab?

In each folder, we will provide a Google Colab notebook with which you can test the full flow and inference on a T4 GPU

Manual deployment

Check the README.md of the following directories:

🔥 Optimizations

Benchmarks

Setup

For hardware, we used 1x40GB A100 GPU with CUDA 11.6 and the results are reported by averaging 50 runs.

The following arguments were used for image generation for all the benchmarks:

{
  'max_seq_length': 64,
  'num_inference_steps': 50, 
  'image_size': (512, 512) 
}

Online results

For batch_size 1, these are the latency results:

A100 GPU

A100_GPU_graph

project Latency (s) GPU VRAM (GB)
PyTorch fp16 5.77 10.3
nvFuser fp16 3.15 ---
FlashAttention fp16 2.80 7.5
TensorRT fp16 1.68 8.1
AITemplate fp16 1.38 4.83
ONNX (CUDA) 7.26 13.3

T4 GPU

Note: AITemplate might not support T4 GPU yet. Check support here

T4_GPU_graph

project Latency (s)
PyTorch fp16 16.2
nvFuser fp16 19.3
FlashAttention fp16 13.7
TensorRT fp16 9.3

Batched results - A100 GPU

The following results were obtained by varying batch_size from 1 to 24.

A100_GPU_batch_size

project \ bs 1 4 8 16 24
Pytorch fp16 5.77s/10.3GB 19.2s/18.5GB 36s/26.7GB OOM
FlashAttention fp16 2.80s/7.5GB 9.1s/17GB 17.7s/29.5GB OOM
TensorRT fp16 1.68s/8.1GB OOM
AITemplate fp16 1.38s/4.83GB 4.25s/8.5GB 7.4s/14.5GB 15.7s/25GB 23.4s/36GB
ONNX (CUDA) 7.26s/13.3GB OOM OOM OOM OOM

Note: TensorRT fails to convert UNet model from ONNX to TensorRT due to memory issues.

Sample images generated

Click here to view the complete list of generated images

Optimization \ Prompt Super Mario learning to fly in an airport, Painting by Leonardo Da Vinci The Easter bunny riding a motorcycle in New York City Drone flythrough of a tropical jungle convered in snow
PyTorch fp16 pytorch_stable-diffusion_mario pytorch_stable-diffusion_bunny pytorch_stable-diffusion_bunny
nvFuser fp16 nvFuser_stable-diffusion_mario nvFuser_stable-diffusion_bunny nvFuser_stable-diffusion_bunny
FlashAttention fp16 FlashAttention_stable-diffusion_mario FlashAttention_stable-diffusion_bunny FlashAttention_stable-diffusion_bunny
TensorRT fp16 TensorRT_stable-diffusion_mario TensorRT_stable-diffusion_bunny TensorRT_stable-diffusion_bunny
AITemplate fp16 AITemplate_stable-diffusion_mario AITemplate_stable-diffusion_bunny AITemplate_stable-diffusion_bunny

References

🌎 Join our community

🌎 Contributing

As an open source project in a rapidly evolving field, we welcome contributions of all kinds, including new features and better documentation. Please read our contributing guide to learn how you can get involved.


For managed hosting on our cloud or on your private cloud Contact us →