106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Fail-fast acceptance test for the isolated VBench scoring environment."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
MAIN_DIMENSIONS = (
|
|
"subject_consistency",
|
|
"background_consistency",
|
|
"motion_smoothness",
|
|
"dynamic_degree",
|
|
"aesthetic_quality",
|
|
"imaging_quality",
|
|
"object_class",
|
|
"multiple_objects",
|
|
"human_action",
|
|
"color",
|
|
"spatial_relationship",
|
|
"scene",
|
|
"temporal_style",
|
|
"appearance_style",
|
|
"overall_consistency",
|
|
)
|
|
|
|
REQUIRED_WEIGHTS = {
|
|
"/root/.cache/vbench/amt_model/amt-s.pth": 10_000_000,
|
|
"/root/.cache/vbench/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth": 500_000_000,
|
|
"/root/.cache/vbench/grit_model/grit_b_densecap_objectdet.pth": 400_000_000,
|
|
"/root/.cache/vbench/caption_model/tag2text_swin_14m.pth": 4_000_000_000,
|
|
"/root/.cache/vbench/ViCLIP/ViClip-InternVid-10M-FLT.pth": 1_500_000_000,
|
|
"/root/.cache/vbench/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth": 100_000_000,
|
|
"/root/.cache/vbench/raft_model/models/raft-things.pth": 20_000_000,
|
|
"/root/.cache/torch/hub/checkpoints/dino_vitbase16_pretrain.pth": 300_000_000,
|
|
"/root/.cache/clip/ViT-B-32.pt": 300_000_000,
|
|
"/root/.cache/clip/ViT-L-14.pt": 900_000_000,
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--vbench-repo", default=os.environ.get("VBENCH_REPO", "/data/wxy/VBench"))
|
|
parser.add_argument("--skip-dino-load", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
repo = Path(args.vbench_repo).resolve()
|
|
if not (repo / "vbench" / "VBench_full_info.json").is_file():
|
|
raise RuntimeError(f"invalid VBench repository: {repo}")
|
|
sys.path.insert(0, str(repo))
|
|
os.chdir(repo)
|
|
|
|
report: dict[str, object] = {
|
|
"python": sys.version,
|
|
"executable": sys.executable,
|
|
"repo": str(repo),
|
|
"imports": {},
|
|
"weights": {},
|
|
}
|
|
|
|
import torch
|
|
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("torch.cuda.is_available() is false")
|
|
x = torch.randn(512, 512, device="cuda")
|
|
torch.cuda.synchronize()
|
|
report["cuda"] = {
|
|
"torch": torch.__version__,
|
|
"runtime": torch.version.cuda,
|
|
"capability": list(torch.cuda.get_device_capability()),
|
|
"probe_shape": list((x @ x).shape),
|
|
}
|
|
|
|
for dimension in MAIN_DIMENSIONS + ("temporal_flickering",):
|
|
importlib.import_module(f"vbench.{dimension}")
|
|
report["imports"][dimension] = "ok"
|
|
|
|
for path_string, minimum_size in REQUIRED_WEIGHTS.items():
|
|
path = Path(path_string)
|
|
size = path.stat().st_size if path.is_file() else 0
|
|
if size < minimum_size:
|
|
raise RuntimeError(f"missing or truncated weight: {path} ({size} bytes)")
|
|
report["weights"][path_string] = size
|
|
|
|
if not args.skip_dino_load:
|
|
model = torch.hub.load(
|
|
"facebookresearch/dino:main",
|
|
"dino_vitb16",
|
|
source="github",
|
|
trust_repo=True,
|
|
verbose=True,
|
|
)
|
|
report["dino_parameters"] = sum(parameter.numel() for parameter in model.parameters())
|
|
|
|
print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|