Blame for summary_server.py

"""Summaryformer streaming server: one trunk, one stage, full-res 24x42.

Inputs: ONE DIRECTORY of checkpoints + ONE SAE checkpoint. Every `*.pt` in `checkpoints_dir` is
servable, named per `streaming_server_utils.ckpt_name_schema` (`gta_model.pt`; a name may also
carry a recipe token, `<labels>-<K><r|nr>.pt`). The stem is the model_id the Next.js client lists
(GET /models) and sends back (POST /z0 `model_id`). A recipe token in the name is checked against
the saved config; a disagreement is refused at startup.

One generator is built and compiled once. Architecture (dim / depth / X/Y/W) must
agree across the directory; weights are swapped in place. Anything the checkpoint does not state
is refused at startup, not defaulted.

The served contract is the ckpt's config (`format` = CKPT_FORMAT):
    timesteps, timestep_shift, noise_schedule, shifted_cosine_ratio,
    cache_mode "recache" | "last_step_writes", recency, sink_gap_max (one int),
    summary_tokens X / raw_window_frames Y / summary_window_frames W, focus_region,
    profile, jump_sprint, dim / ffn_dim / num_heads / num_layers.
A missing key is refused with a message naming it.

Cache: "recache" (KR: K passes + one t=0 writer on x0) or "last_step_writes"
(KNR: the K-th pass's record is committed; there is no tail). z0 is always a clean t=0 warm.
Recency (a per-key bias on the SUM region of the one buffer; Z0 and RAW rows stay 0) and the
sink chase (`SinkChaser`; one Z0 slot, one gap) come from the ckpt. A None gap is refused (an
unbounded stream needs the sink chased). Focus: when the ckpt says `focus_region` the z0 warm gets
a centred diamond at SERVE_FOCUS_COVERAGE; body frames never see it.

Memory: one K/V buffer per layer (Z0 + RAW + SUM regions); 4 forwards per frame at 1104 tokens
under 3R (3 NFE + recache).

Split: this file owns HTTP/WS, the owner lock, the upload encode, H.264 and pacing.
`summary_streaming_inference` owns everything that touches the model (one buffer, the denoising
loop, SAE decode, gaussianization).

Weight dtype (`generator_dtype_mode`):
    "bf16_with_fp32_islands"    bulk bf16; the conditioning MLPs the model runs under
                                autocast(enabled=False) (time / aspect / action / fps) stay fp32.
    "fp32_master_bf16_autocast" all fp32, autocast bf16 on the fly (slower).
The session always computes under bf16 autocast with fp32 latents.

Client contract:
    GET  /models                 {"server_kind": "unified", "models": [id...], "default_model_id",
                                  "active_model_id", "details": {id: {recipe, steps, cache_mode,
                                  timesteps, summary_tokens (X), raw_window_frames (Y),
                                  summary_window_frames (W), sink_gap_max, recency, focus, ...}},
                                  "nfe": {override, mode, modes, ckpt_steps, active_steps, timesteps, max},
                                  "sphere": {mode: off|decode|all, modes}}
                                 `nfe` is the body NFE override state (slider).
    POST /z0/{user_id}           form: z0_image, model_id ("" = keep the active), nfe
                                 (""/0 = ckpt body steps, N = override for the game), nfe_mode. A
                                 different id SWITCHES the served checkpoint (weights + session;
                                 + compile warmup when the session geometry changes; tens of
                                 seconds; 423 if the GPU is busy; 400 if unknown).
    POST /switch                 form: model_id. Same switch as /z0, no encode. 423 if the GPU is busy.
    GET  /create-game/{user_id}
    WS   /ws/{user_id}           binary H.264 Annex-B 384x672 (0x01 key / 0x00 delta), text in =
                                 controls, text out =
                                 {"type":"log","tag","msg","ts"} and, ws_session_limit_s after the
                                 first frame, {"type":"session_over","message","limit_s"} followed by
                                 close 1000 (reason = message). Full-res decode, no aux stream.
    GET  /status                 {ok, busy, ws_active, server_id, tickets_required, session_limit_s,
                                 active_model_id, active_profile}. No ticket needed.
Busy: HTTP 423 / WS 1013. One player, one CUDA thread.
Profiles (`ckpt_profile`): each ckpt's config decides what the player gets. A ckpt with
`jump_sprint` is served with space -> jump / shift -> sprint LIVE in the 5-d action. Everything
else has jump/sprint forced 0.
GET /models `details[id].profile` = {name, jump_sprint, controls}; the client shows the extra keys from it.
Tickets (optional, off by default): with LUCID_DEMO_TICKET_SECRET + LUCID_DEMO_SERVER_ID set, every
endpoint above except /models and /status needs a `uid.server_id.exp.hmac` ticket (`ticket` form field on
POST, `?ticket=` on GET /create-game and WS /ws) or answers 401 / WS close 1008. See _ticket_check.
"""
from __future__ import annotations

import asyncio
import gc
import glob
import hashlib
import hmac
import io
import json
import os
import queue
import re
import sys
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from dataclasses import dataclass, replace
from fractions import Fraction

import numpy as np
import torch
from fastapi import FastAPI, File, Form, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import PlainTextResponse
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from inference_common import num_layers_from_ckpt_config, resolve_sampling, wrapper_ctor_kwargs_from_ckpt_config
from sae import Decoder, SAEConfig, ShakeAE, _apply_channels_last, _assert_channels_last
from streaming_server_utils.ckpt_name_schema import CkptName, check_name_against_config, parse_ckpt_name, trained_recency
from streaming_server_utils.controls import nextjs_client_actions_to_world_model_5_dim_actions as client_actions_to_5d
from summary_streaming_inference import (CACHE_MODES, FOCUS_DEFAULT, SPHERE_MODES, StageSpec,
                                         build_session_from_runtime)
from summaryformer import (Z0_FOCUS_REGION_EMBED_ALLOWED_MISSING, Z0_FOCUS_REGION_EMBED_NAME,
                                set_streaming_attention_backend)
from summaryformer_utils.recency_bias import RecencyBiasConfig
from summaryformer_utils.streaming_sink import (SinkChaser, extend_rope_table_for_streaming, sink_gaps_describe,
                                                     sink_gaps_from_config)
from summaryformer_utils.z0_focus_mask import centered_focus_region
from summaryformer_wrapper import SummaryformerWrapper
# Fraction of the first frame covered by the centred diamond focus region.
SERVE_FOCUS_COVERAGE = 0.20
# Action value that means "no action for this frame"; the model skips its action input for such frames.
ACTION_SENTINEL = -69.0

