orbis2_app_engine.py

"""
Persistent, in-process rollout engine for app2.py.

app2.py used to shell out to `python evaluate/rollout_demo_v2.py` on every click of
"Generate rollouts", which re-read the config, re-instantiated the model, re-loaded
the ~20GB checkpoint from disk, and re-established torch.compile from scratch on
every single request. This module keeps one loaded (and, after the first request,
compiled) model alive for the life of the process instead: `load()` does the
one-time CPU setup, `ensure_ready()` moves to GPU and compiles exactly once, and
`roll_out()` is the cheap per-request call.

This relies on the Hugging Face ZeroGPU worker that runs `@spaces.GPU`-decorated
calls being reused across requests rather than restarted per call -- if the worker
is ever recycled (idle timeout, scale-to-zero), the next request after that pays
the load/compile cost again, same as today's cold start, just no longer on *every*
request.

`orbis2/evaluate/rollout_demo_v2.py` is vendored from another repo and can be
overwritten by a future sync, so nothing app-critical lives there: this module only
imports its small, generic, model-agnostic helpers (context sampling, trajectory
math, config resolution) and owns model construction / the rollout body itself.
"""

import contextlib
import logging
import os
import sys
import threading
from pathlib import Path

import torch
from omegaconf import OmegaConf
from pytorch_lightning import seed_everything
from torchvision.utils import save_image

_REPO_DIR = Path(__file__).parent / "orbis2"
if str(_REPO_DIR) not in sys.path:
    sys.path.insert(0, str(_REPO_DIR))

from evaluate.rollout_demo_v2 import (  # noqa: E402
    STEERING_FORMAT,
    _L1L2FrameIndexer,
    get_rollout_future_frame_count,
    load_l1_l2_context,
    load_trajectory_points,
    make_unconditional_steering,
    maybe_apply_condition_preprocessor_scales,
    maybe_apply_l2_nfe,
    overlay_trajectory_on_images,
    require_l1l2_model,
    resample_trajectory_by_arclength,
    resolve_l2_frame_rate,
    resolve_video_backend,
    trajectory_to_speed_yawrate,
)
from modules import fm_samplers  # noqa: E402
from util import instantiate_from_config  # noqa: E402

logger = logging.getLogger(__name__)


def _describe_sampler(sampler, nfe, eta):
    """Human-readable solver name + init params (e.g. FlowMatchingSamplerHeunPlusEuler
    with its timescale/integration_t_eps/step_schedule/...), plus this call's NFE/eta --
    the latter two are sample()-time args, not stored on the sampler itself."""
    if sampler is None:
        return "none"
    params = {k: v for k, v in vars(sampler).items() if k != "module"}
    return f"{type(sampler).__name__}(NFE={nfe}, eta={eta}, {params})"


@contextlib.contextmanager
def _l1_nfe_override(model, num_steps):
    """Let the "L1 sampler steps" UI control actually change L1's step count.

    FlowMatchingSamplerHeunPlusEuler._resolve_schedule ignores NFE whenever the
    sampler has an explicit `step_schedule` configured (see fm_samplers.py) --
    step_schedule always wins, so a fixed-schedule L1 config (like this app's
    config_distill.yaml) would otherwise make num_steps/l1_steps a no-op. Since
    that's vendored code we don't touch, this instead mutates the already-
    constructed sampler's `step_schedule` attribute directly -- same pattern as
    maybe_apply_l2_nfe's post-init override of l2_pred_NFE, which relies on the
    same fact: these are read at sample time, not at construction.

    Uses the config's tuned step_schedule as-is whenever the requested
    num_steps matches its own step count -- that's the "default config steps"
    case. Only when a genuinely different l1_nfe is requested does this clear
    step_schedule for the call, so the sampler falls back to its own NFE-driven
    schedule (NFE-1 uniform Heun steps + 1 uniform Euler step). Restores the
    original schedule afterward either way.
    """
    sampler = getattr(model, "sampler", None)
    is_fixed_schedule = (
        isinstance(sampler, fm_samplers.FlowMatchingSamplerHeunPlusEuler)
        and sampler.step_schedule is not None
        and len(sampler.step_schedule) != int(num_steps)
    )
    if not is_fixed_schedule:
        yield
        return

    original_schedule = sampler.step_schedule
    print(f"[solver] L1 step_schedule has {len(original_schedule)} steps; requested "
          f"l1_steps={num_steps} differs -- using {int(num_steps) - 1} Heun step(s) + "
          "1 Euler step for this request instead.")
    sampler.step_schedule = None
    try:
        yield
    finally:
        sampler.step_schedule = original_schedule


