UmbraAudio/umbraaudio/voices/qwen_tts.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import numpy as np
import soundfile as sf

logger = logging.getLogger(__name__)

DEFAULT_VOICEDESIGN_MODEL = "/mnt/models/qwen3-tts/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
DEFAULT_BASE_MODEL = "/mnt/models/qwen3-tts/Qwen3-TTS-12Hz-1.7B-Base"
DEFAULT_TOKENIZER = "/mnt/models/qwen3-tts/Qwen3-TTS-Tokenizer-12Hz"

REMOTE_TIMEOUT = 900

# Rainbow Passage — phonetically balanced text used to generate all voice references.
# Passed as ref_text to enable ICL (In-Context Learning) mode for voice cloning,
# which captures full voice identity instead of just a coarse x-vector embedding.
VOICE_REFERENCE_TEXT = (
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. "
    "These take the shape of a long round arch, with its path high above, "
    "and its two ends apparently beyond the horizon. "
    "There is, according to legend, a boiling pot of gold at one end. "
    "People look, but no one ever finds it. "
    "When a man looks for something beyond his reach, "
    "his friends say he is looking for the pot of gold at the end of the rainbow. "
    "When the sunlight strikes raindrops in the air, they act as a prism and form a rainbow. "
    "The rainbow is a division of white light into many beautiful colors. "
    "These take the shape of a long round arch, with its path high above, "
    "and its two ends apparently beyond the horizon. "
    "There is, according to legend, a boiling pot of gold at one end. "
    "People look, but no one ever finds it. "
    "When a man looks for something beyond his reach, "
    "his friends say he is looking for the pot of gold at the end of the rainbow."
)


class QwenTTSUnavailableError(Exception):
    """Raised when the Qwen TTS server is unreachable or returns an error."""


