From f42b7e80a25bf514fcf0c5e2f0e7a11f3a7d15e8 Mon Sep 17 00:00:00 2001
From: tazlin <tazlin.on.github@gmail.com>
Date: Tue, 5 Mar 2024 10:03:54 -0500
Subject: [PATCH] feat: support more model load instructions

See https://github.com/Haidra-Org/horde-sdk/pull/155 for more info.
---
 .pre-commit-config.yaml                       |  4 +--
 bridgeData_template.yaml                      | 13 +++++--
 horde-bridge.cmd                              |  2 +-
 horde_worker_regen/__init__.py                |  2 +-
 horde_worker_regen/bridge_data/load_config.py |  8 +++++
 pyproject.toml                                |  2 +-
 requirements.txt                              |  4 +--
 tests/test_bridge_data.py                     | 34 ++++++++++++++++++-
 8 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 4b5319b7..bb316665 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -41,6 +41,6 @@ repos:
         - torch
         - ruamel.yaml
         - hordelib==2.6.2
-        - horde_sdk==0.8.0
-        - horde_model_reference==0.6.2
+        - horde_sdk==0.8.1
+        - horde_model_reference==0.6.3
         - semver
diff --git a/bridgeData_template.yaml b/bridgeData_template.yaml
index afb0b221..920d3c63 100644
--- a/bridgeData_template.yaml
+++ b/bridgeData_template.yaml
@@ -151,9 +151,16 @@ disable_disk_cache: false # Currently unused in reGen
 #   "ALL"  - means load all possible models. Expect this to take over 1TB of space!
 #   "TOP n"  - load the top "N" most popular models, use for example, "top 5" or "top 3", etc.
 #   "BOTTOM n"  - load the top "N" most popular models, use for example, "top 5" or "top 3", etc.
+#
+#   "ALL SD15 MODELS" - All Stable Diffusion 1.5 models
+#   "ALL SD21 MODELS" - All Stable Diffusion 2.0/2.1 models
+#   "ALL SDXL MODELS" - All Stable Diffusion XL models
+#   "ALL INPAINTING MODELS"  -  All models marked as being for inpainting
+#
+#   "ALL SFW MODELS"  -  All models marked as being SFW
+#   "ALL NSFW MODELS"  -  All models marked as being NSFW
+#
 #   (not currently supported) "ALL <style> MODELS"  -  For example, "all anime models", styles are: generalist, artistic, realistic, anime, furry, other
-#   (not currently supported) "ALL SFW MODELS"  -  All models marked as being SFW
-#   (not currently supported) "ALL NSFW MODELS"  -  All models marked as being NSFW
 #
 # The official model reference is (in json format) found at https://github.com/Haidra-Org/AI-Horde-image-model-reference/blob/main/stable_diffusion.json.
 # Several front ends also a list of model names.
@@ -163,6 +170,7 @@ models_to_load:
   - "top 2"
   #- "ALL MODELS"
   #- "TOP 3"
+  #- "ALL SFW"
   #- "stable_diffusion_2.1"
   #- "stable_diffusion"
   #- "Anything Diffusion"
@@ -185,6 +193,7 @@ models_to_skip:
   - "pix2pix" # Not currently supported
   - "SDXL_beta::stability.ai#6901" # Do not remove this, as this model would never work
   - A to Zovya RPG # This model is known to cause problems with reGen
+  # - ALL NSFW MODELS
   #- "stable_diffusion_inpainting"  # Inpainting is generally quite heavy along with other models for smaller GPUs.
   #- "stable_diffusion_2.1"  # Stable diffusion 2.1 has bigger memory requirements than 1.5, so if your card cannot lift, it, disable it
   #- "stable_diffusion_2.0"  # Same as Stable diffusion 2.1
diff --git a/horde-bridge.cmd b/horde-bridge.cmd
index 14f4af0a..e6a1f42f 100644
--- a/horde-bridge.cmd
+++ b/horde-bridge.cmd
@@ -4,7 +4,7 @@ cd /d %~dp0
 : This first call to runtime activates the environment for the rest of the script
 call runtime python -s -m pip -V
 
