feat: beta

.github/workflows/maintests.yml

@@ -12,6 +12,10 @@ on:
       - '.github/workflows/prtests.yml'
       - '.github/workflows/release.yml'
       - '.pre-commit-config.yaml'
+      - 'bridgeData_template.yaml'
+      - 'requirements.txt'
+      - 'requirements.dev.txt'
+      - 'requirements.docs.txt'
 jobs:
   build:
     runs-on: ubuntu-latest

.github/workflows/prtests.yml

@@ -15,6 +15,10 @@ on:
       - '.github/workflows/prtests.yml'
       - '.github/workflows/release.yml'
       - '.pre-commit-config.yaml'
+      - 'bridgeData_template.yaml'
+      - 'requirements.txt'
+      - 'requirements.dev.txt'
+      - 'requirements.docs.txt'
 jobs:
   build:
     runs-on: ubuntu-latest

bridgeData_template.yaml

@@ -0,0 +1,175 @@
+## Common for all worker Types
+
+# The horde url
+horde_url: "https://aihorde.net"
+# Give a cool name to your instance
+worker_name: "An Awesome AI Horde Worker"
+# The api_key identifies a unique user in the horde
+# Visit https://stablehorde.net/register to create one before you can join
+api_key: "0000000000"
+# Put other users whose prompts you want to prioritize.
+# The owner's username is always included so you don't need to add it here,
+priority_usernames: []
+# The amount of parallel jobs to pick up for the horde. Each running job will consume the amount of RAM needed to run each model, and will also affect the speed of other running jobs
+# so make sure you have enough VRAM to load models in parallel, and that the speed of fulfilling requests is not too slow
+# Expected limit per VRAM size: <6 VRAM: 1, <=8 VRAM: 2, <=12 VRAM:3, <=14 VRAM: 4
+# But remember that the speed of your gens will also be affected for each parallel job
+max_threads: 1
+# We will keep this many requests in the queue so we can start working as soon as a thread is available
+# Recommended to keep no higher than 1
+queue_size: 0
+# If set to True, this worker will not only pick up jobs where the user has the required kudos upfront.
+# Effectively this will exclude all anonymous accounts, and registered accounts who haven't contributed.
+# Users in priority_usernames and trusted users will bypass this restriction
+require_upfront_kudos: false
+# If you set this to True, the worker will detect the most popular models and load them automatically ( Defaults to True if missing )
+# Note this ultimately overrides the models_to_load list
+
+## Dreamer (Stable Diffusion Worker)
+
+# The name to use when running a dreamer instance. Will default to `worker_name` if not set
+dreamer_name: "An Awesome Dreamer"
+# The amount of power your system can handle
+# 8 means 512*512. Each increase increases the possible resoluion by 64 pixes
+# So if you put this to 2 (the minimum, your SD can only generate 64x64 pixels
+# If you put this to 32, it is equivalent to 1024x1024 pixels
+max_power: 8
+# Set this to false, if you do not want your worker to receive requests for NSFW generations
+nsfw: true
+# Set this to True if you want your worker to censor NSFW generations. This will only be active is horde_nsfw == False
+censor_nsfw: false
+# A list of words which you do not want to your worker to accept
+blacklist: []
+# A list of words for which you always want to allow the NSFW censor filter, even when this worker is in NSFW mode
+censorlist: []
+# If set to False, this worker will no longer pick img2img jobs
+allow_img2img: true
+# If set to True, this worker will can pick inpainting jobs
+allow_painting: true
+# If set to False, this worker will no longer pick img2img jobs from unsafe IPs
+allow_unsafe_ip: true
+# If set to False, this worker will not load post-processors like Codeformers and will not pick up jobs which require post-processing
+# In the future this will be adjusted so that post-processing can be split from image generation
+allow_post_processing: true
+# If set to True, this worker start picking up ControlNet jobs
+# ControlNet is really heavy and requires a good GPU with at least 12G VRAM to run.
+# If your controlnet jobs crash by running out of CUDA VRAM, set this to false
+allow_controlnet: false # needs at least 12G VRAM
+# If set to True, this worker start picking up jobs requesting LoRas
+# Your worker will download the top 10Gb of non-character LoRas
+# and then will ad-hoc download any LoRa requested which you do not have, and cache that for a number a days
+allow_lora: false
+# Use this setting to control how much extra space LoRas can take after you downloaded the Top
+# If a new Lora would exceed this space, an old lora you've downloaded previously will be deleted
+# Note THIS IS ON TOP OF THE CURATED LORAS, so plan around +5G more than this
+max_lora_cache_size: 10 # In gigabytes. Min is 10.
+# Set to False to prevent this worker from reading the Horde model queue and loading models which are under load
+dynamic_models: false
+# Adjust how many models to load into memory. In future this will likely be an argument for memory size or may disappear, but for right now, I'm lazy
+number_of_dynamic_models: 0
+# The maximum amount of models to download dynamically for this worker. Increase this amount of you have plenty of space. Keep it low if you do not
+# When the amount of models downloaded reaches this amount, the dynamic list will only use dynamic models already downloaded
+# Therefore make sure you put some generalist and popular models in your models_to_load list if this number is small!
+max_models_to_download: 10
+# The frequency (in seconds) to output worker summary stats, such as kudos per hour.
+# Set to zero to disable stats output completely.
+stats_output_frequency: 30
+# The location in which stable diffusion ckpt models are stored
+cache_home: "./"
+# Always download models when required without prompting
+always_download: true
+# The location of the temp directory, also used for the model cache
+temp_dir: "./tmp"
+# Disable the terminal GUI, which displays information about the worker and the horde.
+disable_terminal_ui: false
+# VRAM to leave unused, as a percentage or in MB. VRAM the worker can use will be used to load and cache models.
+# Note this NOT the amount of VRAM to use, it's the amount to KEEP FREE. So if something else starts using
+# VRAM the worker will attempt to release it to allow the other software to use it.
+# Don't set this too high, or you will run out of vram when it can't be released fast enough.
+vram_to_leave_free: "80%"
+# RAM to leave unused, as a percentage or in MB. RAM the worker can use will be used to cache models. Same
+# notes as for VRAM.
+# Don't set this too high or your OS will likely start using lots of swap space and everything will slow down.
+ram_to_leave_free: "80%"
+# Disable the disk cache. By default if RAM and VRAM are filled (up to the limits above) then models will
+# spill over in to a disk cache. If you don't want this to happen you can disable it here. Note that if you
+# disable disk cache and specify more models to load than will fit in memory your worker will endlessly cycle
+# loading and unloading models.
+disable_disk_cache: false
+
+
+# The models to use. You can select a different main model, or select more than one.
+# With you can easily load 20 of these models with 32Gb RAM and 6G VRAM.
+# Adjust how many models you load based on how much RAM (not VRAM) you have available.
+# The last model in this list takes priority when the client accepts more than 1
+# if you do not know which models you can add here, use the below command
+# python show_available_models.py
+## WARNING: In case you have dynamic models this list is instead specifying models to always load along with the dynamic models!
+# So your total list will be your specific models + the dynamic models
+# in that case, keep this list short, preferrably to only a few more obscure models you'd like to always see available
+# Instead of a model name you may use of any of the following magic constants:
+#   "ALL MODELS"  - 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.
+#   "ALL <style> MODELS"  -  For example, "all anime models", styles are: generalist, artistic, realistic, anime, furry, other
+#   "ALL SFW MODELS"  -  All models marked as being SFW
+#   "ALL NSFW MODELS"  -  All models marked as being NSFW
+models_to_load:
+  - "top 2"
+  #- "ALL MODELS"
+  #- "TOP 3"
+  #- "stable_diffusion_2.1"
+  #- "stable_diffusion"
+  #- "Anything Diffusion"
+  #- "Yiffy"
+  #- "waifu_diffusion"
+  #- "Arcane Diffusion"
+  #- "Spider-Verse Diffusion"
+  #- "Elden Ring Diffusion"
+  #- "Robo-Diffusion"
+  #- "mo-di-diffusion"
+  #- "Knollingcase"
+  #- "stable_diffusion_inpainting"
+
+# This is used when dynamic_models == True or TOP n models are selected in models_to_load
+# The models in this list will not be loaded when they exist in the top models
+# This is to avoid loading models which you do not want either due to VRAM constraints, or due to NSFW content
+models_to_skip:
+  - "pix2pix"
+  #- "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
+  ## Popular NSFW models:
+  #- "Zeipher Female Model"
+  #- "Hentai Diffusion"
+
+# If you are getting messages about jobs taking too long, you can change this to true if you no longer want to see them
+# Please note, that if you *are* getting these messages, you are serving jobs substantially slower than is ideal,
+# and you very likely would get more kudos/hr if you just lower your max_power.
+suppress_speed_warnings: false
+
+## Scribe (LLM Worker)
+
+# The name to use when running a scribe worker. Will default to `worker_name` if not set
+scribe_name: "An Awesome Scribe"
+# The KoboldAI Client API URL
+kai_url: "http://localhost:5000"
+# The max amount of tokens to generate with this worker
+max_length: 80
+# The max tokens to use from the prompt
+max_context_length: 1024
+# When set to true, the horde alias behind the API key will be appended to the model that is advertised to the horde
+# This will prevent the model from being used from the shared pool, but will ensure that no other worker
+# can pretend to serve it
+branded_model: true
+
+## Alchemist (Image interrogation and post-processing)
+
+# The name to use when running an alchemist worker. Will default to `worker_name` if not set
+alchemist_name: "An Awesome Alchemist"
+# The alchemy forms this worker can serve.
+forms:
+  - "caption"
+  - "nsfw" # uses CPU
+  # Heavier than the others, but rewards more kudos
+  - "interrogation"
+  - "post-process"

