[Feature Request] fx2ait example for stable diffusion

[hemildesai]

Hi all,

It would be great to have a fx2ait example for stable diffusion. It will help provide a tutorial on how to use fx2ait for complex pipelines.

The part where I'm confused is for complex models like the SD Unet which has both attention and conv blocks, how do we permute inputs dynamically?

Happy to contribute if possible.

[ipiszy]

cc fx2ait owner @wushirong @frank-wei

FYI we have plans to use aten2ait to support the StableDiffusion model by the end of this half, but you're very welcome to contribute the fx2ait part.

wrt permute inputs dynamically: could you elaborate? For now we only support fixed shaped inputs for SD. We are also working on symbolic shape support in AIT, and after it's done we'll have better support for dynamic inputs for SD.

[hemildesai]

Thanks for the response. In fx2ait vision example, it is mentioned that

Notice that because AIT supports channel last, while pytorch supports channel first operation, we need to permute the input

In stable diffusion, there are both Conv and Attention blocks so I'm confused as to how the input permutation will be handled.