Add VAEImageDecoder for StableDiffusionV3

[james77777778]

Numerics check:
https://colab.research.google.com/drive/1YsWvZ0NBINDgdqipsldso1Y1NKJUDkUf?usp=sharing

Future works:

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

@divyashreepathihalli @mattdangerw @SamanehSaadat

[divyashreepathihalli]

Thanks for the PR!! LGTM!

[mattdangerw]

lgtm! minor nits and a few design notes I don't think we need to solve in this PR

[mattdangerw]

Note that with pali gemma, our "backbone" contains all the weights needed from a pre-trained model. So in that case the image encoder and text decoder collectively form a single backbone class.

We should discuss the high level flows that we want as we go, but our current approach is...

• StableDiffusionBackbone should contain all the pretrained weights for using the entire model without a specific task setup. This can come from stitching other backbones/sub models together. No preprocessing.
• StableDiffusion[TaskName] would wrap the backbone with a setup for a particular task. Preprocessing included. Ideally allowing both find-tuning and inference, but that would depend on the task at hand. For stable diffusion the main task is definitely text to image, though I'm not sure what we should call that. StableDiffusionImageGenerator is kinda long.

[james77777778]

I was unsure of how we wanted to assemble these encoders and the decoder, so I made them as a Backbone first.

We should discuss the high level flows that we want as we go, but our current approach is...

Got it. Will make the encoders and decoder as a keras.Model to follow that pattern.

I think the task name, ImageGenerator, is a bit ambiguous. Maybe we should call it TextToImage instead?
It is also possible to use SD3 for ImageToImage and Inpaint tasks.
https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/stable_diffusion_3

[mattdangerw]

TextToImage sounds fine to me. Shorter.

Got it. Will make the encoders and decoder as a keras.Model to follow that pattern.

I suspect we still have more to figure out here. For these big "composite models" with lots of sub-components, it would be good if we allowed loading sub models individually some how. E.g. load the text encoder of a T5 model, or just the image encoder of PaliGemma. That's a valid use case, that fit's with the flexibility we'd like to shoot for, and we don't support it today. But a probably for another PR I think.