horde_worker_regen/__init__.py

@@ -0,0 +1,13 @@
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from pathlib import Path  # noqa: E402
+
+ASSETS_FOLDER_PATH = Path(__file__).parent / "assets"
+
+from horde_worker_regen.process_management.main_entry_point import start_working  # noqa: E402
+
+__all__ = [
+    "start_working",
+]

horde_worker_regen/bridge_data/data_model.py

@@ -0,0 +1,38 @@
+import os
+
+from horde_sdk.ai_horde_worker.bridge_data import CombinedHordeBridgeData
+from loguru import logger
+from pydantic import Field
+
+from horde_worker_regen.locale_info.regen_bridge_data_fields import BRIDGE_DATA_FIELD_DESCRIPTIONS
+
+
+class reGenBridgeData(CombinedHordeBridgeData):
+    disable_terminal_ui: bool = Field(
+        value=True,
+    )
+
+    def load_env_vars(self) -> None:
+        """Load the environment variables into the config model."""
+        if self.models_folder_parent:
+            os.environ["AIWORKER_CACHE_HOME"] = self.models_folder_parent
+        if self.horde_url:
+            if os.environ.get("AI_HORDE_URL"):
+                logger.warning(
+                    "AI_HORDE_URL environment variable already set. This will override the value for `horde_url` in "
+                    "the config file.",
+                )
+            else:
+                if os.environ.get("AI_HORDE_DEV_URL"):
+                    logger.warning(
+                        "AI_HORDE_DEV_URL environment variable already set. This will override the value for "
+                        "`horde_url` in the config file.",
+                    )
+                else:
+                    os.environ["AI_HORDE_URL"] = self.horde_url
+
+
+# Dynamically add descriptions to the fields of the model
+for field_name, field in reGenBridgeData.model_fields.items():
+    if field_name in BRIDGE_DATA_FIELD_DESCRIPTIONS:
+        field.description = BRIDGE_DATA_FIELD_DESCRIPTIONS[field_name]