RGB_HW = (384, 672)                                             # H.264 size; SAE /16 -> 24x42
SAE_SPATIAL = 16
LATENT_HW = (RGB_HW[0] // SAE_SPATIAL, RGB_HW[1] // SAE_SPATIAL)  # (24, 42)
DTYPE_MODES = ("bf16_with_fp32_islands", "fp32_master_bf16_autocast")
FP8_MODES = ("off", "row", "tensor")   # CONFIG["fp8_linears"]; the module lives below apply_dtype_mode
COMPILE_GENERATOR_MODES = ("default", "max-autotune-no-cudagraphs", "max-autotune+cudagraphs")
SAE_DECODER_DTYPES = {"bf16": torch.bfloat16, "fp32": torch.float32}
# Conditioning MLPs the model runs under autocast(enabled=False); stay fp32 in the islands dtype mode.
_COND_ISLANDS = ("timestep_mlp", "timestep_to_film", "aspect_ratio_proj", "aspect_ratio_proj_z0",
                 "action_proj", "fps_embedder_mlp", "fps_embedder_outs")

# The checkpoint config this server reads. A checkpoint is {"model": state_dict, "config": {...}, "step": int}.
CKPT_FORMAT = "summaryformer-inference-1"
# Model facts every ckpt config must carry. Nothing here has a server-side default.
CKPT_MODEL_KEYS = (
    "format", "timesteps", "timestep_shift", "noise_schedule", "shifted_cosine_ratio",
    "cache_mode", "recency", "sink_gap_max", "summary_tokens", "raw_window_frames",
    "summary_window_frames", "focus_region", "profile", "jump_sprint",
    "dim", "ffn_dim", "num_heads", "num_layers",
)
# The subset that is the ARCHITECTURE: one generator is built and compiled once, weights are swapped
# in place, so every checkpoint in the directory must agree on these.
CKPT_ARCH_KEYS = (
    "dim", "ffn_dim", "num_heads", "num_layers",
    "summary_tokens", "raw_window_frames", "summary_window_frames",
)

CONFIG = {
    # Every *.pt here is a servable checkpoint, named per streaming_server_utils/ckpt_name_schema
    # (`gta_model.pt`). Same architecture, anything else may differ. The stem is the model_id.
    "checkpoints_dir": "./serve_ckpts",
    # Which model_id is loaded at startup (its compile warmup runs before the port opens). "latest" =
    # the ckpt with the highest saved step (ties: last in sorted order); None = the first id in sorted
    # order; else an id. An id that is not in checkpoints_dir fails at startup with the list of what is
    # there. Env LUCID_DEMO_DEFAULT_MODEL=<model_id> overrides.
    "default_model_id": (os.environ.get("LUCID_DEMO_DEFAULT_MODEL") or "").strip() or None,
    # The autoencoder: {"enc": state_dict, "dec": state_dict}.
    "sae_ckpt": "./sae_compact.pt",
    # Diag-gauss stats: per-channel mean / std of the autoencoder's latent, one file for z0 and one for
    # z_rest (full-res 24x42; expect_scale="fine"). They are the space the model was trained in.
    "sae_diag_gauss_stats_path": "./latent_stats_z0.pt",
    "sae_diag_gauss_rest_stats_path": "./latent_stats_z_rest.pt",
    # "bf16_with_fp32_islands" = bulk weights bf16, the cond MLPs stay fp32.
    # "fp32_master_bf16_autocast" = fp32 weights under bf16 autocast (slower).
    "generator_dtype_mode": "bf16_with_fp32_islands",
    # Absolute ids keep growing, so the model's 1024-row RoPE table is extended to `rope_table_len`
    # at load (validated over the largest W + sink_gap_max in the directory via
    # SinkChaser.max_relative_frames); the session HARD-STOPS the game at that row.
    "rope_table_len": 4096,
    # Memory window W in frames. W is cache capacity only (no parameter is sized by it), so every ckpt here
    # is served at THIS W: scan_checkpoints_dir rewrites saved["summary_window_frames"] to it (logged when it
    # differs) before anything reads the config -- model, KV cache, sink chaser, recency floor check, warmup.
    # None = serve each ckpt at its own saved W.
    "serve_window_frames": 16,
    # Recency bias (summaryformer_utils/recency_bias.py). A ckpt whose config carries its own bias is served
    # with that one. A ckpt without is served with no bias unless its file name carries the `rec` token
    # (ckpt_name_schema), in which case THIS bias is applied. decay_end must be <= W+1 of the served ckpt (the
    # cache refuses more): 17 at W=16. Flat for 5 frames, then a smooth ramp to the floor at eviction.
    "serve_recency": {"flat_until": 5, "decay_end": 17, "beta": 6.0, "gamma": 2.0},
    # False = no recency bias on any session, whatever the ckpt config or the file name says.
    # True = the two rules above (ckpt's own bias; `rec` token -> serve_recency; else off).
    "recency_enabled": False,
    # FORWARD RE-PREFILL (SummaryStreamSession.reprefill): every N body frames the KV buffer is rebuilt
    # from z0 + the last W clean x0 (sequential t=0 writes, ids renumbered 0..n) instead of carrying the
    # in-place-rotated sink and stale SUM rows forward. Costs ~(W+1)/(K+1) frame-times per rebuild on the
    # game thread (eager + compiled body passes, no recompile). 0 = off. Form field `reprefill_every` on /z0
    # overrides per game; /models reports it. Needs cache_mode == "recache" (the session refuses otherwise).
    "reprefill_every": 0,
    "noise_seed": None,            # int for reproducible demos
    # fps conditioning: the frame rate the model is told the latent stream runs at. Start value; the client
    # overrides it live with `fps_embedding` in the controls payload (number in [1, 120], or null / "none"
    # for no fps conditioning).
    "fps_embedding": 12,
    # torch.compile: ONE compiled forward_body from body frame 1 (device-end form: the frame id is
    # a persistent 0-d int64 device tensor the session increments in place). z0 warm is eager.
    # STATIC compile (dynamic=False): nothing about the frame reaches the graph as a Python int.
    "use_compiled_generator_in_steady_state": True,
    # "default"                    plain inductor: compiles in seconds, no autotune, no cudagraphs
    # "max-autotune-no-cudagraphs"  proven path (minutes of autotune); launches from C++
    # "max-autotune+cudagraphs"     autotune + CUDA graphs via options. Needs >= 3 compiled frames
    #                               per graph before it REPLAYS (warmup run, recording run, then
    #                               replay) -- the warmup below covers that.
    "compile_generator_mode": "max-autotune+cudagraphs",
    "compile_dynamic": False,
    "compile_sae_decode": True,
    "compile_sae_decode_mode": "max-autotune-no-cudagraphs",
    # SAE decoder weights: "bf16" (default) or "fp32" (A/B only).
    "sae_decoder_dtype": "bf16",
    "use_compiled_sae_in_steady_state": False,  # False = compiled SAE decode from body frame 1
    "dynamo_recompile_limit": 64,
    # Pay the compile (autotune = minutes) at startup: a throwaway game of ring + warmup_compiled_frames
    # runs on the CUDA thread before the server accepts /ws. The session is reset afterwards.
    "warmup_compile_at_startup": True,
    "warmup_compiled_frames": 8,
    # KERNEL PROFILE of the steady state. N > 0 -> after the warmup's steady frames, N more frames run
    # under torch.profiler (CUDA + CPU activities) on the SAME hot graphs, then every CUDA kernel is bucketed
    # (flex attention / GEMM / fused elementwise / conv=SAE / memcpy / other), printed as ms-per-latent and %
    # with the top kernels, and a Chrome trace is written next to the ckpt dir (profile_<model_id>.json).
    # Env SUMMARY_PROFILE=N overrides; SUMMARY_PROFILE_EXIT=1 exits after the report (bench mode).
    # Under compile_generator_mode "max-autotune+cudagraphs" the replayed graph's kernels are still recorded
    # per kernel by CUPTI (CUDA 12); if the report shows one opaque "cudaGraphLaunch" bucket, rerun with
    # compile_generator_mode="max-autotune-no-cudagraphs" (same kernels, launched eagerly).
    "profile_steady_frames": int(os.environ.get("SUMMARY_PROFILE", "0") or 0),
    # fp8 (e4m3) GEMMs for the per-step trunk linears (visual qkv/o, visual SwiGLU proj/out) via
    # torch._scaled_mm, no extra deps. The trunk is compute-bound on small cards (4060 Ti: 1.22 TFLOP/forward,
    # ~1100 FLOP/byte vs a ~360 ridge), and Ada/Hopper fp8 tensor cores run at 2x the bf16 rate -- so this halves
    # the GEMM math; it does nothing for flex attention or the SAE. "row": per-token activation scale and
    # per-output-channel weight scale (sm89+). "tensor": one scale each (fallback if the rowwise kernel is missing).
    # The bf16 `weight` Parameter stays (ckpt switches copy_ into it); the fp8 copy is requantized after every
    # load. Inference-only approximation: eyeball the rollout before trusting it. Env SUMMARY_FP8=row|tensor|off.
    "fp8_linears": str(os.environ.get("SUMMARY_FP8", "off") or "off").lower(),
    "cuda_device": 0,              # one GPU, one player; the other GPUs on the box are untouched
    # Transport. 1 latent = 4 RGB frames at the SAE encode fps. `stream_fps` is the WIRE rate.
    # It is NOT `fps_embedding` (the FiLM value). The client may override it live.
    "encode_fps": 12, "stream_fps": 35, "stream_h264_encoder": "libx264", "stream_h264_bitrate": 4_000_000, "stream_h264_gop": 120,
    # Latency. The generator waits until at most `frames_low_watermark` RGB frames are still queued
    # for the sender, THEN reads the newest controls and makes the next latent (4 frames).
    "frames_low_watermark": 2,
    "enable_ws_log_streaming": True, "host": "0.0.0.0", "port": 3002,
    # Per-player gameplay cap. The clock starts at the FIRST frame sent (warmup / compile do not count); at the
    # limit the server sends {"type": "session_over", "message", "limit_s"} as text, then closes the socket with
    # code 1000 and the same message as the close reason (<= 123 bytes), freeing the lease for the next player.
    # 0 or None = unlimited.
    "ws_session_limit_s": 180,
    "ws_session_limit_message": "Thanks for playing! That's the limit for you.",
    # Browser origins allowed to call the HTTP endpoints (the WS is not CORS-gated).
    # Exact origins (scheme + host, no path, no trailing slash); add your own page's origin here.
    "cors_allow_origins": ["http://localhost:3000", "http://127.0.0.1:3000",
                           "https://lucidml.ai", "https://www.lucidml.ai"],
}


def _assert_config_alive(cfg: dict) -> None:
    if not str(cfg["checkpoints_dir"]).strip():
        raise RuntimeError("[summary-server] checkpoints_dir is empty")
    if cfg["default_model_id"] is not None and not str(cfg["default_model_id"]).strip():
        raise RuntimeError("[summary-server] default_model_id must be None, 'latest' or a model_id (a ckpt stem)")
    sw = cfg["serve_window_frames"]
    if sw is not None and (isinstance(sw, bool) or not isinstance(sw, int) or not 1 <= int(sw) <= 128):
        raise RuntimeError(f"[summary-server] serve_window_frames={sw!r} must be None (each ckpt's saved W) or an int in [1, 128]")
    if not isinstance(cfg.get("recency_enabled", True), bool):
        raise RuntimeError(f"[summary-server] recency_enabled must be a bool; got {cfg.get('recency_enabled')!r}")
    if RecencyBiasConfig.from_dict(cfg["serve_recency"]) is None:
        raise RuntimeError("[summary-server] serve_recency must be a recency dict (it is what a `rec`-named ckpt "
                           "is served with); to serve a ckpt without the bias, drop the token from its file name")
    if not cfg["recency_enabled"]:
        print("[summary-server] recency_enabled=False: no recency bias on any session "
              "(ckpt recency / serve_recency / `rec` file tokens ignored)", flush=True)
    if cfg["sae_decoder_dtype"] not in SAE_DECODER_DTYPES:
        raise RuntimeError(f"[summary-server] sae_decoder_dtype={cfg['sae_decoder_dtype']!r} not in {tuple(SAE_DECODER_DTYPES)}")
    if cfg["generator_dtype_mode"] not in DTYPE_MODES:
        raise RuntimeError(f"[summary-server] generator_dtype_mode={cfg['generator_dtype_mode']!r} not in {DTYPE_MODES}")
    if cfg["fp8_linears"] not in FP8_MODES:
        raise RuntimeError(f"[summary-server] fp8_linears={cfg['fp8_linears']!r} not in {FP8_MODES} (env SUMMARY_FP8)")
    if cfg["fp8_linears"] != "off" and cfg["generator_dtype_mode"] == "fp32_master_bf16_autocast":
        raise RuntimeError("[summary-server] fp8_linears needs generator_dtype_mode='bf16_with_fp32_islands'")
    if not 1.0 <= float(cfg["stream_fps"]) <= 120.0:
        raise RuntimeError(f"[summary-server] stream_fps={cfg['stream_fps']} must be in [1, 120]")
    if not 2048 <= int(cfg["rope_table_len"]) <= 5000:
        raise RuntimeError("[summary-server] rope_table_len must be in [2048, 5000]")
    if cfg["fps_embedding"] is not None and not 1 <= float(cfg["fps_embedding"]) <= 120:
        raise RuntimeError("[summary-server] fps_embedding must be None or in [1, 120]")
    if cfg["compile_generator_mode"] not in COMPILE_GENERATOR_MODES:
        raise RuntimeError(f"[summary-server] compile_generator_mode={cfg['compile_generator_mode']!r} not in {COMPILE_GENERATOR_MODES}")
    if not 0 <= int(cfg["frames_low_watermark"]) <= 4:
        raise RuntimeError("[summary-server] frames_low_watermark must be in [0, 4] (more than one latent of backlog is just latency)")
    if int(cfg["cuda_device"]) < 0:
        raise RuntimeError("[summary-server] cuda_device must be >= 0")
    origins = cfg.get("cors_allow_origins")
    if not (isinstance(origins, (list, tuple)) and origins and all(
            isinstance(o, str) and o.startswith(("http://", "https://")) and not o.endswith("/") for o in origins)):
        raise RuntimeError("[summary-server] cors_allow_origins must be a non-empty list of 'http(s)://host[:port]' origins, no trailing slash")
    if int(cfg["dynamo_recompile_limit"]) < 8:
        raise RuntimeError("[summary-server] dynamo_recompile_limit < 8 cannot hold the compiled graphs")
    if not 3 <= int(cfg["warmup_compiled_frames"]) <= 64:
        raise RuntimeError("[summary-server] warmup_compiled_frames must be in [3, 64] (compile + cudagraph warmup/record + >= 1 steady frame)")
    if not 0 <= int(cfg["profile_steady_frames"]) <= 64:
        raise RuntimeError("[summary-server] profile_steady_frames must be in [0, 64]")
    for k in ("timesteps", "timestep_shift", "cache_mode", "num_layers", "latent_grid"):
        if k in cfg:
            raise RuntimeError(f"[summary-server] CONFIG[{k!r}] is not a server knob: model facts come from the checkpoint's config")


_assert_config_alive(CONFIG)


# ----------------------------------------------------------------------------------------------
# Checkpoint: contract + weights
# ----------------------------------------------------------------------------------------------
def _must_file(path, label: str) -> str:
    p = os.path.abspath(str(path))
    if not os.path.isfile(p):
        raise FileNotFoundError(f"[summary-server] {label} missing: {p}")
    return p


def _noise_schedule_str(saved: dict) -> str:
    return f"{saved['noise_schedule']}:{float(saved['shifted_cosine_ratio'])}"


def _require_ckpt_keys(saved: dict, *, path: str) -> None:
    missing = [k for k in CKPT_MODEL_KEYS if k not in saved]
    if missing:
        raise RuntimeError(f"[summary-server] {path!r}: ckpt config lacks {missing[0]!r}"
                           + (f" (also missing {missing[1:]})" if len(missing) > 1 else ""))


def _sink_gaps_of(saved: dict) -> int:
    """One sink gap from the ckpt (`sink_gap_max: int` in [2, 64]).

    None is refused: an unbounded stream needs the sink chased.
    """
    gap = sink_gaps_from_config(saved["sink_gap_max"])
    if gap is None:
        raise RuntimeError("[summary-server] sink_gap_max is None: an unbounded stream needs the sink chased")
    if not 2 <= int(gap) <= 64:
        raise RuntimeError(f"[summary-server] sink_gap_max={gap} outside [2, 64]")
    return int(gap)


def _max_relative_frames_of(saved: dict) -> int:
    """Largest |query - key| this ckpt can produce: W + sink_gap_max (one Z0 slot, one chase)."""
    return SinkChaser(raw_window=int(saved["raw_window_frames"]),
                      summary_window=int(saved["summary_window_frames"]),
                      gaps=_sink_gaps_of(saved)).max_relative_frames()



def _validate_saved_config(saved: dict, *, path: str) -> None:
    """An inference checkpoint config (CKPT_FORMAT) with every model fact the server needs."""
    if saved.get("format") != CKPT_FORMAT:
        raise RuntimeError(f"[summary-server] {path!r} has config format={saved.get('format')!r}; "
                           f"this server reads {CKPT_FORMAT!r}")
    _require_ckpt_keys(saved, path=path)
    if int(saved["num_layers"]) < 1:
        raise RuntimeError(f"[summary-server] {path!r}: num_layers={saved['num_layers']!r} must be >= 1")
    ts = [int(t) for t in saved["timesteps"]]
    if len(ts) < 1 or any(ts[i] <= ts[i + 1] for i in range(len(ts) - 1)):
        raise RuntimeError(f"[summary-server] {path!r}: timesteps={ts} must be a non-empty list, high -> low")
    if saved["cache_mode"] not in CACHE_MODES:
        raise RuntimeError(f"[summary-server] {path!r}: cache_mode={saved['cache_mode']!r} not in {CACHE_MODES}")
    for key in ("focus_region", "jump_sprint"):
        if not isinstance(saved[key], bool):
            raise RuntimeError(f"[summary-server] {path!r}: {key}={saved[key]!r} must be a bool")
    _sink_gaps_of(saved)   # one int, in [2, 64]


def focus_region_trained(saved: dict) -> bool:
    """The model knows the z0 focus marker (a region of the seed frame flagged as the subject to hold)."""
    return bool(saved["focus_region"])


def ckpt_profile(saved: dict) -> dict:
    """What the player can do with this checkpoint, read from its config:
      name        the config's `profile` (e.g. "gta5").
      jump_sprint the config's `jump_sprint`: dims 2 / 3 of the action vector are live, so the client's
                  space / shift pass through (controls.allow_jump_sprint). False: they are forced to 0.
    Exposed per model in GET /models `details[id].profile` and as `active_profile` in GET /status."""
    jump_sprint = bool(saved["jump_sprint"])
    controls = ["w", "a", "s", "d", "left", "right"] + (["space", "shift"] if jump_sprint else [])
    return {"name": str(saved["profile"]), "jump_sprint": jump_sprint, "controls": controls}


def focus_region_for_ckpt(saved: dict) -> torch.Tensor | None:
    """The default focus region: the server has no subject box, so it marks one centred diamond covering
    SERVE_FOCUS_COVERAGE of the start image. None when the ckpt has no focus operator."""
    if not focus_region_trained(saved):
        return None
    h, w = LATENT_HW
    return centered_focus_region(h, w, shape="diamond", coverage=SERVE_FOCUS_COVERAGE)


def _focus_region_desc(saved: dict, region: torch.Tensor | None) -> str:
    return "diamond" if region is not None else "none"


@dataclass(frozen=True)
class CkptEntry:
    """One servable checkpoint: the file, its name-derived recipe, and its saved config (weights NOT held)."""
    name: CkptName
    path: str
    saved: dict
    step: int | None

    @property
    def model_id(self) -> str:
        return self.name.model_id

    def details(self) -> dict:
        """What GET /models shows per model (JSON-safe)."""
        s = self.saved
        rec, src = _recency_of(self)
        gap = _sink_gaps_of(s)
        region_on = focus_region_trained(s)
        return {
            "recipe": self.name.recipe, "steps": self.name.steps, "cache_mode": str(s["cache_mode"]),
            "timesteps": [int(t) for t in s["timesteps"]],
            "summary_tokens": int(s["summary_tokens"]), "raw_window_frames": int(s["raw_window_frames"]),
            "summary_window_frames": int(s["summary_window_frames"]),
            "sink_gap_max": int(gap), "recency": rec, "recency_source": src,
            "noise_schedule": _noise_schedule_str(s), "train_step": self.step,
            "file": os.path.basename(self.path), "labels": list(self.name.labels),
            "focus": "diamond" if region_on else "none",
            "profile": ckpt_profile(s),
        }


def _recency_of(entry: CkptEntry) -> tuple[dict | None, str]:
    """(bias dict or None, where it came from): "trained" = the ckpt's own saved bias; "serve" =
    CONFIG["serve_recency"] applied because the file name carries `rec`; "off" = trained without
    and no token; "off (disabled)" = CONFIG["recency_enabled"] is False (every ckpt, whatever it
    was trained with)."""
    if not bool(CONFIG.get("recency_enabled", True)):
        return None, "off (disabled)"
    rec = trained_recency(entry.saved)
    if rec is not None:
        rec = dict(rec)
        # decay_end=None in a ckpt config means W+1.
        if rec.get("decay_end") is None:
            rec["decay_end"] = int(entry.saved["summary_window_frames"]) + 1
        return rec, "trained"
    if entry.name.serve_recency:
        return dict(CONFIG["serve_recency"]), "serve"
    return None, "off"


def _recency_describe(rec: dict | None, src: str) -> str:
    if rec is None:
        return "off"
    cfg = RecencyBiasConfig.from_dict(rec)
    body = cfg.describe() if cfg is not None else str(rec)
    return f"{body} ({src})"


def read_ckpt_config(path: str) -> tuple[dict, int | None]:
    """The saved config + train step WITHOUT materialising the weights (mmap): cheap enough to do for a whole directory."""
    ck = torch.load(path, map_location="cpu", weights_only=False, mmap=True)
    saved = ck.get("config") if isinstance(ck, dict) else None
    if not isinstance(saved, dict) or "model" not in ck:
        raise RuntimeError(f"[summary-server] {path!r} is not a {{'model', 'config', ...}} checkpoint")
    step = ck.get("step")
    return saved, (int(step) if isinstance(step, (int, float)) else None)


def load_ckpt_weights(path: str) -> dict:
    """The full state_dict (wrapper key space, 'module.' stripped). Called at startup and on every switch."""
    ck = torch.load(_must_file(path, "checkpoint"), map_location="cpu", weights_only=False)
    return {(k[7:] if k.startswith("module.") else k): v for k, v in ck["model"].items()}


def scan_checkpoints_dir(dir_path: str) -> dict[str, CkptEntry]:
    """Every *.pt in `dir_path`. Name-parsed, config-validated, name-vs-config checked,
    architecture-agreed. Refuses an empty directory, an off-schema name, a mislabeled
    file, an unknown config format or a mixed-architecture set: the operator fixes the directory."""
    root = os.path.abspath(str(dir_path))
    if not os.path.isdir(root):
        raise FileNotFoundError(f"[summary-server] checkpoints_dir is not a directory: {root}")
    paths = sorted(glob.glob(os.path.join(root, "*.pt")))
    if not paths:
        raise FileNotFoundError(f"[summary-server] no *.pt in checkpoints_dir {root} "
                                "(e.g. gta_model.pt)")
    entries: dict[str, CkptEntry] = {}
    for p in paths:
        saved, step = read_ckpt_config(p)
        if saved.get("format") != CKPT_FORMAT:
            raise RuntimeError(f"[summary-server] {p!r} has config format={saved.get('format')!r}: checkpoints_dir "
                               f"takes {CKPT_FORMAT!r} checkpoints only")
        try:
            name = parse_ckpt_name(p, saved)
        except ValueError as e:
            raise RuntimeError(f"{e}  (a ckpt in checkpoints_dir must be schema-named)") from e
        if name.model_id in entries:
            raise RuntimeError(f"[summary-server] two files map to model_id {name.model_id!r} in {root}")
        if CONFIG["serve_window_frames"] is not None and "summary_window_frames" in saved:
            w_ckpt, w_serve = int(saved["summary_window_frames"]), int(CONFIG["serve_window_frames"])
            if w_serve < int(saved.get("raw_window_frames", 0)):
                raise RuntimeError(f"[summary-server] {os.path.basename(p)}: serve_window_frames={w_serve} < "
                                   f"raw_window_frames={int(saved['raw_window_frames'])}")
            if w_ckpt != w_serve:
                print(f"[summary-server] {os.path.basename(p)}: summary_window_frames ckpt={w_ckpt} -> served W={w_serve} "
                      "(sliding window; cache capacity, no weights)", flush=True)
            saved = {**saved, "summary_window_frames": w_serve}
        _validate_saved_config(saved, path=p)
        check_name_against_config(name, saved)
        entry = CkptEntry(name=name, path=p, saved=saved, step=step)
        rec, src = _recency_of(entry)
        if rec is not None:
            cfg = RecencyBiasConfig.from_dict(rec)
            w = int(saved["summary_window_frames"])
            if cfg.decay_end > w + 1:
                raise RuntimeError(f"[summary-server] {os.path.basename(p)}: recency ({src}) decay_end={cfg.decay_end} > "
                                   f"summary_window_frames + 1 = {w + 1}; the bias never reaches its floor before eviction")
        entries[name.model_id] = entry
    if not entries:
        raise FileNotFoundError(f"[summary-server] no servable checkpoint in checkpoints_dir {root}")
    first = next(iter(entries.values()))
    for e in entries.values():
        diff = [k for k in CKPT_ARCH_KEYS if _arch_val(e.saved[k]) != _arch_val(first.saved[k])]
        if diff:
            raise RuntimeError(f"[summary-server] {e.model_id} and {first.model_id} differ in architecture keys {diff}: "
                               "one compiled generator serves the directory, so every ckpt must share the architecture "
                               "(move the odd one out)")
    for e in entries.values():
        d = e.details()
        print(f"[summary-server] ckpt {e.model_id}: t={d['timesteps']} "
              f"cache_mode={d['cache_mode']} recency={_recency_describe(d['recency'], d['recency_source'])} "
              f"{sink_gaps_describe(d['sink_gap_max'])} focus={d['focus']} "
              f"X={d['summary_tokens']} Y={d['raw_window_frames']} W={d['summary_window_frames']} "
              f"schedule={d['noise_schedule']} step={d['train_step']} "
              f"profile={d['profile']['name']}"
              f"{' jump+sprint' if d['profile']['jump_sprint'] else ''} ({d['file']})", flush=True)
    return entries


def _arch_val(v):
    return os.path.basename(v) if isinstance(v, str) and v.endswith(".pt") else (tuple(v) if isinstance(v, (list, tuple)) else v)


def build_wrapper(saved: dict, device) -> SummaryformerWrapper:
    """The ckpt config IS the architecture (`wrapper_ctor_kwargs_from_ckpt_config`)."""
    return SummaryformerWrapper(**wrapper_ctor_kwargs_from_ckpt_config(saved)).to(device)


# Params the model declares but that a checkpoint may legitimately lack. The z0 focus-region embedding
# is zero-init: a ckpt without it loads as the identical model (and is served with no region). A switch
# between the two kinds must overwrite the parameter either way -- load_weights_in_place zeroes it when
# the file lacks it.
INERT_AT_INFERENCE_KEYS = frozenset(Z0_FOCUS_REGION_EMBED_ALLOWED_MISSING)
_FOCUS_EMBED_KEY = f"model.{Z0_FOCUS_REGION_EMBED_NAME}"


def load_weights_in_place(wrapper, sd: dict, saved: dict) -> None:
    """copy_ into existing storages. One trunk: the file holds the whole model."""
    depth, want_depth = int(wrapper.model.num_layers), num_layers_from_ckpt_config(saved)
    if want_depth != depth:
        raise RuntimeError(f"[summary-server] ckpt has {want_depth} blocks != generator {depth}")
    want = set(wrapper.state_dict())
    missing, unexpected = want - set(sd), set(sd) - want
    if unexpected or (missing - INERT_AT_INFERENCE_KEYS):
        raise RuntimeError(f"[summary-server] ckpt key set != generator: missing={sorted(missing - INERT_AT_INFERENCE_KEYS)[:5]} "
                           f"unexpected={sorted(unexpected)[:5]}")
    wrapper.load_state_dict(sd, strict=False)
    params = dict(wrapper.named_parameters())
    if _FOCUS_EMBED_KEY in missing and _FOCUS_EMBED_KEY in params:
        with torch.no_grad():
            params[_FOCUS_EMBED_KEY].zero_()
    if missing:
        print(f"[summary-server] ckpt lacks {sorted(missing)}: inert at inference (the focus embedding is zeroed, "
              f"so a focus region is a no-op)", flush=True)
    if _FOCUS_EMBED_KEY in sd:
        n = float(sd[_FOCUS_EMBED_KEY].float().norm())
        print(f"[summary-server] ckpt carries the z0 focus embedding ({Z0_FOCUS_REGION_EMBED_NAME}): norm={n:.4f} "
              f"{'(trained, served with a centred diamond)' if n > 0 else '(exactly zero: never trained)'}", flush=True)
    print(f"[summary-server] weights loaded in place: {len(sd)} tensors", flush=True)


def apply_dtype_mode(wrapper, mode: str) -> None:
    if mode == "fp32_master_bf16_autocast":
        wrapper.to(torch.float32)
    else:
        wrapper.to(torch.bfloat16)
        for name in _COND_ISLANDS:
            getattr(wrapper.model, name).to(torch.float32)
    n_bf16 = sum(p.numel() for p in wrapper.parameters() if p.dtype == torch.bfloat16)
    n_fp32 = sum(p.numel() for p in wrapper.parameters() if p.dtype == torch.float32)
    print(f"[dtype][generator] mode={mode} bf16_params={n_bf16:,} fp32_params={n_fp32:,}"
          + (f" islands={_COND_ISLANDS}" if n_bf16 else ""), flush=True)


# ----------------------------------------------------------------------------------------------
# fp8 trunk linears (CONFIG["fp8_linears"], env SUMMARY_FP8)
# ----------------------------------------------------------------------------------------------
_FP8_E4M3_MAX = 448.0
# Per-step trunk linears, by attribute path suffix inside each TransformerBlock. The FiLM/cond islands are
# fp32 and excluded: no speed in them, precision for free.
_FP8_LINEAR_SUFFIXES = ("visual_self_attention.visual_qkv", "visual_self_attention.visual_out",
                        "visual_feed_forward.proj", "visual_feed_forward.out")


class Fp8Linear(torch.nn.Module):
    """nn.Linear drop-in: y = scaled_mm(q(x), q(W)^T) + b with e4m3 operands and fp32 scales.
    `weight`/`bias` are the ORIGINAL Parameters (same storage -> state_dict keys and load_weights_in_place's copy_
    keep working); `w8`/`w_scale` are non-persistent buffers rebuilt by requantize() after every weight load."""

    def __init__(self, lin: torch.nn.Linear, mode: str):
        super().__init__()
        if mode not in ("row", "tensor"):
            raise ValueError(mode)
        if lin.in_features % 16 or lin.out_features % 16:
            raise RuntimeError(f"[fp8] {lin.in_features}x{lin.out_features} not a multiple of 16")
        self.mode = mode
        self.in_features, self.out_features = lin.in_features, lin.out_features
        self.weight = lin.weight
        self.bias = lin.bias
        self.register_buffer("w8", torch.empty(0), persistent=False)
        self.register_buffer("w_scale", torch.empty(0), persistent=False)
        self.requantize()

    @torch.no_grad()
    def requantize(self) -> None:
        w = self.weight.detach().float()                                   # (N, K)
        if self.mode == "row":
            amax = w.abs().amax(dim=1, keepdim=True)                        # (N, 1)
        else:
            amax = w.abs().amax().reshape(1, 1)
        scale = (amax.clamp_min(1e-12) / _FP8_E4M3_MAX)
        self.w8 = (w / scale).clamp_(-_FP8_E4M3_MAX, _FP8_E4M3_MAX).to(torch.float8_e4m3fn).contiguous()
        self.w_scale = scale.reshape(1, -1).contiguous() if self.mode == "row" else scale.reshape(()).contiguous()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shp = x.shape
        x2 = x.reshape(-1, shp[-1])
        if self.mode == "row":
            amax = x2.abs().amax(dim=1, keepdim=True).float()               # (M, 1) per-token
            s_a = amax.clamp_min(1e-12) / _FP8_E4M3_MAX
            s_b = self.w_scale                                              # (1, N)
        else:
            amax = x2.abs().amax().float()
            s_a = (amax.clamp_min(1e-12) / _FP8_E4M3_MAX).reshape(())
            s_b = self.w_scale
        x8 = (x2.float() / s_a).clamp_(-_FP8_E4M3_MAX, _FP8_E4M3_MAX).to(torch.float8_e4m3fn)
        y = torch._scaled_mm(x8, self.w8.t(), scale_a=s_a, scale_b=s_b, out_dtype=torch.bfloat16)
        if self.bias is not None:
            y = y + self.bias.to(y.dtype)
        return y.reshape(*shp[:-1], self.out_features)

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}, fp8={self.mode}"


