Skip to content

Commit

Permalink
Merge pull request #76 from Haidra-Org/main
Browse files Browse the repository at this point in the history
feat: patch in lora cache size env var support
  • Loading branch information
tazlin authored Oct 4, 2023
2 parents 98074b0 + 2f8369b commit 2c2d528
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions hordelib/model_manager/lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ class DOWNLOAD_SIZE_CHECK(StrEnum):

TESTS_ONGOING = os.getenv("TESTS_ONGOING", "0") == "1"

AIWORKER_LORA_CACHE_SIZE_DEFAULT = 10 * 1024 # 10GB


class LoraModelManager(BaseModelManager):
LORA_DEFAULTS = "https://raw.githubusercontent.com/Haidra-Org/AI-Horde-image-model-reference/main/lora.json"
Expand All @@ -49,12 +51,20 @@ def __init__(
self,
download_reference=False,
allowed_top_lora_storage=10240 if not TESTS_ONGOING else 1024,
allowed_adhoc_lora_storage=1024,
allowed_adhoc_lora_storage=AIWORKER_LORA_CACHE_SIZE_DEFAULT,
download_wait=False,
):
self.max_adhoc_disk = allowed_adhoc_lora_storage
try:
AIWORKER_LORA_CACHE_SIZE = os.getenv("AIWORKER_LORA_CACHE_SIZE")
if AIWORKER_LORA_CACHE_SIZE is not None:
self.max_adhoc_disk = int(AIWORKER_LORA_CACHE_SIZE)
logger.debug(f"AIWORKER_LORA_CACHE_SIZE is {self.max_adhoc_disk}")
except (ValueError, TypeError):
self.max_adhoc_disk = AIWORKER_LORA_CACHE_SIZE_DEFAULT

self._max_top_disk = allowed_top_lora_storage

self.max_adhoc_disk = allowed_adhoc_lora_storage
self._data = None
self._next_page_url = None
self._mutex = threading.Lock()
Expand Down

0 comments on commit 2c2d528

Please sign in to comment.