server/sync_tts_v2.py

#!/usr/bin/env python3
"""
TTS Alignment Tool v2 - Fixed sample rate issues
"""

import subprocess
import asyncio
import json
import os
import sys
import wave
import tempfile

WHISPER_CLI = "/workspace/MuseTalk1.5/vendor/whisper/build/bin/whisper-cli"
WHISPER_MODEL = "/workspace/MuseTalk1.5/vendor/whisper/models/ggml-large-v3-turbo-q5_0.bin"
SAMPLE_RATE = 22050  # Standard rate for TTS

def get_audio_duration_ms(filepath):
    result = subprocess.run([
        'ffprobe', '-v', 'error', '-show_entries', 'format=duration',
        '-of', 'default=noprint_wrappers=1:nokey=1', filepath
    ], capture_output=True, text=True)
    try:
        return int(float(result.stdout.strip()) * 1000)
    except:
        return 0

def normalize_audio(input_path, output_path):
    """Normalize audio to standard format: mono, 22050Hz, 16-bit"""
    subprocess.run([
        'ffmpeg', '-y', '-v', 'quiet', '-i', input_path,
        '-ar', str(SAMPLE_RATE), '-ac', '1', '-sample_fmt', 's16', output_path
    ], capture_output=True)

def extract_word_timestamps(audio_path):
    # Whisper needs 16kHz
    wav16k = audio_path + '.16k.wav'
    subprocess.run([
        'ffmpeg', '-y', '-v', 'quiet', '-i', audio_path, '-ar', '16000', '-ac', '1', wav16k
    ], capture_output=True)
    
    subprocess.run([
        WHISPER_CLI, '-m', WHISPER_MODEL, '-f', wav16k,
        '--output-json', '--max-len', '1', '-l', 'pt', '--no-prints'
    ], capture_output=True, text=True)
    
    json_path = wav16k + '.json'
    words = []
    if os.path.exists(json_path):
        with open(json_path) as f:
            data = json.load(f)
        os.remove(json_path)
        
        for item in data.get('transcription', []):
            text = item['text'].strip()
            if text and text not in [',', '.', '?', '!', '...', ';', ':']:
                words.append({
                    'text': text,
                    'start': item['offsets']['from'],
                    'end': item['offsets']['to']
                })
    
    if os.path.exists(wav16k):
        os.remove(wav16k)
    return words

def generate_silence_wav(output_path, duration_ms):
    if duration_ms <= 0:
        duration_ms = 1
    num_samples = int(SAMPLE_RATE * duration_ms / 1000)
    with wave.open(output_path, 'w') as f:
        f.setnchannels(1)
        f.setsampwidth(2)
        f.setframerate(SAMPLE_RATE)
        f.writeframes(b'\x00\x00' * num_samples)

def extract_segment(input_path, output_path, start_ms, end_ms):
    start_sec = max(0, start_ms / 1000)
    dur_sec = max(0.01, (end_ms - start_ms) / 1000)
    subprocess.run([
        'ffmpeg', '-y', '-v', 'quiet', '-i', input_path,
        '-ss', str(start_sec), '-t', str(dur_sec),
        '-ar', str(SAMPLE_RATE), '-ac', '1', output_path
    ], capture_output=True)

def time_stretch(input_path, output_path, ratio):
    ratio = max(0.5, min(2.0, ratio))
    if abs(ratio - 1.0) < 0.05:
        subprocess.run(['cp', input_path, output_path], capture_output=True)
    else:
        temp_out = output_path + '.tmp.wav'
        subprocess.run(['rubberband', '-t', str(ratio), input_path, temp_out], capture_output=True)
        # Normalize back to standard rate
        normalize_audio(temp_out, output_path)
        if os.path.exists(temp_out):
            os.remove(temp_out)

def concat_audio(input_files, output_path):
    existing = [f for f in input_files if os.path.exists(f) and get_audio_duration_ms(f) > 0]
    if existing:
        # Use ffmpeg for concatenation (handles different formats better)
        list_file = output_path + '.list'
        with open(list_file, 'w') as f:
            for path in existing:
                f.write(f"file '{path}'\n")
        
        subprocess.run([
            'ffmpeg', '-y', '-v', 'quiet', '-f', 'concat', '-safe', '0',
            '-i', list_file, '-ar', str(SAMPLE_RATE), '-ac', '1', output_path
        ], capture_output=True)
        
        os.remove(list_file)

def generate_espeak(text, output_path):
    temp_out = output_path + '.tmp.wav'
    subprocess.run([
        "espeak-ng", "-v", "pt-br", "-w", temp_out, text
    ], capture_output=True)
    normalize_audio(temp_out, output_path)
    if os.path.exists(temp_out):
        os.remove(temp_out)

async def generate_elevenlabs(text, output_path):
    import httpx
    ELEVENLABS_API_KEY = "sk_857e9e6f2412ddf3ff5334b736e4b571641d26225c0d8d62"
    ELEVENLABS_VOICE_ID = "21m00Tcm4TlvDq8ikWAM"
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(
            f"https://api.elevenlabs.io/v1/text-to-speech/{ELEVENLABS_VOICE_ID}/stream",
            headers={"xi-api-key": ELEVENLABS_API_KEY, "Content-Type": "application/json"},
            json={"text": text, "model_id": "eleven_flash_v2_5", "voice_settings": {"stability": 0.5, "similarity_boost": 0.75}}
        )
        temp_mp3 = output_path + '.mp3'
        with open(temp_mp3, "wb") as f:
            f.write(response.content)
        normalize_audio(temp_mp3, output_path)
        os.remove(temp_mp3)

