adept-sampler-v6/scripts/adept_sampler_v5.py

"""
Adept Sampler v5 for Automatic1111 WebUI
Complete port with ALL custom samplers from ComfyUI

Version: 5.0
"""

import torch
import numpy as np
import math
import inspect
import functools
import os
from tqdm import trange
from modules import scripts, shared, script_callbacks
import gradio as gr
import k_diffusion.sampling

# Try import torchvision for detail enhancement
try:
    from torchvision.transforms.functional import gaussian_blur
    TORCHVISION_AVAILABLE = True
except ImportError:
    TORCHVISION_AVAILABLE = False
    print("⚠️ torchvision not available - detail enhancement disabled")

# ============================================================================
# GLOBAL STATE
# ============================================================================
ADEPT_STATE = {
    "enabled": False,
    "scale": 1.0,
    "shift": 0.0,
    "start_pct": 0.0,
    "end_pct": 1.0,
    "eta": 1.0,
    "s_noise": 1.0,
    "adaptive_eta": False,
    "scheduler": "Standard",
    "vae_reflection": False,
    
    # Custom sampler settings
    "use_custom_sampler": False,
    "custom_sampler": "Akashic Solver v2",
    "adaptive_noise_scale": False,
    "tau": 0.5,
    "phase_strength": 0.5,
    "solver_order": 2,
    "use_corrector": True,
    "phase_noise": False,
    "enhanced_derivative": False,
    "smea_strength": 0.0,
    "ndb_strength": 0.0,
    "eqvae_mode": "Off",

    # Mirror Correction Euler controls
    "mirror_correction_phase": 0.5,
    "mirror_smooth_phase": False,

    # CFG enhancement settings
    "cfg_drift_enabled": False,
    "cfg_drift_method": "mean",
    "cfg_drift_intensity": 0.5,
    "spectral_cfg_enabled": False,
    "spectral_multiplier": 1.0,
    "spectral_percentile": 5.0,
    "phase_cfg_enabled": False,
    "phase_cfg_alpha": 2.0,
    "phase_cfg_beta": 2.0,
    "cfg_runtime_mode": "off",   # off | a1111-postcfg | a1111-monkeypatch | native-hook

    # Internal bookkeeping for phase-aware CFG progress tracking
    "_cfg_step_idx": 0,
    "_cfg_total_steps": 1,

    # Content-Aware Pacing (Euler Ancestral only): auto composition->detail
    # phase switch based on derivative variance, plus optional high-frequency
    # detail enhancement layered on top of either phase.
    "use_content_aware_pacing": False,
    "pacing_coherence_sensitivity": 0.75,
    "use_detail_enhancement_phase": False,
    "detail_enhancement_strength": 0.05,
    "detail_separation_radius": 0.5,
}

# Store original samplers
ORIGINAL_SAMPLERS = {}
# key -> wrapper function currently installed by us (specialized OR generic).
# Needed so unpatch can tell "still ours" from "someone else wrapped us" for
# the full dynamic sampler set, not just the 6 that used to be hardcoded.
_ADEPT_ACTIVE_FUNCS = {}

# VAE Reflection state
_vae_reflection_active = False
_vae_original_padding_modes = {}

# CFG hook / callback runtime state
_ADEPT_CFG_AFTER_CB    = None
_ADEPT_CFG_DENOISER_CB = None
_ADEPT_NATIVE_CFG_HOOK_ACTIVE = False

# CFGDenoiser monkey-patch runtime state
_CFGD_ORIG_COMBINE       = None
_CFGD_ORIG_COMBINE_EDIT  = None
_CFGD_ORIG_FORWARD       = None
_CFGD_MONKEYPATCH_ACTIVE = False
_ADEPT_CFGDENOISER_CTX_ATTR = "_adept_cfg_ctx"

# ============================================================================
# BASIC UTILITY FUNCTIONS (from v3)
# ============================================================================

def to_d(x, sigma, denoised):
    """Convert denoised prediction to derivative."""
    diff = x - denoised
    safe_sigma = torch.clamp(sigma, min=1e-4)
    derivative = diff / safe_sigma
    
    sigma_adaptive_threshold = 1000.0 * (1.0 + sigma / 10.0)
    derivative_max = torch.abs(derivative).max()
    if derivative_max > sigma_adaptive_threshold:
        derivative = torch.clamp(derivative, -sigma_adaptive_threshold, sigma_adaptive_threshold)
    
    return derivative


def get_ancestral_step(sigma, sigma_next, eta=1.0):
    """Calculate ancestral step sizes."""
    if sigma_next == 0:
        return 0.0, 0.0
    sigma_up = min(sigma_next, eta * (sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2) ** 0.5)
    sigma_down = (sigma_next ** 2 - sigma_up ** 2) ** 0.5
    return sigma_down, sigma_up


def compute_dynamic_scale(step_idx, total_steps, base_scale, start_pct, end_pct):
    """
    Compute weight scale for the current step with smooth fade-in/fade-out.

    The fade ramp is proportional to the active window width rather than a
    fixed 0.1 absolute value.  For a narrow window (e.g. start=0.4, end=0.5)
    a hard-coded 0.1 ramp would consume the entire window and produce jarring
    or contradictory behaviour; normalising to 20 % of window width keeps the
    envelope sensible at any window size.

    Returns 1.0 (no-op) outside [start_pct, end_pct].
    """
    # Clamp / validate inputs so callers can pass raw UI values safely.
    start_pct  = max(0.0, min(float(start_pct), 1.0))
    end_pct    = max(start_pct, min(float(end_pct), 1.0))
    total_steps = max(total_steps, 1)
    progress    = step_idx / max(total_steps - 1, 1)

    if progress < start_pct or progress > end_pct:
        return 1.0

    window = end_pct - start_pct
    if window < 1e-6:
        # Degenerate window — treat as fully active for that single step.
        return float(base_scale)

    # Fade ramp = 20 % of window, capped at 0.05 so it never feels sluggish
    # on very wide windows and never feels jarring on narrow ones.
    ramp = min(0.20 * window, 0.05)

    if ramp > 0 and progress < start_pct + ramp:
        fade = (progress - start_pct) / ramp
        return 1.0 + (base_scale - 1.0) * fade
    elif ramp > 0 and progress > end_pct - ramp:
        fade = (end_pct - progress) / ramp
        return 1.0 + (base_scale - 1.0) * fade
    else:
        return float(base_scale)


def default_noise_sampler(x):
    """Simple noise sampler fallback."""
    def sampler(sigma, sigma_next):
        return torch.randn_like(x)
    return sampler


def get_noise_sampler(x):
    """Get noise sampler for the given tensor."""
    return default_noise_sampler(x)


# ============================================================================
# ADVANCED UTILITY FUNCTIONS
# ============================================================================

def to_d_enhanced_ancestral(x, sigma, denoised, eta, progress, generator=None):
    """Enhanced derivative for ancestral sampling."""
    diff = x - denoised
    safe_sigma = torch.clamp(sigma, min=1e-4)
    base_derivative = diff / safe_sigma

    def safe_randn_like(tensor, generator=None):
        if generator is None:
            return torch.randn_like(tensor)
        try:
            return torch.randn(tensor.shape, device=tensor.device, dtype=tensor.dtype, generator=generator)
        except (TypeError, AttributeError):
            return torch.randn_like(tensor)

    if eta > 1.0:
        eta_correction = 0.02 * (eta - 1.0) * safe_randn_like(diff, generator) * progress
        base_derivative = base_derivative + eta_correction
    elif eta < 1.0:
        eta_correction = 0.015 * (1.0 - eta) * safe_randn_like(diff, generator) * (1.0 - progress)
        base_derivative = base_derivative - eta_correction

    if progress < 0.3:
        phase_correction = 0.01 * safe_randn_like(diff, generator)
        base_derivative = base_derivative + phase_correction
    elif progress > 0.7:
        phase_correction = 0.008 * safe_randn_like(diff, generator)
        base_derivative = base_derivative - phase_correction

    sigma_adaptive_threshold = 500.0 * (1.0 + sigma / 10.0)
    derivative_max = torch.abs(base_derivative).max()
    if derivative_max > sigma_adaptive_threshold:
        base_derivative = torch.clamp(base_derivative, -sigma_adaptive_threshold, sigma_adaptive_threshold)

    return base_derivative


def apply_dynamic_thresholding(x, percentile=0.995, clamp_range=1.0):
    """Dynamic thresholding for high CFG."""
    if percentile >= 1.0:
        return x
    
    try:
        batch_size = x.shape[0]
        x_flat = x.view(batch_size, -1)
        
        abs_max = torch.abs(x_flat).max(dim=1, keepdim=True)[0]
        if abs_max.max() < 5.0:
            return x
        
        k = max(1, int(x_flat.shape[1] * (1.0 - percentile)))
        topk_vals = torch.topk(torch.abs(x_flat), k=k, dim=1, largest=True)[0]
        s = topk_vals[:, -1:].clamp(min=1.0)
        
        threshold = s * 2.5
        mask = torch.abs(x_flat) > threshold
        x_flat = torch.where(mask, torch.sign(x_flat) * threshold, x_flat)
        x_flat = x_flat * 0.98
        
        return x_flat.view(x.shape)
    except Exception:
        return x


def compute_compensation_ratio(r, step_idx, total_steps, base_ratio=1.0):
    """DC-Solver compensation."""
    progress = step_idx / max(total_steps - 1, 1)
    if progress < 0.3:
        phase_weight = 1.5
    elif progress < 0.7:
        phase_weight = 1.0
    else:
        phase_weight = 1.3
    return base_ratio * phase_weight * (1.0 + 0.1 * math.tanh(r - 1.0))


def compute_tau_eqvae(progress, base_tau=0.5, phase_strength=0.5):
    """Phase-aware tau for standard VAE."""
    if progress < 0.30:
        phase_factor = 1.0 + 0.2 * phase_strength
    elif progress < 0.60:
        phase_factor = 1.0 - 0.15 * phase_strength
    else:
        phase_factor = 1.0 - 0.3 * phase_strength
    return min(1.0, max(0.0, base_tau * phase_factor))


def compute_eqvae_tau(progress, base_tau, phase_strength):
    """EQ-VAE tau with shifted phases."""
    if progress < 0.25:
        phase_factor = 1.0 + 0.10 * phase_strength
    elif progress < 0.55:
        phase_factor = 1.0 - 0.10 * phase_strength
    else:
        phase_factor = 1.0 - 0.20 * phase_strength
    return min(1.0, max(0.0, base_tau * phase_factor))


def compute_eqvae_noise_scale(base_s_noise, progress):
    """EQ-VAE noise scale."""
    eqvae_base_factor = 0.88
    if progress < 0.25:
        phase_factor = 1.0 + 0.05 * (1.0 - progress / 0.25)
    elif progress < 0.60:
        phase_factor = 1.0 - 0.05 * ((progress - 0.25) / 0.35)
    else:
        phase_factor = 0.95
    return base_s_noise * eqvae_base_factor * phase_factor


def compute_eqvae_ndb(progress, ndb_strength):
    """Native Detail Boost for EQ-VAE."""
    if ndb_strength <= 0:
        return 0.5, 0.0
    
    blur_sigma = 0.6
    if progress < 0.30:
        phase_progress = progress / 0.30
        high_freq_boost = 0.03 * ndb_strength * phase_progress
    elif progress < 0.60:
        phase_progress = (progress - 0.30) / 0.30
        high_freq_boost = (0.03 + 0.07 * phase_progress) * ndb_strength
    else:
        phase_progress = (progress - 0.60) / 0.40
        high_freq_boost = (0.10 + 0.10 * phase_progress) * ndb_strength
    return blur_sigma, high_freq_boost


def compute_native_detail_boost(progress, ndb_strength=0.0):
    """Native Detail Boost for standard VAE."""
    if ndb_strength <= 0:
        return 1.0, 0.0
    
    if progress < 0.30:
        phase_progress = progress / 0.30
        high_freq_boost = 0.03 * ndb_strength * phase_progress
    elif progress < 0.60:
        phase_progress = (progress - 0.30) / 0.30
        high_freq_boost = (0.03 + 0.07 * phase_progress) * ndb_strength
    else:
        phase_progress = (progress - 0.60) / 0.40
        high_freq_boost = (0.10 + 0.08 * phase_progress) * ndb_strength
    return 1.0, high_freq_boost


def compute_smea_factor(progress, smea_strength=0.5):
    """SMEA coherency."""
    if smea_strength <= 0:
        return 1.0
    smea_interp = 0.5 * (1 + math.sin(math.pi * (progress - 0.5)))
    return 1.0 - smea_strength * (1.0 - smea_interp)




# ============================================================================
# ADVANCED CFG TECHNIQUES from reForge
# ============================================================================

def apply_spectral_modulation_clybius(noise_pred, multiplier=1.0, percentile=5.0):
    """
    Clybius Spectral Modulation: Apply frequency-domain corrections to noise prediction.
    
    This is the correct implementation based on ComfyUI-Latent-Modifiers.
    It should be applied to noise_pred (cond - uncond), NOT to denoised latent.
    
    Args:
        noise_pred: The noise prediction tensor (cond - uncond)
        multiplier: Modulation strength (0=none, 1=full Clybius effect). Default: 1.0
        percentile: Upper/lower percentile threshold. Default: 5.0
    
    Returns:
        Spectrally modulated noise prediction
    """
    if multiplier == 0 or percentile <= 0:
        return noise_pred
    
    try:
        # FFT
        fourier = torch.fft.fft2(noise_pred, dim=(-2, -1))
        
        # Log amplitude (with small epsilon for numerical stability)
        log_amp = torch.log(torch.sqrt(fourier.real ** 2 + fourier.imag ** 2) + 1e-8)
        
        # Compute quantiles on absolute log amplitude
        log_amp_flat = log_amp.abs().flatten(2)
        quantile_low = torch.quantile(log_amp_flat, percentile * 0.01, dim=2)
        quantile_high = torch.quantile(log_amp_flat, 1 - percentile * 0.01, dim=2)
        
        # Expand quantiles back to log_amp shape
        quantile_low = quantile_low.unsqueeze(-1).unsqueeze(-1).expand(log_amp.shape)
        quantile_high = quantile_high.unsqueeze(-1).unsqueeze(-1).expand(log_amp.shape)
        
        # Create masks (Clybius approach)
        # mask_low: boost values below low threshold (range 1.0 to 1.5)
        # mask_high: reduce values above high threshold (range 0.5 to 1.0)
        mask_low = ((log_amp < quantile_low).float() + 1).clamp_(max=1.5)
        mask_high = ((log_amp < quantile_high).float()).clamp_(min=0.5)
        
        # Apply modulation via exponentiation
        filtered_fourier = fourier * ((mask_low * mask_high) ** multiplier)
        
        # Inverse FFT
        result = torch.fft.ifft2(filtered_fourier, dim=(-2, -1)).real
        
        return result
        
    except Exception as e:
        print(f"⚠️ Spectral modulation failed: {e}")
        return noise_pred



