Update tts.py

tts.py

@@ -1,24 +1,25 @@
 """
-TTS module - Local Orpheus TTS (finetuned) with transformers + SNAC.
-Uses named voices from the finetuned model.
-Official settings from orpheus_tts source code.
-No vLLM dependency - uses HuggingFace transformers for compatibility.
+TTS module - Local Orpheus TTS with transformers + SNAC.
+Supports voice cloning from reference audio + transcript (pretrained model).
+Falls back to named voice with finetuned model if no reference available.
+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
 
-# Voice options: tara, leah, jess, leo, dan, mia, zac, zoe
-# Leo is a natural male voice
-VOICE = os.environ.get("ORPHEUS_VOICE", "leo")
+# Fallback voice for finetuned model (if no reference audio)
+FALLBACK_VOICE = os.environ.get("ORPHEUS_VOICE", "leo")
 
 # Special token IDs for Orpheus
 TOKEN_START_OF_HUMAN = 128259
@@ -35,6 +36,9 @@ AUDIO_CODE_OFFSET = 128266
 _model = None
 _tokenizer = None
 _snac_model = None
+_ref_tokens = None       # Pre-encoded reference audio SNAC tokens
+_ref_transcript = None   # Transcript of reference audio (REQUIRED for voice cloning)
+_use_voice_cloning = False
 _initialized = False
 
 
@@ -49,77 +53,233 @@ def ensure_temp_dir():
 
 def initialize():
     """
-    Load Orpheus finetuned model, tokenizer, and SNAC decoder.
+    Load Orpheus model, tokenizer, and SNAC decoder.
+    If reference audio + transcript found, uses pretrained model for voice cloning.
+    Otherwise, uses finetuned model with named voices.
     Call once at startup.
     """
-    global _model, _tokenizer, _snac_model, _initialized
+    global _model, _tokenizer, _snac_model, _ref_tokens, _ref_transcript
+    global _use_voice_cloning, _initialized
 
     if _initialized:
         return
 
     t0 = time.time()
-    logger.info(f"[{_ts()}] [TTS] Initializing Orpheus TTS (finetuned, voice: {VOICE})...")
+    logger.info(f"[{_ts()}] [TTS] Initializing local Orpheus TTS...")
 
     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]} - bfloat16")
+            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - using bfloat16")
         else:
             dtype = torch.float16
-            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - float16")
+            logger.info(f"[{_ts()}] [TTS] GPU compute {capability[0]}.{capability[1]} - using float16 (T4)")
     else:
         dtype = torch.float32
-        logger.info(f"[{_ts()}] [TTS] No GPU - float32 (will be slow)")
+        logger.info(f"[{_ts()}] [TTS] No GPU - using float32 (will be slow)")
 
-    # 1. Load tokenizer from finetuned model (NOT gated, no auth needed)
+    # 1. Load SNAC decoder first (needed for reference audio encoding)
     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 finetuned model (NOT gated)
+    _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()-t1:.2f}s")
+
+    # 2. Check for reference audio + transcript
+    _ref_transcript = _load_transcript()
+    ref_audio_path = _find_reference_audio()
+
+    if ref_audio_path and _ref_transcript:
+        # Voice cloning mode: use pretrained model
+        _use_voice_cloning = True
+        model_name = "canopylabs/orpheus-3b-0.1-pretrained"
+        hf_token = os.environ.get("HF_TOKEN", None)
+        if not hf_token:
+            logger.warning(f"[{_ts()}] [TTS] No HF_TOKEN! Pretrained model is gated, falling back to finetuned.")
+            _use_voice_cloning = False
+            model_name = "canopylabs/orpheus-3b-0.1-ft"
+        else:
+            logger.info(f"[{_ts()}] [TTS] Voice cloning mode: ref audio + transcript found")
+            logger.info(f"[{_ts()}] [TTS] Transcript: \"{_ref_transcript[:80]}...\"")
+    else:
+        # Named voice mode: use finetuned model
+        _use_voice_cloning = False
+        model_name = "canopylabs/orpheus-3b-0.1-ft"
+        if not ref_audio_path:
+            logger.info(f"[{_ts()}] [TTS] No reference audio found")
+        if not _ref_transcript:
+            logger.info(f"[{_ts()}] [TTS] No transcript found (voice_transcript.txt)")
+        logger.info(f"[{_ts()}] [TTS] Using finetuned model with voice: {FALLBACK_VOICE}")
+
+    # 3. Load tokenizer
     t2 = time.time()
+    token_kwargs = {}
+    if _use_voice_cloning:
+        hf_token = os.environ.get("HF_TOKEN")
+        token_kwargs["token"] = hf_token
+    _tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
+    logger.info(f"[{_ts()}] [TTS] Tokenizer loaded: {time.time()-t2:.2f}s")
+
+    # 4. Load model
+    t3 = time.time()
     _model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=dtype,
         device_map="auto",
+        **token_kwargs,
     )
     _model.eval()
-    logger.info(f"[{_ts()}] [TTS] Model loaded: {time.time()-t2:.2f}s")
+    logger.info(f"[{_ts()}] [TTS] Model loaded ({model_name.split('/')[-1]}): {time.time()-t3:.2f}s")
 
-    # 3. Load SNAC decoder
-    t3 = time.time()
-    _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 loaded: {time.time()-t3:.2f}s")
+    # 5. Encode reference audio if voice cloning
+    if _use_voice_cloning and ref_audio_path:
+        _ref_tokens = _encode_reference_audio(ref_audio_path)
+        if _ref_tokens is None:
+            logger.warning(f"[{_ts()}] [TTS] Reference encoding failed, falling back to named voice")
+            _use_voice_cloning = False
 
     _initialized = True
-    logger.info(f"[{_ts()}] [TTS] ✓ Ready in {time.time()-t0:.2f}s (voice: {VOICE})")
+    mode = f"voice cloning ({len(_ref_tokens)} ref tokens)" if _use_voice_cloning else f"named voice ({FALLBACK_VOICE})"
+    logger.info(f"[{_ts()}] [TTS] ✓ Orpheus ready in {time.time()-t0:.2f}s | mode: {mode}")
+
+
+def _load_transcript():
+    """Load transcript of reference audio from voice_transcript.txt."""
+    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):
+                text = open(path).read().strip()
+                if text:
+                    return text
+    return None
+
+
+def _find_reference_audio():
+    """Find reference audio file."""
+    for d in [VOICE_REF_DIR, "/app"]:
+        for name in ["voice.mp4", "voice.wav", "voice.mp3"]:
+            path = os.path.join(d, name)
+            if os.path.exists(path) and os.path.getsize(path) > 1024:
+                return path
+    return None
+
+
+def _encode_reference_audio(ref_path):
+    """Encode reference audio to SNAC tokens for voice cloning."""
+    import librosa
+
+    logger.info(f"[{_ts()}] [TTS] Encoding reference: {ref_path}")
+    ensure_temp_dir()
+
+    # Extract to wav if needed
+    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)
+        except Exception as e:
+            logger.error(f"[{_ts()}] [TTS] ffmpeg failed: {e}")
+            return None
+
+    try:
+        t0 = time.time()
+        audio_array, sr = librosa.load(wav_path, sr=SAMPLE_RATE)
+
+        # Limit to 15 seconds
+        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")
+
+        # Encode with SNAC (matching official Colab code exactly)
+        waveform = torch.from_numpy(audio_array).unsqueeze(0)
+        waveform = waveform.to(dtype=torch.float32)
+        waveform = waveform.unsqueeze(0)
+
+        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
 
 
 def _build_prompt_tokens(text):
     """
