Skip to content

Commit

Permalink
feat: add --purge-unused-loras to download_models.py
Browse files Browse the repository at this point in the history
This allows a user-prompted cleaning up of the LoRa folder
  • Loading branch information
tazlin committed Dec 19, 2023
1 parent b39f7f2 commit 883abac
Showing 1 changed file with 33 additions and 2 deletions.
35 changes: 33 additions & 2 deletions download_models.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,20 @@
"""Contains the code to download all models specified in the config file. Executable as a standalone script."""

from load_env_vars import load_env_vars

load_env_vars()

import argparse
import time

from horde_model_reference.model_reference_manager import ModelReferenceManager
from loguru import logger

from horde_worker_regen.bridge_data.load_config import BridgeDataLoader, reGenBridgeData
from horde_worker_regen.consts import BRIDGE_CONFIG_FILENAME


def download_all_models() -> None:
def download_all_models(purge_unused_loras: bool = False) -> None:
"""Download all models specified in the config file."""
horde_model_reference_manager = ModelReferenceManager(
download_and_convert_legacy_dbs=True,
Expand Down Expand Up @@ -45,6 +49,24 @@ def download_all_models() -> None:

SharedModelManager.load_model_managers()

if purge_unused_loras:
logger.info("Purging unused LORAs...")
if SharedModelManager.manager.lora is None:
logger.error("Failed to load LORA model manager")
exit(1)
deleted_loras = SharedModelManager.manager.lora.delete_unused_loras(30)
logger.success(f"Purged {len(deleted_loras)} unused LORAs.")

if bridge_data.allow_lora:
if SharedModelManager.manager.lora is None:
logger.error("Failed to load LORA model manager")
exit(1)
SharedModelManager.manager.lora.download_default_loras()

while SharedModelManager.manager.lora.are_downloads_complete() is False:
logger.info("Waiting for LORA downloads to complete...")
time.sleep(8)

if bridge_data.allow_controlnet:
if SharedModelManager.manager.controlnet is None:
logger.error("Failed to load controlnet model manager")
Expand Down Expand Up @@ -92,4 +114,13 @@ def download_all_models() -> None:


if __name__ == "__main__":
download_all_models()
parser = argparse.ArgumentParser(description="Download all models specified in the config file.")
parser.add_argument(
"--purge-unused-loras",
action="store_true",
help="Purge unused LORAs from the cache",
)

args = parser.parse_args()

download_all_models(purge_unused_loras=args.purge_unused_loras)

0 comments on commit 883abac

Please sign in to comment.