[Community] Make safety model end-to-end compileable - Inference time of JAX / Flax pipeline

[ngaer]

I have a question about JAX / Flax pipeline. We run it on a v3-8 Cloud TPU VM and it works pretty great. The first inference as expected takes around a minute because of the compilation. And the next ones take around 5-6 seconds, which is really impressive. Either way, if we change one of the input parameters other than prompt (width, height, scale, steps), the inference takes again about a minute, so it looks like it compiles it over again.

So the question: is that expected or it still should be 5-6 seconds and we're doing something wrong? If this is expected behavior, are there any ways we can improve inference time when changing these parameters?

[pcuenca]

Hi @ngaer! Yes, that is expected because all the other parameters are compiled using pmap as compile-time constants, so they get inlined as part of the compilation: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py#L303

It's potentially possible to compile them as normal inputs by turning them into tensors, but we haven't done that yet. If you want to explore this avenue, you can follow the steps in the pipeline source code but doing the pmap yourselves, then passing jit=False in the inference call (since you already jitted the code :)) Let us know if that's something you'd like to pursue! We may work on it ourselves, but unfortunately we can't provide a time estimate at this point.

[ngaer]

Hi @pcuenca, thanks for the details. But I'm afraid that I lack any knowledge in JAX and it would really take a while to do it myself. We're ok to wait for the official implementation of these things, and I believe it would be really useful for other users as well.

Anyway, thanks for a great job, you guys doing out there.

[patrickvonplaten]

If anybody of the community would be interested in giving this issue a try, I'd be more than happy to help!

In short, the issue is about "Making JAX stable diffusion end-to-end compilable - including the safety checker"

[soma2000-lang]

@patrickvonplaten I would like to give this a try.

[patrickvonplaten]

Cool!

[soma2000-lang]

@patrickvonplaten Can you please help me get started with this?

[patrickvonplaten]

Sorry actually I just noticed that this feature is already completed in #825