-    Build input tokens for Orpheus finetuned model.
-    Format: [SOH] {voice}: {text} [EOT EOH]
+    Build input tokens for Orpheus generation.
+
+    Voice cloning (pretrained model) - matches official Colab exactly:
+        [SOH] [transcript_tokens] [EOT EOH SOA AUDIO_START] [ref_audio_tokens] [AUDIO_END EOAI]
+        [SOH] [text_tokens] [EOT EOH]
+
+    Named voice (finetuned model):
+        [SOH] {voice}: {text} [EOT EOH]
     """
-    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)
+    if _use_voice_cloning and _ref_tokens is not None:
+        # Voice cloning: pretrained model format (from official Colab)
+        start_tokens = torch.tensor([[TOKEN_START_OF_HUMAN]], dtype=torch.int64)
+        end_tokens = torch.tensor([[TOKEN_END_OF_TEXT, TOKEN_END_OF_HUMAN, TOKEN_START_OF_AI, TOKEN_AUDIO_START]], dtype=torch.int64)
+        final_tokens = torch.tensor([[TOKEN_AUDIO_END, TOKEN_END_OF_AI]], dtype=torch.int64)
+
+        # Transcript tokens (REQUIRED - this is what we were missing before)
+        transcript_ids = _tokenizer(_ref_transcript, return_tensors="pt").input_ids
+
+        # Build reference conditioning: [SOH][transcript][EOT EOH SOA AUDIO_START][audio][AUDIO_END EOAI]
+        ref_input = torch.cat([
+            start_tokens,
+            transcript_ids,
+            end_tokens,
+            torch.tensor([_ref_tokens], dtype=torch.int64),
+            final_tokens
+        ], dim=1)
+
+        # Generation prompt: [SOH][text][EOT EOH]
+        gen_ids = _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_tokens, gen_ids, gen_end], dim=1)
+    else:
+        # Named voice: finetuned model format
+        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)
+
     return input_ids.to(_model.device)
 
 
 def _decode_audio_tokens(generated_ids):
     """
     Extract audio tokens from generated output and decode with SNAC.
-    Returns numpy audio array at 24kHz.
+    Matches official Colab decode logic exactly.
     """
     # Find last AUDIO_START token, crop after it
     indices = (generated_ids == TOKEN_AUDIO_START).nonzero(as_tuple=True)
@@ -170,7 +330,8 @@ def _decode_audio_tokens(generated_ids):
 
 def generate_audio(text: str, output_filename: str = None) -> str:
     """
-    Generate speech from text. Returns path to wav file, or None on failure.
+    Generate speech audio from text using local Orpheus.
+    Returns path to wav file, or None on failure.
     """
     if not text or not text.strip():
         return None
@@ -198,18 +359,32 @@ 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 - official settings from orpheus_tts source code
+        # Generate with settings matched to mode
         t1 = time.time()
+        if _use_voice_cloning:
+            # Official pretrained model settings from Colab
+            gen_kwargs = dict(
+                max_new_tokens=990,
+                temperature=0.5,
+                top_p=0.9,
+                repetition_penalty=1.1,
+            )
+        else:
+            # Official finetuned model settings
+            gen_kwargs = dict(
+                max_new_tokens=1200,
+                temperature=0.6,
+                top_p=0.95,
+                repetition_penalty=1.1,
+            )
+
         with torch.no_grad():
             generated_ids = _model.generate(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
-                max_new_tokens=1200,
                 do_sample=True,
-                temperature=0.6,
-                top_p=0.8,
-                repetition_penalty=1.3,
                 eos_token_id=TOKEN_AUDIO_END,
+                **gen_kwargs,
             )
         t2 = time.time()
         new_tokens = generated_ids.shape[1] - input_ids.shape[1]