Upload 3 files
ec18254 verified 2 months ago
29.8 kB
import math
import torch

import comfy.model_sampling
import comfy.samplers
from comfy.utils import model_trange as trange

_EPS = 1.0e-6


def _get_model_sampling(model):
    try:
        return model.inner_model.model_patcher.get_model_object("model_sampling")
    except Exception as exc:
        raise RuntimeError(
            "ChebyFlow could not access ComfyUI model_sampling. "
            "Use the normal SamplerCustomAdvanced / KSAMPLER path."
        ) from exc



def _validate_flow_model(model):
    ms = _get_model_sampling(model)
    if not isinstance(ms, comfy.model_sampling.CONST):
        raise ValueError(
            "ChebyFlow requires a CONST/rectified-flow model. "
            f"Got {type(ms).__name__}."
        )
    return ms



def _sigma_list(sigmas):
    vals = [float(v) for v in sigmas.detach().float().cpu().tolist()]
    for i, s in enumerate(vals):
        if s < -1.0e-7 or s > 1.0 + 1.0e-5:
            raise ValueError(f"Flow sigma at index {i} is {s:.8g}, outside [0, 1].")
    for i in range(max(0, len(vals) - 1)):
        if vals[i + 1] > vals[i] + 1.0e-7:
            raise ValueError("Sigmas must be monotonically non-increasing.")
    return vals



def _flow_lambda(sigma):
    s = min(max(float(sigma), _EPS), 1.0 - _EPS)
    return math.log1p(-s) - math.log(s)



def _flow_sigma_from_lambda(lam):
    lam = float(lam)
    if lam >= 0.0:
        e = math.exp(-lam)
        return e / (1.0 + e)
    e = math.exp(lam)
    return 1.0 / (1.0 + e)



def _coordinate_value(sigma, coordinate):
    if coordinate == "sigma":
        return float(sigma)
    if coordinate == "flow_lambda":
        return _flow_lambda(sigma)
    raise ValueError(f"Unknown coordinate: {coordinate}")



def _sigma_from_coordinate(q, coordinate):
    if coordinate == "sigma":
        return min(max(float(q), 0.0), 1.0)
    if coordinate == "flow_lambda":
        return min(max(_flow_sigma_from_lambda(q), 0.0), 1.0)
    raise ValueError(f"Unknown coordinate: {coordinate}")



def _flow_velocity(x, denoised, sigma):
    if float(sigma) <= _EPS:
        return torch.zeros_like(x)
    return (x - denoised) / float(sigma)



def _flow_rhs(x, denoised, sigma, coordinate):
    v = _flow_velocity(x, denoised, sigma)
    if coordinate == "sigma":
        return v
    if coordinate == "flow_lambda":
        s = float(sigma)
        return (-s * (1.0 - s)) * v
    raise ValueError(f"Unknown coordinate: {coordinate}")



def _rms_scalar(t):
    return torch.sqrt(torch.mean(t.float() * t.float()) + 1.0e-20)



def _channel_reduce_dims(t):
    return tuple(i for i in range(t.ndim) if i != 1)



def _metric_components_global(a, b):
    af = a.float()
    bf = b.float()
    diff = bf - af

    rms_a = _rms_scalar(af)
    rms_b = _rms_scalar(bf)
    residual = _rms_scalar(diff) / torch.clamp(rms_a, min=1.0e-12)

    dot = torch.mean(af * bf)
    denom = torch.clamp(rms_a * rms_b, min=1.0e-12)
    cosine = torch.clamp(dot / denom, -1.0, 1.0)
    angle = 1.0 - cosine

    magnitude = torch.abs(torch.log(torch.clamp(rms_b, min=1.0e-12) / torch.clamp(rms_a, min=1.0e-12)))
    return residual, angle, magnitude



