image_generation.py

"""
Image Generation and Sampling Module
Handles the core image generation, decoding, and batch management
"""

import time
import os
import re
import json
import torch
import nodes
import numpy as np
from PIL import Image


def generate_image(
    patched_model,
    seed,
    steps,
    cfg,
    sampler_name,
    scheduler,
    positive_conditioning,
    negative_conditioning,
    latent_input,
    denoise,
    attention_mode="default",
    model_sampling_override="none",
    model_sampling_shift=1.73,
    model_sampling_flux_max_shift=1.15,
    model_sampling_flux_base_shift=0.5,
    use_advanced_sampling=False,
    advanced_guider="cfg_guider",
    advanced_scheduler="basic",
    flux_guidance_value=0.0,
    width=1024,
    height=1024
):
    """
    Generate a single image using ComfyUI's KSampler or SamplerCustomAdvanced pipeline.

    Args:
        patched_model: Model (potentially with LoRAs applied)
        seed: Random seed
        steps: Number of sampling steps
        cfg: CFG scale
        sampler_name: Sampler name (e.g., "euler", "dpmpp_2m")
        scheduler: Scheduler name (e.g., "normal", "karras")
        positive_conditioning: Positive conditioning tensor
        negative_conditioning: Negative conditioning tensor
        latent_input: Input latent tensor dict
        denoise: Denoise strength (0.0-1.0)
        attention_mode: Attention implementation to use ("default", "xformers", "pytorch",
                       "flash", "sage", "sage3", "sub_quad", "split")
        model_sampling_override: Model sampling patch type ("none", "aura_flow", "flux", "sd3")
        model_sampling_shift: Shift value for AuraFlow/SD3 model sampling
        model_sampling_flux_max_shift: Max shift for Flux model sampling
        model_sampling_flux_base_shift: Base shift for Flux model sampling
        use_advanced_sampling: Whether to use SamplerCustomAdvanced pipeline
        advanced_guider: Guider type ("cfg_guider", "basic_guider")
        advanced_scheduler: Scheduler type for advanced sampling ("basic", "flux2")
        flux_guidance_value: Flux guidance value (0.0 = disabled)
        width: Image width (used by Flux model sampling and Flux2Scheduler)
        height: Image height (used by Flux model sampling and Flux2Scheduler)

    Returns:
        tuple: (result_latent_dict, generation_duration_seconds)
    """
    # Apply attention mode override if not "default"
    original_attn_override = None
    if attention_mode and attention_mode != "default":
        try:
            from comfy.ldm.modules.attention import get_attention_function
            attn_func = get_attention_function(attention_mode, default=None)
            if attn_func is not None:
                # Save original override (if any) and set the new one
                original_attn_override = patched_model.model_options.get("transformer_options", {}).get("optimized_attention_override")
                if "transformer_options" not in patched_model.model_options:
                    patched_model.model_options["transformer_options"] = {}
                patched_model.model_options["transformer_options"]["optimized_attention_override"] = attn_func
            else:
                print(f"[GridTester] ⚠️ Attention mode '{attention_mode}' not available, using default")
        except Exception as e:
            print(f"[GridTester] ⚠️ Could not set attention mode '{attention_mode}': {e}")

    # === Model Sampling Override ===
    # Clone model and patch model_sampling if an override is requested.
    # This modifies the model's internal noise schedule for specific model families.
    if model_sampling_override and model_sampling_override != "none":
        import comfy.model_sampling
        patched_model = patched_model.clone()

        if model_sampling_override == "aura_flow":
            # AuraFlow/Qwen Image: discrete flow with shift, multiplier=1.0
            sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow
            sampling_type = comfy.model_sampling.CONST
            class ModelSamplingAdvanced(sampling_base, sampling_type):
                pass
            model_sampling = ModelSamplingAdvanced(patched_model.model.model_config)
            model_sampling.set_parameters(shift=float(model_sampling_shift), multiplier=1.0)
            patched_model.add_object_patch("model_sampling", model_sampling)
            print(f"[GridTester] 🔧 Applied ModelSamplingAuraFlow (shift={model_sampling_shift})")

        elif model_sampling_override == "sd3":
            # SD3: discrete flow with shift, multiplier=1000
            sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow
            sampling_type = comfy.model_sampling.CONST
            class ModelSamplingAdvanced(sampling_base, sampling_type):
                pass
            model_sampling = ModelSamplingAdvanced(patched_model.model.model_config)
            model_sampling.set_parameters(shift=float(model_sampling_shift), multiplier=1000)
            patched_model.add_object_patch("model_sampling", model_sampling)
            print(f"[GridTester] 🔧 Applied ModelSamplingSD3 (shift={model_sampling_shift})")

        elif model_sampling_override == "flux":
            # Flux: dynamic shift computed from image dimensions
            sampling_base = comfy.model_sampling.ModelSamplingFlux
            sampling_type = comfy.model_sampling.CONST
            class ModelSamplingAdvanced(sampling_base, sampling_type):
                pass
            max_s = float(model_sampling_flux_max_shift)
            base_s = float(model_sampling_flux_base_shift)
            x1, x2 = 256, 4096
            mm = (max_s - base_s) / (x2 - x1)
            b = base_s - mm * x1
            shift = (width * height / (8 * 8 * 2 * 2)) * mm + b
            model_sampling = ModelSamplingAdvanced(patched_model.model.model_config)
            model_sampling.set_parameters(shift=shift)
            patched_model.add_object_patch("model_sampling", model_sampling)
            print(f"[GridTester] 🔧 Applied ModelSamplingFlux (max_shift={max_s}, base_shift={base_s}, computed_shift={shift:.4f})")

        elif model_sampling_override == "flux2":
            # Flux2: uses ModelSamplingFlux with fixed shift (default 2.02, matching ComfyUI's Flux2 supported_models)
            sampling_base = comfy.model_sampling.ModelSamplingFlux
            sampling_type = comfy.model_sampling.CONST
            class ModelSamplingAdvanced(sampling_base, sampling_type):
                pass
            shift = float(model_sampling_shift) if model_sampling_shift else 2.02
            model_sampling = ModelSamplingAdvanced(patched_model.model.model_config)
            model_sampling.set_parameters(shift=shift)
            patched_model.add_object_patch("model_sampling", model_sampling)
            print(f"[GridTester] 🔧 Applied ModelSamplingFlux2 (shift={shift})")

    # === Flux Guidance ===
    # Modify positive conditioning with guidance value (used by Flux 1 models)
    if flux_guidance_value and float(flux_guidance_value) > 0:
        import node_helpers
        positive_conditioning = node_helpers.conditioning_set_values(
            positive_conditioning, {"guidance": float(flux_guidance_value)}
        )
        print(f"[GridTester] 🔧 Applied FluxGuidance (guidance={flux_guidance_value})")

    t0 = time.time()

    try:
        if use_advanced_sampling:
            # === Advanced Sampling Pipeline (SamplerCustomAdvanced) ===
            import comfy.samplers
            import comfy.sample
            import comfy.sampler_helpers
            import comfy.model_management
            import comfy.utils
            import latent_preview

            # Noise
            class _Noise_RandomNoise:
                def __init__(self, seed):
                    self.seed = seed
                def generate_noise(self, input_latent):
                    latent_image = input_latent["samples"]
                    batch_inds = input_latent.get("batch_index")
                    return comfy.sample.prepare_noise(latent_image, self.seed, batch_inds)

            noise = _Noise_RandomNoise(seed)

            # Guider
            if advanced_guider == "basic_guider":
                guider = comfy.samplers.CFGGuider(patched_model)
                guider.inner_set_conds({"positive": comfy.sampler_helpers.convert_cond(positive_conditioning)})
                guider.set_cfg(1.0)
                print(f"[GridTester] 🔧 Advanced sampling: BasicGuider (no CFG)")
            else:
                guider = comfy.samplers.CFGGuider(patched_model)
                guider.set_conds(positive_conditioning, negative_conditioning)
                guider.set_cfg(cfg)
                print(f"[GridTester] 🔧 Advanced sampling: CFGGuider (cfg={cfg})")

            # Sampler object
            sampler_obj = comfy.samplers.sampler_object(sampler_name)

            # Sigmas
            if advanced_scheduler == "flux2":
                import math

                def _flux2_generalized_time_snr_shift(t, mu, sigma):
                    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)

                def _flux2_compute_empirical_mu(image_seq_len, num_steps):
                    a1, b1 = 8.73809524e-05, 1.89833333
                    a2, b2 = 0.00016927, 0.45666666
                    if image_seq_len > 4300:
                        return float(a2 * image_seq_len + b2)
                    m_200 = a2 * image_seq_len + b2
                    m_10 = a1 * image_seq_len + b1
                    a = (m_200 - m_10) / 190.0
                    b_val = m_200 - 200.0 * a
                    return float(a * num_steps + b_val)

                seq_len = round(width * height / (16 * 16))
                mu = _flux2_compute_empirical_mu(seq_len, steps)
                timesteps = torch.linspace(1, 0, steps + 1)
                sigmas = torch.FloatTensor([_flux2_generalized_time_snr_shift(t.item(), mu, 1.0) for t in timesteps])
                print(f"[GridTester] 🔧 Advanced sampling: Flux2Scheduler (steps={steps}, {width}x{height}, mu={mu:.4f})")
            else:
                # BasicScheduler
                total_steps = steps
                if denoise < 1.0:
                    if denoise <= 0.0:
                        sigmas = torch.FloatTensor([])
                    else:
                        total_steps = int(steps / denoise)
                model_sampling_obj = patched_model.get_model_object("model_sampling")
                sigmas = comfy.samplers.calculate_sigmas(model_sampling_obj, scheduler, total_steps).cpu()
                sigmas = sigmas[-(steps + 1):]
                print(f"[GridTester] 🔧 Advanced sampling: BasicScheduler ({scheduler}, {steps} steps)")

            # Execute SamplerCustomAdvanced
            latent = latent_input.copy()
            latent_image = latent["samples"]
            latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image)
            latent["samples"] = latent_image

            noise_mask = latent.get("noise_mask")

            x0_output = {}
            callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output)

            disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
            samples = guider.sample(
                noise.generate_noise(latent), latent_image, sampler_obj, sigmas,
                denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar,
                seed=noise.seed
            )
            samples = samples.to(comfy.model_management.intermediate_device())

            out = latent.copy()
            out["samples"] = samples
            result = ({"samples": out["samples"]},)

        else:
            # === Standard KSampler path ===
            result = nodes.common_ksampler(
                model=patched_model,
                seed=seed,
                steps=steps,
                cfg=cfg,
                sampler_name=sampler_name,
                scheduler=scheduler,
                positive=positive_conditioning,
                negative=negative_conditioning,
                latent=latent_input,
                denoise=denoise
            )
    finally:
        # Restore original attention override
        if attention_mode and attention_mode != "default":
            if original_attn_override is not None:
                patched_model.model_options["transformer_options"]["optimized_attention_override"] = original_attn_override
            elif "transformer_options" in patched_model.model_options and "optimized_attention_override" in patched_model.model_options["transformer_options"]:
                del patched_model.model_options["transformer_options"]["optimized_attention_override"]

    duration = round(time.time() - t0, 3)

    return result[0], duration


