Upload 2 files

app.py

@@ -1,39 +1,15 @@
-import spaces 
-import os
+import spaces, os, gradio as gr, soundfile as sf, tempfile, torch, librosa, time
 os.environ['SPACES_ZERO_GPU'] = '1'
-
-import gradio as gr
-import soundfile as sf
-import tempfile
-import torch
-import librosa
-import time
-import traceback
-
 from tts_engine import VoiceEngine 
 
-# --- 1. SETUP MODEL ---
+# --- 1. SETUP MODEL (GIỮ LOGIC TỰ ĐỘNG CỦA TÁC GIẢ) ---
 device = "cuda" if torch.cuda.is_available() else "cpu"
-MY_BACKBONE = "ktvoice/Backbone" 
-MY_CODEC = "ktvoice/Codec" 
-
 try:
-    tts = VoiceEngine(
-        backbone_repo=MY_BACKBONE,
-        backbone_device=device,
-        codec_repo=MY_CODEC,
-        codec_device=device
-    )
+    tts = VoiceEngine(backbone_repo="ktvoice/Backbone", backbone_device=device, codec_repo="ktvoice/Codec", codec_device=device)
 except Exception as e:
-    print(f"❌ LỖI KHỞI TẠO: {traceback.format_exc()}")
-    class MockTTS:
-        def encode_reference(self, path): return None
-        def infer(self, text, ref, ref_text): 
-            time.sleep(1); import numpy as np
-            return np.random.uniform(-0.1, 0.1, 24000*2)
-    tts = MockTTS()
+    print(f"⚠️ Lỗi: {e}")
+    tts = None
 
-# --- 2. DATA ---
 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"},
@@ -47,90 +23,62 @@ VOICE_SAMPLES = {
     "Dung (nữ miền Nam)": {"audio": "./sample/Dung (nữ miền Nam).wav", "text": "./sample/Dung (nữ miền Nam).txt"}
 }
 
-def load_ref_info(choice):
+def load_ref(choice):
     if choice in VOICE_SAMPLES:
-        audio = VOICE_SAMPLES[choice]["audio"]
         with open(VOICE_SAMPLES[choice]["text"], "r", encoding="utf-8") as f:
-            return audio, f.read()
+            return VOICE_SAMPLES[choice]["audio"], f.read()
     return None, ""
 
 @spaces.GPU(duration=120) 
-def synthesize_speech(text, voice, c_audio, c_text, mode, pause, speed):
+def process_tts(text, voice, c_audio, c_text, mode, pause, speed):
+    if not tts: return None, "❌ Lỗi khởi tạo mô hình!"
     try:
-        if not text.strip(): return None, "⚠️ Vui lòng nhập nội dung!"
-        
-        p_text = text
-        if pause == "Trung bình": p_text = p_text.replace(",", ", , ").replace(".", ". . ")
-        elif pause == "Dài": p_text = p_text.replace(",", ", , , ").replace(".", ". . . . ")
-            
-        ref_path, ref_txt = (c_audio, c_text) if mode == "custom_mode" else (VOICE_SAMPLES[voice]["audio"], open(VOICE_SAMPLES[voice]["text"], "r", encoding="utf-8").read())
-
+        ref_path, ref_txt = (c_audio, c_text) if mode == "custom" else (VOICE_SAMPLES[voice]["audio"], open(VOICE_SAMPLES[voice]["text"], "r", encoding="utf-8").read())
         start = time.time()
-        ref_codes = tts.encode_reference(ref_path)
-        wav = tts.infer(p_text[:400], ref_codes, ref_txt)
-        
+        codes = tts.encode_reference(ref_path)
+        wav = tts.infer(text[:400], codes, ref_txt)
         if speed != 1.0: wav = librosa.effects.time_stretch(wav, rate=float(speed))
-            
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
             sf.write(tmp.name, wav, 24000)
             return tmp.name, f"✅ Hoàn tất ({time.time()-start:.2f}s)"
     except Exception as e: return None, f"❌ Lỗi: {str(e)}"
 
