From e3cc57979fa555c728563f8fe83cdadf886afa66 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Fri, 16 Dec 2022 05:21:09 -0500
Subject: [PATCH] Model check script

---
 .../softvc/model_check_hifigan.py             | 72 +++++++++++++++++++
 .../softvc/model_check_softvc_acoustic.py     |  3 +
 2 files changed, 75 insertions(+)
 create mode 100755 voice_conversion/softvc/model_check_hifigan.py
 create mode 100644 voice_conversion/softvc/model_check_softvc_acoustic.py

diff --git a/voice_conversion/softvc/model_check_hifigan.py b/voice_conversion/softvc/model_check_hifigan.py
new file mode 100755
index 0000000..403a854
--- /dev/null
+++ b/voice_conversion/softvc/model_check_hifigan.py
@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+
+"""
+HiFi-GAN (SoftVC) Model Check
+"""
+
+import argparse
+import json
+import os
+import sys
+
+import importlib.util
+
+# ML
+import torch
+
+# Hifi-Gan
+sys.path.append("./hifigan/")
+from hifigan.generator import HifiganGenerator
+from hifigan.generator import consume_prefix_in_state_dict_if_present
+
+from gpu_utils import print_gpu_info
+
+print_gpu_info()
+
+parser = argparse.ArgumentParser(description='Check that the HiFi-GAN model is valid')
+
+parser.add_argument('--checkpoint_path', type=str, help='path the hifigan (softvc) model', required=True)
+parser.add_argument('--output_metadata_filename', type=str, help='where to save extra metadata', required=True)
+
+
+def load_hifigan_model(hifigan_checkpoint_path):
+    # NB(bt, 2022-12-07): Not sure why this "spec loading" is needed.
+    spec = importlib.util.spec_from_file_location("hifigan.utils", "./hifigan/hifigan/utils.py")
+    u = importlib.util.module_from_spec(spec)
+    sys.modules["hifigan.utils"] = u
+    spec.loader.exec_module(u)
+
+    spec = importlib.util.spec_from_file_location("hifigan.generator", "./hifigan/hifigan/generator.py")
+    g = importlib.util.module_from_spec(spec)
+    sys.modules["hifigan.generator"] = g
+    spec.loader.exec_module(g)
+
+    checkpoint = torch.load(hifigan_checkpoint_path)["generator"]["model"]
+    hifigan = g.HifiganGenerator()
+    consume_prefix_in_state_dict_if_present(checkpoint, "module.")
+    hifigan.load_state_dict(checkpoint)
+    hifigan.remove_weight_norm()
+    hifigan.eval()
+
+    hifigan.cuda()
+    return hifigan
+
+
+args = parser.parse_args()
+
+print('Checking model validity...', flush=True)
+
+# Check that the model loads
+_ = load_hifigan_model(args.checkpoint_path)
+
+print('Model is valid', flush=True)
+
+file_size_bytes = os.path.getsize(args.checkpoint_path)
+
+metadata = {
+    'file_size_bytes': file_size_bytes,
+}
+
+with open(args.output_metadata_filename, 'w') as json_file:
+    json.dump(metadata, json_file)
+    
\ No newline at end of file
diff --git a/voice_conversion/softvc/model_check_softvc_acoustic.py b/voice_conversion/softvc/model_check_softvc_acoustic.py
new file mode 100644
index 0000000..c019686
--- /dev/null
+++ b/voice_conversion/softvc/model_check_softvc_acoustic.py
@@ -0,0 +1,3 @@
+#!/usr/bin/env python
+
+# TODO...