import time
import numpy as np
import soundfile as sf
import sounddevice as sd
import torch
from qwen_tts import Qwen3TTSModel

model = Qwen3TTSModel.from_pretrained(
    "Qwen/Qwen3-TTS-12Hz-0.6B-Base",
    device_map="cpu",
    dtype=torch.float32,
)

prompt = model.create_voice_clone_prompt(
    ref_audio="../../assets/manito.wav",
    ref_text="¡Hola! Me llamo Manito y me encanta aprender contigo. ¿Jugamos un rato? Vamos a descubrir colores, formas y números juntos.",
)

texto = "¡Hola! Me llamo Manito y me encanta aprender contigo. ¿Jugamos un rato? Vamos a descubrir colores, formas y números juntos."
t0 = time.time()
primero = None
chunks = []

for chunk, sr in model.stream_generate_voice_clone(
    text=texto,
    language="Spanish",
    voice_clone_prompt=prompt,
    emit_every_frames=4,
    decode_window_frames=80,
    overlap_samples=0,
):
    if primero is None:
        primero = time.time() - t0
    chunks.append(np.asarray(chunk, dtype=np.float32))

total = time.time() - t0
audio = np.concatenate(chunks)
dur = len(audio) / sr
sf.write("../../outputs/qwen3-tts-6b/qwen_stream.wav", audio, sr)
print(f"primer chunk {primero:.2f}s | total {total:.2f}s | audio {dur:.2f}s | RTF {total/dur:.2f} | chunks {len(chunks)}")
sd.play(audio, sr)
sd.wait()