def create_spectral_modulation_cfg_hook(multiplier=1.0, percentile=5.0):
    """
    Create a CFG hook that applies Clybius spectral modulation to noise prediction.
    
    This hooks into reForge's set_model_sampler_cfg_function to intercept
    the CFG calculation and apply spectral modulation at the correct point.
    
    Args:
        multiplier: Modulation strength (0=none, 1=full). Default: 1.0
        percentile: Frequency percentile threshold. Default: 5.0
    
    Returns:
        A hook function to pass to set_model_sampler_cfg_function
    """
    def spectral_cfg_hook(args):
        cond = args["cond"]
        uncond = args["uncond"]
        cond_scale = args["cond_scale"]
        sigma = args["sigma"]
        x_orig = args["input"]
        
        # Reshape sigma for broadcasting
        sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
        
        # Convert to v-pred space (from RescaleCFG reference)
        x = x_orig / (sigma * sigma + 1.0)
        cond_v = ((x - (x_orig - cond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma)
        uncond_v = ((x - (x_orig - uncond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma)
        
        # Compute noise prediction
        noise_pred = cond_v - uncond_v
        
        # Apply Clybius spectral modulation to noise prediction
        noise_pred_modulated = apply_spectral_modulation_clybius(noise_pred, multiplier, percentile)
        
        # Compute CFG with modified noise prediction
        x_cfg = uncond_v + cond_scale * noise_pred_modulated
        
        # Convert back from v-pred space
        return x_orig - (x - x_cfg * sigma / (sigma * sigma + 1.0) ** 0.5)
    
    return spectral_cfg_hook



def apply_combat_cfg_drift(latent, method='mean', intensity=1.0):
    """
    Combat CFG Drift: Reduce mean drift from high CFG values.

    Based on ComfyUI-Latent-Modifiers.

    As CFG increases, the latent mean can drift away from 0, which causes
    color shifts and other artifacts. This technique reduces the drift
    proportionally based on intensity.

    Args:
        latent: The latent tensor to correct
        method: 'mean' or 'median'. Default: 'mean'
        intensity: How much drift to remove (0=none, 1=full). Default: 1.0

    Returns:
        Drift-corrected latent
    """
    if intensity <= 0:
        return latent

    try:
        if method == 'median':
            # Compute global median per batch (across all channels and spatial dims)
            center = latent.view(latent.shape[0], -1).median(dim=-1, keepdim=True)[0]
            center = center.view(latent.shape[0], 1, 1, 1)
        else:
            # Compute global mean per batch (across all channels and spatial dims)
            # This matches ComfyUI's PostCFGsubtractMeanNode implementation
            center = latent.mean(dim=(1, 2, 3), keepdim=True)

        # Remove drift proportionally based on intensity
        # intensity=1.0 removes all drift, intensity=0.5 removes half
        return latent - center * intensity

    except Exception as e:
        print(f"⚠️ Combat CFG drift failed: {e}")
        return latent



def compute_phase_aware_cfg_scale(base_scale, progress, alpha=2.0, beta=2.0):
    """
    Phase-Aware CFG Scaling: Adjust CFG scale based on sampling progress.

    Inspired by β-CFG (arXiv:2502.10574).

    CFG effectiveness varies by sampling phase:
    - Early: Lower CFG allows manifold exploration
    - Middle: Higher CFG for prompt adherence
    - Late: Lower CFG to stay on data manifold

    Args:
        base_scale: The user-specified CFG scale
        progress: Sampling progress (0.0 to 1.0)
        alpha: Beta distribution alpha parameter. Default: 2.0
        beta: Beta distribution beta parameter. Default: 2.0

    Returns:
        Adjusted CFG scale for the current step
    """
    try:
        # Use a simple polynomial approximation of beta distribution
        # Beta(2,2) peaks at 0.5 with a smooth curve
        # f(x) = 6 * x * (1-x) for Beta(2,2), normalized to peak at 1
        if alpha == 2.0 and beta == 2.0:
            # Simple case: symmetric peak at 0.5
            scale_factor = 4.0 * progress * (1.0 - progress)  # Peaks at 1.0 when progress=0.5
            scale_factor = 0.7 + 0.6 * scale_factor  # Range: 0.7 to 1.3
        else:
            # General case: use polynomial approximation
            # Mode of Beta(a,b) is at (a-1)/(a+b-2)
            mode = (alpha - 1.0) / (alpha + beta - 2.0) if (alpha + beta) > 2 else 0.5
            # Create a smooth curve that peaks at the mode
            dist_from_mode = abs(progress - mode)
            scale_factor = 1.0 - 0.3 * dist_from_mode * 2  # Simple linear falloff
            scale_factor = max(0.7, min(1.3, scale_factor))

        return base_scale * scale_factor

    except Exception as e:
        print(f"⚠️ Phase-aware CFG scaling failed: {e}")
        return base_scale



# apply_cfg_techniques() removed — was using legacy keys (akashic_combat_cfg_drift /
# akashic_combat_drift_intensity) that no longer match the live CFG runtime, which
# operates through configure_cfg_runtime() / adept_after_cfg_callback instead.


# ============================================================================
# DUAL-MODE CFG RUNTIME (A1111 callbacks + optional native hook)
# ============================================================================

def create_phase_aware_native_cfg_hook(base_hook=None, alpha=2.0, beta=2.0):
    """
    Native CFG hook for Forge/reForge-like backends that support
    set_model_sampler_cfg_function(). Applies phase-aware CFG scaling,
    then optionally delegates to a downstream hook (e.g. spectral modulation).
    """
    def hook(args):
        cond       = args["cond"]
        uncond     = args["uncond"]
        cond_scale = float(args["cond_scale"])
        sigma      = args["sigma"]
        x_orig     = args["input"]

        total_steps = max(int(ADEPT_STATE.get("_cfg_total_steps", 1)), 1)
        step_idx    = int(ADEPT_STATE.get("_cfg_step_idx", 0))
        progress    = min(max(step_idx / max(total_steps - 1, 1), 0.0), 1.0)

        phased_scale = compute_phase_aware_cfg_scale(cond_scale, progress,
                                                     alpha=alpha, beta=beta)
        patched_args = dict(args)
        patched_args["cond_scale"] = phased_scale

        if base_hook is not None:
            return base_hook(patched_args)

        # Vanilla CFG combine with phased scale
        sigma_b = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
        x       = x_orig / (sigma_b * sigma_b + 1.0)
        cond_v  = ((x - (x_orig - cond))  * (sigma_b ** 2 + 1.0) ** 0.5) / sigma_b
        uncond_v = ((x - (x_orig - uncond)) * (sigma_b ** 2 + 1.0) ** 0.5) / sigma_b
        x_cfg   = uncond_v + phased_scale * (cond_v - uncond_v)
        return x_orig - (x - x_cfg * sigma_b / (sigma_b * sigma_b + 1.0) ** 0.5)
    return hook


def create_combined_native_cfg_hook():
    """
    Build one composite native hook from whatever CFG features are enabled.
    Layer order: phase-aware scale → spectral modulation.
    Returns None if nothing is enabled (caller should clear the hook).
    """
    base_hook = None
    if ADEPT_STATE.get("spectral_cfg_enabled", False):
        base_hook = create_spectral_modulation_cfg_hook(
            multiplier=ADEPT_STATE.get("spectral_multiplier", 1.0),
            percentile=ADEPT_STATE.get("spectral_percentile", 5.0),
        )
    if ADEPT_STATE.get("phase_cfg_enabled", False):
        return create_phase_aware_native_cfg_hook(
            base_hook=base_hook,
            alpha=ADEPT_STATE.get("phase_cfg_alpha", 2.0),
            beta=ADEPT_STATE.get("phase_cfg_beta", 2.0),
        )
    return base_hook


def adept_cfg_denoiser_callback(params):
    """
    Official A1111 on_cfg_denoiser callback. Used only to track step
    counters for phase-aware progress bookkeeping; the public API here
    doesn't expose cond/uncond predictions so we can't do CFG math.
    """
    ADEPT_STATE["_cfg_step_idx"]    = int(getattr(params, "sampling_step",        0))
    ADEPT_STATE["_cfg_total_steps"] = int(getattr(params, "total_sampling_steps", 1))


def adept_after_cfg_callback(params):
    """
    Official A1111 on_cfg_after_cfg callback.
    Combat CFG Drift is the only technique that maps cleanly here,
    because AfterCFGCallbackParams only provides (x, sampling_step,
    total_sampling_steps) — no raw cond/uncond tensors.
    """
    if not ADEPT_STATE.get("enabled", False):
        return
    if not ADEPT_STATE.get("cfg_drift_enabled", False):
        return
    try:
        params.x = apply_combat_cfg_drift(
            params.x,
            method=ADEPT_STATE.get("cfg_drift_method", "mean"),
            intensity=ADEPT_STATE.get("cfg_drift_intensity", 0.5),
        )
    except Exception as e:
        print(f"⚠️ Adept post-CFG drift callback failed: {e}")


def uninstall_a1111_cfg_callbacks():
    global _ADEPT_CFG_AFTER_CB, _ADEPT_CFG_DENOISER_CB
    for cb in (_ADEPT_CFG_AFTER_CB, _ADEPT_CFG_DENOISER_CB):
        if cb is not None:
            try:
                script_callbacks.remove_callbacks_for_function(cb)
            except Exception:
                pass
    _ADEPT_CFG_AFTER_CB    = None
    _ADEPT_CFG_DENOISER_CB = None


def install_a1111_cfg_callbacks():
    global _ADEPT_CFG_AFTER_CB, _ADEPT_CFG_DENOISER_CB
    uninstall_a1111_cfg_callbacks()
    _ADEPT_CFG_DENOISER_CB = adept_cfg_denoiser_callback
    _ADEPT_CFG_AFTER_CB    = adept_after_cfg_callback
    script_callbacks.on_cfg_denoiser(  _ADEPT_CFG_DENOISER_CB, name="adept_cfg_denoiser")
    script_callbacks.on_cfg_after_cfg( _ADEPT_CFG_AFTER_CB,    name="adept_after_cfg")


def _get_native_cfg_hook_target():
    """
    Locate a Forge/reForge-like model object that supports
    set_model_sampler_cfg_function(), if one exists.
    """
    sd = getattr(shared, "sd_model", None)
    candidates = [
        sd,
        getattr(sd, "forge_objects", None),
        getattr(sd, "model", None),
        getattr(getattr(sd, "model", None), "model", None) if sd else None,
    ]
    for obj in candidates:
        if obj is not None and hasattr(obj, "set_model_sampler_cfg_function"):
            return obj
    return None


def uninstall_native_cfg_hook():
    global _ADEPT_NATIVE_CFG_HOOK_ACTIVE
    target = _get_native_cfg_hook_target()
    if target is not None:
        try:
            target.set_model_sampler_cfg_function(None)
        except Exception:
            pass
    _ADEPT_NATIVE_CFG_HOOK_ACTIVE = False


def install_native_cfg_hook():
    global _ADEPT_NATIVE_CFG_HOOK_ACTIVE
    target = _get_native_cfg_hook_target()
    if target is None:
        _ADEPT_NATIVE_CFG_HOOK_ACTIVE = False
        return False
    hook = create_combined_native_cfg_hook()
    try:
        target.set_model_sampler_cfg_function(hook)   # None clears it if nothing enabled
        _ADEPT_NATIVE_CFG_HOOK_ACTIVE = (hook is not None)
        return True
    except Exception as e:
        print(f"⚠️ Adept native CFG hook install failed: {e}")
        _ADEPT_NATIVE_CFG_HOOK_ACTIVE = False
        return False


# ============================================================================
# CFGDenoiser MONKEY-PATCH (stock A1111 fallback for spectral/phase CFG)
# ============================================================================

def _adept_cfg_progress_from_denoiser(denoiser):
    """Compute sampling progress [0,1] from CFGDenoiser step counters."""
    total_steps = max(int(getattr(denoiser, "total_steps", 1) or 1), 1)
    step_idx    = int(getattr(denoiser, "step", 0))
    return min(max(step_idx / max(total_steps - 1, 1), 0.0), 1.0)


def _adept_nativeish_cfg_term(x_i, sigma_i, cond_i, uncond_i, scale):
    """
    Approximate native hook behavior for one cond/uncond pair.
    Feeds x/sigma/cond/uncond into the same composite hook builder used in
    native-hook mode so stock A1111 gets as close as possible to Forge parity.
    Falls back to plain weighted delta if hook is unavailable or errors.
    """
    if abs(float(scale)) < 1e-12:
        return torch.zeros_like(uncond_i)
    hook = create_combined_native_cfg_hook()
    if hook is None:
        return (cond_i - uncond_i) * float(scale)
    try:
        combined = hook({
            "cond":       cond_i,
            "uncond":     uncond_i,
            "cond_scale": float(scale),
            "sigma":      sigma_i,
            "input":      x_i,
        })
        return combined - uncond_i
    except Exception as e:
        print(f"⚠️ Adept native-ish CFG term fallback: {e}")
        return (cond_i - uncond_i) * float(scale)


def patch_cfg_denoiser():
    """
    Stock A1111 fallback for Spectral Modulation + Phase-Aware CFG.

    Strategy:
      1. Thin forward() wrapper that stashes x / sigma on the instance so the
         combine methods can reach them — original forward logic is untouched.
      2. Patched combine_denoised() uses those values to call the same composite
         hook builder as native-hook mode, giving near-parity behaviour.
      3. Patched combine_denoised_for_edit_model() does the same for pix2pix.

    This is intentionally safer than a full forward() rewrite: upstream A1111
    changes to refiner/masking/skip-uncond logic remain unaffected.
    """
    global _CFGD_ORIG_COMBINE, _CFGD_ORIG_COMBINE_EDIT, _CFGD_ORIG_FORWARD, _CFGD_MONKEYPATCH_ACTIVE
    try:
        from modules import sd_samplers_cfg_denoiser as sd_cfg
    except Exception as e:
        print(f"⚠️ Adept CFGDenoiser patch import failed: {e}")
        _CFGD_MONKEYPATCH_ACTIVE = False
        return False

    cls = sd_cfg.CFGDenoiser

    if _CFGD_ORIG_COMBINE is None:
        _CFGD_ORIG_COMBINE = cls.combine_denoised
    if _CFGD_ORIG_COMBINE_EDIT is None:
        _CFGD_ORIG_COMBINE_EDIT = cls.combine_denoised_for_edit_model
    if _CFGD_ORIG_FORWARD is None:
        _CFGD_ORIG_FORWARD = cls.forward

    # --- forward wrapper: stash x/sigma, then run original ---
    def adept_forward(self, x, sigma, uncond, cond, cond_scale, s_min_uncond, image_cond):
        setattr(self, _ADEPT_CFGDENOISER_CTX_ATTR, {
            "x":          x,
            "sigma":      sigma,
            "uncond":     uncond,
            "cond":       cond,
            "cond_scale": float(cond_scale),
        })
        try:
            return _CFGD_ORIG_FORWARD(self, x, sigma, uncond, cond,
                                      cond_scale, s_min_uncond, image_cond)
        finally:
            try:
                delattr(self, _ADEPT_CFGDENOISER_CTX_ATTR)
            except AttributeError:
                pass

    # --- combine_denoised: per-cond native-ish or plain path ---
    def adept_combine_denoised(self, x_out, conds_list, uncond, cond_scale):
        denoised_uncond = x_out[-uncond.shape[0]:]
        denoised = torch.clone(denoised_uncond)
        ctx      = getattr(self, _ADEPT_CFGDENOISER_CTX_ATTR, None)
        x_ctx    = None if ctx is None else ctx.get("x",     None)
        sigma_ctx = None if ctx is None else ctx.get("sigma", None)
        progress = _adept_cfg_progress_from_denoiser(self)

        eff_scale = float(cond_scale)
        if ADEPT_STATE.get("phase_cfg_enabled", False):
            eff_scale = compute_phase_aware_cfg_scale(
                eff_scale, progress,
                alpha=ADEPT_STATE.get("phase_cfg_alpha", 2.0),
                beta =ADEPT_STATE.get("phase_cfg_beta",  2.0),
            )

        use_nativeish = (
            (ADEPT_STATE.get("spectral_cfg_enabled", False) or
             ADEPT_STATE.get("phase_cfg_enabled",    False))
            and x_ctx is not None and sigma_ctx is not None
        )

        for i, conds in enumerate(conds_list):
            for cond_index, weight in conds:
                cond_i  = x_out[cond_index:cond_index + 1]
                uncond_i = denoised_uncond[i:i + 1]

                if use_nativeish:
                    term = _adept_nativeish_cfg_term(
                        x_i     = x_ctx[i:i + 1],
                        sigma_i  = sigma_ctx[i:i + 1],
                        cond_i   = cond_i,
                        uncond_i = uncond_i,
                        scale    = float(weight) * eff_scale,
                    )
                    denoised[i:i + 1] += term
                else:
                    delta = cond_i - uncond_i
                    if ADEPT_STATE.get("spectral_cfg_enabled", False):
                        delta = apply_spectral_modulation_clybius(
                            delta,
                            multiplier=ADEPT_STATE.get("spectral_multiplier", 1.0),
                            percentile=ADEPT_STATE.get("spectral_percentile", 5.0),
                        )
                    denoised[i:i + 1] += delta * (float(weight) * eff_scale)

        return denoised

    # --- combine_denoised_for_edit_model: pix2pix / instruct path ---
    def adept_combine_denoised_for_edit_model(self, x_out, cond_scale):
        out_cond, out_img_cond, out_uncond = x_out.chunk(3)
        ctx       = getattr(self, _ADEPT_CFGDENOISER_CTX_ATTR, None)
        x_ctx     = None if ctx is None else ctx.get("x",     None)
        sigma_ctx = None if ctx is None else ctx.get("sigma", None)
        progress  = _adept_cfg_progress_from_denoiser(self)

        eff_scale = float(cond_scale)
        if ADEPT_STATE.get("phase_cfg_enabled", False):
            eff_scale = compute_phase_aware_cfg_scale(
                eff_scale, progress,
                alpha=ADEPT_STATE.get("phase_cfg_alpha", 2.0),
                beta =ADEPT_STATE.get("phase_cfg_beta",  2.0),
            )

        # Native-ish path when context is available
        if (ADEPT_STATE.get("spectral_cfg_enabled", False) or
                ADEPT_STATE.get("phase_cfg_enabled", False)):
            if x_ctx is not None and sigma_ctx is not None:
                try:
                    hook = create_combined_native_cfg_hook()
                    if hook is not None:
                        base = hook({
                            "cond":       out_cond,
                            "uncond":     out_img_cond,
                            "cond_scale": eff_scale,
                            "sigma":      sigma_ctx,
                            "input":      x_ctx,
                        })
                        return base + self.image_cfg_scale * (out_img_cond - out_uncond)
                except Exception as e:
                    print(f"⚠️ Adept edit-model native-ish fallback: {e}")

        # Plain path (no context or hook failed)
        delta = out_cond - out_img_cond
        if ADEPT_STATE.get("spectral_cfg_enabled", False):
            delta = apply_spectral_modulation_clybius(
                delta,
                multiplier=ADEPT_STATE.get("spectral_multiplier", 1.0),
                percentile=ADEPT_STATE.get("spectral_percentile", 5.0),
            )
        return out_uncond + eff_scale * delta + self.image_cfg_scale * (out_img_cond - out_uncond)

    try:
        cls.forward                       = adept_forward
        cls.combine_denoised              = adept_combine_denoised
        cls.combine_denoised_for_edit_model = adept_combine_denoised_for_edit_model
        _CFGD_MONKEYPATCH_ACTIVE = True
        return True
    except Exception as e:
        print(f"⚠️ Adept CFGDenoiser patch failed: {e}")
        _CFGD_MONKEYPATCH_ACTIVE = False
        return False


def unpatch_cfg_denoiser():
    global _CFGD_MONKEYPATCH_ACTIVE
    try:
        from modules import sd_samplers_cfg_denoiser as sd_cfg
    except Exception:
        _CFGD_MONKEYPATCH_ACTIVE = False
        return False
    cls = sd_cfg.CFGDenoiser
    try:
        if _CFGD_ORIG_FORWARD is not None:
            cls.forward = _CFGD_ORIG_FORWARD
        if _CFGD_ORIG_COMBINE is not None:
            cls.combine_denoised = _CFGD_ORIG_COMBINE
        if _CFGD_ORIG_COMBINE_EDIT is not None:
            cls.combine_denoised_for_edit_model = _CFGD_ORIG_COMBINE_EDIT
        _CFGD_MONKEYPATCH_ACTIVE = False
        return True
    except Exception as e:
        print(f"⚠️ Adept CFGDenoiser unpatch failed: {e}")
        _CFGD_MONKEYPATCH_ACTIVE = False
        return False


def configure_cfg_runtime():
    """
    Select and activate the right CFG runtime mode:
      off                 – nothing enabled; all hooks/callbacks/patches cleared
      a1111-postcfg       – stock A1111; Combat CFG Drift only via official callback
      a1111-monkeypatch   – stock A1111; Spectral + Phase-Aware via CFGDenoiser patch
      native-hook         – Forge/reForge-like backend; all three via sampler CFG hook

    Returns the mode string so process() can log it.
    """
    # If the extension is globally disabled, always tear down and return off.
    if not ADEPT_STATE.get("enabled", False):
        uninstall_a1111_cfg_callbacks()
        uninstall_native_cfg_hook()
        unpatch_cfg_denoiser()
        ADEPT_STATE["cfg_runtime_mode"] = "off"
        return "off"

    drift    = ADEPT_STATE.get("cfg_drift_enabled",   False)
    spectral = ADEPT_STATE.get("spectral_cfg_enabled", False)
    phase    = ADEPT_STATE.get("phase_cfg_enabled",    False)

    # Always tear down everything first for a clean slate
    uninstall_a1111_cfg_callbacks()
    uninstall_native_cfg_hook()
    unpatch_cfg_denoiser()

    if not (drift or spectral or phase):
        ADEPT_STATE["cfg_runtime_mode"] = "off"
        return "off"

    # Prefer native hook if backend supports it
    native_target = _get_native_cfg_hook_target()
    if native_target is not None:
        install_a1111_cfg_callbacks()   # keeps drift working in native mode too
        install_native_cfg_hook()
        ADEPT_STATE["cfg_runtime_mode"] = "native-hook"
        return "native-hook"

    # Stock A1111: always install callbacks (drift)
    install_a1111_cfg_callbacks()

    if spectral or phase:
        ok = patch_cfg_denoiser()
        if ok:
            ADEPT_STATE["cfg_runtime_mode"] = "a1111-monkeypatch"
            print("✅ Adept: stock A1111 CFGDenoiser monkey-patch active (spectral/phase + drift enabled)")
            return "a1111-monkeypatch"
        print("⚠️ Adept: CFGDenoiser monkey-patch failed; falling back to post-CFG drift only")

    ADEPT_STATE["cfg_runtime_mode"] = "a1111-postcfg"
    return "a1111-postcfg"


def sa_solver_step(x, d_history, sigma, sigma_next, tau, s_noise=1.0, noise_sampler=None, 
                   order=2, ndb_strength=0.0, progress=0.0, eqvae_mode=False, eqvae_blur_sigma=None):
    """SA-Solver step - CRITICAL for Akashic Solver."""
    dt = sigma_next - sigma
    
    if len(d_history) >= 2 and order >= 2:
        sigma_cur, d_cur = d_history[-1]
        sigma_prev, d_prev = d_history[-2]
        
        h_prev = sigma_cur - sigma_prev
        r = abs(dt / (h_prev + 1e-8)) if abs(h_prev) > 1e-8 else 1.0
        r = min(r, 2.0)
        
        if len(d_history) >= 3 and order >= 3:
            sigma_0, d_0 = d_history[-3]
            h_0 = sigma_prev - sigma_0
            h_1 = h_prev
            
            if abs(h_0) > 1e-6 and abs(h_1) > 1e-6:
                r0 = min(abs(h_1 / h_0), 2.0)
                r1 = min(abs(dt / (h_1 + 1e-8)), 2.0)
                
                tau_blend = 1.0 - tau
                c0_ab3 = 1.0 + (1.0 + r0) * r1 / 2.0
                c1_ab3 = -(1.0 + r0) * r1 / 2.0
                c2_ab3 = r0 * r1 / 2.0
                c0 = tau_blend * c0_ab3 + (1.0 - tau_blend) * 1.0
                c1 = tau_blend * c1_ab3
                c2 = tau_blend * c2_ab3
                
                c_sum = c0 + c1 + c2
                if abs(c_sum) > 1e-8:
                    c0 /= c_sum
                    c1 /= c_sum
                    c2 /= c_sum
                else:
                    c0, c1, c2 = 1.0, 0.0, 0.0
                
                d_interp = c0 * d_cur + c1 * d_prev + c2 * d_0
            else:
                tau_blend = 1.0 - tau
                c1_ab2 = 1.0 + 0.5 * r
                c2_ab2 = -0.5 * r
                c1 = tau_blend * c1_ab2 + (1.0 - tau_blend) * 1.0
                c2 = tau_blend * c2_ab2
                c_sum = c1 + c2
                if abs(c_sum) > 1e-8:
                    c1 /= c_sum
                    c2 /= c_sum
                d_interp = c1 * d_cur + c2 * d_prev
        else:
            tau_blend = 1.0 - tau
            c1_ab2 = 1.0 + 0.5 * r
            c2_ab2 = -0.5 * r
            c1 = tau_blend * c1_ab2 + (1.0 - tau_blend) * 1.0
            c2 = tau_blend * c2_ab2
            c_sum = c1 + c2
            if abs(c_sum) > 1e-8:
                c1 /= c_sum
                c2 /= c_sum
            d_interp = c1 * d_cur + c2 * d_prev
    elif len(d_history) >= 1:
        d_interp = d_history[-1][1]
    else:
        d_interp = torch.zeros_like(x)
    
    # Compute sigma_up based on tau (controls stochasticity)
    sigma_up = 0.0
    if tau > 0 and sigma_next > 0 and noise_sampler is not None:
        sigma_ancestral_sq = sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / (sigma ** 2 + 1e-8)
        sigma_ancestral = sigma_ancestral_sq ** 0.5 if sigma_ancestral_sq > 0 else 0.0
        sigma_up = tau * sigma_ancestral
        
        sigma_down = (sigma_next ** 2 - sigma_up ** 2) ** 0.5
        dt_adjusted = sigma_down - sigma
        
        x_det = x + d_interp * dt_adjusted
        noise = noise_sampler(sigma, sigma_next) * s_noise * sigma_up
        
        # Apply Native Detail Boost if enabled
        if ndb_strength > 0 and TORCHVISION_AVAILABLE:
            # Use EQ-VAE optimized NDB parameters if in EQ-VAE mode
            if eqvae_mode:
                blur_sigma, high_freq_boost = compute_eqvae_ndb(progress, ndb_strength)
            else:
                _, high_freq_boost = compute_native_detail_boost(progress, ndb_strength)
                blur_sigma = 0.5  # Default blur sigma
            
            # Override blur_sigma if explicitly provided
            if eqvae_blur_sigma is not None:
                blur_sigma = eqvae_blur_sigma
            
            # Extract high-frequency component from noise using Gaussian blur
            try:
                low_freq_noise = gaussian_blur(noise, kernel_size=3, sigma=blur_sigma)
                high_freq_noise = noise - low_freq_noise
                noise = noise + high_freq_noise * high_freq_boost
            except Exception:
                pass  # Fallback: use original noise if blur fails
        
        x_next = x_det + noise
    else:
        x_next = x + d_interp * dt
    
    return x_next, sigma_up


def create_detail_enhanced_model(model, x, sigmas, settings):
    # NOTE: Detail Enhancement is currently an internal/experimental path.
    # It is not wired into the UI and callers always pass
    # use_detail_enhancement=False, so this function is never invoked at
    # runtime.  Kept for future re-integration; do not rely on it.
    """Detail enhancement wrapper."""
    if not TORCHVISION_AVAILABLE:
        return model
    
    base_strength = settings.get('detail_enhancement_strength', 0.05)
    radius = settings.get('detail_separation_radius', 0.5)
    total_steps = len(sigmas) - 1
    
    class DetailEnhancer:
        def __init__(self):
            self.current_step = 0
        
        def __call__(self, x_current, sigma, **kwargs):
            denoised = model(x_current, sigma, **kwargs)
            
            try:
                low_freq = gaussian_blur(denoised, kernel_size=3, sigma=radius)
                high_freq = denoised - low_freq
                
                progress = min(self.current_step / max(total_steps, 1), 1.0)
                strength = base_strength * (0.5 + progress)
                
                enhanced = denoised + high_freq * strength
                self.current_step += 1
                return enhanced
            except Exception:
                return denoised
    
    return DetailEnhancer()


# ============================================================================
# ============================================================================
# ============================================================================
# CUSTOM ADVANCED SAMPLERS (Complete port from ComfyUI)
# ============================================================================

@torch.no_grad()
def sample_adept_solver(model, x, sigmas, extra_args=None, callback=None, disable=None, 
                        order=2, use_corrector=True, use_detail_enhancement=False, settings=None):
    """
    Adept Solver: A unified training-free diffusion solver synthesizing improvements from:
    - DPM-Solver++ (data prediction, dynamic thresholding)
    - UniPC (unified predictor-corrector framework)
    - DEIS (exponential integrator)
    - DC-Solver (dynamic compensation)
    """
    extra_args = {} if extra_args is None else extra_args
    settings = settings or {}
    s_in = x.new_ones([x.shape[0]])
    
    order = max(1, min(order, 3))
    
    print(f"🚀 Adept Solver active (Order: {order}, Corrector: {'On' if use_corrector else 'Off'})")
    
    active_model = model
    # use_detail_enhancement is always False from current call-sites;
    # the block below is preserved for future re-integration but is not
    # currently reachable via the UI.
    if use_detail_enhancement and TORCHVISION_AVAILABLE:
        active_model = create_detail_enhanced_model(model, x, sigmas, settings)

    model_outputs = []
    
    for i in range(len(sigmas) - 1):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        denoised = active_model(x, sigma * s_in, **extra_args)
        
        if extra_args.get('cond_scale', 1.0) > 7.0:
            denoised = apply_dynamic_thresholding(denoised, percentile=0.995)
        
        d = to_d(x, sigma, denoised)

        derivative_max = torch.abs(d).max()
        sigma_adaptive_threshold = 1000.0 * (1.0 + sigma / 10.0)
        if torch.isnan(d).any() or torch.isinf(d).any() or derivative_max > sigma_adaptive_threshold:
            print(f"⚠️ Extreme derivative detected at step {i}/{len(sigmas)-1}. Clamping for stability.")
            d = torch.clamp(d, -sigma_adaptive_threshold, sigma_adaptive_threshold)
            if torch.isnan(d).any() or torch.isinf(d).any():
                d = torch.zeros_like(d)

        model_outputs.append((sigma, d))
        if len(model_outputs) > order:
            model_outputs.pop(0)
        
        dt = sigma_next - sigma
        
        if len(model_outputs) == 1 or order == 1:
            x_pred = x + d * dt
        elif len(model_outputs) == 2 and order >= 2:
            sigma_prev, d_prev = model_outputs[-2]
            d_cur = model_outputs[-1][1]
            
            h = sigma - sigma_prev
            compensation_ratio = compute_compensation_ratio(h.item() if torch.is_tensor(h) else float(h), i, len(sigmas))
            
            d_interp = d_cur + compensation_ratio * (d_cur - d_prev)
            x_pred = x + d_interp * dt
        else:
            sigma_0, d_0 = model_outputs[-3]
            sigma_1, d_1 = model_outputs[-2]
            sigma_2, d_2 = model_outputs[-1]
            
            h_0 = sigma_2 - sigma_1
            h_1 = sigma_1 - sigma_0
            
            h_0_val = h_0.item() if torch.is_tensor(h_0) else float(h_0)
            h_1_val = h_1.item() if torch.is_tensor(h_1) else float(h_1)
            
            if abs(h_1_val) < 1e-6:
                compensation_ratio = compute_compensation_ratio(h_0_val, i, len(sigmas))
                d_interp = d_2 + compensation_ratio * (d_2 - d_1)
            else:
                r0 = h_0_val / h_1_val
                c0 = 1.0 + r0 / 2.0
                c1 = -r0 / 2.0
                c2 = 0.0
                
                c_sum = c0 + c1 + c2
                c0 /= c_sum
                c1 /= c_sum
                c2 = 1.0 - c0 - c1
                
                d_interp = c0 * d_2 + c1 * d_1 + c2 * d_0
            
            x_pred = x + d_interp * dt
        
        if use_corrector and i < len(sigmas) - 2:
            denoised_pred = active_model(x_pred, sigma_next * s_in, **extra_args)
            
            if extra_args.get('cond_scale', 1.0) > 7.0:
                denoised_pred = apply_dynamic_thresholding(denoised_pred, percentile=0.995)

            d_pred = to_d(x_pred, sigma_next, denoised_pred)

            if torch.isnan(d_pred).any() or torch.isinf(d_pred).any() or torch.abs(d_pred).max() > 1000.0:
                d_pred = torch.clamp(d_pred, -100.0, 100.0)
                if torch.isnan(d_pred).any() or torch.isinf(d_pred).any():
                    d_pred = torch.zeros_like(d_pred)
            
            dt = sigma_next - sigma
            x = x + (d + d_pred) * dt * 0.5
        else:
            x = x_pred
        
        if torch.isnan(x).any() or torch.isinf(x).any():
            print(f"❌ CRITICAL: NaN/Inf detected at step {i}/{len(sigmas)-1}!")
            if i == 0:
                raise RuntimeError("NaN/Inf on first step - check model/inputs")
            
            print("   Attempting recovery with conservative Euler step...")
            denoised_safe = active_model(x, sigma * s_in, **extra_args)
            if torch.isnan(denoised_safe).any():
                raise RuntimeError("Model producing NaN - check CFG scale and model")
            
            d_safe = to_d(x, sigma, denoised_safe)
            dt_safe = (sigma_next - sigma) * 0.5
            x = x + d_safe * dt_safe
            use_corrector = False
            print("   Recovery successful. Corrector disabled for stability.")
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})
    
    return x


@torch.no_grad()
def compute_adaptive_noise_scale(excess_samples, base_s_noise,
                                  correction_power=0.5,
                                  dampen_floor=0.80, boost_ceiling=1.15):
    """
    Adaptive Noise Scale: calibrated sigma-relative excess correction.

    After collecting excess ratio samples during a warmup window inside the
    texture phase (0.5 < sigma < 5.0), computes a single global correction
    factor based on the median excess. The excess ratio measures how much
    noise is slowing down convergence:

        excess = (change_i / change_{i-1}) / (sigma_{i+1} / sigma_i)

    excess > 1.0 -> model converges slower than sigma predicts -> reduce noise.
    excess < 1.0 -> model converges faster -> noise could be boosted.

    Returns (correction_factor, median_excess).
    """
    if not excess_samples:
        return 1.0, 0.0

    sorted_samples = sorted(excess_samples)
    median_excess = sorted_samples[len(sorted_samples) // 2]

    if median_excess > 0:
        correction = 1.0 / (median_excess ** correction_power)
    else:
        correction = 1.0

    correction = max(dampen_floor, min(boost_ceiling, correction))
    return correction, median_excess


def compute_binned_corrections(excess_bins, global_correction,
                                correction_power=0.5, dampen_floor=0.80,
                                boost_ceiling=1.15, min_samples=3):
    """Per-phase (structural/texture/cleanup) corrections. Bins with fewer
    than min_samples samples fall back to the global correction."""
    bin_corrections = {}
    for bin_name, samples in excess_bins.items():
        if len(samples) >= min_samples:
            correction, _ = compute_adaptive_noise_scale(
                samples, 1.0, correction_power, dampen_floor, boost_ceiling
            )
            bin_corrections[bin_name] = correction
        else:
            bin_corrections[bin_name] = global_correction
    return bin_corrections


def get_phase_correction(sigma_val, bin_corrections, global_correction):
    """Look up the correction for the current sigma phase."""
    if bin_corrections is None:
        return global_correction
    if sigma_val > 5.0:
        return bin_corrections.get('structural', global_correction)
    elif sigma_val > 0.5:
        return bin_corrections.get('texture', global_correction)
    else:
        return bin_corrections.get('cleanup', global_correction)


def adaptive_noise_step(denoised_raw, prev_denoised_raw, prev_change_norm,
                         sigma, sigma_next, excess_samples, adaptive_correction,
                         base_s_noise, step_idx, total_steps, warmup=5,
                         binned_enabled=False, excess_bins=None,
                         adaptive_bin_corrections=None):
    """
    Process one step of adaptive noise calibration.

    Call after getting the raw model output and before noise injection.
    Handles warmup collection, calibration trigger, and correction application.

    Returns:
        (effective_s_noise, change_norm, adaptive_correction, should_restart,
         adaptive_bin_corrections)
    """
    if prev_denoised_raw is not None:
        change_norm = torch.norm((denoised_raw - prev_denoised_raw).flatten(1), dim=1).mean().item()
    else:
        change_norm = None

    effective_s_noise = base_s_noise
    should_restart = False

    if sigma_next > 0:
        sigma_val = sigma.item() if torch.is_tensor(sigma) else float(sigma)
        sigma_next_val = sigma_next.item() if torch.is_tensor(sigma_next) else float(sigma_next)
        in_texture_phase = 0.5 < sigma_val < 5.0

        if adaptive_correction is not None:
            # Post-restart: apply the already-computed correction.
            if binned_enabled and adaptive_bin_corrections is not None:
                correction = get_phase_correction(sigma_val, adaptive_bin_corrections, adaptive_correction)
                effective_s_noise *= correction
            else:
                effective_s_noise *= adaptive_correction
        elif change_norm is not None and prev_change_norm is not None:
            change_ratio = change_norm / (prev_change_norm + 1e-8)
            sigma_ratio = sigma_next_val / (sigma_val + 1e-8)
            excess = change_ratio / (sigma_ratio + 1e-8)

            if binned_enabled and excess_bins is not None:
                if sigma_val > 5.0:
                    excess_bins['structural'].append(excess)
                elif sigma_val > 0.5:
                    excess_bins['texture'].append(excess)
                else:
                    excess_bins['cleanup'].append(excess)

            if in_texture_phase:
                excess_samples.append(excess)
                if len(excess_samples) >= warmup:
                    adaptive_correction, median_excess = compute_adaptive_noise_scale(
                        excess_samples, effective_s_noise
                    )
                    if binned_enabled and excess_bins is not None:
                        adaptive_bin_corrections = compute_binned_corrections(
                            excess_bins, adaptive_correction
                        )
                        bin_info = {k: f"{v:.3f}" for k, v in adaptive_bin_corrections.items()}
                        print(f"   🎯 Adaptive Noise Scale calibrated: global={adaptive_correction:.3f}, binned={bin_info}")
                    else:
                        print(f"   🎯 Adaptive Noise Scale calibrated: correction={adaptive_correction:.3f}")
                    should_restart = True

    return effective_s_noise, change_norm, adaptive_correction, should_restart, adaptive_bin_corrections


class _CalibrationComplete(Exception):
    """Raised by CalibrationCallback once enough samples have been collected,
    so the calibration pass can be aborted early instead of running to completion."""
    pass


class CalibrationCallback:
    """
    Wraps the user's/A1111's own step callback to collect ANS excess-ratio
    samples on the fly, from whatever native sampler is currently running.

    This is what makes Adaptive Noise Scale work on *any* k-diffusion
    sampler rather than only the 3 hand-written custom solvers: nearly every
    sampler calls back({'i':, 'sigma':, 'denoised':, ...}) each step (that's
    how A1111 draws the live preview), so we can observe convergence without
    needing to understand that sampler's internals.

    Raises _CalibrationComplete once `warmup` texture-phase samples are in.
    """

    def __init__(self, original_callback, sigmas, warmup=5):
        self.original_callback = original_callback
        self.sigmas = sigmas
        self.warmup = warmup
        self.prev_denoised = None
        self.prev_change_norm = None
        self.excess_samples = []
        self.excess_bins = {'structural': [], 'texture': [], 'cleanup': []}

    def __call__(self, info):
        if self.original_callback is not None:
            self.original_callback(info)

        i = info.get('i')
        denoised = info.get('denoised')
        sigma = info.get('sigma')
        if i is None or denoised is None or sigma is None:
            return  # non-standard callback payload -- can't calibrate, just pass through

        sigma_val = sigma.item() if torch.is_tensor(sigma) else float(sigma)

        if i + 1 >= len(self.sigmas):
            return

        sigma_next_val = self.sigmas[i + 1].item() if torch.is_tensor(self.sigmas[i + 1]) else float(self.sigmas[i + 1])

        if self.prev_denoised is not None:
            change_norm = torch.norm(
                (denoised - self.prev_denoised).flatten(1), dim=1
            ).mean().item()

            if self.prev_change_norm is not None and sigma_next_val > 0:
                change_ratio = change_norm / (self.prev_change_norm + 1e-8)
                sigma_ratio = sigma_next_val / (sigma_val + 1e-8)
                excess = change_ratio / (sigma_ratio + 1e-8)

                if sigma_val > 5.0:
                    self.excess_bins['structural'].append(excess)
                elif sigma_val > 0.5:
                    self.excess_bins['texture'].append(excess)
                else:
                    self.excess_bins['cleanup'].append(excess)

                if 0.5 < sigma_val < 5.0:
                    self.excess_samples.append(excess)
                    if len(self.excess_samples) >= self.warmup:
                        raise _CalibrationComplete()

            self.prev_change_norm = change_norm

        self.prev_denoised = denoised.detach().clone()


def run_ans_two_pass(original_fn, name, model, x, sigmas, extra_args, callback, disable, kwargs):
    """
    Universal Adaptive Noise Scale for any native sampler, via two full runs:

      Pass 1 (calibration): run original_fn wrapped in a CalibrationCallback,
      which aborts the run early (via _CalibrationComplete) once 5 texture-
      phase excess samples are collected.

      Pass 2 (corrected rerun): rerun original_fn from x_initial (the true
      starting noise) with either a wrapped noise_sampler (per-phase binned
      correction, for solvers that accept one) or a scaled scalar s_noise
      (global correction, for solvers that only take s_noise).

    Only called after the caller has already confirmed original_fn's
    signature has 's_noise' or 'noise_sampler' -- deterministic solvers with
    neither (e.g. DPM++ 2M, LMS) are skipped by the caller before this runs.
    """
    x_initial = x.clone()

    print(f"   Adaptive Noise Scale: calibration pass ({name.replace('sample_', '')})")
    cal_callback = CalibrationCallback(callback, sigmas)
    calibration_result = None
    try:
        calibration_result = original_fn(model, x.clone(), sigmas.clone(), dict(extra_args) if extra_args else {}, cal_callback, disable, **kwargs)
    except _CalibrationComplete:
        pass

    if len(cal_callback.excess_samples) < 5:
        if calibration_result is not None:
            # Calibration pass ran to full completion without ever collecting
            # 5 texture-phase samples (e.g. the sigma schedule never dips into
            # 0.5<sigma<5.0). Nothing to correct -- reuse this already-valid,
            # already-complete result instead of wastefully redoing the same
            # deterministic work a second time.
            print(f"   Adaptive Noise Scale: texture phase never reached warmup "
                  f"({len(cal_callback.excess_samples)}/5) -- reusing completed pass")
            return calibration_result
        print(f"   Adaptive Noise Scale: not enough texture-phase samples "
              f"({len(cal_callback.excess_samples)}/5), skipping correction")
        return original_fn(model, x_initial, sigmas, extra_args, callback, disable, **kwargs)

    global_correction, median_excess = compute_adaptive_noise_scale(
        cal_callback.excess_samples, 1.0
    )
    bin_corrections = compute_binned_corrections(
        cal_callback.excess_bins, global_correction
    )
    bin_info = {k: f"{v:.3f}" for k, v in bin_corrections.items()}
    print(f"   Adaptive Noise Scale calibrated: global={global_correction:.3f}, "
          f"binned={bin_info}, median_excess={median_excess:.3f}")

    if 0.99 < global_correction < 1.01:
        print("   Adaptive Noise Scale: correction ~1.0, rerunning with original parameters")
        return original_fn(model, x_initial, sigmas, extra_args, callback, disable, **kwargs)

    print(f"   Adaptive Noise Scale: correction pass ({name.replace('sample_', '')})")
    corrected_kwargs = dict(kwargs)

    sig = inspect.signature(original_fn)
    if 'noise_sampler' in sig.parameters:
        base_noise_sampler = kwargs.get('noise_sampler', None)
        if base_noise_sampler is None:
            base_noise_sampler = default_noise_sampler(x_initial)

        def ans_noise_sampler(sigma, sigma_next):
            noise = base_noise_sampler(sigma, sigma_next)
            sigma_val = sigma.item() if torch.is_tensor(sigma) else float(sigma)
            correction = get_phase_correction(sigma_val, bin_corrections, global_correction)
            return noise * correction
        corrected_kwargs['noise_sampler'] = ans_noise_sampler
    else:
        base_s_noise = kwargs.get('s_noise', 1.0)
        corrected_kwargs['s_noise'] = base_s_noise * global_correction

    return original_fn(model, x_initial, sigmas, extra_args, callback, disable, **corrected_kwargs)


@torch.no_grad()
def sample_adept_ancestral_solver(model, x, sigmas, extra_args=None, callback=None, disable=None,
                                   eta=1.0, s_noise=1.0, adaptive_eta=False, phase_noise=False,
                                   phase_strength=0.5, enhanced_derivative=False, 
                                   use_detail_enhancement=False, settings=None,
                                   adaptive_noise_scale=False):
    """
    Enhanced Adept Ancestral Solver: Advanced ancestral sampling with phase-aware adaptations.
    
    Key innovations:
    1. Adaptive ancestral step sizing that changes throughout sampling phases
    2. Phase-aware noise injection (more noise early, less noise late)
    3. Enhanced derivative computation with ancestral-specific corrections
    4. Dynamic eta scheduling for better control
    5. Adaptive Noise Scale: calibrates s_noise from a short warmup window, then
       restarts from x_initial with the correction applied from step 0
    """
    extra_args = {} if extra_args is None else extra_args
    settings = settings or {}
    s_in = x.new_ones([x.shape[0]])
    
    print(f"🚀 Enhanced Adept Ancestral Solver active (η: {eta:.2f}, s_noise: {s_noise:.2f})")
    print(f"   Adaptive Eta: {adaptive_eta}, Phase Noise: {phase_noise}, Enhanced Derivative: {enhanced_derivative}")
    if adaptive_noise_scale:
        print("   Adaptive Noise Scale: ON (auto-calibrating s_noise)")
    
    active_model = model
    # use_detail_enhancement is always False from current call-sites.
    if use_detail_enhancement and TORCHVISION_AVAILABLE:
        active_model = create_detail_enhanced_model(model, x, sigmas, settings)

    noise_sampler = get_noise_sampler(x)
    n_steps = len(sigmas) - 1

    # Adaptive Noise Scale state
    prev_denoised_raw = None
    prev_change_norm = None
    excess_samples = []
    adaptive_correction = None
    excess_bins = {'structural': [], 'texture': [], 'cleanup': []} if adaptive_noise_scale else None
    adaptive_bin_corrections = None
    x_initial = x.clone() if adaptive_noise_scale else None

    i = 0
    while i < n_steps:
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        progress = i / max(n_steps, 1)
        
        if adaptive_eta:
            if progress < 0.3:
                current_eta = eta * 1.08
            elif progress < 0.7:
                current_eta = eta * 0.95
            else:
                current_eta = eta * 1.02
        else:
            current_eta = eta
        
        denoised = active_model(x, sigma * s_in, **extra_args)
        denoised_raw = denoised.clone() if adaptive_noise_scale else None
        
        if extra_args.get('cond_scale', 1.0) > 7.0:
            denoised = apply_dynamic_thresholding(denoised, percentile=0.995)
        
        if enhanced_derivative:
            d = to_d_enhanced_ancestral(x, sigma, denoised, current_eta, progress, None)
        else:
            d = to_d(x, sigma, denoised)

        derivative_max = torch.abs(d).max()
        sigma_adaptive_threshold = 1000.0 * (1.0 + sigma / 10.0)
        if torch.isnan(d).any() or torch.isinf(d).any() or derivative_max > sigma_adaptive_threshold:
            d = torch.clamp(d, -sigma_adaptive_threshold, sigma_adaptive_threshold)
            if torch.isnan(d).any() or torch.isinf(d).any():
                d = torch.zeros_like(d)

        if sigma_next > 0:
            sigma_up = min(sigma_next, current_eta * (sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2) ** 0.5)
            sigma_down = (sigma_next ** 2 - sigma_up ** 2) ** 0.5
        else:
            sigma_up = 0.0
            sigma_down = 0.0
        
        dt = sigma_down - sigma
        x_pred = x + d * dt

        should_restart = False
        if sigma_next > 0:
            if phase_noise:
                if progress < 0.25:
                    target_multiplier = 1.0 + (0.05 * min(progress / 0.25, 1.0))
                elif progress < 0.6:
                    target_multiplier = 1.0 - (0.02 * min((progress - 0.25) / 0.35, 1.0))
                else:
                    target_multiplier = 1.0 - (0.05 * min((progress - 0.6) / 0.4, 1.0))

                noise_multiplier = 1.0 + (target_multiplier - 1.0) * phase_strength
                adaptive_s_noise = s_noise * noise_multiplier
            else:
                adaptive_s_noise = s_noise

            if adaptive_noise_scale:
                adaptive_s_noise, change_norm, adaptive_correction, should_restart, adaptive_bin_corrections = adaptive_noise_step(
                    denoised_raw, prev_denoised_raw, prev_change_norm,
                    sigma, sigma_next, excess_samples, adaptive_correction,
                    adaptive_s_noise, i, n_steps,
                    binned_enabled=True, excess_bins=excess_bins,
                    adaptive_bin_corrections=adaptive_bin_corrections
                )
                if not should_restart:
                    prev_change_norm = change_norm
                    prev_denoised_raw = denoised_raw

            if should_restart and x_initial is not None:
                x = x_initial.clone()
                prev_denoised_raw = None
                prev_change_norm = None
                excess_bins = {'structural': [], 'texture': [], 'cleanup': []}
                x_initial = None
                i = 0
                continue

            noise = noise_sampler(sigma, sigma_next) * adaptive_s_noise * sigma_up
            x = x_pred + noise
        else:
            x = x_pred
        
        if torch.isnan(x).any() or torch.isinf(x).any():
            print(f"❌ CRITICAL: NaN/Inf detected at step {i}/{n_steps}!")
            if i == 0:
                raise RuntimeError("NaN/Inf on first step - check model/inputs")
            
            print("   Attempting recovery...")
            denoised_safe = active_model(x, sigma * s_in, **extra_args)
            if torch.isnan(denoised_safe).any():
                raise RuntimeError("Model producing NaN - check CFG scale and model")
            
            d_safe = to_d(x, sigma, denoised_safe)
            dt_safe = (sigma_next - sigma) * 0.5
            x = x + d_safe * dt_safe
            print("   Recovery successful.")
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})

        i += 1
    
    return x


@torch.no_grad()
def sample_mirror_correction_euler(model, x, sigmas, extra_args=None, callback=None, disable=None,
                                    eta=1.0, s_noise=1.0, correction_phase=0.5, smooth_phase=False,
                                    adaptive_noise_scale=False):
    """
    Mirror Correction Euler: Euler Ancestral with a semantic reflection probe.

    In the first `correction_phase` fraction of steps, uses a 3-call Heun correction:
      x_probe = 2*D(x) - x  (reflection of x through its own denoised prediction)
    The probe lies on the denoising trajectory, giving a curvature estimate for the
    Heun correction. Remaining steps: standard 1-call Euler Ancestral.

    Args:
        eta: Ancestral noise coefficient. 0=deterministic, 1=full ancestral. Default: 1.0
        s_noise: Noise scale multiplier. Default: 1.0
        correction_phase: Fraction of steps that receive the 3-call correction. Default: 0.5
        smooth_phase: Use continuous log-sigma weighting instead of a binary cutoff. Default: False
        adaptive_noise_scale: Calibrate s_noise from a short warmup window inside the
            texture phase (0.5 < sigma < 5.0), then restart the run from x_initial with
            the correction applied from step 0. Default: False
    """
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])

    print(f"🔮 Mirror Correction Euler active (η: {eta:.2f}, s_noise: {s_noise:.2f})")
    print(f"   Correction Phase: {correction_phase:.2f}, Smooth Phase: {smooth_phase}")
    if adaptive_noise_scale:
        print("   Adaptive Noise Scale: ON (auto-calibrating s_noise)")

    noise_sampler = get_noise_sampler(x)
    n_steps = len(sigmas) - 1

    log_sigma_phase = None
    log_sigma_max = None
    smooth_denom = 1e-6
    if smooth_phase and n_steps > 0:
        sigma_max_val = sigmas[0].clamp(min=1e-6)
        phase_idx = min(int(correction_phase * n_steps), n_steps - 1)
        sigma_phase_val = sigmas[phase_idx].clamp(min=1e-6)
        log_sigma_max = torch.log(sigma_max_val).item()
        log_sigma_phase = torch.log(sigma_phase_val).item()
        smooth_denom = max(log_sigma_max - log_sigma_phase, 1e-6)

    # Adaptive Noise Scale state
    prev_denoised_raw = None
    prev_change_norm = None
    excess_samples = []
    adaptive_correction = None
    excess_bins = {'structural': [], 'texture': [], 'cleanup': []} if adaptive_noise_scale else None
    adaptive_bin_corrections = None
    x_initial = x.clone() if adaptive_noise_scale else None

    i = 0
    while i < n_steps:
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        progress = i / max(n_steps - 1, 1)

        denoised = model(x, sigma * s_in, **extra_args)
        denoised_raw = denoised.clone() if adaptive_noise_scale else None
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})

        d = to_d(x, sigma, denoised)

        if sigma_next > 0:
            sigma_up = min(sigma_next, eta * (sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2) ** 0.5)
            sigma_down = (sigma_next ** 2 - sigma_up ** 2) ** 0.5
        else:
            sigma_up = 0.0
            sigma_down = 0.0
        dt = sigma_down - sigma

        if smooth_phase and log_sigma_phase is not None:
            log_sig = torch.log(sigma.clamp(min=1e-6)).item()
            t = max(0.0, min(1.0, (log_sig - log_sigma_phase) / smooth_denom))
            correction_weight = t ** 0.5

            if correction_weight > 1e-3 and sigma_next > 0:
                x_probe = 2 * denoised - x
                d_probe = to_d(x_probe, sigma, model(x_probe, sigma * s_in, **extra_args))

                d_diff_norm = (d - d_probe).norm()
                d_scale = (d.norm() + d_probe.norm()) / 2 + 1e-6
                gradient_agreement = max(0.0, 1.0 - (d_diff_norm / d_scale).item())
                effective_weight = correction_weight * gradient_agreement

                if effective_weight > 1e-3:
                    x3 = x + ((d + d_probe) / 2) * dt
                    d3 = to_d(x3, sigma, model(x3, sigma * s_in, **extra_args))
                    d_heun = (d + d3) / 2
                    if not (torch.isnan(d_heun).any() or torch.isinf(d_heun).any()):
                        d = d + effective_weight * (d_heun - d)
        else:
            if progress < correction_phase and sigma_next > 0:
                x_probe = 2 * denoised - x
                d_probe = to_d(x_probe, sigma, model(x_probe, sigma * s_in, **extra_args))
                x3 = x + ((d + d_probe) / 2) * dt
                d3 = to_d(x3, sigma, model(x3, sigma * s_in, **extra_args))
                d = (d + d3) / 2
                if torch.isnan(d).any() or torch.isinf(d).any():
                    d = torch.zeros_like(d)

        x = x + d * dt

        effective_s_noise = s_noise
        if adaptive_noise_scale:
            effective_s_noise, change_norm, adaptive_correction, should_restart, adaptive_bin_corrections = adaptive_noise_step(
                denoised_raw, prev_denoised_raw, prev_change_norm,
                sigma, sigma_next, excess_samples, adaptive_correction,
                effective_s_noise, i, n_steps,
                binned_enabled=True, excess_bins=excess_bins,
                adaptive_bin_corrections=adaptive_bin_corrections
            )
            if should_restart and x_initial is not None:
                x = x_initial.clone()
                prev_denoised_raw = None
                prev_change_norm = None
                excess_bins = {'structural': [], 'texture': [], 'cleanup': []}
                x_initial = None
                i = 0
                continue
            prev_change_norm = change_norm
            prev_denoised_raw = denoised_raw

        if sigma_next > 0:
            x = x + noise_sampler(sigma, sigma_next) * effective_s_noise * sigma_up

        i += 1

    return x


