Add MLX / Apple Silicon (MPS) backend support

.gitignore

@@ -21,4 +21,11 @@ wheels/
 .idea/
 .vscode/
 venv/
-env/
+env/
+
+# example/*.py里默认的模型checkpoints路径
+Qwen/
+
+# example 脚本生成的音频输出目录
+outputs/
+*.wav

README.md

@@ -83,23 +83,25 @@ During model loading in the qwen-tts package or vLLM, model weights will be auto
 ```bash
 # Download through ModelScope (recommended for users in Mainland China)
 pip install -U modelscope
-modelscope download --model Qwen/Qwen3-TTS-Tokenizer-12Hz  --local_dir ./Qwen3-TTS-Tokenizer-12Hz 
-modelscope download --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --local_dir ./Qwen3-TTS-12Hz-1.7B-CustomVoice
-modelscope download --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign --local_dir ./Qwen3-TTS-12Hz-1.7B-VoiceDesign
-modelscope download --model Qwen/Qwen3-TTS-12Hz-1.7B-Base --local_dir ./Qwen3-TTS-12Hz-1.7B-Base
-modelscope download --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --local_dir ./Qwen3-TTS-12Hz-0.6B-CustomVoice
-modelscope download --model Qwen/Qwen3-TTS-12Hz-0.6B-Base --local_dir ./Qwen3-TTS-12Hz-0.6B-Base
+modelscope download --model Qwen/Qwen3-TTS-Tokenizer-12Hz  --local_dir ./Qwen/Qwen3-TTS-Tokenizer-12Hz 
+modelscope download --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --local_dir ./Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice
+modelscope download --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign --local_dir ./Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign
+modelscope download --model Qwen/Qwen3-TTS-12Hz-1.7B-Base --local_dir ./Qwen/Qwen3-TTS-12Hz-1.7B-Base
+modelscope download --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --local_dir ./Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+modelscope download --model Qwen/Qwen3-TTS-12Hz-0.6B-Base --local_dir ./Qwen/Qwen3-TTS-12Hz-0.6B-Base
 
 # Download through Hugging Face
 pip install -U "huggingface_hub[cli]"
-huggingface-cli download Qwen/Qwen3-TTS-Tokenizer-12Hz --local-dir ./Qwen3-TTS-Tokenizer-12Hz
-huggingface-cli download Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --local-dir ./Qwen3-TTS-12Hz-1.7B-CustomVoice
-huggingface-cli download Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign --local-dir ./Qwen3-TTS-12Hz-1.7B-VoiceDesign
-huggingface-cli download Qwen/Qwen3-TTS-12Hz-1.7B-Base --local-dir ./Qwen3-TTS-12Hz-1.7B-Base
-huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --local-dir ./Qwen3-TTS-12Hz-0.6B-CustomVoice
-huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-Base --local-dir ./Qwen3-TTS-12Hz-0.6B-Base
+huggingface-cli download Qwen/Qwen3-TTS-Tokenizer-12Hz --local-dir ./Qwen/Qwen3-TTS-Tokenizer-12Hz
+huggingface-cli download Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --local-dir ./Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice
+huggingface-cli download Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign --local-dir ./Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign
+huggingface-cli download Qwen/Qwen3-TTS-12Hz-1.7B-Base --local-dir ./Qwen/Qwen3-TTS-12Hz-1.7B-Base
+huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --local-dir ./Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-Base --local-dir ./Qwen/Qwen3-TTS-12Hz-0.6B-Base
 ```
 
+The `--local_dir` above keeps the `Qwen/` prefix so the downloaded layout (e.g. `./Qwen/Qwen3-TTS-12Hz-1.7B-Base`) matches the default `Qwen/<model-name>` paths used in the code snippets below and in `examples/*.py` / `finetuning/*.py`. When you run from the repo root, the scripts pick up the local weights directly with no path changes.
+
 
 ## Quickstart
 

examples/test_model_12hz_base.py

