Xinyun-Liu's picture
Prefer H264 PAI playback derivatives over cached HEVC originals
8f0efb9 verified
Raw
History Blame Contribute Delete
12.9 kB
#!/usr/bin/env python3
"""Public review service for the WorldModelBench Hugging Face Space."""
from __future__ import annotations
import io
import json
import os
import threading
import uuid
from contextlib import closing
from datetime import datetime, timezone
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import parse_qs, urlparse
from urllib.request import Request, urlopen
from huggingface_hub import CommitOperationAdd, HfApi, get_token
ROOT = Path(__file__).resolve().parent
STATIC_ROOT = ROOT / "static"
MEDIA_CACHE_ROOT = ROOT / "media-cache"
CATALOG_PATH = STATIC_ROOT / "data" / "worldmodelbench-catalog.js"
DATASET_REPO = os.environ.get("WORLDMODELBENCH_REPO", "Xinyun-Liu/WorldModelBench")
HF_TOKEN = os.environ.get("HF_TOKEN") or get_token() or ""
MAX_BODY_BYTES = 32_000
MIN_SUBMISSION_INTERVAL = 0.8
last_submission_by_ip: dict[str, float] = {}
rate_lock = threading.Lock()
def read_catalog() -> dict[str, dict[str, str]]:
raw = CATALOG_PATH.read_text(encoding="utf-8")
payload = raw.split("=", 1)[1].strip().rstrip(";")
rows = json.loads(payload)
return {row["videoPath"]: row for row in rows if isinstance(row, dict) and row.get("videoPath")}
CATALOG = read_catalog()
def cached_media_path(video_path: str) -> Path | None:
"""Return a cache entry only when it stays beneath the cache root."""
cache_root = MEDIA_CACHE_ROOT.resolve()
try:
candidate = (cache_root / video_path).resolve()
candidate.relative_to(cache_root)
except ValueError:
return None
return candidate if candidate.is_file() else None
def utc_now() -> str:
return datetime.now(timezone.utc).isoformat(timespec="seconds")
def clean_string(value: object, field: str, maximum: int, required: bool = False) -> str:
if not isinstance(value, str):
raise ValueError(f"{field}_must_be_text")
value = value.strip()
if required and not value:
raise ValueError(f"{field}_is_required")
if len(value) > maximum:
raise ValueError(f"{field}_is_too_long")
return value
def validate_submission(payload: object) -> dict[str, str]:
if not isinstance(payload, dict):
raise ValueError("request_must_be_a_json_object")
annotator = clean_string(payload.get("annotator"), "annotator", 80, required=True)
video_path = clean_string(payload.get("video_path"), "video_path", 500, required=True)
row = CATALOG.get(video_path)
if row is None:
raise ValueError("video_path_is_not_in_the_catalog")
decision = clean_string(payload.get("decision", "pending"), "decision", 20, required=True)
if decision not in {"pending", "keep", "reject"}:
raise ValueError("invalid_decision")
return {
"annotator": annotator,
"video_path": video_path,
"sample_id": clean_string(row.get("id", ""), "sample_id", 300, required=True),
"data_source": clean_string(row.get("sourceDataset", ""), "data_source", 200, required=True),
"decision": decision,
"track": clean_string(payload.get("track", ""), "track", 200),
"scene": clean_string(payload.get("scene", ""), "scene", 300),
"spatial_ability": clean_string(payload.get("spatial_ability", ""), "spatial_ability", 500),
"perspective": clean_string(payload.get("perspective", ""), "perspective", 300),
"text_caption": clean_string(payload.get("text_caption", ""), "text_caption", 4000),
}
class ReviewHandler(SimpleHTTPRequestHandler):
server_version = "WorldModelBenchReview/1.0"
def __init__(self, *args, **kwargs):
super().__init__(*args, directory=str(STATIC_ROOT), **kwargs)
def end_headers(self) -> None:
if urlparse(self.path).path == "/review.html":
self.send_header("Cache-Control", "no-store")
super().end_headers()
def send_json(self, status: int, value: dict[str, object]) -> None:
body = json.dumps(value, ensure_ascii=False).encode("utf-8")
self.send_response(status)
self.send_header("Content-Type", "application/json; charset=utf-8")
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
self.wfile.write(body)
def proxy_asset(self, url: str, send_body: bool = True) -> None:
"""Relay an HF resolve URL as a same-origin, range-capable response."""
headers = {"User-Agent": "WorldModelBench-Review/1.0"}
requested_range = self.headers.get("Range")
if requested_range:
headers["Range"] = requested_range
try:
with closing(urlopen(Request(url, headers=headers), timeout=90)) as upstream:
self.send_response(upstream.status)
for name in ("Content-Type", "Content-Length", "Content-Range", "Accept-Ranges", "Last-Modified", "ETag"):
value = upstream.headers.get(name)
if value:
self.send_header(name, value)
self.send_header("Cache-Control", "public, max-age=300")
self.end_headers()
if send_body:
while chunk := upstream.read(1024 * 1024):
self.wfile.write(chunk)
except HTTPError as error:
self.send_json(error.code, {"error": "upstream_asset_unavailable"})
except (OSError, URLError, TimeoutError):
self.send_json(502, {"error": "upstream_asset_connection_failed"})
def serve_cached_media(self, asset: Path, send_body: bool = True) -> None:
"""Serve a locally cached original MP4 with browser Range support."""
total_size = asset.stat().st_size
start, end, partial = 0, total_size - 1, False
requested_range = self.headers.get("Range", "")
if requested_range:
try:
unit, spec = requested_range.split("=", 1)
if unit.strip() != "bytes" or "," in spec:
raise ValueError
first, last = spec.strip().split("-", 1)
if first:
start = int(first)
end = int(last) if last else total_size - 1
elif last:
suffix = int(last)
if suffix <= 0:
raise ValueError
start = max(total_size - suffix, 0)
else:
raise ValueError
if start < 0 or start >= total_size or end < start:
raise ValueError
end = min(end, total_size - 1)
partial = True
except (ValueError, TypeError):
self.send_response(416)
self.send_header("Content-Range", f"bytes */{total_size}")
self.send_header("Content-Length", "0")
self.end_headers()
return
length = end - start + 1
self.send_response(206 if partial else 200)
self.send_header("Content-Type", "video/mp4")
self.send_header("Content-Length", str(length))
self.send_header("Accept-Ranges", "bytes")
if partial:
self.send_header("Content-Range", f"bytes {start}-{end}/{total_size}")
self.send_header("Cache-Control", "public, max-age=300")
self.end_headers()
if not send_body:
return
try:
with asset.open("rb") as stream:
stream.seek(start)
remaining = length
while remaining:
chunk = stream.read(min(1024 * 1024, remaining))
if not chunk:
break
self.wfile.write(chunk)
remaining -= len(chunk)
except BrokenPipeError:
return
def proxy_catalog_asset(self, kind: str, query: dict[str, list[str]], send_body: bool = True) -> None:
paths = query.get("path", [])
if len(paths) != 1:
self.send_json(400, {"error": "exactly_one_video_path_is_required"})
return
# Reload the small catalog at request time so a catalog update never
# leaves the media proxy serving a stale playback URL.
row = read_catalog().get(paths[0])
if row is None:
self.send_json(404, {"error": "video_path_is_not_in_the_catalog"})
return
if kind == "media":
cached_asset = cached_media_path(paths[0])
# A cached source original must not override a catalog entry that
# deliberately points at a browser-compatible derivative.
cache_matches_catalog = isinstance(row.get("videoUrl"), str) and row["videoUrl"].endswith(f"/{paths[0]}")
if cached_asset is not None and cache_matches_catalog:
self.serve_cached_media(cached_asset, send_body=send_body)
return
url = row.get("videoUrl") if kind == "media" else row.get("initialFrameUrl")
if not isinstance(url, str) or not url:
self.send_json(404, {"error": "preview_is_not_available"})
return
self.proxy_asset(url, send_body=send_body)
def do_GET(self) -> None:
parsed = urlparse(self.path)
path = parsed.path
if path == "/api/health":
self.send_json(200 if HF_TOKEN else 503, {"ok": bool(HF_TOKEN), "catalog_videos": len(CATALOG), "dataset_repo": DATASET_REPO})
return
if path == "/media":
self.proxy_catalog_asset("media", parse_qs(parsed.query, keep_blank_values=True))
return
if path == "/preview":
self.proxy_catalog_asset("preview", parse_qs(parsed.query, keep_blank_values=True))
return
if path == "/":
self.path = "/review.html"
super().do_GET()
def do_HEAD(self) -> None:
parsed = urlparse(self.path)
if parsed.path == "/media":
self.proxy_catalog_asset("media", parse_qs(parsed.query, keep_blank_values=True), send_body=False)
return
if parsed.path == "/preview":
self.proxy_catalog_asset("preview", parse_qs(parsed.query, keep_blank_values=True), send_body=False)
return
if parsed.path == "/":
self.path = "/review.html"
super().do_HEAD()
def do_POST(self) -> None:
if urlparse(self.path).path != "/api/submissions":
self.send_json(404, {"error": "not_found"})
return
if not HF_TOKEN:
self.send_json(503, {"error": "submission_service_not_configured"})
return
client_ip = self.client_address[0]
now = datetime.now(timezone.utc).timestamp()
with rate_lock:
previous = last_submission_by_ip.get(client_ip, 0.0)
if now - previous < MIN_SUBMISSION_INTERVAL:
self.send_json(429, {"error": "please_wait_before_submitting_again"})
return
last_submission_by_ip[client_ip] = now
try:
content_length = int(self.headers.get("Content-Length", "0"))
if content_length <= 0 or content_length > MAX_BODY_BYTES:
raise ValueError("invalid_request_size")
payload = json.loads(self.rfile.read(content_length).decode("utf-8"))
record = validate_submission(payload)
submission_id = uuid.uuid4().hex
created_at = utc_now()
record.update({
"schema_version": "worldmodelbench-web-review-v1",
"submission_id": submission_id,
"submitted_at": created_at,
"source": "public_review_space",
})
path_in_repo = f"web-review/submissions/{created_at[:10]}/{submission_id}.json"
content = json.dumps(record, ensure_ascii=False, indent=2).encode("utf-8")
HfApi(token=HF_TOKEN).create_commit(
repo_id=DATASET_REPO,
repo_type="dataset",
operations=[CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(content))],
commit_message=f"Add web review annotation {submission_id[:12]}",
)
self.send_json(201, {"ok": True, "submission_id": submission_id, "submitted_at": created_at})
except ValueError as error:
self.send_json(400, {"error": str(error)})
except Exception:
self.send_json(502, {"error": "huggingface_submission_failed"})
def main() -> None:
port = int(os.environ.get("PORT", "7860"))
server = ThreadingHTTPServer(("0.0.0.0", port), ReviewHandler)
print(f"WorldModelBench review service listening on {port}", flush=True)
server.serve_forever()
if __name__ == "__main__":
main()