Use legacy ONNX exporter (dynamo=False) and skip onnxsim

The dynamo exporter in PyTorch 2.13 produces complex shape formulas
(Min, TruncToInt, CEIL_DIV) that TensorRT cannot resolve for dynamic
shapes. The legacy TorchScript exporter produces clean symbolic dims
(batch, height, width). Also skip onnxsim which folds constants based
on dummy input size, breaking TRT large profile builds.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
huchukato
2026-08-28 05:13:16 +02:00
co-authored by Devin
parent fcb1f332e6
commit 38214ffb2a
+11 -9
View File
@@ -139,10 +139,11 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
interpolation_model.load_state_dict(filtered_state_dict, strict=True)
interpolation_model.eval().to(TORCH_DEVICE)
# # dummy data — use 1080x1080 so onnxsim doesn't fix intermediate
# dimensions to 512px (which breaks TRT build for large profiles)
img0 = torch.randn(1, 3, 1080, 1080).to(TORCH_DEVICE)
img1 = torch.randn(1, 3, 1080, 1080).to(TORCH_DEVICE)
# # dummy data — must be divisible by 16 (RIFE 4.25 scale_list=[16,8,4,2,1])
# and large enough that onnxsim doesn't fix intermediate dims to small values.
# 1024 is divisible by 16/32/64 and is a clean power of 2.
img0 = torch.randn(1, 3, 1024, 1024).to(TORCH_DEVICE)
img1 = torch.randn(1, 3, 1024, 1024).to(TORCH_DEVICE)
timestep = torch.tensor([0.5], dtype=torch.float32).to(TORCH_DEVICE)
# result = (interpolation_model(img0, img1, timestep))
@@ -167,7 +168,8 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
do_constant_folding=True,
input_names=['img0', 'img1', 'timestep'],
output_names=['output'],
dynamic_axes=dynamic_axes
dynamic_axes=dynamic_axes,
dynamo=False, # Use legacy exporter for simpler shape annotations
)
print(f"ONNX model exported to: {onnx_save_path}")
@@ -179,10 +181,10 @@ def export_onnx(ckpt_name, ensemble, scale_factor):
# print(onnx.helper.printable_graph(onnx_model.graph))
sim_model_path = os.path.join("models", f"{ckpt_name.split('.')[0]}_ensemble_{ensemble}_scale_{scale_factor}_sim.onnx")
print("=> ONNX simplify start!")
sim_onnx_model, check = simplify(onnx_model)
onnx.save(sim_onnx_model, sim_model_path)
print("=> ONNX simplify done!")
print("=> Skipping ONNX simplify (onnxsim folds constants based on dummy input size, which breaks TRT dynamic shapes)")
# Just save the raw ONNX as the _sim version (skip simplify)
onnx.save(onnx_model, sim_model_path)
print(f"=> Saved raw ONNX (no simplify) to {sim_model_path}")
sim_model = onnx.load(sim_model_path)
onnx.checker.check_model(sim_model)