lektor-ai/docs/architecture.md

# Architecture

How `lektor_ai_comfy.py` turns subtitles into a finished video, and why each
stage works the way it does. Most of the non-obvious choices here exist because
the straightforward version was tried first and broke.

```
SRT / embedded track
   │  parse, clean, split dialogue
   ▼
subtitle blocks ──────────────► <job>/tts_input.srt
   │  one HTTP job per block
   ▼
ComfyUI + Higgs Audio V3
   │  per-clip atempo + fade
   ▼
<job>/clip_NNNN.flac
   │  adelay + amix, batched
   ▼
tts_assembled.flac
   │  pass 1: mix with original audio
   │  pass 2: stream-copy mux
   ▼
{stem}_lektor.mkv
```

## 1. Subtitle parsing

`parse_subtitle_blocks()` reads the SRT with `utf-8-sig` (tolerating a BOM),
splits on blank lines, and keeps blocks whose start falls inside the requested
range. Times become **relative to `--start`**, which is why a later `--rebuild`
must repeat the same `--start`/`--end`.

`_clean_text()` then normalises each block, in this order:

1. **Strip formatting tags.** `{...}` ASS overrides (`{\an8}`, `{\pos(x,y)}`)
   and `<...>` HTML tags (`<i>`, `<font …>`).
2. **Detect dialogue.** Lines beginning `- ` mean two speakers and are kept
   separate; ordinary wrapped lines are joined with a space.
3. **Title-case ALL-CAPS blocks**, which the model would otherwise shout.
   Every word is capitalised rather than lowercasing all but the first letter,
   because the obvious version also lowercases the proper nouns inside a
   caption: `NASA – MARGO MADISON` became `Nasa – margo madison`, which is wrong
   to read aloud and invites the number normalizer to "correct" the name (it
   produced `marga`). A/B clips showed the TTS reads both casings the same, so
   keeping the names costs nothing.
4. **Capitalise a lowercase first letter**, since a fragment continuing from the
   previous subtitle confuses the model.
5. **Apply the replacement dictionary** — `km/h` → `kilometrów na godzinę` and
   anything else the model reads out wrong (see below).
6. **Add terminal punctuation** so the model produces a natural falling tail.

Step 1 must come first. `<i>ALL CAPS</i>` looks mixed-case while the tag is
still present — the `i` in `<i>` is a lowercase letter — so the caps detection
in step 3 would silently miss it.

Step 5 must come after step 3 for the same reason in reverse: a lowercase
replacement dropped into `JEDZIE 200 KM/H` would make the block look mixed-case
and the shouting would survive. It comes before step 6 so a rule can supply
punctuation of its own.

### The replacement dictionary

`Dictionary.load()` reads `dictionary.txt` (`[paths] dictionary`, or
`--dictionary`) into `pattern = replacement` rules — literal and
case-insensitive, or a regular expression when the pattern is wrapped in
slashes. Every rule is compiled into **one alternation** applied in a single
`re.sub` pass, which buys three properties worth keeping:

- Replacement text is never rescanned, so rules cannot chain or loop.
- Literals are ordered longest-first inside the alternation, so `km/h` beats a
  plain `km` rule regardless of the order in the file.
- Rules cost one pass over the text no matter how many there are.

Literal patterns get `(?<!\w)` / `(?!\w)` guards only on edges that are word
characters, which is what lets `km/h` and `np.` match on their own while `znp.`
does not — a guard on a `/` or `.` edge would stop them matching at all.

Matching is case-sensitive and the replacement is verbatim. The alternative —
match loosely, then infer how to re-case the replacement — was tried and
removed: it silently rewrote `NASA` as `Nasa`, and a rule that fires on a case
you did not write is hard to reason about. The cost is that the first letter of
a block is capitalised before the dictionary runs, so a lowercase rule cannot
fire there; on the seven subtitle files here that changed nothing, but it is why
a word that can open a subtitle needs both forms.

A missing file, a line without a separator and an invalid regex are all
warnings, never fatal — the same rule as `settings.ini`, for the same reason: a
typo should not cost a long run.

`split_dialogue_blocks()` then divides a two-speaker block into one clip per
line, apportioning the block's duration by character count.

### Number normalization (last stage)

`normalize_block_numbers()` rewrites digits as Polish words — `1247` →
`tysiąc dwieście czterdzieści siedem` — because Higgs Audio misreads long digit
strings and sometimes emits garbled non-words. It is **last** for two reasons:
the dictionary may have produced the words a numeral has to agree with (`16%` →
`16 procent` → `szesnaście procent`), and dialogue splitting decides which clip
a number ends up in.

The model that does it needs `transformers>=5`, which cannot be installed beside
ComfyUI's pinned dependencies, so it lives in `normalizer/` with its own venv and
runs as a **separate process**. `NumberNormalizer` is the client half:

- One worker per run, not per block — the 1.1 GB load dominates everything else.
  It is started only if some block actually contains a number worth expanding,
  and closed before step 2, so it is not holding VRAM while Higgs generates.
