# download the upscale models & place inside models/upscaler_models
# edit model paths accordingly 

import onnx
import io
import torch
import folder_paths
from spandrel import ModelLoader, ImageModelDescriptor
from onnxsim import simplify

model_name     = "4x_foolhardy_Remacri.pth"
onnx_save_path = "./4x_foolhardy_Remacri.onnx"

#Load the model and weights
model_path  = folder_paths.get_full_path_or_raise("upscale_models", model_name)
model       = ModelLoader().load_from_file(model_path).model.eval().cuda()

dummy_input = torch.rand(1, 3, 64, 64).cuda()

dynamic_axes = {
                "input": {0: "batch_size", 2: "width", 3: "height"},
                "output": {0: "batch_size", 2: "width", 3: "height"},
               }
               
dynamic_shapes = {
                  "x": {
                            0: torch.export.Dim("batch", min=1, max=8),
                            2: torch.export.Dim("width", min=16, max=4096),
                            3: torch.export.Dim("height", min=16, max=4096),
                           }
                 }

onnx_buffer = io.BytesIO()

torch.onnx.export(model,
                  dummy_input,
                  onnx_buffer,
                  verbose=True,
                  input_names=['input'],
                  output_names=['output'],
                  opset_version=18,
                  export_params=True,
                  dynamic_shapes=dynamic_shapes
                  )
                  
print("Simplifying ONNX ")
onnx_buffer.seek(0)
model_onnx        = onnx.load(onnx_buffer)
model_simp, check = simplify(model_onnx)

assert check, "Simplified ONNX model could not be validated"

onnx.save(model_simp, onnx_save_path)
print("Saved simplified ONNX to:", onnx_save_path)