[openvino] Add audio embeddings export support for gemma4_unified

[lcretan]

Summary

Add audio embeddings export support for gemma4_unified models (e.g., google/gemma-4-12B-it).

Background

Gemma4UnifiedForConditionalGeneration has an encoder-free audio embedder (embed_audio)
consisting of RMSNorm + Linear(640 → 3840) — identical in structure to the vision embedder.
While _OVGemma4UnifiedForCausalLM already supports text and vision modalities, audio was
explicitly excluded with the comment "audio is not exported".

This PR adds full audio export and inference support.

Changes

New: DummyGemma4UnifiedAudioInputGenerator (input_generators.py)

• Generates dummy input_features: [batch, n_frames, 640] for tracing
• No data-dependent operations in embed_audio — trace is straightforward

New: Gemma4UnifiedAudioEmbeddingsModelPatcher (model_patcher.py)

• Wraps model.model.embed_audio.forward() for OV export
• Exports as openvino_audio_embeddings_model.xml/bin

Updated: Gemma4UnifiedOpenVINOConfig (model_configs.py)

• SUPPORTED_BEHAVIORS now includes "audio_embeddings"
• with_behavior("audio_embeddings") returns audio export config
• get_model_for_behavior() returns the audio sub-module
• patch_model_for_export() uses Gemma4UnifiedAudioEmbeddingsModelPatcher
• inputs/outputs properties handle audio shape annotations

Updated: _OVGemma4UnifiedForCausalLM (modeling_visual_language.py)

• additional_parts = ["audio_embeddings"] (with backward-compat from_pretrained)
• get_audio_embeddings(input_features, input_features_mask) — calls OV audio model, strips padding
• merge_audio_text_embeddings() — scatters audio embeddings at audio_token_id positions

Audio Architecture

Raw waveform (16 kHz PCM)
    │
    ▼ Gemma4UnifiedAudioFeatureExtractor
input_features: [batch, n_frames, 640]  (640 samples/frame = 40ms/token)
    │
    ▼ openvino_audio_embeddings_model  ← NEW (2.4 MB)
    │  embed_audio = RMSNorm + Linear(640 → 3840)
    ▼
[batch, n_frames, 3840]  → strip padding → scatter at <|audio|> positions
    │
    ▼ openvino_language_model (existing)
    text output

Backward Compatibility

• Models exported without audio (no openvino_audio_embeddings_model.xml) load normally
• from_pretrained() detects audio model file and skips gracefully if absent

Testing

# Export with audio
optimum-cli export openvino \
  --model google/gemma-4-12B-it \
  --task image-text-to-text \
  --weight-format int8 \
  ./gemma-4-12B-it-openvino-int8-audio

# Verify audio model was created
ls gemma-4-12B-it-openvino-int8-audio/openvino_audio_embeddings_model.*

# Confirmed output on Intel Core Ultra 9 185H (Arc iGPU):
# openvino_audio_embeddings_model.bin  2.4 MB
# openvino_audio_embeddings_model.xml  (topology)

References

• Related Issue [OpenVINO] Support Gemma 4 12B with model type gemma4_unified #1764: Gemma 4 audio support
• Transformers: Gemma4UnifiedAudioConfig, Gemma4UnifiedMultimodalEmbedder
• Audio token IDs: audio_token_id=258881, boa_token_id=256000, eoa_token_index=258883

[rkazants]

@lcretan, please provide code for inference, tests, and share code snippets for inference in PR description

[Copilot]

Pull request overview

Adds OpenVINO export + runtime support for Gemma 4 Unified audio embeddings by introducing an audio dummy input generator, an audio export model patcher, wiring a new "audio_embeddings" export behavior into Gemma4UnifiedOpenVINOConfig, and extending the runtime VLM wrapper to accept audio inputs and merge audio soft tokens into the text embedding stream.

Changes:

• Add an OpenVINO export path for Gemma4 Unified audio embeddings (embed_audio) including dummy tracing inputs and a dedicated model patcher.
• Extend Gemma4 Unified OpenVINO config behaviors to include "audio_embeddings" and produce an openvino_audio_embeddings_model.* artifact.
• Add runtime-side audio preprocessing/forward plumbing and embedding scatter logic in _OVGemma4UnifiedForCausalLM.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
optimum/intel/openvino/modeling_visual_language.py	Adds optional audio submodel loading, audio embedding inference, and audio token scatter/merge support for gemma4_unified.
optimum/exporters/openvino/model_patcher.py	Introduces a patcher to export embed_audio with a trace-friendly forward signature.
optimum/exporters/openvino/model_configs.py	Adds "audio_embeddings" behavior for gemma4_unified export and configures dummy inputs / patching / IO shapes.
optimum/exporters/openvino/input_generators.py	Adds a dummy input generator producing [batch, n_frames, 640] audio frames for tracing.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[lcretan]

Fixed in c3718cb — from_pretrained now checks os.path.isdir(model_id) before treating the path as local, so Hub IDs like google/gemma-4-12B-it fall through to the parent's hf_hub_download optional-key handling instead of being incorrectly stripped of audio. Covered by TestFromPretrainedBackwardCompat::test_hub_id_passes_through_unchanged.

[lcretan]

Addressed in 601a348 — replaced prepare_inputs_and_embeddings() with an override of get_multimodal_embeddings() that calls super().get_multimodal_embeddings() then scatters audio, so the audio path is now in the base call chain and input_features is no longer ignored at inference.

[lcretan]

