| |
| """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 |
| |
| |
| 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]) |
| |
| |
| 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() |
|
|