# Copyright 2026 Viggle AI. Licensed under the Apache License, Version 2.0 (see LICENSE-CODE).
# SPDX-License-Identifier: Apache-2.0
"""Re-camera a video from the command line: source clip + an authored camera move -> the model's video.

    # 3-forward DMD turbo (the default: stock MiniMax-H3 + both of this repo's adapters):
    python inference/sample.py --video clip.mp4 --yaw 15 --sweep --out out/clip_yaw15
    # 50-step teacher (the recam adapter alone):
    python inference/sample.py --video clip.mp4 --yaw 15 --sweep --lora teacher_lora --steps 50 --flow-shift 12 --out out/t

The clip is letterboxed into a 1280x1280 frame (`FULL`), reconstructed by one VGGT-Omega pass, and re-rendered
from a second camera rigidly attached to the source camera: `c2w_dst[t] = c2w_src[t] @ delta(t)`, where `delta`
orbits about the pivot at frame 0's median depth (`--yaw`, degrees) and/or trucks sideways (`--truck`, in units of
that depth). `--sweep` ramps `delta` from identity at frame 0 to its full value at the last frame, so a static
source camera turns into a moving one. `--freeze F:N` is bullet time: the window is `start..F-1`, then source
frame `F` held for `N` frames while `delta` ramps from identity to its full value, then `F+1..` for whatever is
left of the window; the frozen frames share `F`'s geometry, so the render is a moving camera over a static cloud.
The render uses the source's own per-frame intrinsics -- there is no target clip to take them from.

`--video` may also be a still image (PNG), which with `--freeze 0:73` makes the whole window that one frame.

Writes to `--out`: out.mp4, out_audio.mp4 (source soundtrack of the same window muxed back on; not for `--freeze`),
render.mp4, source.mp4, grid.mp4 = [source | render | out] at the target canvas, last.png = out's final frame, and
cams.npz with the source and target cameras.
"""

import argparse
import json
import math
import os
import subprocess
import sys
import time

import av
import numpy as np
import torch
from diffusers import AutoencoderKLMiniMaxH3, MiniMaxH3Transformer3DModel
from diffusers.utils.export_utils import encode_video as write_mp4
from PIL import Image

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
import recam.geometry as geo  # noqa: E402
from recam.geometry import FULL, LENGTHS, NUM_FRAMES, RES, reconstruct, resize_u8, to_input, vggt, warp  # noqa: E402
from recam.h3 import FPS, bucket, decode_video, denoise, encode_video, pack  # noqa: E402
from recam.path import plan_path  # noqa: E402

ASSETS = f"{ROOT}/assets"