@@ -19,19 +19,20 @@
 import soundfile as sf
 
 from qwen_tts import Qwen3TTSModel
+from qwen_tts.inference.backend_utils import default_device, normalize_attn_implementation, normalize_dtype_for_device, synchronize_device
 
 
 def ensure_dir(d: str):
     os.makedirs(d, exist_ok=True)
 
 
 def run_case(tts: Qwen3TTSModel, out_dir: str, case_name: str, call_fn):
-    torch.cuda.synchronize()
+    synchronize_device(str(tts.device))
     t0 = time.time()
 
     wavs, sr = call_fn()
 
-    torch.cuda.synchronize()
+    synchronize_device(str(tts.device))
     t1 = time.time()
     print(f"[{case_name}] time: {t1 - t0:.3f}s, n_wavs={len(wavs)}, sr={sr}")
 
@@ -40,16 +41,17 @@ def run_case(tts: Qwen3TTSModel, out_dir: str, case_name: str, call_fn):
 
 
 def main():
-    device = "cuda:0"
+    device = default_device()
+    dtype = normalize_dtype_for_device(torch.bfloat16, device)
     MODEL_PATH = "Qwen/Qwen3-TTS-12Hz-1.7B-Base/"
-    OUT_DIR = "qwen3_tts_test_voice_clone_output_wav"
+    OUT_DIR = os.path.join("outputs", os.path.splitext(os.path.basename(__file__))[0])
     ensure_dir(OUT_DIR)
 
     tts = Qwen3TTSModel.from_pretrained(
         MODEL_PATH,
         device_map=device,
-        dtype=torch.bfloat16,
-        attn_implementation="flash_attention_2",
+        dtype=dtype,
+        attn_implementation=normalize_attn_implementation("flash_attention_2", device),
     )
 
     # Reference audio(s)

examples/test_model_12hz_custom_voice.py

@@ -13,26 +13,31 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 import time
 import torch
 import soundfile as sf
 
 from qwen_tts import Qwen3TTSModel
+from qwen_tts.inference.backend_utils import default_device, normalize_attn_implementation, normalize_dtype_for_device, synchronize_device
 
 
 def main():
-    device = "cuda:0"
+    device = default_device()
+    dtype = normalize_dtype_for_device(torch.bfloat16, device)
     MODEL_PATH = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice/"
+    OUT_DIR = os.path.join("outputs", os.path.splitext(os.path.basename(__file__))[0])
+    os.makedirs(OUT_DIR, exist_ok=True)
 
     tts = Qwen3TTSModel.from_pretrained(
         MODEL_PATH,
         device_map=device,
-        dtype=torch.bfloat16,
-        attn_implementation="flash_attention_2",
+        dtype=dtype,
+        attn_implementation=normalize_attn_implementation("flash_attention_2", device),
     )
 
     # -------- Single (with instruct) --------
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t0 = time.time()
 
     wavs, sr = tts.generate_custom_voice(
@@ -42,19 +47,19 @@ def main():
         instruct="用特别愤怒的语气说",
     )
 
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t1 = time.time()
     print(f"[CustomVoice Single] time: {t1 - t0:.3f}s")
 
-    sf.write("qwen3_tts_test_custom_single.wav", wavs[0], sr)
+    sf.write(os.path.join(OUT_DIR, "single.wav"), wavs[0], sr)
 
     # -------- Batch (some empty instruct) --------
     texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
     languages = ["Chinese", "English"]
     speakers = ["Vivian", "Ryan"]
     instructs = ["", "Very happy."]
 
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t0 = time.time()
 
     wavs, sr = tts.generate_custom_voice(
@@ -65,12 +70,12 @@ def main():
         max_new_tokens=2048,
     )
 
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t1 = time.time()
     print(f"[CustomVoice Batch] time: {t1 - t0:.3f}s")
 
     for i, w in enumerate(wavs):