def apply_fp8_linears(wrapper, mode: str) -> int:
    """Swap the per-step trunk linears for Fp8Linear. Returns the count. Idempotent."""
    if mode == "off":
        return 0
    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] * 10 + torch.cuda.get_device_capability()[1] < 89:
        raise RuntimeError("[fp8] needs an sm89+ GPU (Ada / Hopper / Blackwell)")
    # Probe the exact kernel once so a missing rowwise kernel fails loudly at startup, not inside compile.
    dev = next(wrapper.parameters()).device
    probe = Fp8Linear(torch.nn.Linear(64, 32, device=dev, dtype=torch.bfloat16), mode)
    probe(torch.randn(16, 64, device=dev, dtype=torch.bfloat16))
    n = 0
    for blk in wrapper.model.blocks:
        for suffix in _FP8_LINEAR_SUFFIXES:
            parent_path, attr = suffix.rsplit(".", 1)
            parent = blk
            for p in parent_path.split("."):
                parent = getattr(parent, p, None)
                if parent is None:
                    break
            if parent is None:
                continue
            lin = getattr(parent, attr, None)
            if isinstance(lin, Fp8Linear):
                n += 1
            elif isinstance(lin, torch.nn.Linear):
                setattr(parent, attr, Fp8Linear(lin, mode))
                n += 1
    if n == 0:
        raise RuntimeError("[fp8] found no trunk linears to swap (block attribute names changed?)")
    print(f"[fp8] mode={mode}: {n} trunk linears -> Fp8Linear (e4m3 operands, fp32 scales, bf16 out)", flush=True)
    return n


def requantize_fp8_linears(wrapper) -> int:
    n = 0
    for m in wrapper.modules():
        if isinstance(m, Fp8Linear):
            m.requantize()
            n += 1
    return n


# ----------------------------------------------------------------------------------------------
# StageSpec from the ckpt config (renoise schedule + full-res diag-gauss stats)
# ----------------------------------------------------------------------------------------------
def _load_diag_stats(path: str, *, expect_scale: str, tag: str) -> dict:
    p = path if os.path.isabs(path) else os.path.join(os.path.dirname(os.path.abspath(__file__)), path)
    art = torch.load(_must_file(p, f"diag-gauss {tag}"), map_location="cpu", weights_only=True)
    if int(art["mean"].numel()) != 48 or float(art["std"].min()) <= 1e-3:
        raise RuntimeError(f"[summary-server] diag-gauss {tag} artifact bad (C/std.min): {p!r}")
    got = art.get("scale")
    if not (got == expect_scale or (expect_scale == "fine" and got is None)):
        raise RuntimeError(f"[summary-server] diag-gauss {tag} scale={got!r}, slot needs {expect_scale!r}: {p!r}")
    print(f"[summary-server] diag-gauss {tag} ({expect_scale}): {p}", flush=True)
    return art


# NFE override (the dashboard's "NFE override" slider), two modes:
#   "repeat" (default): K' calls, EVERY one at one of the ckpt's own timesteps. The ckpt's grid is kept
#       and calls are repeated: K'//K each, the remainder going to the LOWEST t's (final refinement), e.g.
#       3r at 4 -> [999, 878, 669, 669], 6 -> [999, 999, 878, 878, 669, 669]. A repeat re-noises x0_hat
#       back to the same t with fresh noise and denoises again. The model only ever sees timesteps it
#       knows; the last call stays at the ckpt's lowest t. Needs K' >= K.
#   "grid": the grid rebuilt from the ckpt's own shift at K' points, `resolve_sampling(saved, steps=K')`
#       (3r at 6 -> [999, 942, 878, 795, 669, 439]). Half the calls, including the LAST one, are at
#       timesteps the model was not made for; visibly worse at 6 and 12. Kept for A/B.
# Frame cost scales ~ K' (+1 recache pass). The override is server-wide state, applied under the lease
# before each body job, and survives a checkpoint switch (absolute step count, not a multiple).
NFE_OVERRIDE_MAX = 32
NFE_MODES = ("repeat", "grid")
NFE_MODE_DEFAULT = "repeat"


def ckpt_nfe(saved: dict) -> int:
    return len(saved["timesteps"])


def parse_nfe_mode_form(raw: str | None) -> str:
    s = str(raw or "").strip().lower() or NFE_MODE_DEFAULT
    if s not in NFE_MODES:
        raise ValueError(f"nfe_mode must be one of {NFE_MODES}; got {raw!r}")
    return s


def repeat_step_plan(trained_spaced: list[int], total: int) -> tuple[int, ...]:
    """Spaced-index plan with ``total`` calls over the trained grid (high -> low): base repeats for
    every point, the remainder on the lowest-t points."""
    k = len(trained_spaced)
    if total < k:
        raise ValueError(f"repeat mode needs nfe >= the trained {k} steps; got {total}")
    base, extra = divmod(int(total), k)
    plan: list[int] = []
    for j, idx in enumerate(trained_spaced):           # trained_spaced is high -> low
        plan += [int(idx)] * (base + (1 if j >= k - extra else 0))
    return tuple(plan)


