app.py

import spaces 
import os
os.environ['SPACES_ZERO_GPU'] = '1'

import gradio as gr
import soundfile as sf
import tempfile
import torch
import librosa
from tts_engine import VoiceEngine
import time

# --- 1. SETUP MODEL (Sử dụng repo cá nhân của bạn) ---
device = "cuda" if torch.cuda.is_available() else "cpu"

# THAY THẾ 'YOUR_USERNAME' bằng tên Hugging Face của bạn
MY_BACKBONE_REPO = "ktvoice/Backbone" 
MY_CODEC_REPO = "ktvoice/Codec"

try:
    tts = VoiceEngine(
        backbone_repo=MY_BACKBONE_REPO,
        backbone_device=device,
        codec_repo=MY_CODEC_REPO,
        codec_device=device
    )
except Exception as e:
    print(f"⚠️ Lỗi khởi tạo: {e}")
    class MockTTS:
        def encode_reference(self, path): return None
        def infer(self, text, ref, ref_text): 
            time.sleep(1.2) 
            import numpy as np
            return np.random.uniform(-0.1, 0.1, 24000*2)
    tts = MockTTS()

# --- 2. DATA (Giữ nguyên danh sách giọng mẫu cục bộ) ---
VOICE_SAMPLES = {
    "Tuyên (nam miền Bắc)": {"audio": "./sample/Tuyên (nam miền Bắc).wav", "text": "./sample/Tuyên (nam miền Bắc).txt"},
    "Thiện Tâm": {"audio": "./sample/thientam.mp3", "text": "./sample/thientam.txt"},
    "Vĩnh (nam miền Nam)": {"audio": "./sample/Vĩnh (nam miền Nam).wav", "text": "./sample/Vĩnh (nam miền Nam).txt"},
    "Bình (nam miền Bắc)": {"audio": "./sample/Bình (nam miền Bắc).wav", "text": "./sample/Bình (nam miền Bắc).txt"},
    "Nguyên (nam miền Nam)": {"audio": "./sample/Nguyên (nam miền Nam).wav", "text": "./sample/Nguyên (nam miền Nam).txt"},
    "Sơn (nam miền Nam)": {"audio": "./sample/Sơn (nam miền Nam).wav", "text": "./sample/Sơn (nam miền Nam).txt"},
    "Đoan (nữ miền Nam)": {"audio": "./sample/Đoan (nữ miền Nam).wav", "text": "./sample/Đoan (nữ miền Nam).txt"},
    "Ngọc (nữ miền Bắc)": {"audio": "./sample/Ngọc (nữ miền Bắc).wav", "text": "./sample/Ngọc (nữ miền Bắc).txt"},
    "Ly (nữ miền Bắc)": {"audio": "./sample/Ly (nữ miền Bắc).wav", "text": "./sample/Ly (nữ miền Bắc).txt"},
    "Dung (nữ miền Nam)": {"audio": "./sample/Dung (nữ miền Nam).wav", "text": "./sample/Dung (nữ miền Nam).txt"}
}

# --- 3. HELPER FUNCTIONS ---
def load_reference_info(voice_choice):
    if voice_choice in VOICE_SAMPLES:
        audio_path = VOICE_SAMPLES[voice_choice]["audio"]
        text_path = VOICE_SAMPLES[voice_choice]["text"]
        if os.path.exists(text_path):
            with open(text_path, "r", encoding="utf-8") as f:
                ref_text = f.read()
            return audio_path, ref_text
    return None, ""

@spaces.GPU(duration=120) 
def synthesize_speech(text, voice_choice, custom_audio, custom_text, mode_tab, pause_level, speed_value):
    try:
        if not text or text.strip() == "":
            return None, "⚠️ Vui lòng nhập nội dung!"
        
        # Tiền xử lý văn bản để tăng độ nghỉ
        processed_text = text
        if pause_level == "Trung bình":
            processed_text = processed_text.replace(",", ", , ").replace(".", ". . ")
        elif pause_level == "Dài":
            processed_text = processed_text.replace(",", ", , , ").replace(".", ". . . . ")
            
        if len(processed_text) > 400:
             processed_text = processed_text[:400]

        # Lấy Reference Data
        if mode_tab == "custom_mode": 
            if custom_audio is None or not custom_text:
                return None, "⚠️ Thiếu Audio mẫu hoặc Text mẫu."
            ref_audio_path = custom_audio
            ref_text_raw = custom_text
        else:
            ref_audio_path = VOICE_SAMPLES[voice_choice]["audio"]
            with open(VOICE_SAMPLES[voice_choice]["text"], "r", encoding="utf-8") as f:
                ref_text_raw = f.read()

        # Inference
        start_time = time.time()
        ref_codes = tts.encode_reference(ref_audio_path)
        wav = tts.infer(processed_text, ref_codes, ref_text_raw)
        
        # Điều chỉnh tốc độ (Time-stretching)
        if speed_value != 1.0:
            wav = librosa.effects.time_stretch(wav, rate=float(speed_value))
            
        process_time = time.time() - start_time
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
            sf.write(tmp_file.name, wav, 24000)
            output_path = tmp_file.name
        
        return output_path, f"⚡ Thành công: {process_time:.2f}s | Tốc độ: {speed_value}x"
    except Exception as e:
        return None, f"❌ Lỗi: {str(e)}"

