📝 Add docstrings to codex/fix-facedetailer-latent-mask-mixup

Docstrings generation was requested by @xmarre.

* https://github.com/xmarre/ComfyUI-ScaleLockedResidualDiffusion/pull/10#issuecomment-4062975146

The following files were modified:

* `nodes.py`

README.md

@@ -55,7 +55,17 @@ 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 with a provider-first adapter surface (`get_custom_sampler()` and related aliases) plus a `pre_ksample(...)` fallback path. The goal is to run the full SLRD bundle per crop: planner pass, anchor recording, nested-noise construction, and scale-locked final sample. This has been validated for import/compile in this repo, but not against a live current Impact Pack checkout in this environment.
+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 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:
+
+- residual lock strength and cutoffs,
+- optional manifold companding controls,
+- optional `lock_mask` / `manifold_mask`, which are intersected with the FaceDetailer 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.
+
+This has been validated for import/compile in this repo, but not against a live current Impact Pack checkout in this environment.
 
 ### 6. Scale-Locked Nested Noise Preview
 
@@ -75,13 +85,14 @@ That path uses the same SLRD runtime pieces as the all-in-one sampler instead of
 
 ## FaceDetailer / Impact Pack note
 
-A public guider node alone does not make FaceDetailer use SLRD. Detailer crop sampling needs the full planner-plus-noise bundle, not only the denoiser correction. `Scale-Locked Detailer Hook Provider` is the experimental integration point for that path.
+A public guider node alone does not make FaceDetailer use SLRD. `Scale-Locked Detailer Hook Provider` is the experimental integration point for applying the masked residual/manifold correction inside the FaceDetailer crop lifecycle.
 
 The current hook implementation is duck-typed rather than source-verified against a live Impact Pack checkout:
 
-- provider-first via `get_custom_sampler()` / provider aliases,
-- fallback interception via `pre_ksample(...)`,
-- crop execution delegated to the same shared SLRD runtime as the regular sampler nodes.
+- FaceDetailer mask capture via `post_upscale(...)`,
+- self-anchoring via `post_encode(...)`,
+- alias-tolerant request capture via `pre_ksample(...)`,
+- masked latent correction via `pre_decode(...)`.
 
 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.
 

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.0"
+version = "1.0.21"
 license = { file = "LICENSE" }
 
 [project.urls]

slrd_core.py

@@ -155,7 +155,7 @@ def build_nested_noise(
     hf = hf - hf_low_up
 
     std = hf.std(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
-    hf = hf / std
+    hf = (hf / std).to(device=device, non_blocking=True)
 
     out = base + float(hf_strength) * hf
     out_std = out.std(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
@@ -210,6 +210,286 @@ def residual_lock_multiband(
     return base_high + torch.lerp(base_low, anchor_low, low_strength) + torch.lerp(base_mid, anchor_mid, mid_strength)
 
 
+def _base_grid(batch: int, height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    yy, xx = torch.meshgrid(
+        torch.linspace(-1.0, 1.0, height, device=device, dtype=dtype),
+        torch.linspace(-1.0, 1.0, width, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    return torch.stack([xx, yy], dim=-1).unsqueeze(0).expand(batch, -1, -1, -1).contiguous()
+
+
+def warp_4d_tensor(x: torch.Tensor, flow_xy_norm: torch.Tensor, mode: str = "bilinear") -> torch.Tensor:
+    if x.ndim != 4:
+        raise ValueError(f"Expected BCHW tensor, got shape {tuple(x.shape)}")
+    batch, _, height, width = x.shape
+    expected = (batch, height, width, 2)
+    if tuple(flow_xy_norm.shape) != expected:
+        raise ValueError(f"Expected flow shape {expected}, got {tuple(flow_xy_norm.shape)}")
+    base_grid = _base_grid(batch, height, width, x.device, x.dtype)
+    sample_grid = (base_grid + flow_xy_norm).clamp(-1.25, 1.25)
+    return F.grid_sample(
+        x,
+        sample_grid,
+        mode=mode,
+        padding_mode="border",
+        align_corners=True,
+    )
+
+
+def _expand_mask_channels(mask: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
+    if mask.ndim != 4:
+        raise ValueError(f"Expected mask tensor BCHW, got shape {tuple(mask.shape)}")
+    if mask.shape[0] < like.shape[0]:
+        repeat = math.ceil(like.shape[0] / max(1, mask.shape[0]))
+        mask = mask.repeat(repeat, 1, 1, 1)[: like.shape[0]]
+    elif mask.shape[0] > like.shape[0]:
+        mask = mask[: like.shape[0]]
+    if mask.shape[1] == like.shape[1]:
+        return mask
+    if mask.shape[1] == 1:
+        return mask.expand(like.shape[0], like.shape[1], like.shape[-2], like.shape[-1]).contiguous()
+    return mask.mean(dim=1, keepdim=True).expand(like.shape[0], like.shape[1], like.shape[-2], like.shape[-1]).contiguous()
+
+
+def _latent_activity_map(x: torch.Tensor, mask_1ch: torch.Tensor) -> torch.Tensor:
+    mass = mask_1ch.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
+    mean = (x * mask_1ch).sum(dim=(-2, -1), keepdim=True) / mass
+    centered = x - mean
+    activity = centered.square().mean(dim=1, keepdim=True)
+    activity = activity * mask_1ch
+    activity = activity + mask_1ch * 1e-8
+    return activity
+
+
+def _masked_spatial_stats(x: torch.Tensor, mask_1ch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch, _, height, width = x.shape
+    yy, xx = torch.meshgrid(
+        torch.linspace(-1.0, 1.0, height, device=x.device, dtype=x.dtype),
+        torch.linspace(-1.0, 1.0, width, device=x.device, dtype=x.dtype),
+        indexing="ij",
+    )
+    xx = xx.view(1, 1, height, width).expand(batch, -1, -1, -1)
+    yy = yy.view(1, 1, height, width).expand(batch, -1, -1, -1)
+
+    activity = _latent_activity_map(x, mask_1ch)
+    mass = activity.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
+    cx = (activity * xx).sum(dim=(-2, -1), keepdim=True) / mass
+    cy = (activity * yy).sum(dim=(-2, -1), keepdim=True) / mass
+
+    dx = xx - cx
+    dy = yy - cy
+    rx = torch.sqrt((activity * dx.square()).sum(dim=(-2, -1), keepdim=True) / mass).clamp_min(1e-4)
+    ry = torch.sqrt((activity * dy.square()).sum(dim=(-2, -1), keepdim=True) / mass).clamp_min(1e-4)
+    return cx, cy, rx, ry
+
+
+def estimate_latent_compaction_flow(
+    anchor_low: torch.Tensor,
+    base_low: torch.Tensor,
+    mask_1ch: torch.Tensor,
+    strength: float,
+    radial_strength: float,
+    anisotropy: float,
+    translation_strength: float,
+    max_shift_px: float,
+) -> torch.Tensor:
+    batch, _, height, width = base_low.shape
+    anchor_cx, anchor_cy, anchor_rx, anchor_ry = _masked_spatial_stats(anchor_low, mask_1ch)
+    base_cx, base_cy, base_rx, base_ry = _masked_spatial_stats(base_low, mask_1ch)
+
+    ratio_x = (base_rx / anchor_rx.clamp_min(1e-5)).clamp(0.85, 1.25)
+    ratio_y = (base_ry / anchor_ry.clamp_min(1e-5)).clamp(0.85, 1.25)
+    outward_x = (ratio_x - 1.0).clamp(min=0.0)
+    outward_y = (ratio_y - 1.0).clamp(min=0.0)
+
+    yy, xx = torch.meshgrid(
+        torch.linspace(-1.0, 1.0, height, device=base_low.device, dtype=base_low.dtype),
+        torch.linspace(-1.0, 1.0, width, device=base_low.device, dtype=base_low.dtype),
+        indexing="ij",
+    )
+    xx = xx.view(1, 1, height, width).expand(batch, -1, -1, -1)
+    yy = yy.view(1, 1, height, width).expand(batch, -1, -1, -1)
+
+    dx = xx - base_cx
+    dy = yy - base_cy
+    ex = dx / base_rx.clamp_min(1e-5)
+    ey = dy / base_ry.clamp_min(1e-5)
+    radius = torch.sqrt(ex.square() + ey.square() + 1e-8)
+    edge_envelope = torch.clamp(radius / 1.25, 0.0, 1.0)
+
+    smooth_mask = lowpass_latent(mask_1ch, 0.35).clamp(0.0, 1.0)
+    mean_outward = 0.5 * (outward_x + outward_y)
+    axis_x = mean_outward * float(radial_strength) + (outward_x - mean_outward) * float(anisotropy)
+    axis_y = mean_outward * float(radial_strength) + (outward_y - mean_outward) * float(anisotropy)
+
+    shift_x = ex * edge_envelope * smooth_mask * float(strength) * axis_x
+    shift_y = ey * edge_envelope * smooth_mask * float(strength) * axis_y
+
+    trans_x = (base_cx - anchor_cx) * smooth_mask * float(strength) * float(translation_strength)
+    trans_y = (base_cy - anchor_cy) * smooth_mask * float(strength) * float(translation_strength)
+
+    max_shift_norm_x = 2.0 * float(max_shift_px) / max(width - 1, 1)
+    max_shift_norm_y = 2.0 * float(max_shift_px) / max(height - 1, 1)
+
+    shift_x = (shift_x + trans_x).clamp(-max_shift_norm_x, max_shift_norm_x)
+    shift_y = (shift_y + trans_y).clamp(-max_shift_norm_y, max_shift_norm_y)
+    return torch.stack([shift_x[:, 0], shift_y[:, 0]], dim=-1)
+
+
+def _weighted_channel_stats(x: torch.Tensor, mask_1ch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    mass = mask_1ch.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
+    mean = (x * mask_1ch).sum(dim=(-2, -1), keepdim=True) / mass
+    var = (((x - mean).square()) * mask_1ch).sum(dim=(-2, -1), keepdim=True) / mass
+    std = torch.sqrt(var.clamp_min(1e-8))
+    return mean, std
+
+
+def _weighted_global_std(x: torch.Tensor, mask_1ch: torch.Tensor, mean: Optional[torch.Tensor] = None) -> torch.Tensor:
+    if mean is None:
+        mean, _ = _weighted_channel_stats(x, mask_1ch)
+    mass = mask_1ch.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
+    denom = mass * float(max(1, x.shape[1]))
+    var = (((x - mean).square()) * mask_1ch).sum(dim=(1, 2, 3), keepdim=True) / denom
+    return torch.sqrt(var.clamp_min(1e-8))
+
+
+def _bounded_ratio(target: torch.Tensor, source: torch.Tensor, gain_cap: float) -> torch.Tensor:
+    gain_cap = float(max(1.0, gain_cap))
+    lo = 1.0 / gain_cap
+    hi = gain_cap
+    return (target / source.clamp_min(1e-5)).clamp(lo, hi)
+
+
+def tether_latent_low_frequency_energy(
+    base_low: torch.Tensor,
+    anchor_low: torch.Tensor,
+    mask_1ch: torch.Tensor,
+    energy_tether: float,
+    channel_tether: float,
+    gain_cap: float,
+) -> torch.Tensor:
+    energy_tether = float(max(0.0, min(1.0, energy_tether)))
+    channel_tether = float(max(0.0, min(1.0, channel_tether)))
+    if energy_tether <= 0.0 and channel_tether <= 0.0:
+        return base_low
+
+    base_mean, _ = _weighted_channel_stats(base_low, mask_1ch)
+    anchor_mean, anchor_std = _weighted_channel_stats(anchor_low, mask_1ch)
+
+    regulated = base_low
+    if energy_tether > 0.0:
+        base_global_std = _weighted_global_std(regulated, mask_1ch, mean=base_mean)
+        anchor_global_std = _weighted_global_std(anchor_low, mask_1ch, mean=anchor_mean)
+        global_gain = _bounded_ratio(anchor_global_std, base_global_std, gain_cap)
+        global_gain = torch.lerp(torch.ones_like(global_gain), global_gain, energy_tether)
+        regulated = base_mean + (regulated - base_mean) * global_gain
+
+    if channel_tether > 0.0:
+        regulated_mean, regulated_std = _weighted_channel_stats(regulated, mask_1ch)
+        channel_gain = _bounded_ratio(anchor_std, regulated_std, gain_cap)
+        channel_gain = torch.lerp(torch.ones_like(channel_gain), channel_gain, channel_tether)
+        regulated = regulated_mean + (regulated - regulated_mean) * channel_gain
+
+    return regulated
+
+
+def restore_latent_low_frequency_stats(
+    warped_low: torch.Tensor,
+    anchor_low: torch.Tensor,
+    mask_1ch: torch.Tensor,
+    anchor_mix: float,
+    mean_anchor_mix: float,
+    contrast_restore: float,
+) -> torch.Tensor:
+    warped_mean, warped_std = _weighted_channel_stats(warped_low, mask_1ch)
+    anchor_mean, anchor_std = _weighted_channel_stats(anchor_low, mask_1ch)
+
+    mean_target = torch.lerp(warped_mean, anchor_mean, float(mean_anchor_mix))
+    contrast_gain = torch.lerp(
+        torch.ones_like(anchor_std),
+        anchor_std / warped_std.clamp_min(1e-5),
+        float(contrast_restore),
+    )
+    restored = mean_target + (warped_low - warped_mean) * contrast_gain
+    return torch.lerp(restored, anchor_low, float(anchor_mix))
+
+
+def latent_manifold_compand(
+    base_denoised: torch.Tensor,
+    anchor_denoised: torch.Tensor,
+    mask: Optional[torch.Tensor],
+    strength: float,
+    cutoff: float,
+    radial_strength: float,
+    anisotropy: float,
+    translation_strength: float,
+    anchor_mix: float,
+    mean_anchor_mix: float,
+    contrast_restore: float,
+    energy_tether: float,
+    channel_tether: float,
+    energy_gain_cap: float,
+    max_shift_px: float,
+) -> torch.Tensor:
+    strength = float(max(0.0, min(1.0, strength)))
+    if strength <= 0.0:
+        return base_denoised
+
+    cutoff = float(max(0.05, min(1.0, cutoff)))
+    work_dtype = base_denoised.dtype
+    compute_dtype = torch.float32
+
+    base = base_denoised.to(dtype=compute_dtype)
+    anchor = anchor_denoised.to(device=base.device, dtype=compute_dtype)
+    if tuple(anchor.shape[-2:]) != tuple(base.shape[-2:]):
+        anchor = resize_4d_tensor(anchor, tuple(base.shape[-2:]))
+
+    if mask is None:
+        mask_1ch = torch.ones((base.shape[0], 1, base.shape[-2], base.shape[-1]), device=base.device, dtype=compute_dtype)
+    else:
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+        elif mask.ndim == 4 and mask.shape[1] != 1:
+            mask = mask.mean(dim=1, keepdim=True)
+        mask_1ch = _expand_mask_channels(mask.to(device=base.device, dtype=compute_dtype), base)[:, :1].clamp(0.0, 1.0)
+
+    base_low = lowpass_latent(base, cutoff)
+    anchor_low = lowpass_latent(anchor, cutoff)
+    base_high = base - base_low
+
+    regulated_low = tether_latent_low_frequency_energy(
+        base_low,
+        anchor_low,
+        mask_1ch,
+        energy_tether=float(max(0.0, min(1.0, energy_tether))) * strength,
+        channel_tether=float(max(0.0, min(1.0, channel_tether))) * strength,
+        gain_cap=float(max(1.0, energy_gain_cap)),
+    )
+
+    flow = estimate_latent_compaction_flow(
+        anchor_low=anchor_low,
+        base_low=regulated_low,
+        mask_1ch=mask_1ch,
+        strength=strength,
+        radial_strength=radial_strength,
+        anisotropy=anisotropy,
+        translation_strength=translation_strength,
+        max_shift_px=max_shift_px,
+    )
+    warped_low = warp_4d_tensor(regulated_low, flow, mode="bilinear")
+    restored_low = restore_latent_low_frequency_stats(
+        warped_low,
+        anchor_low,
+        mask_1ch,
+        anchor_mix=float(max(0.0, min(1.0, anchor_mix))) * strength,
+        mean_anchor_mix=float(max(0.0, min(1.0, mean_anchor_mix))) * strength,
+        contrast_restore=float(max(0.0, min(1.0, contrast_restore))) * strength,
+    )
+    corrected = base_high + restored_low
+    return corrected.to(dtype=work_dtype)
+
+
 def clamp(value: float, lo: float, hi: float) -> float:
     return float(max(lo, min(hi, value)))
 
@@ -322,6 +602,25 @@ class ScaleLockedCFGGuider:
         mid_schedule_power: float,
         mid_schedule_hold: float,
         spatial_mask: Optional[torch.Tensor],
+        manifold_enabled: bool = False,
+        manifold_strength: float = 0.0,
+        manifold_strength_start: float = 1.0,
+        manifold_strength_end: float = 0.0,
+        manifold_schedule: str = "ease_out",
+        manifold_schedule_power: float = 2.0,
+        manifold_schedule_hold: float = 0.0,
+        manifold_cutoff: float = 0.18,
+        manifold_radial_strength: float = 1.0,
+        manifold_anisotropy: float = 0.15,
+        manifold_translation_strength: float = 1.0,
+        manifold_anchor_mix: float = 0.18,
+        manifold_mean_anchor_mix: float = 0.12,
+        manifold_contrast_restore: float = 0.10,
+        manifold_energy_tether: float = 0.0,
+        manifold_channel_tether: float = 0.0,
+        manifold_energy_gain_cap: float = 1.75,
+        manifold_max_shift_px: float = 3.0,
+        manifold_spatial_mask: Optional[torch.Tensor] = None,
     ) -> None:
         init_scale_lock_state(
             self,
@@ -343,6 +642,25 @@ class ScaleLockedCFGGuider:
             mid_schedule_power=mid_schedule_power,
             mid_schedule_hold=mid_schedule_hold,
             spatial_mask=spatial_mask,
+            manifold_enabled=manifold_enabled,
+            manifold_strength=manifold_strength,
+            manifold_strength_start=manifold_strength_start,
+            manifold_strength_end=manifold_strength_end,
+            manifold_schedule=manifold_schedule,
+            manifold_schedule_power=manifold_schedule_power,
+            manifold_schedule_hold=manifold_schedule_hold,
+            manifold_cutoff=manifold_cutoff,
+            manifold_radial_strength=manifold_radial_strength,
+            manifold_anisotropy=manifold_anisotropy,
+            manifold_translation_strength=manifold_translation_strength,
+            manifold_anchor_mix=manifold_anchor_mix,
+            manifold_mean_anchor_mix=manifold_mean_anchor_mix,
+            manifold_contrast_restore=manifold_contrast_restore,
+            manifold_energy_tether=manifold_energy_tether,
+            manifold_channel_tether=manifold_channel_tether,
+            manifold_energy_gain_cap=manifold_energy_gain_cap,
+            manifold_max_shift_px=manifold_max_shift_px,
+            manifold_spatial_mask=manifold_spatial_mask,
         )
 
     def _slrd_resolve_step_index(self, timestep: torch.Tensor | float | int) -> int:
@@ -357,6 +675,12 @@ class ScaleLockedCFGGuider:
     def _slrd_mask_for(self, like: torch.Tensor) -> Optional[torch.Tensor]:
         return scale_lock_mask_for(self, like)
 
+    def _slrd_manifold_strength_for_step(self, idx: int) -> float:
+        return scale_lock_manifold_strength_for_step(self, idx)
+
+    def _slrd_manifold_mask_for(self, like: torch.Tensor) -> Optional[torch.Tensor]:
+        return scale_lock_manifold_mask_for(self, like)
+
 
 def init_scale_lock_state(
     guider,
@@ -379,6 +703,25 @@ def init_scale_lock_state(
     mid_schedule_power: float,
     mid_schedule_hold: float,
     spatial_mask: Optional[torch.Tensor],
+    manifold_enabled: bool = False,
+    manifold_strength: float = 0.0,
+    manifold_strength_start: float = 1.0,
+    manifold_strength_end: float = 0.0,
+    manifold_schedule: str = "ease_out",
+    manifold_schedule_power: float = 2.0,
+    manifold_schedule_hold: float = 0.0,
+    manifold_cutoff: float = 0.18,
+    manifold_radial_strength: float = 1.0,
+    manifold_anisotropy: float = 0.15,
+    manifold_translation_strength: float = 1.0,
+    manifold_anchor_mix: float = 0.18,
+    manifold_mean_anchor_mix: float = 0.12,
+    manifold_contrast_restore: float = 0.10,
+    manifold_energy_tether: float = 0.0,
+    manifold_channel_tether: float = 0.0,
+    manifold_energy_gain_cap: float = 1.75,
+    manifold_max_shift_px: float = 3.0,
+    manifold_spatial_mask: Optional[torch.Tensor] = None,
 ) -> None:
     guider._slrd_model = model
     guider._slrd_anchors_x0_cpu = list(anchors_x0_cpu)
@@ -401,6 +744,25 @@ def init_scale_lock_state(
     guider._slrd_prev_match_idx = 0
     guider._slrd_spatial_mask = spatial_mask
     guider._slrd_last_sigma = None
+    guider._slrd_manifold_enabled = bool(manifold_enabled)
+    guider._slrd_manifold_strength = float(manifold_strength)
+    guider._slrd_manifold_strength_start = float(manifold_strength_start)
+    guider._slrd_manifold_strength_end = float(manifold_strength_end)
+    guider._slrd_manifold_schedule = manifold_schedule
+    guider._slrd_manifold_schedule_power = float(manifold_schedule_power)
+    guider._slrd_manifold_schedule_hold = float(manifold_schedule_hold)
+    guider._slrd_manifold_cutoff = float(max(0.05, min(1.0, manifold_cutoff)))
+    guider._slrd_manifold_radial_strength = float(manifold_radial_strength)
+    guider._slrd_manifold_anisotropy = float(manifold_anisotropy)
+    guider._slrd_manifold_translation_strength = float(manifold_translation_strength)
+    guider._slrd_manifold_anchor_mix = float(manifold_anchor_mix)
+    guider._slrd_manifold_mean_anchor_mix = float(manifold_mean_anchor_mix)
+    guider._slrd_manifold_contrast_restore = float(manifold_contrast_restore)
+    guider._slrd_manifold_energy_tether = float(max(0.0, min(1.0, manifold_energy_tether)))
+    guider._slrd_manifold_channel_tether = float(max(0.0, min(1.0, manifold_channel_tether)))
+    guider._slrd_manifold_energy_gain_cap = float(max(1.0, manifold_energy_gain_cap))
+    guider._slrd_manifold_max_shift_px = float(manifold_max_shift_px)
+    guider._slrd_manifold_spatial_mask = manifold_spatial_mask if manifold_spatial_mask is not None else spatial_mask
 
 
 def resolve_scale_lock_step_index(guider, timestep: torch.Tensor | float | int) -> int:
@@ -494,6 +856,21 @@ def scale_lock_strengths_for_step(guider, idx: int) -> tuple[float, float]:
     return low_strength, mid_strength
 
 
+def scale_lock_manifold_strength_for_step(guider, idx: int) -> float:
+    if not getattr(guider, "_slrd_manifold_enabled", False):
+        return 0.0
+    progress = scale_lock_progress_for_step(guider, idx)
+    return _scheduled_strength(
+        guider._slrd_manifold_strength,
+        guider._slrd_manifold_strength_start,
+        guider._slrd_manifold_strength_end,
+        progress,
+        guider._slrd_manifold_schedule,
+        guider._slrd_manifold_schedule_power,
+        guider._slrd_manifold_schedule_hold,
+    )
+
+
 def scale_lock_anchor_for(guider, idx: int, like: torch.Tensor) -> torch.Tensor:
     anchor = guider._slrd_anchors_x0_cpu[idx].to(device=like.device, dtype=like.dtype, non_blocking=True)
     if tuple(anchor.shape[-2:]) != tuple(like.shape[-2:]):
@@ -515,5 +892,20 @@ def scale_lock_mask_for(guider, like: torch.Tensor) -> Optional[torch.Tensor]:
     return mask
 
 
+def scale_lock_manifold_mask_for(guider, like: torch.Tensor) -> Optional[torch.Tensor]:
+    manifold_mask = getattr(guider, "_slrd_manifold_spatial_mask", None)
+    if manifold_mask is None:
+        return None
+    mask = manifold_mask.to(device=like.device, dtype=like.dtype, non_blocking=True)
+    if tuple(mask.shape[-2:]) != tuple(like.shape[-2:]):
+        mask = resize_4d_tensor(mask, tuple(like.shape[-2:]))
+    if mask.shape[0] < like.shape[0]:
+        repeat = math.ceil(like.shape[0] / max(1, mask.shape[0]))
+        mask = mask.repeat(repeat, 1, 1, 1)[: like.shape[0]]
+    elif mask.shape[0] > like.shape[0]:
+        mask = mask[: like.shape[0]]
+    return mask
+
+
 
 

slrd_runtime.py

@@ -19,12 +19,15 @@ from .slrd_core import (
     build_nested_noise,
     clone_latent,
     init_scale_lock_state,
+    latent_manifold_compand,
     latent_target_hw_from_megapixels,
     resolve_scale_lock_step_index,
     resize_latent_dict,
     resize_mask,
     residual_lock_multiband,
     scale_lock_anchor_for,
+    scale_lock_manifold_mask_for,
+    scale_lock_manifold_strength_for_step,
     scale_lock_mask_for,
     scale_lock_strengths_for_step,
 )
@@ -82,6 +85,25 @@ class ScaleLockConfig:
     mid_schedule_power: float
     mid_schedule_hold: float
     spatial_mask: Optional[torch.Tensor] = None
+    manifold_enabled: bool = False
+    manifold_strength: float = 0.0
+    manifold_strength_start: float = 1.0
+    manifold_strength_end: float = 0.0
+    manifold_schedule: str = "ease_out"
+    manifold_schedule_power: float = 2.0
+    manifold_schedule_hold: float = 0.0
+    manifold_cutoff: float = 0.18
+    manifold_radial_strength: float = 1.0
+    manifold_anisotropy: float = 0.15
+    manifold_translation_strength: float = 1.0
+    manifold_anchor_mix: float = 0.18
+    manifold_mean_anchor_mix: float = 0.12
+    manifold_contrast_restore: float = 0.10
+    manifold_energy_tether: float = 0.0
+    manifold_channel_tether: float = 0.0
+    manifold_energy_gain_cap: float = 1.75
+    manifold_max_shift_px: float = 3.0
+    manifold_spatial_mask: Optional[torch.Tensor] = None
 
 
 @dataclass
@@ -186,6 +208,21 @@ def prepare_noise(latent_samples: torch.Tensor, seed: int, batch_inds=None, disa
     return comfy.sample.prepare_noise(latent_samples, seed, batch_inds)
 
 
+def _resolve_sampler_device(model_or_wrap, fallback: torch.device) -> torch.device:
+    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 fix_latent_channels(model, latent_dict: dict) -> dict:
     out = clone_latent(latent_dict)
     ratio = out.get("downscale_ratio_spacial", None)
@@ -248,29 +285,59 @@ def generate_noise_for_latent(noise, latent: dict) -> torch.Tensor:
     return generated
 
 
+def _apply_manifold_compand_to_noise_prediction(guider, working_noise: torch.Tensor, anchor: torch.Tensor, idx: int) -> torch.Tensor:
+    manifold_strength = scale_lock_manifold_strength_for_step(guider, idx)
+    if manifold_strength <= 0.0:
+        return working_noise
+
+    mask = scale_lock_manifold_mask_for(guider, working_noise)
+    corrected = latent_manifold_compand(
+        working_noise,
+        anchor,
+        mask=mask,
+        strength=manifold_strength,
+        cutoff=guider._slrd_manifold_cutoff,
+        radial_strength=guider._slrd_manifold_radial_strength,
+        anisotropy=guider._slrd_manifold_anisotropy,
+        translation_strength=guider._slrd_manifold_translation_strength,
+        anchor_mix=guider._slrd_manifold_anchor_mix,
+        mean_anchor_mix=guider._slrd_manifold_mean_anchor_mix,
+        contrast_restore=guider._slrd_manifold_contrast_restore,
+        energy_tether=guider._slrd_manifold_energy_tether,
+        channel_tether=guider._slrd_manifold_channel_tether,
+        energy_gain_cap=guider._slrd_manifold_energy_gain_cap,
+        max_shift_px=guider._slrd_manifold_max_shift_px,
+    )
+    if mask is not None:
+        corrected = working_noise + mask * (corrected - working_noise)
+    return corrected
+
+
 def apply_scale_lock_to_noise_prediction(guider, base_noise: torch.Tensor, x, timestep):
+    del x
     if not getattr(guider, "_slrd_anchors_x0_cpu", None):
         return base_noise
 
     idx = resolve_scale_lock_step_index(guider, timestep)
-    low_strength, mid_strength = scale_lock_strengths_for_step(guider, idx)
-    if low_strength <= 0.0 and mid_strength <= 0.0:
-        return base_noise
-
     anchor = scale_lock_anchor_for(guider, idx, base_noise)
-    corrected = residual_lock_multiband(
-        base_noise,
-        anchor,
-        low_strength=low_strength,
-        mid_strength=mid_strength,
-        low_cutoff=guider._slrd_cutoff,
-        mid_cutoff=guider._slrd_mid_cutoff,
-    )
 
-    mask = scale_lock_mask_for(guider, base_noise)
-    if mask is not None:
-        corrected = base_noise + mask * (corrected - base_noise)
+    corrected = base_noise
+    low_strength, mid_strength = scale_lock_strengths_for_step(guider, idx)
+    if low_strength > 0.0 or mid_strength > 0.0:
+        corrected = residual_lock_multiband(
+            corrected,
+            anchor,
+            low_strength=low_strength,
+            mid_strength=mid_strength,
+            low_cutoff=guider._slrd_cutoff,
+            mid_cutoff=guider._slrd_mid_cutoff,
+        )
 
+        mask = scale_lock_mask_for(guider, corrected)
+        if mask is not None:
+            corrected = base_noise + mask * (corrected - base_noise)
+
+    corrected = _apply_manifold_compand_to_noise_prediction(guider, corrected, anchor, idx)
     return corrected
 
 
@@ -302,6 +369,10 @@ def clone_guider_for_scale_lock(guider):
 
 def apply_scale_lock_to_guider(guider, runtime: ScaleLockedRuntimeContext, config: ScaleLockConfig):
     spatial_mask = prepare_spatial_lock_mask(config.spatial_mask, runtime.highres_latent["samples"])
+    manifold_spatial_mask = prepare_spatial_lock_mask(
+        config.manifold_spatial_mask if config.manifold_spatial_mask is not None else config.spatial_mask,
+        runtime.highres_latent["samples"],
+    )
     init_scale_lock_state(
         guider,
         model=runtime.model,
@@ -322,6 +393,25 @@ def apply_scale_lock_to_guider(guider, runtime: ScaleLockedRuntimeContext, confi
         mid_schedule_power=config.mid_schedule_power,
         mid_schedule_hold=config.mid_schedule_hold,
         spatial_mask=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=manifold_spatial_mask,
     )
 
     original_predict_noise = getattr(guider, "_slrd_original_predict_noise", guider.predict_noise)
@@ -355,9 +445,24 @@ def _run_lowres_planner_advanced(
     model = guider.model_patcher
     lowres_latent = fix_latent_channels(model, lowres_latent)
     latent_samples = lowres_latent["samples"]
-    noise_tensor = generate_noise_for_latent(noise, lowres_latent).to(device="cpu", dtype=latent_samples.dtype)
+    target_device = _resolve_sampler_device(model, latent_samples.device)
+    target_dtype = latent_samples.dtype
+    latent_samples = latent_samples.to(
+        device=target_device,
+        dtype=target_dtype,
+        non_blocking=True,
+    )
+    lowres_latent["samples"] = latent_samples
+    noise_tensor = generate_noise_for_latent(noise, lowres_latent).to(
+        device=target_device,
+        dtype=target_dtype,
+        non_blocking=True,
+    )
 
     noise_mask = lowres_latent.get("noise_mask", None)
+    if isinstance(noise_mask, torch.Tensor):
+        noise_mask = noise_mask.to(device=target_device, non_blocking=True)
+    sigmas = sigmas.to(device=target_device, non_blocking=True)
     recorder = TrajectoryRecorder(
         store_dtype=_store_dtype_for(latent_samples),
         capture_noisy_xt=False,
@@ -405,10 +510,30 @@ def _run_lowres_planner(
 
     lowres_latent = fix_latent_channels(model, lowres_latent)
     latent_samples = lowres_latent["samples"]
+    target_device = _resolve_sampler_device(model, latent_samples.device)
+    target_dtype = latent_samples.dtype
+    latent_samples = latent_samples.to(
+        device=target_device,
+        dtype=target_dtype,
+        non_blocking=True,
+    )
+    lowres_latent["samples"] = latent_samples
     batch_inds = lowres_latent.get("batch_index", None)
-    planner_noise = prepare_noise(latent_samples, seed=seed, batch_inds=batch_inds, disable_noise=disable_noise)
+    planner_noise = prepare_noise(
+        latent_samples,
+        seed=seed,
+        batch_inds=batch_inds,
+        disable_noise=disable_noise,
+    ).to(
+        device=target_device,
+        dtype=target_dtype,
+        non_blocking=True,
+    )
 
     noise_mask = lowres_latent.get("noise_mask", None)
+    if isinstance(noise_mask, torch.Tensor):
+        noise_mask = noise_mask.to(device=target_device, non_blocking=True)
+    sigmas = sigmas.to(device=target_device, non_blocking=True)
     recorder = TrajectoryRecorder(
         store_dtype=_store_dtype_for(latent_samples),
         capture_noisy_xt=False,
@@ -438,12 +563,14 @@ def _run_lowres_planner(
 
 def _build_highres_noise(highres_latent: dict, lowres_noise: torch.Tensor, seed: int, hf_strength: float):
     highres_samples = highres_latent["samples"]
+    target_device = highres_samples.device
+    target_dtype = highres_samples.dtype
     if torch.count_nonzero(lowres_noise).item() == 0:
         return torch.zeros(
             highres_samples.size(),
-            dtype=highres_samples.dtype,
+            dtype=target_dtype,
             layout=highres_samples.layout,
-            device="cpu",
+            device=target_device,
         )
 
     return build_nested_noise(
@@ -451,7 +578,7 @@ def _build_highres_noise(highres_latent: dict, lowres_noise: torch.Tensor, seed:
         target_shape=tuple(highres_samples.shape),
         seed=seed,
         hf_strength=hf_strength,
-    ).to(device="cpu", dtype=highres_samples.dtype)
+    ).to(device=target_device, dtype=target_dtype, non_blocking=True)
 
 
 def build_runtime_context_from_advanced(
@@ -467,6 +594,18 @@ def build_runtime_context_from_advanced(
 ) -> ScaleLockedRuntimeContext:
     model = guider.model_patcher
     highres_latent = fix_latent_channels(model, latent_image)
+    target_device = _resolve_sampler_device(model, highres_latent["samples"].device)
+    target_dtype = highres_latent["samples"].dtype
+    highres_latent["samples"] = highres_latent["samples"].to(
+        device=target_device,
+        dtype=target_dtype,
+        non_blocking=True,
+    )
+    if isinstance(highres_latent.get("noise_mask"), torch.Tensor):
+        highres_latent["noise_mask"] = highres_latent["noise_mask"].to(
+            device=target_device,
+            non_blocking=True,
+        )
     lowres_latent = make_lowres_latent(highres_latent, target_megapixels)
 
     if sigmas.numel() == 0:
@@ -475,10 +614,13 @@ def build_runtime_context_from_advanced(
             highres_latent=highres_latent,
             lowres_latent=lowres_latent,
             lowres_out=clean_latent(lowres_latent),
-            sigmas=sigmas,
+            sigmas=sigmas.to(device=target_device, non_blocking=True),
             anchors_x0=[],
             planner_sigmas=[],
-            highres_noise=torch.zeros_like(highres_latent["samples"], device="cpu"),
+            highres_noise=torch.zeros_like(
+                highres_latent["samples"],
+                device=target_device,
+            ),
             noise_seed=noise_seed(noise),
         )
 
@@ -504,7 +646,7 @@ def build_runtime_context_from_advanced(
         highres_latent=highres_latent,
         lowres_latent=lowres_latent,
         lowres_out=lowres_out,
-        sigmas=sigmas,
+        sigmas=sigmas.to(device=target_device, non_blocking=True),
         anchors_x0=anchors_x0,
         planner_sigmas=planner_sigmas,
         highres_noise=highres_noise,