fix: Gemma 4 audio — mel preprocessing, weight loading, feature extractor (#931)

* fix: Gemma 4 audio — mel preprocessing, weight loading, feature extractor

Fix four bugs preventing Gemma 4 audio from working:

1. Missing semicausal left-padding in audio feature extractor.
   The HF reference prepends frame_length//2 (160) zero samples before
   the unfold, centering the first frame at t=0. Without this, the mel
   spectrogram is misaligned and the frame count is wrong, which also
   causes the broadcast shapes error (issue #923).

2. Wrong Hann window formula. Used cos(2*pi*(n+0.5)/N) instead of the
   correct periodic Hann cos(2*pi*n/N). The +0.5 phase shift produces
   meaningfully different spectral values from what the model was
   trained on.

3. sanitize() double-nests language_model weights (issue #912).
   HF keys like model.language_model.model.embed_tokens.weight become
   language_model.model.embed_tokens.weight after stripping model.,
   which already matches the MLX path. The unconditional insertion of
   .model. created language_model.model.model.*, so all LM weights
   loaded as zero.

4. Feature extractor not instantiated (issue #903). Only created when
   processor_config.json contains a "feature_extractor" key, which
   standard HF checkpoints don't include. Now instantiates with USM
   defaults unconditionally.

Fixes #903, #912, #923

* format

* format

* Update audio feature extractor in Gemma4 model to match hf

---------

Co-authored-by: Stephen Cox <stephencoxmail@gmail.com>
Co-authored-by: Prince Canuma <prince.gdt@gmail.com>

Author: 3 people

mlx_vlm/models/gemma4/audio_feature_extractor.py

@@ -126,7 +126,7 @@ def __init__(
         max_frequency: float = 8000.0,
         preemphasis: float = 0.0,
         preemphasis_htk_flavor: bool = True,
-        fft_overdrive: bool = True,
+        fft_overdrive: bool = False,
         dither: float = 0.0,
         input_scale_factor: float = 1.0,
         mel_floor: float = 1e-3,
@@ -153,12 +153,14 @@ def __init__(
             fft_length *= 2
         self.fft_length = fft_length
 
-        # Hanning window (non-zero at endpoints)
-        arg = math.pi * 2.0 / self.frame_length
-        window = 0.5 - (
-            0.5 * np.cos(arg * (np.arange(self.frame_length, dtype=np.float32) + 0.5))
+        # Periodic Hann window: w[n] = 0.5 - 0.5 * cos(2*pi*n / frame_length)
+        # Matches HuggingFace Transformers (signal.hann_window with periodic=True)
+        self.window = 0.5 - 0.5 * np.cos(
+            2.0
+            * np.pi
+            * np.arange(self.frame_length, dtype=np.float32)
+            / self.frame_length
         )
-        self.window = window.astype(np.float32)
 
         # Mel filter bank
         try:
@@ -209,6 +211,14 @@ def _extract_spectrogram(
         if self.input_scale_factor != 1.0:
             waveform = waveform * self.input_scale_factor
 
+        # Semicausal left-padding: prepend frame_length // 2 zeros so that
+        # the first frame is centered at t=0, matching HuggingFace Transformers
+        pad_left = self.frame_length // 2
+        waveform = np.pad(waveform, ((0, 0), (pad_left, 0)), mode="constant")
+        attention_mask = np.pad(
+            attention_mask, (pad_left, 0), mode="constant", constant_values=0
+        )
+
         frame_size_for_unfold = self.frame_length + 1
 
         frames_to_process = _unfold(
@@ -239,7 +249,7 @@ def _extract_spectrogram(
 
         magnitude_spec = np.abs(stft)
         mel_spec = np.matmul(magnitude_spec, self.mel_filters)
-        log_mel_spec = np.log(np.maximum(mel_spec, self.mel_floor))
+        log_mel_spec = np.log(mel_spec + self.mel_floor)
 
         if self.per_bin_mean is not None:
             log_mel_spec = log_mel_spec - self.per_bin_mean
@@ -248,8 +258,13 @@ def _extract_spectrogram(
             log_mel_spec = log_mel_spec / self.per_bin_stddev
 
         mel_spectrogram = log_mel_spec.squeeze(0)
-        mask = attention_mask[:: self.hop_length].astype(bool)
-        return mel_spectrogram, mask[: mel_spectrogram.shape[0]]
+        num_mel_frames = mel_spectrogram.shape[0]
+
+        frame_end_indices = (
+            np.arange(num_mel_frames) * self.hop_length + frame_size_for_unfold - 1
+        )
+        mask = attention_mask[frame_end_indices].astype(bool)
+        return mel_spectrogram, mask
 
     def _pad_waveforms(self, waveforms, max_length=None, pad_to_multiple_of=None):
         """Pad a list of waveforms to equal length."""
@@ -341,6 +356,12 @@ def __call__(
             prepared_speech.append(spec.astype(np.float32))
             prepared_speech_mask.append(spec_mask)
 
+        # Zero out padded spectrogram positions, matching HuggingFace Transformers
+        prepared_speech = [
+            spec * m[..., None]
+            for spec, m in zip(prepared_speech, prepared_speech_mask)
+        ]
+
         return {
             "input_features": prepared_speech,
             "input_features_mask": prepared_speech_mask,

mlx_vlm/models/gemma4/gemma4.py

@@ -195,7 +195,9 @@ def sanitize(self, weights):
             else:
                 new_key = k
 
-            if new_key.startswith("language_model."):
+            if new_key.startswith("language_model.") and not new_key.startswith(
+                "language_model.model."
+            ):
                 rest = new_key[len("language_model.") :]
                 new_key = "language_model.model." + rest
 

mlx_vlm/models/gemma4/processing_gemma4.py

@@ -548,20 +548,22 @@ def _load_json(path):
 
         image_processor = Gemma4ImageProcessor(**ip_config)
 
-        # Load audio feature extractor if config available
+        # Load audio feature extractor.
+        # The standard HF checkpoint does not include a "feature_extractor" key
+        # in processor_config.json, so we instantiate with defaults when the
+        # config is missing — the USM parameters are fixed for all Gemma 4 models.
         feature_extractor = None
-        if fe_config:
-            try:
-                from .audio_feature_extractor import Gemma4AudioFeatureExtractor
+        try:
+            from .audio_feature_extractor import Gemma4AudioFeatureExtractor
 
-                feature_extractor = Gemma4AudioFeatureExtractor(**fe_config)
-            except ImportError:
-                try:
-                    from transformers import Gemma4AudioFeatureExtractor
+            feature_extractor = Gemma4AudioFeatureExtractor(**(fe_config or {}))
+        except ImportError:
+            try:
+                from transformers import Gemma4AudioFeatureExtractor
 
-                    feature_extractor = Gemma4AudioFeatureExtractor(**fe_config)
-                except (ImportError, Exception):
-                    pass
+                feature_extractor = Gemma4AudioFeatureExtractor(**(fe_config or {}))
+            except (ImportError, Exception):
+                pass
 
         image_seq_length = ip_config.get("max_soft_tokens", 280)
         audio_seq_length = proc_config.get("audio_seq_length", 750)