From dcd4eafdf1d4e30631f870f6a4f4e71b55e7d109 Mon Sep 17 00:00:00 2001
From: capacap <simon.sorkin@outlook.com>
Date: Tue, 21 Apr 2026 01:11:14 +0200
Subject: [PATCH] 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.
---
 __init__.py                 |  3 ++
 selective_sigma_detailer.py | 83 +++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/__init__.py b/__init__.py
index 5ad3a5e..5b5fd71 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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",
 }
 
diff --git a/selective_sigma_detailer.py b/selective_sigma_detailer.py
index 665390c..3a0aca8 100644
--- a/selective_sigma_detailer.py
+++ b/selective_sigma_detailer.py
@@ -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 "