diff --git a/basics/base_binarizer.py b/basics/base_binarizer.py index 41091a26..210f6f57 100644 --- a/basics/base_binarizer.py +++ b/basics/base_binarizer.py @@ -245,18 +245,22 @@ def postprocess(_item): lengths.append(aug_item['length']) total_sec += aug_item['seconds'] - if num_workers > 0: - # code for parallel processing - for item in tqdm( - chunked_multiprocess_run(self.process_item, args, num_workers=num_workers), - total=len(list(self.meta_data_iterator(prefix))) - ): - postprocess(item) - else: - # code for single cpu processing - for a in tqdm(args): - item = self.process_item(*a) - postprocess(item) + try: + if num_workers > 0: + # code for parallel processing + for item in tqdm( + chunked_multiprocess_run(self.process_item, args, num_workers=num_workers), + total=len(list(self.meta_data_iterator(prefix))) + ): + postprocess(item) + else: + # code for single cpu processing + for a in tqdm(args): + item = self.process_item(*a) + postprocess(item) + except KeyboardInterrupt: + builder.finalize() + exit(-1) builder.finalize() with open(self.binary_data_dir / f'{prefix}.lengths', 'wb') as f: diff --git a/scripts/export.py b/scripts/export.py index 3328773a..29b4f3f5 100644 --- a/scripts/export.py +++ b/scripts/export.py @@ -130,7 +130,10 @@ def acoustic( export_spk=export_spk_mix, freeze_spk=freeze_spk_mix ) - exporter.export(out) + try: + exporter.export(out) + except KeyboardInterrupt: + exit(-1) @main.command(help='Export DiffSinger variance model to ONNX format.') @@ -177,7 +180,10 @@ def variance( export_spk=export_spk_mix, freeze_spk=freeze_spk_mix ) - exporter.export(out) + try: + exporter.export(out) + except KeyboardInterrupt: + exit(-1) @main.command(help='Export NSF-HiFiGAN vocoder model to ONNX format.') @@ -211,7 +217,10 @@ def nsf_hifigan( model_path=Path(hparams['vocoder_ckpt']).resolve(), model_name=name ) - exporter.export(out) + try: + exporter.export(out) + except KeyboardInterrupt: + exit(-1) if __name__ == '__main__': diff --git a/scripts/infer.py b/scripts/infer.py index fbb855e2..d69622b1 100644 --- a/scripts/infer.py +++ b/scripts/infer.py @@ -123,10 +123,13 @@ def acoustic( infer_ins = DiffSingerAcousticInfer(load_model=not mel, ckpt_steps=ckpt) print(f'| Model: {type(infer_ins.model)}') - infer_ins.run_inference( - params, out_dir=out, title=name, num_runs=num, - spk_mix=spk_mix, seed=seed, save_mel=mel - ) + try: + infer_ins.run_inference( + params, out_dir=out, title=name, num_runs=num, + spk_mix=spk_mix, seed=seed, save_mel=mel + ) + except KeyboardInterrupt: + exit(-1) @main.command(help='Run DiffSinger variance model inference') @@ -208,10 +211,13 @@ def variance( infer_ins = DiffSingerVarianceInfer(ckpt_steps=ckpt, predictions=set(predict)) print(f'| Model: {type(infer_ins.model)}') - infer_ins.run_inference( - params, out_dir=out, title=name, - num_runs=num, seed=seed - ) + try: + infer_ins.run_inference( + params, out_dir=out, title=name, + num_runs=num, seed=seed + ) + except KeyboardInterrupt: + exit(-1) if __name__ == '__main__': diff --git a/scripts/vocode.py b/scripts/vocode.py index 517cdd75..55861866 100644 --- a/scripts/vocode.py +++ b/scripts/vocode.py @@ -84,4 +84,7 @@ def run_vocoder(path: pathlib.Path): os.makedirs(out, exist_ok=True) -run_vocoder(out / (name + '.wav')) +try: + run_vocoder(out / (name + '.wav')) +except KeyboardInterrupt: + exit(-1) diff --git a/utils/multiprocess_utils.py b/utils/multiprocess_utils.py index 41eb1314..236d94a1 100644 --- a/utils/multiprocess_utils.py +++ b/utils/multiprocess_utils.py @@ -18,7 +18,9 @@ def chunked_worker_run(map_func, args, results_queue=None): try: res = map_func(*a) results_queue.put(res) - except: + except KeyboardInterrupt: + break + except Exception: traceback.print_exc() results_queue.put(None)