keras-hub-nightly 0.16.1.dev202410200345__py3-none-any.whl → 0.19.0.dev202412070351__py3-none-any.whl

keras_hub/src/models/clip/clip_backbone.py

@@ -0,0 +1,286 @@
+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
+
+
+@keras_hub_export("keras_hub.models.CLIPBackbone")
+class CLIPBackbone(Backbone):
+    """CLIP core network with hyperparameters.
+
+    This backbone implements the base architecture for Contrastive
+    Language-Image Pretraining (CLIP) model. It includes a vision and text
+    encoders and the corresponding projection layers. This backbone will output
+    the final logit scores corresponding to each image and token input. These
+    values are cosine similarities between the corresponding image and text
+    features.
+
+    The default constructor gives a fully customizable, randomly initialized
+    CLIP model with any number of layers, heads, and embedding dimensions. To
+    load preset architectures and weights, use the `from_preset` constructor.
+
+    Args:
+        vision_encoder: The CLIP vision encoder for encoding the input images.
+        text_encoder: The CLIP text encoder for encoding the input tokens.
+        projection_dim: int. The size of the projection layer.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for the models computations and weights. Note that some
+            computations, such as softmax and layer normalization will always
+            be done a float32 precision regardless of dtype.
+
+    Example:
+    ```python
+    input_data = {
+        "images": np.ones(shape=(1, 224, 224, 3), dtype="float32"),
+        "token_ids": np.ones(shape=(1, 12), dtype="int32"),
+    }
+
+    # Pretrained CLIP model.
+    model = keras_hub.models.CLIPBackbone.from_preset("clip_vit_base_patch32")
+    model(input_data)
+
+    # Randomly initialized CLIP model with custom config.
+    vision_encoder = keras_hub.models.CLIPVisionEncoder(
+        patch_size=32,
+        hidden_dim=768,
+        num_layers=8,
+        num_heads=8,
+        intermediate_dim=2048,
+        image_shape=(384, 384, 3),
+    )
+    text_encoder = keras_hub.models.CLIPTextEncoder(
+        vocabulary_size=49408,
+        embedding_dim=768,
+        hidden_dim=768,
+        num_layers=8,
+        num_heads=8,
+        intermediate_dim=2048,
+    )
+    model = keras_hub.models.CLIPBackbone(
+        vision_encoder=vision_encoder,
+        text_encoder=text_encoder,
+        projection_dim=256,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vision_encoder,
+        text_encoder,
+        projection_dim,
+        dtype=None,
+        name=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.vision_encoder = vision_encoder
+        self.text_encoder = text_encoder
+        self.vision_pooler = CLIPVisionPooler(dtype=dtype, name="vision_pooler")
+        self.text_pooler = CLIPTextPooler(dtype=dtype, name="text_pooler")
+        self.vision_projection = layers.Dense(
+            projection_dim,
+            use_bias=False,
+            dtype=dtype,
+            name="vision_projection",
+        )
+        self.text_projection = layers.Dense(
+            projection_dim,
+            use_bias=False,
+            dtype=dtype,
+            name="text_projection",
+        )
+        self.clip_head = CLIPHead(dtype=dtype, name="clip_head")
+
+        # === Functional Model ===
+        image_input = layers.Input(
+            shape=self.vision_encoder.image_shape, name="images"
+        )
+        token_id_input = layers.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        vision_embeddings = self.get_vision_embeddings(image_input)
+        text_embeddings = self.get_text_embeddings(token_id_input)
+        vision_logits, text_logits = self.clip_head(
+            vision_embeddings, text_embeddings
+        )
+
+        super().__init__(
+            inputs={
+                "images": image_input,
+                "token_ids": token_id_input,
+            },
+            outputs={
+                "vision_logits": vision_logits,
+                "text_logits": text_logits,
+            },
+            dtype=dtype,
+            name=name,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.projection_dim = projection_dim
+
+    def get_vision_embeddings(self, images):
+        """Get the embeddings from the vision encoder.
+
+        Args:
+            images: The input tensor for the vision encoder.
+
+        Returns:
+            The output embeddings obtained by applying projection layer to the
+            pooled output of the vision encoder.
+        """
+        vision_outputs = self.vision_encoder({"images": images})
+        vision_outputs = self.vision_pooler(vision_outputs)
+        return self.vision_projection(vision_outputs)
+
+    def get_text_embeddings(self, token_ids):
+        """Get the embeddings from the text encoder.
+
+        Args:
+            token_ids: The input int tensor for the text encoder.
+
+        Returns:
+            The output embeddings obtained by applying projection layer to the
+            pooled output of the text encoder.
+        """
+        text_outputs = self.text_encoder({"token_ids": token_ids})
+        text_outputs = self.text_pooler(text_outputs, token_ids)
+        return self.text_projection(text_outputs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vision_encoder": layers.serialize(self.vision_encoder),
+                "text_encoder": layers.serialize(self.text_encoder),
+                "projection_dim": self.projection_dim,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        config = config.copy()
+
+        # Propagate `dtype` to submodels if needed.
+        if "dtype" in config and config["dtype"] is not None:
+            dtype_config = config["dtype"]
+            if "dtype" not in config["vision_encoder"]["config"]:
+                config["vision_encoder"]["config"]["dtype"] = dtype_config
+            if "dtype" not in config["text_encoder"]["config"]:
+                config["text_encoder"]["config"]["dtype"] = dtype_config
+
+        # We expect submodels to be instantiated.
+        config["vision_encoder"] = layers.deserialize(
+            config["vision_encoder"], custom_objects=custom_objects
+        )
+        config["text_encoder"] = layers.deserialize(
+            config["text_encoder"], custom_objects=custom_objects
+        )
+        return cls(**config)

keras_hub/src/models/clip/clip_encoder_block.py

@@ -7,6 +7,16 @@ def quick_gelu(x):
     return x * ops.sigmoid(1.702 * x)
 
 
+# TODO: Deprecate this in favor of `keras.layers.MultiHeadAttention` once the
+# dtype compatibility issue is resolved.
+class CLIPMultiHeadAttention(layers.MultiHeadAttention):
+    def _masked_softmax(self, attention_scores, attention_mask=None):
+        attention_scores = super()._masked_softmax(
+            attention_scores, attention_mask
+        )
+        return ops.cast(attention_scores, self._value_dense.compute_dtype)
+
+
 class CLIPEncoderBlock(layers.Layer):
     def __init__(
         self,
@@ -14,6 +24,7 @@ class CLIPEncoderBlock(layers.Layer):
         num_heads,
         intermediate_dim,
         intermediate_activation="quick_gelu",
+        use_causal_mask=True,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -26,21 +37,22 @@ class CLIPEncoderBlock(layers.Layer):
         self.num_heads = num_heads
         self.intermediate_dim = intermediate_dim
         self.intermediate_activation = intermediate_activation
+        self.use_causal_mask = use_causal_mask
 
         if intermediate_activation == "quick_gelu":
             intermediate_activation = quick_gelu
 
         self.layer_norm_1 = layers.LayerNormalization(
-            epsilon=1e-5, dtype="float32", name="layer_norm_1"
+            epsilon=1e-5, dtype=self.dtype_policy, name="layer_norm_1"
         )
-        self.attention = layers.MultiHeadAttention(
+        self.attention = CLIPMultiHeadAttention(
             num_heads,
             hidden_dim // num_heads,
             dtype=self.dtype_policy,
             name="attention",
         )
         self.layer_norm_2 = layers.LayerNormalization(
-            epsilon=1e-5, dtype="float32", name="layer_norm_2"
+            epsilon=1e-5, dtype=self.dtype_policy, name="layer_norm_2"
         )
         self.dense_1 = layers.Dense(
             self.intermediate_dim, dtype=self.dtype_policy, name="dense_1"
@@ -73,7 +85,9 @@ class CLIPEncoderBlock(layers.Layer):
     def call(self, x, training=None):
         residual = x
         x = self.layer_norm_1(x)
-        x = self.attention(x, x, x, training=training, use_causal_mask=True)
+        x = self.attention(
+            x, x, x, training=training, use_causal_mask=self.use_causal_mask
+        )
         x = ops.add(residual, x)
 
         residual = x
@@ -91,6 +105,7 @@ class CLIPEncoderBlock(layers.Layer):
                 "num_heads": self.num_heads,
                 "intermediate_dim": self.intermediate_dim,
                 "intermediate_activation": self.intermediate_activation,
+                "use_causal_mask": self.use_causal_mask,
             }
         )
         return config

keras_hub/src/models/clip/clip_image_converter.py

@@ -0,0 +1,8 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.clip.clip_backbone import CLIPBackbone
+
+
+@keras_hub_export("keras_hub.layers.CLIPImageConverter")
+class CLIPImageConverter(ImageConverter):
+    backbone_cls = CLIPBackbone

keras_hub/src/models/clip/clip_presets.py

@@ -0,0 +1,93 @@
+"""CLIP model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "clip_vit_base_patch16": {
+        "metadata": {
+            "description": (
+                "150 million parameter, 12-layer for vision and 12-layer for "
+                "text, patch size of 16, CLIP model."
+            ),
+            "params": 149620934,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_base_patch16/1",
+    },
+    "clip_vit_base_patch32": {
+        "metadata": {
+            "description": (
+                "151 million parameter, 12-layer for vision and 12-layer for "
+                "text, patch size of 32, CLIP model."
+            ),
+            "params": 151277363,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_base_patch32/1",
+    },
+    "clip_vit_large_patch14": {
+        "metadata": {
+            "description": (
+                "428 million parameter, 24-layer for vision and 12-layer for "
+                "text, patch size of 14, CLIP model."
+            ),
+            "params": 427616770,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_large_patch14/1",
+    },
+    "clip_vit_large_patch14_336": {
+        "metadata": {
+            "description": (
+                "428 million parameter, 24-layer for vision and 12-layer for "
+                "text, patch size of 14, image size of 336, CLIP model."
+            ),
+            "params": 427944770,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_large_patch14_336/1",
+    },
+    "clip_vit_b_32_laion2b_s34b_b79k": {
+        "metadata": {
+            "description": (
+                "151 million parameter, 12-layer for vision and 12-layer for "
+                "text, patch size of 32, Open CLIP model."
+            ),
+            "params": 151277363,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_b_32_laion2b_s34b_b79k/1",
+    },
+    "clip_vit_h_14_laion2b_s32b_b79k": {
+        "metadata": {
+            "description": (
+                "986 million parameter, 32-layer for vision and 24-layer for "
+                "text, patch size of 14, Open CLIP model."
+            ),
+            "params": 986109698,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_h_14_laion2b_s32b_b79k/1",
+    },
+    "clip_vit_g_14_laion2b_s12b_b42k": {
+        "metadata": {
+            "description": (
+                "1.4 billion parameter, 40-layer for vision and 24-layer for "
+                "text, patch size of 14, Open CLIP model."
+            ),
+            "params": 1366678530,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_g_14_laion2b_s12b_b42k/1",
+    },
+    "clip_vit_bigg_14_laion2b_39b_b160k": {
+        "metadata": {
+            "description": (
+                "2.5 billion parameter, 48-layer for vision and 32-layer for "
+                "text, patch size of 14, Open CLIP model."
+            ),
+            "params": 2539567362,
+            "path": "clip",
+        },
+        "kaggle_handle": "kaggle://keras/clip/keras/clip_vit_bigg_14_laion2b_39b_b160k/1",
+    },
+}

keras_hub/src/models/clip/clip_text_encoder.py

@@ -1,5 +1,6 @@
 from keras import layers
 
+from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
@@ -7,6 +8,7 @@ from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.clip.clip_encoder_block import CLIPEncoderBlock
 
 
+@keras_hub_export("keras_hub.models.CLIPTextEncoder")
 class CLIPTextEncoder(Backbone):
     """CLIP text core network with hyperparameters.
 
@@ -80,7 +82,7 @@ class CLIPTextEncoder(Backbone):
             for i in range(num_layers)
         ]
         self.layer_norm = layers.LayerNormalization(
-            epsilon=1e-6, dtype="float32", name=f"{prefix}layer_norm"
+            epsilon=1e-6, dtype=dtype, name=f"{prefix}layer_norm"
         )
 
         # === Functional Model ===
@@ -106,6 +108,7 @@ class CLIPTextEncoder(Backbone):
         super().__init__(
             inputs={"token_ids": token_id_input},
             outputs=outputs,
+            dtype=dtype,
             name=name,
             **kwargs,
         )

keras_hub/src/models/clip/clip_tokenizer.py

@@ -1,4 +1,5 @@
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.clip.clip_backbone import CLIPBackbone
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_hub.src.tokenizers.byte_pair_tokenizer import convert_to_ragged_batch
 from keras_hub.src.tokenizers.byte_pair_tokenizer import split_strings_for_bpe
@@ -39,11 +40,25 @@ class CLIPTokenizer(BytePairTokenizer):
             should have one merge rule per line. Every merge rule contains
             merge entities separated by a space.
         pad_with_end_token: bool. Whether to pad the output with `end_token`.
-    """
 
-    # TODO: Add example and `backbone_cls` once we have a CLIP model.
+    Examples:
+
+    ```python
+    # Unbatched input.
+    tokenizer = keras_hub.models.CLIPTokenizer.from_preset(
+        "clip_vit_base_patch32"
+    )
+    tokenizer("The quick brown fox jumped.")
+
+    # Batched input.
+    tokenizer(["The quick brown fox jumped.", "The fox slept."])
+
+    # Detokenization.
+    tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
+    ```
+    """
 
-    backbone_cls = None
+    backbone_cls = CLIPBackbone
 
     def __init__(
         self,

keras_hub/src/models/clip/clip_vision_embedding.py

@@ -0,0 +1,101 @@
+from keras import layers
+from keras import ops
+
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class CLIPVisionEmbedding(layers.Layer):
+    def __init__(
+        self,
+        hidden_dim,
+        patch_size,
+        image_size,
+        data_format=None,
+        dtype=None,
+        **kwargs
+    ):
+        super().__init__(dtype=dtype, **kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.patch_size = int(patch_size)
+        self.image_size = int(image_size)
+        data_format = standardize_data_format(data_format)
+        self.data_format = data_format
+        num_patches = (image_size // patch_size) ** 2
+        self.num_positions = num_patches + 1
+
+        self.patch_embedding = layers.Conv2D(
+            hidden_dim,
+            kernel_size=patch_size,
+            strides=patch_size,
+            data_format=data_format,
+            use_bias=False,
+            dtype=dtype,
+            name="patch_embedding",
+        )
+        self.position_embedding = layers.Embedding(
+            num_patches + 1, hidden_dim, dtype=dtype, name="position_embedding"
+        )
+
+    def build(self, input_shape):
+        self.class_embedding = self.add_weight(
+            shape=(self.hidden_dim,),
+            initializer="random_normal",
+            dtype=self.variable_dtype,
+            name="class_embedding",
+        )
+        self.position_ids = self.add_weight(
+            shape=(1, self.num_positions),
+            initializer="zeros",
+            # Let the backend determine the int dtype. For example, tf
+            # requires int64 for correct device placement, whereas jax and torch
+            # don't.
+            dtype=int,
+            trainable=False,
+            name="position_ids",
+        )
+        self.patch_embedding.build(input_shape)
+        self.position_embedding.build(self.position_ids.shape)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        batch_size = ops.shape(x)[0]
+        patch_embeddings = self.patch_embedding(x, training=training)
+        if self.data_format == "channels_last":
+            patch_embeddings = ops.reshape(
+                patch_embeddings, (batch_size, -1, self.hidden_dim)
+            )
+        else:
+            patch_embeddings = ops.reshape(
+                patch_embeddings, (batch_size, self.hidden_dim, -1)
+            )
+            patch_embeddings = ops.transpose(patch_embeddings, (0, 2, 1))
+        class_embeddings = ops.expand_dims(self.class_embedding, axis=(0, 1))
+        class_embeddings = ops.tile(class_embeddings, (batch_size, 1, 1))
+        position_embeddings = self.position_embedding(self.position_ids)
+        embeddings = ops.concatenate(
+            [class_embeddings, patch_embeddings], axis=1
+        )
+        return ops.add(embeddings, position_embeddings)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "patch_size": self.patch_size,
+                "image_size": self.image_size,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        output_shape = [input_shape[0], None, self.hidden_dim]
+        if self.data_format == "channels_last":
+            if input_shape[1] is not None and input_shape[2] is not None:
+                patch_num = input_shape[1] // self.patch_size
+                output_shape[1] = patch_num**2 + 1
+        else:
+            if input_shape[2] is not None and input_shape[3] is not None:
+                patch_num = input_shape[2] // self.patch_size
+                output_shape[1] = patch_num**2 + 1
+        return output_shape