XLA compilation even further for speed up ⚡️

[sayakpaul]

I am trying to compile the sub-models of Stable Diffusion with XLA in isolation. This includes the core diffusion process too. See here.

I am able to obtain a 25% speedup during inference with this.

However, the text encoder doesn't compile successfully: tensorflow/tensorflow#59818. Do you know anyone that could help?

You can find the compilation details in this repository: https://github.com/sayakpaul/xla-benchmark-sd

@LukeWood @mattdangerw

Also cc: @deep-diver

[jbischof]

@sayakpaul it looks like you're still working with the TF team but let me know if you get stuck.

[sayakpaul]

Yeah sure. This is anyway a project for the keras sprint. So, very much line I guess.

[jbischof]

I would love to refactor the Stable Diffusion codebase to make it more modular and hackable instead of a monolithic thing. Might be worth a discussion down the line!

[sayakpaul]

You might want to check out how we (@deep-diver and I) decomposed the Stable Diffusion class into intuitive modules in that regard.

[jbischof]

in keras-cv or elsewhere?

[sayakpaul]

If you look at https://github.com/sayakpaul/xla-benchmark-sd, you will notice that it decomposes the modules involved in the StableDiffusion class of KerasCV.