Unified Pixtral (#164)

mattjcly · web-flow · commit ecc2cf48e449 · 2025-05-19T11:12:25.000-04:00
diff --git a/mlx_engine/model_kit/model_kit.py b/mlx_engine/model_kit/model_kit.py
@@ -12,6 +12,7 @@
 from mlx_engine.logging import log_info, log_warn
 from mlx_engine.model_kit.vision_add_ons.base import BaseVisionAddOn
 from mlx_engine.model_kit.vision_add_ons.gemma3 import Gemma3VisionAddOn
+from mlx_engine.model_kit.vision_add_ons.pixtral import PixtralVisionAddOn
 from mlx_engine.utils.kv_cache_quantization import get_kv_cache_quantization_params
 from mlx_engine.utils.prompt_processing import process_prompt_text_only
 
@@ -33,6 +34,7 @@ class ModelKit:
 
     VISION_ADD_ON_MAP = {
         "gemma3": Gemma3VisionAddOn,
+        "pixtral": PixtralVisionAddOn,
     }
 
     # model state tracking
diff --git a/mlx_engine/model_kit/vision_add_ons/gemma3.py b/mlx_engine/model_kit/vision_add_ons/gemma3.py
@@ -1,8 +1,8 @@
-import glob
-import json
 from typing import List
+from pathlib import Path
 
 from mlx import nn
+import mlx.core as mx
 
 from mlx_vlm.models.gemma3 import (
     VisionModel as Gemma3VisionTower,
@@ -12,76 +12,36 @@
     Model as Gemma3CombinedModel,  # for prepare_inputs_for_multimodal
 )
 from mlx_vlm.models.gemma3.gemma3 import Gemma3MultiModalProjector
-from mlx_vlm.utils import sanitize_weights, load_processor, get_class_predicate
 
-from pathlib import Path
-import mlx.core as mx
-
-from mlx_engine.logging import log_info
+from mlx_engine.model_kit.vision_add_ons.base import BaseVisionAddOn
 from mlx_engine.model_kit.vision_add_ons.process_prompt_with_images import (
     common_process_prompt_with_images,
 )
-from mlx_engine.model_kit.vision_add_ons.base import BaseVisionAddOn
+from mlx_engine.model_kit.vision_add_ons.load_utils import load_vision_addon
 
 
-class Gemma3VisionAddOn(BaseVisionAddOn, nn.Module):
+class Gemma3VisionAddOn(BaseVisionAddOn):
     """
-    Vision add-on for Gemma3 model. Uses mlx-vlm vision components of Gemma3
+    Vision add-on for Gemma3 model. Uses mlx-vlm vision components of Gemma3.
     """
 
     GEMMA3_LOG_PREFIX = "Gemma3VisionAddOn"
 
     def __init__(self, model_path: Path):
+        """Initialize Gemma3VisionAddOn with vision components loaded from the given path."""
         super().__init__()
-        config_dict = json.loads((model_path / "config.json").read_text())
-        self.config = Gemma3ModelConfig.from_dict(config_dict)
-        self.config.vision_config = Gemma3VisionConfig.from_dict(
-            self.config.vision_config
-        )
-        self.config.text_config = Gemma3TextConfig.from_dict(self.config.text_config)
-        self.vision_tower = Gemma3VisionTower(self.config.vision_config)
-        self.multi_modal_projector = Gemma3MultiModalProjector(self.config)
-        self.processor = load_processor(model_path=model_path, add_detokenizer=True)
-        # load the weights for the vision tower
-        # ref: https://github.com/Blaizzy/mlx-vlm/blob/d2391123cabac313729f9a2a8d57d396e2592f20/mlx_vlm/utils.py#L147
-        # and https://github.com/Blaizzy/mlx-vlm/blob/d2391123cabac313729f9a2a8d57d396e2592f20/mlx_vlm/models/gemma3/gemma3.py#L86-L87
-        weight_files = glob.glob(str(model_path / "*.safetensors"))
-        if not weight_files:
-            raise FileNotFoundError(
-                f"Failed to load Gemma3 vision model: {model_path} does not contain any safetensors files"
-            )
-        weights = {}
-        for wf in weight_files:
-            weights.update(mx.load(wf))
-        # filter out everything but weights with keys that start with "vision_tower" or "multi_modal_projector"
-        weights = {
-            k: v
-            for k, v in weights.items()
-            if k.startswith("vision_tower") or k.startswith("multi_modal_projector")
-        }
-        weights = sanitize_weights(
-            Gemma3VisionTower, weights, self.config.vision_config
-        )
-        # perform jit quantization if needed
-        if (quantization := config_dict.get("quantization", None)) is not None:
-            class_predicate = get_class_predicate(skip_vision=False, weights=weights)
-            nn.quantize(
-                self,
-                **quantization,
-                class_predicate=class_predicate,
-            )
 
-        # load weights using nn.Module method
-        self.load_weights(list(weights.items()))
-        # hardcode lazy loading to false for now, always load weights to memory here
-        lazy = False
-        if not lazy:
-            mx.eval(self.parameters())
-
-        self.eval()
-        log_info(
-            prefix=self.GEMMA3_LOG_PREFIX,
-            message=f"Gemma3 vision model loaded successfully from {model_path}",
+        # Load vision model components, configuration, and processor
+        self.vision_tower, self.multi_modal_projector, self.config, self.processor = (
+            load_vision_addon(
+                model_path=model_path,
+                model_config_class=Gemma3ModelConfig,
+                vision_config_class=Gemma3VisionConfig,
+                text_config_class=Gemma3TextConfig,
+                vision_tower_class=Gemma3VisionTower,
+                multi_modal_projector_class=Gemma3MultiModalProjector,
+                log_prefix=self.GEMMA3_LOG_PREFIX,
+            )
         )
 
     def compute_embeddings(
@@ -90,6 +50,7 @@ def compute_embeddings(
         prompt_tokens: mx.array,
         images_b64: List[str],
     ) -> mx.array:
+        """Compute embeddings for text with images."""
         input_ids, pixel_values, attention_mask, other_model_inputs = (
             common_process_prompt_with_images(
                 prompt_tokens=prompt_tokens,
@@ -105,6 +66,7 @@ def compute_embeddings(
             pixel_values.transpose(0, 2, 3, 1).astype(input_embeddings.dtype),
             output_hidden_states=True,
         )
+
         # Format image features
         image_features = hidden_state.astype(pixel_values.dtype)
         image_features = self.multi_modal_projector(image_features)
diff --git a/mlx_engine/model_kit/vision_add_ons/load_utils.py b/mlx_engine/model_kit/vision_add_ons/load_utils.py
@@ -0,0 +1,114 @@
+import glob
+import json
+from pathlib import Path
+from typing import Any, Tuple, Type
+
+import mlx.core as mx
+from mlx import nn
+
+from mlx_vlm.utils import sanitize_weights, load_processor, get_class_predicate
+from mlx_engine.logging import log_info
+
+
+def load_vision_addon(
+    model_path: Path,
+    model_config_class: Any,
+    vision_config_class: Any,
+    text_config_class: Any,
+    vision_tower_class: Type[nn.Module],
+    multi_modal_projector_class: Type[nn.Module],
+    log_prefix: str,
+) -> Tuple[nn.Module, nn.Module, Any, Any]:
+    """
+    Load vision add-on components, configuration, and processor.
+
+    Args:
+        model_path: Path to the model directory
+        model_config_class: Configuration class for the model
+        vision_config_class: Configuration class for vision component
+        text_config_class: Configuration class for text component
+        vision_tower_class: The vision tower model class
+        multi_modal_projector_class: The multi-modal projector class
+        log_prefix: Prefix for logging messages
+
+    Returns:
+        Tuple containing:
+            - The vision tower module
+            - The multi-modal projector module
+            - The model configuration
+            - The processor for handling images and text
+    """
+    # Load and parse configuration
+    config_path = model_path / "config.json"
+    if not config_path.exists():
+        raise FileNotFoundError(f"Configuration file not found at {config_path}")
+
+    config_dict = json.loads(config_path.read_text())
+    config = model_config_class.from_dict(config_dict)
+    config.vision_config = vision_config_class.from_dict(config.vision_config)
+    config.text_config = text_config_class.from_dict(config.text_config)
+
+    # Create model components
+    vision_tower = vision_tower_class(config.vision_config)
+    multi_modal_projector = multi_modal_projector_class(config)
+
+    # Combine components into a container module for loading weights
+    class VisionComponents(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.vision_tower = vision_tower
+            self.multi_modal_projector = multi_modal_projector
+
+    components = VisionComponents()
+
+    # Load processor
+    processor = load_processor(model_path=model_path, add_detokenizer=True)
+
+    # Load model weights
+    weight_files = glob.glob(str(model_path / "*.safetensors"))
+    if not weight_files:
+        raise FileNotFoundError(
+            f"Failed to load vision add-on: {model_path} does not contain any safetensors files"
+        )
+
+    # Load and filter weights
+    weights = {}
+    for wf in weight_files:
+        weights.update(mx.load(wf))
+
+    # Filter only vision-related weights
+    vision_weights = {
+        k: v
+        for k, v in weights.items()
+        if k.startswith("vision_tower") or k.startswith("multi_modal_projector")
+    }
+
+    # Sanitize weights for vision tower
+    vision_weights = sanitize_weights(
+        vision_tower_class, vision_weights, config.vision_config
+    )
+
+    # Apply quantization if specified in config
+    if (quantization := config_dict.get("quantization", None)) is not None:
+        class_predicate = get_class_predicate(skip_vision=False, weights=vision_weights)
+        nn.quantize(
+            components,
+            **quantization,
+            class_predicate=class_predicate,
+        )
+
+    # Load weights into the model
+    components.load_weights(list(vision_weights.items()))
+
+    # Always load weights to memory here
+    mx.eval(components.parameters())
+
+    # Set model to evaluation mode
+    components.eval()
+
+    log_info(
+        prefix=log_prefix,
+        message=f"Vision add-on loaded successfully from {model_path}",
+    )
+
+    return vision_tower, multi_modal_projector, config, processor
diff --git a/mlx_engine/model_kit/vision_add_ons/pixtral.py b/mlx_engine/model_kit/vision_add_ons/pixtral.py
@@ -0,0 +1,88 @@
+from typing import List
+from pathlib import Path
+
+from mlx import nn
+import mlx.core as mx
+
+from mlx_vlm.models.pixtral import (
+    VisionModel as PixtralVisionTower,
+    ModelConfig as PixtralModelConfig,
+    VisionConfig as PixtralVisionConfig,
+    TextConfig as PixtralTextConfig,
+    Model as PixtralCombinedModel,  # for merge_input_ids_with_image_features
+)
+from mlx_vlm.models.pixtral.pixtral import (
+    LlavaMultiModalProjector as PixtralMultiModalProjector,
+)
+
+from mlx_engine.model_kit.vision_add_ons.base import BaseVisionAddOn
+from mlx_engine.model_kit.vision_add_ons.process_prompt_with_images import (
+    common_process_prompt_with_images,
+)
+from mlx_engine.model_kit.vision_add_ons.load_utils import load_vision_addon
+
+
+class PixtralVisionAddOn(BaseVisionAddOn):
+    """
+    Vision add-on for Pixtral model. Uses mlx-vlm vision components of Pixtral.
+    """
+
+    PIXTRAL_LOG_PREFIX = "PixtralVisionAddOn"
+
+    def __init__(self, model_path: Path):
+        """Initialize PixtralVisionAddOn with vision components loaded from the given path."""
+        super().__init__()
+
+        # Load vision model components, configuration, and processor
+        self.vision_tower, self.multi_modal_projector, self.config, self.processor = (
+            load_vision_addon(
+                model_path=model_path,
+                model_config_class=PixtralModelConfig,
+                vision_config_class=PixtralVisionConfig,
+                text_config_class=PixtralTextConfig,
+                vision_tower_class=PixtralVisionTower,
+                multi_modal_projector_class=PixtralMultiModalProjector,
+                log_prefix=self.PIXTRAL_LOG_PREFIX,
+            )
+        )
+
+    def compute_embeddings(
+        self,
+        text_model: nn.Module,
+        prompt_tokens: mx.array,
+        images_b64: List[str],
+    ) -> mx.array:
+        """Compute embeddings for text with images."""
+        input_ids, pixel_values, attention_mask, other_model_inputs = (
+            common_process_prompt_with_images(
+                prompt_tokens=prompt_tokens,
+                images_b64=images_b64,
+                processor=self.processor,
+                config=self.config,
+            )
+        )
+        input_embeddings = text_model.language_model.model.embed_tokens(input_ids)
+
+        if isinstance(pixel_values, list):
+            pixel_values = mx.concatenate(
+                [mx.array(pv)[None, ...] for pv in pixel_values], axis=0
+            )
+        if pixel_values.ndim == 3:
+            pixel_values = pixel_values[None, ...]
+
+        # Process image through vision tower
+        *_, hidden_states = self.vision_tower(
+            pixel_values.transpose(0, 2, 3, 1),
+            output_hidden_states=True,
+        )
+        # Select the hidden states from the desired layer
+        selected_image_feature = hidden_states[self.config.vision_feature_layer]
+
+        # Pass image features through the multi-modal projector
+        image_features = self.multi_modal_projector(selected_image_feature)
+
+        # Insert special image tokens in the input_ids
+        final_inputs_embeds = PixtralCombinedModel.merge_input_ids_with_image_features(
+            self.config.image_token_index, image_features, input_embeddings, input_ids
+        )
+        return final_inputs_embeds.squeeze(0)  # remove batch dimension
diff --git a/requirements.txt b/requirements.txt
@@ -22,8 +22,8 @@ jsonschema-specifications==2024.10.1
 jsonschema==4.23.0
 lark==1.2.2
 markupsafe==2.1.5
-mlx-lm @ git+https://github.com/ml-explore/mlx-lm.git@f93589cb
-mlx-vlm==0.1.26
+mlx-lm @ git+https://github.com/ml-explore/mlx-lm.git@064c75d
+mlx-vlm @ git+https://github.com/Blaizzy/mlx-vlm.git@51eecac
 mlx==0.25.2
 mpmath==1.3.0
 multidict==6.1.0