horde_worker_regen/bridge_data/load_config.py

@@ -0,0 +1,82 @@
+import json
+from enum import auto
+from pathlib import Path
+
+import yaml
+from strenum import StrEnum
+
+from horde_worker_regen.bridge_data.data_model import reGenBridgeData
+
+
+class UnsupportedConfigFormat(Exception):
+    def __init__(self, file_path: str | Path, file_format: str) -> None:
+        super().__init__(f"Unsupported config file format: {file_format} ({file_path})")
+
+
+class ConfigFormat(StrEnum):
+    yaml = auto()
+    json = auto()
+
+
+class BridgeDataLoader:
+    @staticmethod
+    def _infer_format(file_path: str | Path) -> ConfigFormat:
+        """Infer the config file format from the file extension.
+
+        Args:
+            file_path (str | Path): The path to the config file.
+
+        Returns:
+            ConfigFormat: The config file format.
+
+        Raises:
+            UnsupportedConfigFormat: If the config file format is not supported.
+        """
+        file_path = Path(file_path)
+
+        if file_path.suffix == ".yaml" or file_path.suffix == ".yml":
+            return ConfigFormat.yaml
+
+        if file_path.suffix == ".json":
+            return ConfigFormat.json
+
+        raise UnsupportedConfigFormat(file_path, file_path.suffix)
+
+    @staticmethod
+    def load(
+        file_path: str | Path,
+        *,
+        file_format: ConfigFormat | None = None,
+    ) -> reGenBridgeData:
+        """Load the config file and validate it.
+
+        Args:
+            file_path (str | Path): The path to the config file.
+            file_format (ConfigFormat | None, optional): The config file format. Defaults to None.
+            The file format will be inferred from the file extension if not provided.
+
+        Returns:
+            reGenBridgeData: The validated config file.
+
+        Raises:
+            ValidationError: If the config file is invalid.
+            UnsupportedConfigFormat: If the config file format is not supported.
+
+        """
+        # Infer the file format if not provided
+        if not file_format:
+            file_format = BridgeDataLoader._infer_format(file_path)
+
+        if file_format == ConfigFormat.yaml:
+            with open(file_path) as f:
+                config = yaml.safe_load(f)
+
+            return reGenBridgeData.model_validate(config)
+
+        if file_format == ConfigFormat.json:
+            with open(file_path) as f:
+                config = json.load(f)
+
+            return reGenBridgeData.model_validate(config)
+
+        raise UnsupportedConfigFormat(file_path, file_format)

