From 345825bbca46b64ea8bb459fb2fea04b6b1b4620 Mon Sep 17 00:00:00 2001 From: tazlin Date: Thu, 25 Jan 2024 18:01:43 -0500 Subject: [PATCH] feat: civitai token --- .pre-commit-config.yaml | 4 ++-- bridgeData_template.yaml | 8 ++++++++ horde_worker_regen/bridge_data/data_model.py | 5 +++++ horde_worker_regen/process_management/process_manager.py | 2 +- load_env_vars.py | 8 ++++++++ requirements.txt | 5 +++-- 6 files changed, 27 insertions(+), 5 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 00d729dc..f7920fdf 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/bridgeData_template.yaml b/bridgeData_template.yaml index d4156310..53c8e9ab 100644 --- a/bridgeData_template.yaml +++ b/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) ## diff --git a/horde_worker_regen/bridge_data/data_model.py b/horde_worker_regen/bridge_data/data_model.py index 7e5dfcb4..cf29575d 100644 --- a/horde_worker_regen/bridge_data/data_model.py +++ b/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: diff --git a/horde_worker_regen/process_management/process_manager.py b/horde_worker_regen/process_management/process_manager.py index 1d84c3d2..f2e3d21f 100644 --- a/horde_worker_regen/process_management/process_manager.py +++ b/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, diff --git a/load_env_vars.py b/load_env_vars.py index 6b82f024..c6eb744c 100644 --- a/load_env_vars.py +++ b/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.", + ) diff --git a/requirements.txt b/requirements.txt index acd5d06e..9f4ebaf6 100644 --- a/requirements.txt +++ b/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