app.py

import os
import re
import time
import numpy as np
import soundfile as sf
import matplotlib.pyplot as plt
import librosa
import gradio as gr
from scipy.signal import fftconvolve
from model import StyleTTModel

SPEAKER_WAV_PATH = "speakers/example_female.wav"
OUTPUT_FILENAME = "output.wav"
SAMPLE_RATE = 24000

# Global model variable
model = None

def initialize_model():
    """Initialize the StyleTTS model with error handling"""
    global model
    try:
        # Check if speaker reference file exists
        if not os.path.exists(SPEAKER_WAV_PATH):
            raise FileNotFoundError(f"Không tìm thấy file giọng nói tham chiếu tại: {SPEAKER_WAV_PATH}. "
                                    "Vui lòng tạo thư mục và đặt file .wav của bạn vào đó.")

        print("Bắt đầu khởi tạo StyleTTS2 Model...")
        model = StyleTTModel(speaker_wav=SPEAKER_WAV_PATH)
        print("Đang tải model StyleTTS2. Quá trình này có thể mất vài phút...")
        start_time = time.time()
        model.load()
        end_time = time.time()
        print(f"Model đã được tải thành công sau {end_time - start_time:.2f} giây.")
        return True
    except Exception as e:
        print(f"Lỗi khi khởi tạo model: {e}")
        model = None
        return False

# Initialize model on startup
model_loaded = initialize_model()

# ---------------------------
# Load HF TTS model (hexgrad/styletts2)
# ---------------------------
SR_OUT = 24000
# tts_pipe = pipeline("text-to-speech", model="hexgrad/styletts2")

# ---------------------------
# Audio helpers
# ---------------------------
def load_wav(path, sr_target=SR_OUT):
    wav, sr = sf.read(path)
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    if sr != sr_target:
        wav = librosa.resample(wav.astype(np.float32), orig_sr=sr, target_sr=sr_target)
        sr = sr_target
    return wav.astype(np.float32), sr

def apply_reverb(wav, ir_path):
    """Apply reverb effect using impulse response"""
    try:
        if not os.path.exists(ir_path):
            print(f"Cảnh báo: Không tìm thấy file impulse response: {ir_path}")
            return wav
        ir, _ = load_wav(ir_path, sr_target=SR_OUT)
        return fftconvolve(wav, ir, mode="full")
    except Exception as e:
        print(f"Lỗi khi áp dụng reverb: {e}")
        return wav

def add_noise(wav, noise_path, snr_db=10):
    """Add background noise to audio"""
    try:
        if not os.path.exists(noise_path):
            print(f"Cảnh báo: Không tìm thấy file noise: {noise_path}")
            return wav
        noise, _ = load_wav(noise_path, sr_target=SR_OUT)
        if len(noise) < len(wav):
            noise = np.tile(noise, int(len(wav)/len(noise)) + 1)
        noise = noise[:len(wav)]
        sig_power = np.mean(wav**2)
        noise_power = np.mean(noise**2)
        if noise_power == 0:
            return wav
        scale = np.sqrt(sig_power / (10**(snr_db/10) * noise_power))
        return wav + noise * scale
    except Exception as e:
        print(f"Lỗi khi thêm noise: {e}")
        return wav

def bandlimit_phone(wav, sr=SR_OUT):
    """Apply phone-like band limiting"""
    try:
        return librosa.effects.preemphasis(wav)
    except Exception as e:
        print(f"Lỗi khi áp dụng band limiting: {e}")
        return wav

def plot_waveforms(clean, processed, sr=SR_OUT):
    """Create waveform comparison plot"""
    try:
        fig, axes = plt.subplots(2, 1, figsize=(10, 4), sharex=True)
        t_clean = np.arange(len(clean)) / sr
        t_proc = np.arange(len(processed)) / sr
        
        axes[0].plot(t_clean, clean, color="blue", linewidth=0.8)
        axes[0].set_title("🎤 Waveform gốc (StyleTTS2)")
        axes[0].set_ylabel("Amplitude")
        axes[0].grid(True, alpha=0.3)
        
        axes[1].plot(t_proc, processed, color="red", linewidth=0.8)
        axes[1].set_title("🎵 Waveform có hiệu ứng môi trường")
        axes[1].set_xlabel("Thời gian (s)")
        axes[1].set_ylabel("Amplitude")
        axes[1].grid(True, alpha=0.3)
        
        fig.tight_layout()
        return fig
    except Exception as e:
        print(f"Lỗi khi tạo biểu đồ: {e}")
        # Return a simple error plot
        fig, ax = plt.subplots(1, 1, figsize=(10, 2))
        ax.text(0.5, 0.5, "Không thể tạo biểu đồ", ha='center', va='center', transform=ax.transAxes)
        ax.set_title("Lỗi tạo biểu đồ")
        return fig