@torch.no_grad()
def sample_akashic_solver(model, x, sigmas, extra_args=None, callback=None, disable=None,
                           tau=0.5, eta=1.0, s_noise=1.0, adaptive_eta=True, phase_strength=0.5,
                           order=2, smea_strength=0.0, ndb_strength=0.0,
                           use_detail_enhancement=False, settings=None, eqvae_mode='Off',
                           adaptive_noise_scale=False):
    """
    AkashicSolver v2 [EXPERIMENTAL]: Advanced sampler optimized for EQ-VAE models.

    Combines:
    1. SA-SOLVER BASE: Multi-step Adams-Bashforth integration with tau function
    2. PHASE-AWARE SAMPLING: Three-phase approach with adaptive parameters
    3. SMEA COHERENCY: Sine-based interpolation for high-resolution coherency

    Args:
        eqvae_mode: EQ-VAE optimization mode ('Off' or 'Balanced')
        adaptive_noise_scale: Calibrate s_noise from a short warmup window inside the
            texture phase, then restart from x_initial with the correction applied
            from step 0 (also clears the multi-step d_history on restart)
    """
    extra_args = {} if extra_args is None else extra_args
    settings = settings or {}
    s_in = x.new_ones([x.shape[0]])

    if isinstance(eqvae_mode, bool):
        eqvae_enabled = eqvae_mode
    else:
        eqvae_enabled = eqvae_mode == 'Balanced'

    if eqvae_enabled:
        print("🌀 AkashicSolver v2 [EQ-VAE BALANCED] active")
        print("   Optimized for EQ-VAE's cleaner latent space")
    else:
        print("🌀 AkashicSolver v2 [EXPERIMENTAL] active")
    print(f"   τ (tau): {tau:.2f}, η (eta): {eta:.2f}, s_noise: {s_noise:.2f}")
    print(f"   Order: {order}, Adaptive Eta: {adaptive_eta}, Phase Strength: {phase_strength:.2f}")
    if smea_strength > 0:
        print(f"   SMEA: {smea_strength:.2f} (high-res coherency)")
    if ndb_strength > 0:
        print(f"   Native Detail Boost: {ndb_strength:.2f} (detail enhancement)")
    if adaptive_noise_scale:
        print("   Adaptive Noise Scale: ON (auto-calibrating s_noise)")
    if not eqvae_enabled:
        print("   ⚠️ Use external rescaleCFG (e.g., 0.7) for EQ-VAE models")
    
    active_model = model
    # use_detail_enhancement is always False from current call-sites.
    if use_detail_enhancement and TORCHVISION_AVAILABLE:
        active_model = create_detail_enhanced_model(model, x, sigmas, settings)

    noise_sampler = get_noise_sampler(x)
    total_steps = len(sigmas) - 1
    d_history = []

    # Adaptive Noise Scale state
    prev_denoised_raw = None
    prev_change_norm = None
    excess_samples = []
    adaptive_correction = None
    excess_bins = {'structural': [], 'texture': [], 'cleanup': []} if adaptive_noise_scale else None
    adaptive_bin_corrections = None
    x_initial = x.clone() if adaptive_noise_scale else None

    i = 0
    while i < total_steps:
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        progress = i / max(total_steps - 1, 1)
        
        if adaptive_eta:
            if eqvae_enabled:
                current_tau = compute_eqvae_tau(progress, tau, phase_strength)
            else:
                current_tau = compute_tau_eqvae(progress, tau, phase_strength)
        else:
            current_tau = tau

        if adaptive_eta:
            if eqvae_enabled:
                if progress < 0.25:
                    current_eta = eta * (1.0 + 0.03 * phase_strength)
                elif progress < 0.55:
                    current_eta = eta * (1.0 - 0.03 * phase_strength)
                else:
                    current_eta = eta * (1.0 + 0.02 * phase_strength)
            else:
                if progress < 0.30:
                    current_eta = eta * (1.0 + 0.08 * phase_strength)
                elif progress < 0.60:
                    current_eta = eta * (1.0 - 0.05 * phase_strength)
                else:
                    current_eta = eta * (1.0 + 0.02 * phase_strength)
        else:
            current_eta = eta
        
        smea_factor = compute_smea_factor(progress, smea_strength)
        
        denoised = active_model(x, sigma * s_in, **extra_args)
        denoised_raw = denoised.clone() if adaptive_noise_scale else None
        
        cfg_scale = extra_args.get('cond_scale', 1.0)
        if cfg_scale > 7.0:
            denoised = apply_dynamic_thresholding(denoised, percentile=0.995)
        
        d = to_d(x, sigma, denoised)
        
        derivative_max = torch.abs(d).max()
        sigma_adaptive_threshold = 1000.0 * (1.0 + sigma / 10.0)
        if torch.isnan(d).any() or torch.isinf(d).any() or derivative_max > sigma_adaptive_threshold:
            d = torch.clamp(d, -sigma_adaptive_threshold, sigma_adaptive_threshold)
            if torch.isnan(d).any() or torch.isinf(d).any():
                d = torch.zeros_like(d)
        
        d_history.append((sigma, d))
        if len(d_history) > order:
            d_history.pop(0)
        
        effective_tau = current_tau

        if eqvae_enabled:
            effective_s_noise = compute_eqvae_noise_scale(s_noise * current_eta, progress) * smea_factor
        else:
            effective_s_noise = s_noise * current_eta * smea_factor

            if progress < 0.30:
                noise_multiplier = 1.0 + 0.03 * phase_strength
            elif progress < 0.60:
                noise_multiplier = 1.0 - 0.01 * phase_strength
            else:
                noise_multiplier = 1.0 - 0.02 * phase_strength

            effective_s_noise *= noise_multiplier

        should_restart = False
        if adaptive_noise_scale:
            effective_s_noise, change_norm, adaptive_correction, should_restart, adaptive_bin_corrections = adaptive_noise_step(
                denoised_raw, prev_denoised_raw, prev_change_norm,
                sigma, sigma_next, excess_samples, adaptive_correction,
                effective_s_noise, i, total_steps,
                binned_enabled=True, excess_bins=excess_bins,
                adaptive_bin_corrections=adaptive_bin_corrections
            )
            if not should_restart:
                prev_change_norm = change_norm
                prev_denoised_raw = denoised_raw

        if should_restart and x_initial is not None:
            x = x_initial.clone()
            d_history = []
            prev_denoised_raw = None
            prev_change_norm = None
            excess_bins = {'structural': [], 'texture': [], 'cleanup': []}
            x_initial = None
            i = 0
            continue

        if eqvae_enabled and ndb_strength > 0:
            eqvae_blur_sigma, _ = compute_eqvae_ndb(progress, ndb_strength)
        else:
            eqvae_blur_sigma = None

        x, sigma_up = sa_solver_step(
            x=x,
            d_history=d_history,
            sigma=sigma,
            sigma_next=sigma_next,
            tau=effective_tau,
            s_noise=effective_s_noise,
            noise_sampler=noise_sampler,
            order=order,
            ndb_strength=ndb_strength,
            progress=progress,
            eqvae_mode=eqvae_enabled,
            eqvae_blur_sigma=eqvae_blur_sigma
        )
        
        if torch.isnan(x).any() or torch.isinf(x).any():
            print(f"❌ AkashicSolver v2: NaN/Inf detected at step {i}/{total_steps}!")
            
            if i == 0:
                raise RuntimeError("NaN/Inf on first step - check model/inputs")
            
            print("   Attempting recovery...")
            denoised_safe = active_model(x, sigma * s_in, **extra_args)
            if torch.isnan(denoised_safe).any():
                raise RuntimeError("Model producing NaN - reduce CFG scale or check model")
            
            d_safe = to_d(x, sigma, denoised_safe)
            dt_safe = (sigma_next - sigma) * 0.5
            x = x + d_safe * dt_safe
            
            d_history.clear()
            print("   Recovery successful. Multi-step history cleared.")
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})

        i += 1
    
    return x


