app.py

import gc
import hashlib
import math
import json
import os
import subprocess
import tempfile
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

# Disable hf_transfer in subprocess environments (uvx/whisperx can trip over it).
os.environ.pop("HF_HUB_ENABLE_HF_TRANSFER", None)
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0")

# Headless rendering for matplotlib (static brain plots).
os.environ.setdefault("MPLBACKEND", "Agg")

# Some plotting stacks try to initialize GL/display; keep this Space headless-friendly.
os.environ.setdefault("PYVISTA_OFF_SCREEN", "true")
os.environ.setdefault("DISPLAY", "")
os.environ.setdefault("VTK_DEFAULT_RENDER_WINDOW_OFFSCREEN", "true")

import gradio as gr

try:
    import spaces
except Exception:  # pragma: no cover
    # Local dev fallback: run without ZeroGPU decorator.
    class _SpacesFallback:
        @staticmethod
        def GPU(*_args, **_kwargs):
            def _wrap(fn):
                return fn

            return _wrap

    spaces = _SpacesFallback()


CACHE_DIR = Path("./cache")
CACHE_DIR.mkdir(parents=True, exist_ok=True)

TRIBEV2_MODEL_ID = os.environ.get("TRIBEV2_MODEL_ID", "facebook/tribev2")
OPENVLA_MODEL_ID = os.environ.get("OPENVLA_MODEL_ID", "openvla/openvla-7b")
DEPTH_MODEL_ID = os.environ.get("DEPTH_MODEL_ID", "depth-anything/Depth-Anything-V2-Small-hf")
DINO_MODEL_ID = os.environ.get("DINO_MODEL_ID", "facebook/dinov2-small")
SAM_MODEL_ID = os.environ.get("SAM_MODEL_ID", "facebook/sam3")
SAM_FALLBACK_MODEL_ID = os.environ.get("SAM_FALLBACK_MODEL_ID", "facebook/sam2-hiera-tiny")
MAP_ANYTHING_MODEL_ID = os.environ.get("MAP_ANYTHING_MODEL_ID", "facebook/map-anything-apache")
DEFAULT_ROBOT_OUTPUT_MODE = os.environ.get("ROBOT_OUTPUT_MODE", "normalized").strip().lower()

SAMPLE_VIDEO_URL = "https://download.blender.org/durian/trailer/sintel_trailer-480p.mp4"
WORLD_SCOUT_CACHE: Dict[str, Dict[str, Any]] = {}
WORLD_EMBED_CACHE: list[Dict[str, Any]] = []


def _now_ms() -> int:
    return int(time.time() * 1000)


def _cuda_cleanup() -> None:
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _best_effort_login_hf() -> None:
    token = os.environ.get("HF_TOKEN")
    if not token:
        return
    try:
        from huggingface_hub import login

        login(token=token, add_to_git_credential=False)
    except Exception:
        # If login fails, downstream calls may still work via env var; keep going.
        pass


def _download_url(url: str, dst_path: Path) -> Path:
    import urllib.request

    dst_path.parent.mkdir(parents=True, exist_ok=True)
    if dst_path.exists() and dst_path.stat().st_size > 0:
        return dst_path
    with urllib.request.urlopen(url) as r, open(dst_path, "wb") as f:
        f.write(r.read())
    return dst_path


def _as_filepath(value: Any) -> str:
    """
    Gradio component values can be strings, dict payloads, or tuples depending on version.
    Normalize to a local filesystem path string when possible.
    """
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        for k in ("path", "video", "name", "filepath"):
            v = value.get(k)
            if isinstance(v, str) and v:
                return v
        # Some payloads store a nested dict.
        for v in value.values():
            p = _as_filepath(v)
            if p:
                return p
        return ""
    if isinstance(value, (list, tuple)) and value:
        return _as_filepath(value[0])
    try:
        return str(value)
    except Exception:
        return ""


def _probe_video(video_path: str) -> Tuple[float, float, int]:
    import cv2

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return 0.0, 0.0, 0
    fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    cap.release()
    duration = (frame_count / fps) if fps > 0 and frame_count > 0 else 0.0
    return duration, fps, frame_count


def _extract_frame(video_path: str, ts_s: float) -> "PIL.Image.Image":
    import cv2
    import numpy as np
    from PIL import Image

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError("Could not open video")

    fps = float(cap.get(cv2.CAP_PROP_FPS) or 0.0) or 30.0
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    duration = (frame_count / fps) if frame_count > 0 else 0.0

    ts_s = float(ts_s or 0.0)
    if duration > 0:
        ts_s = max(0.0, min(ts_s, max(0.0, duration - 1e-3)))

    frame_idx = int(ts_s * fps)
    if frame_count > 0:
        frame_idx = max(0, min(frame_idx, frame_count - 1))

    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
    ok, frame = cap.read()
    if not ok:
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        ok, frame = cap.read()
    cap.release()

    if not ok or frame is None:
        raise RuntimeError("Could not read frame from video")

    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame = np.asarray(frame)
    return Image.fromarray(frame)


def _image_to_short_video(image_path: str, duration_s: float = 2.0, fps: int = 12) -> str:
    import cv2
    import numpy as np

    from PIL import Image

    image = Image.open(image_path).convert("RGB")
    frame = np.array(image)

    h, w = frame.shape[:2]
    out_path = CACHE_DIR / f"img_{_now_ms()}.mp4"
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(str(out_path), fourcc, float(fps), (w, h))
    n_frames = max(1, int(duration_s * fps))
    bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    for _ in range(n_frames):
        writer.write(bgr)
    writer.release()
    return str(out_path)


def _trim_video_ffmpeg(video_path: str, max_duration_s: float = 10.0) -> str:
    out_path = CACHE_DIR / f"trim_{_now_ms()}.mp4"
    cmd = [
        "ffmpeg",
        "-y",
        "-i",
        str(video_path),
        "-t",
        str(float(max_duration_s)),
        "-vf",
        "scale=480:-2",
        "-c:v",
        "libx264",
        "-preset",
        "veryfast",
        "-crf",
        "28",
        "-c:a",
        "aac",
        "-b:a",
        "96k",
        str(out_path),
    ]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
        if proc.returncode != 0:
            return video_path
        if not out_path.exists() or out_path.stat().st_size == 0:
            return video_path
        return str(out_path)
    except Exception:
        return video_path


def _file_hash(path: str, limit_bytes: int = 64 * 1024 * 1024) -> str:
    digest = hashlib.sha256()
    read_bytes = 0
    with open(path, "rb") as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
            read_bytes += len(chunk)
            if read_bytes >= limit_bytes:
                break
    stat = os.stat(path)
    digest.update(str(stat.st_size).encode("utf-8"))
    return digest.hexdigest()[:16]


def _pil_to_gallery_item(image: "PIL.Image.Image", label: str):
    return image, label


def _sample_video_keyframes(video_path: str, max_frames: int = 8) -> Dict[str, Any]:
    import cv2
    import numpy as np
    from PIL import Image

    video_path = _as_filepath(video_path)
    if not video_path:
        raise RuntimeError("No video path provided")

    duration, fps, frame_count = _probe_video(video_path)
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError("Could not open video")

    if frame_count <= 0:
        frame_count = int((duration or 10.0) * (fps or 30.0))
    if fps <= 0:
        fps = 30.0
    if duration <= 0 and frame_count > 0:
        duration = frame_count / fps

    n = max(1, min(int(max_frames or 8), 16))
    frame_indices = np.linspace(0, max(0, frame_count - 1), n, dtype=int).tolist()

    frames = []
    gray_frames = []
    rows = []
    width = 0
    height = 0

    for idx in frame_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ok, frame_bgr = cap.read()
        if not ok or frame_bgr is None:
            continue
        height, width = frame_bgr.shape[:2]
        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        frame_pil = Image.fromarray(frame_rgb)
        ts = float(idx / fps) if fps else 0.0
        frames.append({"image": frame_pil, "ts": ts, "idx": int(idx)})
        gray = cv2.resize(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY), (96, 54))
        gray_frames.append(gray.astype("float32") / 255.0)
        rows.append([round(ts, 2), int(idx), f"{width}x{height}"])

    cap.release()

    diffs = []
    for i in range(1, len(gray_frames)):
        diffs.append(float(np.mean(np.abs(gray_frames[i] - gray_frames[i - 1]))))

    motion_score = float(np.mean(diffs) * 100.0) if diffs else 0.0
    scene_cut_rows = []
    if diffs:
        ranked = sorted(enumerate(diffs, start=1), key=lambda x: x[1], reverse=True)[:3]
        for i, diff in ranked:
            if i < len(frames):
                scene_cut_rows.append([round(float(frames[i]["ts"]), 2), round(diff * 100.0, 2)])

    return {
        "duration_s": float(duration or 0.0),
        "fps": float(fps or 0.0),
        "frame_count": int(frame_count or 0),
        "resolution": f"{int(width)}x{int(height)}" if width and height else "unknown",
        "motion_score": round(motion_score, 2),
        "frames": frames,
        "keyframe_rows": rows,
        "scene_cut_rows": scene_cut_rows,
    }


def world_scout_video(video_path: str, max_frames: int = 8):
    video_path = _as_filepath(video_path)
    if not video_path:
        return None, [], [], [], "Upload a video first."

    try:
        file_id = _file_hash(video_path)
        cache_key = f"{file_id}:{int(max_frames or 8)}"
        cache_hit = cache_key in WORLD_SCOUT_CACHE
        if not cache_hit:
            scout = _sample_video_keyframes(video_path, max_frames=max_frames)
            scout["file_hash"] = file_id
            WORLD_SCOUT_CACHE[cache_key] = scout
        else:
            scout = WORLD_SCOUT_CACHE[cache_key]

        gallery = [
            _pil_to_gallery_item(frame["image"], f"{frame['ts']:.2f}s")
            for frame in scout.get("frames", [])
        ]
        metadata = {
            "file_hash": scout.get("file_hash"),
            "duration_s": scout.get("duration_s"),
            "fps": scout.get("fps"),
            "frame_count": scout.get("frame_count"),
            "resolution": scout.get("resolution"),
            "motion_score": scout.get("motion_score"),
            "cache": "hit" if cache_hit else "miss",
        }
        status = (
            f"**World Scout ready.** Motion energy `{metadata['motion_score']}` | "
            f"duration `{metadata['duration_s']:.2f}s` | hash `{metadata['file_hash']}`"
        )
        return metadata, gallery, scout.get("keyframe_rows", []), scout.get("scene_cut_rows", []), status
    except Exception as e:
        return None, [], [], [], f"```text\n{type(e).__name__}: {e}\n```"


@dataclass
class _Runtime:
    tribe_model: Any = None
    openvla_model: Any = None
    openvla_processor: Any = None
    depth_model: Any = None
    depth_processor: Any = None
    dino_model: Any = None
    dino_processor: Any = None
    sam_pipeline: Any = None
    hf_logged_in: bool = False

    def ensure_hf_login(self) -> None:
        if self.hf_logged_in:
            return
        _best_effort_login_hf()
        self.hf_logged_in = True

    def unload_tribe(self) -> None:
        self.tribe_model = None
        gc.collect()
        _cuda_cleanup()

    def unload_openvla(self) -> None:
        self.openvla_model = None
        self.openvla_processor = None
        gc.collect()
        _cuda_cleanup()

    def unload_world_models(self) -> None:
        self.depth_model = None
        self.depth_processor = None
        self.dino_model = None
        self.dino_processor = None
        self.sam_pipeline = None
        gc.collect()
        _cuda_cleanup()


RUNTIME = _Runtime()


# --- Brain atlas / region masks (CPU cached) ---

SCORE_REGIONS: Dict[str, list[str]] = {
    "attention": [
        "G_front_sup",
        "G_front_middle",
        "S_front_sup",
        "G_parietal_sup",
        "G_pariet_inf-Supramar",
        "S_intrapariet_and_P_trans",
        "G_front_inf-Opercular",
    ],
    "emotion": [
        "G_insular_short",
        "S_circular_insula_ant",
        "G_cingul-Post-dorsal",
        "G_cingul-Post-ventral",
        "G_temp_sup-G_T_transv",
        "Pole_temporal",
        "G_front_inf-Triangul",
    ],
    "memory": [
        "G_oc-temp_med-Parahip",
        "G_temp_sup-Plan_tempo",
        "S_collat_transv_ant",
        "G_precuneus",
        "S_parieto_occipital",
        "G_oc-temp_med-Lingual",
    ],
    "reward": [
        "G_orbital",
        "S_orbital_lateral",
        "S_orbital-H_Shaped",
        "G_rectus",
        "S_suborbital",
        "G_subcallosal",
        "G_cingul-Post-ventral",
    ],
}


def _get_destrieux() -> Dict[str, Any]:
    from nilearn import datasets
    import numpy as np

    fsaverage5 = datasets.fetch_surf_fsaverage(mesh="fsaverage5")
    destrieux = datasets.fetch_atlas_surf_destrieux()

    labels_lh = destrieux["map_left"]
    labels_rh = destrieux["map_right"]
    label_names = [l.decode() if isinstance(l, bytes) else l for l in destrieux["labels"]]
    all_labels = np.concatenate([labels_lh, labels_rh])

    return {
        "fsaverage5": fsaverage5,
        "labels_lh": labels_lh,
        "labels_rh": labels_rh,
        "label_names": label_names,
        "all_labels": all_labels,
        "n_vertices_lh": len(labels_lh),
    }


_DESTRIEUX_CACHE: Optional[Dict[str, Any]] = None
_MASKS_CACHE: Optional[Dict[str, Any]] = None


def _ensure_masks() -> Tuple[Dict[str, Any], Dict[str, Any]]:
    global _DESTRIEUX_CACHE, _MASKS_CACHE
    if _DESTRIEUX_CACHE is None:
        _DESTRIEUX_CACHE = _get_destrieux()
    if _MASKS_CACHE is None:
        import numpy as np

        all_labels = _DESTRIEUX_CACHE["all_labels"]
        label_names = _DESTRIEUX_CACHE["label_names"]

        def region_mask(region_substrings: list[str]) -> "np.ndarray":
            mask = np.zeros(len(all_labels), dtype=bool)
            for rname in region_substrings:
                for idx, lname in enumerate(label_names):
                    if rname in lname:
                        mask |= all_labels == idx
            return mask

        _MASKS_CACHE = {k: region_mask(v) for k, v in SCORE_REGIONS.items()}
    return _DESTRIEUX_CACHE, _MASKS_CACHE


