451 lines
16 KiB
Python
Executable File
451 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Compute paired, frame-aligned YUV420 SSIM for MiniMax-H3 runs.
|
|
|
|
The candidate and reference directories must each contain the throughput
|
|
client ``results.jsonl`` files. Rows are paired by ``request_id`` and checked
|
|
for matching prompt, seed, task, geometry, duration, and aspect ratio before
|
|
the decoded videos are compared with FFmpeg's native ``ssim`` filter.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import shutil
|
|
import statistics
|
|
import subprocess
|
|
import tempfile
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
PAIR_FIELDS = (
|
|
"task",
|
|
"short_edge",
|
|
"prompt_index",
|
|
"prompt",
|
|
"seed",
|
|
"duration_seconds",
|
|
"aspect_ratio",
|
|
)
|
|
|
|
|
|
def percentile(values: list[float], q: float) -> float:
|
|
if not values:
|
|
return 0.0
|
|
ordered = sorted(values)
|
|
pos = (len(ordered) - 1) * q
|
|
lo, hi = math.floor(pos), math.ceil(pos)
|
|
if lo == hi:
|
|
return ordered[lo]
|
|
return ordered[lo] * (hi - pos) + ordered[hi] * (pos - lo)
|
|
|
|
|
|
def resolve_executable(explicit: str | None, name: str) -> str:
|
|
if explicit:
|
|
path = Path(explicit)
|
|
if path.is_file():
|
|
return str(path)
|
|
resolved = shutil.which(explicit)
|
|
if resolved:
|
|
return resolved
|
|
raise SystemExit(f"{name} executable not found: {explicit}")
|
|
|
|
resolved = shutil.which(name)
|
|
if resolved:
|
|
return resolved
|
|
candidates = [
|
|
Path("/root/.miniconda3/envs/deploy/bin") / name,
|
|
Path("/root/.miniconda3/envs/vllm/bin") / name,
|
|
Path("/root/.miniconda3/envs/bbj/bin") / name,
|
|
]
|
|
for path in candidates:
|
|
if path.is_file():
|
|
return str(path)
|
|
raise SystemExit(f"{name} is required; pass --{name} explicitly")
|
|
|
|
|
|
def run_checked(command: list[str]) -> subprocess.CompletedProcess[str]:
|
|
result = subprocess.run(command, text=True, capture_output=True, check=False)
|
|
if result.returncode:
|
|
rendered = " ".join(command)
|
|
raise RuntimeError(
|
|
f"command failed ({result.returncode}): {rendered}\n{result.stderr[-4000:]}"
|
|
)
|
|
return result
|
|
|
|
|
|
def resolve_video_path(root: Path, raw_path: str) -> Path:
|
|
if not raw_path:
|
|
raise ValueError("file_path is empty")
|
|
path = Path(raw_path)
|
|
if path.is_file():
|
|
return path.resolve()
|
|
matches = [candidate for candidate in root.rglob(path.name) if candidate.is_file()]
|
|
if len(matches) == 1:
|
|
return matches[0].resolve()
|
|
if not matches:
|
|
raise ValueError(f"video does not exist: {path}")
|
|
raise ValueError(
|
|
f"video path {path} is stale and filename is ambiguous under {root}: "
|
|
+ ", ".join(str(match) for match in matches[:10])
|
|
)
|
|
|
|
|
|
def load_rows(root: Path) -> dict[str, dict[str, Any]]:
|
|
rows: dict[str, dict[str, Any]] = {}
|
|
files = sorted(root.glob("client_*/results.jsonl"))
|
|
if not files:
|
|
files = sorted(root.rglob("client_*/results.jsonl"))
|
|
if not files:
|
|
raise ValueError(f"no client_*/results.jsonl found under {root}")
|
|
|
|
for path in files:
|
|
for line_number, line in enumerate(
|
|
path.read_text(encoding="utf-8").splitlines(), start=1
|
|
):
|
|
if not line.strip():
|
|
continue
|
|
row = json.loads(line)
|
|
if not row.get("success"):
|
|
continue
|
|
request_id = str(row.get("request_id") or "")
|
|
if not request_id:
|
|
raise ValueError(f"missing request_id: {path}:{line_number}")
|
|
if request_id in rows:
|
|
raise ValueError(
|
|
f"duplicate request_id {request_id!r} under {root}; "
|
|
"pass one task/topology phase rather than a whole matrix"
|
|
)
|
|
try:
|
|
video_path = resolve_video_path(root, str(row.get("file_path") or ""))
|
|
except ValueError as error:
|
|
raise ValueError(f"video for {request_id!r}: {error}") from error
|
|
row["_resolved_file_path"] = str(video_path)
|
|
rows[request_id] = row
|
|
return rows
|
|
|
|
|
|
def check_pair(candidate: dict[str, Any], reference: dict[str, Any]) -> None:
|
|
mismatches = []
|
|
for field in PAIR_FIELDS:
|
|
if candidate.get(field) != reference.get(field):
|
|
mismatches.append(
|
|
f"{field}: candidate={candidate.get(field)!r} "
|
|
f"reference={reference.get(field)!r}"
|
|
)
|
|
if mismatches:
|
|
raise ValueError("pair metadata mismatch: " + "; ".join(mismatches))
|
|
|
|
|
|
def probe_video(ffprobe: str, path: Path) -> dict[str, Any]:
|
|
result = run_checked(
|
|
[
|
|
ffprobe,
|
|
"-v",
|
|
"error",
|
|
"-select_streams",
|
|
"v:0",
|
|
"-count_frames",
|
|
"-show_entries",
|
|
"stream=width,height,pix_fmt,r_frame_rate,avg_frame_rate,nb_frames,nb_read_frames",
|
|
"-of",
|
|
"json",
|
|
str(path),
|
|
]
|
|
)
|
|
payload = json.loads(result.stdout)
|
|
streams = payload.get("streams") or []
|
|
if len(streams) != 1:
|
|
raise ValueError(f"expected one video stream in {path}, got {len(streams)}")
|
|
stream = streams[0]
|
|
frame_text = stream.get("nb_read_frames") or stream.get("nb_frames")
|
|
if frame_text in (None, "N/A"):
|
|
raise ValueError(f"could not determine decoded frame count for {path}")
|
|
return {
|
|
"width": int(stream["width"]),
|
|
"height": int(stream["height"]),
|
|
"pix_fmt": stream.get("pix_fmt"),
|
|
"r_frame_rate": stream.get("r_frame_rate"),
|
|
"avg_frame_rate": stream.get("avg_frame_rate"),
|
|
"frames": int(frame_text),
|
|
}
|
|
|
|
|
|
def check_video_contract(candidate: dict[str, Any], reference: dict[str, Any]) -> None:
|
|
fields = ("width", "height", "r_frame_rate", "frames")
|
|
mismatches = [
|
|
f"{field}: candidate={candidate[field]!r} reference={reference[field]!r}"
|
|
for field in fields
|
|
if candidate[field] != reference[field]
|
|
]
|
|
if mismatches:
|
|
raise ValueError("decoded video mismatch: " + "; ".join(mismatches))
|
|
|
|
|
|
def parse_ffmpeg_stats(path: Path) -> list[dict[str, float | int]]:
|
|
frames: list[dict[str, float | int]] = []
|
|
for line in path.read_text(encoding="utf-8").splitlines():
|
|
values: dict[str, str] = {}
|
|
for token in line.split():
|
|
if ":" in token:
|
|
key, value = token.split(":", 1)
|
|
values[key] = value
|
|
if not {"n", "Y", "U", "V", "All"}.issubset(values):
|
|
continue
|
|
frames.append(
|
|
{
|
|
"frame": int(values["n"]),
|
|
"y": float(values["Y"]),
|
|
"u": float(values["U"]),
|
|
"v": float(values["V"]),
|
|
"all": float(values["All"]),
|
|
}
|
|
)
|
|
if not frames:
|
|
raise ValueError(f"FFmpeg emitted no per-frame SSIM metrics: {path}")
|
|
return frames
|
|
|
|
|
|
def compare_video_pair(
|
|
ffmpeg: str,
|
|
ffprobe: str,
|
|
candidate_path: Path,
|
|
reference_path: Path,
|
|
stats_path: Path,
|
|
) -> tuple[dict[str, Any], list[dict[str, float | int]]]:
|
|
candidate_probe = probe_video(ffprobe, candidate_path)
|
|
reference_probe = probe_video(ffprobe, reference_path)
|
|
check_video_contract(candidate_probe, reference_probe)
|
|
|
|
filter_graph = (
|
|
"[0:v]setpts=PTS-STARTPTS,format=yuv420p[candidate];"
|
|
"[1:v]setpts=PTS-STARTPTS,format=yuv420p[reference];"
|
|
f"[candidate][reference]ssim=stats_file={stats_path}"
|
|
)
|
|
run_checked(
|
|
[
|
|
ffmpeg,
|
|
"-hide_banner",
|
|
"-nostdin",
|
|
"-loglevel",
|
|
"error",
|
|
"-i",
|
|
str(candidate_path),
|
|
"-i",
|
|
str(reference_path),
|
|
"-filter_complex",
|
|
filter_graph,
|
|
"-an",
|
|
"-f",
|
|
"null",
|
|
"-",
|
|
]
|
|
)
|
|
frames = parse_ffmpeg_stats(stats_path)
|
|
if len(frames) != candidate_probe["frames"]:
|
|
raise ValueError(
|
|
f"SSIM frame count mismatch: metrics={len(frames)} "
|
|
f"decoded={candidate_probe['frames']}"
|
|
)
|
|
return candidate_probe, frames
|
|
|
|
|
|
def metric_summary(values: list[float]) -> dict[str, float]:
|
|
return {
|
|
"mean": statistics.fmean(values),
|
|
"p10": percentile(values, 0.10),
|
|
"min": min(values),
|
|
"max": max(values),
|
|
}
|
|
|
|
|
|
def compare_command(args: argparse.Namespace) -> int:
|
|
if not 0.0 <= args.threshold <= 1.0:
|
|
raise SystemExit("--threshold must be between 0 and 1")
|
|
ffmpeg = resolve_executable(args.ffmpeg, "ffmpeg")
|
|
ffprobe = resolve_executable(args.ffprobe, "ffprobe")
|
|
candidates = load_rows(args.candidate_dir)
|
|
references = load_rows(args.reference_dir)
|
|
|
|
candidate_ids = set(candidates)
|
|
reference_ids = set(references)
|
|
if candidate_ids != reference_ids:
|
|
missing_reference = sorted(candidate_ids - reference_ids)
|
|
missing_candidate = sorted(reference_ids - candidate_ids)
|
|
raise SystemExit(
|
|
"request sets do not match: "
|
|
f"missing_reference={missing_reference[:20]} "
|
|
f"missing_candidate={missing_candidate[:20]}"
|
|
)
|
|
|
|
selected_ids = sorted(candidate_ids)
|
|
if args.limit is not None:
|
|
if args.limit < 1:
|
|
raise SystemExit("--limit must be at least 1")
|
|
selected_ids = selected_ids[: args.limit]
|
|
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
pair_rows: list[dict[str, Any]] = []
|
|
all_frame_rows: list[dict[str, Any]] = []
|
|
with tempfile.TemporaryDirectory(prefix="h3-paired-ssim-") as temporary:
|
|
temporary_root = Path(temporary)
|
|
for index, request_id in enumerate(selected_ids, start=1):
|
|
candidate = candidates[request_id]
|
|
reference = references[request_id]
|
|
check_pair(candidate, reference)
|
|
probe, frames = compare_video_pair(
|
|
ffmpeg,
|
|
ffprobe,
|
|
Path(candidate["_resolved_file_path"]),
|
|
Path(reference["_resolved_file_path"]),
|
|
temporary_root / f"{index:05d}.stats",
|
|
)
|
|
all_values = [float(frame["all"]) for frame in frames]
|
|
y_values = [float(frame["y"]) for frame in frames]
|
|
u_values = [float(frame["u"]) for frame in frames]
|
|
v_values = [float(frame["v"]) for frame in frames]
|
|
all_summary = metric_summary(all_values)
|
|
pair = {
|
|
"request_id": request_id,
|
|
"task": candidate["task"],
|
|
"short_edge": candidate["short_edge"],
|
|
"prompt_index": candidate["prompt_index"],
|
|
"prompt": candidate["prompt"],
|
|
"seed": candidate["seed"],
|
|
"candidate_file": candidate["_resolved_file_path"],
|
|
"reference_file": reference["_resolved_file_path"],
|
|
"candidate_num_inference_steps": candidate.get("num_inference_steps"),
|
|
"reference_num_inference_steps": reference.get("num_inference_steps"),
|
|
"width": probe["width"],
|
|
"height": probe["height"],
|
|
"fps": probe["r_frame_rate"],
|
|
"frames": len(frames),
|
|
"ssim_all_mean": all_summary["mean"],
|
|
"ssim_all_p10": all_summary["p10"],
|
|
"ssim_all_min": all_summary["min"],
|
|
"ssim_y_mean": statistics.fmean(y_values),
|
|
"ssim_u_mean": statistics.fmean(u_values),
|
|
"ssim_v_mean": statistics.fmean(v_values),
|
|
"threshold": args.threshold,
|
|
"passed": all_summary["mean"] >= args.threshold,
|
|
}
|
|
pair_rows.append(pair)
|
|
for frame in frames:
|
|
all_frame_rows.append({"request_id": request_id, **frame})
|
|
print(
|
|
f"[{index}/{len(selected_ids)}] {request_id} "
|
|
f"mean={all_summary['mean']:.6f} p10={all_summary['p10']:.6f} "
|
|
f"min={all_summary['min']:.6f}",
|
|
flush=True,
|
|
)
|
|
|
|
all_values = [float(row["all"]) for row in all_frame_rows]
|
|
video_means = [float(row["ssim_all_mean"]) for row in pair_rows]
|
|
by_short_edge: dict[str, dict[str, Any]] = {}
|
|
grouped: dict[int, list[float]] = defaultdict(list)
|
|
for row in pair_rows:
|
|
grouped[int(row["short_edge"])].append(float(row["ssim_all_mean"]))
|
|
for short_edge, values in sorted(grouped.items()):
|
|
by_short_edge[str(short_edge)] = {
|
|
"videos": len(values),
|
|
"mean_video_ssim": statistics.fmean(values),
|
|
"min_video_ssim": min(values),
|
|
}
|
|
|
|
frame_summary = metric_summary(all_values)
|
|
summary = {
|
|
"metric": "FFmpeg decoded YUV420 SSIM All",
|
|
"aggregation": {
|
|
"mean_video_ssim": statistics.fmean(video_means),
|
|
"frame_weighted_mean_ssim": frame_summary["mean"],
|
|
"frame_p10_ssim": frame_summary["p10"],
|
|
"min_frame_ssim": frame_summary["min"],
|
|
},
|
|
"candidate_dir": str(args.candidate_dir.resolve()),
|
|
"reference_dir": str(args.reference_dir.resolve()),
|
|
"ffmpeg": ffmpeg,
|
|
"ffprobe": ffprobe,
|
|
"threshold": args.threshold,
|
|
"overall_passed": statistics.fmean(video_means) >= args.threshold,
|
|
"videos": len(pair_rows),
|
|
"videos_passed": sum(bool(row["passed"]) for row in pair_rows),
|
|
"frames": len(all_frame_rows),
|
|
"by_short_edge": by_short_edge,
|
|
"pairs": pair_rows,
|
|
}
|
|
(args.output_dir / "paired_ssim.json").write_text(
|
|
json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
|
|
pair_columns = [
|
|
"request_id",
|
|
"task",
|
|
"short_edge",
|
|
"prompt_index",
|
|
"seed",
|
|
"width",
|
|
"height",
|
|
"fps",
|
|
"frames",
|
|
"ssim_all_mean",
|
|
"ssim_all_p10",
|
|
"ssim_all_min",
|
|
"ssim_y_mean",
|
|
"ssim_u_mean",
|
|
"ssim_v_mean",
|
|
"threshold",
|
|
"passed",
|
|
"candidate_file",
|
|
"reference_file",
|
|
]
|
|
with (args.output_dir / "paired_ssim.tsv").open("w", encoding="utf-8") as handle:
|
|
handle.write("\t".join(pair_columns) + "\n")
|
|
for row in pair_rows:
|
|
handle.write("\t".join(str(row[column]) for column in pair_columns) + "\n")
|
|
|
|
frame_columns = ("request_id", "frame", "y", "u", "v", "all")
|
|
with (args.output_dir / "paired_ssim_frames.tsv").open(
|
|
"w", encoding="utf-8"
|
|
) as handle:
|
|
handle.write("\t".join(frame_columns) + "\n")
|
|
for row in all_frame_rows:
|
|
handle.write("\t".join(str(row[column]) for column in frame_columns) + "\n")
|
|
|
|
print(json.dumps(summary["aggregation"], ensure_ascii=False), flush=True)
|
|
if args.fail_below_threshold and not summary["overall_passed"]:
|
|
return 2
|
|
return 0
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
compare = subparsers.add_parser("compare")
|
|
compare.add_argument("--candidate-dir", type=Path, required=True)
|
|
compare.add_argument("--reference-dir", type=Path, required=True)
|
|
compare.add_argument("--output-dir", type=Path, required=True)
|
|
compare.add_argument("--threshold", type=float, default=0.90)
|
|
compare.add_argument(
|
|
"--limit",
|
|
type=int,
|
|
help="compare only the first N matched requests (intended for smoke tests)",
|
|
)
|
|
compare.add_argument("--ffmpeg")
|
|
compare.add_argument("--ffprobe")
|
|
compare.add_argument("--fail-below-threshold", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.command == "compare":
|
|
return compare_command(args)
|
|
raise AssertionError(args.command)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|