Merge branch 'dev' of github.com:litagin02/Style-Bert-VITS2

style_bert_vits2/constants.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 
 # Style-Bert-VITS2 のバージョン
-VERSION = "2.3.1"
+VERSION = "2.4"
 
 # Style-Bert-VITS2 のベースディレクトリ
 BASE_DIR = Path(__file__).parent.parent

style_bert_vits2/text_processing/japanese/g2p.py

@@ -1,13 +1,15 @@
-import pyopenjtalk
 import re
 
 from style_bert_vits2.constants import Languages
 from style_bert_vits2.logging import logger
 from style_bert_vits2.text_processing import bert_models
+from style_bert_vits2.text_processing.japanese import pyopenjtalk_worker as pyopenjtalk
 from style_bert_vits2.text_processing.japanese.mora_list import MORA_KATA_TO_MORA_PHONEMES
 from style_bert_vits2.text_processing.japanese.normalizer import replace_punctuation
 from style_bert_vits2.text_processing.symbols import PUNCTUATIONS
 
+pyopenjtalk.initialize()
+
 
 def g2p(
     norm_text: str,

style_bert_vits2/text_processing/japanese/pyopenjtalk_worker/__init__.py

@@ -0,0 +1,126 @@
+"""
+Run the pyopenjtalk worker in a separate process
+to avoid user dictionary access error
+"""
+
+from typing import Any, Optional
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.text_processing.japanese.pyopenjtalk_worker.worker_client import WorkerClient
+from style_bert_vits2.text_processing.japanese.pyopenjtalk_worker.worker_common import WORKER_PORT
+
+
+WORKER_CLIENT: Optional[WorkerClient] = None
+
+
+# pyopenjtalk interface
+# g2p(): not used
+
+
+def run_frontend(text: str) -> list[dict[str, Any]]:
+    assert WORKER_CLIENT
+    ret = WORKER_CLIENT.dispatch_pyopenjtalk("run_frontend", text)
+    assert isinstance(ret, list)
+    return ret
+
+
+def make_label(njd_features: Any) -> list[str]:
+    assert WORKER_CLIENT
+    ret = WORKER_CLIENT.dispatch_pyopenjtalk("make_label", njd_features)
+    assert isinstance(ret, list)
+    return ret
+
+
+def mecab_dict_index(path: str, out_path: str, dn_mecab: Optional[str] = None):
+    assert WORKER_CLIENT
+    WORKER_CLIENT.dispatch_pyopenjtalk("mecab_dict_index", path, out_path, dn_mecab)
+
+
+def update_global_jtalk_with_user_dict(path: str):
+    assert WORKER_CLIENT
+    WORKER_CLIENT.dispatch_pyopenjtalk("update_global_jtalk_with_user_dict", path)
+
+
+def unset_user_dict():
+    assert WORKER_CLIENT
+    WORKER_CLIENT.dispatch_pyopenjtalk("unset_user_dict")
+
+
+# initialize module when imported
+
+
+def initialize(port: int = WORKER_PORT) -> None:
+    import time
+    import socket
+    import sys
+    import atexit
+    import signal
+
+    logger.debug("initialize")
+    global WORKER_CLIENT
+    if WORKER_CLIENT:
+        return
+
+    client = None
+    try:
+        client = WorkerClient(port)
+    except (socket.timeout, socket.error):
+        logger.debug("try starting pyopenjtalk worker server")
+        import os
+        import subprocess
+
+        worker_pkg_path = os.path.relpath(
+            os.path.dirname(__file__), os.getcwd()
+        ).replace(os.sep, ".")
+        args = [sys.executable, "-m", worker_pkg_path, "--port", str(port)]
+        # new session, new process group
+        if sys.platform.startswith("win"):
+            cf = subprocess.CREATE_NEW_CONSOLE | subprocess.CREATE_NEW_PROCESS_GROUP  # type: ignore
+            si = subprocess.STARTUPINFO()  # type: ignore
+            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW  # type: ignore
+            si.wShowWindow = subprocess.SW_HIDE  # type: ignore
+            subprocess.Popen(args, creationflags=cf, startupinfo=si)
+        else:
+            # align with Windows behavior
+            # start_new_session is same as specifying setsid in preexec_fn
+            subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)  # type: ignore
+
+        # wait until server listening
+        count = 0
+        while True:
+            try:
+                client = WorkerClient(port)
+                break
+            except socket.error:
+                time.sleep(1)
+                count += 1
+                # 10: max number of retries
+                if count == 10:
+                    raise TimeoutError("サーバーに接続できませんでした")
+
+    WORKER_CLIENT = client
+    atexit.register(terminate)
+
+    # when the process is killed
+    def signal_handler(signum: int, frame: Any):
+        terminate()
+
+    signal.signal(signal.SIGTERM, signal_handler)
+
+
+# top-level declaration
+def terminate() -> None:
+    logger.debug("terminate")
+    global WORKER_CLIENT
+    if not WORKER_CLIENT:
+        return
+
+    # repare for unexpected errors
+    try:
+        if WORKER_CLIENT.status() == 1:
+            WORKER_CLIENT.quit_server()
+    except Exception as e:
+        logger.error(e)
+
+    WORKER_CLIENT.close()
+    WORKER_CLIENT = None