def _compute_scores(preds: "np.ndarray") -> Dict[str, int]:
    import numpy as np

    _, masks = _ensure_masks()

    preds = np.asarray(preds)
    if preds.ndim != 2 or preds.shape[0] == 0:
        raise ValueError("Invalid preds array")

    avg_activation = np.mean(preds, axis=0)
    peak_activation = np.max(preds, axis=0)

    raw: Dict[str, float] = {}
    for category, mask in masks.items():
        if int(mask.sum()) == 0:
            raw[category] = 0.0
            continue
        avg_score = float(np.mean(np.abs(avg_activation[mask])))
        peak_score = float(np.mean(np.abs(peak_activation[mask])))
        raw[category] = 0.4 * avg_score + 0.6 * peak_score

    whole_mean = float(np.mean(np.abs(avg_activation)))
    whole_std = float(np.std(np.abs(avg_activation)))

    normalized: Dict[str, int] = {}
    for cat, r in raw.items():
        z = (r - whole_mean) / whole_std if whole_std > 1e-8 else 0.0
        normalized[cat] = int(np.clip(50.0 + 25.0 * z, 0.0, 100.0))

    normalized["overall"] = int(
        0.30 * normalized["attention"]
        + 0.30 * normalized["emotion"]
        + 0.20 * normalized["memory"]
        + 0.20 * normalized["reward"]
    )
    return normalized


def _verdict(overall_score: int, threshold: int = 60) -> str:
    return "Brainrot" if int(overall_score) >= int(threshold) else "Not Brainrot"


def _render_brain(preds: "np.ndarray") -> "PIL.Image.Image":
    import numpy as np
    from PIL import Image

    from nilearn import plotting
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    destrieux, _ = _ensure_masks()
    fsaverage5 = destrieux["fsaverage5"]
    n_lh = int(destrieux["n_vertices_lh"])

    avg_pred = np.mean(np.asarray(preds), axis=0)
    lh = avg_pred[:n_lh]
    rh = avg_pred[n_lh:]

    fig = plt.figure(figsize=(10, 4))
    ax1 = fig.add_subplot(1, 2, 1, projection="3d")
    ax2 = fig.add_subplot(1, 2, 2, projection="3d")

    plotting.plot_surf_stat_map(
        fsaverage5["pial_left"],
        lh,
        hemi="left",
        view="lateral",
        bg_map=fsaverage5.get("sulc_left"),
        colorbar=False,
        axes=ax1,
        title="LH",
    )
    plotting.plot_surf_stat_map(
        fsaverage5["pial_right"],
        rh,
        hemi="right",
        view="lateral",
        bg_map=fsaverage5.get("sulc_right"),
        colorbar=False,
        axes=ax2,
        title="RH",
    )
    fig.tight_layout()

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        fig.savefig(tmp.name, dpi=160, bbox_inches="tight")
        tmp_path = tmp.name
    plt.close(fig)

    img = Image.open(tmp_path).convert("RGB")
    try:
        os.remove(tmp_path)
    except Exception:
        pass
    return img


def _format_brainrot_markdown(scores: Dict[str, int], threshold: int = 60) -> Tuple[str, str]:
    overall = int(scores.get("overall", 0))
    label = _verdict(overall, threshold=threshold)
    headline = f"## {label}\n\n**Overall**: `{overall}/100`  (threshold: `{threshold}`)"

    details = (
        "### Region Breakdown\n"
        + "\n".join(
            [
                f"- **{k.capitalize()}**: `{int(scores.get(k, 0))}/100`"
                for k in ["attention", "emotion", "memory", "reward"]
            ]
        )
        + "\n\n"
        "### Interpretation (Demo Layer)\n"
        "- This score is computed from TRIBE v2 cortical predictions using a simple region-weighted heuristic.\n"
        "- It is *not* an official TRIBE metric and should not be used for medical or high-stakes decisions.\n"
    )
    return headline, details


def _tribe_duration(*_args, **_kwargs) -> int:
    # Cold-starts (weights + WhisperX) can take minutes.
    return 420 if RUNTIME.tribe_model is None else 180


@spaces.GPU(duration=_tribe_duration)
def tribe_brainrot_from_video(video_path: str, trim_to_s: int = 10, threshold: int = 60):
    import numpy as np

    # Keep only one heavy stack alive to reduce GPU/CPU pressure.
    RUNTIME.unload_openvla()
    RUNTIME.unload_world_models()

    RUNTIME.ensure_hf_login()

    video_path = _as_filepath(video_path)
    if not video_path:
        return "## Error", None, [], "Upload a video first."

    try:
        from tribev2.demo_utils import TribeModel

        if RUNTIME.tribe_model is None:
            RUNTIME.tribe_model = TribeModel.from_pretrained(TRIBEV2_MODEL_ID, cache_folder=str(CACHE_DIR))

        vp = _trim_video_ffmpeg(video_path, max_duration_s=float(trim_to_s))
        df = RUNTIME.tribe_model.get_events_dataframe(video_path=vp)
        preds, _ = RUNTIME.tribe_model.predict(events=df)
        preds = np.asarray(preds)

        scores = _compute_scores(preds)
        brain_img = _render_brain(preds)
        headline, details = _format_brainrot_markdown(scores, threshold=int(threshold))

        table = [[k, int(scores[k])] for k in ["attention", "emotion", "memory", "reward", "overall"]]
        return headline, brain_img, table, details
    except Exception as e:
        return "## Error", None, [], f"```text\n{type(e).__name__}: {e}\n```"
    finally:
        _cuda_cleanup()


@spaces.GPU(duration=_tribe_duration)
def tribe_brainrot_from_text(text: str, threshold: int = 60):
    import numpy as np

    RUNTIME.unload_openvla()
    RUNTIME.unload_world_models()
    RUNTIME.ensure_hf_login()

    if not text or not text.strip():
        return "## Error", None, [], "Paste some text first."

    try:
        from tribev2.demo_utils import TribeModel

        if RUNTIME.tribe_model is None:
            RUNTIME.tribe_model = TribeModel.from_pretrained(TRIBEV2_MODEL_ID, cache_folder=str(CACHE_DIR))

        text_path = CACHE_DIR / "input_text.txt"
        text_path.write_text(text.strip(), encoding="utf-8")

        df = RUNTIME.tribe_model.get_events_dataframe(text_path=str(text_path))
        preds, _ = RUNTIME.tribe_model.predict(events=df)
        preds = np.asarray(preds)

        scores = _compute_scores(preds)
        brain_img = _render_brain(preds)
        headline, details = _format_brainrot_markdown(scores, threshold=int(threshold))

        table = [[k, int(scores[k])] for k in ["attention", "emotion", "memory", "reward", "overall"]]
        return headline, brain_img, table, details
    except Exception as e:
        hint = ""
        if os.environ.get("HF_TOKEN") is None:
            hint = (
                "\n\n**Hint:** `HF_TOKEN` is not set. TRIBE text mode may require gated model access (LLaMA family). "
                "Set `HF_TOKEN` in Space secrets if you have access."
            )
        return "## Error", None, [], f"```text\n{type(e).__name__}: {e}\n```{hint}"
    finally:
        _cuda_cleanup()


@spaces.GPU(duration=_tribe_duration)
def tribe_brainrot_from_image(image_path: str, threshold: int = 60):
    import numpy as np

    RUNTIME.unload_openvla()
    RUNTIME.unload_world_models()
    RUNTIME.ensure_hf_login()

    if not image_path:
        return "## Error", None, [], "Upload an image first."

    try:
        from tribev2.demo_utils import TribeModel

        if RUNTIME.tribe_model is None:
            RUNTIME.tribe_model = TribeModel.from_pretrained(TRIBEV2_MODEL_ID, cache_folder=str(CACHE_DIR))

        vp = _image_to_short_video(image_path)
        df = RUNTIME.tribe_model.get_events_dataframe(video_path=vp)
        preds, _ = RUNTIME.tribe_model.predict(events=df)
        preds = np.asarray(preds)

        scores = _compute_scores(preds)
        brain_img = _render_brain(preds)
        headline, details = _format_brainrot_markdown(scores, threshold=int(threshold))

        table = [[k, int(scores[k])] for k in ["attention", "emotion", "memory", "reward", "overall"]]
        return headline, brain_img, table, details
    except Exception as e:
        return "## Error", None, [], f"```text\n{type(e).__name__}: {e}\n```"
    finally:
        _cuda_cleanup()


