feat: Allow providing custom models

hordelib/model_manager/base.py

@@ -117,6 +117,17 @@ def load_model_database(self) -> None:
         for attempt in range(3):
             try:
                 self.model_reference = json.loads((self.models_db_path).read_text())
+                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]
             except json.decoder.JSONDecodeError as e:
                 if attempt <= 2:
                     logger.warning(
@@ -420,15 +431,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 +759,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/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
@@ -87,13 +89,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:
+            check_path = Path(ckpt_name)
+            if check_path.is_absolute():
+                ckpt_path = ckpt_name
         with torch.no_grad():
             result = comfy.sd.load_checkpoint_guess_config(
                 ckpt_path,

tests/conftest.py

@@ -111,6 +111,12 @@ def stable_cascade_base_model_name(shared_model_manager: type[SharedModelManager
     return "Stable Cascade 1.0"
 
 
+@pytest.fixture(scope="session")
+def custom_model_name_for_testing(shared_model_manager: type[SharedModelManager]) -> str:
+    # https://civitai.com/models/338712/pvc-style-modelmovable-figure-model-xl?modelVersionId=413807
+    return "Movable figure model XL"
+
+
 @pytest.fixture(scope="session")
 def db0_test_image() -> PIL.Image.Image:
     return PIL.Image.open("images/test_db0.jpg")

tests/test_horde_inference_custom_model.py

@@ -0,0 +1,51 @@
+# test_horde.py
+
+from PIL import Image
+
+from hordelib.horde import HordeLib
+
+
+class TestHordeInference:
+    def test_custom_model_text_to_image(
+        self,
+        hordelib_instance: HordeLib,
+        custom_model_name_for_testing: str,
+    ):
+        data = {
+            "sampler_name": "k_euler_a",
+            "cfg_scale": 7.5,
+            "denoising_strength": 1.0,
+            "seed": 1312,
+            "height": 1024,
+            "width": 1024,
+            "karras": False,
+            "tiling": False,
+            "hires_fix": False,
+            "clip_skip": 2,
+            "control_type": None,
+            "image_is_control": False,
+            "return_control_map": False,
+            "prompt": (
+                "surreal,amazing quality,masterpiece,best quality,awesome,inspiring,cinematic composition"
+                ",soft shadows,Film grain,shallow depth of field,highly detailed,high budget,cinemascope,epic,"
+                "OverallDetail,color graded cinematic,atmospheric lighting,imperfections,natural,shallow dof,"
+                "1girl,solo,looking at viewer,kurumi_ebisuzawa,twin tails,hair ribbon,leather jacket,leather pants,"
+                "black jacket,tight pants,black chocker,zipper,fingerless gloves,biker clothes,spikes,unzipped,"
+                "shoulder spikes,multiple belts,shiny clothes,(graffiti:1.2),brick wall,dutch angle,crossed arms,"
+                "arms under breasts,anarchist mask,v-shaped eyebrows"
+            ),
+            "ddim_steps": 30,
+            "n_iter": 1,
+            "model": custom_model_name_for_testing,
+        }
+        pil_image = hordelib_instance.basic_inference_single_image(data).image
+        assert pil_image is not None
+        assert isinstance(pil_image, Image.Image)
+
+        img_filename = "custom_model_text_to_image.png"
+        pil_image.save(f"images/{img_filename}", quality=100)
+
+        # assert check_single_inference_image_similarity(
+        #     f"images_expected/{img_filename}",
+        #     pil_image,
+        # )

tox.ini

@@ -76,6 +76,7 @@ passenv =
     AIWORKER_CACHE_HOME
     TESTS_ONGOING
     HORDELIB_SKIP_SIMILARITY_FAIL
+    HORDELIB_CUSTOM_MODELS
     CIVIT_API_TOKEN
     HORDE_MODEL_REFERENCE_GITHUB_BRANCH