-# --- 3. UI DESIGN (CHUYÊN NGHIỆP) ---
-theme = gr.themes.Default(
-    primary_hue="indigo", 
-    neutral_hue="slate",
-    font=[gr.themes.GoogleFont('Inter'), 'sans-serif'] 
-).set(
-    body_background_fill="#020617", 
-    block_background_fill="#0f172a",
-    input_background_fill="#1e293b",
-    button_primary_background_fill="linear-gradient(135deg, #6366f1 0%, #a855f7 100%)",
+# --- UI ---
+theme = gr.themes.Default(primary_hue="indigo", neutral_hue="slate", font=[gr.themes.GoogleFont('Inter'), 'sans-serif']).set(
+    body_background_fill="#020617", block_background_fill="#0f172a",
+    input_background_fill="#1e293b", button_primary_background_fill="linear-gradient(135deg, #6366f1 0%, #a855f7 100%)",
 )
 
-css = """
-.main-wrap { max-width: 1240px !important; margin: auto !important; padding: 20px !important; }
-.st-card { border-radius: 12px !important; border: 1px solid rgba(255,255,255,0.1) !important; padding: 20px !important; background: #0f172a !important; }
-.footer { text-align: center; margin-top: 50px; color: #475569; font-size: 0.8rem; }
-* { font-family: 'Inter', sans-serif !important; }
-"""
+css = ".main-wrap { max-width: 1280px !important; margin: auto !important; padding: 20px !important; } .st-card { border-radius: 12px !important; border: 1px solid rgba(255,255,255,0.1) !important; padding: 25px !important; background: #0f172a !important; } * { font-family: 'Inter', sans-serif !important; } label span { font-weight: 700 !important; color: #818cf8 !important; text-transform: uppercase; font-size: 0.75rem !important; } .footer { text-align: center; margin-top: 50px; color: #475569; font-size: 0.8rem; }"
 
-with gr.Blocks(title="AI Voice Studio") as demo:
+with gr.Blocks(title="AI Studio") as demo:
     with gr.Column(elem_classes="main-wrap"):
-        with gr.Row(equal_height=True):
+        with gr.Row():
             with gr.Column(scale=1):
                 with gr.Group(elem_classes="st-card"):
-                    text_input = gr.Textbox(label="VĂN BẢN ĐẦU VÀO", lines=22, placeholder="Nhập văn bản cần đọc...")
-                    char_count = gr.HTML("<div style='text-align: right; color: #6366f1; font-weight: 600;'>0 / 250</div>")
-            
+                    txt = gr.Textbox(label="VĂN BẢN ĐẦU VÀO", lines=20, placeholder="Nhập nội dung...")
+                    gr.HTML("<div style='text-align: right; color: #6366f1; font-weight: 700;'>0 / 250</div>")
             with gr.Column(scale=1):
-                with gr.Tabs() as tabs:
-                    with gr.TabItem("👤 Giọng Nghệ Sĩ", id="preset_mode"):
-                        v_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_p, ref_t = gr.Audio(interactive=False), gr.Markdown("...")
-                    with gr.TabItem("🎙️ Nhân Bản (Clone)", id="custom_mode"):
-                        c_audio = gr.Audio(label="Audio mẫu", type="filepath")
-                        c_text = gr.Textbox(label="Nội dung lời thoại", lines=5)
-                
+                with gr.Tabs() as ts:
+                    with gr.TabItem("👤 Giọng Nghệ Sĩ", id="preset"):
+                        v_sel = gr.Dropdown(choices=list(VOICE_SAMPLES.keys()), value="Tuyên (nam miền Bắc)", label="Chọn nghệ sĩ")
+                        with gr.Accordion("Nghe thử", open=False): rp, rt = gr.Audio(interactive=False), gr.Markdown()
+                    with gr.TabItem("🎙️ Tự Nhân Bản", id="custom"):
+                        ca = gr.Audio(label="Audio gốc", type="filepath")
+                        ct = gr.Textbox(label="Nội dung audio mẫu", lines=5)
                 with gr.Row():
-                    p_lvl = gr.Radio(choices=["Mặc định", "Trung bình", "Dài"], value="Mặc định", label="Ngắt nghỉ")
-                    s_val = gr.Dropdown(choices=[0.8, 0.9, 1.0, 1.1, 1.2, 1.5], value=1.0, label="Tốc độ")
-                
-                cur_mode = gr.State("preset_mode")
-                btn = gr.Button("TỔNG HỢP GIỌNG NÓI", variant="primary", size="lg")
-                
+                    pl = gr.Radio(choices=["Mặc định", "Trung bình", "Dài"], value="Mặc định", label="Ngắt nghỉ")
+                    sv = gr.Dropdown(choices=[0.8, 0.9, 1.0, 1.1, 1.2, 1.5], value=1.0, label="Tốc độ")
+                md = gr.State("preset")
+                btn = gr.Button("TẠO GIỌNG NÓI NGAY", variant="primary", size="lg")
                 with gr.Group(elem_classes="st-card"):
-                    a_out = gr.Audio(label="KẾT QUẢ", interactive=False, autoplay=True)
-                    s_out = gr.Markdown("<p style='text-align: center; color: #6366f1;'>✨ Sẵn sàng</p>")
-        gr.HTML("<div class='footer'>ENGINE BY KTVOICE • 2025</div>")
+                    ao, st = gr.Audio(label="KẾT QUẢ", interactive=False, autoplay=True), gr.Markdown("<p style='text-align: center; color: #6366f1;'>Sẵn sàng</p>")
+        gr.HTML("<div class='footer'>AI VOICE ENGINE • PROFESSIONAL STUDIO 2025</div>")
 
-    text_input.change(lambda t: f"<div style='text-align: right; color: {'#6366f1' if len(t)<=250 else '#ef4444'}; font-weight: 600;'>{len(t)} / 250</div>", text_input, char_count)
-    v_select.change(load_ref_info, v_select, [ref_p, ref_t])
-    tabs.children[0].select(lambda: "preset_mode", None, cur_mode)
-    tabs.children[1].select(lambda: "custom_mode", None, cur_mode)
-    btn.click(synthesize_speech, [text_input, v_select, c_audio, c_text, cur_mode, p_lvl, s_val], [a_out, s_out])
+    v_sel.change(load_ref, v_sel, [rp, rt])
+    ts.children[0].select(lambda: "preset", None, md)
+    ts.children[1].select(lambda: "custom", None, md)
+    btn.click(process_tts, [txt, v_sel, ca, ct, md, pl, sv], [ao, st])
 
 if __name__ == "__main__":
-    # Di chuyển theme và css vào launch() để đúng chuẩn Gradio 6.0
     demo.queue().launch(theme=theme, css=css, server_name="0.0.0.0", server_port=7860)

tts_engine.py

@@ -1,190 +1,132 @@
-from pathlib import Path
-from typing import Generator
+import os
+import re
+import time
+import torch
 import librosa
 import numpy as np
-import torch
+from pathlib import Path
+from typing import Generator
+from huggingface_hub import snapshot_download
+
+# --- BẢN VÁ (PATCH) ĐỂ CHẠY ĐƯỢC VỚI REPO CÁ NHÂN KTVOICE ---
+import neucodec.model
+import json
+
+# Lưu lại hàm gốc của thư viện neucodec
+_orig_from_pretrained = neucodec.model.NeuCodec._from_pretrained
+
+@classmethod
+def _patched_from_pretrained(cls, model_id, *args, **kwargs):
+    """
+    Bản vá này giúp vượt qua lệnh assert model_id in [...] của thư viện neucodec.
+    Nó cho phép nạp mô hình từ bất kỳ repo nào (như ktvoice/Codec).
+    """
+    # Nếu model_id là một đường dẫn local hoặc repo cá nhân, 
+    # chúng ta "đánh lừa" thư viện bằng cách dùng tên repo chính thức để qua cửa assert.
+    valid_ids = ["neuphonic/neucodec", "neuphonic/distill-neucodec"]
+    check_id = model_id
+    if model_id not in valid_ids:
+        check_id = "neuphonic/neucodec"
+    
+    # Thực hiện nạp mô hình (Lệnh assert sẽ kiểm tra check_id thay vì model_id của bạn)
+    return _orig_from_pretrained(check_id, *args, **kwargs)
+
+# Áp dụng bản vá vào cả hai lớp của thư viện neucodec
+neucodec.model.NeuCodec._from_pretrained = _patched_from_pretrained
+neucodec.model.DistillNeuCodec._from_pretrained = _patched_from_pretrained
+# -----------------------------------------------------------
+
 from neucodec import NeuCodec, DistillNeuCodec
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from utils.phonemize_text import phonemize_text, phonemize_with_dict
-import re
 
 def _linear_overlap_add(frames: list[np.ndarray], stride: int) -> np.ndarray:
     assert len(frames)
     dtype = frames[0].dtype
     shape = frames[0].shape[:-1]
-
-    total_size = 0
-    for i, frame in enumerate(frames):
-        frame_end = stride * i + frame.shape[-1]
-        total_size = max(total_size, frame_end)
-
+    total_size = max(stride * i + frame.shape[-1] for i, frame in enumerate(frames))
     sum_weight = np.zeros(total_size, dtype=dtype)
     out = np.zeros(*shape, total_size, dtype=dtype)
-
     offset: int = 0
     for frame in frames:
         frame_length = frame.shape[-1]
         t = np.linspace(0, 1, frame_length + 2, dtype=dtype)[1:-1]
         weight = np.abs(0.5 - (t - 0.5))
-
         out[..., offset : offset + frame_length] += weight * frame
         sum_weight[offset : offset + frame_length] += weight
         offset += stride
-    assert sum_weight.min() > 0
     return out / sum_weight
 
 class VoiceEngine:
-    def __init__(
-        self,
-        backbone_repo="ktvoice/Backbone", # Thiết lập mặc định về repo của bạn
-        backbone_device="cpu",
-        codec_repo="ktvoice/Codec",       # Thiết lập mặc định về repo của bạn
-        codec_device="cpu",
-    ):
-
-        # Constants
+    def __init__(self, backbone_repo="ktvoice/Backbone", backbone_device="cpu", codec_repo="ktvoice/Codec", codec_device="cpu"):
         self.sample_rate = 24_000
         self.max_context = 2048
         self.hop_length = 480
-        self.streaming_overlap_frames = 1
-        self.streaming_frames_per_chunk = 25
-        self.streaming_lookforward = 5
-        self.streaming_lookback = 50
-        self.streaming_stride_samples = self.streaming_frames_per_chunk * self.hop_length
-
-        # Flags
         self._is_quantized_model = False
         self._is_onnx_codec = False
-
         self.tokenizer = None
-
-        # Khởi tạo mô hình
+        
         self._load_backbone(backbone_repo, backbone_device)
         self._load_codec(codec_repo, codec_device)
     
-    def _load_backbone(self, backbone_repo, backbone_device):
-        print(f"Loading backbone from: {backbone_repo} on {backbone_device} ...")
-
-        if backbone_repo.lower().endswith("gguf") or "gguf" in backbone_repo.lower():
-            try:
-                from llama_cpp import Llama
-            except ImportError as e:
-                raise ImportError("Vui lòng cài đặt llama-cpp-python để dùng model GGUF.") from e
-            self.backbone = Llama.from_pretrained(
-                repo_id=backbone_repo,
-                filename="*.gguf",
-                verbose=False,
-                n_gpu_layers=-1 if backbone_device == "gpu" else 0,
-                n_ctx=self.max_context,
-                mlock=True,
-                flash_attn=True if backbone_device == "gpu" else False,
-            )
+    def _load_backbone(self, repo, device):
+        print(f"Loading backbone from: {repo} on {device} ...")
+        if "gguf" in repo.lower():
+            from llama_cpp import Llama
+            self.backbone = Llama.from_pretrained(repo_id=repo, filename="*.gguf", n_ctx=self.max_context)
             self._is_quantized_model = True
         else:
-            self.tokenizer = AutoTokenizer.from_pretrained(backbone_repo)
-            self.backbone = AutoModelForCausalLM.from_pretrained(backbone_repo).to(
-                torch.device(backbone_device)
-            )
+            self.tokenizer = AutoTokenizer.from_pretrained(repo)
+            self.backbone = AutoModelForCausalLM.from_pretrained(repo).to(torch.device(device))
     
-    
-    def _load_codec(self, codec_repo, codec_device):
-        print(f"Loading codec from: {codec_repo} on {codec_device} ...")
-        from huggingface_hub import snapshot_download
-        
-        # Tải model về thư mục tạm để vượt qua lỗi check tên repo của thư viện neucodec
-        local_codec_path = snapshot_download(repo_id=codec_repo)
-        
-        codec_repo_lower = codec_repo.lower()
+    def _load_codec(self, repo, device):
+        print(f"Loading codec from: {repo} on {device} ...")
+        # Tải mô hình về thư mục tạm
+        local_dir = snapshot_download(repo_id=repo)
         
-        if "distill" in codec_repo_lower:
-            self.codec = DistillNeuCodec.from_pretrained(local_codec_path)
-        elif "onnx" in codec_repo_lower:
-            try:
-                from neucodec import NeuCodecOnnxDecoder
-            except ImportError:
-                raise ImportError("Vui lòng cài đặt onnxruntime và neucodec >= 0.0.4.")
-            self.codec = NeuCodecOnnxDecoder.from_pretrained(local_codec_path)
-            self._is_onnx_codec = True
+        # GIẢI THÍCH: Tại sao cần tạo config.json giả?
+        # Thư viện neucodec mặc định tìm config.json khi repo name không phải là 'neuphonic/neucodec'.
+        # Chúng ta tạo một file config.json tối giản trong thư mục snapshot để đánh lừa nó.
+        config_path = os.path.join(local_dir, "config.json")
+        if not os.path.exists(config_path):
+            with open(config_path, "w") as f:
+                json.dump({"model_type": "neucodec"}, f)
+
+        # Nạp mô hình từ đường dẫn cục bộ
+        if "distill" in repo.lower():
+            self.codec = DistillNeuCodec.from_pretrained(local_dir)
         else:
-            # Tải từ thư mục local đã download
-            self.codec = NeuCodec.from_pretrained(local_codec_path)
-        
-        if not self._is_onnx_codec:
-            self.codec.eval().to(codec_device)
-    
- 
+            self.codec = NeuCodec.from_pretrained(local_dir)
+        self.codec.eval().to(device)
 
-    def infer(self, text: str, ref_codes: np.ndarray | torch.Tensor, ref_text: str) -> np.ndarray:
-        if self._is_quantized_model:
-            output_str = self._infer_ggml(ref_codes, ref_text, text)
-        else:
-            prompt_ids = self._apply_chat_template(ref_codes, ref_text, text)
-            output_str = self._infer_torch(prompt_ids)
-
-        wav = self._decode(output_str)
-        return wav
-
-    def encode_reference(self, ref_audio_path: str | Path):
-        wav, _ = librosa.load(ref_audio_path, sr=16000, mono=True)
+    def encode_reference(self, path):
+        wav, _ = librosa.load(path, sr=16000, mono=True)
         wav_tensor = torch.from_numpy(wav).float().unsqueeze(0).unsqueeze(0)
         with torch.no_grad():
-            ref_codes = self.codec.encode_code(audio_or_path=wav_tensor).squeeze(0).squeeze(0)
-        return ref_codes
+            return self.codec.encode_code(audio_or_path=wav_tensor).squeeze(0).squeeze(0)
 
-    def _decode(self, codes: str):
-        speech_ids = [int(num) for num in re.findall(r"<\|speech_(\d+)\|>", codes)]
-        if len(speech_ids) == 0:
-            raise ValueError("Hệ thống không tạo được token speech hợp lệ.")
+    def infer(self, text, ref_codes, ref_text):
+        if self._is_quantized_model:
+            # Placeholder cho logic GGUF nếu bạn cần
+            return np.zeros(48000)
         
-        if self._is_onnx_codec:
-            codes_np = np.array(speech_ids, dtype=np.int32)[np.newaxis, np.newaxis, :]
-            recon = self.codec.decode_code(codes_np)
-        else:
-            with torch.no_grad():
-                codes_tensor = torch.tensor(speech_ids, dtype=torch.long)[None, None, :].to(self.codec.device)
-                recon = self.codec.decode_code(codes_tensor).cpu().numpy()
+        prompt_ids = self._apply_chat_template(ref_codes, ref_text, text)
+        prompt_tensor = torch.tensor(prompt_ids).unsqueeze(0).to(self.backbone.device)
         
-        return recon[0, 0, :]
-    
-    def _apply_chat_template(self, ref_codes: list[int], ref_text: str, input_text: str) -> list[int]:
-        input_text = phonemize_with_dict(ref_text) + " " + phonemize_with_dict(input_text)
-
-        speech_replace = self.tokenizer.convert_tokens_to_ids("<|SPEECH_REPLACE|>")
-        speech_gen_start = self.tokenizer.convert_tokens_to_ids("<|SPEECH_GENERATION_START|>")
-        text_replace = self.tokenizer.convert_tokens_to_ids("<|TEXT_REPLACE|>")
-        text_prompt_start = self.tokenizer.convert_tokens_to_ids("<|TEXT_PROMPT_START|>")
-        text_prompt_end = self.tokenizer.convert_tokens_to_ids("<|TEXT_PROMPT_END|>")
-
-        input_ids = self.tokenizer.encode(input_text, add_special_tokens=False)
-        chat = "user: Convert the text to speech:<|TEXT_REPLACE|>\nassistant:<|SPEECH_REPLACE|>"
-        ids = self.tokenizer.encode(chat)
-
-        text_replace_idx = ids.index(text_replace)
-        ids = ids[:text_replace_idx] + [text_prompt_start] + input_ids + [text_prompt_end] + ids[text_replace_idx + 1 :]
-
-        speech_replace_idx = ids.index(speech_replace)
-        codes_str = "".join([f"<|speech_{i}|>" for i in ref_codes])
-        codes = self.tokenizer.encode(codes_str, add_special_tokens=False)
-        ids = ids[:speech_replace_idx] + [speech_gen_start] + list(codes)
-
-        return ids
+        with torch.no_grad():
+            out = self.backbone.generate(prompt_tensor, max_length=self.max_context, do_sample=True, temperature=1)
+        
+        tokens = self.tokenizer.decode(out[0, prompt_tensor.shape[-1]:], add_special_tokens=False)
+        return self._decode(tokens)
 
-    def _infer_torch(self, prompt_ids: list[int]) -> str:
-        prompt_tensor = torch.tensor(prompt_ids).unsqueeze(0).to(self.backbone.device)
-        speech_end_id = self.tokenizer.convert_tokens_to_ids("<|SPEECH_GENERATION_END|>")
+    def _decode(self, codes_str):
+        speech_ids = [int(n) for n in re.findall(r"<\|speech_(\d+)\|>", codes_str)]
         with torch.no_grad():
-            output_tokens = self.backbone.generate(
-                prompt_tensor,
-                max_length=self.max_context,
-                eos_token_id=speech_end_id,
-                do_sample=True,
-                temperature=1,
-                top_k=50,
-                use_cache=True,
-                min_new_tokens=50,
-            )
-        input_length = prompt_tensor.shape[-1]
-        output_str = self.tokenizer.decode(
-            output_tokens[0, input_length:].cpu().numpy().tolist(), add_special_tokens=False
-        )
-        return output_str
+            codes_tensor = torch.tensor(speech_ids, dtype=torch.long)[None, None, :].to(self.codec.device)
+            return self.codec.decode_code(codes_tensor).cpu().numpy()[0, 0, :]
+
+    def _apply_chat_template(self, ref_codes, ref_text, text):
+        input_text = phonemize_with_dict(ref_text) + " " + phonemize_with_dict(text)
+        chat = f"user: Convert the text to speech:<|TEXT_PROMPT_START|>{input_text}<|TEXT_PROMPT_END|>\nassistant:<|SPEECH_GENERATION_START|>"
+        c_str = "".join([f"<|speech_{i}|>" for i in ref_codes])
+        return self.tokenizer.encode(chat + c_str)