def _world_duration(*_args, **_kwargs) -> int:
    cold = RUNTIME.depth_model is None or RUNTIME.dino_model is None
    return 180 if cold else 90


def _ensure_depth_model():
    from transformers import AutoImageProcessor, AutoModelForDepthEstimation
    import torch

    if RUNTIME.depth_model is None or RUNTIME.depth_processor is None:
        RUNTIME.depth_processor = AutoImageProcessor.from_pretrained(DEPTH_MODEL_ID)
        RUNTIME.depth_model = AutoModelForDepthEstimation.from_pretrained(
            DEPTH_MODEL_ID,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        )
        if torch.cuda.is_available():
            RUNTIME.depth_model = RUNTIME.depth_model.to("cuda")
        RUNTIME.depth_model.eval()
    return RUNTIME.depth_processor, RUNTIME.depth_model


def _ensure_dino_model():
    from transformers import AutoImageProcessor, AutoModel
    import torch

    if RUNTIME.dino_model is None or RUNTIME.dino_processor is None:
        RUNTIME.dino_processor = AutoImageProcessor.from_pretrained(DINO_MODEL_ID)
        RUNTIME.dino_model = AutoModel.from_pretrained(
            DINO_MODEL_ID,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        )
        if torch.cuda.is_available():
            RUNTIME.dino_model = RUNTIME.dino_model.to("cuda")
        RUNTIME.dino_model.eval()
    return RUNTIME.dino_processor, RUNTIME.dino_model


def _depth_to_image(depth: "np.ndarray") -> "PIL.Image.Image":
    import numpy as np
    from PIL import Image
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    d = np.asarray(depth, dtype="float32")
    d = d - float(np.nanmin(d))
    denom = float(np.nanmax(d)) + 1e-8
    d = d / denom
    rgba = plt.get_cmap("magma")(d)
    rgb = (rgba[:, :, :3] * 255.0).astype("uint8")
    return Image.fromarray(rgb)


def _image_edge_density(image: "PIL.Image.Image") -> float:
    import cv2
    import numpy as np

    arr = np.array(image.convert("RGB"))
    gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 80, 160)
    return float((edges > 0).mean())


def _infer_scene_label(depth_std: float, edge_density: float, motion_score: float) -> str:
    if depth_std > 0.24 and edge_density > 0.10:
        return "cluttered / high-affordance workspace"
    if motion_score > 18.0:
        return "dynamic scene / moving camera or objects"
    if depth_std < 0.12 and edge_density < 0.06:
        return "flat or low-structure scene"
    return "structured workspace"


def _cosine(a: list[float], b: list[float]) -> float:
    import numpy as np

    aa = np.asarray(a, dtype="float32")
    bb = np.asarray(b, dtype="float32")
    denom = float(np.linalg.norm(aa) * np.linalg.norm(bb)) + 1e-8
    return float(np.dot(aa, bb) / denom)


def _nearest_cached_embeddings(vector: list[float], current_hash: str, current_ts: float) -> list[list[Any]]:
    rows = []
    for item in WORLD_EMBED_CACHE[-128:]:
        if item.get("file_hash") == current_hash and abs(float(item.get("ts", 0.0)) - current_ts) < 1e-3:
            continue
        rows.append(
            [
                item.get("label", "cached frame"),
                item.get("file_hash", ""),
                round(float(item.get("ts", 0.0)), 2),
                round(_cosine(vector, item["embedding"]), 4),
            ]
        )
    rows.sort(key=lambda row: row[3], reverse=True)
    return rows[:5]


@spaces.GPU(duration=_world_duration)
def fast_world_state(video_path: str, ts_s: float):
    import numpy as np
    import torch
    import torch.nn.functional as F

    RUNTIME.unload_tribe()
    RUNTIME.unload_openvla()
    RUNTIME.ensure_hf_login()

    video_path = _as_filepath(video_path)
    if not video_path:
        return None, None, [], [], "Upload a video first."

    try:
        file_id = _file_hash(video_path)
        frame = _extract_frame(video_path, float(ts_s or 0.0)).convert("RGB")
        duration, _fps, _n = _probe_video(video_path)
        scout_key = f"{file_id}:8"
        motion_score = 0.0
        if scout_key in WORLD_SCOUT_CACHE:
            motion_score = float(WORLD_SCOUT_CACHE[scout_key].get("motion_score", 0.0))

        depth_processor, depth_model = _ensure_depth_model()
        dino_processor, dino_model = _ensure_dino_model()

        device = "cuda" if torch.cuda.is_available() else "cpu"

        depth_inputs = depth_processor(images=frame, return_tensors="pt").to(device)
        with torch.inference_mode():
            depth_outputs = depth_model(**depth_inputs)
            predicted_depth = depth_outputs.predicted_depth
            prediction = F.interpolate(
                predicted_depth.unsqueeze(1),
                size=frame.size[::-1],
                mode="bicubic",
                align_corners=False,
            ).squeeze()
        depth = prediction.float().detach().cpu().numpy()
        depth_img = _depth_to_image(depth)

        dino_inputs = dino_processor(images=frame, return_tensors="pt").to(device)
        with torch.inference_mode():
            dino_outputs = dino_model(**dino_inputs)
        embedding = dino_outputs.last_hidden_state[:, 0].float().detach().cpu().numpy()[0]
        embedding = embedding / (np.linalg.norm(embedding) + 1e-8)
        embedding_list = embedding.astype("float32").tolist()

        depth_range = float(np.nanpercentile(depth, 95) - np.nanpercentile(depth, 5))
        depth_std = float(np.nanstd(depth) / (abs(float(np.nanmean(depth))) + 1e-8))
        edge_density = _image_edge_density(frame)
        scene_label = _infer_scene_label(depth_std, edge_density, motion_score)

        nearest = _nearest_cached_embeddings(embedding_list, file_id, float(ts_s or 0.0))
        WORLD_EMBED_CACHE.append(
            {
                "label": scene_label,
                "file_hash": file_id,
                "ts": float(ts_s or 0.0),
                "embedding": embedding_list,
            }
        )

        metrics = {
            "file_hash": file_id,
            "timestamp_s": round(float(ts_s or 0.0), 2),
            "duration_s": round(float(duration or 0.0), 2),
            "depth_model": DEPTH_MODEL_ID,
            "ssl_model": DINO_MODEL_ID,
            "depth_range_p95_p5": round(depth_range, 4),
            "depth_structure_score": round(min(depth_std * 100.0, 100.0), 2),
            "edge_density": round(edge_density, 4),
            "motion_score": round(motion_score, 2),
            "scene_label": scene_label,
            "embedding_dim": int(len(embedding_list)),
        }
        rows = [
            ["depth_structure", metrics["depth_structure_score"]],
            ["edge_density", metrics["edge_density"]],
            ["motion_score", metrics["motion_score"]],
            ["embedding_dim", metrics["embedding_dim"]],
        ]
        md = (
            f"### Fast World State\n"
            f"- **Scene label:** `{scene_label}`\n"
            f"- **Depth structure:** `{metrics['depth_structure_score']}/100`\n"
            f"- **Nearest cached SSL matches:** `{len(nearest)}`"
        )
        return depth_img, metrics, rows, nearest, md
    except Exception as e:
        return None, None, [], [], f"```text\n{type(e).__name__}: {e}\n```"
    finally:
        _cuda_cleanup()


