Command line script to check HifiGan model validity

tts/vocodes_model_check_hifigan.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+
+"""
+Vocodes HiFi-GAN Model Check
+"""
+
+import argparse
+import json
+import os
+import sys
+
+# ML
+import torch
+
+# Hifi-Gan
+from hifigan.env import AttrDict
+from hifigan.models import Generator
+from hifigan.denoiser import Denoiser as HifiDenoiser
+
+parser = argparse.ArgumentParser(description='Check that the HiFi-GAN model is valid')
+
+parser.add_argument('--checkpoint_path', type=str, help='path the hifigan model', required=True)
+parser.add_argument('--output_metadata_filename', type=str, help='where to save extra metadata', required=True)
+
+print('========================================')
+print('Python interpreter', sys.executable)
+print('PyTorch version', torch.__version__)
+print('CUDA Available?', torch.cuda.is_available())
+print('CUDA Device count', torch.cuda.device_count())
+print('========================================', flush=True)
+
+def load_hifigan_model(checkpoint_path):
+    conf = 'hifigan/config_v1.json'
+    with open(conf) as f:
+        json_config = json.loads(f.read())
+    h = AttrDict(json_config)
+    torch.manual_seed(h.seed)
+    hifigan = Generator(h).to(torch.device("cuda"))
+    state_dict_g = torch.load(checkpoint_path, map_location=torch.device("cuda"))
+    hifigan.load_state_dict(state_dict_g["generator"])
+    hifigan.eval()
+    hifigan.remove_weight_norm()
+    denoiser = HifiDenoiser(hifigan, mode="normal")
+    return hifigan, h, denoiser
+
+args = parser.parse_args()
+
+print('Checking model validity...', flush=True)
+
+# Check that the model loads
+hifigan, h, denoiser = load_hifigan_model(args.checkpoint_path)
+
+_ = hifigan.cuda().eval().half()
+_ = denoiser.cuda().eval().half()
+
+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)