# ============================================================================
# ============================================================================
# ============================================================================
# This file continues from PART1 + PART1B
# ============================================================================

# ============================================================================
# WEIGHT PATCHER (from v3 - unchanged)
# ============================================================================

def should_patch_weights(unet_model, scale, shift):
    """Return True if weight patching is actually needed for these parameters."""
    return (
        unet_model is not None and
        (abs(scale - 1.0) > 1e-6 or abs(shift) > 1e-6)
    )


class AdeptWeightPatcher:
    """Temporary weight scaling for UNet."""
    
    def __init__(self, unet_model, scale=1.0, shift=0.0):
        self.unet_model = unet_model
        self.scale = scale
        self.shift = shift
        self.backups = {}
        self.target_layers = []
    
    def __enter__(self):
        if self.unet_model is None or (abs(self.scale - 1.0) < 1e-6 and abs(self.shift) < 1e-6):
            return self
        
        self.target_layers.clear()
        self.backups.clear()
        
        try:
            for name, module in self.unet_model.named_modules():
                if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear):
                    if hasattr(module, 'weight') and module.weight is not None:
                        self.target_layers.append((name, module))
                        self.backups[name] = module.weight.data.clone()
                        module.weight.data = module.weight.data * self.scale + self.shift
        except Exception as e:
            print(f"❌ Weight patcher failed: {e}")
            self.__exit__(None, None, None)
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            for name, module in self.target_layers:
                if name in self.backups:
                    module.weight.data.copy_(self.backups[name])
            self.backups.clear()
            self.target_layers.clear()
        except Exception as e:
            print(f"❌ CRITICAL: Failed to restore weights: {e}")
            for name, backup_data in self.backups.items():
                try:
                    for n, m in self.target_layers:
                        if n == name:
                            m.weight.data.copy_(backup_data)
                except:
                    pass
        
        return False


# ============================================================================
# VAE REFLECTION PATCHER (from v3 - unchanged)
# ============================================================================

class VAEReflectionPatcher:
    """Context manager for VAE reflection padding."""
    
    def __init__(self, vae_model):
        self.vae_model = vae_model
        self.backups = {}
    
    def __enter__(self):
        global _vae_reflection_active
        
        if _vae_reflection_active or self.vae_model is None:
            return self
        
        _vae_original_padding_modes.clear()
        patched_count = 0
        
        try:
            for name, module in self.vae_model.named_modules():
                if isinstance(module, torch.nn.Conv2d):
                    _vae_original_padding_modes[name] = module.padding_mode
                    module.padding_mode = 'reflect'
                    patched_count += 1
            
            _vae_reflection_active = True
            print(f"🪞 VAE Reflection: Patched {patched_count} Conv2d layers")
        except Exception as e:
            print(f"❌ VAE Reflection failed: {e}")
            self.__exit__(None, None, None)
        
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        global _vae_reflection_active
        
        if self.vae_model is None:
            _vae_reflection_active = False
            _vae_original_padding_modes.clear()
            return False
        
        restored_count = 0
        try:
            for name, module in self.vae_model.named_modules():
                if isinstance(module, torch.nn.Conv2d) and name in _vae_original_padding_modes:
                    module.padding_mode = _vae_original_padding_modes[name]
                    restored_count += 1
            
            _vae_reflection_active = False
            _vae_original_padding_modes.clear()
            print(f"🔄 VAE Reflection: Restored {restored_count} layers")
        except Exception as e:
            print(f"⚠️ VAE Reflection restore warning: {e}")
        
        return False


def force_restore_vae_reflection():
    """
    Emergency / unload-path restore for VAE padding modes.
    Safe to call at any time — does nothing if VAE reflection was not active.
    """
    global _vae_reflection_active
    if not _vae_reflection_active and not _vae_original_padding_modes:
        return
    try:
        sd = getattr(shared, "sd_model", None)
        vae = getattr(sd, "first_stage_model", None) if sd else None
        if vae is not None:
            restored = 0
            for name, module in vae.named_modules():
                if isinstance(module, torch.nn.Conv2d) and name in _vae_original_padding_modes:
                    module.padding_mode = _vae_original_padding_modes[name]
                    restored += 1
            if restored:
                print(f"🔄 VAE Reflection: force-restored {restored} layers")
    except Exception as e:
        print(f"⚠️ VAE Reflection force-restore warning: {e}")
    finally:
        _vae_reflection_active = False
        _vae_original_padding_modes.clear()


# ALL SCHEDULERS (18 types)
# ============================================================================

def create_aos_v_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """AOS-V (Anime-Optimized Schedule for v-prediction models)."""
    rho = 7.0
    
    p1_steps = int(num_steps * 0.2)
    p2_steps = int(num_steps * 0.6)
    
    ramp = torch.empty(num_steps, device=device, dtype=torch.float32)
    
    if p1_steps > 0:
        torch.linspace(0, 1, p1_steps, out=ramp[:p1_steps])
        ramp[:p1_steps].pow_(0.5).mul_(0.6)
    
    if p2_steps > p1_steps:
        torch.linspace(0.6, 0.9, p2_steps - p1_steps, out=ramp[p1_steps:p2_steps])
    
    if num_steps > p2_steps:
        torch.linspace(0, 1, num_steps - p2_steps, out=ramp[p2_steps:])
        ramp[p2_steps:].pow_(3).mul_(0.1).add_(0.9)
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    ramp.mul_(min_inv_rho - max_inv_rho).add_(max_inv_rho).pow_(rho)
    
    return torch.cat([ramp, torch.zeros(1, device=device)])


