path: root/app/converter/clients/qwen.py

"""Client for the qwen-tts Gradio demo servers (CustomVoice / Base / VoiceDesign)."""

import io
import logging
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

from .. import config
from ..audio import concat_audio_files
from ..chunking import split_into_chunks
from .base import (BaseTTSClient, ConversionCancelled, resolve_request_seed,
                   VOICE_MODE_CLONE, VOICE_MODE_CUSTOM, VOICE_MODE_DESIGN,
                   VOICE_MODES)
from .languages import normalize_language
from .speakers import QWEN3_TTS_SPEAKERS, speaker_display_name_for

logger = logging.getLogger(__name__)

# Fixed model facts: the demos run the 1.7B model (each takes its full
# HuggingFace id), and the 12Hz codec outputs 24 kHz audio.
MODEL_SIZE = "1.7B"
CUSTOM_VOICE_MODEL_ID = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"


class QwenTTSClient(BaseTTSClient):
    """Generates audio chunks through a Qwen3-TTS demo server."""

    def __init__(self, chunks_dir: Path,
                 voice_mode: str = "custom_voice", voice_clone_ref_audio: Optional[str] = None,
                 voice_clone_ref_text: Optional[str] = None, skip_transcription: bool = False,
                 language: Optional[str] = None, api_url: Optional[str] = None,
                 instructions: Optional[str] = None, quiet: bool = False,
                 voice: Optional[str] = None):
        super().__init__(chunks_dir, quiet=quiet)
        if voice_mode not in VOICE_MODES:
            raise ValueError(
                f"Unknown voice mode: {voice_mode!r} (expected one of {VOICE_MODES})"
            )
        self.voice_mode = voice_mode
        self.voice_clone_ref_audio = voice_clone_ref_audio
        self.voice_clone_ref_text = (voice_clone_ref_text or "").strip()
        self.skip_transcription = skip_transcription
        # Built-in CustomVoice speaker (VOICE_MODE_CUSTOM): the --voice
        # value / the Generate form's Speaker pick. Required there — there
        # is no configured default speaker.
        self.speaker = (voice or "").strip() or None
        if voice_mode == VOICE_MODE_CUSTOM and not self.speaker:
            raise ValueError(
                "CustomVoice mode requires a speaker: pass --voice SPEAKER "
                f"(one of {', '.join(QWEN3_TTS_SPEAKERS)})")
        # Voice design / style instruction (VoiceDesign mode): describes the
        # voice to design. Required there (validated by the converter).
        self.instructions = (instructions or "").strip()
        # api_url overrides the configured endpoint for the active voice mode
        # (used by the hub's "[remote]" backend entries and --api-url).
        self.api_url = (api_url or "").strip() or None
        # Seed sent with every request: config.SEED as-is, or (with
        # CONSTANT_SEED and SEED < 0) one random value drawn per run and
        # reused for every request so the voice stays consistent across
        # chunk boundaries. Without CONSTANT_SEED, -1 is forwarded so the
        # server re-samples the voice on every generation.
        self._seed = resolve_request_seed()
        if language is None:
            language = config.LANGUAGE
        # Validate before connecting so bad values fail fast without a server.
        self.language = normalize_language(language)
        self.client = None
        self.api_info: Dict[str, Any] = {}
        self.clone_client = None
        self.clone_api_info: Dict[str, Any] = {}
        self._ref_audio_filedata: Optional[Dict[str, Any]] = None
        self._connect()

    # ------------------------------------------------------------------
    # Connection
    # ------------------------------------------------------------------

    def _connect(self) -> None:
        # One demo server runs at a time on the configured port, hosting
        # whichever model this run selected (CustomVoice / Base / VoiceDesign).
        api_url = self.api_url or config.QWEN_API_URL
        try:
            if self.voice_mode == VOICE_MODE_CLONE:
                # Voice clone talks to the Base-model demo.
                self._init_client(api_url, clone=True)
                self._report(f"[OK] Connected to Voice Clone API at {api_url}")
                self._resolve_reference_text()
            else:
                self._init_client(api_url, clone=False)
                if self.voice_mode == VOICE_MODE_DESIGN:
                    self._report(f"[OK] Connected to Voice Design API at {api_url}")
                else:
                    self._report("[OK] Connected to Qwen API")
        except Exception as exc:
            raise RuntimeError(
                f"Qwen API initialization failed at {api_url}: {exc}. "
                "Make sure the Qwen demo server is running and reachable, and that your "
                "installed Qwen3-TTS version matches this converter's API expectations "
                "(voice clone requires the Base-model demo: Qwen/Qwen3-TTS-12Hz-1.7B-Base; "
                "voice design requires: Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign)."
            ) from exc

    def _resolve_reference_text(self) -> None:
        """Resolve the reference transcript: explicit text, then local
        transcription, then x-vector-only mode."""
        if not self.voice_clone_ref_text and self.voice_clone_ref_audio:
            if self.skip_transcription:
                self._report("[INFO] Skipping reference audio transcription (--no-transcription).")
            else:
                self._report("[INFO] Transcribing reference audio for voice cloning...")
                self.voice_clone_ref_text = self.transcribe_audio(self.voice_clone_ref_audio) or ""
        if not self.voice_clone_ref_text:
            self._report("[WARNING] No reference text available; using "
                         "x-vector-only clone mode (lower quality).")
            self._report('          Pass --transcription "..." for higher-quality in-context cloning.')
        else:
            self._report(f"[OK] Reference text:\n{self.voice_clone_ref_text}")

    def _init_client(self, url: str, clone: bool = False) -> None:
        """Initialize a Gradio client and store its API metadata.

        gradio_client prints its usage info directly to stdout while the
        client is created and its API metadata loaded, so stdout is swapped
        for a buffer for the whole process; the captured text is re-emitted
        at DEBUG level for troubleshooting.
        """
        from gradio_client import Client

        logger.info("Connecting to Qwen API at %s...", url)
        old_stdout = sys.stdout
        captured = io.StringIO()
        sys.stdout = captured
        try:
            try:
                client = Client(url, httpx_kwargs={"timeout": config.API_TIMEOUT})
            except TypeError:
                # Older gradio_client versions don't support httpx_kwargs.
                client = Client(url)
            if clone:
                self.clone_client = client
                self.clone_api_info = self._load_api_info(client)
            else:
                self.client = client
                self.api_info = self._load_api_info(client)
        finally:
            sys.stdout = old_stdout
        usage_info = captured.getvalue().strip()
        if usage_info:
            logger.debug("Gradio client output for %s:\n%s", url, usage_info)
        logger.info("Connected to Qwen API")

    @staticmethod
    def _load_api_info(client) -> Dict[str, Any]:
        """Load available API metadata from the Gradio app."""
        try:
            return client.view_api(return_format="dict")
        except Exception as exc:
            logger.warning("Unable to read API metadata: %s", exc)
            return {}

    def _resolve_api_name(self, *candidates: str, api_info: Optional[Dict[str, Any]] = None) -> str:
        """Return the first available api_name from candidate list."""
        info = api_info if api_info is not None else self.api_info
        named_endpoints = info.get("named_endpoints", {})
        for candidate in candidates:
            if candidate in named_endpoints:
                return candidate
        return candidates[0]

    def _endpoint_accepts_param(self, api_name: str, param_name: str,
                                api_info: Optional[Dict[str, Any]] = None) -> bool:
        """Check whether endpoint input schema includes the given parameter."""
        info = api_info if api_info is not None else self.api_info
        endpoint = info.get("named_endpoints", {}).get(api_name, {})
        parameters = endpoint.get("parameters", [])
        return any(parameter.get("parameter_name") == param_name for parameter in parameters)

    # ------------------------------------------------------------------
    # Reference audio transcription (voice clone)
    # ------------------------------------------------------------------

    def transcribe_audio(self, audio_path: str) -> Optional[str]:
        """Transcribe reference audio locally using an optional Whisper backend."""
        from .transcribe import transcribe_reference_audio
        return transcribe_reference_audio(audio_path)

    # ------------------------------------------------------------------
    # Chunk generation
    # ------------------------------------------------------------------

    def generate_chunk(self, text: str, chunk_num: int) -> Optional[str]:
        """Generate one audio chunk; returns its path in the chunks folder.

        The text is split into sub-requests of at most
        ``config.CHUNK_SIZE`` words each (the book-level chunker
        normally guarantees this already; the split is defense in depth
        against pathological input such as a punctuation-free run of
        text), and the audio files returned for the sub-requests are
        concatenated into one chunk file.
        """
        try:
            sub_texts = split_into_chunks(text, max_words=config.CHUNK_SIZE)
            if not sub_texts:
                raise RuntimeError("No text to synthesize")

            output_path: Optional[Path] = None
            with tempfile.TemporaryDirectory(prefix="tts_parts_") as parts_dir, \
                    self._chunk_heartbeat(chunk_num):
                part_paths = [
                    self._generate_sub_request(sub_text, parts_dir, sub_num,
                                                len(sub_texts), chunk_num)
                    for sub_num, sub_text in enumerate(sub_texts, 1)
                ]
                if len(part_paths) == 1:
                    suffix = part_paths[0].suffix or ".wav"
                    output_path = self._chunk_path(chunk_num, suffix)
                    shutil.copy2(part_paths[0], output_path)
                else:
                    output_path = self._chunk_path(chunk_num, ".wav")
                    concat_audio_files(part_paths, output_path)

            logger.debug("Chunk %d generated successfully (%d sub-request(s))",
                         chunk_num, len(sub_texts))
            return str(output_path)

        except ConversionCancelled:
            raise
        except Exception as exc:
            logger.error("Qwen chunk processing failed for chunk %d: %s", chunk_num, exc)
            return None

    def _generate_sub_request(self, text: str, parts_dir: str, sub_num: int,
                              sub_total: int, chunk_num: int) -> Path:
        """Run one API generation for ``text``; returns the downloaded audio."""
        if sub_total > 1:
            logger.info("Chunk %d: oversized input split into %d requests "
                        "(sub-request %d/%d)", chunk_num, sub_total, sub_num, sub_total)
        if self.voice_mode == VOICE_MODE_CUSTOM:
            result = self._generate_custom_voice(text)
        elif self.voice_mode == VOICE_MODE_CLONE:
            result = self._generate_voice_clone(text)
        elif self.voice_mode == VOICE_MODE_DESIGN:
            result = self._generate_voice_design(text)
        else:
            raise ValueError(f"Unknown voice mode: {self.voice_mode}")

        if not isinstance(result, (tuple, list)) or not result:
            raise RuntimeError("Qwen API returned an invalid result")

        audio_path = result[0]  # First element is the audio file path
        if not isinstance(audio_path, (str, Path)) or not audio_path:
            raise RuntimeError("Qwen API did not return an audio file path")

        source = Path(audio_path)
        if not source.exists():
            raise RuntimeError(f"Generated audio file not found: {audio_path}")

        destination = Path(parts_dir) / f"part_{sub_num:02d}{source.suffix or '.wav'}"
        shutil.copy2(source, destination)

        return destination

    # ------------------------------------------------------------------
    # API payloads
    # ------------------------------------------------------------------

    def _generate_custom_voice(self, text: str) -> Tuple:
        """Generate audio using CustomVoice mode with the run's speaker."""
        custom_api = self._resolve_api_name("/run_instruct", "/run_custom_voice", "/generate_custom_voice")
        if custom_api == "/run_instruct":
            payload = dict(
                text=text,
                lang_disp=self.language,
                spk_disp=speaker_display_name_for(self.speaker),
            )
        else:
            payload = dict(
                text=text,
                language=self.language,
                speaker=self.speaker,
            )
            if self._endpoint_accepts_param(custom_api, "model_id_cv"):
                payload["model_id_cv"] = CUSTOM_VOICE_MODEL_ID
            elif self._endpoint_accepts_param(custom_api, "model_size"):
                payload["model_size"] = MODEL_SIZE

            if self._endpoint_accepts_param(custom_api, "seed"):
                payload["seed"] = self._seed

        return self.client.predict(**payload, api_name=custom_api)

    def _generate_voice_design(self, text: str) -> Tuple:
        """Generate audio using VoiceDesign mode (described-voice model)."""
        design_api = self._resolve_api_name("/run_voice_design")
        # The demo's field is named "design"; older builds may call it
        # "instruct" instead.
        design_field = ("instruct"
                        if self._endpoint_accepts_param(design_api, "instruct")
                        and not self._endpoint_accepts_param(design_api, "design")
                        else "design")
        payload = dict(
            text=text,
            lang_disp=self.language,
            **{design_field: self.instructions},
        )
        if self._endpoint_accepts_param(design_api, "seed"):
            payload["seed"] = self._seed

        return self.client.predict(**payload, api_name=design_api)

    def _ref_audio_payload(self) -> Dict[str, Any]:
        """Gradio file payload for the reference audio (built once, reused)."""
        if self._ref_audio_filedata is None:
            from gradio_client import handle_file
            self._ref_audio_filedata = handle_file(self.voice_clone_ref_audio)
        return self._ref_audio_filedata

    def _generate_voice_clone(self, text: str) -> Tuple:
        """Generate audio using Voice Clone mode."""
        if not Path(self.voice_clone_ref_audio).exists():
            raise FileNotFoundError(f"Reference audio not found: {self.voice_clone_ref_audio}")

        if self.clone_client is None:
            raise RuntimeError("Voice Clone client is not initialized. Is the Base-model demo running?")

        clone_api = self._resolve_api_name("/run_voice_clone", "/generate_voice_clone",
                                           api_info=self.clone_api_info)
        use_xvector = not self.voice_clone_ref_text

        if clone_api == "/run_voice_clone":
            payload = dict(
                ref_aud=self._ref_audio_payload(),
                ref_txt=self.voice_clone_ref_text,
                use_xvec=use_xvector,
                text=text,
                lang_disp=self.language,
            )
        else:
            payload = dict(
                ref_audio=self._ref_audio_payload(),
                ref_text=self.voice_clone_ref_text,
                target_text=text,
                language=self.language,
                use_xvector_only=use_xvector,
            )
            optional_params = {
                "model_size": MODEL_SIZE,
                "seed": self._seed,
            }
            for name, value in optional_params.items():
                if self._endpoint_accepts_param(clone_api, name, api_info=self.clone_api_info):
                    payload[name] = value

        return self.clone_client.predict(**payload, api_name=clone_api)