parser = argparse.ArgumentParser()
parser.add_argument("--video", required=True)
parser.add_argument("--camera-path", help="JSON with a 'path' key: recam.path keys in frame-0 camera coordinates / pivot depth; real-time source only, overrides parametric camera offsets")
parser.add_argument("--ckpt", default=None, help="a local diffusers `transformer/` dir; default is `transformer/` inside --model-dir. Meridian is the adapters, not a transformer -- this is stock MiniMax-H3")
parser.add_argument("--model-dir", default="MiniMaxAI/MiniMax-H3", help="the base MiniMax-H3 repo or a local copy of it, for `vae/`")
parser.add_argument("--start", type=int, default=0, help="first source frame of the 73-frame window")
parser.add_argument("--yaw", type=float, default=0.0, help="orbit about the frame-0 median-depth pivot, degrees; + moves the camera left")
parser.add_argument("--yaw-from", type=float, default=0.0, help="with --sweep/--freeze: ramp the yaw from this value instead of 0 (the lead-in holds it)")
parser.add_argument("--truck", type=float, default=0.0, help="sideways camera shift in units of the pivot depth; + moves right")
parser.add_argument("--sweep", action="store_true", help="ramp the offset linearly from 0 at frame 0 to its full value at the last frame; with --freeze the live lead-in and tail orbit too, at --live-speed")
parser.add_argument("--freeze", default=None, help="F:N -- hold source frame F for N frames and ramp the offset over them (bullet time)")
parser.add_argument("--live-speed", type=float, default=0.33, help="with --freeze --sweep: angular speed of the live lead-in and tail relative to the frozen frames")
parser.add_argument("--swing", action="store_true", help="sine ramp 0 -> 1 -> 0 -> -1 -> 0: orbit to one side, back through the source camera, out to the other side and back")
parser.add_argument("--ease", action="store_true", help="cosine ease-in-out on the ramp (start and end at rest)")
parser.add_argument("--bounce", action="store_true", help="there-and-back ramp 0 -> 1 -> 0 (cosine), so the chunk starts and ends at the source camera")
parser.add_argument("--dolly", type=float, default=1.0, help="orbit radius as a fraction of the pivot depth, reached along the ramp; focal scales with r so the pivot plane keeps its size (dolly zoom: closer and wider)")
parser.add_argument("--zoom", type=float, default=0.0, help="final focal multiplier, reached along the ramp; overrides the dolly's automatic focal scaling. --dolly 0.6 --zoom 1 is a true push-in (parallax, subject grows); --dolly 1 --zoom 1.6 is a pure optical zoom")
parser.add_argument("--boom", type=float, default=0.0, help="vertical camera shift in units of the pivot depth; + raises the camera (crane up)")
parser.add_argument("--pivot", default=None, help="fx,fy -- put the orbit pivot at the depth seen there (fractions of the crop box, frame 0); default: median depth")
parser.add_argument("--aim", action="store_true", help="after boom/truck/dolly, rotate the camera to put the pivot back where it was on screen (a crane that keeps looking at the subject); a no-op for pure --yaw with --pivot-lock")
parser.add_argument("--pivot-to", default=None, help="fx,fy -- a second picked pixel: with --aim, the camera pans/tilts off --pivot and ends up looking AT this point (it lands in frame centre). Pure rotation; orbit/boom/truck still use --pivot")
parser.add_argument("--pivot-lock", action="store_true", help="with --pivot: orbit about that 3D point instead of the optical axis, so an off-centre subject keeps its screen position")
parser.add_argument("--follow", action="store_true", help="replay the input clip's OWN estimated camera path over the geometry of frame --start alone: one VGGT pass gives both, so the path and the cloud share a gauge. Ignores --yaw/--truck/--dolly/... -- the trajectory comes from the video")
parser.add_argument("--smooth", type=float, default=8.0, help="with --follow: Gaussian sigma in frames used to low-pass VGGT's per-frame poses (translation, rotation and focal). VGGT estimates every frame independently, so the raw path jitters; 0 disables")
parser.add_argument("--cull", action="store_true", help="drop splats the target camera sees from behind (depth-map normal oriented to the source camera), so a 180-degree view is a hole, not the mirrored front")
parser.add_argument("--fast-back", type=float, default=1.0, help="K>1: sweep the middle half of the yaw range (the unobserved back) K times faster than the two observed quarters")
parser.add_argument("--seed", type=int, default=1234)
parser.add_argument("--gauge-only", action="store_true", help="stop after the geometry diagnostics (fly-through gauge, pivot depth, coverage) -- no VAE, no transformer, no render")
parser.add_argument("--preview-only", action="store_true", help="export the geometry and source previews -- full-resolution, plus the cond_*.mp4 pair at the condition canvas a sampler elsewhere needs -- then stop before loading the VAE or transformer")
parser.add_argument("--canvas", default="", help="render at an explicit WxH (multiples of 32) instead of the 768-class bucket. The rotary grid is normalised by sqrt(area), so the same aspect at a larger canvas keeps the position ids in distribution -- only the sampling density changes. Raise --full with it or the warp is upsampled from a 1280 source")
parser.add_argument("--full", type=int, default=0, help="side of the square the source is letterboxed into, default 1280 (the corpus's). Raise it alongside --canvas so the point cloud is unprojected and resampled at the output resolution")
parser.add_argument("--steps", type=int, default=4, help="scheduler timesteps: 4 = the student's 3 forwards")
parser.add_argument("--lora", nargs="+", default=[f"{ROOT}/teacher_lora", f"{ROOT}/turbo_lora"],
                    help="adapter dirs applied together on --ckpt: the recam teacher, then the DMD turbo. Both is the default and its grid is "
                         "--steps 4 --flow-shift 3; the teacher alone is --lora teacher_lora --steps 50 --flow-shift 12")
