将自己训练的diffusion model转成onnx

模型训练过程参考上上篇diffusers库训练diffusion model

模型转换参考：diffusers官方教程

torch转onnx多输入参考：pytorch转onnx模型多输入问题（如：Bert）

安装依赖：

pip install -q optimum["onnxruntime"]

转换代码：

from diffusers import UNet2DModel, DDPMScheduler, DDPMPipeline
import torch
from PIL import Image
import numpy as np

model = UNet2DModel.from_pretrained(r"xxx/ddpm-anime-faces-64/unet", use_safetensors=True)

sample_size = model.config.sample_size
example_input = torch.randn(1, 3, sample_size, sample_size)
t = torch.tensor(1000)
print(t)
save_onnx_path = r"xxx/ddpm-anime-faces-64/unet/unet.onnx"
with torch.no_grad():
    # 注意：模型forward有两个输入
    torch.onnx.export(model, (example_input, t), save_onnx_path)