Update app.py

app.py

@@ -3,36 +3,46 @@ import random
 import numpy as np
 import torch
 from pathlib import Path
+
+# Đảm bảo 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 huggingface_hub import snapshot_download
 from chatterbox.src.chatterbox.tts import ChatterboxTTS
 import gradio as gr
 
-# --- CẤU HÌNH ---
-MODEL_REPO    = "LTTEAM/TTS_Pro"
-LOCAL_MODEL_DIR = Path("models") / "tts_pro"
+# --- CẤU HÌNH MODEL TỪ HUGGINGFACE ---
+MODEL_REPO = "LTTEAM/TTS_Pro"
+LOCAL_MODEL_DIR = Path(os.getcwd()) / "models" / "tts_pro"
 MODELS = {}
 
-# 1) Download model từ HF Hub (lần đầu) vào thư mục models/tts_pro
+# Download model một lần (cache)
 if not LOCAL_MODEL_DIR.exists():
-    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
+    print(f"📥 Đang tải model từ HuggingFace repo {MODEL_REPO} …")
     snapshot_download(
         repo_id=MODEL_REPO,
         repo_type="model",
         local_dir=str(LOCAL_MODEL_DIR),
         local_dir_use_symlinks=False
     )
+    print(f"✅ Đã tải xong vào {LOCAL_MODEL_DIR}")
 
 def get_or_load_model(device_str: str):
     """
-    Load ChatterboxTTS trên 'cpu' hoặc 'cuda'.
-    device_str: 'cpu' hoặc 'gpu'
+    Lấy hoặc load 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"
+    device = "cuda" if (device_str == "gpu" and torch.cuda.is_available()) else "cpu"
     if device not in MODELS:
         print(f"📂 Loading model lên {device} …")
         model = ChatterboxTTS.from_local(str(LOCAL_MODEL_DIR), device)
         MODELS[device] = model
-        print(f"✅ Model đã load trên {device}")
+        print(f"✅ Model đã được load lên {device}")
     return MODELS[device]
 
 def set_seed(seed: int):
@@ -44,17 +54,17 @@ def set_seed(seed: int):
     np.random.seed(seed)
 
 def chunk_text(text: str, chunk_size: int = 300):
-    """Chia text dài thành các đoạn ≤ chunk_size, giữ nguyên từ."""
+    """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, cur = [], ""
+    chunks, current = [], ""
     for w in words:
-        if len(cur)+len(w)+1 > chunk_size:
-            chunks.append(cur.strip())
-            cur = w
+        if len(current) + len(w) + 1 > chunk_size:
+            chunks.append(current.strip())
+            current = w
         else:
-            cur = f"{cur} {w}".strip()
-    if cur:
-        chunks.append(cur.strip())
+            current = f"{current} {w}".strip()
+    if current:
+        chunks.append(current.strip())
     return chunks
 
 def generate_tts_audio(
@@ -66,15 +76,19 @@ def generate_tts_audio(
     seed_num: int,
     cfg_weight: float
 ):
+    """
+    Sinh audio từ văn bản không giới hạn: chia thành chunk, generate từng chunk, ghép nối.
+    Trả về (sample_rate, numpy.ndarray).
+    """
     model = get_or_load_model(device_choice)
-    if seed_num!=0:
+    if seed_num != 0:
         set_seed(int(seed_num))
 
-    # chia thành chunks và sinh từng phần
-    chunks = chunk_text(text_input, 300)
+    chunks = chunk_text(text_input, chunk_size=300)
     waves, sr = [], model.sr
-    for i, chunk in enumerate(chunks, 1):
-        print(f"🔊 Sinh đoạn {i}/{len(chunks)}...")
+
+    for idx, chunk in enumerate(chunks, start=1):
+        print(f"🔊 Sinh đoạn {idx}/{len(chunks)} trên {model.device}")
         wav = model.generate(
             chunk,
             audio_prompt_path=audio_prompt_path,
@@ -84,49 +98,69 @@ def generate_tts_audio(
         )
         waves.append(wav.squeeze(0).cpu().numpy())
 
-    full = np.concatenate(waves, axis=0)
-    return sr, full
+    full_wave = np.concatenate(waves, axis=0)
+    print("✅ Hoàn thành sinh toàn bộ audio.")
+    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**  
-    Chuyển văn bản thành giọng nói chất lượng cao, không giới hạn độ dài.
-    """)
+    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(
-                ["cpu","gpu"],
+                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=6,
-                placeholder="Nhập hoặc dán văn bản..."
+                lines=8,
+                placeholder="Dán hoặc nhập văn bản vào đây..."
             )
             ref_wav = gr.Audio(
-                sources=["upload","microphone"],
+                sources=["upload", "microphone"],
                 type="filepath",
                 label="Âm thanh mẫu (tùy chọn)"
             )
-            exaggeration = gr.Slider(0.25,2,step=0.05,value=0.5,label="Mức nhấn nhá")
-            cfg_weight    = gr.Slider(0.2,1,step=0.05,value=0.5,label="CFG/Pace weight")
+            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(0.05,5,step=0.05,value=0.8,label="Nhiệt độ")
+                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 = gr.Audio(label="Kết quả")
+            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],
+        outputs=[out_audio],
     )
 
 if __name__ == "__main__":
-    # Trên Spaces, chỉ cần host 0.0.0.0, port mặc định
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Phát hiện Colab qua biến môi trường
+    is_colab = "COLAB_GPU" in os.environ
+    if is_colab:
+        demo.launch(share=True)
+    else:
+        # Dùng host/port để hỗ trợ HuggingFace Spaces
+        demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))