From cb0aa2be61d85dee2b18b6aff6bd401861f1d71f Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Wed, 4 Nov 2026 00:06:37 -0800 Subject: [PATCH] feat(metrics): wire metrics into decode/flush, remote VAE, upscale paths and run context Hooks: decode_latent_with_vae metrics_sink, flush_batch_with_vae, remote VAE worker, orchestrator upscale/Florence2/SeedVR2 inserts, dashboard upscale thread (explicit spec, never inherits), florence2 builders drop metrics. sampler_node wraps the run in a token-based begin_run/end_run context. Adapters mechanism is tested with fake nodes. Co-Authored-By: Claude Fable 5.1 --- florence2_hires.py | 2 +- generation_orchestrator.py | 21 ++- image_generation.py | 20 ++- image_metrics.py | 57 +++++-- remote_vae.py | 7 + sampler_node.py | 32 ++-- tests/test_image_metrics.py | 4 +- tests/test_metrics_hooks.py | 302 ++++++++++++++++++++++++++++++++++++ upscale_runner.py | 19 ++- 9 files changed, 426 insertions(+), 38 deletions(-) create mode 100644 tests/test_metrics_hooks.py diff --git a/florence2_hires.py b/florence2_hires.py index c984eee..7527103 100644 --- a/florence2_hires.py +++ b/florence2_hires.py @@ -850,7 +850,7 @@ def build_florence2_manifest_entry( if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration", "upscale_source", "upscale_pipeline", "upscale_mode", - "upscale_ratio", "upscale_denoise", "upscale_model") + "upscale_ratio", "upscale_denoise", "upscale_model", "metrics") } base.update({ "id": upscale_id, diff --git a/generation_orchestrator.py b/generation_orchestrator.py index 12f7a82..782456c 100644 --- a/generation_orchestrator.py +++ b/generation_orchestrator.py @@ -34,6 +34,7 @@ from .html_generator import get_html_template from .conditioning_cache import ConditioningCache from .remote_vae import RemoteVAEDecodeWorker, is_remote_vae_available, get_endpoint_names, INSTALL_INSTRUCTIONS from . import distribution +from . import image_metrics try: from server import PromptServer @@ -575,8 +576,13 @@ def run_deferred_upscales( is_last_combo = combo == combos[-1] is_final_output = is_last_step and is_last_combo + up_metrics_sink = {} + up_metrics_img = None if isinstance(upscale_result, dict) and "samples" in upscale_result: - upscaled_pil = decode_latent_with_vae(loaded_vae, upscale_result["samples"]) + upscaled_pil = decode_latent_with_vae( + loaded_vae, upscale_result["samples"], + metrics_sink=up_metrics_sink if is_final_output else None) + up_metrics_img = upscaled_pil up_w, up_h = upscaled_pil.size if is_final_output: upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000) @@ -586,6 +592,7 @@ def run_deferred_upscales( format="WEBP", quality=95 ) elif isinstance(upscale_result, PILImage.Image): + up_metrics_img = upscale_result up_w, up_h = upscale_result.size if is_final_output: upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000) @@ -623,6 +630,7 @@ def run_deferred_upscales( upscaled_meta["hires_prompt_behavior"] = hires_prompt_behavior_rt upscaled_meta["hires_prompt_text"] = hires_prompt_text_rt + image_metrics.finalize(upscaled_meta, up_metrics_img, up_metrics_sink) existing_data["items"].append(upscaled_meta) upscale_combo_idx += 1 @@ -2031,6 +2039,7 @@ def run_generation_loop( "florence2_bbox": step_result.get("florence2_bbox"), "hires_denoise": f2_config["hires_denoise"], }) + image_metrics.finalize(upscaled_meta, step_result["image_pil"]) # Insert at beginning + set upscale_produced (matches L2255-2256 # canonical pattern). Without upscale_produced=True, the base # image at L2337 would ALSO get saved as a duplicate. @@ -2087,6 +2096,7 @@ def run_generation_loop( "upscale_mode": "seedvr2", "upscale_model": sv_config.get("dit_model", ""), }) + image_metrics.finalize(upscaled_meta, result_pil) existing_data["items"].append(upscaled_meta) save_manifest(paths["manifest"], existing_data) if PromptServer is not None: @@ -2177,10 +2187,15 @@ def run_generation_loop( # Generate timestamp-based ID (same convention as normal images) upscale_id = int(time.time() * 100000) + upscale_random.randint(0, 1000) upscaled_filename = f"img_{upscale_id}_upscaled.webp" + up_metrics_sink = {} + up_metrics_img = None if isinstance(upscale_result, dict) and "samples" in upscale_result: # HiRes modes return latent — decode to get dimensions - upscaled_pil = decode_latent_with_vae(loaded_vae, upscale_result["samples"]) + upscaled_pil = decode_latent_with_vae( + loaded_vae, upscale_result["samples"], + metrics_sink=up_metrics_sink if is_final_output else None) + up_metrics_img = upscaled_pil up_w, up_h = upscaled_pil.size # Only save to disk if this is a final output if is_final_output: @@ -2194,6 +2209,7 @@ def run_generation_loop( pipe_h = up_h elif isinstance(upscale_result, PILImage.Image): # Model-only returns PIL directly + up_metrics_img = upscale_result up_w, up_h = upscale_result.size # Only save to disk if this is a final output if is_final_output: @@ -2288,6 +2304,7 @@ def run_generation_loop( upscaled_meta["file"] = f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images" upscaled_meta["rejected"] = False + image_metrics.finalize(upscaled_meta, up_metrics_img, up_metrics_sink) # Insert at beginning (same order as flush_batch_with_vae) existing_data["items"].insert(0, upscaled_meta) upscale_produced = True diff --git a/image_generation.py b/image_generation.py index f7ea9f5..60e0b6d 100644 --- a/image_generation.py +++ b/image_generation.py @@ -53,6 +53,9 @@ import nodes import numpy as np from PIL import Image +from . import image_metrics +from .image_metrics import _CURRENT as _METRICS_RUN + def _latent_ratio(patched_model): """Pixels per latent cell for this model (8 for SD/SDXL/Flux.1, 16 for the Flux.2 family).""" @@ -820,13 +823,17 @@ def upscale_image(result_latent, vae, patched_model, upscaling_config, config, p # @torch.inference_mode() — see module docstring. VAE decode is a model forward # pass and needs inference_mode to avoid pinning intermediate activations. @torch.inference_mode() -def decode_latent_with_vae(vae, latent_samples): +def decode_latent_with_vae(vae, latent_samples, metrics_sink=None, metrics_spec=_METRICS_RUN): """ Decode latent samples to pixel space using VAE. Args: vae: VAE model latent_samples: Latent tensor to decode + metrics_sink: optional dict; when given (and metrics are on) the image + metrics are computed on the decoded GPU tensor and written into it + metrics_spec: omitted -> the active generation run's spec; an explicit + spec (or None) for threads that must not read generation-run state Returns: PIL.Image: Decoded image @@ -855,6 +862,9 @@ def decode_latent_with_vae(vae, latent_samples): # already non-grad). Kept for safety in case this function is ever called # without the decorator, but inference_mode is the primary defense — see # module docstring for why. + if metrics_sink is not None: + image_metrics.sink_update(metrics_sink, decoded, metrics_spec) + img_np = decoded.detach().cpu().float().numpy() # Remove extra dimensions (handle shapes like (1, 1, H, W, C) or (1, H, W, C)) @@ -1048,7 +1058,8 @@ def flush_batch_with_vae(pending_batch, vae, img_dir, existing_data, session_nam for latent_samples, meta in pending_batch: # Decode latent to image - image = decode_latent_with_vae(vae, latent_samples) + metrics_sink = {} + image = decode_latent_with_vae(vae, latent_samples, metrics_sink=metrics_sink) # Generate ID using same format as remote_vae ts = int(time.time() * 100000) + random.randint(0, 1000) @@ -1075,6 +1086,11 @@ def flush_batch_with_vae(pending_batch, vae, img_dir, existing_data, session_nam "file": f"/view?filename={filename}&type=output&subfolder=benchmarks/{session_name}/images", "rejected": False }) + + # Optional image metrics (builtins came from the decode sink; adapters run here) + if metrics_sink: + meta["metrics"] = metrics_sink + image_metrics.attach_adapters(meta, image) # Update manifest - insert at beginning (like remote_vae) existing_data["items"].insert(0, meta) diff --git a/image_metrics.py b/image_metrics.py index 57dd662..2f4bfed 100644 --- a/image_metrics.py +++ b/image_metrics.py @@ -314,7 +314,7 @@ class _RunContext: _ctx_lock = threading.Lock() -_current = None # the active generation run's context (or None) +_stack = [] # contexts of runs still open; the newest is the "current" one _CURRENT = object() # sentinel: "use the active run's spec" @@ -326,10 +326,9 @@ def begin_run(spec): meanwhile. The dashboard upscale thread does not use this context at all (it passes an explicit spec to attach()). """ - global _current ctx = _RunContext(normalize_spec(spec)) with _ctx_lock: - _current = ctx + _stack.append(ctx) return ctx @@ -339,14 +338,13 @@ def end_run(token=None): With a token only that context is closed (and `current` cleared only if it still points at it). Without a token the active context is closed. """ - global _current with _ctx_lock: - ctx = token if token is not None else _current + ctx = token if token is not None else (_stack[-1] if _stack else None) if ctx is None or ctx.closed: return ctx.closed = True - if _current is ctx: - _current = None + if ctx in _stack: + _stack.remove(ctx) with ctx.lock: times = list(ctx.times) ctx.loader_cache.clear() @@ -356,13 +354,13 @@ def end_run(token=None): def current_spec(): - ctx = _current - return ctx.spec if ctx is not None and not ctx.closed else None + ctx = _active_ctx() + return ctx.spec if ctx is not None else None def _active_ctx(): - ctx = _current - return ctx if ctx is not None and not ctx.closed else None + with _ctx_lock: + return _stack[-1] if _stack else None def _warn_once(ctx, key, msg): @@ -434,15 +432,42 @@ def attach(meta, image, spec=_CURRENT, builtins=True, adapters=True): if adapters and ctx is not None and spec.get("adapters"): metrics.update(_run_adapters(ctx, spec["adapters"], image, meta)) if metrics: - existing = meta.get("metrics") - if isinstance(existing, dict): - existing.update(metrics) - else: - meta["metrics"] = metrics + meta["metrics"] = metrics # replaces: a derived image never inherits its source's metrics except Exception as e: _warn_once(ctx, "skip", f"metrics skipped: {e}") +def spec_from_item(item): + """Builtin-only spec from a source item's stored metrics (None when it has none). + + Used by the dashboard upscale thread, which has no session settings and must + not read generation-run state: an upscale of a measured image is measured the + same way, an upscale of an unmeasured image is not. + """ + m = item.get("metrics") if isinstance(item, dict) else None + if not isinstance(m, dict): + return None + keys = [k for k in BUILTIN_METRICS if k in m] + return {"builtins": keys, "adapters": []} if keys else None + + +def finalize(meta, image, sink=None, spec=_CURRENT): + """Metrics for a newly saved image. `sink` = builtins already computed at decode time. + + Always drops any metrics copied from a source item first (never inherit). + spec omitted -> the active run's context; explicit spec/None -> exactly that + (never touches generation-run state). Never raises. + """ + meta.pop("metrics", None) + explicit = spec is not _CURRENT + if sink: + meta["metrics"] = dict(sink) + if not explicit: + attach_adapters(meta, image) + elif image is not None: + attach(meta, image, spec=spec) + + def attach_adapters(meta, image): """Run only the active run's adapters (builtins already came from the decode sink).""" ctx = _active_ctx() diff --git a/remote_vae.py b/remote_vae.py index 3cf7204..2bd8210 100644 --- a/remote_vae.py +++ b/remote_vae.py @@ -6,6 +6,8 @@ import threading import time import numpy as np from PIL import Image + +from . import image_metrics try: from diffusers.image_processor import VaeImageProcessor except ImportError as e: @@ -152,6 +154,8 @@ class RemoteVAEDecodeWorker: if fmt not in ("webp", "png", "jpg", "jpeg"): fmt = "webp" self.image_format = fmt + # Captured now so a late worker thread is unaffected by end_run() ordering + self._metrics_spec = image_metrics.current_spec() self.queue = queue.Queue() self.thread = threading.Thread(target=self._worker, daemon=True) self.thread.start() @@ -224,6 +228,9 @@ class RemoteVAEDecodeWorker: "rejected": False }) + # Optional image metrics from the postprocessed [B,C,H,W] 0-1 tensor + image_metrics.attach(meta, result[:1], spec=self._metrics_spec) + # Update manifest import json self.existing_data["items"].insert(0, meta) diff --git a/sampler_node.py b/sampler_node.py index 1f1bdcd..42141b0 100644 --- a/sampler_node.py +++ b/sampler_node.py @@ -44,6 +44,7 @@ except ImportError: PromptServer = None from .html_generator import get_html_template +from . import image_metrics class SamplerGridTester: @@ -327,17 +328,26 @@ class SamplerGridTester: except (TypeError, ValueError): pass - return run_generation_loop( - self, - ckpt_name, positive_text, negative_text, seed, denoise, vae_batch_size, - overwrite_existing, flush_batch_every, configs_json, resolutions_json, - session_name, unique_id, add_random_seeds_to_gens, lora_triggerwords_mode, - remote_vae_endpoint, save_conditioning_cache_to_file, enable_model_cache, - optional_model, optional_clip, optional_vae, - optional_positive, optional_negative, optional_latent, - distribution_config=dist_config, # Extracted from configs_json - session_settings=session_settings # Extracted from configs_json - ) + # Optional per-image metrics (Builder "Compute image metrics"). The run context is + # token based so end_run() only ever closes the context THIS run opened. + metrics_spec = image_metrics.normalize_spec(session_settings.get("metrics") if session_settings else None) + if metrics_spec: + print(f"[GridTester] ⚙️ metrics: {metrics_spec['builtins'] + metrics_spec['adapters']}") + metrics_token = image_metrics.begin_run(metrics_spec) + try: + return run_generation_loop( + self, + ckpt_name, positive_text, negative_text, seed, denoise, vae_batch_size, + overwrite_existing, flush_batch_every, configs_json, resolutions_json, + session_name, unique_id, add_random_seeds_to_gens, lora_triggerwords_mode, + remote_vae_endpoint, save_conditioning_cache_to_file, enable_model_cache, + optional_model, optional_clip, optional_vae, + optional_positive, optional_negative, optional_latent, + distribution_config=dist_config, # Extracted from configs_json + session_settings=session_settings # Extracted from configs_json + ) + finally: + image_metrics.end_run(metrics_token) # Node class mappings for ComfyUI diff --git a/tests/test_image_metrics.py b/tests/test_image_metrics.py index 7f07d24..cda9c14 100644 --- a/tests/test_image_metrics.py +++ b/tests/test_image_metrics.py @@ -9,9 +9,7 @@ import pytest _ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) _PKG = os.path.basename(_ROOT) -_spec = importlib.util.spec_from_file_location(f"{_PKG}.image_metrics_real", os.path.join(_ROOT, "image_metrics.py")) -im = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(im) +im = importlib.import_module(f"{_PKG}.image_metrics") # package module: same instance the hooks use ALL = im.BUILTIN_METRICS diff --git a/tests/test_metrics_hooks.py b/tests/test_metrics_hooks.py new file mode 100644 index 0000000..802bd8b --- /dev/null +++ b/tests/test_metrics_hooks.py @@ -0,0 +1,302 @@ +"""Run context, attach/finalize, never-inherit, and optional scorer adapters (fake nodes).""" +import importlib +import os +import sys +import threading + +import numpy as np +import pytest +from PIL import Image + +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +_PKG = os.path.basename(_ROOT) +im = importlib.import_module(f"{_PKG}.image_metrics") + +SPEC = {"builtins": ["entropy", "luma"], "adapters": []} +OTHER = {"builtins": ["contrast"], "adapters": []} + + +def _pil(v=128, size=(32, 24)): + return Image.fromarray(np.full((size[1], size[0], 3), v, np.uint8)) + + +@pytest.fixture(autouse=True) +def _clean_stack(): + im._stack.clear() + yield + im._stack.clear() + + +# ---- run context ----------------------------------------------------------- +def test_no_context_means_no_metrics(): + meta = {} + im.attach(meta, _pil()) + assert "metrics" not in meta and im.current_spec() is None + + +def test_attach_pil_keys_match_spec(): + tok = im.begin_run({"enabled": True, "builtins": ["luma", "entropy"]}) + meta = {} + im.attach(meta, _pil(128)) + im.end_run(tok) + assert set(meta["metrics"]) == {"luma", "entropy"} + assert meta["metrics"]["luma"] == pytest.approx(128 / 255, abs=1e-3) + + +def test_end_run_does_not_clobber_another_runs_context(): + a = im.begin_run(SPEC) + b = im.begin_run(OTHER) + assert im.current_spec() == OTHER + im.end_run(a) # A finishing late must not disturb B + assert im.current_spec() == OTHER + im.end_run(b) + assert im.current_spec() is None + im.end_run(a) # double end is a no-op + + +def test_inner_run_ending_restores_outer(): + a = im.begin_run(SPEC) + b = im.begin_run(OTHER) + im.end_run(b) + assert im.current_spec() == SPEC + im.end_run(a) + + +def test_other_thread_run_does_not_disturb_main_run(): + main = im.begin_run(SPEC) + seen = {} + + def other(): + t = im.begin_run({"builtins": ["saturation"], "adapters": []}) + im.end_run(t) + seen["done"] = True + + th = threading.Thread(target=other) + th.start() + th.join() + assert seen["done"] and im.current_spec() == SPEC + im.end_run(main) + + +def test_explicit_spec_ignores_generation_run_state(): + """The dashboard upscale thread passes spec explicitly: a concurrent run's spec must not apply.""" + tok = im.begin_run(SPEC) + meta = {} + im.attach(meta, _pil(), spec=None) + assert "metrics" not in meta # explicit None = off, even though a run is active + im.attach(meta, _pil(), spec=OTHER) + assert set(meta["metrics"]) == {"contrast"} # exactly the explicit spec, not SPEC + m2 = {} + im.finalize(m2, _pil(), spec=None) + assert "metrics" not in m2 + sink = {} + im.sink_update(sink, object(), spec=None) # explicit None: never consults the run + assert sink == {} + im.end_run(tok) + + +def test_end_run_prints_timing_line(capsys): + tok = im.begin_run(SPEC) + im.attach({}, _pil()) + im.attach({}, _pil()) + im.end_run(tok) + out = capsys.readouterr().out + assert "metrics: 2 images, mean" in out and "max" in out + + +def test_attach_swallows_errors_and_warns_once(monkeypatch, capsys): + tok = im.begin_run(SPEC) + + def boom(*a, **k): + raise ValueError("boom") + monkeypatch.setattr(im, "compute_for_pil", boom) + m1, m2 = {"x": 1}, {} + im.attach(m1, _pil()) + im.attach(m2, _pil()) + im.end_run(tok) + out = capsys.readouterr().out + assert m1 == {"x": 1} and m2 == {} + assert out.count("metrics skipped: boom") == 1 + + +def test_sink_update_without_real_torch_never_raises(): + tok = im.begin_run(SPEC) + sink = {} + im.sink_update(sink, object()) + im.end_run(tok) + assert sink == {} + + +def test_finalize_never_inherits_source_metrics(): + meta = {"metrics": {"luma": 0.9, "aesthetic": 9.0}} + im.finalize(meta, _pil()) # no run -> nothing computed, stale value dropped + assert "metrics" not in meta + tok = im.begin_run(SPEC) + meta = {"metrics": {"luma": 0.9, "aesthetic": 9.0}} + im.finalize(meta, _pil(0)) + im.end_run(tok) + assert set(meta["metrics"]) == {"luma", "entropy"} and meta["metrics"]["luma"] == 0 + m3 = {"metrics": {"stale": 1}} + im.finalize(m3, None, {"luma": 0.5}) + assert m3["metrics"] == {"luma": 0.5} + + +def test_spec_from_item(): + assert im.spec_from_item({}) is None + assert im.spec_from_item({"metrics": {}}) is None + assert im.spec_from_item({"metrics": {"luma": 1, "aesthetic": 2, "zzz": 3}}) == {"builtins": ["luma"], "adapters": []} + + +def test_upscale_builders_do_not_copy_metrics(): + from florence2_hires import build_florence2_manifest_entry + item = {"id": 1, "model": "m", "metrics": {"luma": 0.4}} + sr = {"image_width": 8, "image_height": 8, "duration": 1.0} + e = build_florence2_manifest_entry(sr, item, session_name="s", pipeline_name="p", upscale_id=2, + upscaled_filename="x.webp", current_index=0, hires_denoise=0.4) + assert "metrics" not in e + src = open(os.path.join(_ROOT, "upscale_runner.py"), encoding="utf-8").read() + assert src.count('"upscale_model", "metrics")') == 2 # SeedVR2 + HiRes/model copies + + +# ---- adapters --------------------------------------------------------------- +class _Loader: + calls = 0 + FUNCTION = "load" + + def load(self): + _Loader.calls += 1 + return ("MODEL",) + + +class _Scorer: + FUNCTION = "score" + seen = [] + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"image": ("IMAGE",), "model": ("MODEL",), "mode": (["fast", "slow"],), + "scale": ("FLOAT", {"default": 2.0})}} + + def score(self, image, model, mode, scale): + _Scorer.seen.append((model, mode, scale)) + return (np.float32(6.12345),) + + +class _V3Out: + def __init__(self, *a): + self.args = a + + +def _fake_ad(**over): + ad = {"enabled": True, "node_id": "FakeScore", "image_arg": "image", + "loaders": [{"node_id": "FakeLoad", "kwargs": {}, "pass_as": "model", "output_index": 0}], + "output_index": 0, "pack_hint": "fake-pack"} + ad.update(over) + return ad + + +@pytest.fixture +def fake_nodes(monkeypatch): + nodes_mod = sys.modules["nodes"] + mapping = {"FakeLoad": _Loader, "FakeScore": _Scorer} + monkeypatch.setattr(nodes_mod, "NODE_CLASS_MAPPINGS", mapping, raising=False) + monkeypatch.setattr(im, "ADAPTERS", {"fake": _fake_ad()}) + _Loader.calls = 0 + _Scorer.seen = [] + return mapping + + +def test_adapter_runs_loader_once_per_run_and_floats_result(fake_nodes): + spec = im.normalize_spec({"enabled": True, "builtins": ["luma"], "adapters": ["fake"]}) + assert spec["adapters"] == ["fake"] + tok = im.begin_run(spec) + m1, m2 = {}, {} + im.attach(m1, _pil()) + im.attach(m2, _pil()) + im.end_run(tok) + assert m1["metrics"]["fake"] == pytest.approx(6.1234, abs=2e-4) + assert "luma" in m1["metrics"] and m2["metrics"]["fake"] == m1["metrics"]["fake"] + assert _Loader.calls == 1 + assert _Scorer.seen[0] == ("MODEL", "fast", 2.0) # combo + widget defaults filled + tok = im.begin_run(spec) + im.attach({}, _pil()) + im.end_run(tok) + assert _Loader.calls == 2 # cache dropped at end_run + + +def test_attach_adapters_after_sink(fake_nodes): + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + meta = {"metrics": {"luma": 0.1}} + im.attach_adapters(meta, _pil()) + im.end_run(tok) + assert set(meta["metrics"]) == {"luma", "fake"} + + +def test_missing_node_reported_and_not_called(fake_nodes, capsys): + del fake_nodes["FakeScore"] + st = {a["name"]: a for a in im.adapter_status()} + assert st["fake"]["installed"] is False and st["fake"]["pack_hint"] == "fake-pack" + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + meta = {} + im.attach(meta, _pil()) + im.end_run(tok) + assert "metrics" not in meta and _Loader.calls == 0 + assert "metric adapter 'fake' disabled" in capsys.readouterr().out + + +def test_uncovered_required_input_disables_with_warning(fake_nodes, capsys): + class Strict(_Scorer): + @classmethod + def INPUT_TYPES(cls): + return {"required": {"image": ("IMAGE",), "model": ("MODEL",), "mystery": ("LATENT",)}} + fake_nodes["FakeScore"] = Strict + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + meta = {} + im.attach(meta, _pil()) + im.attach(meta, _pil()) + im.end_run(tok) + out = capsys.readouterr().out + assert "metrics" not in meta and out.count("disabled") == 1 and "mystery" in out + assert _Loader.calls == 0 + + +def test_raising_adapter_disabled_for_rest_of_run(fake_nodes, capsys): + calls = [] + + class Boom(_Scorer): + def score(self, image, model, mode, scale): + calls.append(1) + raise RuntimeError("kaput") + fake_nodes["FakeScore"] = Boom + tok = im.begin_run({"builtins": ["luma"], "adapters": ["fake"]}) + m1, m2 = {}, {} + im.attach(m1, _pil()) + im.attach(m2, _pil()) + im.end_run(tok) + out = capsys.readouterr().out + assert len(calls) == 1 and out.count("kaput") == 1 + assert set(m1["metrics"]) == {"luma"} and set(m2["metrics"]) == {"luma"} # builtins unaffected + + +def test_v3_style_node_and_prompt_arg(monkeypatch): + seen = {} + + class V3: + @classmethod + def execute(cls, image, prompt, model): + seen["prompt"] = prompt + return _V3Out(7.5) + monkeypatch.setattr(sys.modules["nodes"], "NODE_CLASS_MAPPINGS", {"FakeLoad": _Loader, "FakeScore": V3}, raising=False) + monkeypatch.setattr(im, "ADAPTERS", {"fake": _fake_ad(text_arg="prompt")}) + # a non-V3-subclass with execute() and no FUNCTION is dispatched through FUNCTION default "execute" + V3.FUNCTION = "execute" + tok = im.begin_run({"builtins": [], "adapters": ["fake"]}) + meta = {"positive": "a red cat"} + im.attach(meta, _pil()) + im.end_run(tok) + 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"] == [] diff --git a/upscale_runner.py b/upscale_runner.py index 744e38f..cfc94af 100644 --- a/upscale_runner.py +++ b/upscale_runner.py @@ -19,6 +19,7 @@ from .model_loader import load_checkpoint, load_vae_by_name from .batch_encoding import encode_prompt_with_combinators, encode_pair_cached from .image_generation import upscale_image, decode_latent_with_vae, create_image_metadata from .manifest_utils import save_manifest +from . import image_metrics import comfy.model_management as mm @@ -203,6 +204,9 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, pass for img_idx, item in enumerate(target_items): + # Never inherit the source's metrics; recompute them only when the source had + # them. This thread must not read generation-run state, so the spec is explicit. + item_metrics_spec = image_metrics.spec_from_item(item) if job["cancel"].is_set(): print(f"[DashboardUpscale] 🛑 Cancelled — clearing VRAM...") try: @@ -332,7 +336,7 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, k: v for k, v in item.items() if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration", "upscale_source", "upscale_pipeline", "upscale_mode", - "upscale_ratio", "upscale_denoise", "upscale_model") + "upscale_ratio", "upscale_denoise", "upscale_model", "metrics") } upscaled_meta.update({ "id": upscale_id, @@ -347,6 +351,7 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, "upscale_mode": "seedvr2", "upscale_model": sv_config.get("dit_model", ""), }) + image_metrics.finalize(upscaled_meta, result_pil, spec=item_metrics_spec) manifest_data["items"].insert(0, upscaled_meta) upscale_combo_idx += 1 else: @@ -444,6 +449,7 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, current_index=len(manifest_data["items"]), hires_denoise=f2_config["hires_denoise"], ) + image_metrics.finalize(entry, result_pil, spec=item_metrics_spec) manifest_data["items"].insert(0, entry) upscale_combo_idx += 1 else: @@ -494,8 +500,13 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, is_last_combo = combo == combos[-1] is_final_output = is_last_step and is_last_combo + up_metrics_sink = {} + up_metrics_img = None if isinstance(upscale_result, dict) and "samples" in upscale_result: - upscaled_pil = decode_latent_with_vae(loaded_vae, upscale_result["samples"]) + up_metrics_img = upscaled_pil = decode_latent_with_vae( + loaded_vae, upscale_result["samples"], + metrics_sink=up_metrics_sink if (is_final_output and item_metrics_spec) else None, + metrics_spec=item_metrics_spec) up_w, up_h = upscaled_pil.size if is_final_output: upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000) @@ -508,6 +519,7 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, else: upscaled_pil.save(_up_path, format="WEBP", quality=95) elif isinstance(upscale_result, PILImage.Image): + up_metrics_img = upscale_result up_w, up_h = upscale_result.size if is_final_output: upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000) @@ -534,7 +546,7 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, k: v for k, v in item.items() if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration", "upscale_source", "upscale_pipeline", "upscale_mode", - "upscale_ratio", "upscale_denoise", "upscale_model") + "upscale_ratio", "upscale_denoise", "upscale_model", "metrics") } # gen_index: place upscaled images at the end (newest) since they were just created # Use current manifest length so each new upscale gets the next sequential position @@ -561,6 +573,7 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, upscaled_meta["hires_prompt_behavior"] = hires_prompt_behavior_rt upscaled_meta["hires_prompt_text"] = hires_prompt_text_rt + image_metrics.finalize(upscaled_meta, up_metrics_img, up_metrics_sink, spec=item_metrics_spec) manifest_data["items"].insert(0, upscaled_meta) upscale_combo_idx += 1