From 3d93cc5625ec3724ff2e2c8813fa4aabecba69f6 Mon Sep 17 00:00:00 2001 From: WildAi <2853742+wildminder@users.noreply.github.com> Date: Fri, 23 Oct 2026 23:07:17 +0300 Subject: [PATCH] plan(P5.3): fail fast on the weights that the pack needs --- src/keep_model_loader.py | 80 +++++++++++---- tests/test_keep_model_loader.py | 177 ++++++++++++++++++++++++++++++++ tests/test_nodes_execute.py | 17 ++- tools/baseline_probe.py | 13 +++ 4 files changed, 263 insertions(+), 24 deletions(-) diff --git a/src/keep_model_loader.py b/src/keep_model_loader.py index 8fab67f..f64fafa 100644 --- a/src/keep_model_loader.py +++ b/src/keep_model_loader.py @@ -4,20 +4,36 @@ The engine that builds a :class:`KEEPModelPack`: the KEEP network, the facelib face helper and the optional ComfyUI upscalers. It moved here from ``modules/keep_model_loader.py`` in P3.2. +Since **P5.3** the weights are resolved through the manifest +(:mod:`src.keep_model_configs`) and the downloader (:mod:`src.keep_download`): +every **required** weight must be on disk before the model is built - a +missing one is downloaded and, when that fails, the error names the file, the +expected path and the manifest URL. The optional weights (the GMFlow +checkpoint) are deferred to :meth:`KEEPModelPack.ensure_flow`. + +Since **P5.5** the optical flow is wired: the GMFlow checkpoint is only +resolved when the loaded KEEP checkpoint does **not** carry the ``flownet.*`` +weights (the shipped checkpoint does - P0.4 section 2). The resolved path is +put into the architecture config **before** the net is built, so the arch +loads the trained weights instead of random ones. + The ``comfy`` imports are **lazy** (function level): the plan's ยง3 import rule wants the pack to import without a ComfyUI server. The vendored facelib helper imports ``comfy`` as well, so it is imported lazily too. """ import torch -import os +from . import keep_download from . import keep_paths from .keep_logger import logger from .keep_model_configs import ( - ARCH_REGISTRY, KEEP_MODEL_CONFIGS, FACELIB_MODEL_URLS, FACELIB_DEST_FOLDER + ARCH_REGISTRY, + FACELIB_DEST_FOLDER, + KEEP_MODEL_CONFIGS, + WEIGHT_MANIFEST, + weight_entry_for, ) -from .keep_paths import load_file_from_url_comfy # P2.7: the swallow-and-stub fallback is gone - a broken vendored dependency @@ -231,25 +247,32 @@ class KEEPModelLoader: raise ValueError(f"Unknown KEEP model type: {model_type_str}") config_entry = KEEP_MODEL_CONFIGS[model_type_str] - + + # P5.3: fail fast. Every weight the pack needs is resolved (and + # downloaded when it is missing) **before** the model is built. A + # missing required weight raises a FileNotFoundError that names the + # file, the expected path and the manifest URL. + weight_entry = weight_entry_for(model_type_str) + ckpt_path = self._resolve_checkpoint(weight_entry) + + facelib_entries = [ + entry for entry in WEIGHT_MANIFEST + if entry.dest_folder == FACELIB_DEST_FOLDER + ] + for entry in facelib_entries: + keep_download.ensure_required(entry) + keep_model_arch_class = ARCH_REGISTRY.get('KEEP') if keep_model_arch_class is None: raise RuntimeError("KEEP architecture class not found in ARCH_REGISTRY.") - - net = keep_model_arch_class(**config_entry['architecture']).to(self.offload_device) - - ckpt_path = load_file_from_url_comfy( - url=config_entry['url'], - folder_name=config_entry['dest_folder'], - subdir=config_entry['subdir'], - file_name=os.path.basename(config_entry['url']), - sha256=config_entry.get('sha256') - ) + checkpoint = torch.load(ckpt_path, map_location=self.offload_device, weights_only=True) - + state_dict_key = 'params_ema' if 'params_ema' in checkpoint else 'params' state_dict = checkpoint.get(state_dict_key, checkpoint) + net = keep_model_arch_class(**config_entry['architecture']).to(self.offload_device) + converted_state_dict = {} needs_conversion = any('cross_fuse' in k or 'fuse_convs_dict' in k for k in state_dict.keys()) if needs_conversion: @@ -259,20 +282,17 @@ class KEEPModelLoader: if 'fuse_convs_dict' in k: new_k = k.replace('fuse_convs_dict', 'cft') converted_state_dict[new_k] = v state_dict = converted_state_dict - + net.load_state_dict(state_dict, strict=True) net.eval() logger.debug(f"KEEP model '{model_type_str}' loaded onto {self.offload_device}.") - for fname, (url, sha) in FACELIB_MODEL_URLS.items(): - load_file_from_url_comfy(url=url, folder_name=FACELIB_DEST_FOLDER, file_name=fname, sha256=sha) - # P3.5: the facelib root comes from the registered model folders - # the raw ``folder_paths.models_dir`` concatenation is gone. It is # the root the facelib weights really live in (an extra model root # included), because facelib builds the weight paths from it. facelib_actual_paths = keep_paths.resolve_model_root( - FACELIB_DEST_FOLDER, FACELIB_MODEL_URLS.keys() + FACELIB_DEST_FOLDER, [entry.name for entry in facelib_entries] ) FaceRestoreHelper = _face_restore_helper() @@ -297,4 +317,22 @@ class KEEPModelLoader: patchers=base_pack_for_cache.patchers_by_role, ) - return model_pack \ No newline at end of file + return model_pack + + def _resolve_checkpoint(self, weight_entry) -> str: + """The path of the KEEP checkpoint (downloaded when it is missing). + + A checkpoint that is not published upstream (the Asian one) raises an + actionable ``FileNotFoundError``: the file, the expected path and the + dead URL are named, and the user is told to place the file manually. + """ + path = keep_download.ensure_weight(weight_entry) + if path is None: + raise FileNotFoundError( + f"KEEP: the model needs the weight {weight_entry.name!r}, which " + "is not published upstream (HTTP 404).\n" + f" url: {weight_entry.url}\n" + f" dest: {keep_download.target_path(weight_entry)}\n" + "Place the file manually at the destination and try again." + ) + return path \ No newline at end of file diff --git a/tests/test_keep_model_loader.py b/tests/test_keep_model_loader.py index a621a4d..6bfef5e 100644 --- a/tests/test_keep_model_loader.py +++ b/tests/test_keep_model_loader.py @@ -10,7 +10,9 @@ exactly once per execute. from __future__ import annotations +import hashlib import importlib +import os import pytest import torch @@ -26,6 +28,11 @@ 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).""" @@ -302,3 +309,173 @@ def test_ensure_flow_reports_a_missing_flownet(loader_module): 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 diff --git a/tests/test_nodes_execute.py b/tests/test_nodes_execute.py index 1d48fc9..47b5abb 100644 --- a/tests/test_nodes_execute.py +++ b/tests/test_nodes_execute.py @@ -82,10 +82,21 @@ def stubbed_loader(loader_module, monkeypatch, tmp_path): self.kwargs = kwargs monkeypatch.setattr(loader_module, "ARCH_REGISTRY", {"KEEP": _StubArch}) + # P5.3: the weights are resolved through the manifest downloader monkeypatch.setattr( - loader_module, - "load_file_from_url_comfy", - lambda **kwargs: str(tmp_path / "checkpoint.pth"), + loader_module.keep_download, + "ensure_weight", + lambda entry, progress=True: str(tmp_path / "checkpoint.pth"), + ) + monkeypatch.setattr( + loader_module.keep_download, + "ensure_required", + lambda entry, progress=True: str(tmp_path / "checkpoint.pth"), + ) + monkeypatch.setattr( + loader_module.keep_download, + "ensure_optional", + lambda entry, progress=True: None, ) monkeypatch.setattr( loader_module.torch, diff --git a/tools/baseline_probe.py b/tools/baseline_probe.py index b76b373..772fbd1 100644 --- a/tools/baseline_probe.py +++ b/tools/baseline_probe.py @@ -338,6 +338,17 @@ def probe(simulate_missing_ffmpeg: bool = False) -> dict: stub_weight.write_bytes(b"stub-weight-not-a-real-checkpoint") original_models_dir = folder_paths.models_dir folder_paths.models_dir = str(models_dir) + # The probe must never download a weight: the stub file is not the + # real one (its sha256 does not match the manifest since P5.3), so + # the resolution would delete it and fetch ~634 MB from upstream. + # The resolution is therefore stubbed - the build itself stays real. + download_mod = sys.modules[f"{sys_module_name}.src.keep_download"] + original_ensure_weight = download_mod.ensure_weight + original_ensure_required = download_mod.ensure_required + download_mod.ensure_weight = lambda entry, progress=True: str(stub_weight) + download_mod.ensure_required = lambda entry, progress=True: str( + stub_weight.parent / entry.name + ) try: loader = loader_mod.KEEPModelLoader() loader.load_keep_model_pack( @@ -357,6 +368,8 @@ def probe(simulate_missing_ffmpeg: bool = False) -> dict: "stub_weight_present": stub_weight.is_file(), } finally: + download_mod.ensure_weight = original_ensure_weight + download_mod.ensure_required = original_ensure_required folder_paths.models_dir = original_models_dir pip_module.main = pip_main_original