Update tts.py

tts.py

@@ -1,25 +1,24 @@
 """
 TTS module - Local Orpheus TTS (finetuned) with transformers + SNAC.
-Uses finetuned model for better audio quality, with optional voice cloning
-from a reference audio file. Falls back to named voice if no reference found.
-No vLLM dependency - uses HuggingFace transformers for maximum compatibility.
+Uses named voices from the finetuned model.
+Official settings from orpheus_tts source code.
+No vLLM dependency - uses HuggingFace transformers for compatibility.
 """
 import os
 import time
 import wave
 import logging
-import subprocess
 import numpy as np
 import torch
 
 logger = logging.getLogger(__name__)
 
 TEMP_DIR = "/tmp/tts_output"
-VOICE_REF_DIR = "/app/assets"
 SAMPLE_RATE = 24000
 
-# Fallback voice if no reference audio - options: tara, leah, jess, leo, dan, mia, zac, zoe
-FALLBACK_VOICE = os.environ.get("ORPHEUS_VOICE", "leo")
+# Voice options: tara, leah, jess, leo, dan, mia, zac, zoe
+# Leo is a natural male voice
+VOICE = os.environ.get("ORPHEUS_VOICE", "leo")
 
 # Special token IDs for Orpheus
 TOKEN_START_OF_HUMAN = 128259
@@ -36,8 +35,6 @@ AUDIO_CODE_OFFSET = 128266
 _model = None
 _tokenizer = None
 _snac_model = None
-_ref_tokens = None  # Pre-encoded reference audio tokens
-_ref_transcript = ""  # Transcript of reference audio
 _initialized = False
 
 
@@ -52,42 +49,40 @@ def ensure_temp_dir():
 
 def initialize():
     """
-    Load Orpheus finetuned model, tokenizer, SNAC decoder,
-    and optionally pre-encode reference audio for voice cloning.
+    Load Orpheus finetuned model, tokenizer, and SNAC decoder.
     Call once at startup.
     """
-    global _model, _tokenizer, _snac_model, _ref_tokens, _ref_transcript, _initialized
+    global _model, _tokenizer, _snac_model, _initialized
 
     if _initialized:
         return
 
     t0 = time.time()
-    logger.info(f"[{_ts()}] [TTS] Initializing local Orpheus TTS (finetuned)...")
+    logger.info(f"[{_ts()}] [TTS] Initializing Orpheus TTS (finetuned, voice: {VOICE})...")
 
     from transformers import AutoModelForCausalLM, AutoTokenizer
     from snac import SNAC
 
     # Determine dtype based on GPU capability
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if torch.cuda.is_available():
         capability = torch.cuda.get_device_capability()
         if capability[0] >= 8:
             dtype = torch.bfloat16
-            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - using bfloat16")
+            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - bfloat16")
         else:
             dtype = torch.float16
-            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - using float16 (T4)")
+            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - float16")
     else:
         dtype = torch.float32
-        logger.info(f"[{_ts()}] [TTS] No GPU - using float32 (will be slow)")
+        logger.info(f"[{_ts()}] [TTS] No GPU - float32 (will be slow)")
 
-    # 1. Load tokenizer - finetuned model is NOT gated, no auth needed
+    # 1. Load tokenizer from finetuned model (NOT gated, no auth needed)
     t1 = time.time()
     model_name = "canopylabs/orpheus-3b-0.1-ft"
     _tokenizer = AutoTokenizer.from_pretrained(model_name)
     logger.info(f"[{_ts()}] [TTS] Tokenizer loaded: {time.time()-t1:.2f}s")
 