def parse_nfe_form(raw: str | int | None) -> int | None:
    """Form field -> override. "" / None / 0 = the ckpt's own step count (no override)."""
    s = str(raw if raw is not None else "").strip()
    if not s or s == "0":
        return None
    try:
        n = int(s)
    except ValueError as e:
        raise ValueError(f"nfe must be an integer in 1..{NFE_OVERRIDE_MAX} (0 = checkpoint default); got {raw!r}") from e
    if not (1 <= n <= NFE_OVERRIDE_MAX):
        raise ValueError(f"nfe must be in 1..{NFE_OVERRIDE_MAX} (0 = checkpoint default); got {n}")
    return n


def build_stage_spec(saved: dict, *, recency: dict | None, nfe: int | None = None,
                     nfe_mode: str = NFE_MODE_DEFAULT) -> StageSpec:
    """Full-res 24x42 facts from the ckpt config. The diag-gauss files loaded with
    expect_scale='fine' are the full-res pair. ``nfe`` != the ckpt's own step count is the NFE
    override: "repeat" keeps the ckpt's timesteps and repeats calls (plan), "grid" rebuilds the grid
    at that step count on the ckpt's own shift law (see NFE_MODES)."""
    if nfe_mode not in NFE_MODES:
        raise ValueError(f"[summary-server] nfe_mode {nfe_mode!r} not in {NFE_MODES}")
    own_nfe = ckpt_nfe(saved)
    overridden = nfe is not None and int(nfe) != own_nfe
    if overridden and not (1 <= int(nfe) <= NFE_OVERRIDE_MAX):
        raise ValueError(f"[summary-server] nfe override {nfe} not in 1..{NFE_OVERRIDE_MAX}")
    regrid = overridden and nfe_mode == "grid"
    samp = resolve_sampling(saved, steps=int(nfe) if regrid else None,
                            tag="serve" if not regrid else f"serve/nfe{int(nfe)}")
    # The grid is rebuilt from the schedule and the shift, and must land on the ckpt's own timesteps.
    # Under a "grid" NFE override the grid is deliberately not that list; log it instead.
    want_ts = [int(t) for t in saved["timesteps"]]
    body_ts = list(samp.timesteps)
    if overridden and not regrid:
        if body_ts != want_ts:
            raise RuntimeError(f"[summary-server] ckpt timesteps={want_ts} != rebuilt grid {body_ts} (repeat override)")
        plan = repeat_step_plan(list(samp.indices), int(nfe))
        samp = replace(samp, step_plan=plan)
        body_ts = list(samp.timesteps)
        print(f"[summary-server] NFE override (repeat): body t={body_ts} ({int(nfe)} calls, all at the ckpt's t; "
              f"ckpt grid {want_ts})", flush=True)
    elif regrid:
        if len(body_ts) != int(nfe):
            raise RuntimeError(f"[summary-server] nfe override {nfe} produced {len(body_ts)} body steps: {body_ts}")
        print(f"[summary-server] NFE override (grid): body t={body_ts} ({int(nfe)} steps) on the ckpt's shift law; "
              f"ckpt timesteps={want_ts} ({own_nfe} steps) -- off-grid calls", flush=True)
    elif body_ts != want_ts:
        raise RuntimeError(f"[summary-server] ckpt timesteps={want_ts} != rebuilt grid {body_ts} "
                           f"(timestep_shift={saved['timestep_shift']!r}, schedule {_noise_schedule_str(saved)})")
    cache_mode = str(saved["cache_mode"])
    if cache_mode not in CACHE_MODES:
        raise RuntimeError(f"[summary-server] cache_mode={cache_mode!r} not in {CACHE_MODES}")
    z0p, rp = CONFIG["sae_diag_gauss_stats_path"], CONFIG["sae_diag_gauss_rest_stats_path"]
    return StageSpec(
        hw=LATENT_HW,
        stats_z0=_load_diag_stats(z0p, expect_scale="fine", tag="z0"),
        stats_rest=_load_diag_stats(rp, expect_scale="fine", tag="rest"),
        samp=samp,
        cache_mode=cache_mode,
        sink_gap_max=_sink_gaps_of(saved),
        recency=recency,
        focus_region=focus_region_for_ckpt(saved),
    )


def _load_autoencoder(path, device):
    path = _must_file(path, "sae_ckpt")
    sae_cfg = SAEConfig(out_dims_z0=48, out_dims_z_rest=48)
    ck = torch.load(path, map_location="cpu", weights_only=True)
    for key in ("enc", "dec"):
        if key not in ck:
            raise RuntimeError(f"[summary-server] sae_ckpt missing {key!r}: {path}")
    enc = ShakeAE(sae_cfg)
    dec = Decoder(sae_cfg)
    for mod, key in ((enc, "enc"), (dec, "dec")):
        try:
            mod.load_state_dict(ck[key], strict=True)
        except RuntimeError as e:
            raise RuntimeError(f"[summary-server] SAE {key} does not fit this autoencoder ({path}): {e}") from e
    enc.to(device)
    dec.to(device)
    enc.eval().requires_grad_(False)
    dec.eval().requires_grad_(False)
    dec_dtype = SAE_DECODER_DTYPES[str(CONFIG["sae_decoder_dtype"])]
    dec.to(dtype=dec_dtype)
    n3e, n2e = _apply_channels_last(enc)
    n3d, n2d = _apply_channels_last(dec)
    _assert_channels_last(enc, tag="summary-server/enc")
    _assert_channels_last(dec, tag="summary-server/dec")
    print(f"[summary-server] autoencoder C=48 enc+dec: {path} | decoder {dec_dtype} "
          f"({sum(p.numel() for p in dec.parameters()) / 1e6:.1f}M params) "
          f"| channels_last: enc Conv3d={n3e} Conv2d={n2e}, dec Conv3d={n3d} Conv2d={n2d}", flush=True)
    return enc, dec


# ----------------------------------------------------------------------------------------------
# Runtime (one CUDA thread; cudagraphs are thread-affine)
# ----------------------------------------------------------------------------------------------
RUNTIME: dict = {}
_CUDA_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="summary-cuda")
_CUDA_THREAD_ID: int | None = None


def run_on_cuda_thread(fn, *args, **kwargs):
    if _CUDA_THREAD_ID is not None and threading.get_ident() == _CUDA_THREAD_ID:
        return fn(*args, **kwargs)

    def _bound():
        global _CUDA_THREAD_ID
        _CUDA_THREAD_ID = threading.get_ident()
        torch.cuda.set_device(int(CONFIG["cuda_device"]))
        return fn(*args, **kwargs)

    return _CUDA_EXECUTOR.submit(_bound).result()


def _build_generator(saved: dict, sd: dict, device, *, max_relative_frames: int) -> SummaryformerWrapper:
    """The ONE generator: built from the default ckpt's config, compiled once; switches only swap
    weights into it. `max_relative_frames` = the largest SinkChaser.max_relative_frames() over the
    whole directory (W + sink_gap_max; the RoPE table is extended once, after .to(device), before
    compile). One Z0 slot, one chase."""
    gen = build_wrapper(saved, device)
    load_weights_in_place(gen, sd, saved)
    gen.eval().requires_grad_(False)
    apply_dtype_mode(gen, str(CONFIG["generator_dtype_mode"]))
    apply_fp8_linears(gen, str(CONFIG["fp8_linears"]))   # after the dtype cast: quantizes the final bf16 weights
    extend_rope_table_for_streaming(gen.model, table_len=int(CONFIG["rope_table_len"]),
                                    max_relative_frames=int(max_relative_frames))
    return gen


def _compile_runtime(gen: SummaryformerWrapper, dec: Decoder):
    """-> (forward_body compiled | None, decode_steady compiled | None)."""
    import torch._dynamo as dynamo
    try:
        torch._logging.set_logs(recompiles=True, cudagraphs=True, graph_breaks=True)
    except (TypeError, ValueError):
        torch._logging.set_logs(recompiles=True, graph_breaks=True)
    limit = int(CONFIG["dynamo_recompile_limit"])
    for attr in ("recompile_limit", "cache_size_limit", "accumulated_recompile_limit", "accumulated_cache_size_limit"):
        if hasattr(dynamo.config, attr):
            setattr(dynamo.config, attr, max(int(getattr(dynamo.config, attr)), limit))
    print(f"[compile] dynamo recompile budget >= {limit}", flush=True)

    body_x = sae_x = None
    dyn, gmode = bool(CONFIG["compile_dynamic"]), str(CONFIG["compile_generator_mode"])
    if not CONFIG["use_compiled_generator_in_steady_state"]:
        print("[compile] SKIP generator: use_compiled_generator_in_steady_state=False (eager)", flush=True)
    elif gmode == "max-autotune+cudagraphs":
        opts = {"max_autotune": True, "coordinate_descent_tuning": True, "triton.cudagraphs": True}
        print(f"[compile] generator: torch.compile(forward_body, dynamic={dyn}, fullgraph=False, cudagraphs=ON via options)",
              flush=True)
        body_x = torch.compile(gen.forward_body, fullgraph=False, dynamic=dyn, options=opts)
    else:
        print(f"[compile] generator: torch.compile(forward_body, dynamic={dyn}, fullgraph=False, mode={gmode})", flush=True)
        body_x = torch.compile(gen.forward_body, fullgraph=False, dynamic=dyn, mode=gmode)
    if CONFIG["compile_sae_decode"]:
        mode = str(CONFIG["compile_sae_decode_mode"])
        print(f"[compile] sae: torch.compile(decode_steady, mode={mode}, dynamic=False)", flush=True)
        sae_x = torch.compile(dec.decode_steady, mode=mode, dynamic=False)
    else:
        print("[compile] SKIP sae decode: compile_sae_decode=False (eager)", flush=True)
    return body_x, sae_x


def _latest_id(registry: dict) -> str:
    """The entry with the highest saved train step (None counts as -1); ties -> last in sorted id order."""
    return max(sorted(registry), key=lambda k: (registry[k].step if registry[k].step is not None else -1, k))


def _default_model_id(registry: dict[str, CkptEntry]) -> str:
    want = CONFIG["default_model_id"]
    if want is None:
        return sorted(registry)[0]
    if str(want).lower() == "latest":
        return _latest_id(registry)
    if str(want) not in registry:
        raise RuntimeError(f"[summary-server] default_model_id={want!r} not in checkpoints_dir; have {sorted(registry)}")
    return str(want)


def init_runtime_once() -> dict:
    if RUNTIME:
        return RUNTIME
    if not torch.cuda.is_available():
        raise RuntimeError("[summary-server] CUDA only")
    if int(CONFIG["cuda_device"]) >= torch.cuda.device_count():
        raise RuntimeError(f"[summary-server] cuda_device={CONFIG['cuda_device']} but only {torch.cuda.device_count()} GPUs visible")
    device = torch.device(f"cuda:{int(CONFIG['cuda_device'])}")
    torch.cuda.set_device(device)
    torch.set_float32_matmul_precision("high")
    torch.backends.cudnn.allow_tf32 = True
    # Two kernels exist, flex and flash. Flex is the main one (a read with the recency bias is flex regardless),
    # so make every streaming read flex: one kernel, one set of numbers.
    print(f"[startup] streaming attention backend -> {set_streaming_attention_backend('flex')!r}", flush=True)
    print(f"[startup] runtime load on {device} ({torch.cuda.get_device_name(device)}), {torch.cuda.device_count()} GPUs visible: "
          "ckpt + SAE + session (first /z0 or /ws waits)", flush=True)
    registry = scan_checkpoints_dir(str(CONFIG["checkpoints_dir"]))
    mid = _default_model_id(registry)
    entry = registry[mid]
    saved = entry.saved
    span = max(_max_relative_frames_of(e.saved) for e in registry.values())
    gen = run_on_cuda_thread(_build_generator, saved, load_ckpt_weights(entry.path), device, max_relative_frames=span)
    enc, dec = _load_autoencoder(CONFIG["sae_ckpt"], device)
    body_x, sae_x = run_on_cuda_thread(_compile_runtime, gen, dec)
    RUNTIME.update(device=device, registry=registry, generator=gen, sae_encoder=enc, sae_decoder=dec,
                   body_x=body_x, sae_x=sae_x, active_model_id=None, saved=None, session=None)
    run_on_cuda_thread(_activate_checkpoint, entry, weights=None)
    print(f"[startup] runtime ready: active={mid} of {sorted(registry)}", flush=True)
    return RUNTIME


def _cuda_mem(tag: str) -> None:
    if not torch.cuda.is_available():
        return
    free, total = torch.cuda.mem_get_info()
    print(f"[mem] {tag}: alloc={torch.cuda.memory_allocated() / 2**30:.2f}GiB "
          f"reserved={torch.cuda.memory_reserved() / 2**30:.2f}GiB "
          f"free={free / 2**30:.2f}/{total / 2**30:.2f}GiB", flush=True)


def _drop_compiled_graphs() -> None:
    """Forget dynamo + CUDAGraph trees so they release mark_static_address'd KV buffers.
    Compiled wrappers (`body_x`, `sae_x`) stay; the next warmup recompiles / re-records against the new session."""
    try:
        from torch._inductor.cudagraph_trees import reset_cudagraph_trees
        reset_cudagraph_trees()
    except Exception as e:
        print(f"[switch] cudagraph tree reset skipped: {e!r}", flush=True)
    torch.compiler.reset()


def _recency_key(rec) -> tuple | None:
    cfg = RecencyBiasConfig.from_dict(rec)
    if cfg is None:
        return None
    return (int(cfg.flat_until), int(cfg.decay_end), float(cfg.beta), float(cfg.gamma))


def _same_session_geometry(session, spec: StageSpec) -> bool:
    """True if the live session can keep its KV / compiled graphs and only swap weights + retarget.

    What the compiled graph is shaped by: the buffer geometry (X/Y/W, agreed across the directory),
    the cache mode, the sink gap, and WHETHER a recency bias tensor is attached. The bias VALUES
    are not: `token_bias` is a static-address fp32 buffer refreshed in place per frame. Only
    on <-> off rebuilds."""
    if session is None:
        return False
    old = getattr(session, "spec", None)
    if old is None:
        return False
    if str(old.cache_mode) != str(spec.cache_mode):
        return False
    if int(old.sink_gap_max) != int(spec.sink_gap_max):
        return False
    return (old.recency is None) == (_recency_key(spec.recency) is None)


