Jason Hoku
55ac08a73b
Extends the May 2026 Florence2 fix (29GB -> <5GB) to every other out-of-graph
inference entry point in the project. Same root cause everywhere: we call
ComfyUI models from background threads / our orchestrator, bypassing
ComfyUI/execution.py:732 which wraps every node in torch.inference_mode().
Decorated with @torch.inference_mode() in image_generation.py:
- generate_image() — SDXL/diffusion sampling (nodes.common_ksampler)
- upscale_image() — KSampler hires fix + ImageUpscaleWithModel
(ESRGAN-style) + VAE encode/decode
- decode_latent_with_vae() — VAE decode (with NaN float32 retry)
- seedvr2_upscale() — SeedVR2 diffusion-based video upscaler
(DiT + VAE + sampling loop)
Plus prominent module-docstring warning blocks in:
- image_generation.py: full "why this matters / where / what to do if
you add a new entry point" explanation
- ltx_video_generation.py: warning that _call_node()'s inference_mode wrap
is load-bearing and must not be removed
Inline comment above each decorated function points back to the module
docstring so anyone editing one function still sees the rationale.
Updated the stale comment in decode_latent_with_vae explaining .detach()
(was 'tensor may have requires_grad=True' — under inference_mode it can't,
so .detach() is now a defensive no-op; kept for double-safety).
Tests: 118/118 pass.
1199 lines
54 KiB
Python
1199 lines
54 KiB
Python
"""
|
||
Image Generation and Sampling Module
|
||
Handles the core image generation, decoding, and batch management.
|
||
|
||
============================================================================
|
||
🚨 CRITICAL: torch.inference_mode() WRAPPING — DO NOT REMOVE 🚨
|
||
============================================================================
|
||
Every top-level GPU-inference entry point in this module is decorated with
|
||
`@torch.inference_mode()`. This mirrors ComfyUI's PromptExecutor which wraps
|
||
every node execution in inference_mode (ComfyUI/execution.py:732).
|
||
|
||
WHY THIS MATTERS:
|
||
We are an out-of-graph caller — we invoke ComfyUI's sampling / VAE / upscale
|
||
models directly from a background thread (the dashboard upscale runner) and
|
||
from our orchestrator, bypassing the normal prompt executor. Without
|
||
inference_mode wrapping, autograd's version-counter and reference-keeping
|
||
machinery stays active. For diffusion sampling, beam-search generation,
|
||
VAE encode/decode, and upscale-model inference, that means intermediate
|
||
activations and past_key_values can't be released — measured as a 6× VRAM
|
||
blowup on Florence2 (29GB on a 16GB card vs <5GB in the standalone workflow
|
||
on the same image, May 2026). Same pattern applies to SDXL sampling,
|
||
SeedVR2 diffusion, and ESRGAN-style upscalers, just at a smaller magnitude.
|
||
|
||
torch.inference_mode() is STRICTER than torch.no_grad():
|
||
- no_grad: disables gradient tracking, but version counters and autograd
|
||
machinery stay partly active
|
||
- inference_mode: tensors are entirely outside autograd, true read-only —
|
||
this is what lets HF generate() actually release KV cache between beam
|
||
steps and what lets ComfyUI's sample loop reclaim activations between
|
||
diffusion steps
|
||
|
||
WHERE IT'S APPLIED:
|
||
- @torch.inference_mode() on generate_image, upscale_image,
|
||
decode_latent_with_vae, seedvr2_upscale (this file)
|
||
- Inside _call_node() in ltx_video_generation.py (covers all 25 ComfyUI
|
||
node invocations across florence2_hires.py + ltx_video_generation.py)
|
||
|
||
IF YOU ADD A NEW INFERENCE ENTRY POINT:
|
||
Wrap it in @torch.inference_mode() (decorator) or
|
||
`with torch.inference_mode():` (context). Symptom of missing it: GPU OOM
|
||
on workloads that work fine in a standalone ComfyUI workflow, or sudden
|
||
VRAM growth that doesn't match the model size. See
|
||
ComfyUI/execution.py:732 for the canonical pattern.
|
||
============================================================================
|
||
"""
|
||
|
||
import time
|
||
import os
|
||
import re
|
||
import json
|
||
import torch
|
||
import nodes
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
|
||
# NOTE: torch.inference_mode() decorator — see module docstring above for why
|
||
# this is REQUIRED on all out-of-graph inference entry points. Removing it
|
||
# will reintroduce the same 6× VRAM blowup we fixed in May 2026.
|
||
@torch.inference_mode()
|
||
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,
|
||
use_deep_shrink=False,
|
||
deep_shrink_block_number=3,
|
||
deep_shrink_downscale_factor=2.0,
|
||
deep_shrink_start_percent=0.0,
|
||
deep_shrink_end_percent=0.35,
|
||
deep_shrink_downscale_after_skip=True,
|
||
deep_shrink_downscale_method="bicubic",
|
||
deep_shrink_upscale_method="bicubic",
|
||
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)
|
||
use_deep_shrink: Whether to apply Kohya Deep Shrink (PatchModelAddDownscale)
|
||
deep_shrink_block_number: UNet block to downscale at (1-32, default 3)
|
||
deep_shrink_downscale_factor: Downscale factor (0.1-9.0, default 2.0)
|
||
deep_shrink_start_percent: Diffusion start % for the patch (default 0.0)
|
||
deep_shrink_end_percent: Diffusion end % for the patch (default 0.35)
|
||
deep_shrink_downscale_after_skip: Whether to downscale after skip (default True)
|
||
deep_shrink_downscale_method: Downscale interpolation method (default "bicubic")
|
||
deep_shrink_upscale_method: Upscale interpolation method (default "bicubic")
|
||
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})")
|
||
|
||
# === Kohya Deep Shrink (PatchModelAddDownscale) ===
|
||
# Patches the UNet to downscale features at a specific block during the
|
||
# early portion of diffusion. Wraps ComfyUI's built-in
|
||
# PatchModelAddDownscale (comfy_extras/nodes_model_downscale.py).
|
||
# Looked up via NODE_CLASS_MAPPINGS per project convention.
|
||
#
|
||
# Current ComfyUI ships this as a V3 io.ComfyNode subclass with a
|
||
# classmethod execute() returning io.NodeOutput (whose positional outputs
|
||
# live in .args). Older builds had a V1 instance method `.patch()`. We
|
||
# dispatch to either depending on what we find — same pattern as
|
||
# ltx_video_generation._call_node / _unwrap.
|
||
if use_deep_shrink:
|
||
try:
|
||
deep_shrink_cls = nodes.NODE_CLASS_MAPPINGS.get("PatchModelAddDownscale")
|
||
if deep_shrink_cls is not None:
|
||
ds_kwargs = dict(
|
||
model=patched_model,
|
||
block_number=int(deep_shrink_block_number),
|
||
downscale_factor=float(deep_shrink_downscale_factor),
|
||
start_percent=float(deep_shrink_start_percent),
|
||
end_percent=float(deep_shrink_end_percent),
|
||
downscale_after_skip=bool(deep_shrink_downscale_after_skip),
|
||
downscale_method=str(deep_shrink_downscale_method),
|
||
upscale_method=str(deep_shrink_upscale_method),
|
||
)
|
||
# Try V3 first (classmethod execute, NodeOutput return).
|
||
ds_result = None
|
||
try:
|
||
from comfy_api.latest import io as _ds_io
|
||
if isinstance(deep_shrink_cls, type) and issubclass(deep_shrink_cls, _ds_io.ComfyNode):
|
||
ds_result = deep_shrink_cls.execute(**ds_kwargs)
|
||
except (ImportError, TypeError):
|
||
pass
|
||
if ds_result is None:
|
||
# V1 fallback (instance method named by FUNCTION, default "patch").
|
||
fn_name = getattr(deep_shrink_cls, "FUNCTION", "patch")
|
||
ds_result = getattr(deep_shrink_cls(), fn_name)(**ds_kwargs)
|
||
# Unwrap to a single MODEL value.
|
||
if hasattr(ds_result, "args") and isinstance(getattr(ds_result, "args"), tuple):
|
||
patched_model = ds_result.args[0] # V3 NodeOutput
|
||
elif isinstance(ds_result, (tuple, list)):
|
||
patched_model = ds_result[0] # V1 tuple
|
||
else:
|
||
patched_model = ds_result # bare model
|
||
print(f"[GridTester] 🔧 Applied PatchModelAddDownscale (Kohya Deep Shrink): block={deep_shrink_block_number}, factor={deep_shrink_downscale_factor}, range={deep_shrink_start_percent}-{deep_shrink_end_percent}, after_skip={deep_shrink_downscale_after_skip}, down={deep_shrink_downscale_method}, up={deep_shrink_upscale_method}")
|
||
else:
|
||
print("[GridTester] ⚠️ PatchModelAddDownscale node not available in this ComfyUI build — Deep Shrink skipped")
|
||
except Exception as e:
|
||
print(f"[GridTester] ⚠️ Failed to apply Deep Shrink: {e}")
|
||
|
||
# === 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}
|
||
|
||
|
||
# @torch.inference_mode() — see module docstring. Covers KSampler hires fix,
|
||
# ImageUpscaleWithModel (ESRGAN-style), and VAE encode/decode in this function.
|
||
@torch.inference_mode()
|
||
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
|
||
|
||
|
||
# @torch.inference_mode() — see module docstring. VAE decode is a model forward
|
||
# pass and needs inference_mode to avoid pinning intermediate activations.
|
||
@torch.inference_mode()
|
||
def decode_latent_with_vae(vae, latent_samples):
|
||
"""
|
||
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 a defensive no-op under @torch.inference_mode() (tensors are
|
||
# already non-grad). Kept for safety in case this function is ever called
|
||
# without the decorator, but inference_mode is the primary defense — see
|
||
# module docstring for why.
|
||
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, image_format="webp"):
|
||
"""
|
||
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)
|
||
image_format: Image file format — "webp" (default), "png", or "jpg"
|
||
|
||
Returns:
|
||
int: Number of images saved
|
||
"""
|
||
if not pending_batch:
|
||
return 0
|
||
|
||
import random
|
||
|
||
saved_count = 0
|
||
|
||
# Normalise format string
|
||
fmt = (image_format or "webp").lower().strip()
|
||
if fmt not in ("webp", "png", "jpg", "jpeg"):
|
||
fmt = "webp"
|
||
ext = "jpg" if fmt in ("jpg", "jpeg") else fmt
|
||
|
||
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 with the chosen extension
|
||
filename = f"img_{meta['id']}.{ext}"
|
||
|
||
# Save image in chosen format
|
||
filepath = os.path.join(img_dir, filename)
|
||
if fmt == "png":
|
||
image.save(filepath, format="PNG")
|
||
elif fmt in ("jpg", "jpeg"):
|
||
image.save(filepath, format="JPEG", quality=95)
|
||
else:
|
||
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.
|
||
# =============================================================================
|
||
|
||
# @torch.inference_mode() — see module docstring. SeedVR2 is a diffusion-based
|
||
# upscaler with its own iterative sampling loop, identical autograd concerns
|
||
# to SDXL KSampler. The CurrentNodeContext below sets up V3-API execution but
|
||
# does NOT include inference_mode — that's our job at the function boundary.
|
||
@torch.inference_mode()
|
||
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 |