From 52dfb2baa7c99a83929d921f6a55cc9c28e52858 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Jul 2026 18:11:07 +0000 Subject: [PATCH 1/3] 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 1337785cd..7770a09a3 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 40c86a0f0..5abe2674b 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/3] 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 7770a09a3..c99afec2a 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 5abe2674b..4858e7201 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 From b50c19731b4b6ba8a878e268a42dbd0e790c5dd3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 15 Sep 2026 11:40:14 -0700 Subject: [PATCH 3/3] Test Gemma4 unified native processor contract Lock the exported vision and audio graph ABI, processor attributes, runtime model type, semantic input mappings, and multimodal token metadata to the Extensions and GenAI consumer contracts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6e64e522-b07f-4c1c-abf1-3985268beb7e Signed-off-by: Justin Chu --- .../integrations/ort_genai/auto_export.py | 2 +- .../ort_genai/auto_export_test.py | 134 +++++++++++++++++- 2 files changed, 133 insertions(+), 3 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 24af0c2bc..28fad5659 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1340,7 +1340,7 @@ def _write_audio_processor_config( config: Any, output_dir: str, ) -> str | None: - """Write audio_processor.json for models with audio encoders. + """Write the native audio processor config for models with audio encoders. Returns the path if written, None otherwise. """ diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 0a95d0961..8efc14383 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -1004,9 +1004,14 @@ 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["name"] == "gemma4_audio" assert op["type"] == "Gemma4Audio" - assert op["attrs"]["type"] == "raw_frames" - assert op["attrs"]["audio_samples_per_token"] == 640 + assert op["attrs"] == { + "type": "raw_frames", + "audio_samples_per_token": 640, + "sampling_rate": 16000, + "padding_value": 0.0, + } def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path): config = mock.MagicMock() @@ -4080,6 +4085,131 @@ def test_gemma4_genai_config_from_real_model(self, tmp_path): assert data["model"]["vision"]["spatial_merge_size"] == 2 assert data["model"]["vision"]["config_filename"] == "image_processor.json" + def test_gemma4_unified_native_processor_contract(self, tmp_path): + """Export a real unified package and lock its Extensions/GenAI ABI.""" + from mobius._builder import build_from_module + from mobius._configs import Gemma4AudioConfig, Gemma4Config, VisionConfig + from mobius._registry import registry + from mobius.integrations.transformers._config_resolver import ( + _default_task_for_model, + ) + from mobius.tasks import get_task + + config = Gemma4Config( + model_type="gemma4_unified", + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=256, + rms_norm_eps=1e-6, + hidden_act="gelu_pytorch_tanh", + attn_qk_norm=True, + layer_types=["sliding_attention", "full_attention"], + sliding_window=8, + global_head_dim=32, + global_rope_theta=1_000_000.0, + global_partial_rotary_factor=0.25, + final_logit_softcapping=30.0, + hidden_size_per_layer_input=0, + num_global_key_value_heads=1, + attention_k_eq_v=True, + use_bidirectional_attention="vision", + image_token_id=253, + audio_token_id=252, + boa_token_id=251, + bos_token_id=2, + pad_token_id=0, + tie_word_embeddings=True, + vision=VisionConfig( + hidden_size=32, + position_embedding_size=1120, + patch_size=16, + pooling_kernel_size=3, + out_hidden_size=32, + norm_eps=1e-6, + mm_tokens_per_image=280, + ), + audio=Gemma4AudioConfig( + hidden_size=640, + output_proj_dims=640, + audio_token_id=252, + ), + ) + module = registry.get("gemma4_unified")(config) + task = get_task(_default_task_for_model("gemma4_unified")) + pkg = build_from_module(module, config, task=task) + pkg.config = config + + vision_inputs = {value.name: value for value in pkg["vision_encoder"].graph.inputs} + pixel_values = vision_inputs["pixel_values"] + pixel_position_ids = vision_inputs["pixel_position_ids"] + assert pixel_values.dtype == ir.DataType.FLOAT + assert len(pixel_values.shape) == 3 + assert pixel_values.shape[-1] == 48 * 48 * 3 + assert pixel_position_ids.dtype == ir.DataType.INT64 + assert len(pixel_position_ids.shape) == 3 + assert pixel_position_ids.shape[-1] == 2 + + audio_inputs = {value.name: value for value in pkg["audio_encoder"].graph.inputs} + input_features = audio_inputs["input_features"] + input_features_mask = audio_inputs["input_features_mask"] + assert input_features.dtype == ir.DataType.FLOAT + assert len(input_features.shape) == 3 + assert input_features.shape[-1] == 640 + assert input_features_mask.dtype == ir.DataType.BOOL + assert len(input_features_mask.shape) == 2 + + result = write_ort_genai_config(pkg, str(tmp_path)) + with open(result["genai_config"], encoding="utf-8") as f: + data = json.load(f) + + model = data["model"] + assert model["type"] == "gemma4_unified" + assert model["image_token_id"] == 253 + assert model["audio_token_id"] == 252 + assert model["boa_token_id"] == 251 + assert model["vision"]["config_filename"] == "image_processor.json" + assert model["vision"]["inputs"] == { + "pixel_values": "pixel_values", + "pixel_position_ids": "pixel_position_ids", + } + assert model["speech"]["config_filename"] == "audio_feature_extraction.json" + assert model["speech"]["inputs"] == { + "audio_embeds": "input_features", + "attention_mask": "input_features_mask", + } + assert model["speech"]["outputs"] == {"audio_features": "audio_features"} + + with open(result["processor_config"], encoding="utf-8") as f: + image_processor = json.load(f) + image_op = image_processor["processor"]["transforms"][1]["operation"] + assert image_op == { + "name": "gemma4_image_transform", + "type": "Gemma4ImageTransform", + "attrs": { + "patch_size": 48, + "max_soft_tokens": 280, + "pooling_kernel_size": 1, + }, + } + + with open(result["audio_processor"], encoding="utf-8") as f: + audio_processor = json.load(f) + audio_op = audio_processor["feature_extraction"]["sequence"][1]["operation"] + assert audio_op == { + "name": "gemma4_audio", + "type": "Gemma4Audio", + "attrs": { + "type": "raw_frames", + "audio_samples_per_token": 640, + "sampling_rate": 16000, + "padding_value": 0.0, + }, + } + def test_text_only_genai_config_is_decoder_only(self, tmp_path): """text_only gemma4_unified export -> decoder-only genai config.