async def align_espeak_to_target(text, output_path, verbose=True):
    if verbose:
        print(f"\n{'='*60}")
        print(f"ALINHAMENTO: \"{text[:40]}{'...' if len(text) > 40 else ''}\"")
        print('='*60)
    
    with tempfile.TemporaryDirectory() as tmpdir:
        eleven_wav = os.path.join(tmpdir, "eleven.wav")
        espeak_wav = os.path.join(tmpdir, "espeak.wav")
        
        if verbose: print("\n1. Gerando áudios...")
        await generate_elevenlabs(text, eleven_wav)
        generate_espeak(text, espeak_wav)
        
        eleven_dur = get_audio_duration_ms(eleven_wav)
        espeak_dur = get_audio_duration_ms(espeak_wav)
        if verbose:
            print(f"   ElevenLabs: {eleven_dur}ms, espeak: {espeak_dur}ms")
        
        if verbose: print("\n2. Extraindo timestamps...")
        eleven_words = extract_word_timestamps(eleven_wav)
        espeak_words = extract_word_timestamps(espeak_wav)
        
        if not eleven_words or not espeak_words:
            if verbose: print("   ERRO: Timestamps não encontrados")
            return None
        
        if verbose:
            print(f"   ElevenLabs: {len(eleven_words)} palavras")
            print(f"   espeak: {len(espeak_words)} palavras")
        
        if verbose: print("\n3. Alinhando...")
        segments = []
        
        # Leading silence
        if eleven_words[0]['start'] > 0:
            sil = os.path.join(tmpdir, "sil0.wav")
            generate_silence_wav(sil, eleven_words[0]['start'])
            segments.append(sil)
        
        min_words = min(len(eleven_words), len(espeak_words))
        
        for i in range(min_words):
            ew = eleven_words[i]
            sw = espeak_words[i]
            
            target_dur = ew['end'] - ew['start']
            source_dur = sw['end'] - sw['start']
            
            if source_dur < 10 or target_dur < 10:
                continue
            
            ratio = target_dur / source_dur
            
            seg_in = os.path.join(tmpdir, f"w{i}_in.wav")
            seg_out = os.path.join(tmpdir, f"w{i}_out.wav")
            
            extract_segment(espeak_wav, seg_in, sw['start'], sw['end'])
            
            if os.path.exists(seg_in) and get_audio_duration_ms(seg_in) > 0:
                time_stretch(seg_in, seg_out, ratio)
                if os.path.exists(seg_out) and get_audio_duration_ms(seg_out) > 0:
                    segments.append(seg_out)
                    if verbose:
                        print(f"   {sw['text']:<10} {source_dur:>3}ms → {target_dur:>3}ms (x{ratio:.2f})")
            
            if i < min_words - 1:
                next_ew = eleven_words[i + 1]
                gap = next_ew['start'] - ew['end']
                if gap > 5:
                    gap_path = os.path.join(tmpdir, f"gap{i}.wav")
                    generate_silence_wav(gap_path, gap)
                    segments.append(gap_path)
        
        if eleven_words:
            trailing = eleven_dur - eleven_words[-1]['end']
            if trailing > 10:
                sil = os.path.join(tmpdir, "sil_end.wav")
                generate_silence_wav(sil, trailing)
                segments.append(sil)
        
        if verbose: print("\n4. Concatenando...")
        if segments:
            concat_audio(segments, output_path)
            
            if os.path.exists(output_path):
                final_dur = get_audio_duration_ms(output_path)
                diff = abs(eleven_dur - final_dur)
                
                if verbose:
                    print(f"\n   Target:    {eleven_dur}ms")
                    print(f"   Resultado: {final_dur}ms")
                    print(f"   Diferença: {diff}ms ({diff/eleven_dur*100:.1f}%)")
                
                return output_path
    
    return None

async def verify_alignment(text):
    aligned_path = "/tmp/espeak_aligned.wav"
    eleven_path = "/tmp/eleven_verify.wav"
    
    result = await align_espeak_to_target(text, aligned_path)
    
    if not result or not os.path.exists(aligned_path):
        print("\nERRO: Falha no alinhamento")
        return
    
    await generate_elevenlabs(text, eleven_path)
    
    print("\n" + "="*60)
    print("VERIFICAÇÃO")
    print("="*60)
    
    eleven_words = extract_word_timestamps(eleven_path)
    aligned_words = extract_word_timestamps(aligned_path)
    
    print(f"\n{'PALAVRA':<10} {'TARGET':<12} {'ALINHADO':<12} {'DIFF':>6}")
    print("-"*45)
    
    total_diff = 0
    count = 0
    
    for i in range(min(len(eleven_words), len(aligned_words))):
        ew = eleven_words[i]
        aw = aligned_words[i]
        diff = abs(ew['start'] - aw['start'])
        total_diff += diff
        count += 1
        print(f"{ew['text']:<10} {ew['start']:>4}-{ew['end']:<4}ms  {aw['start']:>4}-{aw['end']:<4}ms  {diff:>4}ms")
    
    if count > 0:
        avg_diff = total_diff / count
        print("-"*45)
        print(f"MÉDIA: {avg_diff:.0f}ms")
        
        if avg_diff < 30: print("\n✓ EXCELENTE!")
        elif avg_diff < 50: print("\n✓ MUITO BOM")
        elif avg_diff < 100: print("\n✓ BOM")
        else: print("\n⚠ ACEITÁVEL")

if __name__ == "__main__":
    text = sys.argv[1] if len(sys.argv) > 1 else "Olá, tudo bem? Como posso te ajudar?"
    asyncio.run(verify_alignment(text))