From 2948c5899d783a69fd77ecb502b5c7b6903a4dbd Mon Sep 17 00:00:00 2001
From: db0 <mail@dbzer0.com>
Date: Fri, 8 Dec 2023 17:01:12 +0100
Subject: [PATCH] fix: Handle missing lora IDs

---
 hordelib/model_manager/lora.py       | 10 ++++++++--
 tests/model_managers/test_mm_lora.py | 30 ++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

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(