1114 lines
34 KiB
Python
1114 lines
34 KiB
Python
from __future__ import annotations
|
|
|
|
import abc
|
|
import inspect
|
|
from types import SimpleNamespace
|
|
from typing import Any, Callable
|
|
|
|
import torch
|
|
from comfy import samplers
|
|
|
|
from . import external, noise
|
|
from .noise import NoiseType
|
|
from .noise_generation import scale_noise
|
|
from .sonar import (
|
|
GuidanceConfig,
|
|
GuidanceType,
|
|
HistoryType,
|
|
SonarConfig,
|
|
SonarDPMPPSDE,
|
|
SonarEuler,
|
|
SonarEulerAncestral,
|
|
)
|
|
|
|
|
|
class NoisyLatentLikeNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"noise_type": (tuple(NoiseType.get_names(skip=(NoiseType.BROWNIAN,))),),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"latent": ("LATENT",),
|
|
"multiplier": ("FLOAT", {"default": 1.0}),
|
|
"add_to_latent": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
"mul_by_sigmas_opt": ("SIGMAS",),
|
|
"model_opt": ("MODEL",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
CATEGORY = "latent/noise"
|
|
|
|
FUNCTION = "go"
|
|
|
|
def go(
|
|
self,
|
|
noise_type: str,
|
|
seed: None | int,
|
|
latent: dict,
|
|
multiplier: float = 1.0,
|
|
add_to_latent=False,
|
|
custom_noise_opt: object | None = None,
|
|
mul_by_sigmas_opt: None | torch.Tensor = None,
|
|
model_opt: object | None = None,
|
|
):
|
|
model, sigmas = model_opt, mul_by_sigmas_opt
|
|
if sigmas is not None and len(sigmas) > 0:
|
|
if model is None:
|
|
raise ValueError(
|
|
"NoisyLatentLike requires a model when sigmas are connected!",
|
|
)
|
|
while hasattr(model, "model"):
|
|
model = model.model
|
|
latent_scale_factor = model.latent_format.scale_factor
|
|
max_denoise = samplers.Sampler().max_denoise(
|
|
SimpleNamespace(inner_model=model),
|
|
sigmas,
|
|
)
|
|
multiplier *= (
|
|
float(
|
|
torch.sqrt(1.0 + sigmas[0] ** 2.0) if max_denoise else sigmas[0],
|
|
)
|
|
/ latent_scale_factor
|
|
)
|
|
if sigmas is not None and sigmas.numel() > 1:
|
|
sigma_min, sigma_max = sigmas[0], sigmas[-1]
|
|
sigma, sigma_next = sigmas[0], sigmas[1]
|
|
else:
|
|
sigma_min, sigma_max, sigma, sigma_next = (None,) * 4
|
|
latent_samples = latent["samples"]
|
|
if custom_noise_opt is not None:
|
|
ns = custom_noise_opt.make_noise_sampler(
|
|
latent_samples,
|
|
sigma_min=sigma_min,
|
|
sigma_max=sigma_max,
|
|
)
|
|
else:
|
|
ns = noise.get_noise_sampler(
|
|
NoiseType[noise_type.upper()],
|
|
latent_samples,
|
|
sigma_min,
|
|
sigma_max,
|
|
seed=seed,
|
|
cpu=True,
|
|
)
|
|
randst = torch.random.get_rng_state()
|
|
try:
|
|
torch.random.manual_seed(seed)
|
|
result = ns(sigma, sigma_next)
|
|
finally:
|
|
torch.random.set_rng_state(randst)
|
|
result = scale_noise(result, multiplier, normalized=True)
|
|
if add_to_latent:
|
|
result += latent_samples.to(result)
|
|
return ({"samples": result},)
|
|
|
|
|
|
class SonarCustomNoiseNodeBase(abc.ABC):
|
|
RETURN_TYPES = ("SONAR_CUSTOM_NOISE",)
|
|
CATEGORY = "advanced/noise"
|
|
FUNCTION = "go"
|
|
|
|
@abc.abstractmethod
|
|
def get_item_class(self):
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls, *, include_rescale=True, include_chain=True):
|
|
result = {
|
|
"required": {
|
|
"factor": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": -100.0,
|
|
"max": 100.0,
|
|
"step": 0.001,
|
|
"round": False,
|
|
},
|
|
),
|
|
},
|
|
"optional": {},
|
|
}
|
|
if include_rescale:
|
|
result["required"] |= {
|
|
"rescale": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.001,
|
|
"round": False,
|
|
},
|
|
),
|
|
}
|
|
if include_chain:
|
|
result["optional"] |= {
|
|
"sonar_custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
}
|
|
return result
|
|
|
|
def go(
|
|
self,
|
|
factor=1.0,
|
|
rescale=0.0,
|
|
sonar_custom_noise_opt=None,
|
|
**kwargs: dict[str, Any],
|
|
):
|
|
nis = (
|
|
sonar_custom_noise_opt.clone()
|
|
if sonar_custom_noise_opt
|
|
else noise.CustomNoiseChain()
|
|
)
|
|
if factor != 0:
|
|
nis.add(self.get_item_class()(factor, **kwargs))
|
|
return (nis if rescale == 0 else nis.rescaled(rescale),)
|
|
|
|
|
|
class SonarCustomNoiseNode(SonarCustomNoiseNodeBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES()
|
|
result["required"] |= {
|
|
"noise_type": (tuple(NoiseType.get_names()),),
|
|
}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.CustomNoiseItem
|
|
|
|
|
|
class SonarNormalizeNoiseNodeMixin:
|
|
@staticmethod
|
|
def get_normalize(val: str) -> None | bool:
|
|
return None if val == "default" else val == "forced"
|
|
|
|
|
|
class SonarModulatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
|
|
"modulation_type": (
|
|
(
|
|
"intensity",
|
|
"frequency",
|
|
"spectral_signum",
|
|
"none",
|
|
),
|
|
),
|
|
"dims": ("INT", {"default": 3, "min": 1, "max": 3}),
|
|
"strength": ("FLOAT", {"default": 2.0, "min": -100.0, "max": 100.0}),
|
|
"normalize_result": (("default", "forced", "disabled"),),
|
|
"normalize_noise": (("default", "forced", "disabled"),),
|
|
"normalize_ref": (
|
|
"BOOLEAN",
|
|
{"default": True},
|
|
),
|
|
}
|
|
result["optional"] |= {"ref_latent_opt": ("LATENT",)}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.ModulatedNoise
|
|
|
|
def go(
|
|
self,
|
|
factor,
|
|
sonar_custom_noise,
|
|
modulation_type,
|
|
dims,
|
|
strength,
|
|
normalize_result,
|
|
normalize_noise,
|
|
normalize_ref,
|
|
ref_latent_opt=None,
|
|
):
|
|
if ref_latent_opt is not None:
|
|
ref_latent_opt = ref_latent_opt["samples"].clone()
|
|
return super().go(
|
|
factor,
|
|
noise=sonar_custom_noise,
|
|
modulation_type=modulation_type,
|
|
modulation_dims=dims,
|
|
modulation_strength=strength,
|
|
normalize_result=self.get_normalize(normalize_result),
|
|
normalize_noise=self.get_normalize(normalize_noise),
|
|
normalize_ref=self.get_normalize(normalize_ref),
|
|
ref_latent_opt=ref_latent_opt,
|
|
)
|
|
|
|
|
|
class SonarRepeatedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
|
|
"repeat_length": ("INT", {"default": 8, "min": 1, "max": 100}),
|
|
"max_recycle": ("INT", {"default": 1000, "min": 1, "max": 1000}),
|
|
"normalize": (("default", "forced", "disabled"),),
|
|
"permute": (("enabled", "disabled", "always"),),
|
|
}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.RepeatedNoise
|
|
|
|
def go(
|
|
self,
|
|
factor,
|
|
sonar_custom_noise,
|
|
repeat_length,
|
|
max_recycle,
|
|
normalize,
|
|
permute=True,
|
|
):
|
|
return super().go(
|
|
factor,
|
|
noise=sonar_custom_noise,
|
|
repeat_length=repeat_length,
|
|
max_recycle=max_recycle,
|
|
normalize=self.get_normalize(normalize),
|
|
permute=permute,
|
|
)
|
|
|
|
|
|
class SonarScheduledNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"model": ("MODEL",),
|
|
"sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
|
|
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
|
|
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
|
|
"normalize": (("default", "forced", "disabled"),),
|
|
}
|
|
result["optional"] |= {"fallback_sonar_custom_noise": ("SONAR_CUSTOM_NOISE",)}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.ScheduledNoise
|
|
|
|
def go(
|
|
self,
|
|
model,
|
|
factor,
|
|
sonar_custom_noise,
|
|
start_percent,
|
|
end_percent,
|
|
normalize,
|
|
fallback_sonar_custom_noise=None,
|
|
):
|
|
ms = model.get_model_object("model_sampling")
|
|
start_sigma = ms.percent_to_sigma(start_percent)
|
|
end_sigma = ms.percent_to_sigma(end_percent)
|
|
return super().go(
|
|
factor,
|
|
noise=sonar_custom_noise,
|
|
start_sigma=start_sigma,
|
|
end_sigma=end_sigma,
|
|
normalize=self.get_normalize(normalize),
|
|
fallback_noise=fallback_sonar_custom_noise,
|
|
)
|
|
|
|
|
|
class SonarCompositeNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"sonar_custom_noise_dst": ("SONAR_CUSTOM_NOISE",),
|
|
"sonar_custom_noise_src": ("SONAR_CUSTOM_NOISE",),
|
|
"normalize_dst": (("default", "forced", "disabled"),),
|
|
"normalize_src": (("default", "forced", "disabled"),),
|
|
"normalize_result": (("default", "forced", "disabled"),),
|
|
"mask": ("MASK",),
|
|
}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.CompositeNoise
|
|
|
|
def go(
|
|
self,
|
|
factor,
|
|
sonar_custom_noise_dst,
|
|
sonar_custom_noise_src,
|
|
normalize_src,
|
|
normalize_dst,
|
|
normalize_result,
|
|
mask,
|
|
):
|
|
return super().go(
|
|
factor,
|
|
dst_noise=sonar_custom_noise_dst,
|
|
src_noise=sonar_custom_noise_src,
|
|
normalize_dst=self.get_normalize(normalize_src),
|
|
normalize_src=self.get_normalize(normalize_dst),
|
|
normalize_result=self.get_normalize(normalize_result),
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
class SonarGuidedNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"latent": ("LATENT",),
|
|
"sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
|
|
"method": (("euler", "linear"),),
|
|
"guidance_factor": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.0125,
|
|
"min": -100.0,
|
|
"max": 100.0,
|
|
"step": 0.001,
|
|
"round": False,
|
|
},
|
|
),
|
|
"normalize_noise": (("default", "forced", "disabled"),),
|
|
"normalize_result": (("default", "forced", "disabled"),),
|
|
"normalize_ref": (
|
|
"BOOLEAN",
|
|
{"default": True},
|
|
),
|
|
}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.GuidedNoise
|
|
|
|
def go(
|
|
self,
|
|
factor,
|
|
latent,
|
|
sonar_custom_noise,
|
|
normalize_noise,
|
|
normalize_result,
|
|
normalize_ref=True,
|
|
method="euler",
|
|
guidance_factor=0.5,
|
|
):
|
|
from .sonar import SonarGuidanceMixin
|
|
|
|
return super().go(
|
|
factor,
|
|
ref_latent=scale_noise(
|
|
SonarGuidanceMixin.prepare_ref_latent(latent["samples"].clone()),
|
|
normalized=normalize_ref,
|
|
),
|
|
guidance_factor=guidance_factor,
|
|
noise=sonar_custom_noise.clone(),
|
|
method=method,
|
|
normalize_noise=self.get_normalize(normalize_noise),
|
|
normalize_result=self.get_normalize(normalize_result),
|
|
)
|
|
|
|
|
|
class SonarRandomNoiseNode(SonarCustomNoiseNodeBase, SonarNormalizeNoiseNodeMixin):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
|
|
"mix_count": ("INT", {"default": 1, "min": 1, "max": 100}),
|
|
"normalize": (("default", "forced", "disabled"),),
|
|
}
|
|
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.RandomNoise
|
|
|
|
def go(
|
|
self,
|
|
factor,
|
|
sonar_custom_noise,
|
|
mix_count,
|
|
normalize,
|
|
):
|
|
return super().go(
|
|
factor,
|
|
noise=sonar_custom_noise,
|
|
mix_count=mix_count,
|
|
normalize=self.get_normalize(normalize),
|
|
)
|
|
|
|
|
|
class GuidanceConfigNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"factor": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.01,
|
|
"min": -2.0,
|
|
"max": 2.0,
|
|
"step": 0.001,
|
|
"round": False,
|
|
},
|
|
),
|
|
"guidance_type": (tuple(t.name.lower() for t in GuidanceType),),
|
|
"start_step": ("INT", {"default": 1, "min": 1}),
|
|
"end_step": ("INT", {"default": 9999, "min": 1}),
|
|
"latent": ("LATENT",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SONAR_GUIDANCE_CFG",)
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
FUNCTION = "make_guidance_cfg"
|
|
|
|
def make_guidance_cfg(
|
|
self,
|
|
guidance_type,
|
|
factor,
|
|
start_step,
|
|
end_step,
|
|
latent,
|
|
):
|
|
return (
|
|
GuidanceConfig(
|
|
guidance_type=GuidanceType[guidance_type.upper()],
|
|
factor=factor,
|
|
start_step=start_step,
|
|
end_step=end_step,
|
|
latent=latent.get("samples"),
|
|
),
|
|
)
|
|
|
|
|
|
class SamplerNodeSonarBase:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"momentum": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.95,
|
|
"min": -0.5,
|
|
"max": 2.5,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"momentum_hist": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.75,
|
|
"min": -1.5,
|
|
"max": 1.5,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"momentum_init": (tuple(t.name for t in HistoryType),),
|
|
"direction": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": -30.0,
|
|
"max": 15.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"rand_init_noise_type": (
|
|
tuple(NoiseType.get_names(skip=(NoiseType.BROWNIAN,))),
|
|
),
|
|
},
|
|
"optional": {
|
|
"guidance_cfg_opt": ("SONAR_GUIDANCE_CFG",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
|
|
class SamplerNodeSonarEuler(SamplerNodeSonarBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES()
|
|
result["required"].update(
|
|
{
|
|
"s_noise": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
},
|
|
)
|
|
return result
|
|
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
FUNCTION = "get_sampler"
|
|
|
|
def get_sampler(
|
|
self,
|
|
momentum,
|
|
momentum_hist,
|
|
momentum_init,
|
|
direction,
|
|
rand_init_noise_type,
|
|
s_noise,
|
|
guidance_cfg_opt=None,
|
|
):
|
|
cfg = SonarConfig(
|
|
momentum=momentum,
|
|
init=HistoryType[momentum_init.upper()],
|
|
momentum_hist=momentum_hist,
|
|
direction=direction,
|
|
rand_init_noise_type=NoiseType[rand_init_noise_type.upper()],
|
|
guidance=guidance_cfg_opt,
|
|
)
|
|
return (
|
|
samplers.KSAMPLER(
|
|
SonarEuler.sampler,
|
|
{
|
|
"s_noise": s_noise,
|
|
"sonar_config": cfg,
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
class SamplerNodeSonarEulerAncestral(SamplerNodeSonarEuler):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES()
|
|
result["required"].update(
|
|
{
|
|
"eta": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"noise_type": (tuple(NoiseType.get_names()),),
|
|
},
|
|
)
|
|
result["optional"].update(
|
|
{
|
|
"custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
},
|
|
)
|
|
return result
|
|
|
|
def get_sampler(
|
|
self,
|
|
momentum,
|
|
momentum_hist,
|
|
momentum_init,
|
|
direction,
|
|
rand_init_noise_type,
|
|
noise_type,
|
|
eta,
|
|
s_noise,
|
|
guidance_cfg_opt=None,
|
|
custom_noise_opt=None,
|
|
):
|
|
cfg = SonarConfig(
|
|
momentum=momentum,
|
|
init=HistoryType[momentum_init.upper()],
|
|
momentum_hist=momentum_hist,
|
|
direction=direction,
|
|
rand_init_noise_type=NoiseType[rand_init_noise_type.upper()],
|
|
noise_type=NoiseType[noise_type.upper()],
|
|
custom_noise=custom_noise_opt.clone() if custom_noise_opt else None,
|
|
guidance=guidance_cfg_opt,
|
|
)
|
|
return (
|
|
samplers.KSAMPLER(
|
|
SonarEulerAncestral.sampler,
|
|
{
|
|
"sonar_config": cfg,
|
|
"eta": eta,
|
|
"s_noise": s_noise,
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
class SamplerNodeSonarDPMPPSDE(SamplerNodeSonarEuler):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES()
|
|
result["required"].update(
|
|
{
|
|
"eta": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"noise_type": (tuple(NoiseType.get_names(default=NoiseType.BROWNIAN)),),
|
|
},
|
|
)
|
|
result["optional"].update(
|
|
{
|
|
"custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
},
|
|
)
|
|
return result
|
|
|
|
def get_sampler(
|
|
self,
|
|
momentum,
|
|
momentum_hist,
|
|
momentum_init,
|
|
direction,
|
|
rand_init_noise_type,
|
|
noise_type,
|
|
eta,
|
|
s_noise,
|
|
guidance_cfg_opt=None,
|
|
custom_noise_opt=None,
|
|
):
|
|
cfg = SonarConfig(
|
|
momentum=momentum,
|
|
init=HistoryType[momentum_init.upper()],
|
|
momentum_hist=momentum_hist,
|
|
direction=direction,
|
|
rand_init_noise_type=NoiseType[rand_init_noise_type.upper()],
|
|
noise_type=NoiseType[noise_type.upper()],
|
|
custom_noise=custom_noise_opt.clone() if custom_noise_opt else None,
|
|
guidance=guidance_cfg_opt,
|
|
)
|
|
return (
|
|
samplers.KSAMPLER(
|
|
SonarDPMPPSDE.sampler,
|
|
{
|
|
"sonar_config": cfg,
|
|
"eta": eta,
|
|
"s_noise": s_noise,
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
class SamplerNodeConfigOverride:
|
|
KWARG_OVERRIDES = ("s_noise", "eta", "s_churn", "r", "solver_type")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sampler": ("SAMPLER",),
|
|
"eta": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"s_noise": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"s_churn": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"r": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.5,
|
|
"step": 0.01,
|
|
"round": False,
|
|
},
|
|
),
|
|
"sde_solver": (("midpoint", "heun"),),
|
|
"cpu_noise": ("BOOLEAN", {"default": True}),
|
|
"normalize": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"noise_type": (tuple(NoiseType.get_names()),),
|
|
"custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
FUNCTION = "get_sampler"
|
|
|
|
def get_sampler(
|
|
self,
|
|
sampler,
|
|
eta,
|
|
s_noise,
|
|
s_churn,
|
|
r,
|
|
sde_solver,
|
|
cpu_noise=True,
|
|
noise_type=None,
|
|
custom_noise_opt=None,
|
|
normalize=True,
|
|
):
|
|
return (
|
|
samplers.KSAMPLER(
|
|
self.sampler_function,
|
|
extra_options=sampler.extra_options
|
|
| {
|
|
"override_sampler_cfg": {
|
|
"sampler": sampler,
|
|
"noise_type": NoiseType[noise_type.upper()]
|
|
if noise_type is not None
|
|
else None,
|
|
"custom_noise": custom_noise_opt,
|
|
"s_noise": s_noise,
|
|
"eta": eta,
|
|
"s_churn": s_churn,
|
|
"r": r,
|
|
"solver_type": sde_solver,
|
|
"cpu_noise": cpu_noise,
|
|
"normalize": normalize,
|
|
},
|
|
},
|
|
inpaint_options=sampler.inpaint_options | {},
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
@torch.no_grad()
|
|
def sampler_function(
|
|
cls,
|
|
model,
|
|
x,
|
|
sigmas,
|
|
*args: list[Any],
|
|
override_sampler_cfg: dict[str, Any] | None = None,
|
|
noise_sampler: Callable | None = None,
|
|
extra_args: dict[str, Any] | None = None,
|
|
**kwargs: dict[str, Any],
|
|
):
|
|
if not override_sampler_cfg:
|
|
raise ValueError("Override sampler config missing!")
|
|
if extra_args is None:
|
|
extra_args = {}
|
|
cfg = override_sampler_cfg
|
|
sampler, noise_type, custom_noise, cpu, normalize = (
|
|
cfg["sampler"],
|
|
cfg.get("noise_type"),
|
|
cfg.get("custom_noise"),
|
|
cfg.get("cpu_noise", True),
|
|
cfg.get("normalize", True),
|
|
)
|
|
sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
|
|
seed = extra_args.get("seed")
|
|
if custom_noise is not None:
|
|
noise_sampler = custom_noise.make_noise_sampler(
|
|
x,
|
|
sigma_min,
|
|
sigma_max,
|
|
seed=seed,
|
|
cpu=cpu,
|
|
normalized=normalize,
|
|
)
|
|
elif noise_type is not None:
|
|
noise_sampler = noise.get_noise_sampler(
|
|
noise_type,
|
|
x,
|
|
sigma_min,
|
|
sigma_max,
|
|
seed=seed,
|
|
cpu=cpu,
|
|
normalized=normalize,
|
|
)
|
|
sig = inspect.signature(sampler.sampler_function)
|
|
params = sig.parameters
|
|
kwargs = kwargs | {}
|
|
if "noise_sampler" in params:
|
|
kwargs["noise_sampler"] = noise_sampler
|
|
for k in cls.KWARG_OVERRIDES:
|
|
if k not in params or cfg.get(k) is None:
|
|
continue
|
|
kwargs[k] = cfg[k]
|
|
return sampler.sampler_function(
|
|
model,
|
|
x,
|
|
sigmas,
|
|
*args,
|
|
extra_args=extra_args,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SamplerSonarEuler": SamplerNodeSonarEuler,
|
|
"SamplerSonarEulerA": SamplerNodeSonarEulerAncestral,
|
|
"SamplerSonarDPMPPSDE": SamplerNodeSonarDPMPPSDE,
|
|
"SonarGuidanceConfig": GuidanceConfigNode,
|
|
"SamplerConfigOverride": SamplerNodeConfigOverride,
|
|
"NoisyLatentLike": NoisyLatentLikeNode,
|
|
"SonarCustomNoise": SonarCustomNoiseNode,
|
|
"SonarCompositeNoise": SonarCompositeNoiseNode,
|
|
"SonarModulatedNoise": SonarModulatedNoiseNode,
|
|
"SonarRepeatedNoise": SonarRepeatedNoiseNode,
|
|
"SonarScheduledNoise": SonarScheduledNoiseNode,
|
|
"SonarGuidedNoise": SonarGuidedNoiseNode,
|
|
"SonarRandomNoise": SonarRandomNoiseNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
|
|
if "bleh" in external.MODULES:
|
|
bleh = external.MODULES["bleh"]
|
|
bleh_latentutils = bleh.py.latent_utils
|
|
|
|
class SonarBlendFilterNoiseNode(
|
|
SonarCustomNoiseNodeBase,
|
|
SonarNormalizeNoiseNodeMixin,
|
|
):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
result = super().INPUT_TYPES(include_rescale=False, include_chain=False)
|
|
result["required"] |= {
|
|
"sonar_custom_noise": ("SONAR_CUSTOM_NOISE",),
|
|
"blend_mode": (
|
|
("simple_add", *bleh_latentutils.BLENDING_MODES.keys()),
|
|
),
|
|
"ffilter": (tuple(bleh_latentutils.FILTER_PRESETS.keys()),),
|
|
"ffilter_custom": ("STRING", {"default": ""}),
|
|
"ffilter_scale": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": -100.0, "max": 100.0},
|
|
),
|
|
"ffilter_strength": (
|
|
"FLOAT",
|
|
{"default": 0.0, "min": -100.0, "max": 100.0},
|
|
),
|
|
"ffilter_threshold": (
|
|
"INT",
|
|
{"default": 1, "min": 1, "max": 32},
|
|
),
|
|
"enhance_mode": (("none", *bleh_latentutils.ENHANCE_METHODS),),
|
|
"enhance_strength": (
|
|
"FLOAT",
|
|
{"default": 0.0, "min": -100.0, "max": 100.0},
|
|
),
|
|
"affect": (("result", "noise", "both"),),
|
|
"normalize_result": (("default", "forced", "disabled"),),
|
|
"normalize_noise": (("default", "forced", "disabled"),),
|
|
}
|
|
return result
|
|
|
|
def get_item_class(self):
|
|
return noise.BlendFilterNoise
|
|
|
|
def go(
|
|
self,
|
|
factor,
|
|
sonar_custom_noise,
|
|
blend_mode,
|
|
ffilter,
|
|
ffilter_custom,
|
|
ffilter_scale,
|
|
ffilter_strength,
|
|
ffilter_threshold,
|
|
enhance_mode,
|
|
enhance_strength,
|
|
affect,
|
|
normalize_result,
|
|
normalize_noise,
|
|
):
|
|
import ast
|
|
|
|
ffilter_custom = ffilter_custom.strip()
|
|
normalize_result = (
|
|
None if normalize_result == "default" else normalize_result == "forced"
|
|
)
|
|
normalize_noise = (
|
|
None if normalize_noise == "default" else normalize_noise == "forced"
|
|
)
|
|
if ffilter_custom:
|
|
ffilter = ast.literal_eval(f"[{ffilter_custom}]")
|
|
else:
|
|
ffilter = bleh_latentutils.FILTER_PRESETS[ffilter]
|
|
return super().go(
|
|
factor,
|
|
noise=sonar_custom_noise.clone(),
|
|
blend_mode=blend_mode,
|
|
ffilter=ffilter,
|
|
ffilter_scale=ffilter_scale,
|
|
ffilter_strength=ffilter_strength,
|
|
ffilter_threshold=ffilter_threshold,
|
|
enhance_mode=enhance_mode,
|
|
enhance_strength=enhance_strength,
|
|
affect=affect,
|
|
normalize_noise=self.get_normalize(normalize_noise),
|
|
normalize_result=self.get_normalize(normalize_result),
|
|
)
|
|
|
|
NODE_CLASS_MAPPINGS["SonarBlendFilterNoise"] = SonarBlendFilterNoiseNode
|
|
|
|
if "restart" in external.MODULES:
|
|
rs = external.MODULES["restart"]
|
|
|
|
class KRestartSamplerCustomNoise:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
get_normal_schedulers = getattr(
|
|
rs.nodes,
|
|
"get_supported_normal_schedulers",
|
|
rs.nodes.get_supported_restart_schedulers,
|
|
)
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"add_noise": (["enable", "disable"],),
|
|
"noise_seed": (
|
|
"INT",
|
|
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
|
),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
|
"sampler": ("SAMPLER",),
|
|
"scheduler": (get_normal_schedulers(),),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
|
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
|
|
"return_with_leftover_noise": (["disable", "enable"],),
|
|
"segments": (
|
|
"STRING",
|
|
{
|
|
"default": rs.restart_sampling.DEFAULT_SEGMENTS,
|
|
"multiline": False,
|
|
},
|
|
),
|
|
"restart_scheduler": (rs.nodes.get_supported_restart_schedulers(),),
|
|
"chunked_mode": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "LATENT")
|
|
RETURN_NAMES = ("output", "denoised_output")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
add_noise,
|
|
noise_seed,
|
|
steps,
|
|
cfg,
|
|
sampler,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
start_at_step,
|
|
end_at_step,
|
|
return_with_leftover_noise,
|
|
segments,
|
|
restart_scheduler,
|
|
chunked_mode=False,
|
|
custom_noise_opt=None,
|
|
):
|
|
return rs.restart_sampling.restart_sampling(
|
|
model,
|
|
noise_seed,
|
|
steps,
|
|
cfg,
|
|
sampler,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
segments,
|
|
restart_scheduler,
|
|
disable_noise=add_noise == "disable",
|
|
step_range=(start_at_step, end_at_step),
|
|
force_full_denoise=return_with_leftover_noise != "enable",
|
|
output_only=False,
|
|
chunked_mode=chunked_mode,
|
|
custom_noise=custom_noise_opt.make_noise_sampler
|
|
if custom_noise_opt
|
|
else None,
|
|
)
|
|
|
|
NODE_CLASS_MAPPINGS["KRestartSamplerCustomNoise"] = KRestartSamplerCustomNoise
|
|
|
|
if hasattr(rs.restart_sampling, "RestartSampler"):
|
|
|
|
class RestartSamplerCustomNoise:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sampler": ("SAMPLER",),
|
|
"chunked_mode": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"custom_noise_opt": ("SONAR_CUSTOM_NOISE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("SAMPLER",)
|
|
FUNCTION = "go"
|
|
CATEGORY = "sampling/custom_sampling/samplers"
|
|
|
|
def go(self, sampler, chunked_mode, custom_noise_opt=None):
|
|
restart_options = {
|
|
"restart_chunked": chunked_mode,
|
|
"restart_wrapped_sampler": sampler,
|
|
"restart_custom_noise": None
|
|
if custom_noise_opt is None
|
|
else custom_noise_opt.make_noise_sampler,
|
|
}
|
|
restart_sampler = samplers.KSAMPLER(
|
|
rs.restart_sampling.RestartSampler.sampler_function,
|
|
extra_options=sampler.extra_options | restart_options,
|
|
inpaint_options=sampler.inpaint_options,
|
|
)
|
|
return (restart_sampler,)
|
|
|
|
NODE_CLASS_MAPPINGS["RestartSamplerCustomNoise"] = RestartSamplerCustomNoise
|