decoder-video-gen-src/scripts/render_video_from_cache.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
from bisect import bisect_right
from concurrent.futures import ProcessPoolExecutor, as_completed
import functools
import json
import math
import os
from pathlib import Path
import re
import shutil
import subprocess
import time
import unicodedata

import cv2
import matplotlib
import numpy as np


ROOT = Path(__file__).resolve().parents[1]


def log(msg: str) -> None:
    ts = time.strftime("%H:%M:%S")
    print(f"[{ts}] {msg}", flush=True)


def resolve_path(p: str) -> Path:
    path = Path(p)
    if path.is_absolute():
        return path
    return (ROOT / path).resolve()


def parse_cache_steps(cache_dir: Path) -> tuple[list[str], list[int], list[Path], list[Path]]:
    uv = {}
    c4 = {}
    uv_re = re.compile(r"^uv_step(\d+)\.npy$")
    c4_re = re.compile(r"^c4_step(\d+)\.npy$")

    for p in cache_dir.glob("uv_step*.npy"):
        m = uv_re.match(p.name)
        if m:
            uv[m.group(1)] = p.resolve()
    for p in cache_dir.glob("c4_step*.npy"):
        m = c4_re.match(p.name)
        if m:
            c4[m.group(1)] = p.resolve()

    steps = sorted(set(uv.keys()) & set(c4.keys()), key=lambda s: int(s))
    if len(steps) < 2:
        raise RuntimeError(f"Need >=2 matching cache steps in {cache_dir}, found {len(steps)}")

    step_vals = [int(s) for s in steps]
    uv_paths = [uv[s] for s in steps]
    c4_paths = [c4[s] for s in steps]
    return steps, step_vals, uv_paths, c4_paths


VAL_LOSS_RE = re.compile(r"\bstep:(\d+)/\d+\b.*?\bval_loss:([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)\b")


def parse_stepwise_val_loss_points(log_path: Path) -> tuple[list[int], list[float]]:
    if not log_path.exists():
        return [], []
    lines = log_path.read_text(encoding="utf-8", errors="ignore").splitlines()
    by_step: dict[int, float] = {}
    for line in lines:
        m = VAL_LOSS_RE.search(line)
        if not m:
            continue
        try:
            step = int(m.group(1))
            val = float(m.group(2))
        except ValueError:
            continue
        if math.isfinite(val):
            by_step[step] = float(val)
    if not by_step:
        return [], []
    steps = sorted(by_step.keys())
    vals = [by_step[s] for s in steps]
    return steps, vals


def load_completions(completion_cache_json: Path, step_strs: list[str]) -> list[str]:
    if not completion_cache_json.exists():
        return ["" for _ in step_strs]

    obj = json.loads(completion_cache_json.read_text(encoding="utf-8"))
    cache_steps = obj.get("checkpoint_steps", [])
    cache_vals = obj.get("checkpoint_completions", [])
    if not isinstance(cache_steps, list) or not isinstance(cache_vals, list):
        return ["" for _ in step_strs]

    step_to_completion = {str(s): str(v) for s, v in zip(cache_steps, cache_vals)}
    return [step_to_completion.get(s, "") for s in step_strs]


def shift_subspace_colors(base_rgb: np.ndarray, t01: np.ndarray) -> np.ndarray:
    t = np.clip(t01, 0.0, 1.0).astype(np.float32)
    base = np.broadcast_to(base_rgb.reshape(1, 3).astype(np.float32), (t.shape[0], 3))
    out = np.empty_like(base)
    low = t < 0.5
    if np.any(low):
        dark_scale = 0.35 + 1.30 * t[low]
        out[low] = base[low] * dark_scale[:, None]
    if np.any(~low):
        light_mix = (t[~low] - 0.5) * 1.30
        out[~low] = base[~low] * (1.0 - light_mix[:, None]) + light_mix[:, None]
    return np.clip(out, 0.0, 1.0)


def update_cell_raster(
    cell: np.ndarray,
    uv: np.ndarray,
    c4: np.ndarray,
    *,
    umin: float,
    umax: float,
    vmin: float,
    vmax: float,
    c4min: float,
    c4max: float,
    base_colors: np.ndarray,
    alpha: float = 0.86,
) -> np.ndarray:
    h, w = cell.shape[:2]
    n_sub = int(uv.shape[0])
    acc = np.zeros((h, w, 3), dtype=np.float32)
    cnt = np.zeros((h, w), dtype=np.float32)

    ur = max(float(umax - umin), 1e-8)
    vr = max(float(vmax - vmin), 1e-8)
    c4r = max(float(c4max - c4min), 1e-8)

    for s in range(n_sub):
        us = uv[s, :, 0]
        vs = uv[s, :, 1]
        xs = ((us - umin) / ur) * (w - 1)
        ys = ((vs - vmin) / vr) * (h - 1)
        px = np.round(xs).astype(np.int32)
        py = np.round((h - 1) - ys).astype(np.int32)

        valid = (px >= 0) & (px < w) & (py >= 0) & (py < h)
        if not np.any(valid):
            continue
        px = px[valid]
        py = py[valid]
        t = (c4[s, valid] - c4min) / c4r
        cols = shift_subspace_colors(base_colors[s], t)
        np.add.at(acc, (py, px), cols)
        np.add.at(cnt, (py, px), 1.0)

    mask = cnt > 0
    if np.any(mask):
        avg = np.zeros_like(acc)
        avg[mask] = acc[mask] / cnt[mask][:, None]
        cell[mask] = cell[mask] * (1.0 - alpha) + avg[mask] * alpha
    return cell


@functools.lru_cache(maxsize=4096)
def _text_mask_supersampled(
    text: str,
    font_face: int,
    font_scale: float,
    thickness: int,
    ss: int,
) -> tuple[np.ndarray, int, int]:
    pad = 1
    (w0, h0), b0 = cv2.getTextSize(text, font_face, font_scale, thickness)
    low_w = max(1, int(w0 + 2 * pad))
    low_h = max(1, int(h0 + b0 + 2 * pad))
    hi_w = low_w * ss
    hi_h = low_h * ss
    mask_hi = np.zeros((hi_h, hi_w), dtype=np.uint8)
    cv2.putText(
        mask_hi,
        text,
        (pad * ss, (h0 + pad) * ss),
        font_face,
        float(font_scale) * ss,
        255,
        max(1, int(thickness) * ss),
        cv2.LINE_AA,
    )
    mask = cv2.resize(mask_hi, (low_w, low_h), interpolation=cv2.INTER_AREA)
    return mask, h0, pad


def draw_text_antialiased(
    frame: np.ndarray,
    text: str,
    org: tuple[int, int],
    font_face: int,
    font_scale: float,
    color_bgr: tuple[int, int, int],
    thickness: int = 1,
    ss: int = 4,
) -> None:
    text = (
        str(text)
        .replace("\u2018", "'")
        .replace("\u2019", "'")
        .replace("\u201A", "'")
        .replace("\u201B", "'")
        .replace("\u201C", '"')
        .replace("\u201D", '"')
        .replace("\u2013", "-")
        .replace("\u2014", "-")
        .replace("\u2026", "...")
        .replace("\u00A0", " ")
    )
    text = unicodedata.normalize("NFKC", text)
    safe = []
    for ch in text:
        o = ord(ch)
        if 32 <= o <= 126:
            safe.append(ch)
            continue
        approx = unicodedata.normalize("NFKD", ch).encode("ascii", "ignore").decode("ascii")
        safe.append(approx if approx else "?")
    text = "".join(safe)

    x, y = int(org[0]), int(org[1])
    mask, h0, pad = _text_mask_supersampled(text, int(font_face), float(font_scale), int(thickness), max(2, int(ss)))

    x0 = x - pad
    y0 = y - h0 - pad
    x1 = x0 + mask.shape[1]
    y1 = y0 + mask.shape[0]

    h, w = frame.shape[:2]
    cx0 = max(0, x0)
    cy0 = max(0, y0)
    cx1 = min(w, x1)
    cy1 = min(h, y1)
    if cx0 >= cx1 or cy0 >= cy1:
        return

    mx0 = cx0 - x0
    my0 = cy0 - y0
    mx1 = mx0 + (cx1 - cx0)
    my1 = my0 + (cy1 - cy0)

    alpha = (mask[my0:my1, mx0:mx1].astype(np.float32) / 255.0)[:, :, None]
    if not np.any(alpha > 0):
        return
    roi = frame[cy0:cy1, cx0:cx1].astype(np.float32)
    col = np.array(color_bgr, dtype=np.float32).reshape(1, 1, 3)
    out = roi * (1.0 - alpha) + col * alpha
    frame[cy0:cy1, cx0:cx1] = np.clip(out, 0.0, 255.0).astype(np.uint8)