# ---------------------------
# Tag list
# ---------------------------
TAG_LIST = {
    "laugh": "😆 Cười thoải mái",
    "whisper": "🤫 Thì thầm",
    "naughty": "😏 Tinh nghịch",
    "giggle": "😂 Cười rúc rích",
    "tease": "😉 Trêu chọc",
    "smirk": "😼 Đắc ý",
    "surprise": "😲 Ngạc nhiên",
    "shock": "😱 Hoảng hốt",
    "romantic": "❤️ Lãng mạn",
    "shy": "🫣 Bẽn lẽn",
    "excited": "🤩 Phấn khích",
    "curious": "🧐 Tò mò",
    "discover": "✨ Phát hiện",
    "blush": "🌸 Ngượng ngùng",
    "angry": "😡 Giận dữ",
    "sad": "😢 Buồn",
    "happy": "😊 Vui vẻ",
    "fear": "😨 Sợ hãi",
    "confident": "😎 Tự tin",
    "serious": "😐 Nghiêm túc",
    "tired": "🥱 Mệt mỏi",
    "cry": "😭 Khóc",
    "love": "😍 Yêu thương",
    "disgust": "🤢 Ghê tởm",
}
TAG_PATTERN = r"(<\/?(?:" + "|".join(TAG_LIST.keys()) + ")>)"

# ---------------------------
# Core synthesis
# ---------------------------
def synthesize(text, env, snr_db=10, speed=1.0):
    """Synthesize text to speech with environment effects"""
    try:
        # Check if model is loaded
        if model is None:
            print("Lỗi: Model chưa được tải. Vui lòng khởi động lại ứng dụng.")
            return None, None, None
            
        # Parse text and extract segments
        tokens = re.split(TAG_PATTERN, text)
        clean_segments = []

        for tok in tokens:
            if not tok or tok.isspace():
                continue
            if tok.startswith("<") and tok.endswith(">"):
                # Skip tags for now - they're just for text segmentation
                continue
            else:
                # Synthesize each text segment
                try:
                    audio_array = model.synthesize(tok, speed=speed)
                    clean_segments.append(audio_array)
                except Exception as e:
                    print(f"Lỗi khi tổng hợp đoạn '{tok}': {e}")
                    continue

        if not clean_segments:
            return None, None, None

        # Concatenate all audio segments
        clean_audio = np.concatenate(clean_segments, axis=0)
        processed = clean_audio.copy()

        # Apply environment effects
        try:
            if env == "Church":
                processed = apply_reverb(processed, "ir_church.wav")
            elif env == "Hall":
                processed = apply_reverb(processed, "ir_hall.wav")
            elif env == "Cafe":
                processed = add_noise(processed, "noise_cafe.wav", snr_db=snr_db)
            elif env == "Street":
                processed = add_noise(processed, "noise_street.wav", snr_db=snr_db)
            elif env == "Office":
                processed = add_noise(processed, "noise_office.wav", snr_db=snr_db)
            elif env == "Supermarket":
                processed = add_noise(processed, "noise_supermarket.wav", snr_db=snr_db)
            elif env == "Phone":
                processed = bandlimit_phone(processed, sr=SR_OUT)
        except Exception as e:
            print(f"Cảnh báo: Không thể áp dụng hiệu ứng môi trường '{env}': {e}")
            # Continue with clean audio if environment effects fail

        # Create waveform comparison plot
        fig = plot_waveforms(clean_audio, processed, sr=SR_OUT)
        
        return (SR_OUT, processed), fig, (SR_OUT, clean_audio)
        
    except Exception as e:
        print(f"Lỗi trong quá trình tổng hợp: {e}")
        return None, None, None

