Skip to content

Fix physics_iq scoring result reuse across tasks and validate staged artifacts #379

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
JIN-strong wants to merge 1 commit into from
Closed

Fix physics_iq scoring result reuse across tasks and validate staged artifacts #379

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions evaluation/cosmos3/generator/physics_iq/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,11 @@ The notebook walks through both PhysicsIQ task formats:
72-frame conditioning to 73 frames; generate 193 frames total; score
frames 73..192.

Scoring results are namespaced per task directory (`scoring/i2v_full`,
`scoring/v2v_full`, ...), so the scorer's cached CSV can never be reused
across tasks, and staged artifacts are re-staged automatically when they no
longer probe as the expected frame count.

## Files

- `run_with_cosmos_framework.ipynb` — main notebook.
Expand Down
119 changes: 77 additions & 42 deletions evaluation/cosmos3/generator/physics_iq/run_with_cosmos_framework.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,9 +392,9 @@
"- `case_paths(case_id)` — locate the switch frame JPG and 72-frame conditioning MP4 for a case.\n",
"- `build_i2v_input_jsonl(...)` — assemble a one-line JSONL row pointing at the switch frame.\n",
"- `build_v2v_input_jsonl(...)` — assemble a one-line JSONL row pointing at the 73-frame conditioning MP4.\n",
"- `prepare_v2v_conditioning(...)` — clone-pad the official 72-frame conditioning to 73 frames (the format the V2V model expects for Physics-IQ V2V).\n",
"- `stage_i2v_last120(...)` — drop the first frame of a 121-frame I2V output, keep the last 120.\n",
"- `stage_v2v_omit73(...)` — drop the 73-frame conditioning prefix of a 193-frame V2V output, keep the next 120.\n",
"- `prepare_v2v_conditioning(...)` — clone-pad the official 72-frame conditioning to 73 frames (the format the V2V model expects for Physics-IQ V2V). An existing clip is reused only when it still probes as exactly 73 frames.\n",
"- `stage_i2v_last120(...)` — drop the first frame of a 121-frame I2V output, keep the last 120. Existing artifacts are reused only when they still probe as exactly 120 frames; truncated ones are re-staged.\n",
"- `stage_v2v_omit73(...)` — drop the 73-frame conditioning prefix of a 193-frame V2V output, keep the next 120. Existing artifacts are reused only when they still probe as exactly 120 frames; truncated ones are re-staged.\n",
"\n",
"All staging produces a 120-frame mp4 at 24 fps — the exact format the official PhysicsIQ scorer expects."
]
Expand All @@ -406,6 +406,7 @@
"outputs": [],
"source": [
"import json\n",
"import os\n",
"import shutil\n",
"import subprocess\n",
"from pathlib import Path\n",
Expand Down Expand Up @@ -445,6 +446,35 @@
" cap.release()\n",
"\n",
"\n",
"def _stage_validated(dst_mp4: Path, expected_frames: int, produce) -> None:\n",
" \"\"\"Materialize dst_mp4 through a validated temp file.\n",
"\n",
" Re-stages when an existing artifact does not probe back as exactly\n",
" `expected_frames` frames (e.g. a run was killed mid-encode), so a\n",
" truncated clip can never enter scoring or conditioning. The producer\n",
" writes to a temp file that is renamed into place only after the frame\n",
" count checks out.\n",
" \"\"\"\n",
" if dst_mp4.exists():\n",
" try:\n",
" if _ffprobe_frames(dst_mp4) == expected_frames:\n",
" return\n",
" except Exception:\n",
" pass\n",
" dst_mp4.unlink() # truncated or unreadable — rebuild it\n",
" dst_mp4.parent.mkdir(parents=True, exist_ok=True)\n",
" tmp_mp4 = dst_mp4.with_name(f\".{dst_mp4.name}.{os.getpid()}.tmp\")\n",
" try:\n",
" produce(tmp_mp4)\n",
" n_frames = _ffprobe_frames(tmp_mp4)\n",
" assert n_frames == expected_frames, (\n",
" f\"{dst_mp4}: expected {expected_frames} frames, got {n_frames}\"\n",
" )\n",
" os.replace(tmp_mp4, dst_mp4)\n",
" finally:\n",
" tmp_mp4.unlink(missing_ok=True)\n",
"\n",
"\n",
"def case_paths(case_id: str) -> tuple[Path, Path]:\n",
" sample_id = case_id.split(\"_\", 1)[0]\n",
" sw = sorted((PHYSICSIQ_ROOT / \"switch-frames\").glob(f\"{sample_id}_*.jpg\"))\n",
Expand Down Expand Up @@ -486,18 +516,19 @@
" \"\"\"Clone-pad the first frame so the official 72-frame clip becomes 73 frames.\"\"\"\n",
" n_src = _ffprobe_frames(src_72f)\n",
" assert n_src == 72, f\"{src_72f}: expected 72 source frames, got {n_src}\"\n",
" dst_73f.parent.mkdir(parents=True, exist_ok=True)\n",
" start_duration = 1.0 / float(fps)\n",
" _ffmpeg([\n",
" \"-i\", str(src_72f),\n",
" \"-vf\", f\"tpad=start_mode=clone:start_duration={start_duration:.9f}\",\n",
" \"-r\", str(fps),\n",
" \"-frames:v\", \"73\",\n",
" \"-c:v\", \"libx264\", \"-pix_fmt\", \"yuv420p\", \"-crf\", \"18\",\n",
" str(dst_73f),\n",
" ])\n",
" n_dst = _ffprobe_frames(dst_73f)\n",
" assert n_dst == 73, f\"{dst_73f}: expected 73 frames after pad, got {n_dst}\"\n",
"\n",
" def produce(tmp_mp4: Path) -> None:\n",
" _ffmpeg([\n",
" \"-i\", str(src_72f),\n",
" \"-vf\", f\"tpad=start_mode=clone:start_duration={start_duration:.9f}\",\n",
" \"-r\", str(fps),\n",
" \"-frames:v\", \"73\",\n",
" \"-c:v\", \"libx264\", \"-pix_fmt\", \"yuv420p\", \"-crf\", \"18\",\n",
" str(tmp_mp4),\n",
" ])\n",
"\n",
" _stage_validated(dst_73f, 73, produce)\n",
" return dst_73f\n",
"\n",
"\n",
Expand Down Expand Up @@ -533,33 +564,35 @@
" n_src = _ffprobe_frames(src_mp4)\n",
" assert n_src >= 121, f\"{src_mp4}: expected at least 121 frames, got {n_src}\"\n",
" start_frame = n_src - 120\n",
" dst_mp4.parent.mkdir(parents=True, exist_ok=True)\n",
" _ffmpeg([\n",
" \"-i\", str(src_mp4),\n",
" \"-vf\", f\"trim=start_frame={start_frame},setpts=PTS-STARTPTS,fps={fps}\",\n",
" \"-frames:v\", \"120\",\n",
" \"-c:v\", \"libx264\", \"-preset\", \"ultrafast\", \"-crf\", \"23\",\n",
" \"-pix_fmt\", \"yuv420p\", \"-an\",\n",
" str(dst_mp4),\n",
" ])\n",
" assert _ffprobe_frames(dst_mp4) == 120\n",
" def produce(tmp_mp4: Path) -> None:\n",
" _ffmpeg([\n",
" \"-i\", str(src_mp4),\n",
" \"-vf\", f\"trim=start_frame={start_frame},setpts=PTS-STARTPTS,fps={fps}\",\n",
" \"-frames:v\", \"120\",\n",
" \"-c:v\", \"libx264\", \"-preset\", \"ultrafast\", \"-crf\", \"23\",\n",
" \"-pix_fmt\", \"yuv420p\", \"-an\",\n",
" str(tmp_mp4),\n",
" ])\n",
"\n",
" _stage_validated(dst_mp4, 120, produce)\n",
" return dst_mp4\n",
"\n",
"\n",
"def stage_v2v_omit73(src_mp4: Path, dst_mp4: Path, fps: int = 24, omit: int = 73) -> Path:\n",
" \"\"\"Drop the 73-frame conditioning prefix from a 193-frame V2V output, keep the next 120.\"\"\"\n",
" n_src = _ffprobe_frames(src_mp4)\n",
" assert n_src >= omit + 120, f\"{src_mp4}: expected >= {omit + 120} frames, got {n_src}\"\n",
" dst_mp4.parent.mkdir(parents=True, exist_ok=True)\n",
" _ffmpeg([\n",
" \"-i\", str(src_mp4),\n",
" \"-vf\", f\"select='gte(n,{omit})',setpts=N/({fps}*TB)\",\n",
" \"-frames:v\", \"120\",\n",
" \"-r\", str(fps),\n",
" \"-an\", \"-c:v\", \"libx264\", \"-pix_fmt\", \"yuv420p\", \"-crf\", \"18\",\n",
" str(dst_mp4),\n",
" ])\n",
" assert _ffprobe_frames(dst_mp4) == 120\n",
" def produce(tmp_mp4: Path) -> None:\n",
" _ffmpeg([\n",
" \"-i\", str(src_mp4),\n",
" \"-vf\", f\"select='gte(n,{omit})',setpts=N/({fps}*TB)\",\n",
" \"-frames:v\", \"120\",\n",
" \"-r\", str(fps),\n",
" \"-an\", \"-c:v\", \"libx264\", \"-pix_fmt\", \"yuv420p\", \"-crf\", \"18\",\n",
" str(tmp_mp4),\n",
" ])\n",
"\n",
" _stage_validated(dst_mp4, 120, produce)\n",
" return dst_mp4\n",
"\n",
"\n",
Expand Down Expand Up @@ -864,8 +897,7 @@
" for case_id in case_ids:\n",
" _, conditioning_72f = case_paths(case_id)\n",
" dst = v2v_cond_dir / f\"{case_id}.mp4\"\n",
" if not dst.exists():\n",
" prepare_v2v_conditioning(conditioning_72f, dst)\n",
" prepare_v2v_conditioning(conditioning_72f, dst)\n",
"\n",
" all_jsonl = v2v_inputs_dir / \"all_198.jsonl\"\n",
" with all_jsonl.open(\"w\") as fp:\n",
Expand Down Expand Up @@ -984,15 +1016,13 @@
" for raw in sorted(Path(os.environ[\"I2V_FULL_OUTPUT_DIR\"]).rglob(\"vision.mp4\")):\n",
" case_id = raw.parent.name\n",
" dst = Path(os.environ[\"I2V_FULL_OUTPUT_DIR\"]).parent / \"staged\" / f\"{case_id}.mp4\"\n",
" if not dst.exists():\n",
" stage_i2v_last120(raw, dst)\n",
" stage_i2v_last120(raw, dst)\n",
" print(\"I2V staged dir:\", Path(os.environ[\"I2V_FULL_OUTPUT_DIR\"]).parent / \"staged\")\n",
" if \"v2v\" in TASKS_TO_RUN:\n",
" for raw in sorted(Path(os.environ[\"V2V_FULL_OUTPUT_DIR\"]).rglob(\"vision.mp4\")):\n",
" case_id = raw.parent.name\n",
" dst = Path(os.environ[\"V2V_FULL_OUTPUT_DIR\"]).parent / \"staged\" / f\"{case_id}.mp4\"\n",
" if not dst.exists():\n",
" stage_v2v_omit73(raw, dst)\n",
" stage_v2v_omit73(raw, dst)\n",
" print(\"V2V staged dir:\", Path(os.environ[\"V2V_FULL_OUTPUT_DIR\"]).parent / \"staged\")"
]
},
Expand Down Expand Up @@ -1021,7 +1051,12 @@
"# Demo run (single case V2V): $PHYSICSIQ_OUTPUT_ROOT/v2v/<CASE_ID>/staged\n",
"# Full 198-case I2V sweep: $PHYSICSIQ_OUTPUT_ROOT/i2v_full/staged (default)\n",
"STAGED_DIR=\"${STAGED_DIR:-$PHYSICSIQ_OUTPUT_ROOT/i2v_full/staged}\"\n",
"RESULTS_DIR=\"${RESULTS_DIR:-$PHYSICSIQ_OUTPUT_ROOT/scoring}\"\n",
"# Namespace results by the staged folder's parent (i2v_full / v2v_full / the\n",
"# demo case id): the scorer keys its cached CSV by the input folder's\n",
"# basename, and every staged folder here is named \"staged\", so one shared\n",
"# results directory would let the second task silently reuse the first\n",
"# task's cached scores.\n",
"RESULTS_DIR=\"${RESULTS_DIR:-$PHYSICSIQ_OUTPUT_ROOT/scoring/$(basename \"$(dirname \"$STAGED_DIR\")\")}\"\n",
"PHYSICSIQ_REPO=\"$(dirname \"$PHYSICSIQ_ROOT\")/physics-IQ-benchmark-repo\"\n",
"\n",
"if [ ! -d \"$STAGED_DIR\" ]; then\n",
Expand Down Expand Up @@ -1056,4 +1091,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
}
}
153 changes: 153 additions & 0 deletions evaluation/cosmos3/generator/physics_iq/test_run_with_cosmos_framework.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Offline regression tests for the physics_iq notebook staging helpers.

