"""Score staged baseline predictions with the established WorldMem metric contracts.""" import argparse import importlib import importlib.util import json from pathlib import Path import sys import types sys.path.insert(0, str(Path(__file__).resolve().parent / "metric_vendor")) import numpy as np import torch import torch.nn.functional as F import torchmetrics from torchmetrics.functional.image.ssim import structural_similarity_index_measure from torchmetrics.image.fid import FrechetInceptionDistance from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from protocol import case_key, load_re10k_case def load_source(path, name): spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def load_oasis_vae(worldmem_root, checkpoint, device): # A namespace import keeps the authoritative VAE's relative imports intact # without importing WorldMem's training/experiment package. package = types.ModuleType("_benchmark_worldmem_models") package.__path__ = [str(Path(worldmem_root) / "algorithms/worldmem/models")] sys.modules[package.__name__] = package module = importlib.import_module(package.__name__ + ".vae") vae = module.VAE_models["vit-l-20-shallow-encoder"]() payload = torch.load(checkpoint, map_location="cpu", weights_only=False) state = payload.get("state_dict", payload) state = {key.removeprefix("vae."): value for key, value in state.items()} vae.load_state_dict(state, strict=True) return vae.requires_grad_(False).eval().to(device) @torch.inference_mode() def minecraft_reference(row, feature_root, vae, frames, device): relative = Path(row["relative_path"]) feature = Path(feature_root) / relative.parent / (relative.stem + "_vae_feature.npy") start = int(row["target_start"]) latents = np.load(feature, mmap_mode="r")[start:start + frames] decoded = [] for offset in range(0, frames, 16): z = torch.from_numpy(np.array(latents[offset:offset + 16], copy=True)).float().to(device) # The existing cache stores scaled C x18x32 latents; decode matches # DeMemWM.decode, including float32 scaling before fp16 autocast. z = z.flatten(2).transpose(1, 2) / 0.07843137255 with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"): rgb = (vae.decode(z) + 1) / 2 decoded.append(rgb.float().clamp(0, 1).cpu()) return torch.cat(decoded), str(feature) @torch.inference_mode() def frame_metrics(prediction, reference, dataset, lpips, device, batch_size=16): values = {name: [] for name in (["mse", "psnr", "lpips", "ssim"] if dataset == "re10k" else ["mse", "psnr", "lpips"])} for start in range(0, len(prediction), batch_size): pred, real = (video[start:start + batch_size].float().to(device).clamp(0, 1) for video in (prediction, reference)) mse = (pred - real).square().flatten(1).mean(1) epsilon = 1e-12 if dataset == "re10k" else torch.finfo(mse.dtype).eps values["mse"].append(mse.cpu()) values["psnr"].append((-10 * mse.clamp_min(epsilon).log10()).cpu()) values["lpips"].append(lpips.net(pred, real, normalize=True).flatten().cpu()) if dataset == "re10k": values["ssim"].append(structural_similarity_index_measure( pred, real, data_range=1.0, reduction="none").cpu()) return {name: torch.cat(parts) for name, parts in values.items()} @torch.inference_mode() def inception_features(frames, inception, device): # FrechetInceptionDistance(normalize=True).update performs this exact # truncating uint8 conversion before running torch-fidelity Inception. return torch.cat([inception((batch.to(device) * 255).byte()).reshape(len(batch), -1).float().cpu() for batch in frames.split(16)]) @torch.inference_mode() def video_inception_features(path, legacy_fid, inception, count, device): return torch.cat([inception(batch.to(device)).reshape(len(batch), -1).float().cpu() for batch in legacy_fid.frame_batches(path, 0, count, 16)]) def metric_ranges(dataset, frames): width = 25 if dataset == "re10k" else 100 ranges = [(f"frames_{start:03d}_{min(start + width, frames):03d}", start, min(start + width, frames)) for start in range(0, frames, width)] ranges.append(("overall", 0, frames)) if dataset == "re10k": ranges.append(("forward", 0, min(100, frames))) if frames > 100: ranges.append(("return", 100, frames)) return ranges @torch.inference_mode() def score_case(case_dir, dataset, data_root, worldmem_root, vae_checkpoint=None, feature_root=None, keep_tensor=False, device="cuda"): case_dir, data_root, worldmem_root = map(Path, (case_dir, data_root, worldmem_root)) device = torch.device(device) metadata = json.loads((case_dir / "case.json").read_text()) row = metadata["case"] prediction_path = case_dir / "prediction.pt" prediction = torch.load(prediction_path, map_location="cpu", weights_only=True).float() expected_size = (256, 256) if dataset == "re10k" else (360, 640) max_frames = 300 if dataset == "re10k" else 500 errors = [] if torchmetrics.__version__ != "0.11.4": errors.append(f"Legacy contract requires torchmetrics0.11.4, found {torchmetrics.__version__}") if prediction.ndim != 4 or prediction.shape[1] != 3 or not 0 < len(prediction) <= max_frames: errors.append(f"Invalid generated RGB suffix shape {tuple(prediction.shape)}") if not torch.isfinite(prediction).all(): errors.append("Non-finite prediction") if errors: raise ValueError("; ".join(errors)) native_size = list(prediction.shape[-2:]) if tuple(native_size) != expected_size: prediction = torch.cat([F.interpolate(batch.to(device), size=expected_size, mode="bicubic", align_corners=False, antialias=True).cpu() for batch in prediction.split(16)]) prediction.clamp_(0, 1) if dataset == "re10k": reference = load_re10k_case(row, data_root)["rgb"][100:100 + len(prediction)] reference_source = "raw DFoT test_256 RGB at matching source indices" else: vae = load_oasis_vae(worldmem_root, vae_checkpoint, device) reference, reference_source = minecraft_reference( row, feature_root or data_root / "vae_features", vae, len(prediction), device) del vae torch.cuda.empty_cache() if reference.shape != prediction.shape or not torch.isfinite(reference).all(): raise ValueError("Reference frames do not match the complete finite prediction population") lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).eval().to(device) values = frame_metrics(prediction, reference, dataset, lpips, device) del lpips legacy = load_source(worldmem_root / "evaluation/long_horizon.py", "_benchmark_long_horizon") for name, frames in (("prediction", prediction), ("reference", reference)): legacy._write_video(case_dir / f"{name}.mp4", legacy._uint8_video(frames), fps=15) fid = FrechetInceptionDistance(feature=2048, normalize=dataset == "re10k").eval().to(device) if dataset == "minecraft": # calculate_fid imports only the authoritative long_horizon helpers. sys.path.insert(0, str(worldmem_root)) legacy_fid = load_source(worldmem_root / "calculate_fid.py", "_benchmark_legacy_fid") fake_features = video_inception_features(case_dir / "prediction.mp4", legacy_fid, fid.inception, len(prediction), device) real_features = video_inception_features(case_dir / "reference.mp4", legacy_fid, fid.inception, len(reference), device) else: fake_features = inception_features(prediction, fid.inception, device) real_features = inception_features(reference, fid.inception, device) if any(not torch.isfinite(value).all() for value in (*values.values(), fake_features, real_features)): raise ValueError("Non-finite metric or FID feature; preserving prediction.pt for diagnosis") np.savez(case_dir / "frame_metrics.npz", **{key: value.numpy() for key, value in values.items()}) np.save(case_dir / "fake_features.npy", fake_features.numpy()) np.save(case_dir / "real_features.npy", real_features.numpy()) summary = { "case": row, "sample_seed": int(metadata["seed"]), "dataset": dataset, "frames": len(prediction), "native_resolution": native_size, "metric_resolution": list(expected_size), "resize": "none" if tuple(native_size) == expected_size else "bicubic align_corners=False antialias=True", "reference_source": reference_source, "vae_checkpoint": str(vae_checkpoint) if dataset == "minecraft" else None, "torch": torch.__version__, "torchmetrics": torchmetrics.__version__, "lpips": "alex; normalize=True; mean of individual frame values", "psnr_epsilon": 1e-12 if dataset == "re10k" else torch.finfo(torch.float32).eps, "fid": "direct float01 -> truncating uint8" if dataset == "re10k" else "OpenCV decoding of libx264 quality8 yuv420p uint8 videos", "video": {"codec": "libx264", "quality": 8, "pixelformat": "yuv420p", "macro_block_size": 1, "fps": 15}, "ranges": {name: {metric: float(value[start:end].mean()) for metric, value in values.items()} for name, start, end in metric_ranges(dataset, len(prediction))}, } (case_dir / "metrics.json").write_text(json.dumps(summary, indent=2) + "\n") if not keep_tensor: prediction_path.unlink() return summary def main(): parser = argparse.ArgumentParser() parser.add_argument("--case-dir", type=Path, required=True) parser.add_argument("--dataset", choices=["minecraft", "re10k"], required=True) parser.add_argument("--data-root", type=Path, required=True) parser.add_argument("--worldmem-root", type=Path, required=True) parser.add_argument("--vae-checkpoint", type=Path) parser.add_argument("--feature-root", type=Path) parser.add_argument("--keep-tensor", action="store_true") parser.add_argument("--device", default="cuda") args = parser.parse_args() result = score_case(args.case_dir, args.dataset, args.data_root, args.worldmem_root, args.vae_checkpoint, args.feature_root, args.keep_tensor, args.device) print(json.dumps({"case": case_key(result["case"]), "frames": result["frames"], "metrics": result["ranges"]["overall"]}), flush=True) if __name__ == "__main__": main()