tts_engine.py

"""
VoiceClone AI — TTS engine abstraction
โหลดเอนจินโคลนเสียงจริงตามลำดับความพร้อม:
  1) F5-TTS         (รองรับ checkpoint ภาษาไทย เช่น VIZINTZOR/F5-TTS-THAI)
  2) Coqui XTTS-v2  (OpenVoice-style zero-shot cloning ทางเลือก)

เลือกเอนจินได้ด้วย env: TTS_ENGINE = auto | f5 | xtts
ถ้าไม่มีเอนจินใดติดตั้งไว้ get_engine() จะคืน None และ /generate จะตอบ 503
ทำให้ frontend สลับไปใช้ Web Speech โดยอัตโนมัติ
"""
import os
from pathlib import Path


class EngineUnavailable(Exception):
    pass


class BaseEngine:
    name = "base"
    model_name = None

    def clone(self, text, ref_wav, ref_text, out_path,
              speed=1.0, emotion="neutral", language="th"):
        raise NotImplementedError


class F5TTSEngine(BaseEngine):
    """F5-TTS zero-shot voice cloning.
    ตั้ง checkpoint ภาษาไทยผ่าน env:
      F5_MODEL       (ชื่อโมเดลพื้นฐาน, ค่าเริ่มต้น F5TTS_v1_Base)
      F5_CKPT_FILE   (path ไปยัง .safetensors/.pt ของ F5-TTS-THAI)
      F5_VOCAB_FILE  (path ไปยัง vocab.txt ของโมเดลไทย)
    """
    name = "F5-TTS"

    def __init__(self):
        try:
            from f5_tts.api import F5TTS
        except Exception as e:  # pragma: no cover
            raise EngineUnavailable(f"f5-tts ยังไม่ได้ติดตั้ง: {e}")

        here = Path(__file__).parent
        ckpt = os.getenv("F5_CKPT_FILE")
        vocab = os.getenv("F5_VOCAB_FILE")

        # 1) ตรวจหา checkpoint ภาษาไทยที่ดาวน์โหลดไว้ใน backend/models/ (เครื่อง dev)
        if not ckpt:
            cand = here / "models" / "model_1000000.pt"
            if cand.exists():
                ckpt = str(cand)
                vcand = here / "models" / "vocab.txt"
                if not vocab and vcand.exists():
                    vocab = str(vcand)
            else:
                # 2) ไม่มีในเครื่อง → ดาวน์โหลดจาก HF Hub (เหมาะกับ deploy เช่น HF Spaces / Fly)
                repo = os.getenv("F5_HF_REPO", "VIZINTZOR/F5-TTS-THAI")
                if repo:
                    from huggingface_hub import hf_hub_download
                    ckpt = hf_hub_download(repo, os.getenv("F5_HF_CKPT", "model_1000000.pt"))
                    if not vocab:
                        vocab = hf_hub_download(repo, os.getenv("F5_HF_VOCAB", "vocab.txt"))

        # โมเดลไทยอิงสถาปัตยกรรม F5TTS_Base (ไม่ใช่ v1)
        default_arch = "F5TTS_Base" if ckpt else "F5TTS_v1_Base"
        model_name = os.getenv("F5_MODEL", default_arch)

        kwargs = {}
        if ckpt:
            kwargs["ckpt_file"] = ckpt
        if vocab:
            kwargs["vocab_file"] = vocab
        self.model = F5TTS(model=model_name, **kwargs)

        is_thai = bool(ckpt and ("THAI" in ckpt.upper() or "1000000" in ckpt))
        self.name = "F5-TTS-THAI" if is_thai else "F5-TTS"
        self.model_name = (Path(ckpt).name if ckpt else model_name)

    def clone(self, text, ref_wav, ref_text, out_path,
              speed=1.0, emotion="neutral", language="th"):
        self.model.infer(
            ref_file=str(ref_wav),
            ref_text=ref_text or "",
            gen_text=text,
            file_wave=str(out_path),
            speed=float(speed or 1.0),
            nfe_step=32,
            cfg_strength=2.0,
            remove_silence=True,
        )
        return out_path


class XTTSEngine(BaseEngine):
    """Coqui XTTS-v2 — zero-shot cloning จากตัวอย่างเสียงสั้น ๆ"""
    name = "XTTS-v2"

    def __init__(self):
        try:
            from TTS.api import TTS
            import torch
        except Exception as e:  # pragma: no cover
            raise EngineUnavailable(f"coqui TTS ยังไม่ได้ติดตั้ง: {e}")
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        self.tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(dev)
        self.model_name = "xtts_v2"

    def clone(self, text, ref_wav, ref_text, out_path,
              speed=1.0, emotion="neutral", language="th"):
        lang = "th" if language == "th" else "en"
        try:
            self.tts.tts_to_file(text=text, speaker_wav=str(ref_wav),
                                 language=lang, file_path=str(out_path),
                                 speed=float(speed or 1.0))
        except Exception:
            # บางรุ่นของ XTTS ไม่รองรับไทยอย่างเป็นทางการ — ถอยไปใช้ en
            self.tts.tts_to_file(text=text, speaker_wav=str(ref_wav),
                                 language="en", file_path=str(out_path))
        return out_path


def load_engine():
    pref = os.getenv("TTS_ENGINE", "auto").lower()
    order = []
    if pref in ("f5", "f5-tts", "auto"):
        order.append(F5TTSEngine)
    if pref in ("xtts", "coqui", "openvoice", "auto"):
        order.append(XTTSEngine)

    errors = []
    for Engine in order:
        try:
            return Engine()
        except EngineUnavailable as e:
            errors.append(str(e))
    raise EngineUnavailable("ไม่มีเอนจิน TTS พร้อมใช้งาน: " + " | ".join(errors))