Skip to content

Commit

Permalink
Merge pull request #117 from Haidra-Org/main
Browse files Browse the repository at this point in the history
fix: more graceful civitai lora handling
  • Loading branch information
tazlin authored Dec 9, 2023
2 parents 018d060 + 51fff6a commit 9be8121
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 2 deletions.
10 changes: 8 additions & 2 deletions hordelib/model_manager/lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down
30 changes: 30 additions & 0 deletions tests/model_managers/test_mm_lora.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 9be8121

Please sign in to comment.