exporting the diffusion_net model to ONNX or using TensorRT

[wallkop]

Hi @ltzheng

First of all, I would like to thank Memo for this great project. It is the best open-source digital human project I have used.

I have conducted some performance tests. I found that diffusion_net is very slow during diffusion, with most of the time being spent on downsampling and upsampling. I tried converting diffusion_net into an ONNX model. Since the forward method contains data structures like dict and list, I used a wrapper for packaging, but ultimately, I couldn't successfully export the model due to many internal issues, such as repeated parameter passing.

I would like to ask:
（1）Does the author plan to implement quantization or use tools like TensorRT for acceleration?
（2）Is it feasible to export diffusion_net to an ONNX model? (I am not a professional algorithm engineer)

Thanks.

[ltzheng]

Thank you for your kind words! Regarding the first question, we have collaborators working on TensorRT, but I’m unsure about the timeline. As for the second question, I’m also not very familiar with ONNX.