keras-hub-nightly 0.22.0.dev202507150421__py3-none-any.whl → 0.22.0.dev202507170424__py3-none-any.whl

keras_hub/src/models/dinov2/dinov2_layers.py

@@ -0,0 +1,886 @@
+from keras import backend
+from keras import config
+from keras import initializers
+from keras import layers
+from keras import ops
+from keras import random
+
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class DINOV2PatchEmbedding(layers.Layer):
+    """A layer that converts images into patches.
+
+    Args:
+        hidden_dim: int. The number of units in the hidden layers.
+        patch_size: int. The size of one side of each patch.
+        data_format: `None` or str. If specified, either `"channels_last"` or
+            `"channels_first"`. The ordering of the dimensions in the
+            inputs. `"channels_last"` corresponds to inputs with shape
+            `(batch_size, height, width, channels)`
+            while `"channels_first"` corresponds to inputs with shape
+            `(batch_size, channels, height, width)`. It defaults to the
+            `image_data_format` value found in your Keras config file at
+            `~/.keras/keras.json`. If you never set it, then it will be
+            `"channels_last"`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(self, hidden_dim, patch_size, data_format=None, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.patch_size = int(patch_size)
+        self.data_format = standardize_data_format(data_format)
+
+        self.projection = layers.Conv2D(
+            hidden_dim,
+            kernel_size=patch_size,
+            strides=patch_size,
+            data_format=data_format,
+            kernel_initializer=initializers.TruncatedNormal(stddev=0.02),
+            dtype=self.dtype_policy,
+            name="projection",
+        )
+
+    def build(self, input_shape):
+        self.projection.build(input_shape)
+
+    def call(self, inputs, training=None):
+        batch_size = ops.shape(inputs)[0]
+        embeddings = self.projection(inputs, training=training)
+        if self.data_format == "channels_last":
+            embeddings = ops.reshape(
+                embeddings, (batch_size, -1, self.hidden_dim)
+            )
+        else:
+            embeddings = ops.reshape(
+                embeddings, (batch_size, self.hidden_dim, -1)
+            )
+            embeddings = ops.transpose(embeddings, (0, 2, 1))
+        return embeddings
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "patch_size": self.patch_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
+        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
+        return output_shape
+
+
+class DINOV2Embedding(layers.Layer):
+    """A layer that converts images into patches.
+
+    This layer adds all the necessary tokens to the embeddings, inlcuding
+    the class token, register tokens and mask token if specified. Finally, a
+    position embedding will be added.
+
+    This layer supports the interpolation of the position embeddings to enable
+    the model to work with images of different sizes. Please refer to
+    `_interpolate_position_embeddings` for more details.
+
+    The saving and loading of this layer will automatically handle the position
+    embeddings interpolation. Please refer to `save_own_variables` and
+    `load_own_variables` for more details.
+
+    Args:
+        hidden_dim: int. The number of units in the hidden layers.
+        patch_size: int. The size of one side of each patch.
+        image_size: tuple of ints. The (height, width) of the input images.
+        num_register_tokens: int. The number of register tokens to add to the
+            embeddings. Defaults to `0`.
+        use_mask_token: bool. Whether to use a mask token. Defaults to `True`.
+        dropout_rate: float. The dropout rate to use. Defaults to `0.0`.
+        position_embedding_shape: tuple. The original input shape used to
+            train the position embeddings. This is used to interpolate the
+            position embeddings to the actual input shape. Defaults to
+            `(518, 518)`.
+        antialias_in_interpolation: bool. Whether to use antialiasing in the
+            interpolation of the position embeddings. Defaults to `False`.
+        data_format: `None` or str. If specified, either `"channels_last"` or
+            `"channels_first"`. The ordering of the dimensions in the
+            inputs. `"channels_last"` corresponds to inputs with shape
+            `(batch_size, height, width, channels)`
+            while `"channels_first"` corresponds to inputs with shape
+            `(batch_size, channels, height, width)`. It defaults to the
+            `image_data_format` value found in your Keras config file at
+            `~/.keras/keras.json`. If you never set it, then it will be
+            `"channels_last"`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        patch_size,
+        image_shape,
+        num_register_tokens=0,
+        use_mask_token=True,
+        dropout_rate=0.0,
+        position_embedding_shape=(518, 518),
+        antialias_in_interpolation=False,
+        data_format=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.patch_size = int(patch_size)
+        self.image_shape = (int(image_shape[0]), int(image_shape[1]))
+        self.position_embedding_shape = (
+            int(position_embedding_shape[0]),
+            int(position_embedding_shape[1]),
+        )
+        self.num_register_tokens = int(num_register_tokens)
+        self.use_mask_token = bool(use_mask_token)
+        self.dropout_rate = float(dropout_rate)
+        self.antialias_in_interpolation = bool(antialias_in_interpolation)
+        self.data_format = standardize_data_format(data_format)
+        self.interpolated_num_patches = (
+            self.image_shape[0] // self.patch_size
+        ) * (self.image_shape[1] // self.patch_size)
+        self.num_patches = (
+            self.position_embedding_shape[0] // self.patch_size
+        ) * (self.position_embedding_shape[1] // self.patch_size)
+
+        self.patch_embeddings = DINOV2PatchEmbedding(
+            hidden_dim,
+            patch_size,
+            data_format=data_format,
+            dtype=self.dtype_policy,
+            name="patch_embeddings",
+        )
+        self.dropout = layers.Dropout(
+            rate=self.dropout_rate,
+            dtype=self.dtype_policy,
+            name="dropout",
+        )
+
+    def build(self, input_shape):
+        self.cls_token = self.add_weight(
+            shape=(1, 1, self.hidden_dim),
+            initializer=initializers.TruncatedNormal(stddev=0.02),
+            trainable=True,
+            name="cls_token",
+        )
+        if self.use_mask_token:
+            self.mask_token = self.add_weight(
+                shape=(1, self.hidden_dim),
+                initializer="zeros",
+                trainable=True,
+                name="mask_token",
+            )
+        if self.num_register_tokens > 0:
+            self.register_tokens = self.add_weight(
+                shape=(1, self.num_register_tokens, self.hidden_dim),
+                initializer="zeros",
+                trainable=True,
+                name="register_tokens",
+            )
+        self.patch_embeddings.build(input_shape)
+
+        # Note that there are two position embeddings:
+        # `self.interpolated_position_embeddings` is used for the image inputs
+        # during both training and inference.
+        # `self.position_embeddings` is used to load pretrained weights and
+        # remains unchanged during training and inference. It will be updated
+        # during saving once `self.interpolated_position_embeddings` is
+        # modified.
+        self.position_embeddings = self.add_weight(
+            shape=(1, self.num_patches + 1, self.hidden_dim),
+            initializer=initializers.TruncatedNormal(stddev=0.02),
+            trainable=False,
+            name="position_embeddings",
+        )
+        self.interpolated_position_embeddings = self.add_weight(
+            shape=(1, self.interpolated_num_patches + 1, self.hidden_dim),
+            initializer="zeros",  # Will be initialized by interpolation.
+            trainable=True,
+            name="interpolated_position_embeddings",
+        )
+
+        # Initialize the interpolated position embeddings.
+        self.interpolated_position_embeddings.assign(
+            self._interpolate_position_embeddings(
+                self.position_embeddings,
+                patch_size=self.patch_size,
+                source_shape=self.position_embedding_shape,
+                target_shape=self.image_shape,
+                antialias=self.antialias_in_interpolation,
+            )
+        )
+
+    def call(self, inputs, masks=None, training=None):
+        batch_size = ops.shape(inputs)[0]
+        embeddings = self.patch_embeddings(inputs, training=training)
+
+        # Repalce the embeddings with the mask tokens if specified.
+        # Basically, this is only used during training.
+        if masks is not None and self.use_mask_token:
+            masks = ops.expand_dims(masks, axis=-1)
+            mask_token = ops.cast(
+                ops.expand_dims(self.mask_token, axis=0), embeddings.dtype
+            )
+            embeddings = ops.where(masks, mask_token, embeddings)
+
+        # Add the [CLS] token to the embedded patch tokens.
+        cls_tokens = ops.tile(self.cls_token, (batch_size, 1, 1))
+        embeddings = ops.concatenate((cls_tokens, embeddings), axis=1)
+
+        # Add positional encoding to each token.
+        embeddings = ops.add(embeddings, self.interpolated_position_embeddings)
+
+        # Add register tokens if specified.
+        if self.num_register_tokens > 0:
+            register_tokens = ops.tile(self.register_tokens, (batch_size, 1, 1))
+            embeddings = ops.concatenate(
+                (
+                    embeddings[:, :1, ...],
+                    register_tokens,
+                    embeddings[:, 1:, ...],
+                ),
+                axis=1,
+            )
+
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "patch_size": self.patch_size,
+                "image_shape": self.image_shape,
+                "num_register_tokens": self.num_register_tokens,
+                "use_mask_token": self.use_mask_token,
+                "dropout_rate": self.dropout_rate,
+                "position_embedding_shape": self.position_embedding_shape,
+                "antialias_in_interpolation": self.antialias_in_interpolation,
+            }
+        )
+        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
+                # 1 is for cls token.
+                output_shape[1] = 1 + self.num_register_tokens + patch_num**2
+        else:
+            if input_shape[2] is not None and input_shape[3] is not None:
+                patch_num = input_shape[2] // self.patch_size
+                # 1 is for cls token.
+                output_shape[1] = 1 + self.num_register_tokens + patch_num**2
+        return output_shape
+
+    @staticmethod
+    def _interpolate_position_embeddings(
+        position_embeddings,
+        patch_size,
+        source_shape,
+        target_shape,
+        antialias=False,
+    ):
+        """Interpolate position embeddings to match the target image shape.
+
+        Reference:
+            - https://github.com/huggingface/transformers/blob/main/src/transformers/models/dinov2/modeling_dinov2.py
+        """
+        position_embeddings = ops.convert_to_tensor(position_embeddings)
+        patch_size = int(patch_size)
+        source_shape = (int(source_shape[0]), int(source_shape[1]))
+        target_shape = (int(target_shape[0]), int(target_shape[1]))
+        hidden_dim = int(position_embeddings.shape[-1])
+
+        if (
+            source_shape[0] == target_shape[0]
+            and source_shape[1] == target_shape[1]
+        ):
+            # No need to interpolate if the image size is the same as the
+            # position embedding image size.
+            return ops.copy(position_embeddings)
+
+        num_positions = int(position_embeddings.shape[1]) - 1
+
+        # Handle class token and patch embeddings separately.
+        class_position_embeddings = position_embeddings[:, :1, ...]
+        patch_position_embeddings = position_embeddings[:, 1:, ...]
+
+        # Calculate new dimensions
+        new_height = target_shape[0] // patch_size
+        new_width = target_shape[1] // patch_size
+
+        # Reshape for interpolation
+        sqrt_num_positions = int(num_positions**0.5)
+        patch_position_embeddings = ops.reshape(
+            patch_position_embeddings,
+            (1, sqrt_num_positions, sqrt_num_positions, hidden_dim),
+        )
+
+        # Interpolate at float32 precision.
+        original_dtype = backend.standardize_dtype(
+            patch_position_embeddings.dtype
+        )
+        interpolated_patch_position_embeddings = ops.image.resize(
+            ops.cast(patch_position_embeddings, "float32"),
+            size=(new_height, new_width),
+            interpolation="bicubic",
+            antialias=antialias,
+            data_format="channels_last",
+        )
+        interpolated_patch_position_embeddings = ops.cast(
+            interpolated_patch_position_embeddings, original_dtype
+        )
+
+        # Reshape back to the original format
+        interpolated_patch_position_embeddings = ops.reshape(
+            interpolated_patch_position_embeddings, (1, -1, hidden_dim)
+        )
+        interpolated_position_embeddings = ops.concatenate(
+            (class_position_embeddings, interpolated_patch_position_embeddings),
+            axis=1,
+        )
+        return interpolated_position_embeddings
+
+    def _is_interpolated_position_embeddings_updated(self):
+        """Check if the interpolated position embeddings are updated."""
+        original_interpolated_position_embeddings = (
+            self._interpolate_position_embeddings(
+                self.position_embeddings,
+                patch_size=self.patch_size,
+                source_shape=self.position_embedding_shape,
+                target_shape=self.image_shape,
+                antialias=self.antialias_in_interpolation,
+            )
+        )
+        diff = ops.sum(
+            ops.subtract(
+                original_interpolated_position_embeddings,
+                self.interpolated_position_embeddings,
+            )
+        )
+        return ops.cond(
+            ops.greater(diff, config.epsilon()), lambda: True, lambda: False
+        )
+
+    def save_own_variables(self, store):
+        if self._is_interpolated_position_embeddings_updated():
+            self.position_embeddings.assign(
+                self._interpolate_position_embeddings(
+                    self.interpolated_position_embeddings,
+                    patch_size=self.patch_size,
+                    source_shape=self.image_shape,
+                    target_shape=self.position_embedding_shape,
+                    antialias=self.antialias_in_interpolation,
+                )
+            )
+        super().save_own_variables(store)
+
+    def load_own_variables(self, store):
+        all_vars = self._trainable_variables + self._non_trainable_variables
+        for i, v in enumerate(all_vars):
+            if v is self.interpolated_position_embeddings:
+                continue
+            v.assign(store[f"{i}"])
+        self.interpolated_position_embeddings.assign(
+            self._interpolate_position_embeddings(
+                self.position_embeddings,
+                patch_size=self.patch_size,
+                source_shape=self.position_embedding_shape,
+                target_shape=self.image_shape,
+                antialias=self.antialias_in_interpolation,
+            )
+        )
+
+
+class DINOV2Attention(layers.Layer):
+    """A multi-head attention layer with dropout.
+
+    Args:
+        hidden_dim: int. The number of units in the hidden layers.
+        num_heads: int. Number of attention heads.
+        dropout_rate: float. The dropout rate to use. Defaults to `0.0`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(self, hidden_dim, num_heads, dropout_rate=0.0, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.num_heads = int(num_heads)
+        self.dropout_rate = float(dropout_rate)
+
+        self.attention = layers.MultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=self.hidden_dim // self.num_heads,
+            dropout=self.dropout_rate,
+            dtype=self.dtype_policy,
+            name="attention",
+        )
+        self.dropout = layers.Dropout(
+            rate=self.dropout_rate,
+            dtype=self.dtype_policy,
+            name="dropout",
+        )
+
+    def build(self, input_shape):
+        self.attention.build(input_shape, input_shape)
+
+    def call(self, inputs, training=None):
+        attention_output = self.attention(
+            query=inputs,
+            value=inputs,
+            key=inputs,
+            training=training,
+            use_causal_mask=False,
+        )
+        outputs = self.dropout(attention_output, training=training)
+        return outputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "num_heads": self.num_heads,
+                "dropout_rate": self.dropout_rate,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class DINOV2LayerScale(layers.Layer):
+    """A layer scale.
+
+    Args:
+        hidden_dim: int. The number of units in the hidden layers.
+        init_values: float. The initial value for the scale. Defaults to `1.0`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(self, hidden_dim, init_values=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.init_values = float(init_values)
+
+    def build(self, input_shape):
+        self.lambda1 = self.add_weight(
+            shape=(self.hidden_dim,),
+            initializer=initializers.Constant(self.init_values),
+            trainable=True,
+            name="lambda1",
+        )
+
+    def call(self, inputs, training=None):
+        return ops.multiply(inputs, self.lambda1)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"hidden_dim": self.hidden_dim})
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class DINOV2DropPath(layers.Layer):
+    """A drop path layer.
+
+    Args:
+        rate: float. The drop path rate to use. Defaults to `0.0`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(self, rate=0.0, **kwargs):
+        super().__init__(**kwargs)
+        self.rate = float(rate)
+
+    def build(self, input_shape):
+        self.noise_shape = (input_shape[0],) + (1,) * (len(input_shape) - 1)
+
+    def call(self, inputs, training=None):
+        if not training or self.rate == 0.0:
+            return inputs
+
+        keep_prob = 1.0 - self.rate
+        random_tensor = random.uniform(self.noise_shape, dtype=inputs.dtype)
+        random_tensor = ops.add(random_tensor, keep_prob)
+        return ops.multiply(ops.divide(inputs, keep_prob), random_tensor)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"rate": self.rate})
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class DINOV2MLP(layers.Layer):
+    """A DINOV2 MLP block.
+
+    Args:
+        hidden_dim: int. The number of units in the output layer.
+        intermediate_dim: int. The output dimension of the first Dense layer.
+        activation: str of callable. Activation to use in the intermediate
+            layer. Defaults to `"gelu"`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(
+        self, hidden_dim, intermediate_dim, activation="gelu", **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.intermediate_dim = int(intermediate_dim)
+        self.activation = activation
+
+        self.fc1 = layers.Dense(
+            self.intermediate_dim,
+            activation=activation,
+            kernel_initializer=initializers.TruncatedNormal(stddev=0.02),
+            dtype=self.dtype_policy,
+            name="fc1",
+        )
+        self.fc2 = layers.Dense(
+            self.hidden_dim,
+            kernel_initializer=initializers.TruncatedNormal(stddev=0.02),
+            dtype=self.dtype_policy,
+            name="fc2",
+        )
+
+    def build(self, input_shape):
+        self.fc1.build(input_shape)
+        input_shape = self.fc1.compute_output_shape(input_shape)
+        self.fc2.build(input_shape)
+
+    def call(self, inputs, training=None):
+        x = self.fc1(inputs, training=training)
+        x = self.fc2(x, training=training)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        output_shape = list(input_shape)
+        output_shape[-1] = self.hidden_dim
+        return output_shape
+
+
+class DINOV2SwiGLUFFN(layers.Layer):
+    """A DINOV2 SwiGLU Feed-Forward Network layer.
+
+    Please refer to [GLU Variants Improve Transformer](
+    https://arxiv.org/abs/2002.05202) for more details on SwiGLU.
+
+    Args:
+        hidden_dim: int. The number of units in the output layer.
+        intermediate_dim: int. The output dimension of the first Dense layer.
+            Note that this value will be multiplied by `2 / 3` and rounded up to
+            the nearest multiple of `8`. The reason for this is that SwiGLUFFN
+            achieves similar or better performance with fewer parameters
+            compared to the original FFN implementation.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(self, hidden_dim, intermediate_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.intermediate_dim = int(intermediate_dim)
+        self.actual_intermediate_dim = (
+            (int(intermediate_dim * 2 / 3) + 7) // 8 * 8
+        )
+
+        self.weights_in = layers.Dense(
+            2 * self.actual_intermediate_dim,
+            kernel_initializer=initializers.TruncatedNormal(stddev=0.02),
+            dtype=self.dtype_policy,
+            name="weights_in",
+        )
+        self.weights_out = layers.Dense(
+            self.hidden_dim,
+            kernel_initializer=initializers.TruncatedNormal(stddev=0.02),
+            dtype=self.dtype_policy,
+            name="weights_out",
+        )
+
+    def build(self, input_shape):
+        self.weights_in.build(input_shape)
+        input_shape = list(input_shape)
+        input_shape[-1] = self.actual_intermediate_dim
+        self.weights_out.build(input_shape)
+
+    def call(self, inputs, training=None):
+        x = self.weights_in(inputs, training=training)
+        x1, x2 = ops.split(x, 2, axis=-1)
+        x = ops.multiply(ops.silu(x1), x2)
+        x = self.weights_out(x, training=training)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        output_shape = list(input_shape)
+        output_shape[-1] = self.hidden_dim
+        return output_shape
+
+
+class DINOV2Layer(layers.Layer):
+    """A DINOV2 encoder layer.
+
+    Args:
+        hidden_dim: int. The number of units in the hidden layers.
+        num_heads: int. Number of attention heads.
+        layer_scale_init_value: float. The initial value for the scale.
+            Defaults to `1.0`.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        use_swiglu_ffn: bool. Whether to use SwigLUFFN instead of MLP.
+            Defaults to `False`.
+        dropout_rate: float. The dropout rate to use. Defaults to `0.0`.
+        drop_path_rate: float. The drop path rate to use. Defaults to `0.0`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        num_heads,
+        intermediate_dim,
+        layer_scale_init_value=1.0,
+        use_swiglu_ffn=False,
+        dropout_rate=0.0,
+        drop_path_rate=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = int(hidden_dim)
+        self.num_heads = int(num_heads)
+        self.intermediate_dim = int(intermediate_dim)
+        self.layer_scale_init_value = float(layer_scale_init_value)
+        self.use_swiglu_ffn = bool(use_swiglu_ffn)
+        self.dropout_rate = float(dropout_rate)
+        self.drop_path_rate = float(drop_path_rate)
+
+        self.norm1 = layers.LayerNormalization(
+            epsilon=1e-6, dtype=self.dtype_policy, name="norm1"
+        )
+        self.attention = DINOV2Attention(
+            hidden_dim=self.hidden_dim,
+            num_heads=self.num_heads,
+            dropout_rate=self.dropout_rate,
+            dtype=self.dtype_policy,
+            name="attention",
+        )
+        self.layer_scale1 = DINOV2LayerScale(
+            hidden_dim=self.hidden_dim,
+            init_values=self.layer_scale_init_value,
+            dtype=self.dtype_policy,
+            name="layer_scale1",
+        )
+        self.drop_path = (
+            DINOV2DropPath(
+                rate=self.drop_path_rate,
+                dtype=self.dtype_policy,
+                name="drop_path",
+            )
+            if self.drop_path_rate > 0
+            else layers.Identity(dtype=self.dtype_policy, name="drop_path")
+        )
+        self.norm2 = layers.LayerNormalization(
+            epsilon=1e-6, dtype=self.dtype_policy, name="norm2"
+        )
+        if self.use_swiglu_ffn:
+            self.mlp = DINOV2SwiGLUFFN(
+                hidden_dim=self.hidden_dim,
+                intermediate_dim=self.intermediate_dim,
+                dtype=self.dtype_policy,
+                name="mlp",
+            )
+        else:
+            self.mlp = DINOV2MLP(
+                hidden_dim=self.hidden_dim,
+                intermediate_dim=self.intermediate_dim,
+                activation="gelu",
+                dtype=self.dtype_policy,
+                name="mlp",
+            )
+        self.layer_scale2 = DINOV2LayerScale(
+            hidden_dim=self.hidden_dim,
+            init_values=self.layer_scale_init_value,
+            dtype=self.dtype_policy,
+            name="layer_scale2",
+        )
+
+    def build(self, input_shape):
+        self.norm1.build(input_shape)
+        self.attention.build(input_shape)
+        input_shape = self.attention.compute_output_shape(input_shape)
+        self.layer_scale1.build(input_shape)
+        self.drop_path.build(input_shape)
+        self.norm2.build(input_shape)
+        self.mlp.build(input_shape)
+        input_shape = self.mlp.compute_output_shape(input_shape)
+        self.layer_scale2.build(input_shape)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        x = self.norm1(x, training=training)
+        x = self.attention(x, training=training)
+        x = self.layer_scale1(x, training=training)
+
+        # First residual connection.
+        hidden_states = ops.add(self.drop_path(x, training=training), inputs)
+        x = self.norm2(hidden_states, training=training)
+        x = self.mlp(x, training=training)
+        x = self.layer_scale2(x, training=training)
+
+        # Second residual connection.
+        return ops.add(self.drop_path(x, training=training), hidden_states)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "num_heads": self.num_heads,
+                "intermediate_dim": self.intermediate_dim,
+                "layer_scale_init_value": self.layer_scale_init_value,
+                "use_swiglu_ffn": self.use_swiglu_ffn,
+                "dropout_rate": self.dropout_rate,
+                "drop_path_rate": self.drop_path_rate,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+class DINOV2Encoder(layers.Layer):
+    """A DINOV2 encoder.
+
+    Args:
+        num_layers: int. The number of transformer layers.
+        hidden_dim: int. The number of units in the hidden layers.
+        num_heads: int. Number of attention heads.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        layer_scale_init_value: float. The initial value for the scale.
+            Defaults to `1.0`.
+        use_swiglu_ffn: bool. Whether to use SwigLUFFN instead of MLP.
+            Defaults to `False`.
+        dropout_rate: float. The dropout rate to use. Defaults to `0.0`.
+        drop_path_rate: float. The drop path rate to use. Defaults to `0.0`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        hidden_dim,
+        num_heads,
+        intermediate_dim,
+        layer_scale_init_value=1.0,
+        use_swiglu_ffn=False,
+        dropout_rate=0.0,
+        drop_path_rate=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_layers = int(num_layers)
+        self.hidden_dim = int(hidden_dim)
+        self.num_heads = int(num_heads)
+        self.intermediate_dim = int(intermediate_dim)
+        self.layer_scale_init_value = float(layer_scale_init_value)
+        self.use_swiglu_ffn = bool(use_swiglu_ffn)
+        self.dropout_rate = float(dropout_rate)
+        self.drop_path_rate = float(drop_path_rate)
+
+        self.layers = [
+            DINOV2Layer(
+                hidden_dim=self.hidden_dim,
+                num_heads=self.num_heads,
+                intermediate_dim=self.intermediate_dim,
+                layer_scale_init_value=self.layer_scale_init_value,
+                use_swiglu_ffn=self.use_swiglu_ffn,
+                dropout_rate=self.dropout_rate,
+                drop_path_rate=self.drop_path_rate,
+                dtype=self.dtype_policy,
+                name=f"layers_{i}",
+            )
+            for i in range(self.num_layers)
+        ]
+
+    def build(self, input_shape):
+        for layer in self.layers:
+            layer.build(input_shape)
+            input_shape = layer.compute_output_shape(input_shape)
+
+    def call(self, inputs, training=None):
+        x = inputs
+        for layer in self.layers:
+            x = layer(x, training=training)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_layers": self.num_layers,
+                "hidden_dim": self.hidden_dim,
+                "num_heads": self.num_heads,
+                "intermediate_dim": self.intermediate_dim,
+                "layer_scale_init_value": self.layer_scale_init_value,
+                "use_swiglu_ffn": self.use_swiglu_ffn,
+                "dropout_rate": self.dropout_rate,
+                "drop_path_rate": self.drop_path_rate,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape