Add an asyncio tcp server

meissa/main.py

@@ -19,8 +19,19 @@ This module contains all the CLI subcommands and interfaces that Meissa
 provides for user interaction
 """
 import os
+import yaml
+from pathlib import Path
+
 import click
 
+from meissa import utils
+
+
+DEFAULT_MODEL = (
+    "https://coqui.gateway.scarf.sh/english/coqui/v1.0.0-huge-vocab/model.tflite"
+)
+DEFAULT_SCORER = "https://coqui.gateway.scarf.sh/english/coqui/v1.0.0-huge-vocab/huge-vocabulary.scorer"
+
 
 @click.group("CLI")
 @click.pass_context
@@ -30,8 +41,27 @@ def cli(ctx):
 
 @cli.command(help="Seutp the environment")
 @click.pass_context
-def setup(ctx, f, packages, arch):
-    click.echo("here")
+def setup(ctx):
+    config = utils.config(ctx)
+    home = utils.home(ctx)
+
+    model = config.get("model", DEFAULT_MODEL)
+    scorer = config.get("scorer", DEFAULT_SCORER)
+
+    # TODO: Add support for multiple profiles
+    utils.info(f"Downloading the model: {model}")
+    utils.download(model, home, "model")
+    utils.info(f"Downloading the scorer: {scorer}")
+    utils.download(scorer, home, "scorer")
+
+
+@cli.command(help="Starts the Meissa server")
+@click.pass_context
+def start(ctx):
+    try:
+        asyncio.run(start(config, synth(conf)))
+    except KeyboardInterrupt:
+        sys.exit()
 
 
 def main():
@@ -39,13 +69,21 @@ def main():
     The main entry point for Faraday
     """
     user_home = os.environ.get("HOME")
-    meissa_home = os.path.join(user_home, ".meissa")
+    meissa_home = Path(user_home) / ".meissa"
 
-    os.makedirs(meissa_home, exist_ok=True)
+    meissa_home.mkdir(exist_ok=True)
+
+    config_path = meissa_home / "config.yml"
+
+    config = {}
+    if config_path.exists():
+        with open(config_path, "r") as stream:
+            config = yaml.safe_load(stream)
 
     cli(
         obj={
             "home": meissa_home,
+            "config": config,
         }
     )
 

meissa/server.py

@@ -0,0 +1,119 @@
+# Meissa - A trainable and simple text to speech server
+#
+# Copyright (c) 2023 Sameer Rahmani <lxsameer@gnu.org>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 2.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import asyncio
+
+# pylint: disable=redefined-outer-name, unused-argument
+from pathlib import Path
+
+import numpy as np
+
+
+import simpleaudio as sa
+
+from TTS.utils.manage import ModelManager
+from TTS.utils.synthesizer import Synthesizer
+
+from meissa import utils
+
+
+def to_wav_data(wav):
+    wav_norm = np.array(wav) * (32767 / max(0.01, np.max(np.abs(wav))))
+    return wav_norm.astype(np.int16)
+
+
+def play(wav):
+    try:
+        play = sa.play_buffer(to_wav_data(wav), 1, 2, 22050)
+        # Wait for audio playback to finish before exiting
+        play.wait_done()
+
+    finally:
+        play.stop()
+
+
+def synth(ctx):
+    path = Path(__file__).parent / ".models.json"
+    manager = ModelManager(path)
+
+    language_ids_file_path = None
+    vocoder_path = None
+    vocoder_config_path = None
+    encoder_path = None
+    encoder_config_path = None
+
+    model_path, config_path, model_item = manager.download_model(
+        ctx.get("MODEL_NAME", "tts_models/en/ljspeech/tacotron2-DDC")
+    )
+    vocoder_name = model_item["default_vocoder"]
+    vocoder_path, vocoder_config_path, _ = manager.download_model(vocoder_name)
+    speaker_idx = ctx.get("SPEAKER_IDX")
+
+    # load models
+    synthesizer = Synthesizer(
+        model_path,
+        config_path,
+        None,
+        language_ids_file_path,
+        vocoder_path,
+        vocoder_config_path,
+        encoder_path,
+        encoder_config_path,
+        False,
+    )
+
+    async def tcp_handler(reader, writer):
+        while True:
+            # Read till EOF
+            data = await reader.read(4)
+            msg_len = int.from_bytes(data, "big")
+
+            if msg_len != 0:
+                data = await reader.read(msg_len)
+                message = data.decode()
+
+                print(f"Received {message!r}")
+
+                if message == "//close":
+                    break
+
+                wav = synthesizer.tts(message, speaker_idx, "None", None)
+                synthesizer.save_wav(wav, "/tmp/blah.wav")
+                play(wav)
+
+                writer.write(b"Ok")
+                await writer.drain()
+        writer.close()
+
+    return tcp_handler
+
+
+async def start(ctx, fn):
+    """
+    Start a TCP socket and pass the given connection handler function `fn`
+    to it. It uses the `host` and `port` in the config file for the server.
+    """
+    host = utils.config(ctx).get("host", "127.0.0.1")
+    port = utils.config(ctx).get("port", 6666)
+
+    server = await asyncio.start_server(fn, host, port)
+
+    addrs = ", ".join(str(sock.getsockname()) for sock in server.sockets)
+
+    utils.info(f"Serving on {addrs}")
+
+    async with server:
+        await server.serve_forever()

