Spaces:
Build error
Build error
import argparse | |
import asyncio | |
import queue | |
from multiprocessing import Pipe, Process | |
from multiprocessing.connection import Connection | |
from tempfile import NamedTemporaryFile | |
from typing import List, Optional, Tuple | |
import soundfile | |
# FIXME: remove FastAPI dependency | |
from fastapi import HTTPException, Request | |
from .model import AudioQuery | |
from .synthesis_engine import make_synthesis_engines | |
from .utility import get_latest_core_version | |
class CancellableEngine: | |
""" | |
音声合成のキャンセル機能に関するクラス | |
初期化後は、synthesis関数で音声合成できる | |
(オリジナルと比べ引数が増えているので注意) | |
Attributes | |
---------- | |
watch_con_list: List[Tuple[Request, Process]] | |
Requestは接続の監視に使用され、Processは通信切断時のプロセスキルに使用される | |
クライアントから接続があるとListにTupleが追加される | |
接続が切断、もしくは音声合成が終了すると削除される | |
procs_and_cons: queue.Queue[Tuple[Process, Connection]] | |
音声合成の準備が終わっているプロセスのList | |
(音声合成中のプロセスは入っていない) | |
""" | |
def __init__(self, args: argparse.Namespace) -> None: | |
""" | |
変数の初期化を行う | |
また、args.init_processesの数だけプロセスを起動し、procs_and_consに格納する | |
""" | |
self.args = args | |
if not self.args.enable_cancellable_synthesis: | |
raise HTTPException( | |
status_code=404, | |
detail="実験的機能はデフォルトで無効になっています。使用するには引数を指定してください。", | |
) | |
self.watch_con_list: List[Tuple[Request, Process]] = [] | |
self.procs_and_cons: queue.Queue[Tuple[Process, Connection]] = queue.Queue() | |
for _ in range(self.args.init_processes): | |
self.procs_and_cons.put(self.start_new_proc()) | |
def start_new_proc( | |
self, | |
) -> Tuple[Process, Connection]: | |
""" | |
新しく開始したプロセスを返す関数 | |
Returns | |
------- | |
ret_proc: Process | |
新規のプロセス | |
sub_proc_con1: Connection | |
ret_procのプロセスと通信するためのPipe | |
""" | |
sub_proc_con1, sub_proc_con2 = Pipe(True) | |
ret_proc = Process( | |
target=start_synthesis_subprocess, | |
kwargs={ | |
"args": self.args, | |
"sub_proc_con": sub_proc_con2, | |
}, | |
daemon=True, | |
) | |
ret_proc.start() | |
return ret_proc, sub_proc_con1 | |
def finalize_con( | |
self, | |
req: Request, | |
proc: Process, | |
sub_proc_con: Optional[Connection], | |
) -> None: | |
""" | |
接続が切断された時の処理を行う関数 | |
watch_con_listからの削除、プロセスの後処理を行う | |
プロセスが生きている場合はそのままprocs_and_consに加える | |
死んでいる場合は新しく生成したものをprocs_and_consに加える | |
Parameters | |
---------- | |
req: fastapi.Request | |
接続確立時に受け取ったものをそのまま渡せばよい | |
https://fastapi.tiangolo.com/advanced/using-request-directly/ | |
proc: Process | |
音声合成を行っていたプロセス | |
sub_proc_con: Connection, optional | |
音声合成を行っていたプロセスとのPipe | |
指定されていない場合、プロセスは再利用されず終了される | |
""" | |
try: | |
self.watch_con_list.remove((req, proc)) | |
except ValueError: | |
pass | |
try: | |
if not proc.is_alive() or sub_proc_con is None: | |
proc.close() | |
raise ValueError | |
# プロセスが死んでいない場合は再利用する | |
self.procs_and_cons.put((proc, sub_proc_con)) | |
except ValueError: | |
# プロセスが死んでいるので新しく作り直す | |
self.procs_and_cons.put(self.start_new_proc()) | |
def _synthesis_impl( | |
self, | |
query: AudioQuery, | |
speaker_id: int, | |
request: Request, | |
core_version: Optional[str], | |
) -> str: | |
""" | |
音声合成を行う関数 | |
通常エンジンの引数に比べ、requestが必要になっている | |
また、返り値がファイル名になっている | |
Parameters | |
---------- | |
query: AudioQuery | |
speaker_id: int | |
request: fastapi.Request | |
接続確立時に受け取ったものをそのまま渡せばよい | |
https://fastapi.tiangolo.com/advanced/using-request-directly/ | |
core_version: str | |
Returns | |
------- | |
f_name: str | |
生成された音声ファイルの名前 | |
""" | |
proc, sub_proc_con1 = self.procs_and_cons.get() | |
self.watch_con_list.append((request, proc)) | |
try: | |
sub_proc_con1.send((query, speaker_id, core_version)) | |
f_name = sub_proc_con1.recv() | |
except EOFError: | |
raise HTTPException(status_code=422, detail="既にサブプロセスは終了されています") | |
except Exception: | |
self.finalize_con(request, proc, sub_proc_con1) | |
raise | |
self.finalize_con(request, proc, sub_proc_con1) | |
return f_name | |
async def catch_disconnection(self): | |
""" | |
接続監視を行うコルーチン | |
""" | |
while True: | |
await asyncio.sleep(1) | |
for con in self.watch_con_list: | |
req, proc = con | |
if await req.is_disconnected(): | |
try: | |
if proc.is_alive(): | |
proc.terminate() | |
proc.join() | |
proc.close() | |
except ValueError: | |
pass | |
finally: | |
self.finalize_con(req, proc, None) | |
def start_synthesis_subprocess( | |
args: argparse.Namespace, | |
sub_proc_con: Connection, | |
): | |
""" | |
音声合成を行うサブプロセスで行うための関数 | |
pickle化の関係でグローバルに書いている | |
Parameters | |
---------- | |
args: argparse.Namespace | |
起動時に作られたものをそのまま渡す | |
sub_proc_con: Connection | |
メインプロセスと通信するためのPipe | |
""" | |
synthesis_engines = make_synthesis_engines( | |
use_gpu=args.use_gpu, | |
voicelib_dirs=args.voicelib_dir, | |
voicevox_dir=args.voicevox_dir, | |
runtime_dirs=args.runtime_dir, | |
cpu_num_threads=args.cpu_num_threads, | |
enable_mock=args.enable_mock, | |
) | |
assert len(synthesis_engines) != 0, "音声合成エンジンがありません。" | |
latest_core_version = get_latest_core_version(versions=synthesis_engines.keys()) | |
while True: | |
try: | |
query, speaker_id, core_version = sub_proc_con.recv() | |
if core_version is None: | |
_engine = synthesis_engines[latest_core_version] | |
elif core_version in synthesis_engines: | |
_engine = synthesis_engines[core_version] | |
else: | |
# バージョンが見つからないエラー | |
sub_proc_con.send("") | |
continue | |
wave = _engine._synthesis_impl(query, speaker_id) | |
with NamedTemporaryFile(delete=False) as f: | |
soundfile.write( | |
file=f, data=wave, samplerate=query.outputSamplingRate, format="WAV" | |
) | |
sub_proc_con.send(f.name) | |
except Exception: | |
sub_proc_con.close() | |
raise | |