diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index ef906199c..28fad5659 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -99,11 +99,12 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]: "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", @@ -194,12 +195,12 @@ class _DecoderAbi: ) # 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"}) _GLMASR_MODEL_TYPES = frozenset({"glmasr"}) _MINICPM_MODEL_TYPES = frozenset({"minicpmv4_6"}) @@ -937,9 +938,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). - **Gemma3** (``gemma3`` or ``gemma3_text``): Writes ``processor_config.json`` with a 5-step pipeline (DecodeImage → Resize[fixed] → Rescale → Normalize → Permute3D). Uses a @@ -971,16 +973,6 @@ 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 if model_type in _MINICPM_MODEL_TYPES: # MiniCPM needs adaptive slicing and NaViT horizontal patch packing. # ort-extensions has no equivalent transform, so preserving the HF @@ -1007,7 +999,49 @@ def _write_vision_processor_config( 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 = ( @@ -1017,7 +1051,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": [ @@ -1306,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. """ @@ -1317,18 +1351,40 @@ 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: - # Gemma4 USM-style 128-dim log-mel spectrogram. + # 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 + # ``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": { + "sequence": [ + { + "operation": { + "name": "audio_decoder", + "type": "AudioDecoder", + } + }, + { + "operation": { + "name": "gemma4_audio", + "type": "Gemma4Audio", + "attrs": { + "type": "raw_frames", + "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 via the ort-extensions + # ``Gemma4Audio`` op with ``type="log_mel"``. # OrtxCreateSpeechFeatureExtractor requires the feature_extraction.sequence format. processor = { "feature_extraction": { @@ -1341,9 +1397,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 1d3f4da7d..8efc14383 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -155,10 +155,10 @@ def test_qwen25_text_subconfig_maps_to_multimodal_runtime(self): assert _resolve_ort_genai_model_type("qwen2_5_vl_text") == "qwen2_5_vl" 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" @@ -194,12 +194,12 @@ def test_decoder_only_uses_generic_decoder(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): @@ -420,14 +420,29 @@ def test_mage_vl_processor_propagates_trust_remote_code(self, tmp_path): trust_remote_code=True, ) - 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.""" @@ -973,12 +988,30 @@ 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["name"] == "gemma4_audio" + assert op["type"] == "Gemma4Audio" + 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() @@ -996,8 +1029,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 @@ -1934,9 +1968,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 @@ -4050,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. @@ -4134,6 +4294,8 @@ def test_text_only_genai_config_is_decoder_only(self, tmp_path): with open(result["genai_config"]) as f: data = json.load(f) + # Decoder-only package: the generic "decoder" type, NOT the + # multimodal HF gemma4_unified -> gemma4_unified. assert data["model"]["type"] == "decoder" # Decoder-only: input_ids decoder, no multimodal sections. assert "vision" not in data["model"]