def _metric_components_channel(a, b):
    af = a.float()
    bf = b.float()
    dims = _channel_reduce_dims(af)
    diff = bf - af

    rms_a = torch.sqrt(torch.mean(af * af, dim=dims) + 1.0e-20)
    rms_b = torch.sqrt(torch.mean(bf * bf, dim=dims) + 1.0e-20)
    rms_d = torch.sqrt(torch.mean(diff * diff, dim=dims) + 1.0e-20)
    residual = rms_d / torch.clamp(rms_a, min=1.0e-12)

    dot = torch.mean(af * bf, dim=dims)
    denom = torch.clamp(rms_a * rms_b, min=1.0e-12)
    cosine = torch.clamp(dot / denom, -1.0, 1.0)
    angle = 1.0 - cosine

    magnitude = torch.abs(torch.log(torch.clamp(rms_b, min=1.0e-12) / torch.clamp(rms_a, min=1.0e-12)))
    return residual, angle, magnitude



def _combine_metric(residual, angle, magnitude, metric, residual_weight, angle_weight, magnitude_weight):
    if metric == "rms_delta":
        return residual
    if metric == "angle":
        return angle
    if metric == "magnitude":
        return magnitude
    if metric == "combined":
        wr = float(max(residual_weight, 0.0))
        wa = float(max(angle_weight, 0.0))
        wm = float(max(magnitude_weight, 0.0))
        denom = max(wr + wa + wm, 1.0e-12)
        return (wr * residual + wa * angle + wm * magnitude) / denom
    raise ValueError(f"Unknown curvature metric: {metric}")



def _metric_global(a, b, metric, rw, aw, mw):
    parts = _metric_components_global(a, b)
    return _combine_metric(*parts, metric, rw, aw, mw)



def _metric_channel(a, b, metric, rw, aw, mw):
    parts = _metric_components_channel(a, b)
    return _combine_metric(*parts, metric, rw, aw, mw)



def _score_from_metric(metric_value, scale):
    scale = float(max(scale, 1.0e-6))
    return metric_value / (metric_value + scale)



def _broadcast_channel(values, like):
    shape = [1] * like.ndim
    shape[1] = int(values.numel())
    return values.to(device=like.device, dtype=like.dtype).reshape(shape)



def _observer_init(value):
    v = value.detach().float()
    base_dev = torch.clamp(torch.abs(v) * 0.25 + 0.05, min=0.05)
    return {
        "slow": v.clone(),
        "fast": v.clone(),
        "dev": base_dev.clone(),
    }



def _observer_update(obs, value, slow_decay, fast_decay, dev_decay):
    v = value.detach().float()
    if obs is None:
        return _observer_init(v)
    slow = slow_decay * obs["slow"] + (1.0 - slow_decay) * v
    fast = fast_decay * obs["fast"] + (1.0 - fast_decay) * v
    dev = dev_decay * obs["dev"] + (1.0 - dev_decay) * torch.abs(v - slow)
    dev = torch.clamp(dev, min=0.01)
    return {"slow": slow, "fast": fast, "dev": dev}



def _observer_score(obs, value, current_weight, trend_weight, temperature):
    if obs is None:
        return torch.full_like(value.detach().float(), 0.5)
    v = value.detach().float()
    slow = obs["slow"]
    fast = obs["fast"]
    dev = torch.clamp(obs["dev"], min=0.01)
    z_current = (v - slow) / dev
    z_trend = (fast - slow) / dev
    z = current_weight * z_current + trend_weight * z_trend
    return torch.sigmoid(z / float(max(temperature, 1.0e-4)))



def _weighted_score(scores, weights):
    num = None
    denom = 0.0
    for score, weight in zip(scores, weights):
        w = float(max(weight, 0.0))
        if w == 0.0:
            continue
        term = score * w
        num = term if num is None else num + term
        denom += w
    if num is None or denom <= 0.0:
        ref = scores[0]
        return torch.full_like(ref, 0.5)
    return num / denom



def _centered_response(score):
    return 2.0 * torch.clamp(score, 0.0, 1.0) - 1.0



