# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

import os
import torch
import soundfile as sf
from qwen_tts import Qwen3TTSModel

model_path = "/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-Base"
out_dir = "/Volumes/1TB_APL_Storage/UmbraAudio/output"
os.makedirs(out_dir, exist_ok=True)

device = os.environ.get("DEVICE", "mps")
dtype = torch.bfloat16 if device != "cpu" else torch.float32

print(f"Loading model on {device} (dtype={dtype}) ...")
model = Qwen3TTSModel.from_pretrained(
    model_path,
    device_map=device,
    dtype=dtype,
    attn_implementation="eager",
)

ref_audio = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone.wav"
ref_text = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."

print("Generating voice clone ...")
wavs, sr = model.generate_voice_clone(
    text="This is a local test of Qwen three TTS running on my Mac mini. The audio should sound natural and clear.",
    language="English",
    ref_audio=ref_audio,
    ref_text=ref_text,
)
out = os.path.join(out_dir, "test_voice_clone.wav")
sf.write(out, wavs[0], sr)
print(f"DONE: wrote {out} ({sr} Hz, {len(wavs[0])/sr:.2f}s)")