import sys, json, traceback
from pathlib import Path
import numpy as np, soundfile as sf, torch
from qwen_tts import Qwen3TTSModel
B = Path("/marimo/ai_studio")
model = Qwen3TTSModel.from_pretrained(str(B/"models"/"lyra-voice-cantonese"),
                                      device_map="cuda", dtype=torch.bfloat16)
langs = model.get_supported_languages(); spks = model.get_supported_speakers()
print("LANGS:", langs, flush=True)
print("SPEAKERS:", spks, flush=True)
default_spk = (spks[0] if isinstance(spks, (list, tuple)) and spks else spks)
spec = json.loads(Path(sys.argv[1]).read_text())
for it in spec:
    try:
        spk = it.get("speaker") or default_spk
        kw = dict(text=it["text"], speaker=spk)
        if it.get("language"): kw["language"] = it["language"]
        wavs, sr = model.generate_custom_voice(**kw)
        out = B/"assets/audio"/it["name"]
        sf.write(str(out), np.asarray(wavs[0]), sr)
        print("WROTE", out.name, f"{len(np.asarray(wavs[0]))/sr:.2f}s", flush=True)
    except Exception as e:
        print("FAIL", it.get("name"), type(e).__name__, e, flush=True)
        traceback.print_exc()
print("TTS_DONE")