Fix project URLs to match published repository

Add Comfy registry metadata and publish workflow

.github/workflows/publish.yml

@@ -0,0 +1,28 @@
+name: Publish to Comfy registry
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "pyproject.toml"
+
+permissions:
+  issues: write
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+        with:
+          submodules: true
+
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@v1
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.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
-}

README.md

@@ -33,11 +33,15 @@ Main all-in-one node.
 **Important controls**
 - `target_megapixels`: planner resolution in pixel-space megapixels (usually `0.8` to `1.5` for Flux-like native planning)
 - `lock_strength`: global multiplier for the scale lock
-- `lock_strength_start` / `lock_strength_end`: how strongly the lock applies early vs late in denoising
-- `lock_schedule`: linear / cosine / flat interpolation for the lock schedule
+- `lock_strength_start` / `lock_strength_end`: how strongly the low-band lock applies early vs late in denoising
+- `lock_schedule`: low-band schedule shape; progress follows normalized log-sigma position when planner sigmas are available, with fallback to raw step-index progress otherwise; `flat` means constant start-value scheduling and ignores `lock_strength_end`
+- `lock_schedule_hold` / `lock_schedule_power`: knee and curvature controls for schedules such as `hold_then_drop`, `ease_in`, and `ease_out`
 - `coarse_cutoff`: retained spatial fraction for the strongest coarse lock band
 - `mid_band_cutoff`: retained spatial fraction for an additional mid-frequency lock band
-- `mid_band_strength`: how strongly the mid-band is pulled toward the planner relative to the low-band schedule
+- `mid_band_strength`: overall mid-band lock multiplier
+- `mid_band_strength_start` / `mid_band_strength_end`: independent mid-band envelope when `mid_band_schedule` is not `linked`; this envelope multiplies the base `mid_band_strength`
+- `mid_band_schedule`: `linked` preserves the old behavior; other modes decouple the mid band from the low band
+- `mid_band_schedule_hold` / `mid_band_schedule_power`: shape controls for independent mid-band schedules
 - `nested_noise_strength`: amount of zero-mean high-frequency detail noise added on top of the lifted low-res noise
 - `lock_mask` (optional): spatial mask to strengthen the lock only in selected regions (for example body / face / hands)
 - `pin_anchors`: store planner anchors in pinned CPU memory when possible for faster non-blocking transfer during the high-res pass
@@ -67,10 +71,13 @@ For a first test when your high-res target is around 4 MP:
 - `lock_strength = 0.85`
 - `lock_strength_start = 0.95`
 - `lock_strength_end = 0.25`
-- `lock_schedule = cosine`
+- `lock_schedule = hold_then_drop`
+- `lock_schedule_hold = 0.35`
+- `lock_schedule_power = 3.0`
 - `coarse_cutoff = 0.33`
 - `mid_band_cutoff = 0.60`
 - `mid_band_strength = 0.35`
+- `mid_band_schedule = linked`
 - `sampler_guard = warn`
 - `nested_noise_strength = 0.35`
 
@@ -117,6 +124,7 @@ What is not implemented yet:
 - automatic anatomy / pose / segmentation mask extraction,
 - explicit residual-only tiled model execution,
 - sigma-perfect trajectory matching for samplers that perform unusual extra model evaluations,
+- scheduled cutoff animation for coarse or mid bands,
 - multi-stage 1 MP -> 2 MP -> 4 MP progressive ladder inside one node,
 - exact support tuning for every possible exotic custom sampler.
 

docs/ScaleLockedResidualKSampler.md

@@ -12,12 +12,17 @@ A native / low-resolution planner branch is sampled first. Its denoised trajecto
 - `seed`, `steps`, `cfg`, `sampler_name`, `scheduler`, `denoise`: standard sampler controls
 - `target_megapixels`: planner resolution in pixel-space MP
 - `lock_strength`: overall lock multiplier
-- `lock_strength_start`: early-step lock amount
-- `lock_strength_end`: late-step lock amount
-- `lock_schedule`: linear / cosine / flat
+- `lock_strength_start`: early low-band lock amount
+- `lock_strength_end`: late low-band lock amount
+- `lock_schedule`: low-band schedule shape, evaluated against normalized log-sigma position when planner sigmas are available, with raw step-index fallback otherwise; `flat` keeps the start value for the whole run and ignores the end value
+- `lock_schedule_hold`: hold region before `hold_then_drop` releases
+- `lock_schedule_power`: curvature control for power-based schedules
 - `coarse_cutoff`: strongest coarse-band resolution fraction
 - `mid_band_cutoff`: second, looser mid-band resolution fraction
-- `mid_band_strength`: relative strength of the mid-band lock
+- `mid_band_strength`: overall mid-band lock multiplier
+- `mid_band_strength_start` / `mid_band_strength_end`: independent mid-band envelope when `mid_band_schedule` is not `linked`; it multiplies the base `mid_band_strength`
+- `mid_band_schedule`: `linked` for legacy behavior, or an independent curve mode
+- `mid_band_schedule_hold` / `mid_band_schedule_power`: shape controls for the independent mid-band schedule
 - `nested_noise_strength`: amount of extra high-frequency detail noise
 - `pin_anchors`: pinned-memory staging for planner anchors when possible
 - `sampler_guard`: warn / error / off handling for samplers outside the conservative safe set
@@ -33,5 +38,7 @@ A native / low-resolution planner branch is sampled first. Its denoised trajecto
 
 - Lower `coarse_cutoff` = stronger global structure control.
 - Lower `mid_band_cutoff` and higher `mid_band_strength` = tighter control over medium-scale body/shape structure.
+- `hold_then_drop` with a `lock_schedule_hold` around `0.30` to `0.45` gives a stronger early anchor with a later release knee.
+- Use `mid_band_schedule = linked` to preserve the legacy shared curve, or switch it off to let medium structure release earlier than the coarse band.
 - Higher `nested_noise_strength` = more detail freedom, but also more chance of drift.
 - A `lock_mask` is recommended for body-heavy and anatomy-sensitive generations.

js/docs/ScaleLockedResidualKSampler/en.md

@@ -12,12 +12,17 @@ A native / low-resolution planner branch is sampled first. Its denoised trajecto
 - `seed`, `steps`, `cfg`, `sampler_name`, `scheduler`, `denoise`: standard sampler controls
 - `target_megapixels`: planner resolution in pixel-space MP
 - `lock_strength`: overall lock multiplier
-- `lock_strength_start`: early-step lock amount
-- `lock_strength_end`: late-step lock amount
-- `lock_schedule`: linear / cosine / flat
+- `lock_strength_start`: early low-band lock amount
+- `lock_strength_end`: late low-band lock amount
+- `lock_schedule`: low-band schedule shape, evaluated against normalized log-sigma position when planner sigmas are available, with raw step-index fallback otherwise; `flat` keeps the start value for the whole run and ignores the end value
+- `lock_schedule_hold`: hold region before `hold_then_drop` releases
+- `lock_schedule_power`: curvature control for power-based schedules
 - `coarse_cutoff`: strongest coarse-band resolution fraction
 - `mid_band_cutoff`: second, looser mid-band resolution fraction
-- `mid_band_strength`: relative strength of the mid-band lock
+- `mid_band_strength`: overall mid-band lock multiplier
+- `mid_band_strength_start` / `mid_band_strength_end`: independent mid-band envelope when `mid_band_schedule` is not `linked`; it multiplies the base `mid_band_strength`
+- `mid_band_schedule`: `linked` for legacy behavior, or an independent curve mode
+- `mid_band_schedule_hold` / `mid_band_schedule_power`: shape controls for the independent mid-band schedule
 - `nested_noise_strength`: amount of extra high-frequency detail noise
 - `pin_anchors`: pinned-memory staging for planner anchors when possible
 - `sampler_guard`: warn / error / off handling for samplers outside the conservative safe set
@@ -33,5 +38,7 @@ A native / low-resolution planner branch is sampled first. Its denoised trajecto
 
 - Lower `coarse_cutoff` = stronger global structure control.
 - Lower `mid_band_cutoff` and higher `mid_band_strength` = tighter control over medium-scale body/shape structure.
+- `hold_then_drop` with a `lock_schedule_hold` around `0.30` to `0.45` gives a stronger early anchor with a later release knee.
+- Use `mid_band_schedule = linked` to preserve the legacy shared curve, or switch it off to let medium structure release earlier than the coarse band.
 - Higher `nested_noise_strength` = more detail freedom, but also more chance of drift.
 - A `lock_mask` is recommended for body-heavy and anatomy-sensitive generations.

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_strengths_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})
 
@@ -47,9 +56,17 @@ class _ScaleLockedCFGGuiderImpl(ScaleLockedCFGGuider, comfy.samplers.CFGGuider):
         mid_cutoff: float,
         mid_strength: float,
         schedule: str,
+        schedule_power: float,
+        schedule_hold: float,
+        mid_strength_start: float,
+        mid_strength_end: float,
+        mid_schedule: str,
+        mid_schedule_power: float,
+        mid_schedule_hold: float,
         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,
@@ -60,6 +77,13 @@ class _ScaleLockedCFGGuiderImpl(ScaleLockedCFGGuider, comfy.samplers.CFGGuider):
             mid_cutoff=mid_cutoff,
             mid_strength=mid_strength,
             schedule=schedule,
+            schedule_power=schedule_power,
+            schedule_hold=schedule_hold,
+            mid_strength_start=mid_strength_start,
+            mid_strength_end=mid_strength_end,
+            mid_schedule=mid_schedule,
+            mid_schedule_power=mid_schedule_power,
+            mid_schedule_hold=mid_schedule_hold,
             spatial_mask=spatial_mask,
         )
 
@@ -84,28 +108,30 @@ 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)
-        if strength <= 0.0:
+        idx = resolve_scale_lock_step_index(self, timestep)
+        low_strength, mid_strength = scale_lock_strengths_for_step(self, idx)
+        if low_strength <= 0.0 and mid_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,
-            low_strength=strength,
-            mid_strength=min(1.0, strength * self._slrd_mid_strength),
+            low_strength=low_strength,
+            mid_strength=mid_strength,
             low_cutoff=self._slrd_cutoff,
             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
@@ -121,6 +147,21 @@ _CONSERVATIVE_SAFE_SAMPLERS = {
     "dpmpp_2m_cfg_pp",
 }
 
+_LOCK_SCHEDULE_OPTIONS = [
+    "linear",
+    "cosine",
+    "flat",
+    "smoothstep",
+    "smootherstep",
+    "ease_in",
+    "ease_out",
+    "ease_in_out",
+    "hold_then_drop",
+    "fast_drop",
+]
+
+_MID_SCHEDULE_OPTIONS = ["linked", *_LOCK_SCHEDULE_OPTIONS]
+
 
 def _normalize_sampler_name(name: str) -> str:
     return str(name).strip().lower()
@@ -228,6 +269,142 @@ 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)
+    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)
+
+    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,
+    schedule_power: float,
+    schedule_hold: float,
+    mid_strength_start: float,
+    mid_strength_end: float,
+    mid_schedule: str,
+    mid_schedule_power: float,
+    mid_schedule_hold: float,
+    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,
+        schedule_power=schedule_power,
+        schedule_hold=schedule_hold,
+        mid_strength_start=mid_strength_start,
+        mid_strength_end=mid_strength_end,
+        mid_schedule=mid_schedule,
+        mid_schedule_power=mid_schedule_power,
+        mid_schedule_hold=mid_schedule_hold,
+        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,
@@ -314,10 +491,17 @@ class ScaleLockedResidualKSampler:
                 "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"],),
+                "lock_schedule": (_LOCK_SCHEDULE_OPTIONS,),
+                "lock_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
+                "lock_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
                 "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}),
+                "mid_band_strength_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "mid_band_strength_end": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "mid_band_schedule": (_MID_SCHEDULE_OPTIONS,),
+                "mid_band_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
+                "mid_band_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
                 "nested_noise_strength": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 4.0, "step": 0.01, "round": 0.001}),
                 "add_noise": ("BOOLEAN", {"default": True}),
                 "pin_anchors": ("BOOLEAN", {"default": True}),
@@ -349,9 +533,16 @@ class ScaleLockedResidualKSampler:
         lock_strength_start,
         lock_strength_end,
         lock_schedule,
+        lock_schedule_hold,
+        lock_schedule_power,
         coarse_cutoff,
         mid_band_cutoff,
         mid_band_strength,
+        mid_band_strength_start,
+        mid_band_strength_end,
+        mid_band_schedule,
+        mid_band_schedule_hold,
+        mid_band_schedule_power,
         nested_noise_strength,
         add_noise,
         pin_anchors,
@@ -444,6 +635,13 @@ class ScaleLockedResidualKSampler:
             mid_cutoff=mid_band_cutoff,
             mid_strength=mid_band_strength,
             schedule=lock_schedule,
+            schedule_power=lock_schedule_power,
+            schedule_hold=lock_schedule_hold,
+            mid_strength_start=mid_band_strength_start,
+            mid_strength_end=mid_band_strength_end,
+            mid_schedule=mid_band_schedule,
+            mid_schedule_power=mid_band_schedule_power,
+            mid_schedule_hold=mid_band_schedule_hold,
             spatial_mask=_prepare_spatial_lock_mask(lock_mask, highres_samples),
         )
 
@@ -485,6 +683,177 @@ 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": (_LOCK_SCHEDULE_OPTIONS,),
+                "lock_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
+                "lock_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
+                "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}),
+                "mid_band_strength_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "mid_band_strength_end": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
+                "mid_band_schedule": (_MID_SCHEDULE_OPTIONS,),
+                "mid_band_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
+                "mid_band_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
+                "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,
+        lock_schedule_hold,
+        lock_schedule_power,
+        coarse_cutoff,
+        mid_band_cutoff,
+        mid_band_strength,
+        mid_band_strength_start,
+        mid_band_strength_end,
+        mid_band_schedule,
+        mid_band_schedule_hold,
+        mid_band_schedule_power,
+        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,
+            schedule_power=lock_schedule_power,
+            schedule_hold=lock_schedule_hold,
+            mid_strength_start=mid_band_strength_start,
+            mid_strength_end=mid_band_strength_end,
+            mid_schedule=mid_band_schedule,
+            mid_schedule_power=mid_band_schedule_power,
+            mid_schedule_hold=mid_band_schedule_hold,
+            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 +892,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",
 }
 

pyproject.toml

@@ -0,0 +1,16 @@
+[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"
+license = { file = "LICENSE" }
+
+[project.urls]
+Repository = "https://github.com/xmarre/ComfyUI-ScaleLockedResidualDiffusion"
+Documentation = "https://github.com/xmarre/ComfyUI-ScaleLockedResidualDiffusion/blob/main/README.md"
+"Bug Tracker" = "https://github.com/xmarre/ComfyUI-ScaleLockedResidualDiffusion/issues"
+
+[tool.comfy]
+PublisherId = "scalelockedresidualdiffusion"
+DisplayName = "ComfyUI-ScaleLockedResidualDiffusion"
+Icon = ""
+includes = []

slrd_core.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 import math
-from typing import Iterable, Optional
+from typing import Iterable, Optional, Sequence
 
 import torch
 import torch.nn.functional as F
@@ -210,14 +210,67 @@ def residual_lock_multiband(
     return base_high + torch.lerp(base_low, anchor_low, low_strength) + torch.lerp(base_mid, anchor_mid, mid_strength)
 
 
-def schedule_value(start: float, end: float, progress: float, mode: str) -> float:
-    progress = float(max(0.0, min(1.0, progress)))
+def clamp(value: float, lo: float, hi: float) -> float:
+    return float(max(lo, min(hi, value)))
+
+
+def sigma_progress(planner_sigmas: Optional[Sequence[float]], idx: int) -> Optional[float]:
+    if planner_sigmas is None:
+        return None
+
+    if len(planner_sigmas) == 0:
+        return None
+
+    idx = max(0, min(int(idx), len(planner_sigmas) - 1))
+    sigma_hi = max(float(planner_sigmas[0]), 1e-6)
+    sigma_lo = max(float(planner_sigmas[-1]), 1e-6)
+    sigma = max(float(planner_sigmas[idx]), 1e-6)
+
+    hi = math.log(sigma_hi)
+    lo = math.log(sigma_lo)
+    denom = hi - lo
+    if abs(denom) < 1e-6:
+        return None
+
+    cur = math.log(sigma)
+    return clamp((hi - cur) / denom, 0.0, 1.0)
+
+
+def schedule_curve(progress: float, mode: str, power: float = 2.0, hold: float = 0.0) -> float:
+    p = clamp(progress, 0.0, 1.0)
+    power = max(1e-6, float(power))
+    hold = clamp(hold, 0.0, 0.95)
+
+    if mode == "flat":
+        return 0.0
+    if mode == "linear":
+        return p
     if mode == "cosine":
-        t = 0.5 - 0.5 * math.cos(math.pi * progress)
-    elif mode == "flat":
-        t = 0.0
-    else:
-        t = progress
+        return 0.5 - 0.5 * math.cos(math.pi * p)
+    if mode == "smoothstep":
+        return p * p * (3.0 - 2.0 * p)
+    if mode == "smootherstep":
+        return p * p * p * (p * (p * 6.0 - 15.0) + 10.0)
+    if mode == "ease_in":
+        return p ** power
+    if mode == "ease_out":
+        return 1.0 - (1.0 - p) ** power
+    if mode == "ease_in_out":
+        if p < 0.5:
+            return 0.5 * ((2.0 * p) ** power)
+        return 1.0 - 0.5 * ((2.0 * (1.0 - p)) ** power)
+    if mode == "hold_then_drop":
+        if p <= hold:
+            return 0.0
+        u = (p - hold) / max(1e-6, 1.0 - hold)
+        return u ** power
+    if mode == "fast_drop":
+        return p ** (1.0 / power)
+    return p
+
+
+def schedule_value(start: float, end: float, progress: float, mode: str, power: float = 2.0, hold: float = 0.0) -> float:
+    t = schedule_curve(progress, mode, power=power, hold=hold)
     return float(start + (end - start) * t)
 
 
@@ -239,7 +292,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.
 
@@ -261,86 +314,205 @@ class ScaleLockedCFGGuider(torch.nn.Module):
         mid_cutoff: float,
         mid_strength: float,
         schedule: str,
+        schedule_power: float,
+        schedule_hold: float,
+        mid_strength_start: float,
+        mid_strength_end: float,
+        mid_schedule: str,
+        mid_schedule_power: float,
+        mid_schedule_hold: float,
         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,
+            schedule_power=schedule_power,
+            schedule_hold=schedule_hold,
+            mid_strength_start=mid_strength_start,
+            mid_strength_end=mid_strength_end,
+            mid_schedule=mid_schedule,
+            mid_schedule_power=mid_schedule_power,
+            mid_schedule_hold=mid_schedule_hold,
+            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,
+    schedule_power: float,
+    schedule_hold: float,
+    mid_strength_start: float,
+    mid_strength_end: float,
+    mid_schedule: str,
+    mid_schedule_power: float,
+    mid_schedule_hold: float,
+    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_schedule_power = float(schedule_power)
+    guider._slrd_schedule_hold = float(schedule_hold)
+    guider._slrd_mid_strength_start = float(mid_strength_start)
+    guider._slrd_mid_strength_end = float(mid_strength_end)
+    guider._slrd_mid_schedule = mid_schedule
+    guider._slrd_mid_schedule_power = float(mid_schedule_power)
+    guider._slrd_mid_schedule_hold = float(mid_schedule_hold)
+    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:
+    return scale_lock_strengths_for_step(guider, idx)[0]
+
+
+def scale_lock_progress_for_step(guider, idx: int) -> float:
+    sigma_based = sigma_progress(getattr(guider, "_slrd_planner_sigmas", None), idx)
+    if sigma_based is not None:
+        return sigma_based
+    total = max(1, len(guider._slrd_anchors_x0_cpu) - 1)
+    return clamp(idx / total, 0.0, 1.0)
+
+
+def _scheduled_strength(
+    base_strength: float,
+    start: float,
+    end: float,
+    progress: float,
+    mode: str,
+    power: float,
+    hold: float,
+) -> float:
+    scheduled = schedule_value(start, end, progress, mode, power=power, hold=hold)
+    return clamp(base_strength * scheduled, 0.0, 1.0)
+
+
+def scale_lock_strengths_for_step(guider, idx: int) -> tuple[float, float]:
+    progress = scale_lock_progress_for_step(guider, idx)
+    low_strength = _scheduled_strength(
+        guider._slrd_lock_strength,
+        guider._slrd_lock_strength_start,
+        guider._slrd_lock_strength_end,
+        progress,
+        guider._slrd_schedule,
+        guider._slrd_schedule_power,
+        guider._slrd_schedule_hold,
+    )
+    if getattr(guider, "_slrd_mid_schedule", "linked") == "linked":
+        mid_strength = clamp(low_strength * guider._slrd_mid_strength, 0.0, 1.0)
+    else:
+        mid_strength = _scheduled_strength(
+            guider._slrd_mid_strength,
+            guider._slrd_mid_strength_start,
+            guider._slrd_mid_strength_end,
+            progress,
+            guider._slrd_mid_schedule,
+            guider._slrd_mid_schedule_power,
+            guider._slrd_mid_schedule_hold,
+        )
+    return low_strength, mid_strength
+
+
+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