# E-OS v5 — ONNX-Export (opset 17, App-Aufloesung 256x544 -> 512x1088) + Verify
import os
import sys
import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from os5_model import FrameGeneratorV5  # noqa: E402

CKPT = "/home/z/my-project/eos-v5/train/v5_ckpt.pt"
OUTD = "/home/z/my-project/eos-v5/models"
os.makedirs(OUTD, exist_ok=True)

model = FrameGeneratorV5()
ck = torch.load(CKPT, map_location="cpu", weights_only=False)
model.load_state_dict(ck["model"])
model.eval()
print("step:", ck["step"], "hist:", ck["hist"][-3:])

hint = torch.rand(1, 3, 544, 256)
state = torch.rand(1, 32)
with torch.no_grad():
    ref = model(hint, state)

onnx_path = os.path.join(OUTD, "os_glass_v5.onnx")
torch.onnx.export(
    model, (hint, state), onnx_path,
    input_names=["hint", "state"], output_names=["frame"],
    opset_version=17, dynamo=False,
    do_constant_folding=True,
)
print("exported:", onnx_path, os.path.getsize(onnx_path) // 1024, "KB")

import onnxruntime as ort
so = ort.SessionOptions()
sess = ort.InferenceSession(onnx_path, so, providers=["CPUExecutionProvider"])
outs = sess.run(None, {"hint": hint.numpy(), "state": state.numpy()})
diff = np.abs(outs[0] - ref.numpy()).max()
print("ort output shape:", outs[0].shape, "max diff vs torch:", diff)
assert diff < 1e-4, "ONNX mismatch"
print("ONNX_OK")