def _activate_checkpoint(entry: CkptEntry, *, weights: dict | None) -> None:
    """CUDA thread. Make `entry` the served model: weights in place (if given), session retargeted
    or rebuilt. Same cache_mode + sink gap + recency on/off: reuse the live session (weight swap
    only). Different geometry: drop the old KV and the compiled graphs that pin it, then build a
    new session and warmup."""
    rt = RUNTIME
    gen, dec, device = rt["generator"], rt["sae_decoder"], rt["device"]
    saved = entry.saved
    old = rt.get("session")
    rec, rec_src = _recency_of(entry)
    spec = build_stage_spec(saved, recency=rec, nfe=rt.get("nfe_override"),
                            nfe_mode=rt.get("nfe_mode", NFE_MODE_DEFAULT))
    rt["active_model_id"], rt["saved"] = None, None
    if weights is not None:
        load_weights_in_place(gen, weights, saved)
        del weights
        n8 = requantize_fp8_linears(gen)
        if n8:
            print(f"[activate] fp8: requantized {n8} trunk linears from the new weights", flush=True)
    print(f"[activate] {entry.model_id}: {entry.name.recipe} body t={list(spec.timesteps)} "
          f"(ckpt timesteps={[int(t) for t in saved['timesteps']]}, shift={saved['timestep_shift']}) "
          f"cache_mode={spec.cache_mode} recency={_recency_describe(rec, rec_src)} "
          f"{sink_gaps_describe(spec.sink_gap_max)} focus={_focus_region_desc(saved, spec.focus_region)} "
          f"X={int(saved['summary_tokens'])} Y={int(saved['raw_window_frames'])} W={int(saved['summary_window_frames'])}",
          flush=True)

    if _same_session_geometry(old, spec):
        old.retarget(spec)
        rt["session"], rt["saved"] = old, saved
        rt["active_model_id"] = entry.model_id
        print("[activate] reused session (same cache_mode/sink gap and recency "
              f"{'on' if rec else 'off'} in both): weights + spec retargeted in place, no realloc, no recompile",
              flush=True)
        return

    if old is not None:
        _cuda_mem("before release")
        old.release()
        rt["session"] = None
        del old
        _drop_compiled_graphs()
        gc.collect()
        if device.type == "cuda":
            torch.cuda.empty_cache()
        _cuda_mem("after release")

    session = build_session_from_runtime(gen, dec, spec, device, rgb_hw=RGB_HW,
                                         forward_body_compiled=rt["body_x"], decode_steady_compiled=rt["sae_x"],
                                         use_compiled_sae_in_steady_state=bool(CONFIG["use_compiled_sae_in_steady_state"]))
    if bool(CONFIG["warmup_compile_at_startup"]) and (rt["body_x"] is not None or rt["sae_x"] is not None):
        _warmup_compile(session, ring=int(saved["summary_window_frames"]), label=entry.model_id)
    session.set_sphere_mode(rt.get("sphere_mode", SPHERE_MODE_DEFAULT))
    session.reprefill_every = int(rt.get("reprefill_every", int(CONFIG["reprefill_every"] or 0)))
    rt["session"], rt["saved"] = session, saved
    rt["active_model_id"] = entry.model_id
    if device.type == "cuda":
        torch.cuda.empty_cache()
    _cuda_mem("after activate")


# SAE sphere projection (summary_streaming_inference.SPHERE_MODES: off | decode | all). Server-wide,
# like the NFE override: sent as form field `sphere` on /z0,
# applied before the game, reported by /models. Cheap eager ops; no recompile.
SPHERE_MODE_DEFAULT = "off"


def parse_sphere_form(raw: str | None) -> str:
    s = str(raw or "").strip().lower() or SPHERE_MODE_DEFAULT
    if s in ("0", "false", "none"):
        s = "off"
    if s in ("1", "true", "on"):
        s = "all"
    if s not in SPHERE_MODES:
        raise ValueError(f"sphere must be one of {SPHERE_MODES}; got {raw!r}")
    return s


def apply_sphere_mode(mode: str) -> None:
    """Any thread (plain attribute write; the session reads it per call). GPU free or under the lease."""
    rt = init_runtime_once()
    session = rt.get("session")
    if session is None:
        raise RuntimeError("[sphere] no active session")
    session.set_sphere_mode(mode)
    rt["sphere_mode"] = mode


def _sphere_state() -> dict:
    rt = init_runtime_once()
    return {"mode": rt.get("sphere_mode", SPHERE_MODE_DEFAULT), "modes": list(SPHERE_MODES)}


# Forward re-prefill cadence (SummaryStreamSession.reprefill). Server-wide like the sphere mode: form field
# `reprefill_every` on /z0 ("" = keep, 0 = off, N = every N body frames), reported by /models.
REPREFILL_EVERY_MAX = 4096


def parse_reprefill_form(raw: str | None) -> int | None:
    """"" -> None (keep the current cadence); "0"/"off" -> 0; N in [1, REPREFILL_EVERY_MAX]."""
    s = str(raw or "").strip().lower()
    if not s:
        return None
    if s in ("off", "none", "false"):
        return 0
    try:
        n = int(s)
    except ValueError as e:
        raise ValueError(f"reprefill_every must be an int in [0, {REPREFILL_EVERY_MAX}]; got {raw!r}") from e
    if not (0 <= n <= REPREFILL_EVERY_MAX):
        raise ValueError(f"reprefill_every must be in [0, {REPREFILL_EVERY_MAX}]; got {n}")
    return n


def apply_reprefill_every(every: int | None) -> None:
    """Any thread (plain attribute write, read by the session at the top of step()). None = keep."""
    if every is None:
        return
    rt = init_runtime_once()
    session, saved = rt.get("session"), rt.get("saved")
    if session is None or saved is None:
        raise RuntimeError("[reprefill] no active session")
    if every > 0 and str(session.spec.cache_mode) != "recache":
        raise ValueError(f"reprefill needs cache_mode='recache'; the active ckpt is {session.spec.cache_mode!r}")
    session.reprefill_every = int(every)
    rt["reprefill_every"] = int(every)


def _reprefill_state() -> dict:
    rt = init_runtime_once()
    session = rt.get("session")
    return {
        "every": int(rt.get("reprefill_every", int(CONFIG["reprefill_every"] or 0))),
        "max": REPREFILL_EVERY_MAX,
        "supported": session is not None and str(session.spec.cache_mode) == "recache",
        "history_frames": None if session is None else int(session.W),
        "count": 0 if session is None else int(session.n_reprefills),
        "last": None if session is None else session.last_reprefill,
    }


# Custom z0 focus region (replaces the served centred diamond for one game). Form field `focus_mask` on /z0:
#   ""            keep the ckpt's served default (diamond when the ckpt has focus_region, else none)
#   "none"        no region for this game
#   "HxW:<bits>"  a painted token grid: H rows x W cols, `bits` = H*W chars of '0'/'1', row-major, '1' = in the
#                 region. The grid must be LATENT_HW (24x42 for 384x672). The client paints it over the same
#                 cover-crop the server applies to the image.
# The operator saw regions covering 10% to 30% of the tokens in training; coverage outside that is allowed but
# reported so the UI can warn. A ckpt without the operator (embedding zero) refuses a custom mask: it would
# be a no-op and the user would think it worked.
FOCUS_COVERAGE_TRAINED = (0.10, 0.30)
_FOCUS_KEEP = object()


def parse_focus_mask_form(raw: str | None):
    """"" -> _FOCUS_KEEP; "none" -> None; "HxW:bits" -> Bool [1, H, W] (CPU). ValueError on anything else."""
    s = str(raw or "").strip()
    if not s or s.lower() in ("default", "diamond", "keep"):
        return _FOCUS_KEEP
    if s.lower() in ("none", "off", "0", "false"):
        return None
    head, sep, bits = s.partition(":")
    if not sep:
        raise ValueError("focus_mask must be '', 'none' or 'HxW:<H*W bits of 0/1>'")
    try:
        hs, ws = head.lower().split("x")
        h, w = int(hs), int(ws)
    except ValueError as e:
        raise ValueError(f"focus_mask grid {head!r} is not HxW") from e
    if not (1 <= h <= 256 and 1 <= w <= 256):
        raise ValueError(f"focus_mask grid {h}x{w} out of range")
    bits = bits.strip()
    if len(bits) != h * w or set(bits) - {"0", "1"}:
        raise ValueError(f"focus_mask needs exactly {h * w} chars of '0'/'1' for a {h}x{w} grid; got {len(bits)}")
    region = torch.tensor([c == "1" for c in bits], dtype=torch.bool).view(1, h, w)
    if not bool(region.any()):
        raise ValueError("focus_mask is empty; send 'none' to disable the region instead")
    return region


def resolve_focus_region(choice, *, saved: dict, hw: tuple[int, int] | None):
    """(region | None, mode) for a parsed form choice. ``hw`` = the grid the mask must match exactly (the game);
    None = any grid. Custom on a ckpt without the focus operator -> ValueError."""
    if choice is _FOCUS_KEEP:
        region = focus_region_for_ckpt(saved)
        return region, ("diamond" if region is not None else "none")
    if choice is None:
        return None, "none"
    if not focus_region_trained(saved):
        raise ValueError("this checkpoint was trained without the z0 focus operator; a custom focus_mask would be a no-op")
    if hw is not None and tuple(choice.shape[-2:]) != (int(hw[0]), int(hw[1])):
        raise ValueError(f"focus_mask grid {choice.shape[-2]}x{choice.shape[-1]} != the game's latent grid {hw[0]}x{hw[1]}")
    return choice, "custom"


def _focus_summary(region, mode: str) -> dict:
    n = 0 if region is None else int(region.sum())
    total = 0 if region is None else int(region.numel())
    return {"mode": mode, "tokens": n, "grid": None if region is None else [int(region.shape[-2]), int(region.shape[-1])],
            "coverage": (n / total) if total else 0.0}


def _focus_state() -> dict:
    """What /models reports: the game grid the client paints on, whether the active ckpt has the operator, and
    what it serves by default."""
    rt = init_runtime_once()
    saved = rt.get("saved")
    trained = bool(saved) and focus_region_trained(saved)
    return {
        "grid": [int(LATENT_HW[0]), int(LATENT_HW[1])],
        "custom_supported": True,
        "trained": trained,
        "default": "diamond" if trained else "none",
        "default_coverage": float(SERVE_FOCUS_COVERAGE),
        "coverage_trained": [float(FOCUS_COVERAGE_TRAINED[0]), float(FOCUS_COVERAGE_TRAINED[1])],
    }


def _nfe_state() -> dict:
    """What /models and the generators report: the override (None = ckpt law), the ckpt's own step
    count, and the body grid the live session actually steps on."""
    rt = init_runtime_once()
    session, saved = rt.get("session"), rt.get("saved")
    return {
        "override": rt.get("nfe_override"),
        "mode": rt.get("nfe_mode", NFE_MODE_DEFAULT),
        "modes": list(NFE_MODES),
        "ckpt_steps": None if saved is None else ckpt_nfe(saved),
        "active_steps": None if session is None else int(session.spec.nfe),
        "timesteps": [] if session is None else [int(t) for t in session.spec.timesteps],
        "max": NFE_OVERRIDE_MAX,
    }


def apply_nfe_override(nfe: int | None, mode: str = NFE_MODE_DEFAULT) -> None:
    """CUDA thread, GPU free or under the caller's lease. Retarget the live session onto the body plan
    for ``nfe`` (None = the ckpt's own) in ``mode`` (NFE_MODES). No-op when nothing changes. Retarget
    keeps the KV buffer and the compiled graphs (NFE is a loop count, not a graph shape) and resets
    the session."""
    rt = init_runtime_once()
    session, saved = rt.get("session"), rt.get("saved")
    if session is None or saved is None:
        raise RuntimeError("[nfe] no active session")
    if mode not in NFE_MODES:
        raise ValueError(f"[nfe] mode {mode!r} not in {NFE_MODES}")
    want = None if nfe is None or int(nfe) == ckpt_nfe(saved) else int(nfe)
    mode = NFE_MODE_DEFAULT if want is None else mode   # mode is meaningless without an override
    if (want == rt.get("nfe_override") and mode == rt.get("nfe_mode", NFE_MODE_DEFAULT)
            and int(session.spec.nfe) == (want or ckpt_nfe(saved))):
        return
    entry = rt["registry"][rt["active_model_id"]]
    rec, _src = _recency_of(entry)
    spec = build_stage_spec(saved, recency=rec, nfe=want, nfe_mode=mode)
    session.retarget(spec)
    rt["nfe_override"], rt["nfe_mode"] = want, mode
    print(f"[nfe] {entry.model_id}: body t={list(spec.timesteps)} ({spec.nfe} calls"
          f"{'' if want is None else f', {mode} override of the trained {ckpt_nfe(saved)}'})", flush=True)


def switch_checkpoint(model_id: str) -> str:
    """Serve `model_id`. No-op if already active. The caller guarantees the GPU is free
    (OWNER.busy() is False and the engine is stopped): the session is replaced under it."""
    rt = init_runtime_once()
    mid = str(model_id).strip().lower()
    if mid not in rt["registry"]:
        raise KeyError(f"unknown model_id={model_id!r}; valid={sorted(rt['registry'])}")
    if rt.get("active_model_id") == mid:
        return mid
    entry = rt["registry"][mid]
    weights = load_ckpt_weights(entry.path)
    print(f"[switch] {rt.get('active_model_id')!r} -> {mid!r} (loading {os.path.basename(entry.path)})", flush=True)
    t0 = time.perf_counter()
    run_on_cuda_thread(_activate_checkpoint, entry, weights=weights)
    print(f"[switch] active={mid} in {time.perf_counter() - t0:.1f}s", flush=True)
    return mid


_PROFILE_BUCKETS: tuple[tuple[str, str], ...] = (
    # (bucket, regex on the CUDA kernel name; first match wins; order matters)
    # inductor emits flex as a templated kernel named triton_tem_fused_<ops>; the autotune log calls it
    # triton_flex_attention_N. "## Call CompiledFxGraph ##" rows are CPU-side range markers that CUPTI reports
    # with device time spanning the whole replayed graph (double-counts everything) -> dropped in _profile_steady.
    ("flex_attention",   r"flex_attention|triton_tem_fused"),
    ("attention (other)", r"fmha_cutlass|PyTorchMemEffAttention|flash_fwd|fused_attention"),
    ("conv (sae)",       r"conv|cudnn|winograd|implicit_gemm|fprop|fft2d|nchw|nhwc"),
    ("gemm",             r"gemm|cutlass|cublas|nvjet|f8f8|e4m3|rowwise|scaled_mm|Kernel2?x?[0-9]*_[a-z0-9_]*tn|xmma|wgmma|matmul|splitK|sgemm|hgemm|bgemm|ampere_|sm[0-9]+_"),
    ("fused elementwise", r"triton_(poi|red|per|tem|for|unk)|elementwise_kernel|vectorized_|reduce_kernel|index_elementwise|cat_|CatArray|fill_|unrolled_"),
    ("softmax/norm",     r"softmax|layer_norm|LayerNorm|rms_norm|RMSNorm|cunn_"),
    ("memcpy/memset",    r"Memcpy|Memset|memcpy|memset"),
    ("cudagraph replay (opaque)", r"cudaGraphLaunch|CUDA graph"),
)


def _profile_bucket(name: str) -> str:
    import re
    for b, rx in _PROFILE_BUCKETS:
        if re.search(rx, name):
            return b
    return "other"


