diff --git a/hordelib/model_manager/lora.py b/hordelib/model_manager/lora.py index 4caa3b63..d472fb21 100644 --- a/hordelib/model_manager/lora.py +++ b/hordelib/model_manager/lora.py @@ -36,9 +36,9 @@ class DOWNLOAD_SIZE_CHECK(StrEnum): class LoraModelManager(BaseModelManager): LORA_DEFAULTS = "https://raw.githubusercontent.com/Haidra-Org/AI-Horde-image-model-reference/main/lora.json" LORA_API = "https://civitai.com/api/v1/models?types=LORA&sort=Highest%20Rated&primaryFileOnly=true" - MAX_RETRIES = 10 if not TESTS_ONGOING else 1 + MAX_RETRIES = 10 if not TESTS_ONGOING else 3 MAX_DOWNLOAD_THREADS = 3 - RETRY_DELAY = 5 + RETRY_DELAY = 5 if not TESTS_ONGOING else 0.2 """The time to wait between retries in seconds""" REQUEST_METADATA_TIMEOUT = 30 """The time to wait for a response from the server in seconds""" @@ -84,6 +84,7 @@ def __init__( self._stop_all_threads = False self._index_ids = {} self._index_orig_names = {} + self.total_retries_attempted = 0 models_db_path = LEGACY_REFERENCE_FOLDER.joinpath("lora.json").resolve() @@ -182,7 +183,11 @@ def _get_json(self, url): return response.json() except (requests.HTTPError, requests.ConnectionError, requests.Timeout, json.JSONDecodeError): + # CivitAI Errors when the model ID is too long + if response.status_code in [404, 500]: + return None retries += 1 + self.total_retries_attempted += 1 if retries <= self.MAX_RETRIES: time.sleep(self.RETRY_DELAY) else: @@ -379,6 +384,7 @@ def _download_thread(self, thread_number): logger.debug(f"Fatal error downloading {lora['filename']} {e}. Retry {retries}/{self.MAX_RETRIES}") retries += 1 + self.total_retries_attempted += 1 if retries > self.MAX_RETRIES: break # fail diff --git a/tests/model_managers/test_mm_lora.py b/tests/model_managers/test_mm_lora.py index ab53dd12..ce78526e 100644 --- a/tests/model_managers/test_mm_lora.py +++ b/tests/model_managers/test_mm_lora.py @@ -148,6 +148,36 @@ def test_adhoc_non_existing(self): assert not lora_model_manager.is_model_available(lora_name) lora_model_manager.stop_all() + def test_adhoc_non_existing_intstring_small(self): + lora_model_manager = LoraModelManager( + download_wait=False, + allowed_adhoc_lora_storage=1024, + ) + lora_model_manager.download_default_loras() + lora_model_manager.wait_for_downloads(600) + lora_model_manager.wait_for_adhoc_reset(15) + lora_name = "12345" + lora_key = lora_model_manager.fetch_adhoc_lora(lora_name) + assert lora_model_manager.total_retries_attempted == 0 + assert lora_key is None + assert not lora_model_manager.is_model_available(lora_name) + lora_model_manager.stop_all() + + def test_adhoc_non_existing_intstring_large(self): + lora_model_manager = LoraModelManager( + download_wait=False, + allowed_adhoc_lora_storage=1024, + ) + lora_model_manager.download_default_loras() + lora_model_manager.wait_for_downloads(600) + lora_model_manager.wait_for_adhoc_reset(15) + lora_name = "99999999999999" + lora_key = lora_model_manager.fetch_adhoc_lora(lora_name) + assert lora_model_manager.total_retries_attempted == 0 + assert lora_key is None + assert not lora_model_manager.is_model_available(lora_name) + lora_model_manager.stop_all() + ## Disabling this until I can figure out a better way to test these # def test_unused_loras(self): # lora_model_manager = LoraModelManager(