clip.py

#!/usr/bin/env python3

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Sequence

from core.rf_detr_runtime import ensure_python_nvidia_libs_preferred, sync_python_nvidia_runtime_libs_to_system

ensure_python_nvidia_libs_preferred()
sync_python_nvidia_runtime_libs_to_system()

from core.clip_orchestrator import (
    ClipRequest,
    RenderType,
    is_openpilot_render_type,
    is_smear_render_type,
    run_clip,
    supports_driver_face_anonymization,
)
from core.driver_face_swap import (
    canonical_driver_face_profile,
    default_driver_face_donor_bank_dir,
    default_driver_face_source_image,
    default_facefusion_model,
    default_facefusion_root,
)
from core.forward_upon_wide import parse_forward_upon_wide_h
from core.openpilot_bootstrap import bootstrap_openpilot, ensure_openpilot_checkout
from core.openpilot_config import default_local_openpilot_root, default_openpilot_branch, default_openpilot_repo_url
from core.ui_layouts import UI_ALT_VARIANTS


DEMO_ROUTE = "5beb9b58bd12b691|0000010a--a51155e496"
DEMO_START_SECONDS = 90
DEMO_LENGTH_SECONDS = 15
RENDER_TYPES: tuple[RenderType, ...] = (
    "ui",
    "ui-alt",
    "driver-debug",
    "forward",
    "wide",
    "driver",
    "360",
    "360-ui",
    "forward_upon_wide",
    "360_forward_upon_wide",
)
PASSENGER_REDACTION_STYLE_CHOICES = ["blur", "silhouette", "black_silhouette", "ir_tint"]


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Primary local CLI for openpilot replay clipping. Use this for cheap local validation before GCE."
    )
    parser.add_argument("render_type", choices=RENDER_TYPES)
    parser.add_argument("route", nargs="?", help='Comma Connect URL or route id (e.g. "dongle|route")')
    parser.add_argument("--demo", action="store_true", help="Use a known public demo route")
    parser.add_argument("-s", "--start-seconds", type=int, default=None)
    parser.add_argument("-l", "--length-seconds", type=int, default=None)
    parser.add_argument("--smear-seconds", type=int, default=3)
    parser.add_argument("-j", "--jwt-token", default="")
    parser.add_argument("-o", "--output", default="./shared/local-clip.mp4")
    parser.add_argument("--openpilot-dir", default=default_local_openpilot_root())
    parser.add_argument("--openpilot-branch", default=default_openpilot_branch())
    parser.add_argument("--openpilot-repo-url", default=default_openpilot_repo_url())
    parser.add_argument("-m", "--file-size-mb", type=int, default=9)
    parser.add_argument("--file-format", choices=["auto", "h264", "hevc"], default="auto")
    parser.add_argument("--forward-upon-wide-h", type=parse_forward_upon_wide_h, default="auto", help=argparse.SUPPRESS)
    parser.add_argument("--qcam", action="store_true")
    parser.add_argument(
        "--ui-alt-variant",
        choices=UI_ALT_VARIANTS,
        default=None,
        help="Alternate UI composition for `ui-alt`. `device` keeps a single main camera view with telemetry, while the stacked modes require wide video.",
    )
    parser.add_argument("--windowed", action="store_true")
    parser.add_argument("--skip-openpilot-update", action="store_true")
    parser.add_argument("--skip-openpilot-bootstrap", action="store_true")
    parser.add_argument("--data-root", default="./shared/data_dir")
    parser.add_argument("--data-dir", default="", help="Explicit data dir. If unset, uses --data-root/<dongle-id>.")
    parser.add_argument("--skip-download", action="store_true", help="Reuse already-downloaded route data.")
    parser.add_argument("--accel", choices=["auto", "cpu", "videotoolbox", "nvidia"], default="auto")
    parser.add_argument(
        "--driver-face-anonymization",
        choices=["none", "facefusion"],
        default="none",
        help="Optionally replace the face in the backing driver video before rendering a driver-backed clip such as `driver`, `driver-debug`, or the 360 modes.",
    )
    parser.add_argument(
        "--driver-face-profile",
        choices=[
            "driver_unchanged_passenger_hidden",
            "driver_unchanged_passenger_face_swap",
            "driver_face_swap_passenger_unchanged",
            "driver_face_swap_passenger_hidden",
            "driver_face_swap_passenger_face_swap",
        ],
        type=canonical_driver_face_profile,
        default="driver_face_swap_passenger_face_swap",
        help="Seat strategy when driver face anonymization is enabled.",
    )
    parser.add_argument(
        "--passenger-redaction-style",
        choices=PASSENGER_REDACTION_STYLE_CHOICES,
        default="blur",
        help="How to hide the passenger when the selected anonymization profile uses passenger hidden mode.",
    )
    parser.add_argument(
        "--driver-face-source-image",
        default=default_driver_face_source_image(),
        help="Donor/source portrait for driver face anonymization.",
    )
    parser.add_argument(
        "--driver-face-selection",
        choices=["manual", "auto_best_match"],
        default="manual",
        help="Choose a specific donor image manually or auto-pick the best same-tone donor from the donor bank.",
    )
    parser.add_argument(
        "--driver-face-donor-bank-dir",
        default=default_driver_face_donor_bank_dir(),
        help="Directory of donor portraits used when --driver-face-selection auto_best_match is enabled.",
    )
    parser.add_argument(
        "--driver-face-preset",
        choices=["fast", "quality"],
        default="fast",
        help="FaceFusion preset for driver face anonymization.",
    )
    parser.add_argument(
        "--facefusion-root",
        default=default_facefusion_root(),
        help="Path to the local FaceFusion checkout used for driver face anonymization.",
    )
    parser.add_argument(
        "--facefusion-model",
        default=default_facefusion_model(),
        help="FaceFusion model name used for driver face anonymization.",
    )
    return parser


