TTS and ASR function implemented based on openai api

2023-11-22 19:00:41 +08:00
parent 4c80586630
commit b305a963ba
11 changed files with 551 additions and 139 deletions
@@ -0,0 +1,120 @@
+import asyncio
+import io
+import logging
+import os
+from asyncio import Queue
+
+from openai import AsyncOpenAI
+
+from aios_kernel import ComputeNode, ComputeTask, ComputeTaskState, ComputeTaskResult, ComputeTaskType, AIStorage
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAITTSComputeNode(ComputeNode):
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        super().__init__()
+        self.is_start = False
+        self.node_id = "openai_tts_node"
+        self.task_queue = Queue()
+        self.voice_list = {
+            "female": ["nova", "shimmer"],
+            "man": ["alloy", "echo", "fable", "onyx"]
+        }
+        if os.getenv("OPENAI_API_KEY") is not None:
+            self.openai_api_key = os.getenv("OPENAI_API_KEY")
+        else:
+            self.openai_api_key = AIStorage.get_instance().get_user_config().get_value("openai_api_key")
+
+        self.start()
+
+    def start(self):
+        if self.is_start is True:
+            logger.warn("OpenAITTSComputeNode is already start")
+            return
+        self.is_start = True
+
+        async def _run_task_loop():
+            while True:
+                task = await self.task_queue.get()
+                try:
+                    result = await self._run_task(task)
+                    if result is not None:
+                        task.state = ComputeTaskState.DONE
+                        task.result = result
+                except Exception as e:
+                    logger.error(f"openai_tts_node run task error: {e}")
+                    task.state = ComputeTaskState.ERROR
+                    task.result = ComputeTaskResult()
+                    task.result.set_from_task(task)
+                    task.result.worker_id = self.node_id
+                    task.result.result_str = str(e)
+
+        asyncio.create_task(_run_task_loop())
+
+    async def _run_task(self,task: ComputeTask):
+        task.state = ComputeTaskState.RUNNING
+        text = task.params["text"]
+        voice_name = task.params["voice_name"]
+        if voice_name is None:
+            voice_name = "default"
+        gender = task.params["gender"]
+        if gender is None:
+            gender = "female"
+
+        voice_list = self.voice_list[gender]
+        voice = voice_list[hash(voice_name)%len(voice_list)]
+
+        model_name = task.params['model_name']
+        if model_name is None:
+            model_name = 'tts-1'
+
+        client = AsyncOpenAI(api_key=self.openai_api_key)
+
+        response = await client.audio.speech.create(model=model_name, voice=voice, input=text)
+
+        cache = io.BytesIO()
+        async for data in await response.aiter_bytes():
+            cache.write(data)
+
+        cache.seek(0)
+
+        result = ComputeTaskResult()
+        result.set_from_task(task)
+        result.worker_id = self.node_id
+        result.result = cache.read()
+        return result
+
+    async def push_task(self, task: ComputeTask, proiority: int = 0):
+        logger.info(f"openai_tts_node push task: {task.display()}")
+        self.task_queue.put_nowait(task)
+
+    async def remove_task(self, task_id: str):
+        pass
+
+    def get_task_state(self, task_id: str):
+        pass
+
+    def display(self) -> str:
+        return f"OpenAITTSComputeNode: {self.node_id}"
+
+    def get_capacity(self):
+        return 0
+
+    def is_support(self, task: ComputeTask) -> bool:
+        if task.task_type == ComputeTaskType.TEXT_2_VOICE:
+            if task.params['model_name'] is None or task.params['model_name'] == 'tts-1' or task.params['model_name'] == 'tts-1-hd':
+                return True
+        return False
+
+
+    def is_local(self) -> bool:
+        return False