LLMVideoPipeline/pipeline/tts_gen.py

"""
Voiceover Synthesis Module using Kokoro TTS on IntelLLM (port 8003).
Generates crystal-clear narration in English and Spanish with accurate per-segment timestamps.
"""

import os
import sys
import re
import shutil
import subprocess
import logging
import wave
import contextlib
from typing import Dict, Any, List, Tuple
from .remote import RemotePC

logger = logging.getLogger("YTFactory.TTS")

def get_wav_duration(path: str) -> float:
    """Accurately compute WAV audio duration in seconds."""
    try:
        with contextlib.closing(wave.open(path, 'r')) as f:
            frames = f.getnframes()
            rate = f.getframerate()
            return frames / float(rate)
    except Exception:
        # Fallback to ffprobe
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-show_entries", "format=duration",
             "-of", "default=noprint_wrappers=1:nokey=1", path],
            capture_output=True, text=True
        )
        try:
            return float(r.stdout.strip())
        except ValueError:
            return 0.0

import asyncio
try:
    import edge_tts
except ImportError:
    edge_tts = None

def clean_and_format_spanish_for_tts(text: str) -> str:
    """
    Format Spanish narration text to sound 100% natural when spoken:
    - Expands historical abbreviations (a.C., d.C., siglos, etc.)
    - Formats numbers, units, and punctuation for natural human breathing pauses.
    """
    t = text.strip()
    # Normalize historical dates and eras
    t = re.sub(r'\b(?:a\.\s*C\.|a\s*C|A\.\s*C\.|A\s*C)\b', 'antes de Cristo', t)
    t = re.sub(r'\b(?:d\.\s*C\.|d\s*C|D\.\s*C\.|D\s*C)\b', 'después de Cristo', t)
    t = re.sub(r'\bsiglo\s+XXI\b', 'siglo veintiuno', t, flags=re.IGNORECASE)
    t = re.sub(r'\bsiglo\s+XX\b', 'siglo veinte', t, flags=re.IGNORECASE)
    t = re.sub(r'\bsiglo\s+XIX\b', 'siglo diecinueve', t, flags=re.IGNORECASE)
    t = re.sub(r'\bsiglo\s+XVIII\b', 'siglo dieciocho', t, flags=re.IGNORECASE)
    t = re.sub(r'\bsiglo\s+XV\b', 'siglo quince', t, flags=re.IGNORECASE)
    t = re.sub(r'\bEE\.?\s*UU\.?\b', 'Estados Unidos', t)
    t = re.sub(r'\bkm/h\b', 'kilómetros por hora', t)
    t = re.sub(r'(\d+)\s*km\b', r'\1 kilómetros', t)
    t = re.sub(r'(\d+)\s*m\b', r'\1 metros', t)
    t = re.sub(r'(\d+)\s*ha\b', r'\1 hectáreas', t)
    
    # Soften ellipses and replace multiple hyphens with clean commas for natural pausing
    t = re.sub(r'\.{3,}', ', ', t)
    t = re.sub(r'\s*--+\s*', ', ', t)
    t = re.sub(r'\s+', ' ', t)
    return t

def _run_edge_tts(text: str, voice: str, out_wav_path: str, rate: str = "-4%"):
    """Synthesize high-quality natural neural voiceover directly via module or CLI with studio mastering."""
    global edge_tts
    if edge_tts is None:
        try:
            import edge_tts
        except ImportError:
            try:
                subprocess.run([sys.executable, "-m", "pip", "install", "edge-tts"], check=True, capture_output=True)
                import edge_tts
            except Exception as e:
                logger.warning(f"[TTS] Dynamic edge-tts install failed: {e}")

    # Preprocess Spanish text for maximum naturalness
    if voice.startswith("es-") or "Neural" in voice:
        spoken_text = clean_and_format_spanish_for_tts(text)
    else:
        spoken_text = text

    tmp_mp3 = out_wav_path.replace(".wav", "_temp.mp3")
    
    if edge_tts is not None:
        async def _speak():
            communicate = edge_tts.Communicate(spoken_text, voice, rate=rate)
            await communicate.save(tmp_mp3)
        asyncio.run(_speak())
    else:
        # Fallback to sys.executable -m edge_tts
        cmd = [sys.executable, "-m", "edge_tts", "--voice", voice, "--rate", rate, "--text", spoken_text, "--write-media", tmp_mp3]
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            edge_bin = shutil.which("edge-tts") or os.path.expanduser("~/.local/bin/edge-tts") or os.path.expanduser("~/miniconda3/bin/edge-tts")
            if edge_bin and os.path.exists(edge_bin):
                r2 = subprocess.run([edge_bin, "--voice", voice, "--rate", rate, "--text", spoken_text, "--write-media", tmp_mp3], capture_output=True, text=True)
                if r2.returncode != 0:
                    raise RuntimeError(f"edge-tts failed: {r.stderr or r2.stderr}")
            else:
                raise RuntimeError(f"edge-tts execution error: {r.stderr}")

    # Broadcast vocal mastering filter (warm low-mid EQ, subtle presence, transparent leveling)
    vocal_af = (
        "highpass=f=75,"
        "equalizer=f=200:t=q:w=1.2:g=1.8,"
        "equalizer=f=3500:t=q:w=1.5:g=1.0,"
        "compand=attacks=0.03:decays=0.15:points=-80/-80|-40/-30|-20/-14|0/-2:gain=2,"
        "volume=1.05"
    )

    subprocess.run([
        "ffmpeg", "-y", "-i", tmp_mp3,
        "-af", vocal_af,
        "-ar", "44100", "-ac", "2", "-c:a", "pcm_s16le", out_wav_path
    ], check=True, capture_output=True)
    if os.path.exists(tmp_mp3):
        os.remove(tmp_mp3)

