matcha_inference_standalone/api_server.py

import os
import sys
import re
import tempfile
import subprocess
import hashlib
import shutil
import asyncio
from concurrent.futures import ThreadPoolExecutor

import torch
import soundfile as sf
import numpy as np
import uvicorn
from fastapi import FastAPI, HTTPException, Body, BackgroundTasks
from fastapi.responses import FileResponse

# ─── Cấu hình ────────────────────────────────────────────────────────────────
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, "Matcha-TTS"))

CHECKPOINT_PATH = os.path.join(current_dir, "model", "checkpoint_016_int8.pt")
VOCODER_PATH    = os.path.join(current_dir, "model", "generator_v1")
CACHE_DIR       = os.path.join(current_dir, "cache_1x")
SAMPLE_RATE     = 22050

# Số worker = số core CPU muốn dùng song song (mặc định: dùng hết tất cả core)
NUM_WORKERS = int(os.environ.get("TTS_WORKERS", os.cpu_count() or 4))

os.makedirs(CACHE_DIR, exist_ok=True)

# ─── Import model (lazy, chỉ import khi cần) ─────────────────────────────────
from matcha.hifigan.config import v1
from matcha.hifigan.env import AttrDict
from matcha.hifigan.models import Generator as HiFiGAN
from matcha.models.matcha_tts import MatchaTTS
from matcha.text import text_to_sequence
from matcha.utils.utils import intersperse

app = FastAPI(title="Matcha-TTS API – Per-Core Worker Pool")

# ─── Trạng thái toàn cục ─────────────────────────────────────────────────────
_matcha_model  = None
_vocoder_model = None
_executor      = None   # ThreadPoolExecutor
_device        = None

# ─── Khởi động server ────────────────────────────────────────────────────────
@app.on_event("startup")
def startup():
    global _matcha_model, _vocoder_model, _executor, _device

    # Mỗi thread PyTorch chỉ dùng đúng 1 CPU core
    torch.set_num_threads(1)

    _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    mode = "GPU" if _device.type == "cuda" else f"CPU ({NUM_WORKERS} core song song)"
    print(f"[+] Chế độ: {mode}")

    # Tải mô hình một lần, chia sẻ giữa các worker (weights only = read-only)
    print(f"[!] Đang tải Matcha-TTS: {CHECKPOINT_PATH}")
    # Đọc lớp PyTorch 2.6 an toàn
    original_load = torch.load
    def patched_load(*args, **kwargs):
        kwargs['weights_only'] = False
        return original_load(*args, **kwargs)
    torch.load = patched_load

    ckpt_or_model = torch.load(CHECKPOINT_PATH, map_location=_device)
    
    # Hỗ trợ tự động nhận diện cả Model cũ (Float32) và Model nén (INT8)
    if isinstance(ckpt_or_model, dict) and "state_dict" in ckpt_or_model:
        _matcha_model = MatchaTTS(**ckpt_or_model["hyper_parameters"])
        _matcha_model.load_state_dict(ckpt_or_model["state_dict"])
    else:
        _matcha_model = ckpt_or_model
        
    _matcha_model = _matcha_model.to(_device).eval()

    print("[!] Đang tải HiFi-GAN Vocoder...")
    h = AttrDict(v1)
    _vocoder_model = HiFiGAN(h).to(_device)
    _vocoder_model.load_state_dict(
        torch.load(VOCODER_PATH, map_location=_device)["generator"]
    )
    _vocoder_model.eval()
    _vocoder_model.remove_weight_norm()

    # Warm-up Matcha-TTS & Vocoder
    print("[!] Đang chạy warm-up cho mô hình TTS...", flush=True)
    try:
        dummy_text = "Khởi động."
        x = torch.tensor(
            intersperse(text_to_sequence(dummy_text, ["basic_cleaners_vi_female"])[0], 0),
            dtype=torch.long, device=_device,
        )[None]
        x_lengths = torch.tensor([x.shape[-1]], dtype=torch.long, device=_device)
        with torch.inference_mode():
            out = _matcha_model.synthesise(
                x, x_lengths, n_timesteps=2,
                temperature=0.5, spks=None, length_scale=1.0
            )
            _ = _vocoder_model(out["mel"]).clamp(-1, 1).squeeze().cpu().numpy()
        print("[✓] Warm-up mô hình TTS thành công!", flush=True)
    except Exception as e:
        print(f"[⚠️] Lỗi chạy warm-up mô hình TTS: {e}", flush=True)

    # Tạo thread pool: mỗi thread = 1 CPU core worker
    _executor = ThreadPoolExecutor(
        max_workers=NUM_WORKERS,
        thread_name_prefix="tts-worker"
    )
    print(f"[✓] Server sẵn sàng! {NUM_WORKERS} worker(s), mỗi worker = 1 CPU core")

