openvoiceui 1.0.0

package/tts_providers/supertonic_tts.py

@@ -0,0 +1,399 @@
+#!/usr/bin/env python3
+"""
+Supertonic TTS wrapper for OpenVoiceUI.
+
+This module provides a clean interface to the Supertonic Text-to-Speech engine,
+wrapping the helper.py functionality for use in Flask applications.
+
+Author: OpenVoiceUI
+Date: 2026-02-11
+"""
+
+import os
+import sys
+import logging
+from io import BytesIO
+from typing import Optional
+
+import numpy as np
+import soundfile as sf
+
+# Add the Supertonic helper.py directory to the path
+SUPERTONIC_HELPER_PATH = os.environ.get("SUPERTONIC_HELPER_PATH", os.path.expanduser("~/supertonic/py"))
+if SUPERTONIC_HELPER_PATH not in sys.path:
+    sys.path.insert(0, SUPERTONIC_HELPER_PATH)
+
+try:
+    from helper import (
+        load_text_to_speech,
+        load_voice_style,
+        Style,
+    )
+except ImportError as e:
+    logging.error(f"Failed to import Supertonic helper: {e}")
+    logging.error(f"Make sure {SUPERTONIC_HELPER_PATH}/helper.py exists")
+    raise
+
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class SupertonicTTS:
+    """
+    Wrapper class for Supertonic Text-to-Speech engine.
+
+    This class provides a simple interface for generating speech from text
+    using the Supertonic ONNX models. It handles model loading, voice style
+    management, and audio generation.
+
+    Example:
+        >>> tts = SupertonicTTS(
+        ...     onnx_dir="~/supertonic/assets/onnx",
+        ...     voice_style_path="~/supertonic/assets/voice_styles/M1.json"
+        ... )
+        >>> audio_bytes = tts.generate_speech("Hello world, this is a test")
+        >>> # audio_bytes contains WAV format audio data
+    """
+
+    # Default paths (use SUPERTONIC_MODEL_PATH env var or ~/supertonic)
+    DEFAULT_ONNX_DIR = os.environ.get("SUPERTONIC_ONNX_DIR", os.path.expanduser("~/supertonic/assets/onnx"))
+    DEFAULT_VOICE_STYLES_DIR = os.environ.get("SUPERTONIC_VOICE_STYLES_DIR", os.path.expanduser("~/supertonic/assets/voice_styles"))
+
+    # Available voice styles
+    AVAILABLE_VOICE_STYLES = {
+        'M1': 'M1.json',  # Male voice 1
+        'M2': 'M2.json',  # Male voice 2
+        'F1': 'F1.json',  # Female voice 1
+        'F2': 'F2.json',  # Female voice 2
+    }
+
+    def __init__(
+        self,
+        onnx_dir: Optional[str] = None,
+        voice_style_path: Optional[str] = None,
+        voice_style_name: str = 'M1',
+        use_gpu: bool = False
+    ):
+        """
+        Initialize the Supertonic TTS engine.
+
+        Args:
+            onnx_dir: Path to the ONNX models directory. If None, uses DEFAULT_ONNX_DIR.
+            voice_style_path: Full path to the voice style JSON file. If None,
+                             constructs path from voice_style_name.
+            voice_style_name: Name of the voice style (M1, M2, F1, F2). Used only
+                             if voice_style_path is None.
+            use_gpu: Whether to use GPU for inference. Default is False (CPU only).
+
+        Raises:
+            FileNotFoundError: If onnx_dir or voice_style file doesn't exist.
+            RuntimeError: If model loading fails.
+        """
+        # Set paths
+        self.onnx_dir = onnx_dir or self.DEFAULT_ONNX_DIR
+        self.voice_style_name = voice_style_name
+
+        # Validate onnx directory
+        if not os.path.exists(self.onnx_dir):
+            raise FileNotFoundError(
+                f"ONNX models directory not found: {self.onnx_dir}"
+            )
+        logger.info(f"Using ONNX models from: {self.onnx_dir}")
+
+        # Set voice style path
+        if voice_style_path:
+            self.voice_style_path = voice_style_path
+        else:
+            # Construct path from voice style name
+            if voice_style_name not in self.AVAILABLE_VOICE_STYLES:
+                raise ValueError(
+                    f"Invalid voice_style_name: {voice_style_name}. "
+                    f"Available: {list(self.AVAILABLE_VOICE_STYLES.keys())}"
+                )
+            voice_style_file = self.AVAILABLE_VOICE_STYLES[voice_style_name]
+            self.voice_style_path = os.path.join(
+                self.DEFAULT_VOICE_STYLES_DIR, voice_style_file
+            )
+
+        # Validate voice style file
+        if not os.path.exists(self.voice_style_path):
+            raise FileNotFoundError(
+                f"Voice style file not found: {self.voice_style_path}"
+            )
+        logger.info(f"Using voice style: {self.voice_style_path}")
+
+        # Initialize models
+        try:
+            logger.info("Loading Supertonic TTS models...")
+            self.text_to_speech = load_text_to_speech(self.onnx_dir, use_gpu=use_gpu)
+            self.style = load_voice_style([self.voice_style_path], verbose=True)
+            self.sample_rate = self.text_to_speech.sample_rate
+            logger.info(f"TTS models loaded successfully (sample rate: {self.sample_rate}Hz)")
+        except Exception as e:
+            logger.error(f"Failed to load TTS models: {e}")
+            raise RuntimeError(f"TTS model loading failed: {e}")
+
+    def generate_speech(
+        self,
+        text: str,
+        lang: str = 'en',
+        speed: float = 1.0,
+        total_step: int = 15
+    ) -> bytes:
+        """
+        Generate speech from text.
+
+        Args:
+            text: The text to synthesize into speech.
+            lang: Language code ('en', 'ko', 'es', 'pt', 'fr'). Default is 'en'.
+            speed: Speech speed multiplier. Higher values = faster speech.
+                   Recommended range: 0.8 to 1.3. Default is 1.05.
+            total_step: Number of denoising steps for generation. More steps =
+                        better quality but slower. Recommended range: 3-10.
+                        Default is 5 (good balance).
+
+        Returns:
+            bytes: Raw WAV audio data (can be written directly to file or sent
+                   via HTTP with Content-Type: audio/wav).
+
+        Raises:
+            ValueError: If lang is not supported or parameters are invalid.
+            RuntimeError: If speech generation fails.
+
+        Example:
+            >>> audio = tts.generate_speech("Hello world", lang='en', speed=1.05)
+            >>> with open('output.wav', 'wb') as f:
+            ...     f.write(audio)
+        """
+        # Validate inputs
+        if not text or not text.strip():
+            raise ValueError("Text cannot be empty")
+
+        supported_langs = ['en', 'ko', 'es', 'pt', 'fr']
+        if lang not in supported_langs:
+            raise ValueError(
+                f"Unsupported language: {lang}. Supported: {supported_langs}"
+            )
+
+        if speed <= 0 or speed > 3:
+            raise ValueError(f"Invalid speed: {speed}. Must be between 0 and 3")
+
+        if total_step < 1 or total_step > 50:
+            raise ValueError(f"Invalid total_step: {total_step}. Must be between 1 and 50")
+
+        logger.info(f"Generating speech: '{text[:50]}...' (lang={lang}, speed={speed}, steps={total_step})")
+
+        # Maximum character length per chunk to stay under ONNX token limit (~1000 tokens)
+        MAX_CHUNK_LENGTH = 500
+
+        def split_text_into_chunks(text: str, max_length: int) -> list:
+            """Split text into chunks at sentence boundaries."""
+            if len(text) <= max_length:
+                return [text]
+
+            chunks = []
+            # Split on sentence boundaries
+            sentence_endings = ['. ', '! ', '? ', '\n']
+
+            current_chunk = ""
+            # Split by sentences first
+            sentences = [text]
+            for ending in sentence_endings:
+                new_sentences = []
+                for s in sentences:
+                    parts = s.split(ending)
+                    for i, part in enumerate(parts):
+                        if i < len(parts) - 1:
+                            new_sentences.append(part + ending.strip())
+                        elif part.strip():
+                            new_sentences.append(part)
+                sentences = new_sentences if new_sentences else sentences
+
+            # Combine sentences into chunks up to max_length
+            for sentence in sentences:
+                sentence = sentence.strip()
+                if not sentence:
+                    continue
+
+                if len(current_chunk) + len(sentence) + 1 <= max_length:
+                    current_chunk += (" " if current_chunk else "") + sentence
+                else:
+                    if current_chunk:
+                        chunks.append(current_chunk.strip())
+                    # If single sentence is too long, just use it (will be truncated by tokenizer)
+                    if len(sentence) > max_length:
+                        chunks.append(sentence)
+                        current_chunk = ""
+                    else:
+                        current_chunk = sentence
+
+            if current_chunk.strip():
+                chunks.append(current_chunk.strip())
+
+            return chunks
+
+        try:
+            chunks = split_text_into_chunks(text, MAX_CHUNK_LENGTH)
+            logger.info(f"Text split into {len(chunks)} chunk(s)")
+
+            all_audio_chunks = []
+
+            for i, chunk in enumerate(chunks):
+                logger.info(f"Processing chunk {i+1}/{len(chunks)}: '{chunk[:30]}...'")
+
+                # Generate speech using the Supertonic TextToSpeech instance
+                wav, duration = self.text_to_speech(
+                    text=chunk,
+                    lang=lang,
+                    style=self.style,
+                    total_step=total_step,
+                    speed=speed
+                )
+
+                # Extract the audio data (first batch item, trim to actual duration)
+                audio_data = wav[0, :int(self.sample_rate * duration[0].item())]
+                all_audio_chunks.append(audio_data)
+
+            # Concatenate all audio chunks
+            if len(all_audio_chunks) == 1:
+                final_audio = all_audio_chunks[0]
+            else:
+                final_audio = np.concatenate(all_audio_chunks)
+
+            # Write to BytesIO buffer to get raw bytes
+            buffer = BytesIO()
+            sf.write(buffer, final_audio, self.sample_rate, format='WAV')
+            audio_bytes = buffer.getvalue()
+
+            total_duration = len(final_audio) / self.sample_rate
+            logger.info(f"Generated {len(audio_bytes)} bytes of audio ({total_duration:.2f}s)")
+            return audio_bytes
+
+        except Exception as e:
+            logger.error(f"Speech generation failed: {e}")
+            raise RuntimeError(f"Failed to generate speech: {e}")
+
+    def set_voice_style(self, voice_style_name: str) -> None:
+        """
+        Change the voice style.
+
+        Args:
+            voice_style_name: Name of the new voice style (M1, M2, F1, F2).
+
+        Raises:
+            ValueError: If voice_style_name is not available.
+            FileNotFoundError: If the voice style file doesn't exist.
+            RuntimeError: If loading the new style fails.
+        """
+        if voice_style_name == self.voice_style_name:
+            logger.info(f"Already using voice style: {voice_style_name}")
+            return
+
+        if voice_style_name not in self.AVAILABLE_VOICE_STYLES:
+            raise ValueError(
+                f"Invalid voice_style_name: {voice_style_name}. "
+                f"Available: {list(self.AVAILABLE_VOICE_STYLES.keys())}"
+            )
+
+        voice_style_file = self.AVAILABLE_VOICE_STYLES[voice_style_name]
+        new_voice_style_path = os.path.join(
+            self.DEFAULT_VOICE_STYLES_DIR, voice_style_file
+        )
+
+        if not os.path.exists(new_voice_style_path):
+            raise FileNotFoundError(
+                f"Voice style file not found: {new_voice_style_path}"
+            )
+
+        try:
+            self.style = load_voice_style([new_voice_style_path], verbose=True)
+            self.voice_style_name = voice_style_name
+            self.voice_style_path = new_voice_style_path
+            logger.info(f"Voice style changed to: {voice_style_name}")
+        except Exception as e:
+            logger.error(f"Failed to load voice style: {e}")
+            raise RuntimeError(f"Failed to load voice style: {e}")
+
+
+# Singleton instance for use in Flask app
+_tts_instance: Optional[SupertonicTTS] = None
+
+
+def get_tts_instance() -> Optional[SupertonicTTS]:
+    """
+    Get the global TTS instance (singleton).
+
+    Returns:
+        The global SupertonicTTS instance, or None if not initialized.
+
+    This is useful for Flask apps where you want to initialize TTS once
+    at startup and reuse the instance across requests.
+    """
+    global _tts_instance
+    return _tts_instance
+
+
+def initialize_tts(
+    onnx_dir: Optional[str] = None,
+    voice_style_name: str = 'M1',
+    use_gpu: bool = False
+) -> Optional[SupertonicTTS]:
+    """
+    Initialize the global TTS instance.
+
+    Args:
+        onnx_dir: Path to ONNX models directory.
+        voice_style_name: Default voice style to use.
+        use_gpu: Whether to use GPU for inference.
+
+    Returns:
+        The initialized SupertonicTTS instance, or None if initialization fails.
+    """
+    global _tts_instance
+    try:
+        _tts_instance = SupertonicTTS(
+            onnx_dir=onnx_dir,
+            voice_style_name=voice_style_name,
+            use_gpu=use_gpu
+        )
+        logger.info("Global TTS instance initialized")
+        return _tts_instance
+    except Exception as e:
+        logger.error(f"Failed to initialize TTS: {e}")
+        _tts_instance = None
+        return None
+
+
+if __name__ == "__main__":
+    # Simple test when run directly
+    print("Supertonic TTS Wrapper - Direct Test")
+    print("=" * 50)
+
+    try:
+        # Initialize TTS
+        tts = SupertonicTTS(
+            onnx_dir=os.environ.get("SUPERTONIC_ONNX_DIR", os.path.expanduser("~/supertonic/assets/onnx")),
+            voice_style_name="M1"
+        )
+
+        # Generate speech
+        test_text = "Hello world, this is a test of the Supertonic TTS system."
+        audio = tts.generate_speech(test_text, lang='en', speed=1.05)
+
+        # Save to file
+        output_path = "/tmp/supertonic_test_output.wav"
+        with open(output_path, 'wb') as f:
+            f.write(audio)
+
+        print(f"Success! Audio saved to: {output_path}")
+        print(f"Generated {len(audio)} bytes of audio data")
+
+    except Exception as e:
+        print(f"Error: {e}")
+        import traceback
+        traceback.print_exc()