@torch.no_grad()
def _profile_steady(session, *, action, fps, n_frames: int, label: str, out_dir: str) -> None:
    """N steady-state latents under torch.profiler on the hot graphs. Prints ms-per-latent by kernel bucket and the
    top kernels, writes a Chrome trace. Wall is measured with the profiler ON, so compare buckets to each other and
    to the un-profiled steady wall the warmup printed just above, not to it in absolute terms."""
    from torch.profiler import ProfilerActivity, profile
    dev = session.device
    torch.cuda.synchronize(dev)
    walls: list[float] = []
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=False,
                 profile_memory=False, with_stack=False) as prof:
        for _ in range(int(n_frames)):
            t = time.perf_counter()
            session.step(action, fps=fps)
            torch.cuda.synchronize(dev)
            walls.append((time.perf_counter() - t) * 1e3)
    n = float(n_frames)
    ev = [e for e in prof.key_averages() if getattr(e, "device_type", None) is not None
          and str(e.device_type).endswith("CUDA")]
    # only device-side kernel rows (self device time), not the CPU-side launch rows
    rows = []
    for e in ev:
        self_us = float(getattr(e, "self_device_time_total", 0.0) or 0.0)
        if self_us <= 0.0 or "CompiledFxGraph" in str(e.key) or str(e.key).startswith("## "):
            continue
        rows.append((str(e.key), self_us / 1e3 / n, int(e.count) / n))
    total_ms = sum(r[1] for r in rows)
    buckets: dict[str, list[float]] = {}
    for name, ms, cnt in rows:
        b = _profile_bucket(name)
        acc = buckets.setdefault(b, [0.0, 0.0])
        acc[0] += ms
        acc[1] += cnt
    print(f"\n[profile][{label}] ===== {int(n_frames)} steady latents under torch.profiler =====", flush=True)
    print(f"[profile][{label}] wall per latent (profiler ON): min {min(walls):.0f} / mean {sum(walls) / n:.0f} ms | "
          f"CUDA kernel time per latent {total_ms:.1f} ms | kernel launches per latent {sum(r[2] for r in rows):.0f}",
          flush=True)
    print(f"[profile][{label}] {'bucket':<28}{'ms/latent':>11}{'%':>7}{'launches':>10}", flush=True)
    for b, (ms, cnt) in sorted(buckets.items(), key=lambda kv: -kv[1][0]):
        print(f"[profile][{label}] {b:<28}{ms:>11.2f}{100.0 * ms / max(total_ms, 1e-9):>7.1f}{cnt:>10.0f}", flush=True)
    if "cudagraph replay (opaque)" in buckets and buckets["cudagraph replay (opaque)"][0] > 0.25 * total_ms:
        print(f"[profile][{label}] ** most time is inside an opaque CUDA-graph replay: rerun with "
              f"compile_generator_mode='max-autotune-no-cudagraphs' to see the per-kernel split **", flush=True)
    print(f"[profile][{label}] top kernels (self CUDA ms per latent):", flush=True)
    for name, ms, cnt in sorted(rows, key=lambda r: -r[1])[:30]:
        short = name if len(name) <= 96 else name[:93] + "..."
        print(f"[profile][{label}]   {ms:>8.3f} ms  x{cnt:>6.1f}  [{_profile_bucket(name):<17}] {short}", flush=True)
    os.makedirs(out_dir, exist_ok=True)
    trace = os.path.join(out_dir, f"profile_{label}.json")
    try:
        prof.export_chrome_trace(trace)
        print(f"[profile][{label}] chrome trace -> {trace} (chrome://tracing or perfetto.dev)", flush=True)
    except Exception as e:  # noqa: BLE001
        print(f"[profile][{label}] chrome trace export failed: {e!r}", flush=True)
    sys.stdout.flush()


@torch.no_grad()
def _warmup_compile(session, *, ring: int, label: str) -> None:
    """A throwaway game: the generator is compiled on body frame 1, then runs the whole ring fill
    and `warmup_compiled_frames` steady-state frames on the SAME graph, then reset so the first
    player starts clean."""
    n_c = max(3, int(CONFIG["warmup_compiled_frames"]))
    n = int(ring) + n_c
    print(f"[warmup] ===== {label}: {n} frames on random z0 before serving "
          f"(frame 1 = compile, {ring} ring-fill frames + {n_c} steady frames on the compiled graphs) =====",
          flush=True)
    t0 = time.perf_counter()
    g = torch.Generator(device=session.device).manual_seed(0)
    z0 = torch.randn((1, 48, *LATENT_HW), generator=g, device=session.device)
    fps = None if CONFIG["fps_embedding"] is None else float(CONFIG["fps_embedding"])
    session.start(z0, fps=fps)
    action = torch.zeros((1, 1, 5), device=session.device)
    fill_ms, steady_ms, total_ms = [], [], []
    report_at = {0, 1, 2, 3, ring // 2, ring - 1, ring, ring + 1, n - 2, n - 1}
    for i in range(n):
        tf = time.perf_counter()
        session.step(action, fps=fps)
        wall = (time.perf_counter() - tf) * 1e3
        ms = session.last_ms
        (fill_ms if i < ring else steady_ms).append(float(ms.get("model", float("nan"))))
        total_ms.append(wall)
        if i in report_at or wall > 2000.0:
            phase = "compile" if i == 0 else ("cudagraph warm/record" if i in (1, 2) else ("ring fill" if i < ring else "steady"))
            print(f"[warmup][{label}] frame {i + 1}/{n} ({phase}): wall {wall:.0f} ms | "
                  f"model {float(ms.get('model', float('nan'))):.1f} sae {float(ms.get('sae', float('nan'))):.1f} ms"
                  + ("  ** >2 s: a compile happened on this frame; read the [__recompiles] lines above **"
                     if wall > 2000.0 and i > 2 else ""),
                  flush=True)
    n_prof = int(CONFIG["profile_steady_frames"])
    if n_prof > 0:
        _profile_steady(session, action=action, fps=fps, n_frames=n_prof, label=label,
                        out_dir=os.path.abspath(str(CONFIG["checkpoints_dir"])))
        if os.environ.get("SUMMARY_PROFILE_EXIT", "0") == "1":
            print(f"[profile][{label}] SUMMARY_PROFILE_EXIT=1 -> exiting after the report", flush=True)
            os._exit(0)
    session.reset()
    fill_hot = fill_ms[3:] if len(fill_ms) > 3 else fill_ms
    steady_wall = total_ms[ring:]
    print(f"[warmup] ===== {label} done in {time.perf_counter() - t0:.0f}s; compiled graphs are hot =====", flush=True)
    print(f"[warmup][{label}][model ms/frame = all denoising passes] first frame (compile) {fill_ms[0]:.0f} ms | "
          f"ring fill (hot) min/max {min(fill_hot):.1f}/{max(fill_hot):.1f} | steady per frame={[round(x, 1) for x in steady_ms]} "
          f"-> {min(steady_ms):.1f} ms"
          f"{'' if max(fill_hot) < 3 * min(steady_ms) else '  ** a ring-fill frame was >3x steady: recompile in the fill? **'}",
          flush=True)
    steady_min = min(steady_wall)
    print(f"[warmup][{label}] steady wall per latent min {steady_min:.0f} ms -> ~{1e3 / steady_min:.1f} lat/s = "
          f"~{4e3 / steady_min:.0f} rgb fps ceiling (model+sae)", flush=True)


_RGB_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3, 1, 1)
_RGB_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3, 1, 1)


def _sae_encode_z0(enc, x01: torch.Tensor, *, want_hw: tuple[int, int], name: str) -> torch.Tensor:
    z0, _ = enc.encode_bootstrap((x01 - _RGB_MEAN.to(x01.device)) / _RGB_STD.to(x01.device))
    z0 = z0.squeeze(2) if z0.dim() == 5 else z0
    if tuple(z0.shape) != (1, 48, *want_hw):
        raise RuntimeError(f"[summary-server] {name} {tuple(z0.shape)} != {(1, 48, *want_hw)}")
    return z0.contiguous()


def center_crop_resize(im: Image.Image, size_wh: tuple[int, int]) -> Image.Image:
    """Aspect-preserving: scale so the image COVERS size_wh, then center crop. Never stretches (a
    stretched start image is off-distribution for the whole game)."""
    tw, th = (int(v) for v in size_wh)
    w, h = im.size
    s = max(tw / w, th / h)
    nw, nh = max(tw, round(w * s)), max(th, round(h * s))
    im = im.resize((nw, nh), Image.BICUBIC if s > 1 else Image.LANCZOS)
    left, top = (nw - tw) // 2, (nh - th) // 2
    out = im.crop((left, top, left + tw, top + th))
    if out.size != (tw, th):
        raise RuntimeError(f"[summary-server] center_crop_resize produced {out.size}, wanted {(tw, th)}")
    if (w, h) != (nw, nh) or (nw, nh) != (tw, th):
        print(f"[z0] {w}x{h} -> scale {s:.3f} -> {nw}x{nh} -> center crop {tw}x{th} (aspect kept, "
              f"{'sides' if nw > tw else 'top/bottom' if nh > th else 'nothing'} cropped)", flush=True)
    return out