class TTSGen:
    def __init__(self, pcs: Dict[str, RemotePC], voices: Dict[str, str], tts_port: int = 8003):
        self.pcs = pcs
        self.voices = voices or {"en": "af_heart", "es": "es-US-AlonsoNeural"}
        self.tts_port = tts_port

    def generate(self, script_data: Dict[str, Any], job_dir: str, lang: str = "en") -> Dict[str, Any]:
        """Generate voiceover audio for all segments in the specified language."""
        segments = script_data.get("script_segments", [])
        text_key = "text" if lang == "en" else "text_es"
        default_voice = "af_heart" if lang == "en" else "es-MX-JorgeNeural"
        voice = self.voices.get(lang, default_voice)

        # Ensure Spanish uses native neural voice if an old Kokoro voice was configured
        if lang == "es" and voice in ("ef_dora", "em_alex", "ef_clara", "em_santa"):
            voice = "es-MX-JorgeNeural"

        logger.info(f"[TTS:{lang.upper()}] Synthesizing narration with voice '{voice}'...")

        abs_job_dir = os.path.abspath(job_dir)
        os.makedirs(abs_job_dir, exist_ok=True)

        segment_wavs = []
        segment_timings = []
        current_time = 0.0

        is_neural = voice.startswith(("es-", "en-", "fr-", "de-", "it-", "pt-"))

        for i, seg in enumerate(segments):
            text = seg.get(text_key, "").strip()
            if not text:
                text = seg.get("text", f"Segment {i+1}")

            seg_path = os.path.join(abs_job_dir, f"tts_{lang}_{i:03d}.wav")

            if is_neural:
                # Use native neural speech engine
                try:
                    _run_edge_tts(text, voice, seg_path)
                except Exception as e:
                    logger.warning(f"[TTS:{lang.upper()}] Neural TTS failed: {e}, attempting Kokoro fallback...")
                    is_neural = False

            if not is_neural:
                # Use Kokoro on IntelLLM
                pc = self.pcs.get("intelllm")
                if not pc or not pc.check(self.tts_port):
                    raise RuntimeError(f"IntelLLM Kokoro TTS not reachable on port {self.tts_port}")

                payload = {
                    "input": text,
                    "voice": "af_heart" if lang == "en" else "ef_dora",
                    "speed": 1.0,
                    "lang": "e" if lang == "es" else "a"
                }

                wav_data = None
                for attempt in range(1, 4):
                    wav_data = pc.post_binary(self.tts_port, "/v1/audio/speech", payload, timeout=120)
                    if wav_data and len(wav_data) >= 1024:
                        break
                    logger.warning(f"[TTS:{lang.upper()}] Segment {i+1} attempt {attempt} failed, retrying...")

                if not wav_data or len(wav_data) < 1024:
                    raise RuntimeError(f"Failed to generate TTS audio for segment {i+1} ({lang})")

                with open(seg_path, "wb") as f:
                    f.write(wav_data)

            dur = get_wav_duration(seg_path)
            chapter_title = seg.get("chapter_title" if lang == "en" else "chapter_title_es", f"Part {i+1}")
            
            segment_timings.append({
                "index": i,
                "start": current_time,
                "end": current_time + dur,
                "duration": dur,
                "chapter_title": chapter_title,
                "text": text,
                "wav_file": seg_path
            })
            current_time += dur
            segment_wavs.append(seg_path)
            logger.info(f"[TTS:{lang.upper()}] Segment {i+1}/{len(segments)} OK ({dur:.2f}s)")

        # Create concat list for ffmpeg
        concat_txt = os.path.join(abs_job_dir, f"concat_tts_{lang}.txt")
        with open(concat_txt, "w") as f:
            for sw in segment_wavs:
                f.write(f"file '{sw}'\n")

        master_wav = os.path.join(abs_job_dir, f"voiceover_{lang}.wav")
        res = subprocess.run([
            "ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", concat_txt,
            "-ar", "44100", "-ac", "2", "-c:a", "pcm_s16le", master_wav
        ], capture_output=True, text=True)

        if res.returncode != 0:
            logger.error(f"[TTS:{lang.upper()}] FFmpeg concat failed: {res.stderr}")
            raise RuntimeError(f"FFmpeg audio concatenation failed: {res.stderr}")

        if os.path.exists(concat_txt):
            os.remove(concat_txt)

        total_dur = get_wav_duration(master_wav)
        logger.info(f"[TTS:{lang.upper()}] Master voiceover generated: {master_wav} ({total_dur:.2f}s)")

        return {
            "master_wav": master_wav,
            "total_duration": total_dur,
            "segments": segment_timings,
            "voice": voice,
            "language": lang
        }