def dynamic_layer_bounds(
    uv: np.ndarray,
    c4: np.ndarray,
    fallback: dict[str, float],
) -> tuple[float, float, float, float, float, float]:
    u = uv[:, :, 0].reshape(-1)
    v = uv[:, :, 1].reshape(-1)
    c = c4.reshape(-1)

    if u.size == 0 or v.size == 0:
        return (
            float(fallback["umin"]),
            float(fallback["umax"]),
            float(fallback["vmin"]),
            float(fallback["vmax"]),
            float(fallback["c4min"]),
            float(fallback["c4max"]),
        )

    uq0, uq1 = np.quantile(u, [0.01, 0.99])
    vq0, vq1 = np.quantile(v, [0.01, 0.99])
    uc = 0.5 * float(uq0 + uq1)
    vc = 0.5 * float(vq0 + vq1)
    r = 0.52 * max(float(uq1 - uq0), float(vq1 - vq0), 1e-4)
    umin, umax = uc - r, uc + r
    vmin, vmax = vc - r, vc + r

    cq0, cq1 = np.quantile(c, [0.02, 0.98])
    c4min, c4max = float(cq0), float(cq1)

    if not np.isfinite(umin) or not np.isfinite(umax) or abs(umax - umin) < 1e-6:
        umin, umax = float(fallback["umin"]), float(fallback["umax"])
    if not np.isfinite(vmin) or not np.isfinite(vmax) or abs(vmax - vmin) < 1e-6:
        vmin, vmax = float(fallback["vmin"]), float(fallback["vmax"])
    if not np.isfinite(c4min) or not np.isfinite(c4max) or abs(c4max - c4min) < 1e-6:
        c4min, c4max = float(fallback["c4min"]), float(fallback["c4max"])
    return umin, umax, vmin, vmax, c4min, c4max