def create_aos_e_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """AOS-ε (Anime-Optimized Schedule for epsilon-prediction models)."""
    rho = 7.0
    
    p1_frac, p2_frac = 0.35, 0.7
    ramp_p1_val, ramp_p2_val = 0.4, 0.75

    p1_steps = int(num_steps * p1_frac)
    p2_steps = int(num_steps * p2_frac)

    phase1_ramp = torch.linspace(0, 1, p1_steps, device=device) ** 1.5 * ramp_p1_val
    phase2_ramp = torch.linspace(ramp_p1_val, ramp_p2_val, p2_steps - p1_steps, device=device)
    phase3_base = torch.linspace(0, 1, num_steps - p2_steps, device=device) ** 0.7
    phase3_ramp = phase3_base * (1 - ramp_p2_val) + ramp_p2_val
    
    if p1_steps == 0: phase1_ramp = torch.empty(0, device=device)
    if p2_steps - p1_steps == 0: phase2_ramp = torch.empty(0, device=device)
    if num_steps - p2_steps == 0: phase3_ramp = torch.empty(0, device=device)

    ramp = torch.cat([phase1_ramp, phase2_ramp, phase3_ramp])
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_aos_akashic_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """AkashicAOS v2: Detail-Progressive Schedule for EQ-VAE SDXL models."""
    rho = 7.0
    
    u = torch.linspace(0, 1, num_steps, device=device)
    
    detail_power = 0.85
    u_progressive = u ** detail_power
    
    mid_boost_strength = 0.08
    mid_boost = mid_boost_strength * torch.sin(math.pi * u) * (1 - u * 0.5)
    
    u_modulated = u_progressive + mid_boost
    
    u_min, u_max = u_modulated.min(), u_modulated.max()
    if u_max - u_min > 1e-8:
        u_modulated = (u_modulated - u_min) / (u_max - u_min)
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + u_modulated * (min_inv_rho - max_inv_rho)) ** rho
    
    for i in range(1, len(sigmas)):
        if sigmas[i] >= sigmas[i-1]:
            sigmas[i] = sigmas[i-1] * 0.995
        max_ratio = 1.5
        if i > 0 and sigmas[i-1] / sigmas[i] > max_ratio:
            sigmas[i] = sigmas[i-1] / max_ratio
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_entropic_sigmas(sigma_max, sigma_min, num_steps, power=6.0, device='cpu'):
    """Entropic power schedule."""
    rho = 7.0
    
    linear_ramp = torch.linspace(0, 1, num_steps, device=device)
    power_ramp = 1 - torch.linspace(1, 0, num_steps, device=device) ** power
    
    ramp = (linear_ramp + power_ramp) / 2.0
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_snr_optimized_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Schedule optimized around log SNR = 0 region."""
    log_snr_max = 2 * torch.log(sigma_max)
    log_snr_min = 2 * torch.log(sigma_min)
    
    t = torch.linspace(0, 1, num_steps, device=device)
    
    concentration_power = 3.0
    sigmoid_t = torch.sigmoid(concentration_power * (t - 0.5))
    
    linear_t = t
    blend_factor = 0.7
    combined_t = blend_factor * sigmoid_t + (1 - blend_factor) * linear_t
    
    log_snr = log_snr_max + combined_t * (log_snr_min - log_snr_max)
    sigmas = torch.exp(log_snr / 2)
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_constant_rate_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Constant rate of distributional change."""
    rho = 7.0
    
    t = torch.linspace(0, 1, num_steps, device=device)
    corrected_t = t + 0.3 * torch.sin(math.pi * t) * (1 - t)
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + corrected_t * (min_inv_rho - max_inv_rho)) ** rho
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_adaptive_optimized_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Adaptive schedule combining multiple strategies."""
    rho = 7.0
    
    base_t = torch.linspace(0, 1, num_steps, device=device)
    
    strategies = [
        lambda t: t,
        lambda t: t ** 0.8,
        lambda t: t + 0.2 * torch.sin(2 * math.pi * t) * (1 - t),
        lambda t: 1 / (1 + torch.exp(-3 * (t - 0.5))),
    ]
    
    weights = [0.2, 0.3, 0.2, 0.3]
    combined_t = sum(w * s(base_t) for w, s in zip(weights, strategies))
    
    if (combined_t.max() - combined_t.min()) > 1e-6:
        combined_t = (combined_t - combined_t.min()) / (combined_t.max() - combined_t.min())
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + combined_t * (min_inv_rho - max_inv_rho)) ** rho
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_cosine_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Cosine-annealed schedule."""
    rho = 7.0
    u = torch.linspace(0, 1, num_steps, device=device)
    t = (1 - torch.cos(math.pi * u)) / 2
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + t * (min_inv_rho - max_inv_rho)) ** rho
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_logsnr_uniform_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Uniform in log-SNR space."""
    u = torch.linspace(0, 1, num_steps, device=device)
    log_snr_max = 2 * torch.log(sigma_max)
    log_snr_min = 2 * torch.log(sigma_min)
    log_snr = log_snr_max + u * (log_snr_min - log_snr_max)
    sigmas = torch.exp(log_snr / 2)
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_tanh_midboost_sigmas(sigma_max, sigma_min, num_steps, device='cpu', k=4.0):
    """Concentrate steps near mid-range sigmas."""
    rho = 7.0
    u = torch.linspace(0, 1, num_steps, device=device)
    k_tensor = torch.tensor(k, device=device, dtype=u.dtype)
    t = 0.5 * (torch.tanh(k_tensor * (u - 0.5)) / torch.tanh(k_tensor / 2) + 1.0)
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + t * (min_inv_rho - max_inv_rho)) ** rho
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_exponential_tail_sigmas(sigma_max, sigma_min, num_steps, device='cpu', pivot=0.7, gamma=0.8, beta=5.0):
    """Faster early lock-in with extra resolution in final steps."""
    rho = 7.0
    u = torch.linspace(0, 1, num_steps, device=device)
    
    early_mask = u < pivot
    late_mask = ~early_mask
    
    t = torch.empty_like(u)
    t[early_mask] = (u[early_mask] / pivot) ** gamma * pivot
    late_u = u[late_mask]
    t[late_mask] = pivot + (1 - pivot) * (1 - torch.exp(-beta * (late_u - pivot) / (1 - pivot)))
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + t * (min_inv_rho - max_inv_rho)) ** rho
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_jittered_karras_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Karras schedule with controlled jitter."""
    if num_steps <= 0:
        return torch.cat([sigma_max.unsqueeze(0), torch.zeros(1, device=device)])
    
    rho = 7.0
    indices = torch.arange(num_steps, device=device, dtype=torch.float32)
    denom = max(1, num_steps - 1)
    
    base = (indices + 0.5) / denom
    jitter_seed = torch.sin((indices + 1) * 2.3999632)
    jitter_strength = 0.35
    jitter = jitter_seed * jitter_strength / denom
    
    u = torch.clamp(base + jitter, 0.0, 1.0)
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + u * (min_inv_rho - max_inv_rho)) ** rho
    
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_stochastic_sigmas(sigma_max, sigma_min, num_steps, device='cpu', noise_type='brownian', noise_scale=0.3, base_schedule='karras'):
    """Stochastic scheduler with controlled randomness."""
    rho = 7.0
    
    # Base schedule
    if base_schedule == 'karras':
        indices = torch.arange(num_steps, device=device, dtype=torch.float32)
        u = (indices / max(1, num_steps - 1)) ** (1 / rho)
    elif base_schedule == 'cosine':
        u = torch.linspace(0, 1, num_steps, device=device)
        u = (1 - torch.cos(math.pi * u)) / 2
    else:  # uniform
        u = torch.linspace(0, 1, num_steps, device=device)
    
    # Add noise
    if noise_type == 'brownian':
        noise = torch.randn(num_steps, device=device).cumsum(0)
        if noise.numel() > 1:
            noise = noise / noise.std()
        # else: nothing to normalize relative to with a single point --
        # leave as-is (torch.std() of one element is NaN by definition)
    elif noise_type == 'uniform':
        noise = torch.rand(num_steps, device=device) * 2 - 1
    else:  # normal
        noise = torch.randn(num_steps, device=device)
    
    u_noisy = u + noise * noise_scale / num_steps
    u_noisy = torch.clamp(u_noisy, 0, 1)
    
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + u_noisy * (min_inv_rho - max_inv_rho)) ** rho

    # Sort the final sigmas descending so schedule is always noise→clean.
    # Sorting u_noisy descending before the transform gives wrong order
    # because the Karras mapping is monotone-decreasing in u.
    sigmas, _ = torch.sort(sigmas, descending=True)

    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_jys_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """
    JYS (Jump Your Steps) schedule using dynamically computed timestep sequences.
    Strategy: Large jumps early, dense clustering in detail region, fine steps at end.
    Ported from ComfyUI reference implementation.
    """
    # _compute_jys_timesteps returns num_steps entries + a trailing 0.
    # Strip the trailing 0 so we get exactly num_steps timesteps; the
    # explicit zeros(1) terminator is appended below.
    jys_timesteps = _compute_jys_timesteps(num_steps)
    if jys_timesteps and jys_timesteps[-1] == 0:
        jys_timesteps = jys_timesteps[:-1]

    rho = 7.0

    normalized_timesteps = [(1000 - t) / 1000.0 for t in jys_timesteps]
    t_tensor = torch.tensor(normalized_timesteps, device=device, dtype=torch.float32)

    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + t_tensor * (min_inv_rho - max_inv_rho)) ** rho

    sigmas, _ = torch.sort(sigmas, descending=True)
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def _compute_jys_timesteps(num_steps):
    """Dynamically compute optimised JYS timestep sequence (0..1000 scale)."""
    if num_steps <= 0:
        return [0]
    if num_steps == 1:
        return [1000, 0]
    elif num_steps == 2:
        return [1000, 500, 0]
    elif num_steps == 3:
        return [1000, 600, 200, 0]

    early_steps  = max(1, int(num_steps * 0.2))
    final_steps  = max(1, int(num_steps * 0.2))
    middle_steps = max(1, num_steps - early_steps - final_steps)

    early_jump_size = max(50, (1000 - 600) // early_steps)
    early_timesteps = []
    current_t = 1000
    for _ in range(early_steps):
        early_timesteps.append(int(current_t))
        current_t = max(600, current_t - early_jump_size)

    middle_timesteps = []
    structure_steps     = max(1, middle_steps // 2)
    structure_jump_size = max(10, (600 - 300) // structure_steps)
    current_t = 600
    for _ in range(structure_steps):
        middle_timesteps.append(int(current_t))
        current_t = max(300, current_t - structure_jump_size)

    detail_steps = middle_steps - structure_steps
    if detail_steps > 0:
        detail_jump_size = max(5, (300 - 200) // detail_steps)
        current_t = 300
        for _ in range(detail_steps):
            middle_timesteps.append(int(current_t))
            current_t = max(200, current_t - detail_jump_size)

    final_start     = min(middle_timesteps) if middle_timesteps else 200
    final_jump_size = max(5, final_start // final_steps)
    final_timesteps = []
    current_t = final_start
    for _ in range(final_steps):
        final_timesteps.append(int(current_t))
        current_t = max(0, current_t - final_jump_size)

    all_timesteps    = early_timesteps + middle_timesteps + final_timesteps
    unique_timesteps = list(dict.fromkeys(all_timesteps))
    unique_timesteps.sort(reverse=True)

    while len(unique_timesteps) < num_steps:
        for i in range(len(unique_timesteps) - 1):
            mid_point = (unique_timesteps[i] + unique_timesteps[i + 1]) // 2
            if mid_point not in unique_timesteps:
                unique_timesteps.insert(i + 1, mid_point)
                if len(unique_timesteps) >= num_steps:
                    break

    if len(unique_timesteps) > num_steps:
        unique_timesteps = unique_timesteps[:num_steps]

    if unique_timesteps[-1] != 0:
        unique_timesteps.append(0)

    return unique_timesteps


def create_hybrid_jys_karras_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """Hybrid: JYS mid-phase with Karras locks."""
    if num_steps <= 0:
        return torch.cat([sigma_max.unsqueeze(0), torch.zeros(1, device=device)])

    rho = 7.0

    jys_sigmas = create_jys_sigmas(sigma_max, sigma_min, num_steps, device=device)[:-1]

    indices = torch.arange(num_steps, device=device, dtype=torch.float32)
    denom = max(1, num_steps - 1)
    base = (indices + 0.5) / denom
    jitter_seed = torch.sin((indices + 1) * 2.3999632)
    jitter_strength = 0.35
    jitter = jitter_seed * jitter_strength / denom
    u = torch.clamp(base + jitter, 0.0, 1.0)

    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    karras_sigmas = (max_inv_rho + u * (min_inv_rho - max_inv_rho)) ** rho

    positions = torch.linspace(0, 1, num_steps, device=device)
    jys_weight = torch.empty_like(positions)
    early_mask = positions < 0.3
    mid_mask = (positions >= 0.3) & (positions < 0.8)
    late_mask = positions >= 0.8
    jys_weight[early_mask] = 0.2 + 0.4 * (positions[early_mask] / 0.3)
    jys_weight[mid_mask] = 0.6 + 0.3 * ((positions[mid_mask] - 0.3) / 0.5)
    jys_weight[late_mask] = 0.9
    jys_weight = jys_weight.clamp(0.2, 0.9)

    log_jys = torch.log(jys_sigmas.clamp_min(1e-6))
    log_karras = torch.log(karras_sigmas.clamp_min(1e-6))
    log_hybrid = torch.lerp(log_karras, log_jys, jys_weight)

    hybrid = torch.exp(log_hybrid)

    smoothing = 1.0 - 0.05 * (1 - positions) ** 2
    hybrid = hybrid * smoothing

    for i in range(1, hybrid.shape[0]):
        if hybrid[i] > hybrid[i - 1]:
            hybrid[i] = hybrid[i - 1] * 0.999

    return torch.cat([hybrid, torch.zeros(1, device=device)])


def create_ays_sdxl_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """AYS (Align Your Steps) optimized for SDXL."""

    if num_steps <= 1:
        # The interpolation path's normalized_np[-1] = 0.0 (meant to bring
        # the point just before the terminator close to zero across
        # multiple points) degenerates to zeroing the only point when
        # there's just one, producing a no-op [0, 0] schedule. One step
        # should simply go straight from sigma_max to the terminator.
        return torch.cat([
            torch.full((1,), float(sigma_max), device=device, dtype=torch.float32),
            torch.zeros(1, device=device)
        ])
    
    AYS_SCHEDULES = {
        10: [1.0000, 0.8751, 0.7502, 0.6254, 0.5004, 0.3755, 0.2506, 0.1253, 0.0502, 0.0000],
        15: [1.0000, 0.9167, 0.8334, 0.7501, 0.6668, 0.5835, 0.5002, 0.4169, 0.3336, 
             0.2503, 0.1670, 0.0837, 0.0335, 0.0084, 0.0000],
        20: [1.0000, 0.9375, 0.8750, 0.8125, 0.7500, 0.6875, 0.6250, 0.5625, 0.5000,
             0.4375, 0.3750, 0.3125, 0.2500, 0.1875, 0.1250, 0.0625, 0.0313, 0.0156, 
             0.0039, 0.0000],
        25: [1.0000, 0.9500, 0.9000, 0.8500, 0.8000, 0.7500, 0.7000, 0.6500, 0.6000,
             0.5500, 0.5000, 0.4500, 0.4000, 0.3500, 0.3000, 0.2500, 0.2000, 0.1500,
             0.1000, 0.0625, 0.0391, 0.0195, 0.0098, 0.0024, 0.0000],
        30: [1.0000, 0.9583, 0.9167, 0.8750, 0.8333, 0.7917, 0.7500, 0.7083, 0.6667,
             0.6250, 0.5833, 0.5417, 0.5000, 0.4583, 0.4167, 0.3750, 0.3333, 0.2917,
             0.2500, 0.2083, 0.1667, 0.1250, 0.0833, 0.0521, 0.0326, 0.0163, 0.0081,
             0.0041, 0.0010, 0.0000],
    }
    
    if num_steps in AYS_SCHEDULES:
        normalized = torch.tensor(AYS_SCHEDULES[num_steps], device=device, dtype=torch.float32)
    else:
        available_steps = sorted(AYS_SCHEDULES.keys())
        
        if num_steps < available_steps[0]:
            ref_steps = available_steps[0]
        elif num_steps > available_steps[-1]:
            ref_steps = available_steps[-1]
        else:
            ref_steps = min([s for s in available_steps if s >= num_steps], default=available_steps[-1])
        
        ref_schedule = np.array(AYS_SCHEDULES[ref_steps])
        
        t_ref = np.linspace(0, 1, len(ref_schedule))
        t_new = np.linspace(0, 1, num_steps)
        
        log_ref = np.log(ref_schedule + 1e-8)
        log_ref[-1] = log_ref[-2] - 3.0
        
        log_interp = np.interp(t_new, t_ref, log_ref)
        normalized_np = np.exp(log_interp)
        normalized_np[-1] = 0.0
        
        normalized = torch.tensor(normalized_np, device=device, dtype=torch.float32)
    
    sigma_range = sigma_max - sigma_min
    sigmas = normalized * sigma_range + sigma_min
    
    sigmas[0] = sigma_max
    sigmas[-1] = 0.0
    
    for i in range(1, len(sigmas) - 1):
        if sigmas[i] >= sigmas[i-1]:
            sigmas[i] = sigmas[i-1] * 0.999

    # Append zero-terminator so output has num_steps+1 entries like all other schedulers.
    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_aos_akashic_alt_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """
    AkashicAOS Alt: Karras-based schedule with EQ-VAE-tuned warping.
    Stronger detail-progressive bias (power=0.78) and shifted tanh crossover at t=0.55.
    Adaptive rho scales with step count for multi-step solver stability.
    """
    if num_steps <= 0:
        return torch.zeros(1, device=device)

    rho = min(11.0, max(7.0, 7.0 + 2.0 * (20.0 / max(num_steps, 10))))
    u = torch.linspace(0, 1, num_steps, device=device)

    detail_power = 0.78
    u_detail = u ** detail_power

    t_center = 0.55
    beta = 0.07
    gamma = 4.0
    crossover = beta * torch.tanh(gamma * (u - t_center))
    u_modulated = u_detail + crossover

    u_min, u_max = u_modulated.min(), u_modulated.max()
    if u_max - u_min > 1e-8:
        u_modulated = (u_modulated - u_min) / (u_max - u_min)

    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + u_modulated * (min_inv_rho - max_inv_rho)) ** rho

    max_ratio = 1.5
    for i in range(1, len(sigmas)):
        if sigmas[i] >= sigmas[i - 1]:
            sigmas[i] = sigmas[i - 1] * 0.995
        if sigmas[i - 1] / sigmas[i].clamp(min=1e-10) > max_ratio:
            sigmas[i] = sigmas[i - 1] / max_ratio

    return torch.cat([sigmas, torch.zeros(1, device=device)])


def create_akashic_eqflow_sigmas(sigma_max, sigma_min, num_steps, device='cpu'):
    """
    AkashicEQFlow: Robust crossover-focused log-SNR schedule for EQ-VAE models.
    Concentrates steps around the structure-to-detail transition in logSNR space,
    blended with a Karras prior. Adaptive density width + ratio slew-rate limiting.
    """
    if num_steps <= 0:
        return torch.zeros(1, device=device)

    lambda_min = -2.0 * math.log(max(float(sigma_max), 1e-10))
    lambda_max = -2.0 * math.log(max(float(sigma_min), 1e-10))
    lambda_range = max(lambda_max - lambda_min, 1e-8)

    step_factor = min(1.0, max(0.0, (num_steps - 16) / 30.0))
    lambda_center = 0.20 + 0.15 * step_factor
    u_center = (lambda_center - lambda_min) / lambda_range
    u_center = float(min(0.88, max(0.12, u_center)))

    concentration = min(3.2, max(1.35, 1.1 + num_steps / 16.0))
    base_width = min(0.30, max(0.18, 0.31 - 0.0028 * num_steps))
    width_left = base_width * 1.06
    width_right = base_width * 0.94
    detail_side_gain = 1.08 + 0.04 * step_factor

    N = 1200
    t = torch.linspace(0, 1, N, device=device)
    delta = t - u_center
    left_core = torch.exp(-((delta / width_left) ** 2) / 2.0)
    right_core = detail_side_gain * torch.exp(-((delta / width_right) ** 2) / 2.0)
    crossover_core = torch.where(delta <= 0, left_core, right_core)

    detail_floor = 0.08 * (t ** 1.4)
    composition_floor = 0.05 * ((1 - t) ** 1.7)
    density = 1.0 + concentration * crossover_core + detail_floor + composition_floor

    dt_val = 1.0 / (N - 1)
    cdf = torch.zeros(N, device=device)
    cdf[1:] = torch.cumsum((density[:-1] + density[1:]) * 0.5 * dt_val, dim=0)
    cdf = cdf / cdf[-1].clamp(min=1e-12)

    targets = torch.linspace(0, 1, num_steps, device=device)
    indices = torch.searchsorted(cdf, targets).clamp(1, N - 1)
    lo = indices - 1
    hi = indices
    frac = (targets - cdf[lo]) / (cdf[hi] - cdf[lo]).clamp(min=1e-12)
    u_steps = t[lo] + frac * (t[hi] - t[lo])

    lambdas_eqflow = lambda_min + u_steps * lambda_range

    rho = min(10.0, max(7.0, 7.0 + 1.5 * (22.0 / max(num_steps, 12))))
    u_karras = torch.linspace(0, 1, num_steps, device=device)
    min_inv_rho = sigma_min ** (1 / rho)
    max_inv_rho = sigma_max ** (1 / rho)
    sigmas_karras = (max_inv_rho + u_karras * (min_inv_rho - max_inv_rho)) ** rho
    lambdas_karras = -2.0 * torch.log(sigmas_karras.clamp(min=1e-10))

    blend_eqflow = min(0.60, max(0.35, 0.38 + num_steps / 200.0))
    lambdas = (1.0 - blend_eqflow) * lambdas_karras + blend_eqflow * lambdas_eqflow
    sigmas = torch.exp(-lambdas / 2.0)

    if num_steps >= 40:
        max_ratio = 1.50
    elif num_steps >= 28:
        max_ratio = 1.55
    elif num_steps >= 18:
        max_ratio = 1.65
    else:
        max_ratio = 1.85
    ratio_slew = 1.18
    prev_ratio = None

    sigmas[0] = sigma_max
    for i in range(1, len(sigmas)):
        if sigmas[i] >= sigmas[i - 1]:
            sigmas[i] = sigmas[i - 1] * 0.995
        ratio = float((sigmas[i - 1] / sigmas[i].clamp(min=1e-10)).item())
        ratio = min(ratio, max_ratio)
        if prev_ratio is not None:
            ratio = min(ratio, prev_ratio * ratio_slew)
            ratio = max(ratio, prev_ratio / ratio_slew)
        ratio = max(1.001, ratio)
        sigmas[i] = sigmas[i - 1] / ratio
        prev_ratio = ratio

    return torch.cat([sigmas, torch.zeros(1, device=device)])


def apply_custom_scheduler(sigmas, scheduler_type="Standard"):
    """
    Apply a custom sigma schedule.

    sigma_min uses sigmas[-2] (last non-zero step), never the zero-terminator.
    Each scheduler is invoked via a lambda so keyword args with non-standard
    defaults (e.g. Entropic's `power`) are always passed correctly.
    """
    if scheduler_type == "Standard" or len(sigmas) < 2:
        return sigmas

    sigma_max = sigmas[0]
    # Use the last non-zero sigma as sigma_min; sigmas[-1] is always 0.
    sigma_min = sigmas[-2] if len(sigmas) >= 2 else sigmas[0]
    if sigma_min <= 0:
        sigma_min = sigma_max * 1e-3
    num_steps = len(sigmas) - 1
    device    = sigmas.device

    scheduler_map = {
        "AOS-V":              lambda: create_aos_v_sigmas(sigma_max, sigma_min, num_steps, device),
        "AOS-Epsilon":        lambda: create_aos_e_sigmas(sigma_max, sigma_min, num_steps, device),
        "AkashicAOS":         lambda: create_aos_akashic_sigmas(sigma_max, sigma_min, num_steps, device),
        "Entropic":           lambda: create_entropic_sigmas(sigma_max, sigma_min, num_steps, power=6.0, device=device),
        "SNR-Optimized":      lambda: create_snr_optimized_sigmas(sigma_max, sigma_min, num_steps, device),
        "Constant-Rate":      lambda: create_constant_rate_sigmas(sigma_max, sigma_min, num_steps, device),
        "Adaptive-Optimized": lambda: create_adaptive_optimized_sigmas(sigma_max, sigma_min, num_steps, device),
        "Cosine-Annealed":    lambda: create_cosine_sigmas(sigma_max, sigma_min, num_steps, device),
        "LogSNR-Uniform":     lambda: create_logsnr_uniform_sigmas(sigma_max, sigma_min, num_steps, device),
        "Tanh Mid-Boost":     lambda: create_tanh_midboost_sigmas(sigma_max, sigma_min, num_steps, device),
        "Exponential Tail":   lambda: create_exponential_tail_sigmas(sigma_max, sigma_min, num_steps, device),
        "Jittered-Karras":    lambda: create_jittered_karras_sigmas(sigma_max, sigma_min, num_steps, device),
        "Stochastic":         lambda: create_stochastic_sigmas(sigma_max, sigma_min, num_steps, device=device),
        "JYS (Dynamic)":      lambda: create_jys_sigmas(sigma_max, sigma_min, num_steps, device),
        "Hybrid JYS-Karras":  lambda: create_hybrid_jys_karras_sigmas(sigma_max, sigma_min, num_steps, device),
        "AYS-SDXL":           lambda: create_ays_sdxl_sigmas(sigma_max, sigma_min, num_steps, device),
        "AkashicAOS Alt":     lambda: create_aos_akashic_alt_sigmas(sigma_max, sigma_min, num_steps, device),
        "AkashicEQFlow":      lambda: create_akashic_eqflow_sigmas(sigma_max, sigma_min, num_steps, device),
    }

    fn = scheduler_map.get(scheduler_type)
    if fn is not None:
        try:
            result = fn()
            if result is not None and len(result) > 1:
                return result
            print(f"⚠️ Scheduler {scheduler_type} returned empty/None, using standard")
        except Exception as e:
            print(f"⚠️ Scheduler {scheduler_type} failed: {e}, using standard")

    return sigmas


# ============================================================================
# ============================================================================
# K-DIFFUSION SAMPLERS with Custom Sampler Integration
# ============================================================================

@torch.no_grad()
def sample_adept_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., _skip_ans_check=False):
    """Euler sampler with Adept weight scaling OR custom sampler."""
    
    # CUSTOM SAMPLER INTEGRATION
    if ADEPT_STATE.get('enabled', False) and ADEPT_STATE.get('use_custom_sampler', False):
        custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
        print(f"🌀 Redirecting to {custom_type}")
        
        # Apply scheduler to sigmas for custom samplers
        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        if scheduler != "Standard":
            sigmas = apply_custom_scheduler(sigmas, scheduler)
            print(f"   📊 Applied {scheduler} scheduler")
        
        if custom_type == "Akashic Solver v2":
            return sample_akashic_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                tau=ADEPT_STATE.get('tau', 0.5),
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5),
                order=ADEPT_STATE.get('solver_order', 2),
                smea_strength=ADEPT_STATE.get('smea_strength', 0.0),
                ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
                use_detail_enhancement=False,
                settings={},
                eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Adept Solver":
            return sample_adept_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                order=ADEPT_STATE.get('solver_order', 2),
                use_corrector=ADEPT_STATE.get('use_corrector', True),
                use_detail_enhancement=False,
                settings={}
            )
        elif custom_type == "Adept Ancestral Solver":
            return sample_adept_ancestral_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', False),
                phase_noise=ADEPT_STATE.get('phase_noise', False),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5),
                enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
                use_detail_enhancement=False,
                settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Mirror Correction Euler":
            return sample_mirror_correction_euler(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
                smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
    
    # STANDARD K-DIFFUSION MODE (from v3 - unchanged)
    if not ADEPT_STATE.get('enabled', False):
        if 'euler' in ORIGINAL_SAMPLERS:
            return ORIGINAL_SAMPLERS['euler'](model, x, sigmas, extra_args, callback, disable, s_churn, s_tmin, s_tmax, s_noise)
        return _basic_euler(model, x, sigmas, extra_args, callback, disable)

    # Universal Adaptive Noise Scale for plain "Adept Euler" mode (no custom
    # solver active). _skip_ans_check guards the recursive re-entry from
    # run_ans_two_pass against infinite recursion.
    if not _skip_ans_check and ADEPT_STATE.get('adaptive_noise_scale', False):
        _self_partial = functools.partial(sample_adept_euler, _skip_ans_check=True)
        if 's_noise' in inspect.signature(_self_partial).parameters:
            return run_ans_two_pass(
                _self_partial, 'sample_euler', model, x, sigmas, extra_args, callback, disable,
                dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise)
            )

    # Apply custom scheduler deterministically (before the loop, not via p.sampler.model_wrap)
    _sched = ADEPT_STATE.get('scheduler', 'Standard')
    if _sched != 'Standard':
        sigmas = apply_custom_scheduler(sigmas, _sched)

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    base_scale = ADEPT_STATE.get('scale', 1.0)
    shift = ADEPT_STATE.get('shift', 0.0)
    start_pct = ADEPT_STATE.get('start_pct', 0.0)
    end_pct = ADEPT_STATE.get('end_pct', 1.0)
    
    try:
        unet_model = shared.sd_model.model.diffusion_model
    except AttributeError:
        unet_model = None
    
    total_steps = len(sigmas) - 1
    
    for i in range(total_steps):
        sigma = sigmas[i]
        gamma = min(s_churn / total_steps, 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0
        
        current_scale = compute_dynamic_scale(i, total_steps, base_scale, start_pct, end_pct)
        
        with AdeptWeightPatcher(unet_model, current_scale, shift):
            eps = torch.randn_like(x) * s_noise if gamma > 0 else 0
            sigma_hat = sigma * (gamma + 1)
            if gamma > 0:
                x = x + eps * (sigma_hat ** 2 - sigma ** 2) ** 0.5
            
            denoised = model(x, sigma_hat * s_in, **extra_args)
            d = to_d(x, sigma_hat, denoised)
            dt = sigmas[i + 1] - sigma_hat
            x = x + d * dt
            
            if callback is not None:
                callback({'x': x, 'i': i, 'sigma': sigma_hat, 'denoised': denoised})
    
    return x

def apply_progressive_enhancement(base_strength, phase, progress):
    """Ramp the high-frequency detail-enhancement strength across a phase.
    composition: gentle ramp 0.25x->0.75x. detail: stronger ramp 0.75x->1.5x.
    single_phase (pacing disabled but enhancement still on): 0.5x->1.5x."""
    if phase == 'composition':
        return base_strength * (0.25 + 0.5 * progress)
    elif phase == 'detail':
        return base_strength * (0.75 + 0.75 * progress)
    else:
        return base_strength * (0.5 + progress)


def create_pacing_detail_schedule(sigma_at_switch, sigma_min, num_steps, device, scheduler_name):
    """Build a fresh sigma schedule for the detail phase's remaining steps,
    respecting whichever custom scheduler is currently active. Reuses
    apply_custom_scheduler's dispatch table by constructing a probe tensor of
    the right length whose [0] and [-2] entries are the actual sigma_max/min
    it reads -- the intermediate values are ignored by every create_*_sigmas
    function, which only look at (sigma_max, sigma_min, num_steps, device).
    Falls back to a neutral Entropic curve when no custom scheduler is active,
    matching upstream's fallback for the same "Standard" case."""
    if num_steps <= 0:
        return torch.tensor([sigma_at_switch, 0.0], device=device)

    ramp = torch.linspace(float(sigma_at_switch), float(sigma_min), num_steps, device=device)
    probe = torch.cat([ramp, torch.zeros(1, device=device)])

    if scheduler_name == "Standard":
        return create_entropic_sigmas(probe[0], probe[-2], num_steps, power=1.0, device=device)
    return apply_custom_scheduler(probe, scheduler_name)


def _run_content_aware_pacing(model, x, sigmas, extra_args, callback, disable,
                               current_eta, current_s_noise, use_adaptive_eta, base_eta,
                               s_in, noise_sampler, unet_model, base_scale, shift,
                               start_pct, end_pct, total_steps):
    """
    Content-Aware Pacing variant of Adept Euler Ancestral.

    Composition phase: run steps normally while monitoring the variance of
    the derivative (x-denoised)/sigma. Once variance drops below a threshold
    derived from an established baseline (or a step-count-based fallback cap
    is hit), switch to the detail phase.

    Detail phase: regenerate a fresh sigma sub-schedule for the remaining
    steps (respecting the active scheduler) starting from the sigma where
    the switch happened, and continue sampling with it.

    Only called when total_steps >= 26 -- below that the caller falls back
    to the normal single-phase loop, matching upstream's own guard against
    phase-switching being counterproductive on very short runs.
    """
    sensitivity = ADEPT_STATE.get('pacing_coherence_sensitivity', 0.75)
    use_detail_fx = ADEPT_STATE.get('use_detail_enhancement_phase', False) and TORCHVISION_AVAILABLE
    fx_strength = ADEPT_STATE.get('detail_enhancement_strength', 0.05)
    fx_radius = ADEPT_STATE.get('detail_separation_radius', 0.5)
    scheduler_name = ADEPT_STATE.get('scheduler', 'Standard')

    original_sigmas = sigmas
    coherence_check_interval = 1 if total_steps <= 40 else 2
    fallback_step_pct = 0.4 + 0.3 * min(1.0, (total_steps - 20) / 40.0)
    max_composition_steps = max(1, int(total_steps * fallback_step_pct))

    print(f"🧠 Content-Aware Pacing active (sensitivity={sensitivity:.2f}, "
          f"coherence check every {coherence_check_interval} step(s), "
          f"composition budget={max_composition_steps}/{total_steps} steps)")

    def eta_for(progress):
        if use_adaptive_eta:
            if progress < 0.3:
                return base_eta * 1.08
            elif progress < 0.7:
                return base_eta * 0.95
            return base_eta * 1.02
        return current_eta

    def scaled_model_call(overall_idx, x_in, sigma_in):
        scale = compute_dynamic_scale(overall_idx, total_steps, base_scale, start_pct, end_pct)
        if should_patch_weights(unet_model, scale, shift):
            with AdeptWeightPatcher(unet_model, scale, shift):
                return model(x_in, sigma_in * s_in, **extra_args)
        return model(x_in, sigma_in * s_in, **extra_args)

    def detail_fx(x_in, denoised, dt, sigma_in, phase, progress):
        if not use_detail_fx:
            return x_in
        strength = apply_progressive_enhancement(fx_strength, phase, progress)
        low_freq = gaussian_blur(denoised, kernel_size=3, sigma=fx_radius)
        high_freq = denoised - low_freq
        enhancement_amount = dt.abs() / sigma_in.clamp(min=1e-6)
        return x_in + high_freq * enhancement_amount * strength

    # --- Composition phase ---
    initial_variance = None
    is_coherent = False
    composition_steps_taken = 0
    i = 0
    while i < (total_steps - 1) and composition_steps_taken < max_composition_steps:
        composition_steps_taken += 1
        current_sigma = original_sigmas[i]
        next_sigma = original_sigmas[i + 1]
        if current_sigma < next_sigma:
            break

        denoised = scaled_model_call(i, x, current_sigma)
        derivative = to_d(x, current_sigma, denoised)

        if composition_steps_taken % coherence_check_interval == 0 and composition_steps_taken >= 2:
            variance = torch.var(derivative.flatten(1), dim=1).mean().item()
            if initial_variance is None:
                initial_variance = variance
                print(f"   Baseline variance at step {composition_steps_taken}: {variance:.6f}")
            else:
                threshold = initial_variance * (sensitivity * 0.4 + 0.5)
                if variance < threshold:
                    print(f"   Coherence reached at step {composition_steps_taken} "
                          f"(variance={variance:.6f} < threshold={threshold:.6f}) -> detail phase")
                    is_coherent = True
                    break  # discard this iteration's step; detail phase restarts from here

        if callback is not None:
            callback({'x': x, 'i': composition_steps_taken - 1, 'sigma': current_sigma, 'denoised': denoised})

        step_eta = eta_for(i / max(total_steps, 1))
        sigma_down, sigma_up = get_ancestral_step(current_sigma, next_sigma, step_eta)
        dt = sigma_down - current_sigma
        x = x + derivative * dt

        x = detail_fx(x, denoised, dt, current_sigma, 'composition',
                      composition_steps_taken / max_composition_steps)

        if next_sigma > 0:
            x = x + noise_sampler(current_sigma, next_sigma) * current_s_noise * sigma_up

        i += 1

    sigma_idx_at_switch = i
    if not is_coherent:
        print(f"   Composition budget exhausted ({composition_steps_taken} steps) "
              f"without reaching coherence -> switching to detail phase anyway")

    # --- Detail phase ---
    remaining_steps = total_steps - composition_steps_taken
    if remaining_steps <= 0:
        return x

    safe_idx = min(sigma_idx_at_switch, len(original_sigmas) - 2)
    sigma_at_switch = original_sigmas[safe_idx]
    sigma_min = original_sigmas[-2]
    if sigma_at_switch <= sigma_min:
        sigma_at_switch = original_sigmas[len(original_sigmas) // 2]

    print(f"   Detail phase: {remaining_steps} steps, sigma {sigma_at_switch:.4f} -> {sigma_min:.4f}")
    detail_sigmas = create_pacing_detail_schedule(
        sigma_at_switch, sigma_min, remaining_steps, x.device, scheduler_name
    )

    for j in range(len(detail_sigmas) - 1):
        current_sigma = detail_sigmas[j]
        next_sigma = detail_sigmas[j + 1]
        if current_sigma < next_sigma:
            break

        overall_idx = composition_steps_taken + j
        denoised = scaled_model_call(overall_idx, x, current_sigma)
        derivative = to_d(x, current_sigma, denoised)

        if callback is not None:
            callback({'x': x, 'i': overall_idx, 'sigma': current_sigma, 'denoised': denoised})

        step_eta = eta_for(overall_idx / max(total_steps, 1))
        sigma_down, sigma_up = get_ancestral_step(current_sigma, next_sigma, step_eta)
        dt = sigma_down - current_sigma
        x = x + derivative * dt

        x = detail_fx(x, denoised, dt, current_sigma, 'detail', overall_idx / max(total_steps, 1))

        if next_sigma > 0:
            x = x + noise_sampler(current_sigma, next_sigma) * current_s_noise * sigma_up

    return x


def sample_adept_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, 
                                  disable=None, eta=1.0, s_noise=1.0, noise_sampler=None, _skip_ans_check=False):
    """Euler Ancestral with Adept weight scaling."""

    # CUSTOM SAMPLER INTEGRATION
    if ADEPT_STATE.get('enabled', False) and ADEPT_STATE.get('use_custom_sampler', False):
        custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
        print(f"🌀 Redirecting to {custom_type}")
        
        # Apply scheduler to sigmas for custom samplers
        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        if scheduler != "Standard":
            sigmas = apply_custom_scheduler(sigmas, scheduler)
            print(f"   📊 Applied {scheduler} scheduler")
        
        if custom_type == "Akashic Solver v2":
            return sample_akashic_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                tau=ADEPT_STATE.get('tau', 0.5), eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0), adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), order=ADEPT_STATE.get('solver_order', 2),
                smea_strength=ADEPT_STATE.get('smea_strength', 0.0), ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
                use_detail_enhancement=False, settings={}, eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Adept Solver":
            return sample_adept_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                order=ADEPT_STATE.get('solver_order', 2), use_corrector=ADEPT_STATE.get('use_corrector', True),
                use_detail_enhancement=False, settings={}
            )
        elif custom_type == "Adept Ancestral Solver":
            return sample_adept_ancestral_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', False), phase_noise=ADEPT_STATE.get('phase_noise', False),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
                use_detail_enhancement=False, settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Mirror Correction Euler":
            return sample_mirror_correction_euler(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
                smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
    

    
    if not ADEPT_STATE.get('enabled', False):
        if 'euler_ancestral' in ORIGINAL_SAMPLERS:
            return ORIGINAL_SAMPLERS['euler_ancestral'](model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
        return _basic_euler_ancestral(model, x, sigmas, extra_args, callback, disable, eta, s_noise)

    # Universal Adaptive Noise Scale for plain "Adept Euler A" mode (no custom
    # solver active). Deliberately skipped when Content-Aware Pacing is also
    # on: pacing already has its own mid-run restart mechanism, and running
    # two different restart mechanisms back-to-back on the same function
    # adds complexity for a combination neither reference implementation
    # actually supports together. _skip_ans_check guards the recursive
    # re-entry from run_ans_two_pass against infinite recursion.
    if (not _skip_ans_check and ADEPT_STATE.get('adaptive_noise_scale', False)
            and not ADEPT_STATE.get('use_content_aware_pacing', False)):
        _self_partial = functools.partial(sample_adept_euler_ancestral, _skip_ans_check=True)
        if 's_noise' in inspect.signature(_self_partial).parameters:
            return run_ans_two_pass(
                _self_partial, 'sample_euler_ancestral', model, x, sigmas, extra_args, callback, disable,
                dict(eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
            )

    # Apply custom scheduler deterministically (before the loop, not via p.sampler.model_wrap)
    _sched = ADEPT_STATE.get('scheduler', 'Standard')
    if _sched != 'Standard':
        sigmas = apply_custom_scheduler(sigmas, _sched)

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    # Get settings
    base_scale = ADEPT_STATE.get('scale', 1.0)
    shift = ADEPT_STATE.get('shift', 0.0)
    start_pct = ADEPT_STATE.get('start_pct', 0.0)
    end_pct = ADEPT_STATE.get('end_pct', 1.0)
    use_adaptive_eta = ADEPT_STATE.get('adaptive_eta', False)
    current_eta = ADEPT_STATE.get('eta', eta)
    current_s_noise = ADEPT_STATE.get('s_noise', s_noise)
    
    # Get UNet
    try:
        unet_model = shared.sd_model.model.diffusion_model
    except AttributeError:
        unet_model = None
    
    if noise_sampler is None:
        noise_sampler = default_noise_sampler(x)
    
    total_steps = len(sigmas) - 1

    if ADEPT_STATE.get('use_content_aware_pacing', False):
        if total_steps >= 26:
            return _run_content_aware_pacing(
                model, x, sigmas, extra_args, callback, disable,
                current_eta, current_s_noise, use_adaptive_eta, eta,
                s_in, noise_sampler, unet_model, base_scale, shift,
                start_pct, end_pct, total_steps
            )
        print("🧠 Pacing: disabled automatically for low step count (< 26) to ensure quality.")

    print(f"✅ Adept Euler A active: scale={base_scale:.2f}, eta={current_eta:.2f}")
    
    for i in trange(len(sigmas) - 1, disable=disable, desc="Adept Euler A"):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        progress = i / max(total_steps, 1)
        
        # Adaptive eta
        if use_adaptive_eta:
            if progress < 0.3:
                current_eta = eta * 1.08
            elif progress < 0.7:
                current_eta = eta * 0.95
            else:
                current_eta = eta * 1.02
        else:
            current_eta = eta
        
        # Dynamic scale
        current_scale = compute_dynamic_scale(i, total_steps, base_scale, start_pct, end_pct)
        
        # Evaluate model with weight patching
        if should_patch_weights(unet_model, current_scale, shift):
            with AdeptWeightPatcher(unet_model, current_scale, shift):
                denoised = model(x, sigma * s_in, **extra_args)
        else:
            denoised = model(x, sigma * s_in, **extra_args)
        
        # Euler Ancestral step
        sigma_down, sigma_up = get_ancestral_step(sigma, sigma_next, current_eta)
        d = to_d(x, sigma, denoised)
        
        if torch.isnan(d).any() or torch.isinf(d).any():
            d = torch.nan_to_num(d, nan=0.0, posinf=1.0, neginf=-1.0)
        
        dt = sigma_down - sigma
        x = x + d * dt
        
        if sigma_up > 0:
            noise = noise_sampler(sigma, sigma_next) * current_s_noise
            x = x + noise * sigma_up
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})
    
    return x


def _basic_euler(model, x, sigmas, extra_args=None, callback=None, disable=None):
    """Fallback basic Euler (used when ORIGINAL_SAMPLERS has no 'euler' key)."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    for i in trange(len(sigmas) - 1, disable=disable):
        sigma = sigmas[i]
        denoised = model(x, sigma * s_in, **extra_args)
        d = to_d(x, sigma, denoised)
        dt = sigmas[i + 1] - sigma
        x = x + d * dt
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})
    return x


def _basic_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0):
    """Fallback basic Euler Ancestral."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    noise_sampler = default_noise_sampler(x)
    
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta)
        d = to_d(x, sigmas[i], denoised)
        dt = sigma_down - sigmas[i]
        x = x + d * dt
        if sigma_up > 0:
            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})
    
    return x


