diff --git a/ProjectStructure.md b/ProjectStructure.md index 06b7a2d..af97d86 100644 --- a/ProjectStructure.md +++ b/ProjectStructure.md @@ -38,6 +38,8 @@ ComfyUI-Ultimate-Sampler-Grid/ │ ├── generation_orchestrator.py # The "Conductor" (Loop Logic, 2281 lines) ├── image_generation.py # KSampler & VAE Wrappers (748 lines) +├── image_metrics.py # Optional per-image metrics (sharpness/entropy/...), run context, scorer-pack adapters +├── metrics_runner.py # Background job: compute metrics for an existing session (CPU, manifest lock) ├── model_loader.py # Checkpoint/LoRA/VAE/GGUF Loading Logic (805 lines) ├── model_cache.py # 3-Tier Model Caching System (900 lines) ├── remote_vae.py # Remote VAE Offloading (340 lines) @@ -195,6 +197,10 @@ All routes registered via `PromptServer.instance.routes`. | POST | `/configbuilder/lookup_lora_metadata` | `config_builder_node.py` | 768 | Full CivitAI metadata for LoRA `{lora_name}` | | POST | `/configbuilder/lookup_model_metadata` | `config_builder_node.py` | 916 | Full CivitAI metadata for model `{model_name, model_type}` | | GET | `/configbuilder/model_lists` | `config_builder_node.py` | 1029 | All model/sampler/scheduler/VAE/upscale lists | +| POST | `/config_tester/compute_metrics` | `__init__.py` | - | Start the background metrics job for a session (body: `{session_name, overwrite?}`) | +| GET | `/config_tester/metrics_status` | `__init__.py` | - | Job progress (`?job_id=`) | +| POST | `/config_tester/cancel_metrics` | `__init__.py` | - | Cancel a metrics job (body: `{job_id}`) | +| GET | `/configbuilder/metrics_info` | `config_builder_node.py` | - | Built-in metric names + optional scorer-pack install status (Run Settings hint) | | POST | `/configbuilder/refresh_models` | `config_builder_node.py` | 1093 | Signal frontend cache clear | | GET | `/configbuilder/get_lora_triggers` | `config_builder_node.py` | 1115 | Get triggers for single LoRA | | POST | `/configbuilder/save_lora_triggers` | `config_builder_node.py` | 1138 | Save edited triggers for LoRA | @@ -835,6 +841,8 @@ __init__.py 28. **Path Security** — All filesystem endpoints use `_is_path_within()` (line 27 in `__init__.py`) to validate paths stay within `benchmarks/` base directory. Session names sanitized with `re.sub(r'[^\w\-]', '', name)`. +29. **Image Metrics Run Context** - `image_metrics.begin_run(spec)` (wrapped around `run_generation_loop` in `sampler_node.py`) returns a token and `end_run(token)` closes only that context, so a second run or the dashboard upscale thread can never clobber it. Generation-side hooks read the active context (`decode_latent_with_vae(..., metrics_sink=)`, `flush_batch_with_vae`, orchestrator upscale inserts via `image_metrics.finalize`); the dashboard upscale thread (`upscale_runner.py`) must NOT read generation-run state, so it passes an explicit spec derived from the source item (`spec_from_item`). Every place that copies a source item into an upscaled entry excludes `"metrics"` (never inherit); the new entry is re-measured. `metrics` is session-level (`_session_settings.metrics`), not a per-config key. + --- ### **Testing Changes** diff --git a/README.md b/README.md index 5a2d052..a7fb3c4 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,7 @@ If you'd rather build the graph yourself: * **Seed Behavior Settings:** Per-config seed control — `fixed` uses the node seed, `randomize` generates a new seed per image, `full_run_seed` overrides the node seed entirely. * **Configurable Label Mode:** Toggle overlay labels on dashboard cards showing model, sampler, scheduler, steps, CFG, and other parameters directly on each image. * **GPU Cooldown:** Automatic pause between generations to prevent thermal throttling. Configurable interval, duration, and optional VRAM clearing. +* **Image Metrics (optional):** Run Settings checkbox that stores sharpness, entropy, colorfulness, contrast, luma and saturation for every saved image in the manifest (see [Image Metrics](#image-metrics-optional)). ### 🌐 Distributed Processing (Multi-Machine) * **Multi-Worker Distribution:** Split large grid runs across multiple ComfyUI instances. A master node coordinates job distribution while workers process images in parallel and upload results back. @@ -802,6 +803,29 @@ Both live in the **Run Settings** section of the Config Builder and default to t --- +### Image Metrics (Optional) + +Run Settings has a **Compute Image Metrics** checkbox (default OFF). When it is on, every saved still image (base generations and upscales; videos are skipped) gets a `metrics` object in `manifest.json`, computed from the decoded image at save time (about 1-5 ms per image on the GPU, nothing when the box is unchecked): + +| Metric | What it is | Range | +|---|---|---| +| `sharpness` | Variance of the 3x3 Laplacian of luma on the 0-255 scale (OpenCV-style "variance of Laplacian") | 0 = flat, large = lots of edges | +| `entropy` | Shannon entropy of the 256-bin luma histogram | 0-8 bits | +| `colorfulness` | Hasler and Suesstrunk (2003) colorfulness | ~0 grey, >100 very colorful | +| `contrast` | Standard deviation of luma (RMS contrast) | 0-0.5 | +| `luma` | Mean brightness (Rec.601 luma) | 0-1 | +| `saturation` | Mean HSV saturation | 0-1 | + +Notes: +* Metrics are measured on an "analysis image": the long side is block-averaged down to at most 1024 px (by the smallest integer factor), so a 2x upscale stays comparable with its 1024 base image and the cost is capped. The constant is `ANALYSIS_MAX_SIDE` in `image_metrics.py`; changing it would make old sessions incomparable. +* They measure image **signal, not taste**: noise raises sharpness and entropy, and a dark moody image is not a worse one. Use them to rank and filter, then look. +* Upscales are measured again (a derived image never inherits its source's numbers). Dashboard-triggered upscales measure only when the source image had metrics. +* Items produced by remote Distributed workers inside the companion plugin get no metrics yet; master-side saves and upscales do. +* Optional scorer packs (aesthetic / reward models) plug in through `nodes.NODE_CLASS_MAPPINGS` without any import, and their scores are stored next to the built-ins. The shipped entries stay disabled until their node ids are verified against the real packs, so only the built-in metrics are available today. The Builder lists each pack with an installed / not installed hint (text only, nothing is downloaded). +* **In the Dashboard:** once a session has metrics, the Sort popup gains one option per metric (descending, items without the metric sort last, in both the primary and secondary lists), the Logic Filter field list gains a **Metrics** group, and a **Metric range** row adds `>= min` / `<= max` filters in one click (placeholders show the observed range). The label panel has a **Metrics** checkbox that prints `Shp:812 Ent:7.41 Col:41.2 ...` on each card. Logic filters skip items that lack the field, so images without metrics (older items, videos) stay visible; Quick Filter is categorical and does not list metrics, use the Logic Filter for ranges. Sessions without metrics show none of this. +* **Older sessions:** the cog panel's **IMAGE METRICS** section has **Compute metrics for this session**, a background job (CPU only, no VRAM, about 50 ms per image) that measures every image without metrics, skips videos, shows progress and can be cancelled by clicking the button again. It rewrites only the `metrics` field under the manifest lock, so favorites and notes you change meanwhile survive. Do not run it on a session that is still generating. +* The setting is session-level (like GPU Cooldown), so it is NOT a per-config field: it lives in `_session_settings.metrics`, not in `_EXPECTED_CONFIG_KEYS` or `expand_configs`. + ## 🖥️ Dashboard Interface ### Header Bar diff --git a/__init__.py b/__init__.py index f9dc52c..90dbd45 100644 --- a/__init__.py +++ b/__init__.py @@ -1417,6 +1417,44 @@ async def cancel_upscale(request): except Exception as e: return web.json_response({"error": str(e)}, status=500) +# --- Image metrics for an EXISTING session (background job, CPU only) --- +@server.PromptServer.instance.routes.post("/config_tester/compute_metrics") +async def compute_metrics(request): + """Start a job that computes image metrics for the session's items that have none.""" + try: + data = await request.json() + session_name = re.sub(r'[^\w\-]', '', str(data.get("session_name") or "")) + if not session_name: + return web.json_response({"error": "Missing session_name"}, status=400) + base_dir = os.path.join(folder_paths.get_output_directory(), "benchmarks", session_name) + if not _is_path_within(base_dir, _get_benchmarks_base()): + return web.json_response({"error": "Forbidden: path outside benchmarks directory"}, status=403) + from .metrics_runner import start_metrics_job, get_metrics_status + job_id, error = start_metrics_job(session_name, overwrite=bool(data.get("overwrite", False))) + if error: + return web.json_response({"error": error}, status=400) + return web.json_response({"job_id": job_id, "total_images": get_metrics_status(job_id)["total"]}) + except Exception as e: + print(f"[ConfigTester] Error starting metrics job: {e}") + return web.json_response({"error": str(e)}, status=500) + +@server.PromptServer.instance.routes.get("/config_tester/metrics_status") +async def metrics_status(request): + job_id = request.query.get("job_id", "") + if not job_id: + return web.json_response({"error": "Missing job_id"}, status=400) + from .metrics_runner import get_metrics_status + return web.json_response(get_metrics_status(job_id)) + +@server.PromptServer.instance.routes.post("/config_tester/cancel_metrics") +async def cancel_metrics(request): + try: + data = await request.json() + from .metrics_runner import cancel_metrics_job + return web.json_response({"cancelled": cancel_metrics_job(data.get("job_id", ""))}) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) + def _create_image_links(source_dir, link_dir, items): """ diff --git a/config_builder_node.py b/config_builder_node.py index 2854a67..8e0f8dc 100644 --- a/config_builder_node.py +++ b/config_builder_node.py @@ -15,6 +15,7 @@ import server from aiohttp import web import hashlib from .civitai import civitai_fetch_by_hash +from . import image_metrics def safe_print(*args, **kwargs): @@ -905,6 +906,18 @@ class UltimateConfigBuilder: cooldown_data = state.get("cooldown", {}) if cooldown_data and cooldown_data.get("enabled", False): session_settings["cooldown"] = cooldown_data + # Optional per-image metrics (session-level, like cooldown: NOT a per-config key, so it is + # not in _EXPECTED_CONFIG_KEYS / expand_configs). Emitted only when enabled and non-empty. + metrics_data = state.get("metrics") + if isinstance(metrics_data, dict) and metrics_data.get("enabled", False): + _b = metrics_data.get("builtins") or {} + _a = metrics_data.get("adapters") or {} + _b_on = set(_b) if isinstance(_b, (list, tuple)) else {k for k, v in _b.items() if v} + _a_on = set(_a) if isinstance(_a, (list, tuple)) else {k for k, v in _a.items() if v} + _builtins = [k for k in image_metrics.BUILTIN_METRICS if k in _b_on] + _adapters = sorted(k for k in _a_on if k in image_metrics.ADAPTERS) + if _builtins or _adapters: + session_settings["metrics"] = {"enabled": True, "builtins": _builtins, "adapters": _adapters} # Start At Job # (skip to a specific job number) — always emit. try: session_settings["start_at_job"] = int(state.get("start_at_job", 0)) @@ -1342,6 +1355,18 @@ async def get_model_counts_endpoint(request): except Exception as e: return web.json_response({"error": str(e)}, status=500) +# Which image metrics exist and which optional scorer packs are installed (Run Settings hint). +# Separate from /configbuilder/model_lists so the cached model-list shape is unchanged. +@server.PromptServer.instance.routes.get("/configbuilder/metrics_info") +async def get_metrics_info(request): + try: + return web.json_response({ + "builtins": list(image_metrics.BUILTIN_METRICS), + "adapters": image_metrics.adapter_status(), + }) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) + # API endpoint to get all model lists for unified model selector @server.PromptServer.instance.routes.get("/configbuilder/model_lists") async def get_model_lists_endpoint(request): 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 67d30b5..aed4226 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 @@ -588,8 +589,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) @@ -599,6 +605,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) @@ -640,6 +647,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 @@ -2163,6 +2171,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. @@ -2221,6 +2230,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: @@ -2311,10 +2321,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: @@ -2328,6 +2343,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: @@ -2424,6 +2440,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 d8e8a50..8360bf0 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 new file mode 100644 index 0000000..aaa8119 --- /dev/null +++ b/image_metrics.py @@ -0,0 +1,668 @@ +""" +Optional per-image metrics (sharpness, entropy, ...) stored on manifest items. + +Everything here is pure functions plus a small run-scoped context. No new +dependencies: numpy is always present, torch is imported lazily (and only when +it is the real thing -- the unit tests run against a stub). + +Definitions (RGB in [0,1], gamma-encoded as saved). The "analysis image" is the +input with 3 channels, clamped to 0-1, and block-mean downscaled by the smallest +integer factor that brings the long side to <= ANALYSIS_MAX_SIDE, so a 2x +upscale stays comparable with its base and the cost is capped. + + sharpness variance of the 3x3 Laplacian of luma on the 0-255 scale (valid conv) + entropy Shannon entropy (bits, 0..8) of the 256-bin luma histogram + colorfulness Hasler & Suesstrunk 2003 on 0-255 channels + contrast std of luma (RMS contrast, 0..0.5) + luma mean of Y = 0.299 R + 0.587 G + 0.114 B + saturation mean HSV saturation + +These measure signal, not taste: noise raises sharpness and entropy. +ANALYSIS_MAX_SIDE is a versioned constant; changing it breaks comparisons +with metrics stored by older sessions. +""" + +import math +import threading +import time + +import numpy as np + +BUILTIN_METRICS = ("sharpness", "entropy", "colorfulness", "contrast", "luma", "saturation") +ANALYSIS_MAX_SIDE = 1024 + +_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(). +# --------------------------------------------------------------------------- +ADAPTERS = { + "aesthetic": { + "enabled": False, + "label": "Aesthetic score", + "node_id": "CalculateAestheticScore", + "image_arg": "image", + "loaders": [{"node_id": "LoadAesteticModel", "kwargs": {}, "pass_as": "aesthetic_model", "output_index": 0}], + "output_index": 0, + "pack_hint": "ComfyUI-Strimmlarns-Aesthetic-Score", + "note": "node ids not yet verified against the pack", + }, + "imagereward": { + "enabled": False, + "label": "ImageReward", + "node_id": "ImageRewardScore", + "image_arg": "images", + "text_arg": "prompt", + "loaders": [{"node_id": "ImageRewardLoader", "kwargs": {}, "pass_as": "model", "output_index": 0}], + "output_index": 0, + "pack_hint": "ComfyUI-ImageReward", + "note": "node ids not yet verified against the pack", + }, +} + + +# --------------------------------------------------------------------------- +# Spec +# --------------------------------------------------------------------------- +def normalize_spec(raw): + """Session `metrics` setting -> {"builtins": [...], "adapters": [...]} or None. + + None means "feature off": missing, not enabled, nothing selected, or garbage. + Unknown builtin names are dropped; adapters must exist AND be enabled in + the registry. + """ + if not isinstance(raw, dict): + return None + if "enabled" in raw and not raw.get("enabled"): + return None + b = raw.get("builtins") or [] + a = raw.get("adapters") or [] + if isinstance(b, dict): + b = [k for k, v in b.items() if v] + if isinstance(a, dict): + a = [k for k, v in a.items() if v] + if not isinstance(b, (list, tuple)) or not isinstance(a, (list, tuple)): + return None + builtins = [k for k in BUILTIN_METRICS if k in b] + adapters = [k for k in a if isinstance(k, str) and k in ADAPTERS and ADAPTERS[k].get("enabled")] + adapters = sorted(set(adapters)) + if not builtins and not adapters: + return None + return {"builtins": builtins, "adapters": adapters} + + +# --------------------------------------------------------------------------- +# numpy reference implementation (also the CPU fallback) +# --------------------------------------------------------------------------- +def _block_factor(h, w): + side = max(h, w) + return 1 if side <= ANALYSIS_MAX_SIDE else int(math.ceil(side / ANALYSIS_MAX_SIDE)) + + +def _round4(v): + return round(float(v), 4) + + +def _analysis_np(arr): + """H,W,C (float 0-1 or uint8) -> float32 [H,W,3] clamped, block-mean downscaled.""" + a = np.asarray(arr) + while a.ndim > 3: + a = a[0] + if a.ndim == 2: + a = a[..., None] + if a.shape[-1] not in (1, 2, 3, 4) and a.shape[0] in (1, 3, 4): # CHW + a = np.transpose(a, (1, 2, 0)) + if a.dtype == np.uint8: + a = a.astype(np.float32) / 255.0 + else: + a = a.astype(np.float32) + if a.shape[-1] == 1: + a = np.repeat(a, 3, axis=-1) + a = a[..., :3] + a = np.clip(a, 0.0, 1.0) # NaN stays NaN + h, w = a.shape[:2] + f = _block_factor(h, w) + if f > 1: + hh, ww = (h // f) * f, (w // f) * f + a = a[:hh, :ww].reshape(hh // f, f, ww // f, f, 3).mean(axis=(1, 3), dtype=np.float32) + return a + + +def _var64(x): + """Population variance with float64 accumulation, float32 inputs (no float64 copy).""" + m = x.mean(dtype=np.float64) + d = x - np.float32(m) + return float((d * d).mean(dtype=np.float64)) + + +def compute_metrics_numpy(arr, keys=BUILTIN_METRICS): + """Reference implementation. arr: H,W,C float 0-1 or uint8. Returns {key: float}.""" + a = _analysis_np(arr) + if not np.isfinite(a).all(): + return {} + keys = [k for k in BUILTIN_METRICS if k in keys] + r, g, b = a[..., 0], a[..., 1], a[..., 2] + y = _LUMA_W[0] * r + _LUMA_W[1] * g + _LUMA_W[2] * b + out = {} + if "sharpness" in keys and y.shape[0] >= 3 and y.shape[1] >= 3: + y255 = y * np.float32(255.0) + lap = (y255[:-2, 1:-1] + y255[2:, 1:-1] + y255[1:-1, :-2] + y255[1:-1, 2:] - np.float32(4.0) * y255[1:-1, 1:-1]) + out["sharpness"] = float(_var64(lap)) + if "entropy" in keys: + q = np.clip(np.round(y * 255.0), 0, 255).astype(np.int64) + p = np.bincount(q.ravel(), minlength=256).astype(np.float64) / q.size + p = p[p > 0] + out["entropy"] = float(-(p * np.log2(p)).sum()) + if "colorfulness" in keys: + rg = (r - g) * np.float32(255.0) + yb = (np.float32(0.5) * (r + g) - b) * np.float32(255.0) + out["colorfulness"] = float(math.sqrt(_var64(rg) + _var64(yb)) + + 0.3 * math.sqrt(float(rg.mean(dtype=np.float64)) ** 2 + float(yb.mean(dtype=np.float64)) ** 2)) + if "contrast" in keys: + out["contrast"] = math.sqrt(_var64(y)) + if "luma" in keys: + out["luma"] = float(y.mean(dtype=np.float64)) + if "saturation" in keys: + mx = a.max(axis=-1) + mn = a.min(axis=-1) + s = np.where(mx >= 1e-6, (mx - mn) / np.maximum(mx, np.float32(1e-6)), np.float32(0.0)) + out["saturation"] = float(s.mean(dtype=np.float64)) + return {k: _round4(v) for k, v in out.items() if math.isfinite(v)} + + +# --------------------------------------------------------------------------- +# torch implementation (GPU, one sync per batch) +# --------------------------------------------------------------------------- +def _real_torch(): + try: + import torch + except Exception: + return None + return torch if hasattr(torch, "from_numpy") and hasattr(torch, "conv2d") else None + + +def _is_tensor(x): + t = _real_torch() + return t is not None and t.is_tensor(x) + + +def _to_bchw(images): + """Any [B,H,W,C] / [B,C,H,W] / [H,W,C] / 5-D tensor -> float32 [B,3,H,W] clamped 0-1.""" + torch = _real_torch() + x = images.detach().float() + if x.dim() == 2: + x = x[..., None] + if x.dim() == 3: + x = x[None] + if x.dim() > 4: + x = x.reshape(-1, *x.shape[-3:]) + if x.shape[-1] in (1, 3, 4): # BHWC (ComfyUI IMAGE convention) + x = x.permute(0, 3, 1, 2) + elif x.shape[1] not in (1, 3, 4): + raise ValueError(f"unsupported image tensor shape {tuple(images.shape)}") + if x.shape[1] == 1: + x = x.expand(-1, 3, -1, -1) + x = x[:, :3].clamp(0.0, 1.0) + h, w = x.shape[-2:] + f = _block_factor(h, w) + if f > 1: + x = torch.nn.functional.avg_pool2d(x, f) # floor mode == crop to a multiple, then block mean + return x.contiguous() + + +def compute_metrics_torch(images, keys=BUILTIN_METRICS): + """Per-image metrics for a batch tensor. Returns list[dict] (one per image). + + Runs under inference_mode; a single .cpu().tolist() syncs the device once. + """ + torch = _real_torch() + if torch is None: + raise RuntimeError("torch is not available") + with torch.inference_mode(): + keys = [k for k in BUILTIN_METRICS if k in keys] + x = _to_bchw(images) + B, _, H, W = x.shape + finite = torch.isfinite(x).reshape(B, -1).all(dim=1) + x = torch.nan_to_num(x, nan=0.0, posinf=1.0, neginf=0.0) + r, g, b = x[:, 0], x[:, 1], x[:, 2] + y = _LUMA_W[0] * r + _LUMA_W[1] * g + _LUMA_W[2] * b # [B,H,W] + cols = [] + names = [] + if "sharpness" in keys and H >= 3 and W >= 3: + y255 = (y * 255.0)[:, None] + kern = torch.tensor([[0.0, 1.0, 0.0], [1.0, -4.0, 1.0], [0.0, 1.0, 0.0]], + device=y.device, dtype=y.dtype)[None, None] + lap = torch.nn.functional.conv2d(y255, kern) + cols.append(lap.reshape(B, -1).var(dim=1, unbiased=False)) + names.append("sharpness") + if "entropy" in keys: + q = torch.clamp(torch.round(y * 255.0), 0, 255).long().reshape(B, -1) + q = q + (torch.arange(B, device=q.device) * 256)[:, None] + counts = torch.bincount(q.reshape(-1), minlength=256 * B).reshape(B, 256).float() + p = counts / float(H * W) + ent = -torch.where(p > 0, p * torch.log2(p.clamp_min(1e-30)), torch.zeros_like(p)).sum(dim=1) + cols.append(ent) + names.append("entropy") + if "colorfulness" in keys: + rg = ((r - g) * 255.0).reshape(B, -1) + yb = ((0.5 * (r + g) - b) * 255.0).reshape(B, -1) + std = torch.sqrt(rg.var(dim=1, unbiased=False) + yb.var(dim=1, unbiased=False)) + mean = torch.sqrt(rg.mean(dim=1) ** 2 + yb.mean(dim=1) ** 2) + cols.append(std + 0.3 * mean) + names.append("colorfulness") + if "contrast" in keys: + cols.append(y.reshape(B, -1).var(dim=1, unbiased=False).sqrt()) + names.append("contrast") + if "luma" in keys: + cols.append(y.reshape(B, -1).mean(dim=1)) + names.append("luma") + if "saturation" in keys: + mx = x.max(dim=1).values + mn = x.min(dim=1).values + s = torch.where(mx >= 1e-6, (mx - mn) / mx.clamp_min(1e-6), torch.zeros_like(mx)) + cols.append(s.reshape(B, -1).mean(dim=1)) + names.append("saturation") + if not cols: + return [{} for _ in range(B)] + table = torch.stack([c.float() for c in cols], dim=1).cpu().tolist() # the one sync + ok = finite.cpu().tolist() + out = [] + for i in range(B): + d = {} + if ok[i]: + for n, v in zip(names, table[i]): + if math.isfinite(v): + d[n] = _round4(v) + out.append(d) + return out + + +def _comfy_device(): + try: + import comfy.model_management as mm + return mm.get_torch_device() + except Exception: + return None + + +def compute_for_pil(pil, keys=BUILTIN_METRICS, use_torch=True): + """PIL image -> metrics. Torch on the comfy device when torch is real, else numpy.""" + arr = np.asarray(pil.convert("RGB")) + torch = _real_torch() if use_torch else None + if torch is not None: + try: + t = torch.from_numpy(np.ascontiguousarray(arr)).float().div_(255.0)[None] + dev = _comfy_device() + if dev is not None: + t = t.to(dev) + return compute_metrics_torch(t, keys)[0] + except Exception: + pass + return compute_metrics_numpy(arr, keys) + + +# --------------------------------------------------------------------------- +# Run-scoped context +# --------------------------------------------------------------------------- +class _RunContext: + def __init__(self, spec): + self.spec = spec + self.warned = set() + self.loader_cache = {} # adapter name -> {"outputs": {...}} (loader outputs live for the run) + self.disabled_adapters = set() + self.checked_adapters = set() + self.times = [] # seconds per image (builtins only) + self.lock = threading.Lock() + self.closed = False + + +_ctx_lock = threading.Lock() +_stack = [] # contexts of runs still open; the newest is the "current" one +_CURRENT = object() # sentinel: "use the active run's spec" + + +def begin_run(spec): + """Start a run context and return its token. spec: normalized dict or raw setting or None. + + The caller passes the token back to end_run(token) so a run can only close + its own context: end_run never clobbers a context another run started + meanwhile. The dashboard upscale thread does not use this context at all + (it passes an explicit spec to attach()). + """ + ctx = _RunContext(normalize_spec(spec)) + with _ctx_lock: + _stack.append(ctx) + return ctx + + +def end_run(token=None): + """Close a run: print the timing line, release adapter loader outputs. + + 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. + """ + with _ctx_lock: + 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 ctx in _stack: + _stack.remove(ctx) + with ctx.lock: + times = list(ctx.times) + ctx.loader_cache.clear() + 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") + + +def current_spec(): + ctx = _active_ctx() + return ctx.spec if ctx is not None else None + + +def _active_ctx(): + with _ctx_lock: + return _stack[-1] if _stack else None + + +def _warn_once(ctx, key, msg): + if ctx is not None: + with ctx.lock: + if key in ctx.warned: + return + ctx.warned.add(key) + print(f"[GridTester] ⚠️ {msg}") + + +def record_time(seconds, ctx=None): + ctx = ctx or _active_ctx() + if ctx is not None: + with ctx.lock: + ctx.times.append(seconds) + + +def sink_update(sink, decoded, spec=_CURRENT): + """Hook for decode_latent_with_vae: compute builtin metrics on the GPU tensor into `sink`. + + spec defaults to the active run's; pass an explicit spec (or None = off) + from threads that must not read generation-run state. Never raises. + """ + ctx = None + if spec is _CURRENT: + ctx = _active_ctx() + spec = ctx.spec if ctx else None + if not spec or not spec.get("builtins"): + return + try: + t0 = time.perf_counter() + sink.update(compute_metrics_torch(decoded, spec["builtins"])[0]) + record_time(time.perf_counter() - t0, ctx) + except Exception as e: + _warn_once(ctx, "skip", f"metrics skipped: {e}") + + +# --------------------------------------------------------------------------- +# attach (never raises) +# --------------------------------------------------------------------------- +def attach(meta, image, spec=_CURRENT, builtins=True, adapters=True): + """Set meta["metrics"] from `image` (torch tensor | PIL | numpy). Never raises. + + spec: omitted -> the active run's spec; explicit dict/None -> exactly that + (None = do nothing). Explicit specs run builtins only (adapters need a run + context for their loader cache). + """ + ctx = None + explicit = spec is not _CURRENT + if not explicit: + ctx = _active_ctx() + spec = ctx.spec if ctx else None + if not spec: + return + try: + metrics = {} + bkeys = spec.get("builtins") or [] + if builtins and bkeys: + t0 = time.perf_counter() + if _is_tensor(image): + metrics.update(compute_metrics_torch(image, bkeys)[0]) + elif isinstance(image, np.ndarray): + metrics.update(compute_metrics_numpy(image, bkeys)) + else: + metrics.update(compute_for_pil(image, bkeys)) + if ctx is not None: + record_time(time.perf_counter() - t0, ctx) + if adapters and ctx is not None and spec.get("adapters"): + metrics.update(_run_adapters(ctx, spec["adapters"], image, meta)) + if 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() + if ctx is None or not ctx.spec or not ctx.spec.get("adapters"): + return + try: + res = _run_adapters(ctx, ctx.spec["adapters"], image, meta) + if res: + meta.setdefault("metrics", {}).update(res) + except Exception as e: + _warn_once(ctx, "skip", f"metrics skipped: {e}") + + +# --------------------------------------------------------------------------- +# Adapters +# --------------------------------------------------------------------------- +def _node_mappings(): + try: + import nodes + return getattr(nodes, "NODE_CLASS_MAPPINGS", None) or {} + except Exception: + return {} + + +def _is_v3(cls): + try: + from comfy_api.latest import io as _io + return isinstance(cls, type) and issubclass(cls, _io.ComfyNode) + except Exception: + return False + + +def _call_node(cls, kwargs): + """Same dispatch as image_generation._apply_model_patch_node.""" + if _is_v3(cls): + result = cls.execute(**kwargs) + else: + result = getattr(cls(), getattr(cls, "FUNCTION", "execute"))(**kwargs) + return result + + +def _output(result, index): + if hasattr(result, "args") and isinstance(getattr(result, "args"), tuple): + return result.args[index] + if isinstance(result, (tuple, list)): + return result[index] + return result + + +def _required_inputs(cls): + """{name: default-or-_MISSING} for required inputs, or None when it cannot be introspected.""" + _MISSING = object() + try: + if _is_v3(cls): + schema = cls.GET_SCHEMA() if hasattr(cls, "GET_SCHEMA") else cls.define_schema() + out = {} + for i in schema.inputs: + if getattr(i, "optional", False): + continue + d = getattr(i, "default", None) + out[i.id] = _MISSING if d is None else d + return out, _MISSING + spec = cls.INPUT_TYPES().get("required", {}) + out = {} + for name, v in spec.items(): + kind = v[0] if isinstance(v, (tuple, list)) and v else v + opts = v[1] if isinstance(v, (tuple, list)) and len(v) > 1 and isinstance(v[1], dict) else {} + if "default" in opts: + out[name] = opts["default"] + elif isinstance(kind, (list, tuple)) and kind: + out[name] = kind[0] # combo: first option + else: + out[name] = _MISSING + return out, _MISSING + except Exception: + return None + + +def _adapter_defs(name): + ad = ADAPTERS.get(name) + return ad if ad and ad.get("enabled") else 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" + 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) + req = _required_inputs(mappings[ad["node_id"]]) + if req is not None: + req, missing = req + for name, default in req.items(): + if name not in covered and default is missing: + return f"required input '{name}' of '{ad['node_id']}' is not covered" + return None + + +def adapter_status(): + """For the Builder hint route: [{name,label,installed,enabled,pack_hint,note}].""" + mappings = _node_mappings() + out = [] + for name, ad in ADAPTERS.items(): + 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", [])]), + "pack_hint": ad.get("pack_hint", ""), + "note": ad.get("note", ""), + }) + return out + + +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) + 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 + arr = arr[None] if arr.ndim == 3 else arr + return torch.from_numpy(np.ascontiguousarray(arr)) if torch is not None else arr + + +def _to_bhwc_full(t): + """Full-resolution (no analysis downscale) BHWC float view of a tensor.""" + x = t.detach().float() + if x.dim() == 3: + x = x[None] + if x.dim() > 4: + x = x.reshape(-1, *x.shape[-3:]) + if x.shape[-1] not in (1, 3, 4) and x.shape[1] in (1, 3, 4): + x = x.permute(0, 2, 3, 1) + return x[..., :3].clamp(0.0, 1.0).contiguous() + + +def _float_of(v): + if hasattr(v, "item") and not isinstance(v, (list, tuple, str)): + try: + return float(v.reshape(-1)[0].item()) if hasattr(v, "reshape") else float(v.item()) + except Exception: + pass + if isinstance(v, (list, tuple)): + return _float_of(v[0]) + return float(v) + + +def _run_adapters(ctx, names, image, meta): + mappings = _node_mappings() + out = {} + img = None + for name in names: + ad = _adapter_defs(name) + if ad is None or name in ctx.disabled_adapters: + continue + try: + if name not in ctx.checked_adapters: + ctx.checked_adapters.add(name) + problem = _adapter_problem(ad, mappings) + if problem: + raise RuntimeError(problem) + if img is None: + img = _adapter_image(image) + 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", {}))) + outputs[li] = _output(res, loader.get("output_index", 0)) + kwargs[loader["pass_as"]] = outputs[li] + kwargs[ad["image_arg"]] = 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"]] + req = _required_inputs(cls) + if req is not None: + req, missing = req + for k, d in req.items(): + if k not in kwargs and d is not missing: + kwargs[k] = d # widget defaults + res = _call_node(cls, kwargs) + out[name] = _round4(_float_of(_output(res, ad.get("output_index", 0)))) + except Exception as e: + ctx.disabled_adapters.add(name) + _warn_once(ctx, f"adapter:{name}", f"metric adapter '{name}' disabled: {e}") + return out diff --git a/metrics_runner.py b/metrics_runner.py new file mode 100644 index 0000000..159965f --- /dev/null +++ b/metrics_runner.py @@ -0,0 +1,166 @@ +""" +Phase 2 of the image-metrics feature: "Compute metrics for this session". + +A background job (same shape as upscale_runner: job dict + daemon thread + +status/cancel helpers) that measures the items of an EXISTING session that have +no `metrics` yet. Pure CPU (PIL -> numpy), no VRAM, never touches the GPU. + +Manifest writes happen under manifest_utils._manifest_lock: the manifest is +RELOADED from disk, the new `metrics` are set by item id, and the file is +replaced atomically, so a stale full copy is never written back and a user's +concurrent favorites/notes survive. Do not run it on a session that a +generation run is currently writing: that run saves its own in-memory copy of +the manifest and would drop metrics written meanwhile. +""" +import json +import os +import re +import threading +import uuid + +import numpy as np +from PIL import Image + +from . import image_metrics +from . import manifest_utils + +_jobs = {} +FLUSH_EVERY = 25 + + +def _image_path_for(item, images_dir): + """Resolve the item's image file INSIDE images_dir (None for videos/sentinels/anything else).""" + if item.get("media_type") == "video" or item.get("florence2_no_detection"): + return None + name = item.get("filename") + if not name: + m = re.search(r'[?&]filename=([^&]+)', str(item.get("file") or "")) + name = m.group(1) if m else None + if not name or name != os.path.basename(name) or "%" in name: + return None + if not name.lower().endswith((".webp", ".png", ".jpg", ".jpeg")): + return None + root = os.path.realpath(images_dir) + path = os.path.realpath(os.path.join(root, name)) + if os.path.dirname(path) != root or not os.path.isfile(path): + return None + return path + + +def start_metrics_job(session_name, overwrite=False, keys=None): + """Start the job. Returns (job_id, error). session_name must already be sanitized by the caller.""" + import folder_paths + benchmarks_dir = os.path.realpath(os.path.join(folder_paths.get_output_directory(), "benchmarks")) + session_dir = os.path.realpath(os.path.join(benchmarks_dir, session_name)) + if os.path.dirname(session_dir) != benchmarks_dir: + return None, "Invalid session" + manifest_path = os.path.join(session_dir, "manifest.json") + images_dir = os.path.join(session_dir, "images") + if not os.path.exists(manifest_path): + return None, "Manifest not found" + + spec = image_metrics.normalize_spec({"enabled": True, "builtins": list(keys or image_metrics.BUILTIN_METRICS)}) + if not spec: + return None, "No valid metric names" + + with manifest_utils._manifest_lock: + with open(manifest_path, "r", encoding="utf-8") as f: + items = json.load(f).get("items", []) + targets = [it["id"] for it in items if "id" in it and (overwrite or not it.get("metrics")) + and _image_path_for(it, images_dir)] + if not targets: + return None, "Nothing to compute (every image already has metrics, or there are no images)" + for jid, j in list(_jobs.items()): + if j["status"] == "running" and j["session"] == session_name: + return None, "A metrics job is already running for this session" + + job_id = f"metrics_{uuid.uuid4().hex[:12]}" + job = {"id": job_id, "session": session_name, "status": "running", "completed": 0, "total": len(targets), + "failed": 0, "error": None, "cancel": threading.Event()} + _jobs[job_id] = job + threading.Thread(target=_run_metrics_thread, + args=(job, targets, spec["builtins"], overwrite, manifest_path, images_dir, session_name), + daemon=True).start() + return job_id, None + + +def get_metrics_status(job_id): + job = _jobs.get(job_id) + if not job: + return {"status": "not_found"} + return {k: job[k] for k in ("status", "completed", "total", "failed", "error")} + + +def cancel_metrics_job(job_id): + job = _jobs.get(job_id) + if job and job["status"] == "running": + job["cancel"].set() + return True + return False + + +def _flush(manifest_path, results, overwrite): + """Merge `results` ({id: metrics}) into the manifest on disk under the lock; atomic replace.""" + if not results: + return + with manifest_utils._manifest_lock: + with open(manifest_path, "r", encoding="utf-8") as f: + data = json.load(f) + for it in data.get("items", []): + m = results.get(it.get("id")) + if m and (overwrite or not it.get("metrics")): + it["metrics"] = m + tmp = manifest_path + ".metrics.tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(data, f, indent=4) + os.replace(tmp, manifest_path) + + +def _send_progress(job, session_name): + try: + from server import PromptServer + PromptServer.instance.send_sync("ultimate_grid.metrics_progress", { + "session_name": session_name, "job_id": job["id"], "status": job["status"], + "completed": job["completed"], "total": job["total"], + }) + except Exception: + pass + + +def _run_metrics_thread(job, target_ids, keys, overwrite, manifest_path, images_dir, session_name): + pending = {} + try: + with manifest_utils._manifest_lock: + with open(manifest_path, "r", encoding="utf-8") as f: + by_id = {it.get("id"): it for it in json.load(f).get("items", [])} + for tid in target_ids: + if job["cancel"].is_set(): + break + it = by_id.get(tid) + path = _image_path_for(it, images_dir) if it else None + try: + if path is None: + raise ValueError("no image file") + with Image.open(path) as im: + arr = np.asarray(im.convert("RGB")) + m = image_metrics.compute_metrics_numpy(arr, keys) + if m: + pending[tid] = m + else: + job["failed"] += 1 + except Exception as e: + job["failed"] += 1 + print(f"[GridTester] ⚠️ metrics skipped for item {tid}: {e}") + job["completed"] += 1 + if len(pending) >= FLUSH_EVERY: + _flush(manifest_path, pending, overwrite) + pending = {} + _send_progress(job, session_name) + _flush(manifest_path, pending, overwrite) + job["status"] = "cancelled" if job["cancel"].is_set() else "complete" + print(f"[GridTester] 📏 metrics job {job['status']}: {job['completed']}/{job['total']} images, {job['failed']} skipped") + except Exception as e: + job["status"] = "error" + job["error"] = str(e) + print(f"[GridTester] ❌ metrics job failed: {e}") + _send_progress(job, session_name) diff --git a/remote_vae.py b/remote_vae.py index 21ff6c0..216c3dc 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/resources/logic_pipeline.js b/resources/logic_pipeline.js index 92ed7c0..7967cc9 100644 --- a/resources/logic_pipeline.js +++ b/resources/logic_pipeline.js @@ -25,6 +25,10 @@ function refreshIndices() { * Returns null if the field is missing or not parseable as a number. */ function _extractFieldValue(item, field) { + if (typeof field === 'string' && field.startsWith('metric:')) { + const mv = item.metrics ? item.metrics[field.slice(7)] : undefined; + return (typeof mv === 'number' && isFinite(mv)) ? mv : null; + } if (field === 'lora_strength') { // Parse first lora's model strength from "name:0.5:0.7" format const lora = item.lora; @@ -165,6 +169,43 @@ function getFilterKey() { return parts.join('|'); } +// --- METRIC KEY DISCOVERY (optional per-image metrics) --- +var _metricKeysCache = { arr: null, len: -1, keys: [] }; + +/** Union of item.metrics keys over activeData, built-ins first. Cached per (array, length), so live inserts refresh it. */ +function getMetricKeys() { + if (!activeData) return []; + if (_metricKeysCache.arr === activeData && _metricKeysCache.len === activeData.length) return _metricKeysCache.keys; + const seen = new Set(); + for (let i = 0; i < activeData.length; i++) { + const m = activeData[i] && activeData[i].metrics; + if (m && typeof m === 'object') { + for (const k in m) seen.add(k); + } + } + const keys = METRIC_BUILTIN_ORDER.filter(k => seen.has(k)) + .concat(Array.from(seen).filter(k => !METRIC_BUILTIN_ORDER.includes(k)).sort()); + _metricKeysCache = { arr: activeData, len: activeData.length, keys }; + return keys; +} + +function metricLabel(k) { + if (METRIC_LABELS[k]) return METRIC_LABELS[k]; + return String(k).replace(/[_-]+/g, ' ').replace(/\b\w/g, c => c.toUpperCase()); +} + +/** Static sort options plus one descending option per metric present in the data. */ +function getSortOptionsList() { + return sortOptionsList.concat(getMetricKeys().map(k => ({ id: 'metric:' + k, label: metricLabel(k) + ' \u2193' }))); +} + +function _sortLabelFor(id) { + const hit = getSortOptionsList().find(x => x.id === id); + if (hit) return hit.label; + if (typeof id === 'string' && id.startsWith('metric:')) return metricLabel(id.slice(7)) + ' \u2193'; + return undefined; +} + // --- 3. SORT UI LOGIC --- function initSortUI() { @@ -200,7 +241,7 @@ function initSortUI() { container.appendChild(div); } - sortOptionsList.forEach(opt => { + getSortOptionsList().forEach(opt => { const div = document.createElement('div'); div.className = 'sort-option'; div.textContent = opt.label; @@ -231,8 +272,8 @@ function initSortUI() { renderOptions(sList, true); // Update the main button label - const primLabel = sortOptionsList.find(x => x.id === currentSort)?.label || 'Oldest'; - const secLabel = sortOptionsList.find(x => x.id === currentSecondarySort)?.label; + const primLabel = _sortLabelFor(currentSort) || 'Oldest'; + const secLabel = _sortLabelFor(currentSecondarySort); const btnLbl = document.getElementById('sort-primary-label'); if(btnLbl) { @@ -293,6 +334,18 @@ function setSecondarySort(id) { // Unified Comparison Helper function compareItems(a, b, sortKey) { + if (typeof sortKey === 'string' && sortKey.startsWith('metric:')) { + // Metric sort: descending; items without the metric sort LAST (never NaN), ties fall through. + const k = sortKey.slice(7); + const av = a.metrics ? a.metrics[k] : undefined; + const bv = b.metrics ? b.metrics[k] : undefined; + const aOk = typeof av === 'number' && isFinite(av); + const bOk = typeof bv === 'number' && isFinite(bv); + if (aOk && bOk) return bv - av; + if (aOk) return -1; + if (bOk) return 1; + return 0; + } switch (sortKey) { case 'newest': return (b.gen_index ?? b.id) - (a.gen_index ?? a.id); case 'oldest': return (a.gen_index ?? a.id) - (b.gen_index ?? b.id); @@ -493,6 +546,7 @@ let pipelineAbortId = 0; function executePipeline() { const startTime = performance.now(); refreshSupersededFps(); + if (typeof populateMetricFilterUI === 'function') populateMetricFilterUI(); const currentFilterKey = getFilterKey(); const filtersChanged = currentFilterKey !== lastFilterKey; @@ -708,7 +762,8 @@ function loadSortPreference() { const validSortOptions = sortOptionsList.map(o => o.id); - if (savedSort && validSortOptions.includes(savedSort)) { + // A saved metric sort is accepted even before data loads (an unknown metric behaves as "all missing") + if (savedSort && (validSortOptions.includes(savedSort) || savedSort.startsWith('metric:'))) { currentSort = savedSort; } diff --git a/resources/logic_state.js b/resources/logic_state.js index 250bd2e..3542052 100644 --- a/resources/logic_state.js +++ b/resources/logic_state.js @@ -63,6 +63,18 @@ const sortOptionsList = [ { id: 'rating', label: 'Rating' } ]; +// Optional per-image metrics (item.metrics = {sharpness, entropy, ...}); see image_metrics.py. +// Built-ins come first in this order, unknown keys (scorer packs) follow alphabetically. +const METRIC_BUILTIN_ORDER = ['sharpness', 'entropy', 'colorfulness', 'contrast', 'luma', 'saturation']; +const METRIC_LABELS = { + sharpness: 'Sharpness', entropy: 'Entropy', colorfulness: 'Colorful', + contrast: 'Contrast', luma: 'Luma', saturation: 'Saturation', aesthetic: 'Aesthetic' +}; +const METRIC_SHORT = { + sharpness: 'Shp', entropy: 'Ent', colorfulness: 'Col', + contrast: 'Ctr', luma: 'Lum', saturation: 'Sat', aesthetic: 'Aes' +}; + // --- SEARCH FILTER STATE --- // Array of search filter objects: { type: 'model', term: 'mountain' } @@ -129,7 +141,8 @@ var labelMode = labelMode || { upscaleTiling: false, hiresPromptBehavior: false, hiresPromptText: false, - rating: false // Pairwise-review rating badge + rating: false, // Pairwise-review rating badge + metrics: false } }; diff --git a/resources/logic_ui.js b/resources/logic_ui.js index a4cf148..2d7cc0d 100644 --- a/resources/logic_ui.js +++ b/resources/logic_ui.js @@ -302,6 +302,22 @@ function buildLabelOverlay(d) { } } + if (labelMode.fields.metrics && d.metrics && typeof d.metrics === 'object') { + const sig3 = (v) => { + const n = Number(v); + if (!isFinite(n)) return null; + return String(parseFloat(n.toPrecision(3))); + }; + const mkeys = (typeof getMetricKeys === 'function') ? getMetricKeys() : Object.keys(d.metrics); + for (const k of mkeys) { + if (!(k in d.metrics)) continue; + const v = sig3(d.metrics[k]); + if (v === null) continue; + const short = METRIC_SHORT[k] || metricLabel(k).slice(0, 3); + tags.push(`${short}:${v}`); + } + } + if (tags.length === 0) return ''; return `
`; } @@ -327,7 +343,7 @@ function toggleLabelMode(enabled) { function updateLabelFields() { const fields = ['model', 'lora', 'prompt', 'sampler', 'scheduler', 'cfg', 'steps', 'seed', 'denoise', 'upscale', 'upscaleMode', 'upscaleModel', 'upscaleRatio', 'upscaleDenoise', 'upscaleResizeMethod', 'upscaleHiresSteps', 'upscaleTiling', - 'hiresPromptBehavior', 'hiresPromptText', 'rating']; + 'hiresPromptBehavior', 'hiresPromptText', 'rating', 'metrics']; fields.forEach(f => { const cb = document.getElementById(`label-field-${f}`); if (cb) labelMode.fields[f] = cb.checked; @@ -395,7 +411,7 @@ function loadLabelPreferences() { const fields = ['model', 'lora', 'prompt', 'sampler', 'scheduler', 'cfg', 'steps', 'seed', 'denoise', 'upscale', 'upscaleMode', 'upscaleModel', 'upscaleRatio', 'upscaleDenoise', 'upscaleResizeMethod', 'upscaleHiresSteps', 'upscaleTiling', - 'hiresPromptBehavior', 'hiresPromptText', 'rating']; + 'hiresPromptBehavior', 'hiresPromptText', 'rating', 'metrics']; fields.forEach(f => { const cb = document.getElementById(`label-field-${f}`); if (cb) cb.checked = labelMode.fields[f] || false; @@ -1986,6 +2002,86 @@ function addLogicFilter() { console.log(`[Logic Filter] Added: ${field} ${op} ${value}`); } +// --- Metric fields in the Logic Filter (populated from the data; hidden when no item has metrics) --- +var _metricUiSig = null; + +function populateMetricFilterUI() { + if (typeof getMetricKeys !== 'function') return; + const keys = getMetricKeys(); + const sig = keys.join(','); + if (sig === _metricUiSig) return; + _metricUiSig = sig; + + const sel = document.getElementById('logic-filter-field'); + if (sel) { + const prev = sel.value; + const old = document.getElementById('logic-filter-metric-group'); + if (old) old.remove(); + if (keys.length) { + const group = document.createElement('optgroup'); + group.id = 'logic-filter-metric-group'; + group.label = 'Metrics'; + keys.forEach(k => { + const o = document.createElement('option'); + o.value = 'metric:' + k; + o.textContent = metricLabel(k); + group.appendChild(o); + }); + sel.appendChild(group); + } + sel.value = prev; + } + + const row = document.getElementById('metric-range-row'); + const rsel = document.getElementById('metric-range-field'); + if (row && rsel) { + const prev = rsel.value; + rsel.innerHTML = ''; + keys.forEach(k => { + const o = document.createElement('option'); + o.value = k; + o.textContent = metricLabel(k); + rsel.appendChild(o); + }); + if (keys.includes(prev)) rsel.value = prev; + row.style.display = keys.length ? 'flex' : 'none'; + updateMetricRangeHints(); + } +} + +// Placeholders show the observed range of the selected metric +function updateMetricRangeHints() { + const rsel = document.getElementById('metric-range-field'); + const minI = document.getElementById('metric-range-min'); + const maxI = document.getElementById('metric-range-max'); + if (!rsel || !minI || !maxI || !rsel.value) return; + let lo = Infinity, hi = -Infinity; + for (const it of activeData) { + const v = it && it.metrics ? it.metrics[rsel.value] : undefined; + if (typeof v === 'number' && isFinite(v)) { if (v < lo) lo = v; if (v > hi) hi = v; } + } + minI.placeholder = isFinite(lo) ? 'min ' + parseFloat(lo.toPrecision(3)) : 'min'; + maxI.placeholder = isFinite(hi) ? 'max ' + parseFloat(hi.toPrecision(3)) : 'max'; +} + +// "Metric range" ADD: pushes up to two logic filters (>= min, <= max) +function addMetricRange() { + const rsel = document.getElementById('metric-range-field'); + const minI = document.getElementById('metric-range-min'); + const maxI = document.getElementById('metric-range-max'); + if (!rsel || !minI || !maxI || !rsel.value) return; + const field = 'metric:' + rsel.value; + const lo = minI.value.trim() === '' ? null : parseFloat(minI.value); + const hi = maxI.value.trim() === '' ? null : parseFloat(maxI.value); + if ((lo === null || isNaN(lo)) && (hi === null || isNaN(hi))) { alert('Enter a min and/or max'); return; } + if (lo !== null && !isNaN(lo)) logicFilters.push({ field, op: '>=', value: lo }); + if (hi !== null && !isNaN(hi)) logicFilters.push({ field, op: '<=', value: hi }); + minI.value = ''; + maxI.value = ''; + renderLogicFilters(); + updateDataPipeline(); +} + // Remove a logic filter by index function removeLogicFilter(index) { if (index >= 0 && index < logicFilters.length) { @@ -2031,7 +2127,8 @@ function renderLogicFilters() { const tag = document.createElement('div'); tag.className = 'search-filter-tag'; - const label = fieldLabels[filter.field] || filter.field; + const label = fieldLabels[filter.field] || + (typeof filter.field === 'string' && filter.field.startsWith('metric:') ? metricLabel(filter.field.slice(7)) : filter.field); tag.innerHTML = ` ${label} ${filter.op} ${filter.value} @@ -2754,4 +2851,56 @@ function showAnalyticsModal(title, results, type) { overlay.addEventListener('click', (e) => { if (e.target === overlay) overlay.remove(); }); -} \ No newline at end of file +} + + +// ============================================================ +// IMAGE METRICS FOR AN EXISTING SESSION (background job, CPU only) +// ============================================================ +var _sessionMetricsJob = { id: null, timer: null }; + +async function startSessionMetrics() { + const btn = document.getElementById('compute-metrics-btn'); + const out = document.getElementById('compute-metrics-progress'); + if (_sessionMetricsJob.id) { + // Second click while running = cancel + try { + await fetch('/config_tester/cancel_metrics', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ job_id: _sessionMetricsJob.id }) + }); + } catch (e) { /* ignore */ } + if (out) out.textContent = 'Cancelling...'; + return; + } + const sessionName = document.getElementById('session-input')?.value || ''; + if (!sessionName) { alert('No session selected'); return; } + if (out) { out.style.display = 'block'; out.textContent = 'Starting...'; } + try { + const resp = await fetch('/config_tester/compute_metrics', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ session_name: sessionName }) + }); + const result = await resp.json(); + if (result.error) { if (out) out.textContent = result.error; return; } + _sessionMetricsJob.id = result.job_id; + if (btn) btn.textContent = 'CANCEL (click again)'; + _sessionMetricsJob.timer = setInterval(async () => { + try { + const st = await (await fetch('/config_tester/metrics_status?job_id=' + _sessionMetricsJob.id)).json(); + if (out) out.textContent = 'Measuring ' + st.completed + '/' + st.total + '...'; + if (st.status !== 'running') { + clearInterval(_sessionMetricsJob.timer); + _sessionMetricsJob = { id: null, timer: null }; + if (btn) btn.textContent = 'COMPUTE METRICS FOR THIS SESSION'; + if (out) out.textContent = st.status === 'error' ? ('Error: ' + (st.error || 'unknown')) + : (st.status === 'cancelled' ? 'Cancelled at ' : 'Done: ') + st.completed + '/' + st.total + + (st.failed ? ' (' + st.failed + ' skipped)' : '') + (st.status === 'complete' ? '. Reloading...' : ''); + if (st.status === 'complete' && typeof loadSession === 'function') loadSession(); + } + } catch (e) { /* ignore poll errors */ } + }, 1500); + } catch (e) { + if (out) out.textContent = 'Failed to start: ' + e; + } +} diff --git a/resources/report.css b/resources/report.css index 77f97de..f60940b 100644 --- a/resources/report.css +++ b/resources/report.css @@ -661,6 +661,7 @@ body { .label-tag.label-steps { background: rgba(90, 30, 30, 0.6); } .label-tag.label-seed { background: rgba(40, 40, 40, 0.6); } .label-tag.label-denoise { background: rgba(0, 70, 60, 0.6); } +.label-tag.label-metric { background: rgba(30, 60, 100, 0.6); } .label-tag.label-upscale { background: rgba(90, 30, 55, 0.6); } .label-tag.label-hiresPrompt { background: rgba(70, 30, 90, 0.6); } diff --git a/resources/template.html b/resources/template.html index 07dd0d1..b1bb63c 100644 --- a/resources/template.html +++ b/resources/template.html @@ -184,6 +184,10 @@ +