Added chatterbox connections

labs/glitch_odin/chatterbox_worker.py

@@ -0,0 +1,51 @@
+"""Persistent local GPU voice worker. JSON lines on stdin/stdout; logs on stderr."""
+import contextlib
+import json
+from pathlib import Path
+import sys
+
+
+def main():
+    protocol = sys.stdout
+    with contextlib.redirect_stdout(sys.stderr):
+        import numpy as np
+        import soundfile as sf
+        import torch
+        from chatterbox.tts_turbo import ChatterboxTurboTTS
+        torch.set_num_threads(4)
+        if not torch.cuda.is_available():
+            raise RuntimeError("Chatterbox lab requires the configured NVIDIA GPU")
+        model = ChatterboxTurboTTS.from_pretrained(device="cuda")
+        model.prepare_conditionals(sys.argv[1])
+    print(json.dumps({"ready": True}), file=protocol, flush=True)
+    for line in sys.stdin:
+        try:
+            request = json.loads(line)
+            with contextlib.redirect_stdout(sys.stderr):
+                # Sentence-sized chunks avoid Turbo's long-utterance truncation.
+                import re
+                chunks = re.split(r"(?<=[.!?])\s+", request["text"].strip())
+                parts = []
+                for chunk in chunks:
+                    if not chunk:
+                        continue
+                    # Bound even punctuation-free custom hints.
+                    words = chunk.split()
+                    for offset in range(0, len(words), 45):
+                        samples = model.generate(" ".join(words[offset:offset+45])).squeeze().detach().cpu().numpy()
+                        if parts:
+                            parts.append(np.zeros(int(model.sr * .18), dtype=np.float32))
+                        parts.append(samples)
+                if not parts:
+                    raise ValueError("Hint text is empty")
+                output = Path(request["output"])
+                temp = output.with_suffix(".partial")
+                sf.write(temp, np.concatenate(parts), model.sr, format="WAV", subtype="PCM_16")
+                temp.replace(output)
+            print(json.dumps({"ok": True}), file=protocol, flush=True)
+        except Exception as error:
+            print(json.dumps({"ok": False, "error": str(error)}), file=protocol, flush=True)
+
+
+if __name__ == "__main__":
+    main()