"""qwen-tts (Qwen3-TTS 1.7B, fine-tune siriusfreak/qwen3-tts-12hz-1.7b-ru-stress-cf) on a local GPU.

ICL voice clone: the reference clip and its exact text are encoded once; texts are spoken in batches
of tts.batch with torch.manual_seed(seed + batch start), so a seed is repeatable for the same segments
in the same order. fp16 on Volta/Turing, bfloat16 on Ampere and newer (tts.dtype).
"""
import soundfile as sf
import torch


class LocalTTS:
    def __init__(self, cfg, ref_wav, ref_text):
        from qwen_tts import Qwen3TTSModel
        t = cfg["tts"]
        self.model = Qwen3TTSModel.from_pretrained(t.get("model", "siriusfreak/qwen3-tts-12hz-1.7b-ru-stress-cf"),
                                                   device_map=t.get("device", "cuda:0"),
                                                   dtype=getattr(torch, t.get("dtype", "bfloat16")))
        self.batch = int(t.get("batch", 8))
        self.set_voice(ref_wav, ref_text)

    def set_voice(self, ref_wav, ref_text):
        self.prompt = self.model.create_voice_clone_prompt(ref_audio=str(ref_wav), ref_text=ref_text)

    def speak(self, items, seed, log=print):
        items = sorted(items, key=lambda x: len(x[0]))  # similar lengths in one batch: less padding
        for i in range(0, len(items), self.batch):
            b = items[i:i + self.batch]
            torch.manual_seed(seed + i)
            wavs, sr = self.model.generate_voice_clone(text=[t for t, _ in b], language=["Russian"] * len(b),
                                                       voice_clone_prompt=self.prompt * len(b))
            for (_, out), w in zip(b, wavs):
                sf.write(out, w, sr)
            log(f"seed {seed}: {min(i + self.batch, len(items))}/{len(items)}")