【深度学习】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