# --- 4. THEME & CSS (Deep Night Pro) ---
theme = gr.themes.Default(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont('Inter'), 'sans-serif'],
).set(
    body_background_fill="#020617",
    block_background_fill="#0f172a",
    block_border_width="1px",
    input_background_fill="#1e293b",
    input_border_color="#334155",
    button_primary_background_fill="linear-gradient(135deg, #4f46e5 0%, #7c3aed 100%)",
)

css = """

.main-wrap { max-width: 1200px !important; margin: auto !important; padding: 20px !important; }

.st-card { 

    border-radius: 16px !important; 

    border: 1px solid rgba(255,255,255,0.1) !important; 

    box-shadow: 0 4px 20px rgba(0,0,0,0.5) !important;

    padding: 15px;

}

.result-card {

    background: linear-gradient(180deg, rgba(15, 23, 42, 0.8) 0%, rgba(30, 41, 59, 0.8) 100%) !important;

    border: 1px solid rgba(99, 102, 241, 0.2) !important;

    margin-top: 15px;

}

audio { filter: invert(90%) hue-rotate(180deg) brightness(1.5); width: 100%; border-radius: 8px; }

.footer { text-align: center; margin-top: 40px; color: #475569; font-size: 0.8rem; font-weight: 500; }

"""

# --- 5. UI CONSTRUCTION ---
with gr.Blocks(title="AI Voice Studio") as demo:
    
    with gr.Column(elem_classes="main-wrap"):
        with gr.Row(equal_height=True):
            # TRÁI: Nhập văn bản
            with gr.Column(scale=1):
                with gr.Group(elem_classes="st-card"):
                    text_input = gr.Textbox(
                        label="VĂN BẢN CẦN CHUYỂN ĐỔI",
                        placeholder="Chào mừng bạn. Hãy nhập nội dung vào đây...",
                        lines=20, 
                        show_label=True,
                    )
                    char_count = gr.HTML("<div style='text-align: right; color: #6366f1; font-size: 0.85rem; font-weight: bold; padding: 5px;'>0 / 250</div>")

            # PHẢI: Cấu hình
            with gr.Column(scale=1):
                with gr.Tabs() as tabs:
                    with gr.TabItem("👤 Nghệ sĩ đọc", id="preset_mode"):
                        voice_select = gr.Dropdown(
                            choices=list(VOICE_SAMPLES.keys()),
                            value="Tuyên (nam miền Bắc)",
                            label="Chọn giọng đọc",
                        )
                        with gr.Accordion("Nghe thử giọng mẫu", open=False):
                            ref_audio_preview = gr.Audio(interactive=False, show_label=False)
                            ref_text_preview = gr.Markdown("...")

                    with gr.TabItem("🎙️ Nhân bản (Clone)", id="custom_mode"):
                        custom_audio = gr.Audio(label="Audio gốc", type="filepath")
                        custom_text = gr.Textbox(label="NỘI DUNG AUDIO MẪU", lines=4)

                # Cấu hình âm thanh chuyên nghiệp
                with gr.Row():
                    pause_level = gr.Radio(choices=["Mặc định", "Trung bình", "Dài"], value="Mặc định", label="Độ ngắt nghỉ", scale=1)
                    speed_select = gr.Dropdown(choices=[0.8, 0.9, 1.0, 1.1, 1.2, 1.5], value=1.0, label="Tốc độ đọc", scale=1)

                current_mode = gr.State(value="preset_mode")
                gr.Markdown("<br>")
                btn_generate = gr.Button("BẮT ĐẦU TỔNG HỢP", variant="primary", size="lg")
                
                with gr.Group(elem_classes="st-card result-card"):
                    audio_output = gr.Audio(label="KẾT QUẢ", interactive=False, autoplay=True)
                    status_output = gr.Markdown("<p style='text-align: center; color: #818cf8;'>✨ Sẵn sàng</p>")

        gr.HTML("<div class='footer'>ENGINE BY VIENEU-TTS • PROFESSIONAL AI SOLUTIONS 2025</div>")

    # LOGIC
    text_input.change(lambda t: f"<div style='text-align: right; color: {'#6366f1' if len(t)<=250 else '#f43f5e'}; font-size: 0.85rem; font-weight: bold; padding: 5px;'>{len(t)} / 250</div>", text_input, char_count)
    voice_select.change(update_ref_preview, voice_select, [ref_audio_preview, ref_text_preview])
    tabs.children[0].select(fn=lambda: "preset_mode", outputs=current_mode)
    tabs.children[1].select(fn=lambda: "custom_mode", outputs=current_mode)
    btn_generate.click(fn=synthesize_speech, inputs=[text_input, voice_select, custom_audio, custom_text, current_mode, pause_level, speed_select], outputs=[audio_output, status_output])

if __name__ == "__main__":
    demo.queue().launch(theme=theme, css=css, server_name="0.0.0.0", server_port=7860)