AIGC | 十行代码用TensorRT加速Stable Diffusion模型

环境配置:

!docker run --rm -it --gpus all -p 8888:8888 -v $PWD:/workspace nvcr.io/nvidia/tensorrt:23.04-py3 /bin/bash

#进入镜像内安装所需工具

!pip install --upgrade "torch <2.0.0"

!pip install --upgrade "tensorrt>=8.6"

!pip install --extra-index-url https://pypi.ngc.nvidia.com --upgrade "onnx-graphsurgeon" "onnxruntime" "polygraphy"

!pip install --upgrade "accelerate" "diffusers>=0.16" "transformers"

执行代码:

import diffusers

import torch

import tensorrt

from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline

from diffusers import DDIMScheduler

scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1", subfolder="scheduler")

pipe_trt = StableDiffusionPipeline.from_pretrained(

  "stabilityai/stable-diffusion-2-1",

  custom_pipeline="stable_diffusion_tensorrt_txt2img",

  revision='fp16',

  torch_dtype=torch.float16,

  scheduler=scheduler,

  image_height=512,

  image_width=512)

pipe_trt.set_cached_folder("stabilityai/stable-diffusion-2-1", revision='fp16')

pipe_trt = pipe_trt.to("cuda")

prompt = "A hyper realistic avatar of a guy riding on a black KTM duke 390 in leather suit,high detail, high quality,8K,photo realism"

neg_prompt = " flying mirrors,low quality"

image = pipe_trt(prompt, negative_prompt=neg_prompt).images[0]

display(image)

for _ in range(5):

  image = pipe_trt(prompt, negative_prompt=neg_prompt, num_inference_steps=20).images[0]

  display(image)