#!/usr/bin/env python3
# ADOBE CONFIDENTIAL
# Copyright 2026 Adobe
# All Rights Reserved.
# NOTICE: All information contained herein is, and remains
# the property of Adobe and its suppliers, if any. The intellectual
# and technical concepts contained herein are proprietary to Adobe
# and its suppliers and are protected by all applicable intellectual
# property laws, including trade secret and copyright laws.
# Dissemination of this information or reproduction of this material
# is strictly forbidden unless prior written permission is obtained
# from Adobe.

"""TPIPS embedding demo as a Gradio ZeroGPU app.

ZeroGPU only supports the Gradio SDK, and it exposes a GPU *only* for the
duration of a function wrapped with ``@spaces.GPU`` (allocated on demand, freed
right after).  That rules out the Flask + SSE + persistent-worker design of the
non-ZeroGPU demo, so this file rebuilds the same two experiences natively:

* **Image pair** - one reference image compared against any number of candidate
  images across an editable list of aspects, shown as grouped similarity bars.
* **Video** - cosine similarity of every sampled frame against a chosen
  reference frame, over an editable list of aspects, shown as a line plot.

Per-session isolation is handled by ``gr.State`` (Gradio keeps one copy per
browser session automatically).  All GPU compute funnels through
``embed_images`` so ZeroGPU only holds the accelerator while embeddings run.
"""

from __future__ import annotations

import asyncio
import base64
import html as _html
import io
import json
import math
import os
import uuid
from typing import Any

# Gradio's queue creates its asyncio locks via ``asyncio.get_event_loop()`` at
# construction time.  On Python 3.12+ that call raises ``RuntimeError`` when no
# loop is running in the main thread, so Gradio silently stores ``None`` locks
# and every queued request 500s with "NoneType ... asynchronous context
# manager".  Ensure a loop exists up front; no-op on older Pythons and on the
# Space's default interpreter.
try:
    asyncio.get_event_loop()
except RuntimeError:
    asyncio.set_event_loop(asyncio.new_event_loop())

import gradio as gr
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from PIL import Image

# Compatibility shim for a gradio_client bug: its JSON-schema-to-type converter
# assumes every subschema is a dict, but newer pydantic emits boolean subschemas
# (e.g. ``"additionalProperties": false``) for common components (Image, Files,
# Gallery, Video).  Without this, ``get_api_info()`` raises "argument of type
# 'bool' is not a container" and the root page 500s on load.  Guarded so it is a
# no-op if a future gradio_client renames the internal or fixes the bug.
try:
    import gradio_client.utils as _gc_utils

    _orig_schema_to_type = _gc_utils._json_schema_to_python_type

    def _schema_to_type(schema, defs=None):
        if isinstance(schema, bool):
            return "Any" if schema else "None"
        return _orig_schema_to_type(schema, defs)

    _gc_utils._json_schema_to_python_type = _schema_to_type
except Exception:  # pragma: no cover - defensive: never block startup on the shim
    pass

# ``spaces`` provides the ZeroGPU allocator on the Space.  Keep a no-op fallback
# so the same file still runs locally (e.g. on CPU) where it is not installed.
try:
    import spaces

    gpu = spaces.GPU
except Exception:  # pragma: no cover - only exercised off-Space

    def gpu(*args, **kwargs):
        if len(args) == 1 and callable(args[0]) and not kwargs:
            return args[0]

        def decorator(function):
            return function

        return decorator


from tpips import load_model

DEFAULT_FACTOR = "overall"
# Chips offered in the UI; users may also type any custom aspect.
PRESET_FACTORS = [DEFAULT_FACTOR, "lighting", "color palette", "composition", "camera angle"]
# Similarity axis range shared by the bar plot and the video line plot.
BAR_MIN, BAR_MAX = 0.4, 1.0
MAX_VIDEO_SECONDS = float(os.environ.get("TPIPS_MAX_VIDEO_SECONDS", "10"))
# ZeroGPU allocates the accelerator per call with a bounded duration, so cap the
# number of sampled frames instead of decoding at the native frame rate.
MAX_VIDEO_FRAMES = int(os.environ.get("TPIPS_MAX_VIDEO_FRAMES", "60"))
# Images per forward pass. Small keeps peak GPU memory bounded (raise on big GPUs
# like ZeroGPU's H200 for speed; lower if you still OOM on a small card).
EMBED_BATCH = int(os.environ.get("TPIPS_EMBED_BATCH", "8"))
GPU_SECONDS_PER_BATCH = int(os.environ.get("TPIPS_GPU_SECONDS_PER_BATCH", "10"))

MODEL_PATH = os.environ.get("TPIPS_MODEL_PATH", "sywang/TPIPS-Embed-Qwen3VL-8B")
CACHE_DIR = os.environ.get("HF_HOME")
# Keep the active checkpoint on local disk so ZeroGPU can pack its memory-mapped
# tensors quickly. CACHE_DIR still routes the separately loaded backbone through
# the persistent Hugging Face cache mounted at /data.
LOCAL_CHECKPOINT_CACHE = "/home/user/.cache/huggingface"
# On ZeroGPU the ``spaces`` runtime patches CUDA so a cuda placement here is
# resolved inside the GPU worker; set TPIPS_DEVICE=cpu for local CPU testing.
DEVICE = os.environ.get("TPIPS_DEVICE", "cuda")

print(f"Loading TPIPS embedding model '{MODEL_PATH}' for device '{DEVICE}' ...", flush=True)
checkpoint_path = snapshot_download(
    repo_id=MODEL_PATH,
    cache_dir=LOCAL_CHECKPOINT_CACHE,
)
model = load_model("embedding", checkpoint_path, device=DEVICE, cache_dir=CACHE_DIR)


def gpu_duration(images: list[Image.Image], factor: str) -> int:
    """Reserve ten GPU seconds per embedding batch, capped at 50 seconds."""
    del factor
    batches = max(1, math.ceil(len(images) / EMBED_BATCH))
    return min(50, GPU_SECONDS_PER_BATCH * batches)


