diff --git a/melgan-swpark/wrap_melgan_model.py b/melgan-swpark/wrap_melgan_model.py index 4603c9c..7b6af8a 100755 --- a/melgan-swpark/wrap_melgan_model.py +++ b/melgan-swpark/wrap_melgan_model.py @@ -3,110 +3,54 @@ import torch from collections import OrderedDict +import argparse import re +import os +import sys import yaml from model_melgan.generator import load_melgan #torch.set_default_tensor_type('torch.DoubleTensor') -#model_checkpoint_path = '/home/bt/models/melgan-swpark/trumptxlearn183_fa14347_7675.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/ljs-high-pitch_8741385_6975.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/trump-txlearn-2020.05.05_8741385_7525.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/trump-txlearn-2020.05.05_8741385_13675.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_trump_txlearn_2020-05-30_steno_chkpt9800.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_multivoice-1.0_2020.06.15_12d5838_6275.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_multivoice-1.0_2020.06.15_12d5838_6300.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_multivoice-1.0_2020.06.15_12d5838_6475.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_multivoice-1.0_2020.06.15_12d5838_6725.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_6675.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_6900.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_7125.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_7400.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/vocodes1.0_12d5838_7900.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/vocodes1.0_12d5838_9025.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/steno1.5_12d5838_9225.pt' # ugh, forgot the naming scheme -model_checkpoint_path = '/home/bt/models/melgan-swpark/steno1.5_12d5838_9475.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/steno1.5_12d5838_10275.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_12d5838_6300.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_12d5838_6410.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_12d5838_6590.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_2020-06-28_12d5838_6610.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_12d5838_6730.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_2020-06-28_12d5838_6730.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_2020-06-28_12d5838_7050.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_2020-06-28_12d5838_7940.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice1.0_2020-06-28_12d5838_8570.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_8640.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_8810.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_8880.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_8900.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_8940.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_9050.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice2.0_2020-07-08_12d5838_9330.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_9820.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_9860.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_9920.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_9950.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_9960.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_10050.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice3.0_2020-07-17_12d5838_10090.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice5.0_2020-07-23_12d5838_10490.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice5.0_2020-07-23_12d5838_10750.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice5.0_2020-07-23_12d5838_10760.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice6.0_2020-07-27_12d5838_10830.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice6.0_2020-07-27_12d5838_10870.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice6.0_2020-07-27_12d5838_10880.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice6.0_2020-07-27_12d5838_10910.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/john-oliver/melgan_john_oliver_2020-07-29_12d5838_10920.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/john-oliver/melgan_john_oliver_2020-07-29_12d5838_10970.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/john-oliver/melgan_john_oliver_2020-07-29_12d5838_11000.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/john-oliver/melgan_john_oliver_2020-07-29_12d5838_13010.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/john-oliver/melgan_john_oliver_2020-07-29_12d5838_14110.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/bart-simpson/melgan_bart-simpson_58e5387_10950.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/bart-simpson/melgan_bart-simpson_58e5387_14600.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/manyvoice6.0_2020-07-27_12d5838_10910.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/bart-simpson/melgan_bart-simpson_2020.08.03_58e5387_15100.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/spongebob-squarepants/melgan_spongebob_58e5387_11575.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/spongebob-squarepants/melgan_spongebob-squarepants_2020.08.11_58e5387_11725.pt' -model_checkpoint_path = '/home/bt/models/melgan-swpark/spongebob-squarepants/melgan_spongebob_58e5387_18725.pt' +DEFAULT_NAME = 'melgan-vocodes.torchjit' -# renamed for posterity -#model_checkpoint_path = '/home/bt/models/melgan-swpark/melgan_vocodes1.0_multispeaker_2020.06.19_12d5838_experiment-2_chkpt-10400.pt' +def parse_args(): + parser = argparse.ArgumentParser('JITswap the melgan') + parser.add_argument('--model_file', type=str, required=True) # Input + parser.add_argument('--output_file', type=str, required=False) # Explicitly set the filename + parser.add_argument('--copy_name', default=False, action='store_true') # Save as '{model_name}.torchjit' + parser.add_argument('--output_directory', type=str, required=False, default='./OUTPUT/') + + args = parser.parse_args() + + if not args.model_file.endswith('.pt'): + raise Exception('Melgan model filename `{}` does not end with extention `.pt`.'.format(args.model_file)) + + if args.output_file and not args.output_file.endswith('.torchjit'): + raise Exception('(Optional) Melgan JIT output filename `{}` does not end with extention `.torchjit`.'.format(args.output_file)) + + if args.copy_name: + if args.output_file: + raise Exception('Cannot --copy_name and set an explicit --output_file') + basename = os.path.basename(args.model_file) + args.output_file = os.path.splitext(basename)[0] + '.torchjit' + else: + args.output_file = DEFAULT_NAME + + return args + +args = parse_args() + + +model_checkpoint_path = args.model_file +output_filename = os.path.join(args.output_directory, args.output_file) + +print('model file', model_checkpoint_path) +print('output file', output_filename) melgan_model = load_melgan(model_checkpoint_path) -#output_filename = 'shared_melgan_container2.pt' -#output_filename = 'melgan_ljs-high-pitch_container.pt' -output_filename = 'melgan_trump-txlearn-2020.05.05.jit' -output_filename = 'melgan_trump-txlearn-2020.05.05_13675.jit' -output_filename = 'melgan_trump_txlearn_2020-05-30_steno_chkpt9800.torchjit' -output_filename = 'melgan_multivoice-1.0_2020.06.15_12d5838_6275.torchjit' -output_filename = 'melgan_multivoice-1.0_2020.06.15_12d5838_6300.torchjit' -output_filename = 'melgan_multivoice-1.0_2020.06.15_12d5838_6475.torchjit' -output_filename = 'melgan_multivoice-1.0_2020.06.15_12d5838_6725.torchjit' -output_filename = 'melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_6675.torchjit' -output_filename = 'melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_6900.torchjit' -output_filename = 'melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_7125.torchjit' -output_filename = 'melgan_vocodes1.0_multispeaker_2020_06_19_12d5838_7400.torchjit' -output_filename = 'melgan_vocodes1.0_multispeaker_2020.06.19_12d5838_experiment-2_chkpt-10400.torchjit' -output_filename = 'OUTPUT/manyvoice1.0_2020-06-28_12d5838_6610.torchjit' -output_filename = 'OUTPUT/manyvoice1.0_2020-06-28_12d5838_6730.torchjit' -output_filename = 'OUTPUT/manyvoice1.0_2020-06-28_12d5838_7050.torchjit' -output_filename = 'OUTPUT/manyvoice2.0_2020-07-08_12d5838_8940.torchjit' -output_filename = 'OUTPUT/manyvoice2.0_2020-07-08_12d5838_9170.torchjit' -output_filename = 'OUTPUT/manyvoice2.0_2020-07-08_12d5838_9330.torchjit' -output_filename = 'OUTPUT/manyvoice3.0_2020-07-17_12d5838_9920.torchjit' -output_filename = 'OUTPUT/manyvoice5.0_2020-07-23_12d5838_10490.torchjit' -output_filename = 'OUTPUT/manyvoice5.0_2020-07-23_12d5838_10490.torchjit' -output_filename = 'OUTPUT/manyvoice5.0_2020-07-23_12d5838_10760.torchjit' -output_filename = 'OUTPUT/melgan-vocodes-2.torchjit' # Stop renaming them for a bit -output_filename = 'OUTPUT/melgan_bart-simpson_2020.08.03_58e5387_15100.torchjit' -output_filename = 'OUTPUT/melgan_spongebob-squarepants_2020.08.11_58e5387_11725.torchjit' -output_filename = 'OUTPUT/melgan-vocodes.torchjit' # Stop renaming them for a bit - - -#output_filename = 'OUTPUT/melgan-vocodes.torchjit' # Stop renaming them for a bit # NB: Tracing evaluates the model on input and unrolls and hardcodes branching # and loops. Scripting allows these to remain by converting the entire program