def _resolve_route_and_timing(args: argparse.Namespace) -> tuple[str, int, int]:
    if args.demo:
        start_seconds = DEMO_START_SECONDS if args.start_seconds is None else args.start_seconds
        length_seconds = DEMO_LENGTH_SECONDS if args.length_seconds is None else args.length_seconds
        return DEMO_ROUTE, start_seconds, length_seconds
    if not args.route:
        raise SystemExit("route is required unless --demo is used")
    start_seconds = 50 if args.start_seconds is None else args.start_seconds
    length_seconds = 20 if args.length_seconds is None else args.length_seconds
    return args.route, start_seconds, length_seconds


def _prepare_openpilot_if_needed(args: argparse.Namespace) -> str:
    openpilot_path = Path(args.openpilot_dir).expanduser().resolve()
    openpilot_dir = str(openpilot_path)
    needs_openpilot = is_openpilot_render_type(args.render_type) or (
        supports_driver_face_anonymization(args.render_type) and args.driver_face_anonymization != "none"
    )
    if not needs_openpilot:
        return openpilot_dir
    if args.skip_openpilot_update and not openpilot_path.exists():
        raise SystemExit(
            f"Openpilot checkout not found at {openpilot_dir}. Remove --skip-openpilot-update or point --openpilot-dir at an existing checkout."
        )
    if not args.skip_openpilot_update:
        ensure_openpilot_checkout(
            openpilot_path,
            branch=args.openpilot_branch,
            repo_url=args.openpilot_repo_url,
        )
    if not args.skip_openpilot_bootstrap:
        bootstrap_openpilot(openpilot_path)
    elif not (openpilot_path / ".venv/bin/python").exists():
        raise SystemExit(
            f"Openpilot is not bootstrapped at {openpilot_dir}. Remove --skip-openpilot-bootstrap or run bootstrap first."
        )
    return openpilot_dir


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)
    if args.ui_alt_variant is not None and args.render_type != "ui-alt":
        raise SystemExit("--ui-alt-variant is only supported with the `ui-alt` render type")

    route, start_seconds, length_seconds = _resolve_route_and_timing(args)
    openpilot_dir = _prepare_openpilot_if_needed(args)

    try:
        result = run_clip(
            ClipRequest(
                render_type=args.render_type,
                route_or_url=route,
                start_seconds=start_seconds,
                length_seconds=length_seconds,
                target_mb=args.file_size_mb,
                ui_alt_variant=args.ui_alt_variant,
                file_format=args.file_format,
                output_path=args.output,
                smear_seconds=args.smear_seconds if is_smear_render_type(args.render_type) else 0,
                jwt_token=args.jwt_token or None,
                forward_upon_wide_h=args.forward_upon_wide_h,
                explicit_data_dir=args.data_dir or None,
                data_root=args.data_root,
                execution_context="local",
                minimum_length_seconds=1,
                maximum_length_seconds=300,
                local_acceleration=args.accel,
                openpilot_dir=openpilot_dir,
                qcam=args.qcam,
                headless=not args.windowed,
                skip_download=args.skip_download,
                driver_face_anonymization=args.driver_face_anonymization,
                driver_face_profile=args.driver_face_profile,
                passenger_redaction_style=args.passenger_redaction_style,
                driver_face_source_image=args.driver_face_source_image,
                driver_face_preset=args.driver_face_preset,
                facefusion_root=args.facefusion_root,
                facefusion_model=args.facefusion_model,
                driver_face_selection=args.driver_face_selection,
                driver_face_donor_bank_dir=args.driver_face_donor_bank_dir,
            )
        )
    except ModuleNotFoundError as error:
        raise SystemExit(
            f"Missing local dependency: {error.name}. Run `uv sync` and then use `uv run python clip.py ...`."
        ) from error

    print(f"Wrote clip: {result.output_path}")
    if result.acceleration:
        print(f"Acceleration: {result.acceleration}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())