example_inference.py

"""Minimal inference example for qwen3_voice_design_t8.

Install:
    pip install qwen-tts transformers torch soundfile

Run:
    python example_inference.py                                  # loads from local dir (./)
    python example_inference.py --repo macminix/qwen3_voice_design_t8   # or pull from HF

The model is self-contained. No base model download is required.
"""
from __future__ import annotations

import argparse
from pathlib import Path

import soundfile as sf
import torch
from qwen_tts import Qwen3TTSModel


PROMPTS = [
    dict(
        name="uk_male_natural",
        text="The train to Edinburgh departs from platform four.",
        instruct="A man with a British English accent, calm and natural.",
    ),
    dict(
        name="uk_female_warm",
        text="Tea is ready in the kitchen if you would like a cup.",
        instruct="A woman with a British English accent, conversational and unhurried.",
    ),
    dict(
        name="uk_scottish_male",
        text="It is a fair walk from here to the centre of town.",
        instruct="A man with a Scottish accent, calm and conversational.",
    ),
    dict(
        name="neutral_natural_male",
        text="The meeting starts at ten o\u2019clock in the conference room.",
        instruct="A clear, neutral voice reading the sentence.",
    ),
    dict(
        name="subtle_sad_female",
        text="I keep thinking about what I could have done differently.",
        instruct="A woman, quietly sad, in a gentle conversational tone.",
    ),
    dict(
        name="happy_male_subtle",
        text="Come and look at this, you are not going to believe it.",
        instruct="A warm male voice with a happy lift, at a natural pace.",
    ),
]

GEN_KWARGS = dict(
    language="english",
    temperature=0.9,
    top_k=50,
    top_p=1.0,
    repetition_penalty=1.05,
    max_new_tokens=600,
    do_sample=True,
)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default=".", help="HF repo id or local path (default: current dir)")
    ap.add_argument("--out-dir", default="./out", help="where to write wavs")
    ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
    ap.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu",
                    help="torch device (default: cuda:0 if available, else cpu)")
    args = ap.parse_args()

    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]

    print(f"loading model from {args.repo} (device={args.device}, dtype={args.dtype})")
    wrap = Qwen3TTSModel.from_pretrained(args.repo, device_map=args.device, dtype=dtype)

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    for p in PROMPTS:
        wavs, sr = wrap.generate_voice_design(text=p["text"], instruct=p["instruct"], **GEN_KWARGS)
        path = out_dir / f"{p['name']}.wav"
        sf.write(path, wavs[0], sr)
        print(f"  {path}  ({len(wavs[0]) / sr:.1f} s @ {sr} Hz)")

    print("done")


if __name__ == "__main__":
    main()