Demo.py

import os
import random
import numpy as np
import torch

# Monkey‐patch torch.load để luôn map về CPU khi cần
_orig_torch_load = torch.load
def _torch_load_cpu(f, *args, **kwargs):
    if "map_location" not in kwargs:
        kwargs["map_location"] = torch.device("cpu")
    return _orig_torch_load(f, *args, **kwargs)
torch.load = _torch_load_cpu

from pathlib import Path
from chatterbox.src.chatterbox.tts import ChatterboxTTS
import gradio as gr

# --- CẤU HÌNH MODEL OFFLINE ---
LOCAL_MODEL_DIR = Path(os.getcwd()) / "models" / "TTS_Pro"
MODELS = {}

def get_or_load_model(device_str: str):
    """
    Load hoặc lấy lại model ChatterboxTTS trên device 'cpu' hoặc 'cuda'.
    device_str = "cpu" hoặc "gpu".
    """
    device = "cuda" if (device_str == "gpu" and torch.cuda.is_available()) else "cpu"
    if device not in MODELS:
        print(f"📂 Đang load model lên {device} …")
        if not LOCAL_MODEL_DIR.is_dir():
            raise FileNotFoundError(f"Không tìm thấy thư mục model: {LOCAL_MODEL_DIR}")
        model = ChatterboxTTS.from_local(str(LOCAL_MODEL_DIR), device)
        MODELS[device] = model
        print(f"✅ Đã load model trên {device}")
    return MODELS[device]

def set_seed(seed: int):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)

def chunk_text(text: str, chunk_size: int = 300):
    """Chia text dài thành các đoạn tối đa chunk_size kí tự, giữ nguyên từ."""
    words = text.split()
    chunks = []
    current = ""
    for w in words:
        if len(current) + len(w) + 1 > chunk_size:
            chunks.append(current.strip())
            current = w
        else:
            current += " " + w if current else w
    if current:
        chunks.append(current.strip())
    return chunks

def generate_tts_audio(
    device_choice: str,
    text_input: str,
    audio_prompt_path: str,
    exaggeration: float,
    temperature: float,
    seed_num: int,
    cfg_weight: float
):
    """
    Sinh audio từ văn bản không giới hạn bằng cách chia thành chunk nhỏ,
    generate từng chunk và ghép nối kết quả.
    Trả về (sample_rate, numpy.ndarray) toàn bộ audio.
    """
    model = get_or_load_model(device_choice)
    if seed_num != 0:
        set_seed(int(seed_num))

    chunks = chunk_text(text_input, chunk_size=300)
    waves = []
    sr = model.sr

    print(f"🔊 Đang sinh {len(chunks)} đoạn trên {model.device}")
    for idx, chunk in enumerate(chunks, start=1):
        print(f" • Đoạn {idx}/{len(chunks)}: '{chunk[:50]}...'")
        wav = model.generate(
            chunk,
            audio_prompt_path=audio_prompt_path,
            exaggeration=exaggeration,
            temperature=temperature,
            cfg_weight=cfg_weight,
        )
        waves.append(wav.squeeze(0).cpu().numpy())

    full_wave = np.concatenate(waves, axis=0)
    print("✅ Sinh audio hoàn thành.")
    return sr, full_wave

# --- GIAO DIỆN GRADIO TIẾNG VIỆT ---
with gr.Blocks(title="LTTEAM TTS") as demo:
    gr.Markdown(
        """
        # LTTEAM TTS  
        **Phát triển bởi: Lý Trần**  
        Ứng dụng chuyển văn bản thành giọng nói chất lượng cao, hỗ trợ đầu vào không giới hạn.
        """
    )
    with gr.Row():
        with gr.Column():
            device_choice = gr.Radio(
                choices=["cpu", "gpu"],
                value="gpu" if torch.cuda.is_available() else "cpu",
                label="Chọn thiết bị"
            )
            text = gr.Textbox(
                label="Văn bản (không giới hạn độ dài)",
                lines=8,
                placeholder="Dán hoặc nhập văn bản vào đây..."
            )
            ref_wav = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Âm thanh mẫu (tùy chọn)"
            )
            exaggeration = gr.Slider(
                minimum=0.25, maximum=2, step=0.05,
                value=0.5,
                label="Mức nhấn nhá (Exaggeration)"
            )
            cfg_weight = gr.Slider(
                minimum=0.2, maximum=1, step=0.05,
                value=0.5,
                label="Trọng số CFG / Tốc độ"
            )
            with gr.Accordion("Tùy chọn thêm", open=False):
                seed_num = gr.Number(0, label="Seed (0 = random)")
                temperature = gr.Slider(
                    minimum=0.05, maximum=5, step=0.05,
                    value=0.8,
                    label="Nhiệt độ (Temperature)"
                )
            run = gr.Button("Chuyển giọng", variant="primary")

        with gr.Column():
            out_audio = gr.Audio(label="Kết quả âm thanh")

    run.click(
        fn=generate_tts_audio,
        inputs=[device_choice, text, ref_wav, exaggeration, temperature, seed_num, cfg_weight],
        outputs=[out_audio],
    )

if __name__ == "__main__":
    demo.launch()