infer: accept HF Hub repo id for --adapter-dir

infer.py

@@ -0,0 +1,413 @@
+"""Run inference with the anonymous judge LoRA adapter.
+
+The script can either load files from a local directory or pull them
+directly from the Hugging Face Hub. By default it points at the
+companion repository ``anonymouscla/physground-judger9B``:
+
+  # From the Hub (no clone needed):
+  python infer.py --video demo.mp4 --caption "A ball rolls down a ramp." --metric SA
+  python infer.py --video demo.mp4 --caption "A ball rolls down a ramp." --law gravity
+
+  # From a local clone of the model repo:
+  python infer.py --adapter-dir /path/to/local/clone --video demo.mp4 \
+      --caption "A ball rolls down a ramp." --law gravity
+
+It loads:
+  - adapter_config.json to find the base model
+  - adapter_model.safetensors through PEFT
+  - subq+human.yaml to render the scoring prompt
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+import torch
+import yaml
+from peft import PeftModel
+from transformers import AutoProcessor
+
+
+GENERAL_SUB_QUESTIONS: dict[str, list[str]] = {
+    "SA": [
+        "Are the main objects in the caption present in the video?",
+        "Are the key actions or interactions from the caption visible?",
+        "Are important scene attributes and relationships preserved?",
+        "Does the video avoid major contradictions to the caption?",
+    ],
+    "PTV": [
+        "Do causes appear before their effects?",
+        "Do physical events unfold in a plausible temporal order?",
+        "Are motion transitions continuous rather than abrupt jumps or loops?",
+        "Does the sequence avoid impossible reversals or repeated resets?",
+    ],
+    "persistence": [
+        "Do objects maintain consistent existence throughout the video?",
+        "Do objects keep a stable shape, size, color, and texture?",
+        "Do objects avoid disappearing, appearing, or transforming unexpectedly?",
+        "Do objects preserve identity through motion and brief occlusion?",
+    ],
+}
+
+
+PHYSICAL_CRITERIA: dict[str, str] = {
+    "gravity": "Do unsupported objects fall downward? Do thrown objects follow a curved trajectory? Does poured liquid fall with gravity?",
+    "inertia": "Do stationary objects remain still unless acted upon? Do moving objects maintain their motion unless stopped by friction, collision, or an obstacle?",
+    "momentum": "After collision, push, or pull, is the direction of motion reasonable? Ignore speed magnitude.",
+    "impenetrability": "Do objects maintain impenetrability -- no passing through each other?",
+    "collision": "After impact, is there reasonable bounce/shatter/deformation? Does response match impact force?",
+    "material": "Does each material respond according to its properties? (glass shatters, rubber bounces, metal is rigid, cloth deforms softly, etc.)",
+    "buoyancy": "Do dense objects sink? Do wood/plastic float?",
+    "displacement": "When you add more liquid or put an object into it, does the liquid level rise in a realistic way? Does it overflow when full?",
+    "flow_dynamics": "Does the liquid's overall motion behave realistically over time -- flowing along surfaces, spreading, draining naturally?",
+    "boundary_interaction": "When the liquid hits a boundary such as a rock face, container wall, or floor, does it respond realistically? Do local splash, rebound, or split patterns on impact look physically plausible?",
+    "fluid_continuity": "Does the liquid avoid disappearing or appearing out of nowhere? Small splashes that briefly break apart are okay.",
+    "reflection": "Does the reflection roughly match objects and colors in the scene, and avoid completely unrelated content?",
+    "shadow": "Are shadow directions consistent with light source? Do shadows move with objects?",
+}
+
+
+PHYSICAL_SUB_QUESTIONS: dict[str, list[str]] = {
+    "gravity": [
+        "Do unsupported objects or liquids move downward over time?",
+        "Do thrown or falling objects follow a plausible gravity-driven path?",
+        "Does the video avoid objects floating or rising without support?",
+    ],
+    "inertia": [
+        "Do stationary objects remain still unless a visible force acts on them?",
+        "Do moving objects continue plausibly until friction, collision, or an obstacle changes their motion?",
+        "Does the video avoid unexplained starts, stops, or direction changes?",
+    ],
+    "momentum": [
+        "After contact, push, pull, or collision, are motion directions plausible?",
+        "Does the reacting object move in a direction consistent with the interaction?",
+        "Does the video avoid impossible reversals or unrelated motion changes?",
+    ],
+    "impenetrability": [
+        "Do solid objects avoid passing through one another?",
+        "Do contacts and overlaps remain physically plausible?",
+        "Does the video avoid obvious clipping or penetration artifacts?",
+    ],
+    "collision": [
+        "Does impact cause a plausible bounce, break, deformation, or transfer of motion?",
+        "Is the response direction consistent with the collision?",
+        "Does the response avoid being much too weak, too strong, or unrelated to the impact?",
+    ],
+    "material": [
+        "Do objects respond consistently with their apparent material?",
+        "Are rigid, soft, brittle, elastic, or fluid-like objects animated appropriately?",
+        "Does the video avoid material behavior that contradicts the scene?",
+    ],
+    "buoyancy": [
+        "Do objects sink or float in a way consistent with apparent density?",
+        "Does the floating or sinking behavior stay stable over time?",
+        "Does the video avoid unsupported hovering or impossible underwater motion?",
+    ],
+    "displacement": [
+        "Does liquid level rise when volume is added or an object enters it?",
+        "Does overflow happen only when the container is plausibly full?",
+        "Does the liquid volume remain visually plausible?",
+    ],
+    "flow_dynamics": [
+        "Does liquid flow along surfaces, spread, or drain naturally?",
+        "Does the flow direction follow gravity and boundaries?",
+        "Does the video avoid abrupt stops, reversals, or unsupported uphill flow?",
+    ],
+    "boundary_interaction": [
+        "Does liquid react plausibly when hitting a wall, floor, container, or obstacle?",
+        "Are splash, rebound, or split patterns locally plausible?",
+        "Does the liquid remain consistent after interacting with boundaries?",
+    ],
+    "fluid_continuity": [
+        "Does liquid avoid disappearing or appearing without cause?",
+        "Does the amount of liquid remain broadly consistent?",
+        "Are splashes and separations temporary and physically plausible?",
+    ],
+    "reflection": [
+        "Does the reflection match nearby objects, colors, and motion?",
+        "Does the reflected content stay spatially consistent with the scene?",
+        "Does the video avoid unrelated or impossible reflection content?",
+    ],
+    "shadow": [
+        "Are shadows consistent with the apparent light source direction?",
+        "Do shadows move with the objects that cast them?",
+        "Does the video avoid missing, detached, or contradictory shadows?",
+    ],
+}
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    with path.open() as f:
+        return json.load(f)
+
+
+def load_yaml(path: Path) -> dict[str, Any]:
+    with path.open() as f:
+        return yaml.safe_load(f)
+
+
+def questions_block(questions: list[str]) -> str:
+    return "\n".join(f"{idx}. {question}" for idx, question in enumerate(questions, 1))
+
+
+def build_prompt(
+    cfg: dict[str, Any],
+    caption: str,
+    *,
+    metric: str | None = None,
+    law: str | None = None,
+    criteria: str | None = None,
+) -> tuple[str, str, str]:
+    if metric:
+        if metric not in GENERAL_SUB_QUESTIONS:
+            raise ValueError(f"unknown metric: {metric}")
+        prompt = cfg["eval_prompts"][metric].format(
+            prompt=caption,
+            questions_block=questions_block(GENERAL_SUB_QUESTIONS[metric]),
+        )
+        return cfg["system_prompt"], prompt, metric
+
+    if not law:
+        raise ValueError("either --metric or --law is required")
+    if law not in PHYSICAL_CRITERIA:
+        raise ValueError(f"unknown law: {law}")
+    prompt = cfg["physical_template"].format(
+        prompt=caption,
+        law=law,
+        criteria=criteria or PHYSICAL_CRITERIA[law],
+        questions_block=questions_block(PHYSICAL_SUB_QUESTIONS[law]),
+    )
+    return cfg["system_prompt"], prompt, law
+
+
+def load_base_model(base_id: str, dtype: torch.dtype, device_map: str):
+    errors: list[str] = []
+    for class_name in (
+        "AutoModelForImageTextToText",
+        "AutoModelForVision2Seq",
+        "AutoModelForCausalLM",
+    ):
+        try:
+            module = __import__("transformers", fromlist=[class_name])
+            model_cls = getattr(module, class_name)
+            return model_cls.from_pretrained(
+                base_id,
+                torch_dtype=dtype,
+                device_map=device_map,
+                trust_remote_code=True,
+            )
+        except Exception as exc:  # pragma: no cover - depends on local transformers version
+            errors.append(f"{class_name}: {exc}")
+    raise RuntimeError("failed to load base model:\n" + "\n".join(errors))
+
+
+def resolve_adapter_dir(source: str) -> Path:
+    """Return a local directory holding the adapter files.
+
+    If ``source`` is a directory containing ``adapter_config.json`` it is used
+    as-is. Otherwise ``source`` is interpreted as a HF Hub repo id and the
+    snapshot is downloaded into the local cache.
+    """
+    candidate = Path(source)
+    if candidate.is_dir() and (candidate / "adapter_config.json").exists():
+        return candidate
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as exc:
+        raise ImportError(
+            "huggingface_hub is required to fetch the adapter from the Hub. "
+            "Install it with: pip install huggingface_hub"
+        ) from exc
+    return Path(snapshot_download(repo_id=source))
+
+
+def load_model(adapter_source: str, dtype: torch.dtype, device_map: str) -> tuple[Any, Any, Path]:
+    adapter_dir = resolve_adapter_dir(adapter_source)
+    adapter_cfg = load_json(adapter_dir / "adapter_config.json")
+    base_id = adapter_cfg["base_model_name_or_path"]
+    processor = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
+    base = load_base_model(base_id, dtype=dtype, device_map=device_map)
+    model = PeftModel.from_pretrained(base, adapter_dir)
+    model.eval()
+    return processor, model, adapter_dir
+
+
+def build_messages(system_prompt: str, user_prompt: str, video_path: Path) -> list[dict[str, Any]]:
+    return [
+        {"role": "system", "content": system_prompt},
+        {
+            "role": "user",
+            "content": [
+                {"type": "video", "video": str(video_path)},
+                {"type": "text", "text": user_prompt},
+            ],
+        },
+    ]
+
+
+def prepare_inputs(
+    processor: Any,
+    messages: list[dict[str, Any]],
+    device: torch.device,
+    *,
+    fps: float,
+    max_pixels: int,
+) -> dict[str, Any]:
+    text = processor.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    try:
+        from qwen_vl_utils import process_vision_info
+    except ImportError as exc:
+        raise ImportError(
+            "qwen-vl-utils is required for local video inference. "
+            "Install it with: pip install qwen-vl-utils[decord]"
+        ) from exc
+
+    for msg in messages:
+        content = msg.get("content")
+        if isinstance(content, list):
+            for item in content:
+                if item.get("type") == "video":
+                    item.setdefault("fps", fps)
+                    item.setdefault("max_pixels", max_pixels)
+
+    try:
+        image_inputs, video_inputs, video_kwargs = process_vision_info(
+            messages,
+            return_video_kwargs=True,
+        )
+    except TypeError:
+        image_inputs, video_inputs = process_vision_info(messages)
+        video_kwargs = {}
+
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+        **video_kwargs,
+    )
+    return inputs.to(device)
+
+
+def decode_generated(processor: Any, inputs: dict[str, Any], generated_ids: torch.Tensor) -> str:
+    input_len = inputs["input_ids"].shape[1]
+    generated_ids = generated_ids[:, input_len:]
+    return processor.batch_decode(
+        generated_ids,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=False,
+    )[0].strip()
+
+
+def parse_score(text: str, key: str) -> int | None:
+    match = re.search(r"\{.*?\}", text, flags=re.S)
+    if match:
+        try:
+            obj = json.loads(match.group(0))
+            value = obj.get(key)
+            if isinstance(value, int) and 1 <= value <= 5:
+                return value
+        except json.JSONDecodeError:
+            pass
+    match = re.search(rf'"?{re.escape(key)}"?\s*:\s*([1-5])', text)
+    if match:
+        return int(match.group(1))
+    return None
+
+
+def dtype_from_name(name: str) -> torch.dtype:
+    if name == "bfloat16":
+        return torch.bfloat16
+    if name == "float16":
+        return torch.float16
+    if name == "float32":
+        return torch.float32
+    raise ValueError(f"unsupported dtype: {name}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Infer with the anonymous judge adapter.")
+    parser.add_argument(
+        "--adapter-dir",
+        default="anonymouscla/physground-judger9B",
+        help=(
+            "Local directory with adapter_config.json + adapter_model.safetensors "
+            "+ subq+human.yaml, or a HF Hub repo id "
+            "(default: anonymouscla/physground-judger9B)."
+        ),
+    )
+    parser.add_argument("--video", required=True, type=Path)
+    parser.add_argument("--caption", required=True)
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--metric", choices=["SA", "PTV", "persistence"])
+    group.add_argument("--law", choices=sorted(PHYSICAL_CRITERIA))
+    parser.add_argument("--criteria", help="Override physical-law criterion text.")
+    parser.add_argument("--max-new-tokens", type=int, default=64)
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--fps", type=float, default=2.0)
+    parser.add_argument("--max-pixels", type=int, default=360 * 640)
+    parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
+    parser.add_argument("--device-map", default="auto")
+    parser.add_argument("--print-prompt", action="store_true")
+    args = parser.parse_args()
+
+    if not args.video.is_file():
+        raise FileNotFoundError(args.video)
+
+    dtype = dtype_from_name(args.dtype)
+    processor, model, adapter_dir = load_model(
+        args.adapter_dir, dtype=dtype, device_map=args.device_map
+    )
+
+    prompt_cfg = load_yaml(adapter_dir / "subq+human.yaml")
+    system_prompt, user_prompt, score_key = build_prompt(
+        prompt_cfg,
+        args.caption,
+        metric=args.metric,
+        law=args.law,
+        criteria=args.criteria,
+    )
+
+    if args.print_prompt:
+        print("SYSTEM:")
+        print(system_prompt)
+        print("\nUSER:")
+        print(user_prompt)
+        print()
+    device = next(model.parameters()).device
+    messages = build_messages(system_prompt, user_prompt, args.video)
+    inputs = prepare_inputs(
+        processor,
+        messages,
+        device,
+        fps=args.fps,
+        max_pixels=args.max_pixels,
+    )
+
+    generation_kwargs: dict[str, Any] = {
+        "max_new_tokens": args.max_new_tokens,
+        "do_sample": args.temperature > 0,
+        "temperature": args.temperature if args.temperature > 0 else None,
+    }
+    generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}
+
+    with torch.inference_mode():
+        generated_ids = model.generate(**inputs, **generation_kwargs)
+
+    raw = decode_generated(processor, inputs, generated_ids)
+    score = parse_score(raw, score_key)
+    print(json.dumps({"key": score_key, "score": score, "raw": raw}, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()