parser.add_argument("--no-lora", action="store_true", help="load no adapter at all: stock MiniMax-H3, which cannot re-camera")
parser.add_argument("--flow-shift", type=float, default=3.0, help="video scheduler shift: 3 for the student, 12 (MiniMax-H3's) for the teacher")
parser.add_argument("--frames", type=int, default=NUM_FRAMES, choices=LENGTHS, help="output length in frames (the lengths assets/ has a prompt embed for)")
parser.add_argument("--attn-backend", default="_native_cudnn")
parser.add_argument("--vggt", default=None, help="vggt_omega_1b_512.pt; default $VGGT_OMEGA_CKPT (see README)")
parser.add_argument("--vggt-repo", default=None, help="a checkout of facebookresearch/vggt-omega; default $VGGT_OMEGA_DIR")
parser.add_argument("--out", required=True)
args = parser.parse_args()
camera_path = None
if args.camera_path:
    assert not args.freeze and not args.follow, "--camera-path requires advancing source geometry"
    with open(args.camera_path) as f:
        camera_path = json.load(f)["path"]
    assert all(k["src"] == args.start + k["t"] for k in camera_path), "--camera-path currently supports real-time source only"
if args.full:  # `warp` scales VGGT's intrinsics by the module global, so both names must move
    geo.FULL = FULL = args.full
NUM_FRAMES = args.frames
device = torch.device("cuda")
torch.set_grad_enabled(False)
os.makedirs(args.out, exist_ok=True)
torch.manual_seed(args.seed)

# --- the clip, letterboxed into the corpus' square frame -------------------------------------------
# `tmap[t]` is the source frame shown at output frame t; only the distinct frames are decoded and reconstructed.
if args.freeze:
    fz, n = map(int, args.freeze.split(":"))
    tail = NUM_FRAMES - n - (fz - args.start)
    assert fz >= args.start and tail >= 0, f"freeze {fz}x{n} does not fit a {NUM_FRAMES}-frame window from {args.start}"
    tmap = list(range(args.start, fz)) + [fz] * n + list(range(fz + 1, fz + 1 + tail))
    ramp = torch.cat([torch.zeros(fz - args.start), torch.linspace(0, 1, n), torch.ones(tail)])
    if args.sweep:  # the live lead-in and tail orbit too, --live-speed times slower than the frozen frames
        w = torch.tensor([args.live_speed] * (fz - args.start) + [1.0] * n + [args.live_speed] * tail)
        ramp = torch.cumsum(w, 0) - w[0]
        ramp = ramp / ramp[-1]
else:
    tmap = list(range(args.start, args.start + NUM_FRAMES))
    # --bounce/--swing shape a 0->1 ramp.  on the constant `ones` ramp they collapse to
    # identically zero (cos 2pi = 1, sin 2pi = 0) and the camera never moves at all, so
    # they imply the linear base ramp -- there is no useful reading of the other combination.
    ramp = torch.linspace(0, 1, NUM_FRAMES) if args.sweep or args.bounce or args.swing else torch.ones(NUM_FRAMES)
if args.ease:
    ramp = (1 - torch.cos(math.pi * ramp)) / 2
if args.bounce:
    ramp = (1 - torch.cos(2 * math.pi * ramp)) / 2
if args.swing:
    ramp = torch.sin(2 * math.pi * ramp)
if args.fast_back > 1:  # piecewise-linear time->angle map: speed v on the outer quarters, K*v on the middle half
    K, v = args.fast_back, 0.5 * (1 + 1 / args.fast_back)
    t1 = 0.25 / v
    ramp = torch.where(ramp < t1, v * ramp, torch.where(ramp < 1 - t1, 0.25 + K * v * (ramp - t1), 0.75 + v * (ramp - 1 + t1)))
pf = tmap.index(fz) if args.freeze else 0  # the frame whose depth places the pivot
c = av.open(args.video)
frames = np.stack([f.to_ndarray(format="rgb24") for i, f in enumerate(c.decode(video=0)) if tmap[0] <= i <= tmap[-1]])
c.close()
assert len(frames) == tmap[-1] - tmap[0] + 1, f"decoded {len(frames)} frames from {tmap[0]}, need {tmap[-1] - tmap[0] + 1}"
frames = torch.from_numpy(frames).to(device)
h, w = frames.shape[1:3]
s = FULL / max(h, w)
ch, cw = round(h * s), round(w * s)
ox, oy = (FULL - cw) // 2, (FULL - ch) // 2
full = torch.zeros(len(frames), FULL, FULL, 3, dtype=torch.uint8, device=device)
full[:, oy : oy + ch, ox : ox + cw] = resize_u8(frames, (ch, cw))
del frames