class QwenTTSEngine:
    def __init__(
        self,
        device: str = "cuda",
        sample_rate: int = 24000,
        voicedesign_model: str | None = None,
        base_model: str | None = None,
        qwen_tts_url: str | None = None,
        helper_url: str | None = None,
    ) -> None:
        self.device = device
        self.sample_rate = sample_rate
        self._voicedesign_path = voicedesign_model or DEFAULT_VOICEDESIGN_MODEL
        self._base_path = base_model or DEFAULT_BASE_MODEL
        self._qwen_tts_url = qwen_tts_url
        self._helper_url = helper_url
        self._vd_model = None
        self._base_model = None

    @property
    def _candidate_urls(self) -> list[str]:
        """Ordered TTS endpoints: GPU helper first (when configured), local required last."""
        urls: list[str] = []
        if self._helper_url:
            urls.append(self._helper_url.rstrip("/"))
        if self._qwen_tts_url:
            local = self._qwen_tts_url.rstrip("/")
            if local not in urls:
                urls.append(local)
        return urls

    @property
    def _use_remote(self) -> bool:
        return bool(self._candidate_urls)

    def _url_healthy(self, url: str) -> bool:
        import httpx

        try:
            r = httpx.get(f"{url}/health", timeout=5)
            return r.status_code < 500
        except Exception:
            return False

    def _preferred_candidates(self) -> list[str]:
        """Healthy endpoints first, unhealthy ones after (never dropped).

        With a single endpoint (local-only setup) no health probe is made:
        the request itself determines success, so we never add a second
        round-trip to every TTS call.
        """
        urls = self._candidate_urls
        if len(urls) <= 1:
            return urls
        healthy = [u for u in urls if self._url_healthy(u)]
        return healthy + [u for u in urls if u not in healthy]

    def check_health(self) -> bool:
        if not self._use_remote:
            return False
        return any(self._url_healthy(url) for url in self._candidate_urls)

    def load_remote_model(self, model_type: str = "voicedesign") -> None:
        if not self._use_remote:
            return
        import httpx

        for url in self._candidate_urls:
            try:
                r = httpx.post(
                    f"{url}/load",
                    json={"model_type": model_type},
                    timeout=120,
                )
                logger.info(
                    "Qwen TTS remote model load: HTTP %d %s", r.status_code, r.text.strip()[:200]
                )
                import time

                if "already" in r.text.lower():
                    time.sleep(1)
                else:
                    time.sleep(10)
            except Exception as e:
                logger.warning("Qwen TTS remote model load failed (%s): %s", url, e)

    def unload_remote_model(self) -> None:
        if not self._use_remote:
            return
        import httpx

        for url in self._candidate_urls:
            try:
                r = httpx.post(
                    f"{url}/unload",
                    json={"model_type": "all"},
                    timeout=30,
                )
                logger.info("Qwen TTS remote model unload (%s): HTTP %d", url, r.status_code)
            except Exception as e:
                logger.warning("Qwen TTS remote model unload failed (%s): %s", url, e)

    def unload_and_reload(self) -> None:
        self.unload_remote_model()
        self.load_remote_model()

    def load_voicedesign(self) -> None:
        if self._use_remote:
            return
        if self._vd_model is not None:
            return
        self.unload_base()
        import torch
        from qwen_tts import Qwen3TTSModel

        logger.info("Loading Qwen3-TTS VoiceDesign model from %s", self._voicedesign_path)
        self._vd_model = Qwen3TTSModel.from_pretrained(
            self._voicedesign_path,
            device_map=f"{self.device}:0" if self.device == "cuda" else "cpu",
            dtype=torch.bfloat16 if self.device == "cuda" else torch.float32,
        )
        logger.info(
            "VoiceDesign model loaded. VRAM: %.2fGB",
            torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0,
        )

    def load_base(self) -> None:
        if self._use_remote:
            return
        if self._base_model is not None:
            return
        self.unload_voicedesign()
        import torch
        from qwen_tts import Qwen3TTSModel

        logger.info("Loading Qwen3-TTS Base model from %s", self._base_path)
        self._base_model = Qwen3TTSModel.from_pretrained(
            self._base_path,
            device_map=f"{self.device}:0" if self.device == "cuda" else "cpu",
            dtype=torch.bfloat16 if self.device == "cuda" else torch.float32,
        )
        logger.info(
            "Base model loaded. VRAM: %.2fGB",
            torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0,
        )

    def unload_voicedesign(self) -> None:
        if self._use_remote:
            return
        if self._vd_model is not None:
            del self._vd_model
            self._vd_model = None
            import torch

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            logger.info("VoiceDesign model unloaded")

    def unload_base(self) -> None:
        if self._use_remote:
            return
        if self._base_model is not None:
            del self._base_model
            self._base_model = None
            import torch

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            logger.info("Base model unloaded")

    def unload_all(self) -> None:
        self.unload_voicedesign()
        self.unload_base()

    def generate(
        self,
        text: str,
        output_path: Path,
        voice_ref: Path | None = None,
        instruct: str | None = None,
        ref_text: str | None = None,
        prompt_slug: str | None = None,
        temperature: float = 0.9,
        do_sample: bool = True,
        seed: int | None = None,
        **kwargs: Any,
    ) -> Path:
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if self._use_remote:
            wav, sr = self._remote_generate(
                text, voice_ref, instruct,
                ref_text=ref_text, prompt_slug=prompt_slug,
                temperature=temperature, do_sample=do_sample, seed=seed,
            )
        elif voice_ref and Path(voice_ref).exists():
            wav, sr = self._generate_clone(
                text, str(voice_ref), ref_text=ref_text,
                temperature=temperature, do_sample=do_sample, seed=seed,
            )
        elif instruct:
            wav, sr = self._generate_voicedesign(
                text, instruct,
                temperature=temperature, do_sample=do_sample, seed=seed,
            )
        else:
            raise ValueError(
                "Either instruct (voice description) or voice_ref (reference audio) required"
            )

        sf.write(str(output_path), wav, sr)
        logger.info("Generated Qwen3-TTS audio: %s", output_path)
        return output_path

    @staticmethod
    def slug_for_ref(voice_ref: Path) -> str:
        """Derive the prompt slug (cache key) for a voice reference file."""
        if voice_ref.name.endswith("_ref.wav"):
            return voice_ref.name[: -len("_ref.wav")]
        return voice_ref.stem

    def build_voice_prompt(
        self,
        voice_ref: Path,
        ref_text: str | None = None,
        slug: str | None = None,
        x_vector_only: bool = False,
        force: bool = False,
    ) -> str:
        """Capture & cache a speaker's **vmatrix** (the digital voice) once.

        The vmatrix is the set of numbers inside Qwen3-TTS that *is* the voice:
        - the ICL codec-token "8k matrix" (``ref_code`` — the re-encoded
          reference speech tokens the model conditions on), AND
        - the ``ref_spk_embedding`` (the extracted speaker x-vector).

        Uploads the reference WAV to the server's /v1/voice-prompt endpoint,
        which computes the full vmatrix, caches it as a ``.pt`` under PROMPT_DIR
        keyed by the slug, and returns the slug. Callers then pass it as
        ``prompt_slug`` to :meth:`generate`; every utterance reuses the exact
        same cached numbers, so the voice is identical across all utterances and
        chapters (this is the requirement: an approved voice = WAV to listen to
        + the captured vmatrix for consistency).

        x_vector_only=False is the DEFAULT (and intended pipeline mode): capture
        the FULL vmatrix (ref_code + embedding). x_vector_only=True captures only
        the coarse speaker embedding (ref_code=None) — historically the default,
        which is what produced the "inconsistent / changed voice" problem; keep
        it only for experiments or where the extra VRAM of ICL is unavailable.
        """
        if not self._use_remote:
            raise ValueError("build_voice_prompt only works with the remote TTS server")
        import httpx

        slug = slug or self.slug_for_ref(voice_ref)
        data = {"slug": slug, "x_vector_only": 1 if x_vector_only else 0}
        if force:
            data["force"] = 1
        if ref_text:
            data["ref_text"] = ref_text

        last_err: str | None = None
        ref_bytes = Path(voice_ref).read_bytes()
        for url in self._preferred_candidates():
            try:
                r = httpx.post(
                    f"{url}/v1/voice-prompt",
                    data=data,
                    files={"ref_audio": ("ref.wav", ref_bytes, "audio/wav")},
                    timeout=REMOTE_TIMEOUT,
                )
                if r.status_code != 200:
                    last_err = f"Voice-prompt build failed HTTP {r.status_code}: {r.text[:300]}"
                    logger.warning("%s (url=%s)", last_err, url)
                    continue
                return slug
            except Exception as e:
                last_err = f"Voice-prompt build failed: {e}"
                logger.warning("%s (url=%s)", last_err, url)
        raise QwenTTSUnavailableError(last_err or "No Qwen TTS server available")

    def generate_voice_reference(
        self,
        text: str,
        output_path: Path,
        voice_description: str,
        **kwargs: Any,
    ) -> Path:
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if self._use_remote:
            wav, sr = self._remote_generate(text, instruct=voice_description)
        else:
            wav, sr = self._generate_voicedesign(text, voice_description)

        sf.write(str(output_path), wav, sr)
        logger.info("Generated voice reference: %s", output_path)
        return output_path

    def _remote_generate(
        self,
        text: str,
        voice_ref: Path | None = None,
        instruct: str | None = None,
        ref_text: str | None = None,
        prompt_slug: str | None = None,
        temperature: float = 0.9,
        do_sample: bool = True,
        seed: int | None = None,
    ) -> tuple[np.ndarray, int]:
        import httpx

        files = None

        if prompt_slug:
            pass
        elif voice_ref and Path(voice_ref).exists():
            files = {"ref_audio": ("ref.wav", Path(voice_ref).read_bytes(), "audio/wav")}

        last_err: str | None = None
        for url in self._preferred_candidates():
            try:
                return self._remote_generate_at(
                    url=url,
                    text=text,
                    files=files,
                    instruct=instruct,
                    ref_text=ref_text,
                    prompt_slug=prompt_slug,
                    temperature=temperature,
                    do_sample=do_sample,
                    seed=seed,
                )
            except QwenTTSUnavailableError as e:
                last_err = str(e)
                logger.warning("Qwen TTS endpoint %s failed, trying next: %s", url, e)
        raise QwenTTSUnavailableError(last_err or "No Qwen TTS server available")

    def _remote_generate_at(
        self,
        *,
        url: str,
        text: str,
        files: dict[str, Any] | None,
        instruct: str | None,
        ref_text: str | None,
        prompt_slug: str | None,
        temperature: float,
        do_sample: bool,
        seed: int | None,
    ) -> tuple[np.ndarray, int]:
        import httpx

        generate_url = f"{url}/v1/tts"
        payload: dict[str, Any] = {
            "text": text,
            "sample_rate": self.sample_rate,
            "temperature": temperature,
            "do_sample": do_sample,
        }
        if seed is not None:
            payload["seed"] = seed
        if ref_text:
            payload["ref_text"] = ref_text

        if prompt_slug:
            payload["prompt_slug"] = prompt_slug

        # instruct composes with voice clone / ICL: the server prepends it as
        # style direction (emotion/pace/delivery) while keeping the speaker.
        if instruct:
            payload["voice_description"] = instruct

        try:
            r = httpx.post(
                generate_url, data=payload, files=files or {}, timeout=REMOTE_TIMEOUT
            )
        except httpx.ConnectError as e:
            raise QwenTTSUnavailableError(
                f"Cannot connect to Qwen TTS server at {url}. Make sure the server is running."
            ) from e
        except httpx.TimeoutException as e:
            raise QwenTTSUnavailableError(
                f"Qwen TTS server at {url} timed out after {REMOTE_TIMEOUT}s."
            ) from e

        if r.status_code != 200:
            detail = r.text[:500] if r.text else "no detail"
            raise QwenTTSUnavailableError(
                f"Qwen TTS server returned HTTP {r.status_code}: {detail}"
            )

        content_type = r.headers.get("content-type", "")
        logger.debug(
            "TTS response: status=%d content-type=%s body_len=%d",
            r.status_code,
            content_type,
            len(r.content),
        )
        if "json" in content_type:
            data = r.json()
            audio_b64 = data.get("audio") or data.get("audio_base64")
            if audio_b64:
                import base64

                wav_bytes = base64.b64decode(audio_b64)
            else:
                raise QwenTTSUnavailableError(
                    f"Qwen TTS server returned JSON without audio data: {json.dumps(data)[:300]}"
                )
        else:
            wav_bytes = r.content

        import io

        wav, sr = sf.read(io.BytesIO(wav_bytes), dtype="float32")
        return wav, int(sr)

    @staticmethod
    def _load_cached_embedding(voice_ref: Path) -> np.ndarray | None:
        """Load cached speaker embedding for a voice reference.

        Given a voice ref path like .umbra/voices/narrator_ref.wav,
        looks for .umbra/voice_embeddings/narrator.npy.
        """
        ref_path = Path(voice_ref)
        if ref_path.name.endswith("_ref.wav"):
            slug = ref_path.name[:-len("_ref.wav")]
            emb_dir = ref_path.parent.parent / "voice_embeddings"
            emb_path = emb_dir / f"{slug}.npy"
            if emb_path.is_file():
                try:
                    return np.load(str(emb_path)).astype(np.float32)
                except Exception:
                    logger.warning("Failed to load embedding from %s", emb_path)
        return None

    def _generate_voicedesign(
        self, text: str, instruct: str,
        temperature: float = 0.9, do_sample: bool = True, seed: int | None = None,
    ) -> tuple[np.ndarray, int]:
        if seed is not None:
            import torch
            torch.manual_seed(seed)
        self.load_voicedesign()
        wavs, sr = self._vd_model.generate_voice_design(
            text=text,
            instruct=instruct,
            language="english",
            temperature=temperature,
            do_sample=do_sample,
        )
        if not wavs:
            raise ValueError("VoiceDesign model returned no audio")
        return wavs[0], sr

    def _generate_clone(
        self, text: str, ref_audio_path: str, ref_text: str | None = None,
        temperature: float = 0.9, do_sample: bool = True, seed: int | None = None,
    ) -> tuple[np.ndarray, int]:
        if seed is not None:
            import torch
            torch.manual_seed(seed)
        self.load_base()

        import torch

        ref_audio, ref_sr = sf.read(ref_audio_path)
        if ref_sr != self.sample_rate:
            import librosa
            ref_audio = librosa.resample(ref_audio, orig_sr=ref_sr, target_sr=self.sample_rate)
        wavs, sr = self._base_model.generate_voice_clone(
            text=text,
            ref_audio=(ref_audio, self.sample_rate),
            ref_text=ref_text,
            x_vector_only_mode=False,
            language="english",
            temperature=temperature,
            do_sample=do_sample,
        )
        if not wavs:
            raise ValueError("Voice clone model returned no audio")
        return wavs[0], sr

    def extract_embedding(self, audio_path: Path) -> Any:
        return None

    def clone_voice(self, reference_audio: Path, output_path: Path) -> Path:
        return reference_audio


