feat: Replace Matcha with Soprano TTS and add inference pipeline

- Switch from Matcha TTS to Soprano TTS for simpler setup
- Remove Gemma local model (using Groq API instead)
- Add fallback tokenizer loading in dataset creation
- Make text_tokens optional when tokenizer unavailable
- Add inference.py for end-to-end speech-to-speech
- Add TTS service abstractions (soprano, matcha, piper, gemma)
- Add dataset verification script
- Fix version constraints for WhisperX compatibility

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

.gitignore

@@ -25,3 +25,4 @@ temp_*/
 # IDE
 .vscode/
 .idea/
+datasets/samples/

datasets/verify_dataset.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""
+Verify dataset quality by transcribing audio and comparing with original text.
+Uses WhisperX to transcribe TTS-generated audio.
+"""
+import os
+import sys
+import torch
+import numpy as np
+from difflib import SequenceMatcher
+
+# Patch torch.load
+_orig = torch.load
+torch.load = lambda *a, **kw: _orig(*a, **{**kw, 'weights_only': False})
+
+
+def similarity(a: str, b: str) -> float:
+    """Calculate similarity ratio between two strings."""
+    return SequenceMatcher(None, a.lower(), b.lower()).ratio()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Verify dataset audio quality")
+    parser.add_argument("--dataset", type=str, default="./data/test_b200.pt", help="Dataset path")
+    parser.add_argument("--samples", type=int, default=10, help="Number of samples to verify")
+    parser.add_argument("--gpu", type=int, default=0, help="GPU to use")
+    args = parser.parse_args()
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+
+    import torchaudio
+    import whisperx
+
+    print("=" * 60)
+    print("Dataset Audio Verification")
+    print("=" * 60)
+
+    # Load dataset
+    print(f"\nLoading dataset: {args.dataset}")
+    data = torch.load(args.dataset)
+    print(f"Total items: {len(data)}")
+
+    # Load TTS for regenerating audio
+    print("\nLoading TTS...")
+    from soprano import SopranoTTS
+    tts = SopranoTTS(backend="transformers", device="cuda")
+
+    # Load WhisperX for transcription
+    print("Loading WhisperX...")
+    wx_model = whisperx.load_model("large-v3-turbo", "cuda", compute_type="float16", language="en")
+
+    # Select random samples
+    import random
+    indices = random.sample(range(len(data)), min(args.samples, len(data)))
+
+    print(f"\nVerifying {len(indices)} samples...\n")
+    print("-" * 60)
+
+    results = {
+        "question": {"total": 0, "good": 0, "similarities": []},
+        "answer": {"total": 0, "good": 0, "similarities": []}
+    }
+
+    for i, idx in enumerate(indices):
+        item = data[idx]
+        q_text = item["text"]
+        a_text = item["answer"]
+
+        print(f"[{i+1}/{len(indices)}] Sample {idx}")
+
+        # Generate audio with TTS
+        q_audio = tts.infer(q_text)
+        a_audio = tts.infer(a_text)
+
+        # Convert to numpy
+        q_np = q_audio.cpu().numpy() if hasattr(q_audio, 'cpu') else np.array(q_audio)
+        a_np = a_audio.cpu().numpy() if hasattr(a_audio, 'cpu') else np.array(a_audio)
+
+        # Resample to 16kHz for WhisperX
+        q_16k = torchaudio.functional.resample(torch.from_numpy(q_np), 32000, 16000).numpy()
+        a_16k = torchaudio.functional.resample(torch.from_numpy(a_np), 32000, 16000).numpy()
+
+        # Transcribe
+        q_result = wx_model.transcribe(q_16k, batch_size=1)
+        a_result = wx_model.transcribe(a_16k, batch_size=1)
+
+        q_transcribed = " ".join([s["text"].strip() for s in q_result["segments"]])
+        a_transcribed = " ".join([s["text"].strip() for s in a_result["segments"]])
+
+        # Calculate similarity
+        q_sim = similarity(q_text, q_transcribed)
+        a_sim = similarity(a_text, a_transcribed)
+
+        results["question"]["total"] += 1
+        results["answer"]["total"] += 1
+        results["question"]["similarities"].append(q_sim)
+        results["answer"]["similarities"].append(a_sim)
+
+        if q_sim >= 0.8:
+            results["question"]["good"] += 1
+        if a_sim >= 0.8:
+            results["answer"]["good"] += 1
+
+        # Print results
+        q_status = "✓" if q_sim >= 0.8 else "✗"
+        a_status = "✓" if a_sim >= 0.8 else "✗"
+
+        print(f"  Question ({q_sim:.0%}) {q_status}")
+        print(f"    Original:    \"{q_text[:60]}...\"" if len(q_text) > 60 else f"    Original:    \"{q_text}\"")
+        print(f"    Transcribed: \"{q_transcribed[:60]}...\"" if len(q_transcribed) > 60 else f"    Transcribed: \"{q_transcribed}\"")
+
+        print(f"  Answer ({a_sim:.0%}) {a_status}")
+        print(f"    Original:    \"{a_text[:60]}...\"" if len(a_text) > 60 else f"    Original:    \"{a_text}\"")
+        print(f"    Transcribed: \"{a_transcribed[:60]}...\"" if len(a_transcribed) > 60 else f"    Transcribed: \"{a_transcribed}\"")
+        print()
+
+    # Summary
+    print("=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+
+    q_avg = np.mean(results["question"]["similarities"])
+    a_avg = np.mean(results["answer"]["similarities"])
+
+    print(f"\nQuestions:")
+    print(f"  Good (>=80%): {results['question']['good']}/{results['question']['total']}")
+    print(f"  Avg similarity: {q_avg:.1%}")
+
+    print(f"\nAnswers:")
+    print(f"  Good (>=80%): {results['answer']['good']}/{results['answer']['total']}")
+    print(f"  Avg similarity: {a_avg:.1%}")
+
+    overall = (q_avg + a_avg) / 2
+    print(f"\nOverall similarity: {overall:.1%}")
+
+    if overall >= 0.85:
+        print("\n✓ Dataset quality: GOOD")
+    elif overall >= 0.70:
+        print("\n⚠ Dataset quality: ACCEPTABLE")
+    else:
+        print("\n✗ Dataset quality: POOR")
+
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -0,0 +1,324 @@
+#!/usr/bin/env python3
+"""
+Inference script for Speech-to-Speech model.
+
+Usage:
+    python inference.py --checkpoint ./checkpoints/stage2_best.pt --input audio.wav --output response.wav
+    python inference.py --checkpoint ./checkpoints/stage2_best.pt --text "Hello, how are you?"
+"""
+
+import os
+import sys
+import argparse
+import torch
+import torch.nn as nn
+import numpy as np
+
+# SNAC token offsets for Orpheus
+SNAC_BASE_OFFSET = 128266
+EOS_TOKEN = 128009
+
+
+class SpeechAdapter(nn.Module):
+    """Same architecture as training - must match exactly."""
+    def __init__(self, whisper_dim=1280, llm_dim=3072, downsample=5, intermediate_dim=2048):
+        super().__init__()
+        self.downsample = downsample
+        concat_dim = whisper_dim * downsample
+
+        self.ffn = nn.Sequential(
+            nn.Linear(concat_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, llm_dim),
+            nn.LayerNorm(llm_dim)
+        )
+
+    def forward(self, x):
+        B, T, D = x.shape
+        T_new = (T // self.downsample) * self.downsample
+        x = x[:, :T_new]
+        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
+        return self.ffn(x)
+
+
+def decode_snac_tokens(snac_tokens, device="cuda"):
+    """Decode SNAC tokens to audio waveform."""
+    try:
+        from snac import SNAC
+
+        # Load SNAC model
+        snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
+        snac.eval()
+
+        # Remove offsets from tokens
+        raw_tokens = []
+        for i, tok in enumerate(snac_tokens):
+            pos = i % 7
+            offset = SNAC_BASE_OFFSET + pos * 4096
+            raw_tok = tok - offset
+            if 0 <= raw_tok < 4096:
+                raw_tokens.append(raw_tok)
+
+        if len(raw_tokens) == 0:
+            return None, 24000
+
+        # Reshape to SNAC format: 7 tokens per frame
+        num_frames = len(raw_tokens) // 7
+        if num_frames == 0:
+            return None, 24000
+
+        raw_tokens = raw_tokens[:num_frames * 7]
+
+        # SNAC expects [batch, layers, time] - 3 layers with different rates
+        # Layer 0: 1 token/frame, Layer 1: 2 tokens/frame, Layer 2: 4 tokens/frame
+        codes = []
+        for frame_idx in range(num_frames):
+            base = frame_idx * 7
+            codes.append(raw_tokens[base:base+7])
+
+        codes = torch.tensor(codes, device=device)
+
+        # Reorganize into SNAC layer format
+        layer0 = codes[:, 0:1].T  # [1, num_frames]
+        layer1 = codes[:, 1:3].T.reshape(1, -1)  # [1, num_frames*2]
+        layer2 = codes[:, 3:7].T.reshape(1, -1)  # [1, num_frames*4]
+
+        with torch.no_grad():
+            audio = snac.decode([layer0, layer1, layer2])
+
+        return audio.cpu().numpy().squeeze(), 24000
+
+    except Exception as e:
+        print(f"SNAC decode error: {e}")
+        return None, 24000
+
+
+def extract_whisper_features(audio_path, device="cuda"):
+    """Extract Whisper encoder features from audio file."""
+    try:
+        from transformers import WhisperProcessor, WhisperModel
+        import librosa
+
+        # Load audio
+        audio, sr = librosa.load(audio_path, sr=16000)
+
+        # Load Whisper
+        processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
+        model = WhisperModel.from_pretrained("openai/whisper-large-v3").to(device)
+        model.eval()
+
+        # Process
+        inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
+        input_features = inputs.input_features.to(device)
+
+        with torch.no_grad():
+            encoder_outputs = model.encoder(input_features)
+            features = encoder_outputs.last_hidden_state
+
+        return features
+
+    except Exception as e:
+        print(f"Whisper feature extraction error: {e}")
+        return None
+
+
+def generate_response(model, adapter, tokenizer, audio_embeds, device, max_new_tokens=500):
+    """Generate interleaved text+audio response."""
+
+    # Get the base model for generation
+    if hasattr(model, 'get_base_model'):
+        base_model = model.get_base_model()
+    else:
+        base_model = model
+
+    # Start generation from audio embeddings
+    generated_tokens = []
+
+    # Create initial input from audio embeddings
+    current_embeds = audio_embeds
+
+    with torch.no_grad():
+        for step in range(max_new_tokens):
+            # Forward pass
+            outputs = model(inputs_embeds=current_embeds, use_cache=False)
+            logits = outputs.logits
+
+            # Get next token (greedy)
+            next_token_logits = logits[:, -1, :]
+            next_token = torch.argmax(next_token_logits, dim=-1)
+
+            token_id = next_token.item()
+            generated_tokens.append(token_id)
+
+            # Check for EOS
+            if token_id == EOS_TOKEN:
+                break
+
+            # Get embedding for next token
+            if hasattr(base_model, 'model'):
+                next_embed = base_model.model.embed_tokens(next_token.unsqueeze(0))
+            else:
+                next_embed = base_model.embed_tokens(next_token.unsqueeze(0))
+
+            # Append to current embeddings
+            current_embeds = torch.cat([current_embeds, next_embed], dim=1)
+
+            # Truncate if too long (keep last 2048 tokens)
+            if current_embeds.shape[1] > 2048:
+                current_embeds = current_embeds[:, -2048:]
+
+    return generated_tokens
+
+
+def separate_tokens(generated_tokens):
+    """Separate text and audio tokens from interleaved output."""
+    text_tokens = []
+    audio_tokens = []
+
+    for tok in generated_tokens:
+        if tok >= SNAC_BASE_OFFSET:
+            audio_tokens.append(tok)
+        elif tok != EOS_TOKEN:
+            text_tokens.append(tok)
+
+    return text_tokens, audio_tokens
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Speech-to-Speech Inference")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint (stage1 or stage2)")
+    parser.add_argument("--input", type=str, default=None, help="Input audio file")
+    parser.add_argument("--text", type=str, default=None, help="Input text (for testing without audio)")
+    parser.add_argument("--output", type=str, default="./output.wav", help="Output audio file")
+    parser.add_argument("--model_path", type=str, default="canopylabs/3b-es_it-ft-research_release")
+    parser.add_argument("--max_tokens", type=int, default=500)
+    parser.add_argument("--device", type=str, default=None)
+    args = parser.parse_args()
+
+    # Determine device
+    if args.device:
+        device = torch.device(args.device)
+    elif torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    print(f"Device: {device}")
+
+    # Determine dtype
+    torch_dtype = torch.bfloat16 if device.type == 'cuda' else torch.float32
+
+    # Load tokenizer
+    print(f"Loading tokenizer: {args.model_path}")
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+
+    # Load checkpoint
+    print(f"Loading checkpoint: {args.checkpoint}")
+    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+
+    has_lora = "lora" in ckpt
+    print(f"Checkpoint type: {'Stage 2 (Adapter + LoRA)' if has_lora else 'Stage 1 (Adapter only)'}")
+
+    # Load LLM
+    print(f"Loading LLM: {args.model_path}")
+    llm = AutoModelForCausalLM.from_pretrained(
+        args.model_path,
+        torch_dtype=torch_dtype,
+        attn_implementation="sdpa",
+    ).to(device)
+
+    # Apply LoRA if Stage 2
+    if has_lora:
+        from peft import LoraConfig, get_peft_model, TaskType
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+            lora_dropout=0.0,
+            bias="none",
+            task_type=TaskType.CAUSAL_LM
+        )
+        llm = get_peft_model(llm, lora_config)
+        llm.load_state_dict(ckpt["lora"], strict=False)
+        print("LoRA weights loaded")
+
+    llm.eval()
+
+    # Load adapter
+    print("Loading adapter...")
+    adapter = SpeechAdapter(
+        whisper_dim=1280,
+        llm_dim=3072,
+        downsample=5,
+        intermediate_dim=2048
+    ).to(device, dtype=torch_dtype)
+    adapter.load_state_dict(ckpt["adapter"])
+    adapter.eval()
+    print("Adapter loaded")
+
+    # Get input embeddings
+    if args.input:
+        print(f"Processing audio: {args.input}")
+        whisper_features = extract_whisper_features(args.input, device)
+        if whisper_features is None:
+            print("Failed to extract Whisper features")
+            return
+        audio_embeds = adapter(whisper_features.to(torch_dtype))
+
+    elif args.text:
+        print(f"Processing text: {args.text}")
+        # For text input, create dummy audio embeddings (zeros)
+        # This is just for testing the generation pipeline
+        dummy_features = torch.randn(1, 100, 1280, device=device, dtype=torch_dtype)
+        audio_embeds = adapter(dummy_features)
+
+        # Optionally prepend text tokens
+        text_tokens = tokenizer.encode(args.text, add_special_tokens=False)
+        print(f"Text tokens: {text_tokens[:10]}...")
+    else:
+        print("ERROR: Provide --input (audio file) or --text")
+        return
+
+    print(f"Audio embeddings shape: {audio_embeds.shape}")
+
+    # Generate response
+    print(f"Generating response (max {args.max_tokens} tokens)...")
+    generated_tokens = generate_response(
+        llm, adapter, tokenizer, audio_embeds, device,
+        max_new_tokens=args.max_tokens
+    )
+
+    print(f"Generated {len(generated_tokens)} tokens")
+
+    # Separate text and audio
+    text_tokens, audio_tokens = separate_tokens(generated_tokens)
+    print(f"Text tokens: {len(text_tokens)}, Audio tokens: {len(audio_tokens)}")
+
+    # Decode text
+    if text_tokens:
+        decoded_text = tokenizer.decode(text_tokens, skip_special_tokens=True)
+        print(f"\nGenerated text: {decoded_text}")
+
+    # Decode audio
+    if audio_tokens:
+        print(f"\nDecoding {len(audio_tokens)} audio tokens...")
+        audio, sr = decode_snac_tokens(audio_tokens, device)
+
+        if audio is not None:
+            import soundfile as sf
+            sf.write(args.output, audio, sr)
+            print(f"Audio saved: {args.output}")
+        else:
+            print("Failed to decode audio")
+    else:
+        print("No audio tokens generated")
+
+    # Show raw tokens for debugging
+    print(f"\nFirst 20 generated tokens: {generated_tokens[:20]}")
+
+
+if __name__ == "__main__":
+    main()

services/gemma_service.py

@@ -0,0 +1,188 @@
+"""
+Gemma 3 1B Service for Q&A generation using llama.cpp.
+
+Requires: pip install llama-cpp-python
+For GPU: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python --force-reinstall --no-cache-dir
+"""
+import os
+from typing import List, Dict, Optional
+
+# Default model path - Google's official QAT Q4_0 version (smallest, ~600MB)
+DEFAULT_MODEL_PATH = "/tmp/gemma-3-1b-it-q4_0.gguf"
+DEFAULT_MODEL_URL = "https://huggingface.co/google/gemma-3-1b-it-qat-q4_0-gguf/resolve/main/gemma-3-1b-it-q4_0.gguf"
+
+
+class GemmaService:
+    """Service for generating Q&A pairs using Gemma 3 1B via llama.cpp."""
+
+    def __init__(
+        self,
+        model_path: str = DEFAULT_MODEL_PATH,
+        n_gpu_layers: int = -1,  # -1 = all layers on GPU
+        n_ctx: int = 4096,
+        n_batch: int = 512,
+        verbose: bool = False
+    ):
+        self.model_path = model_path
+        self.n_gpu_layers = n_gpu_layers
+        self.n_ctx = n_ctx
+        self.n_batch = n_batch
+        self.verbose = verbose
+        self._llm = None
+
+    def download_model(self) -> bool:
+        """Download the model if not present."""
+        if os.path.exists(self.model_path):
+            print(f"[GemmaService] Model already exists: {self.model_path}")
+            return True
+
+        print(f"[GemmaService] Downloading Gemma 3 1B to {self.model_path}...")
+        import subprocess
+        result = subprocess.run(
+            ["wget", "-q", "--show-progress", "-O", self.model_path, DEFAULT_MODEL_URL],
+            capture_output=False
+        )
+        return result.returncode == 0
+
+    def load(self) -> 'GemmaService':
+        """Load the Gemma model."""
+        from llama_cpp import Llama
+
+        if not os.path.exists(self.model_path):
+            self.download_model()
+
+        print(f"[GemmaService] Loading Gemma 3 1B from {self.model_path}...")
+        print(f"[GemmaService] GPU layers: {self.n_gpu_layers}, Context: {self.n_ctx}")
+
+        self._llm = Llama(
+            model_path=self.model_path,
+            n_gpu_layers=self.n_gpu_layers,
+            n_ctx=self.n_ctx,
+            n_batch=self.n_batch,
+            verbose=self.verbose
+        )
+
+        print(f"[GemmaService] Model loaded successfully")
+        return self
+
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+        stop: Optional[List[str]] = None
+    ) -> str:
+        """Generate text from prompt."""
+        if self._llm is None:
+            self.load()
+
+        response = self._llm(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stop=stop or [],
+            echo=False
+        )
+
+        return response['choices'][0]['text'].strip()
+
+    def generate_qa_pairs(self, count: int) -> List[Dict[str, str]]:
+        """Generate Q&A pairs using Gemma 3 chat template."""
+        if self._llm is None:
+            self.load()
+
+        # Gemma 3 chat template
+        prompt = f"""<start_of_turn>user
+Generate {count} unique question-answer pairs about general knowledge.
+
+REQUIREMENTS:
+- Questions: 5 to 10 words
+- Answers: 5 to 10 words (SHORT answers only!)
+- Topics: science, history, geography, arts, sports
+
+Format each pair EXACTLY as:
+Q: [question]
+A: [answer]
+
+Generate exactly {count} pairs now:<end_of_turn>
+<start_of_turn>model
+"""
+
+        response = self._llm(
+            prompt,
+            max_tokens=count * 60,  # ~60 tokens per QA pair
+            temperature=0.8,
+            stop=["<end_of_turn>", "<start_of_turn>"],
+            echo=False
+        )
+
+        content = response['choices'][0]['text']
+        return self._parse_qa_pairs(content)
+
+    def _parse_qa_pairs(self, content: str) -> List[Dict[str, str]]:
+        """Parse Q&A pairs from generated text."""
+        pairs = []
+        lines = content.split('\n')
+        current_q, current_a = None, None
+
+        for line in lines:
+            line = line.strip()
+            if line.lower().startswith('q:'):
+                current_q = line[2:].strip()
+            elif line.lower().startswith('a:'):
+                current_a = line[2:].strip()
+                if current_q and current_a:
+                    word_count_q = len(current_q.split())
+                    word_count_a = len(current_a.split())
+                    if 3 <= word_count_q <= 15 and 3 <= word_count_a <= 15:
+                        pairs.append({'question': current_q, 'answer': current_a})
+                    current_q, current_a = None, None
+
+        return pairs
+
+    def generate_batch(self, total_count: int, batch_size: int = 20) -> List[Dict[str, str]]:
+        """Generate multiple batches of Q&A pairs."""
+        all_pairs = []
+        remaining = total_count
+
+        while remaining > 0:
+            count = min(batch_size, remaining)
+            pairs = self.generate_qa_pairs(count)
+            all_pairs.extend(pairs)
+            remaining -= len(pairs)
+
+            if not pairs:
+                pairs = self.generate_qa_pairs(count)
+                all_pairs.extend(pairs)
+                remaining -= len(pairs)
+                if not pairs:
+                    break
+
+        return all_pairs[:total_count]
+
+    @staticmethod
+    def install_with_cuda():
+        """Install llama-cpp-python with CUDA support."""
+        import subprocess
+        print("[GemmaService] Installing llama-cpp-python with CUDA...")
+        cmd = 'CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python --force-reinstall --no-cache-dir'
+        return subprocess.run(cmd, shell=True).returncode == 0
+
+
+def test_gemma_service():
+    """Test the Gemma service."""
+    service = GemmaService()
+    service.load()
+
+    print("\nGenerating 5 Q&A pairs with Gemma 3 1B...")
+    pairs = service.generate_qa_pairs(5)
+
+    for i, pair in enumerate(pairs, 1):
+        print(f"\n{i}. Q: {pair['question']}")
+        print(f"   A: {pair['answer']}")
+
+    return pairs
+
+
+if __name__ == "__main__":
+    test_gemma_service()

services/matcha_service.py

@@ -0,0 +1,196 @@
+"""
+Matcha TTS Service for fast audio generation.
+
+Matcha-TTS uses conditional flow matching for fast synthesis.
+Paper: https://arxiv.org/abs/2309.03199 (ICASSP 2024)
+
+Requires:
+- pip install matcha-tts OR git clone + pip install -e .
+- apt-get install espeak-ng (for phonemization)
+
+For GPU: PyTorch with CUDA
+
+Performance on RTX 5060 Ti: ~19 calls/s, 85x real-time
+"""
+import os
+from typing import Optional, Tuple
+import numpy as np
+import torch
+
+# Fix PyTorch 2.6+ weights_only security for Matcha checkpoints
+# Monkey-patch torch.load to force weights_only=False (needed for PyTorch Lightning checkpoints)
+_original_torch_load = torch.load
+def _patched_torch_load(*args, **kwargs):
+    kwargs['weights_only'] = False
+    return _original_torch_load(*args, **kwargs)
+torch.load = _patched_torch_load
+
+
+class MatchaService:
+    """Service for TTS using Matcha-TTS (faster than Piper/diffusion-based TTS)."""
+
+    def __init__(
+        self,
+        use_cuda: bool = True,
+        gpu_id: int = 0,
+        n_timesteps: int = 5,
+        temperature: float = 0.667,
+        length_scale: float = 1.0
+    ):
+        self.use_cuda = use_cuda
+        self.gpu_id = gpu_id
+        self.n_timesteps = n_timesteps
+        self.temperature = temperature
+        self.length_scale = length_scale
+
+        self._model = None
+        self._vocoder = None
+        self._denoiser = None
+        self._device = None
+        self.sample_rate = 22050  # Matcha-TTS default
+
+    def load(self) -> 'MatchaService':
+        """Load the Matcha-TTS model and vocoder."""
+        from matcha.cli import (
+            load_matcha,
+            load_vocoder,
+            MATCHA_URLS,
+            VOCODER_URLS,
+        )
+        from matcha.utils.utils import get_user_data_dir
+
+        # Set device
+        if self.use_cuda and torch.cuda.is_available():
+            self._device = torch.device(f"cuda:{self.gpu_id}")
+        else:
+            self._device = torch.device("cpu")
+
+        print(f"[MatchaService] Loading on {self._device}...")
+
+        # Get default model paths (auto-download if needed)
+        data_dir = get_user_data_dir()
+        os.makedirs(data_dir, exist_ok=True)
+
+        # Load Matcha model (downloads automatically)
+        matcha_path = os.path.join(data_dir, "matcha_ljspeech.ckpt")
+        if not os.path.exists(matcha_path):
+            print(f"[MatchaService] Downloading Matcha model...")
+            import urllib.request
+            urllib.request.urlretrieve(MATCHA_URLS["matcha_ljspeech"], matcha_path)
+
+        # load_matcha(model_name, checkpoint_path, device)
+        self._model = load_matcha("matcha_ljspeech", matcha_path, self._device)
+
+        # Load HiFi-GAN vocoder (downloads automatically - it's a direct file, not a zip)
+        vocoder_path = os.path.join(data_dir, "hifigan_T2_v1")
+        if not os.path.exists(vocoder_path):
+            print(f"[MatchaService] Downloading vocoder...")
+            import urllib.request
+            urllib.request.urlretrieve(VOCODER_URLS["hifigan_T2_v1"], vocoder_path)
+
+        # load_vocoder(vocoder_name, checkpoint_path, device)
+        self._vocoder, self._denoiser = load_vocoder("hifigan_T2_v1", vocoder_path, self._device)
+
+        print(f"[MatchaService] Loaded successfully (sample_rate={self.sample_rate})")
+        return self
+
+    @torch.inference_mode()
+    def synthesize(self, text: str) -> Tuple[np.ndarray, int]:
+        """Synthesize text to audio.
+
+        Returns:
+            Tuple of (samples as float32 numpy array, sample_rate)
+        """
+        if self._model is None:
+            self.load()
+
+        from matcha.cli import process_text, to_waveform
+
+        # Process text to phonemes
+        # process_text(i, text, device) - i is just an index for logging
+        text_processed = process_text(0, text, self._device)
+        x = text_processed["x"]
+        x_lengths = text_processed["x_lengths"]
+
+        # Synthesize mel spectrogram
+        output = self._model.synthesise(
+            x,
+            x_lengths,
+            n_timesteps=self.n_timesteps,
+            temperature=self.temperature,
+            spks=None,
+            length_scale=self.length_scale
+        )
+
+        # Convert mel to waveform
+        waveform = to_waveform(
+            output["mel"],
+            self._vocoder,
+            self._denoiser,
+            denoiser_strength=0.00025
+        )
+
+        # Convert to numpy float32
+        samples = waveform.cpu().numpy().squeeze().astype(np.float32)
+
+        return samples, self.sample_rate
+
+    def synthesize_to_file(self, text: str, output_path: str) -> float:
+        """Synthesize text and save to WAV file.
+
+        Returns:
+            Duration in seconds
+        """
+        import wave
+
+        samples, sr = self.synthesize(text)
+
+        # Convert to int16 for WAV
+        audio_int16 = (samples * 32768).clip(-32768, 32767).astype(np.int16)
+
+        with wave.open(output_path, 'wb') as wav:
+            wav.setnchannels(1)
+            wav.setsampwidth(2)
+            wav.setframerate(sr)
+            wav.writeframes(audio_int16.tobytes())
+
+        return len(samples) / sr
+
+
+def test_matcha_service():
+    """Test the Matcha TTS service."""
+    import time
+
+    service = MatchaService(use_cuda=True, gpu_id=0)
+    service.load()
+
+    text = "Hello, this is a test of the Matcha text to speech system."
+
+    # Warmup
+    print("[Test] Warmup...")
+    for _ in range(3):
+        service.synthesize(text)
+
+    # Benchmark
+    print("[Test] Benchmarking...")
+    start = time.time()
+    n_iterations = 20
+    total_duration = 0
+    for _ in range(n_iterations):
+        samples, sr = service.synthesize(text)
+        total_duration += len(samples) / sr
+    elapsed = time.time() - start
+
+    print(f"{n_iterations} synthesize calls in {elapsed:.2f}s ({n_iterations/elapsed:.1f}/s)")
+    print(f"Average audio duration: {total_duration/n_iterations:.2f}s")
+    print(f"Real-time factor: {total_duration/elapsed:.2f}x")
+
+    # Save a test file
+    service.synthesize_to_file(text, "/tmp/matcha_test.wav")
+    print(f"Saved test audio to /tmp/matcha_test.wav")
+
+    return service
+
+
+if __name__ == "__main__":
+    test_matcha_service()

services/piper_service.py

@@ -0,0 +1,132 @@
+"""
+Piper TTS Service for fast audio generation.
+
+Requires: pip install piper-tts
+For GPU: needs onnxruntime-gpu with CUDA
+"""
+import os
+from typing import Optional, Tuple
+import numpy as np
+
+
+class PiperService:
+    """Service for TTS using Piper (faster than Kokoro)."""
+
+    DEFAULT_MODEL_PATH = "/tmp/piper/en_US-amy-medium.onnx"
+    DEFAULT_MODEL_URL = "https://huggingface.co/rhasspy/piper-voices/resolve/main/en/en_US/amy/medium/en_US-amy-medium.onnx"
+    DEFAULT_JSON_URL = "https://huggingface.co/rhasspy/piper-voices/resolve/main/en/en_US/amy/medium/en_US-amy-medium.onnx.json"
+
+    def __init__(
+        self,
+        model_path: str = DEFAULT_MODEL_PATH,
+        use_cuda: bool = True,
+        gpu_id: int = 0
+    ):
+        self.model_path = model_path
+        self.use_cuda = use_cuda
+        self.gpu_id = gpu_id
+        self._voice = None
+        self.sample_rate = 22050
+
+    def download_model(self) -> bool:
+        """Download the model if not present."""
+        import subprocess
+
+        model_dir = os.path.dirname(self.model_path)
+        os.makedirs(model_dir, exist_ok=True)
+
+        if not os.path.exists(self.model_path):
+            print(f"[PiperService] Downloading model to {self.model_path}...")
+            subprocess.run(["wget", "-q", "-O", self.model_path, self.DEFAULT_MODEL_URL])
+
+        json_path = self.model_path + ".json"
+        if not os.path.exists(json_path):
+            print(f"[PiperService] Downloading config...")
+            subprocess.run(["wget", "-q", "-O", json_path, self.DEFAULT_JSON_URL])
+
+        return os.path.exists(self.model_path)
+
+    def load(self) -> 'PiperService':
+        """Load the Piper model."""
+        # Set GPU before importing
+        if self.use_cuda:
+            os.environ['CUDA_VISIBLE_DEVICES'] = str(self.gpu_id)
+
+        from piper import PiperVoice
+
+        if not os.path.exists(self.model_path):
+            self.download_model()
+
+        print(f"[PiperService] Loading on GPU {self.gpu_id}...")
+        self._voice = PiperVoice.load(self.model_path, use_cuda=self.use_cuda)
+        self.sample_rate = self._voice.config.sample_rate
+        print(f"[PiperService] Loaded (sample_rate={self.sample_rate})")
+
+        return self
+
+    def synthesize(self, text: str) -> Tuple[np.ndarray, int]:
+        """Synthesize text to audio.
+
+        Returns:
+            Tuple of (samples as float32 numpy array, sample_rate)
+        """
+        if self._voice is None:
+            self.load()
+
+        audio_bytes = b''
+        for chunk in self._voice.synthesize(text):
+            audio_bytes += chunk.audio_int16_bytes
+
+        # Convert to float32 numpy array (same format as Kokoro)
+        samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
+
+        return samples, self.sample_rate
+
+    def synthesize_to_file(self, text: str, output_path: str) -> float:
+        """Synthesize text and save to WAV file.
+
+        Returns:
+            Duration in seconds
+        """
+        import wave
+
+        samples, sr = self.synthesize(text)
+
+        # Convert back to int16 for WAV
+        audio_int16 = (samples * 32768).astype(np.int16)
+
+        with wave.open(output_path, 'wb') as wav:
+            wav.setnchannels(1)
+            wav.setsampwidth(2)
+            wav.setframerate(sr)
+            wav.writeframes(audio_int16.tobytes())
+
+        return len(samples) / sr
+
+
+def test_piper_service():
+    """Test the Piper service."""
+    import time
+
+    service = PiperService(use_cuda=True, gpu_id=0)
+    service.load()
+
+    text = "Hello, this is a test of the Piper text to speech system."
+
+    # Warmup
+    service.synthesize(text)
+
+    # Benchmark
+    start = time.time()
+    for _ in range(10):
+        samples, sr = service.synthesize(text)
+    elapsed = time.time() - start
+
+    print(f"10 synthesize calls in {elapsed:.2f}s ({10/elapsed:.1f}/s)")
+    print(f"Audio duration: {len(samples)/sr:.2f}s")
+
+    return service
+
+
+if __name__ == "__main__":
+    test_piper_service()

services/soprano_service.py

@@ -0,0 +1,153 @@
+"""
+Soprano TTS Service for ultra-fast audio generation.
+
+Soprano is an ultra-lightweight TTS model (80M params) with batch support.
+GitHub: https://github.com/ekwek1/soprano
+
+Performance on RTX 5060 Ti: ~121 items/s (8ms/item) with batch=100
+
+Requires:
+- pip install soprano-tts
+- CUDA-enabled GPU
+"""
+from typing import Optional, Tuple, List
+import numpy as np
+
+
+class SopranoService:
+    """Service for TTS using Soprano (fastest open-source TTS)."""
+
+    def __init__(
+        self,
+        use_cuda: bool = True,
+        gpu_id: int = 0,
+        backend: str = "transformers",  # "lmdeploy" or "transformers"
+        temperature: float = 0.3,
+        top_p: float = 0.95,
+    ):
+        self.use_cuda = use_cuda
+        self.gpu_id = gpu_id
+        self.backend = backend
+        self.temperature = temperature
+        self.top_p = top_p
+
+        self._model = None
+        self.sample_rate = 32000  # Soprano outputs 32kHz
+
+    def load(self) -> 'SopranoService':
+        """Load the Soprano TTS model."""
+        import os
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id)
+
+        from soprano import SopranoTTS
+
+        device = "cuda" if self.use_cuda else "cpu"
+        print(f"[SopranoService] Loading on {device} (GPU {self.gpu_id})...")
+
+        self._model = SopranoTTS(backend=self.backend, device=device)
+
+        print(f"[SopranoService] Loaded successfully (sample_rate={self.sample_rate})")
+        return self
+
+    def synthesize(self, text: str) -> Tuple[np.ndarray, int]:
+        """Synthesize text to audio.
+
+        Returns:
+            Tuple of (samples as float32 numpy array, sample_rate)
+        """
+        if self._model is None:
+            self.load()
+
+        audio = self._model.infer(
+            text,
+            temperature=self.temperature,
+            top_p=self.top_p,
+        )
+
+        # Ensure float32 numpy array
+        if hasattr(audio, 'numpy'):
+            audio = audio.numpy()
+        samples = np.asarray(audio, dtype=np.float32)
+
+        return samples, self.sample_rate
+
+    def synthesize_batch(self, texts: List[str]) -> List[Tuple[np.ndarray, int]]:
+        """Synthesize multiple texts in batch (much faster).
+
+        Returns:
+            List of (samples, sample_rate) tuples
+        """
+        if self._model is None:
+            self.load()
+
+        audios = self._model.infer_batch(texts)
+
+        results = []
+        for audio in audios:
+            if hasattr(audio, 'numpy'):
+                audio = audio.numpy()
+            samples = np.asarray(audio, dtype=np.float32)
+            results.append((samples, self.sample_rate))
+
+        return results
+
+    def synthesize_to_file(self, text: str, output_path: str) -> float:
+        """Synthesize text and save to WAV file.
+
+        Returns:
+            Duration in seconds
+        """
+        import wave
+
+        samples, sr = self.synthesize(text)
+
+        # Convert to int16 for WAV
+        audio_int16 = (samples * 32768).clip(-32768, 32767).astype(np.int16)
+
+        with wave.open(output_path, 'wb') as wav:
+            wav.setnchannels(1)
+            wav.setsampwidth(2)
+            wav.setframerate(sr)
+            wav.writeframes(audio_int16.tobytes())
+
+        return len(samples) / sr
+
+
+def test_soprano_service():
+    """Test the Soprano TTS service."""
+    import time
+
+    service = SopranoService(use_cuda=True, gpu_id=0)
+    service.load()
+
+    # Test single
+    print("[Test] Single inference...")
+    text = "Hello, this is a test of the Soprano text to speech system."
+    samples, sr = service.synthesize(text)
+    print(f"Single: {len(samples)} samples, {len(samples)/sr:.2f}s")
+
+    # Test batch
+    print("[Test] Batch inference...")
+    texts = [
+        "What is the capital of France?",
+        "Paris is the capital of France.",
+        "How does rain form?",
+        "Water evaporates and condenses.",
+    ] * 25  # 100 items
+
+    start = time.time()
+    results = service.synthesize_batch(texts)
+    elapsed = time.time() - start
+
+    print(f"Batch of {len(texts)}: {elapsed:.2f}s ({len(texts)/elapsed:.1f}/s)")
+    print(f"Average: {elapsed/len(texts)*1000:.1f}ms/item")
+
+    # Save test file
+    service.synthesize_to_file(text, "/tmp/soprano_test.wav")
+    print(f"Saved test audio to /tmp/soprano_test.wav")
+
+    return service
+
+
+if __name__ == "__main__":
+    test_soprano_service()