nom

This is a gradio issue: https://github.com/gradio-app/gradio/issues/3060. It's fixed on the gradio side with `gradio==3.19`, however the AUTOMATIC1111 webui CSS does not play well with this version.

If I comment out the line `x = x + d * dt` it works fast for every step, so it seems like extra compilation is being done on that...

Thanks, compilation is triggered on this line: https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/models/image_transformer_v2.py#L384 This does not access a tensor value if I understand `einops.rearrange` correctly. I logged the shapes of the variables used in that...

I dumped the HLO/IR graphs, but how can I see which belongs to which step? The dumps are ~1M lines long.

> Can you dump the `PT_XLA_DEBUG=1` output? Want to know how `einops.rearrange` triggers the recompilation. Interestingly it triggers on different lines each time I run it. This is the latest:...

Thanks @yisol. Are you perhaps not doing EMA? Also if you could share a work-in-progress rough train script here, that'd be really helpful - just to get a better understanding...

Hey Nico, don't think I can look at Vercel logs -- let me know if I can help