@spaces.GPU(duration=180)
def deep_world_state(video_path: str, ts_s: float, run_sam: bool, run_cotracker: bool, run_map_anything: bool):
    RUNTIME.unload_tribe()
    RUNTIME.unload_openvla()

    video_path = _as_filepath(video_path)
    if not video_path:
        return None, "Upload a video first."

    frame = None
    try:
        frame = _extract_frame(video_path, float(ts_s or 0.0)).convert("RGB")
    except Exception:
        frame = None

    results: Dict[str, Any] = {
        "sam": {"enabled": bool(run_sam), "model": SAM_MODEL_ID, "status": "skipped"},
        "cotracker": {"enabled": bool(run_cotracker), "model": "facebook/cotracker3", "status": "skipped"},
        "map_anything": {"enabled": bool(run_map_anything), "model": MAP_ANYTHING_MODEL_ID, "status": "skipped"},
        "notes": [],
    }

    if run_sam and frame is not None:
        try:
            from transformers import pipeline
            import torch

            sam_device = 0 if torch.cuda.is_available() else -1

            if RUNTIME.sam_pipeline is None:
                try:
                    RUNTIME.sam_pipeline = pipeline("mask-generation", model=SAM_MODEL_ID, device=sam_device)
                except Exception:
                    RUNTIME.sam_pipeline = pipeline("mask-generation", model=SAM_FALLBACK_MODEL_ID, device=sam_device)
                    results["sam"]["model"] = SAM_FALLBACK_MODEL_ID
            masks = RUNTIME.sam_pipeline(frame)
            count = len(masks.get("masks", [])) if isinstance(masks, dict) else 0
            results["sam"].update({"status": "ok", "mask_count": int(count)})
        except Exception as e:
            results["sam"].update({"status": "failed", "error": f"{type(e).__name__}: {e}"})

    if run_cotracker:
        results["cotracker"].update(
            {
                "status": "dependency-gated",
                "reason": "CoTracker3 is listed as an opt-in deep path; install/runtime needs validation before enabling inside the live ZeroGPU app.",
                "link": "https://hf.co/facebook/cotracker3",
            }
        )

    if run_map_anything:
        results["map_anything"].update(
            {
                "status": "dependency-gated",
                "reason": "Map-Anything Apache is the preferred 3D path, but its custom runtime should be validated separately before default Space inclusion.",
                "link": "https://hf.co/facebook/map-anything-apache",
            }
        )

    md = (
        "### Deep World State\n"
        f"- **SAM:** `{results['sam']['status']}` via `{results['sam'].get('model')}`\n"
        f"- **CoTracker3:** `{results['cotracker']['status']}`\n"
        f"- **Map-Anything:** `{results['map_anything']['status']}`\n\n"
        "Heavy world-model paths stay opt-in so the default ZeroGPU app remains responsive."
    )
    return results, md


def _openvla_duration(*_args, **_kwargs) -> int:
    # OpenVLA 7B cold-start can be large; allow time for first-time downloads.
    return 420 if RUNTIME.openvla_model is None else 180


def _format_action(action_7d: list[float], normalized: bool) -> str:
    def fmt(xs):
        return ", ".join(f"{float(x):+.3f}" for x in xs)

    xyz = action_7d[0:3]
    rpy = action_7d[3:6]
    g = action_7d[6] if len(action_7d) > 6 else 0.0
    mode = "normalized" if normalized else "unnormalized"

    return (
        f"### Action Breakdown ({mode})\n"
        f"- **Δposition (x,y,z)**: `{fmt(xyz)}`\n"
        f"- **Δorientation (roll,pitch,yaw)**: `{fmt(rpy)}`\n"
        f"- **gripper**: `{float(g):+.3f}`\n\n"
        "### Safety / Validity Notes\n"
        "- This is a *policy proposal* for research/demo use. Do not execute on real hardware without calibration and safeguards.\n"
        "- OpenVLA zero-shot only applies to embodiments/domains represented in its pretraining mixture.\n"
    )


