import math import torch import comfy.k_diffusion.sampling as k_diffusion_sampling import comfy.model_sampling import comfy.samplers import comfy.utils MOTION_PULSE_MODES = ["off", "balanced", "aggressive", "custom"] VELOCITY_BIAS_MODES = ["off", "capture", "apply", "capture_and_apply"] VELOCITY_BIAS_TARGETS = (0.90, 0.80) VELOCITY_BIAS_MEMORY = {} def _sigma_fn(t): return t.neg().exp() def _t_fn(sigma): return sigma.log().neg() def _hybrid_ode_step(model, x, sigma, sigma_next, s_in, extra_args, correction_blend, denoised=None): if denoised is None: denoised = model(x, sigma * s_in, **extra_args) if sigma_next == 0: return denoised, denoised d = k_diffusion_sampling.to_d(x, sigma, denoised) dt = sigma_next - sigma x_euler = x + d * dt if correction_blend <= 0.0: return x_euler, denoised t = _t_fn(sigma) t_next = _t_fn(sigma_next) h = t_next - t r = 0.5 s = t + r * h x_mid = (_sigma_fn(s) / _sigma_fn(t)) * x - torch.expm1(-h * r) * denoised denoised_mid = model(x_mid, _sigma_fn(s) * s_in, **extra_args) x_2s = (_sigma_fn(t_next) / _sigma_fn(t)) * x - torch.expm1(-h) * denoised_mid if correction_blend >= 1.0: return x_2s, denoised return x_euler.lerp(x_2s, correction_blend), denoised def _apply_motion_pulse(x, sigma, sigma_next, pulse_noise, noise_sampler): if pulse_noise <= 0.0 or sigma_next >= sigma: return x sigma_delta_sq = max(0.0, float(sigma * sigma - sigma_next * sigma_next)) if sigma_delta_sq <= 0.0: return x sigma_delta = math.sqrt(sigma_delta_sq) return x + noise_sampler(sigma, sigma_next) * (pulse_noise * sigma_delta) def _velocity_bias_enabled(mode, action): mode = (mode or "off").lower() if mode == "capture_and_apply": return action in {"capture", "apply"} return mode == action def _velocity_bias_key(refinement_key, aspect_bucket, target, x): shape = tuple(int(item) for item in getattr(x, "shape", ())) key = str(refinement_key or "default").strip() or "default" aspect = str(aspect_bucket or "any").strip() or "any" return (key, aspect, f"{float(target):.2f}", shape) def _sigma_ratio(sigmas, sigma): try: start = float(sigmas[0].item()) current = float(sigma.item()) except Exception: return None if start <= 0.0: return None return current / start def _velocity_bias_target(sigmas, sigma): ratio = _sigma_ratio(sigmas, sigma) if ratio is None: return None target = min(VELOCITY_BIAS_TARGETS, key=lambda item: abs(float(item) - ratio)) return target if abs(float(target) - ratio) <= 0.065 else None def _capture_velocity_bias(refinement_key, aspect_bucket, target, x, sigma, denoised): if target is None: return try: direction = k_diffusion_sampling.to_d(x, sigma, denoised).detach().float().cpu() except Exception: return key = _velocity_bias_key(refinement_key, aspect_bucket, target, x) slot = VELOCITY_BIAS_MEMORY.setdefault(key, {"count": 0, "direction": None}) count = int(slot.get("count", 0)) previous = slot.get("direction") if count <= 0 or not isinstance(previous, torch.Tensor) or tuple(previous.shape) != tuple(direction.shape): slot["direction"] = direction slot["count"] = 1 return slot["direction"] = (previous * count + direction) / float(count + 1) slot["count"] = min(count + 1, 256) def _apply_velocity_bias(x, refinement_key, aspect_bucket, target, strength): if target is None or strength <= 0.0: return x key = _velocity_bias_key(refinement_key, aspect_bucket, target, x) slot = VELOCITY_BIAS_MEMORY.get(key) if not isinstance(slot, dict) or not isinstance(slot.get("direction"), torch.Tensor): return x direction = slot["direction"] if tuple(direction.shape) != tuple(x.shape): return x try: direction = direction.to(device=x.device, dtype=x.dtype) delta = direction * max(0.0, min(0.35, float(strength))) max_delta = x.detach().float().norm().clamp_min(1e-8) * 0.045 delta_norm = delta.detach().float().norm().clamp_min(1e-8) if delta_norm > max_delta: delta = delta * (max_delta / delta_norm).to(device=x.device, dtype=x.dtype) return x + delta except Exception: return x def _find_schedule_anchor_index(sigmas, total_steps, schedule_progress): if sigmas is None or total_steps <= 1: return 0 schedule_progress = max(0.0, min(1.0, schedule_progress)) return min(total_steps - 1, max(0, int(round(schedule_progress * max(0, total_steps - 1))))) def _resolve_motion_pulse_options(motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength): mode = (motion_pulse_mode or "off").lower() if mode not in MOTION_PULSE_MODES: mode = "off" if mode == "off": return { "enabled": False, "start_pct": 0.30, "count": 0, "spacing_pct": 0.22, "strength": 0.0, "noise": 0.0, } start_pct = 0.30 if motion_pulse_start_pct is None else float(motion_pulse_start_pct) spacing_pct = 0.22 if motion_pulse_spacing_pct is None else float(motion_pulse_spacing_pct) strength = 0.85 if motion_pulse_strength is None else float(motion_pulse_strength) count = 2 if motion_pulse_count is None else int(motion_pulse_count) start_pct = max(0.02, min(0.90, start_pct)) spacing_pct = max(0.04, min(0.45, spacing_pct)) strength = max(0.0, min(1.0, strength)) count = max(1, min(6, count)) if mode == "balanced": count = min(count, 1) strength = 0.55 if motion_pulse_strength is None else min(strength, 0.70) elif mode == "aggressive": count = max(2, count) strength = max(strength, 0.85) noise = 0.10 + strength * 0.55 return { "enabled": True, "start_pct": start_pct, "count": count, "spacing_pct": spacing_pct, "strength": strength, "noise": max(0.0, min(0.80, noise)), } def _get_late_start_index(total_steps, high_quality_pct): high_quality_pct = max(0.0, min(1.0, float(high_quality_pct))) late_steps = max(1, int(math.ceil(total_steps * high_quality_pct))) if high_quality_pct > 0.0 else 0 return max(0, total_steps - late_steps) def _build_motion_pulse_steps(sigmas, total_steps, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength): options = _resolve_motion_pulse_options( motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) if not options["enabled"] or total_steps <= 2: return [], options late_start = _get_late_start_index(total_steps, high_quality_pct) latest_anchor_pct = max(0.04, min(0.92, (late_start - 1) / max(1, total_steps - 1))) pulse_steps = [] used_anchors = set() for pulse_index in range(options["count"]): trigger_pct = options["start_pct"] + options["spacing_pct"] * pulse_index if trigger_pct >= latest_anchor_pct: break anchor_index = _find_schedule_anchor_index(sigmas, total_steps, trigger_pct) anchor_index = min(total_steps - 1, max(0, anchor_index)) if anchor_index in used_anchors or anchor_index >= late_start: continue used_anchors.add(anchor_index) pulse_steps.append({ "step_index": anchor_index, "noise": options["noise"], }) return pulse_steps, options def _prepare_dynamic_sigmas(sigmas, high_quality_pct, motion_pulse_mode="off", motion_pulse_start_pct=0.30, motion_pulse_count=2, motion_pulse_spacing_pct=0.22, motion_pulse_strength=0.85): if sigmas is None or not isinstance(sigmas, torch.Tensor): return None, None, [], 0.0 base_sigmas = sigmas.detach().clone() total_steps = max(0, int(base_sigmas.shape[0]) - 1) if total_steps <= 0: return base_sigmas, None, [], 0.0 late_start = _get_late_start_index(total_steps, high_quality_pct) quality_sigma_start = None if late_start < base_sigmas.shape[0]: quality_sigma_start = float(base_sigmas[late_start].item()) pulse_steps, motion_pulse_options = _build_motion_pulse_steps( base_sigmas, total_steps, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) return base_sigmas, quality_sigma_start, pulse_steps, motion_pulse_options["noise"] 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, quality_sigma_start=None, motion_pulse_mode="off", motion_pulse_start_pct=0.30, motion_pulse_count=2, motion_pulse_spacing_pct=0.22, motion_pulse_strength=0.85, motion_pulse_noise=0.0, motion_pulse_steps=None, velocity_bias_mode="off", velocity_bias_strength=0.0, velocity_refinement_key="default", velocity_aspect_bucket="any"): """ Hybrid sampler: - Early schedule: Euler ancestral for motion/anatomy buildup. - Late schedule: deterministic Euler / DPM-Solver++(2S) ODE refinement for detail. - Motion pulses: optional monotonic noise kicks before the late quality phase. """ if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST): return k_diffusion_sampling.sample_euler_ancestral( model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise ) extra_args = {} if extra_args is None else extra_args seed = extra_args.get("seed", None) noise_sampler = k_diffusion_sampling.default_noise_sampler(x, seed=seed) total_steps = max(0, len(sigmas) - 1) if total_steps <= 0: return x high_quality_pct = max(0.0, min(1.0, float(high_quality_pct))) correction_blend = max(0.0, min(1.0, float(correction_blend))) velocity_bias_mode = (velocity_bias_mode or "off").lower() if velocity_bias_mode not in VELOCITY_BIAS_MODES: velocity_bias_mode = "off" velocity_bias_strength = max(0.0, min(0.35, float(velocity_bias_strength or 0.0))) if not motion_pulse_steps: _, _, motion_pulse_steps, computed_motion_pulse_noise = _prepare_dynamic_sigmas( sigmas, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) motion_pulse_noise = computed_motion_pulse_noise motion_pulse_noise = max(0.0, float(motion_pulse_noise)) motion_step_noise = { int(item.get("step_index", -1)): max(0.0, float(item.get("noise", motion_pulse_noise))) for item in (motion_pulse_steps or []) if isinstance(item, dict) } s_in = x.new_ones([x.shape[0]]) 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] in_quality_phase = quality_sigma_start is not None and float(sigma.item()) <= quality_sigma_start if not in_quality_phase: pulse_noise = motion_step_noise.get(int(i), 0.0) if pulse_noise > 0.0: x = _apply_motion_pulse(x, sigma, sigma_next, pulse_noise, noise_sampler) velocity_target = _velocity_bias_target(sigmas, sigma) if _velocity_bias_enabled(velocity_bias_mode, "apply"): x = _apply_velocity_bias(x, velocity_refinement_key, velocity_aspect_bucket, velocity_target, velocity_bias_strength) denoised = model(x, sigma * s_in, **extra_args) if _velocity_bias_enabled(velocity_bias_mode, "capture"): _capture_velocity_bias(velocity_refinement_key, velocity_aspect_bucket, velocity_target, x, sigma, denoised) if callback is not None: callback({ "x": x, "i": callback_step, "sigma": sigma, "sigma_hat": sigma, "denoised": denoised, }) callback_step += 1 sigma_down, sigma_up = k_diffusion_sampling.get_ancestral_step(sigma, sigma_next, eta=eta) if sigma_down == 0: x = denoised else: d = k_diffusion_sampling.to_d(x, sigma, denoised) dt = sigma_down - sigma x = x + d * dt if sigma_next > 0 and eta > 0 and s_noise > 0: x = x + noise_sampler(sigma, sigma_next) * s_noise * sigma_up else: 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 class FunPackHybridEuler2SSampler: @classmethod def INPUT_TYPES(cls): return { "required": { "eta": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Ancestral stochasticity. Keep at 1.0 for classic ancestral behaviour." }), "s_noise": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Noise scale for ancestral noise injection." }), "high_quality_pct": ("FLOAT", { "default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Fraction of late denoise steps that receive the slower 2S quality correction." }), "correction_blend": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Blend between Euler ancestral (0.0) and late-step 2S correction (1.0)." }), "motion_pulse_mode": (MOTION_PULSE_MODES, { "default": "off", "tooltip": "Adds early/mid anti-stiffness motion pulses. Off preserves legacy sampler behavior." }), "motion_pulse_start_pct": ("FLOAT", { "default": 0.30, "min": 0.0, "max": 0.95, "step": 0.01, "tooltip": "Sampling progress where the first anti-stiffness pulse is applied." }), "motion_pulse_count": ("INT", { "default": 2, "min": 1, "max": 6, "step": 1, "tooltip": "How many early/mid motion pulses to request before the late quality phase." }), "motion_pulse_spacing_pct": ("FLOAT", { "default": 0.22, "min": 0.04, "max": 0.45, "step": 0.01, "tooltip": "Progress spacing between motion pulses." }), "motion_pulse_strength": ("FLOAT", { "default": 0.85, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "How strongly motion pulses add monotonic noise kicks. Higher values push harder against stale image references." }), "velocity_bias_mode": (VELOCITY_BIAS_MODES, { "default": "off", "tooltip": "Experimental: capture/apply averaged early model velocity around normalized sigma 0.9 and 0.8. Off preserves legacy sampler behavior." }), "velocity_bias_strength": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 0.35, "step": 0.01, "tooltip": "Experimental strength for applying captured early velocity bias. Keep low; 0 disables the applied delta." }), "velocity_refinement_key": ("STRING", { "default": "default", "multiline": False, "tooltip": "Memory key used to capture/apply early velocity bias." }), "velocity_aspect_bucket": ("STRING", { "default": "any", "multiline": False, "tooltip": "Optional aspect bucket such as landscape, portrait, square, ultrawide, or vertical." }), }, "optional": { "sigmas": ("SIGMAS",), } } 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 anti-stiffness motion pulses." def get_sampler(self, eta, s_noise, high_quality_pct, correction_blend, motion_pulse_mode="off", motion_pulse_start_pct=0.30, motion_pulse_count=2, motion_pulse_spacing_pct=0.22, motion_pulse_strength=0.85, velocity_bias_mode="off", velocity_bias_strength=0.0, velocity_refinement_key="default", velocity_aspect_bucket="any", sigmas=None): prepared_sigmas, quality_sigma_start, motion_pulse_steps, motion_pulse_noise = _prepare_dynamic_sigmas( sigmas, high_quality_pct, motion_pulse_mode, motion_pulse_start_pct, motion_pulse_count, motion_pulse_spacing_pct, motion_pulse_strength, ) sampler = comfy.samplers.KSAMPLER( sample_funpack_hybrid_euler_2s, extra_options={ "eta": eta, "s_noise": s_noise, "high_quality_pct": high_quality_pct, "correction_blend": correction_blend, "quality_sigma_start": quality_sigma_start, "motion_pulse_mode": motion_pulse_mode, "motion_pulse_start_pct": motion_pulse_start_pct, "motion_pulse_count": motion_pulse_count, "motion_pulse_spacing_pct": motion_pulse_spacing_pct, "motion_pulse_strength": motion_pulse_strength, "motion_pulse_noise": motion_pulse_noise, "motion_pulse_steps": motion_pulse_steps, "velocity_bias_mode": velocity_bias_mode, "velocity_bias_strength": velocity_bias_strength, "velocity_refinement_key": velocity_refinement_key, "velocity_aspect_bucket": velocity_aspect_bucket, } ) return (sampler, prepared_sigmas)