Refactor CLIP and update SD3.

[james77777778]

Description of the change

Hi team,
I tried to add a numeric check for SD3 but was unsuccessful. It seems that the output latent differs slightly from the diffusers impl.

Here’s what I tried:

• Replace LayerNormalization with the custom layer which is the same as diffusers's impl.
• Updated the dtype of text encoders to be float16.
• Updated the scheduler of SD3 to be consistent with diffusers's impl.
• Fixed the input latent for both diffusers's pipeline and keras_hub's pipeline.
• Used diffusers.StableDiffusion3Pipeline.encode_prompt to get the text embeddings and fed them to keras_hub.models.StableDiffusion3TextToImage
• Ensured the inputs of the diffusion model are consistent with diffusers's impl. (timesteps, sigma, text embeddings and input latent)
• Manually ran the denoise_step to verify the latent numerics.

The latent numerics failed to pass the check with atol=1e-1.

However, the generated images look good.

Model	Generated Image
stable_diffusion_3.5_medium
stable_diffusion_3.5_large	TBD
stable_diffusion_3.5_large_turbo	TBD

cc @abheesht17

I think this PR still helps narrow the gap between keras_hub's impl and diffusers's impl.
Uploading new presets is required.

Reference

Colab Notebook

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and works with all backends (TensorFlow, JAX, and PyTorch).
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have followed the Keras Hub Model contribution guidelines in making these changes.
• I have followed the Keras Hub API design guidelines in making these changes.
• I have signed the Contributor License Agreement.

[abheesht17]

Instead of this, can you try using this layer: https://github.com/keras-team/keras/blob/master/keras/src/layers/normalization/rms_normalization.py#L7?

keras.layers.LayerNormalization(rms_scaling=True) was incorrectly implemented, but keras.layers.RMSNormalization is correct. It should work here.

[james77777778]

Unfortunately, keras.layers.RMSNormalization requires keras >= 3.9.

I have added an identical custom layer as a fallback if layers doesn't include it.

[james77777778]

@abheesht17
I have added the numeric checks for the scheduler and the text encoders. Also fixed the float16 + cpu issue in torch by setting dtype="float32" in CLIP's LayerNormalization.

Using JAX + CUDA:

Model	Numeric Checks
stable_diffusion_3.5_medium	🔶 Scheduler difference: 3.569126e-05 🔶 Text embeddings difference: 0.0018353278

EDITED:
Just spotted some upstream Keras issues. Will try to fix them later. Plz don't merge!

[divyashreepathihalli]

@james77777778 we have error tolerance error with TF GPU, can you please check?

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request aims to align the Keras implementation of Stable Diffusion 3 with the diffusers library by updating the layer normalization, scheduler, and text encoder data types. The changes are extensive, touching model definitions, backbones, and the checkpoint conversion script. The introduction of a custom RMSNormalization layer for compatibility and the refactoring of the scheduler logic are significant improvements. The conversion script is also updated to perform numerical validation against diffusers, which is a great step towards ensuring correctness.

My review focuses on improving code clarity and maintainability. I've identified a few opportunities to refactor repetitive code blocks and clarify stateful logic in a custom layer. These changes should make the code easier to read and manage in the future. Overall, this is a solid contribution toward improving the model's fidelity with the reference implementation.

[james77777778]

In this PR:

• CLIP has been refactored to be consistent with SigLIP
  • CLIPPreprocessor now accepts prompts and images instead of only images.
  • Reorganize the CLIP dir to follow other tasks.
• The dtype of text encoders in SD3 defaults to float16
• Ensure the output numeric of FlowMatchEulerDiscreteScheduler is consistent with diffusers's impl.

Caution: We need to upload the updated presets for CLIP and SD3 after the merge.