@contextlib.contextmanager
def _rollout_progress_tqdm(progress_cb):
    """Temporarily wrap fm_samplers.tqdm -- used only by FlowMatchingSampler.roll_out's
    outer per-rollout-step loop (one iteration per autoregressive chunk, i.e. per
    num_gen_frames) -- so it still prints its normal console progress bar while also
    reporting each completed step to progress_cb(step, total). Real per-step progress
    without editing the vendored fm_samplers.py; restores the original tqdm on exit
    either way."""
    if progress_cb is None:
        yield
        return

    real_tqdm = fm_samplers.tqdm

    def _tqdm_with_progress(iterable, *args, **kwargs):
        total = kwargs.get("total")
        if total is None:
            try:
                total = len(iterable)
            except TypeError:
                total = None

        def _gen():
            for i, item in enumerate(real_tqdm(iterable, *args, **kwargs)):
                yield item
                progress_cb(i + 1, total)

        return _gen()

    fm_samplers.tqdm = _tqdm_with_progress
    try:
        yield
    finally:
        fm_samplers.tqdm = real_tqdm


class RolloutEngine:
    """Owns one loaded L1-L2 hierarchical model and serves repeated rollout requests."""

    def __init__(self, exp_dir, config_name, ckpt_name, evaluate_ema=True):
        self.exp_dir = str(exp_dir)
        self.config_path = os.path.join(self.exp_dir, config_name)
        self.ckpt_path = os.path.join(self.exp_dir, ckpt_name)
        self.evaluate_ema = evaluate_ema

        self.model = None
        self._ready = False
        self._lock = threading.Lock()
        self._compile_artifacts = None
        self._artifacts_saved = False

    def load(self):
        """CPU-only one-time setup: read the config, build the model, load the
        checkpoint. Call before any GPU is attached (i.e. outside a @spaces.GPU
        function) -- mirrors the model-construction prefix of
        rollout_demo_v2.generate_images()."""
        config = OmegaConf.load(self.config_path)
        model = instantiate_from_config(config.model)

        ckpt_result = model.load_state_dict(
            torch.load(self.ckpt_path, map_location="cpu")["state_dict"], strict=False
        )
        exempt_prefixes = tuple(getattr(model, "checkpoint_exempt_key_prefixes", ()))
        unexpected_missing = [k for k in ckpt_result.missing_keys if not k.startswith(exempt_prefixes)]
        assert unexpected_missing == [], unexpected_missing
        model.eval()

        require_l1l2_model(model)

        self.model = model
        logger.info(f"Loaded Orbis 2 model from {self.ckpt_path} (CPU, not yet compiled)")
        return self

    def min_context_end_frame(self, l1_frame_rate):
        """Earliest frame index (0-based, in a video encoded at exactly `l1_frame_rate`
        fps) that has enough L1+L2 lookback to be a valid end of the context window.
        CPU-only (only reads model metadata) -- callable right after load(), before
        ensure_ready() puts anything on the GPU.

        `frame_interval=1` below assumes the video's native fps equals l1_frame_rate
        exactly, which is app2.py's re-encode invariant (CONTEXT_FPS == L1_FRAME_RATE);
        under that assumption l1_span == l1_context_frames, matching load_l1_l2_context's
        own math for the same case.
        """
        model = self.model
        l1_context_frames = int(model.vit.num_context_frames)
        l2_context_frames = int(model.condition_preprocessor.num_context_frames)
        l2_frame_rate = resolve_l2_frame_rate(model)

        indexer = _L1L2FrameIndexer(
            frame_interval=1,
            stored_data_frame_rate=l1_frame_rate,
            num_l2_context=l2_context_frames,
            l2_frame_rate=l2_frame_rate,
            l1_context_frames=l1_context_frames,
        )
        return indexer.get_required_l1_start_offset() + l1_context_frames - 1

    @property
    def frames_per_rollout_step(self):
        """Number of output frames each autoregressive rollout step yields
        (model.num_pred_frames). CPU-only (model metadata), callable right after
        load(). Lets callers convert the "rollout steps" the model actually takes
        into seconds of generated video (frames_per_rollout_step / l1_frame_rate)."""
        return int(self.model.num_pred_frames)

    def ensure_ready(self, device="cuda", compile=True, compile_mode="reduce-overhead",
                      compile_artifacts=None, speed_scale=1.0, yaw_rate_scale=1.0):
        """Move the model to `device` and wrap it with torch.compile, exactly once.
        Safe to call on every request: after the first call this is a no-op, because
        the model lives in this persistent process and keeps its GPU/compiled state
        between calls."""
        if self._ready:
            print("[compile] model already on-device and compiled from a prior request -- reusing it as-is.")
            return self.model

        with self._lock:
            if self._ready:
                return self.model

            model = self.model.to(device)

            if compile:
                def _maybe_compile(module, attr):
                    net = getattr(module, attr, None)
                    if net is not None:
                        print(f"[compile] wrapping {type(module).__name__}.{attr} in torch.compile(mode={compile_mode!r})")
                        setattr(module, attr, torch.compile(net, mode=compile_mode))
                        logger.info(f"Compiled {type(module).__name__}.{attr} with mode={compile_mode!r}")

                _maybe_compile(model, "ema_vit" if self.evaluate_ema else "vit")
                l2_predictor = getattr(getattr(model, "condition_preprocessor", None), "l2_predictor", None)
                if l2_predictor is not None:
                    _maybe_compile(l2_predictor, "ema_vit")

                if compile_artifacts:
                    if os.path.exists(compile_artifacts):
                        with open(compile_artifacts, "rb") as f:
                            torch.compiler.load_cache_artifacts(f.read())
                        print(f"[compile] loaded cached compile artifacts from {compile_artifacts!r} "
                              "-- first call reuses these instead of tracing from scratch.")
                        logger.info(f"Loaded compile artifacts from {compile_artifacts!r}")
                    else:
                        print(f"[compile] no compile artifacts found at {compile_artifacts!r} "
                              "-- compiling from scratch; will save artifacts after the first rollout.")
                        logger.info(
                            f"Compile artifacts not found at {compile_artifacts!r}; "
                            "will save after the first rollout."
                        )
                else:
                    print("[compile] no compile_artifacts path given -- compiling from scratch, nothing cached to disk.")
                self._compile_artifacts = compile_artifacts

            maybe_apply_condition_preprocessor_scales(model, speed_scale, yaw_rate_scale)

            self.model = model
            self._ready = True
            logger.info(f"Model moved to {device} and ready -- subsequent requests reuse this instance.")

        return self.model

    @torch.no_grad()
    def roll_out(self, video_path, output_dir, num_gen_frames, num_steps, l1_frame_rate,
                 height, width, seed, num_videos=1, l2_nfe=None, eta=0.0,
                 trajectory_file=None, vis_mode="none", decode_device="cpu", device="cuda",
                 context_end_frame=None, progress_cb=None):
        """One rollout request against the already-loaded, already-compiled model.
        Mirrors rollout_demo_v2.generate_images() from context loading onward -- the
        model-construction prefix runs once in load()/ensure_ready(), not here."""
        if not self._ready:
            raise RuntimeError("RolloutEngine.ensure_ready() must be called before roll_out().")

        model = self.model

        if int(seed) > 0:
            torch.backends.cudnn.enable = False
            torch.backends.cudnn.deterministic = True
            seed_everything(int(seed))

        maybe_apply_l2_nfe(model, l2_nfe)

        print(f"[solver] L1: {_describe_sampler(getattr(model, 'sampler', None), num_steps, eta)}")
        l2_predictor = getattr(model.condition_preprocessor, "l2_predictor", None)
        if l2_predictor is not None:
            l2_nfe = getattr(model.condition_preprocessor, "l2_pred_NFE", None)
            print(f"[solver] L2: {_describe_sampler(getattr(l2_predictor, 'sampler', None), l2_nfe, eta)}")

        height, width = int(height), int(width)
        backend = resolve_video_backend()
        l2_frame_rate = resolve_l2_frame_rate(model)

        l1_context_frames = int(model.vit.num_context_frames)
        l2_context_frames = int(model.condition_preprocessor.num_context_frames)

        start_frame = None if context_end_frame is None else int(context_end_frame) - l1_context_frames + 1

        print(f"[load] reading context from video: L1 {l1_context_frames} frames @ "
              f"{l1_frame_rate}fps, L2 {l2_context_frames} frames @ {l2_frame_rate}fps…")
        l1_tensor, l2_tensor = load_l1_l2_context(
            video_path=video_path,
            start_frame=start_frame,
            l1_frame_rate=l1_frame_rate,
            l2_frame_rate=l2_frame_rate,
            l1_context_frames=l1_context_frames,
            l2_context_frames=l2_context_frames,
            height=height,
            width=width,
            backend=backend,
            device=device,
        )
        print(f"[load] context ready: L1 tensor {tuple(l1_tensor.shape)}, "
              f"L2 tensor {tuple(l2_tensor.shape)}")

        num_future_frames = get_rollout_future_frame_count(model, num_gen_frames)
        frame_rate = torch.tensor(float(l1_frame_rate), device=device)
        data_batch = {"images": l1_tensor, "l2_context": l2_tensor, "frame_rate": frame_rate}

        get_required_steps = getattr(model.condition_preprocessor, "get_required_rollout_odometry_steps", None)
        min_odo_steps = None
        if callable(get_required_steps):
            min_odo_steps = get_required_steps(
                validation_params=None,
                num_condition_frames=l1_context_frames,
                num_gen_frames=num_future_frames,
                rollout_steps=num_gen_frames,
            )

        print(f"[load] preparing {'steering trajectory' if trajectory_file is not None else 'unconditional'} "
              "odometry conditioning…")
        if trajectory_file is not None:
            if min_odo_steps is None:
                raise ValueError(
                    "trajectory_file requires the model's condition_preprocessor to report a "
                    "required odometry length (get_required_rollout_odometry_steps returned None)."
                )
            odometry_steps_per_image_frame = getattr(model.condition_preprocessor, "odometry_steps_per_image_frame", 1)
            dt = 1.0 / (l1_frame_rate * odometry_steps_per_image_frame)
            traj_xy = load_trajectory_points(trajectory_file)
            traj_xy = resample_trajectory_by_arclength(traj_xy, min_odo_steps + 1)
            speed_yawrate = trajectory_to_speed_yawrate(traj_xy, dt)
            data_batch["steering"] = torch.as_tensor(
                speed_yawrate, dtype=l1_tensor.dtype, device=device
            ).unsqueeze(0)
        else:
            data_batch["steering"] = make_unconditional_steering(min_odo_steps, dtype=l1_tensor.dtype, device=device)
        data_batch["steering_format"] = STEERING_FORMAT

        num_videos = max(1, int(num_videos))
        if num_videos > 1:
            def _tile_batch(t):
                return t.repeat(num_videos, *([1] * (t.dim() - 1)))

            l1_tensor = _tile_batch(l1_tensor)
            l2_tensor = _tile_batch(l2_tensor)
            data_batch["images"] = l1_tensor
            data_batch["l2_context"] = l2_tensor
            data_batch["steering"] = _tile_batch(data_batch["steering"])

        condition_kwargs = model.condition_preprocessor.get_condition_kwargs_from_batch(data_batch, split="rollout")

        # The next call encodes the L1 context into latents and runs L2's own upfront
        # forecast (L1's autoregressive loop conditions on it, so it has to happen
        # before L1 can take a single step) -- neither has a progress hook, and on a
        # fresh torch/CUDA/GPU combo this is also where a from-scratch torch.compile
        # can silently eat a minute or more. No further output until this returns.
        print("[load] encoding context + L2 forecast (first call on this GPU/torch/CUDA "
              "combo may compile from scratch here -- can take a while, silently)…")
        autocast_enabled = device.startswith("cuda")
        with torch.autocast(dtype=torch.float16, device_type="cuda", enabled=autocast_enabled), \
             _rollout_progress_tqdm(progress_cb), \
             _l1_nfe_override(model, num_steps):
            _latents, gen_frames = model.roll_out(
                x_0={"images": l1_tensor},
                num_gen_frames=num_gen_frames,
                latent_input=False,
                NFE=num_steps,
                eta=eta,
                sample_with_ema=self.evaluate_ema,
                num_samples=l1_tensor.size(0),
                frame_rate=frame_rate.reshape(1).repeat(l1_tensor.size(0)),
                condition_kwargs=condition_kwargs,
                decode_device=decode_device,
                num_condition_frames=l1_tensor.size(1),
            )

        if vis_mode in {"trajectory", "trajectory_ego"}:
            overlay_trajectory = model.condition_preprocessor.get_rollout_visualization_trajectory(
                condition_kwargs=model.condition_preprocessor.get_condition_kwargs_from_batch(data_batch, split="rollout"),
                num_condition_frames=l1_context_frames,
                num_gen_steps=num_gen_frames,
                num_pred_frames=model.num_pred_frames,
            )
            if overlay_trajectory is not None:
                gen_frames = overlay_trajectory_on_images(gen_frames, overlay_trajectory, mode=vis_mode)

        del _latents, condition_kwargs, l1_tensor, l2_tensor

        os.makedirs(output_dir, exist_ok=True)
        num_out = gen_frames.shape[0]
        num_frames = gen_frames.shape[1]
        for b in range(num_out):
            seq_dir = os.path.join(output_dir, "fake_images", f"sequence_{b:04d}")
            os.makedirs(seq_dir, exist_ok=True)
            for f in range(num_frames):
                save_image(
                    (gen_frames[b, f] + 1.0) / 2.0,
                    os.path.join(seq_dir, f"frame_{f:04d}.jpg"),
                )

        if device.startswith("cuda"):
            logger.info(f"Max memory: {torch.cuda.max_memory_allocated() / 1024**3:.02f} GB")

        if self._compile_artifacts and not self._artifacts_saved and not os.path.exists(self._compile_artifacts):
            artifacts = torch.compiler.save_cache_artifacts()
            if artifacts is not None:
                with open(self._compile_artifacts, "wb") as f:
                    f.write(artifacts[0])
                logger.info(f"Saved compile artifacts to {self._compile_artifacts!r}")
            self._artifacts_saved = True


_engine = None
_engine_lock = threading.Lock()


def get_engine(exp_dir, config_name, ckpt_name, evaluate_ema=True):
    """Return the process-wide RolloutEngine, constructing and CPU-loading it on
    first call. Subsequent calls (with the same or different args -- args are only
    used for the first, initializing call) just return the cached instance."""
    global _engine
    if _engine is None:
        with _engine_lock:
            if _engine is None:
                _engine = RolloutEngine(exp_dir, config_name, ckpt_name, evaluate_ema=evaluate_ema).load()
    return _engine