horde_worker_regen/consts.py

@@ -0,0 +1,4 @@
+BRIDGE_VERSION = 1
+BRIDGE_IDENTIFIER = "reGen"
+RELEASE = f"{BRIDGE_VERSION}.0.0"  # FIXME
+BRIDGE_CONFIG_FILENAME = "bridgeData.yaml"

horde_worker_regen/locale_info/README.md

@@ -0,0 +1 @@
+If you're making changes to the files in this folder, please only **append** entries, and avoid inserting or otherwise altering the order, as it greatly complicates the localizing process.

horde_worker_regen/locale_info/regen_bridge_data_fields.py

@@ -0,0 +1,5 @@
+from horde_worker_regen.localize import _L
+
+BRIDGE_DATA_FIELD_DESCRIPTIONS = {
+    "disable_terminal_ui": _L("Disable the terminal GUI, which displays information about the worker and the horde."),
+}

horde_worker_regen/localize.py

@@ -0,0 +1,6 @@
+def _L(s: str) -> str:
+    """Indicates that the string is displayed to the user and should be localized."""
+    return s
+
+
+# TODO REMOVE THIS?

horde_worker_regen/process_management/_aliased_types.py

@@ -0,0 +1,9 @@
+import multiprocessing.queues
+from typing import TYPE_CHECKING
+
+from horde_worker_regen.process_management.messages import HordeProcessMessage
+
+if TYPE_CHECKING:  # noqa: SIM108 (breaks mypy)
+    ProcessQueue = multiprocessing.Queue[HordeProcessMessage]
+else:
+    ProcessQueue = multiprocessing.Queue