ONNX compatibility issues

[harishanand95]

Thanks for the great library!

I'm looking into generating ONNX version of the stable diffusion models.

I had to make few changes to unet_2d_condition.py to get it to work. You can find some of the changes here.. harishanand95@8dd4e82
Let me know if you prefer a PR for this and if there exist any implementations for return non dict that I can refer to. Thanks!

1. reason for removing broadcast_to :

torch.onnx.symbolic_registry.UnsupportedOperatorError: Exporting the operator ::broadcast_to to ONNX opset version 16 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

2. reason for dict changes:

RuntimeError: Encountering a dict at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for list, use a tuple instead. for dict, use a NamedTuple instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior. 

I was able to generate images with onnxruntime and DirectML, there are some more issues to be fixed.

• Instructions: https://github.com/harishanand95/diffusers/blob/dml/examples/inference/readme.md#instructions-for-onnx-dml-execution-on-windows

• A script that creates onnx files https://github.com/harishanand95/diffusers/blob/dml/examples/inference/save_onnx.py

• A StableDiffusionPipeline implementation that runs on onnxruntime https://github.com/harishanand95/diffusers/blob/dml/examples/inference/dml_onnx.py

Let me know if you have further concerns or ideas, thanks!

[pcuenca]

Regarding the first point, @hollance suggested to use torch.tile; I'm not sure if that's preferable to multiplying by ones or if it doesn't matter. In my opinion tile is more concise, especially if we want to keep the broadcasted dtype as an int for symmetry with the previous version.

Regarding the second point, @anton-l and @patil-suraj suggested to adopt the transformer's style to select the desired output ('dict' or 'tuple') using a flag.

[patrickvonplaten]

Happy to use torch.tile! @harishanand95 or @pcuenca would you like to open a PR? :-)

Regrading the second point, yes I think we should adopt the return_dict=True/False strategy of Transformers (defaulting it to True and allow one to return tuples) - happy to open a PR for this tomorrow or Friday!

[harishanand95]

I'm not sure if the tile is onnx compatible either.. Feel free to check this on your side to confirm (I hope it's not my torch version causing troubles, currently on 1.13.0.dev20220819+cpu ).. Thanks!

import torch.nn as nn
import torch
import torch.nn.functional as F
print(torch.__version__)

class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        # y = x.broadcast_to(3)
        return x.tile(3)#, y

m = Model()
m(torch.rand(1))

traced_model = torch.jit.trace(m, (torch.rand(1),), strict=True)
torch.onnx.export(traced_model, (torch.rand(1),), "test.onnx", input_names=["input"], dynamic_axes={"input": [0]}, opset_version=16)

[pcuenca]

You are right, @harishanand95, tile does not work with ONNX. I'll open a PR with your approach then. Thanks a lot!

[harishanand95]

Thanks!