Ver 1.3.0: Streaming & Prettier ?

Author: AnyaCoder

fish/gui.py

@@ -10,6 +10,7 @@
 from PyQt6.QtGui import QIcon, QPixmap
 from PyQt6.QtMultimedia import QAudioOutput, QMediaPlayer
 from PyQt6.QtWidgets import (
+    QCheckBox,
     QComboBox,
     QFileDialog,
     QGridLayout,
@@ -40,8 +41,9 @@
     FAPTranscribeWidget,
 )
 from fish.input import TextEditorWidget
-from fish.modules.console import ConsoleStream, ConsoleWidget
+from fish.modules.console import ConsoleWidget
 from fish.modules.globals import STOP_BUTTON_QSS
+from fish.modules.log import stderr_stream, stdout_stream
 from fish.modules.registry import widget_registry
 from fish.modules.worker import TTSWorker
 from fish.utils.audio import get_devices
@@ -101,23 +103,14 @@ def __init__(self):
         self.main_layout.setAlignment(Qt.AlignmentFlag.AlignTop)
         self.setup_action_buttons(self.main_layout)
 
-        self.change_theme(self.theme_combo.currentIndex())  # initialize theme for 1st
+        self.change_theme(self.theme_combo.currentIndex())  # initialize theme first
+
+        stdout_stream.new_message.connect(lambda msg: self.update_console(msg, "white"))
+        stderr_stream.new_message.connect(lambda msg: self.update_console(msg, "red"))
 
         # Use size hint to set a reasonable size
         self.setMinimumWidth(800)
 
-        # Redefined Stream
-        self.stdout_stream = ConsoleStream()
-        self.stderr_stream = ConsoleStream()
-        self.stdout_stream.new_message.connect(
-            lambda msg: self.update_console(msg, "white")
-        )
-        self.stderr_stream.new_message.connect(
-            lambda msg: self.update_console(msg, "red")
-        )
-        sys.stdout = self.stdout_stream
-        sys.stderr = self.stderr_stream
-
         # Uploaded ref files
         self.files = []
 
@@ -637,10 +630,14 @@ def setup_action_buttons(self, layout: QVBoxLayout):
         row_layout = QHBoxLayout()
         widget_registry.register(row, "action_widget")
 
-        self.now_audio = QLabel(_t("action.audio").format(audio_name="(null)"))
+        self.now_audio = QLineEdit(_t("action.audio").format(audio_name="(null)"))
+        self.now_audio.setMinimumWidth(200)
+
         row_layout.addWidget(self.now_audio)
-        row_layout.addStretch(1)
+        # row_layout.addStretch(1)
 
+        self.stream = QCheckBox(_t("action.stream"))
+        row_layout.addWidget(self.stream)
         self.start_button = QPushButton(_t("action.start"))
         self.start_button.clicked.connect(self.start_conversion)
         row_layout.addWidget(self.start_button)
