fix: integrity check models too in download_models.py

Haidra-Org · Nov 14, 2023 · 1acbc78 · 1acbc78
1 parent c5da46a
commit 1acbc78
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/download_models.py b/download_models.py
@@ -85,6 +85,10 @@ def download_all_models() -> None:
             logger.error(f"Failed to download model {model}")
             any_model_failed_to_download = True
 
+        if not SharedModelManager.manager.compvis.validate_model(model):
+            SharedModelManager.manager.compvis.taint_model(model)
+            SharedModelManager.manager.compvis.download_model(model)
+
     if any_model_failed_to_download:
         logger.error("Failed to download all models.")
     else:

diff --git a/requirements.txt b/requirements.txt
@@ -3,7 +3,7 @@ torch>=2.1.0
 
 horde_sdk>=0.7.15
 git+https://github.com/Haidra-Org/horde-safety@nsfw-checking-and-anime-handling
-hordelib>=2.0.2
+hordelib>=2.1.2
 horde_model_reference
 
 python-dotenv