-    # 2. Load model
+    # 2. Load finetuned model (NOT gated)
     t2 = time.time()
     _model = AutoModelForCausalLM.from_pretrained(
         model_name,
@@ -102,142 +97,22 @@ def initialize():
     _snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval()
     if torch.cuda.is_available():
         _snac_model = _snac_model.cuda()
-    logger.info(f"[{_ts()}] [TTS] SNAC decoder loaded: {time.time()-t3:.2f}s")
-
-    # 4. Pre-encode reference audio for voice cloning
-    _ref_tokens = _load_reference_audio()
-
-    # 5. Load reference transcript if available (improves voice cloning quality)
-    _ref_transcript = ""
-    if _ref_tokens is not None:
-        for d in [VOICE_REF_DIR, "/app"]:
-            for name in ["voice_transcript.txt", "transcript.txt"]:
-                path = os.path.join(d, name)
-                if os.path.exists(path):
-                    _ref_transcript = open(path).read().strip()
-                    logger.info(f"[{_ts()}] [TTS] Reference transcript: \"{_ref_transcript[:60]}...\"")
-                    break
-            if _ref_transcript:
-                break
-        if not _ref_transcript:
-            logger.info(f"[{_ts()}] [TTS] No transcript file. Add voice_transcript.txt for better cloning.")
+    logger.info(f"[{_ts()}] [TTS] SNAC loaded: {time.time()-t3:.2f}s")
 
     _initialized = True
-    mode = "voice cloning" if _ref_tokens is not None else f"named voice ({FALLBACK_VOICE})"
-    logger.info(f"[{_ts()}] [TTS] ✓ Orpheus ready in {time.time()-t0:.2f}s | mode: {mode}")
-
-
-def _load_reference_audio():
-    """
-    Find and encode reference audio for voice cloning.
-    Looks for voice.mp4/voice.wav in assets and app root.
-    """
-    import librosa
-
-    ref_path = None
-    search_dirs = [VOICE_REF_DIR, "/app"]
-    search_names = ["voice.mp4", "voice.wav", "voice.mp3"]
-
-    for d in search_dirs:
-        for name in search_names:
-            candidate = os.path.join(d, name)
-            if os.path.exists(candidate) and os.path.getsize(candidate) > 1024:
-                ref_path = candidate
-                break
-        if ref_path:
-            break
-
-    if ref_path is None:
-        logger.info(f"[{_ts()}] [TTS] No reference audio found. Using named voice: {FALLBACK_VOICE}")
-        return None
-
-    logger.info(f"[{_ts()}] [TTS] Loading reference audio: {ref_path}")
-    ensure_temp_dir()
-
-    # Extract audio to wav if not already wav
-    wav_path = ref_path
-    if not ref_path.endswith(".wav"):
-        wav_path = os.path.join(TEMP_DIR, "voice_ref.wav")
-        try:
-            subprocess.run([
-                "ffmpeg", "-y", "-i", ref_path,
-                "-ar", str(SAMPLE_RATE), "-ac", "1", "-f", "wav", wav_path
-            ], capture_output=True, check=True)
-            logger.info(f"[{_ts()}] [TTS] Extracted audio from {os.path.basename(ref_path)}")
-        except Exception as e:
-            logger.error(f"[{_ts()}] [TTS] ffmpeg extraction failed: {e}")
-            return None
-
-    try:
-        t0 = time.time()
-        audio_array, sr = librosa.load(wav_path, sr=SAMPLE_RATE)
-
-        # Limit to 15 seconds for reference
-        max_samples = 15 * SAMPLE_RATE
-        if len(audio_array) > max_samples:
-            audio_array = audio_array[:max_samples]
-            logger.info(f"[{_ts()}] [TTS] Trimmed reference to 15s")
-
-        duration = len(audio_array) / SAMPLE_RATE
-        logger.info(f"[{_ts()}] [TTS] Reference audio: {duration:.1f}s")
-
-        waveform = torch.from_numpy(audio_array).unsqueeze(0).unsqueeze(0).float()
-        if torch.cuda.is_available():
-            waveform = waveform.cuda()
-
-        with torch.inference_mode():
-            codes = _snac_model.encode(waveform)
-
-        # Interleave SNAC codes into Orpheus token format
-        all_codes = []
-        for i in range(codes[0].shape[1]):
-            all_codes.append(codes[0][0][i].item() + AUDIO_CODE_OFFSET)
-            all_codes.append(codes[1][0][2 * i].item() + AUDIO_CODE_OFFSET + 4096)
-            all_codes.append(codes[2][0][4 * i].item() + AUDIO_CODE_OFFSET + (2 * 4096))
-            all_codes.append(codes[2][0][(4 * i) + 1].item() + AUDIO_CODE_OFFSET + (3 * 4096))
-            all_codes.append(codes[1][0][(2 * i) + 1].item() + AUDIO_CODE_OFFSET + (4 * 4096))
-            all_codes.append(codes[2][0][(4 * i) + 2].item() + AUDIO_CODE_OFFSET + (5 * 4096))
-            all_codes.append(codes[2][0][(4 * i) + 3].item() + AUDIO_CODE_OFFSET + (6 * 4096))
-
-        logger.info(f"[{_ts()}] [TTS] Reference encoded: {len(all_codes)} tokens ({time.time()-t0:.2f}s)")
-        return all_codes
-
-    except Exception as e:
-        logger.error(f"[{_ts()}] [TTS] Failed to encode reference: {e}", exc_info=True)
-        return None
+    logger.info(f"[{_ts()}] [TTS] ✓ Ready in {time.time()-t0:.2f}s (voice: {VOICE})")
 
 
 def _build_prompt_tokens(text):
     """
-    Build input tokens for Orpheus generation.
-    With voice cloning: reference audio tokens + text prompt.
-    Without: named voice format (e.g. "leo: text").
+    Build input tokens for Orpheus finetuned model.
+    Format: [SOH] {voice}: {text} [EOT EOH]
     """
-    if _ref_tokens is not None:
-        # Voice cloning: [SOH][transcript][EOT EOH SOA AUDIO_START][ref_audio][AUDIO_END EOAI][SOH][gen_text][EOT EOH]
-        start = torch.tensor([[TOKEN_START_OF_HUMAN]], dtype=torch.int64)
-        end = torch.tensor([[TOKEN_END_OF_TEXT, TOKEN_END_OF_HUMAN, TOKEN_START_OF_AI, TOKEN_AUDIO_START]], dtype=torch.int64)
-        final = torch.tensor([[TOKEN_AUDIO_END, TOKEN_END_OF_AI]], dtype=torch.int64)
-
-        ref_text = _tokenizer(_ref_transcript, return_tensors="pt").input_ids
-        ref_input = torch.cat([
-            start, ref_text, end,
-            torch.tensor([_ref_tokens], dtype=torch.int64),
-            final
-        ], dim=1)
-
-        gen_text = _tokenizer(text, return_tensors="pt").input_ids
-        gen_end = torch.tensor([[TOKEN_END_OF_TEXT, TOKEN_END_OF_HUMAN]], dtype=torch.int64)
-
-        input_ids = torch.cat([ref_input, start, gen_text, gen_end], dim=1)
-    else:
-        # Named voice format for finetuned model
-        prompt = f"{FALLBACK_VOICE}: {text}"
-        tokens = _tokenizer(prompt, return_tensors="pt")
-        start = torch.tensor([[TOKEN_START_OF_HUMAN]], dtype=torch.int64)
-        end = torch.tensor([[TOKEN_END_OF_TEXT, TOKEN_END_OF_HUMAN]], dtype=torch.int64)
-        input_ids = torch.cat([start, tokens.input_ids, end], dim=1)
-
+    prompt = f"{VOICE}: {text}"
+    tokens = _tokenizer(prompt, return_tensors="pt")
+    start = torch.tensor([[TOKEN_START_OF_HUMAN]], dtype=torch.int64)
+    end = torch.tensor([[TOKEN_END_OF_TEXT, TOKEN_END_OF_HUMAN]], dtype=torch.int64)
+    input_ids = torch.cat([start, tokens.input_ids, end], dim=1)
     return input_ids.to(_model.device)
 
 
@@ -295,8 +170,7 @@ def _decode_audio_tokens(generated_ids):
 
 def generate_audio(text: str, output_filename: str = None) -> str:
     """
-    Generate speech audio from text using local Orpheus finetuned model.
-    Returns path to wav file, or None on failure.
+    Generate speech from text. Returns path to wav file, or None on failure.
     """
     if not text or not text.strip():
         return None
@@ -324,7 +198,7 @@ def generate_audio(text: str, output_filename: str = None) -> str:
         attention_mask = torch.ones_like(input_ids)
         logger.info(f"[{_ts()}] [TTS] Prompt: {input_ids.shape[1]} tokens")
 
-        # Generate with official Orpheus settings
+        # Generate - official settings from orpheus_tts source code
         t1 = time.time()
         with torch.no_grad():
             generated_ids = _model.generate(
@@ -333,8 +207,8 @@ def generate_audio(text: str, output_filename: str = None) -> str:
                 max_new_tokens=1200,
                 do_sample=True,
                 temperature=0.6,
-                top_p=0.95,
-                repetition_penalty=1.1,
+                top_p=0.8,
+                repetition_penalty=1.3,
                 eos_token_id=TOKEN_AUDIO_END,
             )
         t2 = time.time()