482 lines
17 KiB
Python
482 lines
17 KiB
Python
"""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 hashlib
|
|
import importlib
|
|
import os
|
|
|
|
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(scope="module")
|
|
def configs_module():
|
|
return importlib.import_module("ComfyUI_KEEP.src.keep_model_configs")
|
|
|
|
|
|
@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.
|
|
|
|
The conftest mock answers ``cpu`` by default (``conftest.py``: the
|
|
``get_torch_device`` / ``unet_offload_device`` of the mocked
|
|
``comfy.model_management``), so patching the device to ``cpu`` alone
|
|
would exercise the default path and prove nothing. The CPU leg is
|
|
therefore compared with a **GPU leg**: both devices must produce the
|
|
same residency calls - the pack must not branch on the device. A
|
|
device dependent path is exactly the B7 CPU crash (the pack called
|
|
``torch.cuda.empty_cache()`` and crashed on a CPU-only host).
|
|
"""
|
|
|
|
def residency_calls(device):
|
|
monkeypatch.setattr(model_management, "get_torch_device", lambda: device)
|
|
monkeypatch.setattr(model_management, "unet_offload_device", lambda: device)
|
|
model_management.calls.clear()
|
|
|
|
pack = loader_module.KEEPModelPack(
|
|
_StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
|
|
)
|
|
# the patched device really reaches the pack - the leg is not a no-op
|
|
assert pack.device == device
|
|
assert pack.offload_device == device
|
|
|
|
pack.load_device()
|
|
pack.offload()
|
|
return [call[0] for call in model_management.calls]
|
|
|
|
gpu_calls = residency_calls(torch.device("cuda"))
|
|
cpu_calls = residency_calls(torch.device("cpu"))
|
|
|
|
assert cpu_calls == gpu_calls, (cpu_calls, gpu_calls)
|
|
assert cpu_calls.count("load_models_gpu") == 1, cpu_calls
|
|
assert cpu_calls.count("unload_all_models") == 1, cpu_calls
|
|
assert cpu_calls.count("soft_empty_cache") == 1, cpu_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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P5.3 - the fail-fast weight validation
|
|
# ---------------------------------------------------------------------------
|
|
PAYLOAD = b"not a real checkpoint"
|
|
PAYLOAD_SHA256 = hashlib.sha256(PAYLOAD).hexdigest()
|
|
|
|
|
|
@pytest.fixture
|
|
def synthetic_manifest(loader_module, configs_module, monkeypatch):
|
|
"""The manifest with the hash of the test payload.
|
|
|
|
The real hashes cannot be reproduced by a test, so the manifest the
|
|
loader and the downloader see is replaced by entries that carry the hash
|
|
of the payload the transport double writes. Everything else - the
|
|
resolution order, the fail-fast rule, the error messages - is the
|
|
production code.
|
|
"""
|
|
entries = tuple(
|
|
configs_module.WeightEntry(
|
|
name=entry.name,
|
|
url=entry.url,
|
|
sha256=PAYLOAD_SHA256,
|
|
dest_folder=entry.dest_folder,
|
|
subdir=entry.subdir,
|
|
required=entry.required,
|
|
available=entry.available,
|
|
unverified=entry.unverified,
|
|
sha256_prefix=entry.sha256_prefix,
|
|
)
|
|
for entry in configs_module.WEIGHT_MANIFEST
|
|
)
|
|
by_name = {entry.name: entry for entry in entries}
|
|
monkeypatch.setattr(loader_module, "WEIGHT_MANIFEST", entries)
|
|
monkeypatch.setattr(
|
|
loader_module,
|
|
"weight_entry_for",
|
|
lambda model_type: by_name[
|
|
"KEEP_Asian-4765ebe0.pth" if model_type == "Asian" else "KEEP-b76feb75.pth"
|
|
],
|
|
)
|
|
monkeypatch.setattr(
|
|
loader_module.keep_download, "WEIGHT_MANIFEST", entries
|
|
)
|
|
return by_name
|
|
|
|
|
|
@pytest.fixture
|
|
def transport(monkeypatch):
|
|
"""The transport double: it records the URLs and writes the payload."""
|
|
download_module = importlib.import_module("ComfyUI_KEEP.src.keep_download")
|
|
calls: list[str] = []
|
|
|
|
def _download(url, dst, hash_prefix=None, progress=True):
|
|
calls.append(url)
|
|
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
|
with open(dst, "wb") as handle:
|
|
handle.write(PAYLOAD)
|
|
|
|
monkeypatch.setattr(download_module, "download_url_to_file", _download)
|
|
return calls
|
|
|
|
|
|
@pytest.fixture
|
|
def loader(loader_module, monkeypatch, tmp_path, synthetic_manifest, transport):
|
|
"""A loader whose model build is stubbed out (the weights are real)."""
|
|
|
|
class _StubArch:
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
|
|
def to(self, device):
|
|
return self
|
|
|
|
def load_state_dict(self, state_dict, strict=True):
|
|
self.loaded = state_dict
|
|
return self
|
|
|
|
def eval(self):
|
|
return self
|
|
|
|
class _StubHelper:
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
|
|
monkeypatch.setattr(loader_module, "ARCH_REGISTRY", {"KEEP": _StubArch})
|
|
monkeypatch.setattr(
|
|
loader_module.torch,
|
|
"load",
|
|
lambda *args, **kwargs: {"params_ema": {"weight": torch.zeros(1)}},
|
|
)
|
|
monkeypatch.setattr(loader_module, "_face_restore_helper", lambda: _StubHelper)
|
|
monkeypatch.setattr(
|
|
loader_module.keep_paths, "resolve_model_root", lambda *args, **kwargs: str(tmp_path)
|
|
)
|
|
return loader_module.KEEPModelLoader()
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_a_missing_required_weight_triggers_a_download(loader, transport):
|
|
"""The weights are resolved (and downloaded) before the net is built."""
|
|
pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")
|
|
|
|
assert pack.keep_net is not None
|
|
urls = [url.rsplit("/", 1)[-1] for url in transport]
|
|
assert "KEEP-b76feb75.pth" in urls, urls
|
|
for name in (
|
|
"detection_Resnet50_Final.pth",
|
|
"detection_mobilenet0.25_Final.pth",
|
|
"yolov5n-face.pth",
|
|
"yolov5l-face.pth",
|
|
"parsing_parsenet.pth",
|
|
):
|
|
assert name in urls, (name, urls)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_a_failed_download_raises_the_actionable_message(
|
|
loader_module, loader, monkeypatch, transport
|
|
):
|
|
"""A missing required weight names the file, the path and the URL."""
|
|
download_module = importlib.import_module("ComfyUI_KEEP.src.keep_download")
|
|
|
|
def _fail(url, dst, hash_prefix=None, progress=True):
|
|
raise RuntimeError("the server is not reachable")
|
|
|
|
monkeypatch.setattr(download_module, "download_url_to_file", _fail)
|
|
|
|
with pytest.raises(FileNotFoundError) as excinfo:
|
|
loader.load_keep_model_pack("KEEP", "retinaface_resnet50")
|
|
|
|
message = str(excinfo.value)
|
|
assert "KEEP-b76feb75.pth" in message
|
|
assert "https://github.com/jnjaby/KEEP/releases/download/v1.0.0/KEEP-b76feb75.pth" in message
|
|
assert "keep_models" in message and "KEEP" in message
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_optional_gmflow_weight_is_not_resolved_at_pack_construction(
|
|
loader, transport
|
|
):
|
|
"""The GMFlow checkpoint is deferred to ``ensure_flow()``."""
|
|
loader.load_keep_model_pack("KEEP", "retinaface_resnet50")
|
|
|
|
urls = [url.rsplit("/", 1)[-1] for url in transport]
|
|
assert "gmflow_sintel-0c07dcb3.pth" not in urls, urls
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_unavailable_asian_entry_does_not_break_the_default_pack(loader, transport):
|
|
"""The Asian checkpoint is not required - the default pack still builds."""
|
|
pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")
|
|
|
|
assert pack is not None
|
|
urls = [url.rsplit("/", 1)[-1] for url in transport]
|
|
assert "KEEP_Asian-4765ebe0.pth" not in urls, urls
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_selecting_the_asian_model_raises_an_actionable_error(loader):
|
|
"""The Asian checkpoint is not published upstream (HTTP 404)."""
|
|
with pytest.raises(FileNotFoundError) as excinfo:
|
|
loader.load_keep_model_pack("Asian", "retinaface_resnet50")
|
|
|
|
message = str(excinfo.value)
|
|
assert "KEEP_Asian-4765ebe0.pth" in message
|
|
assert "not published upstream" in message
|
|
assert "KEEP_Asian-4765ebe0.pth" in message
|
|
assert "keep_models" in message
|