From 2b1f5392d30b74311408880ace3e3d049ef2d4e8 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Jul 2026 18:03:16 +0000 Subject: [PATCH 1/3] Support gemma-4-12B unified (encoder-free) multimodal model The gemma-4-12B "unified" variant is encoder-free: it consumes raw 48px merged pixel patches (patch_dim 6912) and raw 640-sample waveform frames directly, rather than the SigLIP image / Conformer log-mel audio contract of the standard gemma-4 (E2B/E4B) model. Register "gemma4_unified" as an MMM model type and route it to Gemma4MultiModalProcessor with a unified flag that adjusts two things: * Vision: no pixel-value trimming. The unified vision graph strips padding patches internally (position == -1), so the full padded (max_soft_tokens, 6912) pixel_values and position_ids are fed as-is instead of being trimmed to actual_soft_tokens * pooling^2. * Audio: each 640-sample frame is exactly one audio soft token, so the audio-token count equals the number of frames (no Conv2d stride-2 subsampling as in the Conformer speech encoder). Everything else (prompt image/audio token expansion, embedding fusion, per-layer-input decoder wiring) is shared with the standard gemma4 path. The paired preprocessing ops live in onnxruntime-extensions (Gemma4ImageTransform at patch_size=48/pooling=1, Gemma4UnifiedAudioFrames). Signed-off-by: Justin Chu --- src/models/gemma4_multimodal_processor.cpp | 28 +++++++++++++++++----- src/models/gemma4_multimodal_processor.h | 1 + src/models/model.cpp | 1 + src/models/model_type.h | 2 +- 4 files changed, 25 insertions(+), 7 deletions(-) diff --git a/src/models/gemma4_multimodal_processor.cpp b/src/models/gemma4_multimodal_processor.cpp index 3c7331a96b..24711cede5 100644 --- a/src/models/gemma4_multimodal_processor.cpp +++ b/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 Gemma4MultiModalProcessor::Process(const Tokenizer std::make_shared(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 audio_sizes_shape = {1}; auto audio_sizes = OrtValue::CreateTensor(allocator, audio_sizes_shape); audio_sizes->GetTensorMutableData()[0] = num_audio_tokens; @@ -320,7 +336,7 @@ std::unique_ptr 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 +381,7 @@ std::unique_ptr 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; diff --git a/src/models/gemma4_multimodal_processor.h b/src/models/gemma4_multimodal_processor.h index f53e4ec061..c9981731ec 100644 --- a/src/models/gemma4_multimodal_processor.h +++ b/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}; }; diff --git a/src/models/model.cpp b/src/models/model.cpp index 06e7729d51..0231f0e6a6 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -974,6 +974,7 @@ MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& sess {"phi4mm", Processor::Create}, {"gemma3", Processor::Create}, {"gemma4", Processor::Create}, + {"gemma4_unified", Processor::Create}, {"mistral3", Processor::Create}, {"fara", Processor::Create}, {"qwen2_5_vl", Processor::Create}, diff --git a/src/models/model_type.h b/src/models/model_type.h index 66f4dff5e2..f1f6704c8b 100644 --- a/src/models/model_type.h +++ b/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 MMM = {"gemma4", "phi4mm"}; + static constexpr std::array MMM = {"gemma4", "gemma4_unified", "phi4mm"}; return std::find(MMM.begin(), MMM.end(), model_type) != MMM.end(); } From 689fa0fe748edba7843aa229c6f3136361f9d272 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Jul 2026 22:38:59 +0000 Subject: [PATCH 2/3] Address review: clarify unified trim comment and add gemma4_unified test coverage * gemma4_multimodal_processor.cpp: rewrite the pixel-values comment to distinguish standard gemma4 (trim to actual teacher patches) from gemma4_unified (feed the full padded grid; the encoder-free graph strips padding via position_ids == -1). * Add test/python/create/create_dummy_gemma4_unified_models.py to derive a gemma4_unified fixture dir from the gemma4 fixtures with the unified I/O contract (pixel_values dim 6912, audio_embeds dim 640, model.type gemma4_unified, unified image/audio processor configs). * Add test/python/models/test_gemma4_unified_models.py covering processor creation and the unified vision/audio contracts (6912-dim pixel_values fed untrimmed; 640-dim audio frames with audio_sizes == frame count). Signed-off-by: Justin Chu --- src/models/gemma4_multimodal_processor.cpp | 11 +- .../create_dummy_gemma4_unified_models.py | 131 ++++++++++++++++ .../models/test_gemma4_unified_models.py | 144 ++++++++++++++++++ 3 files changed, 284 insertions(+), 2 deletions(-) create mode 100644 test/python/create/create_dummy_gemma4_unified_models.py create mode 100644 test/python/models/test_gemma4_unified_models.py diff --git a/src/models/gemma4_multimodal_processor.cpp b/src/models/gemma4_multimodal_processor.cpp index 24711cede5..a821207f3e 100644 --- a/src/models/gemma4_multimodal_processor.cpp +++ b/src/models/gemma4_multimodal_processor.cpp @@ -320,8 +320,15 @@ std::unique_ptr 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(actual_soft_tokens) * kPoolingKernelSize * kPoolingKernelSize; diff --git a/test/python/create/create_dummy_gemma4_unified_models.py b/test/python/create/create_dummy_gemma4_unified_models.py new file mode 100644 index 0000000000..6904347a60 --- /dev/null +++ b/test/python/create/create_dummy_gemma4_unified_models.py @@ -0,0 +1,131 @@ +# 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_unified_audio_frames", + "type": "Gemma4UnifiedAudioFrames", + "attrs": {"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() diff --git a/test/python/models/test_gemma4_unified_models.py b/test/python/models/test_gemma4_unified_models.py new file mode 100644 index 0000000000..ebea0c451c --- /dev/null +++ b/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]}" + ) From 678fdf5255424f2a4ad9b245e2341b8ac47c5846 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 10 Jul 2026 04:10:38 +0000 Subject: [PATCH 3/3] Use consolidated Gemma4Audio op in gemma4_unified fixture Follow the onnxruntime-extensions consolidation: the unified audio config now uses the single Gemma4Audio op with type="raw_frames" instead of the separate Gemma4UnifiedAudioFrames kernel. Signed-off-by: Justin Chu --- .../create/create_dummy_gemma4_unified_models.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/test/python/create/create_dummy_gemma4_unified_models.py b/test/python/create/create_dummy_gemma4_unified_models.py index 6904347a60..52cac59588 100644 --- a/test/python/create/create_dummy_gemma4_unified_models.py +++ b/test/python/create/create_dummy_gemma4_unified_models.py @@ -113,9 +113,14 @@ def main() -> None: {"operation": {"name": "audio_decoder", "type": "AudioDecoder"}}, { "operation": { - "name": "gemma4_unified_audio_frames", - "type": "Gemma4UnifiedAudioFrames", - "attrs": {"audio_samples_per_token": 640, "sampling_rate": 16000, "padding_value": 0.0}, + "name": "gemma4_audio", + "type": "Gemma4Audio", + "attrs": { + "type": "raw_frames", + "audio_samples_per_token": 640, + "sampling_rate": 16000, + "padding_value": 0.0, + }, } }, ]