Merge pull request #13 from xmarre/codex/fix-detailer-mask-for-slrd

Fix detailer hook to use effective guider semantics

README.md

@@ -55,13 +55,15 @@ One-node version of the modular custom-sampling workflow. Internally it now call
 
 Experimental Impact Pack / FaceDetailer integration path.
 
-This node returns a `DETAILER_HOOK` object that preserves the FaceDetailer crop mask, records the encoded crop latent as the anchor in `post_encode(...)`, and applies one-shot residual/manifold correction in `pre_decode(...)`.
+This node returns a `DETAILER_HOOK` object that captures the FaceDetailer crop mask in `post_upscale(...)` and drives the masked residual/manifold correction through the hook's sampler path.
 
-This provider no longer advertises sampler-runtime controls that do not participate in its one-shot hook path. The exposed knobs are the ones that still affect the masked latent correction directly:
+The live hook path is `pre_ksample(...)` plus the custom sampler/runtime integration, not the inert `post_encode(...)` / `pre_decode(...)` pair.
+
+This provider no longer advertises sampler-runtime controls that do not participate in its live sampler-driven hook path. The exposed knobs are the ones that still affect the masked latent correction directly:
 
 - residual lock strength and cutoffs,
 - optional manifold companding controls,
-- optional `lock_mask` / `manifold_mask`, which are intersected with the FaceDetailer mask.
+- optional `lock_mask` / `manifold_mask`, which are intersected with the FaceDetailer support mask.
 
 Compatibility note: this provider's input signature changed when the inert detailer-only sampler/runtime controls were removed. Older saved workflows that used the previous `Scale-Locked Detailer Hook Provider` input surface will need to be re-wired to the current node inputs.
 