# ─── Utility ─────────────────────────────────────────────────────────────────
def cleanup_file(path: str):
    try:
        if os.path.exists(path):
            os.remove(path)
    except Exception:
        pass

def build_ffmpeg_filter(speed: float, volume: float) -> str:
    filters = []
    if volume != 1.0:
        filters.append(f"volume={volume}")
    rem = speed
    while rem > 2.0:
        filters.append("atempo=2.0")
        rem /= 2.0
    while rem < 0.5:
        filters.append("atempo=0.5")
        rem /= 0.5
    if abs(rem - 1.0) > 0.01:
        filters.append(f"atempo={rem}")
    return ",".join(filters) if filters else "anull"

def split_text_into_chunks(text: str, max_words: int = 85) -> list[tuple[str, str]]:
    """Tách văn bản thành các chunk dựa trên các đoạn văn (xuống dòng). 
    Giúp giọng đọc trong một đoạn văn (gồm nhiều câu) được liền mạch, trôi chảy và tự nhiên nhất.
    Nếu một đoạn văn quá dài (> max_words), mới chia nhỏ theo dấu chấm câu.
    """
    text = text.strip()
    if not text:
        return []

    raw_paragraphs = [p.strip() for p in text.split('\n') if p.strip()]
    sentence_boundary = re.compile(r'([.?!]+|\.\.\.)(?=\s|$)')
    final_chunks = []

    for paragraph in raw_paragraphs:
        words = paragraph.split()
        if len(words) <= max_words:
            final_chunks.append((paragraph, '.'))
            continue
        
        # Nếu đoạn văn quá dài, tách theo dấu câu kết thúc câu (. ? !)
        pos = 0
        current_sentence = ""
        for match in sentence_boundary.finditer(paragraph):
            end_punc = match.group(1)
            start_idx, end_idx = match.span()
            chunk_content = paragraph[pos:start_idx].strip()
            
            if chunk_content:
                if current_sentence:
                    current_sentence += " " + chunk_content
                else:
                    current_sentence = chunk_content
            
            if current_sentence:
                final_chunks.append((current_sentence + end_punc, '.'))
                current_sentence = ""
            pos = end_idx
        
        remaining = paragraph[pos:].strip()
        if remaining:
            if current_sentence:
                current_sentence += " " + remaining
            else:
                current_sentence = remaining
        if current_sentence:
            final_chunks.append((current_sentence, '.'))
            
    return [(c.strip(), p) for c, p in final_chunks if c.strip()]

