wildminder-ComfyUI-KEEP/tests/test_keep_model_loader.py

"""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(loader_module, stub_models):
    """``ensure_flow`` resolves once per pack - the second call is a no-op.

    A checkpoint that carries the flownet weights (the shipped one, P0.4
    section 2) needs no resolution at all.
    """
    net, helper = stub_models
    pack = loader_module.KEEPModelPack(
        net, helper, None, None, "KEEP", flownet_from_checkpoint=True
    )
    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_MANIFEST_BY_NAME", by_name)
    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:
        """A stand-in for the vendored arch that behaves like ``nn.Module``.

        The real arch **owns** the flownet as a submodule
        (``keep_arch.py:882``), so the state dict of the net always carries
        the ``flownet.*`` keys - and a checkpoint without them fails the
        **strict** load.  The stub keeps that behaviour (a no-op recorder
        would hide exactly that defect).
        """

        #: The state dict of the stub net (the flownet is part of it).
        STATE = {
            "weight": torch.zeros(1),
            "flownet.model.conv1.weight": torch.ones(1),
        }

        def __init__(self, **kwargs):
            self.kwargs = kwargs
            # the real arch owns the flownet as a submodule (keep_arch.py:882)
            self.flownet = _StubModule("flownet")
            self.state = dict(self.STATE)

        def to(self, device):
            return self

        def state_dict(self):
            return dict(self.state)

        def load_state_dict(self, state_dict, strict=True):
            missing = [key for key in self.state if key not in state_dict]
            unexpected = [key for key in state_dict if key not in self.state]
            if strict and (missing or unexpected):
                raise RuntimeError(
                    "Error(s) in loading state_dict for KEEP: missing "
                    f"{missing}, unexpected {unexpected}"
                )
            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_an_unresolvable_optional_gmflow_weight_does_not_raise(
    loader_module, loader, monkeypatch
):
    """The GMFlow checkpoint is optional - its absence never breaks the pack.

    The flow weights are only *needed* when the checkpoint does not carry
    them; the resolution is deferred to ``ensure_flow()`` (P5.5).
    """
    monkeypatch.setattr(
        loader_module.keep_download,
        "ensure_optional",
        lambda entry, progress=True: None,
    )

    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack is not None
    assert pack.keep_net.kwargs["flownet_path"] is None


@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


# ---------------------------------------------------------------------------
# P5.5 - the flownet wiring
# ---------------------------------------------------------------------------
def _downloaded(transport) -> list[str]:
    return [url.rsplit("/", 1)[-1] for url in transport]


@pytest.mark.unit
def test_the_arch_gets_a_resolved_flownet_path_without_flownet_weights(loader, transport):
    """A checkpoint **without** flownet weights -> the GMFlow file is resolved.

    The resolved path is put into the architecture config **before** the net
    is built, so the vendored arch loads the trained weights instead of
    random ones.
    """
    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    flownet_path = pack.keep_net.kwargs["flownet_path"]
    assert flownet_path is not None, "the GMFlow checkpoint was not resolved"
    assert flownet_path.endswith("gmflow_sintel-0c07dcb3.pth"), flownet_path
    assert "gmflow_sintel-0c07dcb3.pth" in _downloaded(transport)


@pytest.mark.unit
def test_a_checkpoint_without_flownet_weights_loads_strictly(loader):
    """The fallback branch must not break the strict load.

    The real arch **always** owns the flownet (``keep_arch.py:882``), so the
    state dict of the net always carries the ``flownet.*`` keys.  A
    checkpoint without them used to make the strict load raise
    ``RuntimeError: Missing key(s) in state_dict: flownet.*`` - and the
    GMFlow weights the arch had just loaded were thrown away.  The flownet
    tensors of the built net fill the gap, so the load stays strict **and**
    the trained flow weights survive.
    """
    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    # the strict load happened (the stub arch raises on a missing key)
    assert pack.keep_net.loaded is not None
    # the flownet of the built net (the trained GMFlow weights) is in the
    # loaded state dict - it was not discarded
    assert torch.equal(
        pack.keep_net.loaded["flownet.model.conv1.weight"],
        pack.keep_net.state_dict()["flownet.model.conv1.weight"],
    )


@pytest.mark.unit
def test_the_strict_load_still_rejects_an_unexpected_key(loader_module, loader, monkeypatch):
    """The flownet merge does not weaken the strict load."""
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {
            "params_ema": {"weight": torch.zeros(1), "unexpected.key": torch.zeros(1)}
        },
    )

    with pytest.raises(RuntimeError):
        loader.load_keep_model_pack("KEEP", "retinaface_resnet50")


@pytest.mark.unit
def test_a_checkpoint_with_flownet_weights_downloads_no_gmflow(
    loader_module, loader, transport, monkeypatch
):
    """The shipped checkpoint carries the flownet -> nothing is downloaded."""
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {
            "params_ema": {
                "flownet.model.conv1.weight": torch.zeros(1),
                "weight": torch.zeros(1),
            }
        },
    )

    pack = loader.load_keep_model_pack("KEEP", "retinaface_resnet50")

    assert pack.keep_net.kwargs["flownet_path"] is None
    assert "gmflow_sintel-0c07dcb3.pth" not in _downloaded(transport)
    # the flownet is restored from the checkpoint (the strict load)
    assert "flownet.model.conv1.weight" in pack.keep_net.loaded
    assert pack.flownet_from_checkpoint is True
    # the hook is a no-op - the weights are already there
    assert pack.ensure_flow() is True
    assert "gmflow_sintel-0c07dcb3.pth" not in _downloaded(transport)


class _RecordingInner:
    """The inner module of the flownet (it records the strict load)."""

    def __init__(self) -> None:
        self.loaded: tuple | None = None

    def load_state_dict(self, weights, strict=True):
        self.loaded = (weights, strict)
        return self


class _RecordingFlownet:
    """A flownet that records the weights it receives."""

    def __init__(self) -> None:
        self.model = _RecordingInner()


@pytest.mark.unit
def test_ensure_flow_loads_the_gmflow_weights_strictly(
    loader_module, stub_models, monkeypatch, tmp_path
):
    """A net without flow weights gets them from the GMFlow file."""
    net, helper = stub_models
    net.flownet = _RecordingFlownet()
    pack = loader_module.KEEPModelPack(
        net, helper, None, None, "KEEP",
        flownet_from_checkpoint=False, flownet_path=None,
    )

    gmflow_path = tmp_path / "gmflow_sintel-0c07dcb3.pth"
    gmflow_path.write_bytes(b"weights")
    monkeypatch.setattr(
        loader_module.keep_download,
        "ensure_optional",
        lambda entry, progress=True: str(gmflow_path),
    )
    monkeypatch.setattr(
        loader_module.torch,
        "load",
        lambda *args, **kwargs: {"model": {"w": torch.zeros(1)}},
    )

    assert pack.ensure_flow() is True
    weights, strict = net.flownet.model.loaded
    assert strict is True, "the flow weights must be loaded with strict=True"
    assert "w" in weights
    # the resolution is cached - a second call does nothing
    assert pack.ensure_flow() is True


@pytest.mark.mock_integration
def test_the_image_node_does_not_call_ensure_flow(
    loader_module, node_module, monkeypatch
):
    """The image path leaves the flow resolution deferred."""
    net = _StubKeepNet()
    pack = loader_module.KEEPModelPack(
        net, _StubFaceHelper(), None, None, "KEEP"
    )
    calls: list[str] = []
    monkeypatch.setattr(pack, "ensure_flow", lambda: calls.append("ensure_flow"))
    _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,
    )

    assert calls == [], "the image path must not resolve the flow weights"