coderabbitai[bot]
47aa1ddfbe
Docstrings generation was requested by @xmarre. * https://github.com/xmarre/ComfyUI-ScaleLockedResidualDiffusion/pull/10#issuecomment-4062975146 The following files were modified: * `nodes.py`
1627 lines
72 KiB
Python
1627 lines
72 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, replace
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
import math
|
|
import torch
|
|
import comfy.samplers
|
|
|
|
from .slrd_core import build_nested_noise, clone_latent, init_scale_lock_state, resize_4d_tensor
|
|
from .slrd_runtime import (
|
|
LOCK_SCHEDULE_OPTIONS,
|
|
MID_SCHEDULE_OPTIONS,
|
|
ScaleLockConfig,
|
|
apply_scale_lock_to_noise_prediction,
|
|
build_runtime_context_from_advanced,
|
|
clean_latent,
|
|
compat_sampler_names,
|
|
compat_scheduler_names,
|
|
create_cfg_guider,
|
|
fix_latent_channels,
|
|
guard_sampler_alignment,
|
|
make_lowres_latent,
|
|
prepare_noise,
|
|
run_scale_locked_ksampler,
|
|
sample_with_runtime,
|
|
apply_scale_lock_to_guider,
|
|
clone_guider_for_scale_lock,
|
|
)
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ScaleLockedResidualKSampler:
|
|
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": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"steps": ("INT", {"default": 24, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"sampler_name": (compat_sampler_names(),),
|
|
"scheduler": (compat_scheduler_names(),),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"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}),
|
|
"manifold_enabled": ("BOOLEAN", {"default": False}),
|
|
"manifold_strength": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_end": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule": (LOCK_SCHEDULE_OPTIONS,),
|
|
"manifold_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
|
|
"manifold_cutoff": ("FLOAT", {"default": 0.18, "min": 0.05, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_radial_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anisotropy": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 2.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_translation_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anchor_mix": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_mean_anchor_mix": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_contrast_restore": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_channel_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_gain_cap": ("FLOAT", {"default": 1.75, "min": 1.0, "max": 4.0, "step": 0.05, "round": 0.01}),
|
|
"manifold_max_shift_px": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.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}),
|
|
"sampler_guard": (["warn", "error", "off"],),
|
|
},
|
|
"optional": {
|
|
"lock_mask": ("MASK",),
|
|
"manifold_mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
denoise,
|
|
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,
|
|
manifold_enabled,
|
|
manifold_strength,
|
|
manifold_strength_start,
|
|
manifold_strength_end,
|
|
manifold_schedule,
|
|
manifold_schedule_hold,
|
|
manifold_schedule_power,
|
|
manifold_cutoff,
|
|
manifold_radial_strength,
|
|
manifold_anisotropy,
|
|
manifold_translation_strength,
|
|
manifold_anchor_mix,
|
|
manifold_mean_anchor_mix,
|
|
manifold_contrast_restore,
|
|
manifold_energy_tether,
|
|
manifold_channel_tether,
|
|
manifold_energy_gain_cap,
|
|
manifold_max_shift_px,
|
|
nested_noise_strength,
|
|
add_noise,
|
|
pin_anchors,
|
|
sampler_guard,
|
|
lock_mask=None,
|
|
manifold_mask=None,
|
|
):
|
|
result = run_scale_locked_ksampler(
|
|
model=model,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_image,
|
|
seed=seed,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
denoise=denoise,
|
|
target_megapixels=target_megapixels,
|
|
nested_noise_strength=nested_noise_strength,
|
|
add_noise=add_noise,
|
|
pin_anchors=pin_anchors,
|
|
sampler_guard=sampler_guard,
|
|
config=_scale_lock_config(
|
|
lock_strength=lock_strength,
|
|
lock_strength_start=lock_strength_start,
|
|
lock_strength_end=lock_strength_end,
|
|
coarse_cutoff=coarse_cutoff,
|
|
mid_band_cutoff=mid_band_cutoff,
|
|
mid_band_strength=mid_band_strength,
|
|
lock_schedule=lock_schedule,
|
|
lock_schedule_hold=lock_schedule_hold,
|
|
lock_schedule_power=lock_schedule_power,
|
|
mid_band_strength_start=mid_band_strength_start,
|
|
mid_band_strength_end=mid_band_strength_end,
|
|
mid_band_schedule=mid_band_schedule,
|
|
mid_band_schedule_hold=mid_band_schedule_hold,
|
|
mid_band_schedule_power=mid_band_schedule_power,
|
|
manifold_enabled=manifold_enabled,
|
|
manifold_strength=manifold_strength,
|
|
manifold_strength_start=manifold_strength_start,
|
|
manifold_strength_end=manifold_strength_end,
|
|
manifold_schedule=manifold_schedule,
|
|
manifold_schedule_hold=manifold_schedule_hold,
|
|
manifold_schedule_power=manifold_schedule_power,
|
|
manifold_cutoff=manifold_cutoff,
|
|
manifold_radial_strength=manifold_radial_strength,
|
|
manifold_anisotropy=manifold_anisotropy,
|
|
manifold_translation_strength=manifold_translation_strength,
|
|
manifold_anchor_mix=manifold_anchor_mix,
|
|
manifold_mean_anchor_mix=manifold_mean_anchor_mix,
|
|
manifold_contrast_restore=manifold_contrast_restore,
|
|
manifold_energy_tether=manifold_energy_tether,
|
|
manifold_channel_tether=manifold_channel_tether,
|
|
manifold_energy_gain_cap=manifold_energy_gain_cap,
|
|
manifold_max_shift_px=manifold_max_shift_px,
|
|
lock_mask=lock_mask,
|
|
manifold_mask=manifold_mask,
|
|
),
|
|
)
|
|
return (result.output, result.lowres_planner, result.denoised_output)
|
|
|
|
|
|
class ScaleLockedRuntimeContextBuilder:
|
|
CATEGORY = "sampling/scale_locked"
|
|
RETURN_TYPES = ("SLRD_RUNTIME", "NOISE", "LATENT")
|
|
RETURN_NAMES = ("runtime", "prepared_noise", "lowres_planner")
|
|
FUNCTION = "build"
|
|
|
|
@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}),
|
|
"nested_noise_strength": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 4.0, "step": 0.01, "round": 0.001}),
|
|
"pin_anchors": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
def build(self, noise, guider, sampler, sigmas, latent_image, target_megapixels, nested_noise_strength, pin_anchors):
|
|
runtime = build_runtime_context_from_advanced(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=sigmas,
|
|
latent_image=latent_image,
|
|
target_megapixels=target_megapixels,
|
|
nested_noise_strength=nested_noise_strength,
|
|
pin_anchors=pin_anchors,
|
|
)
|
|
return (runtime, runtime.prepared_noise(), clean_latent(runtime.lowres_out))
|
|
|
|
|
|
class ScaleLockedCFGGuiderNode:
|
|
CATEGORY = "sampling/scale_locked"
|
|
RETURN_TYPES = ("GUIDER",)
|
|
RETURN_NAMES = ("guider",)
|
|
FUNCTION = "build"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"guider": ("GUIDER",),
|
|
"runtime": ("SLRD_RUNTIME",),
|
|
"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}),
|
|
"manifold_enabled": ("BOOLEAN", {"default": False}),
|
|
"manifold_strength": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_end": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule": (LOCK_SCHEDULE_OPTIONS,),
|
|
"manifold_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
|
|
"manifold_cutoff": ("FLOAT", {"default": 0.18, "min": 0.05, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_radial_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anisotropy": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 2.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_translation_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anchor_mix": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_mean_anchor_mix": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_contrast_restore": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_channel_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_gain_cap": ("FLOAT", {"default": 1.75, "min": 1.0, "max": 4.0, "step": 0.05, "round": 0.01}),
|
|
"manifold_max_shift_px": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.0, "step": 0.1, "round": 0.01}),
|
|
},
|
|
"optional": {
|
|
"lock_mask": ("MASK",),
|
|
"manifold_mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
def build(
|
|
self,
|
|
guider,
|
|
runtime,
|
|
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,
|
|
manifold_enabled,
|
|
manifold_strength,
|
|
manifold_strength_start,
|
|
manifold_strength_end,
|
|
manifold_schedule,
|
|
manifold_schedule_hold,
|
|
manifold_schedule_power,
|
|
manifold_cutoff,
|
|
manifold_radial_strength,
|
|
manifold_anisotropy,
|
|
manifold_translation_strength,
|
|
manifold_anchor_mix,
|
|
manifold_mean_anchor_mix,
|
|
manifold_contrast_restore,
|
|
manifold_energy_tether,
|
|
manifold_channel_tether,
|
|
manifold_energy_gain_cap,
|
|
manifold_max_shift_px,
|
|
lock_mask=None,
|
|
manifold_mask=None,
|
|
):
|
|
patched_guider = clone_guider_for_scale_lock(guider)
|
|
apply_scale_lock_to_guider(
|
|
patched_guider,
|
|
runtime,
|
|
_scale_lock_config(
|
|
lock_strength=lock_strength,
|
|
lock_strength_start=lock_strength_start,
|
|
lock_strength_end=lock_strength_end,
|
|
coarse_cutoff=coarse_cutoff,
|
|
mid_band_cutoff=mid_band_cutoff,
|
|
mid_band_strength=mid_band_strength,
|
|
lock_schedule=lock_schedule,
|
|
lock_schedule_hold=lock_schedule_hold,
|
|
lock_schedule_power=lock_schedule_power,
|
|
mid_band_strength_start=mid_band_strength_start,
|
|
mid_band_strength_end=mid_band_strength_end,
|
|
mid_band_schedule=mid_band_schedule,
|
|
mid_band_schedule_hold=mid_band_schedule_hold,
|
|
mid_band_schedule_power=mid_band_schedule_power,
|
|
manifold_enabled=manifold_enabled,
|
|
manifold_strength=manifold_strength,
|
|
manifold_strength_start=manifold_strength_start,
|
|
manifold_strength_end=manifold_strength_end,
|
|
manifold_schedule=manifold_schedule,
|
|
manifold_schedule_hold=manifold_schedule_hold,
|
|
manifold_schedule_power=manifold_schedule_power,
|
|
manifold_cutoff=manifold_cutoff,
|
|
manifold_radial_strength=manifold_radial_strength,
|
|
manifold_anisotropy=manifold_anisotropy,
|
|
manifold_translation_strength=manifold_translation_strength,
|
|
manifold_anchor_mix=manifold_anchor_mix,
|
|
manifold_mean_anchor_mix=manifold_mean_anchor_mix,
|
|
manifold_contrast_restore=manifold_contrast_restore,
|
|
manifold_energy_tether=manifold_energy_tether,
|
|
manifold_channel_tether=manifold_channel_tether,
|
|
manifold_energy_gain_cap=manifold_energy_gain_cap,
|
|
manifold_max_shift_px=manifold_max_shift_px,
|
|
lock_mask=lock_mask,
|
|
manifold_mask=manifold_mask,
|
|
),
|
|
)
|
|
return (patched_guider,)
|
|
|
|
|
|
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}),
|
|
"manifold_enabled": ("BOOLEAN", {"default": False}),
|
|
"manifold_strength": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_end": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule": (LOCK_SCHEDULE_OPTIONS,),
|
|
"manifold_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
|
|
"manifold_cutoff": ("FLOAT", {"default": 0.18, "min": 0.05, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_radial_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anisotropy": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 2.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_translation_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anchor_mix": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_mean_anchor_mix": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_contrast_restore": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_channel_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_gain_cap": ("FLOAT", {"default": 1.75, "min": 1.0, "max": 4.0, "step": 0.05, "round": 0.01}),
|
|
"manifold_max_shift_px": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.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",),
|
|
"manifold_mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
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,
|
|
manifold_enabled,
|
|
manifold_strength,
|
|
manifold_strength_start,
|
|
manifold_strength_end,
|
|
manifold_schedule,
|
|
manifold_schedule_hold,
|
|
manifold_schedule_power,
|
|
manifold_cutoff,
|
|
manifold_radial_strength,
|
|
manifold_anisotropy,
|
|
manifold_translation_strength,
|
|
manifold_anchor_mix,
|
|
manifold_mean_anchor_mix,
|
|
manifold_contrast_restore,
|
|
manifold_energy_tether,
|
|
manifold_channel_tether,
|
|
manifold_energy_gain_cap,
|
|
manifold_max_shift_px,
|
|
nested_noise_strength,
|
|
pin_anchors,
|
|
lock_mask=None,
|
|
manifold_mask=None,
|
|
):
|
|
runtime = build_runtime_context_from_advanced(
|
|
noise=noise,
|
|
guider=guider,
|
|
sampler=sampler,
|
|
sigmas=sigmas,
|
|
latent_image=latent_image,
|
|
target_megapixels=target_megapixels,
|
|
nested_noise_strength=nested_noise_strength,
|
|
pin_anchors=pin_anchors,
|
|
)
|
|
result = sample_with_runtime(
|
|
guider=guider,
|
|
sampler=sampler,
|
|
runtime=runtime,
|
|
config=_scale_lock_config(
|
|
lock_strength=lock_strength,
|
|
lock_strength_start=lock_strength_start,
|
|
lock_strength_end=lock_strength_end,
|
|
coarse_cutoff=coarse_cutoff,
|
|
mid_band_cutoff=mid_band_cutoff,
|
|
mid_band_strength=mid_band_strength,
|
|
lock_schedule=lock_schedule,
|
|
lock_schedule_hold=lock_schedule_hold,
|
|
lock_schedule_power=lock_schedule_power,
|
|
mid_band_strength_start=mid_band_strength_start,
|
|
mid_band_strength_end=mid_band_strength_end,
|
|
mid_band_schedule=mid_band_schedule,
|
|
mid_band_schedule_hold=mid_band_schedule_hold,
|
|
mid_band_schedule_power=mid_band_schedule_power,
|
|
manifold_enabled=manifold_enabled,
|
|
manifold_strength=manifold_strength,
|
|
manifold_strength_start=manifold_strength_start,
|
|
manifold_strength_end=manifold_strength_end,
|
|
manifold_schedule=manifold_schedule,
|
|
manifold_schedule_hold=manifold_schedule_hold,
|
|
manifold_schedule_power=manifold_schedule_power,
|
|
manifold_cutoff=manifold_cutoff,
|
|
manifold_radial_strength=manifold_radial_strength,
|
|
manifold_anisotropy=manifold_anisotropy,
|
|
manifold_translation_strength=manifold_translation_strength,
|
|
manifold_anchor_mix=manifold_anchor_mix,
|
|
manifold_mean_anchor_mix=manifold_mean_anchor_mix,
|
|
manifold_contrast_restore=manifold_contrast_restore,
|
|
manifold_energy_tether=manifold_energy_tether,
|
|
manifold_channel_tether=manifold_channel_tether,
|
|
manifold_energy_gain_cap=manifold_energy_gain_cap,
|
|
manifold_max_shift_px=manifold_max_shift_px,
|
|
lock_mask=lock_mask,
|
|
manifold_mask=manifold_mask,
|
|
),
|
|
)
|
|
return (result.output, result.lowres_planner, result.denoised_output)
|
|
|
|
|
|
class ScaleLockedNestedNoisePreview:
|
|
CATEGORY = "sampling/scale_locked"
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("nested_noise_latent", "lowres_reference_latent")
|
|
FUNCTION = "build"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"target_megapixels": ("FLOAT", {"default": 1.0, "min": 0.10, "max": 16.0, "step": 0.05, "round": 0.01}),
|
|
"nested_noise_strength": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 4.0, "step": 0.01, "round": 0.001}),
|
|
}
|
|
}
|
|
|
|
def build(self, latent_image, seed, target_megapixels, nested_noise_strength):
|
|
high = clone_latent(latent_image)
|
|
low = make_lowres_latent(high, target_megapixels)
|
|
|
|
low_samples = low["samples"]
|
|
high_samples = high["samples"]
|
|
low_noise = prepare_noise(low_samples, int(seed), low.get("batch_index", None))
|
|
high_noise = build_nested_noise(
|
|
lowres_noise=low_noise,
|
|
target_shape=tuple(high_samples.shape),
|
|
seed=int(seed) ^ 0x9E3779B97F4A7C15,
|
|
hf_strength=nested_noise_strength,
|
|
)
|
|
|
|
out = clone_latent(high)
|
|
out["samples"] = high_noise
|
|
return (out, low)
|
|
|
|
|
|
@dataclass
|
|
class _ImpactSampleRequest:
|
|
model: Any
|
|
seed: int
|
|
steps: int
|
|
cfg: float
|
|
sampler_name: str
|
|
scheduler: str
|
|
positive: Any
|
|
negative: Any
|
|
latent: dict
|
|
denoise: float
|
|
|
|
|
|
_IMPACT_FIELD_ALIASES = {
|
|
"model": ("model",),
|
|
"seed": ("seed",),
|
|
"steps": ("steps",),
|
|
"cfg": ("cfg",),
|
|
"sampler_name": ("sampler_name", "sampler"),
|
|
"scheduler": ("scheduler",),
|
|
"positive": ("positive", "positive_cond", "cond"),
|
|
"negative": ("negative", "negative_cond", "uncond"),
|
|
"latent": ("latent", "latent_image"),
|
|
"denoise": ("denoise",),
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class _ImpactHookSettings:
|
|
target_megapixels: float
|
|
nested_noise_strength: float
|
|
add_noise: bool
|
|
pin_anchors: bool
|
|
sampler_guard: str
|
|
warn_if_inactive_sampler: bool
|
|
config: ScaleLockConfig
|
|
|
|
|
|
@dataclass
|
|
class _ScaleLockedImpactRuntimeState:
|
|
request: _ImpactSampleRequest
|
|
runtime: Any
|
|
config: ScaleLockConfig
|
|
|
|
|
|
def _expand_mask_for_like(mask: torch.Tensor | None, like: torch.Tensor) -> torch.Tensor | None:
|
|
if mask is None:
|
|
return None
|
|
|
|
m = mask
|
|
if m.ndim == 2:
|
|
m = m.unsqueeze(0).unsqueeze(0)
|
|
elif m.ndim == 3:
|
|
m = m.unsqueeze(1)
|
|
elif m.ndim == 4:
|
|
if m.shape[-1] == 1 and m.shape[1] != 1:
|
|
m = m.permute(0, 3, 1, 2)
|
|
elif m.shape[1] != 1:
|
|
m = m.mean(dim=1, keepdim=True)
|
|
else:
|
|
raise ValueError(f"Unsupported mask rank {m.ndim} for shape {tuple(m.shape)}")
|
|
|
|
m = m.to(device=like.device, dtype=like.dtype, non_blocking=True)
|
|
if tuple(m.shape[-2:]) != tuple(like.shape[-2:]):
|
|
m = resize_4d_tensor(m, tuple(like.shape[-2:]))
|
|
if m.shape[0] < like.shape[0]:
|
|
repeat = math.ceil(like.shape[0] / max(1, m.shape[0]))
|
|
m = m.repeat(repeat, 1, 1, 1)[: like.shape[0]]
|
|
elif m.shape[0] > like.shape[0]:
|
|
m = m[: like.shape[0]]
|
|
if m.shape[1] == 1:
|
|
m = m.expand(like.shape[0], like.shape[1], like.shape[-2], like.shape[-1]).contiguous()
|
|
elif m.shape[1] != like.shape[1]:
|
|
m = m.mean(dim=1, keepdim=True).expand(like.shape[0], like.shape[1], like.shape[-2], like.shape[-1]).contiguous()
|
|
return m.clamp(0.0, 1.0)
|
|
|
|
|
|
def _combine_masks(*masks: torch.Tensor | None) -> torch.Tensor | None:
|
|
combined = None
|
|
for mask in masks:
|
|
if mask is None:
|
|
continue
|
|
combined = mask if combined is None else combined * mask
|
|
return None if combined is None else combined.clamp(0.0, 1.0)
|
|
|
|
|
|
def _canonicalize_impact_noise_mask(mask: Any) -> torch.Tensor | None:
|
|
if not isinstance(mask, torch.Tensor):
|
|
return None
|
|
if mask.ndim == 4:
|
|
if mask.shape[1] == 0:
|
|
return None
|
|
return mask[:, :1, :, :].squeeze(1).contiguous()
|
|
if mask.ndim == 3:
|
|
return mask.contiguous()
|
|
if mask.ndim == 2:
|
|
return mask.unsqueeze(0).contiguous()
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: ignoring unsupported denoise_mask rank %s with shape %s.",
|
|
mask.ndim,
|
|
tuple(mask.shape),
|
|
)
|
|
return None
|
|
|
|
|
|
def _impact_request_tuple_from_kwargs(kwargs: dict[str, Any]) -> tuple[Any, ...]:
|
|
if not kwargs:
|
|
return ()
|
|
|
|
data = {}
|
|
for field, aliases in _IMPACT_FIELD_ALIASES.items():
|
|
for alias in aliases:
|
|
if alias in kwargs:
|
|
data[field] = kwargs[alias]
|
|
break
|
|
if not data:
|
|
return ()
|
|
|
|
names = ("model", "seed", "steps", "cfg", "sampler_name", "scheduler", "positive", "negative", "latent", "denoise")
|
|
return tuple(data.get(name) for name in names)
|
|
|
|
|
|
class _ScaleLockedPlannerNoise:
|
|
def __init__(self, seed: int, disable_noise: bool):
|
|
self.seed = int(seed)
|
|
self.disable_noise = bool(disable_noise)
|
|
|
|
def generate_noise(self, latent: dict) -> torch.Tensor:
|
|
return prepare_noise(
|
|
latent["samples"],
|
|
seed=self.seed,
|
|
batch_inds=latent.get("batch_index", None),
|
|
disable_noise=self.disable_noise,
|
|
)
|
|
|
|
|
|
def _is_impact_guider_like(obj) -> bool:
|
|
return (
|
|
obj is not None
|
|
and hasattr(obj, "sample")
|
|
and hasattr(obj, "model_patcher")
|
|
and (hasattr(obj, "set_conds") or hasattr(obj, "inner_set_conds"))
|
|
)
|
|
|
|
|
|
def _clear_impact_ag_guider(owner) -> None:
|
|
if owner is None or not hasattr(owner, "_ag_detailer_guider"):
|
|
return
|
|
try:
|
|
delattr(owner, "_ag_detailer_guider")
|
|
return
|
|
except Exception as exc:
|
|
try:
|
|
setattr(owner, "_ag_detailer_guider", None)
|
|
except Exception as fallback_exc:
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: failed to clear _ag_detailer_guider on %s (delete=%r, fallback=%r).",
|
|
type(owner).__name__,
|
|
exc,
|
|
fallback_exc,
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: delattr(_ag_detailer_guider) failed on %s; replaced with None instead (%r).",
|
|
type(owner).__name__,
|
|
exc,
|
|
)
|
|
|
|
|
|
def _resolve_impact_ag_guider_template(request: "_ImpactSampleRequest", model_wrap):
|
|
"""
|
|
Locate an existing impact guider template from the provided request or model wrapper and return it.
|
|
|
|
If a candidate guider is found on either model_wrap.model_patcher or request.model, the guider is returned and the internal `_ag_detailer_guider` attribute is cleared from both potential owners as a side effect.
|
|
|
|
Parameters:
|
|
request (_ImpactSampleRequest): Impact request that may contain a `model` owning a guider.
|
|
model_wrap: Model wrapper that may contain a `model_patcher` owning a guider.
|
|
|
|
Returns:
|
|
The found guider instance, or `None` if no guider template is present.
|
|
"""
|
|
owners = (
|
|
getattr(model_wrap, "model_patcher", None),
|
|
getattr(request, "model", None),
|
|
)
|
|
for owner in owners:
|
|
if owner is None:
|
|
continue
|
|
candidate = getattr(owner, "_ag_detailer_guider", None)
|
|
if candidate is None:
|
|
continue
|
|
for clear_owner in owners:
|
|
_clear_impact_ag_guider(clear_owner)
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _resolve_sampler_device(model_or_wrap, fallback: torch.device) -> torch.device:
|
|
"""
|
|
Determine the device to use for sampling by inspecting the provided model or wrapper.
|
|
|
|
Parameters:
|
|
model_or_wrap: An object (model or wrapper) to inspect for a `load_device` attribute; common wrapper fields (`inner_model`, `model`, `model_patcher`) are checked as well.
|
|
fallback (torch.device): Device to return if no `load_device` attribute is found on the inspected objects.
|
|
|
|
Returns:
|
|
torch.device: The first `load_device` found on `model_or_wrap` or its common attributes, otherwise `fallback`.
|
|
"""
|
|
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 _normalize_sampler_extra_args(
|
|
extra_args: dict[str, Any] | None,
|
|
sigmas: torch.Tensor,
|
|
target_device: torch.device,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Ensure the provided extra_args include a `model_options.sigmas` tensor moved to the target device.
|
|
|
|
If `extra_args` contains a `model_options` mapping, this function copies it, sets its `sigmas` entry to `sigmas` converted to `target_device` (non-blocking), and returns a shallow-copied dict with the updated `model_options`. If `extra_args` is None or `model_options` is not a dict, returns a shallow copy of `extra_args` (or an empty dict when None).
|
|
|
|
Parameters:
|
|
extra_args (dict[str, Any] | None): Extra sampler arguments that may include a `model_options` dict.
|
|
sigmas (torch.Tensor): Sigmas tensor to insert into `model_options`.
|
|
target_device (torch.device): Device to which `sigmas` will be moved.
|
|
|
|
Returns:
|
|
dict[str, Any]: A normalized copy of `extra_args` with `model_options.sigmas` set to `sigmas` on `target_device` when applicable.
|
|
"""
|
|
normalized = {} if extra_args is None else dict(extra_args)
|
|
model_options = normalized.get("model_options", None)
|
|
if isinstance(model_options, dict):
|
|
model_options = dict(model_options)
|
|
model_options["sigmas"] = sigmas.to(device=target_device, non_blocking=True)
|
|
normalized["model_options"] = model_options
|
|
return normalized
|
|
|
|
|
|
class _ScaleLockedPlannerGuiderProxy:
|
|
def __init__(self, model_wrap, extra_args: dict[str, Any] | None):
|
|
"""
|
|
Initialize the planner guider proxy which forwards sampling calls to the underlying model wrapper.
|
|
|
|
Parameters:
|
|
model_wrap: An object that exposes a sampler-compatible interface and may provide a `model_patcher` attribute.
|
|
extra_args (dict[str, Any] | None): Optional model/sampler options to be merged and stored for use when sampling; a shallow copy is made if provided.
|
|
"""
|
|
self.model_patcher = getattr(model_wrap, "model_patcher", None)
|
|
self._model_wrap = model_wrap
|
|
self._extra_args = {} if extra_args is None else dict(extra_args)
|
|
|
|
def sample(
|
|
self,
|
|
noise,
|
|
latent_samples,
|
|
sampler,
|
|
sigmas,
|
|
denoise_mask=None,
|
|
callback=None,
|
|
disable_pbar=False,
|
|
seed=None,
|
|
):
|
|
"""
|
|
Prepare inputs on the sampler's device/dtype, normalize extra args, and forward the call to the underlying sampler.
|
|
|
|
Parameters:
|
|
denoise_mask (torch.Tensor | None): Optional mask moved to the sampler device when provided.
|
|
callback (callable | None): Optional progress callback; if None a no-op callback is used.
|
|
seed: Ignored by this proxy.
|
|
|
|
Returns:
|
|
The value returned by the underlying sampler.sample(...) call.
|
|
"""
|
|
del seed
|
|
if callback is None:
|
|
callback = lambda *args, **kwargs: None
|
|
|
|
target_device = _resolve_sampler_device(self._model_wrap, latent_samples.device)
|
|
target_dtype = latent_samples.dtype
|
|
latent_samples = latent_samples.to(
|
|
device=target_device,
|
|
dtype=target_dtype,
|
|
non_blocking=True,
|
|
)
|
|
noise = noise.to(device=target_device, dtype=target_dtype, non_blocking=True)
|
|
sigmas = sigmas.to(device=target_device, non_blocking=True)
|
|
if isinstance(denoise_mask, torch.Tensor):
|
|
denoise_mask = denoise_mask.to(device=target_device, non_blocking=True)
|
|
|
|
sample_extra_args = _normalize_sampler_extra_args(self._extra_args, sigmas, target_device)
|
|
return sampler.sample(
|
|
self._model_wrap,
|
|
sigmas,
|
|
sample_extra_args,
|
|
callback,
|
|
noise,
|
|
latent_image=latent_samples,
|
|
denoise_mask=denoise_mask,
|
|
disable_pbar=disable_pbar,
|
|
)
|
|
|
|
|
|
class _ScaleLockedGuiderProxy:
|
|
def __init__(self, base_guider, runtime, config: ScaleLockConfig):
|
|
"""
|
|
Wraps a base guider and initializes its scale-lock state using the provided runtime context and configuration.
|
|
|
|
Parameters:
|
|
base_guider: The underlying guider object to delegate to; must implement the guider interface used during sampling.
|
|
runtime: Runtime context providing model, anchors_x0, and planner_sigmas required to initialize scale-lock state.
|
|
config (ScaleLockConfig): Configuration object containing all scale-lock parameters (strengths, cutoffs, schedules, manifold controls, spatial masks, etc.) used to initialize the guider's scale-lock behavior.
|
|
"""
|
|
self._base_guider = base_guider
|
|
init_scale_lock_state(
|
|
self,
|
|
model=runtime.model,
|
|
anchors_x0_cpu=runtime.anchors_x0,
|
|
planner_sigmas=runtime.planner_sigmas,
|
|
lock_strength=config.lock_strength,
|
|
lock_strength_start=config.lock_strength_start,
|
|
lock_strength_end=config.lock_strength_end,
|
|
cutoff=config.cutoff,
|
|
mid_cutoff=config.mid_cutoff,
|
|
mid_strength=config.mid_strength,
|
|
schedule=config.schedule,
|
|
schedule_power=config.schedule_power,
|
|
schedule_hold=config.schedule_hold,
|
|
mid_strength_start=config.mid_strength_start,
|
|
mid_strength_end=config.mid_strength_end,
|
|
mid_schedule=config.mid_schedule,
|
|
mid_schedule_power=config.mid_schedule_power,
|
|
mid_schedule_hold=config.mid_schedule_hold,
|
|
spatial_mask=config.spatial_mask,
|
|
manifold_enabled=config.manifold_enabled,
|
|
manifold_strength=config.manifold_strength,
|
|
manifold_strength_start=config.manifold_strength_start,
|
|
manifold_strength_end=config.manifold_strength_end,
|
|
manifold_schedule=config.manifold_schedule,
|
|
manifold_schedule_power=config.manifold_schedule_power,
|
|
manifold_schedule_hold=config.manifold_schedule_hold,
|
|
manifold_cutoff=config.manifold_cutoff,
|
|
manifold_radial_strength=config.manifold_radial_strength,
|
|
manifold_anisotropy=config.manifold_anisotropy,
|
|
manifold_translation_strength=config.manifold_translation_strength,
|
|
manifold_anchor_mix=config.manifold_anchor_mix,
|
|
manifold_mean_anchor_mix=config.manifold_mean_anchor_mix,
|
|
manifold_contrast_restore=config.manifold_contrast_restore,
|
|
manifold_energy_tether=config.manifold_energy_tether,
|
|
manifold_channel_tether=config.manifold_channel_tether,
|
|
manifold_energy_gain_cap=config.manifold_energy_gain_cap,
|
|
manifold_max_shift_px=config.manifold_max_shift_px,
|
|
manifold_spatial_mask=config.manifold_spatial_mask,
|
|
)
|
|
|
|
def __getattr__(self, name: str):
|
|
"""
|
|
Delegate attribute lookup to the wrapped base guider.
|
|
|
|
Parameters:
|
|
name (str): The attribute name to retrieve from the wrapped guider.
|
|
|
|
Returns:
|
|
Any: The value of the requested attribute from the wrapped base guider.
|
|
"""
|
|
return getattr(self._base_guider, name)
|
|
|
|
def __call__(self, x, timestep, model_options=None, seed=None):
|
|
"""
|
|
Run the wrapped guider to produce a noise prediction then apply scale-lock adjustments.
|
|
|
|
Parameters:
|
|
x (torch.Tensor): Input latent tensor for guidance.
|
|
timestep: Timestep value passed to the guider (e.g., scheduler timestep).
|
|
model_options (dict, optional): Additional model options forwarded to the base guider.
|
|
seed (int | None, optional): RNG seed forwarded to the base guider.
|
|
|
|
Returns:
|
|
torch.Tensor: Noise prediction after applying scale-lock modifications.
|
|
"""
|
|
if model_options is None:
|
|
model_options = {}
|
|
base_noise = self._base_guider(x, timestep, model_options=model_options, seed=seed)
|
|
return apply_scale_lock_to_noise_prediction(self, base_noise, x, timestep)
|
|
|
|
|
|
def _sync_impact_guider_conditions(guider, positive, negative) -> None:
|
|
"""
|
|
Synchronizes positive and negative conditioning on a guider object.
|
|
|
|
Attempts to apply the provided positive and negative conditioning to the guider.
|
|
If the guider exposes set_conds, that method is called with (positive, negative).
|
|
If it exposes inner_set_conds, that method is called with {"positive": positive, "negative": negative}.
|
|
If neither method is present, a warning is logged. Failures during the call are caught and logged as warnings.
|
|
|
|
Parameters:
|
|
guider (object): The guider instance to update; may implement `set_conds` or `inner_set_conds`.
|
|
positive: The positive conditioning to apply (type depends on guider implementation).
|
|
negative: The negative conditioning to apply (type depends on guider implementation).
|
|
"""
|
|
try:
|
|
if hasattr(guider, "set_conds"):
|
|
guider.set_conds(positive, negative)
|
|
elif hasattr(guider, "inner_set_conds"):
|
|
guider.inner_set_conds({"positive": positive, "negative": negative})
|
|
else:
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: guider %s does not expose set_conds/inner_set_conds.",
|
|
type(guider).__name__,
|
|
)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: failed to sync guider conditions on %s: %r",
|
|
type(guider).__name__,
|
|
exc,
|
|
)
|
|
|
|
|
|
def _sync_impact_guider_cfg(guider, cfg: float) -> None:
|
|
try:
|
|
if hasattr(guider, "set_scales"):
|
|
w_ag = getattr(guider, "w_ag", None)
|
|
if w_ag is None:
|
|
w_ag = getattr(guider, "w_autoguide", 2.0)
|
|
guider.set_scales(cfg=float(cfg), w_ag=float(w_ag))
|
|
elif hasattr(guider, "set_cfg"):
|
|
guider.set_cfg(float(cfg))
|
|
else:
|
|
guider.cfg = float(cfg)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: failed to sync guider cfg/scales on %s: %r",
|
|
type(guider).__name__,
|
|
exc,
|
|
)
|
|
|
|
|
|
def _apply_impact_model_options(guider, extra_args: dict[str, Any] | None) -> None:
|
|
if not isinstance(extra_args, dict):
|
|
return
|
|
model_options = extra_args.get("model_options", None)
|
|
if not isinstance(model_options, dict):
|
|
return
|
|
try:
|
|
guider.model_options = dict(model_options)
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: failed to copy model_options onto guider %s: %r",
|
|
type(guider).__name__,
|
|
exc,
|
|
)
|
|
|
|
|
|
def _build_impact_effective_guider(
|
|
request: "_ImpactSampleRequest",
|
|
model_wrap,
|
|
extra_args: dict[str, Any] | None,
|
|
guider_template=None,
|
|
):
|
|
if guider_template is None:
|
|
if _is_impact_guider_like(model_wrap):
|
|
guider_template = model_wrap
|
|
else:
|
|
guider_template = _resolve_impact_ag_guider_template(request, model_wrap)
|
|
|
|
if guider_template is None:
|
|
guider = create_cfg_guider(request.model, request.positive, request.negative, float(request.cfg))
|
|
else:
|
|
guider = clone_guider_for_scale_lock(guider_template)
|
|
_sync_impact_guider_conditions(guider, request.positive, request.negative)
|
|
_sync_impact_guider_cfg(guider, float(request.cfg))
|
|
|
|
_apply_impact_model_options(guider, extra_args)
|
|
|
|
return guider
|
|
|
|
|
|
class _ScaleLockedImpactSampler:
|
|
def __init__(self, hook: "_ScaleLockedDetailerHook"):
|
|
self._hook = hook
|
|
|
|
def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
|
|
"""
|
|
Dispatches a scale-locked sampling run through the underlying sampler using a guider proxy and device-aware tensors.
|
|
|
|
Parameters:
|
|
model_wrap: Model wrapper or model object used to resolve device information and guider behavior.
|
|
sigmas (torch.Tensor): Noise schedule tensor for the sampler; will be moved to the target device.
|
|
extra_args (dict | None): Additional sampler/model options; may be normalized to include device-sharded sigmas.
|
|
callback: Progress/callback object passed through to the underlying sampler.
|
|
noise (torch.Tensor): Initial noise tensor for sampling; may be replaced with prepared high-resolution noise if shapes match.
|
|
latent_image (dict | torch.Tensor | None): Optional high-resolution latent to sample from; if a dict is provided its "samples" entry is used.
|
|
denoise_mask (torch.Tensor | None): Optional denoising mask to pass to the sampler.
|
|
disable_pbar (bool): If true, disables progress bar propagation to the underlying sampler.
|
|
|
|
Raises:
|
|
RuntimeError: If no pending impact sampling request was captured before this sampler was invoked.
|
|
|
|
Returns:
|
|
The value returned by the underlying sampler's sample(...) call.
|
|
"""
|
|
request = self._hook._pending_request
|
|
if request is None:
|
|
raise RuntimeError("ScaleLockedDetailerHook: sampler request was not captured before the custom sampler was used.")
|
|
try:
|
|
self._hook._sampler_ran = True
|
|
|
|
planner_latent = clone_latent(request.latent)
|
|
if latent_image is not None:
|
|
planner_latent["samples"] = latent_image
|
|
|
|
planner_mask = _canonicalize_impact_noise_mask(planner_latent.get("noise_mask", None))
|
|
if planner_mask is None:
|
|
planner_mask = _canonicalize_impact_noise_mask(denoise_mask)
|
|
if planner_mask is not None:
|
|
planner_latent["noise_mask"] = planner_mask
|
|
|
|
base_sampler = comfy.samplers.sampler_object(request.sampler_name)
|
|
state = self._hook._prepare_runtime_state_for_sampler(
|
|
request=request,
|
|
model_wrap=model_wrap,
|
|
sampler=base_sampler,
|
|
sigmas=sigmas,
|
|
extra_args=extra_args,
|
|
live_latent=planner_latent,
|
|
)
|
|
|
|
proxy = _ScaleLockedGuiderProxy(model_wrap, state.runtime, state.config)
|
|
|
|
sampling_noise = noise
|
|
fallback_device = latent_image.device if latent_image is not None else noise.device
|
|
target_device = _resolve_sampler_device(model_wrap, fallback_device)
|
|
target_dtype = latent_image.dtype if latent_image is not None else noise.dtype
|
|
|
|
prepared_noise = state.runtime.highres_noise
|
|
if tuple(prepared_noise.shape) == tuple(noise.shape):
|
|
sampling_noise = prepared_noise.to(
|
|
device=target_device,
|
|
dtype=target_dtype,
|
|
non_blocking=True,
|
|
).clone()
|
|
else:
|
|
sampling_noise = sampling_noise.to(
|
|
device=target_device,
|
|
dtype=target_dtype,
|
|
non_blocking=True,
|
|
)
|
|
|
|
sampling_latent = latent_image if latent_image is not None else state.runtime.highres_latent["samples"]
|
|
sampling_latent = sampling_latent.to(
|
|
device=target_device,
|
|
dtype=target_dtype,
|
|
non_blocking=True,
|
|
)
|
|
|
|
sampling_mask = denoise_mask
|
|
sigmas = sigmas.to(device=target_device, non_blocking=True)
|
|
if isinstance(sampling_mask, torch.Tensor):
|
|
sampling_mask = sampling_mask.to(device=target_device, non_blocking=True)
|
|
|
|
sample_extra_args = _normalize_sampler_extra_args(extra_args, sigmas, target_device)
|
|
return base_sampler.sample(
|
|
proxy,
|
|
sigmas,
|
|
sample_extra_args,
|
|
callback,
|
|
sampling_noise,
|
|
latent_image=sampling_latent,
|
|
denoise_mask=sampling_mask,
|
|
disable_pbar=disable_pbar,
|
|
)
|
|
finally:
|
|
self._hook._clear_sampler_state()
|
|
|
|
|
|
class _ScaleLockedDetailerHook:
|
|
def __init__(self, settings: _ImpactHookSettings):
|
|
self._settings = settings
|
|
self._pending_request: _ImpactSampleRequest | None = None
|
|
self._step_info: Any | None = None
|
|
self._upscale_mask: torch.Tensor | None = None
|
|
self._active_runtime: _ScaleLockedImpactRuntimeState | None = None
|
|
self._sampler_ran = False
|
|
self._inactive_sampler_warned = False
|
|
self._custom_sampler = _ScaleLockedImpactSampler(self)
|
|
|
|
def _clear_sampler_state(self):
|
|
self._pending_request = None
|
|
self._active_runtime = None
|
|
self._sampler_ran = False
|
|
self._inactive_sampler_warned = False
|
|
|
|
def _clear_cycle_state(self):
|
|
self._clear_sampler_state()
|
|
self._upscale_mask = None
|
|
|
|
def _warn_if_sampler_inactive(self, stage: str):
|
|
if self._pending_request is None or self._sampler_ran or self._inactive_sampler_warned:
|
|
return
|
|
if not self._settings.warn_if_inactive_sampler:
|
|
return
|
|
self._inactive_sampler_warned = True
|
|
logger.warning(
|
|
"ScaleLockedDetailerHook: custom sampler was not selected for this cycle before %s. "
|
|
"An earlier hook in the Impact chain provided the active sampler, so Scale-Locked sampler logic stayed inert.",
|
|
stage,
|
|
)
|
|
|
|
def _effective_config_for_samples(self, samples: torch.Tensor) -> ScaleLockConfig:
|
|
cfg = self._settings.config
|
|
|
|
face_mask = _expand_mask_for_like(self._upscale_mask, samples)
|
|
lock_mask = _combine_masks(face_mask, _expand_mask_for_like(cfg.spatial_mask, samples))
|
|
manifold_source = cfg.manifold_spatial_mask if cfg.manifold_spatial_mask is not None else cfg.spatial_mask
|
|
manifold_mask = _combine_masks(face_mask, _expand_mask_for_like(manifold_source, samples))
|
|
return replace(cfg, spatial_mask=lock_mask, manifold_spatial_mask=manifold_mask)
|
|
|
|
def _prepare_runtime_state_for_sampler(
|
|
self,
|
|
request: _ImpactSampleRequest,
|
|
model_wrap,
|
|
sampler,
|
|
sigmas,
|
|
extra_args: dict[str, Any] | None,
|
|
live_latent: dict[str, Any] | None = None,
|
|
) -> _ScaleLockedImpactRuntimeState:
|
|
"""
|
|
Builds and registers a scale-locked impact runtime state for the given impact sampling request.
|
|
|
|
Parameters:
|
|
request (_ImpactSampleRequest): Canonicalized impact sampling request containing model, seed, sampler name, and latent.
|
|
model_wrap: Model or model wrapper used to construct the planner guider proxy and to resolve runtime device/dtype.
|
|
sampler: Sampler instance to be used for planning and runtime creation.
|
|
sigmas (torch.Tensor): Noise schedule tensor to use for runtime construction.
|
|
extra_args (dict | None): Optional extra model arguments forwarded into the planner guider proxy (e.g., model_options).
|
|
live_latent (dict | None): Optional override latent to use for planning instead of request.latent.
|
|
|
|
Returns:
|
|
_ScaleLockedImpactRuntimeState: The created runtime state containing the original request, the prepared runtime context, and the effective ScaleLockConfig computed for the planner latent.
|
|
"""
|
|
guard_sampler_alignment(request.sampler_name, self._settings.sampler_guard)
|
|
planner_noise = _ScaleLockedPlannerNoise(request.seed, disable_noise=not self._settings.add_noise)
|
|
planner_latent = request.latent if live_latent is None else live_latent
|
|
planner_guider = _ScaleLockedPlannerGuiderProxy(model_wrap, extra_args)
|
|
runtime = build_runtime_context_from_advanced(
|
|
noise=planner_noise,
|
|
guider=planner_guider,
|
|
sampler=sampler,
|
|
sigmas=sigmas,
|
|
latent_image=planner_latent,
|
|
target_megapixels=self._settings.target_megapixels,
|
|
nested_noise_strength=self._settings.nested_noise_strength,
|
|
pin_anchors=self._settings.pin_anchors,
|
|
)
|
|
state = _ScaleLockedImpactRuntimeState(
|
|
request=request,
|
|
runtime=runtime,
|
|
config=self._effective_config_for_samples(runtime.highres_latent["samples"]),
|
|
)
|
|
self._active_runtime = state
|
|
return state
|
|
|
|
def get_skip_sampling(self):
|
|
return False
|
|
|
|
def set_steps(self, info):
|
|
self._step_info = info
|
|
|
|
def post_crop_region(self, w, h, item_bbox, crop_region):
|
|
del w, h, item_bbox
|
|
return crop_region
|
|
|
|
def post_detection(self, segs):
|
|
return segs
|
|
|
|
def touch_scaled_size(self, w, h):
|
|
return w, h
|
|
|
|
def post_upscale(self, pixels, mask=None):
|
|
self._clear_sampler_state()
|
|
self._upscale_mask = None if mask is None else mask.detach().clone()
|
|
return pixels
|
|
|
|
def post_encode(self, samples):
|
|
return samples
|
|
|
|
def pre_decode(self, samples):
|
|
self._warn_if_sampler_inactive("pre_decode")
|
|
self._clear_sampler_state()
|
|
return samples
|
|
|
|
def post_decode(self, pixels):
|
|
self._warn_if_sampler_inactive("post_decode")
|
|
self._clear_sampler_state()
|
|
return pixels
|
|
|
|
def cycle_latent(self, latent):
|
|
return latent
|
|
|
|
def post_paste(self, image):
|
|
self._warn_if_sampler_inactive("post_paste")
|
|
self._clear_cycle_state()
|
|
return image
|
|
|
|
def get_custom_noise(self, seed, noise, is_touched):
|
|
del seed
|
|
return noise, is_touched
|
|
|
|
def should_retry_patch(self, image):
|
|
del image
|
|
return False
|
|
|
|
def get_custom_sampler(self, *args, **kwargs):
|
|
del args, kwargs
|
|
return self._custom_sampler
|
|
|
|
def get_custom_sampler_provider(self, *args, **kwargs):
|
|
return self.get_custom_sampler(*args, **kwargs)
|
|
|
|
def get_custom_ksampler_provider(self, *args, **kwargs):
|
|
return self.get_custom_sampler(*args, **kwargs)
|
|
|
|
def pre_ksample(self, *args, **kwargs):
|
|
self._clear_sampler_state()
|
|
request = self._remember_request(args, kwargs, strict=False)
|
|
if request is None:
|
|
return _impact_request_tuple_from_kwargs(kwargs) if kwargs else args
|
|
|
|
return _impact_request_tuple_from_kwargs(kwargs) if kwargs else args
|
|
|
|
def post_ksample(self, *args, **kwargs):
|
|
self._clear_cycle_state()
|
|
if kwargs:
|
|
return kwargs
|
|
if len(args) == 1:
|
|
return args[0]
|
|
return args
|
|
|
|
def _remember_request(self, args, kwargs, strict: bool) -> _ImpactSampleRequest | None:
|
|
try:
|
|
request = _coerce_impact_request(args, kwargs, fallback=self._pending_request)
|
|
except TypeError:
|
|
if strict:
|
|
raise
|
|
request = self._pending_request
|
|
if request is not None:
|
|
self._pending_request = request
|
|
return request
|
|
|
|
|
|
class ScaleLockedDetailerHookProvider:
|
|
CATEGORY = "sampling/scale_locked/impact"
|
|
RETURN_TYPES = ("DETAILER_HOOK",)
|
|
RETURN_NAMES = ("detailer_hook",)
|
|
FUNCTION = "build"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
del cls, args, kwargs
|
|
# Force a fresh hook object for each prompt execution.
|
|
return time.monotonic_ns()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"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}),
|
|
"manifold_enabled": ("BOOLEAN", {"default": False}),
|
|
"manifold_strength": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_strength_end": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule": (LOCK_SCHEDULE_OPTIONS,),
|
|
"manifold_schedule_hold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 0.95, "step": 0.01, "round": 0.001}),
|
|
"manifold_schedule_power": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 8.0, "step": 0.1, "round": 0.01}),
|
|
"manifold_cutoff": ("FLOAT", {"default": 0.18, "min": 0.05, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_radial_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anisotropy": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 2.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_translation_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_anchor_mix": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_mean_anchor_mix": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_contrast_restore": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_channel_tether": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
|
"manifold_energy_gain_cap": ("FLOAT", {"default": 1.75, "min": 1.0, "max": 4.0, "step": 0.05, "round": 0.01}),
|
|
"manifold_max_shift_px": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 64.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}),
|
|
"sampler_guard": (["warn", "error", "off"],),
|
|
"warn_if_inactive_sampler": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"lock_mask": ("MASK",),
|
|
"manifold_mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
def build(
|
|
self,
|
|
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,
|
|
manifold_enabled,
|
|
manifold_strength,
|
|
manifold_strength_start,
|
|
manifold_strength_end,
|
|
manifold_schedule,
|
|
manifold_schedule_hold,
|
|
manifold_schedule_power,
|
|
manifold_cutoff,
|
|
manifold_radial_strength,
|
|
manifold_anisotropy,
|
|
manifold_translation_strength,
|
|
manifold_anchor_mix,
|
|
manifold_mean_anchor_mix,
|
|
manifold_contrast_restore,
|
|
manifold_energy_tether,
|
|
manifold_channel_tether,
|
|
manifold_energy_gain_cap,
|
|
manifold_max_shift_px,
|
|
nested_noise_strength,
|
|
add_noise,
|
|
pin_anchors,
|
|
sampler_guard,
|
|
warn_if_inactive_sampler,
|
|
lock_mask=None,
|
|
manifold_mask=None,
|
|
):
|
|
settings = _ImpactHookSettings(
|
|
target_megapixels=target_megapixels,
|
|
nested_noise_strength=nested_noise_strength,
|
|
add_noise=add_noise,
|
|
pin_anchors=pin_anchors,
|
|
sampler_guard=sampler_guard,
|
|
warn_if_inactive_sampler=warn_if_inactive_sampler,
|
|
config=_scale_lock_config(
|
|
lock_strength=lock_strength,
|
|
lock_strength_start=lock_strength_start,
|
|
lock_strength_end=lock_strength_end,
|
|
coarse_cutoff=coarse_cutoff,
|
|
mid_band_cutoff=mid_band_cutoff,
|
|
mid_band_strength=mid_band_strength,
|
|
lock_schedule=lock_schedule,
|
|
lock_schedule_hold=lock_schedule_hold,
|
|
lock_schedule_power=lock_schedule_power,
|
|
mid_band_strength_start=mid_band_strength_start,
|
|
mid_band_strength_end=mid_band_strength_end,
|
|
mid_band_schedule=mid_band_schedule,
|
|
mid_band_schedule_hold=mid_band_schedule_hold,
|
|
mid_band_schedule_power=mid_band_schedule_power,
|
|
manifold_enabled=manifold_enabled,
|
|
manifold_strength=manifold_strength,
|
|
manifold_strength_start=manifold_strength_start,
|
|
manifold_strength_end=manifold_strength_end,
|
|
manifold_schedule=manifold_schedule,
|
|
manifold_schedule_hold=manifold_schedule_hold,
|
|
manifold_schedule_power=manifold_schedule_power,
|
|
manifold_cutoff=manifold_cutoff,
|
|
manifold_radial_strength=manifold_radial_strength,
|
|
manifold_anisotropy=manifold_anisotropy,
|
|
manifold_translation_strength=manifold_translation_strength,
|
|
manifold_anchor_mix=manifold_anchor_mix,
|
|
manifold_mean_anchor_mix=manifold_mean_anchor_mix,
|
|
manifold_contrast_restore=manifold_contrast_restore,
|
|
manifold_energy_tether=manifold_energy_tether,
|
|
manifold_channel_tether=manifold_channel_tether,
|
|
manifold_energy_gain_cap=manifold_energy_gain_cap,
|
|
manifold_max_shift_px=manifold_max_shift_px,
|
|
lock_mask=lock_mask,
|
|
manifold_mask=manifold_mask,
|
|
),
|
|
)
|
|
return (_ScaleLockedDetailerHook(settings),)
|
|
|
|
|
|
def _scale_lock_config(
|
|
*,
|
|
lock_strength,
|
|
lock_strength_start,
|
|
lock_strength_end,
|
|
coarse_cutoff,
|
|
mid_band_cutoff,
|
|
mid_band_strength,
|
|
lock_schedule,
|
|
lock_schedule_hold,
|
|
lock_schedule_power,
|
|
mid_band_strength_start,
|
|
mid_band_strength_end,
|
|
mid_band_schedule,
|
|
mid_band_schedule_hold,
|
|
mid_band_schedule_power,
|
|
lock_mask,
|
|
manifold_enabled=False,
|
|
manifold_strength=0.0,
|
|
manifold_strength_start=1.0,
|
|
manifold_strength_end=0.0,
|
|
manifold_schedule="ease_out",
|
|
manifold_schedule_hold=0.0,
|
|
manifold_schedule_power=2.0,
|
|
manifold_cutoff=0.18,
|
|
manifold_radial_strength=1.0,
|
|
manifold_anisotropy=0.15,
|
|
manifold_translation_strength=1.0,
|
|
manifold_anchor_mix=0.18,
|
|
manifold_mean_anchor_mix=0.12,
|
|
manifold_contrast_restore=0.10,
|
|
manifold_energy_tether=0.0,
|
|
manifold_channel_tether=0.0,
|
|
manifold_energy_gain_cap=1.75,
|
|
manifold_max_shift_px=3.0,
|
|
manifold_mask=None,
|
|
):
|
|
return ScaleLockConfig(
|
|
lock_strength=float(lock_strength),
|
|
lock_strength_start=float(lock_strength_start),
|
|
lock_strength_end=float(lock_strength_end),
|
|
cutoff=float(coarse_cutoff),
|
|
mid_cutoff=float(mid_band_cutoff),
|
|
mid_strength=float(mid_band_strength),
|
|
schedule=lock_schedule,
|
|
schedule_power=float(lock_schedule_power),
|
|
schedule_hold=float(lock_schedule_hold),
|
|
mid_strength_start=float(mid_band_strength_start),
|
|
mid_strength_end=float(mid_band_strength_end),
|
|
mid_schedule=mid_band_schedule,
|
|
mid_schedule_power=float(mid_band_schedule_power),
|
|
mid_schedule_hold=float(mid_band_schedule_hold),
|
|
spatial_mask=lock_mask,
|
|
manifold_enabled=bool(manifold_enabled),
|
|
manifold_strength=float(manifold_strength),
|
|
manifold_strength_start=float(manifold_strength_start),
|
|
manifold_strength_end=float(manifold_strength_end),
|
|
manifold_schedule=manifold_schedule,
|
|
manifold_schedule_power=float(manifold_schedule_power),
|
|
manifold_schedule_hold=float(manifold_schedule_hold),
|
|
manifold_cutoff=float(manifold_cutoff),
|
|
manifold_radial_strength=float(manifold_radial_strength),
|
|
manifold_anisotropy=float(manifold_anisotropy),
|
|
manifold_translation_strength=float(manifold_translation_strength),
|
|
manifold_anchor_mix=float(manifold_anchor_mix),
|
|
manifold_mean_anchor_mix=float(manifold_mean_anchor_mix),
|
|
manifold_contrast_restore=float(manifold_contrast_restore),
|
|
manifold_energy_tether=float(manifold_energy_tether),
|
|
manifold_channel_tether=float(manifold_channel_tether),
|
|
manifold_energy_gain_cap=float(manifold_energy_gain_cap),
|
|
manifold_max_shift_px=float(manifold_max_shift_px),
|
|
manifold_spatial_mask=manifold_mask,
|
|
)
|
|
|
|
|
|
def _coerce_impact_request(args, kwargs, fallback=None) -> _ImpactSampleRequest:
|
|
data = {}
|
|
if fallback is not None:
|
|
data.update(fallback.__dict__)
|
|
|
|
if len(args) >= 10:
|
|
names = ("model", "seed", "steps", "cfg", "sampler_name", "scheduler", "positive", "negative", "latent", "denoise")
|
|
for name, value in zip(names, args[:10]):
|
|
data[name] = value
|
|
|
|
for field, aliases in _IMPACT_FIELD_ALIASES.items():
|
|
for alias in aliases:
|
|
if alias in kwargs:
|
|
data[field] = kwargs[alias]
|
|
break
|
|
|
|
required = ("model", "seed", "steps", "cfg", "sampler_name", "scheduler", "positive", "negative", "latent", "denoise")
|
|
missing = [field for field in required if field not in data or data[field] is None]
|
|
if missing:
|
|
raise TypeError(f"Unable to resolve Impact detailer sample request fields: {', '.join(missing)}")
|
|
|
|
latent = data["latent"]
|
|
if not isinstance(latent, dict) or "samples" not in latent:
|
|
raise TypeError("Impact detailer sample request did not provide a LATENT dict with a 'samples' entry.")
|
|
|
|
return _ImpactSampleRequest(
|
|
model=data["model"],
|
|
seed=int(data["seed"]),
|
|
steps=int(data["steps"]),
|
|
cfg=float(data["cfg"]),
|
|
sampler_name=str(data["sampler_name"]),
|
|
scheduler=str(data["scheduler"]),
|
|
positive=data["positive"],
|
|
negative=data["negative"],
|
|
latent=_snapshot_impact_latent(latent),
|
|
denoise=float(data["denoise"]),
|
|
)
|
|
|
|
|
|
def _snapshot_impact_latent(latent: dict[str, Any]) -> dict[str, Any]:
|
|
snapshot = clone_latent(latent)
|
|
for key, value in snapshot.items():
|
|
if isinstance(value, torch.Tensor):
|
|
snapshot[key] = value.detach().clone()
|
|
elif isinstance(value, list):
|
|
snapshot[key] = list(value)
|
|
elif isinstance(value, dict):
|
|
snapshot[key] = dict(value)
|
|
return snapshot
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ScaleLockedResidualKSampler": ScaleLockedResidualKSampler,
|
|
"ScaleLockedRuntimeContextBuilder": ScaleLockedRuntimeContextBuilder,
|
|
"ScaleLockedCFGGuider": ScaleLockedCFGGuiderNode,
|
|
"ScaleLockedResidualSamplerCustomAdvanced": ScaleLockedResidualSamplerCustomAdvanced,
|
|
"ScaleLockedNestedNoisePreview": ScaleLockedNestedNoisePreview,
|
|
"ScaleLockedDetailerHookProvider": ScaleLockedDetailerHookProvider,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ScaleLockedResidualKSampler": "Scale-Locked Residual KSampler",
|
|
"ScaleLockedRuntimeContextBuilder": "Scale-Locked Runtime Context",
|
|
"ScaleLockedCFGGuider": "Scale-Locked CFG Guider",
|
|
"ScaleLockedResidualSamplerCustomAdvanced": "Scale-Locked Residual SamplerCustomAdvanced",
|
|
"ScaleLockedNestedNoisePreview": "Scale-Locked Nested Noise Preview",
|
|
"ScaleLockedDetailerHookProvider": "Scale-Locked Detailer Hook Provider",
|
|
}
|
|
|