plan(P5.3): fail fast on the weights that the pack needs

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
+        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

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

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,

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