From 52dfb2baa7c99a83929d921f6a55cc9c28e52858 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Jul 2026 18:11:07 +0000 Subject: [PATCH 1/2] ort_genai: emit native processor configs for gemma-4-12B unified The encoder-free gemma-4-12B "unified" model can now be preprocessed natively by ORT GenAI (no HuggingFace processor / set_inputs needed), so stop skipping its processor artifacts and emit the correct configs: * model.type: gemma4_unified now maps to the dedicated "gemma4_unified" ORT GenAI multimodal type (was reusing "gemma4"). * image_processor.json: DecodeImage -> Gemma4ImageTransform configured for the 48px merged-patch contract (patch_size=48, pooling_kernel_size=1, patch_dim=6912). HF builds these via 16px patchify + 3x3 patches_merge, which is identical to a direct 48px patchify, so the existing op is reused with merged geometry. * audio_feature_extraction.json: AudioDecoder -> Gemma4UnifiedAudioFrames (raw 640-sample waveform framing) instead of the 128-dim log-mel op. Requires the companion onnxruntime-extensions (Gemma4UnifiedAudioFrames) and onnxruntime-genai (gemma4_unified processor) support. Updates the unified auto_export tests to assert the emitted configs. Signed-off-by: Justin Chu --- .../integrations/ort_genai/auto_export.py | 129 ++++++++++++------ .../ort_genai/auto_export_test.py | 53 +++++-- 2 files changed, 131 insertions(+), 51 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 1337785c..7770a09a 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -75,11 +75,12 @@ "gemma2": "gemma", "gemma4": "gemma4", "gemma4_text": "gemma4_text", - # gemma-4-12B "unified" (encoder-free) variant reuses the gemma4 ORT GenAI - # pipelines: the multimodal package (decoder taking inputs_embeds + vision - # embedder + embedding fusion) maps to "gemma4"; the standalone text - # backbone maps to "gemma4_text". - "gemma4_unified": "gemma4", + # gemma-4-12B "unified" (encoder-free) variant. The multimodal package + # (decoder taking inputs_embeds + encoder-free vision/audio embedders + + # embedding fusion) maps to the dedicated "gemma4_unified" ORT GenAI type, + # which drives the unified processor (48px merged patches / raw 640-sample + # audio frames). The standalone text backbone maps to "gemma4_text". + "gemma4_unified": "gemma4_unified", "gemma4_unified_text": "gemma4_text", "mistral": "mistral", "mistral3": "mistral3", @@ -99,12 +100,12 @@ ) # Encoder-free gemma-4-12B "unified" variants. Their image/audio inputs are raw # merged pixel patches (48px, 6912-dim) / raw waveform frames (640-dim), NOT the -# SigLIP 16px / 128-dim log-mel contract that the ort-extensions -# ``Gemma4ImageTransform`` / ``Gemma4LogMel`` ops implement. There is no -# genai-native transform for the unified contract, so we deliberately do NOT -# emit image_processor.json / audio_processor.json for these models — callers -# must preprocess with the HuggingFace processor and feed tensors via -# ``Generator.set_inputs`` (see examples/gemma4_unified_ort_genai.py). +# SigLIP 16px / 128-dim log-mel contract of the standard gemma4 (E2B/E4B) model. +# These are produced natively by ort-extensions: ``Gemma4ImageTransform`` with +# patch_size=48 / pooling_kernel_size=1 (the 48px merged patch is identical to a +# direct 48px patchify) and the ``Gemma4UnifiedAudioFrames`` raw-framing op. The +# genai ``gemma4_unified`` processor consumes both (see the companion +# onnxruntime-genai / onnxruntime-extensions support). _GEMMA4_UNIFIED_MODEL_TYPES = frozenset({"gemma4_unified", "gemma4_unified_text"}) _PIXTRAL_MODEL_TYPES = frozenset({"mistral3"}) _QWEN_VL_MODEL_TYPES = frozenset( @@ -427,9 +428,10 @@ def _write_vision_processor_config( - **Gemma4** (``gemma4``, ``gemma4_text``): Writes ``image_processor.json`` with a ``DecodeImage → Gemma4ImageTransform`` pipeline. - - **Gemma4 unified** (``gemma4_unified*``): Returns ``None`` — the - encoder-free model has no matching ort-extensions transform; callers feed - HF-preprocessed pixel_values via ``Generator.set_inputs``. + - **Gemma4 unified** (``gemma4_unified*``): Writes ``image_processor.json`` + with a ``DecodeImage → Gemma4ImageTransform`` pipeline configured for the + encoder-free 48px merged-patch contract (patch_size=48, + pooling_kernel_size=1, patch_dim=6912). - **Pixtral / Mistral3**: Writes ``processor_config.json`` with a 7-step pipeline (DecodeImage → ConvertRGB → Resize → Rescale → Normalize → Permute3D → PixtralImageSizes). @@ -443,21 +445,52 @@ def _write_vision_processor_config( return None model_type = getattr(config, "model_type", "") - if model_type in _GEMMA4_UNIFIED_MODEL_TYPES: - # Encoder-free unified model: no ort-extensions transform matches its - # raw merged-patch contract. Emit no image_processor.json; callers feed - # HF-preprocessed pixel_values via Generator.set_inputs. - logger.info( - "Skipping image_processor.json for encoder-free %s " - "(no native ort-extensions transform; use HF processor + set_inputs)", - model_type, - ) - return None - vision_model_type = getattr(vision, "model_type", None) is_pixtral = vision_model_type == "pixtral" or model_type in _PIXTRAL_MODEL_TYPES - if model_type in _GEMMA4_MODEL_TYPES: + if model_type in _GEMMA4_UNIFIED_MODEL_TYPES: + # Encoder-free unified model: it consumes 48px MERGED patches + # (patch_dim = 48*48*3 = 6912) directly, with no SigLIP tower to pool + # 3x3 teacher patches. HuggingFace produces these via + # (16px patchify -> 3x3 patches_merge), which is provably identical to a + # direct 48px patchify. So we reuse the same ort-extensions + # ``Gemma4ImageTransform`` op with the merged geometry: patch_size = + # patch_size * pooling_kernel_size (48) and pooling_kernel_size = 1. + max_soft_tokens = ( + getattr(vision, "mm_tokens_per_image", None) + or getattr(config, "mm_tokens_per_image", None) + or 280 + ) + patch_size = getattr(vision, "patch_size", None) or 16 + pooling = getattr(vision, "pooling_kernel_size", None) or 3 + model_patch_size = patch_size * pooling # 16 * 3 = 48 + processor_config: dict[str, Any] = { + "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": model_patch_size, + "max_soft_tokens": max_soft_tokens, + "pooling_kernel_size": 1, + }, + } + }, + ], + } + } + path = os.path.join(output_dir, "image_processor.json") + elif model_type in _GEMMA4_MODEL_TYPES: # Gemma4 needs an onnxruntime-extensions format processor config # with a transforms pipeline (DecodeImage -> Gemma4ImageTransform). max_soft_tokens = ( @@ -467,7 +500,7 @@ def _write_vision_processor_config( ) patch_size = getattr(vision, "patch_size", None) or 16 pooling_kernel_size = getattr(vision, "pooling_kernel_size", None) or 3 - processor_config: dict[str, Any] = { + processor_config = { "processor": { "name": "gemma_4_image_processing", "transforms": [ @@ -631,17 +664,37 @@ def _write_audio_processor_config( model_type = getattr(config, "model_type", "") if model_type in _GEMMA4_UNIFIED_MODEL_TYPES: - # Encoder-free unified model: raw 640-dim waveform frames, not the - # 128-dim log-mel Gemma4LogMel contract. Emit no audio_processor.json; - # callers feed HF-preprocessed input_features via Generator.set_inputs. - logger.info( - "Skipping audio_processor.json for encoder-free %s " - "(no native ort-extensions transform; use HF processor + set_inputs)", - model_type, - ) - return None - - if model_type in _GEMMA4_MODEL_TYPES: + # Encoder-free unified model: each audio soft token is a raw chunk of the + # 16 kHz waveform (audio_samples_per_token = audio_embed_dim, 640), not a + # 128-dim log-mel frame. Reproduced natively by the ort-extensions + # ``Gemma4UnifiedAudioFrames`` op (pad to a whole number of frames, + # reshape to (num_tokens, 640)). + samples_per_token = getattr(audio, "hidden_size", None) or 640 + processor = { + "feature_extraction": { + "sequence": [ + { + "operation": { + "name": "audio_decoder", + "type": "AudioDecoder", + } + }, + { + "operation": { + "name": "gemma4_unified_audio_frames", + "type": "Gemma4UnifiedAudioFrames", + "attrs": { + "audio_samples_per_token": samples_per_token, + "sampling_rate": 16000, + "padding_value": 0.0, + }, + } + }, + ] + } + } + proc_filename = "audio_feature_extraction.json" + elif model_type in _GEMMA4_MODEL_TYPES: # Gemma4 USM-style 128-dim log-mel spectrogram. # OrtxCreateSpeechFeatureExtractor requires the feature_extraction.sequence format. processor = { diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 40c86a0f..5abe2674 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -89,10 +89,10 @@ def test_phi4mm_model_types(self): assert _resolve_ort_genai_model_type("phi") == "phi" def test_gemma4_unified_model_types(self): - # The gemma-4-12B unified checkpoint (model_type "gemma4_unified") - # reuses the multimodal "gemma4" ORT GenAI pipeline; its standalone - # text decoder ("gemma4_unified_text") maps to "gemma4_text". - assert _resolve_ort_genai_model_type("gemma4_unified") == "gemma4" + # The gemma-4-12B unified checkpoint (model_type "gemma4_unified") maps + # to the dedicated "gemma4_unified" ORT GenAI multimodal pipeline; its + # standalone text decoder ("gemma4_unified_text") maps to "gemma4_text". + assert _resolve_ort_genai_model_type("gemma4_unified") == "gemma4_unified" assert _resolve_ort_genai_model_type("gemma4_unified_text") == "gemma4_text" # Released gemma4 mappings remain unchanged. assert _resolve_ort_genai_model_type("gemma4") == "gemma4" @@ -115,12 +115,12 @@ def test_decoder_only_prefers_config_type(self): def test_multimodal_keeps_hf_type(self): # Full multimodal export: build() unwraps the composite to its text # sub-config, so config.model_type may be a text type even though the - # package is multimodal. Must keep the HF parent type -> gemma4. + # package is multimodal. Must keep the HF parent type -> gemma4_unified. assert ( _select_ort_model_type( "gemma4_unified_text", "gemma4_unified", is_decoder_only=False ) - == "gemma4" + == "gemma4_unified" ) def test_decoder_only_falls_back_to_hf_when_config_missing(self): @@ -177,14 +177,29 @@ def test_writes_transform_pipeline(self, tmp_path): assert len(norm_attrs["mean"]) == 3 assert len(norm_attrs["std"]) == 3 - def test_gemma4_unified_skips_image_processor(self, tmp_path): - """Encoder-free gemma4_unified has no native transform: no image_processor.json.""" + def test_gemma4_unified_image_processor(self, tmp_path): + """Encoder-free gemma4_unified writes a 48px merged-patch transform.""" vision = mock.MagicMock() vision.model_type = None + vision.patch_size = 16 + vision.pooling_kernel_size = 3 + vision.mm_tokens_per_image = 280 config = mock.MagicMock() config.vision = vision config.model_type = "gemma4_unified" - assert _write_vision_processor_config(config, str(tmp_path)) is None + + path = _write_vision_processor_config(config, str(tmp_path)) + assert path is not None + assert path.endswith("image_processor.json") + with open(path) as f: + data = json.load(f) + + transform = data["processor"]["transforms"][1]["operation"] + assert transform["type"] == "Gemma4ImageTransform" + # 48px merged patches (16*3), no further pooling. + assert transform["attrs"]["patch_size"] == 48 + assert transform["attrs"]["pooling_kernel_size"] == 1 + assert transform["attrs"]["max_soft_tokens"] == 280 def test_pixtral_vision_config(self, tmp_path): """Generates pixtral-specific processor config with 7 transforms.""" @@ -326,12 +341,24 @@ def test_audio_non_gemma4_returns_none(self, tmp_path): config.model_type = "whisper" assert _write_audio_processor_config(config, str(tmp_path)) is None - def test_audio_gemma4_unified_skips_audio_processor(self, tmp_path): - """Encoder-free gemma4_unified has no native transform: no audio_processor.json.""" + def test_audio_gemma4_unified_writes_raw_frames(self, tmp_path): + """Encoder-free gemma4_unified writes a raw-waveform-frame extractor.""" config = mock.MagicMock() config.audio = mock.MagicMock() + config.audio.hidden_size = 640 config.model_type = "gemma4_unified" - assert _write_audio_processor_config(config, str(tmp_path)) is None + + path = _write_audio_processor_config(config, str(tmp_path)) + assert path is not None + assert path.endswith("audio_feature_extraction.json") + with open(path) as f: + data = json.load(f) + + seq = data["feature_extraction"]["sequence"] + assert seq[0]["operation"]["type"] == "AudioDecoder" + op = seq[1]["operation"] + assert op["type"] == "Gemma4UnifiedAudioFrames" + assert op["attrs"]["audio_samples_per_token"] == 640 def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path): config = mock.MagicMock() @@ -1513,7 +1540,7 @@ def test_text_only_genai_config_is_decoder_only(self, tmp_path): data = json.load(f) # ORT-GenAI type resolved from pkg.config.model_type (gemma4_text), - # NOT the multimodal HF gemma4_unified -> gemma4. + # NOT the multimodal HF gemma4_unified -> gemma4_unified. assert data["model"]["type"] == "gemma4_text" # Decoder-only: input_ids decoder, no multimodal sections. assert "vision" not in data["model"] From 4082f1cbc1ef2282c920747da2598c4b27390f87 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 10 Jul 2026 04:12:04 +0000 Subject: [PATCH 2/2] Emit consolidated Gemma4Audio op for gemma-4 audio configs Follow the onnxruntime-extensions consolidation: both gemma4 audio configs now use the single Gemma4Audio op with an explicit type attribute instead of two separate op types. * gemma4 / gemma4_text -> Gemma4Audio type="log_mel" (128-dim USM log-mel) * gemma4_unified* -> Gemma4Audio type="raw_frames" (raw 640-sample frames) Updates the auto_export audio-config tests accordingly. 115 tests pass. Signed-off-by: Justin Chu --- .../integrations/ort_genai/auto_export.py | 17 ++++++++++------- .../integrations/ort_genai/auto_export_test.py | 11 +++++++---- 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 7770a09a..c99afec2 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -667,8 +667,8 @@ def _write_audio_processor_config( # Encoder-free unified model: each audio soft token is a raw chunk of the # 16 kHz waveform (audio_samples_per_token = audio_embed_dim, 640), not a # 128-dim log-mel frame. Reproduced natively by the ort-extensions - # ``Gemma4UnifiedAudioFrames`` op (pad to a whole number of frames, - # reshape to (num_tokens, 640)). + # ``Gemma4Audio`` op with ``type="raw_frames"`` (pad to a whole number of + # frames, reshape to (num_tokens, 640)). samples_per_token = getattr(audio, "hidden_size", None) or 640 processor = { "feature_extraction": { @@ -681,9 +681,10 @@ def _write_audio_processor_config( }, { "operation": { - "name": "gemma4_unified_audio_frames", - "type": "Gemma4UnifiedAudioFrames", + "name": "gemma4_audio", + "type": "Gemma4Audio", "attrs": { + "type": "raw_frames", "audio_samples_per_token": samples_per_token, "sampling_rate": 16000, "padding_value": 0.0, @@ -695,7 +696,8 @@ def _write_audio_processor_config( } proc_filename = "audio_feature_extraction.json" elif model_type in _GEMMA4_MODEL_TYPES: - # Gemma4 USM-style 128-dim log-mel spectrogram. + # Gemma4 USM-style 128-dim log-mel spectrogram via the ort-extensions + # ``Gemma4Audio`` op with ``type="log_mel"``. # OrtxCreateSpeechFeatureExtractor requires the feature_extraction.sequence format. processor = { "feature_extraction": { @@ -708,9 +710,10 @@ def _write_audio_processor_config( }, { "operation": { - "name": "gemma4_log_mel", - "type": "Gemma4LogMel", + "name": "gemma4_audio", + "type": "Gemma4Audio", "attrs": { + "type": "log_mel", "feature_size": 128, "sampling_rate": 16000, "frame_length_ms": 20.0, diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 5abe2674..4858e720 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -357,7 +357,8 @@ def test_audio_gemma4_unified_writes_raw_frames(self, tmp_path): seq = data["feature_extraction"]["sequence"] assert seq[0]["operation"]["type"] == "AudioDecoder" op = seq[1]["operation"] - assert op["type"] == "Gemma4UnifiedAudioFrames" + assert op["type"] == "Gemma4Audio" + assert op["attrs"]["type"] == "raw_frames" assert op["attrs"]["audio_samples_per_token"] == 640 def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path): @@ -376,8 +377,9 @@ def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path): seq = data["feature_extraction"]["sequence"] assert len(seq) == 2 assert seq[0]["operation"]["type"] == "AudioDecoder" - assert seq[1]["operation"]["type"] == "Gemma4LogMel" + assert seq[1]["operation"]["type"] == "Gemma4Audio" attrs = seq[1]["operation"]["attrs"] + assert attrs["type"] == "log_mel" assert attrs["feature_size"] == 128 assert attrs["sampling_rate"] == 16000 assert attrs["frame_length_ms"] == 20.0 # noqa: RUF069 @@ -747,9 +749,10 @@ class FakeConfig: op0 = seq[0]["operation"] assert op0["type"] == "AudioDecoder" - # Second op: Gemma4LogMel with expected attrs + # Second op: Gemma4Audio (type=log_mel) with expected attrs op1 = seq[1]["operation"] - assert op1["type"] == "Gemma4LogMel" + assert op1["type"] == "Gemma4Audio" + assert op1["attrs"]["type"] == "log_mel" assert op1["attrs"]["feature_size"] == 128 assert op1["attrs"]["sampling_rate"] == 16000 assert op1["attrs"]["mel_floor"] == 0.001 # noqa: RUF069