import torch
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'vigs'))

from midas.omnidata import OmnidataModel


def export_onnx(model, dummy, filename):
    model = model.eval().cuda()

    torch.onnx.export(
        model,
        dummy,
        filename,
        dynamo=True,               # Not sure if there is a bit difference between dynamo and not
        input_names=["input"],
        output_names=["output"],
        opset_version=18,          # safe for MiDaS / ViT models
    )

    print(f"[OK] Saved ONNX → {filename}")


# ---- Input shape ----
dummy = torch.randn(1, 3, 512, 512).cuda()


# ---- Depth model ----
depth_model = OmnidataModel(
    'depth',
    'pretrained_models/omnidata_dpt_depth_v2.ckpt',
    device='cuda'
).model
export_onnx(depth_model, dummy, "pretrained_models/omnidata_depth_512.onnx")


# ---- Normal model ----
normal_model = OmnidataModel(
    'normal',
    'pretrained_models/omnidata_dpt_normal_v2.ckpt',
    device='cuda'
).model
export_onnx(normal_model, dummy, "pretrained_models/omnidata_normal_512.onnx")