def tiled_hires_sample(latent_input, patched_model, config, positive_conditioning, negative_conditioning,
                       hires_steps, hires_denoise, tile_width, tile_height, mask_blur, tile_padding,
                       force_uniform, pixel_width, pixel_height):
    """
    Run HiRes fix sampling in tiles to prevent OOM on large images.
    Splits the latent into overlapping tiles, samples each, then blends them back together.

    Args:
        latent_input: Dict with "samples" key — the upscaled latent to denoise
        patched_model: Model for sampling
        config: Generation config (seed, cfg, sampler, scheduler)
        positive_conditioning, negative_conditioning: Conditioning tensors
        hires_steps: Number of sampling steps
        hires_denoise: Denoise strength
        tile_width, tile_height: Tile size in pixels (will be converted to latent space /8)
        mask_blur: Gaussian blur radius for tile seam blending (pixels)
        tile_padding: Extra context padding around each tile (pixels)
        force_uniform: If True, force all tiles to be the same size (may crop edges)
        pixel_width, pixel_height: Full image pixel dimensions (for logging)

    Returns:
        dict: Result latent dict with "samples" key
    """
    import torch

    samples = latent_input["samples"]
    # Convert pixel dimensions to latent space (8x smaller)
    # Handle both 4D [B, C, H, W] and 5D [B, C, T, H, W] latent formats (video VAEs add temporal dim)
    if samples.ndim == 5:
        lat_h, lat_w = samples.shape[3], samples.shape[4]
    else:
        lat_h, lat_w = samples.shape[2], samples.shape[3]
    tw = tile_width // 8
    th = tile_height // 8
    pad = tile_padding // 8
    blur = max(1, mask_blur // 8)  # Blur in latent space

    # Calculate tile grid
    def calc_tiles(total, tile_size, padding, uniform):
        """Calculate tile start positions with overlap = 2 * padding."""
        if total <= tile_size:
            return [(0, total)]
        stride = tile_size - 2 * padding
        if stride <= 0:
            stride = tile_size // 2
        tiles = []
        pos = 0
        while pos < total:
            end = min(pos + tile_size, total)
            if uniform and end == total and end - pos < tile_size:
                # Shift last tile back to maintain uniform size
                pos = max(0, total - tile_size)
                end = total
            tiles.append((pos, end))
            if end == total:
                break
            pos += stride
        return tiles

    x_tiles = calc_tiles(lat_w, tw, pad, force_uniform)
    y_tiles = calc_tiles(lat_h, th, pad, force_uniform)

    total_tiles = len(x_tiles) * len(y_tiles)
    print(f"[GridTester] 🔍 Tiled HiRes sampling: {len(x_tiles)}x{len(y_tiles)} = {total_tiles} tiles "
          f"(tile={tile_width}x{tile_height}px, padding={tile_padding}px, blur={mask_blur}px)")

    # Output accumulator with weighted blending
    result_samples = torch.zeros_like(samples)
    is_5d = samples.ndim == 5

    # Weight map shape matches spatial dims only
    if is_5d:
        weight_map = torch.zeros(1, 1, 1, lat_h, lat_w, device=samples.device)
    else:
        weight_map = torch.zeros(1, 1, lat_h, lat_w, device=samples.device)

    tile_idx = 0
    for yi, (y_start, y_end) in enumerate(y_tiles):
        for xi, (x_start, x_end) in enumerate(x_tiles):
            tile_idx += 1
            # Extract tile from latent (handle both 4D and 5D)
            if is_5d:
                tile_latent = samples[:, :, :, y_start:y_end, x_start:x_end].clone()
            else:
                tile_latent = samples[:, :, y_start:y_end, x_start:x_end].clone()

            print(f"[GridTester] 🔍   Tile {tile_idx}/{total_tiles}: latent region [{y_start}:{y_end}, {x_start}:{x_end}]")

            # Run KSampler on this tile
            tile_result, _ = generate_image(
                patched_model, config.get("seed", 0), hires_steps, config.get("cfg", 7),
                config.get("sampler", "euler"), config.get("scheduler", "normal"),
                positive_conditioning, negative_conditioning,
                {"samples": tile_latent}, hires_denoise,
                width=(x_end - x_start) * 8, height=(y_end - y_start) * 8
            )

            tile_out = tile_result["samples"]
            tile_h = y_end - y_start
            tile_w = x_end - x_start

            # Create feathered weight mask for this tile (higher weight in center, fading at edges)
            mask = torch.ones(tile_h, tile_w, device=samples.device)
            if blur > 0:
                # Feather edges: linear ramp over blur pixels
                for b in range(blur):
                    factor = (b + 1) / (blur + 1)
                    # Top edge
                    if y_start > 0 and b < tile_h:
                        mask[b, :] *= factor
                    # Bottom edge
                    if y_end < lat_h and b < tile_h:
                        mask[tile_h - 1 - b, :] *= factor
                    # Left edge
                    if x_start > 0 and b < tile_w:
                        mask[:, b] *= factor
                    # Right edge
                    if x_end < lat_w and b < tile_w:
                        mask[:, tile_w - 1 - b] *= factor

            # Accumulate weighted results (broadcast mask to match tensor dims)
            if is_5d:
                mask_shaped = mask.unsqueeze(0).unsqueeze(0).unsqueeze(0)  # [1, 1, 1, H, W]
                result_samples[:, :, :, y_start:y_end, x_start:x_end] += tile_out * mask_shaped
                weight_map[:, :, :, y_start:y_end, x_start:x_end] += mask_shaped
            else:
                mask_shaped = mask.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
                result_samples[:, :, y_start:y_end, x_start:x_end] += tile_out * mask_shaped
                weight_map[:, :, y_start:y_end, x_start:x_end] += mask_shaped

    # Normalize by weights to blend overlapping regions
    weight_map = torch.clamp(weight_map, min=1e-6)
    result_samples = result_samples / weight_map

    print(f"[GridTester] 🔍 Tiled HiRes sampling complete ({total_tiles} tiles)")
    return {"samples": result_samples}


def upscale_image(result_latent, vae, patched_model, upscaling_config, config, positive_conditioning, negative_conditioning, width, height):
    """
    Apply upscaling to a generated latent based on upscaling settings.

    Args:
        result_latent: Generated latent dict with "samples" key
        vae: VAE model for encode/decode
        patched_model: Patched model for re-sampling (HiRes fix)
        upscaling_config: Dict with mode, upscale_ratio, hires_denoise, etc.
        config: Current generation config (steps, sampler, scheduler, etc.)
        positive_conditioning: Positive conditioning
        negative_conditioning: Negative conditioning
        width: Original image width
        height: Original image height

    Returns:
        tuple: (result, duration) where result is either a latent dict or PIL Image
    """
    import torch
    import time

    mode = upscaling_config.get("mode", "hires_only")
    upscale_ratio = float(upscaling_config.get("upscale_ratio", 1.5))
    hires_denoise = float(upscaling_config.get("hires_denoise", 0.5))
    hires_steps = int(upscaling_config.get("hires_steps", 0)) or config.get("steps", 20)
    tiled_vae = upscaling_config.get("tiled_vae", False)
    tile_size = int(upscaling_config.get("tile_size", 512))
    tile_overlap = int(upscaling_config.get("tile_overlap", 64))
    temporal_size = int(upscaling_config.get("temporal_size", 512))
    temporal_overlap = int(upscaling_config.get("temporal_overlap", 64))
    upscale_model_name = upscaling_config.get("upscale_model", "")
    upscale_size = float(upscaling_config.get("upscale_size", 2.0))
    resize_method = upscaling_config.get("resize_method", "bilinear")
    hires_tiled_sampling = upscaling_config.get("hires_tiled_sampling", False)
    hires_tile_width = int(upscaling_config.get("hires_tile_width", 512))
    hires_tile_height = int(upscaling_config.get("hires_tile_height", 512))
    hires_mask_blur = int(upscaling_config.get("hires_mask_blur", 8))
    hires_tile_padding = int(upscaling_config.get("hires_tile_padding", 32))
    hires_force_uniform_tiles = upscaling_config.get("hires_force_uniform_tiles", False)

    new_w = int(width * upscale_ratio)
    new_h = int(height * upscale_ratio)

    t0 = time.time()
    print(f"[GridTester] 🔍 Upscaling: mode={mode}, ratio={upscale_ratio}, target={new_w}x{new_h}")

    if mode == "hires_only":
        # Upscale latent in latent space → re-sample with denoise
        import comfy.utils
        latent_samples = result_latent["samples"]
        # Latent space is 8x smaller than pixel space
        upscaled_latent = comfy.utils.common_upscale(
            latent_samples, new_w // 8, new_h // 8, resize_method, "disabled"
        )

        if hires_tiled_sampling:
            hires_latent = tiled_hires_sample(
                {"samples": upscaled_latent}, patched_model, config,
                positive_conditioning, negative_conditioning,
                hires_steps, hires_denoise,
                hires_tile_width, hires_tile_height, hires_mask_blur, hires_tile_padding,
                hires_force_uniform_tiles, new_w, new_h
            )
        else:
            hires_latent, hires_duration = generate_image(
                patched_model, config.get("seed", 0), hires_steps, config.get("cfg", 7),
                config.get("sampler", "euler"), config.get("scheduler", "normal"),
                positive_conditioning, negative_conditioning,
                {"samples": upscaled_latent}, hires_denoise,
                width=new_w, height=new_h
            )

        duration = round(time.time() - t0, 3)
        print(f"[GridTester] 🔍 HiRes fix complete in {duration}s → {new_w}x{new_h}")
        return hires_latent, duration

    elif mode == "model_only":
        # Decode → model upscale → optional resize to target → return as PIL image
        from comfy_extras.nodes_upscale_model import UpscaleModelLoader, ImageUpscaleWithModel
        import numpy as np

        # Use tiled VAE decode if enabled (prevents OOM on large images)
        if tiled_vae:
            from comfy_extras.nodes_post_processing import ImageScaleToTotalPixels
            vae.first_stage_model.tile_sample_min_size = tile_size
            vae.first_stage_model.tile_latent_min_size = tile_size // 8
            vae.first_stage_model.tile_overlap_factor = tile_overlap / tile_size if tile_size > 0 else 0.125
            if hasattr(vae.first_stage_model, 'tile_sample_min_size_temporal'):
                vae.first_stage_model.tile_sample_min_size_temporal = temporal_size
            if hasattr(vae.first_stage_model, 'tile_latent_min_size_temporal'):
                vae.first_stage_model.tile_latent_min_size_temporal = temporal_size // 8
            if hasattr(vae.first_stage_model, 'tile_overlap_factor_temporal'):
                vae.first_stage_model.tile_overlap_factor_temporal = temporal_overlap / temporal_size if temporal_size > 0 else 0.125
        pil_image = decode_latent_with_vae(vae, result_latent["samples"])

        img_np = np.array(pil_image).astype(np.float32) / 255.0
        img_tensor = torch.from_numpy(img_np).unsqueeze(0)  # (1, H, W, C)

        loader = UpscaleModelLoader()
        (up_model,) = loader.load_model(upscale_model_name)
        upscaler = ImageUpscaleWithModel()
        (upscaled_tensor,) = upscaler.upscale(up_model, img_tensor)

        # If upscale_size differs from model's native scale, resize the output
        target_w = int(width * upscale_size)
        target_h = int(height * upscale_size)
        actual_h, actual_w = upscaled_tensor.shape[1], upscaled_tensor.shape[2]
        if abs(actual_w - target_w) > 4 or abs(actual_h - target_h) > 4:
            import comfy.utils
            # Resize from model's native output to user-specified upscale_size
            upscaled_tensor = upscaled_tensor.permute(0, 3, 1, 2)  # NHWC → NCHW
            upscaled_tensor = comfy.utils.common_upscale(upscaled_tensor, target_w, target_h, resize_method, "disabled")
            upscaled_tensor = upscaled_tensor.permute(0, 2, 3, 1)  # NCHW → NHWC

        up_np = upscaled_tensor[0].cpu().float().numpy()
        up_np = np.clip(up_np * 255, 0, 255).astype(np.uint8)
        from PIL import Image as PILImage
        upscaled_image = PILImage.fromarray(up_np)

        duration = round(time.time() - t0, 3)
        print(f"[GridTester] 🔍 Model upscale complete in {duration}s → {upscaled_image.size[0]}x{upscaled_image.size[1]}")
        return upscaled_image, duration

    elif mode == "model_then_hires":
        # Model upscale first → optional resize → encode to latent → HiRes fix
        from comfy_extras.nodes_upscale_model import UpscaleModelLoader, ImageUpscaleWithModel
        import numpy as np

        # Use tiled VAE decode if enabled
        if tiled_vae:
            vae.first_stage_model.tile_sample_min_size = tile_size
            vae.first_stage_model.tile_latent_min_size = tile_size // 8
            vae.first_stage_model.tile_overlap_factor = tile_overlap / tile_size if tile_size > 0 else 0.125
            if hasattr(vae.first_stage_model, 'tile_sample_min_size_temporal'):
                vae.first_stage_model.tile_sample_min_size_temporal = temporal_size
            if hasattr(vae.first_stage_model, 'tile_latent_min_size_temporal'):
                vae.first_stage_model.tile_latent_min_size_temporal = temporal_size // 8
            if hasattr(vae.first_stage_model, 'tile_overlap_factor_temporal'):
                vae.first_stage_model.tile_overlap_factor_temporal = temporal_overlap / temporal_size if temporal_size > 0 else 0.125
        pil_image = decode_latent_with_vae(vae, result_latent["samples"])

        img_np = np.array(pil_image).astype(np.float32) / 255.0
        img_tensor = torch.from_numpy(img_np).unsqueeze(0)

        loader = UpscaleModelLoader()
        (up_model,) = loader.load_model(upscale_model_name)
        upscaler = ImageUpscaleWithModel()
        (upscaled_tensor,) = upscaler.upscale(up_model, img_tensor)

        # If upscale_size differs from model's native scale, resize before HiRes fix
        target_w = int(width * upscale_size)
        target_h = int(height * upscale_size)
        actual_h, actual_w = upscaled_tensor.shape[1], upscaled_tensor.shape[2]
        if abs(actual_w - target_w) > 4 or abs(actual_h - target_h) > 4:
            import comfy.utils
            upscaled_tensor = upscaled_tensor.permute(0, 3, 1, 2)  # NHWC → NCHW
            upscaled_tensor = comfy.utils.common_upscale(upscaled_tensor, target_w, target_h, resize_method, "disabled")
            upscaled_tensor = upscaled_tensor.permute(0, 2, 3, 1)  # NCHW → NHWC

        up_h, up_w = upscaled_tensor.shape[1], upscaled_tensor.shape[2]

        # Encode back to latent space for HiRes fix
        encoded_latent = vae.encode(upscaled_tensor[:, :, :, :3])

        if hires_tiled_sampling:
            hires_latent = tiled_hires_sample(
                {"samples": encoded_latent}, patched_model, config,
                positive_conditioning, negative_conditioning,
                hires_steps, hires_denoise,
                hires_tile_width, hires_tile_height, hires_mask_blur, hires_tile_padding,
                hires_force_uniform_tiles, up_w, up_h
            )
        else:
            hires_latent, hires_duration = generate_image(
                patched_model, config.get("seed", 0), hires_steps, config.get("cfg", 7),
                config.get("sampler", "euler"), config.get("scheduler", "normal"),
                positive_conditioning, negative_conditioning,
                {"samples": encoded_latent}, hires_denoise,
                width=up_w, height=up_h
            )

        duration = round(time.time() - t0, 3)
        print(f"[GridTester] 🔍 Model+HiRes upscale complete in {duration}s → {up_w}x{up_h}")
        return hires_latent, duration

    else:
        print(f"[GridTester] ⚠️ Unknown upscale mode: {mode}")
        return result_latent, 0


def decode_latent_with_vae(vae, latent_samples):
    """
    Decode latent samples to pixel space using VAE.

    Args:
        vae: VAE model
        latent_samples: Latent tensor to decode

    Returns:
        PIL.Image: Decoded image
    """
    import torch

    # Check for NaN in input latent
    if torch.isnan(latent_samples).any():
        print(f"[GridTester] ⚠️ VAE decode: input latent contains NaN! shape={latent_samples.shape} dtype={latent_samples.dtype}")

    decoded = vae.decode(latent_samples)

    # Check for NaN in decoded output and attempt float32 retry
    if torch.isnan(decoded).any():
        print(f"[GridTester] ⚠️ VAE decode produced NaN! decoded shape={decoded.shape} dtype={decoded.dtype}")
        print(f"[GridTester] ⚠️ Input latent: shape={latent_samples.shape} dtype={latent_samples.dtype} device={latent_samples.device}")
        print(f"[GridTester] 🔄 Retrying VAE decode with float32 latent...")
        decoded = vae.decode(latent_samples.to(torch.float32))
        if torch.isnan(decoded).any():
            print(f"[GridTester] ❌ VAE decode still NaN after float32 retry")
        else:
            print(f"[GridTester] ✅ float32 retry succeeded")

    # Convert to PIL Image
    # .detach() is required because the tensor may have requires_grad=True
    # (e.g., when called from distributed worker threads outside ComfyUI's
    # normal execution context where autograd state may differ)
    img_np = decoded.detach().cpu().float().numpy()

    # Remove extra dimensions (handle shapes like (1, 1, H, W, C) or (1, H, W, C))
    while img_np.ndim > 3:
        img_np = img_np[0]

    # Now should be (H, W, C) or (C, H, W)
    img_np = np.clip(img_np * 255, 0, 255).astype(np.uint8)

    # Handle different channel orders
    if img_np.shape[0] == 3 and img_np.ndim == 3:  # CHW format
        img_np = np.transpose(img_np, (1, 2, 0))
    elif img_np.shape[-1] != 3 and img_np.ndim == 3:  # Not HWC format
        img_np = np.transpose(img_np, (1, 2, 0))

    return Image.fromarray(img_np)


def save_image_to_disk(image, output_dir, filename):
    """
    Save PIL Image to disk.
    
    Args:
        image: PIL.Image object
        output_dir: Directory to save in
        filename: Filename (including extension)
        
    Returns:
        str: Full path to saved image
    """
    os.makedirs(output_dir, exist_ok=True)
    filepath = os.path.join(output_dir, filename)
    image.save(filepath)
    return filepath


def create_image_metadata(config, width, height, duration, seed, batch_idx, actual_positive_prompt, actual_negative_prompt, gen_index=None):
    """
    Create metadata dictionary for an image.

    Args:
        config: Configuration dictionary
        width: Image width
        height: Image height
        duration: Generation duration in seconds
        seed: Random seed used
        batch_idx: Batch index
        actual_positive_prompt: Final positive prompt (with triggers)
        actual_negative_prompt: Final negative prompt
        gen_index: Sequential generation index for deterministic sort ordering (optional, backwards-compatible)

    Returns:
        dict: Metadata dictionary
    """
    meta = config.copy()
    
    # Remove global settings that should only be in manifest.meta, not in individual items
    # These are session-wide settings that don't change per-image
    global_settings_to_remove = [
        "lora_triggerwords_append_settings",
        "lora_omit_triggers",
        "seed_behavior",
        "gguf_options",
        "model_prompt_prefix",
        "model_prompt_suffix"
    ]

    for key in global_settings_to_remove:
        meta.pop(key, None)

    # Remove attention_mode if it's "default" (keep manifest backward-compatible)
    if meta.get("attention_mode") == "default":
        meta.pop("attention_mode", None)

    update_dict = {
        "width": width,
        "height": height,
        "duration": duration,
        "seed": seed,
        "batch_idx": batch_idx,
        "positive": actual_positive_prompt,
        "negative": actual_negative_prompt
    }
    if gen_index is not None:
        update_dict["gen_index"] = gen_index
    # Preserve raw config prompts (without trigger words) for dashboard toggle
    meta["config_positive"] = config.get("positive", "")
    meta["config_negative"] = config.get("negative", "")
    meta.update(update_dict)

    return meta


def calculate_eta(job_durations, current_job, total_jobs):
    """
    Calculate ETA based on average job duration.
    
    Args:
        job_durations: List of previous job durations
        current_job: Current job number (1-indexed)
        total_jobs: Total number of jobs
        
    Returns:
        dict: Dictionary with eta info (hours, minutes, seconds, finish_time, finish_formatted)
    """
    if not job_durations:
        return None

    # Use rolling window of last 10 jobs for more responsive ETA
    # This adapts faster when upscaling patterns change (e.g., some configs have upscale, some don't)
    recent_window = job_durations[-10:] if len(job_durations) > 10 else job_durations
    avg_duration = sum(recent_window) / len(recent_window)
    remaining_jobs = total_jobs - current_job
    estimated_seconds = avg_duration * remaining_jobs
    
    eta_hours = int(estimated_seconds // 3600)
    eta_minutes = int((estimated_seconds % 3600) // 60)
    eta_seconds = int(estimated_seconds % 60)
    eta_finish_time = time.time() + estimated_seconds
    eta_finish_formatted = time.strftime("%H:%M:%S", time.localtime(eta_finish_time))
    
    return {
        "hours": eta_hours,
        "minutes": eta_minutes,
        "seconds": eta_seconds,
        "finish_time": eta_finish_time,
        "finish_formatted": eta_finish_formatted,
        "avg_duration": avg_duration
    }


def print_generation_progress(current_job, total_jobs, config, width, height, duration, eta_info):
    """
    Print progress information for current generation.
    
    Args:
        current_job: Current job number
        total_jobs: Total jobs
        config: Configuration dict
        width: Image width
        height: Image height
        duration: Generation duration
        eta_info: ETA info dict from calculate_eta()
    """
    progress_pct = int((current_job / total_jobs) * 100)
    
    print(f"{'='*80}")
    print(f"[GridTester] 📊 Job {current_job}/{total_jobs} ({progress_pct}%)")
    print(f"[GridTester] 🎨 {config['sampler']} @ {config['steps']} steps | {width}x{height}")
    print(f"[GridTester] ⏱️  {duration:.1f}s | Avg: {eta_info['avg_duration']:.1f}s/job")
    
    if eta_info['hours'] > 0:
        print(f"[GridTester] 🕒 ETA: {eta_info['hours']}h {eta_info['minutes']}m (finish ~{eta_info['finish_formatted']})")
    elif eta_info['minutes'] > 0:
        print(f"[GridTester] 🕒 ETA: {eta_info['minutes']}m {eta_info['seconds']}s (finish ~{eta_info['finish_formatted']})")
    else:
        print(f"[GridTester] 🕒 ETA: {eta_info['seconds']}s (finish ~{eta_info['finish_formatted']})")
    
    print(f"{'='*80}")


def flush_batch_with_vae(pending_batch, vae, img_dir, existing_data, session_name, manifest_path=None, unique_id=None):
    """
    Flush a batch of latents by decoding them with VAE and saving.
    
    Args:
        pending_batch: List of (latent_samples, metadata) tuples
        vae: VAE model for decoding
        img_dir: Image output directory
        existing_data: Manifest data dict
        session_name: Session name for filenames
        manifest_path: Path to manifest.json file (optional, enables disk syncing)
        unique_id: Node unique ID for dashboard updates (optional, enables dashboard updates)
        
    Returns:
        int: Number of images saved
    """
    if not pending_batch:
        return 0
    
    import random
    
    saved_count = 0
    
    for latent_samples, meta in pending_batch:
        # Decode latent to image
        image = decode_latent_with_vae(vae, latent_samples)
        
        # Generate ID using same format as remote_vae
        ts = int(time.time() * 100000) + random.randint(0, 1000)
        meta["id"] = ts
        
        # Generate filename using webp format (like remote_vae)
        filename = f"img_{meta['id']}.webp"
        
        # Save image as webp
        filepath = os.path.join(img_dir, filename)
        image.save(filepath, quality=80)
        
        # Normalize denoise to int if it's 1.0 (like remote_vae)
        if meta.get("denoise") == 1.0:
            meta["denoise"] = 1
        
        # Update meta with file path and rejected flag
        meta.update({
            "file": f"/view?filename={filename}&type=output&subfolder=benchmarks/{session_name}/images",
            "rejected": False
        })
        
        # Update manifest - insert at beginning (like remote_vae)
        existing_data["items"].insert(0, meta)
        
        # Sync with disk manifest to preserve tags (like remote_vae) - only if manifest_path provided
        if manifest_path and os.path.exists(manifest_path):
            try:
                with open(manifest_path, "r") as f:
                    disk_manifest = json.load(f)
                
                # Create a lookup map for items currently in memory
                memory_items_map = {
                    i.get("id"): i 
                    for i in existing_data.get("items", []) 
                    if "id" in i
                }

                # Check every item on disk. If it exists in memory, copy the tags over.
                for disk_item in disk_manifest.get("items", []):
                    d_id = disk_item.get("id")
                    if d_id and d_id in memory_items_map:
                        local_item = memory_items_map[d_id]
                        
                        # PRESERVE TAGS: Copy these keys from disk to memory
                        if "favorited" in disk_item:
                            local_item["favorited"] = disk_item["favorited"]
                        if "rejected" in disk_item:
                            local_item["rejected"] = disk_item["rejected"]

            except Exception as e:
                print(f"[GridTester] ⚠️ Error syncing with disk manifest: {e}")

        # Save manifest to disk (like remote_vae) - only if manifest_path provided
        if manifest_path:
            with open(manifest_path, "w") as f:
                json.dump(existing_data, f, indent=4)
        
        # Send update to dashboard (like remote_vae) - only if unique_id provided
        if unique_id:
            try:
                from server import PromptServer
                if PromptServer:
                    # Get meta, use empty dict if not present
                    manifest_meta = existing_data.get("meta", {})
                    
                    PromptServer.instance.send_sync("ultimate_grid.update", {
                        "node": unique_id,
                        "session_name": session_name,
                        "new_items": [meta],
                        "meta": manifest_meta
                    })
            except (ImportError, KeyError):
                # Silently ignore dashboard update errors
                pass
        
        saved_count += 1
    
    print(f"[GridTester] 💾 Flushed {saved_count} images")
    return saved_count


def flush_batch_with_remote_vae(pending_batch, remote_vae_worker, existing_data, session_name):
    """
    Flush a batch of latents by sending them to remote VAE worker.
    
    Args:
        pending_batch: List of (latent_samples, metadata) tuples
        remote_vae_worker: RemoteVAEDecodeWorker instance
        existing_data: Manifest data dict (not used - worker handles manifest)
        session_name: Session name for filenames
        
    Returns:
        int: Number of images queued
    """
    if not pending_batch:
        return 0
    
    import random
    
    queued_count = 0
    
    for latent_samples, meta in pending_batch:
        # Generate ID using same format as old code
        ts = int(time.time() * 100000) + random.randint(0, 1000)
        meta["id"] = ts  # ← Set ID directly on meta dict
        
        # Queue the job - worker will create manifest entry
        # Pass the original meta dict, not a new one
        remote_vae_worker.add_job(
            latent_samples, 
            meta,  # ← Pass original meta, not a new dict
            meta["height"],
            meta["width"]
        )
        queued_count += 1
    
    print(f"[GridTester] 🌐 Queued {queued_count} images for remote VAE decoding")
    return queued_count


# =============================================================================
# SEEDVR2 UPSCALE — Calls ComfyUI-SeedVR2_VideoUpscaler nodes programmatically
# Requires ComfyUI-SeedVR2_VideoUpscaler to be installed as a dependency.
# =============================================================================

def seedvr2_upscale(pil_image, seedvr2_config):
    """
    Upscale an image using SeedVR2 diffusion-based upscaler.

    Args:
        pil_image: PIL Image (RGB)
        seedvr2_config: dict with SeedVR2 options (dit_model, resolution, seed, etc.)

    Returns:
        tuple: (pil_result, width, height, duration)
    """
    import time
    import torch
    import numpy as np

    t0 = time.time()
    sv = seedvr2_config

    # Find SeedVR2 nodes from ComfyUI's node registry.
    # SeedVR2 uses V3 API (ComfyExtension) so we scan NODE_CLASS_MAPPINGS directly
    # rather than trying to import the custom_nodes package.
    SeedVR2LoadDiTModel = None
    SeedVR2LoadVAEModel = None
    SeedVR2VideoUpscaler = None
    try:
        import nodes
        for name, cls in nodes.NODE_CLASS_MAPPINGS.items():
            if name == "SeedVR2LoadDiTModel":
                SeedVR2LoadDiTModel = cls
            elif name == "SeedVR2LoadVAEModel":
                SeedVR2LoadVAEModel = cls
            elif name == "SeedVR2VideoUpscaler":
                SeedVR2VideoUpscaler = cls
    except Exception:
        pass

    if not SeedVR2LoadDiTModel or not SeedVR2LoadVAEModel or not SeedVR2VideoUpscaler:
        raise RuntimeError(
            "SeedVR2 upscale requires ComfyUI-SeedVR2_VideoUpscaler to be installed.\n"
            "Install from: https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler\n"
            "SeedVR2 nodes not found in ComfyUI's node registry."
        )

    print(f"[GridTester] 🎬 SeedVR2 upscale: model={sv.get('dit_model', '3b_fp8')}, "
          f"resolution={sv.get('resolution', 1080)}, seed={sv.get('seed', 42)}")

    # Set up V3 API execution context — SeedVR2 nodes require this
    # (they call get_executing_context().node_id internally)
    from comfy_execution.utils import CurrentNodeContext
    import uuid
    dummy_prompt_id = str(uuid.uuid4())

    # Convert PIL image to ComfyUI IMAGE tensor [1, H, W, 3] float32 [0,1]
    img_np = np.array(pil_image.convert("RGB")).astype(np.float32) / 255.0
    img_tensor = torch.from_numpy(img_np).unsqueeze(0)  # [1, H, W, 3]

    # All SeedVR2 execute() calls run inside a mock execution context
    # (V3 API nodes call get_executing_context().node_id internally)
    with CurrentNodeContext(prompt_id=dummy_prompt_id, node_id="uscg_seedvr2", list_index=0):
        # Load DiT model
        dit_result = SeedVR2LoadDiTModel.execute(
            model=sv.get("dit_model", "seedvr2_ema_3b_fp8_e4m3fn.safetensors"),
            device="cuda:0",
            blocks_to_swap=sv.get("blocks_to_swap", 0),
            swap_io_components=False,
            offload_device=sv.get("offload_device", "cpu"),
            cache_model=sv.get("cache_model", False),
            attention_mode=sv.get("attention_mode", "sdpa")
        )
        dit = dit_result.output[0] if hasattr(dit_result, 'output') else dit_result[0]

        # Load VAE model
        vae_result = SeedVR2LoadVAEModel.execute(
            model="ema_vae_fp16.safetensors",
            device="cuda:0",
            encode_tiled=sv.get("encode_tiled", False),
            encode_tile_size=sv.get("encode_tile_size", 1024),
            encode_tile_overlap=sv.get("encode_tile_overlap", 128),
            decode_tiled=sv.get("decode_tiled", False),
            decode_tile_size=sv.get("decode_tile_size", 1024),
            decode_tile_overlap=sv.get("decode_tile_overlap", 128),
            offload_device=sv.get("vae_offload_device", "none"),
            cache_model=sv.get("vae_cache_model", False)
        )
        vae = vae_result.output[0] if hasattr(vae_result, 'output') else vae_result[0]

        # Run upscale
        upscale_result = SeedVR2VideoUpscaler.execute(
            image=img_tensor,
            dit=dit,
            vae=vae,
            seed=sv.get("seed", 42),
            resolution=sv.get("resolution", 1080),
            max_resolution=sv.get("max_resolution", 0),
            batch_size=sv.get("batch_size", 1),
            uniform_batch_size=False,
            color_correction=sv.get("color_correction", "lab"),
            input_noise_scale=sv.get("input_noise_scale", 0.0),
            latent_noise_scale=sv.get("latent_noise_scale", 0.0),
            offload_device=sv.get("offload_device", "cpu")
        )
        result_tensor = upscale_result.output[0] if hasattr(upscale_result, 'output') else upscale_result[0]

    # Convert back to PIL: [1, H', W', 3] float32 → PIL
    from PIL import Image
    result_np = (result_tensor[0].cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
    result_pil = Image.fromarray(result_np, "RGB")
    up_w, up_h = result_pil.size

    duration = round(time.time() - t0, 2)
    print(f"[GridTester] 🎬 SeedVR2 complete: {up_w}x{up_h} in {duration}s")

    return result_pil, up_w, up_h, duration