keras-hub-nightly 0.22.0.dev202508170419__py3-none-any.whl → 0.24.0.dev202511090424__py3-none-any.whl

keras_hub/src/models/dinov3/dinov3_presets.py

@@ -0,0 +1,4 @@
+"""DINOV3 model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {}

keras_hub/src/models/gemma/gemma_backbone.py

@@ -114,7 +114,6 @@ class GemmaBackbone(Backbone):
                 scale=1.0,
                 mode="fan_in",
                 distribution="untruncated_normal",
-                seed=None,
             ),
             dtype=dtype,
             logit_soft_cap=final_logit_soft_cap,

keras_hub/src/models/gemma/gemma_presets.py

@@ -198,4 +198,34 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_27b_en/2",
     },
+    "vault_gemma_1b_en": {
+        "metadata": {
+            "description": "1 billion parameter, 26-layer, VaultGemma model.",
+            "params": 1038741120,
+            "path": "gemma",
+        },
+        "kaggle_handle": "kaggle://keras/vaultgemma/keras/vault_gemma_1b_en/2",
+    },
+    "c2s_scale_gemma_2_2b_en": {
+        "metadata": {
+            "description": (
+                "A 2 billion parameter, single-cell biology-aware model "
+                "built on the Gemma-2 architecture."
+            ),
+            "params": 2614341888,
+            "path": "gemma",
+        },
+        "kaggle_handle": "kaggle://keras/cell2sentence/keras/c2s_scale_gemma_2_2b_en/1",
+    },
+    "c2s_scale_gemma_2_27b_en": {
+        "metadata": {
+            "description": (
+                "A 27 billion parameter, single-cell biology-aware model "
+                "built on the Gemma-2 architecture."
+            ),
+            "params": 27227128320,
+            "path": "gemma",
+        },
+        "kaggle_handle": "kaggle://keras/cell2sentence/keras/c2s_scale_gemma_2_27b_en/1",
+    },
 }

keras_hub/src/models/gemma3/gemma3_attention.py

@@ -46,6 +46,7 @@ class CachedGemma3Attention(keras.layers.Layer):
         layer_norm_epsilon=1e-6,
         rope_wavelength=10_000.0,
         rope_scaling_factor=1.0,
+        use_bidirectional_attention=False,
         dropout=0,
         **kwargs,
     ):
@@ -61,6 +62,7 @@ class CachedGemma3Attention(keras.layers.Layer):
         self.layer_norm_epsilon = layer_norm_epsilon
         self.rope_wavelength = rope_wavelength
         self.rope_scaling_factor = rope_scaling_factor
+        self.use_bidirectional_attention = use_bidirectional_attention
         self.dropout = dropout
 
         self._kernel_initializer = keras.initializers.get(
@@ -240,12 +242,58 @@ class CachedGemma3Attention(keras.layers.Layer):
         results = ops.einsum("bkgts,bskh->btkgh", attention_softmax, v)
         return ops.reshape(results, (b, q_len, self.num_query_heads, h))
 
+    def _compute_bidirectional_sliding_mask(self, batch_size, sequence_length):
+        """Computes a bidirectional sliding window attention mask.
+
+        A token can attend to any other token if their absolute distance is
+        within  half the sliding window size. This mask is used in embedding
+        models like `EmbeddingGemma`.
+
+        Args:
+            batch_size: The batch size for the mask.
+            sequence_length: The length of the sequence.
+
+        Returns:
+            A boolean attention mask with shape
+            `(batch_size, sequence_length, sequence_length)`.
+        """
+        i = keras.ops.expand_dims(
+            keras.ops.arange(sequence_length, dtype="int32"), axis=1
+        )
+        j = keras.ops.arange(sequence_length, dtype="int32")
+
+        # If sliding window size is 4, the token in question attends to 1
+        # token before and 2 tokens after.
+        w_right = self.sliding_window_size // 2
+        w_left = self.sliding_window_size - w_right - 1
+
+        # Calculate the relative distance.
+        distance = i - j
+
+        mask = keras.ops.logical_and(distance <= w_left, distance >= -w_right)
+
+        mask = keras.ops.expand_dims(mask, axis=0)
+        return keras.ops.broadcast_to(
+            mask, (batch_size, sequence_length, sequence_length)
+        )
+
     def _mask_sliding_window(
         self,
         attention_mask,
         cache_update_index=0,
     ):
         batch_size, query_len, key_len = ops.shape(attention_mask)
+
+        if self.use_bidirectional_attention:
+            bidirectional_sliding_mask = (
+                self._compute_bidirectional_sliding_mask(
+                    batch_size=batch_size,
+                    # `query_len = key_len` for embedding models
+                    sequence_length=query_len,
+                )
+            )
+            return ops.logical_and(attention_mask, bidirectional_sliding_mask)
+
         # Compute the sliding window for square attention.
         all_ones = ops.ones((key_len, key_len), "bool")
         if keras.config.backend() == "tensorflow":

keras_hub/src/models/gemma3/gemma3_backbone.py

@@ -196,6 +196,7 @@ class Gemma3Backbone(Backbone):
         global_rope_scaling_factor=1.0,
         vision_encoder=None,
         layer_norm_epsilon=1e-6,
+        use_bidirectional_attention=False,
         dropout=0,
         dtype=None,
         **kwargs,
@@ -209,7 +210,6 @@ class Gemma3Backbone(Backbone):
                 scale=1.0,
                 mode="fan_in",
                 distribution="untruncated_normal",
-                seed=None,
             ),
             dtype=dtype,
             logit_soft_cap=final_logit_soft_cap,
@@ -251,6 +251,7 @@ class Gemma3Backbone(Backbone):
                 sliding_window_size=sliding_window_size,
                 rope_wavelength=rope_wavelength,
                 rope_scaling_factor=rope_scaling_factor,
+                use_bidirectional_attention=use_bidirectional_attention,
                 dropout=dropout,
                 dtype=dtype,
                 name=f"decoder_block_{i}",
@@ -357,6 +358,7 @@ class Gemma3Backbone(Backbone):
         self.sliding_window_size = sliding_window_size
         self.local_rope_scaling_factor = local_rope_scaling_factor
         self.global_rope_scaling_factor = global_rope_scaling_factor
+        self.use_bidirectional_attention = use_bidirectional_attention
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
 
@@ -396,6 +398,7 @@ class Gemma3Backbone(Backbone):
                 "vision_encoder": None
                 if self.vision_encoder is None
                 else keras.layers.serialize(self.vision_encoder),
+                "use_bidirectional_attention": self.use_bidirectional_attention,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,
             }

keras_hub/src/models/gemma3/gemma3_decoder_block.py

@@ -45,6 +45,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
         layer_norm_epsilon=1e-6,
         rope_wavelength=10_000.0,
         rope_scaling_factor=1.0,
+        use_bidirectional_attention=False,
         dropout=0,
         **kwargs,
     ):
@@ -66,6 +67,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
         self.layer_norm_epsilon = layer_norm_epsilon
         self.rope_wavelength = rope_wavelength
         self.rope_scaling_factor = rope_scaling_factor
+        self.use_bidirectional_attention = use_bidirectional_attention
         self.dropout = dropout
 
         self.pre_attention_norm = RMSNormalization(
@@ -93,6 +95,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
             rope_wavelength=rope_wavelength,
             rope_scaling_factor=rope_scaling_factor,
             dropout=dropout,
+            use_bidirectional_attention=use_bidirectional_attention,
             dtype=self.dtype_policy,
             name="attention",
         )
@@ -209,6 +212,14 @@ class Gemma3DecoderBlock(keras.layers.Layer):
         if cache is not None:
             input_length = ops.shape(cache)[2]
 
+        if self.use_bidirectional_attention:
+            # `output_length` and `input_length` will be the same in this case
+            # because we use bidirectional attention for models like
+            # `EmbeddingGemma` which aren't used for text generation.
+            mask_1 = decoder_mask
+            mask_2 = ops.transpose(mask_1, (0, 2, 1))
+            return mask_1 * mask_2
+
         causal_mask = compute_causal_mask(
             batch_size=batch_size,
             input_length=input_length,
@@ -304,6 +315,7 @@ class Gemma3DecoderBlock(keras.layers.Layer):
                 "dropout": self.dropout,
                 "rope_wavelength": self.rope_wavelength,
                 "rope_scaling_factor": self.rope_scaling_factor,
+                "use_bidirectional_attention": self.use_bidirectional_attention,
             }
         )
         return config

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -181,4 +181,43 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_270m/4",
     },
+    "medgemma_instruct_4b": {
+        "metadata": {
+            "description": (
+                "A 4 billion parameter model based on Gemma 3. "
+                "This model is trained for performance on medical text"
+                "and image comprehension and is optimized for medical"
+                "applications that involve a text generation component."
+            ),
+            "params": 4300079472,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_4b/1",
+    },
+    "medgemma_instruct_27b": {
+        "metadata": {
+            "description": (
+                "A 27 billion parameter model based on Gemma 3. "
+                "This model trained for performance on medical text "
+                "and image comprehension and is optimized for medical "
+                "applications that involve a text generation component."
+            ),
+            "params": 27432406640,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_27b/1",
+    },
+    "medgemma_instruct_27b_text": {
+        "metadata": {
+            "description": (
+                "A 27 billion parameter text-only model based on Gemma 3. "
+                "This model is trained for performance on medical text "
+                "comprehension and is optimized for medical applications "
+                "that involve a text generation component."
+            ),
+            "params": 27009002240,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_27b_text/1",
+    },
 }

keras_hub/src/models/hgnetv2/hgnetv2_backbone.py

@@ -157,7 +157,10 @@ class HGNetV2Backbone(Backbone):
             if stage_name in self.out_features
         }
         super().__init__(
-            inputs=pixel_values, outputs=feature_maps_output, **kwargs
+            inputs=pixel_values,
+            outputs=feature_maps_output,
+            dtype=dtype,
+            **kwargs,
         )
 
         # === Config ===

keras_hub/src/models/hgnetv2/hgnetv2_encoder.py

@@ -56,9 +56,10 @@ class HGNetV2Encoder(keras.layers.Layer):
         use_learnable_affine_block,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.stage_in_channels = stage_in_channels
         self.stage_mid_channels = stage_mid_channels
         self.stage_out_channels = stage_out_channels
@@ -90,7 +91,7 @@ class HGNetV2Encoder(keras.layers.Layer):
                 name=f"{self.name}_stage_{stage_idx}"
                 if self.name
                 else f"stage_{stage_idx}",
-                dtype=self.dtype,
+                dtype=dtype,
             )
             self.stages_list.append(stage_layer)
 

keras_hub/src/models/hgnetv2/hgnetv2_layers.py

@@ -17,8 +17,8 @@ class HGNetV2LearnableAffineBlock(keras.layers.Layer):
         **kwargs: Additional keyword arguments passed to the parent class.
     """
 