def _draw_affordance_overlay(frame: "PIL.Image.Image", action_7d: list[float], label: str = "OpenVLA proposal") -> "PIL.Image.Image":
    from PIL import ImageDraw, ImageFont

    img = frame.convert("RGB").copy()
    draw = ImageDraw.Draw(img)
    w, h = img.size
    cx, cy = w // 2, h // 2

    dx = float(action_7d[0]) if len(action_7d) > 0 else 0.0
    dy = float(action_7d[1]) if len(action_7d) > 1 else 0.0
    dz = float(action_7d[2]) if len(action_7d) > 2 else 0.0
    gripper = float(action_7d[6]) if len(action_7d) > 6 else 0.0

    scale = min(w, h) * 0.22
    ex = int(cx + max(-1.0, min(1.0, dx)) * scale)
    ey = int(cy - max(-1.0, min(1.0, dy)) * scale)

    draw.line((cx, cy, ex, ey), fill=(56, 189, 248), width=max(3, w // 180))
    arrow_angle = math.atan2(ey - cy, ex - cx)
    head = max(10, min(w, h) // 25)
    for delta in (2.55, -2.55):
        hx = int(ex - head * math.cos(arrow_angle + delta))
        hy = int(ey - head * math.sin(arrow_angle + delta))
        draw.line((ex, ey, hx, hy), fill=(56, 189, 248), width=max(3, w // 180))

    radius = max(7, min(w, h) // 55)
    color = (34, 197, 94) if gripper > 0 else (244, 63, 94)
    draw.ellipse((ex - radius, ey - radius, ex + radius, ey + radius), fill=color)

    z_text = "toward camera" if dz > 0 else "away / down"
    g_text = "open" if gripper > 0 else "close"
    box = (12, 12, min(w - 12, 520), 86)
    draw.rounded_rectangle(box, radius=8, fill=(8, 13, 24), outline=(56, 189, 248), width=1)
    text = f"{label}\nxy arrow, z={z_text}, gripper={g_text}"
    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 16)
    except Exception:
        font = None
    draw.multiline_text((24, 22), text, fill=(240, 249, 255), font=font, spacing=4)
    return img


def _ensure_openvla_model():
    from transformers import AutoModelForVision2Seq, AutoProcessor
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("OpenVLA 7B requires a ZeroGPU CUDA allocation; run this inside the Space GPU event.")

    if RUNTIME.openvla_model is None or RUNTIME.openvla_processor is None:
        RUNTIME.openvla_processor = AutoProcessor.from_pretrained(
            OPENVLA_MODEL_ID, trust_remote_code=True
        )
        RUNTIME.openvla_model = AutoModelForVision2Seq.from_pretrained(
            OPENVLA_MODEL_ID,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        ).to("cuda")
        RUNTIME.openvla_model.eval()
    return RUNTIME.openvla_processor, RUNTIME.openvla_model


def _predict_openvla_action(frame: "PIL.Image.Image", instruction: str, output_mode: str) -> Tuple[list[float], bool]:
    import numpy as np
    import torch

    processor, model = _ensure_openvla_model()
    prompt = f"In: What action should the robot take to {instruction.strip()}?\nOut:"
    inputs = processor(prompt, frame, return_tensors="pt")
    inputs = {k: v.to("cuda") for k, v in inputs.items()}

    predict_kwargs: Dict[str, Any] = {"do_sample": False}
    normalized = True
    if output_mode == "bridge_orig":
        predict_kwargs["unnorm_key"] = "bridge_orig"
        normalized = False

    with torch.inference_mode():
        action = model.predict_action(**inputs, **predict_kwargs)

    if hasattr(action, "detach"):
        action = action.detach().float().cpu().numpy()
    action = np.asarray(action).reshape(-1).astype("float32").tolist()
    return action, normalized


@spaces.GPU(duration=_openvla_duration)
def openvla_action_from_video(video_path: str, instruction: str, ts_s: float, output_mode: str):
    from PIL import Image

    RUNTIME.unload_tribe()
    RUNTIME.unload_world_models()
    RUNTIME.ensure_hf_login()

    video_path = _as_filepath(video_path)
    if not video_path:
        return None, None, "Upload a robot clip first."
    if not instruction or not instruction.strip():
        return None, None, "Provide an instruction first."

    output_mode = (output_mode or DEFAULT_ROBOT_OUTPUT_MODE or "normalized").strip().lower()

    try:
        frame: Image.Image = _extract_frame(video_path, float(ts_s or 0.0))
        action, normalized = _predict_openvla_action(frame, instruction, output_mode)
        overlay = _draw_affordance_overlay(frame, action + [0.0] * max(0, 7 - len(action)))

        if len(action) != 7:
            return (
                overlay,
                {"action": action, "note": "Unexpected action shape; expected length 7."},
                _format_action(action + [0.0] * max(0, 7 - len(action)), normalized=normalized),
            )

        return overlay, {"action_7d": action, "output_mode": output_mode}, _format_action(action, normalized=normalized)
    except Exception as e:
        return None, None, f"```text\n{type(e).__name__}: {e}\n```"
    finally:
        _cuda_cleanup()


@spaces.GPU(duration=_openvla_duration)
def openvla_robustness_from_video(video_path: str, instruction: str, output_mode: str):
    import numpy as np

    RUNTIME.unload_tribe()
    RUNTIME.unload_world_models()
    RUNTIME.ensure_hf_login()

    video_path = _as_filepath(video_path)
    if not video_path:
        return None, [], "Upload a robot clip first."
    if not instruction or not instruction.strip():
        return None, [], "Provide an instruction first."

    output_mode = (output_mode or DEFAULT_ROBOT_OUTPUT_MODE or "normalized").strip().lower()
    try:
        duration, _fps, _n = _probe_video(video_path)
        if duration <= 0:
            duration = 6.0
        timestamps = [max(0.0, duration * r) for r in (0.25, 0.5, 0.75)]
        prompt_variants = [
            instruction.strip(),
            f"Carefully execute the robot task: {instruction.strip()}",
        ]

        rows = []
        actions = []
        for ts in timestamps:
            frame = _extract_frame(video_path, ts).convert("RGB")
            for prompt_idx, prompt in enumerate(prompt_variants, start=1):
                action, normalized = _predict_openvla_action(frame, prompt, output_mode)
                padded = action + [0.0] * max(0, 7 - len(action))
                rows.append(
                    [
                        round(float(ts), 2),
                        prompt_idx,
                        *[round(float(v), 4) for v in padded[:7]],
                    ]
                )
                actions.append(padded[:7])

        arr = np.asarray(actions, dtype="float32")
        mean_action = arr.mean(axis=0).tolist()
        std_action = arr.std(axis=0).tolist()
        uncertainty = float(np.linalg.norm(arr.std(axis=0)))
        stability = float(max(0.0, 100.0 - min(100.0, uncertainty * 100.0)))
        payload = {
            "output_mode": output_mode,
            "samples": int(len(actions)),
            "mean_action_7d": [round(float(v), 5) for v in mean_action],
            "std_action_7d": [round(float(v), 5) for v in std_action],
            "uncertainty_norm": round(uncertainty, 5),
            "stability_score_0_100": round(stability, 2),
        }
        md = (
            "### Counterfactual VLA Debugger\n"
            f"- **Samples:** `{len(actions)}` across 3 timestamps and 2 prompt variants\n"
            f"- **Action stability:** `{stability:.1f}/100`\n"
            f"- **Uncertainty norm:** `{uncertainty:.4f}`"
        )
        return payload, rows, md
    except Exception as e:
        return None, [], f"```text\n{type(e).__name__}: {e}\n```"
    finally:
        _cuda_cleanup()


def ui_download_sample() -> str:
    try:
        dst = CACHE_DIR / "sample_sintel.mp4"
        _download_url(SAMPLE_VIDEO_URL, dst)
        return str(dst)
    except Exception as e:
        raise gr.Error(f"Failed to download sample video: {e}")


def ui_robot_video_changed(video_path: str):
    video_path = _as_filepath(video_path)
    if not video_path:
        return (
            gr.update(maximum=10.0, value=0.0),
            None,
            "Upload a robot clip (mp4/mkv/avi). Then select a timestamp to extract a frame.",
        )
    duration, _fps, _n = _probe_video(video_path)
    if duration <= 0:
        duration = 10.0
    default_ts = min(duration * 0.5, 2.0)
    try:
        frame = _extract_frame(video_path, default_ts)
    except Exception:
        frame = None
    return (
        gr.update(maximum=float(max(0.1, duration)), value=float(default_ts)),
        frame,
        f"Detected duration: `{duration:.2f}s`",
    )


def ui_robot_ts_changed(video_path: str, ts_s: float):
    video_path = _as_filepath(video_path)
    if not video_path:
        return None
    try:
        return _extract_frame(video_path, float(ts_s or 0.0))
    except Exception:
        return None


def _coerce_json(value: Any) -> Dict[str, Any]:
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
            return parsed if isinstance(parsed, dict) else {}
        except Exception:
            return {}
    return {}


def _overall_from_region_df(region_df: Any) -> float:
    try:
        if hasattr(region_df, "to_dict"):
            records = region_df.to_dict("records")
            for row in records:
                if str(row.get("region", "")).lower() == "overall":
                    return float(row.get("score_0_100", 0.0))
        if isinstance(region_df, list):
            for row in region_df:
                if len(row) >= 2 and str(row[0]).lower() == "overall":
                    return float(row[1])
    except Exception:
        pass
    return 0.0


def brain_robot_fusion(region_df: Any, scout_json: Any, world_json: Any, action_json: Any, robustness_json: Any):
    scout = _coerce_json(scout_json)
    world = _coerce_json(world_json)
    action = _coerce_json(action_json)
    robustness = _coerce_json(robustness_json)

    brain = _overall_from_region_df(region_df)
    depth_structure = float(world.get("depth_structure_score", 0.0) or 0.0)
    motion = float((world or scout).get("motion_score", scout.get("motion_score", 0.0)) or 0.0)
    stability = float(robustness.get("stability_score_0_100", 0.0) or 0.0)
    action_present = 100.0 if action.get("action_7d") else 0.0

    actionability = (
        0.28 * min(depth_structure, 100.0)
        + 0.24 * min(stability, 100.0)
        + 0.18 * min(motion * 3.0, 100.0)
        + 0.18 * min(brain, 100.0)
        + 0.12 * action_present
    )
    payload = {
        "brain_score": round(brain, 2),
        "depth_structure": round(depth_structure, 2),
        "motion_signal": round(motion, 2),
        "vla_stability": round(stability, 2),
        "has_action": bool(action.get("action_7d")),
        "fusion_actionability_0_100": round(float(actionability), 2),
        "interpretation": "high" if actionability >= 70 else "medium" if actionability >= 45 else "low",
    }
    md = (
        "### Brain / Robot Fusion\n"
        f"- **Actionability:** `{payload['fusion_actionability_0_100']}/100` (`{payload['interpretation']}`)\n"
        f"- **Brain score:** `{payload['brain_score']}/100`\n"
        f"- **Depth structure:** `{payload['depth_structure']}/100`\n"
        f"- **VLA stability:** `{payload['vla_stability']}/100`\n\n"
        "This fusion layer is a demo heuristic that combines cortical engagement, geometry, motion, and VLA consistency."
    )
    return payload, md


CSS = """
:root {
  --bg0: #0b1220;
  --bg1: #0f172a;
  --card: rgba(255,255,255,0.06);
  --stroke: rgba(255,255,255,0.10);
  --text: rgba(255,255,255,0.92);
  --muted: rgba(255,255,255,0.70);
  --accent: #38bdf8;
  --accent2: #a78bfa;
  --warn: #fbbf24;
}

body { background: radial-gradient(1200px 600px at 20% 10%, rgba(56,189,248,0.18), transparent),
                 radial-gradient(1000px 500px at 70% 0%, rgba(167,139,250,0.16), transparent),
                 linear-gradient(180deg, var(--bg0), var(--bg1)) !important; }

#hero {
  border: 1px solid var(--stroke);
  background: linear-gradient(135deg, rgba(56,189,248,0.08), rgba(167,139,250,0.08));
  border-radius: 16px;
  padding: 18px 18px;
}
.hero-title {
  font-size: 22px;
  font-weight: 700;
  letter-spacing: 0.02em;
  color: var(--text);
}
.hero-sub {
  color: var(--muted);
  margin-top: 6px;
}
.pill {
  display: inline-block;
  font-size: 12px;
  padding: 3px 10px;
  border-radius: 999px;
  border: 1px solid var(--stroke);
  background: rgba(255,255,255,0.05);
  color: var(--muted);
  margin-right: 6px;
  margin-top: 10px;
}
.notice {
  border: 1px solid rgba(251,191,36,0.35);
  background: rgba(251,191,36,0.10);
  padding: 10px 12px;
  border-radius: 12px;
  color: var(--text);
}
.card {
  border: 1px solid var(--stroke);
  background: var(--card);
  border-radius: 16px;
  padding: 14px;
}
"""


with gr.Blocks(css=CSS, title="Brain x Robot x World Model Lab") as demo:
    gr.HTML(
        """
        <div id="hero">
          <div class="hero-title">Brain x Robot x World Model Lab</div>
          <div class="hero-sub">ZeroGPU demo: TRIBE v2 cortical predictions, SSL world-state probes, and OpenVLA action proposals</div>
          <div>
            <span class="pill">World Scout: CPU fast path</span>
            <span class="pill">ZeroGPU: H200 on-demand</span>
            <span class="pill">TRIBE v2 (CC BY-NC)</span>
            <span class="pill">OpenVLA 7B (MIT)</span>
          </div>
        </div>
        """
    )
    gr.Markdown(
        """
<div class="notice">
<strong>Cold start warning:</strong> first run may take minutes (model downloads, transcription tooling).
This Space tries to be robust on ZeroGPU, but timeouts and queues can happen.
</div>
        """
    )

    with gr.Tabs():
        with gr.Tab("World Scout (SSL)"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes=["card"]):
                    gr.Markdown(
                        """
**Fast path:** keyframes, scene cuts, motion energy, depth, and DINOv2 SSL signatures before expensive TRIBE/OpenVLA calls.
                        """
                    )
                    world_video = gr.Video(label="World / robot clip")
                    sample_world_btn = gr.Button("Load sample video (Sintel)", variant="secondary")
                    scout_max_frames = gr.Slider(
                        minimum=3,
                        maximum=16,
                        value=8,
                        step=1,
                        label="Scout keyframes",
                    )
                    run_scout = gr.Button("Run World Scout", variant="primary")
                    world_ts = gr.Slider(
                        minimum=0.0,
                        maximum=10.0,
                        value=0.0,
                        step=0.05,
                        label="Fast State timestamp (seconds)",
                    )
                    run_fast_world = gr.Button("Run Fast World State", variant="primary")
                    with gr.Accordion("Deep World State", open=False):
                        run_sam = gr.Checkbox(value=False, label="Try SAM3/SAM2 segmentation")
                        run_cotracker = gr.Checkbox(value=False, label="Report CoTracker3 path")
                        run_map = gr.Checkbox(value=False, label="Report Map-Anything path")
                        run_deep_world = gr.Button("Run Deep World State", variant="secondary")
                    with gr.Accordion("Model Cards / License Notes", open=False):
                        gr.Markdown(
                            """
- Fast defaults: [Depth Anything V2 Small](https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf) and [DINOv2 small](https://huggingface.co/facebook/dinov2-small).
- Optional deep paths: [SAM3](https://huggingface.co/facebook/sam3), [SAM2 tiny](https://huggingface.co/facebook/sam2-hiera-tiny), [CoTracker3](https://huggingface.co/facebook/cotracker3), and [Map-Anything Apache](https://huggingface.co/facebook/map-anything-apache).
- Linked world-generation experiments: [HunyuanWorld-1](https://huggingface.co/tencent/HunyuanWorld-1), [HY-World 2.0](https://huggingface.co/tencent/HY-World-2.0), and [HY-WorldPlay](https://huggingface.co/tencent/HY-WorldPlay). These are not in the default ZeroGPU request path.
- License warning: Meta and Tencent research models may have non-commercial or custom terms. Verify upstream model cards before commercial use.
                            """
                        )

                with gr.Column(scale=2, elem_classes=["card"]):
                    world_probe_md = gr.Markdown()
                    world_frame_preview = gr.Image(label="Selected frame", type="pil")
                    scout_json = gr.JSON(label="Scout metadata")
                    scout_gallery = gr.Gallery(label="Keyframes", columns=4, height=360)
                    keyframe_df = gr.Dataframe(
                        headers=["timestamp_s", "frame_idx", "resolution"],
                        datatype=["number", "number", "str"],
                        label="Keyframe table",
                    )
                    scene_cut_df = gr.Dataframe(
                        headers=["timestamp_s", "diff_score"],
                        datatype=["number", "number"],
                        label="Coarse scene cuts",
                    )
                    scout_md = gr.Markdown()

                    depth_img = gr.Image(label="Depth Anything V2 Small map", type="pil")
                    world_json = gr.JSON(label="Fast world-state metrics")
                    world_metrics_df = gr.Dataframe(
                        headers=["metric", "value"],
                        datatype=["str", "number"],
                        label="World metrics",
                    )
                    ssl_neighbors_df = gr.Dataframe(
                        headers=["cached_label", "file_hash", "timestamp_s", "cosine"],
                        datatype=["str", "str", "number", "number"],
                        label="DINOv2 nearest cached frames",
                    )
                    world_md = gr.Markdown()
                    deep_json = gr.JSON(label="Deep world-state status")
                    deep_md = gr.Markdown()

            sample_world_btn.click(fn=ui_download_sample, inputs=[], outputs=[world_video])
            world_video.change(
                fn=ui_robot_video_changed,
                inputs=[world_video],
                outputs=[world_ts, world_frame_preview, world_probe_md],
                show_progress="minimal",
            )
            world_ts.change(
                fn=ui_robot_ts_changed,
                inputs=[world_video, world_ts],
                outputs=[world_frame_preview],
                show_progress="minimal",
            )
            run_scout.click(
                fn=world_scout_video,
                inputs=[world_video, scout_max_frames],
                outputs=[scout_json, scout_gallery, keyframe_df, scene_cut_df, scout_md],
                show_progress="minimal",
            )
            run_fast_world.click(
                fn=fast_world_state,
                inputs=[world_video, world_ts],
                outputs=[depth_img, world_json, world_metrics_df, ssl_neighbors_df, world_md],
                show_progress="full",
            )
            run_deep_world.click(
                fn=deep_world_state,
                inputs=[world_video, world_ts, run_sam, run_cotracker, run_map],
                outputs=[deep_json, deep_md],
                show_progress="full",
            )

        with gr.Tab("Brainrot Lab (TRIBE v2)"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes=["card"]):
                    gr.Markdown(
                        """
**Inputs:** text, image, or video.  
We run `facebook/tribev2` and then compute a region-weighted engagement score.
                        """
                    )
                    threshold = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=60,
                        step=1,
                        label="Brainrot Threshold",
                    )

                    with gr.Accordion("Limitations", open=False):
                        gr.Markdown(
                            """
- The **Brainrot** label is a demo-layer heuristic, not an official TRIBE metric.
- TRIBE **text** mode may require gated model access; set `HF_TOKEN` in Space secrets.
- This is not medical advice.
                            """
                        )

                with gr.Column(scale=2, elem_classes=["card"]):
                    verdict_md = gr.Markdown()
                    brain_img = gr.Image(label="Cortical Activation (static surface plot)", type="pil")
                    region_df = gr.Dataframe(
                        headers=["region", "score_0_100"],
                        datatype=["str", "number"],
                        row_count=(5, "fixed"),
                        col_count=(2, "fixed"),
                        label="Region Scores",
                    )
                    details_md = gr.Markdown()

            with gr.Tabs():
                with gr.Tab("Video"):
                    video_in = gr.Video(label="Upload a video")
                    with gr.Row():
                        sample_btn = gr.Button("Load sample video (Sintel)", variant="secondary")
                        trim_to = gr.Slider(minimum=3, maximum=20, value=10, step=1, label="Trim to first N seconds")
                        run_video = gr.Button("Run TRIBE v2", variant="primary")

                    sample_btn.click(fn=ui_download_sample, inputs=[], outputs=[video_in])
                    run_video.click(
                        fn=tribe_brainrot_from_video,
                        inputs=[video_in, trim_to, threshold],
                        outputs=[verdict_md, brain_img, region_df, details_md],
                        show_progress="full",
                    )

                with gr.Tab("Image"):
                    image_in = gr.Image(label="Upload an image", type="filepath")
                    run_img = gr.Button("Run TRIBE v2", variant="primary")
                    run_img.click(
                        fn=tribe_brainrot_from_image,
                        inputs=[image_in, threshold],
                        outputs=[verdict_md, brain_img, region_df, details_md],
                        show_progress="full",
                    )

                with gr.Tab("Text"):
                    text_in = gr.Textbox(
                        label="Text",
                        placeholder="Paste a tweet, post, script, or caption...",
                        lines=6,
                    )
                    run_txt = gr.Button("Run TRIBE v2", variant="primary")
                    run_txt.click(
                        fn=tribe_brainrot_from_text,
                        inputs=[text_in, threshold],
                        outputs=[verdict_md, brain_img, region_df, details_md],
                        show_progress="full",
                    )

        with gr.Tab("Robot Action Lab (OpenVLA)"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes=["card"]):
                    gr.Markdown(
                        """
**Input:** a video clip of a robot workspace + a task instruction.  
We extract one frame at a chosen timestamp, then run `openvla/openvla-7b`.
                        """
                    )
                    robot_video = gr.Video(label="Robot clip")
                    sample_btn2 = gr.Button("Load sample video (Sintel)", variant="secondary")
                    ts_slider = gr.Slider(
                        minimum=0.0,
                        maximum=10.0,
                        value=0.0,
                        step=0.05,
                        label="Timestamp (seconds)",
                    )
                    probe_md = gr.Markdown()

                    instruction = gr.Textbox(
                        label="Instruction",
                        placeholder="e.g. pick up the block and place it in the bin",
                        lines=2,
                    )
                    output_mode = gr.Dropdown(
                        choices=["normalized", "bridge_orig"],
                        value="normalized" if DEFAULT_ROBOT_OUTPUT_MODE not in ("bridge_orig",) else "bridge_orig",
                        label="Output mode",
                        info="Default returns normalized actions. bridge_orig is an opt-in unnormalization for Bridge/WidowX.",
                    )
                    run_openvla = gr.Button("Run OpenVLA 7B", variant="primary")
                    run_robustness = gr.Button("Run VLA Robustness Debugger", variant="secondary")

                    with gr.Accordion("Limitations", open=False):
                        gr.Markdown(
                            """
- Output is a **proposal**, not a safe controller.
- OpenVLA does not zero-shot generalize to unseen embodiments outside its training mixture.
- `bridge_orig` unnormalization is only meaningful for the BridgeV2/WidowX domain.
                            """
                        )

                with gr.Column(scale=2, elem_classes=["card"]):
                    frame_preview = gr.Image(label="Extracted frame", type="pil")
                    affordance_img = gr.Image(label="Robot affordance overlay", type="pil")
                    action_json = gr.JSON(label="Predicted action")
                    action_md = gr.Markdown()
                    robust_json = gr.JSON(label="VLA robustness summary")
                    robust_df = gr.Dataframe(
                        headers=["timestamp_s", "prompt_variant", "dx", "dy", "dz", "droll", "dpitch", "dyaw", "gripper"],
                        datatype=["number", "number", "number", "number", "number", "number", "number", "number", "number"],
                        label="Counterfactual actions",
                    )
                    robust_md = gr.Markdown()

            sample_btn2.click(fn=ui_download_sample, inputs=[], outputs=[robot_video])
            robot_video.change(
                fn=ui_robot_video_changed,
                inputs=[robot_video],
                outputs=[ts_slider, frame_preview, probe_md],
                show_progress="minimal",
            )
            ts_slider.change(
                fn=ui_robot_ts_changed,
                inputs=[robot_video, ts_slider],
                outputs=[frame_preview],
                show_progress="minimal",
            )
            run_openvla.click(
                fn=openvla_action_from_video,
                inputs=[robot_video, instruction, ts_slider, output_mode],
                outputs=[affordance_img, action_json, action_md],
                show_progress="full",
            )
            run_robustness.click(
                fn=openvla_robustness_from_video,
                inputs=[robot_video, instruction, output_mode],
                outputs=[robust_json, robust_df, robust_md],
                show_progress="full",
            )

        with gr.Tab("Fusion Report"):
            with gr.Row():
                with gr.Column(scale=1, elem_classes=["card"]):
                    gr.Markdown(
                        """
Combine the latest TRIBE score, World Scout/Fast State metrics, and VLA robustness into a single actionability report.
                        """
                    )
                    run_fusion = gr.Button("Fuse Latest Results", variant="primary")
                with gr.Column(scale=2, elem_classes=["card"]):
                    fusion_json = gr.JSON(label="Fusion payload")
                    fusion_md = gr.Markdown()

            run_fusion.click(
                fn=brain_robot_fusion,
                inputs=[region_df, scout_json, world_json, action_json, robust_json],
                outputs=[fusion_json, fusion_md],
                show_progress="minimal",
            )

    gr.Markdown(
        """
---
	**References**
	- TRIBE v2: https://huggingface.co/facebook/tribev2  | https://github.com/facebookresearch/tribev2
	- OpenVLA: https://huggingface.co/openvla/openvla-7b | arXiv:2406.09246
	- Fast world state: https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf | https://huggingface.co/facebook/dinov2-small
	- Deep/experimental world models: SAM3/SAM2, CoTracker3, Map-Anything, VGGT, HY-World/HunyuanWorld. Validate upstream licenses before commercial use.
	        """
	    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(ssr_mode=False)