handler.py

import io
import base64
import logging
import subprocess
import sys

import numpy as np
import soundfile as sf

logger = logging.getLogger(__name__)


def _install_deps():
    """Install OmniVoice and voicetut-tts with --no-deps to avoid conflicts
    with the HF Inference Endpoint base image.

    Then upgrade transformers (WITH deps) because OmniVoice needs
    HiggsAudioV2TokenizerModel which only exists in transformers>=4.50.
    The HF Inference Toolkit already loaded successfully with the old
    transformers at this point, so upgrading is safe."""
    # 1. Install OmniVoice and voicetut-tts without their dep trees
    pkgs = [
        "git+https://github.com/k2-fsa/OmniVoice.git",
        "voicetut-tts>=0.1.0",
    ]
    for pkg in pkgs:
        try:
            subprocess.check_call(
                [sys.executable, "-m", "pip", "install", "--no-deps", pkg],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
            logger.info(f"Installed {pkg}")
        except subprocess.CalledProcessError as e:
            stderr = e.stderr.decode() if e.stderr else ""
            if "already satisfied" not in stderr.lower():
                logger.warning(f"Failed to install {pkg}: {stderr[:200]}")

    # 2. Upgrade transformers to get HiggsAudioV2TokenizerModel (needed by OmniVoice)
    try:
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "--upgrade",
             "transformers>=4.50.0", "tokenizers>=0.21.0"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
        )
        logger.info("Upgraded transformers for OmniVoice compatibility")
    except subprocess.CalledProcessError as e:
        stderr = e.stderr.decode() if e.stderr else ""
        logger.warning(f"Failed to upgrade transformers: {stderr[:200]}")


# Install deps before importing voicetut_tts
_install_deps()

from voicetut_tts import VoiceTutTTS, GenerationParams


class EndpointHandler:
    def __init__(self, model_dir=""):
        logger.info(f"Loading VoiceTut-TTS from {model_dir or 'mohammedaly22/VoiceTut-TTS'}...")
        self.tts = VoiceTutTTS.from_pretrained(
            model_dir or "mohammedaly22/VoiceTut-TTS",
            dtype="float16",
        )
        self.sampling_rate = self.tts.sampling_rate
        logger.info(f"VoiceTut-TTS loaded. Sampling rate: {self.sampling_rate}")

    def __call__(self, data: dict) -> dict:
        inputs = data.get("inputs", data)

        text = inputs.get("text", "").strip()
        if not text:
            return {"error": "text is required"}

        voice = inputs.get("voice", "Mohamed")
        language = inputs.get("language", "arz")
        normalize = inputs.get("normalize", True)
        num_step = int(inputs.get("num_step", 32))
        guidance_scale = float(inputs.get("guidance_scale", 2.0))
        speed = float(inputs.get("speed", 1.0))

        params = GenerationParams(
            num_step=num_step,
            guidance_scale=guidance_scale,
            speed=speed,
        )

        try:
            wav = self.tts.synthesize(
                text,
                speaker=voice,
                language=language,
                normalize=normalize,
                params=params,
            )
        except Exception as e:
            logger.exception("Synthesis failed")
            return {"error": str(e)}

        buf = io.BytesIO()
        sf.write(buf, wav, self.sampling_rate, format="WAV")
        wav_bytes = buf.getvalue()
        wav_b64 = base64.b64encode(wav_bytes).decode("utf-8")

        return {
            "audio_base64": wav_b64,
            "sampling_rate": self.sampling_rate,
            "duration_sec": float(len(wav) / self.sampling_rate),
        }