-    def __init__(self, scale_value=1.0, bias_value=0.0, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, scale_value=1.0, bias_value=0.0, dtype=None, **kwargs):
+        super().__init__(dtype=dtype, **kwargs)
         self.scale_value = scale_value
         self.bias_value = bias_value
 
@@ -87,9 +87,10 @@ class HGNetV2ConvLayer(keras.layers.Layer):
         use_learnable_affine_block=False,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.kernel_size = kernel_size
@@ -104,6 +105,7 @@ class HGNetV2ConvLayer(keras.layers.Layer):
             padding=((pad, pad), (pad, pad)),
             data_format=self.data_format,
             name=f"{self.name}_pad" if self.name else None,
+            dtype=self.dtype_policy,
         )
         self.convolution = keras.layers.Conv2D(
             filters=self.out_channels,
@@ -156,7 +158,8 @@ class HGNetV2ConvLayer(keras.layers.Layer):
             )
         else:
             self.lab = keras.layers.Identity(
-                name=f"{self.name}_identity_lab" if self.name else None
+                name=f"{self.name}_identity_lab" if self.name else None,
+                dtype=self.dtype_policy,
             )
 
     def build(self, input_shape):
@@ -230,9 +233,10 @@ class HGNetV2ConvLayerLight(keras.layers.Layer):
         use_learnable_affine_block=False,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.kernel_size = kernel_size
