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)