【深度学习】SDXL tensorRT 推理，Stable Diffusion 转onnx，转TensorRT

文章目录

1 sdxl 转 diffusers

juggernautXL_version6Rundiffusion.safetensors文件是pth pytroch文件，需要先转为diffusers 的文件结构。

def convert_sdxl_to_diffusers(pretrained_ckpt_path, output_diffusers_path): import os os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" # 设置 HF 镜像源（国内用户使用） os.environ["CUDA_VISIBLE_DEVICES"] = "1" # 设置 GPU 所使用的节点 import torch from diffusers import StableDiffusionXLPipeline pipe = StableDiffusionXLPipeline.from_single_file(pretrained_ckpt_path, torch_dtype=torch.float16).to("cuda") pipe.save_pretrained(output_diffusers_path, variant="fp16") if __name__ == '__main__': convert_sdxl_to_diffusers("/ssd/wangmiaojun/tensorRT_test/juggernautXL_version6Rundiffusion.safetensors", "/ssd/wangmiaojun/tensorRT_test/mj_onnx") 

FP16在后面不好操作，所以最好先是FP32:

```bash def convert_sdxl_to_diffusers(pretrained_ckpt_path, output_diffusers_path): import os os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" # 设置 HF 镜像源（国内用户使用） os.environ["CUDA_VISIBLE_DEVICES"] = "1" # 设置 GPU 所使用的节点 import torch from diffusers import StableDiffusionXLPipeline pipe = StableDiffusionXLPipeline.from_single_file(pretrained_ckpt_path).to("cuda") pipe.save_pretrained(output_diffusers_path) if __name__ == '__main__': convert_sdxl_to_diffusers("/ssd/wangmiaojun/tensorRT_test/juggernautXL_version6Rundiffusion.safetensors", "/ssd/wangmiaojun/tensorRT_test/mj_onnx32") 

2 转onnx

有了diffusers 的文件结构，就可以转onnx文件。

项目：https://huggingface.co/docs/diffusers/optimization/onnx

pip install -q optimum["onnxruntime"] 

optimum-cli export onnx --model /data/xiedong/fooocus_tensorRT/juggernautXL_version6Rundiffusion_onnx/ --task stable-diffusion-xl juggernautXL_version6Rundiffusion_onnx_optinmum 

3 转TensorRT

博客里第2个步骤转onnx不是必须的，因为在这个步骤里，tensorRT官方代码里是有带转onnx的。可以第1个步骤后就走这个步骤。

stabilityai/stable-diffusion-xl-1.0-tensorrt

项目：https://huggingface.co/stabilityai/stable-diffusion-xl-1.0-tensorrt

TensorRT环境：

git clone https://github.com/rajeevsrao/TensorRT.git cd TensorRT git checkout release/9.2 

stabilityai/stable-diffusion-xl-1.0-tensorrt项目

git lfs install git clone https://huggingface.co/stabilityai/stable-diffusion-xl-1.0-tensorrt cd stable-diffusion-xl-1.0-tensorrt git lfs pull cd .. 

进入容器：

docker run --gpus all -it --ipc=host --ulimit memlock=-1 --ulimit