Added chatterbox connections

labs/glitch_odin/chatterbox_client.py

@@ -0,0 +1,100 @@
+"""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()