diffusers icon indicating copy to clipboard operation
diffusers copied to clipboard

Flax support for Stable Diffusion 2

Open pcuenca opened this issue 3 years ago • 5 comments

pcuenca avatar Nov 25 '22 16:11 pcuenca

The documentation is not available anymore as the PR was closed or merged.

thanks @pcuenca ♥ it is working 🎉

camenduru avatar Nov 26 '22 02:11 camenduru

I converted to flax I put it https://huggingface.co/flax/stable-diffusion-2 if someone needs

camenduru avatar Nov 26 '22 03:11 camenduru

Super cool!

Should we maybe add one slow test for SD-2?

patrickvonplaten avatar Nov 28 '22 09:11 patrickvonplaten

Should we maybe add one slow test for SD-2?

I added a couple of integration tests to ensure that the output from the Flax UNet is close enough to the output from PyTorch. This works within a tolerance of 1e-2 in my hardware, despite PyTorch values being generated in float16 and Flax ones in bfloat16. This sounds like the UNet implementation appears to be correct.

The tests may fail on the real testing hardware though (V100 and TPU). I'll adapt if that's the case.

pcuenca avatar Nov 28 '22 16:11 pcuenca

Merging then to deploy the backend. Thanks!

pcuenca avatar Nov 29 '22 11:11 pcuenca