tts_client.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Hugging Face Text-to-Speech (TTS) client for Sushruta Patient 360."""

import logging
import requests
import hashlib
from typing import Tuple, Optional
from config import HF_TOKEN, GENERATE_SPEECH
from cache_manager import intake_cache

logger = logging.getLogger(__name__)


def synthesize_tts(text: str, voice: str = "af_heart") -> Tuple[Optional[bytes], Optional[str]]:
    """Synthesize text to speech using Hugging Face's serverless inference API.

    Uses `hexgrad/Kokoro-82M` or a fallback TTS model.
    Returns:
        A tuple of (audio_bytes, mime_type), or (None, None) if failed or disabled.
    """
    if not text or not text.strip():
        return None, None

    # Calculate cache key
    key_str = f"tts:{voice}:{text.strip()}"
    key_hash = hashlib.sha256(key_str.encode("utf-8")).hexdigest()

    # Check cache first
    try:
        cached_audio = intake_cache.get(key_hash)
        if cached_audio is not None:
            logger.info("TTS Cache hit for text: '%s...'", text[:30])
            return cached_audio, "audio/mpeg"
    except Exception as e:
        logger.warning("Error reading TTS cache: %s", e)

    # If speech generation is disabled and cache missed, return None
    if not GENERATE_SPEECH:
        logger.debug("TTS generation disabled (GENERATE_SPEECH=false), skipping API call.")
        return None, None

    # Call HF serverless API for TTS (Kokoro-82M)
    api_url = "https://api-inference.huggingface.co/models/hexgrad/Kokoro-82M"
    headers = {}
    if HF_TOKEN:
        headers["Authorization"] = f"Bearer {HF_TOKEN}"

    # Kokoro-82M might accept speakers/voice in the payload, but standard inputs is required
    payload = {
        "inputs": text,
        "parameters": {
            "voice": voice
        }
    }

    try:
        logger.info("Requesting TTS from Hugging Face: %s", api_url)
        response = requests.post(api_url, headers=headers, json=payload, timeout=30)
        
        # Check if the model is still loading
        if response.status_code == 503:
            logger.warning("HF TTS model is loading. Retrying with simple inputs...")
            response = requests.post(api_url, headers=headers, json={"inputs": text}, timeout=30)

        response.raise_for_status()
        audio_bytes = response.content
        mime_type = response.headers.get("Content-Type", "audio/mpeg")

        # Save to cache
        try:
            intake_cache.set(key_hash, audio_bytes)
        except Exception as e:
            logger.warning("Failed to save TTS to cache: %s", e)

        return audio_bytes, mime_type

    except Exception as e:
        logger.error("Failed to generate speech via Hugging Face API: %s", e)
        # Try a fallback TTS model if Kokoro fails
        fallback_api_url = "https://api-inference.huggingface.co/models/facebook/mms-tts-eng"
        try:
            logger.info("Attempting fallback TTS via %s", fallback_api_url)
            response = requests.post(fallback_api_url, headers=headers, json={"inputs": text}, timeout=30)
            response.raise_for_status()
            audio_bytes = response.content
            mime_type = response.headers.get("Content-Type", "audio/wav")
            
            # Save fallback to cache
            try:
                intake_cache.set(key_hash, audio_bytes)
            except Exception as ce:
                logger.warning("Failed to cache fallback TTS: %s", ce)
                
            return audio_bytes, mime_type
        except Exception as fallback_err:
            logger.error("Fallback TTS also failed: %s", fallback_err)
            return None, None