Perform graceful exit on KeyboardInterrupt

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:

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__':

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__':

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)

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)