diff --git a/image_metrics.py b/image_metrics.py index aaa8119..c80be10 100644 --- a/image_metrics.py +++ b/image_metrics.py @@ -35,32 +35,44 @@ _LUMA_W = (0.299, 0.587, 0.114) # --------------------------------------------------------------------------- -# Optional scorer packs (adapters). Entries stay DISABLED until the node ids -# have been confirmed against the real pack: nothing here can be verified -# without the pack installed. A disabled entry is reported to the Builder UI -# but never runs and is dropped by normalize_spec(). +# Optional scorer packs (adapters). Each entry is checked against the pack's +# real node signatures (pinned shas below). A disabled entry is reported to the +# Builder UI but never runs and is dropped by normalize_spec(). Keys: +# loaders[].prefer = {combo_input: [preferred values]} (first one present wins), +# scale = multiplier applied to the node's output, max_batch = images per call. # --------------------------------------------------------------------------- ADAPTERS = { + # strimmlarn/ComfyUI-Strimmlarns-Aesthetic-Score@42bed3fd78488dda50d83bdc39b5cff6772919c9 + # (LoadAesteticModel: combo model_name; CalculateAestheticScore: image + aesthetic_model, + # scores image[0] only, returns int(score * 100)) "aesthetic": { - "enabled": False, - "label": "Aesthetic score", + "enabled": True, + "label": "Aesthetic score (LAION, ~1-10)", "node_id": "CalculateAestheticScore", "image_arg": "image", - "loaders": [{"node_id": "LoadAesteticModel", "kwargs": {}, "pass_as": "aesthetic_model", "output_index": 0}], + "loaders": [{"node_id": "LoadAesteticModel", "kwargs": {}, "pass_as": "aesthetic_model", "output_index": 0, + "prefer": {"model_name": ["sac+logos+ava1-l14-linearMSE.pth", "chadscorer.pth"]}}], "output_index": 0, - "pack_hint": "ComfyUI-Strimmlarns-Aesthetic-Score", - "note": "node ids not yet verified against the pack", + "scale": 0.01, + "max_batch": 1, + "vram_mb": 1500, + "pack_hint": "ComfyUI-Strimmlarns-Aesthetic-Score (+ pip pytorch-lightning, openai CLIP; model in models/aesthetic/)", + "note": "verified against strimmlarn@42bed3f; the node returns score x100 (int), rescaled", }, + # ZaneA/ComfyUI-ImageReward@5e9936ded972d8447c26ea2670a92e999036821d + # (ImageRewardLoader: STRING model; ImageRewardScore: model + prompt + images, FLOAT mean) "imagereward": { "enabled": False, - "label": "ImageReward", + "label": "ImageReward (prompt-aware, ~-2..2)", "node_id": "ImageRewardScore", "image_arg": "images", "text_arg": "prompt", - "loaders": [{"node_id": "ImageRewardLoader", "kwargs": {}, "pass_as": "model", "output_index": 0}], + "loaders": [{"node_id": "ImageRewardLoader", "kwargs": {"model": "ImageReward-v1.0"}, "pass_as": "model", "output_index": 0}], "output_index": 0, - "pack_hint": "ComfyUI-ImageReward", - "note": "node ids not yet verified against the pack", + "max_batch": 1, + "vram_mb": 2500, + "pack_hint": "ComfyUI-ImageReward (pip image-reward==1.4; downloads ~1.8 GB on first use)", + "note": "signatures verified against ZaneA@5e9936d; not live-tested, so disabled", }, } @@ -316,6 +328,7 @@ class _RunContext: self.disabled_adapters = set() self.checked_adapters = set() self.times = [] # seconds per image (builtins only) + self.adapter_times = {} # adapter name -> [seconds per call] self.lock = threading.Lock() self.closed = False @@ -355,6 +368,12 @@ def end_run(token=None): with ctx.lock: times = list(ctx.times) ctx.loader_cache.clear() + adapter_times = {k: list(v) for k, v in ctx.adapter_times.items()} + if adapter_times: + _release_vram(None) + for aname, ts in adapter_times.items(): + if ts: + print(f"[GridTester] 📏 metric adapter '{aname}': {len(ts)} images, mean {sum(ts) * 1000.0 / len(ts):.0f} ms") if times: ms = [t * 1000.0 for t in times] print(f"[GridTester] 📏 metrics: {len(ms)} images, mean {sum(ms) / len(ms):.1f} ms, max {max(ms):.1f} ms") @@ -558,12 +577,68 @@ def _adapter_defs(name): return ad if ad and ad.get("enabled") else None +def _combo_options(cls, name): + """Option list of a V1 combo input (required or optional), or None when `name` is not a combo.""" + try: + types = cls.INPUT_TYPES() + for section in ("required", "optional"): + v = types.get(section, {}).get(name) + if v is not None: + kind = v[0] if isinstance(v, (tuple, list)) and v else v + if isinstance(kind, (list, tuple)): + return list(kind) + return None + except Exception: + pass + return None + + +def _loader_kwargs(cls, loader): + """Kwargs for a loader node: required-input defaults (combo -> first option), then the + `prefer` picks (first preferred value present in the combo), then explicit `kwargs`. + + Returns (kwargs, problem). problem is a short reason when a required input stays uncovered + or a combo has no options at all (e.g. no model file in the pack's folder). + """ + kwargs = {} + req = _required_inputs(cls) + missing = None + if req is not None: + req, missing = req + explicit = dict(loader.get("kwargs", {})) + prefer = loader.get("prefer", {}) or {} + for name, default in (req or {}).items(): + opts = _combo_options(cls, name) + if opts is not None and not opts and name not in explicit: + if loader.get("node_id") == "LoadAesteticModel": + return kwargs, "no model file in models/aesthetic (download sac+logos+ava1-l14-linearMSE.pth)" + return kwargs, f"'{loader['node_id']}' input '{name}' has no options (no model file installed?)" + if default is not missing: + kwargs[name] = default + for name, wanted in prefer.items(): + opts = _combo_options(cls, name) or [] + for w in wanted: + if w in opts: + kwargs[name] = w + break + kwargs.update(explicit) + if req is not None: + for name, default in req.items(): + if name not in kwargs and default is missing: + return kwargs, f"required input '{name}' of '{loader['node_id']}' is not covered" + return kwargs, None + + def _adapter_problem(ad, mappings): """Return None when installed and call-compatible, else a short reason.""" ids = [ad["node_id"]] + [l["node_id"] for l in ad.get("loaders", [])] for nid in ids: if nid not in mappings: return f"node '{nid}' not installed" + for loader in ad.get("loaders", []): + _, problem = _loader_kwargs(mappings[loader["node_id"]], loader) + if problem: + return problem covered = {ad.get("image_arg"), ad.get("text_arg")} | set(ad.get("kwargs", {})) covered |= {l["pass_as"] for l in ad.get("loaders", [])} covered.discard(None) @@ -577,26 +652,49 @@ def _adapter_problem(ad, mappings): def adapter_status(): - """For the Builder hint route: [{name,label,installed,enabled,pack_hint,note}].""" + """For the Builder hint route: [{name,label,installed,enabled,problem,pack_hint,note}].""" mappings = _node_mappings() out = [] for name, ad in ADAPTERS.items(): + installed = all(n in mappings for n in [ad["node_id"]] + [l["node_id"] for l in ad.get("loaders", [])]) + try: + problem = _adapter_problem(ad, mappings) + except Exception as e: + problem = str(e) out.append({ "name": name, "label": ad.get("label", name), "enabled": bool(ad.get("enabled")), - "installed": all(n in mappings for n in [ad["node_id"]] + [l["node_id"] for l in ad.get("loaders", [])]), + "installed": installed, + "problem": problem, "pack_hint": ad.get("pack_hint", ""), "note": ad.get("note", ""), }) return out +def _release_vram(need_mb): + """Ask ComfyUI's memory manager to make room (need_mb) or just empty the cache (None). Never raises.""" + try: + import torch + if not torch.cuda.is_available(): + return + import comfy.model_management as mm + if need_mb: + mm.free_memory(int(need_mb) * 1024 * 1024, mm.get_torch_device()) + else: + import gc + gc.collect() + mm.soft_empty_cache() + except Exception: + pass + + def _adapter_image(image): """Whatever we were given -> [B,H,W,3] float tensor in [0,1] when torch is real (else numpy).""" torch = _real_torch() if _is_tensor(image): - return _to_bhwc_full(image) + return _to_bhwc_full(image).cpu() # some packs call .numpy() on the input arr = np.asarray(image.convert("RGB")) if hasattr(image, "convert") else np.asarray(image) if arr.dtype == np.uint8: arr = arr.astype(np.float32) / 255.0 @@ -643,14 +741,19 @@ def _run_adapters(ctx, names, image, meta): raise RuntimeError(problem) if img is None: img = _adapter_image(image) + t0 = time.perf_counter() outputs = ctx.loader_cache.setdefault(name, {}) kwargs = dict(ad.get("kwargs", {})) for li, loader in enumerate(ad.get("loaders", [])): if li not in outputs: - res = _call_node(mappings[loader["node_id"]], dict(loader.get("kwargs", {}))) + lk, lproblem = _loader_kwargs(mappings[loader["node_id"]], loader) + if lproblem: + raise RuntimeError(lproblem) + res = _call_node(mappings[loader["node_id"]], lk) outputs[li] = _output(res, loader.get("output_index", 0)) kwargs[loader["pass_as"]] = outputs[li] - kwargs[ad["image_arg"]] = img + mb = ad.get("max_batch") + kwargs[ad["image_arg"]] = img[:mb] if mb and getattr(img, "shape", None) is not None and img.shape[0] > mb else img if ad.get("text_arg"): kwargs[ad["text_arg"]] = str(meta.get("positive") or meta.get("config_positive") or "") cls = mappings[ad["node_id"]] @@ -660,8 +763,12 @@ def _run_adapters(ctx, names, image, meta): for k, d in req.items(): if k not in kwargs and d is not missing: kwargs[k] = d # widget defaults + _release_vram(ad.get("vram_mb")) res = _call_node(cls, kwargs) - out[name] = _round4(_float_of(_output(res, ad.get("output_index", 0)))) + out[name] = _round4(_float_of(_output(res, ad.get("output_index", 0))) * ad.get("scale", 1.0)) + _release_vram(None) + with ctx.lock: + ctx.adapter_times.setdefault(name, []).append(time.perf_counter() - t0) except Exception as e: ctx.disabled_adapters.add(name) _warn_once(ctx, f"adapter:{name}", f"metric adapter '{name}' disabled: {e}") diff --git a/tests/test_image_metrics.py b/tests/test_image_metrics.py index cda9c14..05a6757 100644 --- a/tests/test_image_metrics.py +++ b/tests/test_image_metrics.py @@ -90,9 +90,9 @@ def test_normalize_spec(): assert im.normalize_spec({"enabled": True, "builtins": ["bogus"]}) is None assert im.normalize_spec({"enabled": True, "builtins": ["luma", "bogus", "entropy"]}) == \ {"builtins": ["entropy", "luma"], "adapters": []} - # idempotent on a normalized spec; disabled registry entries are dropped - s = im.normalize_spec({"enabled": True, "builtins": ["luma"], "adapters": ["aesthetic"]}) - assert s == {"builtins": ["luma"], "adapters": []} + # idempotent on a normalized spec; disabled registry entries (imagereward) are dropped + s = im.normalize_spec({"enabled": True, "builtins": ["luma"], "adapters": ["aesthetic", "imagereward"]}) + assert s == {"builtins": ["luma"], "adapters": ["aesthetic"]} assert im.normalize_spec(s) == s assert im.normalize_spec({"enabled": True, "builtins": {"luma": True, "entropy": False}}) == \ {"builtins": ["luma"], "adapters": []} diff --git a/tests/test_metrics_hooks.py b/tests/test_metrics_hooks.py index 802bd8b..a3fc3cf 100644 --- a/tests/test_metrics_hooks.py +++ b/tests/test_metrics_hooks.py @@ -297,6 +297,232 @@ def test_v3_style_node_and_prompt_arg(monkeypatch): assert meta["metrics"] == {"fake": 7.5} and seen["prompt"] == "a red cat" -def test_shipped_registry_entries_are_disabled(): - assert im.ADAPTERS and all(not a["enabled"] for a in im.ADAPTERS.values()) - assert im.normalize_spec({"enabled": True, "builtins": ["luma"], "adapters": ["aesthetic"]})["adapters"] == [] +# ---- real pack signatures (copied from the packs' source) -------------------- +# strimmlarn/ComfyUI-Strimmlarns-Aesthetic-Score@42bed3fd78488dda50d83bdc39b5cff6772919c9 my_node.py +_STRIM_FILES = ["sac+logos+ava1-l14-linearMSE.pth"] + + +class _StrimLoader: + calls = 0 + FUNCTION = "load_model" + RETURN_TYPES = ("AESTHETIC_MODEL",) + seen = [] + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"model_name": (list(_STRIM_FILES),)}} + + def load_model(self, model_name): + _StrimLoader.calls += 1 + _StrimLoader.seen.append(model_name) + return ("models/aesthetic/" + model_name,) + + +class _StrimScore: + FUNCTION = "execute" + RETURN_TYPES = ("SCORE",) + calls = 0 + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"image": ("IMAGE",), "aesthetic_model": ("AESTHETIC_MODEL",)}} + + def execute(self, image, aesthetic_model): + _StrimScore.calls += 1 + arr = image[0].numpy() if hasattr(image[0], "numpy") else np.asarray(image[0]) # real node needs CPU + assert image.shape[0] == 1 and arr.ndim == 3 + return (612,) + + +# ZaneA/ComfyUI-ImageReward@5e9936ded972d8447c26ea2670a92e999036821d +class _IRLoader: + FUNCTION = "load_model" + RETURN_TYPES = ("IMAGEREWARD_MODEL",) + seen = [] + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"model": ("STRING", {"default": "ImageReward-v1.0"})}} + + def load_model(self, model): + _IRLoader.seen.append(model) + return ("IRMODEL",) + + +class _IRScore: + FUNCTION = "score_images" + RETURN_TYPES = ("FLOAT", "STRING") + seen = [] + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"model": ("IMAGEREWARD_MODEL",), "prompt": ("STRING", {"multiline": True}), + "images": ("IMAGE",)}} + + def score_images(self, model, prompt, images): + _IRScore.seen.append((model, prompt)) + return (0.42, "0.42") + + +@pytest.fixture +def strim_nodes(monkeypatch): + mapping = {"LoadAesteticModel": _StrimLoader, "CalculateAestheticScore": _StrimScore, + "ImageRewardLoader": _IRLoader, "ImageRewardScore": _IRScore} + monkeypatch.setattr(sys.modules["nodes"], "NODE_CLASS_MAPPINGS", mapping, raising=False) + monkeypatch.setattr(_StrimLoader, "INPUT_TYPES", + classmethod(lambda cls: {"required": {"model_name": (list(_STRIM_FILES),)}})) + _StrimLoader.calls = 0 + _StrimLoader.seen = [] + _StrimScore.calls = 0 + _IRLoader.seen = [] + _IRScore.seen = [] + monkeypatch.setattr(sys.modules[__name__], "_STRIM_FILES", ["sac+logos+ava1-l14-linearMSE.pth"]) + return mapping + + +def test_shipped_registry_entries_match_pack_signatures(strim_nodes): + assert im.ADAPTERS["aesthetic"]["enabled"] is True + assert im.ADAPTERS["imagereward"]["enabled"] is False # not live-tested + assert im._adapter_problem(im.ADAPTERS["aesthetic"], strim_nodes) is None + assert im._adapter_problem(im.ADAPTERS["imagereward"], strim_nodes) is None + assert im.normalize_spec({"enabled": True, "builtins": ["luma"], "adapters": ["aesthetic", "imagereward"]})["adapters"] == ["aesthetic"] + + +def test_real_aesthetic_entry_stores_rescaled_score_and_loads_once(strim_nodes): + tok = im.begin_run({"builtins": ["luma"], "adapters": ["aesthetic"]}) + m1, m2 = {}, {} + im.attach(m1, _pil()) + im.attach(m2, _pil()) + im.end_run(tok) + assert m1["metrics"]["aesthetic"] == 6.12 and m2["metrics"]["aesthetic"] == 6.12 + assert _StrimLoader.calls == 1 and _StrimScore.calls == 2 + tok = im.begin_run({"builtins": [], "adapters": ["aesthetic"]}) + im.attach({}, _pil()) + im.end_run(tok) + assert _StrimLoader.calls == 2 # end_run dropped the cache + + +def test_prefer_picks_first_preferred_file(strim_nodes, monkeypatch): + files = ["chadscorer.pth", "sac+logos+ava1-l14-linearMSE.pth"] + monkeypatch.setattr(_StrimLoader, "INPUT_TYPES", classmethod(lambda cls: {"required": {"model_name": (list(files),)}})) + tok = im.begin_run({"builtins": [], "adapters": ["aesthetic"]}) + im.attach({}, _pil()) + im.end_run(tok) + assert _StrimLoader.seen == ["sac+logos+ava1-l14-linearMSE.pth"] + + +def test_without_preferred_file_first_option_is_used(strim_nodes, monkeypatch): + files = ["zzz.pth", "other.pth"] + monkeypatch.setattr(_StrimLoader, "INPUT_TYPES", classmethod(lambda cls: {"required": {"model_name": (list(files),)}})) + tok = im.begin_run({"builtins": [], "adapters": ["aesthetic"]}) + im.attach({}, _pil()) + im.end_run(tok) + assert _StrimLoader.seen == ["zzz.pth"] + + +def test_empty_model_folder_reports_problem_and_disables_once(strim_nodes, monkeypatch, capsys): + monkeypatch.setattr(_StrimLoader, "INPUT_TYPES", classmethod(lambda cls: {"required": {"model_name": ([],)}})) + st = {a["name"]: a for a in im.adapter_status()} + assert st["aesthetic"]["installed"] is True and "models/aesthetic" in st["aesthetic"]["problem"] + tok = im.begin_run({"builtins": ["luma"], "adapters": ["aesthetic"]}) + m1, m2 = {}, {} + im.attach(m1, _pil()) + im.attach(m2, _pil()) + im.end_run(tok) + out = capsys.readouterr().out + assert _StrimLoader.calls == 0 and _StrimScore.calls == 0 + assert out.count("metric adapter 'aesthetic' disabled") == 1 and "models/aesthetic" in out + assert set(m1["metrics"]) == {"luma"} + + +def test_adapter_status_problem_is_none_when_ok_and_set_when_missing(strim_nodes): + st = {a["name"]: a for a in im.adapter_status()} + assert st["aesthetic"]["problem"] is None + del strim_nodes["CalculateAestheticScore"] + st = {a["name"]: a for a in im.adapter_status()} + assert st["aesthetic"]["installed"] is False and "not installed" in st["aesthetic"]["problem"] + + +def test_imagereward_entry_against_real_signature(strim_nodes, monkeypatch): + monkeypatch.setitem(im.ADAPTERS["imagereward"], "enabled", True) + tok = im.begin_run({"builtins": [], "adapters": ["imagereward"]}) + m1, m2 = {"positive": "a red cat"}, {"config_positive": "fallback text"} + im.attach(m1, _pil()) + im.attach(m2, _pil()) + im.end_run(tok) + assert _IRLoader.seen == ["ImageReward-v1.0"] + assert _IRScore.seen == [("IRMODEL", "a red cat"), ("IRMODEL", "fallback text")] + assert m1["metrics"]["imagereward"] == 0.42 + + +def test_scale_applied_and_defaults_to_one(fake_nodes, monkeypatch): + monkeypatch.setattr(im, "ADAPTERS", {"fake": _fake_ad(scale=0.5)}) + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + m = {} + im.attach(m, _pil()) + im.end_run(tok) + assert m["metrics"]["fake"] == pytest.approx(3.0617, abs=2e-4) + monkeypatch.setattr(im, "ADAPTERS", {"fake": _fake_ad()}) + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + m = {} + im.attach(m, _pil()) + im.end_run(tok) + assert m["metrics"]["fake"] == pytest.approx(6.1234, abs=2e-4) + + +def test_adapter_image_from_pil_and_max_batch(): + img = im._adapter_image(_pil(255)) + assert tuple(img.shape) == (1, 24, 32, 3) and float(img.max()) <= 1.0 and float(img.min()) >= 0.0 + torch = pytest.importorskip("torch") + if getattr(torch, "__file__", None) is None or not hasattr(torch, "cuda"): + pytest.skip("stub torch") + t = torch.rand(2, 8, 8, 3) + out = im._adapter_image(t) + assert out.device.type == "cpu" and tuple(out.shape) == (2, 8, 8, 3) + + +def test_max_batch_slices_tensor_input(monkeypatch, fake_nodes): + shapes = [] + + class Rec(_Scorer): + def score(self, image, model, mode, scale): + shapes.append(tuple(image.shape)) + return (1.0,) + fake_nodes["FakeScore"] = Rec + monkeypatch.setattr(im, "ADAPTERS", {"fake": _fake_ad(max_batch=1)}) + monkeypatch.setattr(im, "_adapter_image", lambda image: np.zeros((2, 4, 4, 3), np.float32)) + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + im.attach({}, _pil()) + im.end_run(tok) + assert shapes == [(1, 4, 4, 3)] + + +def test_release_vram_is_noop_without_comfy(): + im._release_vram(1500) + im._release_vram(None) + + +def test_release_vram_calls_comfy_with_bytes(monkeypatch): + calls = [] + comfy_pkg = type(sys)("comfy") + mm = type(sys)("comfy.model_management") + mm.free_memory = lambda n, dev: calls.append(("free", n, dev)) + mm.get_torch_device = lambda: "cuda:0" + mm.soft_empty_cache = lambda: calls.append(("empty",)) + comfy_pkg.model_management = mm + fake_torch = type(sys)("torch") + fake_torch.cuda = type("C", (), {"is_available": staticmethod(lambda: True)}) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "comfy", comfy_pkg) + monkeypatch.setitem(sys.modules, "comfy.model_management", mm) + im._release_vram(1500) + im._release_vram(None) + assert calls == [("free", 1500 * 2 ** 20, "cuda:0"), ("empty",)] + + +def test_adapter_timing_line_printed_at_end_run(strim_nodes, capsys): + tok = im.begin_run({"builtins": [], "adapters": ["aesthetic"]}) + im.attach({}, _pil()) + im.attach({}, _pil()) + im.end_run(tok) + assert "metric adapter 'aesthetic': 2 images, mean" in capsys.readouterr().out