""" Generation Orchestrator - Main Entry Point Coordinates the entire grid generation workflow """ import os import json import time import re import gc import threading import torch import hashlib import folder_paths import comfy.sd # Required for async workers from comfy.model_management import InterruptProcessingException from .trigger_words import collect_unique_prompts_with_triggers, build_prompt_with_triggers, clear_trigger_caches from .batch_encoding import batch_encode_prompts, encode_prompt_with_combinators from .manifest_utils import load_existing_manifest, save_manifest from .model_loader import ( load_checkpoint, load_loras, cleanup_model_references, get_latent_channels, load_loras_for_preencoding, print_incompatible_loras_summary, load_diffusion_model_and_clip, load_vae_by_name ) from .lora_utils import expand_lora_folder from .image_generation import ( generate_image, flush_batch_with_vae, flush_batch_with_remote_vae, create_image_metadata, decode_latent_with_vae, calculate_eta, print_generation_progress ) from .config_utils import sanitize_session_name from .html_generator import get_html_template from .conditioning_cache import ConditioningCache from .remote_vae import RemoteVAEDecodeWorker, HF_ENDPOINTS try: from server import PromptServer except ImportError: PromptServer = None def setup_session_directories(session_name): """Create session directories and return paths.""" base_dir = os.path.join(folder_paths.get_output_directory(), "benchmarks", session_name) img_dir = os.path.join(base_dir, "images") manifest_path = os.path.join(base_dir, "manifest.json") os.makedirs(base_dir, exist_ok=True) os.makedirs(img_dir, exist_ok=True) return { "base": base_dir, "images": img_dir, "manifest": manifest_path } def initialize_remote_vae(remote_vae_endpoint, img_dir, manifest_path, existing_data, session_name, unique_id): """Initialize remote VAE worker if enabled.""" if not remote_vae_endpoint or remote_vae_endpoint == "None": return None if remote_vae_endpoint in ["SD", "SDXL", "Flux", "HunyuanVideo"]: actual_endpoint = HF_ENDPOINTS.get(remote_vae_endpoint) print(f"[GridTester] ๐ŸŒ Using {remote_vae_endpoint} endpoint: {actual_endpoint}") elif remote_vae_endpoint == "Auto (Experimental)": print(f"[GridTester] ๐ŸŒ Auto mode selected - worker will initialize on first flush") return None else: actual_endpoint = remote_vae_endpoint print(f"[GridTester] ๐ŸŒ Using custom endpoint: {actual_endpoint}") worker = RemoteVAEDecodeWorker( endpoint=actual_endpoint, img_dir=img_dir, manifest_path=manifest_path, existing_data=existing_data, session_name=session_name, unique_id=unique_id ) print(f"[GridTester] ๐ŸŒ Remote VAE worker started") return worker REMOTE_VAE_PREFIX = "remote:" def is_remote_vae(vae_string): """Check if a VAE value is a per-config remote URL (e.g. 'remote:http://...').""" return isinstance(vae_string, str) and vae_string.startswith(REMOTE_VAE_PREFIX) def extract_remote_vae_url(vae_string): """Extract the URL from a remote VAE string like 'remote:http://...'.""" url = vae_string[len(REMOTE_VAE_PREFIX):] if not url: raise ValueError( "[GridTester] Per-config remote VAE URL is empty.\n" "Please provide a valid endpoint URL in the config builder's VAE section." ) return url def _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, manifest_path, unique_id, config_overrides_vae=False): """Flush pending batch using the appropriate VAE decode method. Handles three-way dispatch: 1. Per-config remote VAE (current_vae_is_remote) โ€” uses per-config worker 2. Global remote VAE (use_remote_vae) โ€” uses global remote_vae_worker ONLY when config_overrides_vae is False (config VAE is "Default") 3. Local VAE โ€” uses loaded_vae for local decode Config Builder VAE settings take priority over the sampler node's remote_vae_endpoint when a config explicitly sets a non-Default VAE. """ if not pending_batch: return if current_vae_is_remote and current_remote_vae_url: worker = per_config_remote_workers.get(current_remote_vae_url) if worker: flush_batch_with_remote_vae(pending_batch, worker, existing_data, session_name) else: print(f"[GridTester] โš ๏ธ No remote worker for {current_remote_vae_url}, falling back to local VAE") flush_batch_with_vae(pending_batch, loaded_vae, paths["images"], existing_data, session_name, manifest_path, unique_id) elif not config_overrides_vae and use_remote_vae and remote_vae_worker: flush_batch_with_remote_vae(pending_batch, remote_vae_worker, existing_data, session_name) else: flush_batch_with_vae(pending_batch, loaded_vae, paths["images"], existing_data, session_name, manifest_path, unique_id) def _cleanup_per_config_remote_workers(per_config_remote_workers): """Shut down all per-config remote VAE workers.""" for url, worker in per_config_remote_workers.items(): try: print(f"[GridTester] ๐ŸŒ Waiting for per-config remote VAE ({url})...") worker.wait_completion() worker.stop() except Exception as e: print(f"[GridTester] โš ๏ธ Error stopping per-config remote worker ({url}): {e}") def calculate_clip_hash(clip_model): """Calculate a hash of the CLIP model for cache validation.""" try: if hasattr(clip_model, 'state_dict'): state_dict = clip_model.state_dict() model_signature = str([(k, tuple(v.shape)) for k, v in list(state_dict.items())[:10]]) elif hasattr(clip_model, 'cond_stage_model'): model_signature = str(type(clip_model.cond_stage_model)) else: model_signature = str(type(clip_model)) return hashlib.md5(model_signature.encode()).hexdigest()[:16] except: return "unknown" def check_if_job_completed(existing_items, conf, seed, width, height, batch_idx, positive_prompt, negative_prompt, has_optional_inputs=False): """Independent check to see if a specific generation job already exists. When has_optional_inputs is True, we cannot reliably detect changes from upstream nodes (models, LoRAs, conditionings, latents), so we disable skip detection entirely and return -1 (no match) to force re-generation. """ # When optional inputs are connected, we can't reliably match on model, # lora, prompts, or latents because those may come from upstream nodes # whose changes we can't detect. Disable skip detection entirely. if has_optional_inputs: return -1 FLOAT_TOLERANCE = 0.0001 for idx, item in enumerate(existing_items): if item.get("seed") != seed: continue if item.get("width") != width: continue if item.get("height") != height: continue if item.get("batch_idx", 0) != batch_idx: continue if item.get("sampler") != conf["sampler"]: continue if item.get("scheduler") != conf["scheduler"]: continue # Check attention mode (default = no attention_mode key or "default") item_attn = item.get("attention_mode", "default") conf_attn = conf.get("attention_mode", "default") if item_attn != conf_attn: continue try: if abs(float(item.get("steps")) - float(conf["steps"])) > FLOAT_TOLERANCE: continue if abs(float(item.get("cfg")) - float(conf["cfg"])) > FLOAT_TOLERANCE: continue if abs(float(item.get("denoise")) - float(conf["denoise"])) > FLOAT_TOLERANCE: continue except (ValueError, TypeError): continue # Standard matching - check model, lora, and prompts if item.get("model") != conf["model"]: continue if item.get("positive", "").strip() != positive_prompt.strip(): continue if item.get("negative", "").strip() != negative_prompt.strip(): continue item_lora = item.get("lora", "None") conf_lora = conf.get("lora_expanded", "None") if item_lora != conf_lora: continue # Check model_type and text_encoders โ€” different text encoders produce # different conditioning even with the same model file, so they must # NOT be considered duplicate jobs item_model_type = item.get("model_type", "checkpoint") conf_model_type = conf.get("model_type", "checkpoint") if item_model_type != conf_model_type: continue item_te = item.get("text_encoders", []) conf_te = conf.get("text_encoders", []) if item_te != conf_te: continue # Check clip_type โ€” different clip types produce different conditioning item_clip_type = item.get("clip_type", "stable_diffusion") conf_clip_type = conf.get("clip_type", "stable_diffusion") if item_clip_type != conf_clip_type: continue # Check clip_skip โ€” different clip_skip values produce different conditioning item_clip_skip = item.get("clip_skip", 0) conf_clip_skip = conf.get("clip_skip", 0) if item_clip_skip != conf_clip_skip: continue # Check VAE โ€” different VAEs produce different decoded images item_vae = item.get("vae", "Default") conf_vae = conf.get("vae", "Default") if item_vae != conf_vae: continue return idx return -1 def get_model_cache_key(conf): """Generate a cache key that uniquely identifies the model+clip combination.""" model_type = conf.get("model_type", "checkpoint") if model_type == "checkpoint": return conf["model"] else: te_key = "|".join(sorted(conf.get("text_encoders", []))) return f"{conf['model']}::{model_type}::{te_key}" def load_model_by_type(conf, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, loaded_clip, loaded_vae, model_cache): """Dispatch to correct loader based on model_type in config.""" model_type = conf.get("model_type", "checkpoint") target = conf["model"] if model_type == "checkpoint": return load_checkpoint( target, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, loaded_clip, loaded_vae, model_cache=model_cache ) else: # diffusion_model or gguf return load_diffusion_model_and_clip( model_name=target, model_type=model_type, text_encoder_paths=conf.get("text_encoders", []), clip_type_str=conf.get("clip_type", "stable_diffusion"), gguf_options=conf.get("gguf_options"), use_remote_vae=use_remote_vae, optional_model=optional_model, optional_clip=optional_clip, optional_vae=optional_vae, model_cache=model_cache ) def run_deferred_upscales( deferred_queue, upscale_settings, session_settings, loaded_vae, patched_model, patched_clip, conditioning_cache, paths, existing_data, session_name, unique_id, PromptServer, config_overrides_vae=None ): """ Process deferred upscale jobs after generation completes. Runs all pipeline chains for each queued image. Called BEFORE finalization cleanup so VAE/model/CLIP are still loaded. """ if not deferred_queue: return from .image_generation import upscale_image, decode_latent_with_vae, create_image_metadata import itertools as up_itertools import random as up_random from PIL import Image as PILImage import comfy.model_management as mm import numpy as np import torch pipelines = upscale_settings.get("pipelines", []) if not pipelines: return total_jobs = len(deferred_queue) print(f"\n[GridTester] ๐Ÿ”„ Starting deferred upscale phase: {total_jobs} images") upscale_durations = [] for job_idx, job in enumerate(deferred_queue): # Check for interrupt if mm.processing_interrupted(): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during deferred upscaling") break conf = job["config"] pipe_w = job["width"] pipe_h = job["height"] current_seed = job["seed"] actual_positive_prompt = job["actual_positive_prompt"] actual_negative_prompt = job["actual_negative_prompt"] # Find the base image in the manifest by matching identity base_item = None for item in existing_data.get("items", []): if (item.get("seed") == current_seed and item.get("width") == pipe_w and item.get("height") == pipe_h and not item.get("upscaled") and item.get("positive") == actual_positive_prompt): base_item = item break if not base_item: print(f"[GridTester] โš ๏ธ Deferred upscale #{job_idx+1}: base image not found in manifest, skipping") continue base_filename = base_item.get("filename") if not base_filename: print(f"[GridTester] โš ๏ธ Deferred upscale #{job_idx+1}: no filename in manifest, skipping") continue base_image_path = os.path.join(paths["images"], base_filename) if not os.path.exists(base_image_path): print(f"[GridTester] โš ๏ธ Deferred upscale #{job_idx+1}: file not found: {base_filename}, skipping") continue # Check if already upscaled (resume support) already_upscaled = False for item in existing_data.get("items", []): if (item.get("upscaled") and item.get("seed") == current_seed and item.get("positive") == actual_positive_prompt and item.get("width") != pipe_w): already_upscaled = True break if already_upscaled: print(f"[GridTester] โญ๏ธ Deferred upscale #{job_idx+1}: already upscaled, skipping") continue # Load base image from disk and VAE-encode back to latent try: pil_image = PILImage.open(base_image_path).convert("RGB") img_array = np.array(pil_image).astype(np.float32) / 255.0 img_tensor = torch.from_numpy(img_array).unsqueeze(0) encoded = loaded_vae.encode(img_tensor[:, :, :, :3]) result_latent = {"samples": encoded} except Exception as e: print(f"[GridTester] โš ๏ธ Deferred upscale #{job_idx+1}: failed to load/encode: {e}") continue # Look up conditioning from cache final_positive = conditioning_cache["positive"].get(actual_positive_prompt) final_negative = conditioning_cache["negative"].get(actual_negative_prompt) if final_positive is None or final_negative is None: print(f"[GridTester] โš ๏ธ Deferred upscale #{job_idx+1}: conditioning not in cache, re-encoding") try: clip_skip = conf.get("clip_skip", 0) if final_positive is None: final_positive = encode_prompt_with_combinators(patched_clip, actual_positive_prompt, clip_skip) conditioning_cache["positive"][actual_positive_prompt] = final_positive if final_negative is None: final_negative = encode_prompt_with_combinators(patched_clip, actual_negative_prompt, clip_skip) conditioning_cache["negative"][actual_negative_prompt] = final_negative except Exception as e: print(f"[GridTester] โš ๏ธ Deferred upscale #{job_idx+1}: re-encode failed: {e}, skipping") continue # HiRes prompt adjustment hires_positive_cond = final_positive hires_prompt_active = False hires_prompt_behavior_rt = "" hires_prompt_text_rt = "" if upscale_settings.get("hires_prompt_adjust") and upscale_settings.get("hires_prompt_text", "").strip(): hires_prompt_behavior_rt = upscale_settings.get("hires_prompt_behavior", "append_end") hires_prompt_text_rt = upscale_settings["hires_prompt_text"].strip() if hires_prompt_behavior_rt == "prepend": adjusted_prompt = hires_prompt_text_rt + " " + actual_positive_prompt elif hires_prompt_behavior_rt == "append_end": adjusted_prompt = actual_positive_prompt + " " + hires_prompt_text_rt elif hires_prompt_behavior_rt == "replace": adjusted_prompt = hires_prompt_text_rt else: adjusted_prompt = actual_positive_prompt hires_cond = conditioning_cache["positive"].get(adjusted_prompt) if hires_cond is not None: hires_positive_cond = hires_cond hires_prompt_active = True job_start_time = time.time() upscale_combo_idx = 0 for pipeline_idx, pipeline in enumerate(pipelines): if pipeline.get("active", True) is False: continue pipeline_name = pipeline.get("name", f"Pipeline {pipeline_idx + 1}") pipeline_steps = pipeline.get("steps", []) if not pipeline_steps: continue pipe_latent = result_latent pipe_w_current = pipe_w pipe_h_current = pipe_h expanded_steps = [] for step in pipeline_steps: if step.get("active", True) is False: continue repeat = max(1, int(step.get("repeat", 1))) for _ in range(repeat): expanded_steps.append(step) for step_idx, ucfg in enumerate(expanded_steps): mode = ucfg.get("mode", "hires_only") show_hires = mode in ("hires_only", "model_then_hires") show_model = mode in ("model_only", "model_then_hires") raw_ratios = str(ucfg.get("upscale_ratios", "1.5")) ratios = [float(r.strip()) for r in raw_ratios.split(",") if r.strip()] or [1.5] raw_denoise = str(ucfg.get("hires_denoise", "0.3")) denoises = [float(d.strip()) for d in raw_denoise.split(",") if d.strip()] or [0.3] models = ucfg.get("upscale_models", []) or [""] if show_hires and show_model: combos = list(up_itertools.product(models, ratios, denoises)) elif show_hires: combos = list(up_itertools.product([""], ratios, denoises)) elif show_model: combos = list(up_itertools.product(models, [1.0], [0.0])) else: combos = [] for combo in combos: up_model, up_ratio, up_denoise = combo single_config = { "mode": mode, "upscale_model": up_model, "upscale_ratio": up_ratio, "hires_denoise": up_denoise, "hires_steps": ucfg.get("hires_steps", 0), "tiled_vae": ucfg.get("tiled_vae", False), "tile_size": ucfg.get("tile_size", 512), "upscale_size": ucfg.get("upscale_size", "2.0"), "resize_method": ucfg.get("resize_method", "bilinear"), "hires_tiled_sampling": ucfg.get("hires_tiled_sampling", False), "hires_tile_width": ucfg.get("hires_tile_width", 512), "hires_tile_height": ucfg.get("hires_tile_height", 512), "hires_mask_blur": ucfg.get("hires_mask_blur", 8), "hires_tile_padding": ucfg.get("hires_tile_padding", 32), "hires_force_uniform_tiles": ucfg.get("hires_force_uniform_tiles", False) } up_positive = hires_positive_cond if (show_hires and hires_prompt_active) else final_positive upscale_result, upscale_duration = upscale_image( pipe_latent, loaded_vae, patched_model, single_config, conf, up_positive, final_negative, pipe_w_current, pipe_h_current ) is_last_step = step_idx == len(expanded_steps) - 1 is_last_combo = combo == combos[-1] is_final_output = is_last_step and is_last_combo if isinstance(upscale_result, dict) and "samples" in upscale_result: upscaled_pil = decode_latent_with_vae(loaded_vae, upscale_result["samples"]) up_w, up_h = upscaled_pil.size if is_final_output: upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000) upscaled_filename = f"img_{upscale_id}_upscaled.webp" upscaled_pil.save( os.path.join(paths["images"], upscaled_filename), format="WEBP", quality=95 ) elif isinstance(upscale_result, PILImage.Image): up_w, up_h = upscale_result.size if is_final_output: upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000) upscaled_filename = f"img_{upscale_id}_upscaled.webp" upscale_result.save( os.path.join(paths["images"], upscaled_filename), format="WEBP", quality=95 ) else: continue if not is_final_output: if isinstance(upscale_result, dict) and "samples" in upscale_result: pipe_latent = upscale_result pipe_w_current = up_w pipe_h_current = up_h continue # Final output: create manifest entry upscaled_meta = create_image_metadata( conf, up_w, up_h, upscale_duration, current_seed, job.get("batch_idx", 0), actual_positive_prompt, actual_negative_prompt, gen_index=job.get("gen_index", 0) ) upscaled_meta["id"] = upscale_id upscaled_meta["filename"] = upscaled_filename upscaled_meta["upscaled"] = True upscaled_meta["upscale_pipeline"] = pipeline_name upscaled_meta["upscale_mode"] = mode upscaled_meta["upscale_ratio"] = up_ratio upscaled_meta["upscale_denoise"] = up_denoise if up_model: upscaled_meta["upscale_model"] = up_model if hires_prompt_active: upscaled_meta["hires_prompt_behavior"] = hires_prompt_behavior_rt upscaled_meta["hires_prompt_text"] = hires_prompt_text_rt existing_data["items"].append(upscaled_meta) upscale_combo_idx += 1 if upscale_combo_idx > 0: save_manifest(paths["manifest"], existing_data) if PromptServer is not None: try: PromptServer.instance.send_sync("ultimate_grid.update_data", { "node": unique_id, "session_name": session_name, "new_items": [existing_data["items"][-1]] }) except Exception: pass job_duration = time.time() - job_start_time upscale_durations.append(job_duration) # Progress reporting if PromptServer is not None: try: avg_dur = sum(upscale_durations) / len(upscale_durations) remaining = (total_jobs - (job_idx + 1)) * avg_dur mins = int(remaining // 60) secs = int(remaining % 60) eta_str = f"{mins}m {secs}s" if mins > 0 else f"{secs}s" PromptServer.instance.send_sync("ultimate_grid.progress", { "node": unique_id, "session_name": session_name, "current_job": job_idx + 1, "total_jobs": total_jobs, "progress_pct": int(((job_idx + 1) / total_jobs) * 100), "eta_str": eta_str, "avg_duration": round(avg_dur, 1), "last_duration": round(job_duration, 1), "phase": "upscaling", }) except Exception: pass print(f"[GridTester] ๐Ÿ”„ Deferred upscale {job_idx+1}/{total_jobs} complete ({job_duration:.1f}s)") print(f"[GridTester] โœ… Deferred upscale phase complete: {len(upscale_durations)}/{total_jobs} images upscaled") def _expand_lora_weight_arrays(expanded_configs): """ Expand LoRA weight array bracket notation into Cartesian product of configs. Input format per lora part: "name:[0.5, 0.8]:1.0" or "name:0.5:[0.7, 1.0]" Weight arrays use brackets in the STRENGTH fields (after first colon), while folder random syntax uses brackets in the NAME field (before first colon). If multiple loras in the same config have weight arrays, produces a Cartesian product of all combinations: lora1:[0.5, 1.0] + lora2:[0.7, 1.0] โ†’ 4 configs """ import re from itertools import product as itertools_product result = [] weight_array_pattern = re.compile(r'\[([^\]]+)\]') for conf in expanded_configs: lora_str = conf.get("lora", "None") if lora_str == "None" or ":" not in lora_str: result.append(conf) continue lora_parts = lora_str.split(" + ") # Check each part for weight arrays (brackets in strength fields, not name field) parts_with_arrays = [] has_weight_arrays = False for part in lora_parts: part = part.strip() colon_idx = part.find(":") if colon_idx < 0: # No strength specified โ€” plain lora name parts_with_arrays.append([(part, None, None)]) continue name = part[:colon_idx] strength_str = part[colon_idx + 1:] # If brackets are in the name portion, it's folder syntax โ€” skip if "[" in name: parts_with_arrays.append([(part, None, None)]) continue # Parse strength fields for bracket arrays # Format: "model_str:clip_str" where either can be "[val1, val2]" strength_parts = strength_str.split(":") model_strs = strength_parts[0] if len(strength_parts) > 0 else "1.0" clip_strs = strength_parts[1] if len(strength_parts) > 1 else model_strs # Extract arrays from model strength model_match = weight_array_pattern.search(model_strs) if model_match: model_vals = [v.strip() for v in model_match.group(1).split(",")] has_weight_arrays = True else: model_vals = [model_strs.strip()] # Extract arrays from clip strength clip_match = weight_array_pattern.search(clip_strs) if clip_match: clip_vals = [v.strip() for v in clip_match.group(1).split(",")] has_weight_arrays = True else: clip_vals = [clip_strs.strip()] # Build all combinations for this lora part combos = [] for mv in model_vals: for cv in clip_vals: combos.append((name, mv, cv)) parts_with_arrays.append(combos) if not has_weight_arrays: result.append(conf) continue # Cartesian product across all lora parts for combo in itertools_product(*parts_with_arrays): new_lora_parts = [] for item in combo: name, mv, cv = item if mv is None: # Passthrough (folder syntax or plain name) new_lora_parts.append(name) else: new_lora_parts.append(f"{name}:{mv}:{cv}") new_conf = conf.copy() new_conf["lora"] = " + ".join(new_lora_parts) result.append(new_conf) if len(result) != len(expanded_configs): print(f"[GridTester] ๐ŸŽฏ LoRA weight arrays expanded: {len(expanded_configs)} โ†’ {len(result)} configs") return result def 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=None, session_settings=None ): """Main generation loop orchestrator.""" # Clear stale trigger word caches from previous runs (within same ComfyUI session). # Without this, @lru_cache and _build_prompt_cache can serve outdated results # if loras_tags.json was updated by a CivitAI fetch in a prior run. clear_trigger_caches() from .model_cache import ModelCache # Initialize Model Cache (or disable completely if user requested) if enable_model_cache: model_cache = ModelCache( max_models=1, # Adjust for your VRAM max_lora_sets=2, # Adjust for your VRAM max_lora_files=30, # Adjust for your VRAM enable_preload=True, # Preloading enabled when cache is on async_preload=True, # Always async when preloading cache_device='cpu', verbose=True ) else: # Completely disable entire cache system print("[GridTester] โš ๏ธ Model cache DISABLED - all models will load from disk (slower but saves RAM/VRAM)") model_cache = None # ==== SETUP ==== session_name = sanitize_session_name(session_name) paths = setup_session_directories(session_name) existing_data = load_existing_manifest(paths["manifest"]) existing_data["session_name"] = session_name from .config_utils import ( parse_json_with_error, parse_float_input, parse_prompt_input_nested, expand_configs, prepare_input_jobs ) raw_configs = parse_json_with_error(configs_json, "configs") denoise_values = parse_float_input(denoise) resolutions = parse_json_with_error(resolutions_json, "resolutions") pos_prompts = parse_prompt_input_nested(positive_text) neg_prompts = parse_prompt_input_nested(negative_text) extra_seeds = [] if add_random_seeds_to_gens > 0: import random extra_seeds = [random.randint(0, 2**32 - 1) for _ in range(add_random_seeds_to_gens)] expanded = expand_configs(raw_configs, pos_prompts, neg_prompts, denoise_values, seed, extra_seeds, ckpt_name) expanded.sort(key=lambda x: (x.get('model_type', 'checkpoint'), x['model'], tuple(x.get('text_encoders', [])), x.get('vae', 'Default'), x['lora'], x.get('attention_mode', 'default'), x['positive'], x['negative'])) # ==== EXPAND LORA WEIGHT ARRAYS (Cartesian product) ==== # Format: "lora:[0.5, 0.8]:1.0" โ†’ expand into separate configs per strength combo expanded = _expand_lora_weight_arrays(expanded) # ==== EXPAND LORA FOLDERS ==== print(f"[GridTester] ๐ŸŽฒ Expanding LoRA folders and random selections...") for conf in expanded: if conf["lora"] != "None": lora_parts = conf["lora"].split(" + ") expanded_parts = [] for part in lora_parts: part = part.strip() if "[" in part and "]" in part: expanded_lora = expand_lora_folder(part, seed=conf.get("seed")) if expanded_lora: if isinstance(expanded_lora, list): expanded_parts.extend(expanded_lora) else: expanded_parts.append(expanded_lora) else: expanded_parts.append(part) conf["lora_expanded"] = " + ".join(expanded_parts) if expanded_parts else "None" conf["lora"] = conf["lora_expanded"] else: conf["lora_expanded"] = "None" print(f"[GridTester] โœ… LoRA expansion complete") # ==== VAE VALIDATION FOR NON-CHECKPOINT MODELS ==== # Only error if non-checkpoint model has no VAE source (no optional_vae AND no per-config VAE) use_remote_vae = remote_vae_endpoint and remote_vae_endpoint != "None" needs_vae = any( c.get("model_type", "checkpoint") != "checkpoint" and c.get("vae", "Default") == "Default" for c in expanded ) if needs_vae and not use_remote_vae and optional_vae is None: raise ValueError( "GGUF and diffusion models do not include a bundled VAE.\n" "You must either:\n" " 1. Connect a VAE to the optional_vae input on the sampler node, or\n" " 2. Enable remote VAE decoding via remote_vae_endpoint, or\n" " 3. Set a specific VAE in the config builder for those configs\n" ) # ==== REGISTER SMART CACHE SCHEDULE ==== if model_cache: model_cache.register_schedule(expanded) input_jobs = prepare_input_jobs(optional_latent, resolutions) # Count total jobs: configs with per-config resolutions provide their own # width/height and don't multiply with the global input_jobs. configs_with_res = sum(1 for c in expanded if c.get("resolution") is not None) configs_without_res = len(expanded) - configs_with_res total_jobs = configs_with_res + (configs_without_res * len(input_jobs)) print(f"{'='*80}") print(f"[GridTester] ๐Ÿš€ GENERATION START") if configs_with_res > 0: print(f"[GridTester] ๐Ÿ“‹ {configs_with_res} configs with per-config resolutions + " f"{configs_without_res} configs ร— {len(input_jobs)} resolutions = {total_jobs} total jobs") else: print(f"[GridTester] ๐Ÿ“‹ {len(expanded)} configs ร— {len(input_jobs)} resolutions = {total_jobs} total jobs") print(f"{'='*80}") # ==== OPTIONAL INPUT DETECTION ==== # Track whether any optional inputs are connected - this affects skip/resume logic has_optional_inputs = any([ optional_model is not None, optional_clip is not None, optional_vae is not None, optional_positive is not None, optional_negative is not None, optional_latent is not None ]) if has_optional_inputs and not overwrite_existing: print(f"[GridTester] โš ๏ธ Optional inputs connected โ€” job skip/resume DISABLED.") print(f"[GridTester] โš ๏ธ Changes to upstream nodes (models, LoRAs, prompts, latents) connected via optional inputs") print(f"[GridTester] โš ๏ธ cannot be automatically detected, so all jobs will be re-generated.") print(f"[GridTester] โš ๏ธ To use resume mode, disconnect optional inputs and use the built-in config fields instead.") # ==== DISTRIBUTED PROCESSING BRANCH ==== # If distribution_config is provided and enabled, delegate to distributed processing # which handles the generation loop, then jump to finalization dist_enabled = ( distribution_config and isinstance(distribution_config, dict) and distribution_config.get("enabled") and distribution_config.get("worker_urls") ) if distribution_config: print(f"[GridTester] ๐ŸŒ Distribution check: type={type(distribution_config).__name__}, " f"enabled={distribution_config.get('enabled') if isinstance(distribution_config, dict) else 'N/A'}, " f"workers={len(distribution_config.get('worker_urls', [])) if isinstance(distribution_config, dict) else 'N/A'}, " f"dist_enabled={dist_enabled}") else: print(f"[GridTester] โ„น๏ธ distribution_config is None/empty, running normal generation") if dist_enabled: print(f"[GridTester] ๐ŸŒ ENTERING DISTRIBUTED MODE with {len(distribution_config.get('worker_urls', []))} worker(s)") # Upscaling in distributed mode: workers generate base images, master runs upscales after collection if session_settings and session_settings.get("upscaling", {}).get("enabled", False): print(f"[GridTester] โ„น๏ธ Upscaling enabled โ€” workers will generate base images, master will run upscales after all generation completes.") return _run_distributed_generation( self, distribution_config, expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode, session_name, paths, unique_id, total_jobs, seed, ckpt_name, positive_text, negative_text, vae_batch_size, configs_json, resolutions_json, flush_batch_every, use_remote_vae, remote_vae_endpoint, save_conditioning_cache_to_file, enable_model_cache, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, optional_latent, model_cache, session_settings=session_settings ) # ==== OPTIONAL CONDITIONING SETUP ==== if optional_positive or optional_negative: pos_hash = hashlib.md5(str(optional_positive).encode()).hexdigest()[:16] if optional_positive else None neg_hash = hashlib.md5(str(optional_negative).encode()).hexdigest()[:16] if optional_negative else None else: pos_hash, neg_hash = None, None # ==== REMOTE VAE SETUP (use_remote_vae already defined above for VAE validation) ==== try: if PromptServer is not None: pbar = PromptServer.instance.progress_bar_pool.get_progress_bar(unique_id) else: pbar = None except: pbar = None # ==== STATE VARIABLES ==== loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None cached_model_key = None cached_lora_key = None cached_lora_cache_key = None cached_vae_key = None # Track which VAE is currently loaded default_model_vae = None # Track the model's bundled/default VAE for reverting conditioning_cache = {"positive": {}, "negative": {}} incompatible_loras = {} pending_batch = [] current_job = 0 total_generated = 0 gen_index_offset = len(existing_data.get("items", [])) # Sequential index for deterministic sort ordering skipped_count = 0 start_at_job = int(session_settings.get("start_at_job", 0)) if session_settings else 0 if start_at_job > 0: print(f"[GridTester] โญ๏ธ Start At Job #{start_at_job} โ€” skipping earlier jobs") job_durations = [] deferred_upscale_queue = [] eta_start_time = time.time() # Initialize remote VAE worker remote_vae_worker = None if use_remote_vae and expanded: remote_vae_worker = initialize_remote_vae( remote_vae_endpoint, paths["images"], paths["manifest"], existing_data, session_name, unique_id ) # Per-config remote VAE state per_config_remote_workers = {} # Keyed by URL, reused across configs current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = False # True when config explicitly sets a non-Default VAE (overrides sampler node's remote_vae_endpoint) # ==== PRE-ENCODING STAGE ==== unique_model_keys = set(get_model_cache_key(conf) for conf in expanded) if not (optional_positive and optional_negative) and expanded and len(unique_model_keys) == 1: first_conf = expanded[0] target_model_name = first_conf["model"] print(f"[GridTester] โœ… Single model detected ({target_model_name}) - enabling pre-encoding") loaded_model, loaded_clip, loaded_vae = load_model_by_type( first_conf, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, None, None, model_cache=model_cache ) # VAE fallback for non-checkpoint models if loaded_vae is None and optional_vae is not None: loaded_vae = optional_vae # Remember this model's default VAE for reverting later default_model_vae = loaded_vae # Per-config VAE: if first config specifies a VAE, load it # Config Builder VAE settings take priority over sampler node's remote_vae_endpoint target_vae = first_conf.get("vae", "Default") if target_vae != "Default": config_overrides_vae = True if is_remote_vae(target_vae): url = extract_remote_vae_url(target_vae) current_vae_is_remote = True current_remote_vae_url = url print(f"[GridTester] ๐ŸŒ Using per-config remote VAE: {url}") if url not in per_config_remote_workers: per_config_remote_workers[url] = RemoteVAEDecodeWorker( endpoint=url, img_dir=paths["images"], manifest_path=paths["manifest"], existing_data=existing_data, session_name=session_name, unique_id=unique_id ) else: print(f"[GridTester] ๐ŸŽจ Loading per-config VAE: {target_vae}") loaded_vae = load_vae_by_name(target_vae) cached_vae_key = target_vae else: cached_vae_key = "Default" if first_conf["lora_expanded"] != "None": patched_model, patched_clip = load_loras_for_preencoding( loaded_model, loaded_clip, first_conf["lora_expanded"] ) else: patched_model, patched_clip = loaded_model, loaded_clip clip_hash = calculate_clip_hash(patched_clip) cond_cache = ConditioningCache( cache_dir=paths["base"], clip_hash=clip_hash, enable_disk_cache=save_conditioning_cache_to_file) # Filter configs to only collect prompts from jobs that actually need to run # This avoids wasting time encoding prompts for jobs that will be skipped in resume mode print(f"[GridTester] ๐Ÿง  Collecting unique prompts...") if not overwrite_existing and existing_data.get("items"): configs_needing_work = [] for conf in expanded: conf_positive, _ = build_prompt_with_triggers(conf, lora_triggerwords_mode) conf_negative = conf["negative"] conf_seed = conf["seed"] # Check if ANY resolution/batch needs this config needs_work = False for job in input_jobs: if check_if_job_completed( existing_data["items"], conf, conf_seed, job["width"], job["height"], job["batch_idx"], conf_positive, conf_negative, has_optional_inputs=has_optional_inputs ) == -1: needs_work = True break if needs_work: configs_needing_work.append(conf) skippable = len(expanded) - len(configs_needing_work) if skippable > 0: print(f"[GridTester] โญ๏ธ {skippable}/{len(expanded)} configs already completed, encoding only {len(configs_needing_work)} needed") if not configs_needing_work: print(f"[GridTester] โญ๏ธ All configs already completed, skipping encoding entirely") unique_positives, unique_negatives = set(), set() else: unique_positives, unique_negatives = collect_unique_prompts_with_triggers( configs_needing_work, lora_triggerwords_mode ) else: unique_positives, unique_negatives = collect_unique_prompts_with_triggers( expanded, lora_triggerwords_mode ) # Collect adjusted HiRes prompts for pre-encoding (if hires prompt adjustment is enabled) upscale_settings = session_settings.get("upscaling", {}) if session_settings else {} if upscale_settings.get("enabled") and upscale_settings.get("hires_prompt_adjust") and upscale_settings.get("hires_prompt_text", "").strip(): hires_behavior = upscale_settings.get("hires_prompt_behavior", "append_end") hires_text = upscale_settings["hires_prompt_text"].strip() hires_adjusted_positives = set() for base_prompt in unique_positives: if hires_behavior == "prepend": adjusted = hires_text + " " + base_prompt elif hires_behavior == "append_end": adjusted = base_prompt + " " + hires_text elif hires_behavior == "replace": adjusted = hires_text else: adjusted = base_prompt hires_adjusted_positives.add(adjusted) # Add adjusted prompts to unique set so they get batch-encoded unique_positives = unique_positives | hires_adjusted_positives print(f"[GridTester] ๐Ÿ” Added {len(hires_adjusted_positives)} HiRes-adjusted prompts for pre-encoding") clip_skip = first_conf.get("clip_skip", 0) if clip_skip != 0: print(f"[GridTester] ๐Ÿ”ง Using clip_skip={clip_skip}") try: conditioning_cache = batch_encode_prompts( patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip, enable_disk_cache=save_conditioning_cache_to_file ) except InterruptProcessingException: print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during pre-encoding - Stopping all jobs") loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Re-raise so ComfyUI knows execution was interrupted (not completed) raise cached_model_key = get_model_cache_key(first_conf) cached_lora_key = first_conf["lora_expanded"] if model_cache: cached_lora_cache_key = model_cache._get_lora_cache_key(cached_model_key, first_conf["lora_expanded"]) else: cached_lora_cache_key = None latent_channels = get_latent_channels(loaded_model, optional_latent) elif len(unique_model_keys) > 1: print(f"[GridTester] โš ๏ธ Multiple models detected ({len(unique_model_keys)} different models) - pre-encoding DISABLED") cond_cache = None latent_channels = 4 else: print(f"[GridTester] โ„น๏ธ Using optional conditioning, skipping pre-encoding") cond_cache = None latent_channels = 4 # ==== FULL RUN SEED BEHAVIOR (PRE-RUN) ==== # Apply "random_before" full run seed behavior: randomize seed before the entire session for conf_idx_pre, conf_pre in enumerate(expanded): if conf_pre.get("full_run_seed_behavior") == "random_before": import random conf_pre["seed"] = random.randint(0, 2**63 - 1) print(f"[GridTester] ๐ŸŽฒ Full run random_before: config {conf_idx_pre} seed โ†’ {conf_pre['seed']}") # ==== MAIN GENERATION LOOP ==== print(f"\n{'='*80}\n") for job_idx, job in enumerate(input_jobs): w, h = job["width"], job["height"] batch_idx = job["batch_idx"] for conf_idx, conf in enumerate(expanded): # Per-config resolution override: if this config has its own resolution, # use it instead of the global input_job's resolution. Only process once # (at job_idx == 0) to avoid duplicating work across input_jobs. if conf.get("resolution") is not None: if job_idx > 0: continue # Already processed at job_idx=0 w, h = conf["resolution"] batch_idx = 0 # ==== CHECK FOR INTERRUPT ==== # ==== UPDATE CURRENT STEP ==== if model_cache: model_cache.set_current_step(conf_idx) try: import comfy.model_management as mm if mm.processing_interrupted(): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED - Stopping all jobs") if pending_batch: _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae) pending_batch = [] _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Re-raise so ComfyUI knows execution was interrupted (not completed) raise InterruptProcessingException() except InterruptProcessingException: raise except: pass current_seed = conf["seed"] if conf.get("seed_behavior") == "randomize": import random current_seed = random.randint(0, 2**63 - 1) current_job += 1 if pbar: try: pbar.update_absolute(current_job, total_jobs) except: pass progress_pct = int((current_job / total_jobs) * 100) print(f"[GridTester] ๐Ÿ“Š {current_job}/{total_jobs} ({progress_pct}%) | " f"{conf['sampler']} @ {conf['steps']} steps | {w}x{h}") actual_positive_prompt, lora_triggers = build_prompt_with_triggers( conf, lora_triggerwords_mode ) actual_negative_prompt = conf["negative"] # ==== START AT JOB # (skip earlier jobs) ==== if start_at_job > 0 and current_job < start_at_job: skipped_count += 1 continue # ==== CHECK EXISTING MATCH ==== match_index = check_if_job_completed( existing_data["items"], conf, current_seed, w, h, batch_idx, actual_positive_prompt, actual_negative_prompt, has_optional_inputs=has_optional_inputs ) if match_index != -1: if not overwrite_existing: skipped_count += 1 continue else: old_item = existing_data["items"][match_index] try: old_fname_match = re.search(r'filename=([^&]+)', old_item["file"]) if old_fname_match: old_file_path = os.path.join(paths["images"], old_fname_match.group(1)) if os.path.exists(old_file_path): os.remove(old_file_path) except Exception as e: print(f"[GridTester] โš ๏ธ Warning: Could not delete old file: {e}") existing_data["items"].pop(match_index) # ==== MODEL SWITCHING ==== target_model_name = conf["model"] target_model_key = get_model_cache_key(conf) if target_model_key != cached_model_key: if cached_model_key is not None: patched_model, patched_clip = cleanup_model_references( patched_model, patched_clip, conditioning_cache ) # CRITICAL FIX: Restore dict structure after clear conditioning_cache["positive"] = {} conditioning_cache["negative"] = {} loaded_model, loaded_clip, loaded_vae = load_model_by_type( conf, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, loaded_clip, loaded_vae, model_cache=model_cache ) # VAE fallback for non-checkpoint models if loaded_vae is None and optional_vae is not None: loaded_vae = optional_vae # Remember this model's default VAE for reverting later default_model_vae = loaded_vae cached_vae_key = "Default" current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = False cached_model_key = target_model_key cached_lora_key = None cached_lora_cache_key = None latent_channels = get_latent_channels(loaded_model, optional_latent) model_switched = True else: model_switched = False # ==== PER-CONFIG VAE SWITCHING ==== # Config Builder VAE settings take priority over sampler node's remote_vae_endpoint target_vae = conf.get("vae", "Default") if target_vae != cached_vae_key: # Flush pending batch before switching VAE (they need current VAE for decoding) if pending_batch: _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae) pending_batch = [] if target_vae == "Default": # Revert to model's bundled/default VAE loaded_vae = default_model_vae current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = False print(f"[GridTester] ๐ŸŽจ Reverting to Default VAE") elif is_remote_vae(target_vae): url = extract_remote_vae_url(target_vae) current_vae_is_remote = True current_remote_vae_url = url config_overrides_vae = True loaded_vae = None # No local VAE needed print(f"[GridTester] ๐ŸŒ Using per-config remote VAE: {url}") if url not in per_config_remote_workers: per_config_remote_workers[url] = RemoteVAEDecodeWorker( endpoint=url, img_dir=paths["images"], manifest_path=paths["manifest"], existing_data=existing_data, session_name=session_name, unique_id=unique_id ) else: current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = True print(f"[GridTester] ๐ŸŽจ Loading per-config VAE: {target_vae}") loaded_vae = load_vae_by_name(target_vae) cached_vae_key = target_vae # ==== LORA SWITCHING ==== current_lora_string = conf["lora_expanded"] if model_cache: current_cache_key = model_cache._get_lora_cache_key(target_model_key, current_lora_string) need_to_load = (current_cache_key != cached_lora_cache_key) or patched_model is None else: need_to_load = (current_lora_string != cached_lora_key) or patched_model is None if need_to_load: patched_model, patched_clip, should_skip = load_loras( loaded_model, loaded_clip, current_lora_string, target_model_key, incompatible_loras, model_cache=model_cache ) if should_skip: skipped_count += 1 continue cached_lora_key = current_lora_string if model_cache: cached_lora_cache_key = current_cache_key # Batch encode logic for multi-model switching # Because of short-circuit, if model_switched is True, conditioning_cache['positive'] isn't checked initially # But inside the loop, we access it. if model_switched or not conditioning_cache["positive"]: model_unique_positives = [] model_unique_negatives = [] _seen_positives = set() _seen_negatives = set() for future_idx in range(conf_idx, len(expanded)): future_conf = expanded[future_idx] if get_model_cache_key(future_conf) == target_model_key: future_positive, _ = build_prompt_with_triggers( future_conf, lora_triggerwords_mode ) # Skip encoding prompts for configs that are already completed # (avoids wasting time on prompts only used by skipped jobs) if not overwrite_existing and existing_data.get("items"): future_seed = future_conf["seed"] all_done = True for fj in input_jobs: if check_if_job_completed( existing_data["items"], future_conf, future_seed, fj["width"], fj["height"], fj["batch_idx"], future_positive, future_conf["negative"], has_optional_inputs=has_optional_inputs ) == -1: all_done = False break if all_done: continue if future_positive not in _seen_positives: _seen_positives.add(future_positive) model_unique_positives.append(future_positive) future_neg = future_conf["negative"] if future_neg not in _seen_negatives: _seen_negatives.add(future_neg) model_unique_negatives.append(future_neg) if model_unique_positives: print(f"[GridTester] ๐Ÿง  Batch encoding {len(model_unique_positives)} prompts for {target_model_name}") import comfy.model_management as mm_batch mm_batch.load_models_gpu([patched_clip.patcher]) clip_skip = conf.get("clip_skip", 0) try: for prompt in model_unique_positives: # Check for interrupt before each prompt encoding if mm.processing_interrupted(): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during positive encoding - Stopping all encoding") raise InterruptProcessingException() if prompt not in conditioning_cache["positive"]: conditioning_cache["positive"][prompt] = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) for prompt in model_unique_negatives: # Check for interrupt before each prompt encoding if mm.processing_interrupted(): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during negative encoding - Stopping all encoding") raise InterruptProcessingException() if prompt not in conditioning_cache["negative"]: conditioning_cache["negative"][prompt] = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) except InterruptProcessingException: print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during encoding - Stopping all jobs") if pending_batch: _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae) pending_batch = [] _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Re-raise so ComfyUI knows execution was interrupted (not completed) raise print(f"[GridTester] โœ… Encoded {len(conditioning_cache['positive'])} positive, {len(conditioning_cache['negative'])} negative") model_switched = False if cond_cache: cond_cache.set_lora_config(conf['lora_expanded']) # Get conditioning if optional_positive: final_positive = optional_positive else: full_positive = actual_positive_prompt final_positive = conditioning_cache["positive"].get(full_positive) if final_positive is None: cached_keys = list(conditioning_cache["positive"].keys()) cached_preview = [k[:40] for k in cached_keys[:5]] raise RuntimeError( f"[GridTester] โŒ BUG: Encoding not found for positive prompt: {full_positive[:80]}...\n" f" model_key={target_model_key}, lora={current_lora_string}\n" f" Cache has {len(cached_keys)} entries: {cached_preview}\n" f" model_switched={model_switched}, need_to_load={need_to_load}" ) if optional_negative: final_negative = optional_negative else: final_negative = conditioning_cache["negative"].get(conf["negative"]) if final_negative is None: cached_keys = list(conditioning_cache["negative"].keys()) cached_preview = [k[:40] for k in cached_keys[:5]] raise RuntimeError( f"[GridTester] โŒ BUG: Encoding not found for negative prompt: {conf['negative'][:80]}...\n" f" model_key={target_model_key}, lora={current_lora_string}\n" f" Cache has {len(cached_keys)} entries: {cached_preview}\n" f" model_switched={model_switched}, need_to_load={need_to_load}" ) # ========================================================= # ==== ๐Ÿš€ ASYNC LOOK-AHEAD PRE-FETCHING (CORRECTED) ==== # ========================================================= if model_cache is not None and model_cache.async_preload: current_overall_index = (job_idx * len(expanded)) + conf_idx # Check if next job exists if current_overall_index + 1 < total_jobs: next_idx = current_overall_index + 1 next_job_idx = next_idx // len(expanded) next_conf_idx = next_idx % len(expanded) if next_job_idx < len(input_jobs): next_conf = expanded[next_conf_idx] # RESOLVE cache keys for next conf next_model_key = get_model_cache_key(next_conf) next_lora = next_conf["lora_expanded"] # CASE 1: Same Base, Different LoRA if next_model_key == target_model_key and next_lora != current_lora_string: def _preload_lora_worker(): return load_loras( loaded_model, loaded_clip, next_lora, target_model_name, {}, model_cache=None ) model_cache.preload_lora_model(target_model_key, next_lora, _preload_lora_worker) # CASE 2: Different Base Model elif next_model_key != target_model_key: next_model_type = next_conf.get("model_type", "checkpoint") next_model_name = next_conf["model"] print(f"[GridTester] ๐Ÿ”ฎ Pre-loading Base Model: {next_model_name} ({next_model_type})") if next_model_type == "checkpoint": next_model_resolved = ckpt_name if next_model_name == "Default" else next_model_name def _preload_base_worker(path_to_load=next_model_resolved): try: import comfy.sd ckpt_path = folder_paths.get_full_path("checkpoints", path_to_load) if use_remote_vae: out = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=False, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings") ) return out[0], out[1], None else: out = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings") ) return out[0], out[1], out[2] except Exception as e: print(f"[GridTester] โŒ Async Worker Error: {e}") return None, None, None model_cache.preload_base_model(next_model_key, _preload_base_worker) else: # GGUF/diffusion model preload def _preload_diff_worker(conf_to_load=next_conf): try: return load_diffusion_model_and_clip( model_name=conf_to_load["model"], model_type=conf_to_load.get("model_type"), text_encoder_paths=conf_to_load.get("text_encoders", []), clip_type_str=conf_to_load.get("clip_type", "stable_diffusion"), gguf_options=conf_to_load.get("gguf_options"), use_remote_vae=use_remote_vae, ) except Exception as e: print(f"[GridTester] โŒ Async Worker Error: {e}") return None, None, None model_cache.preload_base_model(next_model_key, _preload_diff_worker) # ========================================================= # Generate image if job["latent"] is not None: latent_in = {"samples": job["latent"]["samples"].clone()} else: latent_in = {"samples": torch.zeros([1, latent_channels, h // 8, w // 8])} result_latent = None try: attention_mode = conf.get("attention_mode", "default") if attention_mode != "default": print(f"[GridTester] ๐Ÿง  Using attention mode: {attention_mode}") # Log active extra options extra_opts = [] if conf.get("model_sampling_override", "none") != "none": extra_opts.append(f"model_sampling={conf['model_sampling_override']}") if conf.get("use_advanced_sampling", False): extra_opts.append(f"advanced({conf.get('advanced_guider','cfg_guider')}/{conf.get('advanced_scheduler','basic')})") if conf.get("flux_guidance_value", 0) and float(conf.get("flux_guidance_value", 0)) > 0: extra_opts.append(f"flux_guidance={conf['flux_guidance_value']}") if extra_opts: print(f"[GridTester] โš™๏ธ Extra options: {', '.join(extra_opts)}") result_latent, duration = generate_image( patched_model, current_seed, conf["steps"], conf["cfg"], conf["sampler"], conf["scheduler"], final_positive, final_negative, latent_in, conf["denoise"], attention_mode=attention_mode, model_sampling_override=conf.get("model_sampling_override", "none"), model_sampling_shift=conf.get("model_sampling_shift", 1.73), model_sampling_flux_max_shift=conf.get("model_sampling_flux_max_shift", 1.15), model_sampling_flux_base_shift=conf.get("model_sampling_flux_base_shift", 0.5), use_advanced_sampling=conf.get("use_advanced_sampling", False), advanced_guider=conf.get("advanced_guider", "cfg_guider"), advanced_scheduler=conf.get("advanced_scheduler", "basic"), flux_guidance_value=conf.get("flux_guidance_value", 0.0), width=w, height=h ) # ==== UPSCALING (pipeline-based with sequential steps) ==== upscale_produced = False save_pre_upscale = False total_upscale_duration = 0 # Track cumulative upscale time for duration metric if session_settings and session_settings.get("upscaling", {}).get("enabled", False) and result_latent is not None: from .image_generation import upscale_image import itertools as upscale_itertools import random as upscale_random from PIL import Image as PILImage upscale_settings_rt = session_settings["upscaling"] upscale_pipelines = upscale_settings_rt.get("pipelines", []) upscale_combo_idx = 0 save_pre_upscale = upscale_settings_rt.get("save_pre_upscale", False) # If deferring upscales to end, skip inline processing and queue the job run_upscales_at_end = upscale_settings_rt.get("run_upscales_at_end", False) if run_upscales_at_end: deferred_upscale_queue.append({ "config": conf.copy(), "width": w, "height": h, "seed": current_seed, "batch_idx": batch_idx, "actual_positive_prompt": actual_positive_prompt, "actual_negative_prompt": actual_negative_prompt, "gen_index": gen_index_offset + total_generated, }) # Clear pipelines so the inline loop below becomes a no-op upscale_pipelines = [] # Compute HiRes-adjusted conditioning if enabled hires_positive_cond = final_positive # Default: same as base hires_prompt_active = False hires_prompt_behavior_rt = "" hires_prompt_text_rt = "" if upscale_settings_rt.get("hires_prompt_adjust") and upscale_settings_rt.get("hires_prompt_text", "").strip(): hires_prompt_behavior_rt = upscale_settings_rt.get("hires_prompt_behavior", "append_end") hires_prompt_text_rt = upscale_settings_rt["hires_prompt_text"].strip() if hires_prompt_behavior_rt == "prepend": adjusted_prompt = hires_prompt_text_rt + " " + actual_positive_prompt elif hires_prompt_behavior_rt == "append_end": adjusted_prompt = actual_positive_prompt + " " + hires_prompt_text_rt elif hires_prompt_behavior_rt == "replace": adjusted_prompt = hires_prompt_text_rt else: adjusted_prompt = actual_positive_prompt # Look up pre-encoded conditioning from cache hires_cond = conditioning_cache["positive"].get(adjusted_prompt) if hires_cond is not None: hires_positive_cond = hires_cond hires_prompt_active = True print(f"[GridTester] ๐Ÿ” Using HiRes-adjusted prompt ({hires_prompt_behavior_rt}): {adjusted_prompt[:80]}...") else: print(f"[GridTester] โš ๏ธ HiRes-adjusted prompt not found in cache, using original prompt") for pipeline_idx, pipeline in enumerate(upscale_pipelines): # Skip inactive pipelines (safety check โ€” should already be filtered) if pipeline.get("active", True) is False: continue pipeline_name = pipeline.get("name", f"Pipeline {pipeline_idx + 1}") pipeline_steps = pipeline.get("steps", []) if not pipeline_steps: continue # Each pipeline starts from the original base image pipe_latent = result_latent pipe_w = w pipe_h = h pipeline_history = [] # Accumulate per-step metadata for compound manifest entries # Flatten steps with repeat: a step with repeat=3 runs 3 times sequentially expanded_steps = [] for step in pipeline_steps: if step.get("active", True) is False: continue repeat = max(1, int(step.get("repeat", 1))) for _ in range(repeat): expanded_steps.append(step) for step_idx, ucfg in enumerate(expanded_steps): mode = ucfg.get("mode", "hires_only") show_hires = mode in ("hires_only", "model_then_hires") show_model = mode in ("model_only", "model_then_hires") # --- SeedVR2 upscale mode --- if mode == "seedvr2": from .image_generation import seedvr2_upscale from PIL import Image as PILImage sv_config = ucfg.get("seedvr2", {}) # Decode current latent to PIL for SeedVR2 input decoded = decode_latent_with_vae(pipe_latent, loaded_vae) pil_img = PILImage.fromarray((decoded[0].cpu().numpy() * 255).clip(0, 255).astype("uint8")) result_pil, up_w, up_h, sv_duration = seedvr2_upscale(pil_img, sv_config) total_upscale_duration += sv_duration is_last_step = step_idx == len(expanded_steps) - 1 is_final_output = is_last_step if is_final_output: # Save the upscaled image and update manifest upscale_id = int(time.time() * 100000) + random.randint(0, 1000) upscaled_filename = f"img_{upscale_id}.webp" filepath = os.path.join(paths["images"], upscaled_filename) result_pil.save(filepath, quality=80) upscaled_meta = dict(meta) upscaled_meta.update({ "id": upscale_id, "gen_index": len(existing_data["items"]), "file": f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images", "filename": upscaled_filename, "width": up_w, "height": up_h, "duration": sv_duration, "upscaled": True, "upscale_source": "inline", "upscale_pipeline": pipeline_name, "upscale_mode": "seedvr2", "upscale_model": sv_config.get("dit_model", ""), }) existing_data["items"].append(upscaled_meta) save_manifest(paths["manifest"], existing_data) if PromptServer is not None: try: PromptServer.instance.send_sync("ultimate_grid.update_data", { "session_name": session_name, "new_items": [upscaled_meta] }) except Exception: pass upscale_combo_idx += 1 else: # Re-encode for next step in chain import torch as seedvr_torch import numpy as seedvr_np result_np = seedvr_np.array(result_pil.convert("RGB")).astype(seedvr_np.float32) / 255.0 result_tensor = seedvr_torch.from_numpy(result_np).unsqueeze(0) pipe_latent = loaded_vae.encode(result_tensor[:, :, :, :3].movedim(-1, 1)) pipe_latent = {"samples": pipe_latent} pipe_w_current = up_w pipe_h_current = up_h continue # Skip the normal combo loop for this step # Parse multi-value fields raw_ratios = str(ucfg.get("upscale_ratios", "1.5")) ratios = [float(r.strip()) for r in raw_ratios.split(",") if r.strip()] or [1.5] raw_denoise = str(ucfg.get("hires_denoise", "0.3")) denoises = [float(d.strip()) for d in raw_denoise.split(",") if d.strip()] or [0.3] models = ucfg.get("upscale_models", []) or [""] # Build Cartesian product of multi-value fields if show_hires and show_model: combos = list(upscale_itertools.product(models, ratios, denoises)) elif show_hires: combos = list(upscale_itertools.product([""], ratios, denoises)) elif show_model: combos = list(upscale_itertools.product(models, [1.0], [0.0])) else: combos = [] for combo in combos: up_model_name, up_ratio, up_denoise = combo # Skip combos requiring a model when none is selected if show_model and not up_model_name: upscale_combo_idx += 1 continue # Build single-value upscaling config for upscale_image() # For model_only mode, use upscale_ratio for upscale_size so the # model's native output is preserved (not resized back to 1x) effective_size = up_ratio if up_ratio > 1.0 else 2.0 single_config = { "mode": mode, "upscale_ratio": up_ratio, "hires_denoise": up_denoise, "hires_steps": ucfg.get("hires_steps", 0), "tiled_vae": ucfg.get("tiled_vae", False), "tile_size": ucfg.get("tile_size", 512), "tile_overlap": ucfg.get("tile_overlap", 64), "temporal_size": ucfg.get("temporal_size", 512), "temporal_overlap": ucfg.get("temporal_overlap", 64), "upscale_model": up_model_name, "upscale_size": effective_size, "resize_method": ucfg.get("resize_method", "bilinear"), "hires_tiled_sampling": ucfg.get("hires_tiled_sampling", False), "hires_tile_width": ucfg.get("hires_tile_width", 512), "hires_tile_height": ucfg.get("hires_tile_height", 512), "hires_mask_blur": ucfg.get("hires_mask_blur", 8), "hires_tile_padding": ucfg.get("hires_tile_padding", 32), "hires_force_uniform_tiles": ucfg.get("hires_force_uniform_tiles", False) } # Steps within a pipeline chain sequentially (each feeds into the next) # Use HiRes-adjusted conditioning for modes that involve hires fix up_positive = hires_positive_cond if (show_hires and hires_prompt_active) else final_positive upscale_result, upscale_duration = upscale_image( pipe_latent, loaded_vae, patched_model, single_config, conf, up_positive, final_negative, pipe_w, pipe_h ) total_upscale_duration += upscale_duration # Determine if this is the final output (only the last step's last combo in the pipeline) is_last_step = step_idx == len(expanded_steps) - 1 is_last_combo = combo == combos[-1] is_final_output = is_last_step and is_last_combo # 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" 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"]) up_w, up_h = upscaled_pil.size # Only save to disk if this is a final output if is_final_output: upscaled_pil.save( os.path.join(paths["images"], upscaled_filename), "WEBP", quality=80 ) # Feed this latent to the next step in the pipeline pipe_latent = upscale_result pipe_w = up_w pipe_h = up_h elif isinstance(upscale_result, PILImage.Image): # Model-only returns PIL directly up_w, up_h = upscale_result.size # Only save to disk if this is a final output if is_final_output: upscale_result.save( os.path.join(paths["images"], upscaled_filename), "WEBP", quality=80 ) # Encode PIL back to latent for next step in pipeline import numpy as np_stack import torch as torch_stack pil_np = np_stack.array(upscale_result).astype(np_stack.float32) / 255.0 pil_tensor = torch_stack.from_numpy(pil_np).unsqueeze(0) pipe_latent = {"samples": loaded_vae.encode(pil_tensor[:, :, :, :3])} pipe_w = up_w pipe_h = up_h else: upscale_combo_idx += 1 continue # Collect this step's upscale info for pipeline history step_info = { "mode": mode, "ratio": up_ratio, "denoise": up_denoise, "model": up_model_name or "", "resize_method": ucfg.get("resize_method", "bilinear"), "hires_steps": ucfg.get("hires_steps", 0), "tiled_vae": ucfg.get("tiled_vae", False), "tile_size": ucfg.get("tile_size", 512), "tiled_sampling": ucfg.get("hires_tiled_sampling", False), "tile_w": ucfg.get("hires_tile_width", 512), "tile_h": ucfg.get("hires_tile_height", 512), } if mode == "model_only": step_info["upscale_size"] = ucfg.get("upscale_size", "2.0") pipeline_history.append(step_info) if is_final_output: upscaled_meta = create_image_metadata( conf, up_w, up_h, duration + total_upscale_duration, current_seed, batch_idx, actual_positive_prompt, actual_negative_prompt, gen_index=gen_index_offset + total_generated ) upscaled_meta["id"] = upscale_id upscaled_meta["upscaled"] = True upscaled_meta["upscale_pipeline"] = pipeline_name # Include HiRes prompt adjustment info in manifest if hires_prompt_active: upscaled_meta["hires_prompt_behavior"] = hires_prompt_behavior_rt upscaled_meta["hires_prompt_text"] = hires_prompt_text_rt # Multiple steps in pipeline: save arrays of all steps if len(pipeline_history) > 1: upscaled_meta["upscale_stacked"] = True upscaled_meta["upscale_mode"] = [s["mode"] for s in pipeline_history] upscaled_meta["upscale_ratio"] = [s["ratio"] for s in pipeline_history] upscaled_meta["upscale_denoise"] = [s["denoise"] for s in pipeline_history] models = [s["model"] for s in pipeline_history if s["model"]] if models: upscaled_meta["upscale_model"] = models upscaled_meta["upscale_resize_method"] = [s["resize_method"] for s in pipeline_history] upscaled_meta["upscale_hires_steps"] = [s["hires_steps"] for s in pipeline_history] # Tiling info: only include steps that have it enabled tiled_vae_steps = [s["tile_size"] for s in pipeline_history if s["tiled_vae"]] if tiled_vae_steps: upscaled_meta["upscale_tiled_vae"] = True upscaled_meta["upscale_tile_size"] = tiled_vae_steps tiled_sampling_steps = [f'{s["tile_w"]}x{s["tile_h"]}' for s in pipeline_history if s["tiled_sampling"]] if tiled_sampling_steps: upscaled_meta["upscale_tiled_sampling"] = True upscaled_meta["upscale_tile_w"] = [s["tile_w"] for s in pipeline_history if s["tiled_sampling"]] upscaled_meta["upscale_tile_h"] = [s["tile_h"] for s in pipeline_history if s["tiled_sampling"]] else: # Single step pipeline upscaled_meta["upscale_mode"] = mode upscaled_meta["upscale_ratio"] = up_ratio upscaled_meta["upscale_denoise"] = up_denoise if up_model_name: upscaled_meta["upscale_model"] = up_model_name upscaled_meta["upscale_resize_method"] = ucfg.get("resize_method", "bilinear") upscaled_meta["upscale_hires_steps"] = ucfg.get("hires_steps", 0) if ucfg.get("tiled_vae", False): upscaled_meta["upscale_tiled_vae"] = True upscaled_meta["upscale_tile_size"] = ucfg.get("tile_size", 512) if ucfg.get("hires_tiled_sampling", False): upscaled_meta["upscale_tiled_sampling"] = True upscaled_meta["upscale_tile_w"] = ucfg.get("hires_tile_width", 512) upscaled_meta["upscale_tile_h"] = ucfg.get("hires_tile_height", 512) if mode == "model_only": upscaled_meta["upscale_size"] = ucfg.get("upscale_size", "2.0") upscaled_meta["file"] = f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images" upscaled_meta["rejected"] = False # Insert at beginning (same order as flush_batch_with_vae) existing_data["items"].insert(0, upscaled_meta) upscale_produced = True # Save manifest immediately so dashboard can pick it up save_manifest(paths["manifest"], existing_data) # Send live update to dashboard (same as flush_batch_with_vae) if PromptServer is not None: try: manifest_meta = existing_data.get("meta", {}) PromptServer.instance.send_sync("ultimate_grid.update", { "node": unique_id, "session_name": session_name, "new_items": [upscaled_meta], "meta": manifest_meta }) except Exception: pass print(f"[GridTester] ๐Ÿ” Saved upscaled image: {upscaled_filename} " f"(pipeline={pipeline_name}, mode={mode}, ratio={up_ratio}, denoise={up_denoise}" f"{', model=' + up_model_name if up_model_name else ''})") else: print(f"[GridTester] ๐Ÿ” Pipeline '{pipeline_name}' intermediate step {step_idx + 1}/{len(expanded_steps)} " f"(mode={mode}, ratio={up_ratio}, denoise={up_denoise}" f"{', model=' + up_model_name if up_model_name else ''})") upscale_combo_idx += 1 # Save manifest after all pipeline combos if upscale_combo_idx > 0: save_manifest(paths["manifest"], existing_data) # Include upscale time in total duration for ETA accuracy total_duration = duration + total_upscale_duration job_durations.append(total_duration) eta_info = calculate_eta(job_durations, current_job, total_jobs) if eta_info: print_generation_progress(current_job, total_jobs, conf, w, h, total_duration, eta_info) # Send progress to dashboard frontend if PromptServer is not None: try: progress_pct = int((current_job / total_jobs) * 100) if eta_info['hours'] > 0: eta_str = f"{eta_info['hours']}h {eta_info['minutes']}m" elif eta_info['minutes'] > 0: eta_str = f"{eta_info['minutes']}m {eta_info['seconds']}s" else: eta_str = f"{eta_info['seconds']}s" PromptServer.instance.send_sync("ultimate_grid.progress", { "node": unique_id, "session_name": session_name, "current_job": current_job, "total_jobs": total_jobs, "progress_pct": progress_pct, "eta_str": eta_str, "finish_time": eta_info['finish_formatted'], "avg_duration": round(eta_info['avg_duration'], 1), "last_duration": round(total_duration, 1) }) except Exception: pass # Skip saving the base (non-upscaled) image when upscaling produced results # โ€” only the upscaled final version(s) should appear in the output # Unless save_pre_upscale is enabled, in which case save both if not upscale_produced or save_pre_upscale: meta = create_image_metadata( conf, w, h, total_duration, current_seed, batch_idx, actual_positive_prompt, actual_negative_prompt, gen_index=gen_index_offset + total_generated ) if pos_hash or neg_hash: meta["conditioning_pos_hash"] = pos_hash meta["conditioning_neg_hash"] = neg_hash pending_batch.append((result_latent["samples"].clone(), meta)) total_generated += 1 # ==== GPU COOLDOWN (if enabled in session settings) ==== if session_settings and session_settings.get("cooldown", {}).get("enabled", False): cooldown_config = session_settings["cooldown"] cooldown_every_n = int(cooldown_config.get("every_n", 1)) if total_generated > 0 and total_generated % cooldown_every_n == 0: cooldown_seconds = int(cooldown_config.get("seconds", 5)) clear_vram = cooldown_config.get("clear_vram", False) print(f"[GridTester] โ„๏ธ GPU Cooldown: pausing {cooldown_seconds}s after {total_generated} generations") if clear_vram: import comfy.model_management as mm_cooldown mm_cooldown.soft_empty_cache() mm_cooldown.unload_all_models() print(f"[GridTester] โ„๏ธ VRAM cleared") cached_model_key = None # Force model reload on next iteration import time as time_cooldown time_cooldown.sleep(cooldown_seconds) print(f"[GridTester] โ„๏ธ Cooldown complete, resuming generation") except Exception as e: import comfy.model_management if isinstance(e, InterruptProcessingException): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during generation - Stopping all jobs") if result_latent is not None: del result_latent result_latent = None if pending_batch: _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae) pending_batch = [] _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Re-raise so ComfyUI knows execution was interrupted (not completed) raise else: print(f"[GridTester] โŒ Generation failed: {e}") if result_latent is not None: del result_latent continue if result_latent is not None: del result_latent result_latent = None del latent_in latent_in = None if current_job % 10 == 0: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() threshold = vae_batch_size if flush_batch_every <= 0 else flush_batch_every if len(pending_batch) >= threshold: _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae) pending_batch = [] # ==== FINALIZATION ==== if pending_batch: _flush_pending_batch(pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae) # ==== DEFERRED UPSCALE PHASE (must run before cleanup so VAE/model/CLIP are still loaded) ==== if deferred_upscale_queue and session_settings and session_settings.get("upscaling", {}).get("run_upscales_at_end", False): upscale_settings_deferred = session_settings["upscaling"] run_deferred_upscales( deferred_upscale_queue, upscale_settings_deferred, session_settings, loaded_vae, patched_model, patched_clip, conditioning_cache, paths, existing_data, session_name, unique_id, PromptServer, config_overrides_vae=config_overrides_vae ) # Shut down per-config remote VAE workers _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: print(f"[GridTester] ๐ŸŒ Waiting for remote VAE...") remote_vae_worker.wait_completion() remote_vae_worker.stop() print_incompatible_loras_summary(incompatible_loras) if skipped_count > 0: print(f"[GridTester] โญ๏ธ Skipped {skipped_count} configs") # ==== FULL RUN SEED BEHAVIOR (POST-RUN) ==== # These modify the seed for the NEXT queue/run, not the current one. # The new seeds are stored in the manifest meta so the user can see/use them. next_run_seeds = {} for conf_idx_post, conf_post in enumerate(expanded): frb = conf_post.get("full_run_seed_behavior", "fixed") if frb == "random_after": import random new_seed = random.randint(0, 2**63 - 1) print(f"[GridTester] ๐ŸŽฒ Full run random_after: config {conf_idx_post} next seed โ†’ {new_seed}") conf_post["seed"] = new_seed next_run_seeds[str(conf_idx_post)] = new_seed elif frb == "increment_after": conf_post["seed"] = conf_post["seed"] + 1 print(f"[GridTester] โž• Full run increment_after: config {conf_idx_post} next seed โ†’ {conf_post['seed']}") next_run_seeds[str(conf_idx_post)] = conf_post["seed"] elif frb == "decrement_after": conf_post["seed"] = conf_post["seed"] - 1 print(f"[GridTester] โž– Full run decrement_after: config {conf_idx_post} next seed โ†’ {conf_post['seed']}") next_run_seeds[str(conf_idx_post)] = conf_post["seed"] # Persist next-run seeds to the node's seed widget via PromptServer if next_run_seeds and PromptServer is not None: try: PromptServer.instance.send_sync("ultimate_grid.next_run_seeds", { "node": unique_id, "seeds": next_run_seeds }) except Exception: pass existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) print(f"[GridTester] ๐Ÿงน Cleaning up...") loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() html = get_html_template(session_name, existing_data, unique_id) if job_durations: total_elapsed = time.time() - eta_start_time total_hours = int(total_elapsed // 3600) total_minutes = int((total_elapsed % 3600) // 60) total_seconds = int(total_elapsed % 60) avg_per_job = sum(job_durations) / len(job_durations) print(f"\n{'='*80}") print(f"[GridTester] ๐ŸŽ‰ COMPLETE!") print(f"[GridTester] โœ… {total_generated} images generated") print(f"[GridTester] โฑ๏ธ {total_hours}h {total_minutes}m {total_seconds}s total") print(f"[GridTester] ๐Ÿ“Š {avg_per_job:.1f}s average per job") print(f"{'='*80}\n") # Send completion event to dashboard if PromptServer is not None: try: PromptServer.instance.send_sync("ultimate_grid.progress", { "node": unique_id, "session_name": session_name, "current_job": total_jobs, "total_jobs": total_jobs, "progress_pct": 100, "eta_str": "Done", "finish_time": time.strftime("%H:%M:%S"), "avg_duration": round(avg_per_job, 1), "last_duration": 0, "complete": True, "total_elapsed": f"{total_hours}h {total_minutes}m {total_seconds}s", "total_generated": total_generated }) except Exception: pass return (html,) def _get_master_url(): """Detect this ComfyUI instance's URL for remote workers to connect to.""" import socket try: # Get local IP by connecting to external (doesn't actually send data) s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) s.connect(("8.8.8.8", 80)) local_ip = s.getsockname()[0] s.close() except Exception: local_ip = "127.0.0.1" port = 8188 try: import comfy.cli_args port = comfy.cli_args.args.port except Exception: pass return f"http://{local_ip}:{port}" # === Conditioning Serialization Helpers === # These convert ComfyUI conditioning tensors to/from JSON-serializable dicts # for transmitting pre-encoded conditionings from master to workers. # Pattern mirrors ConditioningCache._conditioning_to_dict/_dict_to_conditioning # but handles multi-entry conditionings (AND combinator creates multiple entries). def _conditioning_to_serializable(conditioning): """ Convert a ComfyUI conditioning to a JSON-serializable list of dicts. Args: conditioning: ComfyUI format [[cond_tensor, {"pooled_output": pooled_tensor}], ...] May have multiple entries when AND combinator is used. Returns: list of dicts, each with base64-encoded tensor data: [{"cond": {"data": ..., "shape": ..., "dtype": ...}, "pooled": {...}}, ...] """ import base64 import numpy as np entries = [] for cond_tensor, extra in conditioning: cond_np = cond_tensor.cpu().numpy() entry = { "cond": { "data": base64.b64encode(cond_np.tobytes()).decode('utf-8'), "shape": list(cond_np.shape), "dtype": str(cond_np.dtype) } } # Serialize ALL keys from the extra dict, not just pooled_output. # Models like Wan/T5 rely on attention_mask and other keys for proper sampling. if isinstance(extra, dict): extra_serialized = {} for key, val in extra.items(): if val is None: extra_serialized[key] = {"type": "none"} elif isinstance(val, torch.Tensor): val_np = val.cpu().numpy() extra_serialized[key] = { "type": "tensor", "data": base64.b64encode(val_np.tobytes()).decode('utf-8'), "shape": list(val_np.shape), "dtype": str(val_np.dtype) } else: # Scalar or other JSON-serializable value extra_serialized[key] = {"type": "value", "data": val} entry["extra"] = extra_serialized entries.append(entry) return entries def _serializable_to_conditioning(entries): """ Convert a serialized list of dicts back to ComfyUI conditioning format. Args: entries: list of dicts with base64-encoded tensor data (from _conditioning_to_serializable) Returns: ComfyUI conditioning format [[cond_tensor, {"pooled_output": pooled_tensor}], ...] """ import base64 import numpy as np conditioning = [] for entry in entries: # Reconstruct main conditioning tensor cond_data = entry["cond"] cond_bytes = base64.b64decode(cond_data["data"]) cond_np = np.frombuffer(cond_bytes, dtype=np.dtype(cond_data["dtype"])).reshape( tuple(cond_data["shape"]) ) # .copy() is required because np.frombuffer returns a read-only array, # and torch.from_numpy cannot send read-only arrays to GPU cond_tensor = torch.from_numpy(cond_np.copy()) # Reconstruct ALL extra dict keys (attention_mask, pooled_output, etc.) extra_dict = {} if "extra" in entry: for key, val_info in entry["extra"].items(): if val_info.get("type") == "none": extra_dict[key] = None elif val_info.get("type") == "tensor": val_bytes = base64.b64decode(val_info["data"]) val_np = np.frombuffer(val_bytes, dtype=np.dtype(val_info["dtype"])).reshape( tuple(val_info["shape"]) ) extra_dict[key] = torch.from_numpy(val_np.copy()) elif val_info.get("type") == "value": extra_dict[key] = val_info["data"] elif "pooled" in entry: # Backwards compatibility with old serialization format pooled_data = entry["pooled"] pooled_bytes = base64.b64decode(pooled_data["data"]) pooled_np = np.frombuffer(pooled_bytes, dtype=np.dtype(pooled_data["dtype"])).reshape( tuple(pooled_data["shape"]) ) extra_dict["pooled_output"] = torch.from_numpy(pooled_np.copy()) conditioning.append([cond_tensor, extra_dict]) return conditioning def _preencode_all_conditionings( expanded, lora_triggerwords_mode, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, model_cache ): """ Pre-encode all unique prompts across all (model, LoRA) combinations. Groups expanded configs by (model_key, lora_key), loads model+LoRA for each group, encodes all unique prompts, serializes the conditionings using base64 encoding, then cleans up. Returns: dict: { "model_key|lora_key": { "positive": {prompt_text: serialized_entries_list, ...}, "negative": {prompt_text: serialized_entries_list, ...} }, ... } Empty dict if encoding fails or no configs to encode. """ from collections import defaultdict from .model_loader import load_loras_for_preencoding # Group configs by (model_key, lora_key) groups = defaultdict(list) for conf in expanded: model_key = get_model_cache_key(conf) lora_key = conf.get("lora_expanded", conf.get("lora", "None")) group_key = f"{model_key}|{lora_key}" groups[group_key].append(conf) print(f"[Distribution] ๐Ÿง  Found {len(groups)} unique (model, LoRA) group(s) to pre-encode") encoded_all = {} for group_key, configs in groups.items(): parts = group_key.split("|", 1) lora_key = parts[1] if len(parts) > 1 else "None" first_conf = configs[0] print(f"[Distribution] ๐Ÿง  Pre-encoding group: model={first_conf['model']}, lora={lora_key}") # Collect unique prompts for this group unique_positives = set() unique_negatives = set() for conf in configs: try: actual_positive, _ = build_prompt_with_triggers(conf, lora_triggerwords_mode) except Exception: actual_positive = conf.get("positive", "") actual_negative = conf.get("negative", "") unique_positives.add(actual_positive) unique_negatives.add(actual_negative) print(f"[Distribution] ๐Ÿง  {len(unique_positives)} unique positive, " f"{len(unique_negatives)} unique negative prompt(s)") try: # Load model + CLIP for this group loaded_model, loaded_clip, loaded_vae = load_model_by_type( first_conf, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, None, None, model_cache=model_cache ) # Apply LoRA to get patched CLIP if lora_key != "None": patched_model, patched_clip = load_loras_for_preencoding( loaded_model, loaded_clip, lora_key ) else: patched_model, patched_clip = loaded_model, loaded_clip # Force CLIP onto GPU for encoding import comfy.model_management as mm_enc mm_enc.load_models_gpu([patched_clip.patcher]) clip_skip = first_conf.get("clip_skip", 0) serialized_positive = {} serialized_negative = {} for prompt in unique_positives: cond = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) serialized_positive[prompt] = _conditioning_to_serializable(cond) for prompt in unique_negatives: cond = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) serialized_negative[prompt] = _conditioning_to_serializable(cond) encoded_all[group_key] = { "positive": serialized_positive, "negative": serialized_negative } print(f"[Distribution] ๐Ÿง  Encoded {len(serialized_positive)} positive, " f"{len(serialized_negative)} negative conditioning(s)") # Cleanup patched references del patched_model, patched_clip except Exception as e: print(f"[Distribution] โš ๏ธ Pre-encoding failed for group {group_key}: {e}") import traceback traceback.print_exc() continue # Final cleanup gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() return encoded_all def _run_distributed_generation( self, distribution_config, expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode, session_name, paths, unique_id, total_jobs, seed, ckpt_name, positive_text, negative_text, vae_batch_size, configs_json, resolutions_json, flush_batch_every, use_remote_vae, remote_vae_endpoint, save_conditioning_cache_to_file, enable_model_cache, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, optional_latent, model_cache, session_settings=None ): """ Run generation in distributed mode. The master acts as both coordinator and local worker: - Creates a DistributionManager with the full job list - Notifies remote workers to start pulling jobs via API - Processes jobs locally by claiming from its own manager - Waits for all remote workers to finish - Returns the final HTML dashboard """ from .distribution_manager import DistributionManager, JobState from .distribution_routes import ( set_distribution_manager, notify_workers_to_start, stop_all_workers, _send_distribution_status ) worker_urls = distribution_config.get("worker_urls", []) claim_timeout = distribution_config.get("claim_timeout", 600) # Inject trigger word mode into configs so workers can use it for conf in expanded: conf["_lora_triggerwords_mode"] = lora_triggerwords_mode # === Create and populate distribution manager === manager = DistributionManager( session_name=session_name, paths=paths, unique_id=unique_id, existing_data=existing_data, claim_timeout_seconds=claim_timeout ) # Pass session-level settings (cooldown, etc.) to manager for worker passthrough if session_settings: manager._session_settings = session_settings manager.populate_jobs( expanded, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode ) # Make manager accessible to API endpoints set_distribution_manager(manager) _send_distribution_status(manager) # === Master pre-encoding phase (if enabled) === # When "Use Master Text Encoding" is on, pre-encode ALL unique prompts for all # (model, LoRA) combinations. Workers will receive these conditionings with their # job claims and skip CLIP encoding entirely. use_master_encoding = distribution_config.get("use_master_encoding", False) if use_master_encoding and not (optional_positive and optional_negative): print(f"[Distribution] ๐Ÿง  MASTER PRE-ENCODING: Encoding all unique prompts for workers...") encoded_conditionings = _preencode_all_conditionings( expanded, lora_triggerwords_mode, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, model_cache ) if encoded_conditionings: manager.set_encoded_conditionings(encoded_conditionings) print(f"[Distribution] ๐Ÿง  Pre-encoding complete, conditionings will be sent with job claims") else: print(f"[Distribution] โš ๏ธ Pre-encoding produced no results, workers will encode locally") elif use_master_encoding and (optional_positive and optional_negative): print(f"[Distribution] โš ๏ธ Master encoding disabled: optional conditioning inputs override prompts") # Early exit if all jobs were already completed (resume with nothing to do) initial_status = manager.get_status() if initial_status["total"] == 0: print(f"[Distribution] โญ๏ธ All jobs already completed, nothing to distribute") manager.deactivate() set_distribution_manager(None) existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) html = get_html_template(session_name, existing_data, unique_id) return (html,) # === Notify remote workers === master_url = _get_master_url() print(f"[Distribution] ๐ŸŒ Master URL: {master_url}") print(f"[Distribution] ๐ŸŒ Notifying {len(worker_urls)} worker(s)...") sync_models = distribution_config.get("sync_models_to_workers", False) if sync_models: manager.sync_models_to_workers = True print(f"[Distribution] โ˜๏ธ Model sync enabled โ€” workers will download missing models from master") worker_results = notify_workers_to_start(worker_urls, master_url, session_name, sync_models_to_workers=sync_models) successful_workers = sum(1 for _, ok, _ in worker_results if ok) print(f"[Distribution] โœ… {successful_workers}/{len(worker_urls)} workers started") # === Master local processing state === loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None cached_model_key = None cached_lora_key = None cached_vae_key = None default_model_vae = None conditioning_cache = {"positive": {}, "negative": {}} incompatible_loras = {} pending_batch = [] pending_job_ids = [] # Track job_ids for items in pending_batch total_generated = 0 gen_index_offset = len(existing_data.get("items", [])) job_durations = [] eta_start_time = time.time() master_jobs_processed = 0 # Remote VAE setup for master's local processing remote_vae_worker = None if use_remote_vae and expanded: remote_vae_worker = initialize_remote_vae( remote_vae_endpoint, paths["images"], paths["manifest"], existing_data, session_name, unique_id ) per_config_remote_workers = {} current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = False # True when config explicitly sets a non-Default VAE (overrides sampler node's remote_vae_endpoint) try: if PromptServer is not None: pbar = PromptServer.instance.progress_bar_pool.get_progress_bar(unique_id) else: pbar = None except: pbar = None dist_total = manager.get_status()["total"] # Ensure last_prompt_id exists for sampling callbacks (some ComfyUI versions # may not have it set when running inside the distributed code path) try: if PromptServer is not None and PromptServer.instance is not None: if not hasattr(PromptServer.instance, 'last_prompt_id'): PromptServer.instance.last_prompt_id = f"dist_master_{unique_id}" except Exception: pass # === Master local processing loop === # Master claims jobs from its own distribution manager and processes them # using the same model loading, encoding, and generation code paths as the # normal (non-distributed) generation loop. try: while True: # Check for interrupt try: import comfy.model_management as mm if mm.processing_interrupted(): print(f"\n[Distribution/Master] ๐Ÿ›‘ INTERRUPTED") raise InterruptProcessingException() except InterruptProcessingException: raise except: pass # Claim next job from manager batch = manager.claim_batch_for_local(1) if not batch: break # No more pending jobs for master claimed_job = batch[0] conf = claimed_job["config"] input_job = claimed_job["input_job"] job_id = claimed_job["job_id"] w = input_job["width"] h = input_job["height"] batch_idx = input_job.get("batch_idx", 0) current_seed = conf["seed"] if conf.get("seed_behavior") == "randomize": import random current_seed = random.randint(0, 2**63 - 1) master_jobs_processed += 1 completed_total = manager.total_completed print(f"[Distribution/Master] ๐Ÿ“Š Processing job {job_id} | " f"{conf['sampler']} @ {conf['steps']} steps | {w}x{h} | " f"Progress: {completed_total}/{dist_total}") if pbar: try: pbar.update_absolute(completed_total, dist_total) except: pass # Build prompts with trigger words actual_positive_prompt, _ = build_prompt_with_triggers(conf, lora_triggerwords_mode) actual_negative_prompt = conf["negative"] # --- Model Switching --- target_model_key = get_model_cache_key(conf) if target_model_key != cached_model_key: if cached_model_key is not None: patched_model, patched_clip = cleanup_model_references( patched_model, patched_clip, conditioning_cache ) conditioning_cache = {"positive": {}, "negative": {}} loaded_model, loaded_clip, loaded_vae = load_model_by_type( conf, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, loaded_clip, loaded_vae, model_cache=model_cache ) if loaded_vae is None and optional_vae is not None: loaded_vae = optional_vae default_model_vae = loaded_vae cached_model_key = target_model_key cached_lora_key = None cached_vae_key = "Default" current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = False # --- VAE Switching --- # Config Builder VAE settings take priority over sampler node's remote_vae_endpoint target_vae = conf.get("vae", "Default") if target_vae != cached_vae_key: # Flush pending batch before switching VAE if pending_batch: _flush_pending_batch( pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae ) for jid in pending_job_ids: manager.complete_job(jid) pending_batch = [] pending_job_ids = [] if target_vae == "Default": loaded_vae = default_model_vae current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = False elif is_remote_vae(target_vae): url = extract_remote_vae_url(target_vae) current_vae_is_remote = True current_remote_vae_url = url config_overrides_vae = True loaded_vae = None if url not in per_config_remote_workers: per_config_remote_workers[url] = RemoteVAEDecodeWorker( endpoint=url, img_dir=paths["images"], manifest_path=paths["manifest"], existing_data=existing_data, session_name=session_name, unique_id=unique_id ) else: current_vae_is_remote = False current_remote_vae_url = None config_overrides_vae = True loaded_vae = load_vae_by_name(target_vae) cached_vae_key = target_vae # --- LoRA Switching --- current_lora_string = conf.get("lora_expanded", conf.get("lora", "None")) if current_lora_string != cached_lora_key or patched_model is None: patched_model, patched_clip, should_skip = load_loras( loaded_model, loaded_clip, current_lora_string, target_model_key, incompatible_loras, model_cache=model_cache ) if should_skip: manager.fail_job(job_id, "LoRA incompatible") continue cached_lora_key = current_lora_string # --- Prompt Encoding --- if optional_positive: final_positive = optional_positive else: if actual_positive_prompt not in conditioning_cache["positive"]: clip_skip = conf.get("clip_skip", 0) conditioning_cache["positive"][actual_positive_prompt] = \ encode_prompt_with_combinators(patched_clip, actual_positive_prompt, clip_skip) final_positive = conditioning_cache["positive"][actual_positive_prompt] if optional_negative: final_negative = optional_negative else: if actual_negative_prompt not in conditioning_cache["negative"]: clip_skip = conf.get("clip_skip", 0) conditioning_cache["negative"][actual_negative_prompt] = \ encode_prompt_with_combinators(patched_clip, actual_negative_prompt, clip_skip) final_negative = conditioning_cache["negative"][actual_negative_prompt] # --- Create Latent --- latent_channels = get_latent_channels(loaded_model, optional_latent) if optional_latent is not None: latent_in = {"samples": optional_latent["samples"].clone()} else: latent_in = {"samples": torch.zeros([1, latent_channels, h // 8, w // 8])} # --- Generate Image --- try: attention_mode = conf.get("attention_mode", "default") # Log active extra options extra_opts = [] if conf.get("model_sampling_override", "none") != "none": extra_opts.append(f"model_sampling={conf['model_sampling_override']}") if conf.get("use_advanced_sampling", False): extra_opts.append(f"advanced({conf.get('advanced_guider','cfg_guider')}/{conf.get('advanced_scheduler','basic')})") if conf.get("flux_guidance_value", 0) and float(conf.get("flux_guidance_value", 0)) > 0: extra_opts.append(f"flux_guidance={conf['flux_guidance_value']}") if extra_opts: print(f"[GridTester] โš™๏ธ Extra options: {', '.join(extra_opts)}") result_latent, duration = generate_image( patched_model, current_seed, conf["steps"], conf["cfg"], conf["sampler"], conf["scheduler"], final_positive, final_negative, latent_in, conf["denoise"], attention_mode=attention_mode, model_sampling_override=conf.get("model_sampling_override", "none"), model_sampling_shift=conf.get("model_sampling_shift", 1.73), model_sampling_flux_max_shift=conf.get("model_sampling_flux_max_shift", 1.15), model_sampling_flux_base_shift=conf.get("model_sampling_flux_base_shift", 0.5), use_advanced_sampling=conf.get("use_advanced_sampling", False), advanced_guider=conf.get("advanced_guider", "cfg_guider"), advanced_scheduler=conf.get("advanced_scheduler", "basic"), flux_guidance_value=conf.get("flux_guidance_value", 0.0), width=w, height=h ) job_durations.append(duration) eta_info = calculate_eta(job_durations, completed_total + 1, dist_total) if eta_info: print_generation_progress( completed_total + 1, dist_total, conf, w, h, duration, eta_info ) meta = create_image_metadata( conf, w, h, duration, current_seed, batch_idx, actual_positive_prompt, actual_negative_prompt, gen_index=gen_index_offset + total_generated ) pending_batch.append((result_latent["samples"].clone(), meta)) pending_job_ids.append(job_id) total_generated += 1 del result_latent, latent_in except InterruptProcessingException: raise except Exception as e: print(f"[Distribution/Master] โŒ Generation failed: {e}") manager.fail_job(job_id, str(e)) del latent_in continue # --- Flush batch at threshold --- threshold = vae_batch_size if flush_batch_every <= 0 else flush_batch_every if len(pending_batch) >= threshold: _flush_pending_batch( pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae ) for jid in pending_job_ids: manager.complete_job(jid) pending_batch = [] pending_job_ids = [] _send_distribution_status(manager) # Periodic cleanup if master_jobs_processed % 10 == 0: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Flush remaining master batch if pending_batch: _flush_pending_batch( pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae ) for jid in pending_job_ids: manager.complete_job(jid) pending_batch = [] pending_job_ids = [] _send_distribution_status(manager) except InterruptProcessingException: # Flush any remaining batch before cleanup if pending_batch: _flush_pending_batch( pending_batch, current_vae_is_remote, current_remote_vae_url, per_config_remote_workers, use_remote_vae, remote_vae_worker, loaded_vae, paths, existing_data, session_name, paths["manifest"], unique_id, config_overrides_vae=config_overrides_vae ) for jid in pending_job_ids: manager.complete_job(jid) # Stop remote workers and deactivate manager stop_all_workers(worker_urls) manager.deactivate() set_distribution_manager(None) _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Re-raise so ComfyUI knows execution was interrupted (not completed) raise # === Wait for remote workers to finish === # Max wait timeout prevents infinite loops if workers crash without cleanup. # After timeout, remaining claimed jobs are reclaimed and processed on master. dist_wait_timeout = distribution_config.get("wait_timeout", 600) # Default 10 minutes if manager.has_pending_or_claimed: print(f"[Distribution] โณ Waiting for remote workers to finish (timeout: {dist_wait_timeout}s)...") wait_start = time.time() while manager.has_pending_or_claimed: # Release timed-out jobs from dead workers (was only called during claim_job) manager.release_timed_out_jobs() # Check max wait timeout โ€” reclaim remaining jobs for master processing elapsed_wait = time.time() - wait_start if elapsed_wait > dist_wait_timeout: status = manager.get_status() remaining = status['pending'] + status['claimed'] print(f"[Distribution] โฐ Wait timeout ({dist_wait_timeout}s) reached with {remaining} jobs remaining โ€” reclaiming for master") # Force-reclaim all remaining claimed jobs with manager._lock: for job in manager._jobs.values(): if job.state == JobState.CLAIMED and job.worker_id != "master": job.state = JobState.PENDING job.worker_id = None job.claimed_at = None manager._pending_queue.append(job.job_id) break # Check for interrupt while waiting try: import comfy.model_management as mm if mm.processing_interrupted(): print(f"\n[Distribution] ๐Ÿ›‘ INTERRUPTED while waiting for workers") stop_all_workers(worker_urls) manager.deactivate() set_distribution_manager(None) _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() raise InterruptProcessingException() except InterruptProcessingException: raise except: pass time.sleep(5) _send_distribution_status(manager) status = manager.get_status() print(f"[Distribution] โณ Pending: {status['pending']}, " f"Claimed: {status['claimed']}, " f"Completed: {status['completed']}/{status['total']}") # === Handle remaining jobs after wait timeout === if manager.has_pending_or_claimed: status = manager.get_status() remaining = status['pending'] + status['claimed'] print(f"[Distribution] โš ๏ธ {remaining} job(s) still incomplete after wait timeout โ€” marking as failed and proceeding to finalization") # Mark remaining jobs as failed so we don't block forever with manager._lock: for job in manager._jobs.values(): if job.state in (JobState.PENDING, JobState.CLAIMED): job.state = JobState.COMPLETED # Mark done to unblock print(f"[Distribution] โš ๏ธ Abandoned job {job.job_id} (was {job.state})") # === Finalization === stop_all_workers(worker_urls) manager.deactivate() # Don't set_distribution_manager(None) immediately โ€” a slow worker may still # be in the middle of submitting a result. The submit_result route only needs # manager != None (it does NOT check is_active). Deactivation already prevents # new job claims (claim_job returns 503). Delay the cleanup with a daemon thread # that waits for workers to go silent before clearing the manager. def _delayed_manager_cleanup(): # Wait up to 5 minutes for all workers to stop submitting. # Check every 10s whether any registered workers still have recent heartbeats. for _ in range(30): # 30 ร— 10s = 5 minutes max time.sleep(10) if not manager or not manager._workers: break now = time.time() any_alive = False with manager._lock: for w in manager._workers.values(): if now - w.get("last_heartbeat", 0) < 30: any_alive = True break if not any_alive: break set_distribution_manager(None) threading.Thread(target=_delayed_manager_cleanup, daemon=True).start() _cleanup_per_config_remote_workers(per_config_remote_workers) if remote_vae_worker: print(f"[Distribution] ๐ŸŒ Waiting for remote VAE...") remote_vae_worker.wait_completion() remote_vae_worker.stop() print_incompatible_loras_summary(incompatible_loras) # === Run deferred upscales on master after all distributed generation is complete === if session_settings and session_settings.get("upscaling", {}).get("enabled", False): upscale_settings = session_settings.get("upscaling", {}) print(f"[Distribution] ๐Ÿ” Running upscales on master for distributed results...") # Build deferred upscale queue from all generated items deferred_upscale_queue = [] for item_meta in existing_data.get("items", []): if not item_meta.get("upscaled"): # Skip items already upscaled deferred_upscale_queue.append({ "meta": item_meta, "config": { "model": item_meta.get("model", ckpt_name), "steps": item_meta.get("steps", 20), "sampler": item_meta.get("sampler", "euler"), "scheduler": item_meta.get("scheduler", "normal"), "cfg": item_meta.get("cfg", 7.0), } }) if deferred_upscale_queue: run_deferred_upscales( deferred_upscale_queue, upscale_settings, existing_data, session_name, paths, unique_id, ckpt_name, conditioning_cache, model_cache, use_remote_vae, remote_vae_endpoint ) else: print(f"[Distribution] โ„น๏ธ No items to upscale") existing_data["meta"] = { "positive": positive_text, "negative": negative_text, "model": ckpt_name, "seed": seed, "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } save_manifest(paths["manifest"], existing_data) print(f"[Distribution] ๐Ÿงน Cleaning up...") loaded_model, loaded_clip, loaded_vae = None, None, None patched_model, patched_clip = None, None conditioning_cache.clear() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() html = get_html_template(session_name, existing_data, unique_id) final_status = manager.get_status() if job_durations: total_elapsed = time.time() - eta_start_time total_hours = int(total_elapsed // 3600) total_minutes = int((total_elapsed % 3600) // 60) total_seconds = int(total_elapsed % 60) avg_per_job = sum(job_durations) / len(job_durations) print(f"\n{'='*80}") print(f"[Distribution] ๐ŸŽ‰ DISTRIBUTED GENERATION COMPLETE!") print(f"[Distribution] โœ… {final_status.get('completed', total_generated)} total images generated") print(f"[Distribution] ๐Ÿ“Š Master processed: {total_generated} images ({avg_per_job:.1f}s avg)") print(f"[Distribution] โฑ๏ธ {total_hours}h {total_minutes}m {total_seconds}s total") if final_status.get('failed', 0) > 0: print(f"[Distribution] โŒ {final_status['failed']} jobs failed permanently") # Per-worker breakdown workers = final_status.get("workers", {}) if workers: print(f"[Distribution] ๐Ÿ‘ฅ Worker breakdown:") print(f"[Distribution] ๐Ÿ“Œ master: {total_generated} jobs completed") for wid, winfo in workers.items(): completed = winfo.get("jobs_completed", 0) failed = winfo.get("jobs_failed", 0) duplicated = winfo.get("jobs_duplicated", 0) parts = [] if failed > 0: parts.append(f"{failed} failed") if duplicated > 0: parts.append(f"{duplicated} late/duplicate") extra_str = f" ({', '.join(parts)})" if parts else "" print(f"[Distribution] ๐Ÿ“Œ {wid}: {completed} jobs completed{extra_str}") print(f"{'='*80}\n") # Send completion event to dashboard if PromptServer is not None: try: PromptServer.instance.send_sync("ultimate_grid.progress", { "node": unique_id, "session_name": session_name, "current_job": dist_total, "total_jobs": dist_total, "progress_pct": 100, "eta_str": "Done", "finish_time": time.strftime("%H:%M:%S"), "avg_duration": round(avg_per_job, 1), "last_duration": 0, "complete": True, "total_elapsed": f"{total_hours}h {total_minutes}m {total_seconds}s", "total_generated": final_status.get("completed", total_generated) }) except Exception: pass return (html,)