@torch.no_grad()

def sample_adept_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., _skip_ans_check=False):
    """Heun sampler with Adept weight scaling."""

    # CUSTOM SAMPLER INTEGRATION
    if ADEPT_STATE.get('enabled', False) and ADEPT_STATE.get('use_custom_sampler', False):
        custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
        print(f"🌀 Redirecting to {custom_type}")
        
        # Apply scheduler to sigmas for custom samplers
        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        if scheduler != "Standard":
            sigmas = apply_custom_scheduler(sigmas, scheduler)
            print(f"   📊 Applied {scheduler} scheduler")
        
        if custom_type == "Akashic Solver v2":
            return sample_akashic_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                tau=ADEPT_STATE.get('tau', 0.5), eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0), adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), order=ADEPT_STATE.get('solver_order', 2),
                smea_strength=ADEPT_STATE.get('smea_strength', 0.0), ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
                use_detail_enhancement=False, settings={}, eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Adept Solver":
            return sample_adept_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                order=ADEPT_STATE.get('solver_order', 2), use_corrector=ADEPT_STATE.get('use_corrector', True),
                use_detail_enhancement=False, settings={}
            )
        elif custom_type == "Adept Ancestral Solver":
            return sample_adept_ancestral_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', False), phase_noise=ADEPT_STATE.get('phase_noise', False),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
                use_detail_enhancement=False, settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Mirror Correction Euler":
            return sample_mirror_correction_euler(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
                smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
    

    
    if not ADEPT_STATE.get('enabled', False):
        if 'heun' in ORIGINAL_SAMPLERS:
            return ORIGINAL_SAMPLERS['heun'](model, x, sigmas, extra_args, callback, disable, s_churn, s_tmin, s_tmax, s_noise)
        return _basic_heun(model, x, sigmas, extra_args, callback, disable)

    if not _skip_ans_check and ADEPT_STATE.get('adaptive_noise_scale', False):
        _self_partial = functools.partial(sample_adept_heun, _skip_ans_check=True)
        if 's_noise' in inspect.signature(_self_partial).parameters:
            return run_ans_two_pass(
                _self_partial, 'sample_heun', model, x, sigmas, extra_args, callback, disable,
                dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise)
            )

    # Apply custom scheduler deterministically (before the loop, not via p.sampler.model_wrap)
    _sched = ADEPT_STATE.get('scheduler', 'Standard')
    if _sched != 'Standard':
        sigmas = apply_custom_scheduler(sigmas, _sched)

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    # Get settings
    base_scale = ADEPT_STATE.get('scale', 1.0)
    shift = ADEPT_STATE.get('shift', 0.0)
    start_pct = ADEPT_STATE.get('start_pct', 0.0)
    end_pct = ADEPT_STATE.get('end_pct', 1.0)
    
    # Get UNet
    try:
        unet_model = shared.sd_model.model.diffusion_model
    except AttributeError:
        unet_model = None
    
    total_steps = len(sigmas) - 1
    print(f"✅ Adept Heun active: scale={base_scale:.2f}")
    
    for i in trange(len(sigmas) - 1, disable=disable, desc="Adept Heun"):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        # Dynamic scale
        current_scale = compute_dynamic_scale(i, total_steps, base_scale, start_pct, end_pct)
        
        # First evaluation
        if should_patch_weights(unet_model, current_scale, shift):
            with AdeptWeightPatcher(unet_model, current_scale, shift):
                denoised = model(x, sigma * s_in, **extra_args)
        else:
            denoised = model(x, sigma * s_in, **extra_args)
        
        d = to_d(x, sigma, denoised)
        
        if torch.isnan(d).any() or torch.isinf(d).any():
            d = torch.nan_to_num(d, nan=0.0, posinf=1.0, neginf=-1.0)
        
        dt = sigma_next - sigma
        
        if sigma_next == 0:
            # Last step
            x = x + d * dt
        else:
            # Heun's method: two-stage
            x_2 = x + d * dt
            
            # Second evaluation
            if should_patch_weights(unet_model, current_scale, shift):
                with AdeptWeightPatcher(unet_model, current_scale, shift):
                    denoised_2 = model(x_2, sigma_next * s_in, **extra_args)
            else:
                denoised_2 = model(x_2, sigma_next * s_in, **extra_args)
            
            d_2 = to_d(x_2, sigma_next, denoised_2)
            
            if torch.isnan(d_2).any() or torch.isinf(d_2).any():
                d_2 = torch.nan_to_num(d_2, nan=0.0, posinf=1.0, neginf=-1.0)
            
            # Average
            d_prime = (d + d_2) / 2
            x = x + d_prime * dt
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})
    
    return x


def _basic_heun(model, x, sigmas, extra_args=None, callback=None, disable=None):
    """Fallback basic Heun."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        d = to_d(x, sigmas[i], denoised)
        dt = sigmas[i + 1] - sigmas[i]
        
        if sigmas[i + 1] == 0:
            x = x + d * dt
        else:
            x_2 = x + d * dt
            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
            d_prime = (d + d_2) / 2
            x = x + d_prime * dt
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})
    
    return x


@torch.no_grad()

def sample_adept_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
    """DPM++ 2M sampler with Adept weight scaling."""

    # CUSTOM SAMPLER INTEGRATION
    if ADEPT_STATE.get('enabled', False) and ADEPT_STATE.get('use_custom_sampler', False):
        custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
        print(f"🌀 Redirecting to {custom_type}")
        
        # Apply scheduler to sigmas for custom samplers
        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        if scheduler != "Standard":
            sigmas = apply_custom_scheduler(sigmas, scheduler)
            print(f"   📊 Applied {scheduler} scheduler")
        
        if custom_type == "Akashic Solver v2":
            return sample_akashic_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                tau=ADEPT_STATE.get('tau', 0.5), eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0), adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), order=ADEPT_STATE.get('solver_order', 2),
                smea_strength=ADEPT_STATE.get('smea_strength', 0.0), ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
                use_detail_enhancement=False, settings={}, eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Adept Solver":
            return sample_adept_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                order=ADEPT_STATE.get('solver_order', 2), use_corrector=ADEPT_STATE.get('use_corrector', True),
                use_detail_enhancement=False, settings={}
            )
        elif custom_type == "Adept Ancestral Solver":
            return sample_adept_ancestral_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', False), phase_noise=ADEPT_STATE.get('phase_noise', False),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
                use_detail_enhancement=False, settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Mirror Correction Euler":
            return sample_mirror_correction_euler(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
                smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
    

    
    if not ADEPT_STATE.get('enabled', False):
        if 'dpmpp_2m' in ORIGINAL_SAMPLERS:
            return ORIGINAL_SAMPLERS['dpmpp_2m'](model, x, sigmas, extra_args, callback, disable)
        return _basic_dpmpp_2m(model, x, sigmas, extra_args, callback, disable)

    if ADEPT_STATE.get('adaptive_noise_scale', False):
        print("   Adaptive Noise Scale: skipped (deterministic solver: dpmpp_2m, no s_noise/noise_sampler)")

    # Apply custom scheduler deterministically (before the loop, not via p.sampler.model_wrap)
    _sched = ADEPT_STATE.get('scheduler', 'Standard')
    if _sched != 'Standard':
        sigmas = apply_custom_scheduler(sigmas, _sched)

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    # Get settings
    base_scale = ADEPT_STATE.get('scale', 1.0)
    shift = ADEPT_STATE.get('shift', 0.0)
    start_pct = ADEPT_STATE.get('start_pct', 0.0)
    end_pct = ADEPT_STATE.get('end_pct', 1.0)
    
    # Get UNet
    try:
        unet_model = shared.sd_model.model.diffusion_model
    except AttributeError:
        unet_model = None
    
    total_steps = len(sigmas) - 1
    print(f"✅ Adept DPM++ 2M active: scale={base_scale:.2f}")
    
    old_denoised = None
    
    for i in trange(len(sigmas) - 1, disable=disable, desc="Adept DPM++ 2M"):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        # Dynamic scale
        current_scale = compute_dynamic_scale(i, total_steps, base_scale, start_pct, end_pct)
        
        # Evaluate model with weight patching
        if should_patch_weights(unet_model, current_scale, shift):
            with AdeptWeightPatcher(unet_model, current_scale, shift):
                denoised = model(x, sigma * s_in, **extra_args)
        else:
            denoised = model(x, sigma * s_in, **extra_args)
        
        # DPM++ 2M step
        t, t_next = sigma, sigma_next
        h = t_next - t
        
        if old_denoised is None or sigma_next == 0:
            # First step (Euler)
            x = (sigma_next / sigma) * x - (-h).expm1() * denoised
        else:
            # Second order
            h_last = t - sigmas[i - 1]
            r = h_last / h
            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
            x = (sigma_next / sigma) * x - (-h).expm1() * denoised_d
        
        old_denoised = denoised
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})
    
    return x


def _basic_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
    """Fallback basic DPM++ 2M."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    old_denoised = None
    
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        t, t_next = sigmas[i], sigmas[i + 1]
        h = t_next - t
        
        if old_denoised is None or sigmas[i + 1] == 0:
            x = (t_next / t) * x - (-h).expm1() * denoised
        else:
            h_last = t - sigmas[i - 1]
            r = h_last / h
            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
            x = (t_next / t) * x - (-h).expm1() * denoised_d
        
        old_denoised = denoised
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})
    
    return x


@torch.no_grad()

