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 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 _run_edge_tts(text: str, voice: str, out_wav_path: str):
    """Synthesize high-quality natural neural voiceover directly via module or CLI."""
    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}")

    tmp_mp3 = out_wav_path.replace(".wav", "_temp.mp3")
    
    if edge_tts is not None:
        async def _speak():
            communicate = edge_tts.Communicate(text, voice)
            await communicate.save(tmp_mp3)
        asyncio.run(_speak())
    else:
        # Fallback to sys.executable -m edge_tts
        cmd = [sys.executable, "-m", "edge_tts", "--voice", voice, "--text", 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, "--text", 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}")

    subprocess.run([
        "ffmpeg", "-y", "-i", tmp_mp3,
        "-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-MX-JorgeNeural"}
        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
        }