def zero_gpu_identity_status(request: gr.Request) -> str:
    """Report whether the Hub forwarded a ZeroGPU request identity."""
    headers = {
        str(key).lower(): str(value)
        for key, value in dict(getattr(request, "headers", {})).items()
    }
    if headers.get("x-ip-token"):
        return (
            "✅ **ZeroGPU identity forwarding detected.** Hugging Face will apply "
            "the caller's account tier when scheduling GPU work."
        )
    return (
        "⚠️ **ZeroGPU identity forwarding not detected.** Sign in on Hugging Face "
        "and use the app through its public Space page to use account quota."
    )


@gpu(duration=gpu_duration)
def embed_images(images: list[Image.Image], factor: str) -> torch.Tensor:
    """Return L2-normalized embeddings for ``images`` under ``factor``.

    This is the only GPU entry point; ZeroGPU holds the accelerator for the
    length of this call and releases it on return.

    Images are embedded in small sub-batches so peak GPU memory stays bounded
    regardless of how many candidates/frames are scored at once (embedding all
    60 video frames in one forward pass OOMs an 8B model on a 24 GB card).
    """
    outputs: list[torch.Tensor] = []
    with torch.inference_mode():
        for start in range(0, len(images), EMBED_BATCH):
            chunk = images[start : start + EMBED_BATCH]
            result = model.embed(chunk, factor=factor, normalized=True)
            tensor = result.detach() if isinstance(result, torch.Tensor) else torch.as_tensor(result)
            outputs.append(F.normalize(tensor.float().cpu(), dim=-1))
    return torch.cat(outputs, dim=0)


def normalize_factors(factors: Any) -> list[str]:
    """Clean, de-duplicate, and cap the aspect list, falling back to 'overall'."""
    if isinstance(factors, str):
        factors = [factors]
    cleaned: list[str] = []
    for raw in factors or []:
        name = " ".join(str(raw).strip().split())[:120]
        if name and name not in cleaned:
            cleaned.append(name)
    return cleaned or [DEFAULT_FACTOR]


# --------------------------------------------------------------------------- #
# Image pair tab
#
# Faithful to the local demo: a reference plus any number of candidate cards in
# a scrollable row, each candidate showing one green->red gradient bar per aspect
# (the bar width maps the 0.4..1.0 similarity range).  Scores recompute
# automatically whenever the reference, candidates, or aspects change.
# --------------------------------------------------------------------------- #
def _bar_row(name: str, value: float | None) -> str:
    label = _html.escape(name)
    if value is None:
        right = "<span class='tp-status'>&mdash;</span>"
    else:
        pct = max(0.0, min(1.0, (value - BAR_MIN) / (BAR_MAX - BAR_MIN)))
        red = round(255 * (1 - pct))
        green = round(255 * pct)
        right = (
            f"<div class='tp-bar-bg'><div class='tp-bar-fill' "
            f"style='width:{pct * 100:.1f}%;background:rgb({red},{green},0)'></div></div>"
            f"<span class='tp-score-val'>{value:.3f}</span>"
        )
    return (
        f"<div class='tp-score-line'><span class='tp-score-name' title='{label}'>"
        f"{label}</span>{right}</div>"
    )


def candidate_bars_html(scores: dict[str, float], aspects: list[str]) -> str:
    if not aspects:
        return "<div class='tp-empty'>No aspects yet.</div>"
    rows = "".join(_bar_row(a, (scores or {}).get(a)) for a in aspects)
    return f"<div class='tp-cand-scores'>{rows}</div>"


def recompute_image_scores(
    ref_path: str | None,
    candidates: list[dict[str, Any]] | None,
    aspects: list[str] | None,
) -> list[dict[str, Any]]:
    """Return candidates with each one's per-aspect cosine similarity filled in."""
    candidates = [dict(c, scores={}) for c in (candidates or [])]
    if not ref_path or not candidates or not aspects:
        return candidates
    ref = Image.open(ref_path).convert("RGB")
    imgs = [Image.open(c["path"]).convert("RGB") for c in candidates]
    for factor in aspects:
        embeddings = embed_images([ref] + imgs, factor)
        sims = F.cosine_similarity(
            embeddings[0:1].expand_as(embeddings[1:]), embeddings[1:]
        ).tolist()
        for c, value in zip(candidates, sims):
            c["scores"][factor] = round(value, 4)
    return candidates


# --------------------------------------------------------------------------- #
# Video tab
#
# Faithful to the local demo: the GPU embeds every sampled frame once per aspect;
# from those (L2-normalized) embeddings we build the full N x N cosine matrix per
# aspect and ship it to the browser with small frame previews.  The interactive
# viewer (dual reference/candidate frames, play button, ref+frame sliders, a
# Chart.js curve with a moving scrubber, and live per-aspect bars with show/hide)
# then runs entirely client-side: moving the reference frame just indexes a row of
# the cached matrix, so it never touches the GPU (mirrors the local demo).
# --------------------------------------------------------------------------- #
def extract_video_frames(path: str) -> tuple[list[Image.Image], list[float], float]:
    """Evenly sample up to ``MAX_VIDEO_FRAMES`` frames within the time limit."""
    try:
        from moviepy import VideoFileClip
    except ImportError:  # moviepy 1.x compatibility
        from moviepy.editor import VideoFileClip

    clip = VideoFileClip(path)
    try:
        fps = float(clip.fps or 0.0)
        duration = float(clip.duration or 0.0)
        if fps <= 0 or duration <= 0:
            raise ValueError("video has no readable frames or duration")
        processed = min(duration, MAX_VIDEO_SECONDS)
        native = max(1, int(math.ceil(processed * fps - 1e-9)))
        count = min(native, MAX_VIDEO_FRAMES)
        # Sample within the processed window; keep true timestamps for the x axis.
        step = processed / count
        timestamps = [round(index * step, 3) for index in range(count)]
        frames = [
            Image.fromarray(clip.get_frame(min(ts, processed)).astype("uint8")).convert("RGB")
            for ts in timestamps
        ]
        return frames, timestamps, duration
    finally:
        clip.close()


