101 lines
4.0 KiB
Python
101 lines
4.0 KiB
Python
"""Nonblocking, cancellable requests to a resident Chatterbox worker in WSL/Linux."""
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import queue
|
|
import shlex
|
|
import subprocess
|
|
import threading
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
REFERENCE = ROOT / ".speech/voices/teddy-chatterbox-reference.wav"
|
|
|
|
|
|
def linux_path(path):
|
|
path = Path(path).resolve()
|
|
if os.name != "nt":
|
|
return str(path)
|
|
return "/mnt/" + path.drive[0].lower() + path.as_posix()[2:]
|
|
|
|
|
|
def voice_id():
|
|
return "teddy-chatterbox-turbo-v1-" + hashlib.sha256(REFERENCE.read_bytes()).hexdigest()[:16]
|
|
|
|
|
|
class Job:
|
|
def __init__(self, text, output):
|
|
self.text, self.output = text, output
|
|
self.cancelled = threading.Event()
|
|
self.returncode = None
|
|
def poll(self):
|
|
return self.returncode
|
|
def terminate(self):
|
|
# In-flight GPU work may finish and populate cache, but must never play.
|
|
self.cancelled.set()
|
|
def wait(self, timeout=None):
|
|
return self.returncode
|
|
|
|
|
|
class Engine:
|
|
def __init__(self):
|
|
self.jobs = queue.Queue()
|
|
self.ready = threading.Event()
|
|
self.error = None
|
|
self.process = None
|
|
self.thread = threading.Thread(target=self._run, daemon=True)
|
|
self.thread.start()
|
|
def submit(self, text, output):
|
|
job = Job(text, output)
|
|
self.jobs.put(job)
|
|
return job
|
|
def close(self):
|
|
self.jobs.put(None)
|
|
def _run(self):
|
|
current = None
|
|
try:
|
|
home = os.environ.get("TEDDY_CHATTERBOX_HOME", "~/teddy-chatterbox")
|
|
# Expand the conventional home prefix in the Linux shell only.
|
|
home_arg = '"$HOME"' + shlex.quote(home[1:]) if home.startswith("~/") else shlex.quote(home)
|
|
script = linux_path(ROOT / "chatterbox_worker.py")
|
|
command = (f"export HF_HOME={home_arg}/models; export HF_HUB_DISABLE_XET=1; "
|
|
f"exec {home_arg}/venv/bin/python -u {shlex.quote(script)} {shlex.quote(linux_path(REFERENCE))}")
|
|
args = (["wsl", "-d", "Ubuntu", "--", "bash", "-lc", command] if os.name == "nt"
|
|
else ["bash", "-lc", command])
|
|
log = ROOT / ".speech/chatterbox.log"
|
|
log.parent.mkdir(parents=True, exist_ok=True)
|
|
with log.open("a", encoding="utf-8") as logs:
|
|
self.process = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
|
stderr=logs, text=True, encoding="utf-8",
|
|
creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0)
|
|
if not json.loads(self.process.stdout.readline()).get("ready"):
|
|
raise RuntimeError("Chatterbox did not become ready")
|
|
self.ready.set()
|
|
while True:
|
|
current = self.jobs.get()
|
|
if current is None:
|
|
break
|
|
if current.cancelled.is_set():
|
|
current.returncode = 1
|
|
continue
|
|
self.process.stdin.write(json.dumps({"text": current.text,
|
|
"output": linux_path(current.output)}) + "\n")
|
|
self.process.stdin.flush()
|
|
result = json.loads(self.process.stdout.readline())
|
|
current.returncode = 0 if result.get("ok") else 1
|
|
if not result.get("ok"):
|
|
print("Chatterbox synthesis: " + result.get("error", "unknown error"), flush=True)
|
|
current = None
|
|
except Exception as error:
|
|
self.error = str(error)
|
|
if current is not None:
|
|
current.returncode = 1
|
|
print("Chatterbox worker failed: " + self.error, flush=True)
|
|
finally:
|
|
if self.process:
|
|
self.process.stdin.close()
|
|
try:
|
|
self.process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
self.process.terminate()
|