diff --git "a/adept-sampler-v6/scripts/adept_sampler_v6.py" "b/adept-sampler-v6/scripts/adept_sampler_v6.py" new file mode 100644--- /dev/null +++ "b/adept-sampler-v6/scripts/adept_sampler_v6.py" @@ -0,0 +1,4717 @@ +""" +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 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("

" + "⚠️ Weight Scale / Shift / Start–End apply to the 6 patched k-diffusion samplers only. " + "Custom samplers (Akashic, Adept, Mirror) use their own internal parameters.

") + + 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("
") + gr.HTML("

🧠 Content-Aware Pacing (Euler Ancestral only)

") + gr.HTML("

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.

") + 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("
") + gr.HTML("

🌀 Custom Advanced Samplers

") + gr.HTML("

Enable to use Akashic/Adept/Ancestral samplers instead of k-diffusion

") + + 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("

" + "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<σ<5.0 range), then restarts the generation once with the " + "auto-computed correction applied from step 0 — total model calls for that " + "run will be higher than the step count.

") + + 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("

Active only when Custom Sampler = Mirror Correction Euler

") + 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("
") + gr.HTML("

🎛️ CFG Enhancements

") + gr.HTML("

" + "Combat CFG Drift works in stock A1111 via official callback. " + "Spectral Modulation & 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).

") + + + 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") +