""" Generation Orchestrator - Main Entry Point Coordinates the entire grid generation workflow """ import os import json import time import re import gc import torch import hashlib import folder_paths from .trigger_words import collect_unique_prompts_with_triggers, build_prompt_with_triggers from .batch_encoding import batch_encode_prompts 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 ) 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, 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. Returns: dict: Paths dictionary with 'base', 'images', 'manifest' """ # base_dir = os.path.join("benchmarks", session_name) # img_dir = os.path.join(base_dir, "images") 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. Args: remote_vae_endpoint: Remote VAE endpoint URL OR model type (SD/SDXL/Flux/HunyuanVideo) img_dir: Image output directory manifest_path: Path to manifest file existing_data: Existing manifest data session_name: Session name unique_id: Unique ID for this node Returns: RemoteVAEDecodeWorker or None """ if not remote_vae_endpoint or remote_vae_endpoint == "None": return None # ==== FIX: Convert model type to actual endpoint URL ==== # If user selected a model type (SD/SDXL/Flux/HunyuanVideo), convert to URL 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)": # For Auto mode, we can't initialize worker yet - return None # Worker will be initialized lazily in flush_batch when model type is detected print(f"[GridTester] ๐ŸŒ Auto mode selected - worker will initialize on first flush") return None else: # Assume it's a direct URL actual_endpoint = remote_vae_endpoint print(f"[GridTester] ๐ŸŒ Using custom endpoint: {actual_endpoint}") # ==== Create worker with actual endpoint URL ==== worker = RemoteVAEDecodeWorker( endpoint=actual_endpoint, # โ† FIXED: Use resolved URL, not model type string 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 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 run_generation_loop( node_instance, 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, lookup_and_append_lora_triggerwords, remote_vae_endpoint, optional_model=None, optional_clip=None, optional_vae=None, optional_positive=None, optional_negative=None, optional_latent=None ): """ Main generation loop orchestrator. Coordinates: - Session setup - Config expansion - Model/LoRA loading - Batch encoding - Image generation - Progress tracking - Manifest management Returns: tuple: (html_dashboard,) """ # ==== 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 # Parse configs and expand 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['model'], x['lora'], x['positive'], x['negative'])) # ==== EXPAND LORA FOLDERS AND RANDOM SELECTIONS ==== # Process all LoRA random syntax ONCE before generation loop # This ensures consistent LoRA selection and allows trigger words to be included 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() # Check if this part uses random syntax: "folder/[count,strength]" or "folder/[count,strength,random]" if "[" in part and "]" in part: # Expand using the config's seed for reproducibility (unless 'random' keyword is used) expanded_lora = expand_lora_folder(part, seed=conf.get("seed")) if expanded_lora: # expand_lora_folder returns a string or list if isinstance(expanded_lora, list): expanded_parts.extend(expanded_lora) else: expanded_parts.append(expanded_lora) else: # Regular LoRA path - keep as-is expanded_parts.append(part) # Store expanded version in config conf["lora_expanded"] = " + ".join(expanded_parts) if expanded_parts else "None" # Also update the original lora field so trigger word lookup uses expanded version conf["lora"] = conf["lora_expanded"] else: conf["lora_expanded"] = "None" print(f"[GridTester] โœ… LoRA expansion complete") input_jobs = prepare_input_jobs(optional_latent, resolutions) total_jobs = len(expanded) * len(input_jobs) print(f"\n{'='*80}") print(f"[GridTester] ๐Ÿš€ GENERATION START") print(f"[GridTester] ๐Ÿ“‹ {len(expanded)} configs ร— {len(input_jobs)} resolutions = {total_jobs} total jobs") print(f"{'='*80}\n") # ==== 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 MATCH_KEYS = [ "sampler", "scheduler", "steps", "cfg", "lora", "denoise", "seed", "width", "height", "batch_idx", "model", "conditioning_pos_hash", "conditioning_neg_hash" ] else: pos_hash, neg_hash = None, None MATCH_KEYS = [ "sampler", "scheduler", "steps", "cfg", "lora", "denoise", "seed", "width", "height", "positive", "negative", "batch_idx", "model" ] # ==== REMOTE VAE SETUP ==== use_remote_vae = remote_vae_endpoint and remote_vae_endpoint != "None" # ==== PROGRESS BAR ==== 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 conditioning_cache = {"positive": {}, "negative": {}} incompatible_loras = {} pending_batch = [] current_job = 0 total_generated = 0 skipped_count = 0 job_durations = [] eta_start_time = time.time() # Initialize remote VAE worker if needed (before pre-encoding stage) # This needs to happen regardless of pre-encoding status 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 ) if remote_vae_worker: print(f"[GridTester] ๐ŸŒ Remote VAE initialized") # ==== PRE-ENCODING STAGE ==== # CRITICAL: Only pre-encode if all configs use the SAME model # Different models have different CLIP models, so we can't reuse encodings unique_models = set(conf["model"] for conf in expanded) if not (optional_positive and optional_negative) and expanded and len(unique_models) == 1: first_conf = expanded[0] target_model_name = first_conf["model"] print(f"[GridTester] โœ… Single model detected ({target_model_name}) - enabling pre-encoding") # Load model for pre-encoding loaded_model, loaded_clip, loaded_vae = load_checkpoint( target_model_name, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, None, None ) # Load LoRAs if needed (use expanded version) 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 # Initialize conditioning cache clip_hash = calculate_clip_hash(patched_clip) cond_cache = ConditioningCache(paths["base"], clip_hash) # Collect unique prompts with triggers print(f"[GridTester] ๐Ÿง  Collecting unique prompts...") unique_positives, unique_negatives = collect_unique_prompts_with_triggers( expanded, lookup_and_append_lora_triggerwords ) # Batch encode all prompts with clip_skip from first config clip_skip = first_conf.get("clip_skip", 0) if clip_skip != 0: print(f"[GridTester] ๐Ÿ”ง Using clip_skip={clip_skip}") conditioning_cache = batch_encode_prompts( patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip ) cached_model_key = target_model_name cached_lora_key = first_conf["lora_expanded"] latent_channels = get_latent_channels(loaded_model, optional_latent) elif len(unique_models) > 1: print(f"[GridTester] โš ๏ธ Multiple models detected ({len(unique_models)} different models) - pre-encoding DISABLED") print(f"[GridTester] โ„น๏ธ Each model has a different CLIP - encoding will happen per-generation") print(f"[GridTester] โ„น๏ธ This is slower but ensures correct CLIP encodings for each model") cond_cache = None latent_channels = 4 else: print(f"[GridTester] โ„น๏ธ Using optional conditioning, skipping pre-encoding") cond_cache = None latent_channels = 4 # ==== MAIN GENERATION LOOP ==== print(f"\n{'='*80}\n") for job in input_jobs: w, h = job["width"], job["height"] batch_idx = job["batch_idx"] for conf_idx, conf in enumerate(expanded): # ==== CHECK FOR INTERRUPT AT START OF EACH ITERATION ==== try: import comfy.model_management as mm if mm.processing_interrupted(): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED - Stopping all jobs") print(f"[GridTester] โœ… Completed {total_generated}/{total_jobs} images before interrupt") # Flush any pending images if pending_batch: if use_remote_vae: 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, paths["manifest"], unique_id) pending_batch = [] # Wait for remote VAE if active if remote_vae_worker: print(f"[GridTester] ๐ŸŒ Waiting for remote VAE...") remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, # Base positive prompt "negative": negative_text, # Base negative prompt "model": ckpt_name, # Base model "seed": seed, # Base seed "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } # Save manifest save_manifest(paths["manifest"], existing_data) # Cleanup 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() # Generate HTML with what we have html = get_html_template(session_name, existing_data, unique_id) return (html,) except: pass # If interrupt checking fails, continue normally current_seed = conf["seed"] current_job += 1 # Update progress bar if pbar: try: pbar.update_absolute(current_job, total_jobs) except: pass # Console progress 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}") # Build prompt with triggers actual_positive_prompt, lora_triggers = build_prompt_with_triggers( conf, lookup_and_append_lora_triggerwords ) actual_negative_prompt = conf["negative"] # Check for existing match (overwrite mode) if overwrite_existing: match_index = node_instance.find_existing_match( existing_data["items"], conf, w, h, current_seed, batch_idx, MATCH_KEYS ) if match_index != -1: 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: pass existing_data["items"].pop(match_index) continue # Skip generation, we're overwriting # Load model if switching target_model_name = conf["model"] if target_model_name != cached_model_key: if cached_model_key is not None: patched_model, patched_clip = cleanup_model_references( patched_model, patched_clip, conditioning_cache ) loaded_model, loaded_clip, loaded_vae = load_checkpoint( target_model_name, ckpt_name, use_remote_vae, optional_model, optional_clip, optional_vae, optional_positive, optional_negative, loaded_clip, loaded_vae ) cached_model_key = target_model_name cached_lora_key = None latent_channels = get_latent_channels(loaded_model, optional_latent) # Flag to trigger batch encoding after LoRAs are loaded model_switched = True else: model_switched = False # Load LoRAs if switching (use expanded version) current_lora_key = conf["lora_expanded"] if current_lora_key != cached_lora_key or patched_model is None: patched_model, patched_clip, should_skip = load_loras( loaded_model, loaded_clip, conf["lora_expanded"], target_model_name, incompatible_loras ) if should_skip: skipped_count += 1 continue cached_lora_key = current_lora_key # ==== FIX: Batch encode for THIS model when switching models/LoRAs ==== # Even in multi-model mode, batch encode all prompts for the current # model to avoid encoding during the hot path (which causes CLIP thrashing) if model_switched or not conditioning_cache["positive"]: # Collect all unique prompts for THIS model from remaining configs model_unique_positives = set() model_unique_negatives = set() # Look ahead at remaining configs for this model for future_idx in range(conf_idx, len(expanded)): future_conf = expanded[future_idx] if future_conf["model"] == target_model_name: # Build prompt with triggers future_positive, _ = build_prompt_with_triggers( future_conf, lookup_and_append_lora_triggerwords ) model_unique_positives.add(future_positive) model_unique_negatives.add(future_conf["negative"]) if model_unique_positives: print(f"[GridTester] ๐Ÿง  Batch encoding {len(model_unique_positives)} prompts for {target_model_name}") # Keep CLIP in VRAM during batch encoding import comfy.model_management as mm_batch mm_batch.load_models_gpu([patched_clip.patcher], force_patch_weights=True) # Get clip_skip from current config clip_skip = conf.get("clip_skip", 0) # Batch encode all prompts for this model for prompt in model_unique_positives: if prompt not in conditioning_cache["positive"]: # Apply clip_skip if needed original_layer = None if clip_skip != 0 and hasattr(patched_clip.cond_stage_model, 'clip_layer'): original_layer = patched_clip.cond_stage_model.clip_layer patched_clip.cond_stage_model.set_clip_options({"layer": clip_skip}) tokens = patched_clip.tokenize(prompt) cond, pooled = patched_clip.encode_from_tokens(tokens, return_pooled=True) conditioning_cache["positive"][prompt] = [[cond, {"pooled_output": pooled}]] # Restore layer if original_layer is not None: patched_clip.cond_stage_model.set_clip_options({"layer": original_layer}) for prompt in model_unique_negatives: if prompt not in conditioning_cache["negative"]: # Apply clip_skip if needed original_layer = None if clip_skip != 0 and hasattr(patched_clip.cond_stage_model, 'clip_layer'): original_layer = patched_clip.cond_stage_model.clip_layer patched_clip.cond_stage_model.set_clip_options({"layer": clip_skip}) tokens = patched_clip.tokenize(prompt) cond, pooled = patched_clip.encode_from_tokens(tokens, return_pooled=True) conditioning_cache["negative"][prompt] = [[cond, {"pooled_output": pooled}]] # Restore layer if original_layer is not None: patched_clip.cond_stage_model.set_clip_options({"layer": original_layer}) 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 (should always be in cache after batch encoding) 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: # This should never happen since we batch encoded all prompts raise RuntimeError(f"[GridTester] โŒ BUG: Encoding not found for: {full_positive[:50]}...") if optional_negative: final_negative = optional_negative else: final_negative = conditioning_cache["negative"].get(conf["negative"]) if final_negative is None: # This should never happen since we batch encoded all prompts raise RuntimeError(f"[GridTester] โŒ BUG: Encoding not found for: {conf['negative'][:50]}...") # 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])} # ==== GENERATION WITH PROPER INTERRUPT HANDLING ==== result_latent = None try: 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"] ) # Track ETA job_durations.append(duration) eta_info = calculate_eta(job_durations, current_job, total_jobs) if eta_info: print_generation_progress(current_job, total_jobs, conf, w, h, duration, eta_info) # Create metadata meta = create_image_metadata( conf, w, h, duration, current_seed, batch_idx, actual_positive_prompt, actual_negative_prompt ) if pos_hash or neg_hash: meta["conditioning_pos_hash"] = pos_hash meta["conditioning_neg_hash"] = neg_hash # Clone tensor to break reference chain pending_batch.append((result_latent["samples"].clone(), meta)) total_generated += 1 # ==== FIX: CATCH INTERRUPT SEPARATELY ==== except Exception as e: # Check if this is an interrupt exception import comfy.model_management if isinstance(e, comfy.model_management.InterruptProcessingException): print(f"\n[GridTester] ๐Ÿ›‘ INTERRUPTED during generation - Stopping all jobs") # Clean up current generation if result_latent is not None: del result_latent result_latent = None # Flush pending batch if pending_batch: if use_remote_vae: 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, paths["manifest"], unique_id) pending_batch = [] # Wait for remote VAE if remote_vae_worker: print(f"[GridTester] ๐ŸŒ Waiting for remote VAE...") remote_vae_worker.wait_completion() remote_vae_worker.stop() existing_data["meta"] = { "positive": positive_text, # Base positive prompt "negative": negative_text, # Base negative prompt "model": ckpt_name, # Base model "seed": seed, # Base seed "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } # Save manifest save_manifest(paths["manifest"], existing_data) # Cleanup all models 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() # Generate HTML and return - this stops everything html = get_html_template(session_name, existing_data, unique_id) return (html,) else: # Regular error - just skip this config print(f"[GridTester] โŒ Generation failed: {e}") if result_latent is not None: del result_latent continue # ==== END FIX ==== # Clean up after each generation if result_latent is not None: del result_latent result_latent = None del latent_in latent_in = None # Periodic garbage collection if current_job % 10 == 0: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Flush batch if needed threshold = vae_batch_size if flush_batch_every <= 0 else flush_batch_every if len(pending_batch) >= threshold: if use_remote_vae: 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, paths["manifest"], unique_id) pending_batch = [] # ==== FINALIZATION ==== # Flush remaining if pending_batch: if use_remote_vae: 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, paths["manifest"], unique_id) # Wait for remote VAE if remote_vae_worker: print(f"[GridTester] ๐ŸŒ Waiting for remote VAE...") remote_vae_worker.wait_completion() remote_vae_worker.stop() # Print summaries print_incompatible_loras_summary(incompatible_loras) if skipped_count > 0: print(f"[GridTester] โญ๏ธ Skipped {skipped_count} configs") existing_data["meta"] = { "positive": positive_text, # Base positive prompt "negative": negative_text, # Base negative prompt "model": ckpt_name, # Base model "seed": seed, # Base seed "vae_batch_size": vae_batch_size, "configs_json": configs_json, "resolutions_json": resolutions_json } # Save manifest save_manifest(paths["manifest"], existing_data) # Cleanup 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() # Generate HTML html = get_html_template(session_name, existing_data, unique_id) # Final summary 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") return (html,)