nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add sigma IO to hybrid restart sampler

Browse Source
This commit is contained in:
digital-garbage
2026-04-30 21:10:27 +03:00
parent 3b291831f2
commit ca3ca3ce4c
2 changed files with 104 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/FunPackHybridEuler2SSampler.md
+23 -7
View File
@@ -1,6 +1,6 @@
# FunPack Hybrid Euler 2S Sampler
This node outputs a ComfyUI `SAMPLER` object that can be connected directly into `CustomSamplerAdvanced` / `SamplerCustomAdvanced`.
This node outputs a ComfyUI `SAMPLER` and can also take optional `SIGMAS` input and output modified `SIGMAS` for use with `CustomSamplerAdvanced` / `SamplerCustomAdvanced`.
## Purpose
@@ -11,11 +11,23 @@ It is designed as a quality/speed compromise between:
This sampler keeps classic Euler ancestral for the early "structure building" stage where motion, anatomy and main composition are forming. On the late denoise steps it switches to a deterministic Euler / DPM-Solver++(2S) ODE refinement path where fine detail and cleanup usually matter most.
It can also run an optional paper-style Restart pass during that late quality stage:
It can also build an optional paper-style Restart schedule during that late quality stage:
- re-noise the current latent up to a higher sigma
- replay a short late-step sigma interval with the deterministic refinement path
- repeat that replay multiple times if requested
- detect a restart anchor inside the late quality phase
- insert upward sigma jumps back to a higher sigma
- replay the chosen sigma interval one or more times
Unlike the older implementation, Restart is now encoded directly into the outgoing `SIGMAS` schedule, so the visible schedule passed into `SamplerCustomAdvanced` is the actual schedule that gets executed.
## Recommended wiring
Use this node with `SamplerCustomAdvanced` or `CustomSamplerAdvanced`:
- connect your base scheduler output into this node's optional `sigmas` input
- connect this node's `sampler` output into the sampler input
- connect this node's `sigmas` output into the sigmas input
This node is not intended for `KSampler`.
## Parameters
@@ -31,7 +43,7 @@ Example: `0.35` means only the last 35% of steps use the ODE refinement path.
- `0.0` = pure late-step Euler ODE
- `1.0` = full late-step DPM++(2S)-style correction
**restart_steps**: Target size of the restarted sigma interval. The sampler now interprets this against the actual late-phase sigma curve, so with custom sigmas it may choose a slightly different number of replay transitions in order to cover a comparable sigma span.
**restart_steps**: Target size of the restarted sigma interval. The node interprets this against the actual late-phase sigma curve and expands the outgoing `SIGMAS` schedule accordingly.
**restart_repeats**: How many Restart loops to run. `0` disables Restart.
@@ -39,6 +51,8 @@ Example: `0.35` means only the last 35% of steps use the ODE refinement path.
**restart_noise**: Noise strength used when re-noising the latent up to the restart interval's higher sigma.
**sigmas**: Optional incoming sigma schedule. If connected, the node outputs a restart-expanded sigma schedule that matches the sampler settings. If not connected, the sampler output still works, but there is no schedule for the node to rewrite.
## Recommended starting values
- `eta = 1.0`
@@ -56,7 +70,7 @@ For a first Restart test, try:
- `restart_trigger_pct = 0.85`
- `restart_noise = 1.0`
Example: with `8` base sampling steps, `restart_steps = 3`, and `restart_repeats = 1`, a normal evenly spaced late schedule will usually add about `2` extra replay denoise steps. With custom sigmas, the exact replay step count can shift slightly because Restart is now anchored to the real sigma values rather than raw step indices.
Example: with `8` base sampling steps, `restart_steps = 3`, and `restart_repeats = 1`, a normal evenly spaced late schedule will usually expand into a schedule with about `2` extra replay denoise transitions. With custom sigmas, the exact replay step count can shift slightly because Restart is anchored to the real sigma values rather than raw step indices.
## Expected behavior
@@ -69,6 +83,8 @@ Compared to plain `euler_ancestral`, this sampler should usually:
With Restart enabled, it may further improve late cleanup and detail consistency, but it will increase runtime and may be too aggressive for some video workflows if overused.
Because Restart is now represented in the outgoing `SIGMAS`, the total number of executed transitions should match what the expanded schedule shows in the downstream custom sampler.
## Limitation
This node improves the **sampler-side quality/speed tradeoff**. It does **not** reduce the underlying cost of CFG++ guidance itself. If the goal becomes specifically "CFG++ quality at lower guidance cost", that likely requires a custom `GUIDER`, not just a custom `SAMPLER`.
samplers.py
+81 -70
View File
@@ -142,46 +142,53 @@ def _build_restart_sigmas(sigmas, late_start, total_steps, restart_trigger_pct,
return sigmas[restart_from:anchor_index + 1], anchor_index
def _run_restart_replay(model, x, restart_sigmas, s_in, extra_args, correction_blend,
callback=None, callback_step_start=0):
def _expand_restart_sigmas(sigmas, high_quality_pct, restart_steps, restart_repeats, restart_trigger_pct):
if sigmas is None or not isinstance(sigmas, torch.Tensor):
return None, None
expanded_sigmas = sigmas.detach().clone()
total_steps = max(0, int(expanded_sigmas.shape[0]) - 1)
if total_steps <= 0:
return expanded_sigmas, None
high_quality_pct = max(0.0, min(1.0, float(high_quality_pct)))
restart_steps = max(2, int(restart_steps))
restart_repeats = max(0, int(restart_repeats))
restart_trigger_pct = max(0.0, min(1.0, float(restart_trigger_pct)))
late_steps = max(1, int(math.ceil(total_steps * high_quality_pct))) if high_quality_pct > 0.0 else 0
late_start = max(0, total_steps - late_steps)
quality_sigma_start = None
if late_start < expanded_sigmas.shape[0]:
quality_sigma_start = float(expanded_sigmas[late_start].item())
if restart_repeats <= 0:
return expanded_sigmas, quality_sigma_start
restart_sigmas, anchor_index = _build_restart_sigmas(
expanded_sigmas,
late_start,
total_steps,
restart_trigger_pct,
restart_steps,
)
if restart_sigmas.numel() <= 1:
return x, callback_step_start
return expanded_sigmas, quality_sigma_start
callback_step = int(callback_step_start)
for idx in range(int(restart_sigmas.shape[0]) - 1):
sigma = restart_sigmas[idx]
sigma_next = restart_sigmas[idx + 1]
denoised = model(x, sigma * s_in, **extra_args)
if callback is not None:
callback({
"x": x,
"i": callback_step,
"sigma": sigma,
"sigma_hat": sigma,
"denoised": denoised,
})
callback_step += 1
x, _ = _hybrid_ode_step(
model,
x,
sigma,
sigma_next,
s_in,
extra_args,
correction_blend,
denoised=denoised,
)
return x, callback_step
prefix = expanded_sigmas[:anchor_index + 1]
suffix = expanded_sigmas[anchor_index + 1:]
repeated = [restart_sigmas.clone() for _ in range(restart_repeats)]
expanded_sigmas = torch.cat([prefix] + repeated + [suffix], dim=0)
return expanded_sigmas, quality_sigma_start
def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=None,
disable=None, eta=1.0, s_noise=1.0,
high_quality_pct=0.35, correction_blend=1.0,
restart_steps=3, restart_repeats=0,
restart_trigger_pct=0.85, restart_noise=1.0):
restart_trigger_pct=0.85, restart_noise=1.0,
quality_sigma_start=None):
"""
Hybrid sampler:
- Early schedule: Euler ancestral for motion/anatomy buildup.
@@ -208,24 +215,39 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
restart_repeats = max(0, int(restart_repeats))
restart_trigger_pct = max(0.0, min(1.0, float(restart_trigger_pct)))
restart_noise = max(0.0, float(restart_noise))
late_steps = max(1, int(math.ceil(total_steps * high_quality_pct))) if high_quality_pct > 0.0 else 0
late_start = max(0, total_steps - late_steps)
s_in = x.new_ones([x.shape[0]])
restart_sigmas, restart_anchor = _build_restart_sigmas(
sigmas,
late_start,
total_steps,
restart_trigger_pct,
restart_steps,
)
restart_has_run = False
callback_step = 0
if quality_sigma_start is None:
late_steps = max(1, int(math.ceil(total_steps * high_quality_pct))) if high_quality_pct > 0.0 else 0
late_start = max(0, total_steps - late_steps)
if late_start < sigmas.shape[0]:
quality_sigma_start = float(sigmas[late_start].item())
else:
quality_sigma_start = float(quality_sigma_start)
for i in comfy.utils.model_trange(total_steps, disable=disable):
sigma = sigmas[i]
sigma_next = sigmas[i + 1]
is_restart_jump = sigma_next > sigma
in_quality_phase = quality_sigma_start is not None and float(sigma.item()) <= quality_sigma_start
if i < late_start:
if is_restart_jump:
denoised = None
if callback is not None:
denoised = model(x, sigma * s_in, **extra_args)
callback({
"x": x,
"i": callback_step,
"sigma": sigma,
"sigma_hat": sigma,
"denoised": denoised,
})
callback_step += 1
x = _renoise_to_sigma(x, sigma, sigma_next, restart_noise, noise_sampler)
continue
if not in_quality_phase:
denoised = model(x, sigma * s_in, **extra_args)
if callback is not None:
@@ -250,28 +272,6 @@ def sample_funpack_hybrid_euler_2s(model, x, sigmas, extra_args=None, callback=N
if sigma_next > 0 and eta > 0 and s_noise > 0:
x = x + noise_sampler(sigma, sigma_next) * s_noise * sigma_up
else:
if restart_repeats > 0 and not restart_has_run and i >= restart_anchor:
if restart_sigmas.numel() > 1 and restart_sigmas[0] > restart_sigmas[-1]:
for _ in range(restart_repeats):
x = _renoise_to_sigma(
x,
restart_sigmas[-1],
restart_sigmas[0],
restart_noise,
noise_sampler,
)
x, callback_step = _run_restart_replay(
model,
x,
restart_sigmas,
s_in,
extra_args,
correction_blend,
callback=callback,
callback_step_start=callback_step,
)
restart_has_run = True
denoised = model(x, sigma * s_in, **extra_args)
if callback is not None:
@@ -359,17 +359,27 @@ class FunPackHybridEuler2SSampler:
"step": 0.01,
"tooltip": "Noise strength used when re-noising up to the Restart interval's higher sigma."
}),
},
"optional": {
"sigmas": ("SIGMAS",),
}
}
RETURN_TYPES = ("SAMPLER",)
RETURN_NAMES = ("sampler",)
RETURN_TYPES = ("SAMPLER", "SIGMAS")
RETURN_NAMES = ("sampler", "sigmas")
FUNCTION = "get_sampler"
CATEGORY = "FunPack/Sampling"
DESCRIPTION = "Hybrid sampler: early Euler ancestral for motion, late DPM-Solver++(2S) ODE refinement for quality, with optional paper-style Restart replay."
DESCRIPTION = "Hybrid sampler: early Euler ancestral for motion, late DPM-Solver++(2S) ODE refinement for quality, with optional sigma-driven Restart replay."
def get_sampler(self, eta, s_noise, high_quality_pct, correction_blend,
restart_steps, restart_repeats, restart_trigger_pct, restart_noise):
restart_steps, restart_repeats, restart_trigger_pct, restart_noise, sigmas=None):
expanded_sigmas, quality_sigma_start = _expand_restart_sigmas(
sigmas,
high_quality_pct,
restart_steps,
restart_repeats,
restart_trigger_pct,
)
sampler = comfy.samplers.KSAMPLER(
sample_funpack_hybrid_euler_2s,
extra_options={
@@ -381,6 +391,7 @@ class FunPackHybridEuler2SSampler:
"restart_repeats": restart_repeats,
"restart_trigger_pct": restart_trigger_pct,
"restart_noise": restart_noise,
"quality_sigma_start": quality_sigma_start,
}
)
return (sampler,)
return (sampler, expanded_sigmas)