diff --git a/hordelib/model_manager/lora.py b/hordelib/model_manager/lora.py index 9a27c45..30568c7 100644 --- a/hordelib/model_manager/lora.py +++ b/hordelib/model_manager/lora.py @@ -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() @@ -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 @@ -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 @@ -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( @@ -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: