nodes/wildminder-ComfyUI-KEEP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

plan(P4.5): arbitrate the model residency through ModelPatcher

Browse Source
This commit is contained in:
WildAi
2026-10-23 21:23:07 +03:00
parent 7d84b9f879
commit e31cae3c4e
2 changed files with 412 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/keep_model_loader.py
+126 -29
View File
@@ -48,8 +48,50 @@ def __getattr__(name: str):
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
def _build_patchers(keep_net, face_helper, device, offload_device) -> dict:
"""The three patchers of a pack (P4.5).
* ``keep`` - the KEEP net; it **owns the flownet** as a submodule
(``keep_arch.py:882``), so the flownet must **not** get a patcher of
its own: it would be counted twice in the VRAM accounting.
* ``detector`` - the face detector of the facelib helper.
* ``parser`` - the face parser of the facelib helper.
The patchers are built **once** per model build and are shared by every
pack that uses these models (the loader caches them).
"""
from comfy.model_patcher import ModelPatcher # lazy: the pack imports without ComfyUI
patchers: dict = {}
if keep_net is not None:
patchers["keep"] = ModelPatcher(
keep_net, load_device=device, offload_device=offload_device
)
detector = getattr(face_helper, "face_detector", None)
if detector is not None:
patchers["detector"] = ModelPatcher(
detector, load_device=device, offload_device=offload_device
)
parser = getattr(face_helper, "face_parse", None)
if parser is not None:
patchers["parser"] = ModelPatcher(
parser, load_device=device, offload_device=offload_device
)
return patchers
class KEEPModelPack:
def __init__(self, keep_net, face_helper, bg_upscale_model, face_upscale_model, model_type_str):
"""The models of the pack and their residency (P4.5).
The models are wrapped in three ``ModelPatcher`` instances (the KEEP net
with its flownet, the face detector, the face parser) and ComfyUI
arbitrates the VRAM: :meth:`load_device` asks
``model_management.load_models_gpu`` for the patchers, :meth:`offload`
gives the memory back.
"""
def __init__(self, keep_net, face_helper, bg_upscale_model, face_upscale_model,
model_type_str, patchers=None):
from comfy import model_management # lazy: the pack imports without ComfyUI
self.keep_net = keep_net
@@ -60,53 +102,99 @@ class KEEPModelPack:
self.device = model_management.get_torch_device()
self.offload_device = model_management.unet_offload_device()
def load_device(self):
"""Move all models to the main compute device."""
if self.keep_net is not None:
# the flownet is a submodule of the KEEP net - it belongs to the
# patcher of the net (no second patcher, no double accounting)
self.flownet = getattr(keep_net, "flownet", None)
if patchers is None:
patchers = _build_patchers(
keep_net, face_helper, self.device, self.offload_device
)
self.patchers_by_role = dict(patchers)
self.keep_patcher = self.patchers_by_role.get("keep")
self.detector_patcher = self.patchers_by_role.get("detector")
self.parser_patcher = self.patchers_by_role.get("parser")
self.patchers = list(self.patchers_by_role.values())
# P4.5 / P5.5: the state of the flow resolution (see ensure_flow)
self._flow_checked = False
self._flow_available = False
def load_device(self, memory_required: int = 0):
"""Move the models to the main compute device.
ComfyUI arbitrates the VRAM: the patchers are handed to
``model_management.load_models_gpu``, which decides what fits.
"""
from comfy import model_management # lazy: the pack imports without ComfyUI
if self.patchers:
model_management.load_models_gpu(
self.patchers, memory_required=memory_required
)
elif self.keep_net is not None:
# no patcher (a pack without models) - move the net directly
self.keep_net.to(self.device)
# The spandrel model descriptor has a 'model' attribute that needs to be moved
# The optional ComfyUI upscalers are not patchers - they are moved
# directly (they are small and they are not part of the KEEP models).
if self.bg_upscale_model is not None:
self.bg_upscale_model.model.to(self.device)
if self.face_upscale_model is not None:
self.face_upscale_model.model.to(self.device)
if self.face_helper is not None:
self.face_helper.device = self.device
if hasattr(self.face_helper, 'face_detector'):
self.face_helper.face_detector.to(self.device)
if hasattr(self.face_helper, 'face_parse'):
self.face_helper.face_parse.to(self.device)
def offload(self):
"""Move all models to the offload device (CPU)."""
"""Move the models back to the offload device and free the VRAM."""
from comfy import model_management # lazy: the pack imports without ComfyUI
if self.keep_net is not None:
self.keep_net.to(self.offload_device)
for patcher in self.patchers:
self._unload_patcher(patcher)
if self.bg_upscale_model is not None:
self.bg_upscale_model.model.to(self.offload_device)
if self.face_upscale_model is not None:
self.face_upscale_model.model.to(self.offload_device)
if self.face_helper is not None:
self.face_helper.device = self.offload_device
if hasattr(self.face_helper, 'face_detector'):
self.face_helper.face_detector.to(self.offload_device)
if hasattr(self.face_helper, 'face_parse'):
self.face_helper.face_parse.to(self.offload_device)
model_management.unload_all_models()
model_management.soft_empty_cache()
def _unload_patcher(self, patcher):
"""Move the model of a patcher back to its offload device.
The real ``ModelPatcher`` has **no** ``unload()`` method (ComfyUI
v0.42.0) - the model is moved to the offload device of the patcher
and ComfyUI does the rest (``unload_all_models``).
"""
model = getattr(patcher, "model", None)
if model is None:
return
offload_device = getattr(patcher, "offload_device", None) or self.offload_device
try:
model.to(offload_device)
except Exception as exc: # noqa: BLE001 - the offload must not raise
logger.debug(f"KEEP: a model could not be offloaded ({exc})")
def ensure_flow(self) -> bool:
"""Make sure the optical flow net of the KEEP net is usable.
The method is **idempotent** - it may be called before every
sequence run. P5.5 implements the conditional resolution of the
GMFlow checkpoint behind it: the checkpoint is only downloaded when
the loaded KEEP checkpoint carries no ``flownet.*`` keys. Until then
the hook is a no-op that reports ``False`` (nothing was resolved).
The method is **idempotent** - the resolution happens at most once
per pack. P5.5 implements the conditional resolution of the GMFlow
checkpoint behind it: the checkpoint is only downloaded when the
loaded KEEP checkpoint carries no ``flownet.*`` keys. Until then the
flownet comes from the KEEP checkpoint itself (it is a submodule of
the net) and nothing has to be resolved.
"""
return False
if self._flow_checked:
return self._flow_available
self._flow_checked = True
self._flow_available = self.flownet is not None
return self._flow_available
class KEEPModelLoader:
@@ -132,7 +220,10 @@ class KEEPModelLoader:
cached_pack.face_helper,
bg_upscale_model,
face_upscale_model,
model_type_str
model_type_str,
# P4.5: the patchers belong to the models - they are built
# once and shared by every pack that uses these models.
patchers=cached_pack.patchers_by_role,
)
return model_pack
@@ -195,9 +286,15 @@ class KEEPModelLoader:
logger.error(f"Error initializing FaceRestoreHelper: {e}")
raise
model_pack = KEEPModelPack(net, face_helper, bg_upscale_model, face_upscale_model, model_type_str)
# P4.5: the patchers are built once (with the base pack) and shared
# with the returned pack - a second set of patchers for the same
# models would confuse the VRAM accounting of ComfyUI.
base_pack_for_cache = KEEPModelPack(net, face_helper, None, None, model_type_str)
self.loaded_models[cache_key] = base_pack_for_cache
model_pack = KEEPModelPack(
net, face_helper, bg_upscale_model, face_upscale_model, model_type_str,
patchers=base_pack_for_cache.patchers_by_role,
)
return model_pack
tests/test_keep_model_loader.py
+286
View File
@@ -0,0 +1,286 @@
"""P4.5 - the model residency of the pack (ModelPatcher) and ``ensure_flow``.
The pack wraps its three model components (the KEEP net with its flownet, the
face detector, the face parser) in ``ModelPatcher`` instances and lets ComfyUI
arbitrate the VRAM. The tests run against the **mocked**
``comfy.model_management`` of the conftest: every residency call is recorded,
so the tests can assert what the pack asks ComfyUI to do - and that it asks
exactly once per execute.
"""
from __future__ import annotations
import importlib
import pytest
import torch
@pytest.fixture(scope="module")
def loader_module():
return importlib.import_module("ComfyUI_KEEP.src.keep_model_loader")
@pytest.fixture(scope="module")
def node_module():
return importlib.import_module("ComfyUI_KEEP.nodes.keep_face_upscale_image")
@pytest.fixture
def model_management():
"""The mocked ``comfy.model_management`` (the calls are reset per test)."""
module = importlib.import_module("comfy.model_management")
module.calls.clear()
return module
class _StubModule:
"""A minimal stand-in for a torch module (it only has to move)."""
def __init__(self, name: str = "module") -> None:
self.name = name
self.moved_to: list = []
def to(self, device):
self.moved_to.append(device)
return self
class _StubKeepNet(_StubModule):
"""The KEEP net - it owns the flownet as a submodule."""
def __init__(self) -> None:
super().__init__("keep_net")
self.flownet = _StubModule("flownet")
class _StubFaceHelper:
"""The facelib helper - it owns the detector and the parser."""
def __init__(self) -> None:
self.face_detector = _StubModule("detector")
self.face_parse = _StubModule("parser")
self.device = None
@pytest.fixture
def stub_models():
return _StubKeepNet(), _StubFaceHelper()
@pytest.fixture
def pack(loader_module, stub_models):
"""A real ``KEEPModelPack`` around the stub models."""
keep_net, face_helper = stub_models
return loader_module.KEEPModelPack(
keep_net=keep_net,
face_helper=face_helper,
bg_upscale_model=None,
face_upscale_model=None,
model_type_str="KEEP",
)
# ---------------------------------------------------------------------------
# the patchers
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_the_pack_builds_three_patchers(pack):
"""Three patchers: the KEEP net, the face detector and the face parser."""
assert len(pack.patchers) == 3, pack.patchers
assert pack.keep_patcher.model is pack.keep_net
assert pack.detector_patcher.model is pack.face_helper.face_detector
assert pack.parser_patcher.model is pack.face_helper.face_parse
@pytest.mark.unit
def test_exactly_one_patcher_contains_the_flownet(pack):
"""The flownet belongs to the KEEP net - it must not be patched twice.
A second patcher for the flownet would count its VRAM twice (the plan
P4.5).
"""
assert pack.flownet is pack.keep_net.flownet
assert pack.keep_patcher.model.flownet is pack.flownet
# the flownet is not the model of a patcher of its own
assert all(patcher.model is not pack.flownet for patcher in pack.patchers)
assert pack.detector_patcher.model is not pack.flownet
assert pack.parser_patcher.model is not pack.flownet
@pytest.mark.unit
def test_the_patchers_are_shared_by_the_packs_of_the_same_models(
loader_module, stub_models
):
"""The patchers belong to the models - they are built once per build."""
keep_net, face_helper = stub_models
first = loader_module.KEEPModelPack(keep_net, face_helper, None, None, "KEEP")
second = loader_module.KEEPModelPack(
keep_net, face_helper, None, None, "KEEP", patchers=first.patchers_by_role
)
assert second.keep_patcher is first.keep_patcher
assert second.patchers == first.patchers
# ---------------------------------------------------------------------------
# the residency
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_load_device_lets_comfyui_arbitrate_the_vram(pack, model_management):
"""The pack hands all its patchers to ``load_models_gpu`` - once."""
pack.load_device()
load_calls = [
call for call in model_management.calls if call[0] == "load_models_gpu"
]
assert len(load_calls) == 1, model_management.calls
_, models, memory_required = load_calls[0]
assert list(models) == list(pack.patchers)
assert memory_required == 0
@pytest.mark.unit
def test_offload_moves_the_models_back_and_frees_the_vram(pack, model_management):
"""The release path: the models go back to the offload device and ComfyUI
frees the VRAM."""
pack.offload()
names = [call[0] for call in model_management.calls]
assert "unload_all_models" in names, model_management.calls
assert "soft_empty_cache" in names, model_management.calls
# the models were moved back to the offload device
assert pack.keep_net.moved_to[-1] == pack.offload_device
assert pack.face_helper.face_detector.moved_to[-1] == pack.offload_device
assert pack.face_helper.face_parse.moved_to[-1] == pack.offload_device
@pytest.mark.unit
def test_the_cpu_mode_makes_the_same_calls(
loader_module, model_management, monkeypatch
):
"""``get_torch_device()`` -> cpu: the same residency calls are made."""
monkeypatch.setattr(
model_management, "get_torch_device", lambda: torch.device("cpu")
)
pack = loader_module.KEEPModelPack(
_StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
)
assert pack.device == torch.device("cpu")
pack.load_device()
pack.offload()
names = [call[0] for call in model_management.calls]
assert names.count("load_models_gpu") == 1, model_management.calls
assert names.count("unload_all_models") == 1, model_management.calls
assert names.count("soft_empty_cache") == 1, model_management.calls
# ---------------------------------------------------------------------------
# the node around the pack (mock-integration)
# ---------------------------------------------------------------------------
def _patch_the_node(monkeypatch, node_module, processor_cls):
"""Stub the heavy parts of the single image node."""
monkeypatch.setattr(node_module, "KEEPFaceProcessor", processor_cls)
monkeypatch.setattr(node_module, "comfy_image_to_cv2", lambda image: image)
monkeypatch.setattr(node_module, "cv2_to_comfy_image", lambda image: image)
monkeypatch.setattr(node_module, "report_progress", lambda *args: True)
class _StubProcessor:
"""A processor that returns a stub image."""
def __init__(self, pack) -> None:
self.pack = pack
def process_image(self, **kwargs):
return _StubModule("processed")
class _FailingProcessor:
"""A processor that fails (the node must give the VRAM back)."""
def __init__(self, pack) -> None:
self.pack = pack
def process_image(self, **kwargs):
raise RuntimeError("the processing failed")
@pytest.mark.mock_integration
def test_one_execute_asks_for_the_vram_once_and_gives_it_back(
node_module, loader_module, model_management, monkeypatch
):
"""One execute -> one ``load_models_gpu`` with all the pack patchers and
the release afterwards."""
pack = loader_module.KEEPModelPack(
_StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
)
_patch_the_node(monkeypatch, node_module, _StubProcessor)
node_module.KEEP_FaceUpscaleImageNode.execute(
image=torch.zeros((1, 8, 8, 3)),
keep_model=pack,
final_upscale_factor=1.0,
has_aligned_face=False,
only_center_face=True,
draw_bounding_box=False,
)
names = [call[0] for call in model_management.calls]
assert names.count("load_models_gpu") == 1, model_management.calls
assert names.count("unload_all_models") == 1, model_management.calls
assert names.count("soft_empty_cache") == 1, model_management.calls
load_call = [c for c in model_management.calls if c[0] == "load_models_gpu"][0]
assert list(load_call[1]) == list(pack.patchers)
@pytest.mark.mock_integration
def test_the_node_gives_the_vram_back_when_the_processing_fails(
node_module, loader_module, model_management, monkeypatch
):
"""The release happens in the ``finally`` path - also on a failure."""
pack = loader_module.KEEPModelPack(
_StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
)
_patch_the_node(monkeypatch, node_module, _FailingProcessor)
with pytest.raises(RuntimeError):
node_module.KEEP_FaceUpscaleImageNode.execute(
image=torch.zeros((1, 8, 8, 3)),
keep_model=pack,
final_upscale_factor=1.0,
has_aligned_face=False,
only_center_face=True,
draw_bounding_box=False,
)
names = [call[0] for call in model_management.calls]
assert names.count("load_models_gpu") == 1, model_management.calls
assert names.count("unload_all_models") == 1, model_management.calls
assert names.count("soft_empty_cache") == 1, model_management.calls
# ---------------------------------------------------------------------------
# ensure_flow
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_ensure_flow_is_idempotent(pack):
"""``ensure_flow`` resolves once per pack - the second call is a no-op."""
assert pack.ensure_flow() is True, "the flownet comes with the KEEP net"
# the flownet is gone -> the cached result is still returned
pack.flownet = None
assert pack.ensure_flow() is True
@pytest.mark.unit
def test_ensure_flow_reports_a_missing_flownet(loader_module):
"""A pack without a flownet reports that nothing is available (yet).
P5.5 resolves the GMFlow checkpoint behind this hook.
"""
net = _StubKeepNet()
net.flownet = None
pack = loader_module.KEEPModelPack(net, _StubFaceHelper(), None, None, "KEEP")
assert pack.ensure_flow() is False