Update app.py

app.py

@@ -1,13 +1,15 @@
 import gradio as gr
-import whisper
+from transformers import pipeline
 import os
 import pandas as pd
 from pydub import AudioSegment
 import zipfile
 import shutil
+import torch
 
 # Tạo thư mục tạm để xử lý
 TEMP_DIR = "temp_dataset"
+MODEL_ID = "kelvinbksoh/whisper-large-v2-vietnamese-lyrics-transcription"
 
 def format_audio_for_piper(audio_segment):
     """
@@ -18,7 +20,7 @@ def format_audio_for_piper(audio_segment):
     audio_segment = audio_segment.set_sample_width(2) # 16-bit
     return audio_segment
 
-def create_dataset(audio_path, model_size, speaker_name):
+def create_dataset(audio_path, speaker_name):
     # Kiểm tra đầu vào
     if audio_path is None:
         return None, 0, pd.DataFrame()
@@ -31,17 +33,28 @@ def create_dataset(audio_path, model_size, speaker_name):
     wavs_dir = os.path.join(TEMP_DIR, "wavs")
     os.makedirs(wavs_dir, exist_ok=True)
 
-    print(f"Đang tải model Whisper: {model_size}...")
+    print(f"Đang tải model: {MODEL_ID}...")
     try:
-        # Load model Whisper
-        model = whisper.load_model(model_size)
+        # Load model bằng Transformers Pipeline
+        # Sử dụng CPU (device=-1) cho Space Free. Nếu có GPU thì chỉnh device=0
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipe = pipeline(
+            "automatic-speech-recognition", 
+            model=MODEL_ID, 
+            device=device,
+            chunk_length_s=30
+        )
     except Exception as e:
-        return None, 0, pd.DataFrame(columns=["Lỗi"], data=[[str(e)]])
+        return None, 0, pd.DataFrame(columns=["Lỗi"], data=[[f"Lỗi tải model: {str(e)}"]])
 
-    print("Đang xử lý audio...")
-    # Transcribe audio để lấy text và timestamps
-    result = model.transcribe(audio_path, language="vi")
-    segments = result["segments"]
+    print("Đang xử lý audio (Model Large rất nặng, vui lòng chờ)...")
+    
+    # Transcribe audio: return_timestamps=True là bắt buộc để cắt
+    try:
+        output = pipe(audio_path, return_timestamps=True, generate_kwargs={"language": "vi"})
+        chunks = output["chunks"] # Lấy danh sách các đoạn đã nhận diện
+    except Exception as e:
+        return None, 0, pd.DataFrame(columns=["Lỗi"], data=[[f"Lỗi xử lý audio: {str(e)}"]])
     
     # Load audio gốc bằng pydub
     original_audio = AudioSegment.from_file(audio_path)
@@ -49,37 +62,47 @@ def create_dataset(audio_path, model_size, speaker_name):
     metadata_rows = []
     preview_data = []
     
-    print(f"Tìm thấy {len(segments)} câu. Đang cắt và lưu...")
+    print(f"Tìm thấy {len(chunks)} câu. Đang cắt và lưu...")
     
-    for i, segment in enumerate(segments):
-        start_time = segment["start"] * 1000 # pydub dùng mili giây
-        end_time = segment["end"] * 1000
-        text = segment["text"].strip()
+    file_counter = 0
+    
+    for chunk in chunks:
+        # Transformers trả về timestamp dạng tuple (start, end) đơn vị giây
+        timestamp = chunk["timestamp"]
+        text = chunk["text"].strip()
+        
+        # Kiểm tra timestamp hợp lệ (đôi khi model trả về None ở cuối)
+        if timestamp[0] is None or timestamp[1] is None:
+            continue
+            
+        start_time = timestamp[0] * 1000 
+        end_time = timestamp[1] * 1000
         
-        # Bỏ qua các đoạn quá ngắn hoặc không có chữ
+        # Bỏ qua các đoạn quá ngắn (dưới 0.5s) hoặc không có chữ
         if (end_time - start_time) < 500 or len(text) < 2:
             continue
 
         # Cắt audio
-        chunk = original_audio[start_time:end_time]
+        audio_chunk = original_audio[start_time:end_time]
         
         # Chuẩn hóa format cho Piper
-        chunk = format_audio_for_piper(chunk)
+        audio_chunk = format_audio_for_piper(audio_chunk)
         
-        # Đặt tên file: 0001.wav, 0002.wav...
-        filename = f"{i:04d}.wav"
+        # Đặt tên file
+        filename = f"{file_counter:04d}.wav"
         save_path = os.path.join(wavs_dir, filename)
         
-        chunk.export(save_path, format="wav")
+        audio_chunk.export(save_path, format="wav")
         
         # Format Piper: filename|speaker|text
         file_id = filename.replace(".wav", "")
         row_str = f"{file_id}|{speaker_name}|{text}"
         metadata_rows.append(row_str)
         
-        # Data cho bảng xem trước (chỉ lấy 10 dòng đầu để hiển thị nhanh)
         if len(preview_data) < 10:
             preview_data.append([file_id, speaker_name, text])
+            
+        file_counter += 1
 
     # Ghi file metadata.csv
     csv_path = os.path.join(TEMP_DIR, "metadata.csv")
@@ -98,27 +121,25 @@ def create_dataset(audio_path, model_size, speaker_name):
     return zip_filename, len(metadata_rows), pd.DataFrame(preview_data, columns=["ID", "Speaker", "Text"])
 
 # Giao diện Gradio
-with gr.Blocks(title="Auto Piper Dataset Maker") as demo:
-    gr.Markdown("## 🎧 Tự động tạo Dataset cho Piper TTS")
-    gr.Markdown("Upload file âm thanh dài -> Tự động cắt câu & nhận dạng chữ (ASR) -> Xuất file ZIP chuẩn Piper.")
+with gr.Blocks(title="Auto Piper Dataset (Large Model)") as demo:
+    gr.Markdown(f"## 🎧 Auto Piper Dataset Maker")
+    gr.Markdown(f"Model sử dụng: `{MODEL_ID}` (Whisper Large V2 Fine-tuned)")
+    gr.Markdown("⚠️ **Lưu ý:** Model này rất nặng. Quá trình xử lý sẽ chậm hơn nhiều so với model Small.")
     
     with gr.Row():
         with gr.Column():
             audio_input = gr.Audio(type="filepath", label="Upload Audio Gốc")
             speaker_name = gr.Textbox(value="vner", label="Tên người đọc (Speaker Name)")
-            # Mặc định chọn small hoặc medium để cân bằng tốc độ trên CPU
-            model_size = gr.Dropdown(choices=["base", "small", "medium"], value="small", label="Whisper Model Size")
-            btn_submit = gr.Button("🚀 Xử lý & Tạo Dataset", variant="primary")
+            btn_submit = gr.Button("🚀 Xử lý (Chờ khoảng 2-5 phút)", variant="primary")
         
         with gr.Column():
             output_zip = gr.File(label="Tải về Dataset (.zip)")
             info_text = gr.Number(label="Số lượng câu đã cắt")
-            # Đã sửa lỗi: Bỏ tham số max_rows
             preview_table = gr.Dataframe(label="Xem trước Metadata (10 câu đầu)", headers=["ID", "Speaker", "Text"])
 
     btn_submit.click(
         fn=create_dataset,
-        inputs=[audio_input, model_size, speaker_name],
+        inputs=[audio_input, speaker_name],
         outputs=[output_zip, info_text, preview_table]
     )