From b54b69fda50636088131038c419f8f74abb19085 Mon Sep 17 00:00:00 2001 From: tazlin Date: Fri, 5 Apr 2024 19:14:08 -0400 Subject: [PATCH 1/2] feat: update to comfyui `a7dd82e` --- hordelib/consts.py | 2 +- hordelib/model_manager/base.py | 25 ++++++++++- hordelib/nodes/node_model_loader.py | 6 +++ tests/conftest.py | 6 +++ tests/test_horde_inference_custom_model.py | 51 ++++++++++++++++++++++ tox.ini | 1 + 6 files changed, 88 insertions(+), 3 deletions(-) create mode 100644 tests/test_horde_inference_custom_model.py diff --git a/hordelib/consts.py b/hordelib/consts.py index 6e1a0c1f..d834d377 100644 --- a/hordelib/consts.py +++ b/hordelib/consts.py @@ -6,7 +6,7 @@ from hordelib.config_path import get_hordelib_path -COMFYUI_VERSION = "a28a9dc83684624ee2167c0b92d976bb68f2c606" +COMFYUI_VERSION = "a7dd82e668bfaf7fac365a4e73a1ba1acf224fbb" """The exact version of ComfyUI version to load.""" REMOTE_PROXY = "" diff --git a/hordelib/model_manager/base.py b/hordelib/model_manager/base.py index 3e264c9a..d3390fd3 100644 --- a/hordelib/model_manager/base.py +++ b/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 diff --git a/hordelib/nodes/node_model_loader.py b/hordelib/nodes/node_model_loader.py index 111f2cee..f656330f 100644 --- a/hordelib/nodes/node_model_loader.py +++ b/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 @@ -94,6 +96,10 @@ def load_checkpoint( 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, diff --git a/tests/conftest.py b/tests/conftest.py index f8962842..32aba561 100644 --- a/tests/conftest.py +++ b/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") diff --git a/tests/test_horde_inference_custom_model.py b/tests/test_horde_inference_custom_model.py new file mode 100644 index 00000000..720bac48 --- /dev/null +++ b/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, + # ) diff --git a/tox.ini b/tox.ini index 15a0309d..67540290 100644 --- a/tox.ini +++ b/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 From 7f1e69bfebbeff97958c586c4d30b477c9fe7a85 Mon Sep 17 00:00:00 2001 From: db0 Date: Mon, 8 Apr 2024 01:12:17 +0200 Subject: [PATCH 2/2] fix: handle full model name --- hordelib/model_manager/base.py | 12 ++++++------ hordelib/nodes/node_model_loader.py | 5 ++++- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/hordelib/model_manager/base.py b/hordelib/model_manager/base.py index d3390fd3..25219f1a 100644 --- a/hordelib/model_manager/base.py +++ b/hordelib/model_manager/base.py @@ -122,12 +122,12 @@ def load_model_database(self) -> None: 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] + 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( diff --git a/hordelib/nodes/node_model_loader.py b/hordelib/nodes/node_model_loader.py index f656330f..f67630e3 100644 --- a/hordelib/nodes/node_model_loader.py +++ b/hordelib/nodes/node_model_loader.py @@ -89,7 +89,10 @@ 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