52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
"""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()
|