diff --git a/.gitignore b/.gitignore index 1f6b14cc..d1ddb33f 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,11 @@ wheels/ .idea/ .vscode/ venv/ -env/ \ No newline at end of file +env/ + +# example/*.py里默认的模型checkpoints路径 +Qwen/ + +# example 脚本生成的音频输出目录 +outputs/ +*.wav \ No newline at end of file diff --git a/README.md b/README.md index 8287bac4..a3d0a615 100644 --- a/README.md +++ b/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/` 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 diff --git a/examples/test_model_12hz_base.py b/examples/test_model_12hz_base.py index ec7ea4b5..0ecd2e80 100644 --- a/examples/test_model_12hz_base.py +++ b/examples/test_model_12hz_base.py @@ -19,6 +19,7 @@ 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): @@ -26,12 +27,12 @@ def ensure_dir(d: str): 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) diff --git a/examples/test_model_12hz_custom_voice.py b/examples/test_model_12hz_custom_voice.py index 560ecd13..93488fef 100644 --- a/examples/test_model_12hz_custom_voice.py +++ b/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,11 +47,11 @@ 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."] @@ -54,7 +59,7 @@ def main(): 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__": diff --git a/examples/test_model_12hz_voice_design.py b/examples/test_model_12hz_voice_design.py index 7e755391..934015a3 100644 --- a/examples/test_model_12hz_voice_design.py +++ b/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__": diff --git a/examples/test_tokenizer_12hz.py b/examples/test_tokenizer_12hz.py index 595809e3..b9beabcc 100644 --- a/examples/test_tokenizer_12hz.py +++ b/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) \ No newline at end of file +sf.write(os.path.join(OUT_DIR, "decoded_numpy_12hz.wav"), wavs3[0], out_sr3) \ No newline at end of file diff --git a/qwen_tts/cli/demo.py b/qwen_tts/cli/demo.py index e267861d..47d37e0f 100644 --- a/qwen_tts/cli/demo.py +++ b/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, ) diff --git a/qwen_tts/inference/backend_utils.py b/qwen_tts/inference/backend_utils.py new file mode 100644 index 00000000..dfca4edf --- /dev/null +++ b/qwen_tts/inference/backend_utils.py @@ -0,0 +1,89 @@ +# coding=utf-8 +# Copyright 2026 The Alibaba Qwen team. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any, Dict, Optional + +import torch + + +_MLX_ALIASES = {"mlx", "mlx:0", "apple-mlx"} +_MPS_ALIASES = {"mps", "mps:0"} +_CUDA_PREFIXES = ("cuda",) + + +def _has_mps() -> bool: + return bool(torch.backends.mps.is_built() and torch.backends.mps.is_available()) + + +def default_device() -> str: + if torch.cuda.is_available(): + return "cuda:0" + if _has_mps(): + return "mps" + return "cpu" + + +def device_supports_bfloat16(device: str) -> bool: + device = normalize_device_name(device) + if device.startswith("cuda"): + return True + if device == "cpu": + return True + return False + + +def normalize_device_name(device: Optional[str]) -> str: + raw = (device or "").strip().lower() + if not raw or raw == "auto": + return default_device() + if raw in _MLX_ALIASES: + return "mps" if _has_mps() else "cpu" + if raw in _MPS_ALIASES: + return "mps" if _has_mps() else "cpu" + if raw.startswith(_CUDA_PREFIXES): + return raw if torch.cuda.is_available() else ("mps" if _has_mps() else "cpu") + if raw == "cpu": + return "cpu" + return raw + + +def normalize_dtype_for_device(dtype: Optional[torch.dtype], device: str) -> Optional[torch.dtype]: + if dtype is None: + return None + if dtype == torch.bfloat16 and not device_supports_bfloat16(device): + return torch.float16 if device != "cpu" else torch.float32 + return dtype + + +def normalize_attn_implementation(attn_implementation: Optional[str], device: str) -> Optional[str]: + device = normalize_device_name(device) + if attn_implementation == "flash_attention_2" and not device.startswith("cuda"): + return None + return attn_implementation + + +def resolve_model_load_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: + resolved = dict(kwargs) + device = normalize_device_name(resolved.get("device_map")) + resolved["device_map"] = device + + dtype_key = "dtype" if "dtype" in resolved else "torch_dtype" if "torch_dtype" in resolved else None + if dtype_key is not None: + resolved[dtype_key] = normalize_dtype_for_device(resolved.get(dtype_key), device) + + if "attn_implementation" in resolved: + resolved["attn_implementation"] = normalize_attn_implementation(resolved.get("attn_implementation"), device) + + return resolved + + +def synchronize_device(device: Optional[str]) -> None: + name = normalize_device_name(device) + if name.startswith("cuda") and torch.cuda.is_available(): + torch.cuda.synchronize() + return + if name == "mps" and hasattr(torch, "mps") and torch.backends.mps.is_available(): + torch.mps.synchronize() diff --git a/qwen_tts/inference/qwen3_tts_model.py b/qwen_tts/inference/qwen3_tts_model.py index f4d33bf3..5ce00675 100644 --- a/qwen_tts/inference/qwen3_tts_model.py +++ b/qwen_tts/inference/qwen3_tts_model.py @@ -26,6 +26,7 @@ import torch from transformers import AutoConfig, AutoModel, AutoProcessor +from .backend_utils import resolve_model_load_kwargs from ..core.models import Qwen3TTSConfig, Qwen3TTSForConditionalGeneration, Qwen3TTSProcessor AudioLike = Union[ @@ -109,7 +110,8 @@ def from_pretrained( AutoModel.register(Qwen3TTSConfig, Qwen3TTSForConditionalGeneration) AutoProcessor.register(Qwen3TTSConfig, Qwen3TTSProcessor) - model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs) + model_load_kwargs = resolve_model_load_kwargs(kwargs) + model = AutoModel.from_pretrained(pretrained_model_name_or_path, **model_load_kwargs) if not isinstance(model, Qwen3TTSForConditionalGeneration): raise TypeError( f"AutoModel returned {type(model)}, expected Qwen3TTSForConditionalGeneration. " diff --git a/qwen_tts/inference/qwen3_tts_tokenizer.py b/qwen_tts/inference/qwen3_tts_tokenizer.py index 24f8768e..384e1529 100644 --- a/qwen_tts/inference/qwen3_tts_tokenizer.py +++ b/qwen_tts/inference/qwen3_tts_tokenizer.py @@ -26,6 +26,7 @@ from torch.nn.utils.rnn import pad_sequence from transformers import AutoConfig, AutoFeatureExtractor, AutoModel +from .backend_utils import resolve_model_load_kwargs from ..core import ( Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model, @@ -85,7 +86,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs) -> "Qwen3 AutoModel.register(Qwen3TTSTokenizerV2Config, Qwen3TTSTokenizerV2Model) inst.feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path) - inst.model = AutoModel.from_pretrained(pretrained_model_name_or_path, **kwargs) + model_load_kwargs = resolve_model_load_kwargs(kwargs) + inst.model = AutoModel.from_pretrained(pretrained_model_name_or_path, **model_load_kwargs) inst.config = inst.model.config inst.device = getattr(inst.model, "device", None)