render.py

#!/usr/bin/env python3
"""
Build ONCE end to end.

  python3 render.py            # everything: segments, score, cut, master
  python3 render.py segments   # just re-render the picture
  python3 render.py assemble   # just re-cut from existing segments

Movements render as independent processes — six cores, six movements.
"""
from __future__ import annotations

import os
import subprocess
import sys
import time
from multiprocessing import Process

from once import core as C
from once import run as R
from once import score as S

WORK = "work"
OUT = "out"
SEG = os.path.join(WORK, "seg_{}.mkv")
SCORE = os.path.join(WORK, "score.wav")
MASTER = os.path.join(OUT, "ONCE_1080p.mp4")


def _one(name: str, seconds: float):
    os.environ.setdefault("NUMBA_NUM_THREADS", "2")
    mod = R.load(name)
    n = C.frames_for(seconds)
    mv = mod.build(n)
    post_kw = dict(getattr(mod, "POST", {}))
    t0 = time.time()
    with C.Encoder(SEG.format(name), lossless=True) as enc:
        for i in range(n):
            enc.write(C.post(mv.frame(i), frame_idx=i, **post_kw))
            if i % 90 == 0:
                el = time.time() - t0
                r = (i + 1) / max(el, 1e-6)
                print(f"  [{name:9s}] {i+1:4d}/{n}  {r:4.2f}fps  "
                      f"eta {(n-i-1)/max(r,1e-6):5.0f}s", flush=True)
    print(f"  [{name:9s}] COMPLETE {time.time()-t0:.0f}s", flush=True)


def segments():
    print(f"rendering {len(C.MOVEMENTS)} movements in parallel\n", flush=True)
    t0 = time.time()
    procs = [Process(target=_one, args=(n, d)) for n, d in C.MOVEMENTS]
    for p in procs:
        p.start()
    for p in procs:
        p.join()
    bad = [p.exitcode for p in procs if p.exitcode != 0]
    if bad:
        raise SystemExit(f"a movement failed (exit codes {bad})")
    print(f"\nall segments done in {time.time()-t0:.0f}s", flush=True)


def assemble():
    """Chain the segments with crossfades, then marry the score."""
    names = [n for n, _ in C.MOVEMENTS]
    durs = [d for _, d in C.MOVEMENTS]

    ins = []
    for n in names:
        ins += ["-i", SEG.format(n)]

    # xfade offset is measured against the accumulated output, which shortens
    # by one xfade at every join
    parts, prev, acc = [], "[0:v]", durs[0]
    for k in range(1, len(names)):
        off = acc - C.XFADE
        lbl = f"[v{k}]"
        parts.append(f"{prev}[{k}:v]xfade=transition=fade:"
                     f"duration={C.XFADE}:offset={off:.4f}{lbl}")
        prev = lbl
        acc = acc + durs[k] - C.XFADE
    parts.append(f"{prev}format=yuv420p[vout]")
    fc = ";".join(parts)

    os.makedirs(OUT, exist_ok=True)
    cmd = (["ffmpeg", "-y", "-hide_banner", "-loglevel", "error"] + ins +
           ["-i", SCORE, "-filter_complex", fc, "-map", "[vout]",
            "-map", f"{len(names)}:a",
            "-c:v", "libx264", "-preset", "veryslow", "-crf", "16",
            "-pix_fmt", "yuv420p", "-profile:v", "high", "-level", "4.2",
            "-x264-params", "aq-mode=3:aq-strength=1.05:psy-rd=1.0:deblock=-1,-1",
            "-c:a", "aac", "-b:a", "320k", "-ar", "48000",
            "-movflags", "+faststart", "-shortest", MASTER])
    print(f"cutting -> {MASTER}", flush=True)
    t0 = time.time()
    subprocess.run(cmd, check=True)
    print(f"cut in {time.time()-t0:.0f}s", flush=True)


def main():
    what = sys.argv[1] if len(sys.argv) > 1 else "all"
    os.makedirs(WORK, exist_ok=True)
    if what in ("all", "segments"):
        segments()
    if what in ("all", "score"):
        S.build(C.MOVEMENTS, C.XFADE, SCORE)
    if what in ("all", "assemble"):
        assemble()
    if os.path.exists(MASTER):
        mb = os.path.getsize(MASTER) / 1e6
        print(f"\nONCE — {mb:.1f} MB — {MASTER}")


if __name__ == "__main__":
    main()