def build_layer_bounds(uv_paths: list[Path], c4_paths: list[Path], n_layers: int) -> list[dict[str, float]]:
    bounds = [
        {"umin": np.inf, "umax": -np.inf, "vmin": np.inf, "vmax": -np.inf, "c4min": np.inf, "c4max": -np.inf}
        for _ in range(n_layers)
    ]
    for uv_path, c4_path in zip(uv_paths, c4_paths):
        uv = np.load(uv_path, mmap_mode="r")
        c4 = np.load(c4_path, mmap_mode="r")
        for i in range(n_layers):
            uvi = uv[i].astype(np.float32, copy=False)
            c4i = c4[i].astype(np.float32, copy=False)
            b = bounds[i]
            b["umin"] = min(float(b["umin"]), float(np.min(uvi[:, :, 0])))
            b["umax"] = max(float(b["umax"]), float(np.max(uvi[:, :, 0])))
            b["vmin"] = min(float(b["vmin"]), float(np.min(uvi[:, :, 1])))
            b["vmax"] = max(float(b["vmax"]), float(np.max(uvi[:, :, 1])))
            b["c4min"] = min(float(b["c4min"]), float(np.min(c4i)))
            b["c4max"] = max(float(b["c4max"]), float(np.max(c4i)))
    for b in bounds:
        if not np.isfinite(b["umin"]) or not np.isfinite(b["umax"]):
            b["umin"], b["umax"] = -1.0, 1.0
        if not np.isfinite(b["vmin"]) or not np.isfinite(b["vmax"]):
            b["vmin"], b["vmax"] = -1.0, 1.0
        if not np.isfinite(b["c4min"]) or not np.isfinite(b["c4max"]):
            b["c4min"], b["c4max"] = -1.0, 1.0
        if abs(b["umax"] - b["umin"]) < 1e-6:
            b["umin"] -= 1.0
            b["umax"] += 1.0
        if abs(b["vmax"] - b["vmin"]) < 1e-6:
            b["vmin"] -= 1.0
            b["vmax"] += 1.0
        if abs(b["c4max"] - b["c4min"]) < 1e-6:
            b["c4min"] -= 1.0
            b["c4max"] += 1.0
    return bounds


