xmarre-ComfyUI-ScaleLockedResidualDiffusion/slrd_runtime.py

from __future__ import annotations

import copy
import logging
import types
from dataclasses import dataclass
from typing import Any, Optional

import torch

import comfy.model_management
import comfy.sample
import comfy.samplers
import comfy.utils
import latent_preview

from .slrd_core import (
    TrajectoryRecorder,
    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,
)


_LOGGER = logging.getLogger(__name__)



class _NullPreviewCallback:
    def __call__(self, step, x0, x, total_steps):
        del step, x0, x, total_steps


_CONSERVATIVE_SAFE_SAMPLERS = {
    "ddim",
    "euler",
    "euler_cfg_pp",
    "heun",
    "lcm",
    "dpmpp_2m",
    "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]


@dataclass
class ScaleLockConfig:
    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
    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
class ScaleLockedRuntimeContext:
    model: Any
    highres_latent: dict
    lowres_latent: dict
    lowres_out: dict
    sigmas: torch.Tensor
    anchors_x0: list[torch.Tensor]
    planner_sigmas: list[float]
    highres_noise: torch.Tensor
    noise_seed: int

    def prepared_noise(self) -> "ScaleLockedPreparedNoise":
        return ScaleLockedPreparedNoise(self.highres_noise, self.noise_seed)


@dataclass
class ScaleLockedSampleResult:
    output: dict
    lowres_planner: dict
    denoised_output: dict


class ScaleLockedPreparedNoise:
    def __init__(self, noise_tensor: torch.Tensor, seed: int):
        self._noise_tensor = noise_tensor.detach().to(device="cpu").contiguous().clone()
        self.seed = int(seed)

    def generate_noise(self, latent: dict) -> torch.Tensor:
        latent_samples = latent["samples"]
        if tuple(latent_samples.shape) != tuple(self._noise_tensor.shape):
            raise ValueError(
                "ScaleLockedPreparedNoise expected latent shape "
                f"{tuple(self._noise_tensor.shape)} but received {tuple(latent_samples.shape)}."
            )
        return self._noise_tensor.clone()


def _normalize_sampler_name(name: str) -> str:
    return str(name).strip().lower()


def guard_sampler_alignment(sampler_name: str, mode: str) -> None:
    mode = str(mode).strip().lower()
    if mode == "off":
        return

    normalized = _normalize_sampler_name(sampler_name)
    if normalized in _CONSERVATIVE_SAFE_SAMPLERS:
        return

    msg = (
        "ScaleLockedResidualKSampler: sampler "
        f"'{sampler_name}' is outside the conservative SLRD alignment-safe allowlist. "
        "The node will still work in many cases, but planner/high-res anchor matching is less trustworthy "
        "for samplers with more complex internal evaluation patterns."
    )
    if mode == "error":
        raise ValueError(msg)
    _LOGGER.warning(msg)


def compat_sampler_names():
    return getattr(comfy.samplers, "SAMPLER_NAMES", comfy.samplers.KSampler.SAMPLERS)


def compat_scheduler_names():
    return getattr(comfy.samplers, "SCHEDULER_NAMES", comfy.samplers.KSampler.SCHEDULERS)


def clean_latent(latent: dict) -> dict:
    out = clone_latent(latent)
    out.pop("downscale_ratio_spacial", None)
    return out


def _model_sampling_obj(model):
    if hasattr(model, "get_model_object"):
        return model.get_model_object("model_sampling")
    if hasattr(model, "model") and hasattr(model.model, "model_sampling"):
        return model.model.model_sampling
    raise AttributeError("Unable to resolve model_sampling from the ComfyUI model patcher.")


def calculate_sigmas(model, scheduler: str, steps: int, denoise: float) -> torch.Tensor:
    total_steps = int(steps)
    if denoise < 1.0:
        if denoise <= 0.0:
            return torch.FloatTensor([])
        total_steps = int(steps / denoise)
    sigmas = comfy.samplers.calculate_sigmas(_model_sampling_obj(model), scheduler, total_steps).cpu()
    if denoise < 1.0:
        sigmas = sigmas[-(steps + 1) :]
    return sigmas


def prepare_noise(latent_samples: torch.Tensor, seed: int, batch_inds=None, disable_noise: bool = False) -> torch.Tensor:
    if disable_noise:
        return torch.zeros(latent_samples.size(), dtype=latent_samples.dtype, layout=latent_samples.layout, device="cpu")
    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)
    out["samples"] = comfy.sample.fix_empty_latent_channels(model, out["samples"], ratio)
    return out


def make_lowres_latent(latent: dict, target_megapixels: float) -> dict:
    low_hw = latent_target_hw_from_megapixels(latent["samples"], target_megapixels)
    return resize_latent_dict(latent, low_hw)


def _store_dtype_for(x: torch.Tensor) -> torch.dtype:
    if x.dtype in (torch.float32, torch.float16, torch.bfloat16):
        return x.dtype
    return torch.float32


def prepare_spatial_lock_mask(mask: Optional[torch.Tensor], latent_samples: torch.Tensor) -> Optional[torch.Tensor]:
    if mask is None:
        return None
    return resize_mask(mask, tuple(latent_samples.shape[-2:]), latent_samples.shape[0], latent_samples.shape[1])


