Jason HokuandClaude Fable 5.1
e1936bfdad
metrics_runner computes metrics for an existing session straight to disk; a generation run's in-memory copy of those items had no metrics and would have wiped them on its next save. Disk wins for an item's metrics when the memory item has none (merge_manifest_user_changes, flush_batch_with_vae's disk sync, the remote VAE worker's disk sync). Fresh in-memory metrics are kept. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1322 lines
60 KiB
Python
1322 lines
60 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
|
||
|
||
from . import image_metrics
|
||
from .image_metrics import _CURRENT as _METRICS_RUN
|
||
|
||
|
||
def _latent_ratio(patched_model):
|
||
"""Pixels per latent cell for this model (8 for SD/SDXL/Flux.1, 16 for the Flux.2 family)."""
|
||
from .model_loader import get_latent_downscale_ratio
|
||
return get_latent_downscale_ratio(patched_model)
|
||
|
||
|
||
def _vae_spatial_compression(vae):
|
||
"""Pixels per latent cell for a VAE (what VAEDecodeTiled divides tile_size by); 8 if unknown."""
|
||
try:
|
||
ratio = vae.spacial_compression_decode()
|
||
if isinstance(ratio, int) and ratio > 0:
|
||
return ratio
|
||
except Exception:
|
||
pass
|
||
return 8
|
||
|
||
|
||
def _vae_temporal_compression(vae):
|
||
"""Frames per latent frame for a video VAE; 8 (the historical divisor here) if unknown."""
|
||
try:
|
||
ratio = vae.temporal_compression_decode()
|
||
if isinstance(ratio, int) and ratio > 0:
|
||
return ratio
|
||
except Exception:
|
||
pass
|
||
return 8
|
||
|
||
|
||
def _apply_model_patch_node(node_id, model, label=None, **kwargs):
|
||
"""
|
||
Run a ComfyUI model-patching node (MODEL in -> MODEL out) by its node id and
|
||
return the patched model, or the input model if the node is missing/fails.
|
||
|
||
Looked up via nodes.NODE_CLASS_MAPPINGS per project convention. Current
|
||
ComfyUI ships these as V3 io.ComfyNode subclasses (classmethod execute()
|
||
returning io.NodeOutput whose positional outputs live in .args); older
|
||
builds used V1 instance methods named by FUNCTION. Same dispatch as the
|
||
inline Deep Shrink block in generate_image and ltx_video_generation._call_node.
|
||
"""
|
||
node_cls = nodes.NODE_CLASS_MAPPINGS.get(node_id)
|
||
if node_cls is None:
|
||
print(f"[GridTester] ⚠️ {node_id} node not available in this ComfyUI build — skipped")
|
||
return model
|
||
try:
|
||
result = None
|
||
try:
|
||
from comfy_api.latest import io as _io
|
||
if isinstance(node_cls, type) and issubclass(node_cls, _io.ComfyNode):
|
||
result = node_cls.execute(model=model, **kwargs)
|
||
except (ImportError, TypeError):
|
||
pass
|
||
if result is None:
|
||
fn_name = getattr(node_cls, "FUNCTION", "patch")
|
||
result = getattr(node_cls(), fn_name)(model=model, **kwargs)
|
||
if hasattr(result, "args") and isinstance(getattr(result, "args"), tuple):
|
||
patched = result.args[0] # V3 NodeOutput
|
||
elif isinstance(result, (tuple, list)):
|
||
patched = result[0] # V1 tuple
|
||
else:
|
||
patched = result # bare model
|
||
print(f"[GridTester] 🔧 Applied {label or node_id}")
|
||
return patched
|
||
except Exception as e:
|
||
print(f"[GridTester] ⚠️ Failed to apply {node_id}: {e}")
|
||
return model
|
||
|
||
|
||
# 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",
|
||
use_slg=False,
|
||
slg_double_layers="",
|
||
slg_single_layers="22",
|
||
slg_start_percent=0.0,
|
||
slg_end_percent=0.5,
|
||
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")
|
||
use_slg: Whether to apply Skip Layer Guidance (SkipLayerGuidanceDiTSimple)
|
||
slg_double_layers: Comma-separated double-block indices the uncond pass skips ("" = none)
|
||
slg_single_layers: Comma-separated single-block indices the uncond pass skips (MOSS-Image template: "22")
|
||
slg_start_percent: Sampling percent where the skip starts (default 0.0)
|
||
slg_end_percent: Sampling percent where the skip ends (default 0.5)
|
||
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}")
|
||
|
||
# === Skip Layer Guidance (SkipLayerGuidanceDiTSimple) ===
|
||
# DiT models only: between start/end percent the uncond (negative) pass
|
||
# skips the listed transformer blocks, which sharpens structure. The
|
||
# official MOSS-Image workflow ships this on single block 22 for the first
|
||
# half of sampling. Wraps comfy_extras/nodes_slg.py; the node clones the
|
||
# model, so the LoRA-patched model handed in is left untouched.
|
||
if use_slg:
|
||
double_s = str(slg_double_layers or "").strip()
|
||
single_s = str(slg_single_layers or "").strip()
|
||
if not double_s and not single_s:
|
||
print("[GridTester] ⚠️ Skip Layer Guidance enabled but no layers listed — skipped")
|
||
else:
|
||
patched_model = _apply_model_patch_node(
|
||
"SkipLayerGuidanceDiTSimple", patched_model,
|
||
label=f"SkipLayerGuidanceDiTSimple: double=[{double_s}] single=[{single_s}] range={float(slg_start_percent)}-{float(slg_end_percent)}",
|
||
double_layers=double_s,
|
||
single_layers=single_s,
|
||
start_percent=float(slg_start_percent),
|
||
end_percent=float(slg_end_percent),
|
||
)
|
||
|
||
# === 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 (8 px per cell for SD/SDXL/Flux.1,
|
||
# 16 for the Flux.2 family incl. MOSS-Image; read from the model's latent format)
|
||
# Handle both 4D [B, C, H, W] and 5D [B, C, T, H, W] latent formats (video VAEs add temporal dim)
|
||
ratio = _latent_ratio(patched_model)
|
||
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 = max(1, tile_width // ratio)
|
||
th = max(1, tile_height // ratio)
|
||
pad = tile_padding // ratio
|
||
blur = max(1, mask_blur // ratio) # 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 cells per pixel come from the model's latent format (8 or 16)
|
||
ratio = _latent_ratio(patched_model)
|
||
upscaled_latent = comfy.utils.common_upscale(
|
||
latent_samples, new_w // ratio, new_h // ratio, 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 // _vae_spatial_compression(vae)
|
||
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 // _vae_temporal_compression(vae)
|
||
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 // _vae_spatial_compression(vae)
|
||
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 // _vae_temporal_compression(vae)
|
||
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, metrics_sink=None, metrics_spec=_METRICS_RUN):
|
||
"""
|
||
Decode latent samples to pixel space using VAE.
|
||
|
||
Args:
|
||
vae: VAE model
|
||
latent_samples: Latent tensor to decode
|
||
metrics_sink: optional dict; when given (and metrics are on) the image
|
||
metrics are computed on the decoded GPU tensor and written into it
|
||
metrics_spec: omitted -> the active generation run's spec; an explicit
|
||
spec (or None) for threads that must not read generation-run state
|
||
|
||
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.
|
||
if metrics_sink is not None:
|
||
image_metrics.sink_update(metrics_sink, decoded, metrics_spec)
|
||
|
||
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
|
||
metrics_sink = {}
|
||
image = decode_latent_with_vae(vae, latent_samples, metrics_sink=metrics_sink)
|
||
|
||
# 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
|
||
})
|
||
|
||
# Optional image metrics (builtins came from the decode sink; adapters run here)
|
||
if metrics_sink:
|
||
meta["metrics"] = metrics_sink
|
||
image_metrics.attach_adapters(meta, image)
|
||
|
||
# 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)
|
||
|
||
# Pairwise-review votes live on disk (disk wins), never drop them on this save
|
||
from .manifest_utils import preserve_disk_comparisons
|
||
preserve_disk_comparisons(disk_manifest, existing_data)
|
||
|
||
# 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"]
|
||
if "metrics" in disk_item and "metrics" not in local_item:
|
||
local_item["metrics"] = disk_item["metrics"] # computed later for an existing session
|
||
|
||
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 |