unet export to onnx issue

[tanayvarshney]

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.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[patrickvonplaten]

@anton-l any updates here? Can we close this?

[tanayvarshney]

@patrickvonplaten I got around the issue by using the external weights, as @anton-l mentioned. Conversion to fp16 works fine. ONNX still makes a fuzz about it, but the model runs without any major issues.

[anton-l]

#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.