Add Llama Nemotron VL multimodal embedding model support (#54)

* Add Llama Bidirectional embedding model support

Add support for bidirectional Llama-based embedding models (e.g. nvidia/llama-embed-nemotron-8b).

Key changes:
- New model file: mlx_embeddings/models/llama_bidirec.py
- Imports TransformerBlock from mlx_lm for the Llama backbone
- Bidirectional attention via full (non-causal) attention masking
- Average pooling + L2 normalization for text embeddings
- Unit test with synthetic config
- README updated with new architecture

* Add input_embeddings parameter to LlamaBidirectionalModel

Allows passing pre-computed embeddings (e.g. with vision tokens injected)
instead of input_ids. Required for the upcoming llama_nemotron_vl model
which injects vision features into the token embedding stream.

Backward compatible — existing callers are unaffected.

* Add Llama Nemotron VL multimodal embedding model support

Add support for nvidia/llama-nemotron-embed-vl-1b-v2, a 1.7B multimodal
embedding model combining SigLIP vision encoder with bidirectional Llama.

Key changes:
- New model package: mlx_embeddings/models/llama_nemotron_vl/
- Reuses SiglipVisionTransformer from siglip.py for vision encoding
- Reuses LlamaBidirectionalModel from llama_bidirec.py for text encoding
- pixel_shuffle + MLP1 projection for vision-to-LLM token injection
- Average pooling + L2 normalization for 2048-dim embeddings
- Follows mlx-vlm InternVL reference patterns for pixel_shuffle and token merge
- Unit test with synthetic config
- README updated with new architecture

Benchmarked against HF reference:
- VL 1B text embeddings: 0.9999 cosine (near-identical)
- Conversion works (regular 3.1GB + quantized 1.5GB)
- Rankings match 5/5 across all test texts

---------

Co-authored-by: Joe Schlesinger <157418651+joe-verkada@users.noreply.github.com>
Co-authored-by: Prince Canuma <prince.gdt@gmail.com>

Author: 3 people

README.md

@@ -20,6 +20,7 @@ MLX-Embeddings supports a variety of model architectures for text embedding task
 - Qwen3 (Qwen3's embedding model)
 - Qwen3-VL (multimodal Qwen3-VL embedding and reranking model)
 - Llama Bidirectional (Llama-based bidirectional embedding models, e.g. NVIDIA NV-Embed)
+- Llama Nemotron VL (multimodal vision-language embedding model with SigLIP vision + bidirectional Llama)
 
 We're continuously working to expand our support for additional model architectures. Check our GitHub repository or documentation for the most up-to-date list of supported models and their specific versions.
 

mlx_embeddings/models/llama_nemotron_vl/__init__.py

@@ -0,0 +1,6 @@
+from .model import Model, ModelArgs
+
+__all__ = [
+    "Model",
+    "ModelArgs",
+]

mlx_embeddings/models/llama_nemotron_vl/model.py

@@ -0,0 +1,283 @@
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from ..base import BaseModelArgs, BaseModelOutput, mean_pooling, normalize_embeddings
+from ..llama_bidirec import LlamaBidirectionalModel
+from ..llama_bidirec import ModelArgs as LlamaBidirectModelArgs
+from ..siglip import SiglipVisionTransformer
+from ..siglip import VisionConfig as SiglipVisionConfig
+
+
+@dataclass
+class VisionConfig(BaseModelArgs):
+    model_type: str = "siglip_vision_model"
+    hidden_size: int = 1152
+    intermediate_size: int = 4304
+    num_hidden_layers: int = 27
+    num_attention_heads: int = 16
+    num_channels: int = 3
+    image_size: int = 512
+    patch_size: int = 16
+    layer_norm_eps: float = 1e-6
+    attention_dropout: float = 0.0
+    hidden_act: str = "gelu_pytorch_tanh"
+
+
+@dataclass
+class TextConfig(BaseModelArgs):
+    model_type: str = "llama_bidirec"
+    hidden_size: int = 2048
+    num_hidden_layers: int = 16
+    intermediate_size: int = 8192
+    num_attention_heads: int = 32
+    num_key_value_heads: Optional[int] = 8
+    head_dim: Optional[int] = 64
+    max_position_embeddings: int = 131072
+    vocab_size: int = 128266
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 500000.0
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    rope_traditional: bool = False
+    tie_word_embeddings: bool = True
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    layer_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.head_dim is None:
+            self.head_dim = self.hidden_size // self.num_attention_heads
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "llama_nemotron_vl"
+    vision_config: Optional[VisionConfig] = None
+    llm_config: Optional[TextConfig] = None
+    downsample_ratio: float = 0.5
+    force_image_size: Optional[int] = 512
+    img_context_token_id: int = 128258
+    select_layer: int = -1
+    hidden_size: int = 2048
+
+    def __post_init__(self):
+        if isinstance(self.vision_config, dict):
+            self.vision_config = VisionConfig.from_dict(self.vision_config)
+        elif self.vision_config is None:
+            self.vision_config = VisionConfig()
+        if isinstance(self.llm_config, dict):
+            self.llm_config = TextConfig.from_dict(self.llm_config)
+        elif self.llm_config is None:
+            self.llm_config = TextConfig()
+        self.hidden_size = self.llm_config.hidden_size
+
+
+def pixel_shuffle(input_tensor: mx.array, shuffle_ratio: float) -> mx.array:
+    batch_size, num_patches, channels = input_tensor.shape
+    patch_size = int(math.sqrt(num_patches))
+
+    input_tensor = input_tensor.reshape(batch_size, patch_size, patch_size, -1)
+    batch_size, height, width, channels = input_tensor.shape
+
+    reshaped = input_tensor.reshape(
+        batch_size, height, int(width * shuffle_ratio), int(channels / shuffle_ratio)
+    )
+    reshaped = reshaped.transpose(0, 2, 1, 3)
+    reshaped = reshaped.reshape(
+        batch_size,
+        int(height * shuffle_ratio),
+        int(width * shuffle_ratio),
+        int(channels / (shuffle_ratio**2)),
+    )
+    reshaped = reshaped.transpose(0, 2, 1, 3)
+
+    return reshaped.reshape(batch_size, -1, reshaped.shape[-1])
+
+
+class VisionEncoder(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        siglip_config = SiglipVisionConfig(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            num_hidden_layers=config.num_hidden_layers,
+            num_attention_heads=config.num_attention_heads,
+            num_channels=config.num_channels,
+            image_size=config.image_size,
+            patch_size=config.patch_size,
+            vision_use_head=True,
+        )
+        self.vision_model = SiglipVisionTransformer(siglip_config)
+
+    def __call__(self, pixel_values: mx.array) -> mx.array:
+        x, _ = self.vision_model(pixel_values)
+        return x
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+
+        self.vision_model = VisionEncoder(config.vision_config)
+
+        llm_args = LlamaBidirectModelArgs.from_dict(
+            {
+                "model_type": "llama_bidirec",
+                "hidden_size": config.llm_config.hidden_size,
+                "num_hidden_layers": config.llm_config.num_hidden_layers,
+                "intermediate_size": config.llm_config.intermediate_size,
+                "num_attention_heads": config.llm_config.num_attention_heads,
+                "num_key_value_heads": config.llm_config.num_key_value_heads,
+                "head_dim": config.llm_config.head_dim,
+                "max_position_embeddings": config.llm_config.max_position_embeddings,
+                "vocab_size": config.llm_config.vocab_size,
+                "rms_norm_eps": config.llm_config.rms_norm_eps,
+                "rope_theta": config.llm_config.rope_theta,
+                "rope_scaling": config.llm_config.rope_scaling,
+                "rope_traditional": config.llm_config.rope_traditional,
+                "tie_word_embeddings": config.llm_config.tie_word_embeddings,
+                "attention_bias": config.llm_config.attention_bias,
+                "mlp_bias": config.llm_config.mlp_bias,
+                "layer_types": config.llm_config.layer_types,
+            }
+        )
+        self.language_model = LlamaBidirectionalModel(llm_args)
+
+        vit_hidden = config.vision_config.hidden_size
+        ds = config.downsample_ratio
+        mlp_input_size = int(vit_hidden * (1 / ds) ** 2)
+        llm_hidden = config.llm_config.hidden_size
+
+        self.mlp1 = [
+            nn.LayerNorm(mlp_input_size),
+            nn.Linear(mlp_input_size, llm_hidden),
+            nn.GELU(),
+            nn.Linear(llm_hidden, llm_hidden),
+        ]
+
+        self.downsample_ratio = ds
+        self.img_context_token_id = config.img_context_token_id
+
+    def get_extended_attention_mask(self, attention_mask):
+        if attention_mask.ndim == 2:
+            extended = attention_mask[:, None, None, :]
+            extended = mx.repeat(extended, attention_mask.shape[-1], -2)
+        elif attention_mask.ndim == 3:
+            extended = attention_mask[:, None, :, :]
+        else:
+            raise ValueError(
+                f"Wrong shape for attention_mask (shape {attention_mask.shape})"
+            )
+        return extended
+
+    def extract_feature(self, pixel_values: mx.array) -> mx.array:
+        vit_embeds = self.vision_model(pixel_values)
+        vit_embeds = pixel_shuffle(vit_embeds, shuffle_ratio=self.downsample_ratio)
+
+        for layer in self.mlp1:
+            vit_embeds = layer(vit_embeds)
+
+        return vit_embeds
+
+    def _merge_input_ids_with_image_features(
+        self, image_features, inputs_embeds, input_ids
+    ):
+        B, N, C = inputs_embeds.shape
+
+        image_positions = input_ids == self.img_context_token_id
+        image_indices = np.where(image_positions)[1].tolist()
+
+        image_features = image_features.reshape(-1, image_features.shape[-1])
+        inputs_embeds[:, image_indices, :] = image_features
+
+        return inputs_embeds.reshape(B, N, C)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        attention_mask: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+    ):
+        if attention_mask is None:
+            attention_mask = mx.ones(input_ids.shape)
+
+        input_embeds = self.language_model.embed_tokens(input_ids)
+
+        if pixel_values is not None:
+            if pixel_values.ndim == 5:
+                pixel_values = pixel_values[0]
+
+            dtype = (
+                self.vision_model.vision_model.embeddings.patch_embedding.weight.dtype
+            )
+            pixel_values = pixel_values.astype(dtype)
+
+            vit_embeds = self.extract_feature(pixel_values)
+            input_embeds = self._merge_input_ids_with_image_features(
+                vit_embeds, input_embeds, input_ids
+            )
+
+        extended_mask = self.get_extended_attention_mask(attention_mask)
+        extended_mask = mx.where(
+            extended_mask.astype(mx.bool_),
+            0.0,
+            -mx.inf,
+        )
+        extended_mask = extended_mask.astype(
+            self.language_model.embed_tokens.weight.dtype
+        )
+
+        out = self.language_model(
+            input_ids, extended_mask, input_embeddings=input_embeds
+        )
+
+        text_embeds = mean_pooling(out, attention_mask)
+        text_embeds = normalize_embeddings(text_embeds)
+
+        return BaseModelOutput(
+            last_hidden_state=out,
+            text_embeds=text_embeds,
+            pooler_output=None,
+        )
+
+    def sanitize(self, weights):
+        sanitized = {}
+        for k, v in weights.items():
+            if "rotary_emb.inv_freq" in k:
+                continue
+            if "lm_head" in k:
+                continue
+            if "head.attention.in_proj_weight" in k:
+                sanitized[k.replace("in_proj_weight", "in_proj.weight")] = v
+                continue
+            if "head.attention.in_proj_bias" in k:
+                sanitized[k.replace("in_proj_bias", "in_proj.bias")] = v
+                continue
+
+            new_key = k
+
+            if k.startswith("language_model.model."):
+                new_key = k.replace("language_model.model.", "language_model.")
+
+            if "patch_embedding.weight" in k and v.ndim == 4:
+                # HF: (out, in, kH, kW) -> MLX: (out, kH, kW, in)
+                # Skip if already in MLX format (last dim is smallest = in_channels)
+                if v.shape[1] < v.shape[2]:
+                    v = v.transpose(0, 2, 3, 1)
+
+            sanitized[new_key] = v
+        return sanitized
+
+    @property
+    def layers(self):
+        return self.language_model.layers

mlx_embeddings/tests/test_models.py

@@ -421,6 +421,50 @@ def test_siglip2_model(self):
         self.assertEqual(outputs.logits_per_image.shape, (batch_size, batch_size))
         self.assertEqual(outputs.logits_per_text.shape, (batch_size, batch_size))
 
+    def test_llama_nemotron_vl_text_only(self):
+        from mlx_embeddings.models.llama_nemotron_vl.model import (
+            Model as NemotronVLModel,
+        )
+        from mlx_embeddings.models.llama_nemotron_vl.model import (
+            ModelArgs as NemotronVLModelArgs,
+        )
+        from mlx_embeddings.models.llama_nemotron_vl.model import (
+            TextConfig,
+            VisionConfig,
+        )
+
+        config = NemotronVLModelArgs(
+            model_type="llama_nemotron_vl",
+            hidden_size=64,
+            vision_config=VisionConfig(
+                hidden_size=64,
+                intermediate_size=128,
+                num_hidden_layers=2,
+                num_attention_heads=2,
+                image_size=32,
+                patch_size=16,
+            ),
+            llm_config=TextConfig(
+                hidden_size=64,
+                num_hidden_layers=2,
+                intermediate_size=128,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=32,
+                max_position_embeddings=128,
+                vocab_size=1000,
+            ),
+            downsample_ratio=0.5,
+            img_context_token_id=999,
+        )
+        model = NemotronVLModel(config)
+
+        self.model_test_runner(
+            model,
+            config.model_type,
+            config.llm_config.num_hidden_layers,
+        )
+
     def test_llama_bidirec_model(self):
         from mlx_embeddings.models import llama_bidirec