def _frame_data_uri(image: Image.Image, max_dim: int = 360, quality: int = 72) -> str:
    """Small JPEG data URI so the whole frame strip travels in one JSON payload."""
    width, height = image.size
    scale = min(1.0, max_dim / max(width, height))
    if scale < 1.0:
        image = image.resize((max(1, round(width * scale)), max(1, round(height * scale))))
    buffer = io.BytesIO()
    image.save(buffer, format="JPEG", quality=quality)
    return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")


def _empty_video_state() -> dict[str, Any]:
    return {"loaded": False, "frames": [], "matrices": {}, "nonce": 0}


def build_video_payload(state: dict[str, Any] | None, aspects: list[str]) -> str:
    """Serialize the client-side viewer's data: previews + per-aspect NxN cosine."""
    if not state or not state.get("loaded"):
        return json.dumps({"loaded": False, "nonce": (state or {}).get("nonce", 0)})
    matrices = state.get("matrices", {})
    ordered = [a for a in aspects if a in matrices]
    state["nonce"] = int(state.get("nonce", 0)) + 1
    return json.dumps(
        {
            "loaded": True,
            "nonce": state["nonce"],
            "n_frames": len(state["frames"]),
            "fps": state["fps"],
            "processed_duration": state["processed"],
            "original_duration": state["duration"],
            "notification": state.get("notification", ""),
            "frames": state["previews"],
            "timestamps": state["timestamps"],
            "factors": ordered,
            "matrices": {a: matrices[a] for a in ordered},
        }
    )


def load_video_state(path: str | None):
    """Upload handler: extract frames (no GPU) and reset per-aspect embeddings."""
    if not path:
        state = _empty_video_state()
        return state, build_video_payload(state, [])
    frames, timestamps, duration = extract_video_frames(path)
    processed = min(duration, MAX_VIDEO_SECONDS)
    note = ""
    if duration > MAX_VIDEO_SECONDS + 1e-6:
        note = (
            f"⚠️ Only the first {MAX_VIDEO_SECONDS:g}s were analyzed "
            f"(processed {processed:.2f}s of {duration:.2f}s). "
        )
    state = {
        "loaded": True,
        "frames": frames,
        "previews": [_frame_data_uri(f) for f in frames],
        "timestamps": timestamps,
        "fps": len(frames) / max(processed, 1e-6),
        "processed": processed,
        "duration": duration,
        "notification": note.strip(),
        "matrices": {},
        "nonce": 0,
    }
    return state, build_video_payload(state, [])


def score_video_state(state: dict[str, Any] | None, aspects: Any):
    """Embed every frame once per aspect (GPU) and cache the full NxN cosine matrix.

    Reference-frame moves happen client-side by indexing a row of the matrix, so
    this only runs the model when a brand-new aspect is added.
    """
    if not state or not state.get("loaded"):
        state = state or _empty_video_state()
        return state, build_video_payload(state, [])
    aspects = normalize_factors(aspects)
    frames = state["frames"]
    matrices = state.setdefault("matrices", {})
    for aspect in aspects:
        if aspect not in matrices:
            embeddings = embed_images(frames, aspect)  # (N, D), L2-normalized
            cosine = (embeddings @ embeddings.T).clamp_(-1.0, 1.0)
            matrices[aspect] = [[round(v, 4) for v in row] for row in cosine.tolist()]
    return state, build_video_payload(state, aspects)


