Update tts.py

tts.py

@@ -1,22 +1,24 @@
 """
-TTS module - Local Orpheus TTS with transformers + SNAC.
-Supports voice cloning from a reference audio file.
+TTS module - Local Orpheus TTS (finetuned) with transformers + SNAC.
+Uses named voices (leo, tara, etc.) from the finetuned model.
 No vLLM dependency - uses HuggingFace transformers for maximum 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
 
+# Default voice - options: tara, leah, jess, leo, dan, mia, zac, zoe
+# Leo is the most natural-sounding male voice per Orpheus docs
+VOICE = os.environ.get("ORPHEUS_VOICE", "leo")
+
 # Special token IDs for Orpheus
 TOKEN_START_OF_HUMAN = 128259
 TOKEN_END_OF_TEXT = 128009
@@ -32,8 +34,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
 
 
@@ -48,17 +48,17 @@ def ensure_temp_dir():
 
 def initialize():
     """
-    Load Orpheus model, tokenizer, and SNAC decoder.
-    Also 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, _initialized
+    global _model, _tokenizer, _snac_model, _initialized
 
     if _initialized:
         return
 
     t0 = time.time()
-    logger.info(f"[{_ts()}] [TTS] Initializing local Orpheus TTS...")
+    logger.info(f"[{_ts()}] [TTS] Initializing local Orpheus TTS (finetuned)...")
+    logger.info(f"[{_ts()}] [TTS] Voice: {VOICE}")
 
     from transformers import AutoModelForCausalLM, AutoTokenizer
     from snac import SNAC
@@ -77,15 +77,10 @@ def initialize():
         dtype = torch.float32
         logger.info(f"[{_ts()}] [TTS] No GPU - using float32 (will be slow)")
 
-    # 1. Load tokenizer
+    # 1. Load tokenizer - finetuned model is NOT gated, no auth needed
     t1 = time.time()
-    model_name = "canopylabs/orpheus-3b-0.1-pretrained"
-    hf_token = os.environ.get("HF_TOKEN", None)
-    if hf_token:
-        logger.info(f"[{_ts()}] [TTS] HF_TOKEN found, will use for gated model auth")
-    else:
-        logger.warning(f"[{_ts()}] [TTS] No HF_TOKEN found! Gated model download will fail.")
-    _tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
+    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
@@ -94,7 +89,6 @@ def initialize():
         model_name,
         torch_dtype=dtype,
         device_map="auto",
-        token=hf_token,
     )
     _model.eval()
     logger.info(f"[{_ts()}] [TTS] Model loaded: {time.time()-t2:.2f}s")
@@ -106,140 +100,20 @@ def initialize():
         _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)
-    global _ref_transcript
-    _ref_transcript = ""
-    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 loaded: \"{_ref_transcript[:60]}...\"")
-                break
-        if _ref_transcript:
-            break
-    if not _ref_transcript and _ref_tokens:
-        logger.info(f"[{_ts()}] [TTS] No transcript file found. Voice cloning may be less consistent.")
-        logger.info(f"[{_ts()}] [TTS] Add voice_transcript.txt with the spoken text for better results.")
-
     _initialized = True
-    logger.info(f"[{_ts()}] [TTS] ✓ Orpheus ready in {time.time()-t0:.2f}s (voice cloning: {_ref_tokens is not None})")
-
-
-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.warning(f"[{_ts()}] [TTS] No reference audio found. Will use named voice fallback.")
-        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 (keeps prompt length manageable)
-        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] ✓ Orpheus 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.
+    Build input tokens for Orpheus finetuned model.
+    Format: [SOH] {voice}: {text} [EOT EOH]
     """
-    if _ref_tokens is not None:
-        # Voice cloning: [SOH][text][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"tara: {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)
 
 
@@ -297,7 +171,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.
+    Generate speech audio from text using local Orpheus finetuned model.
     Returns path to wav file, or None on failure.
     """
     if not text or not text.strip():
@@ -326,17 +200,17 @@ 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
+        # Generate with official Orpheus settings
         t1 = time.time()
         with torch.no_grad():
             generated_ids = _model.generate(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
-                max_new_tokens=2048,
+                max_new_tokens=1200,
                 do_sample=True,
                 temperature=0.6,
-                top_p=0.9,
-                repetition_penalty=1.2,
+                top_p=0.95,
+                repetition_penalty=1.1,
                 eos_token_id=TOKEN_AUDIO_END,
             )
         t2 = time.time()