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,63 @@ 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
+    if mask.ndim == 4:
+        if mask.shape[1] == 0:
+            return None
+        return mask[:, :1, :, :].squeeze(1).contiguous()
+    if mask.ndim == 3:
+        return mask.contiguous()
+    if mask.ndim == 2:
+        return mask.unsqueeze(0).contiguous()
+    logger.warning(
+        "ScaleLockedDetailerHook: ignoring unsupported denoise_mask rank %s with shape %s.",
+        mask.ndim,
+        tuple(mask.shape),
+    )
+    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 ()
@@ -719,6 +778,18 @@ def _clear_impact_ag_guider(owner) -> None:
 
 
 def _resolve_impact_ag_guider_template(request: "_ImpactSampleRequest", model_wrap):
+    """
+    Locate an existing impact guider template from the provided request or model wrapper and return it.
+    
+    If a candidate guider is found on either model_wrap.model_patcher or request.model, the guider is returned and the internal `_ag_detailer_guider` attribute is cleared from both potential owners as a side effect.
+    
+    Parameters:
+        request (_ImpactSampleRequest): Impact request that may contain a `model` owning a guider.
+        model_wrap: Model wrapper that may contain a `model_patcher` owning a guider.
+    
+    Returns:
+        The found guider instance, or `None` if no guider template is present.
+    """
     owners = (
         getattr(model_wrap, "model_patcher", None),
         getattr(request, "model", None),
@@ -734,8 +805,195 @@ def _resolve_impact_ag_guider_template(request: "_ImpactSampleRequest", model_wr
         return candidate
     return None
 
+def _normalize_sampler_extra_args(
+    extra_args: dict[str, Any] | None,
+    sigmas: torch.Tensor,
+    target_device: torch.device,
+) -> dict[str, Any]:
+    """
+    Ensure the provided extra_args include a `model_options.sigmas` tensor moved to the target device.
+    
+    If `extra_args` contains a `model_options` mapping, this function copies it, sets its `sigmas` entry to `sigmas` converted to `target_device` (non-blocking), and returns a shallow-copied dict with the updated `model_options`. If `extra_args` is None or `model_options` is not a dict, returns a shallow copy of `extra_args` (or an empty dict when None).
+    
+    Parameters:
+        extra_args (dict[str, Any] | None): Extra sampler arguments that may include a `model_options` dict.
+        sigmas (torch.Tensor): Sigmas tensor to insert into `model_options`.
+        target_device (torch.device): Device to which `sigmas` will be moved.
+    
+    Returns:
+        dict[str, Any]: A normalized copy of `extra_args` with `model_options.sigmas` set to `sigmas` on `target_device` when applicable.
+    """
+    normalized = {} if extra_args is None else dict(extra_args)
+    model_options = normalized.get("model_options", None)
+    if isinstance(model_options, dict):
+        model_options = dict(model_options)
+        model_options["sigmas"] = sigmas.to(device=target_device, non_blocking=True)
+        normalized["model_options"] = model_options
+    return normalized
+
+
+class _ScaleLockedPlannerGuiderProxy:
+    def __init__(self, model_wrap, extra_args: dict[str, Any] | None):
+        """
+        Initialize the planner guider proxy which forwards sampling calls to the underlying model wrapper.
+        
+        Parameters:
+            model_wrap: An object that exposes a sampler-compatible interface and may provide a `model_patcher` attribute.
+            extra_args (dict[str, Any] | None): Optional model/sampler options to be merged and stored for use when sampling; a shallow copy is made if provided.
+        """
+        self.model_patcher = getattr(model_wrap, "model_patcher", None)
+        self._model_wrap = model_wrap
+        self._extra_args = {} if extra_args is None else dict(extra_args)
+
+    def sample(
+        self,
+        noise,
+        latent_samples,
+        sampler,
+        sigmas,
+        denoise_mask=None,
+        callback=None,
+        disable_pbar=False,
+        seed=None,
+    ):
+        """
+        Prepare inputs on the sampler's device/dtype, normalize extra args, and forward the call to the underlying sampler.
+        
+        Parameters:
+            denoise_mask (torch.Tensor | None): Optional mask moved to the sampler device when provided.
+            callback (callable | None): Optional progress callback; if None a no-op callback is used.
+            seed: Ignored by this proxy.
+        
+        Returns:
+            The value returned by the underlying sampler.sample(...) call.
+        """
+        del seed
+        if callback is None:
+            def callback(*args, **kwargs):
+                return None
+
+        target_device = _resolve_sampler_device(self._model_wrap, latent_samples.device)
+        target_dtype = latent_samples.dtype
+        latent_samples = latent_samples.to(
+            device=target_device,
+            dtype=target_dtype,
+            non_blocking=True,
+        )
+        noise = noise.to(device=target_device, dtype=target_dtype, non_blocking=True)
+        sigmas = sigmas.to(device=target_device, non_blocking=True)
+        if isinstance(denoise_mask, torch.Tensor):
+            denoise_mask = denoise_mask.to(device=target_device, non_blocking=True)
+
+        sample_extra_args = _normalize_sampler_extra_args(self._extra_args, sigmas, target_device)
+        return sampler.sample(
+            self._model_wrap,
+            sigmas,
+            sample_extra_args,
+            callback,
+            noise,
+            latent_image=latent_samples,
+            denoise_mask=denoise_mask,
+            disable_pbar=disable_pbar,
+        )
+
+
+class _ScaleLockedGuiderProxy:
+    def __init__(self, base_guider, runtime, config: ScaleLockConfig):
+        """
+        Wraps a base guider and initializes its scale-lock state using the provided runtime context and configuration.
+        
+        Parameters:
+            base_guider: The underlying guider object to delegate to; must implement the guider interface used during sampling.
+            runtime: Runtime context providing model, anchors_x0, and planner_sigmas required to initialize scale-lock state.
+            config (ScaleLockConfig): Configuration object containing all scale-lock parameters (strengths, cutoffs, schedules, manifold controls, spatial masks, etc.) used to initialize the guider's scale-lock behavior.
+        """
+        self._base_guider = base_guider
+        init_scale_lock_state(
+            self,
+            model=runtime.model,
+            anchors_x0_cpu=runtime.anchors_x0,
+            planner_sigmas=runtime.planner_sigmas,
+            lock_strength=config.lock_strength,
+            lock_strength_start=config.lock_strength_start,
+            lock_strength_end=config.lock_strength_end,
+            cutoff=config.cutoff,
+            mid_cutoff=config.mid_cutoff,
+            mid_strength=config.mid_strength,
+            schedule=config.schedule,
+            schedule_power=config.schedule_power,
+            schedule_hold=config.schedule_hold,
+            mid_strength_start=config.mid_strength_start,
+            mid_strength_end=config.mid_strength_end,
+            mid_schedule=config.mid_schedule,
+            mid_schedule_power=config.mid_schedule_power,
+            mid_schedule_hold=config.mid_schedule_hold,
+            spatial_mask=config.spatial_mask,
+            manifold_enabled=config.manifold_enabled,
+            manifold_strength=config.manifold_strength,
+            manifold_strength_start=config.manifold_strength_start,
+            manifold_strength_end=config.manifold_strength_end,
+            manifold_schedule=config.manifold_schedule,
+            manifold_schedule_power=config.manifold_schedule_power,
+            manifold_schedule_hold=config.manifold_schedule_hold,
+            manifold_cutoff=config.manifold_cutoff,
+            manifold_radial_strength=config.manifold_radial_strength,
+            manifold_anisotropy=config.manifold_anisotropy,
+            manifold_translation_strength=config.manifold_translation_strength,
+            manifold_anchor_mix=config.manifold_anchor_mix,
+            manifold_mean_anchor_mix=config.manifold_mean_anchor_mix,
+            manifold_contrast_restore=config.manifold_contrast_restore,
+            manifold_energy_tether=config.manifold_energy_tether,
+            manifold_channel_tether=config.manifold_channel_tether,
+            manifold_energy_gain_cap=config.manifold_energy_gain_cap,
+            manifold_max_shift_px=config.manifold_max_shift_px,
+            manifold_spatial_mask=config.manifold_spatial_mask,
+        )
+
+    def __getattr__(self, name: str):
+        """
+        Delegate attribute lookup to the wrapped base guider.
+        
+        Parameters:
+            name (str): The attribute name to retrieve from the wrapped guider.
+        
+        Returns:
+            Any: The value of the requested attribute from the wrapped base guider.
+        """
+        return getattr(self._base_guider, name)
+
+    def __call__(self, x, timestep, model_options=None, seed=None):
+        """
+        Run the wrapped guider to produce a noise prediction then apply scale-lock adjustments.
+        
+        Parameters:
+            x (torch.Tensor): Input latent tensor for guidance.
+            timestep: Timestep value passed to the guider (e.g., scheduler timestep).
+            model_options (dict, optional): Additional model options forwarded to the base guider.
+            seed (int | None, optional): RNG seed forwarded to the base guider.
+        
+        Returns:
+            torch.Tensor: Noise prediction after applying scale-lock modifications.
+        """
+        if model_options is None:
+            model_options = {}
+        base_noise = self._base_guider(x, timestep, model_options=model_options, seed=seed)
+        return apply_scale_lock_to_noise_prediction(self, base_noise, x, timestep)
+
 
 def _sync_impact_guider_conditions(guider, positive, negative) -> None:
+    """
+    Synchronizes positive and negative conditioning on a guider object.
+    
+    Attempts to apply the provided positive and negative conditioning to the guider.
+    If the guider exposes set_conds, that method is called with (positive, negative).
+    If it exposes inner_set_conds, that method is called with {"positive": positive, "negative": negative}.
+    If neither method is present, a warning is logged. Failures during the call are caught and logged as warnings.
+    
+    Parameters:
+        guider (object): The guider instance to update; may implement `set_conds` or `inner_set_conds`.
+        positive: The positive conditioning to apply (type depends on guider implementation).
+        negative: The negative conditioning to apply (type depends on guider implementation).
+    """
     try:
         if hasattr(guider, "set_conds"):
             guider.set_conds(positive, negative)
@@ -818,6 +1076,25 @@ class _ScaleLockedImpactSampler:
         self._hook = hook
 
     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 an effective guider and device-aware tensors.
+        
+        Parameters:
+            model_wrap: Model wrapper or model object used to resolve device information and guider behavior.
+            sigmas (torch.Tensor): Noise schedule tensor for the sampler; will be moved to the target device.
+            extra_args (dict | None): Additional sampler/model options; may be normalized to include device-sharded sigmas.
+            callback: Progress/callback object passed through to the underlying sampler.
+            noise (torch.Tensor): Initial noise tensor for sampling; may be replaced with prepared high-resolution noise if shapes match.
+            latent_image (dict | torch.Tensor | None): Optional high-resolution latent to sample from; if a dict is provided its "samples" entry is used.
+            denoise_mask (torch.Tensor | None): Optional denoising mask to pass to the sampler.
+            disable_pbar (bool): If true, disables progress bar propagation to the underlying sampler.
+        
+        Raises:
+            RuntimeError: If no pending impact sampling request was captured before this sampler was invoked.
+        
+        Returns:
+            The value returned by the underlying sampler's sample(...) call.
+        """
         request = self._hook._pending_request
         if request is None:
             raise RuntimeError("ScaleLockedDetailerHook: sampler request was not captured before the custom sampler was used.")
@@ -827,66 +1104,74 @@ class _ScaleLockedImpactSampler:
             planner_latent = clone_latent(request.latent)
             if latent_image is not None:
                 planner_latent["samples"] = latent_image
-            if isinstance(denoise_mask, torch.Tensor):
-                planner_latent["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)
-            guider_template = (
-                model_wrap
-                if _is_impact_guider_like(model_wrap)
-                else _resolve_impact_ag_guider_template(request, model_wrap)
-            )
-            planner_guider = _build_impact_effective_guider(
-                request,
-                model_wrap,
-                extra_args,
-                guider_template=guider_template,
-            )
+            effective_guider = _build_impact_effective_guider(request, model_wrap, extra_args)
             state = self._hook._prepare_runtime_state_for_sampler(
                 request=request,
-                guider=planner_guider,
+                guider=effective_guider,
                 sampler=base_sampler,
                 sigmas=sigmas,
                 live_latent=planner_latent,
+                runtime_mask=planner_mask,
             )
 
-            sample_guider = _build_impact_effective_guider(
-                request,
-                model_wrap,
-                extra_args,
-                guider_template=guider_template,
+            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
+            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 = noise.to(
+                    device=target_device,
+                    dtype=target_dtype,
+                    non_blocking=True,
+                )
+
+            sampling_latent = latent_image if latent_image is not None else state.runtime.highres_latent["samples"]
+            sampling_latent = sampling_latent.to(
+                device=target_device,
+                dtype=target_dtype,
+                non_blocking=True,
             )
-            apply_scale_lock_to_guider(sample_guider, state.runtime, state.config)
 
-            target_device = state.runtime.highres_latent["samples"].device
-            target_dtype = state.runtime.highres_latent["samples"].dtype
-            sampling_noise = state.runtime.highres_noise
-            sampling_latent = state.runtime.highres_latent["samples"]
-            sampling_mask = state.runtime.highres_latent.get("noise_mask", denoise_mask)
-
-            if tuple(sampling_noise.shape) != tuple(noise.shape):
-                sampling_noise = noise
-                if latent_image is not None:
-                    sampling_latent = latent_image
-                else:
-                    sampling_latent = planner_latent["samples"]
-                sampling_mask = denoise_mask
-
-            sampling_noise = sampling_noise.to(device=target_device, dtype=target_dtype, non_blocking=True)
-            sampling_latent = sampling_latent.to(device=target_device, dtype=target_dtype, non_blocking=True)
+            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)
 
-            return sample_guider.sample(
-                sampling_noise,
-                sampling_latent,
-                base_sampler,
-                state.runtime.sigmas,
-                denoise_mask=sampling_mask,
-                callback=callback,
-                disable_pbar=disable_pbar,
-                seed=state.runtime.noise_seed,
-            )
+            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()
 
@@ -924,13 +1209,30 @@ 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(
@@ -940,7 +1242,22 @@ class _ScaleLockedDetailerHook:
         sampler,
         sigmas,
         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.
+            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.
+            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.
+        """
         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
@@ -957,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.18"
+version = "1.0.32"
 license = { file = "LICENSE" }
 
 [project.urls]