ofa-sys / diffusion-deploy

Diffusion deployment

This repository contains scripts for the deployment of diffusion models (based on diffusers) on both GPU (Nvidia) and CPU (Intel). The aim is to significantly speed up the inference of diffusion models. It provides a ~12x speedup on CPUs and a ~4x speedup on GPUs. Integrated with small-stable-diffusion-v0, it could generate an image in just 5s on the CPU.

CPU speedup

We develop the diffusion deployment on CPU based on Intel OpenVINO. The pipeline OpenVINOStableDiffusionPipeline is modified from OnnxStableDiffusionPipeline. The code used here is in pipeline_openvino_stable_diffusion.py.

Results

Here are some experimental results on stable diffusion v1.4 for comparison with the default Pytorch CPU and Onnx pipeline.

Pipeline	Pytorch CPU	Onnx	OpenVINO
Time Cost	397s	77s ± 2.56 s	33.9 s ± 247 ms
Speedup	1	5.2	11.7

Test setting: CPU Intel(R) Xeon(R) Platinum 8369B CPU @ 2.90GHz / PNDM scheduler 50 steps)

Prerequisites

There are several limitations of OpenVINO now, Therefore, we only support the following platforms for CPU speedup.

• Ubuntu 18.04, 20.04, RHEL(CPU only) or Windows 10 - 64 bit
• Python 3.7, 3.8 or 3.9 for Linux and only Python3.9 for Windows

Requirements

• diffuers
• transformers
• openvino runtime

To install openvino runtime, you could simply use pip install onnxruntime-openvino==1.13.0.

Usage

To use this deployment, you could follow the following code:

# Load a onnx pipeline firstly.  
from diffusers import OnnxStableDiffusionPipeline
onnx_pipe = OnnxStableDiffusionPipeline.from_pretrained(
    "OFA-Sys/small-stable-diffusion-v0",
    revision="onnx",
    provider="CPUExecutionProvider",
)
# Convert it to OpenVINO pipeline.  
import pipeline_openvino_stable_diffusion
openvino_pipe = pipeline_openvino_stable_diffusion.OpenVINOStableDiffusionPipeline.from_onnx_pipeline(onnx_pipe)

# Generate images.
images = openvino_pipe("an apple, 4k")  

GPU speedup

We develop the deployment on GPU based on TensorRT and its plugins.

Comparison

Here are some experimental results of stable-diffusion-v1-4 and small-stable-diffusion-v0.

Model\Pipeline	Pytorch GPU	TensorRT	TensorRT Plugin
Stable diffusion	3.94	1.44s	1.07s
Small stable diffusion	2.7s	1.01s	0.65s

Prerequisites

See the gpu_requirements.txt for requirments. To use plugins, we need tensorrt>=8.5. If you use tensorrt==8.4, you could run it by deleting trt.init_libnvinfer_plugins(TRT_LOGGER, '') in gpu-trt-infer-demo.py and not adding PLUGIN_LIBS to LD_PRELOAD.

Usage

export PLUGIN_LIBS="/path/to/libnvinfer_plugin.so.8.5.1"
export HF_TOKEN="Your_HF_TOKEN"

mkdir -p onnx engine output
LD_PRELOAD=${PLUGIN_LIBS} python3 demo-diffusion.py "a beautiful photograph of Mt. Fuji during cherry blossom" --enable-preview-features --hf-token=$HF_TOKEN -v

Contributions

Contributions to this repository are welcome. If you would like to contribute, please open a pull request and make sure to follow the existing code style.

diffusion-deploy's Issues

Models.py is missing

Nice work! Thanks for sharing your deploying code.

For gpu tensorrt, models.py containing CLIP, VAE, and UNet is missing. The models.py from the official Tensorrt GitHub did not work.

Here, requiring the models.py.

请问这个CPU用的是多少线程？需要多少内存，谢谢！

很不错的工作，想要跑一下玩玩

Error: No module named 'diffusers.onnx_utils'

I tried to launch pipeline_openvino_stable_diffusion.py but got this error:

C:\opvin>pipeline_openvino_stable_diffusion.py
Traceback (most recent call last):
File "C:\opvin\pipeline_openvino_stable_diffusion.py", line 21, in
from diffusers.onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel
ModuleNotFoundError: No module named 'diffusers.onnx_utils'