@@ -90,9 +92,8 @@ A public guider node alone does not make FaceDetailer use SLRD. `Scale-Locked De
 The current hook implementation is duck-typed rather than source-verified against a live Impact Pack checkout:
 
 - FaceDetailer mask capture via `post_upscale(...)`,
-- self-anchoring via `post_encode(...)`,
 - alias-tolerant request capture via `pre_ksample(...)`,
-- masked latent correction via `pre_decode(...)`.
+- masked latent correction via the custom sampler/runtime path.
 
 Because Impact Pack's internal contracts can move, this integration should be treated as unverified runtime glue until it is exercised against the current Impact Pack source.
 

nodes.py

@@ -9,11 +9,12 @@ import math
 import torch
 import comfy.samplers
 
-from .slrd_core import build_nested_noise, clone_latent, init_scale_lock_state, resize_4d_tensor
+from .slrd_core import build_nested_noise, clone_latent, init_scale_lock_state, resize_4d_tensor, resize_mask
 from .slrd_runtime import (
     LOCK_SCHEDULE_OPTIONS,
     MID_SCHEDULE_OPTIONS,
     ScaleLockConfig,
+    _resolve_sampler_device,
     apply_scale_lock_to_noise_prediction,
     build_runtime_context_from_advanced,
     clean_latent,
@@ -28,6 +29,7 @@ from .slrd_runtime import (
     sample_with_runtime,
     apply_scale_lock_to_guider,
     clone_guider_for_scale_lock,
+    restore_original_predict_noise,
 )
 
 
@@ -654,6 +656,21 @@ def _combine_masks(*masks: torch.Tensor | None) -> torch.Tensor | None:
     return None if combined is None else combined.clamp(0.0, 1.0)
 
 
+_DETAILER_SUPPORT_MASK_EPS = 1e-3
+
+
+def _detailer_support_mask_for_like(mask: torch.Tensor | None, like: torch.Tensor) -> torch.Tensor | None:
+    expanded = _expand_mask_for_like(mask, like)
+    if expanded is None:
+        return None
+
+    # Impact detailer masks are compositing weights. Preserving their soft latent-cell
+    # variation turns them into a spatially varying SLRD strength field, which causes
+    # blotchy low-frequency corrections inside smooth regions. SLRD should see a support
+    # mask here: on inside the selected region, off outside it.
+    return expanded.gt(_DETAILER_SUPPORT_MASK_EPS).to(dtype=like.dtype)
+
+
 def _canonicalize_impact_noise_mask(mask: Any) -> torch.Tensor | None:
     if not isinstance(mask, torch.Tensor):
         return None
@@ -673,6 +690,29 @@ def _canonicalize_impact_noise_mask(mask: Any) -> torch.Tensor | None:
     return None
 
 
+def _resolve_impact_sampling_mask(
+    latent_samples: torch.Tensor,
+    latent_mask: Any,
+    denoise_mask: Any,
+) -> torch.Tensor | None:
+    target_hw = tuple(latent_samples.shape[-2:])
+    batch = latent_samples.shape[0]
+
+    def _normalize(mask: Any) -> torch.Tensor | None:
+        canonical = _canonicalize_impact_noise_mask(mask)
+        if canonical is None:
+            return None
+        resized = resize_mask(canonical, target_hw, batch, 1)
+        if resized is None:
+            return None
+        return resized[:, :1, :, :].squeeze(1).contiguous()
+
+    sampler_mask = _normalize(denoise_mask)
+    if sampler_mask is not None:
+        return sampler_mask
+    return _normalize(latent_mask)
+
+
 def _impact_request_tuple_from_kwargs(kwargs: dict[str, Any]) -> tuple[Any, ...]:
     if not kwargs:
         return ()
@@ -765,32 +805,6 @@ def _resolve_impact_ag_guider_template(request: "_ImpactSampleRequest", model_wr
         return candidate
     return None
 
-
-def _resolve_sampler_device(model_or_wrap, fallback: torch.device) -> torch.device:
-    """
-    Determine the device to use for sampling by inspecting the provided model or wrapper.
-    
-    Parameters:
-        model_or_wrap: An object (model or wrapper) to inspect for a `load_device` attribute; common wrapper fields (`inner_model`, `model`, `model_patcher`) are checked as well.
-        fallback (torch.device): Device to return if no `load_device` attribute is found on the inspected objects.
-    
-    Returns:
-        torch.device: The first `load_device` found on `model_or_wrap` or its common attributes, otherwise `fallback`.
-    """
-    for obj in (
-        model_or_wrap,
-        getattr(model_or_wrap, "inner_model", None),
-        getattr(model_or_wrap, "model", None),
-        getattr(model_or_wrap, "model_patcher", None),
-    ):
-        if obj is None:
-            continue
-        device = getattr(obj, "load_device", None)
-        if device is not None:
-            return device
-    return fallback
-
-
 def _normalize_sampler_extra_args(
     extra_args: dict[str, Any] | None,
     sigmas: torch.Tensor,
@@ -855,7 +869,8 @@ class _ScaleLockedPlannerGuiderProxy:
         """
         del seed
         if callback is None:
-            callback = lambda *args, **kwargs: None
+            def callback(*args, **kwargs):
+                return None
 
         target_device = _resolve_sampler_device(self._model_wrap, latent_samples.device)
         target_dtype = latent_samples.dtype
@@ -1062,7 +1077,7 @@ class _ScaleLockedImpactSampler:
 
     def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
         """
-        Dispatches a scale-locked sampling run through the underlying sampler using a guider proxy and device-aware tensors.
+        Dispatches a scale-locked sampling run through the underlying sampler using an effective guider and device-aware tensors.
         
         Parameters:
             model_wrap: Model wrapper or model object used to resolve device information and guider behavior.
@@ -1090,38 +1105,43 @@ class _ScaleLockedImpactSampler:
             if latent_image is not None:
                 planner_latent["samples"] = latent_image
 
-            planner_mask = _canonicalize_impact_noise_mask(planner_latent.get("noise_mask", None))
-            if planner_mask is None:
-                planner_mask = _canonicalize_impact_noise_mask(denoise_mask)
+            planner_mask = _resolve_impact_sampling_mask(
+                planner_latent["samples"],
+                planner_latent.get("noise_mask", None),
+                denoise_mask,
+            )
             if planner_mask is not None:
                 planner_latent["noise_mask"] = planner_mask
+            else:
+                planner_latent.pop("noise_mask", None)
 
             base_sampler = comfy.samplers.sampler_object(request.sampler_name)
+            effective_guider = _build_impact_effective_guider(request, model_wrap, extra_args)
             state = self._hook._prepare_runtime_state_for_sampler(
                 request=request,
-                model_wrap=model_wrap,
+                guider=effective_guider,
                 sampler=base_sampler,
                 sigmas=sigmas,
-                extra_args=extra_args,
                 live_latent=planner_latent,
+                runtime_mask=planner_mask,
             )
 
-            proxy = _ScaleLockedGuiderProxy(model_wrap, state.runtime, state.config)
-
-            sampling_noise = noise
             fallback_device = latent_image.device if latent_image is not None else noise.device
             target_device = _resolve_sampler_device(model_wrap, fallback_device)
             target_dtype = latent_image.dtype if latent_image is not None else noise.dtype
+            sampling_mask = planner_mask
 
             prepared_noise = state.runtime.highres_noise
-            if tuple(prepared_noise.shape) == tuple(noise.shape):
+            use_prepared_noise = tuple(prepared_noise.shape) == tuple(noise.shape)
+
+            if use_prepared_noise:
                 sampling_noise = prepared_noise.to(
                     device=target_device,
                     dtype=target_dtype,
                     non_blocking=True,
                 ).clone()
             else:
-                sampling_noise = sampling_noise.to(
+                sampling_noise = noise.to(
                     device=target_device,
                     dtype=target_dtype,
                     non_blocking=True,
@@ -1134,22 +1154,24 @@ class _ScaleLockedImpactSampler:
                 non_blocking=True,
             )
 
-            sampling_mask = denoise_mask
             sigmas = sigmas.to(device=target_device, non_blocking=True)
             if isinstance(sampling_mask, torch.Tensor):
                 sampling_mask = sampling_mask.to(device=target_device, non_blocking=True)
 
-            sample_extra_args = _normalize_sampler_extra_args(extra_args, sigmas, target_device)
-            return base_sampler.sample(
-                proxy,
-                sigmas,
-                sample_extra_args,
-                callback,
-                sampling_noise,
-                latent_image=sampling_latent,
-                denoise_mask=sampling_mask,
-                disable_pbar=disable_pbar,
-            )
+            apply_scale_lock_to_guider(effective_guider, state.runtime, state.config)
+            try:
+                return effective_guider.sample(
+                    sampling_noise,
+                    sampling_latent,
+                    base_sampler,
+                    sigmas,
+                    denoise_mask=sampling_mask,
+                    callback=callback,
+                    disable_pbar=disable_pbar,
+                    seed=state.runtime.noise_seed,
+                )
+            finally:
+                restore_original_predict_noise(effective_guider)
         finally:
             self._hook._clear_sampler_state()
 
@@ -1187,34 +1209,51 @@ class _ScaleLockedDetailerHook:
             stage,
         )
 
-    def _effective_config_for_samples(self, samples: torch.Tensor) -> ScaleLockConfig:
+    def _effective_config_for_samples(
+        self,
+        samples: torch.Tensor,
+        runtime_mask: torch.Tensor | None = None,
+    ) -> ScaleLockConfig:
         cfg = self._settings.config
 
-        face_mask = _expand_mask_for_like(self._upscale_mask, samples)
-        lock_mask = _combine_masks(face_mask, _expand_mask_for_like(cfg.spatial_mask, samples))
-        manifold_source = cfg.manifold_spatial_mask if cfg.manifold_spatial_mask is not None else cfg.spatial_mask
-        manifold_mask = _combine_masks(face_mask, _expand_mask_for_like(manifold_source, samples))
+        detailer_support = _detailer_support_mask_for_like(
+            runtime_mask if runtime_mask is not None else self._upscale_mask,
+            samples,
+        )
+
+        # In the detailer hook path, SLRD must not extend beyond the region that the
+        # detailer actually denoises/composites for this crop. Treat the live sampler mask
+        # or captured upscale mask as the default support mask, then let explicit user masks
+        # further restrict that support.
+        explicit_lock_mask = _expand_mask_for_like(cfg.spatial_mask, samples)
+        lock_mask = _combine_masks(detailer_support, explicit_lock_mask)
+
+        explicit_manifold_source = (
+            cfg.manifold_spatial_mask if cfg.manifold_spatial_mask is not None else cfg.spatial_mask
+        )
+        explicit_manifold_mask = _expand_mask_for_like(explicit_manifold_source, samples)
+        manifold_mask = _combine_masks(detailer_support, explicit_manifold_mask)
         return replace(cfg, spatial_mask=lock_mask, manifold_spatial_mask=manifold_mask)
 
     def _prepare_runtime_state_for_sampler(
         self,
         request: _ImpactSampleRequest,
-        model_wrap,
+        guider,
         sampler,
         sigmas,
-        extra_args: dict[str, Any] | None,
         live_latent: dict[str, Any] | None = None,
+        runtime_mask: torch.Tensor | None = None,
     ) -> _ScaleLockedImpactRuntimeState:
         """
         Builds and registers a scale-locked impact runtime state for the given impact sampling request.
         
         Parameters:
             request (_ImpactSampleRequest): Canonicalized impact sampling request containing model, seed, sampler name, and latent.
-            model_wrap: Model or model wrapper used to construct the planner guider proxy and to resolve runtime device/dtype.
+            guider: Effective guider used for the planner runtime build.
             sampler: Sampler instance to be used for planning and runtime creation.
             sigmas (torch.Tensor): Noise schedule tensor to use for runtime construction.
-            extra_args (dict | None): Optional extra model arguments forwarded into the planner guider proxy (e.g., model_options).
             live_latent (dict | None): Optional override latent to use for planning instead of request.latent.
+            runtime_mask (torch.Tensor | None): Optional latent-space mask actually used by the detailer sampler for this cycle.
         
         Returns:
             _ScaleLockedImpactRuntimeState: The created runtime state containing the original request, the prepared runtime context, and the effective ScaleLockConfig computed for the planner latent.
@@ -1222,10 +1261,9 @@ class _ScaleLockedDetailerHook:
         guard_sampler_alignment(request.sampler_name, self._settings.sampler_guard)
         planner_noise = _ScaleLockedPlannerNoise(request.seed, disable_noise=not self._settings.add_noise)
         planner_latent = request.latent if live_latent is None else live_latent
-        planner_guider = _ScaleLockedPlannerGuiderProxy(model_wrap, extra_args)
         runtime = build_runtime_context_from_advanced(
             noise=planner_noise,
-            guider=planner_guider,
+            guider=guider,
             sampler=sampler,
             sigmas=sigmas,
             latent_image=planner_latent,
@@ -1236,7 +1274,7 @@ class _ScaleLockedDetailerHook:
         state = _ScaleLockedImpactRuntimeState(
             request=request,
             runtime=runtime,
-            config=self._effective_config_for_samples(runtime.highres_latent["samples"]),
+            config=self._effective_config_for_samples(runtime.highres_latent["samples"], runtime_mask=runtime_mask),
         )
         self._active_runtime = state
         return state

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "scale-locked-residual-diffusion"
 description = "A ComfyUI custom node pack implementing Scale-Locked Residual Diffusion for high-resolution composition and anatomy stability."
-version = "1.0.21"
+version = "1.0.32"
 license = { file = "LICENSE" }
 
 [project.urls]