import sys, torch
sys.path.insert(0, ".")
from rife_net import IFNet
W, H = 768, 640
net = IFNet().eval()
state = torch.load("train_log/flownet.pkl", map_location="cpu", weights_only=True)
state = {k.replace("module.", ""): v for k, v in state.items()}
missing, unexpected = net.load_state_dict(state, strict=False)
print("missing", missing[:6], "unexpected", [k for k in unexpected if not k.startswith(("teacher", "caltime"))][:6])
torch.manual_seed(0)
a = torch.rand(1, 3, H, W); b = torch.roll(a, shifts=(3, 5), dims=(2, 3))
with torch.no_grad():
    ref = net(a, b)
prev = torch.load("probe.pt", weights_only=True)["ref"] if __import__("os").path.exists("probe.pt") else None
if prev is not None: print("same as scale-factor version:", float((prev - ref).abs().max()))
torch.save({"a": a, "b": b, "ref": ref}, "probe.pt")
print("torch out", tuple(ref.shape), round(float(ref.min()), 3), round(float(ref.max()), 3))
class Middle(torch.nn.Module):
    """The in-between at t = 0.5 from two frames: all the export takes."""
    def __init__(self, net):
        super().__init__()
        self.net = net

    def forward(self, img0, img1):
        return self.net(img0, img1)

torch.onnx.export(Middle(net).eval(), (a, b), "rife425lite_576p.onnx", input_names=["img0", "img1"], output_names=["middle"], opset_version=17)
print("onnx written")