diff --git a/finetune/compact.py b/finetune/compact.py new file mode 100644 index 0000000..71c7ac3 --- /dev/null +++ b/finetune/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.")