style_bert_vits2/text_processing/japanese/pyopenjtalk_worker/__main__.py

@@ -0,0 +1,16 @@
+import argparse
+
+from style_bert_vits2.text_processing.japanese.pyopenjtalk_worker.worker_common import WORKER_PORT
+from style_bert_vits2.text_processing.japanese.pyopenjtalk_worker.worker_server import WorkerServer
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=WORKER_PORT)
+    args = parser.parse_args()
+    server = WorkerServer()
+    server.start_server(port=args.port)
+
+
+if __name__ == "__main__":
+    main()

style_bert_vits2/text_processing/japanese/pyopenjtalk_worker/worker_client.py

@@ -0,0 +1,63 @@
+import socket
+from typing import Any, cast
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.text_processing.japanese.pyopenjtalk_worker.worker_common import RequestType, receive_data, send_data
+
+
+class WorkerClient:
+    """ pyopenjtalk worker client """
+
+
+    def __init__(self, port: int) -> None:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        # 5: timeout
+        sock.settimeout(5)
+        sock.connect((socket.gethostname(), port))
+        self.sock = sock
+
+
+    def __enter__(self) -> "WorkerClient":
+        return self
+
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+
+    def close(self) -> None:
+        self.sock.close()
+
+
+    def dispatch_pyopenjtalk(self, func: str, *args: Any, **kwargs: Any) -> Any:
+        data = {
+            "request-type": RequestType.PYOPENJTALK,
+            "func": func,
+            "args": args,
+            "kwargs": kwargs,
+        }
+        logger.trace(f"client sends request: {data}")
+        send_data(self.sock, data)
+        logger.trace("client sent request successfully")
+        response = receive_data(self.sock)
+        logger.trace(f"client received response: {response}")
+        return response.get("return")
+
+
+    def status(self) -> int:
+        data = {"request-type": RequestType.STATUS}
+        logger.trace(f"client sends request: {data}")
+        send_data(self.sock, data)
+        logger.trace("client sent request successfully")
+        response = receive_data(self.sock)
+        logger.trace(f"client received response: {response}")
+        return cast(int, response.get("client-count"))
+
+
+    def quit_server(self) -> None:
+        data = {"request-type": RequestType.QUIT_SERVER}
+        logger.trace(f"client sends request: {data}")
+        send_data(self.sock, data)
+        logger.trace("client sent request successfully")
+        response = receive_data(self.sock)
+        logger.trace(f"client received response: {response}")

style_bert_vits2/text_processing/japanese/pyopenjtalk_worker/worker_common.py

@@ -0,0 +1,45 @@
+import json
+import socket
+from enum import IntEnum, auto
+from typing import Any, Final
+
+
+WORKER_PORT: Final[int] = 7861
+HEADER_SIZE: Final[int] = 4
+
+
+class RequestType(IntEnum):
+    STATUS = auto()
+    QUIT_SERVER = auto()
+    PYOPENJTALK = auto()
+
+
+class ConnectionClosedException(Exception):
+    pass
+
+
+# socket communication
+
+
+def send_data(sock: socket.socket, data: dict[str, Any]):
+    json_data = json.dumps(data).encode()
+    header = len(json_data).to_bytes(HEADER_SIZE, byteorder="big")
+    sock.sendall(header + json_data)
+
+
+def __receive_until(sock: socket.socket, size: int):
+    data = b""
+    while len(data) < size:
+        part = sock.recv(size - len(data))
+        if part == b"":
+            raise ConnectionClosedException("接続が閉じられました")
+        data += part
+
+    return data
+
+
+def receive_data(sock: socket.socket) -> dict[str, Any]:
+    header = __receive_until(sock, HEADER_SIZE)
+    data_length = int.from_bytes(header, byteorder="big")
+    body = __receive_until(sock, data_length)
+    return json.loads(body.decode())

style_bert_vits2/text_processing/japanese/pyopenjtalk_worker/worker_server.py