meissa/utils.py

@@ -0,0 +1,56 @@
+# Meissa - A trainable and simple text to speech server
+#
+# Copyright (c) 2023 Sameer Rahmani <lxsameer@gnu.org>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 2.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+from pathlib import Path
+
+import requests
+import click
+
+
+def log(level: str, msg: str):
+    click.echo(f"[{level}]: {msg}")
+
+
+def info(msg: str):
+    log(click.style("INFO", fg="green"), msg)
+
+
+def warn(msg: str):
+    log(click.style("WARN", fg="yellow"), msg)
+
+
+def error(msg: str):
+    log(click.style("ERR", fg="red"), msg)
+
+
+def home(ctx) -> Path:
+    return Path(ctx.obj["home"])
+
+
+def config(ctx) -> dict:
+    return ctx.obj["config"]
+
+
+def download(url: str, path: Path, fname: str = None) -> Path:
+    local_filename = url.split("/")[-1] if not fname else fname
+
+    output = path / local_filename
+    with requests.get(url, stream=True, timeout=120) as r:
+        r.raise_for_status()
+        with open(output, "wb") as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+    return output

poetry.lock

@@ -2723,6 +2723,23 @@ docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "pygments-g
 testing = ["build[virtualenv]", "filelock (>=3.4.0)", "flake8 (<5)", "flake8-2020", "ini2toml[lite] (>=0.9)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pip (>=19.1)", "pip-run (>=8.8)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-timeout", "pytest-xdist", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"]
 testing-integration = ["build[virtualenv]", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"]
 
+[[package]]
+name = "simpleaudio"
+version = "1.0.4"
+description = "Simple, asynchronous audio playback for Python 3."
+category = "main"
+optional = false
+python-versions = "*"
+files = [
+    {file = "simpleaudio-1.0.4-cp37-cp37m-macosx_10_6_intel.whl", hash = "sha256:05b63da515f5fc7c6f40e4d9673d22239c5e03e2bda200fc09fd21c185d73713"},
+    {file = "simpleaudio-1.0.4-cp37-cp37m-win32.whl", hash = "sha256:f1a4fe3358429b2ea3181fd782e4c4fff5c123ca86ec7fc29e01ee9acd8a227a"},
+    {file = "simpleaudio-1.0.4-cp37-cp37m-win_amd64.whl", hash = "sha256:86f1b0985629852afe67259ac6c24905ca731cb202a6e96b818865c56ced0c27"},
+    {file = "simpleaudio-1.0.4-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:f68820297ad51577e3a77369e7e9b23989d30d5ae923bf34c92cf983c04ade04"},
+    {file = "simpleaudio-1.0.4-cp38-cp38-win32.whl", hash = "sha256:67348e3d3ccbae73bd126beed7f1e242976889620dbc6974c36800cd286430fc"},
+    {file = "simpleaudio-1.0.4-cp38-cp38-win_amd64.whl", hash = "sha256:f346a4eac9cdbb1b3f3d0995095b7e86c12219964c022f4d920c22f6ca05fb4c"},
+    {file = "simpleaudio-1.0.4.tar.gz", hash = "sha256:691c88649243544db717e7edf6a9831df112104e1aefb5f6038a5d071e8cf41d"},
+]
+
 [[package]]
 name = "six"
 version = "1.16.0"
@@ -3424,4 +3441,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<3.11"
-content-hash = "bc0667c6cd290a2486cc1ffa2e92374d97f68692666b1559be91634de602e0ee"
+content-hash = "8d72955bdf3e95e400eb660e7f877c3e2abd29e209a8864979cc0981ad0fda3a"

pyproject.toml

@@ -11,6 +11,7 @@ python = ">=3.10,<3.11"
 tts = "^0.13.3"
 click = "^8.1.3"
 msgpack = "^1.0.5"
+simpleaudio = "^1.0.4"
 
 
 [tool.poetry.group.dev.dependencies]