The tests extract the real function definitions from
run_with_cosmos_framework.ipynb (no logic is duplicated here) and drive them
with fake ffmpeg/ffprobe pairs, so no GPU, network, or multimedia stack is
needed.

Run: python3 test_run_with_cosmos_framework.py
"""

import json
import os
import subprocess
import types
import unittest
from pathlib import Path

NOTEBOOK = Path(__file__).resolve().parent / "run_with_cosmos_framework.ipynb"


def load_staging_namespace() -> dict:
"""Exec the notebook cell that defines the staging helpers and return its namespace."""
nb = json.loads(NOTEBOOK.read_text())
for cell in nb["cells"]:
source = "".join(cell["source"])
if "def stage_i2v_last120(" in source:
func_source = source[source.index("def _ffmpeg("):]
namespace = {
"Path": Path,
"os": os,
"FFMPEG": "/usr/bin/true",
"shutil": types.SimpleNamespace(which=lambda _p: "/usr/bin/true"),
"subprocess": subprocess,
}
exec(compile(func_source, str(NOTEBOOK), "exec"), namespace)
return namespace
raise AssertionError("staging helper cell not found in notebook")


class FakeMedia:
"""Stand-in for ffmpeg/ffprobe: the 'frame count' is embedded in the file bytes."""

def __init__(self, namespace: dict):
self.calls: list[list[str]] = []
self._namespace = namespace
namespace["_ffmpeg"] = self._ffmpeg
namespace["_ffprobe_frames"] = self._ffprobe_frames

@staticmethod
def frames_from_args(args: list[str]) -> int:
return int(args[args.index("-frames:v") + 1])

def _ffmpeg(self, args: list[str]) -> None:
self.calls.append([str(a) for a in args])
out_path = Path(self.calls[-1][-1])
out_path.write_bytes(f"FRAMES={self.frames_from_args(self.calls[-1])}".encode())

@staticmethod
def _ffprobe_frames(path) -> int:
data = Path(path).read_bytes()
return int(data.removeprefix(b"FRAMES=").strip())

@staticmethod
def write_source(path: Path, frames: int) -> Path:
path.write_bytes(f"FRAMES={frames}".encode())
return path


class StagingHelperTest(unittest.TestCase):
def setUp(self):
import tempfile

self._tmp = tempfile.TemporaryDirectory()
self.addCleanup(self._tmp.cleanup)
self.tmp = Path(self._tmp.name)
self.ns = load_staging_namespace()
self.media = FakeMedia(self.ns)

def last_output_arg(self) -> Path:
return Path(self.media.calls[-1][-1])

def test_stage_i2v_last120_writes_through_temp_file(self):
src = self.media.write_source(self.tmp / "raw.mp4", 121)
dst = self.tmp / "staged" / "case.mp4"
self.ns["stage_i2v_last120"](src, dst)
self.assertEqual(self.media._ffprobe_frames(dst), 120)
# the producer must write to a temp path, never directly to the final one
self.assertTrue(self.last_output_arg().name.startswith(".case.mp4."))
self.assertTrue(self.last_output_arg().name.endswith(".tmp"))
self.assertNotEqual(self.last_output_arg(), dst)

def test_stage_i2v_last120_reuses_valid_artifact_without_rerun(self):
src = self.media.write_source(self.tmp / "raw.mp4", 121)
dst = self.tmp / "staged" / "case.mp4"
self.ns["stage_i2v_last120"](src, dst)
calls_after_first = len(self.media.calls)
self.ns["stage_i2v_last120"](src, dst)
self.assertEqual(len(self.media.calls), calls_after_first)
self.assertEqual(self.media._ffprobe_frames(dst), 120)

def test_stage_i2v_last120_restages_corrupted_artifact(self):
src = self.media.write_source(self.tmp / "raw.mp4", 121)
dst = self.tmp / "staged" / "case.mp4"
self.ns["stage_i2v_last120"](src, dst)
# simulate a run killed mid-encode: the artifact is short
dst.write_bytes(b"FRAMES=100")
self.ns["stage_i2v_last120"](src, dst)
self.assertEqual(self.media._ffprobe_frames(dst), 120)

def test_stage_v2v_omit73_frames_and_reuse(self):
src = self.media.write_source(self.tmp / "raw.mp4", 193)
dst = self.tmp / "staged" / "case.mp4"
self.ns["stage_v2v_omit73"](src, dst)
self.assertEqual(self.media._ffprobe_frames(dst), 120)
calls_after_first = len(self.media.calls)
self.ns["stage_v2v_omit73"](src, dst)
self.assertEqual(len(self.media.calls), calls_after_first)

def test_stage_v2v_omit73_restages_corrupted_artifact(self):
src = self.media.write_source(self.tmp / "raw.mp4", 193)
dst = self.tmp / "staged" / "case.mp4"
self.ns["stage_v2v_omit73"](src, dst)
dst.write_bytes(b"FRAMES=64")
self.ns["stage_v2v_omit73"](src, dst)
self.assertEqual(self.media._ffprobe_frames(dst), 120)

def test_prepare_conditioning_pads_72_to_73_and_reuses(self):
src = self.media.write_source(self.tmp / "cond72.mp4", 72)
dst = self.tmp / "cond73.mp4"
self.ns["prepare_v2v_conditioning"](src, dst)
self.assertEqual(self.media._ffprobe_frames(dst), 73)
calls_after_first = len(self.media.calls)
self.ns["prepare_v2v_conditioning"](src, dst)
self.assertEqual(len(self.media.calls), calls_after_first)

def test_prepare_conditioning_restages_corrupted_artifact(self):
src = self.media.write_source(self.tmp / "cond72.mp4", 72)
dst = self.tmp / "cond73.mp4"
self.ns["prepare_v2v_conditioning"](src, dst)
dst.write_bytes(b"FRAMES=59")
self.ns["prepare_v2v_conditioning"](src, dst)
self.assertEqual(self.media._ffprobe_frames(dst), 73)

def test_prepare_conditioning_rejects_unexpected_source(self):
src = self.media.write_source(self.tmp / "cond71.mp4", 71)
with self.assertRaises(AssertionError):
self.ns["prepare_v2v_conditioning"](src, self.tmp / "cond73.mp4")


if __name__ == "__main__":
unittest.main(verbosity=2)