chumtts2 / gui.py
chumdz97's picture
Upload folder using huggingface_hub
2271dee verified
Raw
History Blame Contribute Delete
9.68 kB
import sys
import json
import re
import numpy as np
import os
import torch
import soundfile as sf
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
QLabel, QComboBox, QSlider, QTextEdit, QPushButton, QMessageBox, QProgressBar)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from omnivoice import OmniVoice, OmniVoiceGenerationConfig
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
class WorkerThread(QThread):
progress = pyqtSignal(int, str)
finished = pyqtSignal(str)
error = pyqtSignal(str)
def __init__(self, model, voice_prompts_cache, text, voice_name, num_step, guidance_scale, speed):
super().__init__()
self.model = model
self.voice_prompts_cache = voice_prompts_cache
self.text = text
self.voice_name = voice_name
self.num_step = num_step
self.guidance_scale = guidance_scale
self.speed = speed
def run(self):
try:
if not self.text.strip():
self.error.emit("Vui lòng nhập văn bản cần đọc!")
return
prompt = self.voice_prompts_cache.get(self.voice_name)
if not prompt:
self.error.emit(f"Không tìm thấy dữ liệu đã cache cho giọng: {self.voice_name}")
return
config = OmniVoiceGenerationConfig(num_step=self.num_step, guidance_scale=self.guidance_scale)
sentences = [s.strip() for s in re.split(r'(?<=[.!?\n])\s+', self.text) if s.strip()]
audio_chunks = []
for i, sentence in enumerate(sentences):
self.progress.emit(int((i / len(sentences)) * 100), f"Đang đọc câu {i+1}/{len(sentences)}...")
audio = self.model.generate(
text=sentence,
language="vietnamese",
voice_clone_prompt=prompt,
generation_config=config,
speed=self.speed,
)
audio_chunks.append(audio[0])
silence = np.zeros(int(24000 * 0.1), dtype=np.float32)
audio_chunks.append(silence)
self.progress.emit(100, "Đang ghép file audio...")
final_audio = np.concatenate(audio_chunks)
output_path = os.path.join(BASE_DIR, "output_gui.wav")
sf.write(output_path, final_audio, 24000)
self.finished.emit(output_path)
except Exception as e:
self.error.emit(str(e))
class OmniVoiceGUI(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("OmniVoice Vietnamese - PyQt5 GUI")
self.resize(800, 600)
self.model = None
self.voice_prompts_cache = {}
self.init_ui()
self.load_model_and_voices()
def init_ui(self):
central_widget = QWidget()
self.setCentralWidget(central_widget)
main_layout = QHBoxLayout()
# Cột bên trái: Settings
left_layout = QVBoxLayout()
left_layout.addWidget(QLabel("🧑‍🎤 Chọn giọng đọc:"))
self.voice_combo = QComboBox()
left_layout.addWidget(self.voice_combo)
left_layout.addSpacing(20)
left_layout.addWidget(QLabel("⚙️ Num Steps (Độ mượt):"))
self.step_label = QLabel("32")
left_layout.addWidget(self.step_label)
self.step_slider = QSlider(Qt.Horizontal)
self.step_slider.setMinimum(8)
self.step_slider.setMaximum(50)
self.step_slider.setValue(32)
self.step_slider.valueChanged.connect(lambda v: self.step_label.setText(str(v)))
left_layout.addWidget(self.step_slider)
left_layout.addSpacing(10)
left_layout.addWidget(QLabel("⚙️ Guidance Scale (Chuẩn xác/Robot):"))
self.cfg_label = QLabel("5.0")
left_layout.addWidget(self.cfg_label)
self.cfg_slider = QSlider(Qt.Horizontal)
self.cfg_slider.setMinimum(10) # 1.0
self.cfg_slider.setMaximum(100) # 10.0
self.cfg_slider.setValue(50)
self.cfg_slider.valueChanged.connect(lambda v: self.cfg_label.setText(str(v/10.0)))
left_layout.addWidget(self.cfg_slider)
left_layout.addSpacing(10)
left_layout.addWidget(QLabel("⚙️ Tốc độ đọc (Speed):"))
self.speed_label = QLabel("1.0")
left_layout.addWidget(self.speed_label)
self.speed_slider = QSlider(Qt.Horizontal)
self.speed_slider.setMinimum(5) # 0.5
self.speed_slider.setMaximum(20) # 2.0
self.speed_slider.setValue(10)
self.speed_slider.valueChanged.connect(lambda v: self.speed_label.setText(str(v/10.0)))
left_layout.addWidget(self.speed_slider)
left_layout.addStretch()
# Cột bên phải: Text và Button
right_layout = QVBoxLayout()
right_layout.addWidget(QLabel("📝 Nhập nội dung cần đọc:"))
self.text_input = QTextEdit()
self.text_input.setPlaceholderText("Nhập văn bản dài thoải mái, hệ thống sẽ tự chia nhỏ theo câu...")
right_layout.addWidget(self.text_input)
self.progress_bar = QProgressBar()
self.progress_bar.setValue(0)
right_layout.addWidget(self.progress_bar)
self.status_label = QLabel("Trạng thái: Sẵn sàng")
right_layout.addWidget(self.status_label)
self.generate_btn = QPushButton("🚀 TẠO GIỌNG NÓI")
self.generate_btn.setMinimumHeight(50)
self.generate_btn.clicked.connect(self.start_generation)
right_layout.addWidget(self.generate_btn)
main_layout.addLayout(left_layout, 1)
main_layout.addLayout(right_layout, 2)
central_widget.setLayout(main_layout)
def load_model_and_voices(self):
self.status_label.setText("Trạng thái: Đang load OmniVoice Model... (Vui lòng đợi vài chục giây)")
QApplication.processEvents()
try:
self.model = OmniVoice.from_pretrained(
BASE_DIR,
device_map="cuda:0",
dtype=torch.float16,
)
VOICES_FILE = os.path.join(BASE_DIR, "voice", "voices.json")
if os.path.exists(VOICES_FILE):
with open(VOICES_FILE, "r", encoding="utf-8") as f:
voices_config = json.load(f)
self.status_label.setText("Trạng thái: Đang cache các giọng mẫu...")
QApplication.processEvents()
for name, data in voices_config.items():
audio_rel_path = data.get("audio")
audio_path = os.path.join(BASE_DIR, "voice", audio_rel_path) if audio_rel_path else None
ref_text = data.get("text")
if audio_path and os.path.exists(audio_path) and ref_text:
self.voice_prompts_cache[name] = self.model.create_voice_clone_prompt(
ref_audio=audio_path,
ref_text=ref_text,
)
self.voice_combo.addItem(name)
self.status_label.setText("Trạng thái: Load model thành công. Sẵn sàng!")
self.progress_bar.setValue(100)
except Exception as e:
QMessageBox.critical(self, "Lỗi", f"Không thể load model: {e}")
self.status_label.setText("Trạng thái: Lỗi Load Model!")
def start_generation(self):
if not self.model:
QMessageBox.warning(self, "Cảnh báo", "Model chưa được load thành công!")
return
text = self.text_input.toPlainText()
voice_name = self.voice_combo.currentText()
num_step = self.step_slider.value()
guidance_scale = self.cfg_slider.value() / 10.0
speed = self.speed_slider.value() / 10.0
self.generate_btn.setEnabled(False)
self.progress_bar.setValue(0)
self.status_label.setText("Trạng thái: Bắt đầu xử lý...")
self.worker = WorkerThread(self.model, self.voice_prompts_cache, text, voice_name, num_step, guidance_scale, speed)
self.worker.progress.connect(self.update_progress)
self.worker.finished.connect(self.on_finished)
self.worker.error.connect(self.on_error)
self.worker.start()
def update_progress(self, percent, msg):
self.progress_bar.setValue(percent)
self.status_label.setText(f"Trạng thái: {msg}")
def on_finished(self, output_path):
self.progress_bar.setValue(100)
self.status_label.setText(f"Trạng thái: Đã lưu thành công tại {output_path}")
self.generate_btn.setEnabled(True)
# Tự động mở file vừa xuất
if sys.platform == "win32":
os.startfile(output_path)
elif sys.platform == "darwin":
os.system(f"open {output_path}")
else:
os.system(f"xdg-open {output_path} &")
def on_error(self, err_msg):
QMessageBox.critical(self, "Lỗi", err_msg)
self.status_label.setText("Trạng thái: Lỗi khi xử lý!")
self.generate_btn.setEnabled(True)
if __name__ == "__main__":
app = QApplication(sys.argv)
# Optional: Thiết lập style cho giống app xịn xò xíu
app.setStyle("Fusion")
gui = OmniVoiceGUI()
gui.show()
sys.exit(app.exec_())