CUSTOM_CSS = """
.tp-compare-row { flex-wrap: nowrap !important; overflow-x: auto !important; gap: 14px; padding-bottom: 10px; align-items: flex-start !important; }
.tp-card { min-width: 300px !important; max-width: 340px !important; flex: 0 0 auto !important; }
.tp-head { font-weight: 600; margin: 2px 0 6px; font-size: 0.95rem; }
.tp-cand-scores { display: flex; flex-direction: column; gap: 8px; margin-top: 8px; }
.tp-score-line { display: flex; align-items: center; gap: 8px; }
.tp-score-name { flex: 1; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 0.86rem; opacity: 0.8; }
.tp-bar-bg { flex: 1.6; height: 12px; background: rgba(128,128,128,0.25); border-radius: 3px; overflow: hidden; }
.tp-bar-fill { height: 100%; width: 0; border-radius: 3px; transition: width .3s ease, background .3s ease; }
.tp-score-val { width: 48px; text-align: right; font-variant-numeric: tabular-nums; font-size: 0.86rem; }
.tp-status { opacity: 0.5; }
.tp-empty { opacity: 0.6; font-size: 0.85rem; padding: 8px 0; }
/* aspects sidebar */
.tp-aspect-add { align-items: center !important; gap: 6px !important; }
.tp-chip-row { flex-wrap: wrap !important; gap: 6px !important; }
.tp-chip { flex: 0 0 auto !important; min-width: 0 !important; }
.tp-chip button, button.tp-chip { padding: 3px 12px !important; font-size: 0.78rem !important; border-radius: 14px !important; min-width: 0 !important; }
.tp-aspect-row { align-items: center !important; gap: 8px !important; margin-bottom: 4px !important; flex-wrap: nowrap !important; }
.tp-aspect-name { flex: 1; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 0.9rem; padding: 4px 2px; }
.tp-x { flex: 0 0 auto !important; }
.tp-x button, button.tp-x { min-width: 34px !important; padding: 2px 6px !important; }

/* video viewer (client-side, mirrors demo.html) */
.tpv-main { display: flex; flex-direction: row; gap: 18px; }
.tpv-left { flex: 1; min-width: 0; }
.tpv-right { width: 260px; flex-shrink: 0; }
.tpv-frames-row { display: flex; gap: 14px; }
.tpv-frame-box { flex: 1; min-width: 0; }
.tpv-frame-box h4 { margin: 0 0 6px; font-size: 0.8rem; text-transform: uppercase; letter-spacing: .05em; opacity: .7; }
.tpv-slot { position: relative; height: 300px; border: 1px solid rgba(128,128,128,0.3); border-radius: 8px; overflow: hidden; background: rgba(128,128,128,0.08); display: flex; align-items: center; justify-content: center; }
.tpv-slot img { position: absolute; inset: 0; width: 100%; height: 100%; object-fit: contain; }
.tpv-empty-slot { color: rgba(128,128,128,0.9); font-size: 0.85rem; padding: 12px; text-align: center; }
.tpv-frame-time { z-index: 2; position: absolute; right: 8px; bottom: 8px; padding: 2px 7px; border-radius: 5px; background: rgba(15,28,24,0.85); color: #fff; font-family: monospace; font-size: 0.72rem; }
.tpv-play { z-index: 2; position: absolute; left: 10px; bottom: 10px; width: 42px; height: 42px; border-radius: 50%; border: none; cursor: pointer; background: #2563eb; color: #fff; display: flex; align-items: center; justify-content: center; box-shadow: 0 2px 10px rgba(0,0,0,.35); }
.tpv-play:hover { opacity: .9; }
.tpv-slider-row { display: flex; align-items: center; gap: 10px; margin: 8px 0 4px; }
.tpv-slider-row .tpv-lbl { width: 40px; flex-shrink: 0; opacity: .7; font-size: 0.8rem; }
.tpv-slider-row input[type=range] { flex: 1; cursor: pointer; }
.tpv-frame-label { font-family: monospace; font-size: 0.8rem; white-space: nowrap; min-width: 70px; text-align: right; opacity: .8; }
.tpv-frame-label.ref { color: #2563eb; opacity: 1; }
.tpv-meta { opacity: .65; font-size: 0.78rem; margin: 6px 0 4px; }
.tpv-chart-wrap { border: 1px solid rgba(128,128,128,0.3); border-radius: 8px; padding: 12px 14px; margin-top: 10px; }
.tpv-chart-wrap h4 { margin: 0 0 8px; font-size: 0.8rem; text-transform: uppercase; letter-spacing: .05em; opacity: .7; }
.tpv-chart-holder { position: relative; height: 250px; }
.tpv-factors-head { display: flex; align-items: center; justify-content: space-between; margin-bottom: 6px; }
.tpv-factors-head h4 { margin: 0; font-size: 0.8rem; text-transform: uppercase; letter-spacing: .05em; opacity: .7; }
.tpv-btn-sm { border: 1px solid rgba(128,128,128,0.4); background: transparent; color: inherit; border-radius: 6px; padding: 3px 8px; font-size: 0.7rem; cursor: pointer; margin-left: 6px; }
.tpv-btn-sm:hover { border-color: #2563eb; }
.tpv-hint { opacity: .6; font-size: 0.75rem; margin-bottom: 12px; }
.tpv-factor-item { border: 1px solid rgba(128,128,128,0.3); border-radius: 8px; padding: 10px 11px; margin-bottom: 8px; }
.tpv-factor-head { display: flex; justify-content: space-between; align-items: center; gap: 8px; margin-bottom: 6px; }
.tpv-fname-wrap { display: flex; align-items: center; gap: 6px; min-width: 0; flex: 1; }
.tpv-fname { font-size: 0.85rem; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
.tpv-eye { background: none; border: none; cursor: pointer; padding: 0; opacity: .75; display: flex; align-items: center; }
.tpv-eye.off { opacity: .3; }
.tpv-cosine-row { display: flex; align-items: center; gap: 8px; }
.tpv-bar-bg { flex: 1; height: 12px; background: rgba(128,128,128,0.25); border-radius: 3px; overflow: hidden; }
.tpv-bar-fill { height: 100%; width: 0; border-radius: 3px; transition: width .12s, background .12s; }
.tpv-score-val { width: 46px; text-align: right; font-variant-numeric: tabular-nums; font-size: 0.8rem; }
.tpv-empty { opacity: .6; font-size: 0.82rem; padding: 10px 0; }
@media (max-width: 900px) { .tpv-main { flex-direction: column; } .tpv-right { width: 100%; } }
"""


# Static skeleton for the video viewer. Rendered once as trusted HTML (gr.HTML
# does not sanitize developer content); every element carries a stable id so the
# controller in VIDEO_HEAD can drive it from the JSON payload. No inline scripts
# (gr.HTML would strip them) - all behaviour lives in VIDEO_HEAD.
VIDEO_VIEWER_HTML = """
<div class="tpv-main">
  <div class="tpv-left">
    <div class="tpv-frames-row">
      <div class="tpv-frame-box">
        <h4 id="tpv-ref-title">Reference frame</h4>
        <div class="tpv-slot">
          <div class="tpv-empty-slot" id="tpv-ref-empty">Upload a video and score an aspect to begin.</div>
          <img id="tpv-ref-img" alt="Reference frame" hidden>
          <span class="tpv-frame-time" id="tpv-ref-time" hidden>0.00s</span>
        </div>
        <div class="tpv-slider-row" id="tpv-ref-slider-row" hidden>
          <span class="tpv-lbl">Ref</span>
          <input type="range" id="tpv-ref-slider" min="0" max="0" step="1" value="0">
          <span class="tpv-frame-label ref" id="tpv-ref-label">0.00s</span>
        </div>
      </div>
      <div class="tpv-frame-box">
        <h4 id="tpv-cur-title">Candidate frame</h4>
        <div class="tpv-slot">
          <div class="tpv-empty-slot" id="tpv-cur-empty">The candidate frame plays here.</div>
          <img id="tpv-cur-img" alt="Candidate frame" hidden>
          <button class="tpv-play" id="tpv-play" title="Play / pause" hidden></button>
          <span class="tpv-frame-time" id="tpv-cur-time" hidden>0.00s</span>
        </div>
        <div class="tpv-slider-row" id="tpv-frame-slider-row" hidden>
          <span class="tpv-lbl">Frame</span>
          <input type="range" id="tpv-frame-slider" min="0" max="0" step="1" value="0">
          <span class="tpv-frame-label" id="tpv-frame-label">0.00s</span>
        </div>
      </div>
    </div>
    <div class="tpv-meta" id="tpv-meta" hidden></div>
    <div class="tpv-chart-wrap" id="tpv-chart-wrap" hidden>
      <h4>Cosine similarity over video time</h4>
      <div class="tpv-chart-holder"><canvas id="tpv-chart"></canvas></div>
    </div>
  </div>
  <div class="tpv-right">
    <div class="tpv-factors-head">
      <h4>Scored aspects</h4>
      <div><button class="tpv-btn-sm" id="tpv-show-all">Show all</button><button class="tpv-btn-sm" id="tpv-hide-all">Hide all</button></div>
    </div>
    <div class="tpv-hint">Bars track the candidate frame &mdash; drag the frame slider or hit play. Moving the reference re-reads cached scores (no model rerun).</div>
    <div id="tpv-factors"><div class="tpv-empty">Score an aspect to see per-aspect curves.</div></div>
  </div>
</div>
"""


