Add FlaxCLIPTextModelWithProjection

[pcuenca]

What does this PR do?

FlaxCLIPTextModelWithProjection is necessary to support the Flax port of Stable Diffusion XL: https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/fb6d705fb518524cabc79c77f13a0e7921bcab3a/text_encoder_2/config.json#L3

I can add some tests, if necessary, after this approach is validated.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@patrickvonplaten @patil-suraj @sanchit-gandhi @younesbelkada

[HuggingFaceDocBuilderDev]

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

[patrickvonplaten]

Should we maybe for now just add it in a subfolder of sdxl in diffusers here: https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/stable_diffusion_xl instead of having to rely on transformers here? I'm not 100% convinced this model is really needed for core transformers usage.

Would also not force the user to have to install transformers from main :-)

[pcuenca]

Should we maybe for now just add it in a subfolder of sdxl in diffusers here: https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/stable_diffusion_xl instead of having to rely on transformers here? I'm not 100% convinced this model is really needed for core transformers usage.

The PyTorch version of the same model was added 9 months ago, so I assumed it was ok.

But sure, we can do that. In that case, how do we deal with it?

• Change the library to diffusers here: https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/model_index.json#L15. Unless I'm mistaken, then we'd need to distribute the flax weights separately, or use a branch.
• Create a hack in diffusers to map the library.

Would also not force the user to have to install transformers from main :-)

Yes, of course, this was meant as the long-term solution.

[patrickvonplaten]

Ah yeah good point JAX & PyTorch share the same config - this will become complicated indeed then. Ok let's try to get it merged here. CLIP is important enough to be merged to transformers indeed

[sanchit-gandhi]

Thanks for adding this @pcuenca - looking clean! Good with me to add this Flax model if you require it for SD / SD XL - I'll leave it to you and Patrick to decide on whether it's an appropriate addition, but from a Flax transformers side the design is great.

Feel free to add the model to the testing file such that it is run by the CI, just need to add the model class here:

transformers/tests/models/clip/test_modeling_flax_clip.py

Line 318 in 5ee9693

all_model_classes = (FlaxCLIPTextModel,) if is_flax_available() else ()

[pcuenca]

Thanks a lot for your in-depth review @sanchit-gandhi! 🙌 I couldn't get back to this PR until today, but I think I addressed all your comments and fixed an error in the example docstring. I was getting some seemingly unrelated CI failures so I just merged the latest main to see if they pass.

[sanchit-gandhi]

Looks good - thanks for iterating @pcuenca!

[pcuenca]

cc @patrickvonplaten too, in case we want to consider other alternatives :) Otherwise feel free to merge when appropriate, as I can't do it.

[patrickvonplaten]

Looks good to me!