unet export to onnx issue

Describe the bug

Trying to export the unet model to onnx. Used this script

.onnx file created successfully but export wasn’t successful:

import onnx
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)

Error:

protobuf_string = model if isinstance(model, bytes) else model.SerializeToString()
ValueError: Message onnx.ModelProto exceeds maximum protobuf size of 2GB: 3439066593

Reproduction

No response

Logs

No response

System Info

Python 3.8.13 Pytorch 1.13 ONNX 1.12.0 transformers 4.22.1 diffusers 0.3.0

[anton-l]

Hi @tanayvarshney! The exported UNet is indeed >2GB due to the protobuf size limits, but the directory of the exported model contains external weight tensors. So while this fails the onnx check, onnxruntime is able to load the model just fine 😃

P.S. I’m preparing an update for the onnx exporting script that will save the UNet in fp16, which should fit into 2GB.

[anton-l]

https://github.com/huggingface/diffusers/pull/932 enabled an FP16 conversion, so it should be possible to manually convert and use the checkpoints with size-sensitive onnx tools now. Feel free to ping me here if you have any issues with that workflow! Hardware-independent FP16 checkpoints are still WIP, as they require still-unreleased onnxruntime updates.