def _stage_from_difficulty(base_c, min_c, max_c, difficulty, adaptivity, stage_gain=1.0):
    shift = float(torch.clamp(_centered_response(difficulty) * float(stage_gain), -1.0, 1.0).item())
    if shift >= 0.0:
        target = base_c + shift * (max_c - base_c)
    else:
        target = base_c + shift * (base_c - min_c)
    c = (1.0 - adaptivity) * base_c + adaptivity * target
    return min(max(float(c), min_c), max_c)



def _correction_strength_from_difficulty(difficulty, low, high):
    low = float(low)
    high = float(high)
    shift = torch.clamp(_centered_response(difficulty), -1.0, 1.0)
    hard = torch.clamp(shift, min=0.0)
    easy = torch.clamp(-shift, min=0.0)
    out = torch.ones_like(shift)
    out = out + hard * (low - 1.0)
    out = out + easy * (high - 1.0)
    return out



def _auto_profile(name):
    name = str(name)
    if name == "detail":
        return {
            "weights": (0.50, 0.35, 0.15),
            "slow_decay": 0.88,
            "fast_decay": 0.42,
            "dev_decay": 0.84,
            "current_weight": 0.70,
            "trend_weight": 0.30,
            "temperature": 1.05,
            "stage_gain": 1.10,
            "corr_low": 0.92,
            "corr_high": 1.10,
            "warmup_steps": 2,
            "soft_start_steps": 0,
        }
    if name == "smooth":
        return {
            "weights": (0.40, 0.45, 0.15),
            "slow_decay": 0.93,
            "fast_decay": 0.34,
            "dev_decay": 0.89,
            "current_weight": 0.62,
            "trend_weight": 0.38,
            "temperature": 1.35,
            "stage_gain": 0.88,
            "corr_low": 0.94,
            "corr_high": 1.04,
            "warmup_steps": 2,
            "soft_start_steps": 0,
        }
    if name == "stable":
        return {
            "weights": (0.42, 0.43, 0.15),
            "slow_decay": 0.94,
            "fast_decay": 0.32,
            "dev_decay": 0.90,
            "current_weight": 0.60,
            "trend_weight": 0.40,
            "temperature": 1.45,
            "stage_gain": 0.78,
            "corr_low": 0.95,
            "corr_high": 1.03,
            "warmup_steps": 2,
            "soft_start_steps": 0,
        }
    if name == "dynamic":
        return {
            "weights": (0.35, 0.50, 0.15),
            "slow_decay": 0.87,
            "fast_decay": 0.38,
            "dev_decay": 0.83,
            "current_weight": 0.67,
            "trend_weight": 0.33,
            "temperature": 1.00,
            "stage_gain": 1.18,
            "corr_low": 0.91,
            "corr_high": 1.09,
            "warmup_steps": 2,
            "soft_start_steps": 0,
        }
    # balanced
    return {
        "weights": (0.45, 0.40, 0.15),
        "slow_decay": 0.90,
        "fast_decay": 0.40,
        "dev_decay": 0.86,
        "current_weight": 0.68,
        "trend_weight": 0.32,
        "temperature": 1.18,
        "stage_gain": 1.00,
        "corr_low": 0.90,
        "corr_high": 1.08,
        "warmup_steps": 2,
        "soft_start_steps": 0,
    }