def make_preview_callback(model, steps: int, x0_output: dict):
    try:
        return latent_preview.prepare_callback(model, max(0, steps), x0_output)
    except Exception:
        return _NullPreviewCallback()


def _planner_sigmas_for_recorded_steps(sigmas: torch.Tensor, recorded_steps: int) -> list[float]:
    if recorded_steps <= 0:
        return []

    sigma_values = [float(v) for v in sigmas.detach().flatten().cpu().tolist()]
    if not sigma_values:
        return []

    visible_sigmas = sigma_values[:-1] if len(sigma_values) > 1 else sigma_values
    if not visible_sigmas:
        visible_sigmas = sigma_values

    if recorded_steps <= len(visible_sigmas):
        return visible_sigmas[:recorded_steps]
    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 _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)
    anchor = scale_lock_anchor_for(guider, idx, 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


def create_cfg_guider(model, positive, negative, cfg):
    guider = comfy.samplers.CFGGuider(model)
    guider.set_conds(positive, negative)
    guider.set_cfg(cfg)
    return guider


def clone_guider_for_scale_lock(guider):
    cloned = copy.copy(guider)
    if hasattr(guider, "__dict__"):
        cloned.__dict__ = dict(guider.__dict__)

    original_predict_noise = getattr(cloned, "_slrd_original_predict_noise", None)
    if original_predict_noise is not None:
        if isinstance(original_predict_noise, types.MethodType):
            cloned.predict_noise = types.MethodType(original_predict_noise.__func__, cloned)
        else:
            cloned.predict_noise = original_predict_noise

    stale_keys = [key for key in getattr(cloned, "__dict__", {}) if key.startswith("_slrd_")]
    for key in stale_keys:
        delattr(cloned, key)

    return cloned


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,
        anchors_x0_cpu=runtime.anchors_x0,
        planner_sigmas=runtime.planner_sigmas,
        lock_strength=config.lock_strength,
        lock_strength_start=config.lock_strength_start,
        lock_strength_end=config.lock_strength_end,
        cutoff=config.cutoff,
        mid_cutoff=config.mid_cutoff,
        mid_strength=config.mid_strength,
        schedule=config.schedule,
        schedule_power=config.schedule_power,
        schedule_hold=config.schedule_hold,
        mid_strength_start=config.mid_strength_start,
        mid_strength_end=config.mid_strength_end,
        mid_schedule=config.mid_schedule,
        mid_schedule_power=config.mid_schedule_power,
        mid_schedule_hold=config.mid_schedule_hold,
        spatial_mask=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)
    guider._slrd_original_predict_noise = original_predict_noise

    def _wrapped_predict_noise(self, x, timestep, model_options=None, seed=None):
        if model_options is None:
            model_options = {}
        base_noise = self._slrd_original_predict_noise(x, timestep, model_options=model_options, seed=seed)
        return apply_scale_lock_to_noise_prediction(self, base_noise, x, timestep)

    guider.predict_noise = types.MethodType(_wrapped_predict_noise, guider)
    return guider


def restore_original_predict_noise(guider) -> None:
    original_predict_noise = getattr(guider, "_slrd_original_predict_noise", None)
    if original_predict_noise is not None:
        guider.predict_noise = 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"]
    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,
        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,
    positive,
    negative,
    cfg: float,
    sampler_name: str,
    sigmas: torch.Tensor,
    lowres_latent: dict,
    seed: int,
    disable_noise: bool,
    pin_anchors: bool,
) -> tuple[dict, list[torch.Tensor], list[float], torch.Tensor]:
    sampler_obj = comfy.samplers.sampler_object(sampler_name)
    guider = comfy.samplers.CFGGuider(model)
    guider.set_conds(positive, negative)
    guider.set_cfg(cfg)

    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,
    ).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,
        pin_memory=pin_anchors,
    )

    samples = guider.sample(
        planner_noise,
        latent_samples,
        sampler_obj,
        sigmas,
        denoise_mask=noise_mask,
        callback=recorder.callback,
        disable_pbar=True,
        seed=seed,
    )
    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, planner_noise


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=target_dtype,
            layout=highres_samples.layout,
            device=target_device,
        )

    return build_nested_noise(
        lowres_noise=lowres_noise,
        target_shape=tuple(highres_samples.shape),
        seed=seed,
        hf_strength=hf_strength,
    ).to(device=target_device, dtype=target_dtype, non_blocking=True)


def build_runtime_context_from_advanced(
    *,
    noise,
    guider,
    sampler,
    sigmas: torch.Tensor,
    latent_image,
    target_megapixels: float,
    nested_noise_strength: float,
    pin_anchors: bool,
) -> 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:
        return ScaleLockedRuntimeContext(
            model=model,
            highres_latent=highres_latent,
            lowres_latent=lowres_latent,
            lowres_out=clean_latent(lowres_latent),
            sigmas=sigmas.to(device=target_device, non_blocking=True),
            anchors_x0=[],
            planner_sigmas=[],
            highres_noise=torch.zeros_like(
                highres_latent["samples"],
                device=target_device,
            ),
            noise_seed=noise_seed(noise),
        )

    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_noise = _build_highres_noise(
        highres_latent=highres_latent,
        lowres_noise=lowres_noise,
        seed=noise_seed(noise) ^ 0x9E3779B97F4A7C15,
        hf_strength=nested_noise_strength,
    )
    return ScaleLockedRuntimeContext(
        model=model,
        highres_latent=highres_latent,
        lowres_latent=lowres_latent,
        lowres_out=lowres_out,
        sigmas=sigmas.to(device=target_device, non_blocking=True),
        anchors_x0=anchors_x0,
        planner_sigmas=planner_sigmas,
        highres_noise=highres_noise,
        noise_seed=noise_seed(noise),
    )


def sample_with_runtime(
    *,
    guider,
    sampler,
    runtime: ScaleLockedRuntimeContext,
    config: ScaleLockConfig,
    restore_after: bool = True,
) -> ScaleLockedSampleResult:
    if runtime.sigmas.numel() == 0:
        out = clean_latent(runtime.highres_latent)
        return ScaleLockedSampleResult(output=out, lowres_planner=clean_latent(runtime.lowres_out), denoised_output=out)

    apply_scale_lock_to_guider(guider, runtime, config)

    x0_output = {}
    callback = make_preview_callback(runtime.model, len(runtime.sigmas) - 1, x0_output)
    noise_mask = runtime.highres_latent.get("noise_mask", None)
    disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED

    try:
        samples = guider.sample(
            runtime.highres_noise,
            runtime.highres_latent["samples"],
            sampler,
            runtime.sigmas,
            denoise_mask=noise_mask,
            callback=callback,
            disable_pbar=disable_pbar,
            seed=runtime.noise_seed,
        )
    finally:
        if restore_after:
            restore_original_predict_noise(guider)
    samples = samples.to(comfy.model_management.intermediate_device())

    out = clean_latent(runtime.highres_latent)
    out["samples"] = samples

    if "x0" in x0_output:
        try:
            x0_out = runtime.model.model.process_latent_out(x0_output["x0"].cpu())
        except Exception:
            x0_out = x0_output["x0"].detach().cpu()
        denoised = clean_latent(runtime.highres_latent)
        denoised["samples"] = x0_out
    else:
        denoised = out

    return ScaleLockedSampleResult(
        output=out,
        lowres_planner=clean_latent(runtime.lowres_out),
        denoised_output=denoised,
    )


def run_scale_locked_ksampler(
    *,
    model,
    positive,
    negative,
    latent_image,
    seed,
    steps,
    cfg,
    sampler_name,
    scheduler,
    denoise,
    target_megapixels,
    nested_noise_strength,
    add_noise,
    pin_anchors,
    sampler_guard,
    config: ScaleLockConfig,
) -> ScaleLockedSampleResult:
    if steps < 1:
        raise ValueError("steps must be >= 1")

    highres_latent = fix_latent_channels(model, latent_image)
    lowres_latent = make_lowres_latent(highres_latent, target_megapixels)
    if denoise <= 0.0:
        out = clean_latent(highres_latent)
        return ScaleLockedSampleResult(output=out, lowres_planner=clean_latent(lowres_latent), denoised_output=out)

    guard_sampler_alignment(sampler_name, sampler_guard)
    sigmas = calculate_sigmas(model, scheduler=scheduler, steps=steps, denoise=denoise)
    if sigmas.numel() == 0:
        out = clean_latent(highres_latent)
        return ScaleLockedSampleResult(output=out, lowres_planner=clean_latent(lowres_latent), denoised_output=out)

    disable_noise = not bool(add_noise)
    planner_seed = int(seed)
    detail_seed = int(seed) ^ 0x9E3779B97F4A7C15

    lowres_out, anchors_x0, planner_sigmas, lowres_noise = _run_lowres_planner(
        model=model,
        positive=positive,
        negative=negative,
        cfg=cfg,
        sampler_name=sampler_name,
        sigmas=sigmas,
        lowres_latent=lowres_latent,
        seed=planner_seed,
        disable_noise=disable_noise,
        pin_anchors=pin_anchors,
    )
    if len(anchors_x0) == 0:
        raise RuntimeError("ScaleLockedResidualKSampler: planner pass did not record any x0 anchors.")

    runtime = ScaleLockedRuntimeContext(
        model=model,
        highres_latent=highres_latent,
        lowres_latent=lowres_latent,
        lowres_out=lowres_out,
        sigmas=sigmas,
        anchors_x0=anchors_x0,
        planner_sigmas=planner_sigmas,
        highres_noise=_build_highres_noise(
            highres_latent=highres_latent,
            lowres_noise=lowres_noise,
            seed=detail_seed,
            hf_strength=nested_noise_strength,
        ),
        noise_seed=detail_seed,
    )
    guider = create_cfg_guider(model, positive, negative, cfg)
    sampler = comfy.samplers.sampler_object(sampler_name)
    return sample_with_runtime(guider=guider, sampler=sampler, runtime=runtime, config=config)