155 lines
6.2 KiB
Python
155 lines
6.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the pinned Qwen3-TTS CustomVoice checkpoint on native Windows CPU.
|
|
|
|
This uses Qwen's PyTorch wrapper as a standalone backend feasibility check.
|
|
It does not prove Omni streaming, playback, or speech intelligibility.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import platform
|
|
import time
|
|
from pathlib import Path
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as stream:
|
|
for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--output-wav", type=Path, required=True)
|
|
parser.add_argument("--output-report", type=Path, required=True)
|
|
parser.add_argument("--text", default="Hello from the local computer.")
|
|
parser.add_argument("--speaker", default="Ryan")
|
|
parser.add_argument("--dtype", choices=("float16", "bfloat16", "float32"), default="bfloat16")
|
|
parser.add_argument("--max-new-tokens", type=int, default=64)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--warmups", type=int, default=0)
|
|
parser.add_argument("--repeats", type=int, default=1)
|
|
args = parser.parse_args()
|
|
if args.max_new_tokens <= 0 or args.threads <= 0 or args.warmups < 0 or args.repeats <= 0:
|
|
parser.error("Token/thread/repeat counts must be positive and warmups nonnegative")
|
|
checkpoint = args.model_dir / "model.safetensors"
|
|
tokenizer_weights = args.model_dir / "speech_tokenizer" / "model.safetensors"
|
|
if not checkpoint.is_file() or not tokenizer_weights.is_file():
|
|
raise FileNotFoundError("The exact talker and speech-tokenizer weights are required")
|
|
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
import transformers
|
|
from qwen_tts import Qwen3TTSModel
|
|
|
|
torch.set_num_threads(args.threads)
|
|
started = time.perf_counter()
|
|
model = Qwen3TTSModel.from_pretrained(
|
|
str(args.model_dir.resolve()),
|
|
device_map="cpu",
|
|
dtype=getattr(torch, args.dtype),
|
|
attn_implementation="sdpa",
|
|
local_files_only=True,
|
|
)
|
|
startup_s = time.perf_counter() - started
|
|
model_devices = sorted({parameter.device.type for parameter in model.model.parameters()})
|
|
if model_devices != ["cpu"]:
|
|
raise RuntimeError(f"TTS model was not wholly on CPU: {model_devices}")
|
|
if args.speaker.lower() not in model.get_supported_speakers():
|
|
raise ValueError("The requested speaker is not supported by this checkpoint")
|
|
|
|
def generate_one(*, save_wav: bool) -> dict:
|
|
torch.manual_seed(42)
|
|
started = time.perf_counter()
|
|
wavs, sample_rate = model.generate_custom_voice(
|
|
text=args.text,
|
|
language="English",
|
|
speaker=args.speaker,
|
|
max_new_tokens=args.max_new_tokens,
|
|
do_sample=True,
|
|
non_streaming_mode=True,
|
|
)
|
|
request_s = time.perf_counter() - started
|
|
if len(wavs) != 1 or sample_rate <= 0:
|
|
raise RuntimeError("Expected one waveform with a valid sample rate")
|
|
wav = np.asarray(wavs[0], dtype=np.float32).reshape(-1)
|
|
if not wav.size or not np.isfinite(wav).all():
|
|
raise RuntimeError("TTS returned an empty or non-finite waveform")
|
|
rms = float(np.sqrt(np.mean(wav.astype(np.float64) ** 2)))
|
|
if rms <= 1e-5:
|
|
raise RuntimeError("TTS waveform is effectively silent")
|
|
if save_wav:
|
|
args.output_wav.parent.mkdir(parents=True, exist_ok=True)
|
|
sf.write(args.output_wav, wav, sample_rate, subtype="PCM_16")
|
|
return {
|
|
"request_s": request_s,
|
|
"sample_rate": sample_rate,
|
|
"samples": int(wav.size),
|
|
"duration_s": float(wav.size / sample_rate),
|
|
"rms": rms,
|
|
"peak_abs": float(np.max(np.abs(wav))),
|
|
"nonzero_samples": int(np.count_nonzero(wav)),
|
|
"float32_sha256": hashlib.sha256(wav.tobytes()).hexdigest(),
|
|
}
|
|
|
|
warmups = [generate_one(save_wav=False) for _ in range(args.warmups)]
|
|
measured = [generate_one(save_wav=index == 0) for index in range(args.repeats)]
|
|
first = measured[0]
|
|
latencies = sorted(row["request_s"] for row in measured)
|
|
report = {
|
|
"scope": "one complete native Windows CPU text-to-WAV request via Qwen's standalone PyTorch wrapper; no Omni streaming or speech-quality claim",
|
|
"platform": platform.platform(),
|
|
"python": platform.python_version(),
|
|
"torch": torch.__version__,
|
|
"transformers": transformers.__version__,
|
|
"model_dir": str(args.model_dir.resolve()),
|
|
"talker_sha256": _sha256(checkpoint),
|
|
"speech_tokenizer_sha256": _sha256(tokenizer_weights),
|
|
"model_devices": model_devices,
|
|
"dtype": args.dtype,
|
|
"attention": "sdpa",
|
|
"threads": args.threads,
|
|
"text": args.text,
|
|
"speaker": args.speaker,
|
|
"language": "English",
|
|
"max_new_tokens": args.max_new_tokens,
|
|
"seed": 42,
|
|
"startup_s": startup_s,
|
|
"request_s": first["request_s"],
|
|
"sample_rate": first["sample_rate"],
|
|
"samples": first["samples"],
|
|
"duration_s": first["duration_s"],
|
|
"rms": first["rms"],
|
|
"peak_abs": first["peak_abs"],
|
|
"nonzero_samples": first["nonzero_samples"],
|
|
"wav_sha256": _sha256(args.output_wav),
|
|
"warmups": warmups,
|
|
"measured": measured,
|
|
"profile": {
|
|
"repeats": args.repeats,
|
|
"concurrency": 1,
|
|
"p50_request_s": latencies[math.ceil(0.5 * len(latencies)) - 1],
|
|
"p95_request_s": latencies[math.ceil(0.95 * len(latencies)) - 1],
|
|
"all_same_float32_audio": len({row["float32_sha256"] for row in measured}) == 1,
|
|
},
|
|
"status": "passed",
|
|
}
|
|
args.output_report.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output_report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(report, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|