環(huán)境配置:

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

#進入鏡像內(nèi)安裝所需工具

!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"

執(zhí)行代碼:

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)