class QwenTTSClient:
    def __init__(
        self,
        base_url: str,
        timeout: int = REMOTE_TIMEOUT,
        helper_url: str | None = None,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.helper_url = (helper_url or "").rstrip("/") or None
        self.timeout = timeout

    def check_health(self) -> bool:
        import httpx

        for url in (self.helper_url, self.base_url):
            if not url:
                continue
            try:
                r = httpx.get(f"{url}/health", timeout=5)
                if r.status_code < 500:
                    return True
            except Exception:
                continue
        return False

    def generate_voice_reference(
        self, text: str, voice_description: str, output_path: Path
    ) -> Path:
        engine = QwenTTSEngine(qwen_tts_url=self.base_url, helper_url=self.helper_url)
        return engine.generate_voice_reference(text, output_path, voice_description)

    def generate(
        self,
        text: str,
        output_path: Path,
        voice_ref: Path | None = None,
        instruct: str | None = None,
        temperature: float = 0.9,
        do_sample: bool = True,
        seed: int | None = None,
    ) -> Path:
        engine = QwenTTSEngine(qwen_tts_url=self.base_url, helper_url=self.helper_url)
        return engine.generate(
            text, output_path, voice_ref=voice_ref, instruct=instruct,
            temperature=temperature, do_sample=do_sample, seed=seed,
        )


def build_voice_description(
    character_name: str,
    profile: dict[str, Any] | None = None,
    gender: str | None = None,
    energy: str | None = None,
    timbre: str | None = None,
    pitch_register: str | None = None,
    age: str | None = None,
    accent: str | None = None,
    extra_traits: str | None = None,
    texture: str | None = None,
    formality: str | None = None,
    emotional_baseline: str | None = None,
    personality: str | None = None,
    distinguishing_trait: str | None = None,
) -> str:
    if profile:
        gender = gender or profile.get("gender")
        energy = energy or profile.get("energy")
        timbre = timbre or profile.get("timbre")
        pitch_register = pitch_register or profile.get("pitch_register")
        age = age or profile.get("age")
        accent = accent or profile.get("accent")
        texture = texture or profile.get("texture")
        formality = formality or profile.get("formality")
        emotional_baseline = emotional_baseline or profile.get("emotional_baseline")
        personality = personality or profile.get("personality")
        distinguishing_trait = distinguishing_trait or profile.get("distinguishing_trait")
        extra_traits = extra_traits or profile.get("voice_description")

    if extra_traits:
        return extra_traits

    parts: list[str] = []

    if gender:
        gender_map = {
            "male": "A male voice",
            "female": "A female voice",
            "child": "A young child's voice",
            "robot": "A synthetic robotic voice",
        }
        parts.append(gender_map.get(gender.lower(), f"A {gender} voice"))
    else:
        parts.append("A voice")

    if age:
        age_map = {
            "child": "sounding like a young child",
            "young": "sounding young",
            "adult": "sounding adult",
            "middle-aged": "sounding middle-aged",
            "middle_aged": "sounding middle-aged",
            "elderly": "sounding elderly",
            "old": "sounding elderly",
        }
        parts.append(age_map.get(age.lower(), age))

    timbre_map = {
        "warm": "warm and rich",
        "breathy": "soft and breathy",
        "nasal": "slightly nasal",
        "gravelly": "rough and gravelly",
        "resonant": "deep and resonant",
        "thin": "thin and reedy",
        "clear": "clear and articulate",
        "metallic": "metallic and synthetic",
        "velvety": "smooth and velvety",
    }
    if timbre and timbre.lower() in timbre_map:
        parts.append(f"with a {timbre_map[timbre.lower()]} timbre")

    if pitch_register:
        pitch_map = {
            "very_low": "with a very low pitch",
            "low": "with a low pitch",
            "mid_low": "with a slightly low pitch",
            "mid": "with a medium pitch",
            "mid_high": "with a slightly high pitch",
            "high": "with a high pitch",
            "very_high": "with a very high pitch",
            "ultra_high": "with an extremely high pitch",
        }
        parts.append(pitch_map.get(pitch_register.lower(), pitch_register))

    energy_map = {
        "whispered": "barely audible, hushed whisper",
        "subdued": "speaking softly and calmly",
        "soft_spoken": "speaking gently",
        "measured": "speaking in a measured, even tone",
        "assertive": "speaking with confidence",
        "commanding": "speaking with authority and power",
        "volatile": "speaking with intense emotion",
        "explosive": "speaking with sudden, thunderous force",
    }
    if energy and energy.lower() in energy_map:
        parts.append(energy_map[energy.lower()])

    texture_map = {
        "smooth": "smooth and polished",
        "raspy": "with a raspy, textured quality",
        "clipped": "with clipped, precise articulation",
        "melodic": "with a melodic, singing cadence",
        "monotone": "in a flat, uninflected monotone",
        "staccato": "in a staccato, choppy rhythm",
        "lilting": "with a lilting, bouncy cadence",
    }
    if texture and texture.lower() in texture_map:
        parts.append(texture_map[texture.lower()])

    formality_map = {
        "casual": "casual and informal",
        "conversational": "conversational and natural",
        "formal": "formal and measured",
        "archaic": "archaic and old-fashioned",
        "military": "military and clipped",
        "academic": "academic and precise",
    }
    if formality and formality.lower() in formality_map:
        parts.append(f"speaking {formality_map[formality.lower()]}")

    if accent:
        parts.append(f"with a {accent} accent")

    if personality:
        parts.append(f" personality: {personality[:150]}")

    if distinguishing_trait:
        parts.append(f" Most distinctive: {distinguishing_trait}")

    desc = " ".join(parts) + "."
    return desc