# Chart.js + the client-side controller for the video viewer. Injected into
# <head>; ``window.tpvInit(payload)`` is invoked from the payload textbox's
# .change (see the Video tab wiring) each time the server ships new scores.
VIDEO_HEAD = """
<script src="https://cdn.jsdelivr.net/npm/chart.js@4"></script>
<script>
(function(){
  if (window.__tpvReady) return;
  window.__tpvReady = true;
  var COLORS=['#2563eb','#059669','#dc2626','#d97706','#7c3aed','#db2777','#0891b2','#65a30d','#ca8a04','#e11d48'];
  var color=function(i){return COLORS[i%COLORS.length];};
  var BAR_MIN=0.4, BAR_MAX=1.0;
  var barPct=function(v){return Math.max(0,Math.min(1,(v-BAR_MIN)/(BAR_MAX-BAR_MIN)));};
  var EYE_OPEN='<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M1 12s4-8 11-8 11 8 11 8-4 8-11 8-11-8-11-8z"/><circle cx="12" cy="12" r="3"/></svg>';
  var EYE_CLOSED='<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M17.94 17.94A10.07 10.07 0 0 1 12 20c-7 0-11-8-11-8a18.45 18.45 0 0 1 5.06-5.94M9.9 4.24A9.12 9.12 0 0 1 12 4c7 0 11 8 11 8a18.5 18.5 0 0 1-2.16 3.19m-6.72-1.07a3 3 0 1 1-4.24-4.24"/><line x1="1" y1="1" x2="23" y2="23"/></svg>';
  var ICON_PLAY='<svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor"><path d="M8 5v14l11-7z"/></svg>';
  var ICON_PAUSE='<svg width="16" height="16" viewBox="0 0 24 24" fill="currentColor"><rect x="6" y="5" width="4" height="14" rx="1"/><rect x="14" y="5" width="4" height="14" rx="1"/></svg>';
  var $=function(id){return document.getElementById(id);};
  function esc(s){var n=document.createElement('span');n.textContent=String(s);return n.innerHTML;}
  var VP=null, refIdx=0, curIdx=0, playing=false, raf=null, chart=null, wired=false;
  var hidden={};

  function clamp(){ if(!VP)return; var m=VP.n_frames-1; if(refIdx>m)refIdx=m; if(refIdx<0)refIdx=0; if(curIdx>m)curIdx=m; if(curIdx<0)curIdx=0; }
  function curveFor(f){ return (VP&&VP.matrices[f])?VP.matrices[f][refIdx]:null; }

  function showFrames(){
    if(!VP||!VP.frames.length)return; clamp();
    if($('tpv-ref-img'))$('tpv-ref-img').src=VP.frames[refIdx];
    if($('tpv-cur-img'))$('tpv-cur-img').src=VP.frames[curIdx];
    if($('tpv-ref-time'))$('tpv-ref-time').textContent=VP.timestamps[refIdx].toFixed(2)+'s';
    if($('tpv-cur-time'))$('tpv-cur-time').textContent=VP.timestamps[curIdx].toFixed(2)+'s';
    if($('tpv-ref-label'))$('tpv-ref-label').textContent=VP.timestamps[refIdx].toFixed(2)+'s';
    if($('tpv-frame-label'))$('tpv-frame-label').textContent=VP.timestamps[curIdx].toFixed(2)+'s';
    if($('tpv-ref-title'))$('tpv-ref-title').textContent='Reference frame \\u00b7 '+refIdx;
    if($('tpv-cur-title'))$('tpv-cur-title').textContent='Candidate frame \\u00b7 '+curIdx;
  }
  function updateBars(){
    if(!VP)return;
    var items=document.querySelectorAll('#tpv-factors .tpv-factor-item');
    for(var k=0;k<items.length;k++){
      var item=items[k], f=item.getAttribute('data-factor'), curve=curveFor(f);
      var bar=item.querySelector('.tpv-bar-fill'), val=item.querySelector('.tpv-score-val');
      if(!curve||curve[curIdx]===undefined)continue;
      var vv=curve[curIdx],pct=barPct(vv),r=Math.round(255*(1-pct)),g=Math.round(255*pct);
      if(bar){bar.style.width=(pct*100).toFixed(1)+'%';bar.style.background='rgb('+r+','+g+',0)';}
      if(val)val.textContent=vv.toFixed(3);
    }
  }
  function playIcon(){var b=$('tpv-play');if(b)b.innerHTML=playing?ICON_PAUSE:ICON_PLAY;}
  function stopPlay(){playing=false;if(raf)cancelAnimationFrame(raf);raf=null;playIcon();}
  function startPlay(){
    if(!VP||VP.n_frames<2)return; playing=true; playIcon();
    var max=VP.n_frames-1, fps=Math.max(1,VP.fps||12);
    var pos=curIdx>=max?0:curIdx, last=performance.now();
    var tick=function(now){
      if(!playing)return;
      pos+=(now-last)/1000*fps; last=now; if(pos>max)pos=0;
      curIdx=Math.min(max,Math.round(pos));
      var fs=$('tpv-frame-slider'); if(fs)fs.value=curIdx;
      showFrames(); updateBars();
      if(chart){chart._scrub=curIdx;chart.update('none');}
      raf=requestAnimationFrame(tick);
    };
    raf=requestAnimationFrame(tick);
  }
  function togglePlay(){ playing?stopPlay():startPlay(); }

  var scrubberPlugin={id:'tpvscrub',afterDraw:function(c){
    var idx=c._scrub; if(idx===null||idx===undefined)return; var pt=null;
    for(var i=0;i<c.data.datasets.length;i++){var m=c.getDatasetMeta(i);if(!m.hidden&&m.data[idx]){pt=m.data[idx];break;}}
    if(!pt)return; var ctx=c.ctx, top=c.chartArea.top, bottom=c.chartArea.bottom;
    ctx.save();ctx.strokeStyle='rgba(120,120,120,0.7)';ctx.lineWidth=1;ctx.setLineDash([4,3]);
    ctx.beginPath();ctx.moveTo(pt.x,top);ctx.lineTo(pt.x,bottom);ctx.stroke();ctx.restore();
  }};
  function buildChart(){
    if(!window.Chart)return; var cv=$('tpv-chart'); if(!cv)return;
    if(chart){chart.destroy();chart=null;}
    chart=new Chart(cv.getContext('2d'),{type:'line',data:{datasets:[]},
      options:{animation:false,responsive:true,maintainAspectRatio:false,interaction:{mode:'index',intersect:false},
        scales:{x:{type:'linear',min:0,title:{display:true,text:'Time (s)'},ticks:{callback:function(v){return (+v).toFixed(1);}}},
                y:{min:BAR_MIN,max:BAR_MAX,title:{display:true,text:'Cosine similarity'}}},
        plugins:{legend:{labels:{boxWidth:12,font:{size:11}},onClick:function(e,item,legend){
          Chart.defaults.plugins.legend.onClick(e,item,legend);
          var ci=legend.chart, di=ci.data.datasets.findIndex(function(d){return d.label===item.text;});
          if(di<0)return; hidden[item.text]=!ci.isDatasetVisible(di); renderFactors();
        }}}},
      plugins:[scrubberPlugin]});
    chart._scrub=curIdx;
  }
  function syncChart(){
    if(!chart||!VP)return;
    chart.options.scales.x.max=Math.max(VP.processed_duration,0.001);
    chart.data.datasets=VP.factors.map(function(f,i){
      var curve=curveFor(f); if(!curve)return null;
      return {label:f,data:curve.map(function(y,j){return {x:VP.timestamps[j],y:y};}),
        borderColor:color(i),backgroundColor:'transparent',borderWidth:1.5,pointRadius:0,tension:0.25,hidden:hidden[f]===true};
    }).filter(Boolean);
    chart._scrub=curIdx; chart.update('none');
  }
  function renderFactors(){
    var host=$('tpv-factors'); if(!host)return;
    if(!VP||!VP.factors.length){host.innerHTML='<div class="tpv-empty">Score an aspect to see per-aspect curves.</div>';return;}
    host.innerHTML=VP.factors.map(function(f,i){
      var c=color(i), vis=hidden[f]!==true;
      return '<div class="tpv-factor-item" data-factor="'+esc(f)+'">'+
        '<div class="tpv-factor-head"><span class="tpv-fname-wrap">'+
        '<button class="tpv-eye '+(vis?'':'off')+'" data-eye="'+esc(f)+'">'+(vis?EYE_OPEN:EYE_CLOSED)+'</button>'+
        '<span class="tpv-fname" style="color:'+c+'">'+esc(f)+'</span></span></div>'+
        '<div class="tpv-cosine-row"><div class="tpv-bar-bg"><div class="tpv-bar-fill"></div></div><span class="tpv-score-val">&mdash;</span></div></div>';
    }).join('');
    if(chart)syncChart();
    updateBars();
  }
  function setVis(loaded){
    ['tpv-ref-img','tpv-cur-img','tpv-ref-time','tpv-cur-time','tpv-play','tpv-ref-slider-row','tpv-frame-slider-row','tpv-meta','tpv-chart-wrap'].forEach(function(id){var el=$(id);if(el)el.hidden=!loaded;});
    ['tpv-ref-empty','tpv-cur-empty'].forEach(function(id){var el=$(id);if(el)el.style.display=loaded?'none':'';});
  }
  function wire(){
    if(wired)return; wired=true;
    document.addEventListener('click',function(e){
      var t=e.target; if(!t.closest)return;
      if(t.closest('#tpv-play')){togglePlay();return;}
      if(t.closest('#tpv-show-all')){if(VP)VP.factors.forEach(function(f){hidden[f]=false;});renderFactors();return;}
      if(t.closest('#tpv-hide-all')){if(VP)VP.factors.forEach(function(f){hidden[f]=true;});renderFactors();return;}
      var eye=t.closest('[data-eye]'); if(eye){var f=eye.getAttribute('data-eye');hidden[f]=hidden[f]!==true;renderFactors();return;}
    });
    document.addEventListener('input',function(e){
      if(e.target.id==='tpv-frame-slider'){stopPlay();curIdx=+e.target.value;showFrames();updateBars();if(chart){chart._scrub=curIdx;chart.update('none');}}
      else if(e.target.id==='tpv-ref-slider'){refIdx=+e.target.value;showFrames();syncChart();updateBars();}
    });
  }
  window.tpvInit=function(payloadStr){
    try{
      wire();
      var p; try{p=JSON.parse(payloadStr||'{}');}catch(_){p={};}
      if(!p||!p.loaded){VP=null;setVis(false);var h=$('tpv-factors');if(h)h.innerHTML='<div class="tpv-empty">Score an aspect to see per-aspect curves.</div>';return;}
      VP=p; setVis(true); clamp();
      var rs=$('tpv-ref-slider'), fs=$('tpv-frame-slider');
      if(rs){rs.max=VP.n_frames-1;rs.value=refIdx;}
      if(fs){fs.max=VP.n_frames-1;fs.value=curIdx;}
      if($('tpv-meta'))$('tpv-meta').textContent=VP.n_frames+' frames \\u00b7 '+VP.fps.toFixed(2)+' fps \\u00b7 analyzed '+VP.processed_duration.toFixed(2)+'s of '+VP.original_duration.toFixed(2)+'s'+(VP.notification?(' \\u00b7 '+VP.notification):'');
      if(!chart)buildChart();
      playIcon(); showFrames(); renderFactors(); syncChart(); updateBars();
    }catch(err){console.error('tpvInit',err);}
  };
  var tries=0;
  var waitChart=function(){ if(window.Chart){ if(VP&&!chart){buildChart();syncChart();} } else if(tries++<80){ setTimeout(waitChart,150); } };
  waitChart();
})();
</script>
"""


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
with gr.Blocks(
    title="TPIPS interactive demo", theme=gr.themes.Soft(), css=CUSTOM_CSS, head=VIDEO_HEAD
) as demo:
    gr.Markdown(
        "# TPIPS interactive demo\n"
        "Embedding-based perceptual similarity for image pairs and short videos. "
        "Pick or type any aspect to score against. "
        "More details on our [project page](https://peterwang512.github.io/TPIPS)."
    )
    identity_status = gr.Markdown()
    demo.load(zero_gpu_identity_status, outputs=identity_status)

    # Per-session image-tab state (Gradio keeps one copy per browser session).
    img_ref_state = gr.State(None)  # reference image filepath
    img_cands_state = gr.State([])  # list of {"id", "path", "scores": {aspect: value}}
    img_aspects_state = gr.State([DEFAULT_FACTOR])
    # Per-session video-tab state (frames, previews, cached per-aspect NxN cosine).
    video_state = gr.State(_empty_video_state())
    video_aspects_state = gr.State([DEFAULT_FACTOR])

    with gr.Tabs():
        with gr.Tab("Image Pair"):
            gr.Markdown(
                "Upload a **reference** and add any number of **candidates**. Each "
                "candidate shows a green→red similarity bar per aspect; bars update "
                "automatically as you add candidates or aspects."
            )
            with gr.Row():
                with gr.Column(scale=3):
                    # Reference + candidates + add tile all live inside one render so
                    # they are flex siblings of a single horizontally scrollable row.
                    @gr.render(inputs=[img_ref_state, img_cands_state, img_aspects_state])
                    def render_compare(ref_path, cands, aspects):
                        with gr.Row(elem_classes="tp-compare-row"):
                            with gr.Column(min_width=300, elem_classes="tp-card"):
                                gr.HTML("<div class='tp-head'>Reference Image</div>")
                                ref_img = gr.Image(
                                    value=ref_path,
                                    type="filepath",
                                    sources=["upload", "clipboard"],
                                    height=280,
                                    show_label=False,
                                )

                                def _set_reference(path, cs, asp):
                                    return path, recompute_image_scores(path, cs, asp)

                                ref_img.change(
                                    _set_reference,
                                    [ref_img, img_cands_state, img_aspects_state],
                                    [img_ref_state, img_cands_state],
                                )

                            for index, cand in enumerate(cands):
                                with gr.Column(min_width=300, elem_classes="tp-card"):
                                    gr.HTML(f"<div class='tp-head'>Candidate {index + 1}</div>")
                                    gr.Image(
                                        value=cand["path"],
                                        interactive=False,
                                        height=280,
                                        show_label=False,
                                    )
                                    gr.HTML(candidate_bars_html(cand.get("scores", {}), aspects))
                                    remove_btn = gr.Button("✕ Remove", size="sm")

                                    def _remove(cs, asp, ref, cid=cand["id"]):
                                        cs = [c for c in cs if c["id"] != cid]
                                        return recompute_image_scores(ref, cs, asp)

                                    remove_btn.click(
                                        _remove,
                                        [img_cands_state, img_aspects_state, img_ref_state],
                                        img_cands_state,
                                    )

                            with gr.Column(min_width=260, elem_classes="tp-card"):
                                gr.HTML("<div class='tp-head'>&nbsp;</div>")
                                add_btn = gr.UploadButton(
                                    "＋ Add candidate",
                                    file_count="multiple",
                                    file_types=["image"],
                                )

                                def _add(files, cs, asp, ref):
                                    paths = [getattr(f, "name", f) for f in (files or [])]
                                    cs = list(cs) + [
                                        {"id": uuid.uuid4().hex, "path": p, "scores": {}}
                                        for p in paths
                                    ]
                                    return recompute_image_scores(ref, cs, asp)

                                add_btn.upload(
                                    _add,
                                    [add_btn, img_cands_state, img_aspects_state, img_ref_state],
                                    img_cands_state,
                                )

                with gr.Column(scale=1, min_width=240):
                    gr.Markdown("### Comparison aspects")
                    with gr.Row(elem_classes="tp-aspect-add"):
                        aspect_box = gr.Textbox(
                            placeholder="e.g. material texture",
                            show_label=False,
                            container=False,
                            scale=4,
                            max_lines=1,
                        )
                        aspect_add = gr.Button("Add", scale=1, size="sm")
                    with gr.Row(elem_classes="tp-chip-row"):
                        chip_btns = [
                            gr.Button(p, size="sm", elem_classes="tp-chip") for p in PRESET_FACTORS
                        ]

                    @gr.render(inputs=[img_aspects_state])
                    def render_aspects(aspects):
                        if not aspects:
                            gr.HTML("<div class='tp-empty'>Add an aspect to compare.</div>")
                        for name in aspects:
                            with gr.Row(elem_classes="tp-aspect-row"):
                                gr.HTML(f"<div class='tp-aspect-name'>{_html.escape(name)}</div>")
                                remove_aspect = gr.Button(
                                    "✕", size="sm", scale=0, min_width=34, elem_classes="tp-x"
                                )

                                def _remove_aspect(asp, cs, ref, target=name):
                                    asp = [a for a in asp if a != target]
                                    return asp, recompute_image_scores(ref, cs, asp)

                                remove_aspect.click(
                                    _remove_aspect,
                                    [img_aspects_state, img_cands_state, img_ref_state],
                                    [img_aspects_state, img_cands_state],
                                )

                    def _add_aspect(text, asp, cs, ref):
                        name = " ".join(str(text or "").strip().split())[:120]
                        asp = list(asp)
                        if name and name not in asp:
                            asp = asp + [name]
                        return "", asp, recompute_image_scores(ref, cs, asp)

                    aspect_add.click(
                        _add_aspect,
                        [aspect_box, img_aspects_state, img_cands_state, img_ref_state],
                        [aspect_box, img_aspects_state, img_cands_state],
                    )
                    aspect_box.submit(
                        _add_aspect,
                        [aspect_box, img_aspects_state, img_cands_state, img_ref_state],
                        [aspect_box, img_aspects_state, img_cands_state],
                    )

                    for chip, preset in zip(chip_btns, PRESET_FACTORS):
                        def _add_chip(asp, cs, ref, name=preset):
                            asp = list(asp)
                            if name not in asp:
                                asp = asp + [name]
                            return asp, recompute_image_scores(ref, cs, asp)

                        chip.click(
                            _add_chip,
                            [img_aspects_state, img_cands_state, img_ref_state],
                            [img_aspects_state, img_cands_state],
                        )

        with gr.Tab("Video"):
            gr.Markdown(
                f"Upload a clip (first {MAX_VIDEO_SECONDS:g}s analyzed), pick a reference "
                "frame, then play the candidate — the curve shows how similar every frame "
                "is to the reference along each aspect, and the bars update live as it plays."
            )
            # Hidden channel: server writes the JSON payload here, and its .change
            # forwards it to the client-side controller (window.tpvInit). Declared
            # early because the aspect-remove buttons rendered below target it.
            video_payload = gr.Textbox(visible=False, elem_id="tpv-payload")
            with gr.Row():
                with gr.Column(scale=3):
                    video_input = gr.Video(
                        label=f"Video (first {MAX_VIDEO_SECONDS:g}s analyzed)", height=220
                    )
                with gr.Column(scale=1, min_width=240):
                    gr.Markdown("### Video aspects")
                    with gr.Row(elem_classes="tp-aspect-add"):
                        v_aspect_box = gr.Textbox(
                            placeholder="e.g. camera motion",
                            show_label=False,
                            container=False,
                            scale=4,
                            max_lines=1,
                        )
                        v_aspect_add = gr.Button("Add", scale=1, size="sm")
                    with gr.Row(elem_classes="tp-chip-row"):
                        v_chip_btns = [
                            gr.Button(p, size="sm", elem_classes="tp-chip") for p in PRESET_FACTORS
                        ]

                    @gr.render(inputs=[video_aspects_state])
                    def render_video_aspects(aspects):
                        if not aspects:
                            gr.HTML("<div class='tp-empty'>Add an aspect to embed frames.</div>")
                        for name in aspects:
                            with gr.Row(elem_classes="tp-aspect-row"):
                                gr.HTML(
                                    f"<div class='tp-aspect-name'>{_html.escape(name)}</div>"
                                )
                                v_remove = gr.Button(
                                    "✕", size="sm", scale=0, min_width=34, elem_classes="tp-x"
                                )

                                def _remove_v_aspect(asp, st, target=name):
                                    asp = [a for a in asp if a != target]
                                    st, payload = score_video_state(st, asp)
                                    return asp, st, payload

                                v_remove.click(
                                    _remove_v_aspect,
                                    [video_aspects_state, video_state],
                                    [video_aspects_state, video_state, video_payload],
                                )

            # The interactive viewer (static skeleton; behaviour is in VIDEO_HEAD).
            gr.HTML(VIDEO_VIEWER_HTML)

            video_payload.change(
                None, video_payload, None, js="(p) => { if (window.tpvInit) window.tpvInit(p); }"
            )

            # Upload -> extract frames (no GPU) -> auto-score the current aspects.
            video_input.change(
                load_video_state, [video_input], [video_state, video_payload]
            ).then(
                score_video_state,
                [video_state, video_aspects_state],
                [video_state, video_payload],
            )

            def _add_v_aspect(text, asp, st):
                name = " ".join(str(text or "").strip().split())[:120]
                asp = list(asp)
                if name and name not in asp:
                    asp = asp + [name]
                st, payload = score_video_state(st, asp)
                return "", asp, st, payload

            v_aspect_add.click(
                _add_v_aspect,
                [v_aspect_box, video_aspects_state, video_state],
                [v_aspect_box, video_aspects_state, video_state, video_payload],
            )
            v_aspect_box.submit(
                _add_v_aspect,
                [v_aspect_box, video_aspects_state, video_state],
                [v_aspect_box, video_aspects_state, video_state, video_payload],
            )

            for v_chip, v_preset in zip(v_chip_btns, PRESET_FACTORS):
                def _add_v_chip(asp, st, name=v_preset):
                    asp = list(asp)
                    if name not in asp:
                        asp = asp + [name]
                    st, payload = score_video_state(st, asp)
                    return asp, st, payload

                v_chip.click(
                    _add_v_chip,
                    [video_aspects_state, video_state],
                    [video_aspects_state, video_state, video_payload],
                )


if __name__ == "__main__":
    demo.queue(max_size=32).launch(
        # Bind all interfaces by default (required on the Space); override with
        # GRADIO_SERVER_NAME=127.0.0.1 for local testing on hosts where
        # ``localhost`` resolves only to IPv6 and Gradio's self-check would fail.
        server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.environ.get("PORT", 7860)),
    )