fix(api): add an option to skip certain models during conversion
This commit is contained in:
parent
dab2cc7752
commit
556d5b84d6
|
@ -350,36 +350,50 @@ def main() -> int:
|
||||||
parser = ArgumentParser(
|
parser = ArgumentParser(
|
||||||
prog='onnx-web model converter',
|
prog='onnx-web model converter',
|
||||||
description='convert checkpoint models to ONNX')
|
description='convert checkpoint models to ONNX')
|
||||||
|
|
||||||
|
# model groups
|
||||||
parser.add_argument('--diffusers', action='store_true', default=True)
|
parser.add_argument('--diffusers', action='store_true', default=True)
|
||||||
parser.add_argument('--gfpgan', action='store_true', default=False)
|
parser.add_argument('--gfpgan', action='store_true', default=False)
|
||||||
parser.add_argument('--resrgan', action='store_true', default=False)
|
parser.add_argument('--resrgan', action='store_true', default=False)
|
||||||
parser.add_argument(
|
parser.add_argument('--skip', nargs='*', type=str, default=[])
|
||||||
'--opset',
|
|
||||||
default=14,
|
# export options
|
||||||
type=int,
|
|
||||||
help="The version of the ONNX operator set to use.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
'--half',
|
'--half',
|
||||||
action='store_true',
|
action='store_true',
|
||||||
default=False,
|
default=False,
|
||||||
help='Export models for half precision, faster on some Nvidia cards'
|
help='Export models for half precision, faster on some Nvidia cards'
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
'--opset',
|
||||||
|
default=14,
|
||||||
|
type=int,
|
||||||
|
help="The version of the ONNX operator set to use.",
|
||||||
|
)
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
print(args)
|
print(args)
|
||||||
|
|
||||||
if args.diffusers:
|
if args.diffusers:
|
||||||
for source in sources.get('diffusers'):
|
for source in sources.get('diffusers'):
|
||||||
convert_diffuser(*source, args.opset, args.half)
|
if source[0] in args.skip:
|
||||||
|
print('Skipping model: %s' % source[0])
|
||||||
|
else:
|
||||||
|
convert_diffuser(*source, args.opset, args.half)
|
||||||
|
|
||||||
if args.resrgan:
|
if args.resrgan:
|
||||||
for source in sources.get('real_esrgan'):
|
for source in sources.get('real_esrgan'):
|
||||||
convert_real_esrgan(*source, args.opset)
|
if source[0] in args.skip:
|
||||||
|
print('Skipping model: %s' % source[0])
|
||||||
|
else:
|
||||||
|
convert_real_esrgan(*source, args.opset)
|
||||||
|
|
||||||
if args.gfpgan:
|
if args.gfpgan:
|
||||||
for source in sources.get('gfpgan'):
|
for source in sources.get('gfpgan'):
|
||||||
convert_gfpgan(*source, args.opset)
|
if source[0] in args.skip:
|
||||||
|
print('Skipping model: %s' % source[0])
|
||||||
|
else:
|
||||||
|
convert_gfpgan(*source, args.opset)
|
||||||
|
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue