Merge pull request #247 from Haidra-Org/main

feat: custom models, comfyui `16eabdf7`; project rename to `horde-engine`

.github/workflows/maintests.yml

@@ -1,4 +1,4 @@
-name: Tests
+name: Main Tests
 
 on:
   push:
@@ -39,7 +39,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install --upgrade -r requirements.dev.txt
-      - name: Check build_helper.py hordelib imports have no breaking dependency changes
+      - name: Check build_helper.py horde-engine imports have no breaking dependency changes
         run: tox -e test-build-helper
       - name: Build unit test environment, confirm CUDA is available on host
         run: tox -e tests -- -k test_cuda

.github/workflows/prtests.yml

@@ -1,4 +1,4 @@
-name: Unstable Tests
+name: Pull Request Tests
 
 on:
   pull_request_target:
@@ -44,7 +44,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install --upgrade -r requirements.dev.txt
-      - name: Check build_helper.py hordelib imports have no breaking dependency changes
+      - name: Check build_helper.py horde-engine imports have no breaking dependency changes
         run: tox -e test-build-helper
       - name: Build unit test environment, confirm CUDA is available on host
         run: tox -e tests -- -k test_cuda

.github/workflows/release.yml

@@ -1,4 +1,4 @@
-name: hordelib build and publish
+name: horde-engine build and publish
 
 on:
   push:
@@ -11,8 +11,13 @@ permissions:
 
 jobs:
   build-n-publish:
-    name: hordelib build and publish
+    name: horde-engine build and publish
     runs-on: ubuntu-latest
+
+    environment: # see https://github.com/pypa/gh-action-pypi-publish/tree/release/v1/
+      name: pypi
+      url: https://pypi.org/p/horde-engine/
+
     steps:
 
     # Version bumps the project and creates a tag by creating
@@ -46,8 +51,8 @@ jobs:
         python-version: "3.10"
 
     # Install build deps
-    # NOTE If any hordelib imports used by build_helper.py are changed, or the specific modules
-    # imported from hordelib depend on a package not included here, running build_helper.py later on will fail.
+    # NOTE If any horde-engine imports used by build_helper.py are changed, or the specific modules
+    # imported from horde-engine depend on a package not included here, running build_helper.py later on will fail.
     # See `build_helper.py` for more information.
     - name: "🛠 Install pypa/build"
       if: ${{ steps.release.outputs.version != '' }}
@@ -70,8 +75,6 @@ jobs:
       with:
         add: 'CHANGELOG.md'
         message: 'ci: update changelog'
-        committer_name: GitHub Actions
-        committer_email: [email protected]
 
     # Patches our requirements.txt and pyproject.toml
     # Build a changelog
@@ -104,9 +107,9 @@ jobs:
         curl -X PURGE
         https://camo.githubusercontent.com/769edfb1778d4cbc3f93bc5ad0be9597bbd2d9c162cc1e9fb44172a5b660af01/68747470733a2f2f706570792e746563682f62616467652f686f7264656c6962
 
-    # - name: "Inform with Discord Webhook"
-    #   if: ${{ steps.release.outputs.version != '' }}
-    #   uses: tsickert/[email protected]
-    #   with:
-    #     webhook-url: ${{ secrets.DISCORD_WEBHOOK_URL }}
-    #     content: "New version of hordelib has been published to pypi: ${{ steps.release.outputs.version }}. You can update your worker by running `./update-runtime.sh --hordelib` (Linux) or `update-runtime.cmd --hordelib` (Windows). Changelog: https://t.ly/z2vQ"
+    - name: "Inform with Discord Webhook"
+      if: ${{ steps.release.outputs.version != '' }}
+      uses: tsickert/[email protected]
+      with:
+        webhook-url: ${{ secrets.DISCORD_WEBHOOK_URL }}
+        content: "New version of horde-engine has been published to pypi: ${{ steps.release.outputs.version }}. Changelog: https://t.ly/z2vQ"

.pre-commit-config.yaml

@@ -1,21 +1,21 @@
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
 -   repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.4.2
     hooks:
     -   id: black
         exclude: ^hordelib/nodes/.*\..*$
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.3.3
+    rev: v0.4.3
     hooks:
     -    id: ruff
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: 'v1.9.0'
+    rev: 'v1.10.0'
     hooks:
     -   id: mypy
         exclude: ^examples/.*$ # FIXME

README.md

@@ -1,4 +1,4 @@
-# hordelib
+# horde-engine
 
 [![PyPI Version][pypi-image]][pypi-url]
 [![Downloads][downloads-image]][downloads-url]
@@ -10,11 +10,13 @@
 [![All Models][all-model-images]][all-model-url]
 [![Release Changelog][changelog-image]][changelog-url]
 
-`hordelib` is a wrapper around [ComfyUI](https://github.com/comfyanonymous/ComfyUI) primarily to enable the [AI Horde](https://aihorde.net/) to run inference pipelines designed visually in the ComfyUI GUI.
+> Note: This project was formerly known as `hordelib`. The project namespace will be changed in the near future to reflect this change.
 
-The developers of `hordelib` can be found in the AI Horde Discord server: [https://discord.gg/3DxrhksKzn](https://discord.gg/3DxrhksKzn)
+`horde-engine` is a wrapper around [ComfyUI](https://github.com/comfyanonymous/ComfyUI) primarily to enable the [AI Horde](https://aihorde.net/) to run inference pipelines designed visually in the ComfyUI GUI.
 
-`hordelib` has been the default inference backend library of the [AI Horde](https://aihorde.net/) since `hordelib` v1.0.0.
+The developers of `horde-engine` can be found in the AI Horde Discord server: [https://discord.gg/3DxrhksKzn](https://discord.gg/3DxrhksKzn)
+
+Note that `horde-engine` (previously known as `hordelib`) has been the default inference backend library of the [AI Horde](https://aihorde.net/) since `hordelib` v1.0.0.
 
 ## Purpose
 
@@ -45,19 +47,25 @@ If you only have 16GB of RAM you will also need swap space. So if you typically
 Horde payloads can be processed simply with (for example):
 
 ```python
-import os
+# import os
+# Wherever your models are
+# os.environ["AIWORKER_CACHE_HOME"] = "f:/ai/models" # Defaults to `models/` in the current working directory
+
 import hordelib
-hordelib.initialise()
+
+hordelib.initialise()  # This must be called before any other hordelib functions
 
 from hordelib.horde import HordeLib
 from hordelib.shared_model_manager import SharedModelManager
 
-# Wherever your models are
-os.environ["AIWORKER_CACHE_HOME"] = "f:/ai/models"
-
 generate = HordeLib()
-SharedModelManager.loadModelManagers(compvis=True)
-SharedModelManager.manager.load("Deliberate")
+
+if SharedModelManager.manager.compvis is None:
+    raise Exception("Failed to load compvis model manager")
+
+SharedModelManager.manager.compvis.download_model("Deliberate")
+SharedModelManager.manager.compvis.validate_model("Deliberate")
+
 
 data = {
     "sampler_name": "k_dpmpp_2m",
@@ -79,7 +87,12 @@ data = {
     "model": "Deliberate",
 }
 pil_image = generate.basic_inference_single_image(data).image
+
+if pil_image is None:
+    raise Exception("Failed to generate image")
+
 pil_image.save("test.png")
+
 ```
 
 Note that `hordelib.initialise()` will erase all command line arguments from argv. So make sure you parse them before you call that.
@@ -88,7 +101,7 @@ See `tests/run_*.py` for more standalone examples.
 
 ### Logging
 
-If you don't want `hordelib` to setup and control the logging configuration initialise with:
+If you don't want `hordelib` to setup and control the logging configuration (we use [loguru](https://loguru.readthedocs.io/en/stable/)) initialise with:
 
 ```python
 import hordelib
@@ -111,14 +124,6 @@ Custom nodes for ComfyUI providing Controlnet preprocessing capability. Licened
 
 Custom nodes for ComfyUI providing face restoration.
 
-## DMCA Abuse
-
-On 26th May 2023 an [individual](https://github.com/hlky) issued a [DMCA takedown notice](https://github.com/github/dmca/blob/master/2023/05/2023-05-26-nataili.md) to Github against `hordelib` which claimed their name had been removed from the copyright header in the AGPL license in the 7 files listed in the takedown notice. This claim was true, and this attribution had been removed by a `hordelib` contributor prior to being committed into the `hordelib` repository.
-
-Unfortunately, it appears the individual making the DMCA claim was acting in bad faith, and even when their name was restored to the copyright attribution in the files, they persisted to press the DMCA takedown claim, which, due to the nature of the Github process, resulted in hordelib being subject to a DMCA takedown on Github.
-
-This version of `hordelib` has the 7 files mentioned in the DMCA takedown removed and replaced with alternatives and the functionality required to run the AI Horde Worker was restored on 7th June 2023.
-
 ## Development
 
 Requirements:
@@ -177,9 +182,9 @@ In this example we install the dependencies in the OS default environment. When
 
 `pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu118 --upgrade`
 
-Ensure ComfyUI is installed and patched, one way is running the tests:
+Ensure ComfyUI is installed, one way is running the tests:
 
-`tox`
+`tox -- -k test_comfy_install`
 
 From then on to run ComfyUI:
 
@@ -245,32 +250,10 @@ The `images/` directory should have our test images.
 
 ### Updating the embedded version of ComfyUI
 
-We use a ComfyUI version pinned to a specific commit, see `hordelib/consts.py:COMFYUI_VERSION`
-
-To test if the latest version works and upgrade to it, from the project root simply:
-
-1. `cd ComfyUI` _Change CWD to the embedded comfy_
-1. `git checkout master` _Switch to master branch_
-1. `git pull` _Get the latest comfyui code_
-1. `git rev-parse HEAD` _Update the hash in `hordelib.consts:COMFYUI_VERSION`_
-1. `cd ..` _Get back to the hordelib project root_
-1. `tox` _See if everything still works_
-
-Now ComfyUI is pinned to a new version.
-
-### ComfyUI Patching
-
-We patch the ComfyUI source code to:
-
-1. Modify the model manager to allow us to dynamically move models between VRAM, RAM and disk cache.
-2. Allow make ComfyUI output some handy JSON we need for development purposes.
-
-To create a patch file:
-- Make the required changes to a clean install of ComfyUI and then run `git diff > yourfile.patch` then move the patch file to wherever you want to save it.
+- Change the value in `consts.py` to the desired ComfyUI version.
+- Run the test suite via `tox`
 
-Note that the patch file _really_ needs to be in UTF-8 format and some common terminals, e.g. Powershell, won't do this by default. In Powershell to create a patch file use: `git diff | Set-Content -Encoding utf8 -Path yourfile.patch`
 
-Patches can be applied with the `hordelib.install_comfyui.Installer` classes `apply_patch()` method.
 
 <!-- Badges: -->
 

hordelib/consts.py

@@ -6,7 +6,7 @@
 
 from hordelib.config_path import get_hordelib_path
 
-COMFYUI_VERSION = "a7dd82e668bfaf7fac365a4e73a1ba1acf224fbb"
+COMFYUI_VERSION = "16eabdf70dbdb64dc4822908f0fe455c56d11ec3"
 """The exact version of ComfyUI version to load."""
 
 REMOTE_PROXY = ""

hordelib/model_manager/base.py

@@ -420,15 +420,24 @@ def is_file_available(self, file_path: str | Path) -> bool:
         Returns True if the file exists, False otherwise
         """
         parsed_full_path = Path(f"{self.model_folder_path}/{file_path}")
+        is_custom_model = False
+        if isinstance(file_path, str):
+            check_path = Path(file_path)
+            if check_path.is_absolute():
+                parsed_full_path = Path(file_path)
+                is_custom_model = True
+        if isinstance(file_path, Path) and file_path.is_absolute():
+            parsed_full_path = Path(file_path)
+            is_custom_model = True
         if parsed_full_path.suffix == ".part":
             logger.debug(f"File {file_path} is a partial download, skipping")
             return False
         sha_file_path = Path(f"{self.model_folder_path}/{parsed_full_path.stem}.sha256")
 
-        if parsed_full_path.exists() and not sha_file_path.exists():
+        if parsed_full_path.exists() and not sha_file_path.exists() and not is_custom_model:
             self.get_file_sha256_hash(parsed_full_path)
 
-        return parsed_full_path.exists() and sha_file_path.exists()
+        return parsed_full_path.exists() and (sha_file_path.exists() or is_custom_model)
 
     def download_file(
         self,
@@ -739,6 +748,7 @@ def is_model_available(self, model_name: str) -> bool:
         model_files = self.get_model_filenames(model_name)
         for file_entry in model_files:
             if not self.is_file_available(file_entry["file_path"]):
+                logger.debug([file_entry["file_path"], self.is_file_available(file_entry["file_path"])])
                 return False
         return True
 

hordelib/model_manager/compvis.py

@@ -1,3 +1,9 @@
+import json
+import os
+from pathlib import Path
+
+from loguru import logger
+
 from hordelib.consts import MODEL_CATEGORY_NAMES
 from hordelib.model_manager.base import BaseModelManager
 
@@ -14,3 +20,31 @@ def __init__(
             download_reference=download_reference,
             **kwargs,
         )
+
+    def load_model_database(self) -> None:
+        super().load_model_database()
+
+        num_custom_models = 0
+
+        try:
+            extra_models_path_str = os.getenv("HORDELIB_CUSTOM_MODELS")
+            if extra_models_path_str:
+                extra_models_path = Path(extra_models_path_str)
+                if extra_models_path.exists():
+                    extra_models = json.loads((extra_models_path).read_text())
+                    for mname in extra_models:
+                        # Avoid cloberring
+                        if mname in self.model_reference:
+                            continue
+                        # Merge all custom models into our new model reference
+                        self.model_reference[mname] = extra_models[mname]
+                        if self.is_model_available(mname):
+                            self.available_models.append(mname)
+
+                    num_custom_models += len(extra_models)
+
+        except json.decoder.JSONDecodeError as e:
+            logger.error(f"Custom model database {self.models_db_path} is not valid JSON: {e}")
+            raise
+
+        logger.info(f"Loaded {num_custom_models} models from {self.models_db_path}")

hordelib/nodes/node_model_loader.py

@@ -1,6 +1,8 @@
 # node_model_loader.py
 # Simple proof of concept custom node to load models.
 
+from pathlib import Path
+
 import comfy.model_management
 import comfy.sd
 import folder_paths  # type: ignore
@@ -46,6 +48,9 @@ def load_checkpoint(
         logger.debug(f"Will load Loras: {will_load_loras}, seamless tiling: {seamless_tiling_enabled}")
         if ckpt_name:
             logger.debug(f"Checkpoint name: {ckpt_name}")
+            # Check if the checkpoint name is a path
+            if Path(ckpt_name).is_absolute():
+                logger.debug("Checkpoint name is an absolute path.")
 
         if preloading:
             logger.debug("Preloading model.")
@@ -87,13 +92,20 @@ def load_checkpoint(
                 else:
                     # If there's no file_type passed, we follow the previous approach and pick the first file
                     # (There should be only one)
-                    ckpt_name = file_entry["file_path"].name
+                    if file_entry["file_path"].is_absolute():
+                        ckpt_name = str(file_entry["file_path"])
+                    else:
+                        ckpt_name = file_entry["file_path"].name
                     break
 
         # Clear references so comfy can free memory as needed
         SharedModelManager.manager._models_in_ram = {}
 
-        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+        if ckpt_name is not None and Path(ckpt_name).is_absolute():
+            ckpt_path = ckpt_name
+        else:
+            ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+
         with torch.no_grad():
             result = comfy.sd.load_checkpoint_guess_config(
                 ckpt_path,