fix: improve code formatting across multiple files

Author: Li-dongyang

mlx_audio/codec/models/moss_audio_tokenizer/__init__.py

@@ -5,7 +5,12 @@
 Model = MossAudioTokenizer
 
 DETECTION_HINTS = {
-    "config_keys": {"quantizer_kwargs", "decoder_kwargs", "downsample_rate", "code_dim"},
+    "config_keys": {
+        "quantizer_kwargs",
+        "decoder_kwargs",
+        "downsample_rate",
+        "code_dim",
+    },
     "architectures": {"MossAudioTokenizerModel"},
     "path_patterns": {"moss_audio_tokenizer", "mossaudiotokenizer", "speech_tokenizer"},
 }

mlx_audio/server.py

@@ -24,7 +24,6 @@
 import numpy as np
 import uvicorn
 import webrtcvad
-from mlx_audio.audio_io import sf_read, sf_write
 from fastapi import (
     FastAPI,
     File,
@@ -41,10 +40,10 @@
 from pydantic import BaseModel
 
 from mlx_audio.audio_io import read as audio_read
+from mlx_audio.audio_io import sf_read, sf_write
 from mlx_audio.audio_io import write as audio_write
 from mlx_audio.utils import load_model
 
-
 SpeechConversation = List[Dict[str, Any]]
 
 

mlx_audio/tests/test_server.py

@@ -429,12 +429,12 @@ def test_realtime_ws_streaming_disabled_fallback(client, mock_model_provider):
     # Should have legacy-format messages (no 'type' field), not delta/complete
     deltas = [m for m in messages if m.get("type") == "delta"]
     completes = [m for m in messages if m.get("type") == "complete"]
-    assert len(deltas) == 0, (
-        f"Expected no delta messages when streaming disabled, got: {deltas}"
-    )
-    assert len(completes) == 0, (
-        f"Expected no complete messages when streaming disabled, got: {completes}"
-    )
+    assert (
+        len(deltas) == 0
+    ), f"Expected no delta messages when streaming disabled, got: {deltas}"
+    assert (
+        len(completes) == 0
+    ), f"Expected no complete messages when streaming disabled, got: {completes}"
 
     # Should have at least one legacy text message
     text_msgs = [m for m in messages if "text" in m and "type" not in m]

mlx_audio/tts/models/moss_tts/__init__.py

@@ -1,4 +1,4 @@
-from .moss_tts import Model
 from .config import ModelConfig
+from .moss_tts import Model
 
 __all__ = ["Model", "ModelConfig"]

mlx_audio/tts/models/moss_tts/processor.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
-# pyright: reportMissingImports=false
-
+import re
 from dataclasses import dataclass
 from functools import lru_cache
-import re
 from typing import Any, Iterable, List, Optional, Sequence, Tuple
 
 import mlx.core as mx
 import numpy as np
 
+# pyright: reportMissingImports=false
+
 
 AUDIO_PLACEHOLDER = "<|audio|>"
 
@@ -462,7 +462,11 @@ def _normalize_message(message: dict) -> dict:
         ref = message.get("reference")
         if isinstance(ref, list):
             audio_codes_list = [r for r in ref if isinstance(r, mx.array)]
-        return {"role": "user", "content": content, "audio_codes_list": audio_codes_list}
+        return {
+            "role": "user",
+            "content": content,
+            "audio_codes_list": audio_codes_list,
+        }
     if role == "assistant":
         return build_assistant_message(
             audio_codes_list=message.get("audio_codes_list", []),
@@ -485,12 +489,16 @@ def prepare_conversation_input(
 
     if mode == "generation":
         if len(conversation) % 2 == 0:
-            raise ValueError("generation mode requires odd number of messages (last must be user)")
+            raise ValueError(
+                "generation mode requires odd number of messages (last must be user)"
+            )
         if conversation[-1]["role"] != "user":
             raise ValueError("generation mode requires last message to be user")
     else:
         if len(conversation) % 2 != 0:
-            raise ValueError("continuation mode requires even number of messages (last must be assistant)")
+            raise ValueError(
+                "continuation mode requires even number of messages (last must be assistant)"
+            )
         if conversation[-1]["role"] != "assistant":
             raise ValueError("continuation mode requires last message to be assistant")
 
@@ -541,7 +549,9 @@ def build_user_message(self, **kwargs: Any) -> str:
     def build_assistant_message(
         self, audio_codes_list: List[mx.array], content: str = AUDIO_PLACEHOLDER
     ) -> dict:
-        return build_assistant_message(audio_codes_list=audio_codes_list, content=content)
+        return build_assistant_message(
+            audio_codes_list=audio_codes_list, content=content
+        )
 
     def prepare_conversation_input(
         self,

mlx_audio/tts/tests/test_moss_tts.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
-# pyright: reportMissingImports=false
-
-import os
 import inspect
+import os
 import unittest
 from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, cast
@@ -15,33 +13,41 @@
 
 from mlx_audio.codec.models.moss_audio_tokenizer import MossAudioTokenizer
 from mlx_audio.tts.models.moss_tts.config import ModelConfig
-from mlx_audio.utils import apply_quantization
 from mlx_audio.tts.models.moss_tts.moss_tts import (
     Model,
     _apply_experimental_quantization,
     _apply_mixed_precision_rescue,
+    _encode_reference_audio_with_codec,
     _get_experimental_quant_mode,
     _get_experimental_quant_patterns,
-    _get_mixed_precision_rescue_patterns,
-    _requantize_module_with_mode,
-    _encode_reference_audio_with_codec,
     _get_generated_audio_history,
+    _get_mixed_precision_rescue_patterns,
     _normalize_reference_audio_for_codec,
     _path_matches_rescue_pattern,
+    _requantize_module_with_mode,
     _suppress_token_ids,
     find_last_equal,
 )
-from mlx_audio.tts.models.moss_tts.qwen3 import Qwen3Attention
 from mlx_audio.tts.models.moss_tts.processor import (
     AUDIO_PLACEHOLDER,
-    MossTTSProcessor,
     AssistantMessage,
+    MossTTSProcessor,
+)
+from mlx_audio.tts.models.moss_tts.processor import (
     apply_de_delay_pattern as processor_apply_de_delay_pattern,
+)
+from mlx_audio.tts.models.moss_tts.processor import (
     apply_delay_pattern as processor_apply_delay_pattern,
+)
+from mlx_audio.tts.models.moss_tts.processor import (
     build_user_message,
     parse_output,
     prepare_generation_input,
 )
+from mlx_audio.tts.models.moss_tts.qwen3 import Qwen3Attention
+from mlx_audio.utils import apply_quantization
+
+# pyright: reportMissingImports=false
 
 
 class TestConfig(unittest.TestCase):
@@ -931,20 +937,20 @@ def test_generate_smoke_if_local_weights_exist(self):
         if not os.path.isdir(model_dir):
             self.skipTest(f"model dir not found: {model_dir} (set MOSS_TTS_MODEL_DIR)")
 
-        from mlx_audio.tts.utils import load_model
-
         from pathlib import Path
 
+        from mlx_audio.tts.utils import load_model
+
         model = load_model(Path(model_dir))
         assert model.generate is not None
         results = list(model.generate("Hello from MOSS", max_tokens=64, verbose=False))
         self.assertGreaterEqual(len(results), 1)
 
     def test_generate_ref_audio_runs_with_fixture_q8(self):
-        from mlx_audio.tts.utils import load_model
-
         from pathlib import Path
 
+        from mlx_audio.tts.utils import load_model
+
         model_dir = os.environ.get("MOSS_TTS_MODEL_DIR", "./moss-tts-8bit")
         if not os.path.isdir(model_dir):
             self.skipTest(f"model dir not found: {model_dir} (set MOSS_TTS_MODEL_DIR)")