fix: avoid crash on long lora names

Haidra-Org · Mar 27, 2024 · ab57d9c · ab57d9c
1 parent 6637f76
commit ab57d9c
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git a/hordelib/model_manager/hyper.py b/hordelib/model_manager/hyper.py
@@ -181,9 +181,7 @@ def init_model_managers(
             self.active_model_managers.append(
                 resolve_manager_to_load_type(
                     multiprocessing_lock=multiprocessing_lock,
-                    civitai_api_token=(
-                        os.environ["CIVIT_API_TOKEN"] if "CIVIT_API_TOKEN" in os.environ.keys() else None
-                    ),
+                    civitai_api_token=os.getenv("CIVIT_API_TOKEN", None),
                 ),
             )
 

diff --git a/hordelib/model_manager/lora.py b/hordelib/model_manager/lora.py
@@ -393,7 +393,7 @@ def _parse_civitai_lora_data(self, item, adhoc=False):
         for file in version.get("files", {}):
             if file.get("primary", False) and file.get("name", "").endswith(".safetensors"):
                 sanitized_name = Sanitizer.sanitise_model_name(lora_name)
-                lora_filename = f'{Sanitizer.sanitise_filename(lora_name)}_{version.get("id", 0)}'
+                lora_filename = f'{Sanitizer.sanitise_filename(lora_name)[0:128]}_{version.get("id", 0)}'
                 lora_key = sanitized_name.lower().strip()
                 lora["name"] = lora_key
                 lora["orig_name"] = lora_name

diff --git a/tests/model_managers/test_mm_lora.py b/tests/model_managers/test_mm_lora.py
@@ -148,6 +148,22 @@ def test_fetch_specific_lora_version(self):
         assert lora_model_manager.get_lora_filename("22591") == "GAG-RPGPotionsLoRaXL_197256.safetensors"
         lora_model_manager.stop_all()
 
+    def test_lora_long_filename(self):
+        lora_model_manager = LoraModelManager(
+            download_wait=False,
+            allowed_adhoc_lora_storage=1024,
+            civitai_api_token=os.getenv("CIVIT_API_TOKEN", None),
+        )
+        lora_model_manager.download_default_loras()
+        lora_model_manager.wait_for_downloads(600)
+        lora_model_manager.wait_for_adhoc_reset(15)
+
+        lora_model_manager.ensure_lora_deleted(189973)
+        lora_model_manager.fetch_adhoc_lora("372094", is_version=True)
+        lora_dict = lora_model_manager.get_model_reference_info("189973")
+        assert lora_model_manager.find_latest_version(lora_dict) == "372094"
+        lora_model_manager.stop_all()
+
     def test_reject_adhoc_nsfw_lora(self):
         lora_model_manager = LoraModelManager(
             download_wait=False,