def _apply_turbo_profile(profile):
    """Return a low-step controller variant without changing the preset's character."""
    p = dict(profile)
    # Faster online calibration: fewer historical steps exist, so both baselines must
    # move sooner. Keep slow > fast so trend information remains meaningful.
    p["slow_decay"] = max(0.68, float(p["slow_decay"]) - 0.15)
    p["fast_decay"] = max(0.16, float(p["fast_decay"]) - 0.15)
    p["dev_decay"] = max(0.68, float(p["dev_decay"]) - 0.12)

    # A little more weight on the current interval and a lower temperature make the
    # normalized controller useful after only one baseline interval.
    p["current_weight"] = min(0.78, float(p["current_weight"]) + 0.05)
    p["trend_weight"] = max(0.0, 1.0 - p["current_weight"])
    p["temperature"] = max(0.75, float(p["temperature"]) * 0.78)
    p["stage_gain"] = float(p["stage_gain"]) * 1.08

    # Low-step intervals are large: move the probe decisively, but keep algebraic
    # correction closer to 1.0 so one interval cannot redirect the whole trajectory.
    p["corr_low"] = 1.0 + (float(p["corr_low"]) - 1.0) * 0.50
    p["corr_high"] = 1.0 + (float(p["corr_high"]) - 1.0) * 0.50

    p["warmup_steps"] = 1
    p["soft_start_steps"] = 1
    p["stage_base_shift"] = 0.05
    p["stage_min_shift"] = 0.07
    p["stage_max_shift"] = 0.01
    return p


def _controller_confidence(step_index, warmup_steps, soft_start_steps):
    """0 during warmup, then a short confidence ramp before full adaptation."""
    i = int(step_index)
    warm = int(max(warmup_steps, 0))
    soft = int(max(soft_start_steps, 0))
    if i < warm:
        return 0.0
    if soft <= 0:
        return 1.0
    j = i - warm
    if j < soft:
        return float(j + 1) / float(soft + 1)
    return 1.0


@torch.no_grad()
def sample_h3_adaptive_chebyshev_advanced(
    model,
    x,
    sigmas,
    extra_args=None,
    callback=None,
    disable=None,
    coordinate="flow_lambda",
    base_stage_position=0.70,
    min_stage_position=0.55,
    max_stage_position=0.95,
    adapt_strength=1.0,
    curvature_smoothing=0.50,
    stage_curvature_scale=0.12,
    curvature_metric="combined",
    residual_weight=1.0,
    angle_weight=1.0,
    magnitude_weight=0.5,
    correction_mode="curvature_down",
    correction_granularity="per_channel",
    correction_min=0.88,
    correction_max=1.08,
    correction_curvature_scale=0.12,
):
    """
    Advanced adaptive 2-stage Chebyshev/GS sampler for H3 CONST flow models.
    This is the manually tunable version. Defaults are intentionally set to a strong
    "all three improvements enabled" configuration:
      - combined curvature metric,
      - curvature-dependent correction damping,
      - per-channel correction strengths.
    """
    extra_args = {} if extra_args is None else extra_args
    _validate_flow_model(model)
    sigma_vals = _sigma_list(sigmas)
    if len(sigma_vals) <= 1:
        return x

    base_c = min(max(float(base_stage_position), 0.05), 0.99)
    min_c = min(max(float(min_stage_position), 0.05), 0.99)
    max_c = min(max(float(max_stage_position), min_c + 1.0e-4), 0.99)
    adapt = float(max(adapt_strength, 0.0))
    smooth = float(min(max(curvature_smoothing, 0.0), 0.99))
    stage_scale = float(max(stage_curvature_scale, 1.0e-6))

    corr_min = float(min(correction_min, correction_max))
    corr_max = float(max(correction_min, correction_max))
    corr_scale = float(max(correction_curvature_scale, 1.0e-6))

    c_state = base_c
    global_metric_ema = None
    channel_metric_ema = None
    s_in = x.new_ones([x.shape[0]])

    for i in trange(len(sigmas) - 1, disable=disable):
        sigma_s = sigma_vals[i]
        sigma_t = sigma_vals[i + 1]
        denoised_s = model(x, sigmas[i] * s_in, **extra_args)

        if callback is not None:
            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised_s})

        if sigma_t <= _EPS:
            x = denoised_s
            break

        c = min(max(float(c_state), min_c), max_c)
        b_stage = 1.0 / (2.0 * c)
        b_start = 1.0 - b_stage

        q_s = _coordinate_value(sigma_s, coordinate)
        q_t = _coordinate_value(sigma_t, coordinate)
        h = q_t - q_s
        if abs(h) < 1.0e-12:
            continue

        rhs_s = _flow_rhs(x, denoised_s, sigma_s, coordinate)
        q_c = q_s + c * h
        sigma_c = _sigma_from_coordinate(q_c, coordinate)
        x_c = x + (c * h) * rhs_s
        denoised_c = model(x_c, x.new_full((x.shape[0],), sigma_c), **extra_args)
        rhs_c = _flow_rhs(x_c, denoised_c, sigma_c, coordinate)

        x_first = x + h * rhs_s
        x_cheb = x + h * (b_start * rhs_s + b_stage * rhs_c)
        correction = x_cheb - x_first

        g_metric = _metric_global(rhs_s, rhs_c, curvature_metric, residual_weight, angle_weight, magnitude_weight)
        global_metric_ema = g_metric if global_metric_ema is None else (
            smooth * global_metric_ema + (1.0 - smooth) * g_metric
        )
        g_score = _score_from_metric(global_metric_ema, stage_scale)
        target_c = min_c + (max_c - min_c) * float(g_score.item())
        c_state = (1.0 - adapt) * base_c + adapt * target_c
        c_state = min(max(c_state, min_c), max_c)

        if correction_mode == "constant":
            strength = corr_max
        else:
            if correction_granularity == "per_channel":
                c_metric = _metric_channel(rhs_s, rhs_c, curvature_metric, residual_weight, angle_weight, magnitude_weight)
                channel_metric_ema = c_metric if channel_metric_ema is None else (
                    smooth * channel_metric_ema + (1.0 - smooth) * c_metric
                )
                score = _score_from_metric(channel_metric_ema, corr_scale)
                if correction_mode == "curvature_down":
                    strength = corr_max - (corr_max - corr_min) * score
                elif correction_mode == "curvature_up":
                    strength = corr_min + (corr_max - corr_min) * score
                else:
                    raise ValueError(f"Unknown correction mode: {correction_mode}")
                strength = _broadcast_channel(strength, correction)
            else:
                score = _score_from_metric(global_metric_ema, corr_scale)
                if correction_mode == "curvature_down":
                    strength = corr_max - (corr_max - corr_min) * score
                elif correction_mode == "curvature_up":
                    strength = corr_min + (corr_max - corr_min) * score
                else:
                    raise ValueError(f"Unknown correction mode: {correction_mode}")
                strength = float(strength.item())

        x = x_first + correction * strength

    return x


@torch.no_grad()
def sample_h3_adaptive_chebyshev_auto(
    model,
    x,
    sigmas,
    extra_args=None,
    callback=None,
    disable=None,
    coordinate="flow_lambda",
    base_stage_position=0.72,
    min_stage_position=0.58,
    max_stage_position=0.96,
    adaptivity=1.0,
    quality_bias="balanced",
    turbo_mode=False,
    correction_granularity="per_channel",
):
    """
    Self-calibrating adaptive Chebyshev sampler.
    No extra model evaluations are added. The sampler self-calibrates online during the
    current generation using only the two evaluations it already performs per interval.
    turbo_mode is an orthogonal numerical regime, not a quality preset. It keeps the
    selected quality_bias character but accelerates observer dynamics, shortens startup,
    biases the internal stage later, and narrows correction excursions for large low-step
    intervals such as ~6-10 step distilled/Turbo-LoRA workflows.
    """
    extra_args = {} if extra_args is None else extra_args
    _validate_flow_model(model)
    sigma_vals = _sigma_list(sigmas)
    if len(sigma_vals) <= 1:
        return x

    profile = _auto_profile(quality_bias)
    if bool(turbo_mode):
        profile = _apply_turbo_profile(profile)

    # User values remain the baseline. Turbo mode shifts the operating envelope later
    # without requiring separate turbo presets. Keep the internal stage strictly below 1.
    base_c = float(base_stage_position)
    min_c = float(min_stage_position)
    max_c = float(max_stage_position)
    if bool(turbo_mode):
        base_c += float(profile.get("stage_base_shift", 0.0))
        min_c += float(profile.get("stage_min_shift", 0.0))
        max_c += float(profile.get("stage_max_shift", 0.0))

    base_c = min(max(base_c, 0.05), 0.99)
    min_c = min(max(min_c, 0.05), 0.99)
    max_c = min(max(max_c, min_c + 1.0e-4), 0.99)
    base_c = min(max(base_c, min_c), max_c)

    adaptivity = float(max(adaptivity, 0.0))
    weights = profile["weights"]
    warmup_steps = int(profile["warmup_steps"])
    soft_start_steps = int(profile.get("soft_start_steps", 0))

    obs_g_res = obs_g_ang = obs_g_mag = None
    obs_c_res = obs_c_ang = obs_c_mag = None
    c_state = base_c
    s_in = x.new_ones([x.shape[0]])

    for i in trange(len(sigmas) - 1, disable=disable):
        sigma_s = sigma_vals[i]
        sigma_t = sigma_vals[i + 1]
        denoised_s = model(x, sigmas[i] * s_in, **extra_args)

        if callback is not None:
            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised_s})

        if sigma_t <= _EPS:
            x = denoised_s
            break

        confidence = _controller_confidence(i, warmup_steps, soft_start_steps)
        c = base_c if confidence <= 0.0 else c_state
        c = min(max(float(c), min_c), max_c)
        b_stage = 1.0 / (2.0 * c)
        b_start = 1.0 - b_stage

        q_s = _coordinate_value(sigma_s, coordinate)
        q_t = _coordinate_value(sigma_t, coordinate)
        h = q_t - q_s
        if abs(h) < 1.0e-12:
            continue

        rhs_s = _flow_rhs(x, denoised_s, sigma_s, coordinate)
        q_c = q_s + c * h
        sigma_c = _sigma_from_coordinate(q_c, coordinate)
        x_c = x + (c * h) * rhs_s
        denoised_c = model(x_c, x.new_full((x.shape[0],), sigma_c), **extra_args)
        rhs_c = _flow_rhs(x_c, denoised_c, sigma_c, coordinate)

        x_first = x + h * rhs_s
        x_cheb = x + h * (b_start * rhs_s + b_stage * rhs_c)
        correction = x_cheb - x_first

        # Current raw local difficulty components.
        g_res, g_ang, g_mag = _metric_components_global(rhs_s, rhs_c)

        if obs_g_res is None:
            g_difficulty = torch.full_like(g_res, 0.5)
        else:
            g_res_s = _observer_score(obs_g_res, g_res, profile["current_weight"], profile["trend_weight"], profile["temperature"])
            g_ang_s = _observer_score(obs_g_ang, g_ang, profile["current_weight"], profile["trend_weight"], profile["temperature"])
            g_mag_s = _observer_score(obs_g_mag, g_mag, profile["current_weight"], profile["trend_weight"], profile["temperature"])
            g_difficulty = _weighted_score((g_res_s, g_ang_s, g_mag_s), weights)

        # Global stage controller for NEXT interval. During turbo soft-start, the
        # confidence factor ramps this from fixed-base behavior to full adaptation.
        c_state = _stage_from_difficulty(
            base_c,
            min_c,
            max_c,
            g_difficulty,
            adaptivity * confidence,
            profile["stage_gain"],
        )

        c_res = c_ang = c_mag = None
        if correction_granularity == "per_channel":
            c_res, c_ang, c_mag = _metric_components_channel(rhs_s, rhs_c)

        # Correction controller for CURRENT interval. The automatic strength is blended
        # toward 1.0 by controller confidence, giving Turbo mode a 50% soft-start step.
        if confidence <= 0.0:
            strength = 1.0
        else:
            if correction_granularity == "per_channel":
                if obs_c_res is None:
                    c_difficulty = torch.full_like(c_res, 0.5)
                else:
                    c_res_s = _observer_score(obs_c_res, c_res, profile["current_weight"], profile["trend_weight"], profile["temperature"])
                    c_ang_s = _observer_score(obs_c_ang, c_ang, profile["current_weight"], profile["trend_weight"], profile["temperature"])
                    c_mag_s = _observer_score(obs_c_mag, c_mag, profile["current_weight"], profile["trend_weight"], profile["temperature"])
                    c_difficulty = _weighted_score((c_res_s, c_ang_s, c_mag_s), weights)
                auto_strength = _correction_strength_from_difficulty(
                    c_difficulty, profile["corr_low"], profile["corr_high"]
                )
                auto_strength = 1.0 + confidence * (auto_strength - 1.0)
                strength = _broadcast_channel(auto_strength, correction)
            else:
                auto_strength = _correction_strength_from_difficulty(
                    g_difficulty, profile["corr_low"], profile["corr_high"]
                )
                auto_strength = 1.0 + confidence * (auto_strength - 1.0)
                strength = float(auto_strength.item())

        x = x_first + correction * strength

        # Update observers AFTER using them: current difficulty is judged against the
        # previously learned baseline, rather than being normalized away immediately.
        obs_g_res = _observer_update(obs_g_res, g_res, profile["slow_decay"], profile["fast_decay"], profile["dev_decay"])
        obs_g_ang = _observer_update(obs_g_ang, g_ang, profile["slow_decay"], profile["fast_decay"], profile["dev_decay"])
        obs_g_mag = _observer_update(obs_g_mag, g_mag, profile["slow_decay"], profile["fast_decay"], profile["dev_decay"])

        if correction_granularity == "per_channel":
            obs_c_res = _observer_update(obs_c_res, c_res, profile["slow_decay"], profile["fast_decay"], profile["dev_decay"])
            obs_c_ang = _observer_update(obs_c_ang, c_ang, profile["slow_decay"], profile["fast_decay"], profile["dev_decay"])
            obs_c_mag = _observer_update(obs_c_mag, c_mag, profile["slow_decay"], profile["fast_decay"], profile["dev_decay"])

    return x


