Views
No views yet
UVR-MDX-NET-Inst_HQ_4.pt was obtained w/ the following code:1import argparse
2
3import onnx
4import torch
5from onnx2torch import convert
6
7
8def try_forward(m, shape):
9 x = torch.randn(*shape, dtype=torch.float32)
10 with torch.no_grad():
11 m(x)
12 return True
13
14
15def main(onnx_path, out_prefix):
16 model_onnx = onnx.load(onnx_path)
17 model_torch = convert(model_onnx).eval()
18
19 candidates = [
20 (1, 4, 2560, 256),
21 (1, 4, 2560, 320),
22 (1, 4, 3072, 256),
23 (1, 4, 3072, 320),
24 ]
25
26 ok_shape = None
27 for shape in candidates:
28 try:
29 try_forward(model_torch, shape)
30 ok_shape = shape
31 break
32 except Exception:
33 pass
34 if ok_shape is None:
35 raise RuntimeError("Could not find a working input shape for this ONNX model.")
36
37 try:
38 scripted = torch.jit.script(model_torch)
39 print("Scripted model")
40 except Exception:
41 scripted = torch.jit.trace(model_torch, torch.randn(*ok_shape), strict=False)
42 print("Traced model")
43
44 out_pt = f"{out_prefix}.pt"
45 torch.jit.save(scripted, out_pt)
46
47
48if __name__ == "__main__":
49 ap = argparse.ArgumentParser()
50 ap.add_argument("--onnx", required=True)
51 ap.add_argument("--out-prefix", default="UVR-MDX-NET-Inst_HQ_4")
52 args = ap.parse_args()
53 main(args.onnx, args.out_prefix)
54