Flax: Flip sin to cos in time embeddings

[pcuenca]

This was assumed in the previous implementation, but now (0b61cea) the default is the opposite.

Fixes #1145.

[HuggingFaceDocBuilderDev]

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

[pcuenca]

I'm running the slow tests locally.

[kashif]

ah damn! so sorry! My bad! I believe you are right this is the reason. I do not think this function is called by itself... to be sure shall we set the default of the function to be True?

[pcuenca]

ah damn! so sorry! My bad! I believe you are right this is the reason. I do not think this function is called by itself... to be sure shall we set the default of the function to be True?

No worries! These issues should be easier to catch going forward with the new TPU tests.

Regarding making it the default, I'm not sure because the default is False in the PyTorch version and maybe it's better to keep them consistent. Let's see what the others think :)

[patrickvonplaten]

ah damn! so sorry! My bad! I believe you are right this is the reason. I do not think this function is called by itself... to be sure shall we set the default of the function to be True?

No worries! These issues should be easier to catch going forward with the new TPU tests.

Regarding making it the default, I'm not sure because the default is False in the PyTorch version and maybe it's better to keep them consistent. Let's see what the others think :)

Yes we should indeed mirror the PyTorch behavior here

[kashif]

In that case I can send a PR to add an option to the FlaxUNet2DConditionModel to set the flip_sin_to_cos to True which is the default for the pytorch version?

[patrickvonplaten]

Ah sorry, I didn't do a thoughout check - will take a look now!

[patrickvonplaten]

Actually, let's maybe leave the PR as is for now since it solves the problem. In general, it's nicer to mirror PyTorch API - but we need to be sure that it doesn't break anything and we also shouldn't add features just for the sake of having the same API as PyTorch. So if we put the default to True , we need to make sure that it's set to False for Stable Diffusion.

More generally though, I think it's better to not add new function arguments just for the sake of making Flax API similar to PyTorch. E.g. such a PR: https://github.com/huggingface/diffusers/pull/1081/files only really makes sense if it enables new functionality for Flax/JAX which I think it didn't really there. Given that we have very low usage in Flax, let's not spend too much time on giving the API more features than it needs :-)

I did a really bad job at reviewing #1081, so taking full responsibility here :-)

Will do a patch release now - thanks for the quick fix @pcuenca and @kashif !