From 5e5edac72fb895cfc478c79f93ad11213d5a467e Mon Sep 17 00:00:00 2001 From: Jordan Thompson Date: Mon, 27 Apr 2026 08:25:43 -0700 Subject: [PATCH] Initial commit --- .github/workflows/publish.yml | 22 + LICENSE | 42 +- README.md | 156 ++- __init__.py | 14 + pyproject.toml | 28 + sampling.py | 1927 +++++++++++++++++++++++++++++++++ 6 files changed, 2167 insertions(+), 22 deletions(-) create mode 100644 .github/workflows/publish.yml create mode 100644 __init__.py create mode 100644 pyproject.toml create mode 100644 sampling.py diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..a4fd892 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,22 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + - master + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/LICENSE b/LICENSE index 97b23e7..dfdf465 100644 --- a/LICENSE +++ b/LICENSE @@ -1,21 +1,21 @@ -MIT License - -Copyright (c) 2026 WAS - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. +MIT License + +Copyright (c) 2026 Jordan "WAS" Thompson + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md index c6aefcb..e2ce6d9 100644 --- a/README.md +++ b/README.md @@ -1 +1,155 @@ -# RES4SHO \ No newline at end of file +# High-Frequency Detail Sampling based on Res Sampling + +This is a ComfyUI custom node that enhances fine detail preservation in diffusion model outputs using spectral high-frequency emphasis (HFE). + +## Installation + +Clone or copy this folder into your ComfyUI `custom_nodes` directory: + +``` +ComfyUI/ + custom_nodes/ + RES4SHO/ + __init__.py + sampling.py +``` + +Restart ComfyUI. The new samplers and schedulers will appear in the dropdown menus of any **KSampler** node. + +## Samplers + +All samplers are exponential integrators with phi-function coefficients. The HFE enhancement extracts high-frequency detail from inter-stage correction deltas via a 3x3 spatial high-pass filter and re-injects it with configurable strength. + +### Fixed-Strength Presets + +Each stage count offers 8 strength levels (`s1` = no emphasis, `s8` = maximum potential sharpness): + +| Sampler | Stages | Model Evals/Step | +|---------|--------|-----------------| +| `hfe_s1` .. `hfe_s8` | 2 | 2 | +| `hfe3_s1` .. `hfe3_s8` | 3 | 3 | +| `hfe4_s1` .. `hfe4_s8` | 4 | 4 | +| `hfe5_s1` .. `hfe5_s8` | 5 | 5 | + +Higher stage counts provide better ODE integration accuracy at the cost of more model evaluations per step. + +### Adaptive (Auto) Samplers + +Per-step adaptive `eta` based on sigma envelope and content gating: + +| Sampler | Stages | Description | +|---------|--------|-------------| +| `hfe_auto` | 2 | Variable c2, eta, and kernel per step | +| `hfe3_auto` | 3 | Per-step eta with 3-stage integrator | +| `hfe4_auto` | 4 | Per-step eta with 4-stage integrator | +| `hfe5_auto` | 5 | Per-step eta with 5-stage integrator | + +**How auto adapts:** +- **Sigma envelope** (smoothstep): suppresses emphasis at high noise (early steps), full strength in the detail-forming range +- **Content gate**: reduces emphasis when the model correction is already HF-rich; increases it when the correction is smooth and needs boosting + +### Experimental Samplers (hfx_*) + +Alternative HF extraction methods, all using a 2-stage base: + +| Sampler | Method | +|---------|--------| +| `hfx_lap` | Laplacian pyramid multi-scale (3 bands) | +| `hfx_mom` | Correction momentum (EMA across steps) | +| `hfx_fft` | FFT spectral high-pass with smooth cutoff | +| `hfx_sde` | Stochastic HF noise injection | +| `hfx_spatial` | Spatially-adaptive per-pixel gating | + +**Hybrids** (combine two techniques): +- `hfx_lap_mom` -- Laplacian pyramid + momentum +- `hfx_lap_spatial` -- Laplacian pyramid + spatial gating +- `hfx_fft_spatial` -- FFT spectral + spatial gating + +**Band profile variants:** +- `hfx_lap_fine` -- fine-detail emphasis (edges, texture) +- `hfx_lap_broad` -- even emphasis across frequency bands + +Each experimental mode also has 4 graduated strength presets (`_s1` .. `_s4`), e.g. `hfx_lap_s1`, `hfx_mom_s3`, etc. + +## Schedulers + +Arctangent S-curve schedulers that concentrate step density in the detail-forming sigma range: + +| Scheduler | Description | +|-----------|-------------| +| `atan_gentle` | Mild mid-sigma concentration | +| `atan_focused` | Moderate detail-range concentration | +| `atan_steep` | Aggressive detail-range concentration | +| `karras_tan` | Karras-Tangent hybrid (experimental) | +| `logistic` | Logistic sigmoid S-curve (experimental) | + +An ASCII sigma chart is printed to the console when a scheduler is used. + +## Recommended Combinations + +### Getting Started + +| Goal | Sampler | Scheduler | Notes | +|------|---------|-----------|-------| +| General use | `hfe_auto` | `atan_focused` | Best all-rounder -- adaptive emphasis handles most content | +| Subtle enhancement | `hfe_s3` | `atan_gentle` | Light touch, minimal risk of artifacts | +| Strong detail | `hfe_s6` | `atan_steep` | Noticeably sharper textures and edges | +| Maximum sharpness | `hfe_s7` or `hfe_s8` | `atan_steep` | Aggressive -- inspect for over-sharpening | + +### By Content Type + +| Content | Sampler | Scheduler | Why | +|---------|---------|-----------|-----| +| Portraits / faces | `hfe_auto` | `atan_focused` | Auto gate protects smooth skin while sharpening eyes, hair, pores | +| Landscapes / nature | `hfe_s5` | `atan_gentle` | Fixed mid-strength avoids over-enhancing skies and gradients | +| Architecture / hard surfaces | `hfe_s7` | `atan_steep` | Strong emphasis on edges and geometric detail | +| Text / UI renders | `hfx_lap_fine` | `atan_steep` | Fine-band Laplacian targets glyph edges specifically | +| Fabric / organic texture | `hfx_lap_broad` | `atan_focused` | Even multi-scale emphasis across weave and folds | +| Illustrations / anime | `hfe_s4` | `atan_gentle` | Light emphasis preserves flat shading without adding unwanted texture | + +### High-Accuracy Integrators + +More model evaluations per step for better ODE integration -- useful at low step counts or with difficult models: + +| Sampler | Scheduler | Use Case | +|---------|-----------|----------| +| `hfe3_auto` | `atan_focused` | Good balance of accuracy and speed (3 evals/step) | +| `hfe4_auto` | `atan_focused` | High accuracy for complex prompts (4 evals/step) | +| `hfe5_auto` | `atan_gentle` | Maximum integration accuracy (5 evals/step) | +| `hfe4_s5` | `atan_steep` | Fixed-strength detail + 4-stage accuracy | +| `hfe5_s6` | `karras_tan` | High emphasis + high accuracy + Karras hybrid spacing | + +### Experimental Combinations + +| Sampler | Scheduler | Character | +|---------|-----------|-----------| +| `hfx_lap` | `atan_focused` | Multi-scale detail -- good default experimental choice | +| `hfx_fft` | `atan_steep` | Frequency-domain sharpening -- clean spectral separation | +| `hfx_spatial` | `atan_focused` | Sharpens high-variance regions, leaves smooth areas alone | +| `hfx_mom` | `atan_gentle` | Accumulates detail across steps -- builds up gradually | +| `hfx_sde` | `atan_gentle` | Stochastic texture injection -- adds micro-variation | +| `hfx_lap_mom` | `atan_focused` | Multi-scale + momentum -- rich progressive detail | +| `hfx_lap_spatial` | `atan_steep` | Multi-scale + spatial gating -- targeted sharpening | +| `hfx_fft_spatial` | `atan_focused` | Spectral + spatial -- precise frequency-aware gating | + +### Scheduler Pairings + +| Scheduler | Best With | Character | +|-----------|-----------|-----------| +| `atan_gentle` | Low-strength samplers (`s1`-`s4`), stochastic modes | Mild concentration, safe for all content | +| `atan_focused` | Auto samplers, mid-strength presets (`s4`-`s6`) | Balanced step density in detail range | +| `atan_steep` | High-strength samplers (`s6`-`s8`), architectural content | Aggressive detail-range concentration | +| `karras_tan` | High-stage integrators (`hfe4_*`, `hfe5_*`) | Karras optimal spacing + tangent warp | +| `logistic` | Any -- alternative S-curve shape | Sharper transition through detail range, flatter extremes | + +## How It Works + +**Base integrator:** Multi-stage singlestep exponential integrator (res_Ns) with phi-function coefficients, giving exact treatment of exponential decay and higher-order corrections from intermediate evaluations. + +**HFE enhancement:** The inter-stage correction delta captures what the model reveals at lower noise -- texture, edges, micro-structure. A spatial high-pass (residual after box blur in latent space) extracts the fine detail component, which is re-injected with extra weight `eta`. This compounds across every step. + +**Cost:** One 3x3 `avg_pool` per step for all variants (negligible vs. model evaluation). Auto samplers add a few scalar ops on top. + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..5537097 --- /dev/null +++ b/__init__.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +""" +RES4SHO -- High-Frequency Detail Sampling for ComfyUI + +Custom samplers and schedulers that enhance fine detail preservation +in diffusion model outputs via spectral high-frequency emphasis (HFE). + +Adds new entries to the sampler and scheduler dropdowns in KSampler nodes. +No additional custom nodes are created. +""" + +from .sampling import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..2a377d8 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,28 @@ +[project] +name = "RES4SHO" +version = "1.0.0" +description = "High-Frequency Detail Sampling based on Res Sampling for ComfyUI" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT License" } +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Environment :: GPU :: NVIDIA CUDA", + "Environment :: GPU :: AMD ROCm", + "Environment :: GPU :: Apple Metal" +] +dynamic = ["dependencies"] + +[tool.setuptools.dynamic] +dependencies = { file = ["requirements.txt"] } + +[project.urls] +Repository = "https://github.com/WASasquatch/RES4SHO" +"Bug Tracker" = "https://github.com/WASasquatch/RES4SHO/issues" + +[tool.comfy] +PublisherId = "was" +DisplayName = "RES4SHO" +requires-comfyui = ">=1.0.0" diff --git a/sampling.py b/sampling.py new file mode 100644 index 0000000..57579bb --- /dev/null +++ b/sampling.py @@ -0,0 +1,1927 @@ +# -*- coding: utf-8 -*- +""" +HF-Detail Sampling -- exponential integrator with spectral high-frequency +emphasis (HFE), tuned for realistic detail preservation. + +Samplers +-------- + 2-stage (res_2s base): hfe_s1..s8, hfe_auto + 3-stage (res_3s base): hfe3_s1..s8, hfe3_auto + 4-stage (res_4s base): hfe4_s1..s8, hfe4_auto + 5-stage (res_5s base): hfe5_s1..s8, hfe5_auto + + s1 = no HF emphasis (vanilla integrator) + s8 = maximum sharpness + auto = per-step adaptive eta based on sigma envelope and content gate + + Higher stage counts = more model evaluations per step = higher ODE + integration accuracy. The HFE enhancement is applied the same way + across all stage counts. + +Schedulers (arctangent S-curve, bong_tangent-inspired) +------------------------------------------------------ + atan_gentle -- mild mid-sigma concentration + atan_focused -- moderate detail-range concentration + atan_steep -- aggressive detail-range concentration + +How the HFE sampler works +-------------------------- +Base: 2-stage singlestep exponential integrator (res_2s) with phi-function +coefficients, giving an exact treatment of exponential decay and a second- +order correction from a midpoint evaluation. + +Enhancement: the inter-stage correction delta (denoised_2 - denoised_1) +captures what the model reveals at lower noise -- texture, edges, micro- +structure. A 3x3 spatial high-pass (residual after box blur in latent +space) extracts the fine detail component, which is re-injected with extra +weight ``eta``. This compounds across every step. + +How hfe_auto adapts +-------------------- + eta_effective = eta_peak * sigma_envelope * content_gate + + sigma_envelope: smoothstep from 0 (high noise, no emphasis) to 1 (detail + range, full emphasis). Prevents noise amplification at early steps. + + content_gate: measures HF energy in the correction delta. When the model + correction is already rich in high-frequency content, the gate reduces + emphasis (detail is already there). When the correction is smooth, the + gate opens wider (detail needs boosting). + +Cost: one 3x3 avg_pool per step for all variants (negligible vs model eval). +hfe_auto adds a few scalar ops on top. +""" + +import math +import logging +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from tqdm.auto import trange + +import comfy.samplers as comfy_samplers + +LOGGER = logging.getLogger("HFDetailSampling") + + +# ===================================================================== +# Phi functions (exponential integrator building blocks) +# ===================================================================== +# +# phi1(-h) = (1 - e^{-h}) / h +# phi2(-h) = (e^{-h} - 1 + h) / h^2 +# +# Taylor branches avoid catastrophic cancellation near h ~ 0. + +def _phi1(h: torch.Tensor) -> torch.Tensor: + """phi1(-h) for positive h. Scalar or broadcastable tensor.""" + return torch.where( + h.abs() > 1e-4, + (1.0 - torch.exp(-h)) / h, + 1.0 - h / 2.0 + h * h / 6.0, + ) + + +def _phi2(h: torch.Tensor) -> torch.Tensor: + """phi2(-h) for positive h. Scalar or broadcastable tensor.""" + h2 = h * h + return torch.where( + h.abs() > 1e-4, + (torch.exp(-h) - 1.0 + h) / h2, + 0.5 - h / 6.0 + h2 / 24.0, + ) + + +def _phi3(h: torch.Tensor) -> torch.Tensor: + """phi3(-h) for positive h. Needed by 4-stage and 5-stage integrators.""" + h2 = h * h + h3 = h2 * h + return torch.where( + h.abs() > 1e-4, + (1.0 - torch.exp(-h) - h + h2 / 2.0) / h3, + 1.0 / 6.0 - h / 24.0 + h2 / 120.0, + ) + + +# ===================================================================== +# Spectral detail extraction +# ===================================================================== + +def _extract_hf(t: torch.Tensor, kernel_size: int = 3) -> torch.Tensor: + """ + Spatial high-pass via residual after box blur. + + For a 4D [B,C,H,W] latent tensor this isolates edges, texture, and + micro-structure. Returns zeros for non-4D inputs. + """ + if t.ndim != 4: + return torch.zeros_like(t) + pad = kernel_size // 2 + padded = F.pad(t, [pad, pad, pad, pad], mode='reflect') + low = F.avg_pool2d(padded, kernel_size, stride=1) + return t - low + + +def _extract_hf_pyramid(t: torch.Tensor, levels: int = 3) -> list: + """ + Laplacian pyramid: decompose into multiple frequency bands. + + Returns a list of [fine, medium, coarse] band tensors. + Each band captures progressively lower spatial frequencies via + increasing kernel sizes (3, 5, 7, ...). + """ + if t.ndim != 4: + return [torch.zeros_like(t)] * levels + bands = [] + current = t + for lvl in range(levels): + ks = 3 + 2 * lvl # 3, 5, 7 + pad = ks // 2 + padded = F.pad(current, [pad, pad, pad, pad], mode='reflect') + blurred = F.avg_pool2d(padded, ks, stride=1) + bands.append(current - blurred) + current = blurred + return bands + + +def _extract_hf_fft(t: torch.Tensor, cutoff: float = 0.3) -> torch.Tensor: + """ + FFT high-pass filter with smooth ramp. + + Extracts frequencies above ``cutoff`` (fraction of Nyquist). + Ramps linearly from 0 at cutoff to 1 at 0.5 (Nyquist). + """ + if t.ndim != 4: + return torch.zeros_like(t) + H, W = t.shape[2], t.shape[3] + freq = torch.fft.rfft2(t) + fy = torch.fft.fftfreq(H, device=t.device).unsqueeze(1) + fx = torch.fft.rfftfreq(W, device=t.device).unsqueeze(0) + freq_mag = torch.sqrt(fy ** 2 + fx ** 2) + ramp_width = max(0.5 - cutoff, 1e-6) + mask = torch.clamp((freq_mag - cutoff) / ramp_width, 0.0, 1.0) + return torch.fft.irfft2(freq * mask, s=(H, W)) + + +def _spatial_gate(delta: torch.Tensor, window: int = 7) -> torch.Tensor: + """ + Per-pixel gate based on local energy of the correction delta. + + Returns a [0, 1] spatial map: 1 in high-variance regions (faces, + text, fine objects) where emphasis helps; 0 in smooth areas (sky, + gradients) where emphasis would add noise. + """ + if delta.ndim != 4: + return torch.ones_like(delta) + energy = delta ** 2 + pad = window // 2 + padded = F.pad(energy, [pad, pad, pad, pad], mode='reflect') + local_energy = F.avg_pool2d(padded, window, stride=1) + e_max = local_energy.amax(dim=(-1, -2), keepdim=True).clamp(min=1e-8) + return local_energy / e_max + + +# ===================================================================== +# Console sigma plot +# ===================================================================== + +def _plot_sigmas(sigmas: torch.Tensor, name: str, + width: int = 64, height: int = 16) -> None: + """Render a sigma schedule as an ASCII chart in the console.""" + vals = sigmas.tolist() + if vals and vals[-1] == 0.0: + vals = vals[:-1] + n = len(vals) + if n < 2: + return + + y_hi = max(vals) + y_lo = min(vals) + y_span = y_hi - y_lo + if y_span < 1e-12: + return + + # Build character canvas + grid = [[' '] * width for _ in range(height)] + for i, v in enumerate(vals): + c = int(i * (width - 1) / (n - 1) + 0.5) + r = int((y_hi - v) * (height - 1) / y_span + 0.5) + grid[max(0, min(height - 1, r))][max(0, min(width - 1, c))] = '*' + + # Y-axis label positions (5 evenly spaced) + label_rows = {0, height // 4, height // 2, 3 * height // 4, height - 1} + + out = [ + '', + f' {name} ({n} steps, sigma {y_hi:.2f} -> {y_lo:.4f})', + f' +{"-" * width}+', + ] + for r in range(height): + y = y_hi - r * y_span / (height - 1) + lbl = f'{y:7.2f}' if r in label_rows else ' ' + out.append(f' {lbl} |{"".join(grid[r])}|') + out.append(f' +{"-" * width}+') + + # X-axis labels: 0, midpoint, end + mid_s = str(n // 2) + end_s = str(n) + gap1 = width // 2 - len(mid_s) + gap2 = width - width // 2 - len(end_s) + out.append(f' 0{" " * gap1}{mid_s}{" " * gap2}{end_s}') + out.append(f' {"step":^{width}}') + out.append('') + + print('\n'.join(out)) + + +# ===================================================================== +# Fixed-eta sampler core +# ===================================================================== +# +# Butcher tableau (res_2s exponential, 2-stage singlestep): +# +# 0 | 0 0 +# c2 | c2*phi1(-h*c2) 0 +# ----+----------------------------- +# | phi1(-h) - phi2(-h)/c2 phi2(-h)/c2 +# +# Spectral sharpening: +# delta = e2 - e1 (correction signal) +# delta_hf = high_pass_3x3(delta) (fine spatial detail) +# e2' = e2 + eta * delta_hf (amplified texture/edges) +# x_next = x + h * (b1*e1 + b2*e2') +# +# eta = 0 recovers standard res_2s exactly. + +@torch.no_grad() +def _sample_hfe( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + c2: float = 0.5, + eta: float = 0.0, +) -> torch.Tensor: + """ + Exponential integrator with fixed-strength spectral detail sharpening. + + c2: intermediate evaluation point in (0,1]. + eta: high-frequency amplification strength. 0 = standard res_2s. + """ + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + total_steps = len(sigmas) - 1 + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + # Final step: sigma_next ~ 0, just return denoised + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + h = torch.log(sigma / sigma_next) + phi1 = _phi1(h) + phi2 = _phi2(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Intermediate point --- + sigma_mid = sigma * torch.exp(-c2 * h) + hc2 = h * c2 + a21 = c2 * _phi1(hc2) + X_2 = x + h * a21 * eps_1 + + # --- Stage 2 --- + denoised_2 = model(X_2, sigma_mid * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Spectral HF sharpening --- + # Sigma warmup: suppress at high noise (first ~25% of steps) to + # prevent amplifying noise. Full strength from ~55% onward. + if eta > 0.0: + progress = i / max(total_steps - 1, 1) + sigma_gate = max(0.0, min(1.0, (progress - 0.25) / 0.30)) + eta_step = eta * sigma_gate + + if eta_step > 1e-3: + delta = eps_2 - eps_1 + delta_hf = _extract_hf(delta) + eps_2 = eps_2 + eta_step * delta_hf + + # --- Output weights (standard res_2s) --- + b2 = phi2 / c2 + b1 = phi1 - b2 + + x = x + h * (b1 * eps_1 + b2 * eps_2) + + # Guard against NaN/inf from numerical instability + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE step %d: NaN/inf detected, falling back to " + "standard res_2s for remaining steps.", i) + eta = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_2, + "x": x, + }) + + return x + + +# ===================================================================== +# Adaptive sampler (hfe_auto) +# ===================================================================== +# +# Three things adapt per step: +# +# 1. c2 (Butcher tableau): ramps from c2_start (conservative, high sigma) +# to c2_end (aggressive, low sigma). This changes the actual ODE +# solver weights each step -- not just a scaling knob. +# +# 2. eta (HF emphasis): eta_peak * sigma_envelope * content_gate +# - sigma_envelope: smoothstep, 0 at high noise, 1 in detail range. +# - content_gate: 0 when correction is already HF-rich (model is +# producing detail on its own), 1 when smooth (needs sharpening). +# Full [0, 1] range -- no floor, so it can fully shut off. +# +# 3. kernel_size: 3x3 at low sigma (fine texture), 5x5 at high sigma +# (coarser structural detail). + +@torch.no_grad() +def _sample_hfe_auto( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta_peak: float = 0.55, + c2_start: float = 0.45, + c2_end: float = 0.85, +) -> torch.Tensor: + """ + Fully adaptive HFE sampler. + + eta_peak: maximum HF amplification (reached at low sigma with smooth correction). + c2_start: intermediate eval point at high sigma (conservative). + c2_end: intermediate eval point at low sigma (aggressive detail capture). + """ + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + sigma_max = float(sigmas[0]) + total_steps = len(sigmas) - 1 + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + # --- Adaptive c2: ramps from conservative to aggressive --- + progress = 1.0 - float(sigma) / sigma_max # 0 at start, 1 at end + c2 = c2_start + progress * (c2_end - c2_start) + + h = torch.log(sigma / sigma_next) + phi1 = _phi1(h) + phi2 = _phi2(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Intermediate point (c2 varies per step) --- + sigma_mid = sigma * torch.exp(-c2 * h) + hc2 = h * c2 + a21 = c2 * _phi1(hc2) + X_2 = x + h * a21 * eps_1 + + # --- Stage 2 --- + denoised_2 = model(X_2, sigma_mid * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Adaptive eta --- + # Adaptive kernel: 5x5 early (coarser detail), 3x3 late (fine texture) + ks = 5 if progress < 0.5 else 3 + delta = eps_2 - eps_1 + delta_hf = _extract_hf(delta, kernel_size=ks) + + # Sigma envelope: smoothstep, suppresses at high noise + envelope = progress * progress * (3.0 - 2.0 * progress) + + # Content gate: full range [0, 1] + # 0 = correction already HF-rich (model producing detail on its own) + # 1 = correction is smooth (detail needs boosting) + hf_energy = float((delta_hf ** 2).mean()) + total_energy = float((delta ** 2).mean()) + hf_ratio = hf_energy / (total_energy + 1e-8) + content_gate = max(0.0, 1.0 - 2.0 * hf_ratio) + + eta_step = eta_peak * envelope * content_gate + + if eta_step > 1e-3: + eps_2 = eps_2 + eta_step * delta_hf + + # --- Output weights (change every step with adaptive c2) --- + b2 = phi2 / c2 + b1 = phi1 - b2 + + x = x + h * (b1 * eps_1 + b2 * eps_2) + + # Guard against NaN/inf + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE auto step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta_peak = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_2, + "x": x, + }) + + return x + + +def sample_hfe_auto(model, x, sigmas, extra_args=None, callback=None, disable=False): + """Adaptive HFE -- variable c2, eta, and kernel per step.""" + LOGGER.info(">>> hfe_auto sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfe_auto( + model, x, sigmas, extra_args, callback, disable, + eta_peak=0.55, c2_start=0.45, c2_end=0.85, + ) + + +# ===================================================================== +# 3-stage exponential integrator with HFE (hfe3_*) +# ===================================================================== +# +# Butcher tableau (res_3s, c2=1/2, c3=1): +# +# 0 | 0 0 0 +# 1/2 | a2_1 0 0 +# 1 | a3_1 a3_2 0 +# ----+-------------------------------- +# | b1 b2 b3 +# +# gamma = (3*c3^3 - 2*c3) / (c2*(2 - 3*c2)) = 4 +# 3 model evaluations per step. + +_3S_C2 = 0.5 +_3S_C3 = 1.0 +_3S_GAMMA = 4.0 + + +@torch.no_grad() +def _sample_hfe_3s( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta: float = 0.0, +) -> torch.Tensor: + """3-stage exponential integrator with fixed-strength HFE.""" + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + total_steps = len(sigmas) - 1 + c2, c3, gamma = _3S_C2, _3S_C3, _3S_GAMMA + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + h = torch.log(sigma / sigma_next) + phi1_h = _phi1(h) + phi2_h = _phi2(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 (c2=1/2) --- + hc2 = h * c2 + a2_1 = c2 * _phi1(hc2) + X_2 = x + h * a2_1 * eps_1 + sigma_2 = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_2 * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Stage 3 (c3=1) --- + hc3 = h * c3 + a3_2 = gamma * c2 * _phi2(hc2) + (c3 ** 2 / c2) * _phi2(hc3) + a3_1 = c3 * _phi1(hc3) - a3_2 + X_3 = x + h * (a3_1 * eps_1 + a3_2 * eps_2) + sigma_3 = sigma * torch.exp(-c3 * h) + denoised_3 = model(X_3, sigma_3 * s_in, **extra_args) + eps_3 = denoised_3 - x + + # --- Spectral HF sharpening --- + if eta > 0.0: + progress = i / max(total_steps - 1, 1) + sigma_gate = max(0.0, min(1.0, (progress - 0.25) / 0.30)) + eta_step = eta * sigma_gate + if eta_step > 1e-3: + delta = eps_3 - eps_1 + delta_hf = _extract_hf(delta) + eps_3 = eps_3 + eta_step * delta_hf + + # --- Output weights --- + b3 = phi2_h / (gamma * c2 + c3) + b2 = gamma * b3 + b1 = phi1_h - b2 - b3 + + x = x + h * (b1 * eps_1 + b2 * eps_2 + b3 * eps_3) + + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE 3s step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_3, + "x": x, + }) + + return x + + +@torch.no_grad() +def _sample_hfe_3s_auto( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta_peak: float = 0.55, +) -> torch.Tensor: + """3-stage adaptive HFE -- per-step eta based on sigma and content.""" + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + sigma_max = float(sigmas[0]) + total_steps = len(sigmas) - 1 + c2, c3, gamma = _3S_C2, _3S_C3, _3S_GAMMA + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + progress = 1.0 - float(sigma) / sigma_max + h = torch.log(sigma / sigma_next) + phi1_h = _phi1(h) + phi2_h = _phi2(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 --- + hc2 = h * c2 + a2_1 = c2 * _phi1(hc2) + X_2 = x + h * a2_1 * eps_1 + sigma_2 = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_2 * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Stage 3 --- + hc3 = h * c3 + a3_2 = gamma * c2 * _phi2(hc2) + (c3 ** 2 / c2) * _phi2(hc3) + a3_1 = c3 * _phi1(hc3) - a3_2 + X_3 = x + h * (a3_1 * eps_1 + a3_2 * eps_2) + sigma_3 = sigma * torch.exp(-c3 * h) + denoised_3 = model(X_3, sigma_3 * s_in, **extra_args) + eps_3 = denoised_3 - x + + # --- Adaptive eta --- + ks = 5 if progress < 0.5 else 3 + delta = eps_3 - eps_1 + delta_hf = _extract_hf(delta, kernel_size=ks) + envelope = progress * progress * (3.0 - 2.0 * progress) + hf_energy = float((delta_hf ** 2).mean()) + total_energy = float((delta ** 2).mean()) + hf_ratio = hf_energy / (total_energy + 1e-8) + content_gate = max(0.0, 1.0 - 2.0 * hf_ratio) + eta_step = eta_peak * envelope * content_gate + + if eta_step > 1e-3: + eps_3 = eps_3 + eta_step * delta_hf + + # --- Output weights --- + b3 = phi2_h / (gamma * c2 + c3) + b2 = gamma * b3 + b1 = phi1_h - b2 - b3 + + x = x + h * (b1 * eps_1 + b2 * eps_2 + b3 * eps_3) + + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE 3s auto step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta_peak = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_3, + "x": x, + }) + + return x + + +def sample_hfe3_auto(model, x, sigmas, extra_args=None, callback=None, disable=False): + """3-stage adaptive HFE.""" + LOGGER.info(">>> hfe3_auto sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfe_3s_auto(model, x, sigmas, extra_args, callback, disable, + eta_peak=0.55) + + +# ===================================================================== +# 4-stage exponential integrator with HFE (hfe4_*) +# ===================================================================== +# +# Butcher tableau (Strehmel-Weiner, c2=1/2, c3=1/2, c4=1): +# +# 0 | 0 0 0 0 +# 1/2 | a2_1 0 0 0 +# 1/2 | a3_1 a3_2 0 0 +# 1 | a4_1 a4_2 a4_3 0 +# ----+-------------------------------------------- +# | b1 b2 b3 b4 +# +# 4 model evaluations per step. Weak 4th order accuracy. + +_4S_C2 = 0.5 +_4S_C3 = 0.5 +_4S_C4 = 1.0 + + +@torch.no_grad() +def _sample_hfe_4s( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta: float = 0.0, +) -> torch.Tensor: + """4-stage exponential integrator with fixed-strength HFE.""" + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + total_steps = len(sigmas) - 1 + c2, c3, c4 = _4S_C2, _4S_C3, _4S_C4 + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + h = torch.log(sigma / sigma_next) + hc2 = h * c2 + hc3 = h * c3 + phi1_h = _phi1(h) + phi2_h = _phi2(h) + phi3_h = _phi3(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 (c2=1/2) --- + a2_1 = c2 * _phi1(hc2) + X_2 = x + h * a2_1 * eps_1 + sigma_2 = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_2 * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Stage 3 (c3=1/2) --- + a3_2 = c3 * _phi2(hc3) + a3_1 = c3 * _phi1(hc3) - a3_2 + X_3 = x + h * (a3_1 * eps_1 + a3_2 * eps_2) + sigma_3 = sigma * torch.exp(-c3 * h) + denoised_3 = model(X_3, sigma_3 * s_in, **extra_args) + eps_3 = denoised_3 - x + + # --- Stage 4 (c4=1) --- + a4_2 = -2.0 * phi2_h + a4_3 = 4.0 * phi2_h + a4_1 = phi1_h - a4_2 - a4_3 + X_4 = x + h * (a4_1 * eps_1 + a4_2 * eps_2 + a4_3 * eps_3) + sigma_4 = sigma * torch.exp(-c4 * h) + denoised_4 = model(X_4, sigma_4 * s_in, **extra_args) + eps_4 = denoised_4 - x + + # --- Spectral HF sharpening --- + if eta > 0.0: + progress = i / max(total_steps - 1, 1) + sigma_gate = max(0.0, min(1.0, (progress - 0.25) / 0.30)) + eta_step = eta * sigma_gate + if eta_step > 1e-3: + delta = eps_4 - eps_1 + delta_hf = _extract_hf(delta) + eps_4 = eps_4 + eta_step * delta_hf + + # --- Output weights (Strehmel-Weiner, b2=0) --- + b3 = 4.0 * phi2_h - 8.0 * phi3_h + b4 = -phi2_h + 4.0 * phi3_h + b1 = phi1_h - b3 - b4 + + x = x + h * (b1 * eps_1 + b3 * eps_3 + b4 * eps_4) + + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE 4s step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_4, + "x": x, + }) + + return x + + +@torch.no_grad() +def _sample_hfe_4s_auto( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta_peak: float = 0.55, +) -> torch.Tensor: + """4-stage adaptive HFE -- per-step eta based on sigma and content.""" + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + sigma_max = float(sigmas[0]) + total_steps = len(sigmas) - 1 + c2, c3, c4 = _4S_C2, _4S_C3, _4S_C4 + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + progress = 1.0 - float(sigma) / sigma_max + h = torch.log(sigma / sigma_next) + hc2 = h * c2 + hc3 = h * c3 + phi1_h = _phi1(h) + phi2_h = _phi2(h) + phi3_h = _phi3(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 --- + a2_1 = c2 * _phi1(hc2) + X_2 = x + h * a2_1 * eps_1 + sigma_2 = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_2 * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Stage 3 --- + a3_2 = c3 * _phi2(hc3) + a3_1 = c3 * _phi1(hc3) - a3_2 + X_3 = x + h * (a3_1 * eps_1 + a3_2 * eps_2) + sigma_3 = sigma * torch.exp(-c3 * h) + denoised_3 = model(X_3, sigma_3 * s_in, **extra_args) + eps_3 = denoised_3 - x + + # --- Stage 4 --- + a4_2 = -2.0 * phi2_h + a4_3 = 4.0 * phi2_h + a4_1 = phi1_h - a4_2 - a4_3 + X_4 = x + h * (a4_1 * eps_1 + a4_2 * eps_2 + a4_3 * eps_3) + sigma_4 = sigma * torch.exp(-c4 * h) + denoised_4 = model(X_4, sigma_4 * s_in, **extra_args) + eps_4 = denoised_4 - x + + # --- Adaptive eta --- + ks = 5 if progress < 0.5 else 3 + delta = eps_4 - eps_1 + delta_hf = _extract_hf(delta, kernel_size=ks) + envelope = progress * progress * (3.0 - 2.0 * progress) + hf_energy = float((delta_hf ** 2).mean()) + total_energy = float((delta ** 2).mean()) + hf_ratio = hf_energy / (total_energy + 1e-8) + content_gate = max(0.0, 1.0 - 2.0 * hf_ratio) + eta_step = eta_peak * envelope * content_gate + + if eta_step > 1e-3: + eps_4 = eps_4 + eta_step * delta_hf + + # --- Output weights --- + b2 = 0.0 + b3 = 4.0 * phi2_h - 8.0 * phi3_h + b4 = -phi2_h + 4.0 * phi3_h + b1 = phi1_h - b3 - b4 + + x = x + h * (b1 * eps_1 + b3 * eps_3 + b4 * eps_4) + + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE 4s auto step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta_peak = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_4, + "x": x, + }) + + return x + + +def sample_hfe4_auto(model, x, sigmas, extra_args=None, callback=None, disable=False): + """4-stage adaptive HFE.""" + LOGGER.info(">>> hfe4_auto sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfe_4s_auto(model, x, sigmas, extra_args, callback, disable, + eta_peak=0.55) + + +# ===================================================================== +# 5-stage exponential integrator with HFE (hfe5_*) +# ===================================================================== +# +# Butcher tableau (c2=1/2, c3=1/2, c4=1, c5=1/2): +# +# 0 | 0 0 0 0 0 +# 1/2 | a2_1 0 0 0 0 +# 1/2 | a3_1 a3_2 0 0 0 +# 1 | a4_1 a4_2 a4_3 0 0 +# 1/2 | a5_1 a5_2 a5_3 a5_4 0 +# ----+------------------------------------ +# | b1 b2 b3 b4 b5 +# +# 5 model evaluations per step. Non-monotonic node placement (c5=1/2). + +_5S_C2 = 0.5 +_5S_C3 = 0.5 +_5S_C4 = 1.0 +_5S_C5 = 0.5 + + +@torch.no_grad() +def _sample_hfe_5s( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta: float = 0.0, +) -> torch.Tensor: + """5-stage exponential integrator with fixed-strength HFE.""" + LOGGER.info(">>> _sample_hfe_5s called with eta=%.4f", eta) + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + total_steps = len(sigmas) - 1 + c2, c3, c4, c5 = _5S_C2, _5S_C3, _5S_C4, _5S_C5 + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + h = torch.log(sigma / sigma_next) + hc2 = h * c2 + hc3 = h * c3 + hc5 = h * c5 + phi1_h = _phi1(h) + phi2_h = _phi2(h) + phi3_h = _phi3(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 (c2=1/2) --- + a2_1 = c2 * _phi1(hc2) + X_2 = x + h * a2_1 * eps_1 + sigma_2 = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_2 * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Stage 3 (c3=1/2) --- + a3_2 = _phi2(hc3) + a3_1 = c3 * _phi1(hc3) - a3_2 + X_3 = x + h * (a3_1 * eps_1 + a3_2 * eps_2) + sigma_3 = sigma * torch.exp(-c3 * h) + denoised_3 = model(X_3, sigma_3 * s_in, **extra_args) + eps_3 = denoised_3 - x + + # --- Stage 4 (c4=1) --- + a4_2 = phi2_h + a4_3 = phi2_h + a4_1 = phi1_h - a4_2 - a4_3 + X_4 = x + h * (a4_1 * eps_1 + a4_2 * eps_2 + a4_3 * eps_3) + sigma_4 = sigma * torch.exp(-c4 * h) + denoised_4 = model(X_4, sigma_4 * s_in, **extra_args) + eps_4 = denoised_4 - x + + # --- Stage 5 (c5=1/2, non-monotonic) --- + phi2_hc5 = _phi2(hc5) + phi3_hc5 = _phi3(hc5) + a5_2 = 0.5 * phi2_hc5 - phi3_h + 0.25 * phi2_h - 0.5 * phi3_hc5 + a5_3 = a5_2 + a5_4 = 0.25 * phi2_hc5 - a5_2 + a5_1 = c5 * _phi1(hc5) - a5_2 - a5_3 - a5_4 + X_5 = x + h * (a5_1 * eps_1 + a5_2 * eps_2 + a5_3 * eps_3 + a5_4 * eps_4) + sigma_5 = sigma * torch.exp(-c5 * h) + denoised_5 = model(X_5, sigma_5 * s_in, **extra_args) + eps_5 = denoised_5 - x + + # --- Spectral HF sharpening --- + if eta > 0.0: + progress = i / max(total_steps - 1, 1) + sigma_gate = max(0.0, min(1.0, (progress - 0.25) / 0.30)) + eta_step = eta * sigma_gate + if eta_step > 1e-3: + delta = eps_5 - eps_1 + delta_hf = _extract_hf(delta) + eps_5 = eps_5 + eta_step * delta_hf + + # --- Output weights (b2=0, b3=0) --- + b4 = -phi2_h + 4.0 * phi3_h + b5 = 4.0 * phi2_h - 8.0 * phi3_h + b1 = phi1_h - b4 - b5 + + x = x + h * (b1 * eps_1 + b4 * eps_4 + b5 * eps_5) + + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE 5s step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_5, + "x": x, + }) + + return x + + +@torch.no_grad() +def _sample_hfe_5s_auto( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + eta_peak: float = 0.55, +) -> torch.Tensor: + """5-stage adaptive HFE -- per-step eta based on sigma and content.""" + if extra_args is None: + extra_args = {} + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + sigma_max = float(sigmas[0]) + total_steps = len(sigmas) - 1 + c2, c3, c4, c5 = _5S_C2, _5S_C3, _5S_C4, _5S_C5 + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + progress = 1.0 - float(sigma) / sigma_max + h = torch.log(sigma / sigma_next) + hc2 = h * c2 + hc3 = h * c3 + hc5 = h * c5 + phi1_h = _phi1(h) + phi2_h = _phi2(h) + phi3_h = _phi3(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 --- + a2_1 = c2 * _phi1(hc2) + X_2 = x + h * a2_1 * eps_1 + sigma_2 = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_2 * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Stage 3 --- + a3_2 = _phi2(hc3) + a3_1 = c3 * _phi1(hc3) - a3_2 + X_3 = x + h * (a3_1 * eps_1 + a3_2 * eps_2) + sigma_3 = sigma * torch.exp(-c3 * h) + denoised_3 = model(X_3, sigma_3 * s_in, **extra_args) + eps_3 = denoised_3 - x + + # --- Stage 4 --- + a4_2 = phi2_h + a4_3 = phi2_h + a4_1 = phi1_h - a4_2 - a4_3 + X_4 = x + h * (a4_1 * eps_1 + a4_2 * eps_2 + a4_3 * eps_3) + sigma_4 = sigma * torch.exp(-c4 * h) + denoised_4 = model(X_4, sigma_4 * s_in, **extra_args) + eps_4 = denoised_4 - x + + # --- Stage 5 --- + phi2_hc5 = _phi2(hc5) + phi3_hc5 = _phi3(hc5) + a5_2 = 0.5 * phi2_hc5 - phi3_h + 0.25 * phi2_h - 0.5 * phi3_hc5 + a5_3 = a5_2 + a5_4 = 0.25 * phi2_hc5 - a5_2 + a5_1 = c5 * _phi1(hc5) - a5_2 - a5_3 - a5_4 + X_5 = x + h * (a5_1 * eps_1 + a5_2 * eps_2 + a5_3 * eps_3 + a5_4 * eps_4) + sigma_5 = sigma * torch.exp(-c5 * h) + denoised_5 = model(X_5, sigma_5 * s_in, **extra_args) + eps_5 = denoised_5 - x + + # --- Adaptive eta --- + ks = 5 if progress < 0.5 else 3 + delta = eps_5 - eps_1 + delta_hf = _extract_hf(delta, kernel_size=ks) + envelope = progress * progress * (3.0 - 2.0 * progress) + hf_energy = float((delta_hf ** 2).mean()) + total_energy = float((delta ** 2).mean()) + hf_ratio = hf_energy / (total_energy + 1e-8) + content_gate = max(0.0, 1.0 - 2.0 * hf_ratio) + eta_step = eta_peak * envelope * content_gate + + if eta_step > 1e-3: + eps_5 = eps_5 + eta_step * delta_hf + + # --- Output weights (b2=0, b3=0) --- + b4 = -phi2_h + 4.0 * phi3_h + b5 = 4.0 * phi2_h - 8.0 * phi3_h + b1 = phi1_h - b4 - b5 + + x = x + h * (b1 * eps_1 + b4 * eps_4 + b5 * eps_5) + + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFE 5s auto step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", i) + eta_peak = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_5, + "x": x, + }) + + return x + + +def sample_hfe5_auto(model, x, sigmas, extra_args=None, callback=None, disable=False): + """5-stage adaptive HFE.""" + LOGGER.info(">>> hfe5_auto sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfe_5s_auto(model, x, sigmas, extra_args, callback, disable, + eta_peak=0.55) + + +# ===================================================================== +# Graduated fixed-strength presets (hfe_s1..s8, hfe3_s1..s8, etc.) +# ===================================================================== +# +# eta follows a power-1.5 curve from 0.00 (s1) to 0.48 (s8) so that +# the perceptual jump between adjacent levels feels roughly even. +# +# 2-stage presets also vary c2 from 0.45 to 0.80. +# 3/4/5-stage presets use fixed c values (from reference tableaux) +# and only vary eta. + +_HFE_LEVELS = 8 +_HFE_C2_MIN = 0.45 +_HFE_C2_MAX = 0.80 +_HFE_ETA_MAX = 0.48 + +# Map stage count -> (core function, name prefix) +_STAGE_CORES = { + 2: (_sample_hfe, "hfe"), + 3: (_sample_hfe_3s, "hfe3"), + 4: (_sample_hfe_4s, "hfe4"), + 5: (_sample_hfe_5s, "hfe5"), +} + + +def _make_hfe_preset(level: int, stages: int = 2): + """Factory: create a fixed-strength HFE sampler for the given level and stage count.""" + core_fn, prefix = _STAGE_CORES[stages] + t = level / (_HFE_LEVELS - 1) + eta = _HFE_ETA_MAX * (t ** 1.5) + + if stages == 2: + # 2-stage: also vary c2 + c2 = _HFE_C2_MIN + t * (_HFE_C2_MAX - _HFE_C2_MIN) + def sampler(model, x, sigmas, extra_args=None, callback=None, disable=False, + _c2=c2, _eta=eta): + return core_fn(model, x, sigmas, extra_args, callback, disable, + c2=_c2, eta=_eta) + sampler.__doc__ = (f"HFE {stages}s strength {level + 1}/{_HFE_LEVELS}" + f" -- c2={c2:.3f}, eta={eta:.3f}") + else: + # 3/4/5-stage: fixed c values, only vary eta + def sampler(model, x, sigmas, extra_args=None, callback=None, disable=False, + _eta=eta, _prefix=prefix, _level=level): + LOGGER.info(">>> %s_s%d preset invoked, passing eta=%.4f", + _prefix, _level + 1, _eta) + return core_fn(model, x, sigmas, extra_args, callback, disable, + eta=_eta) + sampler.__doc__ = (f"HFE {stages}s strength {level + 1}/{_HFE_LEVELS}" + f" -- eta={eta:.3f}") + + name = f"{prefix}_s{level + 1}" + sampler.__name__ = f"sample_{name}" + sampler.__qualname__ = sampler.__name__ + return name, sampler + + +# Generate all presets: hfe_s1..s8, hfe3_s1..s8, hfe4_s1..s8, hfe5_s1..s8 +_HFE_PRESETS = {} +for _stages in (2, 3, 4, 5): + for _lvl in range(_HFE_LEVELS): + _name, _fn = _make_hfe_preset(_lvl, _stages) + _HFE_PRESETS[_name] = _fn + + +# ===================================================================== +# Experimental HFE samplers (hfx_*) +# ===================================================================== +# +# Each variant modifies HOW high-frequency detail is extracted and/or +# applied, using a shared 2-stage exponential integrator base. +# +# All use fixed moderate strength (c2=0.65, eta=0.25) for direct +# comparison against hfe_s5. + +_HFX_C2 = 0.65 +_HFX_ETA = 0.25 +_HFX_SDE_STRENGTH = 0.08 +_HFX_MOM_BETA = 0.7 +_HFX_FFT_CUTOFF = 0.3 +_HFX_LAP_WEIGHTS = (1.5, 1.0, 0.5) + + +@torch.no_grad() +def _sample_hfx( + model: Any, + x: torch.Tensor, + sigmas: torch.Tensor, + extra_args: Optional[Dict[str, Any]] = None, + callback: Optional[Any] = None, + disable: bool = False, + *, + c2: float = _HFX_C2, + eta: float = _HFX_ETA, + mode: str = 'lap', + # Per-mode overrides (use module defaults when None) + lap_weights: Optional[tuple] = None, + mom_beta: Optional[float] = None, + fft_cutoff: Optional[float] = None, + sde_strength: Optional[float] = None, + spatial_window: Optional[int] = None, +) -> torch.Tensor: + """ + Generic experimental HFE sampler. + + mode: + 'lap' -- Laplacian pyramid multi-scale (3 bands, weighted) + 'mom' -- correction momentum (EMA across steps) + 'fft' -- FFT spectral high-pass with smooth cutoff + 'sde' -- stochastic HF noise injection after update + 'spatial' -- spatially-adaptive per-pixel gating + Hybrid modes (combine two techniques): + 'lap_mom' -- Laplacian pyramid + momentum accumulation + 'lap_spatial' -- Laplacian pyramid + spatial gating + 'fft_spatial' -- FFT spectral + spatial gating + """ + if extra_args is None: + extra_args = {} + + # Resolve per-mode defaults + _lap_w = lap_weights or _HFX_LAP_WEIGHTS + _mom_b = mom_beta if mom_beta is not None else _HFX_MOM_BETA + _fft_c = fft_cutoff if fft_cutoff is not None else _HFX_FFT_CUTOFF + _sde_s = sde_strength if sde_strength is not None else _HFX_SDE_STRENGTH + _sp_win = spatial_window if spatial_window is not None else 7 + + s_in = x.new_ones([x.shape[0]]) + sigmas = sigmas.to(device=x.device, dtype=x.dtype) + total_steps = len(sigmas) - 1 + + # Per-mode state + momentum_buf = None + + for i in trange(total_steps, disable=disable): + sigma = sigmas[i] + sigma_next = sigmas[i + 1] + + if sigma_next < 1e-6: + denoised = model(x, sigma * s_in, **extra_args) + x = denoised + if callback is not None: + callback({"i": i, "sigma": 0.0, "denoised": denoised, "x": x}) + break + + h = torch.log(sigma / sigma_next) + phi1_h = _phi1(h) + phi2_h = _phi2(h) + + # --- Stage 1 --- + denoised_1 = model(x, sigma * s_in, **extra_args) + eps_1 = denoised_1 - x + + # --- Stage 2 --- + hc2 = h * c2 + a21 = c2 * _phi1(hc2) + X_2 = x + h * a21 * eps_1 + sigma_mid = sigma * torch.exp(-c2 * h) + denoised_2 = model(X_2, sigma_mid * s_in, **extra_args) + eps_2 = denoised_2 - x + + # --- Experimental HFE --- + progress = i / max(total_steps - 1, 1) + sigma_gate = max(0.0, min(1.0, (progress - 0.25) / 0.30)) + eta_step = eta * sigma_gate + + if eta_step > 1e-3 and mode != 'sde': + delta = eps_2 - eps_1 + + if mode == 'lap': + bands = _extract_hf_pyramid(delta, levels=3) + correction = sum(w * b for w, b in zip(_lap_w, bands)) + eps_2 = eps_2 + eta_step * correction + + elif mode == 'mom': + delta_hf = _extract_hf(delta) + if momentum_buf is None: + momentum_buf = delta_hf.clone() + else: + momentum_buf = (_mom_b * momentum_buf + + (1.0 - _mom_b) * delta_hf) + eps_2 = eps_2 + eta_step * momentum_buf + + elif mode == 'fft': + delta_hf = _extract_hf_fft(delta, cutoff=_fft_c) + eps_2 = eps_2 + eta_step * delta_hf + + elif mode == 'spatial': + delta_hf = _extract_hf(delta) + gate = _spatial_gate(delta, window=_sp_win) + eps_2 = eps_2 + eta_step * gate * delta_hf + + elif mode == 'lap_mom': + bands = _extract_hf_pyramid(delta, levels=3) + correction = sum(w * b for w, b in zip(_lap_w, bands)) + if momentum_buf is None: + momentum_buf = correction.clone() + else: + momentum_buf = (_mom_b * momentum_buf + + (1.0 - _mom_b) * correction) + eps_2 = eps_2 + eta_step * momentum_buf + + elif mode == 'lap_spatial': + bands = _extract_hf_pyramid(delta, levels=3) + correction = sum(w * b for w, b in zip(_lap_w, bands)) + gate = _spatial_gate(delta, window=_sp_win) + eps_2 = eps_2 + eta_step * gate * correction + + elif mode == 'fft_spatial': + delta_hf = _extract_hf_fft(delta, cutoff=_fft_c) + gate = _spatial_gate(delta, window=_sp_win) + eps_2 = eps_2 + eta_step * gate * delta_hf + + # --- Output weights --- + b2 = phi2_h / c2 + b1 = phi1_h - b2 + + x = x + h * (b1 * eps_1 + b2 * eps_2) + + # --- SDE: post-update HF noise injection --- + if mode == 'sde' and eta_step > 1e-3: + noise = torch.randn_like(x) + noise_hf = _extract_hf(noise) + x = x + (_sde_s * float(sigma_next) * sigma_gate * noise_hf) + + # NaN/inf guard + if torch.isnan(x).any() or torch.isinf(x).any(): + LOGGER.warning("HFX %s step %d: NaN/inf detected, disabling " + "emphasis for remaining steps.", mode, i) + eta = 0.0 + x = x.nan_to_num(nan=0.0, posinf=1e4, neginf=-1e4) + + if callback is not None: + callback({ + "i": i, + "sigma": float(sigma_next), + "denoised": denoised_2, + "x": x, + }) + + return x + + +# --- Experimental sampler wrappers (base, no strength suffix) --- + +def sample_hfx_lap(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Laplacian pyramid multi-scale HFE (experimental).""" + LOGGER.info(">>> hfx_lap sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='lap') + + +def sample_hfx_mom(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Correction momentum HFE (experimental).""" + LOGGER.info(">>> hfx_mom sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='mom') + + +def sample_hfx_fft(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """FFT spectral shaping HFE (experimental).""" + LOGGER.info(">>> hfx_fft sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='fft') + + +def sample_hfx_sde(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Stochastic HF injection HFE (experimental).""" + LOGGER.info(">>> hfx_sde sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='sde') + + +def sample_hfx_spatial(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Spatially-adaptive gating HFE (experimental).""" + LOGGER.info(">>> hfx_spatial sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='spatial') + + +# --- Hybrid sampler wrappers --- + +def sample_hfx_lap_mom(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Laplacian pyramid + momentum HFE (experimental hybrid).""" + LOGGER.info(">>> hfx_lap_mom sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='lap_mom') + + +def sample_hfx_lap_spatial(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Laplacian pyramid + spatial gating HFE (experimental hybrid).""" + LOGGER.info(">>> hfx_lap_spatial sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='lap_spatial') + + +def sample_hfx_fft_spatial(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """FFT spectral + spatial gating HFE (experimental hybrid).""" + LOGGER.info(">>> hfx_fft_spatial sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='fft_spatial') + + +# --- Band profile variants for hfx_lap --- + +def sample_hfx_lap_fine(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Laplacian pyramid fine-detail emphasis (experimental).""" + LOGGER.info(">>> hfx_lap_fine sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='lap', lap_weights=(2.5, 0.8, 0.2)) + + +def sample_hfx_lap_broad(model, x, sigmas, extra_args=None, callback=None, + disable=False): + """Laplacian pyramid broad/even emphasis (experimental).""" + LOGGER.info(">>> hfx_lap_broad sampler invoked (%d sigmas)", len(sigmas)) + return _sample_hfx(model, x, sigmas, extra_args, callback, disable, + mode='lap', lap_weights=(1.0, 1.2, 1.0)) + + +# ===================================================================== +# Graduated experimental presets (hfx_*_s1..s4) +# ===================================================================== +# +# 4 strength tiers per mode, sweeping the key parameter for each mode. +# All use c2=0.65 (moderate). + +_HFX_LEVELS = 4 + +# Per-mode sweep definitions: (mode, param_name, values_s1_to_s4) +_HFX_SWEEPS = { + 'lap': { + 'param': 'eta', + 'values': (0.10, 0.20, 0.35, 0.50), + }, + 'mom': { + 'param': 'mom_beta', + 'values': (0.40, 0.55, 0.70, 0.85), + }, + 'fft': { + 'param': 'fft_cutoff', + 'values': (0.15, 0.25, 0.35, 0.45), + }, + 'sde': { + 'param': 'sde_strength', + 'values': (0.03, 0.06, 0.10, 0.15), + }, + 'spatial': { + 'param': 'eta', + 'values': (0.10, 0.20, 0.35, 0.50), + }, +} + + +def _make_hfx_preset(mode: str, level: int): + """Factory: create a graduated experimental sampler. + + For 'eta' sweeps, eta varies and mode-specific param stays default. + For mode-specific param sweeps, eta stays at _HFX_ETA and param varies. + """ + sweep = _HFX_SWEEPS[mode] + param = sweep['param'] + value = sweep['values'][level] + + if param == 'eta': + def sampler(model_fn, x, sigmas, extra_args=None, callback=None, + disable=False, _eta=value, _mode=mode): + return _sample_hfx(model_fn, x, sigmas, extra_args, callback, + disable, mode=_mode, eta=_eta) + desc = f"eta={value:.2f}" + else: + kwarg = {param: value} + def sampler(model_fn, x, sigmas, extra_args=None, callback=None, + disable=False, _mode=mode, _kw=kwarg): + return _sample_hfx(model_fn, x, sigmas, extra_args, callback, + disable, mode=_mode, **_kw) + desc = f"{param}={value}" + + name = f"hfx_{mode}_s{level + 1}" + sampler.__name__ = f"sample_{name}" + sampler.__qualname__ = sampler.__name__ + sampler.__doc__ = f"HFX {mode} strength {level + 1}/{_HFX_LEVELS} -- {desc}" + return name, sampler + + +_HFX_PRESETS = {} +for _mode in _HFX_SWEEPS: + for _lvl in range(_HFX_LEVELS): + _name, _fn = _make_hfx_preset(_mode, _lvl) + _HFX_PRESETS[_name] = _fn + + +# ===================================================================== +# Tangent S-curve scheduler (bong_tangent-inspired) +# ===================================================================== +# +# sigma(i) from an arctangent S-curve that concentrates steps around a +# pivot point. Higher slope = sharper bend = more step density at pivot. +# +# slope_adj = slope / (steps / 40) [normalization for step count] + +def _tangent_sigmas( + steps: int, + sigma_max: float, + sigma_min: float, + slope: float, + pivot_frac: float, +) -> torch.Tensor: + n = steps + if n < 1: + return torch.zeros(1, dtype=torch.float32) + + pivot = pivot_frac * (n - 1) + slope_adj = slope / max(n / 40.0, 0.1) + + xs = torch.arange(n, dtype=torch.float64) + raw = ((2.0 / math.pi) * torch.atan(-slope_adj * (xs - pivot)) + 1.0) / 2.0 + + r_max = raw[0].item() + r_min = raw[-1].item() + r_range = r_max - r_min + + if r_range < 1e-12: + normalized = torch.linspace(1.0, 0.0, n, dtype=torch.float64) + else: + normalized = (raw - r_min) / r_range + + sigmas = normalized * (sigma_max - sigma_min) + sigma_min + sigmas = torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]) + return sigmas.float() + + +def _tangent_schedule( + model_sampling: Any, + steps: int, + slope: float, + pivot_frac: float, + name: str = '', +) -> torch.Tensor: + sigma_max = float(model_sampling.sigma_max) + sigma_min = float(model_sampling.sigma_min) + sigmas = _tangent_sigmas(steps, sigma_max, sigma_min, slope, pivot_frac) + if name: + _plot_sigmas(sigmas, name) + return sigmas + + +def scheduler_atan_gentle(model_sampling: Any, steps: int) -> torch.Tensor: + return _tangent_schedule(model_sampling, steps, slope=0.7, pivot_frac=0.35, + name='atan_gentle') + + +def scheduler_atan_focused(model_sampling: Any, steps: int) -> torch.Tensor: + return _tangent_schedule(model_sampling, steps, slope=1.1, pivot_frac=0.40, + name='atan_focused') + + +def scheduler_atan_steep(model_sampling: Any, steps: int) -> torch.Tensor: + return _tangent_schedule(model_sampling, steps, slope=1.6, pivot_frac=0.45, + name='atan_steep') + + +# ===================================================================== +# Experimental schedulers +# ===================================================================== + +def _karras_tangent_sigmas( + steps: int, + sigma_max: float, + sigma_min: float, + rho: float = 7.0, + bend: float = 0.35, + pivot_frac: float = 0.40, +) -> torch.Tensor: + """ + Karras-Tangent hybrid schedule. + + Base: Karras optimal spacing (rho=7). + Enhancement: warp the time ramp with an arctangent bend to concentrate + more steps in the detail-forming sigma range. bend=0 gives pure Karras, + bend=1 gives pure tangent warp. + """ + n = steps + if n < 1: + return torch.zeros(1, dtype=torch.float32) + + t_lin = torch.linspace(0.0, 1.0, n, dtype=torch.float64) + + # Tangent warp of the time ramp + pivot = pivot_frac + slope = 1.2 / max(n / 40.0, 0.1) + raw = ((2.0 / math.pi) + * torch.atan(-slope * (t_lin * (n - 1) - pivot * (n - 1))) + + 1.0) / 2.0 + r_max = raw[0].item() + r_min = raw[-1].item() + r_range = r_max - r_min + if r_range < 1e-12: + t_tan = 1.0 - t_lin + else: + t_tan = (raw - r_min) / r_range # [1, 0] normalized + + # Blend linear descent [1->0] with tangent warp + t_blend = (1.0 - bend) * (1.0 - t_lin) + bend * t_tan + + # Karras formula: sigma = (sig_min^(1/rho) + t*(sig_max^(1/rho)-sig_min^(1/rho)))^rho + inv_rho = 1.0 / rho + lo = sigma_min ** inv_rho + hi = sigma_max ** inv_rho + sigmas = (lo + t_blend * (hi - lo)) ** rho + + sigmas = torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]) + return sigmas.float() + + +def _logistic_sigmas( + steps: int, + sigma_max: float, + sigma_min: float, + steepness: float = 8.0, + midpoint: float = 0.4, +) -> torch.Tensor: + """ + Logistic (sigmoid) S-curve schedule. + + Exponential tails (vs algebraic for atan) give a sharper transition + through the detail range with flatter extremes. + """ + n = steps + if n < 1: + return torch.zeros(1, dtype=torch.float32) + + t = torch.linspace(0.0, 1.0, n, dtype=torch.float64) + + # Sigmoid: 1 / (1 + exp(k*(t - m))) + raw = 1.0 / (1.0 + torch.exp(steepness * (t - midpoint))) + + r_max = raw[0].item() + r_min = raw[-1].item() + r_range = r_max - r_min + + if r_range < 1e-12: + normalized = torch.linspace(1.0, 0.0, n, dtype=torch.float64) + else: + normalized = (raw - r_min) / r_range + + sigmas = normalized * (sigma_max - sigma_min) + sigma_min + sigmas = torch.cat([sigmas, torch.zeros(1, dtype=torch.float64)]) + return sigmas.float() + + +def scheduler_karras_tan(model_sampling: Any, steps: int) -> torch.Tensor: + """Karras-Tangent hybrid schedule (experimental).""" + sigma_max = float(model_sampling.sigma_max) + sigma_min = float(model_sampling.sigma_min) + sigmas = _karras_tangent_sigmas(steps, sigma_max, sigma_min) + _plot_sigmas(sigmas, 'karras_tan') + return sigmas + + +def scheduler_logistic(model_sampling: Any, steps: int) -> torch.Tensor: + """Logistic sigmoid S-curve schedule (experimental).""" + sigma_max = float(model_sampling.sigma_max) + sigma_min = float(model_sampling.sigma_min) + sigmas = _logistic_sigmas(steps, sigma_max, sigma_min) + _plot_sigmas(sigmas, 'logistic') + return sigmas + + +# ===================================================================== +# Registration +# ===================================================================== + +_SAMPLERS: Dict[str, Any] = {} +_SAMPLERS.update(_HFE_PRESETS) # hfe_s1..s8, hfe3_s1..s8, hfe4_s1..s8, hfe5_s1..s8 +_SAMPLERS["hfe_auto"] = sample_hfe_auto # 2-stage adaptive +_SAMPLERS["hfe3_auto"] = sample_hfe3_auto # 3-stage adaptive +_SAMPLERS["hfe4_auto"] = sample_hfe4_auto # 4-stage adaptive +_SAMPLERS["hfe5_auto"] = sample_hfe5_auto # 5-stage adaptive +# Experimental base samplers +_SAMPLERS["hfx_lap"] = sample_hfx_lap +_SAMPLERS["hfx_mom"] = sample_hfx_mom +_SAMPLERS["hfx_fft"] = sample_hfx_fft +_SAMPLERS["hfx_sde"] = sample_hfx_sde +_SAMPLERS["hfx_spatial"] = sample_hfx_spatial +# Hybrid combinators +_SAMPLERS["hfx_lap_mom"] = sample_hfx_lap_mom +_SAMPLERS["hfx_lap_spatial"] = sample_hfx_lap_spatial +_SAMPLERS["hfx_fft_spatial"] = sample_hfx_fft_spatial +# Band profile variants +_SAMPLERS["hfx_lap_fine"] = sample_hfx_lap_fine +_SAMPLERS["hfx_lap_broad"] = sample_hfx_lap_broad +# Graduated experimental presets (hfx_*_s1..s4 for each mode) +_SAMPLERS.update(_HFX_PRESETS) + +_SCHEDULERS = { + "atan_gentle": scheduler_atan_gentle, + "atan_focused": scheduler_atan_focused, + "atan_steep": scheduler_atan_steep, + "karras_tan": scheduler_karras_tan, + "logistic": scheduler_logistic, +} + +# Old names from all previous versions +_OLD_NAMES = [ + "euler_hfdetail", "hfdetail_power", + "hfdetail_soft", "hfdetail", "hfdetail_strong", + "res_2s_soft", "res_2s_sharp", "res_2s_crisp", + "tangent_soft", "tangent_sharp", "tangent_crisp", + "hfe_soft", "hfe_sharp", "hfe_crisp", +] + + +def _unregister_old() -> None: + """Remove entries from previous versions.""" + for attr in ("KSAMPLER_NAMES", "SAMPLER_NAMES", "SCHEDULER_NAMES"): + names = getattr(comfy_samplers, attr, None) + if isinstance(names, (list, tuple)): + names = list(names) + changed = False + for old in _OLD_NAMES: + if old in names: + names.remove(old) + changed = True + if changed: + setattr(comfy_samplers, attr, names) + + KSampler = getattr(comfy_samplers, "KSampler", None) + if KSampler is not None and hasattr(KSampler, "SAMPLERS"): + samplers = list(getattr(KSampler, "SAMPLERS")) + changed = False + for old in _OLD_NAMES: + if old in samplers: + samplers.remove(old) + changed = True + if changed: + KSampler.SAMPLERS = samplers + + kdiff = getattr(comfy_samplers, "k_diffusion_sampling", None) + if kdiff is not None: + for old in _OLD_NAMES: + attr = f"sample_{old}" + if hasattr(kdiff, attr): + delattr(kdiff, attr) + + handlers = getattr(comfy_samplers, "SCHEDULER_HANDLERS", None) + if isinstance(handlers, dict): + for old in _OLD_NAMES: + handlers.pop(old, None) + + +def _register_samplers() -> None: + kdiff = getattr(comfy_samplers, "k_diffusion_sampling", None) + + for name, func in _SAMPLERS.items(): + ksampler_names = getattr(comfy_samplers, "KSAMPLER_NAMES", None) + if isinstance(ksampler_names, (list, tuple)): + kl = list(ksampler_names) + if name not in kl: + kl.append(name) + comfy_samplers.KSAMPLER_NAMES = kl + + sampler_names = getattr(comfy_samplers, "SAMPLER_NAMES", []) + if not isinstance(sampler_names, list): + sampler_names = list(sampler_names) + if name not in sampler_names: + sampler_names.append(name) + comfy_samplers.SAMPLER_NAMES = sampler_names + + KSampler = getattr(comfy_samplers, "KSampler", None) + if KSampler is not None and hasattr(KSampler, "SAMPLERS"): + sl = list(getattr(KSampler, "SAMPLERS")) + if name not in sl: + sl.append(name) + KSampler.SAMPLERS = sl + + if kdiff is not None: + attr = f"sample_{name}" + setattr(kdiff, attr, func) + + LOGGER.info("HFE samplers registered: %s", list(_SAMPLERS.keys())) + + +def _register_schedulers() -> None: + handlers = getattr(comfy_samplers, "SCHEDULER_HANDLERS", None) + + for name, func in _SCHEDULERS.items(): + if isinstance(handlers, dict) and len(handlers) > 0: + any_handler = next(iter(handlers.values())) + HandlerType = type(any_handler) + handlers[name] = HandlerType(handler=func, use_ms=True) + + names = getattr(comfy_samplers, "SCHEDULER_NAMES", []) + if not isinstance(names, list): + names = list(names) + if name not in names: + names.append(name) + comfy_samplers.SCHEDULER_NAMES = names + + LOGGER.info("HFE schedulers registered: %s", list(_SCHEDULERS.keys())) + + +# ===================================================================== +# Initialization +# ===================================================================== + +def initialize_hfdetail_extension() -> None: + try: + _unregister_old() + except Exception: + LOGGER.debug("Old HFDetail entries cleanup skipped.", exc_info=True) + + try: + _register_samplers() + except Exception: + LOGGER.error("Failed to register HFE samplers.", exc_info=True) + + try: + _register_schedulers() + except Exception: + LOGGER.error("Failed to register HFE schedulers.", exc_info=True) + + +initialize_hfdetail_extension() + +NODE_CLASS_MAPPINGS: Dict[str, Any] = {} +NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {}