Refactor scale lock init helper

Add advanced scale-locked sampler node and guider init fix

.gitignore

@@ -1,2 +1,3 @@
 __pycache__/
 *.py[cod]
+.vs/

.vs/ComfyUI-ScaleLockedResidualDiffusion/v17/DocumentLayout.json

@@ -1,23 +0,0 @@
-{
-  "Version": 1,
-  "WorkspaceRootPath": "C:\\Users\\marre\\source\\repos\\ComfyUI-ScaleLockedResidualDiffusion\\",
-  "Documents": [],
-  "DocumentGroupContainers": [
-    {
-      "Orientation": 0,
-      "VerticalTabListWidth": 256,
-      "DocumentGroups": [
-        {
-          "DockedWidth": 200,
-          "SelectedChildIndex": -1,
-          "Children": [
-            {
-              "$type": "Bookmark",
-              "Name": "ST:0:0:{aa2115a1-9712-457b-9047-dbb71ca2cdd2}"
-            }
-          ]
-        }
-      ]
-    }
-  ]
-}

.vs/VSWorkspaceState.json

@@ -1,7 +0,0 @@
-{
-  "ExpandedNodes": [
-    ""
-  ],
-  "SelectedNode": "\\C:\\Users\\marre\\Source\\Repos\\ComfyUI-ScaleLockedResidualDiffusion",
-  "PreviewInSolutionExplorer": false
-}

nodes.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import types
 from typing import Optional
 
 import torch
@@ -16,10 +17,15 @@ from .slrd_core import (
     TrajectoryRecorder,
     build_nested_noise,
     clone_latent,
+    init_scale_lock_state,
     latent_target_hw_from_megapixels,
+    resolve_scale_lock_step_index,
     resize_latent_dict,
     resize_mask,
     residual_lock_multiband,
+    scale_lock_anchor_for,
+    scale_lock_mask_for,
+    scale_lock_strength_for_step,
 )
 
 
@@ -28,6 +34,9 @@ _LOGGER = logging.getLogger(__name__)
 
 # Attach our mixin-style state logic to the real CFG guider base.
 class _ScaleLockedCFGGuiderImpl(ScaleLockedCFGGuider, comfy.samplers.CFGGuider):
+    def __init__(self, model_patcher):
+        comfy.samplers.CFGGuider.__init__(self, model_patcher)
+
     def set_conds(self, positive, negative):
         self.inner_set_conds({"positive": positive, "negative": negative})
 
@@ -49,7 +58,8 @@ class _ScaleLockedCFGGuiderImpl(ScaleLockedCFGGuider, comfy.samplers.CFGGuider):
         schedule: str,
         spatial_mask: Optional[torch.Tensor],
     ):
-        self._init_scale_lock_state(
+        init_scale_lock_state(
+            self,
             model=model,
             anchors_x0_cpu=anchors_x0_cpu,
             planner_sigmas=planner_sigmas,
@@ -84,12 +94,12 @@ class _ScaleLockedCFGGuiderImpl(ScaleLockedCFGGuider, comfy.samplers.CFGGuider):
         if not getattr(self, "_slrd_anchors_x0_cpu", None):
             return base_cfg
 
-        idx = self._slrd_resolve_step_index(timestep)
-        strength = self._slrd_strength_for_step(idx)
+        idx = resolve_scale_lock_step_index(self, timestep)
+        strength = scale_lock_strength_for_step(self, idx)
         if strength <= 0.0:
             return base_cfg
 
-        anchor = self._slrd_anchor_for(idx, base_cfg)
+        anchor = scale_lock_anchor_for(self, idx, base_cfg)
         corrected = residual_lock_multiband(
             base_cfg,
             anchor,
@@ -99,13 +109,15 @@ class _ScaleLockedCFGGuiderImpl(ScaleLockedCFGGuider, comfy.samplers.CFGGuider):
             mid_cutoff=self._slrd_mid_cutoff,
         )
 
-        mask = self._slrd_mask_for(base_cfg)
+        mask = scale_lock_mask_for(self, base_cfg)
         if mask is not None:
             corrected = base_cfg + mask * (corrected - base_cfg)
 
         return corrected
 
 
+
+
 class _NullPreviewCallback:
     def __call__(self, step, x0, x, total_steps):
         del step, x0, x, total_steps
@@ -228,6 +240,128 @@ def _planner_sigmas_for_recorded_steps(sigmas: torch.Tensor, recorded_steps: int
     return visible_sigmas + [visible_sigmas[-1]] * (recorded_steps - len(visible_sigmas))
 
 
+def _noise_seed(noise) -> int:
+    try:
+        return int(getattr(noise, "seed", 0))
+    except Exception:
+        return 0
+
+
+def _generate_noise_for_latent(noise, latent: dict) -> torch.Tensor:
+    generated = noise.generate_noise(latent)
+    if not isinstance(generated, torch.Tensor):
+        raise TypeError("ScaleLockedResidualSamplerCustomAdvanced currently supports tensor noise only.")
+    return generated
+
+def _scale_lock_predict_noise(guider, base_noise: torch.Tensor, x, timestep):
+    if not getattr(guider, "_slrd_anchors_x0_cpu", None):
+        return base_noise
+
+    idx = resolve_scale_lock_step_index(guider, timestep)
+    strength = scale_lock_strength_for_step(guider, idx)
+    if strength <= 0.0:
+        return base_noise
+
+    anchor = scale_lock_anchor_for(guider, idx, base_noise)
+    corrected = residual_lock_multiband(
+        base_noise,
+        anchor,
+        low_strength=strength,
+        mid_strength=min(1.0, strength * guider._slrd_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)
+
+    return corrected
+
+
+def _patch_guider_with_scale_lock(
+    guider,
+    *,
+    model,
+    anchors_x0_cpu,
+    planner_sigmas,
+    lock_strength: float,
+    lock_strength_start: float,
+    lock_strength_end: float,
+    cutoff: float,
+    mid_cutoff: float,
+    mid_strength: float,
+    schedule: str,
+    spatial_mask: Optional[torch.Tensor],
+):
+    init_scale_lock_state(
+        guider,
+        model=model,
+        anchors_x0_cpu=anchors_x0_cpu,
+        planner_sigmas=planner_sigmas,
+        lock_strength=lock_strength,
+        lock_strength_start=lock_strength_start,
+        lock_strength_end=lock_strength_end,
+        cutoff=cutoff,
+        mid_cutoff=mid_cutoff,
+        mid_strength=mid_strength,
+        schedule=schedule,
+        spatial_mask=spatial_mask,
+    )
+
+    original_predict_noise = guider.predict_noise
+
+    def _wrapped_predict_noise(self, x, timestep, model_options={}, seed=None):
+        base_noise = original_predict_noise(x, timestep, model_options=model_options, seed=seed)
+        return _scale_lock_predict_noise(self, base_noise, x, timestep)
+
+    guider.predict_noise = types.MethodType(_wrapped_predict_noise, guider)
+    return original_predict_noise
+
+
+
+def _run_lowres_planner_advanced(
+    *,
+    guider,
+    sampler,
+    sigmas: torch.Tensor,
+    lowres_latent: dict,
+    noise,
+    pin_anchors: bool,
+) -> tuple[dict, list[torch.Tensor], list[float], torch.Tensor]:
+    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)
+
+    noise_mask = lowres_latent.get("noise_mask", None)
+    recorder = TrajectoryRecorder(
+        store_dtype=_store_dtype_for(latent_samples),
+        capture_noisy_xt=False,
+        pin_memory=pin_anchors,
+    )
+
+    samples = guider.sample(
+        noise_tensor,
+        latent_samples,
+        sampler,
+        sigmas,
+        denoise_mask=noise_mask,
+        callback=recorder.callback,
+        disable_pbar=True,
+        seed=_noise_seed(noise),
+    )
+    samples = samples.to(comfy.model_management.intermediate_device())
+
+    out = clone_latent(lowres_latent)
+    out.pop("downscale_ratio_spacial", None)
+    out["samples"] = samples
+
+    planner_sigmas = _planner_sigmas_for_recorded_steps(sigmas, len(recorder.x0_steps))
+    recorder.step_sigmas = planner_sigmas
+    return out, recorder.x0_steps, planner_sigmas, noise_tensor
+
+
 def _run_lowres_planner(
     *,
     model,
@@ -485,6 +619,156 @@ class ScaleLockedResidualKSampler:
         return (out, lowres_out_clean, denoised)
 
 
+class ScaleLockedResidualSamplerCustomAdvanced:
+    CATEGORY = "sampling/scale_locked"
+    RETURN_TYPES = ("LATENT", "LATENT", "LATENT")
+    RETURN_NAMES = ("output", "lowres_planner", "denoised_output")
+    FUNCTION = "sample"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "noise": ("NOISE",),
+                "guider": ("GUIDER",),
+                "sampler": ("SAMPLER",),
+                "sigmas": ("SIGMAS",),
+                "latent_image": ("LATENT",),
+                "target_megapixels": ("FLOAT", {"default": 1.0, "min": 0.10, "max": 16.0, "step": 0.05, "round": 0.01}),
+                "lock_strength": ("FLOAT", {"default": 0.85, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "lock_strength_start": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "lock_strength_end": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "lock_schedule": (["linear", "cosine", "flat"],),
+                "coarse_cutoff": ("FLOAT", {"default": 0.33, "min": 0.05, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "mid_band_cutoff": ("FLOAT", {"default": 0.60, "min": 0.05, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "mid_band_strength": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 2.0, "step": 0.01, "round": 0.001}),
+                "nested_noise_strength": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 4.0, "step": 0.01, "round": 0.001}),
+                "pin_anchors": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "lock_mask": ("MASK",),
+            },
+        }
+
+    @staticmethod
+    def _preview_callback(model, steps: int, x0_output: dict):
+        return _make_preview_callback(model, steps, x0_output)
+
+    def sample(
+        self,
+        noise,
+        guider,
+        sampler,
+        sigmas,
+        latent_image,
+        target_megapixels,
+        lock_strength,
+        lock_strength_start,
+        lock_strength_end,
+        lock_schedule,
+        coarse_cutoff,
+        mid_band_cutoff,
+        mid_band_strength,
+        nested_noise_strength,
+        pin_anchors,
+        lock_mask=None,
+    ):
+        model = guider.model_patcher
+        highres_latent = _fix_latent_channels(model, latent_image)
+        lowres_latent = _make_lowres_latent(highres_latent, target_megapixels)
+
+        if sigmas.numel() == 0:
+            out = clone_latent(highres_latent)
+            out.pop("downscale_ratio_spacial", None)
+            lowres_out = clone_latent(lowres_latent)
+            lowres_out.pop("downscale_ratio_spacial", None)
+            return (out, lowres_out, out)
+
+        lowres_out, anchors_x0, planner_sigmas, lowres_noise = _run_lowres_planner_advanced(
+            guider=guider,
+            sampler=sampler,
+            sigmas=sigmas,
+            lowres_latent=lowres_latent,
+            noise=noise,
+            pin_anchors=pin_anchors,
+        )
+
+        if len(anchors_x0) == 0:
+            raise RuntimeError("ScaleLockedResidualSamplerCustomAdvanced: planner pass did not record any x0 anchors.")
+
+        highres_samples = highres_latent["samples"]
+        if torch.count_nonzero(lowres_noise).item() == 0:
+            highres_noise = torch.zeros(
+                highres_samples.size(),
+                dtype=highres_samples.dtype,
+                layout=highres_samples.layout,
+                device="cpu",
+            )
+        else:
+            highres_noise = build_nested_noise(
+                lowres_noise=lowres_noise,
+                target_shape=tuple(highres_samples.shape),
+                seed=_noise_seed(noise) ^ 0x9E3779B97F4A7C15,
+                hf_strength=nested_noise_strength,
+            )
+        highres_noise = highres_noise.to(device="cpu", dtype=highres_samples.dtype)
+
+        original_predict_noise = _patch_guider_with_scale_lock(
+            guider,
+            model=model,
+            anchors_x0_cpu=anchors_x0,
+            planner_sigmas=planner_sigmas,
+            lock_strength=lock_strength,
+            lock_strength_start=lock_strength_start,
+            lock_strength_end=lock_strength_end,
+            cutoff=coarse_cutoff,
+            mid_cutoff=mid_band_cutoff,
+            mid_strength=mid_band_strength,
+            schedule=lock_schedule,
+            spatial_mask=_prepare_spatial_lock_mask(lock_mask, highres_samples),
+        )
+
+        x0_output = {}
+        callback = self._preview_callback(model, len(sigmas) - 1, x0_output)
+        noise_mask = highres_latent.get("noise_mask", None)
+        disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+
+        try:
+            samples = guider.sample(
+                highres_noise,
+                highres_samples,
+                sampler,
+                sigmas,
+                denoise_mask=noise_mask,
+                callback=callback,
+                disable_pbar=disable_pbar,
+                seed=_noise_seed(noise),
+            )
+        finally:
+            guider.predict_noise = original_predict_noise
+        samples = samples.to(comfy.model_management.intermediate_device())
+
+        out = clone_latent(highres_latent)
+        out.pop("downscale_ratio_spacial", None)
+        out["samples"] = samples
+
+        if "x0" in x0_output:
+            try:
+                x0_out = model.model.process_latent_out(x0_output["x0"].cpu())
+            except Exception:
+                x0_out = x0_output["x0"].detach().cpu()
+            denoised = clone_latent(highres_latent)
+            denoised.pop("downscale_ratio_spacial", None)
+            denoised["samples"] = x0_out
+        else:
+            denoised = out
+
+        lowres_out_clean = clone_latent(lowres_out)
+        lowres_out_clean.pop("downscale_ratio_spacial", None)
+
+        return (out, lowres_out_clean, denoised)
+
+
 class ScaleLockedNestedNoisePreview:
     CATEGORY = "sampling/scale_locked"
     RETURN_TYPES = ("LATENT", "LATENT")
@@ -523,11 +807,13 @@ class ScaleLockedNestedNoisePreview:
 
 NODE_CLASS_MAPPINGS = {
     "ScaleLockedResidualKSampler": ScaleLockedResidualKSampler,
+    "ScaleLockedResidualSamplerCustomAdvanced": ScaleLockedResidualSamplerCustomAdvanced,
     "ScaleLockedNestedNoisePreview": ScaleLockedNestedNoisePreview,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ScaleLockedResidualKSampler": "Scale-Locked Residual KSampler",
+    "ScaleLockedResidualSamplerCustomAdvanced": "Scale-Locked Residual SamplerCustomAdvanced",
     "ScaleLockedNestedNoisePreview": "Scale-Locked Nested Noise Preview",
 }
 

slrd_core.py

@@ -239,7 +239,7 @@ class TrajectoryRecorder:
             self.xt_steps.append(_stage_cpu_tensor(x, self.store_dtype, self.pin_memory))
 
 
-class ScaleLockedCFGGuider(torch.nn.Module):
+class ScaleLockedCFGGuider:
     """
     A custom ComfyUI guider that applies the scale lock in denoised-latent space.
 
@@ -263,84 +263,131 @@ class ScaleLockedCFGGuider(torch.nn.Module):
         schedule: str,
         spatial_mask: Optional[torch.Tensor],
     ) -> None:
-        self._slrd_model = model
-        self._slrd_anchors_x0_cpu = list(anchors_x0_cpu)
-        self._slrd_planner_sigmas = [float(x) for x in planner_sigmas] if planner_sigmas is not None else None
-        self._slrd_lock_strength = float(lock_strength)
-        self._slrd_lock_strength_start = float(lock_strength_start)
-        self._slrd_lock_strength_end = float(lock_strength_end)
-        self._slrd_cutoff = float(cutoff)
-        self._slrd_mid_cutoff = float(max(cutoff, mid_cutoff))
-        self._slrd_mid_strength = float(mid_strength)
-        self._slrd_schedule = schedule
-        self._slrd_seen_sigmas: list[float] = []
-        self._slrd_prev_match_idx: int = 0
-        self._slrd_spatial_mask = spatial_mask
-        self._slrd_last_sigma: Optional[float] = None
+        init_scale_lock_state(
+            self,
+            model=model,
+            anchors_x0_cpu=anchors_x0_cpu,
+            planner_sigmas=planner_sigmas,
+            lock_strength=lock_strength,
+            lock_strength_start=lock_strength_start,
+            lock_strength_end=lock_strength_end,
+            cutoff=cutoff,
+            mid_cutoff=mid_cutoff,
+            mid_strength=mid_strength,
+            schedule=schedule,
+            spatial_mask=spatial_mask,
+        )
 
     def _slrd_resolve_step_index(self, timestep: torch.Tensor | float | int) -> int:
-        sigma = _sigma_scalar(timestep)
-
-        if self._slrd_planner_sigmas:
-            best_idx = 0
-            best_dist = float("inf")
-            for i, planner_sigma in enumerate(self._slrd_planner_sigmas):
-                dist = abs(planner_sigma - sigma)
-                if dist < best_dist:
-                    best_idx = i
-                    best_dist = dist
-            best_idx = max(best_idx, self._slrd_prev_match_idx)
-            best_idx = min(best_idx, len(self._slrd_anchors_x0_cpu) - 1)
-            self._slrd_prev_match_idx = best_idx
-            return best_idx
-
-        if self._slrd_last_sigma is None:
-            self._slrd_last_sigma = sigma
-            step_index = 0
-            self._slrd_seen_sigmas = [sigma]
-        else:
-            tol = 1e-6 * max(1.0, abs(self._slrd_last_sigma), abs(sigma))
-            if abs(sigma - self._slrd_last_sigma) > tol:
-                self._slrd_last_sigma = sigma
-                self._slrd_seen_sigmas.append(sigma)
-            unique_sigmas = []
-            for seen_sigma in self._slrd_seen_sigmas:
-                if all(
-                    abs(seen_sigma - unique_sigma) > (1e-6 * max(1.0, abs(seen_sigma), abs(unique_sigma)))
-                    for unique_sigma in unique_sigmas
-                ):
-                    unique_sigmas.append(seen_sigma)
-            step_index = max(0, len(unique_sigmas) - 1)
-            step_index = max(step_index, self._slrd_prev_match_idx)
-            self._slrd_prev_match_idx = step_index
-        if not self._slrd_anchors_x0_cpu:
-            return 0
-        return min(step_index, len(self._slrd_anchors_x0_cpu) - 1)
+        return resolve_scale_lock_step_index(self, timestep)
 
     def _slrd_strength_for_step(self, idx: int) -> float:
-        total = max(1, len(self._slrd_anchors_x0_cpu) - 1)
-        progress = idx / total
-        scheduled = schedule_value(self._slrd_lock_strength_start, self._slrd_lock_strength_end, progress, self._slrd_schedule)
-        return float(max(0.0, min(1.0, self._slrd_lock_strength * scheduled)))
+        return scale_lock_strength_for_step(self, idx)
 
     def _slrd_anchor_for(self, idx: int, like: torch.Tensor) -> torch.Tensor:
-        anchor = self._slrd_anchors_x0_cpu[idx].to(device=like.device, dtype=like.dtype, non_blocking=True)
-        if tuple(anchor.shape[-2:]) != tuple(like.shape[-2:]):
-            anchor = resize_4d_tensor(anchor, tuple(like.shape[-2:]))
-        return anchor
+        return scale_lock_anchor_for(self, idx, like)
 
     def _slrd_mask_for(self, like: torch.Tensor) -> Optional[torch.Tensor]:
-        if self._slrd_spatial_mask is None:
-            return None
-        mask = self._slrd_spatial_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
+        return scale_lock_mask_for(self, like)
+
+
+def init_scale_lock_state(
+    guider,
+    *,
+    model,
+    anchors_x0_cpu: Iterable[torch.Tensor],
+    planner_sigmas: Optional[Iterable[float]],
+    lock_strength: float,
+    lock_strength_start: float,
+    lock_strength_end: float,
+    cutoff: float,
+    mid_cutoff: float,
+    mid_strength: float,
+    schedule: str,
+    spatial_mask: Optional[torch.Tensor],
+) -> None:
+    guider._slrd_model = model
+    guider._slrd_anchors_x0_cpu = list(anchors_x0_cpu)
+    guider._slrd_planner_sigmas = [float(x) for x in planner_sigmas] if planner_sigmas is not None else None
+    guider._slrd_lock_strength = float(lock_strength)
+    guider._slrd_lock_strength_start = float(lock_strength_start)
+    guider._slrd_lock_strength_end = float(lock_strength_end)
+    guider._slrd_cutoff = float(cutoff)
+    guider._slrd_mid_cutoff = float(max(cutoff, mid_cutoff))
+    guider._slrd_mid_strength = float(mid_strength)
+    guider._slrd_schedule = schedule
+    guider._slrd_seen_sigmas = []
+    guider._slrd_prev_match_idx = 0
+    guider._slrd_spatial_mask = spatial_mask
+    guider._slrd_last_sigma = None
+
+
+def resolve_scale_lock_step_index(guider, timestep: torch.Tensor | float | int) -> int:
+    sigma = _sigma_scalar(timestep)
+
+    if guider._slrd_planner_sigmas:
+        best_idx = 0
+        best_dist = float("inf")
+        for i, planner_sigma in enumerate(guider._slrd_planner_sigmas):
+            dist = abs(planner_sigma - sigma)
+            if dist < best_dist:
+                best_idx = i
+                best_dist = dist
+        best_idx = max(best_idx, guider._slrd_prev_match_idx)
+        best_idx = min(best_idx, len(guider._slrd_anchors_x0_cpu) - 1)
+        guider._slrd_prev_match_idx = best_idx
+        return best_idx
+
+    if guider._slrd_last_sigma is None:
+        guider._slrd_last_sigma = sigma
+        step_index = 0
+        guider._slrd_seen_sigmas = [sigma]
+    else:
+        tol = 1e-6 * max(1.0, abs(guider._slrd_last_sigma), abs(sigma))
+        if abs(sigma - guider._slrd_last_sigma) > tol:
+            guider._slrd_last_sigma = sigma
+            guider._slrd_seen_sigmas.append(sigma)
+        unique_sigmas = []
+        for seen_sigma in guider._slrd_seen_sigmas:
+            if all(
+                abs(seen_sigma - unique_sigma) > (1e-6 * max(1.0, abs(seen_sigma), abs(unique_sigma)))
+                for unique_sigma in unique_sigmas
+            ):
+                unique_sigmas.append(seen_sigma)
+        step_index = max(0, len(unique_sigmas) - 1)
+        step_index = max(step_index, guider._slrd_prev_match_idx)
+        guider._slrd_prev_match_idx = step_index
+    if not guider._slrd_anchors_x0_cpu:
+        return 0
+    return min(step_index, len(guider._slrd_anchors_x0_cpu) - 1)
+
+
+def scale_lock_strength_for_step(guider, idx: int) -> float:
+    total = max(1, len(guider._slrd_anchors_x0_cpu) - 1)
+    progress = idx / total
+    scheduled = schedule_value(guider._slrd_lock_strength_start, guider._slrd_lock_strength_end, progress, guider._slrd_schedule)
+    return float(max(0.0, min(1.0, guider._slrd_lock_strength * scheduled)))
+
+
+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:]):
+        anchor = resize_4d_tensor(anchor, tuple(like.shape[-2:]))
+    return anchor
+
+
+def scale_lock_mask_for(guider, like: torch.Tensor) -> Optional[torch.Tensor]:
+    if guider._slrd_spatial_mask is None:
+        return None
+    mask = guider._slrd_spatial_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