"""P0.1 baseline probe - records how the *current* ComfyUI-KEEP package behaves. This tool changes no production code. It reproduces, in a throw-away interpreter, the two things the modernization plan needs a verified baseline for: 1. the package **imports** through ComfyUI's own custom-node loader (``importlib.util.spec_from_file_location`` + alias, no server), and 2. the **model build fails**, at the eager basicsr arch scan, with ``ImportError: cannot import name 'T' from 're'`` - swallowed by the swallow-and-stub blocks of B9, which is why the import itself still succeeds. Run it with the embedded interpreter: python tools/baseline_probe.py Useful flags: --json PATH write the machine-readable report to PATH --comfyui-root PATH ComfyUI source tree to put on sys.path (default: $COMFYUI_ROOT) --simulate-missing-ffmpeg block ``import ffmpeg`` and record (instead of perform) the ``pip.main([...])`` call of the vendored ``basicsr/utils/video_util.py`` import-time shim. Exit code is 0 when the probe ran; the *content* of the report (not the exit code) is the baseline evidence. """ from __future__ import annotations import argparse import asyncio import importlib import importlib.util import json import os import sys import tempfile import traceback from pathlib import Path REPO_ROOT = Path(__file__).resolve().parent.parent def resolve_comfyui_root(explicit: str | None) -> str | None: """The ComfyUI source tree to probe against. The tree comes from the explicit argument, else from ``COMFYUI_ROOT``. There is no machine-local fallback: the probe runs against a tree the caller names. """ candidates = [] if explicit: candidates.append(explicit) env_root = os.environ.get("COMFYUI_ROOT") if env_root: candidates.append(env_root) for candidate in candidates: if candidate and (Path(candidate) / "folder_paths.py").is_file(): return str(Path(candidate)) return None def _exc_info(exc: BaseException) -> dict: return { "type": type(exc).__name__, "message": str(exc), "traceback": "".join( traceback.format_exception(type(exc), exc, exc.__traceback__) ), } def _symbol_origin(module, name: str) -> str | None: """``__module__`` of ``module.name`` (``None`` when the symbol is absent).""" symbol = getattr(module, name, None) return getattr(symbol, "__module__", None) def _is_local_stub(module, name: str) -> bool: """``True`` when ``module.name`` is a stub defined in ``module`` itself. A symbol that does not exist at all (P2.1 removed the vendored ``VideoReader`` / ``VideoWriter``) is **not** a stub. """ origin = _symbol_origin(module, name) return origin is not None and origin == module.__name__ class _BlockFFmpegFinder: """meta_path finder that makes ``import ffmpeg`` fail (simulates a host without ``ffmpeg-python``) so the vendored pip shim can be observed.""" def find_module(self, fullname, path=None): # pragma: no cover - py2 API return self.find_spec(fullname, path) def find_spec(self, fullname, path=None, target=None): if fullname == "ffmpeg" or fullname.startswith("ffmpeg."): raise ImportError(f"blocked by probe: {fullname} is unavailable") return None def _install_pip_main_recorder(): """Replace ``pip.main`` with a recorder so the shim can be *observed* without actually installing anything into the environment.""" calls = [] import pip # noqa: PLC0415 - optional dependency, imported defensively original = pip.main def recorder(args, *a, **kw): calls.append(list(args)) return 0 pip.main = recorder return calls, original, pip def load_package_like_comfyui(package_dir: Path): """Load the package exactly the way ``ComfyUI/nodes.py:load_custom_node`` does. ComfyUI (v0.42.0, ``nodes.py:2255-2276``) uses, for a directory package: sys_module_name = module_path.replace(".", "_x_") module_spec = spec_from_file_location(sys_module_name, /__init__.py) module = importlib.util.module_from_spec(module_spec) sys.modules[sys_module_name] = module module_spec.loader.exec_module(module) """ module_path = str(package_dir) sys_module_name = module_path.replace(".", "_x_") module_spec = importlib.util.spec_from_file_location( sys_module_name, os.path.join(module_path, "__init__.py") ) module = importlib.util.module_from_spec(module_spec) sys.modules[sys_module_name] = module module_spec.loader.exec_module(module) return sys_module_name, module def _v3_nodes(package) -> tuple[list[str], dict[str, str]]: """The node ids and display names of a **V3** package (P4.4). The extension is loaded the way ``nodes.py:2316`` does it: through the module level ``comfy_entrypoint``. An empty result means the package is not a V3 extension. """ entrypoint = getattr(package, "comfy_entrypoint", None) if entrypoint is None: return [], {} async def collect(): extension = await entrypoint() await extension.on_load() node_list = await extension.get_node_list() return [cls.GET_SCHEMA() for cls in node_list] try: schemas = asyncio.run(collect()) except Exception: # noqa: BLE001 - the probe reports everything return [], {} return ( sorted(schema.node_id for schema in schemas), {schema.node_id: schema.display_name for schema in schemas}, ) def probe(simulate_missing_ffmpeg: bool = False) -> dict: report: dict = { "repo_root": str(REPO_ROOT), "python": sys.version, "python_executable": sys.executable, "simulate_missing_ffmpeg": simulate_missing_ffmpeg, } comfyui_root = resolve_comfyui_root(None) report["comfyui_root"] = comfyui_root if comfyui_root is None: report["fatal"] = ( "No ComfyUI source tree found. Set COMFYUI_ROOT or pass --comfyui-root." ) return report sys.path.insert(0, comfyui_root) # Record pip.main calls *before* the package import so the vendored # import-time shim cannot install anything into the shared environment. pip_calls, pip_main_original, pip_module = _install_pip_main_recorder() if simulate_missing_ffmpeg: sys.meta_path.insert(0, _BlockFFmpegFinder()) report["pip_main_calls"] = pip_calls # ------------------------------------------------------------------ # 1. import through ComfyUI's loader # ------------------------------------------------------------------ package_dir = REPO_ROOT try: sys_module_name, package = load_package_like_comfyui(package_dir) report["import"] = {"ok": True, "error": None, "sys_module_name": sys_module_name} report["node_class_mappings"] = sorted( getattr(package, "NODE_CLASS_MAPPINGS", {}).keys() ) report["node_display_name_mappings"] = dict( getattr(package, "NODE_DISPLAY_NAME_MAPPINGS", {}) ) # P4.4: the pack registers as a V3 extension - the node ids live in # the schemas of the node classes, not in a V1 mapping. v3_ids, v3_display_names = _v3_nodes(package) report["v3_node_ids"] = v3_ids report["v3_display_names"] = v3_display_names except BaseException as exc: # noqa: BLE001 - the probe reports everything report["import"] = {"ok": False, "error": _exc_info(exc)} pip_module.main = pip_main_original return report utils_mod = sys.modules[f"{sys_module_name}.src.keep_model_configs"] loader_mod = sys.modules[f"{sys_module_name}.src.keep_model_loader"] processor_mod = sys.modules[f"{sys_module_name}.src.keep_processor"] # ------------------------------------------------------------------ # 2. which swallow-and-stub block swallowed which failure (B9) # ------------------------------------------------------------------ arch_registry = getattr(utils_mod, "ARCH_REGISTRY", None) registry_type = type(arch_registry).__name__ if arch_registry is not None else None try: has_keep = arch_registry.get("KEEP") is not None except Exception as exc: # noqa: BLE001 has_keep = f"raised {type(exc).__name__}: {exc}" report["arch_registry"] = { "class_name": registry_type, "module": type(arch_registry).__module__ if arch_registry is not None else None, "is_stub": registry_type == "DummyArchRegistry", "get_keep": has_keep, } report["stubs"] = { "utils_arch_registry_dummy": registry_type == "DummyArchRegistry", "utils_load_file_from_url_fallback": ( # P4.7: the dead import is gone - an absent symbol is not a stub. getattr(utils_mod, "load_file_from_url", None) is not None and getattr(utils_mod, "load_file_from_url").__module__ == utils_mod.__name__ ), "processor_img2tensor_is_dummy": ( getattr(processor_mod, "img2tensor").__module__ == processor_mod.__name__ ), "processor_tensor2img_is_dummy": ( getattr(processor_mod, "tensor2img").__module__ == processor_mod.__name__ ), "processor_is_gray_is_dummy": ( getattr(processor_mod, "is_gray").__module__ == processor_mod.__name__ ), # P2.1 deleted the vendored ``video_util`` module (its import-time pip # shim); ``VideoReader`` / ``VideoWriter`` are therefore *absent* from # the processor instead of being dummy stubs. An absent symbol is not # a stub - the dead ``process_video`` path raises instead. "processor_video_reader_is_dummy": _is_local_stub( processor_mod, "VideoReader" ), "processor_video_writer_is_dummy": _is_local_stub( processor_mod, "VideoWriter" ), "loader_face_restore_helper_is_dummy": ( getattr(loader_mod, "FaceRestoreHelper").__module__ == loader_mod.__name__ ), } report["stub_summary"] = { "fired": sorted(k for k, v in report["stubs"].items() if v), "not_fired": sorted(k for k, v in report["stubs"].items() if not v), } report["symbol_origins"] = { "img2tensor": getattr(processor_mod, "img2tensor").__module__, "tensor2img": getattr(processor_mod, "tensor2img").__module__, "is_gray": getattr(processor_mod, "is_gray").__module__, "VideoReader": _symbol_origin(processor_mod, "VideoReader"), "FaceRestoreHelper": getattr(loader_mod, "FaceRestoreHelper").__module__, } # ------------------------------------------------------------------ # 3. state of the interpreter *right after* the package import: # the pip shim (B10) and the global sys.path shadowing (B4). # Captured before the cleanup below, because the eager arch scan # test deliberately purges the partially imported 'basicsr'. # ------------------------------------------------------------------ deps_path = str(REPO_ROOT / "modules" / "deps") try: ffmpeg_importable = importlib.util.find_spec("ffmpeg") is not None except ImportError: ffmpeg_importable = False report["pip_shim"] = { "video_util_imported": any( name.endswith("basicsr.utils.video_util") for name in sys.modules ), "ffmpeg_importable": ffmpeg_importable, "pip_main_calls": list(pip_calls), } report["sys_path_shadowing"] = { "deps_path_on_sys_path": deps_path in sys.path, "basicsr_in_sys_modules": "basicsr" in sys.modules, "basicsr_utils_in_sys_modules": "basicsr.utils" in sys.modules, "facelib_in_sys_modules": "facelib" in sys.modules, } # ------------------------------------------------------------------ # 4. the *real* (unmasked) import of the vendored KEEP arch (B1/B2) # ------------------------------------------------------------------ try: importlib.import_module( f"{sys_module_name}.modules.deps.basicsr.archs.keep_arch" ) report["arch_scan"] = {"ok": True, "error": None} except BaseException as exc: # noqa: BLE001 report["arch_scan"] = {"ok": False, "error": _exc_info(exc)} # ------------------------------------------------------------------ # 5. the diffusers import (B2) - only symbol source in the tree # ------------------------------------------------------------------ try: importlib.import_module("diffusers.models.attention") report["diffusers"] = {"ok": True, "error": None} except BaseException as exc: # noqa: BLE001 report["diffusers"] = {"ok": False, "error": _exc_info(exc)} # ------------------------------------------------------------------ # 6. a real model build, with a stubbed weight file present # ------------------------------------------------------------------ import folder_paths # noqa: PLC0415 - needs the ComfyUI root on sys.path with tempfile.TemporaryDirectory(prefix="keep-baseline-weights-") as tmp: models_dir = Path(tmp) stub_weight = models_dir / "keep_models" / "KEEP" / "KEEP-b76feb75.pth" stub_weight.parent.mkdir(parents=True, exist_ok=True) 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( model_type_str="KEEP", detection_model_str="retinaface_resnet50" ) report["model_build"] = { "ok": True, "error": None, "stub_weight_path": str(stub_weight), "stub_weight_present": stub_weight.is_file(), } except BaseException as exc: # noqa: BLE001 report["model_build"] = { "ok": False, "error": _exc_info(exc), "stub_weight_path": str(stub_weight), "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 return report def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--json", dest="json_path", default=None) parser.add_argument("--comfyui-root", dest="comfyui_root", default=None) parser.add_argument( "--simulate-missing-ffmpeg", action="store_true", help="block 'import ffmpeg' and record the vendored pip.main shim call", ) args = parser.parse_args(argv) if args.comfyui_root: os.environ["COMFYUI_ROOT"] = args.comfyui_root report = probe(simulate_missing_ffmpeg=args.simulate_missing_ffmpeg) text = json.dumps(report, indent=2, sort_keys=True) if args.json_path: Path(args.json_path).write_text(text + "\n", encoding="utf-8") print(text) return 0 if "fatal" not in report else 2 if __name__ == "__main__": raise SystemExit(main())