@torch.no_grad()
def encode_upload_z0(image: Image.Image, rgb_hw: tuple[int, int] = RGB_HW) -> torch.Tensor:
    """PIL -> center crop to ``rgb_hw`` (H, W) [0,1] -> z0 [1,48,H/16,W/16] RAW (CPU). The session
    gaussianizes. Default is the live game grid (384x672 -> 24x42)."""
    rt = init_runtime_once()
    rh, rw = int(rgb_hw[0]), int(rgb_hw[1])
    if rh % SAE_SPATIAL or rw % SAE_SPATIAL:
        raise ValueError(f"[z0] rgb_hw {(rh, rw)} is not divisible by SAE {SAE_SPATIAL}")
    want_hw = (rh // SAE_SPATIAL, rw // SAE_SPATIAL)
    im = center_crop_resize(image.convert("RGB"), (rw, rh))
    x01 = torch.from_numpy(np.asarray(im)).to(device=rt["device"], dtype=torch.float32).permute(2, 0, 1).div_(255).view(1, 1, 3, rh, rw)
    z0 = _sae_encode_z0(rt["sae_encoder"], x01, want_hw=want_hw, name="z0")
    return z0.cpu()


# ----------------------------------------------------------------------------------------------
# OwnerLock, log tee, H264, Engine
# ----------------------------------------------------------------------------------------------
BUSY_HTTP = "Another player is currently using the live demo. Please try again in a moment."
BUSY_WS = "Another player is using the live demo. Try again shortly."
_OWNER_STALE_SECONDS = 20.0
_LOG_QUEUE: queue.Queue[str] = queue.Queue(maxsize=1000)
_TAG_RE = re.compile(r"^\[([^\]]+)\]")


@dataclass(frozen=True)
class Lease:
    token: str
    user_id: str


class OwnerLock:
    """One GPU, one player. The hold is a lease token, not the user id. The same user without a live
    WS: claim() replaces the reservation. The same user with a live WS: None (an HTTP call cannot
    invalidate the stream). Another user: only if the lease is stale and has no live WS.

    A lease without a WS is reclaimable after `_OWNER_STALE_SECONDS`."""

    def __init__(self, stale_seconds: float = _OWNER_STALE_SECONDS):
        self._lock, self._lease, self._ts, self._ws, self._stale = (
            threading.Lock(), None, 0.0, False, float(stale_seconds))

    def _issue(self, user_id: str, now: float, *, why: str) -> Lease:
        self._lease = Lease(token=uuid.uuid4().hex, user_id=str(user_id))
        self._ts, self._ws = now, False
        print(f"[ws][lock] claim granted user={user_id!r} lease={self._lease.token[:8]} why={why}")
        return self._lease

    def claim(self, user_id: str) -> Lease | None:
        now, uid = time.monotonic(), str(user_id)
        with self._lock:
            cur = self._lease
            if cur is None:
                return self._issue(uid, now, why="free")
            if cur.user_id == uid and not self._ws:
                return self._issue(uid, now, why="replace-pre-ws")
            if cur.user_id != uid and not self._ws and (now - self._ts) > self._stale:
                return self._issue(uid, now, why="stale-reclaim")
            print(f"[ws][lock] claim denied user={uid!r} owner={cur.user_id!r} ws={self._ws}")
            return None

    def mark_ws(self, lease: Lease) -> bool:
        with self._lock:
            if self._lease is None or self._lease.token != lease.token:
                return False
            self._ws, self._ts = True, time.monotonic()
            return True

    def release(self, lease: Lease) -> bool:
        with self._lock:
            if self._lease is None or self._lease.token != lease.token:
                return False
            print(f"[ws][lock] released {lease.user_id!r} lease={lease.token[:8]}")
            self._lease, self._ws = None, False
            return True

    def owner_user_id(self) -> str | None:
        with self._lock:
            return None if self._lease is None else self._lease.user_id

    def ws_active(self) -> bool:
        """True while a game WS holds the lease."""
        with self._lock:
            return bool(self._ws)

    def busy(self) -> bool:
        """True while a game is live. HTTP routes 423 on this."""
        with self._lock:
            return bool(self._ws)


OWNER = OwnerLock()


# ── Tickets (optional) ────────────────────────────────────────────────────────────────────────────
# For a server behind a queue. With both env vars set, every GPU-touching endpoint requires a ticket
# `uid.server_id.exp.hmac` (`?ticket=` on GET / WS, `ticket` form field on POST): no ticket -> 401 /
# WS close 1008. The ticket is bound to the user id in the path, to this server's id and to an
# expiry. Both env vars unset -> open server (your own box), announced once at startup. `/models` and
# `/status` stay open.
_TICKET_SECRET = (os.environ.get("LUCID_DEMO_TICKET_SECRET") or "").strip()
_TICKET_SERVER_ID = (os.environ.get("LUCID_DEMO_SERVER_ID") or "").strip()
if bool(_TICKET_SECRET) != bool(_TICKET_SERVER_ID):
    raise RuntimeError("[summary-server] set BOTH LUCID_DEMO_TICKET_SECRET and LUCID_DEMO_SERVER_ID, or neither")
TICKETS_REQUIRED = bool(_TICKET_SECRET)
_TICKET_ERR = "This server only accepts players sent by the lucidml demo queue."


def _ticket_check(user_id: str, ticket: str | None) -> str | None:
    """None when `ticket` admits `user_id` to this server now; else a short reason (never the secret)."""
    if not TICKETS_REQUIRED:
        return None
    parts = (ticket or "").strip().split(".")
    if len(parts) != 4:
        return "missing or malformed ticket"
    uid, sid, exp_s, sig = parts
    if uid != str(user_id):
        return "ticket is for another user"
    if sid != _TICKET_SERVER_ID:
        return "ticket is for another server"
    try:
        exp = int(exp_s)
    except ValueError:
        return "malformed ticket expiry"
    if exp < int(time.time()):
        return "ticket expired"
    want = hmac.new(_TICKET_SECRET.encode(), f"{uid}.{sid}.{exp_s}".encode(), hashlib.sha256).hexdigest()
    if not hmac.compare_digest(want, sig.lower()):
        return "bad ticket signature"
    return None


def _require_ticket_http(user_id: str, ticket: str | None) -> None:
    why = _ticket_check(user_id, ticket)
    if why is not None:
        print(f"[ticket] 401 user={user_id!r}: {why}")
        raise HTTPException(status_code=401, detail=f"{_TICKET_ERR} ({why})")


# ── Occupancy for /status ────────────────────────────────────────────────────────────────────────
# OWNER is held only from /create-game to the end of the game. /status must also report the ramp
# as busy: a checkpoint switch (tens of seconds), the /z0 encode, and the gap between a successful
# /z0 and the /ws open (the client is connecting). `_INFLIGHT` counts GPU-touching requests in
# progress; `_PENDING_Z0` is uid -> time of the last successful /z0 not yet consumed by /ws,
# honoured for `_PENDING_Z0_GRACE_S`.
_OCC_LOCK = threading.Lock()
_INFLIGHT = 0
_PENDING_Z0: dict[str, float] = {}
_PENDING_Z0_GRACE_S = 45.0


class _inflight:
    def __enter__(self):
        global _INFLIGHT
        with _OCC_LOCK:
            _INFLIGHT += 1
        return self

    def __exit__(self, *exc):
        global _INFLIGHT
        with _OCC_LOCK:
            _INFLIGHT -= 1
        return False


def _pending_z0_set(user_id: str) -> None:
    with _OCC_LOCK:
        _PENDING_Z0[str(user_id)] = time.monotonic()


def _pending_z0_clear(user_id: str) -> None:
    with _OCC_LOCK:
        _PENDING_Z0.pop(str(user_id), None)


def _occupied() -> dict:
    """What /status reports. `busy` = any of: owner lock held, a request on the GPU, a fresh z0 waiting
    for its /ws."""
    now = time.monotonic()
    with _OCC_LOCK:
        for uid, ts in list(_PENDING_Z0.items()):
            if now - ts > _PENDING_Z0_GRACE_S:
                del _PENDING_Z0[uid]
        inflight, pending = _INFLIGHT, len(_PENDING_Z0)
    owner = OWNER.busy()
    return {"busy": bool(owner or inflight > 0 or pending > 0), "owner": owner,
            "inflight": inflight, "pending_z0": pending}


class _TeeWriter:
    """print() fragments -> full lines -> bounded queue (drop when full). Locked: FastAPI and the CUDA thread both print."""

    def __init__(self, original):
        self._orig, self._buf, self._lock = original, "", threading.Lock()

    def write(self, text: str) -> int:
        with self._lock:
            self._orig.write(text)
            self._buf += str(text)
            while "\n" in self._buf:
                line, self._buf = self._buf.split("\n", 1)
                if line.strip():
                    try:
                        _LOG_QUEUE.put_nowait(line.strip())
                    except queue.Full:
                        pass
            return len(text)

    def flush(self):
        with self._lock:
            self._orig.flush()

    def __getattr__(self, name):
        return getattr(self._orig, name)


def install_log_tee() -> None:
    if not isinstance(sys.stdout, _TeeWriter):
        sys.stdout = _TeeWriter(sys.stdout)


def pop_log_line() -> str | None:
    try:
        return _LOG_QUEUE.get_nowait()
    except queue.Empty:
        return None


def log_ws_payload(line: str) -> dict:
    m = _TAG_RE.match(line)
    return {"type": "log", "tag": m.group(1) if m else "misc", "msg": line, "ts": time.time()}


class H264StreamEncoder:
    """Annex-B, no muxer, size = RGB_HW. Missing PyAV raises."""

    def __init__(self):
        import av
        height, width = RGB_HW
        fps, name = int(round(float(CONFIG["stream_fps"]))), str(CONFIG["stream_h264_encoder"])
        ctx = av.CodecContext.create(name, "w")
        ctx.width, ctx.height, ctx.pix_fmt = int(width), int(height), "yuv420p"
        ctx.framerate, ctx.time_base = Fraction(fps, 1), Fraction(1, fps)
        ctx.bit_rate, ctx.gop_size, ctx.max_b_frames = int(CONFIG["stream_h264_bitrate"]), int(CONFIG["stream_h264_gop"]), 0
        if name == "libx264":
            ctx.options = {"preset": "ultrafast", "tune": "zerolatency", "profile": "baseline"}
        elif name == "h264_nvenc":
            ctx.options = {"preset": "p1", "tune": "ull", "profile": "baseline", "delay": "0", "zerolatency": "1"}
        self._av, self._ctx, self._pts = av, ctx, 0
        print(f"[ws][video] H264 {name} {width}x{height} fps={fps} bitrate={int(CONFIG['stream_h264_bitrate'])}")

    def encode(self, frame_rgb: np.ndarray) -> list[tuple[bool, bytes]]:
        if frame_rgb.shape != (*RGB_HW, 3):
            raise RuntimeError(f"[summary-server] H264 frame {frame_rgb.shape} != {(*RGB_HW, 3)}")
        vf = self._av.VideoFrame.from_ndarray(frame_rgb, format="rgb24").reformat(format="yuv420p")
        vf.pts, vf.time_base = self._pts, self._ctx.time_base
        self._pts += 1
        return [(bool(pkt.is_keyframe), bytes(pkt)) for pkt in self._ctx.encode(vf)]

    def close(self) -> None:
        list(self._ctx.encode(None))


def pack_h264_au(is_key: bool, au: bytes) -> bytes:
    return (b"\x01" if is_key else b"\x00") + au


class FramePacer:
    """Sender-side pacing at the wire rate: the engine enqueues 4 frames per tick, we send one per
    dt. The rate can be changed live (client `stream_fps`); a change re-anchors the clock so it
    never bursts to "catch up" to the old schedule."""

    def __init__(self):
        self.fps = float(CONFIG["stream_fps"])
        self.dt, self._t = 1.0 / self.fps, None

    def set_fps(self, fps: float) -> None:
        """Pace-only: the H264 codec context keeps CONFIG['stream_fps'] as its framerate/time_base for the
        life of the socket, so a live change alters arrival timing, not the PTS stride inside the stream.
        Arrival-time players (the Next.js client) follow it; PTS-strict players would not."""
        if fps == self.fps:
            return
        self.fps, self.dt, self._t = float(fps), 1.0 / float(fps), None

    def next_delay(self) -> float:
        now = time.monotonic()
        if self._t is None:
            self._t = now
            return 0.0
        self._t += self.dt
        return max(0.0, self._t - now)


class Engine:
    """Controls -> session.start / step / reset, all on the CUDA thread. Frames: uint8 HWC."""

    def __init__(self):
        rt = init_runtime_once()
        self.device = rt["device"]
        self.frames: queue.Queue[np.ndarray] = queue.Queue(maxsize=8)
        self._running, self._future = False, None
        self._ctrl_lock, self._controls = threading.Lock(), {}
        self._stream_fps = float(CONFIG["stream_fps"])
        self._paced = 0

    @property
    def session(self):
        """The ACTIVE checkpoint's session (replaced by switch_checkpoint between games), never cached on the engine."""
        s = RUNTIME.get("session")
        if s is None:
            raise RuntimeError("[summary-server] no active checkpoint session (a switch is in progress or failed)")
        return s

    @property
    def running(self) -> bool:
        """True from start() until stop() clears the game loop."""
        return bool(self._running)

    @staticmethod
    def _stream_fps_of(payload: dict) -> float | None:
        """Client `stream_fps`: the wire rate, number in [1, 120]. Absent -> unchanged. Garbage is an error, not a fallback.
        Distinct from `fps_embedding` (the FiLM value): this one only tells the sender how fast to push frames."""
        if "stream_fps" not in payload:
            return None
        v = payload["stream_fps"]
        try:
            f = float(v)
        except (TypeError, ValueError):
            raise RuntimeError(f"[summary-server] controls.stream_fps={v!r} is not a number")
        if not 1.0 <= f <= 120.0:
            raise RuntimeError(f"[summary-server] controls.stream_fps={f} outside [1, 120]")
        return f

    @property
    def stream_fps(self) -> float:
        with self._ctrl_lock:
            return self._stream_fps

    def set_controls(self, payload) -> None:
        if isinstance(payload, dict):
            fps = self._stream_fps_of(payload)
            with self._ctrl_lock:
                self._controls = payload
                if fps is not None and fps != self._stream_fps:
                    print(f"[ws][video] stream_fps {self._stream_fps:g} -> {fps:g} (client)", flush=True)
                    self._stream_fps = fps

    def _take(self) -> dict:
        with self._ctrl_lock:
            return dict(self._controls)

    @staticmethod
    def _fps_of(state: dict) -> float | None:
        """Client `fps_embedding`: number in [1, 120], or null / "none" / "off" for no fps conditioning.
        Key absent -> CONFIG start value. Garbage is an error, not a fallback."""
        if "fps_embedding" not in state:
            v = CONFIG["fps_embedding"]
            return None if v is None else float(v)
        v = state["fps_embedding"]
        if v is None or (isinstance(v, str) and v.strip().lower() in ("none", "off", "null", "")):
            return None
        try:
            f = float(v)
        except (TypeError, ValueError):
            raise RuntimeError(f"[summary-server] controls.fps_embedding={v!r} is not a number / null")
        if not 1.0 <= f <= 120.0:
            raise RuntimeError(f"[summary-server] controls.fps_embedding={f} outside [1, 120]")
        return f

    def _wait_for_room(self) -> None:
        """Real-time pacing, BEFORE the controls are read: sleep until the sender has drained the queue down to
        `frames_low_watermark`. The action is then sampled as late as possible and its latent lands just as the
        previous frames run out, so action -> pixel is ~ watermark/wire_fps + one model step, not the whole queue."""
        low = int(CONFIG["frames_low_watermark"])
        waited = False
        while self._running and self.frames.qsize() > low:
            waited = True
            time.sleep(0.002)
        if waited:
            self._paced += 1
            if self._paced == 1:
                print(f"[engine][pace] model outruns the wire ({self.stream_fps:g} fps): waiting for queue <= {low} frames "
                      "before each step (real-time pacing, not a stall)", flush=True)

    @staticmethod
    def _to_u8_frames(rgb: torch.Tensor) -> np.ndarray:
        if rgb.dim() != 5 or int(rgb.size(0)) != 1 or int(rgb.size(2)) != 3:
            raise RuntimeError(f"[summary-server] rgb must be [1,T,3,H,W]; got {tuple(rgb.shape)}")
        return rgb[0].permute(0, 2, 3, 1).mul(255.0).clamp(0, 255).to(torch.uint8).cpu().numpy()

    def _put_rgb(self, rgb: torch.Tensor) -> None:
        """Hand 1 or 4 RGB frames to the sender. Blocks if the queue is somehow full (should not happen after
        `_wait_for_room`); never drops (a dropped frame detaches the player's action from what they see)."""
        u8 = self._to_u8_frames(rgb)
        for i in range(int(u8.shape[0])):
            frame = np.ascontiguousarray(u8[i])
            while self._running:
                try:
                    self.frames.put(frame, timeout=0.05)
                    break
                except queue.Full:
                    print(f"[engine][pace][WARN] frame queue full ({self.frames.maxsize}) after _wait_for_room; blocking",
                          flush=True)

    def _run(self, z0: torch.Tensor, focus_region) -> None:
        session = self.session   # pinned for the whole game: a switch is refused while a WS is live
        try:
            # Per-ckpt profile (ckpt_profile): jump_sprint -> space/shift live.
            prof = ckpt_profile(RUNTIME["saved"])
            print(f"[engine] profile={prof['name']} jump_sprint={prof['jump_sprint']}", flush=True)
            state = self._take()
            z0_dev = z0.to(device=self.device, dtype=torch.float32)
            self._put_rgb(session.start(z0_dev, fps=self._fps_of(state), focus_region=focus_region))
            while self._running:
                self._wait_for_room()
                state = self._take()
                fps = self._fps_of(state)
                # A real 5-vector (all zeros = standing still) when there is an fps clock. With no fps
                # conditioning the action is ACTION_SENTINEL and the model skips its action input.
                vec = ([ACTION_SENTINEL] * 5 if fps is None
                       else client_actions_to_5d(state, allow_jump_sprint=prof["jump_sprint"]))
                print(f"[action-input] schema={str(state.get('schemaVersion', 'legacy')).strip()} "
                      f"vec5={[round(float(x), 3) for x in vec]} fps_embedding={fps} "
                      f"stream_fps={self.stream_fps:g} queued_rgb={self.frames.qsize()}", flush=True)
                # Host tensor: the session stages it through pinned memory (async copy).
                self._put_rgb(session.step(torch.tensor(vec, dtype=torch.float32).view(1, 1, 5), fps=fps))
        except BaseException:
            import traceback
            print("[engine][FATAL] generation loop died:\n" + traceback.format_exc(), flush=True)
            raise
        finally:
            session.reset()
            print("[engine] generation loop exited.", flush=True)

    def start(self, z0: torch.Tensor, *, focus_region=FOCUS_DEFAULT) -> None:
        """``focus_region``: FOCUS_DEFAULT (ckpt's served region) | None | Bool [1, h, w] -> session.start."""
        focus_desc = ("default" if focus_region is FOCUS_DEFAULT else "none" if focus_region is None
                      else f"custom:{int(focus_region.sum())}/{int(focus_region.numel())}")
        print(f"[engine] start z0={tuple(z0.shape)} focus={focus_desc}")
        self.stop()
        while not self.frames.empty():
            self.frames.get_nowait()
        with self._ctrl_lock:
            self._controls = {}
            self._stream_fps = float(CONFIG["stream_fps"])
        self._paced = 0
        self._running = True
        self._future = _CUDA_EXECUTOR.submit(self._run, z0, focus_region)

    def stop(self) -> None:
        self._running = False
        fut, self._future = self._future, None
        if fut is None:
            return
        if _CUDA_THREAD_ID is not None and threading.get_ident() == _CUDA_THREAD_ID:
            raise RuntimeError("[summary-server] Engine.stop on the CUDA thread deadlocks")
        fut.result()

    async def stop_async(self) -> None:
        self._running = False
        fut, self._future = self._future, None
        if fut is not None:
            await asyncio.wrap_future(fut)

    def worker_state(self) -> tuple[bool, BaseException | None]:
        """(done, exception)."""
        fut = self._future
        if fut is None or not fut.done():
            return False, None
        return True, fut.exception()


ENGINE: Engine | None = None


def get_engine() -> Engine:
    global ENGINE
    if ENGINE is None:
        ENGINE = Engine()
    return ENGINE


# ----------------------------------------------------------------------------------------------
# Routes
# ----------------------------------------------------------------------------------------------
STASH: dict[str, dict] = {}
_VIDEO_EXT = {".mp4", ".webm", ".mkv", ".mov", ".avi"}


def _image_from_bytes(data: bytes, filename: str = "") -> Image.Image:
    ext = os.path.splitext(str(filename or "").lower())[1]
    if ext in _VIDEO_EXT:
        import tempfile
        import imageio.v2 as imageio
        with tempfile.NamedTemporaryFile(delete=True, suffix=ext) as tmp:
            tmp.write(data)
            tmp.flush()
            reader = imageio.get_reader(tmp.name)
            frame = reader.get_data(0)
            reader.close()
        return Image.fromarray(np.asarray(frame)[:, :, :3].astype(np.uint8), mode="RGB")
    img = Image.open(io.BytesIO(data))
    img.load()
    return img.convert("RGB")


@asynccontextmanager
async def lifespan(app: FastAPI):
    if bool(CONFIG["enable_ws_log_streaming"]):
        install_log_tee()
    print(f"[startup] Summaryformer server: default={CONFIG['default_model_id']} "
          f"dtype={CONFIG['generator_dtype_mode']} checkpoints_dir={CONFIG['checkpoints_dir']}")
    await asyncio.to_thread(init_runtime_once)
    yield
    if ENGINE is not None:
        await ENGINE.stop_async()
    STASH.clear()


app = FastAPI(lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=list(CONFIG["cors_allow_origins"]), allow_credentials=True,
                   allow_methods=["*"], allow_headers=["*"])


_SWITCH_LOCK = asyncio.Lock()   # one switch at a time; /z0 callers queue behind it


@app.get("/models")
async def list_models():
    """The selector's data: ids (the ckpt stems), the startup default, the one loaded now, and per-id
    facts (recipe, timesteps, X/Y/W, cache mode, recency, sink gaps, focus) so the UI can label the entries."""
    rt = init_runtime_once()
    reg: dict[str, CkptEntry] = rt["registry"]
    ids = sorted(reg)
    return {
        "models": ids, "server_kind": "unified",
        "default_model_id": _default_model_id(reg), "active_model_id": rt.get("active_model_id"),
        "details": {mid: reg[mid].details() for mid in ids},
        "nfe": _nfe_state(),
        "sphere": _sphere_state(),
        "reprefill": _reprefill_state(),
        "focus": _focus_state(),
    }


@app.get("/status")
async def server_status():
    """Health + occupancy (no ticket, no GPU work). `busy` = owner lock held OR a GPU request in flight OR
    a z0 waiting for its /ws (see _occupied)."""
    occ = _occupied()
    return {"ok": True, "busy": occ["busy"], "occupancy": occ, "ws_active": OWNER.ws_active(),
            "server_id": _TICKET_SERVER_ID or None, "tickets_required": TICKETS_REQUIRED,
            "session_limit_s": float(CONFIG.get("ws_session_limit_s") or 0.0),
            "active_model_id": RUNTIME.get("active_model_id"),
            "active_profile": ckpt_profile(RUNTIME["saved"])["name"] if RUNTIME.get("saved") else None}


async def _switch_if_requested(model_id: str, *, user_id: str) -> str:
    rt = init_runtime_once()
    mid = str(model_id or "").strip().lower() or str(rt.get("active_model_id"))
    if mid not in rt["registry"]:
        raise HTTPException(status_code=400, detail=f"Unknown model_id={mid!r}. Valid ids: {sorted(rt['registry'])}")
    if mid != rt.get("active_model_id"):
        async with _SWITCH_LOCK:
            if OWNER.busy():
                raise HTTPException(status_code=423, detail=BUSY_HTTP)
            if ENGINE is not None:
                await ENGINE.stop_async()
            print(f"[switch] user={user_id!r} requests {mid!r} (active={rt.get('active_model_id')!r}): switching", flush=True)
            await asyncio.to_thread(switch_checkpoint, mid)
    return mid


def _nfe_or_400(raw: str, mode_raw: str = "") -> tuple[int | None, str]:
    try:
        return parse_nfe_form(raw), parse_nfe_mode_form(mode_raw)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e)) from e


