# E-OS v4 — export FrameGenerator to ONNX (opset 17) + onnxruntime verification.
import os, sys, json
import numpy as np
import torch
sys.path.insert(0, "/home/z/my-project/scripts/os4")
from os4_models import FrameGenerator

CKPT = "/home/z/my-project/checkpoints4/gen_ckpt_final.pt"
OUT = "/home/z/my-project/download/e-os-v4/models"
os.makedirs(OUT, exist_ok=True)

m = FrameGenerator()
ck = torch.load(CKPT, map_location="cpu", weights_only=False)
m.load_state_dict(ck["m"])
m.eval()
n = sum(p.numel() for p in m.parameters())
print("params %.3fM (ckpt step %s)" % (n / 1e6, ck.get("step")))

hint = torch.rand(1, 3, 192, 416)
state = torch.rand(1, 32)
onnx_path = os.path.join(OUT, "os_generator.onnx")
torch.onnx.export(
    m, (hint, state), onnx_path, opset_version=17,
    input_names=["hint", "state"], output_names=["frame"],
    do_constant_folding=True)
# torch 2.14 writes external weight data -> inline into ONE file (AI Hub needs it)
import onnx
om = onnx.load(onnx_path, load_external_data=True)
onnx.save(om, onnx_path, save_as_external_data=False)
print("onnx saved:", onnx_path, os.path.getsize(onnx_path) // 1024, "KB")

# verify with onnxruntime
import onnxruntime as ort
so = ort.SessionOptions()
sess = ort.InferenceSession(onnx_path, so, providers=["CPUExecutionProvider"])
h = np.random.rand(1, 3, 192, 416).astype(np.float32)
s = np.random.rand(1, 32).astype(np.float32)
ref = m(torch.from_numpy(h), torch.from_numpy(s)).detach().numpy()
got = sess.run(None, {"hint": h, "state": s})[0]
print("ort maxdiff:", float(np.abs(ref - got).max()))
md = float(np.abs(ref - got).max())
assert md < 1e-3, "ORT mismatch!"
print("EXPORT OK")