Add model-patch variant compatible with KSampler
The existing node outputs a SAMPLER and only works with SamplerCustom. The new SelectiveSigmaDetailerPatchModelNode takes MODEL in and out, attaching the mask-and-blend logic via set_model_unet_function_wrapper so the same behavior works inside any sampler node. Tradeoffs vs the SAMPLER variant: progress is measured in sigma-space (1 - sigma/sigma_max) instead of step index since the full schedule is not visible at the model layer; intensity is calibrated against a fixed 16-active-step reference instead of the actual active-step count; the cfg_scale coupling is dropped since it is not cleanly accessible from the unet wrapper. Behavior should match the SAMPLER variant closely near 20 steps on standard schedulers.
This commit is contained in:
@@ -1,15 +1,18 @@
|
||||
from .selective_sigma_detailer import (
|
||||
SelectiveSigmaDetailerDeltaV2Node,
|
||||
SelectiveSigmaDetailerPatchModelNode,
|
||||
SelectiveSigmaDetailerMaskPreviewNode,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SelectiveSigmaDetailerDeltaV2Node": SelectiveSigmaDetailerDeltaV2Node,
|
||||
"SelectiveSigmaDetailerPatchModelNode": SelectiveSigmaDetailerPatchModelNode,
|
||||
"SelectiveSigmaDetailerMaskPreviewNode": SelectiveSigmaDetailerMaskPreviewNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SelectiveSigmaDetailerDeltaV2Node": "Selective Sigma Detailer",
|
||||
"SelectiveSigmaDetailerPatchModelNode": "Selective Sigma Detailer (Model Patch)",
|
||||
"SelectiveSigmaDetailerMaskPreviewNode": "Selective Sigma Detailer Mask Preview",
|
||||
}
|
||||
|
||||
|
||||
@@ -83,6 +83,89 @@ class SelectiveSigmaDetailerDeltaV2Node:
|
||||
return (ksampler, mask_ref)
|
||||
|
||||
|
||||
_REFERENCE_ACTIVE_STEPS = 16
|
||||
|
||||
|
||||
class SelectiveSigmaDetailerPatchModelNode:
|
||||
DESCRIPTION = (
|
||||
"Model-patch variant. Attaches the mask-and-blend logic to the model "
|
||||
"itself via set_model_unet_function_wrapper, so any sampler node "
|
||||
"(KSampler, KSamplerAdvanced, SamplerCustom) picks it up. Progress "
|
||||
"is measured in sigma-space since the full schedule is not visible "
|
||||
"at the model layer; intensity is calibrated against a 16-active-step "
|
||||
"reference, so magnitude may need retuning on very short or long runs."
|
||||
)
|
||||
CATEGORY = "sampling/custom_sampling/samplers"
|
||||
RETURN_TYPES = ("MODEL", "SSD_MASK_REF")
|
||||
RETURN_NAMES = ("model", "mask_ref")
|
||||
FUNCTION = "patch"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"intensity": ("FLOAT", {"default": 2.0, "min": -100.0, "max": 100.0, "step": 0.1,
|
||||
"tooltip": "Strength of the sigma adjustment on masked regions."}),
|
||||
"coverage": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
|
||||
"tooltip": "0 = empty mask (no effect), 0.5 = normal delta mask, 1.0 = full mask (applied everywhere)."}),
|
||||
}
|
||||
}
|
||||
|
||||
def patch(self, model, intensity, coverage):
|
||||
mask_ref = {}
|
||||
state = {}
|
||||
per_step_adjustment = intensity * 0.1 / _REFERENCE_ACTIVE_STEPS
|
||||
|
||||
def wrapper(apply_model, options):
|
||||
x = options["input"]
|
||||
sigma = options["timestep"]
|
||||
c = options["c"]
|
||||
|
||||
sigma_float = float(sigma.max().detach().cpu())
|
||||
sigma_max = state.get("sigma_max", 0.0)
|
||||
if sigma_float > sigma_max * 1.001 or sigma_max == 0.0:
|
||||
state.clear()
|
||||
state["sigma_max"] = sigma_float
|
||||
sigma_max = sigma_float
|
||||
|
||||
denoised_normal = apply_model(x, sigma, **c)
|
||||
|
||||
progress = 1.0 - sigma_float / sigma_max if sigma_max > 0 else 0.0
|
||||
if progress < _SCHEDULE_START:
|
||||
state["prev_denoised"] = denoised_normal.detach()
|
||||
return denoised_normal
|
||||
|
||||
prev = state.get("prev_denoised")
|
||||
state["prev_denoised"] = denoised_normal.detach()
|
||||
if prev is None or prev.shape != denoised_normal.shape:
|
||||
return denoised_normal
|
||||
|
||||
raw = (denoised_normal - prev).abs().mean(dim=1, keepdim=True)
|
||||
m = normalize_mask(raw, _MASK_CLIP_PERCENTILE)
|
||||
prev_mask = state.get("mask")
|
||||
if prev_mask is not None and prev_mask.shape == m.shape:
|
||||
m = _MASK_EMA * prev_mask + (1 - _MASK_EMA) * m
|
||||
state["mask"] = m
|
||||
|
||||
shift = 2 * (coverage - 0.5)
|
||||
if shift <= -1.0:
|
||||
return denoised_normal
|
||||
if shift != 0:
|
||||
m = (m + shift).clamp(0, 1)
|
||||
|
||||
mask_ref["mask"] = m.detach()
|
||||
|
||||
adjusted_sigma = sigma * max(1e-6, 1.0 - per_step_adjustment)
|
||||
denoised_detailed = apply_model(x, adjusted_sigma, **c)
|
||||
m_match = m.to(denoised_normal)
|
||||
return denoised_normal * (1 - m_match) + denoised_detailed * m_match
|
||||
|
||||
patched = model.clone()
|
||||
patched.set_model_unet_function_wrapper(wrapper)
|
||||
return (patched, mask_ref)
|
||||
|
||||
|
||||
class SelectiveSigmaDetailerMaskPreviewNode:
|
||||
DESCRIPTION = (
|
||||
"Displays the mask captured during sampling. Place between a Selective "
|
||||
|
||||
Reference in New Issue
Block a user