AI Flow: Reverse Engineering

"""The voice and the captions. Each narration line is one take of `adom-tts say`, so a caption starts exactly where its
line starts. When the voice service does not answer, the build goes on with captions only and says so.
"""
import concurrent.futures as cf
import hashlib
import os
import shutil
import subprocess
import wave
import numpy as np
from engine import ffmpeg_exe

VOICE = "en-US-AndrewNeural"
RATE = 48000


class VoiceDown(Exception):
    pass


def _pcm(path):
    out = subprocess.run([ffmpeg_exe(), "-v", "error", "-i", path, "-f", "s16le", "-ac", "1", "-ar", str(RATE), "-"], capture_output=True)
    if out.returncode != 0:
        raise VoiceDown(f"ffmpeg could not read {os.path.basename(path)}: {out.stderr.decode(errors='replace').strip()[-200:]}")
    return np.frombuffer(out.stdout, np.int16)


def _trim(a):
    """The take without the silence the service puts before and after the words, keeping 40 ms each side."""
    loud = np.flatnonzero(np.abs(a.astype(np.int32)) > 260)
    if loud.size == 0:
        return a[:0]
    pad = int(0.04 * RATE)
    return a[max(0, loud[0] - pad):min(a.size, loud[-1] + pad)]


def _say(text, path):
    exe = os.environ.get("ADOM_AIFLOW_TTS") or "adom-tts"
    env = dict(os.environ, ADOM_TTS_NO_HINTS="1")
    last = ""
    for _ in range(2):
        try:
            r = subprocess.run([exe, "say", "--out", path, "--voice", VOICE, "-"], input=text.encode(), capture_output=True, timeout=90, env=env)
        except FileNotFoundError:
            raise VoiceDown("adom-tts is not installed on this machine")
        except subprocess.TimeoutExpired:
            last = "adom-tts did not answer within 90 seconds"
            continue
        if r.returncode == 0 and os.path.isfile(path) and os.path.getsize(path) > 800:
            return
        last = (r.stderr.decode(errors="replace").strip() or r.stdout.decode(errors="replace").strip() or f"exit {r.returncode}")[-300:]
    raise VoiceDown(f"adom-tts say failed: {last}")


def speak(lines, work, enabled=True):
    """Voice every line. Returns (takes, why): takes[i] is the trimmed PCM of line i, or (None, reason) when there is no voice."""
    if not enabled:
        return None, "the voice was turned off for this build (--no-voice)"
    if not shutil.which(os.environ.get("ADOM_AIFLOW_TTS") or "adom-tts"):
        return None, "adom-tts is not installed on this machine"
    os.makedirs(work, exist_ok=True)

    def one(text):
        path = os.path.join(work, hashlib.sha1((VOICE + "|" + text).encode()).hexdigest()[:16] + ".mp3")
        if not (os.path.isfile(path) and os.path.getsize(path) > 800):
            _say(text, path)
        a = _trim(_pcm(path))
        if a.size < RATE * 0.25:
            os.remove(path)
            raise VoiceDown(f"the take for \"{text[:50]}\" came back empty")
        return a

    try:
        with cf.ThreadPoolExecutor(4) as ex:
            return list(ex.map(one, lines)), None
    except VoiceDown as e:
        return None, str(e)


def estimate(text):
    """How long a line takes to read when there is no voice to measure: about 2.6 words a second."""
    return max(1.8, len(text.split()) / 2.6 + 0.35)


def write_wav(path, total, clips):
    """One mono track of `total` seconds with each (start seconds, pcm) clip laid in, levelled to leave 3 dB of headroom."""
    mix = np.zeros(int(total * RATE) + RATE // 10, np.float32)
    for start, a in clips:
        b = a.astype(np.float32)
        n = min(b.size, int(0.006 * RATE))
        if n > 1:
            ramp = np.linspace(0, 1, n, dtype=np.float32)
            b[:n] *= ramp
            b[-n:] *= ramp[::-1]
        i = int(start * RATE)
        b = b[:max(0, mix.size - i)]
        mix[i:i + b.size] += b
    peak = float(np.abs(mix).max()) if mix.size else 0.0
    if peak > 0:
        mix *= (32767 * 0.708) / peak
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(RATE)
        w.writeframes(mix.astype(np.int16).tobytes())


def _stamp(t, sep):
    t = max(0.0, t)
    h, m, s = int(t // 3600), int(t % 3600 // 60), t % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}".replace(".", sep)


def caption_cues(lines, max_chars=84):
    """Split each spoken line into caption cues short enough to read (two lines of 42), of about equal length, timed by
    their share of the letters."""
    cues = []
    for start, end, text in lines:
        words = text.split()
        n = max(1, -(-len(text) // max_chars))
        parts = []
        while n > 1 and len(words) > 1:
            # the break that leaves this cue nearest its even share of what is left
            want = len(" ".join(words)) / n
            k = min(range(1, len(words)), key=lambda k: abs(len(" ".join(words[:k])) - want))
            parts.append(" ".join(words[:k]))
            words = words[k:]
            n -= 1
        parts.append(" ".join(words))
        total = sum(len(p) for p in parts) or 1
        t = start
        for p in parts:
            d = (end - start) * len(p) / total
            cues.append((t, t + d, p))
            t += d
    return cues


def _two_lines(text):
    if len(text) <= 42:
        return text
    words = text.split()
    best = min(range(1, len(words)), key=lambda k: abs(len(" ".join(words[:k])) - len(" ".join(words[k:]))))
    return " ".join(words[:best]) + "\n" + " ".join(words[best:])


def write_captions(base, cues):
    with open(base + ".vtt", "w", encoding="utf-8") as f:
        f.write("WEBVTT\n\n")
        for k, (a, b, t) in enumerate(cues, 1):
            f.write(f"{k}\n{_stamp(a, '.')} --> {_stamp(b, '.')}\n{_two_lines(t)}\n\n")
    with open(base + ".srt", "w", encoding="utf-8") as f:
        for k, (a, b, t) in enumerate(cues, 1):
            f.write(f"{k}\n{_stamp(a, ',')} --> {_stamp(b, ',')}\n{_two_lines(t)}\n\n")