fix: exported model name should match generator class name

EveryVoiceTTS · Nov 25, 2024 · 87f8298 · 87f8298
1 parent 5a8a227
commit 87f8298
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/hfgl/cli.py b/hfgl/cli.py
@@ -103,6 +103,8 @@ def export(
     del vocoder_ckpt["callbacks"]
     del vocoder_ckpt["optimizer_states"]
     del vocoder_ckpt["lr_schedulers"]
+    if "model_info" in vocoder_ckpt:
+        vocoder_ckpt["model_info"]["name"] = "HiFiGANGenerator"
     torch.save(vocoder_ckpt, output_path)
     new_size = sizeof_fmt(os.path.getsize(output_path))
     logger.info(