feat: civitai token

.pre-commit-config.yaml

@@ -41,5 +41,5 @@ repos:
         - horde_safety
         - torch
         - ruamel.yaml
-        - hordelib==2.4.0
-        - horde_sdk==0.7.33
+        - hordelib==2.4.2
+        - horde_sdk==0.7.34

bridgeData_template.yaml

@@ -43,6 +43,14 @@ safety_on_gpu: false
 # Users in priority_usernames and trusted users will bypass this restriction
 require_upfront_kudos: false
 
+# If set, this worker will use this civitai API token when downloading any resources from civitai.
+# This is required in order to provide LoRas/TIs (or other resources)
+# which are marked as requiring a civitai token to download.
+#
+# You can get your civitai API Key from https://civitai.com/user/account (look for 'Add API Key')
+#
+# Remove the # from the line below and add your civitai API token to enable this feature.
+# civitai_api_token:
 
 #######################################
 ## Dreamer (Stable Diffusion Worker) ##

horde_worker_regen/bridge_data/data_model.py

@@ -29,6 +29,11 @@ class reGenBridgeData(CombinedHordeBridgeData):
         default=False,
     )
 
+    CIVIT_API_TOKEN: str | None = Field(
+        default=None,
+        alias="civitai_api_token",
+    )
+
     def load_env_vars(self) -> None:
         """Load the environment variables into the config model."""
         if self.models_folder_parent and os.getenv("AIWORKER_CACHE_HOME") is None:

horde_worker_regen/process_management/process_manager.py

@@ -2094,7 +2094,7 @@ async def api_job_pop(self) -> None:
             job_pop_request = ImageGenerateJobPopRequest(
                 apikey=self.bridge_data.api_key,
                 name=self.bridge_data.dreamer_worker_name,
-                bridge_agent="AI Horde Worker reGen:4.1.0:https://github.com/Haidra-Org/horde-worker-reGen",
+                bridge_agent="AI Horde Worker reGen:4.1.1:https://github.com/Haidra-Org/horde-worker-reGen",
                 models=self.bridge_data.image_models_to_load,
                 nsfw=self.bridge_data.nsfw,
                 threads=self.max_concurrent_inference_processes,

load_env_vars.py

@@ -45,3 +45,11 @@ def load_env_vars() -> None:  # FIXME: there is a dynamic way to do this
                 "AIWORKER_LORA_CACHE_SIZE environment variable already set. "
                 "This will override the value for `max_lora_cache_size` in the config file.",
             )
+    if "civitai_api_token" in config:
+        if os.getenv("CIVIT_API_TOKEN") is None:
+            os.environ["CIVIT_API_TOKEN"] = config["civitai_api_token"]
+        else:
+            print(
+                "CIVIT_API_TOKEN environment variable already set. "
+                "This will override the value for `civitai_api_token` in the config file.",
+            )

requirements.txt

@@ -1,13 +1,14 @@
 --extra-index-url https://download.pytorch.org/whl/cu121
 torch>=2.1.2
 
-horde_sdk~=0.7.33
+horde_sdk~=0.7.36
 horde_safety~=0.2.3
-hordelib~=2.4.0
+hordelib~=2.4.2
 horde_model_reference
 
 python-dotenv
 ruamel.yaml
+semver
 
 python-Levenshtein