Support gemma-4-12B unified (encoder-free) multimodal model

src/models/gemma4_multimodal_processor.cpp

@@ -167,6 +167,14 @@ ProcessGemma4Prompt(const Generators::Tokenizer& tokenizer, const std::string& p
 
 Gemma4MultiModalProcessor::Gemma4MultiModalProcessor(Config& config, const SessionInfo& session_info)
     : pixel_values_type_{session_info.GetInputDataType(config.model.vision.inputs.pixel_values)} {
+  // gemma-4-12B "unified" is encoder-free: it consumes raw 48px merged pixel
+  // patches (patch_dim 6912) and raw 640-sample waveform frames directly. The
+  // preprocessing ops (Gemma4ImageTransform at patch_size=48/pooling=1 and
+  // Gemma4UnifiedAudioFrames) already emit the final contract, so — unlike the
+  // SigLIP/log-mel "gemma4" path — no pixel trimming or conv-subsampling of the
+  // audio-token count is applied.
+  unified_ = config.model.type == "gemma4_unified";
+
   // Query pixel_position_ids type (int32 or int64) if the vision model has this input
   if (session_info.HasInput(config.model.vision.inputs.pixel_position_ids)) {
     pixel_position_ids_type_ = session_info.GetInputDataType(config.model.vision.inputs.pixel_position_ids);
@@ -285,10 +293,18 @@ std::unique_ptr<NamedTensors> Gemma4MultiModalProcessor::Process(const Tokenizer
                              std::make_shared<Tensor>(std::move(mask)));
     }
 
-    // Compute audio_sizes: the speech encoder uses 2-stage Conv2d with stride=2 each
-    int64_t t_after_1 = (time_dim - 1) / 2 + 1;
-    int64_t t_after_2 = (t_after_1 - 1) / 2 + 1;
-    num_audio_tokens = t_after_2;
+    // Compute audio_sizes / audio-token count.
+    // Unified: each 640-sample frame is exactly one audio soft token, so the
+    // token count equals the number of frames (time_dim) — no subsampling.
+    // Standard gemma4: the Conformer speech encoder applies 2-stage Conv2d with
+    // stride=2 each, so the token count is reduced accordingly.
+    if (unified_) {
+      num_audio_tokens = time_dim;
+    } else {
+      int64_t t_after_1 = (time_dim - 1) / 2 + 1;
+      int64_t t_after_2 = (t_after_1 - 1) / 2 + 1;
+      num_audio_tokens = t_after_2;
+    }
     std::array<int64_t, 1> audio_sizes_shape = {1};
     auto audio_sizes = OrtValue::CreateTensor<int64_t>(allocator, audio_sizes_shape);
     audio_sizes->GetTensorMutableData<int64_t>()[0] = num_audio_tokens;
@@ -304,8 +320,15 @@ std::unique_ptr<NamedTensors> Gemma4MultiModalProcessor::Process(const Tokenizer
 
   if (payload.images) {
     // The Gemma4ImageTransform pads pixel_values and position_ids to max_patches.
-    // The vision ONNX model expects the actual (unpadded) number of patches.
-    // Trim the tensors to actual_patches = actual_soft_tokens * pooling_kernel_size².
+    //
+    // Standard gemma4: the SigLIP vision ONNX expects the actual (unpadded)
+    // number of teacher patches, so the tensors are trimmed to
+    // actual_patches = actual_soft_tokens * pooling_kernel_size².
+    //
+    // Unified (gemma4_unified): the encoder-free vision graph consumes the full
+    // padded (max_soft_tokens, 6912) tensors and strips padding internally using
+    // position_ids (== -1 marks padding). No trimming is applied; actual_patches
+    // is computed but the trim branches below are gated on !unified_.
     constexpr int64_t kPoolingKernelSize = 3;
     const int64_t actual_patches = static_cast<int64_t>(actual_soft_tokens) * kPoolingKernelSize * kPoolingKernelSize;
 
@@ -320,7 +343,7 @@ std::unique_ptr<NamedTensors> Gemma4MultiModalProcessor::Process(const Tokenizer
     const int64_t num_padded_patches = (pv_dims == 3) ? pv_shape[1] : pv_shape[0];
     const int64_t patch_dim = (pv_dims == 3) ? pv_shape[2] : pv_shape[1];
 
-    if (actual_patches < num_padded_patches) {
+    if (!unified_ && actual_patches < num_padded_patches) {
       // Trim: copy only the first actual_patches from the padded tensor.
       // The preprocessor outputs float32 data, so we trim using float32 strides,
       // then cast to the model's pixel_values type (float, fp16, or bf16).
@@ -365,7 +388,7 @@ std::unique_ptr<NamedTensors> Gemma4MultiModalProcessor::Process(const Tokenizer
       const int64_t num_padded_pos = (pos_dims == 3) ? pos_shape[1] : pos_shape[0];
       const int64_t pos_last_dim = (pos_dims == 3) ? pos_shape[2] : pos_shape[1];
 
-      if (actual_patches < num_padded_pos) {
+      if (!unified_ && actual_patches < num_padded_pos) {
         // Trim position_ids: for 3D, copy per-batch with correct stride.
         // Detect the element type from the vision model's input to handle both int32 and int64.
         const int64_t pos_batch = (pos_dims == 3) ? pos_shape[0] : 1;

src/models/gemma4_multimodal_processor.h

@@ -20,6 +20,7 @@ struct Gemma4MultiModalProcessor : Processor {
   ONNXTensorElementDataType audio_features_type_;
 
   bool has_speech_{false};
+  bool unified_{false};  // gemma-4-12B encoder-free "unified" variant
   size_t vision_soft_tokens_per_image_{260};
 };
 

src/models/model.cpp

@@ -974,6 +974,7 @@ MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& sess
           {"phi4mm", Processor::Create<PhiMultiModalProcessor>},
           {"gemma3", Processor::Create<GemmaImageProcessor>},
           {"gemma4", Processor::Create<Gemma4MultiModalProcessor>},
+          {"gemma4_unified", Processor::Create<Gemma4MultiModalProcessor>},
           {"mistral3", Processor::Create<Mistral3ImageProcessor>},
           {"fara", Processor::Create<QwenImageProcessor>},
           {"qwen2_5_vl", Processor::Create<QwenImageProcessor>},

src/models/model_type.h

@@ -63,7 +63,7 @@ struct ModelType {
 
   inline static bool IsMMM(const std::string& model_type) {
     // Multi-modal model (MMM)
-    static constexpr std::array<std::string_view, 2> MMM = {"gemma4", "phi4mm"};
+    static constexpr std::array<std::string_view, 3> MMM = {"gemma4", "gemma4_unified", "phi4mm"};
     return std::find(MMM.begin(), MMM.end(), model_type) != MMM.end();
   }
 

test/python/create/create_dummy_gemma4_unified_models.py

@@ -0,0 +1,136 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License
+
+"""Generate the dummy ``gemma4_unified`` test model directory.
+
+The gemma-4-12B "unified" (encoder-free) model shares the gemma4 decoder /
+embedding contract but consumes raw 48px merged pixel patches
+(``pixel_values`` last dim = 48*48*3 = 6912) and raw 640-sample waveform frames
+(``audio_embeds`` last dim = 640) directly, instead of the SigLIP 16px /
+128-dim log-mel contract.
+
+This derives ``test/models/gemma4_unified`` from the existing
+``test/models/gemma4`` fixtures: the embedding / text decoders are copied
+verbatim, the vision / speech dummies get the unified input dims, and the
+genai / processor configs are rewritten for the ``gemma4_unified`` type.
+
+Usage (from the repo root):
+    python test/python/create/create_dummy_gemma4_unified_models.py
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+
+import onnx
+
+_UNIFIED_PIXEL_DIM = 48 * 48 * 3  # 6912
+_UNIFIED_AUDIO_DIM = 640
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_SRC_DIR = _REPO_ROOT / "test" / "models" / "gemma4"
+_DST_DIR = _REPO_ROOT / "test" / "models" / "gemma4_unified"
+
+_TOKENIZER_FILES = [
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+]
+
+
+def _set_input_last_dim(model_path: Path, out_path: Path, input_name: str, last_dim: int) -> None:
+    """Rewrite the last dimension of a named graph input to a fixed value."""
+    model = onnx.load(str(model_path))
+    for inp in model.graph.input:
+        if inp.name == input_name:
+            dims = inp.type.tensor_type.shape.dim
+            dims[-1].ClearField("dim_param")
+            dims[-1].dim_value = last_dim
+            break
+    else:
+        raise ValueError(f"input {input_name!r} not found in {model_path}")
+    onnx.save(model, str(out_path))
+
+
+def main() -> None:
+    if not _SRC_DIR.exists():
+        raise SystemExit(
+            f"Source gemma4 fixtures not found at {_SRC_DIR}. Generate/download the "
+            "gemma4 test model directory first; gemma4_unified is derived from it."
+        )
+    _DST_DIR.mkdir(parents=True, exist_ok=True)
+
+    # Decoder + embedding are identical to gemma4.
+    for name in ("dummy_text.onnx", "dummy_embedding.onnx"):
+        shutil.copyfile(_SRC_DIR / name, _DST_DIR / name)
+
+    # Vision / speech dummies: same trivial constant-output graphs, but declare
+    # the unified input dims so the fixtures document the real contract.
+    _set_input_last_dim(
+        _SRC_DIR / "dummy_vision.onnx", _DST_DIR / "dummy_vision.onnx", "pixel_values", _UNIFIED_PIXEL_DIM
+    )
+    _set_input_last_dim(
+        _SRC_DIR / "dummy_speech.onnx", _DST_DIR / "dummy_speech.onnx", "audio_embeds", _UNIFIED_AUDIO_DIM
+    )
+
+    for name in _TOKENIZER_FILES:
+        shutil.copyfile(_SRC_DIR / name, _DST_DIR / name)
+
+    # genai_config.json: switch the model type and vision processor config file.
+    with open(_SRC_DIR / "genai_config.json") as f:
+        genai_config = json.load(f)
+    genai_config["model"]["type"] = "gemma4_unified"
+    genai_config["model"]["vision"]["config_filename"] = "image_processor.json"
+    with open(_DST_DIR / "genai_config.json", "w") as f:
+        json.dump(genai_config, f, indent=4)
+
+    # image_processor.json: reuse Gemma4ImageTransform with the merged geometry
+    # (patch_size=48, pooling_kernel_size=1) that yields 6912-dim patches.
+    image_processor = {
+        "processor": {
+            "name": "gemma_4_unified_image_processing",
+            "transforms": [
+                {"operation": {"name": "decode_image", "type": "DecodeImage", "attrs": {"color_space": "RGB"}}},
+                {
+                    "operation": {
+                        "name": "gemma4_image_transform",
+                        "type": "Gemma4ImageTransform",
+                        "attrs": {"patch_size": 48, "max_soft_tokens": 280, "pooling_kernel_size": 1},
+                    }
+                },
+            ],
+        }
+    }
+    with open(_DST_DIR / "image_processor.json", "w") as f:
+        json.dump(image_processor, f, indent=4)
+
+    # audio_feature_extraction.json: raw 640-sample waveform framing.
+    audio_config = {
+        "feature_extraction": {
+            "sequence": [
+                {"operation": {"name": "audio_decoder", "type": "AudioDecoder"}},
+                {
+                    "operation": {
+                        "name": "gemma4_audio",
+                        "type": "Gemma4Audio",
+                        "attrs": {
+                            "type": "raw_frames",
+                            "audio_samples_per_token": 640,
+                            "sampling_rate": 16000,
+                            "padding_value": 0.0,
+                        },
+                    }
+                },
+            ]
+        }
+    }
+    with open(_DST_DIR / "audio_feature_extraction.json", "w") as f:
+        json.dump(audio_config, f, indent=4)
+
+    print(f"Wrote gemma4_unified dummy model to {_DST_DIR}")
+
+
+if __name__ == "__main__":
+    main()

test/python/models/test_gemma4_unified_models.py

@@ -0,0 +1,144 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License
+
+"""
+Unit tests for the Gemma4 "unified" (encoder-free, gemma-4-12B) multimodal model.
+
+Unlike the standard gemma4 model, the unified variant consumes raw 48px merged
+pixel patches (``pixel_values`` last dim = 6912) and raw 640-sample waveform
+frames (``audio_embeds`` last dim = 640) directly, and its ``model.type`` is
+``gemma4_unified``. These tests exercise the ``unified_`` branch of
+``Gemma4MultiModalProcessor``: processor creation, and that the produced
+pixel_values / audio_embeds / audio_sizes follow the unified contract.
+
+The dummy model fixtures live in ``test/models/gemma4_unified`` (generated by
+``test/python/create/create_dummy_gemma4_unified_models.py``).
+
+Usage:
+    pytest test_gemma4_unified_models.py --test_models=/path/to/models
+"""
+
+import logging
+import os
+from pathlib import Path
+
+import numpy as np
+import onnxruntime_genai as og
+import pytest
+
+logging.basicConfig(format="%(asctime)s %(name)s [%(levelname)s] - %(message)s", level=logging.DEBUG)
+log = logging.getLogger("gemma4-unified-tests")
+
+GEMMA4_UNIFIED_MODEL_NAME = "gemma4_unified"
+
+# Encoder-free unified contract.
+_UNIFIED_PIXEL_DIM = 48 * 48 * 3  # 6912
+_UNIFIED_AUDIO_DIM = 640
+
+
+def _get_model_path(test_data_path):
+    """Return the gemma4_unified model path, skipping if it doesn't exist."""
+    model_path = os.path.join(test_data_path, GEMMA4_UNIFIED_MODEL_NAME)
+    if not os.path.exists(model_path):
+        pytest.skip(f"gemma4_unified test model not found at {model_path}")
+    return model_path
+
+
+def _load_model_and_processor(test_data_path):
+    model_path = _get_model_path(test_data_path)
+    model = og.Model(model_path)
+    return model, model.create_multimodal_processor()
+
+
+def _to_numpy(tensor):
+    if hasattr(tensor, "as_numpy"):
+        return tensor.as_numpy()
+    if hasattr(tensor, "numpy"):
+        return tensor.numpy()
+    return np.array(tensor)
+
+
+def test_gemma4_unified_model_load(test_data_path):
+    """The gemma4_unified model (model.type == 'gemma4_unified') loads."""
+    model_path = _get_model_path(test_data_path)
+    model = og.Model(model_path)
+    assert model is not None
+
+
+def test_gemma4_unified_processor_creation(test_data_path):
+    """create_multimodal_processor() succeeds for the gemma4_unified type.
+
+    This exercises the processor-factory registration and the unified image /
+    audio ort-extensions configs (Gemma4ImageTransform at patch_size=48 and
+    Gemma4UnifiedAudioFrames).
+    """
+    _, processor = _load_model_and_processor(test_data_path)
+    assert processor is not None
+
+
+def test_gemma4_unified_text_only(test_data_path):
+    """Text-only processing (no images/audio)."""
+    _, processor = _load_model_and_processor(test_data_path)
+    inputs = processor("What is the capital of France?", images=None)
+    assert inputs is not None
+    assert "input_ids" in inputs
+    ids = _to_numpy(inputs["input_ids"])
+    assert len(ids.shape) == 2 and ids.shape[0] == 1
+
+
+@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
+def test_gemma4_unified_vision_contract(test_data_path, relative_image_path):
+    """Unified vision preprocessing produces 6912-dim merged patches (no trim)."""
+    _, processor = _load_model_and_processor(test_data_path)
+
+    image_path = os.fspath(Path(test_data_path) / relative_image_path)
+    if not os.path.exists(image_path):
+        pytest.skip(f"Test image not found at {image_path}")
+    images = og.Images.open(image_path)
+
+    inputs = processor("<|image|>Describe this image", images=images)
+    assert inputs is not None
+    assert "pixel_values" in inputs
+    assert "pixel_position_ids" in inputs
+
+    pixel_values = _to_numpy(inputs["pixel_values"])
+    assert pixel_values.shape[-1] == _UNIFIED_PIXEL_DIM, (
+        f"unified pixel_values feature dim should be {_UNIFIED_PIXEL_DIM}, got {pixel_values.shape[-1]}"
+    )
+    # Unified feeds the full padded patch grid (no trim); the graph strips
+    # padding via position_ids. pixel_position_ids rows must match pixel_values.
+    pos = _to_numpy(inputs["pixel_position_ids"])
+    assert pos.shape[-2] == pixel_values.shape[-2], (
+        f"position_ids ({pos.shape}) and pixel_values ({pixel_values.shape}) patch counts must match"
+    )
+    assert pos.shape[-1] == 2
+
+
+@pytest.mark.parametrize("relative_audio_path", [Path("audios") / "jfk.flac"])
+def test_gemma4_unified_audio_contract(test_data_path, relative_audio_path):
+    """Unified audio preprocessing produces raw 640-sample frames; audio_sizes = frame count."""
+    _, processor = _load_model_and_processor(test_data_path)
+
+    audio_path = os.fspath(Path(test_data_path) / relative_audio_path)
+    if not os.path.exists(audio_path):
+        pytest.skip(f"Test audio file not found at {audio_path}")
+
+    audios = og.Audios.open(audio_path)
+    inputs = processor("<|audio|>Transcribe this audio", audios=audios)
+    assert inputs is not None
+    assert "audio_embeds" in inputs
+    assert "audio_sizes" in inputs
+
+    audio_embeds = _to_numpy(inputs["audio_embeds"])
+    assert audio_embeds.shape[-1] == _UNIFIED_AUDIO_DIM, (
+        f"unified audio_embeds feature dim should be {_UNIFIED_AUDIO_DIM}, got {audio_embeds.shape[-1]}"
+    )
+    assert audio_embeds.dtype == np.float32
+
+    # Unified: each 640-sample frame is exactly one audio token, so audio_sizes
+    # equals the number of frames (no stride-2 subsampling).
+    num_frames = audio_embeds.shape[-2]
+    audio_sizes = _to_numpy(inputs["audio_sizes"])
+    assert audio_sizes[0] == num_frames, (
+        f"unified audio_sizes should equal frame count {num_frames}, got {audio_sizes[0]}"
+    )