torch.compile with stable diffusion

[Madhumitha-MCW]

I would like to know why any inference example with stable diffusion has only unet part of the model passed into torch.compile(). https://pytorch.org/TensorRT/tutorials/_rendered_examples/dynamo/torch_compile_stable_diffusion.html

1. When using torch.compile with model.vae it relatively takes longer time than just compiling unet part. So, is it recommended to compile only unet or what is the exact reason?

2. And if I want to execute in graph mode, will all other uncompiled parts undergo eager execution if I compile only unet part?

[tolgacangoz]

1. Most of the time, torch.compile() is applied to UNet only because it is UNet that is responsible for most of the calculations. You can compile as many components as possible such as VAE, ControlNet, text encoder(s), and even optimizer, sampler function, torchvision.transforms.v2. But the last three may not be stable and proper yet. Are you saying that compiling both UNet and VAE is slower than compiling only UNet? Speeding-up thanks to compiling is observed in relatively new GPUs more.
2. I think so.

[Madhumitha-MCW]

Yes, the time taken to compile vae is more than only compiling unet part. Thanks on your reply!

[tolgacangoz]

Are you talking about the first inference time? Don't count the first inference. Evaluate the next ones. The first inference (the compilation part) could take more time for more compilation tasks.

[Madhumitha-MCW]

Yes, I just was mentioning the first inference time(compilation part). The final inference is not slow.

[tolgacangoz]

My commonsense -> If someone compiles more parts then wouldn't it take more time to compile?

[Madhumitha-MCW]

Yes, it would. I first thought the entire inference time was longer, but when I verified it, it was just the initial part which took longer