"""
Note: can also use external_data=False during export I think
"""

import onnx

# ONNX model blueprint path (the small .onnx file)
INPUT_MODEL_PATH = "/content/lsnet_xl_artist-dynamo-opset20.onnx"

OUTPUT_MODEL_PATH = "/content/lsnet_xl_artist-dynamo-opset20_merged.onnx"

print(f"Loading model from: {INPUT_MODEL_PATH}")
# will automatically load the associated .onnx.data file
model = onnx.load(INPUT_MODEL_PATH)

print(f"Saving single-file model to: {OUTPUT_MODEL_PATH}")
# This function re-embeds the external data into the main .onnx file
onnx.save_model(
    model,
    OUTPUT_MODEL_PATH,
    save_as_external_data=False,
    all_tensors_to_one_file=True
)

print("finished.")