From d1be3b2aa79e4a4962f5e2e8d72817b99785409f Mon Sep 17 00:00:00 2001
From: gdw6463 <gdw6463@gmail.com>
Date: Mon, 20 May 2024 17:39:31 +0800
Subject: [PATCH] Blackbox g2e models updated (using QWen), Blackbox melotts
 updated with local and docker mode

---
 README.md               |  7 +++++-
 src/blackbox/g2e.py     |  4 +--
 src/blackbox/melotts.py | 54 ++++++++++++++++++++++++++++++++---------
 src/configuration.py    | 14 +++++++++--
 4 files changed, 62 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 07fc7c3..a2f274e 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,12 @@ log:
   filename: "D:/Workspace/Logging/jarvis/jarvis-models.log"
 
 melotts:
-  url: http://{IP running docker melotts-api}:18080/convert/tts
+  mode: local # or docker
+  url: http://10.6.44.16:18080/convert/tts
+  speed: 0.9
+  device: 'cuda'
+  language: 'ZH'
+  speaker: 'ZH'
 
 tesou:
   url: http://120.196.116.194:48891/chat/
diff --git a/src/blackbox/g2e.py b/src/blackbox/g2e.py
index de4a8c8..416ab72 100755
--- a/src/blackbox/g2e.py
+++ b/src/blackbox/g2e.py
@@ -66,8 +66,8 @@ class G2E(Blackbox):
             api_key='YOUR_API_KEY',
             base_url=url
         )
-        #model_name = client.models.list().data[0].id
-        model_name = client.models.list().data[1].id
+        model_name = client.models.list().data[0].id
+        #model_name = client.models.list().data[1].id
         print(model_name)
 
         response = client.chat.completions.create(
diff --git a/src/blackbox/melotts.py b/src/blackbox/melotts.py
index 76a07e1..551ed8e 100644
--- a/src/blackbox/melotts.py
+++ b/src/blackbox/melotts.py
@@ -10,14 +10,38 @@ from injector import singleton
 from ..configuration import MeloConf
 from .blackbox import Blackbox
 
+import soundfile
+from melo.api import TTS
+
+import logging
+logging.basicConfig(level=logging.INFO)
 
 @singleton
 class MeloTTS(Blackbox):
-    melotts: str
+    mode: str
+    url: str
+    speed: int
+    device: str
+    language: str
+    speaker: str
 
     @inject
     def __init__(self, melo_config: MeloConf) -> None:
-        self.melotts = melo_config.melotts
+        self.speed = melo_config.speed
+        self.device = melo_config.device
+        self.language = melo_config.language
+        self.speaker = melo_config.speaker
+        self.device = melo_config.device
+        self.url = ''
+        self.mode = melo_config.mode
+        self.melotts = None
+        self.speaker_ids = None
+        if self.mode == 'local':
+            self.melotts = TTS(language=self.language, device=self.device)
+            self.speaker_ids = self.melotts.hps.data.spk2id
+        else:
+            self.url = melo_config.url
+        logging.info('#### Initializing MeloTTS Service in ' + self.device + ' mode...')
 
     def __call__(self, *args, **kwargs):
         return self.processing(*args, **kwargs)
@@ -26,15 +50,23 @@ class MeloTTS(Blackbox):
         text = args[0]
         return isinstance(text, str)
 
-    def processing(self, *args, **kwargs) -> bytes:
+    def processing(self, *args, **kwargs) -> io.BytesIO | bytes:
         text = args[0]
         current_time = time.time()
-        message = {
-            "text": text
-        }
-        response = requests.post(self.melotts, json=message)
-        print("#### MeloTTS Service consume : ", (time.time()-current_time))
-        return response.content
+        if self.mode == 'local':
+            audio = self.melotts.tts_to_file(text, self.speaker_ids[self.speaker], speed=self.speed)
+            f = io.BytesIO()
+            soundfile.write(f, audio, 44100, format='wav')
+            f.seek(0)
+            print("#### MeloTTS Service consume - local : ", (time.time() - current_time))
+            return f.read()
+        else:
+            message = {
+                "text": text
+            }
+            response = requests.post(self.url, json=message)
+            print("#### MeloTTS Service consume - docker : ", (time.time()-current_time))
+            return response.content
 
     async def fast_api_handler(self, request: Request) -> Response:
         try:
@@ -44,6 +76,4 @@ class MeloTTS(Blackbox):
         text = data.get("text")
         if text is None:
             return JSONResponse(content={"error": "text is required"}, status_code=status.HTTP_400_BAD_REQUEST)
-        return Response(content=self.processing(text), media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=audio.wav"})
-
-
+        return Response(content=self.processing(text), media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=audio.wav"})
\ No newline at end of file
diff --git a/src/configuration.py b/src/configuration.py
index bb4e749..30be08f 100644
--- a/src/configuration.py
+++ b/src/configuration.py
@@ -48,11 +48,21 @@ class TesouConf():
 
 
 class MeloConf():
-    melotts: str
+    mode: str
+    url: str
+    speed: int
+    device: str
+    language: str
+    speaker: str
 
     @inject
     def __init__(self, config: Configuration) -> None:
-        self.melotts = config.get("melotts.url")
+        self.mode = config.get("melotts.mode")
+        self.url = config.get("melotts.url")
+        self.speed = config.get("melotts.speed")
+        self.device = config.get("melotts.device")
+        self.language = config.get("melotts.language")
+        self.speaker = config.get("melotts.speaker")
 
 # 'CRITICAL': CRITICAL,
 # 'FATAL': FATAL,