@@ -674,11 +671,19 @@ def change_theme(self, index):
                     }
                 """
             )
+            self.now_audio.setStyleSheet(
+                """
+                QLineEdit {
+                    border: none;
+                    color: black;
+                }
+            """
+            )
 
         else:
             for widget in widget_registry.get_registered_widgets().values():
                 widget.setStyleSheet("")
-
+            self.now_audio.setStyleSheet("")
         save_config()
         qdarktheme.setup_theme(config.theme)
 
@@ -861,7 +866,7 @@ def open_file(self):
         )
         self.set_audio(file_name)
 
-    def set_audio(self, audio_file):
+    def set_audio(self, audio_file: str):
         if Path(audio_file).exists():
             self.player.setSource(QUrl.fromLocalFile(audio_file))
             self.play_button.setText(_t("tts_output.play"))
@@ -916,7 +921,8 @@ def start_conversion(self):
         text = self.text_editor.input_edit.toPlainText()
 
         audio_name = now.strftime("%Y%m%d_%H%M%S")
-        audio_path = Path(self.save_audio_path.text()) / f"{audio_name}.mp3"
+        wav_suffix = "wav" if self.stream.isChecked() else "mp3"
+        audio_path = Path(self.save_audio_path.text()) / f"{audio_name}.{wav_suffix}"
         audio_path.parent.mkdir(parents=True, exist_ok=True)
         self.audio_path = str(audio_path)
         kwargs = dict(
@@ -926,6 +932,7 @@ def start_conversion(self):
             max_new_tokens=self.max_new_tokens_slider.value(),
             temperature=self.temperature_slider.value() / 1000.0,
             mp3_bitrate=int(self.mp3_bitrate_combo.currentText()),
+            stream=self.stream.isChecked(),
         )
         self.tts_worker = TTSWorker(
             ref_files=self.files,
@@ -936,17 +943,22 @@ def start_conversion(self):
             audio_path=str(audio_path),
             **kwargs,
         )
-        self.tts_worker.finished.connect(self.on_conversion_finished)
+        self.tts_worker.finished_signal.connect(self.on_conversion_finished)
+        self.tts_worker.packet_delay.connect(
+            lambda t: self.latency_label.setText(
+                _t("action.latency").format(latency=(t * 1000.0))
+            )
+        )
         self.tts_worker.start()
 
-        self.now_audio.setText(_t("action.audio").format(audio_name=str(audio_path)))
-
     def stop_conversion(self):
         self.tts_worker.stop()
-        self.tts_worker.wait()
+        # self.tts_worker.wait()
         self.start_button.setEnabled(True)
         self.stop_button.setEnabled(False)
 
-    def on_conversion_finished(self):
-        self.stop_conversion()
-        self.set_audio(self.audio_path)()
+    def on_conversion_finished(self, audio_path):
+        self.now_audio.setText(_t("action.audio").format(audio_name=audio_path))
+        self.set_audio(self.audio_path)
+        self.start_button.setEnabled(True)
+        self.stop_button.setEnabled(False)

fish/modules/log.py

@@ -0,0 +1,24 @@
+import logging
+import sys
+
+from fish.modules.console import ConsoleStream
+
+# Redefined Stream
+stdout_stream = ConsoleStream()
+stderr_stream = ConsoleStream()
+
+sys.stdout = stdout_stream
+sys.stderr = stderr_stream
+
+# Global logger
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+stdout_handler = logging.StreamHandler(sys.stdout)
+stdout_handler.setLevel(logging.INFO)
+stderr_handler = logging.StreamHandler(sys.stdout)
+stderr_handler.setLevel(logging.WARNING)
+formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+stdout_handler.setFormatter(formatter)
+stderr_handler.setFormatter(formatter)
+logger.addHandler(stdout_handler)
+logger.addHandler(stderr_handler)

fish/modules/worker.py

@@ -3,13 +3,17 @@
 os.environ["no_proxy"] = "localhost, 127.0.0.1, 0.0.0.0"
 import re
 import subprocess
+import time
+import wave
 from pathlib import Path
 
 import httpx
 import ormsgpack
 import psutil
-from PyQt6.QtCore import QMutex, QMutexLocker, QThread, QWaitCondition, pyqtSignal
+import pyaudio
+from PyQt6.QtCore import QMutex, QMutexLocker, QThread, pyqtSignal
 
+from fish.modules.log import logger
 from fish.utils.audio import ServeReferenceAudio, ServeTTSRequest
 from fish.utils.i18n import _t
 
@@ -106,8 +110,26 @@ def terminate_process(self):
         self.process = None
 
 
-class TTSWorker(QThread):
-    finished = pyqtSignal()
+class TimeWorker(QThread):
+    time_signal = pyqtSignal(float)
+
+    def __init__(self, pause_time=0.1, parent=None):
+        super().__init__(parent)
+        self.start_time = time.time()
+        self._stop_requested = False
+        self.pause_time = pause_time
+
+    def run(self):
+        while not self._stop_requested:
+            time.sleep(self.pause_time)
+            self.time_signal.emit(time.time() - self.start_time)
+
+    def stop(self):
+        self._stop_requested = True
+
+
+class TTSWorker(BaseWorker):
+    packet_delay = pyqtSignal(float)
 
     def __init__(
         self,
@@ -120,26 +142,91 @@ def __init__(
         **kwargs,
     ):
         super().__init__()
-        self.mutex = QMutex()
-        self.wait_condition = QWaitCondition()
-        self._stop_requested = False
+
         self.ref_files = ref_files
         self.ref_id = ref_id if len(ref_id) > 0 else None
         self.backend = backend
         self.text = text
         self.api_key = api_key
         self.audio_path = audio_path
         self.kwargs = kwargs
+        self.time_worker = TimeWorker(pause_time=0.1)
+        self.time_worker.time_signal.connect(self.calc_elapsed)
+        self.elapsed = 0
+
+    def calc_elapsed(self, elapsed):
+        self.elapsed = elapsed
+        self.packet_delay.emit(elapsed)
 
     def run(self):
-        pre_files = [f for f in self.ref_files if not f.endswith(".lab")]
-        audio_files = [
+        self._process_audio_stream()
+
+    def _process_audio_stream(self):
+        pre_files = self.get_pre_files()
+        audio_files = self.filter_audio_files(pre_files)
+        streaming = self.kwargs.get("stream", False)
+        request = self.create_tts_request(audio_files, streaming)
+        frames_per_buffer = 16384
+        first_packet_time = None
+
+        self.time_worker.start()
+
+        if streaming:
+            p, stream = self.initialize_audio_stream(frames_per_buffer)
+            self.p = p
+            self.stream = stream
+            f = wave.open(self.audio_path, "wb")
+            f.setnchannels(1)
+            f.setsampwidth(2)
+            f.setframerate(44100)
+        else:
+            f = open(self.audio_path, "wb")
+
+        self.f = f
+        with httpx.Client() as client:
+            with client.stream(
+                "POST",
+                self.backend,
+                content=ormsgpack.packb(
+                    request, option=ormsgpack.OPT_SERIALIZE_PYDANTIC
+                ),
+                headers={
+                    "authorization": f"Bearer {self.api_key}",
+                    "content-type": "application/msgpack",
+                },
+                timeout=None,
+            ) as response:
+                for chunk in response.iter_bytes(chunk_size=frames_per_buffer):
+                    if first_packet_time is None:
+                        first_packet_time = self.elapsed
+                        self.time_worker.stop()
+
+                    if self.is_interrupted:
+                        return
+
+                    if streaming:
+                        stream.write(chunk)
+                        f.writeframesraw(chunk)
+                    else:
+                        f.write(chunk)
+
+        self.finish()
+
+        if not self.is_interrupted:
+            self.finished_signal.emit(self.audio_path)
+
+    def get_pre_files(self):
+        return [f for f in self.ref_files if not f.endswith(".lab")]
+
+    def filter_audio_files(self, pre_files: list):
+        return [
             f
             for f in pre_files
             if Path(f).exists() and Path(f).with_suffix(".lab").exists()
         ]
 
-        request = ServeTTSRequest(
+    def create_tts_request(self, audio_files: list, streaming: bool):
+        return ServeTTSRequest(
             text=self.text,
             references=[
                 ServeReferenceAudio(
@@ -149,42 +236,40 @@ def run(self):
                 for f in audio_files
             ],
             reference_id=self.ref_id,
-            streaming=False,
-            format="mp3",
-            chunk_length=self.kwargs["chunk_length"],
-            top_p=self.kwargs["top_p"],
-            repetition_penalty=self.kwargs["repetition_penalty"],
-            max_new_tokens=self.kwargs["max_new_tokens"],
-            temperature=self.kwargs["temperature"],
-            mp3_bitrate=self.kwargs["mp3_bitrate"],
+            streaming=streaming,
+            format="wav" if streaming else "mp3",
+            chunk_length=self.kwargs.get("chunk_length"),
+            top_p=self.kwargs.get("top_p"),
+            repetition_penalty=self.kwargs.get("repetition_penalty"),
+            max_new_tokens=self.kwargs.get("max_new_tokens"),
+            temperature=self.kwargs.get("temperature"),
+            mp3_bitrate=self.kwargs.get("mp3_bitrate"),
         )
 
-        with httpx.Client() as client, open(f"{self.audio_path}", "wb") as f:
-            with client.stream(
-                "POST",
-                self.backend,
-                content=ormsgpack.packb(
-                    request, option=ormsgpack.OPT_SERIALIZE_PYDANTIC
-                ),
-                headers={
-                    "authorization": f"Bearer {self.api_key}",
-                    "content-type": "application/msgpack",
-                },
-                timeout=None,
-            ) as response:
-                for chunk in response.iter_bytes():
-                    self.mutex.lock()
-                    if self._stop_requested:
-                        print("TTS is interrupted!")
-                        self.mutex.unlock()
-                        break
-                    self.mutex.unlock()
-                    f.write(chunk)
-
-        self.finished.emit()
+    def initialize_audio_stream(self, frames_per_buffer: int):
+        p = pyaudio.PyAudio()
+        stream = p.open(
+            format=pyaudio.paInt16,
+            channels=1,
+            rate=44100,
+            output=True,
+            frames_per_buffer=frames_per_buffer,
+        )
+        return p, stream
 
     def stop(self):
-        self.mutex.lock()
-        self._stop_requested = True
-        self.mutex.unlock()
-        self.wait_condition.wakeAll()
+        self.is_interrupted = True
+        logger.info("Stop requested!")
+        self.finish()
+
+    def finish(self):
+        streaming = self.kwargs.get("stream", False)
+        if streaming:
+            self.stream.stop_stream()
+            self.stream.close()
+            self.p.terminate()
+            logger.warning("Stop streaming!")
+        self.time_worker.stop()
+        logger.info("Timer off!")
+        self.f.close()
+        logger.info("File closed!")

locales/en_US.yaml

@@ -82,7 +82,7 @@ tts_output:
 
 action:
   audio: "Now playing: {audio_name}"
-  toggle_console: "Open/Close Console"
+  stream: "Streaming"
   start: "Start Text To Speech"
   stop: "Stop Text To Speech"
   latency: "Latency: {latency:.2f} ms"