55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
# 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)
|