- **ComfyUI is asked to unload first** (`free_comfy_vram`). Higgs stays resident
  between runs and leaves too little VRAM; a normalizer that does not fit is
  spilled to RAM by accelerate *silently*, at about 5× the time per block. A
  `POST /free` is cheaper than a restart — the server stays warm — and the wait
  is on `torch_vram_total` in `/system_stats` rather than a fixed sleep, because
  the flag is consumed by ComfyUI's worker loop rather than by the request.
- JSON lines over stdin/stdout, strictly one request per response. stdout is the
  protocol; logs go to stderr and are kept in a 40-line ring for diagnostics.
- Reads are on a queue fed by a reader thread, which is what makes a per-block
  timeout possible at all — a wedged worker cannot block the run.
- **Nothing here is fatal.** Missing venv, dead worker, timeout, out-of-order
  reply: each prints a warning, disables the stage, and lets the digits through
  untouched. A number read badly is a bad clip; a crash here is a lost film.

- **Two models behind one interface.** Plain integers go to num2words for the
  value and a general Polish LLM for the case; anything with a time, date,
  decimal or identifier goes to the 0.6B fine-tune. Measured on 40 real subtitle
  lines, each is clearly better than the other on its half and clearly worse on
  the other's. The worker reports which leg answered, and the run prints the
  split.
- **Guards, not trust.** Both models can emit something worse than nothing — a
  repetition loop to the token ceiling, a dropped minus that inverts a
  temperature, a vanished noun, a number changed from 8250 to 8205. Each answer
  is checked against the input and falls back to num2words when it fails.

`needs_number_normalization()` duplicates the worker's own selection rule so
blocks that need nothing never cross the process boundary. The two must stay in
step; the worker's copy is the authoritative one, and a test in the normalizer
suite imports the caller and compares both on a shared sample.

The result is written to `tts_input.srt` in that video's job folder
(`output/clips/<video name>/`). It is the source of truth for block
numbering: block *N* there is `clip_NNNN.flac` beside it and the
argument to `--regen N`.

Which is why `--regen`, `--rebuild` **and `--resume-from`** all read that file
rather than the original SRT. Everything upstream of it — the dictionary, the
normalizer, dialogue splitting — can produce different text than it did on the
first run, and any difference breaks the pairing with the clips already on disk.
Step 1 is skipped entirely in those modes.

## 2. Speech generation

For each block the script POSTs a workflow graph to `/prompt`, polls
`/history/{id}` until it completes, and reads the resulting file directly from
ComfyUI's output directory.

**Polling cadence is a performance decision, not a detail.** A clip that
finishes just after a poll waits for the next one, so the interval — not the
GPU — sets the floor on throughput. It was originally a flat 3 s, which
quantised every clip to a 3 s boundary: measured across 2075 real clips, 91% of
the gaps landed exactly on a multiple of 3 s.

`wait_for_completion()` now starts at 0.25 s and backs off to 2 s once a job
passes 10 s, so the common case is picked up almost immediately while a stuck
generation does not sit at 4 requests/second indefinitely. Measured over 34
clips:

| | mean | median | 34 clips |
|---|---|---|---|
| Actual generation | 4.75 s | 4.60 s | 162 s |
| Old flat 3 s poll | 6.35 s | 6.00 s | 216 s |

That is 1.60 s wasted per clip, a **25% reduction in TTS wait** — about 18
minutes on a 693-clip episode. The run also now prints real generation time
(`done (4.2s)`), which is what makes this measurable at all.

ComfyUI also pushes completion events over a WebSocket, which would remove
polling entirely. It is not used because Python has no standard-library
WebSocket client, and that would cost this script its only real virtue as a
dependency-free file. Polling at 0.25 s recovers most of the benefit.

The graph has five nodes:

| Node | Class | Role |
|---|---|---|
| 1 | `HiggsAudioV3EngineNode` | Model, temperature, top_p 0.95, top_k 50 |
| 2 | `LoadAudio` | Voice reference WAV from ComfyUI's input dir |
| 3 | `CharacterVoicesNode` | Pairs that audio with its transcript |
| 4 | `UnifiedTTSTextNode` | The text, the seed, `enable_audio_cache: False` |
| 5 | `SaveAudio` | Writes FLAC |

Three details matter:

- **`enable_audio_cache: False`** is essential. With caching on, ComfyUI returns
  a previous result for repeated text and every `--regen` is a no-op.
- **The seed is `base_seed + block_index`**, so a fixed `--seed` is fully
  reproducible while each clip still differs.
- **No language tag is sent.** Higgs Audio V3 detects Polish from the text
  itself; its tokenizer has no `<|lang:pl|>` token, and a `[pl:]` prefix is
  stripped by the character parser before it reaches the model. A `lang: en`
  line in the ComfyUI log is cosmetic.

Each finished clip is optionally head-trimmed (`--strip-leading-ms`), then
sped up with `atempo` and given a short fade-out — **per clip, before
assembly**. Applying `atempo` to the assembled track instead would speed up the
silences between clips and shift every subsequent timestamp.

