Add Mistral3 model to Hub

keras_hub/api/layers/__init__.py

@@ -129,6 +129,9 @@
 from keras_hub.src.models.metaclip_2.metaclip_2_image_converter import (
     MetaCLIP2ImageConverter as MetaCLIP2ImageConverter,
 )
+from keras_hub.src.models.mistral.mistral_image_converter import (
+    Mistral3ImageConverter as Mistral3ImageConverter,
+)
 from keras_hub.src.models.mit.mit_image_converter import (
     MiTImageConverter as MiTImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -517,6 +517,9 @@
 from keras_hub.src.models.mistral.mistral_tokenizer import (
     MistralTokenizer as MistralTokenizer,
 )
+from keras_hub.src.models.mistral.mistral_vision_encoder import (
+    Mistral3VisionEncoder as Mistral3VisionEncoder,
+)
 from keras_hub.src.models.mit.mit_backbone import MiTBackbone as MiTBackbone
 from keras_hub.src.models.mit.mit_image_classifier import (
     MiTImageClassifier as MiTImageClassifier,

keras_hub/src/models/mistral/mistral_backbone.py

@@ -1,15 +1,19 @@
 import keras
 from keras import ops
 from keras.layers import ReversibleEmbedding
+from keras.layers import RMSNormalization
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.backbone import Backbone
-from keras_hub.src.models.mistral.mistral_layer_norm import (
-    MistralLayerNormalization,
-)
 from keras_hub.src.models.mistral.mistral_transformer_decoder import (
     MistralTransformerDecoder,
 )
+from keras_hub.src.models.mistral.mistral_vision_encoder import (
+    Mistral3ImageFeatureExtractor,
+)
+from keras_hub.src.models.mistral.mistral_vision_encoder import (
+    Mistral3ImageTextEmbeddingMerger,
+)
 
 
 def _mistral_kernel_initializer(stddev=0.02):
@@ -59,6 +63,28 @@ class MistralBackbone(Backbone):
             Set explicitly when the model's head size is not equal to
             `hidden_dim // num_query_heads` — e.g. Magistral uses
             `head_dim=128` with `hidden_dim=5120` and `num_query_heads=32`.
+        vision_encoder: A `keras_hub.models.Mistral3VisionEncoder` instance,
+            or `None` (the default) for a text-only model. When set,
+            `multimodal_projector` must also be provided, and `call()`
+            accepts additional `"pixel_values"`, `"image_sizes"`, and
+            `"placeholder_indices"` inputs. `"pixel_values"` has dynamic
+            spatial dimensions (padded to the batch's largest image, as
+            produced by HF's image processor), not a fixed canvas;
+            `"image_sizes"` gives each image's true `(height, width)` for
+            cropping after patchification. `"placeholder_indices"` gives
+            the flat positions of image placeholder tokens in `token_ids`
+            (into the flattened `batch * seq_length` sequence) that get
+            replaced with projected image features — compute it with
+            `mistral_vision_encoder.compute_image_placeholder_indices`
+            before calling the model, since deriving it in-graph would be
+            incompatible with `jax.jit` tracing.
+        multimodal_projector: A `Mistral3MultiModalProjector` instance.
+            Required when `vision_encoder` is set.
+        image_token_index (int, optional): The token ID in `token_ids` that
+            marks image placeholder positions. Defaults to `10`. Unused in
+            text-only mode; used together with
+            `compute_image_placeholder_indices` to build the
+            `"placeholder_indices"` model input.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
             for model computations and weights. Note that some computations,
             such as softmax and layer normalization, will always be done at
@@ -106,9 +132,19 @@ def __init__(
         sliding_window=512,
         head_dim=None,
         dropout=0,
+        vision_encoder=None,
+        multimodal_projector=None,
+        image_token_index=10,
         dtype=None,
         **kwargs,
     ):
+        if (vision_encoder is None) != (multimodal_projector is None):
+            raise ValueError(
+                "`vision_encoder` and `multimodal_projector` must be "
+                "provided together for a multimodal `MistralBackbone`."
+            )
+        text_only_model = vision_encoder is None
+
         # === Layers ===
         self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
@@ -136,11 +172,18 @@ def __init__(
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
-        self.layer_norm = MistralLayerNormalization(
+        self.layer_norm = RMSNormalization(
             epsilon=layer_norm_epsilon,
             dtype=dtype,
             name="sequence_output_layernorm",
         )
+        self.vision_encoder = vision_encoder
+        self.multimodal_projector = multimodal_projector
+        if not text_only_model:
+            self.image_text_embedding_merger = Mistral3ImageTextEmbeddingMerger(
+                dtype=dtype,
+                name="image_text_embedding_merger",
+            )
 
         # === Functional Model ===
         token_id_input = keras.Input(
@@ -150,14 +193,62 @@ def __init__(
             shape=(None,), dtype="int32", name="padding_mask"
         )
         x = self.token_embedding(token_id_input)
+
+        if not text_only_model:
+            # `None` spatial dims: HF's `PixtralImageProcessor` pads each
+            # batch to its own largest image, not to a fixed canvas, so the
+            # input canvas size varies per call. `image_sizes` carries each
+            # image's true (unpadded) `(height, width)` for cropping.
+            pixel_values_input = keras.Input(
+                shape=(vision_encoder.num_channels, None, None),
+                name="pixel_values",
+            )
+            image_sizes_input = keras.Input(
+                shape=(2,), dtype="int32", name="image_sizes"
+            )
+            # Flat positions of image placeholder tokens in the flattened
+            # `(batch * seq_length,)` sequence. Computed on the host (e.g.
+            # via `compute_image_placeholder_indices`) rather than derived
+            # in-graph, since a `nonzero`-style lookup has a data-dependent
+            # output shape that is incompatible with `jax.jit` tracing.
+            placeholder_indices_input = keras.Input(
+                shape=(None,),
+                dtype="int32",
+                name="placeholder_indices",
+            )
+            self.image_feature_extractor = Mistral3ImageFeatureExtractor(
+                vision_encoder,
+                multimodal_projector,
+                dtype=dtype,
+                name="image_feature_extractor",
+            )
+            image_features = self.image_feature_extractor(
+                pixel_values_input,
+                image_sizes_input,
+            )
+            x = self.image_text_embedding_merger(
+                x, image_features, placeholder_indices_input
+            )
+
         for transformer_layer in self.transformer_layers:
             x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
         sequence_output = self.layer_norm(x)
+
+        inputs = {
+            "token_ids": token_id_input,
+            "padding_mask": padding_mask_input,
+        }
+        if not text_only_model:
+            inputs.update(
+                {
+                    "pixel_values": pixel_values_input,
+                    "image_sizes": image_sizes_input,
+                    "placeholder_indices": placeholder_indices_input,
+                }
+            )
+
         super().__init__(
-            inputs={
-                "token_ids": token_id_input,
-                "padding_mask": padding_mask_input,
-            },
+            inputs=inputs,
             outputs=sequence_output,
             dtype=dtype,
             **kwargs,
@@ -176,6 +267,8 @@ def __init__(
         self.head_dim = head_dim
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
+        self.image_token_index = image_token_index
+        self.text_only_model = text_only_model
 
     def get_config(self):
         config = super().get_config()
@@ -193,6 +286,28 @@ def get_config(self):
                 "head_dim": self.head_dim,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,
+                "image_token_index": self.image_token_index,
+                "vision_encoder": None
+                if self.vision_encoder is None
+                else keras.layers.serialize(self.vision_encoder),
+                "multimodal_projector": None
+                if self.multimodal_projector is None
+                else keras.layers.serialize(self.multimodal_projector),
             }
         )
         return config
+
+    @classmethod
+    def from_config(cls, config):
+        config = dict(config)
+        config.update(
+            {
+                "vision_encoder": None
+                if config["vision_encoder"] is None
+                else keras.layers.deserialize(config["vision_encoder"]),
+                "multimodal_projector": None
+                if config["multimodal_projector"] is None
+                else keras.layers.deserialize(config["multimodal_projector"]),
+            }
+        )
+        return super().from_config(config)

keras_hub/src/models/mistral/mistral_backbone_test.py

@@ -1,7 +1,17 @@
+import numpy as np
 import pytest
 from keras import ops
 
 from keras_hub.src.models.mistral.mistral_backbone import MistralBackbone
+from keras_hub.src.models.mistral.mistral_vision_encoder import (
+    Mistral3MultiModalProjector,
+)
+from keras_hub.src.models.mistral.mistral_vision_encoder import (
+    Mistral3VisionEncoder,
+)
+from keras_hub.src.models.mistral.mistral_vision_encoder import (
+    compute_image_placeholder_indices,
+)
 from keras_hub.src.tests.test_case import TestCase
 
 
@@ -87,3 +97,101 @@ def test_all_presets(self):
                 preset=preset,
                 input_data=self.input_data,
             )
+
+
+class MistralMultimodalBackboneTest(TestCase):
+    """Tests for `MistralBackbone` configured with a vision encoder."""
+
+    def setUp(self):
+        self.text_init_kwargs = {
+            "vocabulary_size": 10,
+            "num_layers": 2,
+            "num_query_heads": 8,
+            "num_key_value_heads": 4,
+            "hidden_dim": 16,
+            "intermediate_dim": 8,
+            "sliding_window": 2,
+        }
+
+        vision_encoder = Mistral3VisionEncoder(
+            image_size=8,
+            patch_size=4,
+            hidden_dim=8,
+            num_layers=1,
+            num_heads=2,
+            head_dim=4,
+            intermediate_dim=8,
+        )
+        multimodal_projector = Mistral3MultiModalProjector(
+            vision_hidden_dim=8,
+            text_hidden_dim=self.text_init_kwargs["hidden_dim"],
+            spatial_merge_size=1,
+            patch_size=4,
+            image_size=8,
+        )
+        # Must stay inside `vocabulary_size`: the embedding lookup happens
+        # before the image-text merger overwrites these positions.
+        self.image_token_index = 9
+        self.init_kwargs = {
+            **self.text_init_kwargs,
+            "vision_encoder": vision_encoder,
+            "multimodal_projector": multimodal_projector,
+            "image_token_index": self.image_token_index,
+        }
+        # Two 8x8 images, each a 2x2 patch grid; `spatial_merge_size=1`
+        # makes every patch its own merge window (4 rows per image).
+        token_ids = ops.array(
+            [
+                [self.image_token_index] * 4 + [3],
+                [self.image_token_index] * 4 + [4],
+            ],
+            dtype="int32",
+        )
+        placeholder_indices = compute_image_placeholder_indices(
+            token_ids, image_token_index=self.image_token_index
+        )
+        self.input_data = {
+            "token_ids": token_ids,
+            "padding_mask": ops.ones((2, 5), dtype="int32"),
+            "pixel_values": ops.convert_to_tensor(
+                np.random.rand(2, 3, 8, 8).astype("float32")
+            ),
+            "image_sizes": ops.array([[8, 8], [8, 8]], dtype="int32"),
+            "placeholder_indices": ops.reshape(
+                ops.convert_to_tensor(placeholder_indices), (2, 4)
+            ),
+        }
+
+    def test_backbone_basics(self):
+        self.run_backbone_test(
+            cls=MistralBackbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=(
+                2,
+                5,
+                self.text_init_kwargs["hidden_dim"],
+            ),
+            # Image inputs have no sequence axis to slice, so skip the
+            # default variable-length sweep.
+            variable_length_data=[self.input_data],
+            # `run_quantization_test` rebuilds `vision_encoder`/
+            # `multimodal_projector` as standalone objects to apply a
+            # path-keyed `DTypePolicyMap`, but their sublayer paths change
+            # once they're no longer nested under the backbone -- the same
+            # structural mismatch `gemma3_backbone_test.py` works around for
+            # its own vision-encoder-bearing backbone.
+            run_quantization_check=False,
+        )
+
+    def test_vision_projector_must_be_paired(self):
+        with self.assertRaises(ValueError):
+            MistralBackbone(
+                **self.text_init_kwargs,
+                vision_encoder=self.init_kwargs["vision_encoder"],
+            )
+        with self.assertRaises(ValueError):
+            MistralBackbone(
+                **self.text_init_kwargs,
+                multimodal_projector=self.init_kwargs["multimodal_projector"],
+            )