-        sf.write(f"qwen3_tts_test_custom_batch_{i}.wav", w, sr)
+        sf.write(os.path.join(OUT_DIR, f"batch_{i}.wav"), w, sr)
 
 
 if __name__ == "__main__":

examples/test_model_12hz_voice_design.py

@@ -13,26 +13,31 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 import time
 import torch
 import soundfile as sf
 
 from qwen_tts import Qwen3TTSModel
+from qwen_tts.inference.backend_utils import default_device, normalize_attn_implementation, normalize_dtype_for_device, synchronize_device
 
 
 def main():
-    device = "cuda:0"
+    device = default_device()
+    dtype = normalize_dtype_for_device(torch.bfloat16, device)
     MODEL_PATH = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign/"
+    OUT_DIR = os.path.join("outputs", os.path.splitext(os.path.basename(__file__))[0])
+    os.makedirs(OUT_DIR, exist_ok=True)
 
     tts = Qwen3TTSModel.from_pretrained(
         MODEL_PATH,
         device_map=device,
-        dtype=torch.bfloat16,
-        attn_implementation="flash_attention_2",
+        dtype=dtype,
+        attn_implementation=normalize_attn_implementation("flash_attention_2", device),
     )
 
     # -------- Single --------
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t0 = time.time()
 
     wavs, sr = tts.generate_voice_design(
@@ -41,11 +46,11 @@ def main():
         instruct="体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
     )
 
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t1 = time.time()
     print(f"[VoiceDesign Single] time: {t1 - t0:.3f}s")
 
-    sf.write("qwen3_tts_test_voice_design_single.wav", wavs[0], sr)
+    sf.write(os.path.join(OUT_DIR, "single.wav"), wavs[0], sr)
 
     # -------- Batch --------
     texts = [
@@ -58,7 +63,7 @@ def main():
         "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
     ]
 
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t0 = time.time()
 
     wavs, sr = tts.generate_voice_design(
@@ -68,12 +73,12 @@ def main():
         max_new_tokens=2048,
     )
 
-    torch.cuda.synchronize()
+    synchronize_device(device)
     t1 = time.time()
     print(f"[VoiceDesign Batch] time: {t1 - t0:.3f}s")
 
     for i, w in enumerate(wavs):
-        sf.write(f"qwen3_tts_test_voice_design_batch_{i}.wav", w, sr)
+        sf.write(os.path.join(OUT_DIR, f"batch_{i}.wav"), w, sr)
 
 
 if __name__ == "__main__":

examples/test_tokenizer_12hz.py

@@ -14,52 +14,57 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import io
+import os
 import requests
 import soundfile as sf
 
 from qwen_tts import Qwen3TTSTokenizer
+from qwen_tts.inference.backend_utils import default_device
 
 audio_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/tokenizer_demo_1.wav"
 audio_2 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/tokenizer_demo_2.wav"
 
+OUT_DIR = os.path.join("outputs", os.path.splitext(os.path.basename(__file__))[0])
+os.makedirs(OUT_DIR, exist_ok=True)
+
 # -------- Single input: wav path --------
 tokenizer_12hz = Qwen3TTSTokenizer.from_pretrained(
     "Qwen/Qwen3-TTS-Tokenizer-12Hz",
-    device_map="cuda:0",
+    device_map=default_device(),
 )
 
 enc1 = tokenizer_12hz.encode(audio_1)
 wavs1, out_sr1 = tokenizer_12hz.decode(enc1)
-sf.write("decoded_single_12hz.wav", wavs1[0], out_sr1)
+sf.write(os.path.join(OUT_DIR, "decoded_single_12hz.wav"), wavs1[0], out_sr1)
 
 # -------- Batch input: wav path list --------
 enc2 = tokenizer_12hz.encode([audio_1, audio_2])
 wavs2, out_sr2 = tokenizer_12hz.decode(enc2)
 for i, w in enumerate(wavs2):
