Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Some more bugs on the lora MM #372

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 24 additions & 22 deletions hordelib/model_manager/lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -531,16 +531,16 @@ def _download_thread(self, thread_number):
# Save the hash file
with open(hashfile, "w") as outfile:
outfile.write(f"{sha256} *{lora['versions'][version]['filename']}")
lora["versions"][version]["last_used"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
lora["last_checked"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self._add_lora_to_reference(lora)
# Shout about it
logger.info(
f"Downloaded LORA {lora['versions'][version]['filename']} "
f"({lora['versions'][version]['size_mb']} MB)",
f"({lora['versions'][version]['size_mb']} MB). "
"The current adhoc cache is at "
f"{round(self.calculate_adhoc_loras_cache() / 1024,1)}/{self.max_adhoc_disk}G",
)
# Maybe we're done
# We store as lower to allow case-insensitive search
lora["versions"][version]["last_used"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
lora["last_checked"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self._add_lora_to_reference(lora)
if self.is_adhoc_cache_full():
for _lora_iter in range(self.amount_of_adhoc_loras_to_delete()):
self.delete_oldest_lora()
Expand Down Expand Up @@ -884,16 +884,16 @@ def calculate_adhoc_loras_cache(self):
return self.calculate_downloaded_loras(DOWNLOAD_SIZE_CHECK.adhoc)

def is_default_cache_full(self):
return self.calculate_default_loras_cache() >= self._max_top_disk
return self.calculate_default_loras_cache() >= self._max_top_disk * 1024

def is_adhoc_cache_full(self):
return self.calculate_adhoc_loras_cache() >= self.max_adhoc_disk
return self.calculate_adhoc_loras_cache() >= self.max_adhoc_disk * 1024

def amount_of_adhoc_loras_to_delete(self):
if not self.is_adhoc_cache_full():
return 0
# If we have exceeded our cache, we delete 1 lora + 1 extra lora per 4G over our cache.
return 1 + int((self.calculate_adhoc_loras_cache() - self.max_adhoc_disk) / 4096)
return 1 + int((self.calculate_adhoc_loras_cache() - self.max_adhoc_disk * 1024) / 4096)

def calculate_download_queue(self):
total_queue = 0
Expand All @@ -906,10 +906,16 @@ def find_oldest_adhoc_lora(self) -> str | None:
oldest_datetime: datetime | None = None
for lora in self._adhoc_loras:
for version in self.model_reference[lora]["versions"]:
lora_datetime = datetime.strptime(
self.model_reference[lora]["versions"][version]["last_used"],
"%Y-%m-%d %H:%M:%S",
)
if "last_used" in self.model_reference[lora]["versions"][version]:
lora_datetime = datetime.strptime(
self.model_reference[lora]["versions"][version]["last_used"],
"%Y-%m-%d %H:%M:%S",
)
else:
lora_datetime = datetime.now()
self.model_reference[lora]["versions"][version]["last_used"] = lora_datetime.strftime(
"%Y-%m-%d %H:%M:%S",
)
if not oldest_lora:
oldest_lora = version
oldest_datetime = lora_datetime
Expand Down Expand Up @@ -1034,15 +1040,7 @@ def reset_adhoc_loras(self):
if prevversion.get("adhoc", True) is False:
continue
# If True, it will initiates a redownload and call _add_lora_to_reference() later
if not self._check_for_refresh(prevlora_key):
if "last_used" not in prevversion:
prevversion["last_used"] = now
# We create a temp lora dict holding the just one version (the one we want to keep)
# The _add_lora_to_reference() will anyway merge versions if we keep more than 1
temp_lora = self.model_reference[prevlora_key].copy()
temp_lora["versions"] = {}
temp_lora["versions"][prevversion["version_id"]] = prevversion
self._add_lora_to_reference(temp_lora)
self._check_for_refresh(prevlora_key)
self._adhoc_loras.add(prevlora_key)
self.save_cached_reference_to_disk()
logger.debug(
Expand Down Expand Up @@ -1164,6 +1162,10 @@ def get_lora_last_use(self, lora_name, is_version: bool = False):
version = self.find_latest_version(lora)
if version is None:
return None
if "last_used" not in lora["versions"][version]:
logger.debug("Lora does not appear to have a last_used key. Setting it to now()")
lora["versions"][version]["last_used"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
return datetime.now()
return datetime.strptime(lora["versions"][version]["last_used"], "%Y-%m-%d %H:%M:%S")

def fetch_adhoc_lora(self, lora_name, timeout: int | None = 45, is_version: bool = False) -> str | None:
Expand Down
Loading