fix(vibevoice): preserve quantization metadata in sanitize() for quantized model loading

mlx_audio/tts/models/vibevoice/vibevoice.py

@@ -250,8 +250,9 @@ def transform_key(key: str) -> str:
 
             # Check if key exists in model
             if new_key not in curr_shapes:
-                # Debug: uncomment to see missing keys
-                # print(f"Warning: Key {new_key} (from {k}) not found in model")
+                # Preserve quantization metadata -- model isn't quantized yet at sanitize time
+                if new_key.endswith((".scales", ".biases")):
+                    new_weights[new_key] = v
                 continue
 
             target_shape = curr_shapes[new_key]

mlx_audio/tts/tests/test_models.py

@@ -1318,6 +1318,31 @@ def test_sanitize_huggingface_keys(self):
         self.assertNotIn("model.prediction_head.t_embedder.mlp.0.weight", sanitized)
         self.assertNotIn("model.prediction_head.adaLN_modulation.1.weight", sanitized)
 
+    def test_sanitize_preserves_quantization_metadata(self):
+        """Test that sanitize preserves .scales and .biases for quantized models."""
+        from mlx.utils import tree_flatten
+
+        from mlx_audio.tts.models.vibevoice.vibevoice import Model
+
+        config = self._default_config
+        model = Model(config)
+
+        # Start with the model's own weights
+        weights = dict(tree_flatten(model.parameters()))
+
+        # Add mock quantization metadata for the key from the bug report:
+        # "Expected shape (151936, 896) but received shape (151936, 224)
+        #  for parameter language_model.embed_tokens.weight"
+        quant_key = "language_model.embed_tokens.weight"
+        weights[f"{quant_key}.scales"] = mx.ones((1,))
+        weights[f"{quant_key}.biases"] = mx.ones((1,))
+
+        sanitized = model.sanitize(weights)
+
+        # Quantization metadata must survive sanitization
+        self.assertIn(f"{quant_key}.scales", sanitized)
+        self.assertIn(f"{quant_key}.biases", sanitized)
+
     def test_config_defaults(self):
         """Test VibeVoiceModel uses correct config defaults."""
         from mlx_audio.tts.models.vibevoice.config import ModelConfig