-    sf.write(f"decoded_batch_12hz_{i}.wav", w, out_sr2)
+    sf.write(os.path.join(OUT_DIR, f"decoded_batch_12hz_{i}.wav"), w, out_sr2)
 
 # -------- Decode input as dict (12hz) --------
 # Take the first sample codes and pass as a dict.
 dict_input_12hz = {"audio_codes": enc2.audio_codes[0]}  # torch.Tensor
 wavs_d1, out_sr_d1 = tokenizer_12hz.decode(dict_input_12hz)
-sf.write("decoded_dict_12hz.wav", wavs_d1[0], out_sr_d1)
+sf.write(os.path.join(OUT_DIR, "decoded_dict_12hz.wav"), wavs_d1[0], out_sr_d1)
 
 # -------- Decode input as list[dict] (12hz) --------
 list_dict_input_12hz = [{"audio_codes": c} for c in enc2.audio_codes]  # list of torch.Tensor
 wavs_d2, out_sr_d2 = tokenizer_12hz.decode(list_dict_input_12hz)
 for i, w in enumerate(wavs_d2):
-    sf.write(f"decoded_listdict_12hz_{i}.wav", w, out_sr_d2)
+    sf.write(os.path.join(OUT_DIR, f"decoded_listdict_12hz_{i}.wav"), w, out_sr_d2)
 
 # -------- Decode input as list[dict] with numpy (12hz) --------
 # Convert codes to numpy to simulate "serialized" payload.
 list_dict_numpy_12hz = [{"audio_codes": c.cpu().numpy()} for c in enc2.audio_codes]
 wavs_d3, out_sr_d3 = tokenizer_12hz.decode(list_dict_numpy_12hz)
 for i, w in enumerate(wavs_d3):
-    sf.write(f"decoded_listdict_numpy_12hz_{i}.wav", w, out_sr_d3)
+    sf.write(os.path.join(OUT_DIR, f"decoded_listdict_numpy_12hz_{i}.wav"), w, out_sr_d3)
 
 # -------- Numpy input (must pass sr) --------
 data = requests.get(audio_2, timeout=30).content
 y, sr = sf.read(io.BytesIO(data))
 enc3 = tokenizer_12hz.encode(y, sr=sr)
 wavs3, out_sr3 = tokenizer_12hz.decode(enc3)
-sf.write("decoded_numpy_12hz.wav", wavs3[0], out_sr3)
+sf.write(os.path.join(OUT_DIR, "decoded_numpy_12hz.wav"), wavs3[0], out_sr3)

qwen_tts/cli/demo.py

@@ -28,6 +28,7 @@
 import torch
 
 from .. import Qwen3TTSModel, VoiceClonePromptItem
+from ..inference.backend_utils import default_device, normalize_attn_implementation, normalize_device_name, normalize_dtype_for_device
 
 
 def _title_case_display(s: str) -> str:
@@ -91,8 +92,8 @@ def build_parser() -> argparse.ArgumentParser:
     # Model loading / from_pretrained args
     parser.add_argument(
         "--device",
-        default="cuda:0",
-        help="Device for device_map, e.g. cpu, cuda, cuda:0 (default: cuda:0).",
+        default=default_device(),
+        help="Device for device_map, e.g. cpu, mps, mlx, cuda, cuda:0 (default: auto-detect).",
     )
     parser.add_argument(
         "--dtype",
@@ -602,12 +603,12 @@ def main(argv=None) -> int:
 
     ckpt = _resolve_checkpoint(args)
 
-    dtype = _dtype_from_str(args.dtype)
-    attn_impl = "flash_attention_2" if args.flash_attn else None
+    dtype = normalize_dtype_for_device(_dtype_from_str(args.dtype), normalize_device_name(args.device))
+    attn_impl = normalize_attn_implementation("flash_attention_2" if args.flash_attn else None, args.device)
 
     tts = Qwen3TTSModel.from_pretrained(
         ckpt,
-        device_map=args.device,
+        device_map=normalize_device_name(args.device),
         dtype=dtype,
         attn_implementation=attn_impl,
     )