nodes/wildminder-ComfyUI-KEEP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
bfede01dd58d8f77e718c6abbb2b5c81ef96c54e
wildminder-ComfyUI-KEEP/tests/test_keep_model_loader.py
T
WildAi 3af08f4357 perf(keep): stream the clip through detection/restore and place the face nets on the pack device 
- R1: the face helper is constructed on the pack device (was the offload
  device); load_device() explicitly places the detector and the parser with
  a logged CPU fallback, offload() mirrors the restore, the YOLO .device
  skew is closed (P-D1..P-D5).
- R2: detection streams per max_clip_length group (no full-clip convert-in
  retention), crop is fused into a chunk-bounded restore loop with
  preallocated staging, the paste tail keeps one contractual float32
  buffer, the helper is released after every run; golden fixtures pin the
  outputs bitwise (P-M1..P-M7).
- The diagnosis: the OOM was O(frames) whole-clip host buffers, not the
  faces (the detector/landmark path was already GPU-placed);
  max_clip_length never bounded convert-in or the paste tail.
- Gate: 523 passed, 6 skipped, 7 deselected, 0 failed (baseline 505).
2026-10-28 23:39:53 +03:00

908 lines
33 KiB
Python
Raw Blame History

"""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 logging
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 device placement of the face nets (round 4, D-R1)
# ---------------------------------------------------------------------------
def _pack_on_distinct_devices(loader_module, model_management, monkeypatch):
"""A pack whose pack device is distinct from its offload device.
The conftest mock answers ``cpu`` for both device functions, so a
placement on the wrong device would be indistinguishable from a correct
one. The patch gives the pack a distinct device object - recording
only, no CUDA op ever runs (the CPU/GPU-leg pattern above).
"""
monkeypatch.setattr(
model_management, "get_torch_device", lambda: torch.device("cuda")
)
monkeypatch.setattr(
model_management, "unet_offload_device", lambda: torch.device("cpu")
)
pack = loader_module.KEEPModelPack(
_StubKeepNet(), _StubFaceHelper(), None, None, "KEEP"
)
assert pack.device != pack.offload_device # the leg is not a no-op
return pack
class _YoloDetector(_StubModule):
"""The YOLO detector shape - it stores its own ``device`` attribute.
The vendored YOLO detector reads ``self.device`` when it pushes its
inputs (the device diagnosis D8b), so a moved net with a stale attribute
would run on the wrong device.
"""
def __init__(self) -> None:
super().__init__("detector")
# a sentinel distinct from every leg device - a sync shows up
self.device = torch.device("meta")
class _FailingParser(_StubModule):
"""A parser whose move raises - the placement must degrade cleanly."""
def __init__(self) -> None:
super().__init__("parser")
self.device = torch.device("meta")
def to(self, device):
raise RuntimeError("the VRAM is full")
@pytest.fixture
def keep_records(loader_module, caplog):
"""The records of the pack logger (the pattern of ``test_keep_download``).
Without a ComfyUI the pack logger owns a private handler and does not
propagate, so the records never reach the root logger of pytest: the
capture handler is attached to the pack logger itself.
"""
loader_module.logger.addHandler(caplog.handler)
try:
yield caplog
finally:
loader_module.logger.removeHandler(caplog.handler)
@pytest.mark.unit
def test_the_face_helper_is_constructed_on_the_pack_device(
loader_module, model_management, monkeypatch, loader
):
"""P-D1: the helper is constructed with the pack device.
The vendored helper stores the construction device and builds its
detector and parser nets on it (``face_restoration_helper.py:66-93``) -
constructed on the offload device, the nets were born on the wrong side
of the patcher flow. The ``_StubHelper`` of the ``loader`` fixture
records the constructor kwargs.
"""
monkeypatch.setattr(
model_management, "get_torch_device", lambda: torch.device("cuda")
)
monkeypatch.setattr(
model_management, "unet_offload_device", lambda: torch.device("cpu")
)
# the device functions are read at construction - a loader built after
# the patch sees the distinct pack device (the fixture's loader saw cpu)
fresh = loader_module.KEEPModelLoader()
pack = fresh.load_keep_model_pack("KEEP", "retinaface_resnet50")
assert pack.face_helper.kwargs["device"] == pack.device
assert pack.face_helper.kwargs["device"] != pack.offload_device
@pytest.mark.unit
def test_load_device_places_the_face_nets_on_the_pack_device(
loader_module, model_management, monkeypatch
):
"""P-D2: ``load_device`` places the detector and the parser explicitly.
The patcher flow may silently skip a module under VRAM pressure - the
explicit placement makes the residency of the two nets a fact, not a
hope, and the helper's own ``device`` stays synced.
"""
pack = _pack_on_distinct_devices(loader_module, model_management, monkeypatch)
pack.load_device()
assert pack.face_helper.face_detector.moved_to[-1] == pack.device
assert pack.face_helper.face_parse.moved_to[-1] == pack.device
assert pack.face_helper.device == pack.device
@pytest.mark.unit
def test_load_device_falls_back_cleanly_when_a_net_cannot_move(
loader_module, model_management, keep_records
):
"""P-D3: a net that cannot move degrades cleanly.
A failed move must not raise out of ``load_device``, must be logged as a
warning, and must leave the net's own ``.device`` attribute untouched -
the sync belongs to the success path only (a failed move that still
re-synced the attribute would point the net at a device its weights
never reached). The other net is still placed.
"""
helper = _StubFaceHelper()
failing = _FailingParser()
helper.face_parse = failing
pack = loader_module.KEEPModelPack(
_StubKeepNet(), helper, None, None, "KEEP"
)
pack.load_device() # must not raise
# the detector is still placed, the helper device still synced
assert pack.face_helper.face_detector.moved_to[-1] == pack.device
assert pack.face_helper.device == pack.device
# the failed parser: no move recorded, no attribute re-sync
assert failing.moved_to == []
assert failing.device == torch.device("meta")
# the failure is a logged warning, not a silent skip
warnings = [
record for record in keep_records.records
if record.levelno == logging.WARNING
]
assert len(warnings) == 1, [r.getMessage() for r in keep_records.records]
assert "face parser" in warnings[0].getMessage()
assert "could not be placed" in warnings[0].getMessage()
@pytest.mark.unit
def test_load_device_syncs_a_detector_own_device_attribute(
loader_module, model_management, monkeypatch
):
"""P-D4: the YOLO-style ``.device`` attribute follows the net - both ways.
Synced on ``load_device`` (after a successful move, D-R1.2) and restored
on ``offload`` (D-R1.3) - the attribute must never lag the weights.
"""
monkeypatch.setattr(
model_management, "get_torch_device", lambda: torch.device("cuda")
)
monkeypatch.setattr(
model_management, "unet_offload_device", lambda: torch.device("cpu")
)
helper = _StubFaceHelper()
helper.face_detector = _YoloDetector()
pack = loader_module.KEEPModelPack(
_StubKeepNet(), helper, None, None, "KEEP"
)
detector = pack.face_helper.face_detector
assert detector.device == torch.device("meta") # the sentinel so far
pack.load_device()
assert detector.moved_to[-1] == pack.device
assert detector.device == pack.device
pack.offload()
assert detector.moved_to[-1] == pack.offload_device
assert detector.device == pack.offload_device
@pytest.mark.unit
def test_the_cpu_mode_places_the_nets_on_cpu(
loader_module, model_management, monkeypatch
):
"""P-D5: the CPU leg places the nets on ``cpu`` with the same calls.
The CPU/GPU-leg comparison of the residency (the B7 guard): the
placement is one more residency step that must not branch on the device,
and on the CPU leg it must land on ``cpu`` - the inherent fallback of a
CPU-only install, where ``get_torch_device()`` answers ``cpu``.
"""
def leg(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"
)
assert pack.device == device # the patched device reaches the pack
pack.load_device()
pack.offload()
placements = (
pack.face_helper.face_detector.moved_to,
pack.face_helper.face_parse.moved_to,
)
return [call[0] for call in model_management.calls], placements
gpu_calls, _ = leg(torch.device("cuda"))
cpu_calls, cpu_placements = leg(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
# both nets were placed on cpu (the placement) and moved back (the
# offload) - exactly the two moves, no device branch
for moved_to in cpu_placements:
assert moved_to == [torch.device("cpu"), torch.device("cpu")], moved_to
# ---------------------------------------------------------------------------
# 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_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, 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)
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
# D5: only the **selected** detector and the face parsing net are
# fetched - the other detectors of the manifest are never downloaded.
assert "detection_Resnet50_Final.pth" in urls, urls
assert "parsing_parsenet.pth" in urls, 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"
Reference in New Issue View Git Blame Copy Permalink