-call python -s -m pip install horde_sdk~=0.8.0 horde_model_reference~=0.6.2 hordelib~=2.6.2 -U
+call python -s -m pip install horde_sdk~=0.8.1 horde_model_reference~=0.6.3 hordelib~=2.6.2 -U
 if %ERRORLEVEL% NEQ 0 (
     echo "Please run update-runtime.cmd."
     GOTO END
diff --git a/horde_worker_regen/__init__.py b/horde_worker_regen/__init__.py
index d1fcecd0..b0f7a2d6 100644
--- a/horde_worker_regen/__init__.py
+++ b/horde_worker_regen/__init__.py
@@ -8,4 +8,4 @@
 
 ASSETS_FOLDER_PATH = Path(__file__).parent / "assets"
 
-__version__ = "4.2.4"
+__version__ = "4.2.5"
diff --git a/horde_worker_regen/bridge_data/load_config.py b/horde_worker_regen/bridge_data/load_config.py
index a677420d..f6ed1d06 100644
--- a/horde_worker_regen/bridge_data/load_config.py
+++ b/horde_worker_regen/bridge_data/load_config.py
@@ -130,6 +130,14 @@ def _resolve_meta_instructions(  # FIXME: This should be moved into the SDK
                 AIHordeAPIManualClient(),
             )
 
+        if bridge_data.meta_skip_instructions is not None:
+            skip_models: set[str] = load_resolver.resolve_meta_instructions(
+                list(bridge_data.meta_skip_instructions),
+                AIHordeAPIManualClient(),
+            )
+            existing_skip_models = set(bridge_data.image_models_to_skip)
+            bridge_data.image_models_to_skip = list(existing_skip_models.union(skip_models))
+
         if resolved_models is not None:
             bridge_data.image_models_to_load = list(set(bridge_data.image_models_to_load + list(resolved_models)))
 
diff --git a/pyproject.toml b/pyproject.toml
index 8f5170e9..f6fbe2eb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "horde_worker_regen"
-version = "4.2.4"
+version = "4.2.5"
 description = "Allows you to connect to the AI Horde and generate images for users."
 authors = [
     {name = "tazlin", email = "tazlin.on.github@gmail.com"},
diff --git a/requirements.txt b/requirements.txt
index 284ada85..e540108b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,9 +1,9 @@
 torch>=2.1.2
 
-horde_sdk~=0.8.0
+horde_sdk~=0.8.1
 horde_safety~=0.2.3
 hordelib~=2.6.2
-horde_model_reference~=0.6.2
+horde_model_reference~=0.6.3
 
 python-dotenv
 ruamel.yaml
diff --git a/tests/test_bridge_data.py b/tests/test_bridge_data.py
index bcd262ba..087b5e0f 100644
--- a/tests/test_bridge_data.py
+++ b/tests/test_bridge_data.py
@@ -1,5 +1,8 @@
 # import yaml
+import pathlib
 
+import pytest
+from horde_model_reference.model_reference_manager import ModelReferenceManager
 from horde_sdk.generic_api.consts import ANON_API_KEY
 from ruamel.yaml import YAML
 
@@ -29,13 +32,42 @@ def test_bridge_data_yaml() -> None:
     assert len(parsed_bridge_data.meta_load_instructions) == 1
 
 
-def test_bridge_data_loader_yaml() -> None:
+def test_bridge_data_loader_yaml_template() -> None:
     bridge_data_loader = BridgeDataLoader()
+
+    horde_model_reference_manager = ModelReferenceManager(
+        download_and_convert_legacy_dbs=True,
+        override_existing=True,
+    )
     bridge_data = bridge_data_loader.load(
         file_path="bridgeData_template.yaml",
         file_format=ConfigFormat.yaml,
+        horde_model_reference_manager=horde_model_reference_manager,
     )
 
     assert bridge_data is not None
     assert bridge_data.disable_terminal_ui is False
     assert bridge_data.api_key == ANON_API_KEY
+
+
+def test_bridge_data_loader_yaml_local_if_present() -> None:
+    bridge_data_loader = BridgeDataLoader()
+
+    horde_model_reference_manager = ModelReferenceManager(
+        download_and_convert_legacy_dbs=True,
+        override_existing=True,
+    )
+
+    if pathlib.Path("bridgeData.yaml").is_file():
+        bridge_data = bridge_data_loader.load(
+            file_path="bridgeData.yaml",
+            file_format=ConfigFormat.yaml,
+            horde_model_reference_manager=horde_model_reference_manager,
+        )
+
+        assert bridge_data is not None
+        assert bridge_data.api_key != ANON_API_KEY
+        assert len(bridge_data.image_models_to_load) > 0
+
+    else:
+        pytest.skip("bridgeData.yaml not found")