diff --git a/src/keep_model_loader.py b/src/keep_model_loader.py index b27c422..8fab67f 100644 --- a/src/keep_model_loader.py +++ b/src/keep_model_loader.py @@ -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 \ No newline at end of file diff --git a/tests/test_keep_model_loader.py b/tests/test_keep_model_loader.py new file mode 100644 index 0000000..5bf47dd --- /dev/null +++ b/tests/test_keep_model_loader.py @@ -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