"""Make subtitles for the Movie Lounge on the Kodi box, with no inbound ports.
── WHAT IT DOES ─────────────────────────────────────────────────────────────
The box (playgrn) decides what needs subtitles; this machine has the spare CPU
to make them. Like lounge_agent.py it only ever reaches out: it asks the box
API for a job (a "lease"), does it, and posts the result back. Two kinds:
translate cues in one language -> the same cues in another, with NLLB-200
(distilled 600M, int8) on CTranslate2.
transcribe a [from, to] window of the film's HLS audio -> timed cues, with
faster-whisper (small by default, medium as a setting).
── WHY IT IS CAREFUL WITH THE MACHINE ───────────────────────────────────────
This runs on the same 6-core, 7.4 GB box that plays the film in Kodi. A dropped
frame on the Go Live share is worse than a subtitle arriving late, so:
* the process runs at BELOW_NORMAL priority (IDLE as a setting), with low
I/O priority, so Kodi's threads always win the CPU;
* it uses 4 threads by default, leaving two cores and all SMT siblings free;
* each model is loaded on the first job that needs it, kept for the next
one, and dropped after a stretch with no jobs, so an idle evening does not
hold ~1.5 GB away from Kodi;
* it refuses to load a model when Windows says there is not enough free RAM,
instead of pushing Kodi into the page file.
── THE TOKEN ────────────────────────────────────────────────────────────────
Same pattern as lounge_agent.py: the GRN_LOUNGE_TOKEN user environment variable
(or --token). It is only ever sent to the API host, and never printed or logged.
pythonw lounge_worker.py # the scheduled task: poll forever
python lounge_worker.py --fetch-models # download models next to this file
python lounge_worker.py --bench # measure speed and RAM, then exit
"""
from __future__ import annotations
import argparse
import ctypes
import gc
import io
import json
import logging
import logging.handlers
import os
import random
import re
import socket
import statistics
import sys
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
from urllib.parse import urljoin, urlparse
HOME = Path(__file__).resolve().parent
MODELS = HOME / "models"
SETTINGS_FILE = HOME / "worker_settings.json"
LOG_FILE = HOME / "worker.log"
# GRN_LOUNGE_API points a test copy at a scratch API; normally unset.
API = os.environ.get("GRN_LOUNGE_API") or "https://lounge.playgrn.com/movielounge/api"
USER_AGENT = "GRN-MovieLounge-Worker/1.0"
TOKEN = "" # set by --token or the GRN_LOUNGE_TOKEN environment variable
WORKER_NAME = os.environ.get("COMPUTERNAME") or socket.gethostname()
# Everything a person might want to change on the box lives in
# worker_settings.json next to this file; these are the values used when the
# file or a key is missing.
DEFAULTS = {
"threads": 4, # CPU threads per model; the 5650GE has 6 cores
"whisperModel": "small", # "small" or "medium"
"whisperBeam": 1, # greedy; see the note in transcribe_window()
"translateBeam": 2,
"translateRepetitionPenalty": 1.0,
"batchSize": 24, # cues per translate_batch call
"priority": "below_normal", # or "idle"
"idleUnloadMinutes": 20, # drop a model after this long without a job
"kinds": ["translate", "transcribe"],
}
# Pinned revisions, so a re-upload on Hugging Face cannot change what the box
# runs without anyone noticing. Each entry: folder, repo, revision, files.
MODEL_SOURCES = {
"nllb": [
# A ready int8 CTranslate2 conversion of facebook/nllb-200-distilled-600M
# (622 MB). Loads on ctranslate2 4.x; verified with the bench.
("nllb-200-distilled-600M-ct2-int8", "JustFrederik/nllb-200-distilled-600M-ct2-int8",
"302d78f00e6fdb50a1064059df7c392b735e9d05",
["config.json", "model.bin", "shared_vocabulary.txt"]),
# The tokenizer comes from the original repo rather than the conversion,
# so it is exactly the one the model was trained with.
("nllb-200-distilled-600M-ct2-int8", "facebook/nllb-200-distilled-600M",
"f8d333a098d19b4fd9a8b18f94170487ad3f821d",
["sentencepiece.bpe.model"]),
],
# These are float16; CTranslate2 quantises them to int8 as it loads.
"small": [("faster-whisper-small", "Systran/faster-whisper-small",
"536b0662742c02347bc0e980a01041f333bce120",
["config.json", "model.bin", "tokenizer.json", "vocabulary.txt"])],
"medium": [("faster-whisper-medium", "Systran/faster-whisper-medium",
"08e178d48790749d25932bbc082711ddcfdfbc4f",
["config.json", "model.bin", "tokenizer.json", "vocabulary.txt"])],
}
NLLB_FOLDER = "nllb-200-distilled-600M-ct2-int8"
# Peak private bytes while loading and running, measured with --bench, used
# to decide whether loading is safe. Whisper peaks high because its weights
# are float16 on disk and converted to int8 in memory as they load: medium
# briefly needs ~2.5 GB, which is why small is the default on a 7.4 GB box.
MODEL_RAM_GB = {"nllb": 1.0, "small": 0.8, "medium": 2.5}
# Free RAM to leave for Kodi and Windows after a model is loaded.
RAM_RESERVE_GB = 0.8
log = logging.getLogger("lounge_worker")
class JobError(Exception):
"""A job that cannot be done; the message goes to /worker/fail as is."""
class NetworkError(Exception):
pass
# ── languages ────────────────────────────────────────────────────────────────
# Subtitle language codes as the box sends them -> NLLB-200 codes. Lookup is by
# the full code first (zh-tw), then the primary subtag (es-419 -> es).
NLLB_CODES = {
"en": "eng_Latn", "es": "spa_Latn", "fr": "fra_Latn", "de": "deu_Latn",
# NLLB has a single Portuguese, trained mostly on European and Brazilian
# text together, so pt-BR and pt-PT both map to it and come out the same.
"pt": "por_Latn", "pt-br": "por_Latn", "pt-pt": "por_Latn",
"it": "ita_Latn", "ja": "jpn_Jpan", "nl": "nld_Latn", "pl": "pol_Latn",
"ru": "rus_Cyrl", "uk": "ukr_Cyrl", "tr": "tur_Latn", "ar": "arb_Arab",
"he": "heb_Hebr", "iw": "heb_Hebr", "hi": "hin_Deva", "ko": "kor_Hang",
"zh": "zho_Hans", "zh-cn": "zho_Hans", "zh-hans": "zho_Hans", "zh-sg": "zho_Hans",
"zh-tw": "zho_Hant", "zh-hant": "zho_Hant", "zh-hk": "zho_Hant", "yue": "yue_Hant",
"sv": "swe_Latn", "da": "dan_Latn", "no": "nob_Latn", "nb": "nob_Latn",
"nn": "nno_Latn", "fi": "fin_Latn", "is": "isl_Latn", "cs": "ces_Latn",
"sk": "slk_Latn", "hu": "hun_Latn", "ro": "ron_Latn", "el": "ell_Grek",
"bg": "bul_Cyrl", "hr": "hrv_Latn", "sr": "srp_Cyrl", "bs": "bos_Latn",
"sl": "slv_Latn", "mk": "mkd_Cyrl", "lt": "lit_Latn", "lv": "lvs_Latn",
"et": "est_Latn", "ca": "cat_Latn", "eu": "eus_Latn", "gl": "glg_Latn",
"id": "ind_Latn", "ms": "zsm_Latn", "vi": "vie_Latn", "th": "tha_Thai",
"tl": "tgl_Latn", "fil": "tgl_Latn", "fa": "pes_Arab", "ur": "urd_Arab",
"bn": "ben_Beng", "ta": "tam_Taml", "te": "tel_Telu", "ml": "mal_Mlym",
"af": "afr_Latn", "sw": "swh_Latn", "cy": "cym_Latn", "ga": "gle_Latn",
# three-letter forms some subtitle sites use
"eng": "eng_Latn", "spa": "spa_Latn", "fre": "fra_Latn", "fra": "fra_Latn",
"ger": "deu_Latn", "deu": "deu_Latn", "por": "por_Latn", "ita": "ita_Latn",
"jpn": "jpn_Jpan",
}
# Scripts written without spaces, where a subtitle line is much shorter and a
# line may break between any two characters.
NO_SPACE_SCRIPTS = ("_Jpan", "_Hans", "_Hant")
def language_lookup(code, table):
if not code:
return None
code = str(code).strip().lower().replace("_", "-")
return table.get(code) or table.get(code.split("-")[0])
def nllb_code(code, role):
if code and "_" in str(code) and len(str(code)) == 8:
return str(code) # already an NLLB code such as spa_Latn
found = language_lookup(code, NLLB_CODES)
if not found:
raise JobError("unsupported %s language %r" % (role, code))
return found
def whisper_code(code):
"""'pt-BR' -> 'pt', 'eng' -> 'en'; None when the job did not say."""
if not code:
return None
code = str(code).strip().lower().replace("_", "-")
three = {"eng": "en", "spa": "es", "fre": "fr", "fra": "fr", "ger": "de",
"deu": "de", "por": "pt", "ita": "it", "jpn": "ja", "fil": "tl",
"nb": "no", "nn": "no", "iw": "he"}
code = three.get(code, code)
code = three.get(code.split("-")[0], code.split("-")[0])
from faster_whisper.tokenizer import _LANGUAGE_CODES
if code not in _LANGUAGE_CODES:
raise JobError("whisper does not know language %r" % code)
return code
# ── subtitle text ────────────────────────────────────────────────────────────
STYLE_TAG = re.compile(r"]*)?>", re.I)
ASS_TAG = re.compile(r"\{\\[^}]*\}")
POSITION = re.compile(r"^\s*(\{\\an\d\})")
DASH = re.compile(r"^\s*[-–—]\s*")
LETTER = re.compile(r"[^\W\d_]")
def wrap_two_lines(text, width=42, any_char=False):
"""At most two lines, each about `width` long.
A long cue gets the break that makes the longer line as short as possible,
nudged toward a break after punctuation, which is how subtitlers split.
It never makes three lines: a third line covers the picture, and a
slightly long line is the lesser harm.
"""
text = " ".join(text.split())
if len(text) <= width:
return text
if any_char:
points = range(1, len(text))
else:
points = [i for i, c in enumerate(text) if c == " "]
if not points:
return text
best, best_score = None, None
for point in points:
first, second = text[:point].rstrip(), text[point:].lstrip()
if not first or not second:
continue
score = max(len(first), len(second))
if first[-1] in ",.;:!?…、。，！？":
score -= 4
if best_score is None or score < best_score:
best, best_score = (first, second), score
return "%s\n%s" % best if best else text
def parse_cue(text):
"""Split a cue into the parts worth translating, plus how to put it back.
Returns (position, italic, dashed, parts). Two-speaker cues ("- Hi.\\n
- Hello.") are kept as separate lines, because joining them would merge two
people's words into one sentence. Anything else is joined with spaces so
the model sees whole sentences, and re-wrapped afterwards.
"""
text = (text or "").replace("\r", "")
position = ""
match = POSITION.match(text)
if match:
position = match.group(1)
text = ASS_TAG.sub("", text)
text = re.sub(r"]*)?>", "", text, flags=re.I)
# Italic is kept only when it covers the whole cue (narration, a voice on
# the phone); partial italics are dropped, since the words move around in
# translation and there is no reliable way to know where they went.
without_italic_spans = re.sub(r".*?", "", text, flags=re.I | re.S)
italic = "" in text.lower() and not without_italic_spans.strip()
text = STYLE_TAG.sub("", text)
lines = [line.strip() for line in text.split("\n") if line.strip()]
dashed = len(lines) >= 2 and all(DASH.match(line) for line in lines)
if dashed:
parts = [DASH.sub("", line) for line in lines]
else:
parts = [" ".join(lines)] if lines else []
return position, italic, dashed, parts
def music_wrap(part):
"""Peel off what the model should not see: music notes and the brackets of
a sound description.
'♪ Happy birthday ♪' -> ('♪ ', 'Happy birthday', ' ♪'), '[door slams]' ->
('[', 'door slams', ']'). Sound descriptions are still translated, since
they help viewers, but given to the model in brackets NLLB treats them as
noise and answers with a stock phrase ("[door slams]" -> "¿Qué pasa?").
"""
match = re.match(r"^([♪♫#\s]*)(.*?)([♪♫#\s]*)$", part, re.S)
lead, core, tail = match.group(1), match.group(2), match.group(3)
bracket = re.match(r"^([\[(])\s*([^\[\]()]*?)\s*([\])])$", core)
if bracket:
lead, core, tail = lead + bracket.group(1), bracket.group(2), bracket.group(3) + tail
return lead, core, tail
def undash(source, translated):
"""NLLB, given a very short line, sometimes invents a second speaker:
"That's right." -> "- Es cierto. - Sí." (it saw a lot of dash-dialogue
subtitles in training). When the source had no dash, the dash and anything
after a second one are not a translation of anything, so they go.
"""
if DASH.match(source) or not DASH.match(translated):
return translated
translated = DASH.sub("", translated, count=1)
return re.split(r"\s[-–—]\s", translated, maxsplit=1)[0].strip()
def build_cue(position, italic, dashed, translated, target):
no_spaces = target.endswith(NO_SPACE_SCRIPTS)
width = 16 if no_spaces else 42
if dashed:
body = "\n".join("- " + part for part in translated)
else:
body = wrap_two_lines(" ".join(translated), width, any_char=no_spaces)
if italic and body:
body = "%s" % body
return position + body
# ── models ───────────────────────────────────────────────────────────────────
class Models:
"""Loads each model on first use and keeps it until it has sat idle."""
def __init__(self, settings, ram_guard=True):
self.settings = settings
self.ram_guard = ram_guard
self.translator = None
self.tokenizer = None
self.whisper = None
self.whisper_name = None
self.last_used = {}
def _make_room(self, wanted):
if not self.ram_guard:
return
import psutil
need = MODEL_RAM_GB[wanted] + RAM_RESERVE_GB
free = psutil.virtual_memory().available / 2 ** 30
if free >= need:
return
# Drop whatever else is loaded before giving up; reloading it later
# costs seconds, a paging Kodi costs the evening.
self.unload("nllb" if wanted != "nllb" else "whisper")
free = psutil.virtual_memory().available / 2 ** 30
if free < need:
raise JobError("not enough free RAM to load %s (%.1f GB free, %.1f GB needed)"
% (wanted, free, need))
def nllb(self):
if self.translator is None:
path = MODELS / NLLB_FOLDER
if not (path / "model.bin").exists():
raise JobError("translation model missing; run lounge_worker.py --fetch-models")
self._make_room("nllb")
import ctranslate2
import sentencepiece
started = time.perf_counter()
# inter_threads=1: one batch at a time, every thread on it. Two
# parallel batches would double the RAM for no gain on a box that
# should only use part of its CPU anyway.
self.translator = ctranslate2.Translator(
str(path), device="cpu", compute_type="int8",
intra_threads=self.settings["threads"], inter_threads=1)
self.tokenizer = sentencepiece.SentencePieceProcessor(
model_file=str(path / "sentencepiece.bpe.model"))
log.info("loaded NLLB in %.1fs", time.perf_counter() - started)
self.last_used["nllb"] = time.monotonic()
return self.translator, self.tokenizer
def whisper_model(self, name=None):
name = name or self.settings["whisperModel"]
if self.whisper is not None and self.whisper_name != name:
self.unload("whisper")
if self.whisper is None:
path = MODELS / ("faster-whisper-%s" % name)
if not (path / "model.bin").exists():
raise JobError("whisper %s missing; run lounge_worker.py --fetch-models" % name)
self._make_room(name)
from faster_whisper import WhisperModel
started = time.perf_counter()
self.whisper = WhisperModel(str(path), device="cpu", compute_type="int8",
cpu_threads=self.settings["threads"], num_workers=1)
self.whisper_name = name
log.info("loaded whisper %s in %.1fs", name, time.perf_counter() - started)
self.last_used["whisper"] = time.monotonic()
return self.whisper
def unload(self, which):
if which == "nllb" and self.translator is not None:
self.translator = self.tokenizer = None
log.info("unloaded NLLB")
elif which == "whisper" and self.whisper is not None:
self.whisper = self.whisper_name = None
log.info("unloaded whisper")
self.last_used.pop(which, None)
gc.collect()
def unload_idle(self):
limit = float(self.settings["idleUnloadMinutes"]) * 60
now = time.monotonic()
for which, used in list(self.last_used.items()):
if now - used > limit:
self.unload(which)
# ── translate ────────────────────────────────────────────────────────────────
def translate_cues(models, cues, source, target, settings):
"""[{"i", "text"}] in `source` -> [{"i", "text"}] in `target`."""
src = nllb_code(source, "source")
tgt = nllb_code(target, "target")
if src == tgt:
raise JobError("source and target are both %s" % src)
translator, sp = models.nllb()
# Flatten every cue into the pieces that need the model, remembering where
# each piece goes back. Pieces with no letters (♪, "...", "Hm?") would only
# make the model invent something, so they pass through untouched.
plans, pieces = [], []
for cue in cues:
position, italic, dashed, parts = parse_cue(cue.get("text", ""))
slots = []
for part in parts:
lead, core, tail = music_wrap(part)
if LETTER.search(core):
slots.append((lead, len(pieces), tail))
pieces.append(core)
else:
slots.append((part, None, ""))
plans.append((cue, position, italic, dashed, slots))
out = [None] * len(pieces)
batch_size = max(1, int(settings["batchSize"]))
# Sorting by length keeps a batch of short lines from waiting on one long
# one; the order is undone when the pieces are put back.
order = sorted(range(len(pieces)), key=lambda k: len(pieces[k]))
for start in range(0, len(order), batch_size):
chunk = order[start:start + batch_size]
# NLLB expects the source language token first and last; the
# target language token is forced as the first output token.
sources = [[src] + sp.encode(pieces[k], out_type=str) + [""] for k in chunk]
results = translator.translate_batch(
sources, target_prefix=[[tgt]] * len(chunk),
beam_size=int(settings["translateBeam"]), max_input_length=256,
max_decoding_length=256,
repetition_penalty=float(settings["translateRepetitionPenalty"]))
for k, result in zip(chunk, results):
tokens = result.hypotheses[0]
if tokens and tokens[0] == tgt:
tokens = tokens[1:]
out[k] = undash(pieces[k], sp.decode(tokens).strip())
answer = []
for cue, position, italic, dashed, slots in plans:
translated = []
for lead, index, tail in slots:
translated.append(lead if index is None else (lead + out[index] + tail).strip())
answer.append({"i": cue.get("i"), "text": build_cue(position, italic, dashed, translated, tgt)})
return answer
# ── transcribe ───────────────────────────────────────────────────────────────
def http_get(url, timeout=30, headers=None):
headers = dict(headers or {})
headers.setdefault("User-Agent", USER_AGENT)
# The stream lives on the same host as the API and may need the token; any
# other host never sees it.
if TOKEN and urlparse(url).netloc == urlparse(API).netloc:
headers["Authorization"] = "Bearer " + TOKEN
request = urllib.request.Request(url, headers=headers)
with urllib.request.urlopen(request, timeout=timeout) as response:
return response.read()
def read_media_playlist(url):
"""(segments [(url, start, duration)], ended) for an HLS media playlist."""
text = http_get(url).decode("utf-8", "replace")
if "#EXT-X-STREAM-INF" in text:
# A master playlist: follow the first variant.
lines = text.splitlines()
for index, line in enumerate(lines):
if line.startswith("#EXT-X-STREAM-INF"):
for candidate in lines[index + 1:]:
if candidate.strip() and not candidate.startswith("#"):
return read_media_playlist(urljoin(url, candidate.strip()))
raise JobError("master playlist without variants")
sequence, duration, raw = 0, None, []
for line in text.splitlines():
line = line.strip()
if line.startswith("#EXT-X-MEDIA-SEQUENCE:"):
sequence = int(line.split(":", 1)[1])
elif line.startswith("#EXTINF:"):
duration = float(line[8:].split(",")[0])
elif line and not line.startswith("#"):
raw.append((urljoin(url, line), duration or 0.0))
duration = None
if not raw:
raise JobError("playlist has no segments")
# Segment k starts at the sum of the earlier EXTINF values. If the playlist
# is a sliding window that has already dropped segments, their durations
# are gone; the typical segment length times the media sequence is the best
# estimate left (the segments are cut on a fixed keyframe interval).
start = 0.0
if sequence:
start = sequence * statistics.median(d for _, d in raw)
log.warning("playlist starts at media sequence %d; assuming it begins at %.1fs",
sequence, start)
segments = []
for segment_url, length in raw:
segments.append((segment_url, start, length))
start += length
return segments, "#EXT-X-ENDLIST" in text
def run_whisper(model, audio, language, settings):
"""faster-whisper over 16 kHz mono samples; returns (segments, info)."""
segments, info = model.transcribe(
audio, language=language, task="transcribe",
# Greedy decoding by default: on CPU, beam 5 costs roughly twice the
# time for a small gain, and keeping ahead of the film matters more.
beam_size=int(settings["whisperBeam"]),
# VAD skips music and silence, which is both the slow part and where
# whisper invents lines ("Thanks for watching!").
vad_filter=True, vad_parameters={"min_silence_duration_ms": 500},
# Each window is independent; carrying text over is what makes whisper
# repeat one line for a minute after a long silence.
condition_on_previous_text=False,
# Word times, so cues can be cut where people actually pause rather
# than where whisper's 30 s chunks happen to end.
word_timestamps=True,
# Without earlier text as context whisper drops most punctuation; a
# short punctuated prompt brings it back. English only, because a
# prompt in the wrong language pulls the transcript towards it.
initial_prompt=PUNCTUATED_PROMPT if language == "en" else None)
return list(segments), info
PUNCTUATED_PROMPT = "Hello. Are you ready? Yes, I am! Let's go, then."
CUE_MAX_SECONDS = 7.0
CUE_MAX_CHARS = 84 # two lines of 42
CUE_PAUSE = 0.6
def cues_from_words(segments, offset):
"""Cut whisper's words into subtitle-sized cues: at a sentence end once the
cue has some length, at a pause, or when it would outgrow two lines or
seven seconds. Each cue stays up a moment after its last word, but never
into the next one."""
words = [w for segment in segments for w in (segment.words or [])]
cues, current = [], []
def flush():
text = "".join(w.word for w in current).strip()
if text:
cues.append({"start": current[0].start + offset, "end": current[-1].end + offset,
"text": text})
current.clear()
for word in words:
if current:
text = "".join(w.word for w in current).strip()
gap = word.start - current[-1].end
sentence_done = text.endswith((".", "?", "!", "\u2026")) and \
(len(text) >= 18 or current[-1].end - current[0].start >= 1.2)
if (gap > CUE_PAUSE or sentence_done or len(text) + len(word.word) > CUE_MAX_CHARS
or word.end - current[0].start > CUE_MAX_SECONDS):
flush()
current.append(word)
flush()
for n, cue in enumerate(cues):
end = max(cue["end"] + 0.4, cue["start"] + 1.0)
if n + 1 < len(cues):
end = min(end, cues[n + 1]["start"] - 0.05)
cue["end"] = max(cue["end"], end)
return cues
def transcribe_window(models, lease, settings):
from faster_whisper import decode_audio
window_from, window_to = float(lease["from"]), float(lease["to"])
if window_to <= window_from:
raise JobError("empty window %.1f-%.1f" % (window_from, window_to))
language = whisper_code(lease.get("language"))
segments, ended = read_media_playlist(lease["audioUrl"])
wanted = [s for s in segments if s[1] + s[2] > window_from and s[1] < window_to]
if not wanted:
have = segments[-1][1] + segments[-1][2]
raise JobError("window %.1f-%.1f is not in the playlist (it covers %.1f-%.1f)"
% (window_from, window_to, segments[0][1], have))
covered_to = wanted[-1][1] + wanted[-1][2]
if covered_to < window_to and not ended:
# The film's last window can stop short, since the stream playlist
# never gets an ENDLIST; but if most of it is missing, the audio has
# simply not been made yet, and a partial answer would leave a gap.
if covered_to - window_from < 0.5 * (window_to - window_from):
raise JobError("audio only ready up to %.1fs" % covered_to)
offset = wanted[0][1]
buffer = io.BytesIO()
for segment_url, _, _ in wanted:
try:
buffer.write(http_get(segment_url, timeout=60))
except urllib.error.HTTPError as error:
raise JobError("segment fetch failed: HTTP %d" % error.code)
buffer.seek(0)
# MPEG-TS segments concatenate into one valid stream, so PyAV decodes them
# in one go, straight from memory.
try:
audio = decode_audio(buffer, sampling_rate=16000)
except Exception as error:
raise JobError("could not decode audio: %s" % error)
del buffer
model = models.whisper_model()
found, info = run_whisper(model, audio, language, settings)
cues = []
for cue in cues_from_words(found, offset):
# A cue that starts in the window is kept, even if it runs past the
# end; the next window starts its own cues after `to`.
if window_from <= cue["start"] < min(window_to, covered_to):
cues.append({"start": round(cue["start"], 2), "end": round(cue["end"], 2),
"text": wrap_two_lines(cue["text"])})
return (lease.get("language") or info.language), cues
# ── API ──────────────────────────────────────────────────────────────────────
def api_post(path, body, timeout=30):
"""(status, json) from the box API. Raises NetworkError when unreachable."""
request = urllib.request.Request(
API + path, data=json.dumps(body).encode("utf-8"), method="POST",
headers={"Content-Type": "application/json", "User-Agent": USER_AGENT,
"Authorization": "Bearer " + TOKEN})
try:
with urllib.request.urlopen(request, timeout=timeout) as response:
raw = response.read()
status = response.status
except urllib.error.HTTPError as error:
raw, status = error.read(), error.code
except (urllib.error.URLError, OSError, TimeoutError) as error:
raise NetworkError(str(getattr(error, "reason", error)))
try:
answer = json.loads(raw.decode("utf-8") or "{}")
except ValueError:
answer = {}
return status, answer if isinstance(answer, dict) else {}
def post_with_retries(path, body, attempts=4):
"""The result took minutes of CPU; a blip in the network should not waste it."""
for attempt in range(attempts):
try:
return api_post(path, body)
except NetworkError as error:
if attempt == attempts - 1:
raise
wait = min(60, 5 * 2 ** attempt) * random.uniform(0.5, 1.0)
log.warning("posting %s failed (%s); retrying in %.0fs", path, error, wait)
time.sleep(wait)
def handle(models, lease, settings):
lease_id, kind = lease.get("id"), lease.get("kind")
started = time.perf_counter()
try:
if kind == "translate":
cues = translate_cues(models, lease.get("cues") or [],
lease.get("sourceLanguage"), lease.get("targetLanguage"), settings)
body = {"id": lease_id, "cues": cues}
summary = "%d cues %s->%s" % (len(cues), lease.get("sourceLanguage"),
lease.get("targetLanguage"))
elif kind == "transcribe":
language, cues = transcribe_window(models, lease, settings)
body = {"id": lease_id, "language": language, "cues": cues}
summary = "%d cues from %.0f-%.0fs (%s)" % (len(cues), lease["from"], lease["to"], language)
else:
raise JobError("unknown job kind %r" % kind)
except Exception as error:
message = str(error) if isinstance(error, JobError) else "%s: %s" % (type(error).__name__, error)
log.error("lease %s (%s) failed: %s", lease_id, kind, message)
if not isinstance(error, JobError):
log.debug("traceback", exc_info=True)
try:
post_with_retries("/worker/fail", {"id": lease_id, "error": message[:300]})
except NetworkError:
pass # the lease simply expires on the server
return
elapsed = time.perf_counter() - started
try:
status, answer = post_with_retries("/worker/result", body)
except NetworkError as error:
log.error("lease %s done but the result could not be posted: %s", lease_id, error)
return
if status == 409:
log.warning("lease %s expired or its run changed; dropped (%s, %.1fs)", lease_id, summary, elapsed)
elif status == 200 and answer.get("ok"):
log.info("lease %s %s: %s in %.1fs", lease_id, kind, summary, elapsed)
else:
log.error("lease %s result refused (HTTP %d): %s", lease_id, status, answer.get("error", "?"))
def run_forever(settings):
models = Models(settings)
kinds = settings["kinds"]
backoff = 0.0
log.info("worker %s polling %s for %s (threads %d, whisper %s)",
WORKER_NAME, API, ",".join(kinds), settings["threads"], settings["whisperModel"])
while True:
models.unload_idle()
try:
status, answer = api_post("/worker/lease", {"worker": WORKER_NAME, "kinds": kinds})
except NetworkError as error:
backoff = min(60.0, max(5.0, backoff * 2))
wait = random.uniform(backoff / 2, backoff)
log.warning("API unreachable (%s); next try in %.0fs", error, wait)
time.sleep(wait)
continue
if status >= 500:
backoff = min(60.0, max(5.0, backoff * 2))
wait = random.uniform(backoff / 2, backoff)
log.warning("API answered HTTP %d; next try in %.0fs", status, wait)
time.sleep(wait)
continue
backoff = 0.0
if status in (401, 403):
log.error("API refused the token (HTTP %d); check GRN_LOUNGE_TOKEN. Waiting 5 min", status)
time.sleep(300)
continue
if status == 404:
# Until the server side is deployed the route does not exist; that
# is "nothing to do", not an error worth filling the log with.
time.sleep(60)
continue
if status != 200:
log.warning("lease answered HTTP %d; waiting 30s", status)
time.sleep(30)
continue
lease = answer.get("lease")
if not lease:
try:
wait = float(answer.get("retryAfter", 10))
except (TypeError, ValueError):
wait = 10.0
time.sleep(min(300.0, max(2.0, wait)))
continue
handle(models, lease, settings)
# ── setup helpers ────────────────────────────────────────────────────────────
def load_settings():
settings = dict(DEFAULTS)
try:
settings.update(json.loads(SETTINGS_FILE.read_text(encoding="utf-8")))
except FileNotFoundError:
pass
except Exception as error:
log.warning("ignoring %s: %s", SETTINGS_FILE.name, error)
settings["threads"] = max(1, int(settings["threads"]))
return settings
def lower_priority(level):
"""Below Kodi, always. Threads the models start inherit the class."""
try:
import psutil
process = psutil.Process()
process.nice(psutil.IDLE_PRIORITY_CLASS if level == "idle"
else psutil.BELOW_NORMAL_PRIORITY_CLASS)
try:
process.ionice(psutil.IOPRIO_LOW)
except Exception:
pass
except Exception:
if os.name == "nt":
# IDLE_PRIORITY_CLASS 0x40, BELOW_NORMAL_PRIORITY_CLASS 0x4000
kernel32 = ctypes.windll.kernel32
kernel32.SetPriorityClass(kernel32.GetCurrentProcess(),
0x40 if level == "idle" else 0x4000)
_MUTEX = None
def single_instance():
"""Two copies would lease two jobs and load every model twice.
The scheduled task uses IgnoreNew, but a copy started by hand is not
covered by that, so a named mutex makes the second copy exit instead.
"""
global _MUTEX
if os.name != "nt":
return True
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
kernel32.CreateMutexW.restype = ctypes.c_void_p
_MUTEX = kernel32.CreateMutexW(None, False, "Local\\GRNMovieLoungeSubtitleWorker")
return ctypes.get_last_error() != 183 # ERROR_ALREADY_EXISTS
def setup_logging(console):
log.setLevel(logging.INFO)
form = logging.Formatter("%(asctime)s %(levelname)s %(message)s", "%Y-%m-%d %H:%M:%S")
file_handler = logging.handlers.RotatingFileHandler(
LOG_FILE, maxBytes=1_000_000, backupCount=2, encoding="utf-8")
file_handler.setFormatter(form)
log.addHandler(file_handler)
if console:
stream = logging.StreamHandler(sys.stdout)
stream.setFormatter(form)
log.addHandler(stream)
def download(url, dest):
part = dest.with_name(dest.name + ".part")
request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
with urllib.request.urlopen(request, timeout=60) as response:
total = int(response.headers.get("Content-Length") or 0)
done, shown = 0, -1
with open(part, "wb") as handle_out:
while True:
chunk = response.read(1 << 20)
if not chunk:
break
handle_out.write(chunk)
done += len(chunk)
if total > 50 << 20 and done * 10 // total != shown:
shown = done * 10 // total
print(" %s %d%%" % (dest.name, shown * 10), flush=True)
if total and done != total:
raise RuntimeError("%s: got %d of %d bytes" % (dest.name, done, total))
os.replace(part, dest)
def fetch_models(names):
for name in names:
for folder, repo, revision, files in MODEL_SOURCES[name]:
target = MODELS / folder
target.mkdir(parents=True, exist_ok=True)
for file_name in files:
dest = target / file_name
if dest.exists() and dest.stat().st_size > 0:
continue
print(" %s/%s" % (repo, file_name), flush=True)
download("https://huggingface.co/%s/resolve/%s/%s" % (repo, revision, file_name), dest)
print("models are in %s" % MODELS)
# ── bench ────────────────────────────────────────────────────────────────────
# LibriVox, public domain: The Count of Monte Cristo, chapter 1. Only the first
# 1.6 MB is fetched (a Range request; 128 kbps mono, so ~100 s), and 60 s of
# narration is cut from it, after the LibriVox preamble.
BENCH_AUDIO_URL = ("https://archive.org/download/count_monte_cristo_0711_librivox/"
"count_of_monte_cristo_001_dumas.mp3")
BENCH_AUDIO_BYTES = 1_600_000
BENCH_AUDIO_FROM, BENCH_AUDIO_SECONDS = 20.0, 60.0
# 200 cues of ordinary film dialogue, written for this bench. "|" is a line
# break inside a cue. Tags, sound descriptions and two-speaker cues are mixed
# in at roughly the rate a real English .srt has them.
BENCH_CUES = """\
Where were you last night?|I called you four times.
I was at my sister's. Her car broke down again.
You could have answered the phone.
- My battery died.|- Of course it did.
[door slams]
Don't walk away from me when I'm talking to you.
I'm not walking away. I'm getting a glass of water.
The summer I turned seventeen,|everything changed.
We moved to a town with one traffic light|and two churches.
My father said it would be good for us.
He said that about everything.
Hey, new kid! You want to play or not?
I guess so.
Then hurry up, we don't have all day.
[children laughing]
How long have you lived here?
My whole life. Nobody ever leaves.
That's not true. My brother left.
- Where did he go?|- California. He never writes.
Mom, I'm home!
In the kitchen, sweetheart.
What's for dinner?
Whatever you're making.
Very funny.
Did you finish your homework?
Almost. I just have to read one more chapter.
Then go read it. Dinner's in twenty minutes.
[phone ringing]
Hello?
Is this Daniel Morgan?
Yes. Who's this?
I'm calling about your father.
What about him?
I think you should come down to the station.
I'll be there in ten minutes.
Detective Ruiz. Thanks for coming in.
Is he all right? Is he hurt?
He's fine. He's not hurt.
But we found him wandering near the highway|at three in the morning.
He didn't know where he was.
That's impossible. He's sharp as a tack.
Has he been forgetting things lately?
Little things. Keys, names. Everybody does that.
Maybe. But you should take him to a doctor.
Can I see him?
Down the hall, second door on the left.
Dad?
Danny. What are you doing here?
I came to take you home.
I was going to the store.
At three in the morning?
They're open all night, aren't they?
Come on. Let's go.
[engine starts]
You don't have to look at me like that.
Like what?
Like I'm a broken lamp you can't decide whether to throw out.
That's not what I'm thinking.
Then what are you thinking?
I'm thinking I should have visited more.
Well, you're here now.
Three weeks later, the tests came back.
I'm sorry. I wish I had better news.
How long?
It's hard to say. Every case is different.
Years? Months?
Most likely years.|But it will get harder.
Is there anything we can do?
There are medications that can slow it down.
And there's a lot you can do at home.
Keep him active. Keep him around people.
Routine helps more than you'd think.
Thank you, doctor.
[sighs]
So, what did the man say?
He said you need to stop driving.
Over my dead body.
Dad, please.
I have been driving since I was fourteen years old.
I know.
I drove your mother to the hospital|the night you were born.
Through a snowstorm, in a car with no heater.
You've told me.
Then you know I can drive.
- Fine. We'll talk about it later.|- There's nothing to talk about.
Sarah, it's me. Call me back when you get this.
It's about Dad.
Hi, you've reached Sarah.|Leave a message.
[rain falling]
You look terrible.
Thanks. I haven't slept.
Sit down. I'll make coffee.
How is he?
Some days he's fine.|Some days he calls me by your name.
That must be hard.
It's harder for him.
He knows something is wrong.|He just won't say it.
Have you thought about a nurse?
He'd never accept a stranger in the house.
What about me?
You live four hundred miles away.
I can work from anywhere.
You'd do that?
He's my father too, Danny.
[knocking]
Who's there?
It's Mrs. Patterson from next door.
I brought you some soup.
That's very kind of you.
Your father helped my husband fix the roof|every spring for thirty years.
It's the least I can do.
Would you like to come in?
Only for a minute. I don't want to be a bother.
- You're never a bother.|- Oh, you're sweet.
Did you hear that?
Hear what?
Somebody's in the garage.
[metal clanking]
Stay here.
Dad? What are you doing out here?
Looking for my fishing rod.
It's the middle of the night.
The fish don't care what time it is.
Okay. Let's look in the morning.
Promise?
I promise.
We went to the lake that Saturday.
He caught three fish|and I caught nothing.
You're holding it wrong.
I'm holding it exactly the way you showed me.
Then I showed you wrong.
[both laughing]
Do you remember when we came here with Mom?
Of course I remember.
She fell in the water trying to catch a frog.
She was furious.
She laughed about it for a week.
I miss her.
Me too, kid. Every day.
♪ Take me home, country roads ♪
Turn that down, will you?
You used to love this song.
I still do. I just don't want to hear it right now.
Are you hungry? We could stop for lunch.
There's a diner up ahead. Best pie in the county.
Is it still there?
Only one way to find out.
Two coffees and two slices of apple pie, please.
Coming right up, hon.
Hey, I know you. You're Walt Morgan's boy.
That's right.
Your father taught me how to change a tire.
I was sixteen and scared to death.
He never mentioned it.
He wouldn't. That's the kind of man he is.
Excuse me, I'm looking for the restroom.
Through the back, on the right.
Dad, sit down. I'll be right back.
Where did he go?
He was right here a minute ago.
Did anyone see an old man in a blue jacket?
He went out the side door.
Dad! Dad!
[car horn honking]
Watch out!
Are you crazy? You could have been killed!
I was going home.
Home is two hundred miles that way.
I know where my home is.
Okay. Okay, I'm sorry. Let's go together.
That night, I called Sarah again.
I can't do this alone anymore.
I know. I'm coming.
When?
I've already booked the flight. I'll be there Friday.
Thank you.
Don't thank me. Just hang in there.
Look who's here.
Hi, Daddy.
Sarah. My little girl.
You cut your hair.
Five years ago.
It looks good. It suits you.
What happened to your face?|You look tired.
I'm fine, Dad. It was a long flight.
- Are you staying?|- For as long as you need me.
Then I'll make up the guest room.
I already did, Dad.
Oh. Well, then we're all set.
[clock ticking]
Can't sleep either?
I keep thinking about when we were kids.
He was always so strong.
He still is, in his way.
Do you think he's scared?
I think he's terrified.|He just doesn't want us to see it.
So what do we do?
We stay. We make it easier.
And when it gets worse?
Then we stay anyway.
Good morning, you two.
Morning, Dad. You're up early.
I made pancakes. Don't get used to it.
These are actually good.
Of course they're good. I'm not dead yet.
And for a while, he wasn't.
For a while, he was just our father.
"""
class PeakSampler:
"""Highest RAM use seen while a phase runs (psutil, every 50 ms).
Both the working set (RSS) and private bytes are tracked: when Windows is
short of RAM it trims working sets, so RSS alone can understate what a
model really needs, while private bytes is the committed total.
"""
def __init__(self):
import psutil
self.process = psutil.Process()
self.peak = self.peak_private = 0
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, daemon=True)
def _run(self):
while not self._stop.is_set():
self._sample()
self._stop.wait(0.05)
def _sample(self):
info = self.process.memory_info()
self.peak = max(self.peak, info.rss)
self.peak_private = max(self.peak_private, getattr(info, "private", 0))
def __enter__(self):
self._sample()
self._thread.start()
return self
def __exit__(self, *exc):
self._stop.set()
self._thread.join()
self._sample()
def text(self):
return "peak RSS %.0f MB, private %.0f MB" % (mb(self.peak), mb(self.peak_private))
def mb(value):
return value / 2 ** 20
def bench(settings, whisper_names, skip_translate=False, skip_transcribe=False):
import psutil
process = psutil.Process()
print("bench on %s: %d threads, %d logical CPUs, %.1f GB RAM (%.1f GB free)" % (
WORKER_NAME, settings["threads"], psutil.cpu_count(),
psutil.virtual_memory().total / 2 ** 30, psutil.virtual_memory().available / 2 ** 30))
print("baseline RSS %.0f MB" % mb(process.memory_info().rss))
# No RAM guard here: the bench is run by hand, and the point is to measure
# what the models take, not to protect a film that is not playing.
models = Models(settings, ram_guard=False)
if not skip_translate:
cues = [{"i": n, "text": line.replace("|", "\n")}
for n, line in enumerate(BENCH_CUES.splitlines()) if line.strip()]
with PeakSampler() as sampler:
started = time.perf_counter()
models.nllb()
loaded = time.perf_counter() - started
# one warm-up batch, so the number is the steady speed of a film
translate_cues(models, cues[:8], "en", "es", settings)
started = time.perf_counter()
result = translate_cues(models, cues, "en", "es", settings)
elapsed = time.perf_counter() - started
print("\nNLLB-200 600M int8, en->es, %d cues, batch %d, beam %d" % (
len(cues), settings["batchSize"], settings["translateBeam"]))
print(" load %.1fs translate %.1fs %.1f cues/s %s" % (
loaded, elapsed, len(cues) / elapsed, sampler.text()))
for n in (0, 4, 7, 27, 55, 62, 135, 185):
print(" %-44s -> %s" % (cues[n]["text"].replace("\n", " / ")[:44],
result[n]["text"].replace("\n", " / ")))
models.unload("nllb")
if not skip_transcribe:
from faster_whisper import decode_audio
clip = MODELS / "bench" / "monte_cristo_head.mp3"
if not clip.exists():
clip.parent.mkdir(parents=True, exist_ok=True)
request = urllib.request.Request(BENCH_AUDIO_URL, headers={
"User-Agent": USER_AGENT, "Range": "bytes=0-%d" % (BENCH_AUDIO_BYTES - 1)})
with urllib.request.urlopen(request, timeout=60) as response:
clip.write_bytes(response.read(BENCH_AUDIO_BYTES))
audio = decode_audio(str(clip), sampling_rate=16000)
audio = audio[int(BENCH_AUDIO_FROM * 16000):int((BENCH_AUDIO_FROM + BENCH_AUDIO_SECONDS) * 16000)]
seconds = len(audio) / 16000
for name in whisper_names:
if not (MODELS / ("faster-whisper-%s" % name) / "model.bin").exists():
print("\nwhisper %s: not downloaded, skipped" % name)
continue
with PeakSampler() as sampler:
started = time.perf_counter()
model = models.whisper_model(name)
loaded = time.perf_counter() - started
started = time.perf_counter()
found, info = run_whisper(model, audio, "en", settings)
elapsed = time.perf_counter() - started
text = " ".join(s.text.strip() for s in found)
print("\nwhisper %s int8, %.0fs of speech, beam %d, VAD on" % (
name, seconds, settings["whisperBeam"]))
print(" load %.1fs transcribe %.1fs RTF %.2f (%.1fx real time) %s" % (
loaded, elapsed, elapsed / seconds, seconds / elapsed, sampler.text()))
print(" %d segments: %s..." % (len(found), text[:150]))
models.unload("whisper")
peak = getattr(process.memory_info(), "peak_wset", 0)
if peak:
print("\nprocess peak working set %.0f MB" % mb(peak))
def main():
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--token", default="", help="defaults to GRN_LOUNGE_TOKEN")
parser.add_argument("--api", default="", help="API base, for testing against a mock")
parser.add_argument("--threads", type=int, help="override the threads setting")
parser.add_argument("--whisper", help="override the whisperModel setting (small/medium)")
parser.add_argument("--fetch-models", nargs="*", metavar="NAME",
help="download models (nllb small medium; default: nllb + the whisper setting)")
parser.add_argument("--bench", action="store_true", help="measure speed and RAM, then exit")
parser.add_argument("--bench-whisper", default="small,medium", help="whisper models to bench")
parser.add_argument("--skip-translate", action="store_true")
parser.add_argument("--skip-transcribe", action="store_true")
parser.add_argument("--models", default="", help="model folder (default: models next to this file)")
parser.add_argument("--normal-priority", action="store_true", help="do not lower the priority")
arguments = parser.parse_args()
# pythonw has no console: sys.stdout is None and the first print, or any
# library that writes a warning, would kill the process.
console = sys.stdout is not None
if sys.stdout is None:
sys.stdout = open(os.devnull, "w")
if sys.stderr is None:
sys.stderr = open(os.devnull, "w")
setup_logging(console and not arguments.bench)
settings = load_settings()
if arguments.threads:
settings["threads"] = arguments.threads
if arguments.whisper:
settings["whisperModel"] = arguments.whisper
# Before ctranslate2 is imported: its OpenMP runtime reads this once, and
# otherwise starts a thread per logical CPU for some internal work.
os.environ["OMP_NUM_THREADS"] = str(settings["threads"])
global TOKEN, API, MODELS
if arguments.models:
MODELS = Path(arguments.models).resolve()
TOKEN = arguments.token or os.environ.get("GRN_LOUNGE_TOKEN", "")
if arguments.api:
API = arguments.api.rstrip("/")
if arguments.fetch_models is not None:
fetch_models(arguments.fetch_models or ["nllb", settings["whisperModel"]])
return 0
if not arguments.normal_priority:
lower_priority(settings["priority"])
if arguments.bench:
bench(settings, [n for n in arguments.bench_whisper.split(",") if n],
arguments.skip_translate, arguments.skip_transcribe)
return 0
if not single_instance():
log.info("another lounge_worker is already running; exiting")
return 0
if not TOKEN:
log.error("no token: set the GRN_LOUNGE_TOKEN user environment variable")
return 1
try:
run_forever(settings)
except KeyboardInterrupt:
return 0
if __name__ == "__main__":
sys.exit(main())