@@ -327,9 +331,10 @@ class HGNetV2Embeddings(keras.layers.Layer):
         use_learnable_affine_block,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.stem_channels = stem_channels
         self.hidden_act = hidden_act
         self.use_learnable_affine_block = use_learnable_affine_block
@@ -352,6 +357,7 @@ class HGNetV2Embeddings(keras.layers.Layer):
             padding=((0, 1), (0, 1)),
             data_format=self.data_format,
             name=f"{self.name}_padding1" if self.name else "padding1",
+            dtype=self.dtype_policy,
         )
         self.stem2a_layer = HGNetV2ConvLayer(
             in_channels=self.stem_channels[1],
@@ -370,6 +376,7 @@ class HGNetV2Embeddings(keras.layers.Layer):
             padding=((0, 1), (0, 1)),
             data_format=self.data_format,
             name=f"{self.name}_padding2" if self.name else "padding2",
+            dtype=self.dtype_policy,
         )
         self.stem2b_layer = HGNetV2ConvLayer(
             in_channels=self.stem_channels[1] // 2,
@@ -390,10 +397,12 @@ class HGNetV2Embeddings(keras.layers.Layer):
             padding="valid",
             data_format=self.data_format,
             name=f"{self.name}_pool" if self.name else "pool",
+            dtype=self.dtype_policy,
         )
         self.concatenate_layer = keras.layers.Concatenate(
             axis=self.channel_axis,
             name=f"{self.name}_concat" if self.name else "concat",
+            dtype=self.dtype_policy,
         )
         self.stem3_layer = HGNetV2ConvLayer(
             in_channels=self.stem_channels[1] * 2,
@@ -550,9 +559,10 @@ class HGNetV2BasicLayer(keras.layers.Layer):
         use_learnable_affine_block=False,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.in_channels_arg = in_channels
         self.middle_channels = middle_channels
         self.out_channels = out_channels
@@ -635,23 +645,27 @@ class HGNetV2BasicLayer(keras.layers.Layer):
                 self.drop_path_rate,
                 noise_shape=(None, 1, 1, 1),
                 name=f"{self.name}_drop_path" if self.name else "drop_path",
+                dtype=self.dtype_policy,
             )
         else:
             self.drop_path_layer = keras.layers.Identity(
                 name=f"{self.name}_identity_drop_path"
                 if self.name
-                else "identity_drop_path"
+                else "identity_drop_path",
+                dtype=self.dtype_policy,
             )
 
         self.concatenate_layer = keras.layers.Concatenate(
             axis=self.channel_axis,
             name=f"{self.name}_concat" if self.name else "concat",
+            dtype=self.dtype_policy,
         )
         if self.residual:
             self.add_layer = keras.layers.Add(
                 name=f"{self.name}_add_residual"
                 if self.name
-                else "add_residual"
+                else "add_residual",
+                dtype=self.dtype_policy,
             )
 
     def build(self, input_shape):
