old-books/_pipeline/voice.py

"""Voice client: script segments -> wav pieces, via an OpenAI-compatible /v1/audio/speech endpoint (URL in settings).
Pieces are sentence-sized. Output goes to scratch (work_dir), with a manifest the assembler reads."""
import json
import os
import re
import time
import urllib.error
import urllib.parse
import urllib.request

import bookio
import settings

ABBREV = re.compile(r"^(Dr|Mr|Mrs|Ms|St|Mt|Jr|Sr|vs|[A-Z])\.$")


def sentences(text):
    """Split at . ! ? (optionally followed by a quote) and a space, but not after abbreviations like 'Dr.' or 'E.'."""
    out, start = [], 0
    for m in re.finditer(r"[.!?][\"']?\s+", text):
        before = text[start:m.start() + 1].split()[-1] if text[start:m.start() + 1].split() else ""
        if ABBREV.match(before):
            continue
        out.append(text[start:m.end()].strip())
        start = m.end()
    if text[start:].strip():
        out.append(text[start:].strip())
    return out


def pieces(text, max_chars):
    """Pack whole sentences into pieces of at most max_chars (a longer sentence stays whole)."""
    out, cur = [], ""
    for s in sentences(text):
        if cur and len(cur) + 1 + len(s) > max_chars:
            out.append(cur)
            cur = s
        else:
            cur = (cur + " " + s).strip()
    if cur:
        out.append(cur)
    return out


def plan(cfg, segs):
    """Ordered items: {'text', 'gap_after_s'} for speech, {'cue': 'chime', 'gap_after_s'} for cues."""
    v = cfg["voice"]
    items = []
    for s in segs:
        if s["kind"] == "cue":
            items.append({"cue": "chime", "gap_after_s": v["gap_sentence_s"]})
            continue
        ps = pieces(s["text"], v["max_chunk_chars"])
        for i, p in enumerate(ps):
            items.append({"text": p, "gap_after_s": v["gap_paragraph_s"] if i == len(ps) - 1 else v["gap_sentence_s"]})
    return items


class VoiceError(RuntimeError):
    """The voice endpoint cannot be reached or refused the request. The message says what to check."""


def base_url(cfg):
    u = urllib.parse.urlsplit(cfg["voice"]["endpoint"])
    return "%s://%s" % (u.scheme, u.netloc)


def _get_json(url, timeout=15):
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.load(r)


def health(cfg):
    """Prove the voice server is reachable without rendering anything. Uses the Chatterbox server's own
    read-only routes: GET /api/model-info (is a model loaded) and GET /get_predefined_voices (is our voice there).
    Returns a dict. Raises VoiceError when the server cannot be reached or the voice is missing."""
    base, name = base_url(cfg), cfg["voice"]["voice"]
    try:
        info = _get_json(base + "/api/model-info")
        voices = _get_json(base + "/get_predefined_voices")
    except (urllib.error.URLError, OSError, ValueError) as e:
        raise VoiceError("voice server not reachable at %s (%s). Check that the voice container is running and that "
                         "AUDIOBOOK_VOICE_ENDPOINT / voice.endpoint is right." % (base, e))
    names = [v.get("filename") for v in voices if isinstance(v, dict)]
    if name not in names:
        raise VoiceError("voice server at %s has no voice file '%s' (it has: %s)" % (base, name, ", ".join(map(str, names)) or "none"))
    return {"base": base, "model_loaded": bool(info.get("loaded")), "model": info.get("class_name"),
            "device": info.get("device"), "voice": name}


def ensure_loaded(cfg):
    """Health check, then load the model if the server is up but the model is unloaded (setting voice.load_if_needed)."""
    h = health(cfg)
    if not h["model_loaded"]:
        if not cfg["voice"].get("load_if_needed", True):
            raise VoiceError("voice server is up but its model is not loaded, and voice.load_if_needed is false")
        print("voice model not loaded; asking the server to load it (POST /restart_server)")
        req = urllib.request.Request(h["base"] + "/restart_server", b"{}", {"Content-Type": "application/json"}, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=600) as r:
                r.read()
        except (urllib.error.URLError, OSError) as e:
            raise VoiceError("voice server could not load its model: %s" % e)
        h = health(cfg)
        if not h["model_loaded"]:
            raise VoiceError("voice server still reports no model loaded after /restart_server")
    return h


def unload(cfg):
    req = urllib.request.Request(base_url(cfg) + "/api/unload", b"", method="POST")
    with urllib.request.urlopen(req, timeout=120) as r:
        r.read()


def speak(cfg, text):
    v = cfg["voice"]
    if v.get("api") == "chatterbox_tts":
        # Native Chatterbox server route: the only one that accepts exaggeration, cfg_weight and temperature per request.
        body = {"text": text, "voice_mode": "predefined", "predefined_voice_id": v["voice"], "output_format": "wav"}
    else:
        body = {"model": v["model"], "input": text, "voice": v["voice"], "response_format": "wav"}
    body.update(v.get("extra") or {})
    if "speed_factor" in body:
        raise VoiceError("speed_factor must never be set (it time-stretches after generation and ruins the sound)")
    last = None
    for attempt in range(int(v.get("retries", 3)) + 1):
        req = urllib.request.Request(v["endpoint"], json.dumps(body).encode(), {"Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=v["timeout_s"]) as r:
                data = r.read()
            if data[:4] != b"RIFF":
                raise VoiceError("voice endpoint did not return a wav file")
            return data
        except urllib.error.HTTPError as e:
            last = "HTTP %d %s" % (e.code, e.read()[:200].decode("utf-8", "replace"))
            if e.code < 500:
                break
        except (urllib.error.URLError, OSError) as e:
            last = str(e)
        if attempt < int(v.get("retries", 3)):
            time.sleep(5 * 3 ** attempt)
    raise VoiceError("voice request to %s failed after %d attempts: %s" % (v["endpoint"], int(v.get("retries", 3)) + 1, last))


def work_dir_for(cfg, book_dir, name):
    return os.path.join(settings.resolve(cfg, "work_dir"), os.path.basename(book_dir), "audio", name)


def run(cfg, book_dir, only=None, names=None):
    """names: render only these script names (None = all, or the one in `only`)."""
    book = bookio.read_book(book_dir)
    for name in book["audio_files"]:
        if only and name != only:
            continue
        if names is not None and name not in names:
            continue
        with open(os.path.join(book_dir, "script", name + ".json"), encoding="utf-8") as f:
            items = plan(cfg, json.load(f))
        out = work_dir_for(cfg, book_dir, name)
        os.makedirs(out, exist_ok=True)
        for i, it in enumerate(items):
            if "text" in it:
                it["wav"] = os.path.join(out, "%04d.wav" % i)
                if not os.path.exists(it["wav"]):      # resume: finished pieces are kept
                    data = speak(cfg, it["text"])
                    with open(it["wav"], "wb") as f:
                        f.write(data)
        with open(os.path.join(out, "manifest.json"), "w", encoding="utf-8") as f:
            json.dump(items, f, indent=1, ensure_ascii=False)
        print("%s: %d pieces" % (name, len(items)))