Add CLIP and T5XXL for StableDiffusionV3

[james77777778]

Please refer to these colabs for numerics checks:

• CLIPPreprocessor: https://colab.research.google.com/drive/1vCAUyVXetyya5EJ_o7U8xuXZCnhgDRxH
• T5XXLPreprocessor: https://colab.research.google.com/drive/1UBerWxCEFwyr2QfUjKROhV7rPQObo6P-
• CLIPTextEncoder: https://colab.research.google.com/drive/1tHKAZYgEhu_Isgo3fd3mAK7uMTkOEW3w
• T5XXLTextEncoder: https://colab.research.google.com/drive/1z99c9JSGRJEHmTesHV0wU-K0K9mrcYgC

CLIPTokenizer and T5XXLPreprocessor are verified using transformers library as the reference implementation.

The reference implementations of CLIPTextEncoder and T5XXLTextEncoder come from private code.

Notes for reviewing:

• CLIPTokenizer subclasses BytePairTokenizer with extensive modifications. Tests have been included.
• CLIPPreprocessor subclasses Preprocessor and uses StartEndPacker for adding start & end token.
• CLIPTextEncoder consists of CLIPAttention and CLIPEncoderBlock and subclasses Backbone.
• T5XXLPreprocessor subclasses Preprocessor and uses StartEndPacker for adding start token.
• T5XXLTextEncoder is similar to T5 but excludes the decoder part.

Future works:

• Implement CLIPPreprocessor
• Wrap CLIPTextEncoder and T5XXLTextEncoder for the use in StableDiffusionV3
• Implement VAEImageDecoder (Add VAEImageDecoder for StableDiffusionV3 #1796)
• Implement MMDiT
• Implement StableDiffusionV3 (inference model)

@divyashreepathihalli @mattdangerw @SamanehSaadat

Should I finish all the implementations in this PR, or would it be better to separate them into multiple ones?

[mattdangerw]

@james77777778 I think probably smaller PRs will be easier as a general strategy will be easier. So let's go with "multiple" to your question above.

General context for me, why do we need two tokenizers and text encoders?

[james77777778]

@mattdangerw

General context for me, why do we need two tokenizers and text encoders?

Actually, we need three pairs of them...

https://github.com/huggingface/diffusers/blob/dc07fc29da8a593f68080fbde0b9161a9a68bd36/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L131-L180

In their paper (https://arxiv.org/abs/2403.03206), they mention that:

While the main motivation for using multiple text-encoders is boosting the overall model performance

It should also be possible to have flexible text encoders:

we train our model with three text encoders, with an individual drop-out rate of 46.3%. Hence, at inference time, we can use an arbitrary subset of all three text encoders

However, the reference implementation currently uses all of them. We could implement the flexible design in a future PR.

[mattdangerw]

Just some minor comments. I suspect the preprocessing might need to move around a bit as we figure out our general vision preprocessing story, but no need to block on that I think.

[mattdangerw]

side note, I've cleaned this up on a commit now on master, this will change slightly when i rebase the whole branch

[james77777778]

I've added a TODO for this call

[mattdangerw]

I think we can land all of this in SD3 folder for now.

But we might eventually want clip in it's own folder here. Usable with SD3. WDYT @divyashreepathihalli ?

[james77777778]

The implementation will be cleaner if we have a CLIP model. I can try adding it after landing SD3.

[divyashreepathihalli]

@james77777778, CLIP is in KerasCV - it needs to be added to keras_hub still. That would make this implementation cleaner for sure.

Given that this is only using the text encoder part and not the whole clip model. We can land this as is now and clean up once CLIP is added.

[mattdangerw]

Not necessarily for this PR, but I think we probably should make it possible to instantiate only a t5 encoder model from our actual t5 implementation. Then you could just use that symbol directly here.

[james77777778]

I can try to decompose T5 into encoder and decoder archs after landing SD3

[mattdangerw]

@james77777778

Actually, we need three pairs of them...

So the actually just concat the encode representation together or something like that?

However, the reference implementation currently uses all of them. We could implement the flexible design in a future PR.

Yeah, let's do whatever is easier for now. This quote from the paper makes me wonder if only using the clip encoders might be the better default. Adding 5B parameters just for a slight performance boost is probably not something most of our users will want.

But let's keep building whatever is simplest for now, can figure that out later.

[james77777778]

So the actually just concat the encode representation together or something like that?

As far as I know, yes.
The encoding of CLIP is zero-padded to the size of the T5XXL encoding and then concatenated.

Yeah, let's do whatever is easier for now. This quote from the paper makes me wonder if only using the clip encoders might be the better default. Adding 5B parameters just for a slight performance boost is probably not something most of our users will want.

But let's keep building whatever is simplest for now, can figure that out later.

Sounds good! Will investigate the drop in T5 after finishing SD3.
BTW, work with T5 model is really a pain because I can't fit it on my GPU and can only run it on the CPU, which is very slow.

EDITED:
I made the implementation much clearer in CLIPEncoderBlock by using layers.MultiHeadAttention. The numeric check has been also updated.
CLIPTextEncoder: https://colab.research.google.com/drive/1tHKAZYgEhu_Isgo3fd3mAK7uMTkOEW3w

[divyashreepathihalli]

LGTM

[mattdangerw]

Let's pull this in! I suspect there will be a lot of churn on this model (given potential refactors of clip, t5 and still hazy preprocessing), but I think we should keep moving and reshuffle as we go.