Refactor CLIP and update SD3. (#2316)

* Update SD3 scheduler and the dtype of the text encoders.

* Fix the test.

* Fix torch float16 issues and jax take issue. Add numeric checks for SD3 scheduler and text encoders.

* Fix CLIP test.

* Refactor CLIP models.

* Update CLIP conversion script.

* Update `from_config`.

* Fix tests.

Author: james77777778

keras_hub/src/models/clip/clip_backbone.py

@@ -1,109 +1,10 @@
-import math
-
 from keras import layers
-from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.backbone import Backbone
-
-
-class CLIPVisionPooler(layers.Layer):
-    """The vision pooler layer of CLIP.
-
-    `CLIPVisionPooler` will extracts the first token (index `0`) from the
-    sequence of the vision embeddings as the pooled outputs.
-
-    Call arguments:
-        vision_embeddings: A tensor of shape
-            `(batch_size, sequence_length, hidden_dim)`.
-    """
-
-    def call(self, vision_embeddings):
-        return vision_embeddings[:, 0, :]
-
-    def compute_output_shape(self, input_shape):
-        return (input_shape[0], input_shape[-1])
-
-
-class CLIPTextPooler(layers.Layer):
-    """The text pooler layer of CLIP.
-
-    `CLIPTextPooler` extracts the text embeddings at the positions of EOS tokens
-    as the pooled outputs.
-
-    Call arguments:
-        text_embeddings: A tensor of shape
-            `(batch_size, sequence_length, hidden_dim)`.
-        token_ids: A tensor of shape `(batch_size, max_tokens)`, used to
-            identify the positions of EOS tokens.
-    """
-
-    def call(self, text_embeddings, token_ids):
-        # `keepdims` is not supported in `keras<=3.1`.
-        eos_index = ops.argmax(token_ids, axis=-1)
-        eos_index = ops.expand_dims(eos_index, axis=-1)
-        eos_index = ops.expand_dims(eos_index, axis=-1)
-        pooled_outputs = ops.take_along_axis(text_embeddings, eos_index, axis=1)
-        return ops.squeeze(pooled_outputs, axis=1)
-
-    def compute_output_shape(self, input_shape):
-        return (input_shape[0], input_shape[-1])
-
-
-class CLIPHead(layers.Layer):
-    """The head layer of CLIP.
-
-    `CLIPHead` takes `vision_embedding` and `text_embedding` as inputs to
-    compute the corresponding logits. Both embeddings are L2 normalized and used
-    to compute pairwise cosine similarity. The resulting logits are then scaled
-    by a learnable `logit_scale` parameter.
-
-    Call arguments:
-        vision_embedding: A tensor of shape `(batch_size, hidden_dim)`.
-        text_embedding: A tensor of shape `(batch_size, hidden_dim)`.
-    """
-
-    def build(self, input_shape):
-        self.logit_scale = self.add_weight(
-            shape=(),
-            initializer=lambda *a, **kw: math.log(1 / 0.07),
-            trainable=True,
-            dtype=self.variable_dtype,
-            name="logit_scale",
-        )
-
-    def call(self, vision_embedding, text_embedding):
-        normalized_vision_embedding = ops.sqrt(
-            ops.sum(ops.power(vision_embedding, 2), axis=-1, keepdims=True)
-        )
-        normalized_text_embedding = ops.sqrt(
-            ops.sum(ops.power(text_embedding, 2), axis=-1, keepdims=True)
-        )
-        vision_embedding = vision_embedding / normalized_vision_embedding
-        text_embedding = text_embedding / normalized_text_embedding
-        logit_scale = ops.exp(self.logit_scale)
-        text_logits = (
-            ops.matmul(
-                text_embedding,
-                ops.transpose(vision_embedding),
-            )
-            * logit_scale
-        )
-        vision_logits = ops.transpose(text_logits)
-        return vision_logits, text_logits
-
-    def compute_output_shape(
-        self, vision_embedding_shape, text_embedding_shape
-    ):
-        vision_logits_shape = (
-            vision_embedding_shape[0],
-            text_embedding_shape[0],
-        )
-        text_logits_shape = (
-            text_embedding_shape[0],
-            vision_embedding_shape[0],
-        )
-        return vision_logits_shape, text_logits_shape
+from keras_hub.src.models.clip.clip_layers import CLIPHead
+from keras_hub.src.models.clip.clip_layers import CLIPTextPooler
+from keras_hub.src.models.clip.clip_layers import CLIPVisionPooler
 
 
 @keras_hub_export("keras_hub.models.CLIPBackbone")