feat: support more model load instructions

See Haidra-Org/horde-sdk#155 for more info.

.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

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

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

horde_worker_regen/__init__.py

@@ -8,4 +8,4 @@
 
 ASSETS_FOLDER_PATH = Path(__file__).parent / "assets"
 
-__version__ = "4.2.4"
+__version__ = "4.2.5"

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)))
 

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 = "[email protected]"},

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

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")