class H3AdaptiveChebyshevAdvancedSampler:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "coordinate": (["flow_lambda", "sigma"], {"default": "flow_lambda"}),
                "base_stage_position": ("FLOAT", {"default": 0.70, "min": 0.05, "max": 0.99, "step": 0.01}),
                "min_stage_position": ("FLOAT", {"default": 0.55, "min": 0.05, "max": 0.99, "step": 0.01}),
                "max_stage_position": ("FLOAT", {"default": 0.95, "min": 0.05, "max": 0.99, "step": 0.01}),
                "adapt_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.5, "step": 0.05}),
                "curvature_smoothing": ("FLOAT", {"default": 0.50, "min": 0.0, "max": 0.95, "step": 0.05}),
                "stage_curvature_scale": ("FLOAT", {"default": 0.12, "min": 0.01, "max": 2.0, "step": 0.01}),
                "curvature_metric": (["rms_delta", "angle", "magnitude", "combined"], {"default": "combined"}),
                "residual_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 4.0, "step": 0.05}),
                "angle_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 4.0, "step": 0.05}),
                "magnitude_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 4.0, "step": 0.05}),
                "correction_mode": (["constant", "curvature_down", "curvature_up"], {"default": "curvature_down"}),
                "correction_granularity": (["global", "per_channel"], {"default": "per_channel"}),
                "correction_min": ("FLOAT", {"default": 0.88, "min": 0.0, "max": 1.5, "step": 0.01}),
                "correction_max": ("FLOAT", {"default": 1.08, "min": 0.0, "max": 1.5, "step": 0.01}),
                "correction_curvature_scale": ("FLOAT", {"default": 0.12, "min": 0.01, "max": 2.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("SAMPLER",)
    RETURN_NAMES = ("sampler",)
    FUNCTION = "get_sampler"
    CATEGORY = "sampling/custom_sampling/samplers"

    def get_sampler(
        self,
        coordinate,
        base_stage_position,
        min_stage_position,
        max_stage_position,
        adapt_strength,
        curvature_smoothing,
        stage_curvature_scale,
        curvature_metric,
        residual_weight,
        angle_weight,
        magnitude_weight,
        correction_mode,
        correction_granularity,
        correction_min,
        correction_max,
        correction_curvature_scale,
    ):
        sampler = comfy.samplers.KSAMPLER(
            sample_h3_adaptive_chebyshev_advanced,
            extra_options={
                "coordinate": coordinate,
                "base_stage_position": base_stage_position,
                "min_stage_position": min_stage_position,
                "max_stage_position": max_stage_position,
                "adapt_strength": adapt_strength,
                "curvature_smoothing": curvature_smoothing,
                "stage_curvature_scale": stage_curvature_scale,
                "curvature_metric": curvature_metric,
                "residual_weight": residual_weight,
                "angle_weight": angle_weight,
                "magnitude_weight": magnitude_weight,
                "correction_mode": correction_mode,
                "correction_granularity": correction_granularity,
                "correction_min": correction_min,
                "correction_max": correction_max,
                "correction_curvature_scale": correction_curvature_scale,
            },
        )
        return (sampler,)


class H3AdaptiveChebyshevAutoSampler:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "coordinate": (["flow_lambda", "sigma"], {"default": "flow_lambda"}),
                "base_stage_position": ("FLOAT", {"default": 0.72, "min": 0.05, "max": 0.99, "step": 0.01}),
                "min_stage_position": ("FLOAT", {"default": 0.58, "min": 0.05, "max": 0.99, "step": 0.01}),
                "max_stage_position": ("FLOAT", {"default": 0.96, "min": 0.05, "max": 0.99, "step": 0.01}),
                "adaptivity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.5, "step": 0.05}),
                "quality_bias": (["balanced", "detail", "smooth", "stable", "dynamic"], {"default": "balanced"}),
                "turbo_mode": ("BOOLEAN", {"default": False}),
                "correction_granularity": (["global", "per_channel"], {"default": "per_channel"}),
            }
        }

    RETURN_TYPES = ("SAMPLER",)
    RETURN_NAMES = ("sampler",)
    FUNCTION = "get_sampler"
    CATEGORY = "sampling/custom_sampling/samplers"

    def get_sampler(self, coordinate, base_stage_position, min_stage_position, max_stage_position, adaptivity, quality_bias, turbo_mode, correction_granularity):
        sampler = comfy.samplers.KSAMPLER(
            sample_h3_adaptive_chebyshev_auto,
            extra_options={
                "coordinate": coordinate,
                "base_stage_position": base_stage_position,
                "min_stage_position": min_stage_position,
                "max_stage_position": max_stage_position,
                "adaptivity": adaptivity,
                "quality_bias": quality_bias,
                "turbo_mode": turbo_mode,
                "correction_granularity": correction_granularity,
            },
        )
        return (sampler,)


# Backward-compatible alias for earlier pack versions.
H3AdaptiveChebyshevProSampler = H3AdaptiveChebyshevAdvancedSampler


NODE_CLASS_MAPPINGS = {
    "H3AdaptiveChebyshevAdvancedSampler": H3AdaptiveChebyshevAdvancedSampler,
    "H3AdaptiveChebyshevAutoSampler": H3AdaptiveChebyshevAutoSampler,
    "H3AdaptiveChebyshevProSampler": H3AdaptiveChebyshevAdvancedSampler,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "H3AdaptiveChebyshevAdvancedSampler": "ChebyFlow Advanced Sampler",
    "H3AdaptiveChebyshevAutoSampler": "ChebyFlow Auto Sampler",
    "H3AdaptiveChebyshevProSampler": "ChebyFlow Advanced Sampler",
}