From 83a379e291a93e4c85fe4ab6e9b5ed34784f47ce Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 23 Apr 2026 19:27:28 -0700 Subject: [PATCH 1/4] Add Gemma4 image_processor.json generation with onnxruntime-extensions transforms pipeline --- .../integrations/ort_genai/auto_export.py | 40 ++++++++++++++----- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 37a9ff74..0e4fc0e1 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -173,30 +173,49 @@ def _write_processor_config( model_type = getattr(config, "model_type", "") if model_type in ("gemma4", "gemma4_text"): - # Gemma4 needs a processor wrapper with model-specific fields - tokens_per_image = ( + # Gemma4 needs an onnxruntime-extensions format processor config + # with a transforms pipeline (DecodeImage -> Gemma4ImageTransform). + # The OrtxCreateProcessor API requires this format. + max_soft_tokens = ( getattr(vision, "mm_tokens_per_image", None) or getattr(config, "mm_tokens_per_image", None) or getattr(vision, "max_soft_tokens", None) or 280 ) - image_size = getattr(vision, "image_size", None) or 448 patch_size = getattr(vision, "patch_size", None) or 16 - processor: dict[str, Any] = { + pooling_kernel_size = getattr(vision, "pooling_kernel_size", None) or 3 + processor = { "processor": { - "name": "gemma4_image_processor", - "image_size": image_size, - "patch_size": patch_size, - "tokens_per_image": tokens_per_image, + "name": "gemma_4_image_processing", + "transforms": [ + { + "operation": { + "name": "decode_image", + "type": "DecodeImage", + "attrs": {"color_space": "RGB"}, + } + }, + { + "operation": { + "name": "gemma4_image_transform", + "type": "Gemma4ImageTransform", + "attrs": { + "patch_size": patch_size, + "max_soft_tokens": max_soft_tokens, + "pooling_kernel_size": pooling_kernel_size, + }, + } + }, + ], } } + proc_filename = "image_processor.json" else: processor = { "image_size": getattr(vision, "image_size", None) or 448, "patch_size": getattr(vision, "patch_size", None) or 14, } - - proc_filename = "processor_config.json" + proc_filename = "processor_config.json" path = os.path.join(output_dir, proc_filename) with open(path, "w", encoding="utf-8") as f: @@ -280,6 +299,7 @@ def _write_genai_config( vision_cfg = getattr(config, "vision", None) sms = getattr(vision_cfg, "spatial_merge_size", 2) vision_kwargs["spatial_merge_size"] = sms + vision_kwargs["config_filename"] = "image_processor.json" if vision_input_mapping is not None: vision_kwargs["input_names"] = vision_input_mapping From 2deb57ed6451e2d7ea87ef0f9f8173fe6a1f58ae Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 23 Apr 2026 21:42:23 -0700 Subject: [PATCH 2/4] update tests --- src/mobius/integrations/ort_genai/auto_export_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 60b8c21c..1a0f0f66 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -780,7 +780,7 @@ def test_gemma4_genai_config_from_real_model(self, tmp_path): # Config-level properties are still present assert data["model"]["image_token_id"] == 255999 assert data["model"]["vision"]["spatial_merge_size"] == 2 - assert data["model"]["vision"]["config_filename"] == "processor_config.json" + assert data["model"]["vision"]["config_filename"] == "image_processor.json" def test_auto_export_produces_genai_config(self, tmp_path): """Mock build() to return a tiny package, verify genai_config.""" From bcb3f649657acc659d62fd09b43aad788b185bc2 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 24 Apr 2026 10:42:40 -0700 Subject: [PATCH 3/4] Address PR comments --- .../integrations/ort_genai/auto_export.py | 22 +++++-- .../ort_genai/auto_export_test.py | 63 +++++++++++++++++++ 2 files changed, 80 insertions(+), 5 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 0e4fc0e1..098f78c8 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -158,13 +158,22 @@ def _copy_tokenizer_files_from_local( return copied -def _write_processor_config( +def _write_vision_processor_config( config: Any, output_dir: str, ) -> str | None: - """Write a minimal processor_config.json for VLM models. + """Write the vision processor config file for VLM models. - Returns the path if written, None otherwise. + The output format depends on the model type: + + - **Gemma4** (``gemma4``, ``gemma4_text``): Writes ``image_processor.json`` + in the onnxruntime-extensions transforms pipeline format required by + ``OrtxCreateProcessor``. The pipeline is + ``DecodeImage → Gemma4ImageTransform``. + - **Other models**: Writes ``processor_config.json`` with a minimal + HuggingFace-style schema (``image_size``, ``patch_size``). + + Returns the written file path, or None if the config has no vision section. """ vision = getattr(config, "vision", None) if vision is None: @@ -176,10 +185,13 @@ def _write_processor_config( # Gemma4 needs an onnxruntime-extensions format processor config # with a transforms pipeline (DecodeImage -> Gemma4ImageTransform). # The OrtxCreateProcessor API requires this format. + # + # max_soft_tokens: maps from HF's mm_tokens_per_image (the number of + # vision tokens per image after pooling) into the Gemma4ImageTransform's + # max_soft_tokens attribute, which controls the padded patch budget. max_soft_tokens = ( getattr(vision, "mm_tokens_per_image", None) or getattr(config, "mm_tokens_per_image", None) - or getattr(vision, "max_soft_tokens", None) or 280 ) patch_size = getattr(vision, "patch_size", None) or 16 @@ -470,7 +482,7 @@ def write_ort_genai_config( result[tf] = os.path.join(directory, tf) # Write processor_config.json for VLMs - processor_path = _write_processor_config(config, directory) + processor_path = _write_vision_processor_config(config, directory) if processor_path: result["processor_config"] = processor_path diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 1a0f0f66..113286c7 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -314,6 +314,69 @@ def test_processor_config_not_written_without_vision(self, tmp_path): assert "processor_config" not in result assert not os.path.exists(os.path.join(str(tmp_path), "processor_config.json")) + def test_gemma4_image_processor_json_written(self, tmp_path): + """Gemma4 writes image_processor.json with onnxruntime-extensions transforms pipeline.""" + import dataclasses + + from mobius._model_package import ModelPackage + from mobius.integrations.ort_genai.auto_export import write_ort_genai_config + + @dataclasses.dataclass + class FakeVision: + image_size: int = 448 + patch_size: int = 16 + mm_tokens_per_image: int = 260 + pooling_kernel_size: int = 3 + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "gemma4" + vocab_size: int = 262144 + hidden_size: int = 1536 + num_hidden_layers: int = 35 + num_attention_heads: int = 8 + num_key_value_heads: int = 1 + head_dim: int = 256 + vision: FakeVision = dataclasses.field(default_factory=FakeVision) + + pkg = ModelPackage( + { + "model": mock.MagicMock(), + "vision": mock.MagicMock(), + "embedding": mock.MagicMock(), + }, + config=FakeConfig(), + ) + result = write_ort_genai_config(pkg, str(tmp_path)) + + # Should write image_processor.json, not processor_config.json + assert "processor_config" in result + proc_path = result["processor_config"] + assert proc_path.endswith("image_processor.json") + assert os.path.isfile(proc_path) + assert not os.path.exists(os.path.join(str(tmp_path), "processor_config.json")) + + with open(proc_path) as f: + data = json.load(f) + + # Verify onnxruntime-extensions transforms pipeline structure + assert "processor" in data + assert "transforms" in data["processor"] + transforms = data["processor"]["transforms"] + assert len(transforms) == 2 + + # First op: DecodeImage + op0 = transforms[0]["operation"] + assert op0["type"] == "DecodeImage" + assert op0["attrs"]["color_space"] == "RGB" + + # Second op: Gemma4ImageTransform with correct attrs from config + op1 = transforms[1]["operation"] + assert op1["type"] == "Gemma4ImageTransform" + assert op1["attrs"]["patch_size"] == 16 + assert op1["attrs"]["max_soft_tokens"] == 260 + assert op1["attrs"]["pooling_kernel_size"] == 3 + def test_tokenizer_not_copied_without_model_id(self, tmp_path): """No tokenizer files copied when hf_model_id=None.""" from mobius.integrations.ort_genai.auto_export import write_ort_genai_config From a7ee07f5976e0541301c6a619b5e7f3f7e36def0 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 24 Apr 2026 11:13:47 -0700 Subject: [PATCH 4/4] Fix tests --- src/mobius/integrations/ort_genai/auto_export_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 113286c7..3f18708e 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -18,7 +18,7 @@ _graph_input_names, _resolve_ort_genai_model_type, _write_genai_config, - _write_processor_config, + _write_vision_processor_config, write_ort_genai_config, ) @@ -42,7 +42,7 @@ class TestWriteProcessorConfig: def test_no_vision_returns_none(self, tmp_path): config = mock.MagicMock(spec=[]) del config.vision # ensure no vision attribute - assert _write_processor_config(config, str(tmp_path)) is None + assert _write_vision_processor_config(config, str(tmp_path)) is None def test_writes_vision_config(self, tmp_path): vision = mock.MagicMock() @@ -51,7 +51,7 @@ def test_writes_vision_config(self, tmp_path): config = mock.MagicMock() config.vision = vision - path = _write_processor_config(config, str(tmp_path)) + path = _write_vision_processor_config(config, str(tmp_path)) assert path is not None with open(path) as f: data = json.load(f)