def _sphere_or_400(raw: str) -> str:
    try:
        return parse_sphere_form(raw)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e)) from e


def _reprefill_or_400(raw: str) -> int | None:
    try:
        return parse_reprefill_form(raw)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e)) from e


def _focus_mask_or_400(raw: str):
    try:
        return parse_focus_mask_form(raw)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=f"focus_mask rejected: {e}") from e


@app.post("/z0/{user_id}")
async def upload_z0(user_id: str, z0_image: UploadFile = File(...), model_id: str = Form(default=""),
                    nfe: str = Form(default=""), nfe_mode: str = Form(default=""), sphere: str = Form(default=""),
                    reprefill_every: str = Form(default=""), focus_mask: str = Form(default=""),
                    ticket: str = Form(default="")):
    """z0 image (+ optional model_id, nfe, nfe_mode, sphere, reprefill_every, focus_mask). `model_id` "" =
    keep the active. A different id switches the served checkpoint here, BEFORE the encode. `nfe` "" / 0 = the ckpt's
    own body steps; N retargets the session onto an N-call body plan (NFE override, `nfe_mode` repeat|grid) for the
    game that follows. `sphere` off|decode|all = SAE sphere projection. `reprefill_every` "" = keep, 0 = off,
    N = rebuild the KV from the clean history every N body frames (forward re-prefill). `focus_mask` "" = the ckpt's
    served region, "none" = off, "HxW:bits" = a painted z0 focus region for this game (parse_focus_mask_form).
    Refused with 423 while the GPU is busy (a live game); 401 without a ticket when tickets are on."""
    _require_ticket_http(user_id, ticket)
    with _inflight():
        return await _upload_z0(user_id, z0_image, model_id, nfe, nfe_mode, sphere, reprefill_every, focus_mask)


async def _upload_z0(user_id: str, z0_image: UploadFile, model_id: str, nfe: str, nfe_mode: str,
                     sphere: str, reprefill_every: str, focus_mask: str):
    if OWNER.busy():
        return PlainTextResponse(content=BUSY_HTTP, status_code=423)
    want_nfe, want_mode = _nfe_or_400(nfe, nfe_mode)
    want_sphere = _sphere_or_400(sphere)
    want_reprefill = _reprefill_or_400(reprefill_every)
    want_focus = _focus_mask_or_400(focus_mask)
    try:
        mid = await _switch_if_requested(model_id, user_id=user_id)
    except HTTPException as e:
        if e.status_code == 423:
            return PlainTextResponse(content=BUSY_HTTP, status_code=423)
        raise
    prof = ckpt_profile(init_runtime_once()["saved"])   # the ckpt that is active AFTER the switch
    try:
        focus_region, focus_mode = resolve_focus_region(want_focus, saved=init_runtime_once()["saved"], hw=LATENT_HW)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=f"focus_mask rejected: {e}") from e
    focus_for_engine = FOCUS_DEFAULT if want_focus is _FOCUS_KEEP else focus_region
    try:
        await asyncio.to_thread(run_on_cuda_thread, apply_nfe_override, want_nfe, want_mode)
    except (ValueError, RuntimeError) as e:
        raise HTTPException(status_code=400, detail=f"nfe override rejected: {e}") from e
    apply_sphere_mode(want_sphere)
    try:
        apply_reprefill_every(want_reprefill)
    except (ValueError, RuntimeError) as e:
        raise HTTPException(status_code=400, detail=f"reprefill_every rejected: {e}") from e
    raw = await z0_image.read()
    focus_info = _focus_summary(focus_region, focus_mode)
    print(f"[z0] encode user={user_id!r} bytes={len(raw)} name={z0_image.filename!r} model={mid!r} "
          f"nfe={_nfe_state()['active_steps']} sphere={want_sphere} reprefill_every={_reprefill_state()['every']} "
          f"focus={focus_mode}:{focus_info['tokens']} ({focus_info['coverage']:.0%})")
    z0 = encode_upload_z0(_image_from_bytes(raw, z0_image.filename or ""))
    STASH[str(user_id)] = {"z0": z0, "model_id": RUNTIME.get("active_model_id"),
                           "focus_region": focus_for_engine}
    _pending_z0_set(user_id)   # /status says busy until this user's /ws opens (or 45 s pass)
    return {"success": True, "model_id": RUNTIME.get("active_model_id"), "z0_shape": list(z0.shape),
            "nfe": _nfe_state(), "sphere": _sphere_state(), "reprefill": _reprefill_state(), "focus": focus_info,
            "profile": prof}


@app.post("/switch")
async def switch_model(model_id: str = Form(...), user_id: str = Form(default=""), ticket: str = Form(default="")):
    """Switch the served checkpoint without encoding a z0. 423 if the GPU is busy (a live game)."""
    _require_ticket_http(user_id, ticket)
    with _inflight():
        if OWNER.busy():
            return PlainTextResponse(content=BUSY_HTTP, status_code=423)
        try:
            mid = await _switch_if_requested(model_id, user_id=user_id or "switch")
        except HTTPException as e:
            if e.status_code == 423:
                return PlainTextResponse(content=BUSY_HTTP, status_code=423)
            raise
        return {"success": True, "model_id": mid}


@app.get("/create-game/{user_id}")
async def create_game(user_id: str, ticket: str = ""):
    """Reserve the GPU for this user. 423 while a game already holds it; 401 without a ticket when tickets are on."""
    _require_ticket_http(user_id, ticket)
    if OWNER.busy():
        return PlainTextResponse(content=BUSY_HTTP, status_code=423)
    lease = OWNER.claim(user_id)
    if lease is None:
        return PlainTextResponse(content=BUSY_HTTP, status_code=423)
    if str(user_id) not in STASH:
        OWNER.release(lease)
        raise HTTPException(status_code=412, detail="POST /z0/{user_id} first")
    print(f"[create-game] ok user={user_id!r} lease={lease.token[:8]} (engine starts on /ws)")
    return {"success": True}


@app.websocket("/ws/{user_id}")
async def ws_stream(user_id: str, websocket: WebSocket, ticket: str = ""):
    await websocket.accept()
    why = _ticket_check(user_id, ticket)
    if why is not None:
        print(f"[ticket][ws] reject user={user_id!r}: {why}")
        await websocket.close(code=1008, reason=f"{_TICKET_ERR} ({why})"[:120])
        return
    lease = OWNER.claim(user_id)
    if lease is None or not OWNER.mark_ws(lease):
        await websocket.close(code=1013, reason=BUSY_WS)
        return
    stash = STASH.get(str(user_id))
    if stash is None:
        await websocket.close(code=1008, reason="POST /z0 first")
        OWNER.release(lease)
        return
    active = RUNTIME.get("active_model_id")
    if stash.get("model_id") != active:
        await websocket.close(code=1008, reason=f"server model is now {active!r}, you uploaded for "
                                                f"{stash.get('model_id')!r}; POST /z0 again")
        OWNER.release(lease)
        return

    engine = get_engine()
    _pending_z0_clear(user_id)   # the lease (ws flag) carries occupancy from here
    print(f"[ws] open user={user_id!r} lease={lease.token[:8]}")
    engine.start(stash["z0"], focus_region=stash.get("focus_region", FOCUS_DEFAULT))
    video, pacer, stop = H264StreamEncoder(), FramePacer(), False
    worker_exc: list[BaseException] = []
    first_frame_sent = asyncio.Event()   # the session clock starts here, not at accept()
    limit_s = float(CONFIG.get("ws_session_limit_s") or 0.0)
    limit_msg = str(CONFIG.get("ws_session_limit_message") or "Thanks for playing! That's the limit for you.")
    session_over = False

    async def send_frames():
        while not stop:
            try:
                frame = engine.frames.get_nowait()
            except queue.Empty:
                done, exc = engine.worker_state()
                if done:
                    print("[ws][frames] CUDA worker exited" + (f" with {exc!r}" if exc else ""))
                    if exc is not None:
                        worker_exc.append(exc)
                        await asyncio.sleep(0.3)
                    return
                await asyncio.sleep(0.005)
                continue
            pacer.set_fps(engine.stream_fps)
            delay = pacer.next_delay()
            if delay > 0:
                await asyncio.sleep(delay)
            for is_key, au in await asyncio.to_thread(video.encode, frame):
                await websocket.send_bytes(pack_h264_au(is_key, au))
            first_frame_sent.set()

    async def session_timer():
        """Ends the game `limit_s` seconds after the first frame: one text message the client can show, then
        return -> FIRST_COMPLETED -> the finally block closes the socket with code 1000 + the message."""
        nonlocal session_over
        await first_frame_sent.wait()
        await asyncio.sleep(limit_s)
        session_over = True
        print(f"[ws] session limit {limit_s:g}s reached user={user_id!r}")
        try:
            await websocket.send_text(json.dumps({"type": "session_over", "message": limit_msg, "limit_s": limit_s}))
        except Exception as exc:
            print(f"[ws] session_over send failed user={user_id!r}: {exc!r}")

    async def send_logs():
        while not stop:
            line = pop_log_line()
            if line is None:
                await asyncio.sleep(0.05)
                continue
            await websocket.send_text(json.dumps(log_ws_payload(line)))

    async def recv_inputs():
        while not stop:
            payload = json.loads(await websocket.receive_text())
            if not isinstance(payload, dict):
                continue
            if payload.get("type", "keys") != "keys":
                continue   # only controls payloads ({"type": "keys", ...} or no type) are read
            engine.set_controls(payload)

    tasks = [asyncio.create_task(send_frames()), asyncio.create_task(recv_inputs())]
    if bool(CONFIG["enable_ws_log_streaming"]):
        tasks.append(asyncio.create_task(send_logs()))
    if limit_s > 0:
        tasks.append(asyncio.create_task(session_timer()))
    try:
        done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        for t in done:
            exc = t.exception()
            if exc is not None and not isinstance(exc, WebSocketDisconnect):
                print(f"[ws] task died user={user_id!r} exc={exc!r}")
    finally:
        stop = True
        for t in tasks:
            t.cancel()
        video.close()
        await engine.stop_async()
        OWNER.release(lease)
        if worker_exc:
            try:
                await websocket.close(code=1011, reason=f"server: {type(worker_exc[0]).__name__}: {str(worker_exc[0])[:90]}")
            except Exception:
                pass
        elif session_over:
            try:
                await websocket.close(code=1000, reason=limit_msg[:120])   # normal closure; the reason is the user-facing text
            except Exception:
                pass
        print(f"[ws] closed user={user_id!r} lease={lease.token[:8]}"
              + (f" worker_error={worker_exc[0]!r}" if worker_exc else "")
              + (" session_limit" if session_over else ""))


if __name__ == "__main__":
    import uvicorn
    print(f"[main] uvicorn {CONFIG['host']}:{int(CONFIG['port'])}")
    print(f"[ticket] {'REQUIRED, server_id=' + _TICKET_SERVER_ID if TICKETS_REQUIRED else 'OFF (open server: LUCID_DEMO_TICKET_SECRET / LUCID_DEMO_SERVER_ID unset)'}")
    uvicorn.run(app, host=str(CONFIG["host"]), port=int(CONFIG["port"]), reload=False)