From f0982cbeed897e7a559ff86da4df5b8c1e4d4f34 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 25 Aug 2024 11:49:06 -0400 Subject: [PATCH] fix: rely on comfyui internals for loading cascade series models This is a patent "temporary" fix to prevent extremely large models, such as cascade models, from being forced to fully load. I will track the resolution of this in an issue: https://github.com/Haidra-Org/hordelib/issues/312 --- hordelib/comfy_horde.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/hordelib/comfy_horde.py b/hordelib/comfy_horde.py index 6c913a3..ea80b23 100644 --- a/hordelib/comfy_horde.py +++ b/hordelib/comfy_horde.py @@ -240,6 +240,10 @@ def do_comfy_import( # isort: on +def _model_patcher_has_cascade(model_patcher): + return "cascade" in str(type(model_patcher.model)).lower() + + def _load_models_gpu_hijack(*args, **kwargs): """Intercepts the comfy load_models_gpu function to force full load. @@ -248,7 +252,18 @@ def _load_models_gpu_hijack(*args, **kwargs): and the worker/horde-engine takes responsibility for managing the memory or the problems this may cause. """ + found_cascade = False + for model_patcher in args[0]: + found_cascade = _model_patcher_has_cascade(model_patcher) + if found_cascade: + break + global _comfy_current_loaded_models + if found_cascade: + logger.debug("Not overriding cascade model load") + _comfy_load_models_gpu(*args, **kwargs) + return + if "force_full_load" in kwargs: kwargs.pop("force_full_load") @@ -262,6 +277,13 @@ def _model_patcher_load_hijack(*args, **kwargs): See _load_models_gpu_hijack for more information """ global _comfy_model_patcher_load + + model_patcher = args[0] + if _model_patcher_has_cascade(model_patcher): + logger.debug("Not overriding cascade model load") + _comfy_model_patcher_load(*args, **kwargs) + return + if "full_load" in kwargs: kwargs.pop("full_load")