## 3. Placement

Clips are positioned with a push-forward cursor:

```python
actual_start = max(t_start, cursor)
cursor       = actual_start + duration
```

A clip starts at its subtitle time, or when the previous clip finished if that
is later. Clips are **never trimmed** to their subtitle window — cutting speech
mid-sentence sounds far worse than running a little long. Overrun cascades
forward and then resolves by itself at the next natural gap in the subtitles,
because a gap lets the cursor fall back below the next `t_start`.

The console prints `[+1.20s behind SRT]` when drift occurs. Steady growth means
`--speed` is too low for the material.

## 4. Assembly

Each clip becomes an `adelay` filter input and everything is mixed with `amix`:

```
[0:a]adelay=1340:all=1[a0]; [1:a]adelay=4820:all=1[a1]; …
[a0][a1]…amix=inputs=N:duration=longest:normalize=0[out]
```

`normalize=0` is important — `amix` otherwise divides by the input count, which
would make every clip inaudibly quiet across hundreds of inputs.

Above 500 clips the work is **batched**, because one ffmpeg process cannot hold
thousands of open inputs. Each batch is assembled to its own temporary FLAC and
the batches are then mixed together.

Batches use **relative** delays: the first clip's start time is subtracted from
every clip in the batch, and re-applied when the batches are combined. Without
that subtraction, batch 7 would begin with 40 minutes of silence, wasting time
and making the progress bar useless.

## 5. Muxing

`merge_voiceover()` runs **two separate ffmpeg passes**. This is deliberate and
should not be collapsed back into one.

### Pass 1 — build the lektor mix

```
[0:a:0]<norm>,aformat=channel_layouts=stereo,volume=<orig>[orig_mix];
[1:a]aformat=channel_layouts=stereo,volume=<tts>[tts];
[orig_mix][tts]amix=inputs=2:duration=first:normalize=0[lektor]
```

Written to a temporary `*_mix.m4a`, AAC 192k, and deleted in a `finally` block.

**Stereo, not 5.1.** ffmpeg's AAC encoder does not signal a 5.1 channel layout —
the track comes back tagged `unknown` in every container tested, and players
that cannot route the channels play silence. Stereo is tagged correctly, plays
everywhere, and spends the 192 kbit/s on two channels instead of six. Track 0
keeps the untouched 5.1, so no surround information is lost.

**PTS normalisation.** MKV rips often start their audio at a non-zero PTS (one
source here begins at 4.043 s). `amix` cannot emit output before both inputs
have frames, so the voiceover would be delayed by that offset. When ffprobe
reports an offset the filter becomes:

```
asetpts=PTS-STARTPTS,adelay=<offset_ms>:all=1,aresample=async=1
```

`asetpts` rebases the stream to zero and `adelay` reinserts the offset as real
silence, so the background audio stays where it belongs while `amix` receives
frames from both inputs starting at zero.

`duration=first` with `orig_mix` listed first means the mix runs as long as the
film, not as long as the voiceover.

### Pass 2 — stream-copy mux

Video, original audio, the mix, and every subtitle track are copied into the
final MKV with `-c copy`. Nothing is re-encoded, and it runs at several hundred
times real time.

Two flags carry the weight here:

**`-max_interleave_delta 0`**, placed **after the maps and before the output
path**. It is an output-muxer option; before `-i` it silently configures the
demuxer and does nothing. Without it, sparse subtitle tracks make the muxer
abandon interleaving after its default 10-second window, and the lektor audio
ends up written hundreds of megabytes away from the video it belongs with —
measurably fine to ffmpeg, but silent in any player that reads sequentially.

**MKV output**, not MP4. MKV copies subtitles of any codec without conversion,
carries DTS and TrueHD in track 0, and avoids the `mov_text` conversion that
made large multi-subtitle files abort mid-mux.

### Why two passes

Doing the filtering and the subtitle muxing in one command fails on sources with
many subtitle tracks. The sparse subtitle streams stall the muxer while the
audio filter graph keeps producing, and ffmpeg aborts partway through with:

```
Error while filtering: Invalid argument
Failed to inject frame into filter network: Invalid argument
```

The message points at audio, but audio is not the cause. Splitting the work
keeps the filter graph free of subtitle streams and makes the final mux a pure
copy. It is also about ten times faster.

## Progress reporting

`_run_progress()` adds `-progress pipe:1 -nostats` before the output path, reads
`out_time_us=` lines from stdout, and rewrites a single line with `\r`. A daemon
thread drains stderr concurrently, so ffmpeg cannot deadlock on a full stderr
pipe while the main thread reads stdout.

## Design rules

Worth preserving, since each was learned by breaking it:

1. **Per-clip processing before assembly**, never on the assembled track.
2. **Never trim clips** to the subtitle window.
3. **Keep every clip file** — `--regen` depends on them.
4. **FLAC throughout**, encoding only the final voiceover track.
5. **Two audio tracks**, with the original stream-copied and untouched.
6. **Verify muxed output by byte position**, not only by volume. A track can
   measure perfectly and still be unplayable.