story maker

src/aios_kernel/text_to_speech_function.py

@@ -0,0 +1,101 @@
+import io
+import logging
+import os
+import random
+from typing import Dict
+
+from aios_kernel import ComputeKernel
+from aios_kernel.ai_function import AIFunction
+
+from pydub import AudioSegment
+
+logger = logging.getLogger(__name__)
+
+class TextToSpeechFunction(AIFunction):
+    def __init__(self):
+        self.func_id = "text_to_speech"
+        self.description = "根据输入的剧本生成音频数据"
+
+    def get_name(self) -> str:
+        return self.func_id
+
+    def get_description(self) -> str:
+        return self.description
+
+    def get_parameters(self) -> Dict:
+        return {
+            "type": "object",
+            "properties": {
+                "language": {"type": "string", "description": "演播语言", "enum": ["zh", "en"]},
+                "roles": {"type": "array", "items": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string", "description": "角色名字"},
+                        "gender": {"type": "string", "description": "角色性别", "enum": ["man", "female"]},
+                        "age": {"type": "string", "description": "年龄", "enum": ["child", "adult"]},
+                    }}},
+                "lines": {"type": "array", "items": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string", "description": "角色名字"},
+                        "tone": {"type": "string", "description": "演播情感",
+                                 "enum": ["happy", "sad", "angry", "fear", "disgust", "surprise", "neutral"]},
+                        "text": {"type": "string", "description": "台词"},
+                    }
+                }}
+            }
+        }
+
+    async def execute(self, **kwargs) -> str:
+        logger.info(f"execute text_to_speech function: {kwargs}")
+
+        language = kwargs.get("language")
+        if language is None:
+            language = "zh"
+        roles = kwargs.get("roles")
+        lines = kwargs.get("lines")
+
+        audio = None
+        for line in lines:
+            name = line.get("name")
+            tone = line.get("tone")
+            text = line.get("text")
+            gender = None
+            age = None
+            for role in roles:
+                role_name = role.get("name")
+                if role_name == name:
+                    gender = role.get("gender")
+                    age = role.get("age")
+                    break
+            i = 0
+            while i < 3:
+                try:
+                    data = await ComputeKernel.get_instance().do_text_to_speech(text, language, gender, age, name, tone)
+                    if audio is None:
+                        audio = AudioSegment.from_mp3(io.BytesIO(data))
+                    else:
+                        audio = audio + AudioSegment.from_mp3(io.BytesIO(data))
+                    break
+                except Exception as e:
+                    logger.error(f"do_text_to_speech failed: {e}")
+                    i += 1
+                    continue
+
+        if audio is not None:
+            path = os.path.join(os.curdir, "{}.mp3".format(random.sample('zyxwvutsrqponmlkjihgfedcba', 10)))
+            audio.export(path, format="mp3")
+            return "complete.file path:{}".format(path)
+        else:
+            return "failed"
+
+    def is_local(self) -> bool:
+        return True
+
+    def is_in_zone(self) -> bool:
+        return True
+
+    def is_ready_only(self) -> bool:
+        return False
+
+