def apply_fades(audio: np.ndarray, sample_rate: int = 22050, fade_duration: float = 0.05) -> np.ndarray:
    """Áp dụng fade-in/out tuyến tính ở đầu và cuối audio để triệt tiêu tiếng click/gập biên độ."""
    audio = audio.copy()
    fade_samples = int(fade_duration * sample_rate)
    n_samples = len(audio)
    actual_fade = min(fade_samples, n_samples // 2)
    if actual_fade <= 0:
        return audio
    
    # Fade in
    fade_in = np.linspace(0.0, 1.0, actual_fade)
    audio[:actual_fade] *= fade_in
    
    # Fade out
    fade_out = np.linspace(1.0, 0.0, actual_fade)
    audio[-actual_fade:] *= fade_out
    
    return audio

def concatenate_with_pause_and_smooth_fades(
    chunks: list[np.ndarray], 
    sample_rate: int = 22050, 
    fade_duration: float = 0.15,   # 150ms fade
    silence_duration: float = 0.20  # 200ms natural pause gap
) -> np.ndarray:
    """
    Nối các chunk bằng cách vuốt Smoothstep nhanh về 0 (150ms), chèn khoảng nghỉ ngắn (200ms),
    và vuốt Smoothstep lên từ 0 (150ms). Giúp giữ giọng bình bình ở cuối câu rồi nghỉ tự nhiên.
    """
    if not chunks:
        return np.zeros(0, dtype=np.float32)
    
    fade_samples = int(sample_rate * fade_duration)
    silence_samples = int(sample_rate * silence_duration)
    
    processed_chunks = []
    for i, chunk in enumerate(chunks):
        chunk = chunk.copy()
        n_samples = len(chunk)
        
        # Fade in ở đầu chunk (từ câu thứ 2 trở đi)
        if i > 0 and fade_samples > 0:
            actual_fade = min(fade_samples, n_samples // 2)
            if actual_fade > 0:
                t = np.linspace(0.0, 1.0, actual_fade)
                s = 3 * (t ** 2) - 2 * (t ** 3)
                chunk[:actual_fade] *= s
                
        # Fade out ở cuối chunk (đến câu kế cuối)
        if i < len(chunks) - 1 and fade_samples > 0:
            actual_fade = min(fade_samples, n_samples // 2)
            if actual_fade > 0:
                t = np.linspace(0.0, 1.0, actual_fade)
                s = 3 * (t ** 2) - 2 * (t ** 3)
                chunk[-actual_fade:] *= (1.0 - s)
                
        processed_chunks.append(chunk)
        
        # Chèn khoảng nghỉ ở giữa
        if i < len(chunks) - 1 and silence_samples > 0:
            processed_chunks.append(np.zeros(silence_samples, dtype=np.float32))
            
    return np.concatenate(processed_chunks)

# ─── Công việc inference (chạy trong thread pool, 1 core / request) ──────────
def _do_inference(text: str, cache_path: str, n_timesteps: int, length_scale: float, temperature: float) -> np.ndarray:
    """Chạy Matcha-TTS + HiFi-GAN cho 1 chunk, trả về waveform và ghi vào cache_path."""
    x = torch.tensor(
        intersperse(text_to_sequence(text, ["basic_cleaners_vi_female"])[0], 0),
        dtype=torch.long, device=_device,
    )[None]
    x_lengths = torch.tensor([x.shape[-1]], dtype=torch.long, device=_device)

    with torch.inference_mode():
        out = _matcha_model.synthesise(
            x, x_lengths, n_timesteps=n_timesteps,
            temperature=temperature, spks=None, length_scale=length_scale
        )
        audio = _vocoder_model(out["mel"]).clamp(-1, 1).squeeze().cpu().numpy()

    # Chuẩn hóa âm lượng cho chunk
    max_val = np.max(np.abs(audio))
    if max_val > 0:
        audio = audio / max_val * 0.95
        
    sf.write(cache_path, audio, SAMPLE_RATE)
    return audio

def _do_ffmpeg(input_path: str, speed: float, volume: float) -> str:
    """Chạy FFmpeg để đổi tốc độ/âm lượng, trả về đường dẫn file tạm."""
    fd, out_path = tempfile.mkstemp(suffix=".wav")
    os.close(fd)
    f = build_ffmpeg_filter(speed, volume)
    subprocess.run(
        ["ffmpeg", "-y", "-i", input_path, "-filter:a", f, out_path],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True
    )
    return out_path

# ─── Endpoint chính ───────────────────────────────────────────────────────────
@app.post("/synthesize")
@app.post("/v1/audio/speech")
async def synthesize(
    background_tasks: BackgroundTasks,
    text: str        = Body(None, embed=True),
    input: str       = Body(None, embed=True),
    speed: float     = Body(1.0, embed=True),
    volume: float    = Body(1.0, embed=True),
    steps: int       = Body(10, embed=True),  # Mặc định: 10 steps
    n_timesteps: int = Body(None, embed=True),
    temperature: float = Body(0.50, embed=True),
    length_scale: float = Body(None, embed=True),
    bypass_cache: bool = Body(False, embed=True),
):
    actual_text = text or input
    if not actual_text or not actual_text.strip():
        raise HTTPException(400, "Văn bản không được để trống")

    actual_steps = n_timesteps if n_timesteps is not None else steps
    loop = asyncio.get_event_loop()

    # Luôn sinh mô hình ở tốc độ 1.0x để tối ưu hoá tỷ lệ trúng cache của các chunk
    # Sau đó sẽ dùng FFmpeg để điều chỉnh tốc độ đọc của tệp âm thanh hoàn chỉnh (bao gồm cả khoảng nghỉ)
    if length_scale is not None:
        model_length_scale = length_scale
        ffmpeg_speed = speed / (1.0 / model_length_scale) if speed else 1.0
    else:
        model_length_scale = 1.0
        ffmpeg_speed = speed if speed else 1.0

    # Kiểm tra cache toàn cục cho câu đầy đủ ở tốc độ gốc 1.0x (ls = model_length_scale)
    full_hash = hashlib.sha256(
        f"{actual_text}_steps_{actual_steps}_ls_{model_length_scale:.4f}_temp_{temperature}".encode()
    ).hexdigest()
    full_cached_path = os.path.join(CACHE_DIR, f"full_{full_hash}.wav")

    if not bypass_cache and os.path.exists(full_cached_path):
        print(f"[✓] Cache hit (full 1.0x hash: {full_hash[:8]}) – bỏ qua GPU/CPU")
        cached_path = full_cached_path
    else:
        # Tách văn bản thành các chunk nhỏ tự nhiên
        chunks = split_text_into_chunks(actual_text, max_words=85)
        if not chunks:
            raise HTTPException(400, "Văn bản không hợp lệ")

        print(f"[~] Bắt đầu xử lý {len(chunks)} chunk(s) với {actual_steps} steps (temp: {temperature}, ls: {model_length_scale:.4f})...")
        
        # Xử lý tuần tự từng chunk (caching riêng cho từng chunk ở tốc độ 1.0x)
        chunk_audios = []
        for i, (chunk_text, punc) in enumerate(chunks):
            chunk_hash = hashlib.sha256(
                f"{chunk_text}_steps_{actual_steps}_ls_{model_length_scale:.4f}_temp_{temperature}".encode()
            ).hexdigest()
            chunk_cache_path = os.path.join(CACHE_DIR, f"chunk_{chunk_hash}.wav")

            if not bypass_cache and os.path.exists(chunk_cache_path):
                print(f"  [✓] Chunk {i+1}/{len(chunks)} cache hit (hash: {chunk_hash[:8]})")
                audio, _ = await loop.run_in_executor(None, sf.read, chunk_cache_path)
                chunk_audios.append(audio)
            else:
                print(f"  [→] Chunk {i+1}/{len(chunks)} cache miss – đang sinh (hash: {chunk_hash[:8]})")
                # Gọi inference chạy trên thread pool
                audio = await loop.run_in_executor(
                    _executor, _do_inference, chunk_text, chunk_cache_path, actual_steps, model_length_scale, temperature
                )
                chunk_audios.append(audio)

        # Ghép nối các chunk 1.0x với khoảng lặng mặc định là 300ms (0.3s)
        # Khi dùng FFmpeg tăng tốc lên x1.5, khoảng nghỉ này tự động thu lại còn 200ms
        full_audio = concatenate_with_pause_and_smooth_fades(
            chunk_audios, 
            SAMPLE_RATE, 
            fade_duration=0.15,       # 150ms vuốt nhanh Smoothstep
            silence_duration=0.1      # Mặc định 100ms cho 1.0x
        )

        # Lưu cache file 1.0x
        await loop.run_in_executor(None, sf.write, full_cached_path, full_audio, SAMPLE_RATE)
        cached_path = full_cached_path
        print(f"[✓] Đã lưu cache full 1.0x thành công: {cached_path}")

    # Áp dụng thay đổi tốc độ/âm lượng bằng FFmpeg toàn cục trên file đã ghép (nếu cần)
    try:
        if abs(ffmpeg_speed - 1.0) > 0.01 or abs(volume - 1.0) > 0.01:
            out_path = await loop.run_in_executor(
                _executor, _do_ffmpeg, cached_path, ffmpeg_speed, volume
            )
            background_tasks.add_task(cleanup_file, out_path)
            return FileResponse(out_path, media_type="audio/wav")
        else:
            return FileResponse(cached_path, media_type="audio/wav")
    except Exception as e:
        raise HTTPException(500, f"Lỗi FFmpeg: {e}")

# ─── Endpoint phụ ─────────────────────────────────────────────────────────────
@app.post("/clear_cache")
async def clear_cache():
    shutil.rmtree(CACHE_DIR, ignore_errors=True)
    os.makedirs(CACHE_DIR, exist_ok=True)
    return {"status": "ok", "message": "Đã xóa toàn bộ cache"}

@app.get("/status")
async def status():
    busy = _executor._work_queue.qsize() if _executor else 0
    return {
        "workers": NUM_WORKERS,
        "device": str(_device),
        "threads_per_worker": 1,
        "queue_pending": busy,
    }

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8016))
    print(f"[*] Khởi động server tại http://0.0.0.0:{port}")
    print(f"[*] Để tuỳ chỉnh số core: TTS_WORKERS=2 python api_server.py")
    uvicorn.run(app, host="0.0.0.0", port=port)