def sample_adept_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0, noise_sampler=None, _skip_ans_check=False):
    """DPM++ 2S Ancestral with Adept weight scaling."""

    # CUSTOM SAMPLER INTEGRATION
    if ADEPT_STATE.get('enabled', False) and ADEPT_STATE.get('use_custom_sampler', False):
        custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
        print(f"🌀 Redirecting to {custom_type}")
        
        # Apply scheduler to sigmas for custom samplers
        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        if scheduler != "Standard":
            sigmas = apply_custom_scheduler(sigmas, scheduler)
            print(f"   📊 Applied {scheduler} scheduler")
        
        if custom_type == "Akashic Solver v2":
            return sample_akashic_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                tau=ADEPT_STATE.get('tau', 0.5), eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0), adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), order=ADEPT_STATE.get('solver_order', 2),
                smea_strength=ADEPT_STATE.get('smea_strength', 0.0), ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
                use_detail_enhancement=False, settings={}, eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Adept Solver":
            return sample_adept_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                order=ADEPT_STATE.get('solver_order', 2), use_corrector=ADEPT_STATE.get('use_corrector', True),
                use_detail_enhancement=False, settings={}
            )
        elif custom_type == "Adept Ancestral Solver":
            return sample_adept_ancestral_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', False), phase_noise=ADEPT_STATE.get('phase_noise', False),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
                use_detail_enhancement=False, settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Mirror Correction Euler":
            return sample_mirror_correction_euler(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
                smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
    

    
    if not ADEPT_STATE.get('enabled', False):
        if 'dpmpp_2s_ancestral' in ORIGINAL_SAMPLERS:
            return ORIGINAL_SAMPLERS['dpmpp_2s_ancestral'](model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
        return _basic_dpmpp_2s_ancestral(model, x, sigmas, extra_args, callback, disable, eta, s_noise)

    if not _skip_ans_check and ADEPT_STATE.get('adaptive_noise_scale', False):
        _self_partial = functools.partial(sample_adept_dpmpp_2s_ancestral, _skip_ans_check=True)
        if 's_noise' in inspect.signature(_self_partial).parameters:
            return run_ans_two_pass(
                _self_partial, 'sample_dpmpp_2s_ancestral', model, x, sigmas, extra_args, callback, disable,
                dict(eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
            )

    # Apply custom scheduler deterministically (before the loop, not via p.sampler.model_wrap)
    _sched = ADEPT_STATE.get('scheduler', 'Standard')
    if _sched != 'Standard':
        sigmas = apply_custom_scheduler(sigmas, _sched)

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    # Get settings
    base_scale = ADEPT_STATE.get('scale', 1.0)
    shift = ADEPT_STATE.get('shift', 0.0)
    start_pct = ADEPT_STATE.get('start_pct', 0.0)
    end_pct = ADEPT_STATE.get('end_pct', 1.0)
    current_eta = ADEPT_STATE.get('eta', eta)
    current_s_noise = ADEPT_STATE.get('s_noise', s_noise)
    
    # Get UNet
    try:
        unet_model = shared.sd_model.model.diffusion_model
    except AttributeError:
        unet_model = None
    
    if noise_sampler is None:
        noise_sampler = default_noise_sampler(x)
    
    total_steps = len(sigmas) - 1
    print(f"✅ Adept DPM++ 2S A active: scale={base_scale:.2f}")
    
    for i in trange(len(sigmas) - 1, disable=disable, desc="Adept DPM++ 2S A"):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        
        # Dynamic scale
        current_scale = compute_dynamic_scale(i, total_steps, base_scale, start_pct, end_pct)
        
        # First evaluation
        if should_patch_weights(unet_model, current_scale, shift):
            with AdeptWeightPatcher(unet_model, current_scale, shift):
                denoised = model(x, sigma * s_in, **extra_args)
        else:
            denoised = model(x, sigma * s_in, **extra_args)
        
        # DPM++ 2S step with ancestral noise
        sigma_down, sigma_up = get_ancestral_step(sigma, sigma_next, current_eta)
        
        if sigma_down == 0:
            d = to_d(x, sigma, denoised)
            x = x + d * (sigma_down - sigma)
        else:
            # Midpoint method
            t, t_next = sigma, sigma_down
            h = t_next - t
            s = t + h * 0.5
            
            # Step to midpoint
            x_mid = (s / t) * x - (-(h * 0.5)).expm1() * denoised
            
            # Evaluate at midpoint
            if should_patch_weights(unet_model, current_scale, shift):
                with AdeptWeightPatcher(unet_model, current_scale, shift):
                    denoised_mid = model(x_mid, s * s_in, **extra_args)
            else:
                denoised_mid = model(x_mid, s * s_in, **extra_args)
            
            # Full step using midpoint
            x = (t_next / t) * x - (-h).expm1() * denoised_mid
        
        # Add ancestral noise
        if sigma_up > 0:
            noise = noise_sampler(sigma, sigma_next) * current_s_noise
            x = x + noise * sigma_up
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})
    
    return x


def _basic_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1.0):
    """Fallback basic DPM++ 2S Ancestral."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    noise_sampler = default_noise_sampler(x)
    
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta)
        
        if sigma_down == 0:
            d = to_d(x, sigmas[i], denoised)
            x = x + d * (sigma_down - sigmas[i])
        else:
            t, t_next = sigmas[i], sigma_down
            h = t_next - t
            s = t + h * 0.5
            x_mid = (s / t) * x - (-(h * 0.5)).expm1() * denoised
            denoised_mid = model(x_mid, s * s_in, **extra_args)
            x = (t_next / t) * x - (-h).expm1() * denoised_mid
        
        if sigma_up > 0:
            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})
    
    return x


@torch.no_grad()

def sample_adept_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
    """LMS sampler with Adept weight scaling."""

    # CUSTOM SAMPLER INTEGRATION
    if ADEPT_STATE.get('enabled', False) and ADEPT_STATE.get('use_custom_sampler', False):
        custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
        print(f"🌀 Redirecting to {custom_type}")
        
        # Apply scheduler to sigmas for custom samplers
        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        if scheduler != "Standard":
            sigmas = apply_custom_scheduler(sigmas, scheduler)
            print(f"   📊 Applied {scheduler} scheduler")
        
        if custom_type == "Akashic Solver v2":
            return sample_akashic_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                tau=ADEPT_STATE.get('tau', 0.5), eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0), adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), order=ADEPT_STATE.get('solver_order', 2),
                smea_strength=ADEPT_STATE.get('smea_strength', 0.0), ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
                use_detail_enhancement=False, settings={}, eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Adept Solver":
            return sample_adept_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                order=ADEPT_STATE.get('solver_order', 2), use_corrector=ADEPT_STATE.get('use_corrector', True),
                use_detail_enhancement=False, settings={}
            )
        elif custom_type == "Adept Ancestral Solver":
            return sample_adept_ancestral_solver(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
                adaptive_eta=ADEPT_STATE.get('adaptive_eta', False), phase_noise=ADEPT_STATE.get('phase_noise', False),
                phase_strength=ADEPT_STATE.get('phase_strength', 0.5), enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
                use_detail_enhancement=False, settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
        elif custom_type == "Mirror Correction Euler":
            return sample_mirror_correction_euler(
                model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
                eta=ADEPT_STATE.get('eta', 1.0),
                s_noise=ADEPT_STATE.get('s_noise', 1.0),
                correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
                smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
            )
    

    
    if not ADEPT_STATE.get('enabled', False):
        if 'lms' in ORIGINAL_SAMPLERS:
            return ORIGINAL_SAMPLERS['lms'](model, x, sigmas, extra_args, callback, disable, order)
        return _basic_lms(model, x, sigmas, extra_args, callback, disable, order)

    if ADEPT_STATE.get('adaptive_noise_scale', False):
        print("   Adaptive Noise Scale: skipped (deterministic solver: lms, no s_noise/noise_sampler)")

    # Apply custom scheduler deterministically (before the loop, not via p.sampler.model_wrap)
    _sched = ADEPT_STATE.get('scheduler', 'Standard')
    if _sched != 'Standard':
        sigmas = apply_custom_scheduler(sigmas, _sched)

    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    
    # Get settings
    base_scale = ADEPT_STATE.get('scale', 1.0)
    shift = ADEPT_STATE.get('shift', 0.0)
    start_pct = ADEPT_STATE.get('start_pct', 0.0)
    end_pct = ADEPT_STATE.get('end_pct', 1.0)
    
    # Get UNet
    try:
        unet_model = shared.sd_model.model.diffusion_model
    except AttributeError:
        unet_model = None
    
    total_steps = len(sigmas) - 1
    print(f"✅ Adept LMS active: scale={base_scale:.2f}, order={order}")
    
    ds = []
    
    for i in trange(len(sigmas) - 1, disable=disable, desc="Adept LMS"):
        sigma = sigmas[i]
        
        # Dynamic scale
        current_scale = compute_dynamic_scale(i, total_steps, base_scale, start_pct, end_pct)
        
        # Evaluate model with weight patching
        if should_patch_weights(unet_model, current_scale, shift):
            with AdeptWeightPatcher(unet_model, current_scale, shift):
                denoised = model(x, sigma * s_in, **extra_args)
        else:
            denoised = model(x, sigma * s_in, **extra_args)
        
        d = to_d(x, sigma, denoised)
        ds.append(d)
        
        if len(ds) > order:
            ds.pop(0)
        
        # Linear multistep coefficients
        cur_order = min(i + 1, order)
        coeffs = [1.0]
        
        for j in range(1, cur_order):
            prod = 1.0
            for k in range(cur_order):
                if k != j:
                    prod *= (sigmas[i] - sigmas[i - k]) / (sigmas[i - j] - sigmas[i - k])
            coeffs.append(prod)
        
        # Apply multistep
        d_multistep = sum(c * d_val for c, d_val in zip(coeffs, reversed(ds[-cur_order:])))
        
        dt = sigmas[i + 1] - sigma
        x = x + d_multistep * dt
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigma, 'denoised': denoised})
    
    return x


def _basic_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
    """Fallback basic LMS."""
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    ds = []
    
    for i in trange(len(sigmas) - 1, disable=disable):
        denoised = model(x, sigmas[i] * s_in, **extra_args)
        d = to_d(x, sigmas[i], denoised)
        ds.append(d)
        
        if len(ds) > order:
            ds.pop(0)
        
        cur_order = min(i + 1, order)
        coeffs = [1.0]
        for j in range(1, cur_order):
            prod = 1.0
            for k in range(cur_order):
                if k != j:
                    prod *= (sigmas[i] - sigmas[i - k]) / (sigmas[i - j] - sigmas[i - k])
            coeffs.append(prod)
        
        d_multistep = sum(c * d_val for c, d_val in zip(coeffs, reversed(ds[-cur_order:])))
        dt = sigmas[i + 1] - sigmas[i]
        x = x + d_multistep * dt
        
        if callback is not None:
            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'denoised': denoised})
    
    return x


# ============================================================================
# MONKEY PATCHING
# ============================================================================

# Samplers that get a hand-written implementation with full Adept weight
# scale/shift patching (compute_dynamic_scale + AdeptWeightPatcher). Every
# other k-diffusion sampler falls back to the generic wrapper below, which
# still gets scheduler override + custom-solver redirect, just not the
# per-algorithm weight patching.
_SPECIALIZED_WRAPPERS = {
    'sample_euler': sample_adept_euler,
    'sample_euler_ancestral': sample_adept_euler_ancestral,
    'sample_heun': sample_adept_heun,
    'sample_dpmpp_2m': sample_adept_dpmpp_2m,
    'sample_dpmpp_2s_ancestral': sample_adept_dpmpp_2s_ancestral,
    'sample_lms': sample_adept_lms,
}


def _adept_redirect_to_custom_solver(name, custom_type, model, x, sigmas, extra_args, callback, disable):
    """Shared custom-solver redirect used by the generic wrapper (mirrors the
    logic duplicated inside each specialized sample_adept_* function)."""
    if custom_type == "Akashic Solver v2":
        return sample_akashic_solver(
            model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
            tau=ADEPT_STATE.get('tau', 0.5), eta=ADEPT_STATE.get('eta', 1.0),
            s_noise=ADEPT_STATE.get('s_noise', 1.0), adaptive_eta=ADEPT_STATE.get('adaptive_eta', True),
            phase_strength=ADEPT_STATE.get('phase_strength', 0.5), order=ADEPT_STATE.get('solver_order', 2),
            smea_strength=ADEPT_STATE.get('smea_strength', 0.0), ndb_strength=ADEPT_STATE.get('ndb_strength', 0.0),
            use_detail_enhancement=False, settings={}, eqvae_mode=ADEPT_STATE.get('eqvae_mode', 'Off'),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
        )
    elif custom_type == "Adept Solver":
        return sample_adept_solver(
            model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
            order=ADEPT_STATE.get('solver_order', 2), use_corrector=ADEPT_STATE.get('use_corrector', True),
            use_detail_enhancement=False, settings={}
        )
    elif custom_type == "Adept Ancestral Solver":
        return sample_adept_ancestral_solver(
            model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
            eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
            adaptive_eta=ADEPT_STATE.get('adaptive_eta', False), phase_noise=ADEPT_STATE.get('phase_noise', False),
            phase_strength=ADEPT_STATE.get('phase_strength', 0.5), enhanced_derivative=ADEPT_STATE.get('enhanced_derivative', False),
            use_detail_enhancement=False, settings={},
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
        )
    elif custom_type == "Mirror Correction Euler":
        return sample_mirror_correction_euler(
            model=model, x=x, sigmas=sigmas, extra_args=extra_args, callback=callback, disable=disable,
            eta=ADEPT_STATE.get('eta', 1.0), s_noise=ADEPT_STATE.get('s_noise', 1.0),
            correction_phase=ADEPT_STATE.get('mirror_correction_phase', 0.5),
            smooth_phase=ADEPT_STATE.get('mirror_smooth_phase', False),
            adaptive_noise_scale=ADEPT_STATE.get('adaptive_noise_scale', False)
        )
    return None  # unknown custom_type -> caller falls through to native sampler


def _make_generic_adept_wrapper(attr_name, original_fn):
    """
    Build a wrapper for a k-diffusion sampler that has no hand-written Adept
    variant. Unlike the specialized wrappers, this does NOT reimplement the
    sampling loop — it only does what's algorithm-agnostic:
      1. no-op passthrough when Adept is disabled
      2. optional redirect to one of the 4 custom solvers (Akashic/Adept/
         Adept Ancestral/Mirror), same as the specialized wrappers do
      3. otherwise, sigma-schedule override while preserving the sampler's
         own native step logic (this is what "Adept preserves the native
         solver and only overrides sigma schedule" in the README refers to)

    functools.wraps is required, not cosmetic: A1111 inspects
    inspect.signature(self.func).parameters to decide whether to pass
    sampler-specific extra kwargs (s_churn, s_tmin, s_tmax, r, ...). Without
    wraps(), inspect.signature() would report our wrapper's own **kwargs
    signature instead of the real one, and those extra params would silently
    stop being passed to samplers that use them (e.g. churn-based samplers).
    """
    @functools.wraps(original_fn)
    def generic_wrapper(model, x, sigmas, extra_args=None, callback=None, disable=None, **kwargs):
        if not ADEPT_STATE.get('enabled', False):
            return original_fn(model, x, sigmas, extra_args, callback, disable, **kwargs)

        scheduler = ADEPT_STATE.get('scheduler', 'Standard')
        local_sigmas = sigmas
        if scheduler != 'Standard':
            try:
                local_sigmas = apply_custom_scheduler(local_sigmas, scheduler)
            except Exception as e:
                print(f"⚠️ [{attr_name}] Scheduler override failed: {e}. Using original schedule.")
                local_sigmas = sigmas

        if ADEPT_STATE.get('use_custom_sampler', False):
            custom_type = ADEPT_STATE.get('custom_sampler', 'Akashic Solver v2')
            print(f"🌀 [{attr_name}] Redirecting to {custom_type}")
            result = _adept_redirect_to_custom_solver(
                attr_name, custom_type, model, x, local_sigmas, extra_args, callback, disable
            )
            if result is not None:
                return result
            print(f"⚠️ [{attr_name}] Unknown custom sampler type '{custom_type}', "
                  f"falling back to native sampler with schedule override.")

        if ADEPT_STATE.get('adaptive_noise_scale', False):
            sig_params = inspect.signature(original_fn).parameters
            if 's_noise' in sig_params or 'noise_sampler' in sig_params:
                return run_ans_two_pass(
                    original_fn, attr_name, model, x, local_sigmas, extra_args, callback, disable, kwargs
                )
            print(f"   Adaptive Noise Scale: skipped (deterministic solver: "
                  f"{attr_name.replace('sample_', '')})")

        return original_fn(model, x, local_sigmas, extra_args, callback, disable, **kwargs)

    return generic_wrapper


def patch_k_diffusion():
    """
    Apply monkey patches to ALL k-diffusion samplers (dynamic discovery via
    dir(), same approach as upstream's patch_samplers_globally) instead of a
    fixed list. Fixes: Adept previously only activated for 6 hardcoded base
    samplers; any other sampler (DPM++ SDE, DPM++ 2M SDE, DPM++ 3M SDE, DPM2,
    DPM2 a, Restart, etc.) silently did nothing even with Adept enabled.
    """

    specialized_count = 0
    generic_count = 0
    skipped_names = []

    for attr_name in dir(k_diffusion.sampling):
        if not attr_name.startswith('sample_'):
            continue

        func = getattr(k_diffusion.sampling, attr_name, None)
        if not callable(func):
            continue

        # Skip samplers with an incompatible signature — e.g. dpm_fast /
        # dpm_adaptive take (model, x, sigma_min, sigma_max, n, ...) instead
        # of (model, x, sigmas, ...) and can't be wrapped the same way.
        try:
            if 'sigmas' not in inspect.signature(func).parameters:
                skipped_names.append(attr_name)
                continue
        except (ValueError, TypeError):
            continue

        key = attr_name.replace('sample_', '')

        # Already patched by us in a previous call (e.g. UI reload) — skip.
        if key in _ADEPT_ACTIVE_FUNCS and func is _ADEPT_ACTIVE_FUNCS[key]:
            continue

        if key not in ORIGINAL_SAMPLERS:
            ORIGINAL_SAMPLERS[key] = func

        if attr_name in _SPECIALIZED_WRAPPERS:
            target_func = _SPECIALIZED_WRAPPERS[attr_name]
            specialized_count += 1
        else:
            target_func = _make_generic_adept_wrapper(attr_name, ORIGINAL_SAMPLERS[key])
            generic_count += 1

        setattr(k_diffusion.sampling, attr_name, target_func)
        _ADEPT_ACTIVE_FUNCS[key] = target_func

    total = specialized_count + generic_count
    print(f"✅ Adept Sampler v5: Patched {total} samplers "
          f"({specialized_count} specialized + {generic_count} generic; "
          f"skipped {len(skipped_names)} incompatible)")
    print("   Schedulers: 18 types available")


def unpatch_k_diffusion():
    """
    Restore original k-diffusion samplers.

    Safe-unpatch strategy: before restoring we check whether the live slot
    still holds *our* wrapper. If another extension has wrapped us on top
    (i.e. live_func is neither our wrapper nor the original we saved),
    blindly restoring would silently remove *their* wrapper too. In that
    case we skip the restore for that slot and log a warning so the
    operator knows the coexistence situation.
    """

    restored_count = 0
    skipped_count = 0

    for key, our_func in list(_ADEPT_ACTIVE_FUNCS.items()):
        attr_name = f"sample_{key}"
        if not hasattr(k_diffusion.sampling, attr_name):
            continue

        live_func = getattr(k_diffusion.sampling, attr_name)
        saved_original = ORIGINAL_SAMPLERS.get(key)

        if live_func is our_func:
            # Normal case: we still own the slot — safe to restore.
            setattr(k_diffusion.sampling, attr_name, saved_original)
            restored_count += 1
        elif live_func is saved_original:
            # Already restored somehow — nothing to do.
            restored_count += 1
        else:
            # Another extension wrapped us. Restoring would silently
            # remove their wrapper; skip and warn instead.
            print(f"⚠️ Adept unpatch: {attr_name} is currently owned by another "
                  f"extension ({live_func!r}). Skipping restore to avoid breaking "
                  f"their wrapper — you may need to reload the UI to fully unload.")
            skipped_count += 1

    ORIGINAL_SAMPLERS.clear()
    _ADEPT_ACTIVE_FUNCS.clear()
    print(f"🔄 Adept Sampler: Restored {restored_count} samplers"
          + (f", skipped {skipped_count} (foreign wrappers)" if skipped_count else ""))


# ============================================================================
# A1111 EXTENSION SCRIPT
# ============================================================================

class AdeptSamplerScript(scripts.Script):
    def title(self):
        return "Adept Sampler v5"
    
    def show(self, is_img2img):
        return scripts.AlwaysVisible
    
    def ui(self, is_img2img):
        with gr.Accordion("Adept Sampler v5", open=False):
            enabled = gr.Checkbox(label="Enable Adept Sampler", value=False, elem_id="adept_enabled")
            
            with gr.Row():
                scale = gr.Slider(minimum=0.5, maximum=2.0, step=0.05, value=1.0, label="Weight Scale")
                shift = gr.Slider(minimum=-0.5, maximum=0.5, step=0.01, value=0.0, label="Weight Shift")

            with gr.Row():
                start_pct = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label="Start Percent")
                end_pct = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=1.0, label="End Percent")

            gr.HTML("<p style='color: #888; font-size: 0.85em; margin: 2px 0 10px;'>"
                    "⚠️ Weight Scale / Shift / Start–End apply to the 6 patched k-diffusion samplers only. "
                    "Custom samplers (Akashic, Adept, Mirror) use their own internal parameters.</p>")
            
            with gr.Row():
                eta = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=1.0, label="Eta (Ancestral samplers)")
                s_noise = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=1.0, label="S-Noise")
            
            adaptive_eta = gr.Checkbox(label="Adaptive Eta (dynamic eta during sampling)", value=False)
            
            scheduler = gr.Dropdown(
                choices=["Standard", "AOS-V", "AOS-Epsilon", "AkashicAOS", "Entropic", "SNR-Optimized", 
                         "Constant-Rate", "Adaptive-Optimized", "Cosine-Annealed", "LogSNR-Uniform", 
                         "Tanh Mid-Boost", "Exponential Tail", "Jittered-Karras", "Stochastic", 
                         "JYS (Dynamic)", "Hybrid JYS-Karras", "AYS-SDXL",
                         "AkashicAOS Alt", "AkashicEQFlow"],
                value="Standard", label="Scheduler Type"
            )
            
            vae_reflection = gr.Checkbox(label="Enable VAE Reflection (fixes edge artifacts for EQ-VAE)", value=False)
            
            gr.HTML("<hr style='margin: 15px 0;'>")
            gr.HTML("<h3 style='margin: 10px 0;'>🧠 Content-Aware Pacing (Euler Ancestral only)</h3>")
            gr.HTML("<p style='color: #888; font-size: 0.9em;'>Auto-switches from a composition phase to a detail "
                    "phase once the model's output stabilizes, using a freshly-generated schedule for the "
                    "remaining steps. No effect on custom solvers or other base samplers. Auto-disables below "
                    "26 steps.</p>")
            use_content_aware_pacing = gr.Checkbox(label="Enable Content-Aware Pacing", value=False)
            with gr.Group(visible=False) as pacing_options_group:
                pacing_coherence_sensitivity = gr.Slider(
                    minimum=0.0, maximum=1.0, step=0.05, value=0.75,
                    label="Coherence Sensitivity (higher = switches to detail phase sooner)"
                )
                use_detail_enhancement_phase = gr.Checkbox(label="High-Frequency Detail Enhancement", value=False)
                with gr.Group(visible=False) as detail_fx_options_group:
                    with gr.Row():
                        detail_enhancement_strength = gr.Slider(minimum=0.0, maximum=0.5, step=0.01, value=0.05, label="Detail Strength")
                        detail_separation_radius = gr.Slider(minimum=0.1, maximum=3.0, step=0.1, value=0.5, label="Detail Separation Radius")

            use_content_aware_pacing.change(
                fn=lambda v: gr.update(visible=v),
                inputs=[use_content_aware_pacing], outputs=[pacing_options_group]
            )
            use_detail_enhancement_phase.change(
                fn=lambda v: gr.update(visible=v),
                inputs=[use_detail_enhancement_phase], outputs=[detail_fx_options_group]
            )
            
            gr.HTML("<hr style='margin: 15px 0;'>")
            gr.HTML("<h3 style='margin: 10px 0;'>🌀 Custom Advanced Samplers</h3>")
            gr.HTML("<p style='color: #888; font-size: 0.9em;'>Enable to use Akashic/Adept/Ancestral samplers instead of k-diffusion</p>")
            
            use_custom = gr.Checkbox(label="Use Custom Sampler (overrides k-diffusion)", value=False)
            custom_type = gr.Dropdown(
                choices=["Akashic Solver v2", "Adept Solver", "Adept Ancestral Solver", "Mirror Correction Euler"],
                value="Akashic Solver v2", label="Custom Sampler Type"
            )
            adaptive_noise_scale = gr.Checkbox(
                label="Adaptive Noise Scale (auto-calibrate s_noise, restarts once mid-run)",
                value=False
            )
            gr.HTML("<p style='color: #888; font-size: 0.85em; margin-top: -8px;'>"
                    "Available on Akashic Solver v2, Adept Ancestral Solver, and Mirror Correction Euler. "
                    "No effect on Adept Solver. Runs a short calibration window (5 steps in the "
                    "0.5&lt;&sigma;&lt;5.0 range), then restarts the generation once with the "
                    "auto-computed correction applied from step 0 &mdash; total model calls for that "
                    "run will be higher than the step count.</p>")
            
            with gr.Accordion("⚙️ Akashic Solver Settings", open=False, visible=False) as akashic_settings_group:
                tau = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Tau (0=ODE, 1=SDE)")
                phase_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.5, label="Phase Strength")
                smea = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.0, label="SMEA (high-res coherency)")
                ndb = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label="NDB (detail boost)")
                eqvae = gr.Dropdown(choices=["Off", "Balanced"], value="Off", label="EQ-VAE Mode")
            
            with gr.Accordion("⚙️ Adept Solver Settings", open=False, visible=False) as adept_settings_group:
                solver_order = gr.Slider(minimum=1, maximum=3, step=1, value=2, label="Order (1-3)")
                use_corrector = gr.Checkbox(value=True, label="Use Corrector")
            
            with gr.Accordion("⚙️ Ancestral Solver Settings", open=False, visible=False) as ancestral_settings_group:
                phase_noise = gr.Checkbox(value=False, label="Phase-Aware Noise")
                enhanced_deriv = gr.Checkbox(value=False, label="Enhanced Derivative")

            with gr.Accordion("⚙️ Mirror Correction Euler Settings", open=False, visible=False) as mirror_settings_group:
                gr.HTML("<p style='color: #888; font-size: 0.9em;'>Active only when Custom Sampler = Mirror Correction Euler</p>")
                mirror_correction_phase = gr.Slider(
                    minimum=0.0, maximum=1.0, step=0.05, value=0.5,
                    label="Correction Phase (fraction of steps with 3-call Heun correction)"
                )
                mirror_smooth_phase = gr.Checkbox(
                    value=False,
                    label="Smooth Phase (log-sigma blend instead of binary cutoff)"
                )

            def _adept_on_solver_type_change(use_custom_val, custom_type_val):
                """Only the settings block matching the active custom solver
                (and only when a custom solver is enabled at all) is shown --
                the other three, which have no effect in that configuration,
                stay collapsed out of view instead of just closed."""
                return {
                    akashic_settings_group: gr.update(visible=use_custom_val and custom_type_val == "Akashic Solver v2"),
                    adept_settings_group: gr.update(visible=use_custom_val and custom_type_val == "Adept Solver"),
                    ancestral_settings_group: gr.update(visible=use_custom_val and custom_type_val == "Adept Ancestral Solver"),
                    mirror_settings_group: gr.update(visible=use_custom_val and custom_type_val == "Mirror Correction Euler"),
                }

            _solver_visibility_outputs = [akashic_settings_group, adept_settings_group,
                                           ancestral_settings_group, mirror_settings_group]
            use_custom.change(fn=_adept_on_solver_type_change, inputs=[use_custom, custom_type],
                               outputs=_solver_visibility_outputs)
            custom_type.change(fn=_adept_on_solver_type_change, inputs=[use_custom, custom_type],
                                outputs=_solver_visibility_outputs)

            gr.HTML("<hr style='margin: 15px 0;'>")
            gr.HTML("<h3 style='margin: 10px 0;'>🎛️ CFG Enhancements</h3>")
            gr.HTML("<p style='color: #888; font-size: 0.9em;'>"
                    "Combat CFG Drift works in stock A1111 via official callback. "
                    "Spectral Modulation &amp; Phase-Aware CFG use a native sampler hook on "
                    "Forge/reForge-like backends, or a CFGDenoiser monkey-patch on stock A1111 "
                    "(near-parity; active mode logged to console).</p>")


            with gr.Accordion("⚙️ CFG Enhancement Settings", open=False):
                cfg_drift_enabled = gr.Checkbox(value=False, label="Enable Combat CFG Drift")
                with gr.Row():
                    cfg_drift_method = gr.Dropdown(
                        choices=["mean", "median"], value="mean",
                        label="Drift Method"
                    )
                    cfg_drift_intensity = gr.Slider(
                        minimum=0.0, maximum=1.0, step=0.05, value=0.5,
                        label="Drift Intensity"
                    )
                spectral_cfg_enabled = gr.Checkbox(
                    value=False, label="Enable Spectral Modulation (native hook or A1111 monkey-patch)"
                )
                with gr.Row():
                    spectral_multiplier = gr.Slider(
                        minimum=0.0, maximum=2.0, step=0.05, value=1.0,
                        label="Spectral Multiplier"
                    )
                    spectral_percentile = gr.Slider(
                        minimum=1.0, maximum=25.0, step=0.5, value=5.0,
                        label="Spectral Percentile"
                    )
                phase_cfg_enabled = gr.Checkbox(
                    value=False, label="Enable Phase-Aware CFG (native hook or A1111 monkey-patch)"
                )
                with gr.Row():
                    phase_cfg_alpha = gr.Slider(
                        minimum=1.1, maximum=4.0, step=0.1, value=2.0,
                        label="Phase CFG Alpha"
                    )
                    phase_cfg_beta = gr.Slider(
                        minimum=1.1, maximum=4.0, step=0.1, value=2.0,
                        label="Phase CFG Beta"
                    )

        return [enabled, scale, shift, start_pct, end_pct, eta, s_noise, adaptive_eta, scheduler, vae_reflection,
                use_content_aware_pacing, pacing_coherence_sensitivity, use_detail_enhancement_phase,
                detail_enhancement_strength, detail_separation_radius,
                use_custom, custom_type, adaptive_noise_scale, tau, phase_strength, smea, ndb, eqvae, solver_order, use_corrector,
                phase_noise, enhanced_deriv,
                mirror_correction_phase, mirror_smooth_phase,
                cfg_drift_enabled, cfg_drift_method, cfg_drift_intensity,
                spectral_cfg_enabled, spectral_multiplier, spectral_percentile,
                phase_cfg_enabled, phase_cfg_alpha, phase_cfg_beta]
    
    def process(self, p, enabled, scale, shift, start_pct, end_pct, eta, s_noise, adaptive_eta, scheduler, vae_reflection,
                use_content_aware_pacing, pacing_coherence_sensitivity, use_detail_enhancement_phase,
                detail_enhancement_strength, detail_separation_radius,
                use_custom, custom_type, adaptive_noise_scale, tau, phase_strength, smea, ndb, eqvae, solver_order, use_corrector,
                phase_noise, enhanced_deriv,
                mirror_correction_phase, mirror_smooth_phase,
                cfg_drift_enabled, cfg_drift_method, cfg_drift_intensity,
                spectral_cfg_enabled, spectral_multiplier, spectral_percentile,
                phase_cfg_enabled, phase_cfg_alpha, phase_cfg_beta):

        # --- XYZ Grid overrides for this cell only (see set_value / make_axis_on_xyz_grid) ---
        # locals() reassignment isn't reliable in CPython function scopes, so
        # each field is applied explicitly rather than looped generically.
        xyz = getattr(p, "_adept_xyz", {})
        if xyz:
            if "enabled" in xyz: enabled = xyz["enabled"]
            if "scale" in xyz: scale = xyz["scale"]
            if "shift" in xyz: shift = xyz["shift"]
            if "start_pct" in xyz: start_pct = xyz["start_pct"]
            if "end_pct" in xyz: end_pct = xyz["end_pct"]
            if "eta" in xyz: eta = xyz["eta"]
            if "s_noise" in xyz: s_noise = xyz["s_noise"]
            if "adaptive_eta" in xyz: adaptive_eta = xyz["adaptive_eta"]
            if "scheduler" in xyz: scheduler = xyz["scheduler"]
            if "vae_reflection" in xyz: vae_reflection = xyz["vae_reflection"]
            if "use_content_aware_pacing" in xyz: use_content_aware_pacing = xyz["use_content_aware_pacing"]
            if "pacing_coherence_sensitivity" in xyz: pacing_coherence_sensitivity = xyz["pacing_coherence_sensitivity"]
            if "use_detail_enhancement_phase" in xyz: use_detail_enhancement_phase = xyz["use_detail_enhancement_phase"]
            if "detail_enhancement_strength" in xyz: detail_enhancement_strength = xyz["detail_enhancement_strength"]
            if "detail_separation_radius" in xyz: detail_separation_radius = xyz["detail_separation_radius"]
            if "use_custom" in xyz: use_custom = xyz["use_custom"]
            if "custom_type" in xyz: custom_type = xyz["custom_type"]
            if "adaptive_noise_scale" in xyz: adaptive_noise_scale = xyz["adaptive_noise_scale"]
            if "tau" in xyz: tau = xyz["tau"]
            if "phase_strength" in xyz: phase_strength = xyz["phase_strength"]
            if "smea" in xyz: smea = xyz["smea"]
            if "ndb" in xyz: ndb = xyz["ndb"]
            if "eqvae" in xyz: eqvae = xyz["eqvae"]
            if "solver_order" in xyz: solver_order = xyz["solver_order"]
            if "use_corrector" in xyz: use_corrector = xyz["use_corrector"]
            if "phase_noise" in xyz: phase_noise = xyz["phase_noise"]
            if "enhanced_deriv" in xyz: enhanced_deriv = xyz["enhanced_deriv"]
            if "mirror_correction_phase" in xyz: mirror_correction_phase = xyz["mirror_correction_phase"]
            if "mirror_smooth_phase" in xyz: mirror_smooth_phase = xyz["mirror_smooth_phase"]
            if "cfg_drift_enabled" in xyz: cfg_drift_enabled = xyz["cfg_drift_enabled"]
            if "cfg_drift_method" in xyz: cfg_drift_method = xyz["cfg_drift_method"]
            if "cfg_drift_intensity" in xyz: cfg_drift_intensity = xyz["cfg_drift_intensity"]
            if "spectral_cfg_enabled" in xyz: spectral_cfg_enabled = xyz["spectral_cfg_enabled"]
            if "spectral_multiplier" in xyz: spectral_multiplier = xyz["spectral_multiplier"]
            if "spectral_percentile" in xyz: spectral_percentile = xyz["spectral_percentile"]
            if "phase_cfg_enabled" in xyz: phase_cfg_enabled = xyz["phase_cfg_enabled"]
            if "phase_cfg_alpha" in xyz: phase_cfg_alpha = xyz["phase_cfg_alpha"]
            if "phase_cfg_beta" in xyz: phase_cfg_beta = xyz["phase_cfg_beta"]

        # Gate all sub-features through the master enabled switch.
        # This prevents CFG hooks, native patches, and VAE Reflection from
        # activating when the extension is globally toggled off.
        ADEPT_STATE.update({
            "enabled": enabled,
            "scale": scale,
            "shift": shift,
            "start_pct": start_pct,
            "end_pct": end_pct,
            "eta": eta,
            "s_noise": s_noise,
            "adaptive_eta": adaptive_eta,
            "scheduler": scheduler,
            "vae_reflection": enabled and vae_reflection,   # gated
            "use_content_aware_pacing": use_content_aware_pacing,
            "pacing_coherence_sensitivity": pacing_coherence_sensitivity,
            "use_detail_enhancement_phase": use_detail_enhancement_phase,
            "detail_enhancement_strength": detail_enhancement_strength,
            "detail_separation_radius": detail_separation_radius,
            "use_custom_sampler": use_custom,
            "custom_sampler": custom_type,
            "adaptive_noise_scale": adaptive_noise_scale,
            "tau": tau,
            "phase_strength": phase_strength,
            "smea_strength": smea,
            "ndb_strength": ndb,
            "eqvae_mode": eqvae,
            "solver_order": int(solver_order),
            "use_corrector": use_corrector,
            "phase_noise": phase_noise,
            "enhanced_derivative": enhanced_deriv,
            # Mirror Correction Euler
            "mirror_correction_phase": mirror_correction_phase,
            "mirror_smooth_phase": mirror_smooth_phase,
            # CFG enhancements — all gated through enabled
            "cfg_drift_enabled":    enabled and cfg_drift_enabled,
            "cfg_drift_method":     cfg_drift_method,
            "cfg_drift_intensity":  cfg_drift_intensity,
            "spectral_cfg_enabled": enabled and spectral_cfg_enabled,
            "spectral_multiplier":  spectral_multiplier,
            "spectral_percentile":  spectral_percentile,
            "phase_cfg_enabled":    enabled and phase_cfg_enabled,
            "phase_cfg_alpha":      phase_cfg_alpha,
            "phase_cfg_beta":       phase_cfg_beta,
        })

        # Scheduler is now applied inside each patched sampler function,
        # so p.sampler.model_wrap patching is no longer needed here.

        # Always reconfigure CFG runtime — even when disabled — so any previously
        # installed native hook or A1111 callbacks get cleanly removed.
        runtime_mode = configure_cfg_runtime()

        if enabled:
            info = {
                "Adept Sampler": "v5",
                "Adept Scheduler": scheduler,
                "CFG Runtime": runtime_mode,
            }
            if use_custom:
                info["Adept Custom"] = custom_type
                if custom_type == "Akashic Solver v2":
                    info["Adept Tau"] = tau
                    info["Adept EQ-VAE"] = eqvae
            p.extra_generation_params.update(info)

    def process_batch(self, p, *args, **kwargs):
        """Apply VAE Reflection before batch processing."""
        if ADEPT_STATE.get("enabled", False) and ADEPT_STATE.get("vae_reflection", False):
            try:
                vae_model = shared.sd_model.first_stage_model
                patcher = VAEReflectionPatcher(vae_model)
                patcher.__enter__()
                p.adept_vae_patcher = patcher
            except Exception as e:
                print(f"⚠️ VAE Reflection error: {e}")
    
    def postprocess_batch(self, p, *args, **kwargs):
        """Restore VAE padding modes after batch processing."""
        if hasattr(p, 'adept_vae_patcher'):
            try:
                p.adept_vae_patcher.__exit__(None, None, None)
                delattr(p, 'adept_vae_patcher')
            except Exception as e:
                print(f"⚠️ VAE Reflection restore error: {e}")
                # Safety net: force-restore even if the patcher context failed
                force_restore_vae_reflection()


# ============================================================================
# XYZ GRID INTEGRATION
# ============================================================================
#
# XYZ Grid's per-cell pipeline is: reset p -> call each active axis's apply()
# -> process_images(p) -> our Script.process() runs with the UI's *default*
# values (script_args aren't touched). So set_value() stashes the override on
# p itself, and process() (below, in the ADEPT_STATE.update block) checks for
# it after computing its normal locals but before writing state -- same
# pattern the reForge version of this project uses via p._adept_xyz.

_ADEPT_XYZ_BOOL_FIELDS = {
    "enabled", "adaptive_eta", "vae_reflection", "use_custom", "adaptive_noise_scale",
    "use_corrector", "phase_noise", "enhanced_deriv", "mirror_smooth_phase",
    "cfg_drift_enabled", "spectral_cfg_enabled", "phase_cfg_enabled",
    "use_content_aware_pacing", "use_detail_enhancement_phase",
}
_ADEPT_XYZ_FLOAT_FIELDS = {
    "scale", "shift", "start_pct", "end_pct", "eta", "s_noise", "tau", "phase_strength",
    "smea", "ndb", "mirror_correction_phase", "cfg_drift_intensity", "spectral_multiplier",
    "spectral_percentile", "phase_cfg_alpha", "phase_cfg_beta",
    "pacing_coherence_sensitivity", "detail_enhancement_strength", "detail_separation_radius",
}
_ADEPT_XYZ_INT_FIELDS = {"solver_order"}
_ADEPT_XYZ_STR_FIELDS = {"scheduler", "custom_type", "eqvae", "cfg_drift_method"}


def set_value(p, x, xs, *, field):
    """Generic XYZ Grid apply() target: validate/convert x for `field`, then
    stash it on p so our own process() can pick it up for this cell only."""
    if not hasattr(p, "_adept_xyz"):
        p._adept_xyz = {}
    try:
        if field in _ADEPT_XYZ_BOOL_FIELDS:
            x = str(x).strip().lower() == "true"
        elif field in _ADEPT_XYZ_FLOAT_FIELDS:
            x = float(x)
        elif field in _ADEPT_XYZ_INT_FIELDS:
            x = int(x)
            if x not in (1, 2, 3):
                raise ValueError(f"solver_order must be 1, 2, or 3, got {x}")
        elif field in _ADEPT_XYZ_STR_FIELDS:
            x = str(x)
        p._adept_xyz[field] = x
    except (ValueError, TypeError) as e:
        print(f"⚠️ Adept XYZ Grid: invalid value '{x}' for field '{field}': {e}. Skipping override.")


def make_axis_on_xyz_grid():
    """Register one AxisOption per Adept setting with the real xyz_grid.py
    script, if it's loaded. Safe to call more than once (dedup-guarded)."""
    xyz_grid = None
    for data in scripts.scripts_data:
        if os.path.basename(data.path) == "xyz_grid.py":
            xyz_grid = data.module
            break

    if xyz_grid is None:
        return  # xyz_grid.py disabled or not present -- nothing to hook

    bool_choices = lambda: ["True", "False"]

    axis = [
        xyz_grid.AxisOption("(Adept) Enabled", str, functools.partial(set_value, field="enabled"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Weight Scale", float, functools.partial(set_value, field="scale")),
        xyz_grid.AxisOption("(Adept) Weight Shift", float, functools.partial(set_value, field="shift")),
        xyz_grid.AxisOption("(Adept) Start %", float, functools.partial(set_value, field="start_pct")),
        xyz_grid.AxisOption("(Adept) End %", float, functools.partial(set_value, field="end_pct")),
        xyz_grid.AxisOption("(Adept) Eta", float, functools.partial(set_value, field="eta")),
        xyz_grid.AxisOption("(Adept) S Noise", float, functools.partial(set_value, field="s_noise")),
        xyz_grid.AxisOption("(Adept) Adaptive Eta", str, functools.partial(set_value, field="adaptive_eta"), choices=bool_choices),
        xyz_grid.AxisOption(
            "(Adept) Scheduler", str, functools.partial(set_value, field="scheduler"),
            choices=lambda: ["Standard", "AOS-V", "AOS-Epsilon", "AkashicAOS", "Entropic", "SNR-Optimized",
                              "Constant-Rate", "Adaptive-Optimized", "Cosine-Annealed", "LogSNR-Uniform",
                              "Tanh Mid-Boost", "Exponential Tail", "Jittered-Karras", "Stochastic",
                              "JYS (Dynamic)", "Hybrid JYS-Karras", "AYS-SDXL",
                              "AkashicAOS Alt", "AkashicEQFlow"]
        ),
        xyz_grid.AxisOption("(Adept) VAE Reflection", str, functools.partial(set_value, field="vae_reflection"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Content-Aware Pacing", str, functools.partial(set_value, field="use_content_aware_pacing"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Pacing Coherence Sensitivity", float, functools.partial(set_value, field="pacing_coherence_sensitivity")),
        xyz_grid.AxisOption("(Adept) Pacing Detail Enhancement", str, functools.partial(set_value, field="use_detail_enhancement_phase"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Pacing Detail Strength", float, functools.partial(set_value, field="detail_enhancement_strength")),
        xyz_grid.AxisOption("(Adept) Pacing Detail Radius", float, functools.partial(set_value, field="detail_separation_radius")),
        xyz_grid.AxisOption("(Adept) Use Custom Sampler", str, functools.partial(set_value, field="use_custom"), choices=bool_choices),
        xyz_grid.AxisOption(
            "(Adept) Custom Sampler Type", str, functools.partial(set_value, field="custom_type"),
            choices=lambda: ["Akashic Solver v2", "Adept Solver", "Adept Ancestral Solver", "Mirror Correction Euler"]
        ),
        xyz_grid.AxisOption("(Adept) Adaptive Noise Scale", str, functools.partial(set_value, field="adaptive_noise_scale"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Tau", float, functools.partial(set_value, field="tau")),
        xyz_grid.AxisOption("(Adept) Phase Strength", float, functools.partial(set_value, field="phase_strength")),
        xyz_grid.AxisOption("(Adept) SMEA Strength", float, functools.partial(set_value, field="smea")),
        xyz_grid.AxisOption("(Adept) NDB Strength", float, functools.partial(set_value, field="ndb")),
        xyz_grid.AxisOption("(Adept) EQ-VAE Mode", str, functools.partial(set_value, field="eqvae"), choices=lambda: ["Off", "Balanced"]),
        xyz_grid.AxisOption("(Adept) Solver Order", int, functools.partial(set_value, field="solver_order"), choices=lambda: ["1", "2", "3"]),
        xyz_grid.AxisOption("(Adept) Use Corrector", str, functools.partial(set_value, field="use_corrector"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Ancestral Phase Noise", str, functools.partial(set_value, field="phase_noise"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Ancestral Enhanced Derivative", str, functools.partial(set_value, field="enhanced_deriv"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Mirror Correction Phase", float, functools.partial(set_value, field="mirror_correction_phase")),
        xyz_grid.AxisOption("(Adept) Mirror Smooth Phase", str, functools.partial(set_value, field="mirror_smooth_phase"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Combat CFG Drift Enabled", str, functools.partial(set_value, field="cfg_drift_enabled"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) CFG Drift Method", str, functools.partial(set_value, field="cfg_drift_method"), choices=lambda: ["mean", "median"]),
        xyz_grid.AxisOption("(Adept) CFG Drift Intensity", float, functools.partial(set_value, field="cfg_drift_intensity")),
        xyz_grid.AxisOption("(Adept) Spectral Modulation Enabled", str, functools.partial(set_value, field="spectral_cfg_enabled"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Spectral Multiplier", float, functools.partial(set_value, field="spectral_multiplier")),
        xyz_grid.AxisOption("(Adept) Spectral Percentile", float, functools.partial(set_value, field="spectral_percentile")),
        xyz_grid.AxisOption("(Adept) Phase-Aware CFG Enabled", str, functools.partial(set_value, field="phase_cfg_enabled"), choices=bool_choices),
        xyz_grid.AxisOption("(Adept) Phase CFG Alpha", float, functools.partial(set_value, field="phase_cfg_alpha")),
        xyz_grid.AxisOption("(Adept) Phase CFG Beta", float, functools.partial(set_value, field="phase_cfg_beta")),
    ]

    if not any(getattr(x, "label", "").startswith("(Adept)") for x in xyz_grid.axis_options):
        xyz_grid.axis_options.extend(axis)


# ============================================================================
# INITIALIZATION
# ============================================================================
#
# k-diffusion wrappers are installed via on_before_ui (fires after all
# extensions are imported) rather than at bare module import time.
# This reduces the risk of interacting badly with other extensions that
# also wrap k_diffusion.sampling functions, because our wrappers are put
# on last and therefore sit outermost in the call chain.
# Uninstall happens in on_script_unloaded() as before.

def _adept_deferred_init():
    patch_k_diffusion()
    try:
        make_axis_on_xyz_grid()
    except Exception:
        print("⚠️ Adept Sampler: XYZ Grid integration failed (xyz_grid.py may be disabled)")

try:
    script_callbacks.on_before_ui(_adept_deferred_init)
except Exception:
    # Fallback: if on_before_ui isn't available (older A1111), patch immediately.
    patch_k_diffusion()

def on_script_unloaded():
    try:
        force_restore_vae_reflection()
    except Exception:
        pass
    try:
        uninstall_a1111_cfg_callbacks()
    except Exception:
        pass
    try:
        uninstall_native_cfg_hook()
    except Exception:
        pass
    try:
        unpatch_cfg_denoiser()
    except Exception:
        pass
    try:
        unpatch_k_diffusion()
    except Exception:
        pass

try:
    script_callbacks.on_script_unloaded(on_script_unloaded)
except AttributeError:
    print("⚠️ Script unload callback not available")

print("🚀 Adept Sampler v5 loaded!")
print("   ✨ 4 Custom Samplers: Akashic v2, Adept Solver, Adept Ancestral, Mirror Correction Euler")
print("   ⚡ 6 k-diffusion Samplers with weight scaling")
print("   📅 18 Schedulers (including AkashicAOS Alt, AkashicEQFlow)")
print("   🎨 VAE Reflection")
print("   ✅ A1111 port of ComfyUI-Adept-Sampler")