Addressed in 601a348 — added an explicit count check before masked_scatter(): it compares <audio> token positions against audio_embeds.shape[0] and raises a clear ValueError on mismatch. Covered by TestMergeAudioTextEmbeddings::test_shape_mismatch_raises.

[lcretan]

Fixed in 601a348 — the docstring now states the actual return shape [n_valid_frames_total, hidden_size] (flattened, padding removed via input_features_mask).

[lcretan]

Fixed in 601a348 — input_features_mask now uses constant_tensor() so the dummy input honors the framework argument like the other generators.

[lcretan]

Addressed in 2fb824f — added audio_embeddings_model to the gemma4_unified INT8 expectations in tests/openvino/utils_tests.py, and c3718cb adds tests/openvino/test_gemma4_audio.py covering the new export artifact.

[lcretan]

Thanks for the review, @rkazants!

---

Tests added — tests/openvino/test_gemma4_audio.py (14 tests, all passing)

Class	Coverage
TestMergeAudioTextEmbeddings	scatter correctness, shape-mismatch guard, numpy→torch conversion, mask/position_ids passthrough
TestGetAudioEmbeddings	decode-step skip, ValueError when audio model absent, mask-based frame stripping, no-mask flatten
TestFromPretrainedBackwardCompat	local dir without audio XML skips gracefully; Hub ID passes through unchanged
TestGemma4AudioIntegration	audio model loaded, output shape, dtype, full merge roundtrip (requires local int8-audio model)

Also fixed a Hub-ID bug in from_pretrained: previously os.path.exists("google/gemma-4-12B-it/openvino_audio_embeddings_model.xml") always returned False, causing Hub models to silently strip audio. Fixed by checking os.path.isdir() first — Hub IDs are not local dirs, so they fall through to the parent's hf_hub_download optional-key try/except.

---

Inference code snippet

Export with audio support:

optimum-cli export openvino \
  --model google/gemma-4-12B-it \
  --task image-text-to-text \
  --weight-format int8 \
  ./gemma-4-12B-it-openvino-int8-audio

The exported directory will contain openvino_audio_embeddings_model.xml (~2.4 MB) alongside the existing vision and language model files.

Text-only inference (unchanged):

from optimum.intel import OVModelForVisualCausalLM
from transformers import AutoProcessor

model_dir = "./gemma-4-12B-it-openvino-int8-audio"
processor = AutoProcessor.from_pretrained(model_dir)
model = OVModelForVisualCausalLM.from_pretrained(model_dir, device="CPU")

messages = [{"role": "user", "content": "What is the capital of Japan?"}]
text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
inputs = processor(text=text, return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=64)
print(processor.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
# → "The capital of Japan is Tokyo."

Audio inference:

import numpy as np
import soundfile as sf  # or librosa / scipy

# Load 16 kHz mono audio
audio, sr = sf.read("speech.wav")  # shape: [n_samples]
assert sr == 16000, "Gemma4 audio feature extractor expects 16 kHz"

# Build a multimodal message — processor handles feature extraction internally
messages = [{"role": "user", "content": [
    {"type": "audio", "audio": audio},
    {"type": "text",  "text": "Transcribe the audio above in Japanese."},
]}]

# apply_chat_template with audio data → produces input_features + input_features_mask
inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_tensors="pt",
    return_dict=True,
)
# inputs now contains: input_ids, attention_mask, input_features [B, T, 640], input_features_mask

out = model.generate(**inputs, max_new_tokens=256)
print(processor.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Backward compatibility (models without audio):

# Loading a model exported before audio support works transparently.
# The audio sub-model is simply absent and audio inputs raise a clear ValueError.
model = OVModelForVisualCausalLM.from_pretrained(
    "./gemma-4-12B-it-openvino-int4",  # no openvino_audio_embeddings_model.xml
    device="CPU",
)
# model.audio_embeddings is None; text+vision still work normally

[rkazants]

use existing test suite in test_seq2seq and test_export. No need to implement a new one

[lcretan]

Thanks @rkazants — makes sense, I'll fold this into the existing suites rather than a separate file. Since gemma4_unified is already parametrized in test_export.py and test_seq2seq.py, I'll add the audio-specific coverage on that path (assert the openvino_audio_embeddings_model.* artifact is produced on export, plus an audio-embeddings inference + merge check) and remove tests/openvino/test_gemma4_audio.py.

One note for your call: a few cases in that file are fine-grained unit tests for the audio↔text embedding-merge helpers (scatter correctness, shape-mismatch guard, mask-based frame stripping) that don't map onto the export/inference model tests. My default is to drop the standalone units and keep equivalent coverage as assertions within the gemma4_unified test path — let me know if you'd prefer I retain any of them. Will push the updated tests shortly.

[lcretan]

Pushed 5d59cbc — removed tests/openvino/test_gemma4_audio.py. The gemma4_unified audio coverage now lives entirely in the existing suites:

• Export + artifact + quantization: tests/openvino/utils_tests.py lists audio_embeddings_model in the gemma4_unified expectations, so test_export.py / test_seq2seq.py already export the tiny model and assert the openvino_audio_embeddings_model.* artifact and its INT8 quantization.
• Component validity: test_compare_to_transformers → _check_openvino_model_attributes iterates every component of the exported model — including the audio embeddings model — and asserts each loads as a valid openvino.Model.

So the audio embedder is covered structurally without a dedicated file. The fine-grained merge-helper unit tests are dropped as discussed; a dedicated audio-inference check would need audio preprocessing the image-text-to-text suite doesn't currently do — happy to add one if you'd like it. Thanks again for the steer.