# The 768-class canvas nearest the clip's aspect, and a crop box of exactly that aspect inside the content
# (`crop_box`'s construction, centred instead of drawn): the resize factor `f` stays isotropic, as in training.
canvas, cond_canvas = bucket(w, h)  # `<Video 1>` stays at the 480 class whatever the target is, as in training
if args.canvas:
    canvas = tuple(int(v) for v in args.canvas.split("x"))
a = canvas[0] / canvas[1]
bw, bh = (cw, round(cw / a)) if cw / ch <= a else (round(ch * a), ch)
box = (ox + (cw - bw) // 2, oy + (ch - bh) // 2, bw, bh, canvas[0] / bw)
x0, y0, bw, bh, _ = box
print(f"{w}x{h} -> content {cw}x{ch} in {FULL}^2, box {box[:4]}, canvas {canvas}, cond {cond_canvas}", flush=True)

# --- geometry: one solo pass, then the authored second camera in the same gauge ----------------------
geometry = vggt(args.vggt, args.vggt_repo, device)
t0 = time.time()
S = reconstruct(geometry, to_input(full))
del geometry
idx = torch.tensor(tmap, device=device) - tmap[0]
S = {k: v[idx] for k, v in S.items()}
full = full[idx]
if args.follow:  # keep the clip's per-frame cameras, then pin every frame's geometry to the first
    c2w_f, intr_f = S["c2w"].clone(), S["intr"].clone()
    if args.smooth:  # VGGT poses every frame independently, so the raw path jitters -- Gaussian low-pass it
        n = len(c2w_f)
        k = min(int(3 * args.smooth), n - 1)  # point-reflect the ends (x[-j] = 2x[0] - x[j]): a plain
        pad = lambda v: torch.cat([2 * v[:1] - v[1:k + 1].flip(0), v, 2 * v[-1:] - v[-k - 1:-1].flip(0)])
        ts = torch.arange(n + 2 * k, device=device, dtype=torch.float32)  # truncated window would pull the
        w = torch.exp(-0.5 * ((ts[k:k + n, None] - ts[None, :]) / args.smooth) ** 2)  # endpoints inward and
        w = w / w.sum(1, keepdim=True)  # eat 7% of the travel at sigma 8; reflection keeps the end velocity
        sm = lambda v: torch.einsum("ij,jab->iab", w, pad(v))
        jit = float((c2w_f[:, :3, 3] - sm(c2w_f)[:, :3, 3]).norm(dim=1).mean())
        c2w_f, intr_f = sm(c2w_f), sm(intr_f)
        U, _, Vh = torch.linalg.svd(c2w_f[:, :3, :3])  # blurring leaves R off SO(3); project it back
        U[:, :, 2] *= torch.linalg.det(U @ Vh)[:, None]  # never let the fit flip handedness
        c2w_f[:, :3, :3], c2w_f[:, 3] = U @ Vh, torch.tensor([0.0, 0.0, 0.0, 1.0], device=device)
        print(f"smoothed follow path, sigma {args.smooth} frames, removed {jit:.4f} mean jitter", flush=True)
    z = [0] * NUM_FRAMES
    S = {k: v[z] for k, v in S.items()}
    full = full[z]
if args.pivot and args.pivot != "none":  # median depth in a +-5% window around the picked point, in VGGT's 512 grid
    fx, fy = map(float, args.pivot.split(","))
    r = RES / FULL
    px, py, rw, rh = (x0 + fx * bw) * r, (y0 + fy * bh) * r, 0.05 * bw * r, 0.05 * bh * r
    win = (slice(round(py - rh), round(py + rh)), slice(round(px - rw), round(px + rw)))
    zm = float(S["depth"][pf][win][S["keep"][pf][win]].median())
else:
    zm = float(S["depth"][pf][S["keep"][pf]].median())
piv = torch.tensor([0.0, 0.0, zm], device=device)
if args.pivot and args.pivot != "none" and args.pivot_lock:  # unproject the picked pixel; orbit about it so it holds its screen position
    K = S["intr"][pf]
    piv = torch.tensor([(px - float(K[0, 2])) / float(K[0, 0]) * zm, (py - float(K[1, 2])) / float(K[1, 1]) * zm, zm], device=device)
    print(f"pivot locked at {piv.tolist()}", flush=True)
piv_to = piv
if args.pivot_to:  # a second picked pixel; the aim target slides from `piv` to it along the ramp
    gx, gy = map(float, args.pivot_to.split(","))
    g = RES / FULL
    qx, qy, qw, qh = (x0 + gx * bw) * g, (y0 + gy * bh) * g, 0.05 * bw * g, 0.05 * bh * g
    wn = (slice(round(qy - qh), round(qy + qh)), slice(round(qx - qw), round(qx + qw)))
    zt = float(S["depth"][pf][wn][S["keep"][pf][wn]].median())
    K = S["intr"][pf]
    piv_to = torch.tensor([(qx - float(K[0, 2])) / float(K[0, 0]) * zt, (qy - float(K[1, 2])) / float(K[1, 1]) * zt, zt], device=device)
    print(f"reframe target at {piv_to.tolist()}", flush=True)
c2w, intr_t = (c2w_f, intr_f) if args.follow else (S["c2w"].clone(), S["intr"].clone())
for ti in range(0 if args.follow or camera_path is not None else NUM_FRAMES):
    th = math.radians(args.yaw_from + (args.yaw - args.yaw_from) * float(ramp[ti]))
    r = 1 + (args.dolly - 1) * float(ramp[ti])
    R = torch.tensor([[math.cos(th), 0, math.sin(th)], [0, 1, 0], [-math.sin(th), 0, math.cos(th)]], device=device)
    delta = torch.eye(4, device=device)
    delta[:3, :3] = R
    # orbit about `piv`: sit at r*|piv| from it along the rotated line of sight, then truck/boom in the rotated frame
    delta[:3, 3] = piv - R @ (r * piv) - R @ torch.tensor([-args.truck * zm * float(ramp[ti]), args.boom * zm * float(ramp[ti]), 0.0], device=device)
    if args.aim:  # re-point at the pivot: boom/truck reframe the shot instead of sliding the subject out of frame
        a = piv / piv.norm()
        b = piv + (piv_to - piv) * float(ramp[ti]) - delta[:3, 3]
        b = b / b.norm()
        v, c = torch.cross(a, b, dim=0), float(a @ b)
        sn = float(v.norm())
        K = torch.zeros(3, 3, device=device)
        K[0, 1], K[0, 2], K[1, 0], K[1, 2], K[2, 0], K[2, 1] = -v[2], v[1], v[2], -v[0], -v[1], v[0]
        delta[:3, :3] = torch.eye(3, device=device) + K + K @ K * ((1 - c) / sn ** 2) if sn > 1e-8 else torch.eye(3, device=device)
    c2w[ti] = S["c2w"][ti] @ delta
    f = 1 + (args.zoom - 1) * float(ramp[ti]) if args.zoom else r
    intr_t[ti, 0, 0] *= f
    intr_t[ti, 1, 1] *= f
if camera_path is not None:
    poses, path_tmap, focal, _ = plan_path(camera_path, NUM_FRAMES, zm)
    assert path_tmap == tmap
    c2w = S["c2w"][0][None] @ torch.as_tensor(poses, dtype=c2w.dtype, device=device)
    focal = torch.as_tensor(focal, dtype=intr_t.dtype, device=device)
    intr_t[:, 0, 0] *= focal
    intr_t[:, 1, 1] *= focal
if args.cull:  # per frame: unproject the RES-grid depth, normal from finite differences, oriented toward the source camera
    yy, xx = torch.meshgrid(torch.arange(RES, device=device), torch.arange(RES, device=device), indexing="ij")
    uv1 = torch.stack([xx, yy, torch.ones_like(xx)], -1).float().reshape(-1, 3)
    for ti in range(NUM_FRAMES):
        Xc = (uv1 @ torch.linalg.inv(S["intr"][ti]).T).reshape(RES, RES, 3) * S["depth"][ti][..., None]
        X = Xc @ S["c2w"][ti][:3, :3].T + S["c2w"][ti][:3, 3]
        n = torch.cross(torch.roll(X, -1, 1) - X, torch.roll(X, -1, 0) - X, dim=-1)
        n = n * torch.sign((n * (S["c2w"][ti][:3, 3] - X)).sum(-1, keepdim=True))
        S["keep"][ti] &= ((c2w[ti][:3, 3] - X) * n).sum(-1) > 0
    print(f"cull: keep fraction per frame {S['keep'].float().mean((1, 2)).min():.2f}..{S['keep'].float().mean((1, 2)).max():.2f}", flush=True)
w2c = torch.linalg.inv(c2w)
render, cov = warp(S, w2c, intr_t, full, box, canvas)
yy, xx = torch.meshgrid(torch.arange(RES, device=device), torch.arange(RES, device=device), indexing="ij")
g1 = torch.stack([xx, yy, torch.ones_like(xx)], -1).float().reshape(-1, 3)
near, behind, coll, ahead = [], [], [], []
for ti in range(NUM_FRAMES):  # fly-through gauge: how much of the scene ends up near, behind, or (collide) within a 0.05-pivot-depth ball of the moved camera
    X = ((g1 @ torch.linalg.inv(S["intr"][ti]).T).reshape(RES, RES, 3) * S["depth"][ti][..., None]) @ S["c2w"][ti][:3, :3].T + S["c2w"][ti][:3, 3]
    zd = (X - c2w[ti][:3, 3]) @ c2w[ti][:3, 2]
    near.append(float((zd[S["keep"][ti]] < 0.1 * zm).float().mean()))
    behind.append(float((zd[S["keep"][ti]] < 0).float().mean()))
    coll.append(float(((X - c2w[ti][:3, 3]).norm(dim=-1)[S["keep"][ti]] < 0.05 * zm).float().mean()))
    cb = S["keep"][ti].clone(); cb[: RES // 4] = cb[-(RES // 4):] = False; cb[:, : RES // 4] = cb[:, -(RES // 4):] = False  # central half of the source frame: the subject, not the floor
    ahead.append(float(torch.quantile(zd[cb], 0.05)) / zm if cb.any() else 9.0)
print(f"fly-through gauge: scene within 0.1 pivot-depths of the moved camera, per-frame max {max(near):.3f} at frame {near.index(max(near))}; "
      f"behind the camera max {max(behind):.3f} at frame {behind.index(max(behind))}; "
      f"collide (within a 0.05-pivot-depth ball) max {max(coll):.4f} at frame {coll.index(max(coll))}; "
      f"central-subject 5th-pct depth ahead of the moved camera min {min(ahead):+.3f} pivot-depths at frame {ahead.index(min(ahead))}", flush=True)
print(f"pivot depth {zm:.3f}, camera moved {float((c2w[-1, :3, 3] - S['c2w'][-1, :3, 3]).norm()) / zm:.3f} pivot-depths "
      f"by the last frame, coverage {float(cov.float().mean()):.3f}, geometry in {time.time() - t0:.0f}s", flush=True)
if args.gauge_only:
    sys.exit(0)
if args.preview_only:
    source = resize_u8(full[:, y0 : y0 + bh, x0 : x0 + bw], canvas[::-1])
    write_mp4(render.cpu(), fps=int(FPS), output_path=f"{args.out}/render.mp4")
    write_mp4(source.cpu(), fps=int(FPS), output_path=f"{args.out}/source.mp4")
    # ... and the same pair at the condition canvas, which is what a sampler elsewhere has to be fed
    write_mp4(resize_u8(render, cond_canvas[::-1]).cpu(), fps=int(FPS), output_path=f"{args.out}/cond_render.mp4")
    write_mp4(resize_u8(full[:, y0 : y0 + bh, x0 : x0 + bw], cond_canvas[::-1]).cpu(), fps=int(FPS),
              output_path=f"{args.out}/cond_source.mp4")
    np.savez(f"{args.out}/requested_cameras.npz",
             c2w_src=S["c2w"].cpu().numpy(), c2w_dst=c2w.cpu().numpy(),
             intr_src=S["intr"].cpu().numpy(), intr_dst=intr_t.cpu().numpy(),
             piv=piv.cpu().numpy(), piv_to=piv_to.cpu().numpy(), zm=zm,
             box=box, canvas=canvas, fps=FPS, tmap=tmap, argv=" ".join(sys.argv[1:]))
    print(f"wrote geometry previews to {args.out}; no model video generated", flush=True)
    sys.exit(0)

# --- encode exactly as `build` does; the target latent is a placeholder that only sets the layout --------
render = resize_u8(render, cond_canvas[::-1])
source = resize_u8(full[:, y0 : y0 + bh, x0 : x0 + bw], canvas[::-1])
cond = resize_u8(full[:, y0 : y0 + bh, x0 : x0 + bw], cond_canvas[::-1])
del full
vae = AutoencoderKLMiniMaxH3.from_pretrained(args.model_dir, subfolder="vae").to(device).eval()
d = {"cond": encode_video(vae, cond)[0], "render": encode_video(vae, render)[0], "target": encode_video(vae, source)[1]}

embed = torch.load(f"{ASSETS}/fixed_embed_{NUM_FRAMES}.pt", weights_only=False)
audio_x0 = torch.load(f"{ASSETS}/silence_audio_{NUM_FRAMES}.pt", weights_only=True)["audio_x0"].float()
batch = pack(d["cond"], d["render"], d["target"], embed["prompt_embeds"][0], embed["text_token_tags"], audio_x0)

# --- denoise ----------------------------------------------------------------------------------------
ckpt, sfold = (args.ckpt, None) if args.ckpt else (args.model_dir, "transformer")
transformer = MiniMaxH3Transformer3DModel.from_pretrained(ckpt, subfolder=sfold, torch_dtype=torch.bfloat16)
if not args.no_lora:
    # `prefix=None` + the safetensors name, or the loader looks for a `.bin` / filters for `transformer.` keys and
    # silently loads nothing. The assert is the guard against that silence.
    # Each dir gets its own adapter name; loading alone leaves only the last one active, so `set_adapters` is what
    # sums them at weight 1.0 each -- which is how the turbo was distilled (see the model card).
    names = [f"lora{i}" for i in range(len(args.lora))]
    for name, ldir in zip(names, args.lora):
        transformer.load_lora_adapter(ldir, weight_name="pytorch_lora_weights.safetensors", prefix=None, adapter_name=name)
    assert any("lora_" in n for n, _ in transformer.named_parameters()), "no LoRA weights landed"
    transformer.set_adapters(names, [1.0] * len(names))
    print(f"lora: {len(names)} adapter(s) active {names}", flush=True)
transformer.set_attention_backend(args.attn_backend)
transformer.to(device).eval()
t0 = time.time()
torch.cuda.reset_peak_memory_stats()
rows = denoise(transformer, batch, args.steps, args.flow_shift, device)
print(f"denoised in {time.time() - t0:.0f}s, peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB", flush=True)
del transformer
out = decode_video(vae, rows, d["target"].shape[1:])

# --- write ----------------------------------------------------------------------------------------
hw = canvas[::-1]
for name, fr in ("out", out), ("render", render.cpu()), ("source", source.cpu()):
    write_mp4(fr, fps=int(FPS), output_path=f"{args.out}/{name}.mp4")
write_mp4(torch.cat([source.cpu(), resize_u8(render.cpu(), hw), out], 2), fps=int(FPS), output_path=f"{args.out}/grid.mp4")
Image.fromarray(out[-1].numpy()).save(f"{args.out}/last.png")  # the next link of an autoregressive orbit
if not args.freeze:  # a frozen window has no soundtrack that lines up
    subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", f"{args.out}/out.mp4", "-ss", f"{args.start / FPS:.4f}",
                    "-t", f"{NUM_FRAMES / FPS:.4f}", "-i", args.video, "-map", "0:v", "-map", "1:a?", "-c:v", "copy",
                    "-c:a", "aac", "-shortest", f"{args.out}/out_audio.mp4"], check=True)
# poses last: the mp4s are already on disk, so nothing here can cost a render
_x0, _y0, _bw, _bh, _sx = box  # intr_* live on the RES grid of the FULL^2 letterbox; *_px are in out.mp4 pixels
_M = np.array([[FULL / RES * _sx, 0, -_x0 * _sx], [0, FULL / RES * canvas[1] / _bh, -_y0 * canvas[1] / _bh], [0, 0, 1]])
np.savez(f"{args.out}/cams.npz", c2w_src=S["c2w"].cpu().numpy(), c2w_dst=c2w.cpu().numpy(),
         intr_src=S["intr"].cpu().numpy(), intr_dst=intr_t.cpu().numpy(),
         intr_src_px=_M @ S["intr"].cpu().numpy(), intr_dst_px=_M @ intr_t.cpu().numpy(),
         piv=piv.cpu().numpy(), piv_to=piv_to.cpu().numpy(), zm=float(zm),
         box=np.asarray(box), canvas=np.asarray(canvas), fps=float(FPS), tmap=tmap, argv=" ".join(sys.argv[1:]))
print(f"wrote {args.out}/{{out,render,source,grid}}.mp4 + last.png + cams.npz" + ("" if args.freeze else " + out_audio.mp4"))