232 lines
8.5 KiB
Python
232 lines
8.5 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import comfy
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
from comfy.samplers import KSAMPLER
|
|
|
|
import latent_preview
|
|
|
|
|
|
def get_cosine_similarity(a: torch.Tensor, b: torch.Tensor, dim: int = 1) -> torch.Tensor:
|
|
return F.cosine_similarity(a.flatten(start_dim=dim), b.flatten(start_dim=dim), dim=dim)
|
|
|
|
|
|
def sample_adaptive_custom(
|
|
model,
|
|
x,
|
|
sigmas,
|
|
extra_args=None,
|
|
callback=None,
|
|
disable=None,
|
|
wrapped_sampler=None,
|
|
error_type="cosine",
|
|
base_step_size=0.0004,
|
|
min_step_size=0.01,
|
|
max_step_size=0.2,
|
|
max_steps=100,
|
|
smoothing_coef=0.0,
|
|
sigma_schedule_out=None,
|
|
denoise=1.0,
|
|
error_bias=0.002,
|
|
):
|
|
"""Adaptive sampler that wraps another sampler with adaptive step sizes."""
|
|
|
|
extra_args = {} if extra_args is None else extra_args
|
|
|
|
disable_pbar = disable if disable is not None else not comfy.utils.PROGRESS_BAR_ENABLED
|
|
|
|
model_sampling = model.inner_model.inner_model.model_sampling
|
|
sigma_min = model_sampling.sigma_min
|
|
sigma_max = model_sampling.sigma_max
|
|
|
|
current_sigma = sigma_min + (sigma_max - sigma_min) * denoise
|
|
step_size = min_step_size # First step = smallest
|
|
|
|
total_steps = 0
|
|
error_val = 1.0
|
|
v_prev = None
|
|
|
|
# Track sigma schedule
|
|
sigma_schedule = [current_sigma]
|
|
|
|
denoised_capture = [None]
|
|
|
|
def capture_denoised(inner_callback):
|
|
def wrapper(info):
|
|
denoised_capture[0] = info.get('denoised', None)
|
|
if inner_callback:
|
|
inner_callback(info)
|
|
return wrapper
|
|
|
|
while current_sigma > sigma_min and total_steps < max_steps:
|
|
# Run wrapped sampler for one step
|
|
sigma_target = max(current_sigma - step_size, sigma_min)
|
|
|
|
# Create a single-step sigma tensor for the wrapped sampler
|
|
step_sigmas = torch.tensor([current_sigma, sigma_target], device=x.device, dtype=x.dtype)
|
|
|
|
x_prev = x.clone()
|
|
|
|
# Run wrapped sampler (e.g., euler, heun)
|
|
wrapped_callback = capture_denoised(callback)
|
|
x = wrapped_sampler.sampler_function(model, x, step_sigmas, extra_args, wrapped_callback, disable_pbar)
|
|
|
|
# Get velocity: v = denoised - x
|
|
denoised = denoised_capture[0]
|
|
if denoised is not None:
|
|
v_current = denoised - x
|
|
else:
|
|
v_current = x - x_prev
|
|
|
|
if callback is not None:
|
|
callback({
|
|
'x': x,
|
|
'i': total_steps,
|
|
'sigma': sigma_target,
|
|
'sigma_hat': sigma_target,
|
|
'denoised': denoised if denoised is not None else x
|
|
})
|
|
|
|
if not disable_pbar:
|
|
if total_steps > 0:
|
|
print(f"Step {total_steps + 1}: sigma={sigma_target:.6f}, step_size={step_size:.6f}, error={error_val:.6f}")
|
|
else:
|
|
print(f"Step {total_steps + 1}: sigma={sigma_target:.6f}, step_size={step_size:.6f}")
|
|
|
|
# Compute error from velocity change
|
|
if v_prev is not None:
|
|
if error_type == "cosine":
|
|
error = 1.0 - get_cosine_similarity(v_prev, v_current, dim=1).abs().mean()
|
|
else:
|
|
error = F.mse_loss(v_prev, v_current)
|
|
|
|
error_val = max(error.item(), 1e-10)
|
|
|
|
# Adjust step size for next iteration: base_step_size / error
|
|
new_step_size = base_step_size / (error_val + error_bias)
|
|
step_size = max(min_step_size, min(max_step_size, step_size * smoothing_coef + new_step_size * (1 - smoothing_coef)))
|
|
|
|
v_prev = v_current
|
|
|
|
current_sigma = sigma_target
|
|
sigma_schedule.append(sigma_target)
|
|
total_steps += 1
|
|
|
|
if not disable_pbar:
|
|
print(f"Completed in {total_steps} steps, final sigma: {current_sigma:.6f}")
|
|
|
|
if sigma_schedule_out is not None:
|
|
sigma_schedule_out[0] = torch.tensor(sigma_schedule, device=x.device, dtype=x.dtype)
|
|
|
|
return x
|
|
|
|
|
|
class AdaptiveSamplerCustom:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"noise": ("NOISE",),
|
|
"guider": ("GUIDER",),
|
|
"sampler": ("SAMPLER",),
|
|
"latent_image": ("LATENT",),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Denoising strength (1.0 = full)"}),
|
|
"error_type": (["cosine", "mse"], {"default": "cosine", "tooltip": "cosine: measure direction change, mse: measure magnitude change"}),
|
|
"base_step_size": ("FLOAT", {"default": 0.0004, "min": 0.0001, "max": 1.0, "step": 0.0001, "tooltip": "Base multiplier for step size calculation"}),
|
|
"min_step_size": ("FLOAT", {"default": 0.005, "min": 0.0001, "max": 1.0, "step": 0.0001, "tooltip": "Minimum allowed step size"}),
|
|
"max_step_size": ("FLOAT", {"default": 0.2, "min": 0.001, "max": 1.0, "step": 0.001, "tooltip": "Maximum allowed step size"}),
|
|
"max_steps": ("INT", {"default": 100, "min": 1, "max": 10000, "tooltip": "Maximum number of adaptive steps"}),
|
|
"smoothing_coef": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "EMA coefficient: 0 = new value only, 1 = keep old value"}),
|
|
"error_bias": ("FLOAT", {"default": 0.002, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Bias added to error. Higher values = smaller steps"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "LATENT", "SIGMAS")
|
|
RETURN_NAMES = ("output", "denoised", "sigmas")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling/custom_sampling"
|
|
|
|
def sample(
|
|
self,
|
|
noise,
|
|
guider,
|
|
sampler,
|
|
latent_image,
|
|
denoise=1.0,
|
|
error_type="cosine",
|
|
base_step_size=0.0004,
|
|
min_step_size=0.01,
|
|
max_step_size=0.2,
|
|
max_steps=100,
|
|
smoothing_coef=0.0,
|
|
error_bias=0.002,
|
|
):
|
|
latent = latent_image.copy()
|
|
x = latent["samples"]
|
|
|
|
noise_mask = None
|
|
if "noise_mask" in latent:
|
|
noise_mask = latent["noise_mask"]
|
|
|
|
x = comfy.sample.fix_empty_latent_channels(guider.model_patcher, x, latent.get("downscale_ratio_spacial", None))
|
|
latent["samples"] = x
|
|
|
|
# Create dummy sigmas for the wrapped sampler - we'll ignore them
|
|
model_sampling = guider.model_patcher.get_model_object("model_sampling")
|
|
sigma_min = model_sampling.sigma_min
|
|
sigma_max = model_sampling.sigma_max
|
|
start_sigma = sigma_min + (sigma_max - sigma_min) * denoise
|
|
sigmas = torch.tensor([start_sigma, 0.0], device=x.device, dtype=x.dtype)
|
|
|
|
x0_output = {}
|
|
callback = latent_preview.prepare_callback(guider.model_patcher, max_steps, x0_output)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
|
|
# Store sigma schedule
|
|
sigma_schedule_out = [None]
|
|
|
|
def adaptive_sampler_function(model, x, sigmas, extra_args, callback, disable):
|
|
return sample_adaptive_custom(
|
|
model=model,
|
|
x=x,
|
|
sigmas=sigmas,
|
|
extra_args=extra_args,
|
|
callback=callback,
|
|
disable=disable,
|
|
wrapped_sampler=sampler,
|
|
error_type=error_type,
|
|
base_step_size=base_step_size,
|
|
min_step_size=min_step_size,
|
|
max_step_size=max_step_size,
|
|
max_steps=max_steps,
|
|
smoothing_coef=smoothing_coef,
|
|
sigma_schedule_out=sigma_schedule_out,
|
|
denoise=denoise,
|
|
error_bias=error_bias,
|
|
)
|
|
|
|
sampler_obj = KSAMPLER(adaptive_sampler_function, {})
|
|
|
|
noise_gen = noise.generate_noise(latent)
|
|
samples = guider.sample(noise_gen, x, sampler_obj, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
|
|
if "x0" in x0_output:
|
|
x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
|
|
if samples.is_nested:
|
|
latent_shapes = [s.shape for s in samples.unbind()]
|
|
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = x0_out
|
|
else:
|
|
out_denoised = out
|
|
|
|
return (out, out_denoised, sigma_schedule_out[0])
|