| """Compare a Core ML export with the PyTorch model on a real pair.""" |
| import sys |
| from pathlib import Path |
| import coremltools as ct |
| import numpy as np |
| import torch |
| from PIL import Image |
| sys.path.insert(0, str(Path(__file__).parent)) |
| from liquid_model import IFNet, load_teacher, student_from_teacher |
|
|
| weights, package, a_path, b_path, out_png = sys.argv[1:6] |
| net = student_from_teacher(load_teacher(Path(weights))) if weights.endswith(".safetensors") else IFNet(timestep=True) |
| if not weights.endswith(".safetensors"): |
| net.load_state_dict(torch.load(weights, map_location="cpu")) |
| net.eval() |
| a = Image.open(a_path).convert("RGB").resize((512, 512), Image.LANCZOS) |
| b = Image.open(b_path).convert("RGB").resize((512, 512), Image.LANCZOS) |
| model = ct.models.MLModel(package) |
| coreml = np.asarray(model.predict({"first": a, "second": b, "t": np.full((1, 1, 1, 1), 0.5, np.float32)})["frame"].convert("RGB")).astype(np.float32) |
| ta, tb = (torch.from_numpy(np.asarray(x)).permute(2, 0, 1)[None].float() / 255 for x in (a, b)) |
| with torch.no_grad(): |
| ref = (net(ta, tb, torch.full((1, 1, 1, 1), 0.5)).clamp(0, 1)[0].permute(1, 2, 0).numpy() * 255) |
| mse = ((coreml - ref) ** 2).mean() |
| print(f"Core ML vs PyTorch PSNR: {10 * np.log10(255 ** 2 / max(mse, 1e-9)):.2f} dB") |
| Image.fromarray(np.concatenate([ref, coreml], 1).astype(np.uint8)).save(out_png) |
|
|
