run_native.py

"""Drive the 4DAnyone phases against a native Rerun sink, no Gradio involved.

This is the debugging seam the Gradio app is a wrapper over: the exact
``begin_phase`` / ``source_phase`` / ``motion_phase`` / ``generate_phase`` /
``publish_phase`` implementation runs here with a file or live-viewer sink, so a
rendering problem can be split into "the logging is wrong" (reproduces here)
versus "the gradio/web-viewer plumbing is wrong" (only in the browser).

Nothing forks here, so one recording serves every phase; the app opens a fresh
one per link only because ZeroGPU forks each of its GPU callbacks.

Examples:
    pixi run python run_native.py --smoke --sink save --out /tmp/smoke.rrd
    pixi run python run_native.py --sink spawn            # live native viewer
    pixi run python run_native.py --sink serve            # rerun+http://... to connect
"""

from __future__ import annotations

from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, TypeAlias

import rerun as rr
import tyro

SinkChoice: TypeAlias = Literal["save", "spawn", "serve", "connect"]


@dataclass(frozen=True)
class NativeArgs:
    """CLI arguments for a native, Gradio-free run."""

    video: Path = Path(__file__).parent / "examples" / "jump-rope.mp4"
    """Source clip; defaults to the bundled jump-rope example."""
    start_time: float = 0.0
    """Seconds into the clip where the 121-frame window begins."""
    seed: int = 0
    """Generation seed."""
    smoke: bool = False
    """Stream synthetic data through the real machinery, no GPU."""
    sink: SinkChoice = "save"
    """Where the recording goes: an .rrd file, a spawned viewer, or a gRPC server."""
    out: Path = Path("/tmp/4danyone-native.rrd")
    """Output recording path when ``sink`` is ``save``."""
    addr: str = "rerun+http://127.0.0.1:9876/proxy"
    """Running viewer's proxy URL when ``sink`` is ``connect``."""
    save_also: Path | None = None
    """With ``connect``: additionally write the recording to this .rrd file."""


def main(args: NativeArgs) -> None:
    from fdanyone.rerun_streaming import (
        APPLICATION_ID,
        RunSpec,
        Session,
        begin_phase,
        generate_phase,
        motion_phase,
        new_spec,
        publish_phase,
        smoke_generate_phase,
        smoke_motion_phase,
        smoke_publish_phase,
        source_phase,
    )

    spec: RunSpec = new_spec(args.video, args.start_time, args.seed)
    recording: rr.RecordingStream = rr.RecordingStream(APPLICATION_ID, recording_id=spec.token)
    if args.sink == "save":
        recording.save(str(args.out))
    elif args.sink == "spawn":
        recording.spawn()
    elif args.sink == "connect":
        if args.save_also is not None:
            recording.set_sinks(rr.GrpcSink(url=args.addr), rr.FileSink(str(args.save_also)))
        else:
            recording.connect_grpc(args.addr)
        print(f"[native] connected to {args.addr}", flush=True)
    else:
        server_uri: str = recording.serve_grpc()
        print(f"[native] serving at {server_uri}", flush=True)

    session: Session = Session(spec)
    begin_phase(recording, spec)
    print(f"[native] {source_phase(recording, spec)}", flush=True)
    motion: Iterator[str] = (
        smoke_motion_phase(recording, session) if args.smoke else motion_phase(recording, session)
    )
    for label in motion:
        print(f"[native] {label}", flush=True)
    generate: Iterator[str] = (
        smoke_generate_phase(recording, session)
        if args.smoke
        else generate_phase(recording, session)
    )
    for label in generate:
        print(f"[native] {label}", flush=True)
    assert session.summary is not None, "the generation phase published no result"
    published: str = (
        smoke_publish_phase(recording, spec)
        if args.smoke
        else publish_phase(recording, spec, session.summary)
    )
    print(f"[native] {published}", flush=True)
    recording.flush()
    if args.sink == "save":
        print(f"[native] wrote {args.out}", flush=True)
    print("[native] DONE", flush=True)


if __name__ == "__main__":
    main(tyro.cli(NativeArgs))