# ---------------------------
# Examples
# ---------------------------
EXAMPLES = [
    "Xin chào <whisper> tôi nói nhỏ </whisper> rồi <laugh> bật cười </laugh>.",
    "Tôi cảm thấy <happy> vui </happy> nhưng cũng <sad> buồn </sad>.",
    "Khi <surprise> bất ngờ </surprise> tôi <shock> hoảng hốt </shock>.",
]

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks(
    title="StyleTTS2 Text-to-Speech", 
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        max-width: 1200px !important;
        margin: auto !important;
    }
    .main-header {
        text-align: center;
        margin-bottom: 20px;
    }
    .status-box {
        padding: 10px;
        border-radius: 8px;
        margin: 10px 0;
    }
    .status-success {
        background-color: #d4edda;
        border: 1px solid #c3e6cb;
        color: #155724;
    }
    .status-error {
        background-color: #f8d7da;
        border: 1px solid #f5c6cb;
        color: #721c24;
    }
    """
) as demo:
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🎙️ StyleTTS2 Text-to-Speech với Hiệu ứng Môi trường</h1>
    </div>
    """)
    
    # Model status indicator
    if model_loaded:
        gr.HTML("""
        <div class="status-box status-success">
            <strong>✅ Model đã sẵn sàng</strong> - Bạn có thể bắt đầu tạo giọng nói!
        </div>
        """)
    else:
        gr.HTML("""
        <div class="status-box status-error">
            <strong>❌ Lỗi tải model</strong> - Vui lòng kiểm tra file giọng nói tham chiếu và khởi động lại.
        </div>
        """)
    
    gr.Markdown("""
    ### 📖 Hướng dẫn sử dụng
    - **Nhập văn bản** vào ô text và sử dụng các tags để tạo cảm xúc
    - **Chọn môi trường** để áp dụng hiệu ứng âm thanh phù hợp
    - **Điều chỉnh tốc độ** và mức độ nhiễu theo ý muốn
    - **Nhấn "Tạo giọng nói"** để tạo audio với hiệu ứng
    """)

    with gr.Accordion("📑 Danh sách Tags + Emoji", open=False):
        md = "| Tag | Ý nghĩa |\n|-----|----------|\n"
        for k, v in TAG_LIST.items():
            md += f"| `<{k}>...</{k}>` | {v} |\n"
        gr.Markdown(md)

    # Main content area
    with gr.Row(equal_height=True):
        # Left column - Input controls
        with gr.Column(scale=1, min_width=400):
            with gr.Group():
                gr.Markdown("### ⚙️ Cài đặt đầu vào")
                
                text_in = gr.Textbox(
                    value=EXAMPLES[0], 
                    label="📝 Văn bản cần chuyển đổi", 
                    lines=4,
                    placeholder="Nhập văn bản của bạn ở đây. Sử dụng tags để tạo cảm xúc...",
                    max_lines=6
                )
                
                # Environment and Speed in one row
                with gr.Row():
                    env_in = gr.Dropdown(
                        choices=["Neutral", "Church", "Hall", "Cafe", "Street", "Phone", "Office", "Supermarket"],
                        value="Neutral", 
                        label="🌍 Môi trường âm thanh",
                        info="Chọn môi trường để áp dụng hiệu ứng"
                    )
                    speed_slider = gr.Slider(
                        minimum=0.5,
                        maximum=2.0,
                        value=1.0,
                        step=0.1,
                        label="⚡ Tốc độ nói",
                        info="1.0 = bình thường, < 1.0 = chậm, > 1.0 = nhanh"
                    )
                
                # Noise level slider
                snr_slider = gr.Slider(
                    0, 30, 
                    value=10, 
                    step=1, 
                    label="🔊 Mức độ nhiễu (SNR dB)",
                    info="Chỉ áp dụng cho môi trường có tiếng ồn. Cao hơn = ít nhiễu hơn"
                )
                
                # Generate button
                btn = gr.Button("🎵 Tạo giọng nói", variant="primary", size="lg", scale=1)
            
            # Examples section
            with gr.Group():
                gr.Examples(
                    examples=[[ex] for ex in EXAMPLES], 
                    inputs=[text_in], 
                    label="💡 Ví dụ nhanh"
                )
            
        # Right column - Output results
        with gr.Column(scale=1, min_width=400):
            with gr.Group():
                gr.Markdown("### 🎧 Kết quả")
                
                # Audio outputs
                audio_out = gr.Audio(
                    label="🎵 Âm thanh có hiệu ứng", 
                    type="numpy",
                    info="Phiên bản có áp dụng hiệu ứng môi trường"
                )
                clean_out = gr.Audio(
                    label="🎤 Âm thanh gốc", 
                    type="numpy",
                    info="Phiên bản gốc không có hiệu ứng"
                )
            
            # Waveform plot
            with gr.Group():
                wave_plot = gr.Plot(
                    label="📊 So sánh dạng sóng",
                    info="Biểu đồ so sánh âm thanh gốc và có hiệu ứng",
                    show_label=True
                )

    btn.click(fn=synthesize,
              inputs=[text_in, env_in, snr_slider, speed_slider],
              outputs=[audio_out, wave_plot, clean_out])

    # Footer with additional information
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            ---
            ### ℹ️ Thông tin bổ sung
            
            **🎯 Môi trường âm thanh:**
            - **Neutral**: Không có hiệu ứng
            - **Church/Hall**: Thêm reverb (tiếng vang)
            - **Cafe/Street/Office/Supermarket**: Thêm tiếng ồn nền
            - **Phone**: Giới hạn tần số như điện thoại
            
            **📊 SNR (Signal-to-Noise Ratio):**
            - Giá trị cao = ít nhiễu hơn
            - Giá trị thấp = nhiều nhiễu hơn
            - Chỉ áp dụng cho môi trường có tiếng ồn
            
            **⚡ Tốc độ nói:**
            - 0.5x = Nói chậm một nửa
            - 1.0x = Tốc độ bình thường
            - 2.0x = Nói nhanh gấp đôi
            """)

# Launch the application
if __name__ == "__main__":
    try:
        print("🚀 Đang khởi động ứng dụng StyleTTS2...")
        print("📱 Giao diện sẽ mở tại: http://localhost:7860")
        print("🔄 Để dừng ứng dụng, nhấn Ctrl+C")
        print("-" * 50)
        
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            show_error=True,
            show_tips=True,
            enable_queue=True,
            max_threads=10
        )
    except KeyboardInterrupt:
        print("\n👋 Ứng dụng đã được dừng bởi người dùng.")
    except Exception as e:
        print(f"❌ Lỗi khi khởi động ứng dụng: {e}")
        print("💡 Vui lòng kiểm tra lại:")
        print("   - File giọng nói tham chiếu có tồn tại không")
        print("   - Cổng 7860 có bị chiếm dụng không")
        print("   - Các thư viện đã được cài đặt đầy đủ chưa")