fix: lora cache size issues

download_models.py

@@ -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,
@@ -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(bridge_data.nsfw)
+
+        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")
@@ -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)

load_env_vars.py

@@ -26,3 +26,18 @@ def load_env_vars() -> None:  # FIXME: there is a dynamic way to do this
                 "AIWORKER_CACHE_HOME environment variable already set. "
                 "This will override the value for `cache_home` in the config file.",
             )
+
+    if "max_lora_cache_size" in config:
+        if os.getenv("AIWORKER_LORA_CACHE_SIZE") is None:
+            try:
+                int(config["max_lora_cache_size"])
+            except ValueError as e:
+                raise ValueError(
+                    "max_lora_cache_size must be an integer, but is not.",
+                ) from e
+            os.environ["AIWORKER_LORA_CACHE_SIZE"] = str(config["max_lora_cache_size"])
+        else:
+            print(
+                "AIWORKER_LORA_CACHE_SIZE environment variable already set. "
+                "This will override the value for `max_lora_cache_size` in the config file.",
+            )