@@ -794,9 +808,10 @@ class HGNetV2Stage(keras.layers.Layer):
         drop_path: float = 0.0,
         data_format=None,
         channel_axis=None,
+        dtype=None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(dtype=dtype, **kwargs)
         self.stage_in_channels = stage_in_channels
         self.stage_mid_channels = stage_mid_channels
         self.stage_out_channels = stage_out_channels
@@ -842,7 +857,8 @@ class HGNetV2Stage(keras.layers.Layer):
             self.downsample_layer = keras.layers.Identity(
                 name=f"{self.name}_identity_downsample"
                 if self.name
-                else "identity_downsample"
+                else "identity_downsample",
+                dtype=self.dtype_policy,
             )
 
         self.blocks_list = []

keras_hub/src/models/image_to_image.py

@@ -415,3 +415,8 @@ class ImageToImage(Task):
         # Image-to-image.
         outputs = [generate(*x) for x in inputs]
         return self._normalize_generate_outputs(outputs, input_is_scalar)
+
+    def _post_quantize(self, mode, **kwargs):
+        super()._post_quantize(mode, **kwargs)
+        # Reset the compiled generate function.
+        self.generate_function = None

keras_hub/src/models/inpaint.py

@@ -518,3 +518,8 @@ class Inpaint(Task):
         # Inpaint.
         outputs = [generate(*x) for x in inputs]
         return self._normalize_generate_outputs(outputs, input_is_scalar)
+
+    def _post_quantize(self, mode, **kwargs):
+        super()._post_quantize(mode, **kwargs)
+        # Reset the compiled generate function.
+        self.generate_function = None

keras_hub/src/models/mobilenetv5/__init__.py

@@ -0,0 +1,9 @@
+from keras_hub.src.models.mobilenetv5.mobilenetv5_backbone import (
+    MobileNetV5Backbone,
+)
+from keras_hub.src.models.mobilenetv5.mobilenetv5_presets import (
+    backbone_presets,
+)
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, MobileNetV5Backbone)