def render_frame_png(
    *,
    fi: int,
    transition_frames: int,
    checkpoint_count: int,
    checkpoint_step_values: list[int],
    checkpoint_completions: list[str],
    completion_prompt: str,
    tokens_per_step: float,
    val_loss_steps: list[int],
    val_loss_values: list[float],
    val_loss_label: str,
    total_tokens_b: float,
    uv_cache_paths: list[str],
    c4_cache_paths: list[str],
    layer_bounds: list[dict[str, float]],
    base_colors: np.ndarray,
    width: int,
    height: int,
    cols: int,
    margin_l: int,
    margin_t: int,
    cell_w: int,
    cell_h: int,
    gx: int,
    gy: int,
    out_png: str,
) -> int:
    if fi >= transition_frames:
        i0 = checkpoint_count - 2
        i1 = checkpoint_count - 1
        seg_t = 1.0
        ck_display_idx = checkpoint_count - 1
    else:
        tau = fi / float(max(1, transition_frames - 1))
        p = tau * float(checkpoint_count - 1)
        i0 = int(math.floor(p))
        i0 = min(i0, checkpoint_count - 2)
        i1 = i0 + 1
        seg_t = p - float(i0)
        ck_display_idx = int(round(p))
        ck_display_idx = max(0, min(checkpoint_count - 1, ck_display_idx))

    uv0 = np.load(uv_cache_paths[i0], mmap_mode="r")
    c40 = np.load(c4_cache_paths[i0], mmap_mode="r")
    uv1 = np.load(uv_cache_paths[i1], mmap_mode="r")
    c41 = np.load(c4_cache_paths[i1], mmap_mode="r")

    n_layers = int(uv0.shape[0])
    n_subspaces = int(uv0.shape[1])
    frame = np.full((height, width, 3), 255, dtype=np.uint8)

    for layer in range(n_layers):
        r = layer // cols
        c = layer % cols
        x0 = margin_l + c * (cell_w + gx)
        y0 = margin_t + r * (cell_h + gy)
        uv = uv0[layer].astype(np.float32, copy=False) * (1.0 - seg_t) + uv1[layer].astype(np.float32, copy=False) * seg_t
        c4 = c40[layer].astype(np.float32, copy=False) * (1.0 - seg_t) + c41[layer].astype(np.float32, copy=False) * seg_t

        cell = np.full((cell_h, cell_w, 3), 1.0, dtype=np.float32)
        b = layer_bounds[layer]
        umin, umax, vmin, vmax, c4min, c4max = dynamic_layer_bounds(uv, c4, b)
        cell = update_cell_raster(
            cell,
            uv,
            c4,
            umin=umin,
            umax=umax,
            vmin=vmin,
            vmax=vmax,
            c4min=c4min,
            c4max=c4max,
            base_colors=base_colors,
            alpha=0.86,
        )
        cell_u8 = np.clip(cell * 255.0, 0.0, 255.0).astype(np.uint8)
        frame[y0 : y0 + cell_h, x0 : x0 + cell_w] = cell_u8
        cv2.rectangle(frame, (x0, y0), (x0 + cell_w - 1, y0 + cell_h - 1), (210, 210, 210), 1)
        draw_text_antialiased(frame, f"L{layer}", (x0 + 3, y0 + 11), cv2.FONT_HERSHEY_SIMPLEX, 0.30, (35, 35, 35), 1)

    step_interp = float(checkpoint_step_values[i0]) * (1.0 - seg_t) + float(checkpoint_step_values[i1]) * seg_t
    tokens_b = step_interp * float(tokens_per_step) / 1e9
    display_val_loss = None
    if val_loss_steps:
        idx = bisect_right(val_loss_steps, int(math.floor(step_interp))) - 1
        if idx >= 0:
            display_val_loss = float(val_loss_values[idx])

    draw_text_antialiased(frame, "How do LLM decoders change over training?", (20, 26), cv2.FONT_HERSHEY_SIMPLEX, 0.72, (25, 25, 25), 2)
    tokens_line = f"Tokens trained: {tokens_b:0.3f}B"
    if display_val_loss is not None:
        tokens_line += f"   {val_loss_label}: {display_val_loss:0.3f}"
    draw_text_antialiased(frame, tokens_line, (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.56, (40, 40, 40), 1)

    completion_text = checkpoint_completions[ck_display_idx] if 0 <= ck_display_idx < len(checkpoint_completions) else ""
    completion_line = f'Completion for "{completion_prompt}": {completion_text}'
    draw_text_antialiased(frame, completion_line, (20, 72), cv2.FONT_HERSHEY_SIMPLEX, 0.40, (45, 45, 45), 1)

    lx = width - 420
    ly = 14
    cv2.rectangle(frame, (lx - 8, ly - 6), (width - 16, ly + 56), (245, 245, 245), -1)
    cv2.rectangle(frame, (lx - 8, ly - 6), (width - 16, ly + 56), (160, 160, 160), 1)
    draw_text_antialiased(frame, "Subspace key", (lx, ly + 8), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (30, 30, 30), 1)
    for s in range(n_subspaces):
        cx = lx + (s % 8) * 48
        cy = ly + 22 + (s // 8) * 18
        col = tuple(int(x) for x in (base_colors[s] * 255.0)[::-1])
        cv2.rectangle(frame, (cx, cy - 7), (cx + 12, cy + 5), col, -1)
        cv2.rectangle(frame, (cx, cy - 7), (cx + 12, cy + 5), (0, 0, 0), 1)
        draw_text_antialiased(frame, f"S{s}", (cx + 16, cy + 4), cv2.FONT_HERSHEY_SIMPLEX, 0.35, (40, 40, 40), 1)

    note_lines = [
        f"The charts above show all 128 layers of an 18-million parameter transformer model while being trained on {total_tokens_b:0.1f} billion tokens.",
        "Each layer has 16x K-Splanifolds (fast spline-manifolds) which encode geometric transformations instead of a neural network MLP.",
        "Model architecture uses a byte-level vocab and 8x 16-layer AttnRes blocks.",
    ]
    line_h = 15
    box_pad_top = 7
    box_pad_bottom = 7
    box_h = box_pad_top + box_pad_bottom + line_h * len(note_lines)
    box_top = height - 8 - box_h
    cv2.rectangle(frame, (14, box_top), (width - 14, height - 8), (245, 245, 245), -1)
    cv2.rectangle(frame, (14, box_top), (width - 14, height - 8), (160, 160, 160), 1)
    for i, line in enumerate(note_lines):
        draw_text_antialiased(frame, line, (22, box_top + box_pad_top + 11 + i * line_h), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (35, 35, 35), 1)

    ok = cv2.imwrite(out_png, frame, [cv2.IMWRITE_PNG_COMPRESSION, 2])
    if not ok:
        raise RuntimeError(f"Failed to write frame image: {out_png}")
    return fi


def main() -> None:
    parser = argparse.ArgumentParser(description="Render decoder projection progression video from cached projection arrays.")
    parser.add_argument("--cache-dir", default="data/cache", help="Directory containing uv_stepXXXXXX.npy and c4_stepXXXXXX.npy")
    parser.add_argument("--completion-cache-json", default="data/completions.json", help="Checkpoint completion cache json")
    parser.add_argument("--training-log", default="data/training_log.txt", help="Training log for step-wise val_loss display")
    parser.add_argument("--music", default="data/music/hydrostatic_mind.mp3", help="Optional background music file")
    parser.add_argument("--no-music", action="store_true", help="Disable background music mux")
    parser.add_argument("--music-volume", type=float, default=0.8, help="Music volume multiplier")
    parser.add_argument("--music-fade-sec", type=float, default=5.0, help="Outro fade duration in seconds")
    parser.add_argument("--duration-sec", type=float, default=158.0, help="Transition duration")
    parser.add_argument("--end-hold-sec", type=float, default=5.0, help="Hold last checkpoint")
    parser.add_argument("--fps", type=int, default=60, help="Frames per second")
    parser.add_argument("--width", type=int, default=1920)
    parser.add_argument("--height", type=int, default=1080)
    parser.add_argument("--workers", type=int, default=max(1, os.cpu_count() or 1), help="Parallel frame workers")
    parser.add_argument("--tokens-per-step", type=float, default=500000.0)
    parser.add_argument("--completion-prompt", default="The cat is")
    parser.add_argument("--val-loss-label", default="Val loss (byte vocab, 500k tokens, fineweb)")
    parser.add_argument("--ffmpeg-preset", default="medium")
    parser.add_argument("--ffmpeg-crf", type=int, default=14, help="Used only when --video-bitrate is unset")
    parser.add_argument("--ffmpeg-pix-fmt", default="yuv420p")
    parser.add_argument("--ffmpeg-threads", type=int, default=0)
    parser.add_argument("--video-bitrate", default="14M", help="Set target bitrate (e.g. 14M). Empty disables CBR mode")
    parser.add_argument("--video-maxrate", default=None)
    parser.add_argument("--video-bufsize", default=None)
    parser.add_argument(
        "--keep-frame-pngs",
        dest="keep_frame_pngs",
        action="store_true",
        default=True,
        help="Keep rendered frame PNGs (default: on).",
    )
    parser.add_argument("--cleanup-frame-pngs", dest="keep_frame_pngs", action="store_false")
    parser.add_argument("--output", default="outputs/subspace_projection_progression_video.mp4")
    args = parser.parse_args()

    cache_dir = resolve_path(args.cache_dir)
    completion_json = resolve_path(args.completion_cache_json)
    training_log = resolve_path(args.training_log)
    music_path = resolve_path(args.music)
    output_path = resolve_path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    step_strs, step_vals, uv_paths, c4_paths = parse_cache_steps(cache_dir)
    completions = load_completions(completion_json, step_strs)
    val_loss_steps, val_loss_values = parse_stepwise_val_loss_points(training_log)
    total_tokens_b = float(step_vals[-1]) * float(args.tokens_per_step) / 1e9
    log(f"Checkpoints from cache: {len(step_strs)} ({step_strs[0]}..{step_strs[-1]})")
    if val_loss_steps:
        log(
            "Loaded val_loss points: "
            f"{len(val_loss_steps)} entries, first={val_loss_steps[0]}, "
            f"latest={val_loss_steps[-1]}:{val_loss_values[-1]:0.4f}"
        )

    uv0 = np.load(uv_paths[0], mmap_mode="r")
    n_layers = int(uv0.shape[0])
    n_subspaces = int(uv0.shape[1])
    cols = 16
    rows = int(math.ceil(n_layers / cols))
    layer_bounds = build_layer_bounds(uv_paths, c4_paths, n_layers)

    width = int(args.width)
    height = int(args.height)
    fps = max(1, int(args.fps))
    transition_frames = max(2, int(round(float(args.duration_sec) * fps)))
    hold_frames = max(0, int(round(float(args.end_hold_sec) * fps)))
    frames = transition_frames + hold_frames

    margin_l, margin_r = 24, 24
    margin_t, margin_b = 92, 56
    gx, gy = 2, 2
    cell_w = int((width - margin_l - margin_r - (cols - 1) * gx) / cols)
    cell_h = int((height - margin_t - margin_b - (rows - 1) * gy) / rows)
    base_colors = matplotlib.colormaps["tab20"](np.linspace(0.0, 1.0, n_subspaces, endpoint=False))[:, :3].astype(np.float32)

    frame_dir = output_path.parent / f"{output_path.stem}_frames"
    if frame_dir.exists():
        shutil.rmtree(frame_dir)
    frame_dir.mkdir(parents=True, exist_ok=True)
    log(f"Rendering {frames} frames -> {frame_dir}")

    uv_strs = [str(p) for p in uv_paths]
    c4_strs = [str(p) for p in c4_paths]
    workers = max(1, int(args.workers))
    if workers == 1:
        for fi in range(frames):
            out_png = frame_dir / f"frame_{fi:06d}.png"
            render_frame_png(
                fi=fi,
                transition_frames=transition_frames,
                checkpoint_count=len(step_strs),
                checkpoint_step_values=step_vals,
                checkpoint_completions=completions,
                completion_prompt=str(args.completion_prompt),
                tokens_per_step=float(args.tokens_per_step),
                val_loss_steps=val_loss_steps,
                val_loss_values=val_loss_values,
                val_loss_label=str(args.val_loss_label),
                total_tokens_b=float(total_tokens_b),
                uv_cache_paths=uv_strs,
                c4_cache_paths=c4_strs,
                layer_bounds=layer_bounds,
                base_colors=base_colors,
                width=width,
                height=height,
                cols=cols,
                margin_l=margin_l,
                margin_t=margin_t,
                cell_w=cell_w,
                cell_h=cell_h,
                gx=gx,
                gy=gy,
                out_png=str(out_png),
            )
            if fi % max(1, fps) == 0 or fi == frames - 1:
                log(f"Rendered {fi + 1}/{frames}")
    else:
        futures = []
        done = 0
        with ProcessPoolExecutor(max_workers=workers) as ex:
            for fi in range(frames):
                out_png = frame_dir / f"frame_{fi:06d}.png"
                futures.append(
                    ex.submit(
                        render_frame_png,
                        fi=fi,
                        transition_frames=transition_frames,
                        checkpoint_count=len(step_strs),
                        checkpoint_step_values=step_vals,
                        checkpoint_completions=completions,
                        completion_prompt=str(args.completion_prompt),
                        tokens_per_step=float(args.tokens_per_step),
                        val_loss_steps=val_loss_steps,
                        val_loss_values=val_loss_values,
                        val_loss_label=str(args.val_loss_label),
                        total_tokens_b=float(total_tokens_b),
                        uv_cache_paths=uv_strs,
                        c4_cache_paths=c4_strs,
                        layer_bounds=layer_bounds,
                        base_colors=base_colors,
                        width=width,
                        height=height,
                        cols=cols,
                        margin_l=margin_l,
                        margin_t=margin_t,
                        cell_w=cell_w,
                        cell_h=cell_h,
                        gx=gx,
                        gy=gy,
                        out_png=str(out_png),
                    )
                )
            for fut in as_completed(futures):
                _ = fut.result()
                done += 1
                if done % max(1, fps) == 0 or done == frames:
                    log(f"Rendered {done}/{frames}")

    if shutil.which("ffmpeg") is None:
        raise RuntimeError("ffmpeg is required")

    video_only = output_path.parent / f"{output_path.stem}.video_only.mp4"
    ffmpeg_cmd = [
        "ffmpeg",
        "-y",
        "-framerate",
        str(fps),
        "-i",
        str(frame_dir / "frame_%06d.png"),
        "-c:v",
        "libx264",
        "-preset",
        str(args.ffmpeg_preset),
        "-pix_fmt",
        str(args.ffmpeg_pix_fmt),
    ]
    if args.video_bitrate:
        ffmpeg_cmd += ["-b:v", str(args.video_bitrate)]
        ffmpeg_cmd += ["-maxrate", str(args.video_maxrate or args.video_bitrate)]
        default_bufsize = f"{int(float(str(args.video_bitrate).rstrip('M')) * 2)}M" if str(args.video_bitrate).endswith("M") else str(args.video_bitrate)
        ffmpeg_cmd += ["-bufsize", str(args.video_bufsize or default_bufsize)]
    else:
        ffmpeg_cmd += ["-crf", str(int(args.ffmpeg_crf))]
    if int(args.ffmpeg_threads) > 0:
        ffmpeg_cmd += ["-threads", str(int(args.ffmpeg_threads))]
    ffmpeg_cmd += ["-movflags", "+faststart", str(video_only)]
    log("Encoding video stream")
    subprocess.run(ffmpeg_cmd, check=True)

    if args.no_music or not music_path.exists():
        shutil.move(str(video_only), str(output_path))
    else:
        fade_start = max(0.0, float(args.duration_sec + args.end_hold_sec) - float(args.music_fade_sec))
        mux_cmd = [
            "ffmpeg",
            "-y",
            "-i",
            str(video_only),
            "-i",
            str(music_path),
            "-filter_complex",
            (
                f"[1:a]atrim=0:{float(args.duration_sec + args.end_hold_sec):.6f},"
                f"asetpts=PTS-STARTPTS,volume={float(args.music_volume):.6f},"
                f"afade=t=out:st={fade_start:.6f}:d={float(args.music_fade_sec):.6f}[a]"
            ),
            "-map",
            "0:v:0",
            "-map",
            "[a]",
            "-c:v",
            "copy",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-shortest",
            "-movflags",
            "+faststart",
            str(output_path),
        ]
        log("Muxing music")
        subprocess.run(mux_cmd, check=True)
        video_only.unlink(missing_ok=True)

    if not args.keep_frame_pngs:
        shutil.rmtree(frame_dir, ignore_errors=True)

    meta = {
        "output": str(output_path),
        "cache_dir": str(cache_dir),
        "checkpoint_steps": step_strs,
        "checkpoint_step_values": step_vals,
        "tokens_per_step": float(args.tokens_per_step),
        "completion_cache_json": str(completion_json),
        "training_log": str(training_log),
        "val_loss_steps": [int(x) for x in val_loss_steps],
        "val_loss_values": [float(x) for x in val_loss_values],
        "duration_sec": float(args.duration_sec),
        "end_hold_sec": float(args.end_hold_sec),
        "fps": int(fps),
        "frames": int(frames),
        "resolution": [int(width), int(height)],
        "grid": {"cols": int(cols), "rows": int(rows), "layers": int(n_layers)},
        "subspaces": int(n_subspaces),
        "music": None if args.no_music else str(music_path),
        "music_volume": float(args.music_volume),
        "music_fade_sec": float(args.music_fade_sec),
        "frame_dir": str(frame_dir),
    }
    meta_path = output_path.with_suffix(".json")
    meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
    log(f"Wrote {output_path}")
    log(f"Wrote {meta_path}")


if __name__ == "__main__":
    main()