""" Main SamplerGridTester Node for ComfyUI Orchestrates grid generation across multiple configurations """ import re import torch import json import os import time import random import hashlib import uuid import folder_paths import nodes import comfy.utils import comfy.sd import comfy.samplers import comfy.model_management from PIL import Image import numpy as np # Import from split modules from .remote_vae import ( HF_ENDPOINTS, detect_model_type, remote_decode_hf, RemoteVAEDecodeWorker ) from .lora_utils import load_and_save_tags from .config_utils import ( parse_json_with_error, parse_float_input, parse_string_input, expand_configs, prepare_input_jobs, sanitize_session_name, normalize_str, get_files_from_folder, parse_lora_definition ) try: from server import PromptServer except ImportError: PromptServer = None from .html_generator import get_html_template class SamplerGridTester: @classmethod def INPUT_TYPES(s): return { "required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "Default checkpoint model. Used when config has model set to 'Default'. Configs can also specify GGUF or diffusion models directly via model_type."}), "positive_text": ("STRING", {"multiline": True, "default": "masterpiece, best quality, 1girl"}), "negative_text": ("STRING", {"multiline": True, "default": "bad quality, worst quality, lowres"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "denoise": ("STRING", {"default": "1.0", "multiline": False}), "vae_batch_size": ("INT", {"default": 4, "min": -1, "max": 64}), "configs_json": ("STRING", {"multiline": True, "default": '[{"sampler": "euler", "scheduler": "normal", "steps": 20, "cfg": 7.0}]'}), "resolutions_json": ("STRING", {"default": '[[1024, 1024]]'}), "session_name": ("STRING", {"default": "my_session"}), "overwrite_existing": ("BOOLEAN", {"default": False, "tooltip": "True = Re-run everything. False = Skip already generated images (Resume)."}), "flush_batch_every": ("INT", {"default": 4, "min": 0, "max": 64, "tooltip": "Update dashboard every X images. 0 = Use VAE Batch Size."}), "add_random_seeds_to_gens": ("INT", {"default": 0, "min": 0, "max": 100, "tooltip": "Generate X extra images per config using consistent random seeds."}), "lora_triggerwords_mode": (["None", "Append To End", "Append To Start", "Read From Config"], { "default": "None", "tooltip": "None = Don't fetch/append trigger words. Append To End = Add triggers at end of prompt (default behavior). Append To Start = Add triggers at start of prompt. Read From Config = Use lora_triggerwords_append_settings in config JSON to specify per-lora placement." }), "remote_vae_endpoint": (["None", "Auto (Experimental)", "SD", "SDXL", "Flux", "HunyuanVideo"], { "default": "None", "tooltip": "Offload VAE decoding to HuggingFace remote endpoints. Auto detects model type, or manually select endpoint. Ignores vae_batch_size and flush_batch_every when enabled." }), "save_conditioning_cache_to_file": ("BOOLEAN", { "default": False, "tooltip": "Save CLIP conditioning cache to disk. \n\n Pretty much only helps if you want to use the same prompts without changing models or loras and are planning to experiment with different values AND want to skip the text encoding step after stopping and restarting/resuming a job. \n\n WARNING: Can create very large files in your output/benchmarks folder if you use many prompts with many LoRAs on large runs. Automatically disabled when optional inputs (model/clip/conditioning) are connected, as changes cannot be reliably detected." }), "enable_model_cache": ("BOOLEAN", { "default": False, "tooltip": "\n\n Experimental WIP feature that cuts lora switching loading times. Very useful if you are using lots of loras and then switching out some of them on additional runs. Skips re-loading from disk and loads cached loras from RAM instead. \n\n This is currently hard-coded to work best on my 8gb VRAM and 64gb RAM setup, it should work fine on most systems unless you have more LoRAs than available RAM. \n\n Enable intelligent model/LoRA caching system with async background preloading. Dramatically speeds up generation by caching models and preloading next model in background. Disable for debugging or to reduce RAM/VRAM usage (can be much slower if your text encoding step takes a long time and you're doing lots of starting and stopping with the same prompt). \n\n WILL NOT AUTODELETE THE CACHED ENCODED PROMPTS. You'll need to delete them manually from the benchmark session's folder manually when you're done with them." }), }, "optional": { "optional_model": ("MODEL",), "optional_clip": ("CLIP",), "optional_vae": ("VAE",), "optional_positive": ("CONDITIONING",), "optional_negative": ("CONDITIONING",), "optional_latent": ("LATENT",), }, "hidden": {"unique_id": "UNIQUE_ID"}, } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("dashboard_html",) FUNCTION = "run_tests" CATEGORY = "sampling/testing" @classmethod def IS_CHANGED(cls, **kwargs): """ Force re-execution when optional inputs are connected. ComfyUI's standard IS_CHANGED uses input hash comparisons to decide whether to skip execution. However, for complex objects like MODEL, CLIP, VAE, and CONDITIONING passed via optional inputs, the hash doesn't reliably change when the upstream node's internal state changes (e.g., user swaps model in a loader node). When ANY optional input is connected, we return a unique value each time to force ComfyUI to always re-execute this node, ensuring the grid tester picks up upstream changes to models, LoRAs, prompts, etc. When NO optional inputs are connected, we return a deterministic hash based on the text/numeric inputs so ComfyUI can properly cache. """ optional_keys = [ "optional_model", "optional_clip", "optional_vae", "optional_positive", "optional_negative", "optional_latent" ] # Check if any optional input is connected (not None and present in kwargs) has_optional = any( kwargs.get(key) is not None for key in optional_keys ) if has_optional: # Force re-execution every time - optional inputs can't be reliably hashed return float("NaN") # No optional inputs - return deterministic hash so ComfyUI can cache properly # Hash the text/numeric inputs that we CAN reliably track hash_parts = [] for key in sorted(kwargs.keys()): if key not in optional_keys and key != "unique_id": val = kwargs.get(key) if val is not None: hash_parts.append(f"{key}={val}") return hashlib.md5("|".join(hash_parts).encode()).hexdigest() # --- HELPER METHODS --- def is_float_equal(self, a, b, tolerance=1e-5): """Robust float comparison""" try: return abs(float(a) - float(b)) < tolerance except: return str(a) == str(b) def hash_conditioning(self, conditioning): """Create a hash of conditioning tensor for change detection""" if conditioning is None: return "none" try: tensor = conditioning[0][0] tensor_bytes = tensor.cpu().numpy().tobytes() hash_obj = hashlib.md5(tensor_bytes) return hash_obj.hexdigest()[:16] except Exception as e: print(f"[GridTester] Warning: Could not hash conditioning: {e}") return "unknown" def get_latent_channels(self, model, optional_latent): """Detect the correct number of latent channels for the model""" # First, check if we have an optional_latent to extract from if optional_latent is not None: channels = optional_latent["samples"].shape[1] print(f"[GridTester] 🔍 Detected {channels} latent channels from optional_latent") return channels # Try to detect from model if model is not None: try: if hasattr(model, 'model') and hasattr(model.model, 'latent_format'): latent_format = model.model.latent_format if hasattr(latent_format, 'latent_channels'): channels = latent_format.latent_channels print(f"[GridTester] 🔍 Detected {channels} latent channels from model.latent_format") return channels if hasattr(model, 'model') and hasattr(model.model, 'diffusion_model'): diff_model = model.model.diffusion_model if hasattr(diff_model, 'in_channels'): channels = diff_model.in_channels print(f"[GridTester] 🔍 Detected {channels} latent channels from diffusion_model.in_channels") return channels except Exception as e: print(f"[GridTester] ⚠ī¸ Could not detect latent channels: {e}") # Default to 4 (SD1.5/SDXL) print(f"[GridTester] 🔍 Using default 4 latent channels (SD1.5/SDXL)") return 4 def find_existing_match(self, existing_items, conf, w, h, current_seed, batch_idx, match_keys): """Returns the index of matching item, or -1 if not found""" for idx, item in enumerate(existing_items): is_match = True for k in match_keys: val_conf = conf.get(k) # Override with current job values if k == "width": val_conf = w elif k == "height": val_conf = h elif k == "seed": val_conf = current_seed elif k == "batch_idx": val_conf = batch_idx val_item = item.get(k) # Handle model defaults if k == "model": if val_item is None: if val_conf != "Default": is_match = False break elif val_conf == "Default" and val_item is None: continue # Float comparison if isinstance(val_conf, float) or isinstance(val_item, float): if not self.is_float_equal(val_conf, val_item): is_match = False break # String comparison elif isinstance(val_conf, str) and isinstance(val_item, str): if normalize_str(val_conf) != normalize_str(val_item): is_match = False break # Direct comparison elif val_item != val_conf: is_match = False break if is_match: return idx return -1 def run_tests(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=None, optional_clip=None, optional_vae=None, optional_positive=None, optional_negative=None, optional_latent=None): # Import the generation logic from the separate module from .generation_orchestrator import run_generation_loop # Disable cache saving if any optional inputs are connected # (changes in models/LoRAs cannot be reliably detected from optional inputs) if optional_model is not None or optional_clip is not None or optional_positive is not None or optional_negative is not None: if save_conditioning_cache_to_file: print("[GridTester] ⚠ī¸ save_conditioning_cache_to_file disabled: optional inputs connected (changes cannot be reliably detected)") save_conditioning_cache_to_file = False # Extract distribution config from configs_json (embedded as _distribution key) dist_config = None session_settings = None try: import json parsed = json.loads(configs_json) if isinstance(parsed, dict) and "_distribution" in parsed: dist_config = parsed["_distribution"] session_settings = parsed.get("_session_settings") if session_settings: print(f"[GridTester] ⚙ī¸ Session settings extracted: {list(session_settings.keys())}") # Replace configs_json with just the configs array for downstream configs_json = json.dumps(parsed["configs"], indent=2, ensure_ascii=False) print(f"[GridTester] 🌐 Distribution config extracted: {len(dist_config.get('worker_urls', []))} worker(s), enabled={dist_config.get('enabled')}") elif isinstance(parsed, dict) and "configs" in parsed: # Configs wrapped in dict but no distribution settings — unwrap for downstream session_settings = parsed.get("_session_settings") if session_settings: print(f"[GridTester] ⚙ī¸ Session settings extracted: {list(session_settings.keys())}") configs_json = json.dumps(parsed["configs"], indent=2, ensure_ascii=False) print(f"[GridTester] ℹī¸ No distribution settings in configs_json") else: print(f"[GridTester] ℹī¸ No distribution settings in configs_json") except Exception as e: print(f"[GridTester] ⚠ī¸ Error parsing configs_json for distribution: {e}") # Builder UI Run Settings override widget values when present. # When a user sets these in the Builder UI, _session_settings carries # them through configs_json and they take precedence over what's # wired to the SamplerGridTester widgets. Generator widgets are still # the fallback for backward compat. if session_settings: if "overwrite_existing" in session_settings: overwrite_existing = bool(session_settings["overwrite_existing"]) print(f"[GridTester] ⚙ī¸ overwrite_existing overridden by Builder UI: {overwrite_existing}") if "flush_batch_every" in session_settings: try: flush_batch_every = int(session_settings["flush_batch_every"]) print(f"[GridTester] ⚙ī¸ flush_batch_every overridden by Builder UI: {flush_batch_every}") except (TypeError, ValueError): pass if "lora_triggerwords_mode" in session_settings: lora_triggerwords_mode = str(session_settings["lora_triggerwords_mode"]) print(f"[GridTester] ⚙ī¸ lora_triggerwords_mode overridden by Builder UI: {lora_triggerwords_mode}") if "save_conditioning_cache_to_file" in session_settings: save_conditioning_cache_to_file = bool(session_settings["save_conditioning_cache_to_file"]) print(f"[GridTester] ⚙ī¸ save_conditioning_cache_to_file overridden by Builder UI: {save_conditioning_cache_to_file}") if "enable_model_cache" in session_settings: enable_model_cache = bool(session_settings["enable_model_cache"]) print(f"[GridTester] ⚙ī¸ enable_model_cache overridden by Builder UI: {enable_model_cache}") if "vae_batch_size" in session_settings: try: vae_batch_size = int(session_settings["vae_batch_size"]) print(f"[GridTester] ⚙ī¸ vae_batch_size overridden by Builder UI: {vae_batch_size}") 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 ) # Node class mappings for ComfyUI NODE_CLASS_MAPPINGS = { "SamplerGridTester": SamplerGridTester } NODE_DISPLAY_NAME_MAPPINGS = { "SamplerGridTester": "Sampler Grid Tester" }