From 8b0ade5b9f26029d6180a8cb5eb0f354fa3d7d05 Mon Sep 17 00:00:00 2001
From: erew123 <35898566+erew123@users.noreply.github.com>
Date: Thu, 28 Dec 2023 18:34:49 +0000
Subject: [PATCH] compact routine - separate to finetune

a mini compact routine for anyone who wants to compact their existing model. This is built into the new version of finetune anyway, but is here just to capture those outlier instances.
---
 finetune/put-voice-samples-in-here/compact.py | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)
 create mode 100644 finetune/put-voice-samples-in-here/compact.py

diff --git a/finetune/put-voice-samples-in-here/compact.py b/finetune/put-voice-samples-in-here/compact.py
new file mode 100644
index 0000000..71c7ac3
--- /dev/null
+++ b/finetune/put-voice-samples-in-here/compact.py
@@ -0,0 +1,26 @@
+import torch
+import os
+MODEL_PATH = "best_model.pth"
+# Check if best_model.pth exists
+if not os.path.exists(MODEL_PATH):
+    print(f"\nCannot find \033[94m{MODEL_PATH}\033[0m in the current folder.")
+    
+    # Check if model.pth exists as an alternative
+    ALTERNATIVE_PATH = "model.pth"
+    if os.path.exists(ALTERNATIVE_PATH):
+        print(f"\nIf your model is currently called \033[93m{ALTERNATIVE_PATH}\033[0m please rename it to \033[94mbest_model.pth\033[0m.")
+    else:
+        print(f"Please put \033[91mbest_model.pth\033[0m in the current folder or rename your model file to \033[94mbest_model.pth\033[0m.")
+    exit(1)
+# Attempt to load the model
+try:
+    checkpoint = torch.load(MODEL_PATH, map_location=torch.device("cpu"))
+except Exception as e:
+    print("Error loading checkpoint:", e)
+    raise
+del checkpoint["optimizer"]
+for key in list(checkpoint["model"].keys()):
+    if "dvae" in key:
+        del checkpoint["model"][key]
+torch.save(checkpoint, "model.pth")
+print(f"\nYour compacted model is created \033[93mmodel.pth\033[0m.")