@@ -0,0 +1,126 @@
+import select
+import socket
+import time
+from typing import Any, cast
+
+import pyopenjtalk
+
+from style_bert_vits2.logging import logger
+from style_bert_vits2.text_processing.japanese.pyopenjtalk_worker.worker_common import (
+    ConnectionClosedException,
+    RequestType,
+    receive_data,
+    send_data,
+)
+
+
+# To make it as fast as possible
+# Probably faster than calling getattr every time
+__PYOPENJTALK_FUNC_DICT = {
+    "run_frontend": pyopenjtalk.run_frontend,
+    "make_label": pyopenjtalk.make_label,
+    "mecab_dict_index": pyopenjtalk.mecab_dict_index,
+    "update_global_jtalk_with_user_dict": pyopenjtalk.update_global_jtalk_with_user_dict,
+    "unset_user_dict": pyopenjtalk.unset_user_dict,
+}
+
+
+class WorkerServer:
+    """ pyopenjtalk worker server """
+
+
+    def __init__(self) -> None:
+        self.client_count: int = 0
+        self.quit: bool = False
+
+
+    def handle_request(self, request: dict[str, Any]) -> dict[str, Any]:
+        request_type = None
+        try:
+            request_type = RequestType(cast(int, request.get("request-type")))
+        except Exception:
+            return {
+                "success": False,
+                "reason": "request-type is invalid",
+            }
+
+        response: dict[str, Any] = {}
+        if request_type:
+            if request_type == RequestType.STATUS:
+                response = {
+                    "success": True,
+                    "client-count": self.client_count,
+                }
+            elif request_type == RequestType.QUIT_SERVER:
+                self.quit = True
+                response = {"success": True}
+            elif request_type == RequestType.PYOPENJTALK:
+                func_name = request.get("func")
+                assert isinstance(func_name, str)
+                func = __PYOPENJTALK_FUNC_DICT[func_name]
+                args = request.get("args")
+                kwargs = request.get("kwargs")
+                assert isinstance(args, list)
+                assert isinstance(kwargs, dict)
+                ret = func(*args, **kwargs)
+                response = {"success": True, "return": ret}
+            else:
+                # NOT REACHED
+                response = request
+
+        return response
+
+
+    def start_server(self, port: int, no_client_timeout: int = 30) -> None:
+        logger.info("start pyopenjtalk worker server")
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
+            server_socket.bind((socket.gethostname(), port))
+            server_socket.listen()
+            sockets = [server_socket]
+            no_client_since = time.time()
+            while True:
+                if self.client_count == 0:
+                    if no_client_since is None:
+                        no_client_since = time.time()
+                    elif (time.time() - no_client_since) > no_client_timeout:
+                        logger.info("quit because there is no client")
+                        return
+                else:
+                    no_client_since = None
+
+                ready_sockets, _, _ = select.select(sockets, [], [], 0.1)
+                for sock in ready_sockets:
+                    if sock is server_socket:
+                        logger.info("new client connected")
+                        client_socket, _ = server_socket.accept()
+                        sockets.append(client_socket)
+                        self.client_count += 1
+                    else:
+                        # client
+                        try:
+                            request = receive_data(sock)
+                        except Exception as e:
+                            sock.close()
+                            sockets.remove(sock)
+                            self.client_count -= 1
+                            # unexpected disconnections
+                            if not isinstance(e, ConnectionClosedException):
+                                logger.error(e)
+
+                            logger.info("close connection")
+                            continue
+
+                        logger.trace(f"server received request: {request}")
+
+                        response = self.handle_request(request)
+                        logger.trace(f"server sends response: {response}")
+                        try:
+                            send_data(sock, response)
+                            logger.trace("server sent response successfully")
+                        except Exception:
+                            logger.warning(
+                                "an exception occurred during sending responce"
+                            )
+                        if self.quit:
+                            logger.info("quit pyopenjtalk worker server")
+                            return

style_bert_vits2/text_processing/japanese/user_dict/__init__.py

@@ -13,15 +13,15 @@ from typing import Dict, List, Optional
 from uuid import UUID, uuid4
 
 import numpy as np
-import pyopenjtalk
 from fastapi import HTTPException
 
 from style_bert_vits2.constants import DEFAULT_USER_DICT_DIR
+from style_bert_vits2.text_processing.japanese import pyopenjtalk_worker as pyopenjtalk
 from style_bert_vits2.text_processing.japanese.user_dict.word_model import UserDictWord, WordTypes
-# from ..utility.mutex_utility import mutex_wrapper
-# from ..utility.path_utility import engine_root, get_save_dir
 from style_bert_vits2.text_processing.japanese.user_dict.part_of_speech_data import MAX_PRIORITY, MIN_PRIORITY, part_of_speech_data
 
+pyopenjtalk.initialize()
+
 # root_dir = engine_root()
 # save_dir = get_save_dir()