Upload 12 files

mega_freeu_a1111/README.md

@@ -0,0 +1,30 @@
+# ⚡ Mega FreeU — A1111 Extension
+
+**Mega FreeU** combines the best features from 5 FreeU implementations into a single
+production-ready A1111 extension.
+
+## Sources & Features
+
+| Source | Features taken |
+|--------|---------------|
+| **sd-webui-freeu** | `th.cat` hijack (only correct A1111 approach), V1/V2 backbone, box filter (BUG FIXED), schedule start/stop/smoothness, backbone region (width+offset), presets JSON, PNG metadata, XYZ grid, ControlNet patch |
+| **WAS FreeU_Advanced** | 9 blending modes (lerp/inject/bislerp/colorize/cosine/cuberp/hslerp/stable_slerp/linear_dodge), 13 multi-scale FFT presets, override_scales textarea, Post-CFG Shift (ported to A1111). Note: `target_block` / `input_block` / `middle_block` / `slice_b1/b2` not ported — A1111's `th.cat` hijack operates on output-side skip connections only. |
+| **ComfyUI_FreeU_V2_Advanced** | Gaussian FFT filter (smooth, no ringing), Adaptive Cap loop (MAX_CAP_ITER=3), independent B/S timestep ranges per stage, channel_threshold matching |
+| **FreeU_V2_timestepadd** | b_start/b_end%, s_start/s_end% per-stage step-fraction gating (note: original ComfyUI used `percent_to_sigma`; this port uses `current_step / total_steps` — conceptually equivalent for typical schedulers) |
+| **nrs_kohaku_v3.5** | hf_boost parameter, gaussian on output, on_cpu fallback tracker |
+
+## Bug Fixed
+- `sdwebui-freeU-extension` had Fourier mask applied to ONE quadrant:  
+  `mask[..., crow-t:crow, ccol-t:ccol]`  
+  **Fixed**: `mask[..., crow-t:crow+t, ccol-t:ccol+t]` (symmetric center)
+
+## Installation
+Copy `mega_freeu_a1111/` into `stable-diffusion-webui/extensions/` and restart.
+
+## Recommended Settings (SD1.5 V2 Gaussian + Independent B/S)
+- Stage 1: B=1.2, S=0.9, FFT=gaussian, B End%=0.35, S Start%=0.35
+- Stage 2: B=1.4, S=0.2, FFT=gaussian, B End%=0.35, S Start%=0.35
+
+## Key Concept: Independent B/S Timestep Ranges
+- **B-scaling** is most effective in structure phase (early steps, B Start%=0.0, B End%=0.35)
+- **S-filtering** is most effective in detail phase (late steps, S Start%=0.35, S End%=1.0)

mega_freeu_a1111/install.py

@@ -0,0 +1,3 @@
+# install.py — auto-run by A1111 extension system (no external deps needed)
+import launch
+# All deps (torch, gradio, modules) are already in A1111 environment

mega_freeu_a1111/lib_mega_freeu/__init__.py

@@ -0,0 +1 @@
+# lib_mega_freeu — Mega FreeU for A1111

mega_freeu_a1111/lib_mega_freeu/global_state.py

@@ -0,0 +1,285 @@
+"""
+lib_mega_freeu/global_state.py
+Runtime state, data structures, presets for ⚡ Mega FreeU.
+
+Sources:
+  sd-webui-freeu/lib_free_u/global_state.py  -- StageInfo layout, State, preset JSON, XYZ
+  WAS FreeU_Advanced/nodes.py                -- BLEND_MODE_NAMES, MSCALES
+  ComfyUI_FreeU_V2_Advanced/FreeU_B1B2.py   -- b_start/b_end, channel_threshold
+  ComfyUI_FreeU_V2_Advanced/FreeU_S1S2.py   -- s_start/s_end, adaptive cap
+  nrs_kohaku_enhanced_v3_5.py                -- hf_boost, gaussian standalone
+"""
+import dataclasses
+import json
+import math
+import pathlib
+import re
+import sys
+from typing import Any, Dict, List, Optional, Union
+
+# ─── Blending modes (WAS FreeU_Advanced/nodes.py blending_modes keys) ─────────
+BLEND_MODE_NAMES: List[str] = [
+    "lerp", "inject", "bislerp", "colorize",
+    "cosine interp", "cuberp", "hslerp", "stable_slerp", "linear dodge",
+]
+
+# ─── Multi-scale presets (WAS nodes.py mscales dict -- exact) ─────────────────
+MSCALES: Dict[str, Optional[list]] = {
+    "Default":                None,
+    "Low-Pass":               [(10, 1.0)],
+    "Pass-Through":           [(10, 1.0)],
+    "Gaussian-Blur":          [(10, 0.5)],
+    "Edge-Enhancement":       [(10, 2.0)],
+    "Sharpen":                [(10, 1.5)],
+    "Multi-Bandpass":         [[(5, 0.0),  (15, 1.0), (25, 0.0)]],
+    "Multi-Low-Pass":         [[(5, 1.0),  (10, 0.5), (15, 0.2)]],
+    "Multi-High-Pass":        [[(5, 0.0),  (10, 0.5), (15, 0.8)]],
+    "Multi-Pass-Through":     [[(5, 1.0),  (10, 1.0), (15, 1.0)]],
+    "Multi-Gaussian-Blur":    [[(5, 0.5),  (10, 0.8), (15, 0.2)]],
+    "Multi-Edge-Enhancement": [[(5, 1.2),  (10, 1.5), (15, 2.0)]],
+    "Multi-Sharpen":          [[(5, 1.5),  (10, 2.0), (15, 2.5)]],
+}
+
+ALL_VERSIONS: Dict[str, str] = {"Version 1": "1", "Version 2": "2"}
+REVERSED_VERSIONS: Dict[str, str] = {v: k for k, v in ALL_VERSIONS.items()}
+FFT_TYPES: List[str] = ["gaussian", "box"]
+STAGES_COUNT: int = 3
+
+_shorthand_re = re.compile(r"^([a-z]{1,3})(\d+)$")
+
+# ─── StageInfo ─────────────────────────────────────────────────────────────────
+@dataclasses.dataclass
+class StageInfo:
+    """
+    All per-stage parameters.
+    Fields 1-6: same order as sd-webui-freeu for PNG backwards compat.
+    New fields appended at end.
+    """
+    # sd-webui-freeu compat (DO NOT REORDER first 6)
+    backbone_factor:      float = 1.0
+    skip_factor:          float = 1.0
+    backbone_offset:      float = 0.0
+    backbone_width:       float = 0.5
+    skip_cutoff:          float = 0.0
+    skip_high_end_factor: float = 1.0
+    # WAS blending
+    backbone_blend_mode:  str   = "lerp"
+    backbone_blend:       float = 1.0
+    # ComfyUI V2 independent timestep ranges
+    b_start_ratio:  float = 0.0
+    b_end_ratio:    float = 1.0
+    s_start_ratio:  float = 0.0
+    s_end_ratio:    float = 1.0
+    # FFT
+    fft_type:         str   = "box"
+    fft_radius_ratio: float = 0.07
+    hf_boost:         float = 1.0
+    # Adaptive Cap (FreeU_S1S2)
+    enable_adaptive_cap: bool  = False
+    cap_threshold:       float = 0.35
+    cap_factor:          float = 0.6
+    adaptive_cap_mode:   str   = "adaptive"
+
+    def to_dict(self, include_default=False):
+        default = StageInfo()
+        d = dataclasses.asdict(self)
+        if not include_default:
+            d = {k: v for k, v in d.items() if v != getattr(default, k)}
+        return d
+
+    def copy(self):
+        return StageInfo(**dataclasses.asdict(self))
+
+STAGE_FIELD_NAMES = [f.name for f in dataclasses.fields(StageInfo)]
+STAGE_FIELD_COUNT = len(STAGE_FIELD_NAMES)
+
+# ─── State ─────────────────────────────────────────────────────────────────────
+@dataclasses.dataclass
+class State:
+    enable:                bool  = True
+    start_ratio:           Any   = 0.0
+    stop_ratio:            Any   = 1.0
+    transition_smoothness: float = 0.0
+    version:               str   = "1"
+    multiscale_mode:       str   = "Default"
+    multiscale_strength:   float = 1.0
+    override_scales:       str   = ""
+    channel_threshold:     int   = 96
+    stage_infos: List[Any] = dataclasses.field(
+        default_factory=lambda: [StageInfo() for _ in range(STAGES_COUNT)]
+    )
+    # Post-CFG Shift (WAS_PostCFGShift) — stored in presets & PNG
+    pcfg_enabled:   bool  = False
+    pcfg_steps:     int   = 20
+    pcfg_mode:      str   = "inject"
+    pcfg_blend:     float = 1.0
+    pcfg_b:         float = 1.1
+    pcfg_fourier:   bool  = False
+    pcfg_ms_mode:   str   = "Default"
+    pcfg_ms_str:    float = 1.0
+    pcfg_threshold: int   = 1
+    pcfg_s:         float = 0.5
+    pcfg_gain:      float = 1.0
+    verbose:        bool  = False
+
+    def __post_init__(self):
+        self.stage_infos = self._coerce_stages()
+        self.version = ALL_VERSIONS.get(self.version, self.version)
+
+    def _coerce_stages(self):
+        result, raw = [], list(self.stage_infos)
+        i = 0
+        while i < len(raw) and len(result) < STAGES_COUNT:
+            item = raw[i]
+            if isinstance(item, StageInfo):
+                result.append(item); i += 1
+            elif isinstance(item, dict):
+                known = {k: v for k, v in item.items() if k in STAGE_FIELD_NAMES}
+                result.append(StageInfo(**known)); i += 1
+            else:
+                chunk = raw[i:i+STAGE_FIELD_COUNT]
+                result.append(StageInfo(*chunk))
+                i += STAGE_FIELD_COUNT
+        while len(result) < STAGES_COUNT:
+            result.append(StageInfo())
+        return result
+
+    def to_dict(self):
+        d = dataclasses.asdict(self)
+        d["stage_infos"] = [si.to_dict() for si in self.stage_infos]
+        del d["enable"]
+        return d
+
+    def copy(self):
+        d = dataclasses.asdict(self)
+        d["stage_infos"] = [StageInfo(**s) for s in d["stage_infos"]]
+        return State(**d)
+
+    def update_attr(self, key, value):
+        if m := _shorthand_re.match(key):
+            char, idx = m.group(1), int(m.group(2))
+            if 0 <= idx < STAGES_COUNT:
+                si = self.stage_infos[idx]
+                _MAP = {
+                    "b":"backbone_factor","s":"skip_factor","o":"backbone_offset",
+                    "w":"backbone_width","t":"skip_cutoff","h":"skip_high_end_factor",
+                    "bm":"backbone_blend_mode","bb":"backbone_blend",
+                    "bs":"b_start_ratio","be":"b_end_ratio",
+                    "ss":"s_start_ratio","se":"s_end_ratio",
+                    "ft":"fft_type","fr":"fft_radius_ratio","hfb":"hf_boost",
+                    "cap":"enable_adaptive_cap","ct":"cap_threshold","cf":"cap_factor","acm":"adaptive_cap_mode",
+                }
+                if char in _MAP:
+                    setattr(si, _MAP[char], value); return
+        if hasattr(self, key):
+            setattr(self, key, value)
+
+# ─── Singletons ────────────────────────────────────────────────────────────────
+instance: State = State()
+xyz_attrs: Dict[str, Any] = {}
+current_sampling_step: int = 0
+
+# ─── Preset builders ───────────────────────────────────────────────────────────
+def _v1(*pairs):
+    infos = [StageInfo(backbone_factor=b, skip_factor=s) for b,s in pairs]
+    while len(infos) < STAGES_COUNT: infos.append(StageInfo())
+    return State(version="1", stage_infos=infos)
+
+def _v2g(pairs):
+    infos = []
+    for b, s, r, hfb, bs, be, ss, se in pairs:
+        infos.append(StageInfo(
+            backbone_factor=b, skip_factor=s,
+            fft_type="gaussian", fft_radius_ratio=r, hf_boost=hfb,
+            b_start_ratio=bs, b_end_ratio=be,
+            s_start_ratio=ss, s_end_ratio=se,
+        ))
+    while len(infos) < STAGES_COUNT: infos.append(StageInfo())
+    return State(version="2", stage_infos=infos)
+
+default_presets: Dict[str, State] = {
+    "SD1.4 Recommendations": _v1((1.2,0.9),(1.4,0.2),(1.0,1.0)),
+    "SD2.1 Recommendations": _v1((1.1,0.9),(1.2,0.2),(1.0,1.0)),
+    "SDXL Recommendations":  _v1((1.1,0.6),(1.2,0.4),(1.0,1.0)),
+    "SD1.5 V2 Gaussian": _v2g([
+        (1.2,0.9,0.07,1.0, 0.0,0.35, 0.35,1.0),
+        (1.4,0.2,0.07,1.0, 0.0,0.35, 0.35,1.0),
+    ]),
+    "SD1.5 V2 High Detail": _v2g([
+        (1.4,0.8,0.08,1.2, 0.0,0.35, 0.35,1.0),
+        (1.6,0.1,0.06,1.0, 0.0,0.35, 0.35,1.0),
+    ]),
+    "SDXL V2 Gaussian": _v2g([
+        (1.1,0.6,0.05,1.1, 0.0,0.35, 0.35,1.0),
+        (1.2,0.4,0.05,1.1, 0.0,0.35, 0.35,1.0),
+    ]),
+    "SD1.5 Adaptive Cap": State(version="2", stage_infos=[
+        StageInfo(backbone_factor=1.3,skip_factor=0.9,
+                  fft_type="gaussian",fft_radius_ratio=0.08,hf_boost=1.2,
+                  b_start_ratio=0.0,b_end_ratio=0.35,
+                  s_start_ratio=0.35,s_end_ratio=1.0,
+                  enable_adaptive_cap=True,cap_threshold=0.35,
+                  cap_factor=0.6,adaptive_cap_mode="adaptive"),
+        StageInfo(backbone_factor=1.4,skip_factor=0.2,
+                  fft_type="gaussian",fft_radius_ratio=0.06,hf_boost=1.0,
+                  b_start_ratio=0.0,b_end_ratio=0.35,
+                  s_start_ratio=0.35,s_end_ratio=1.0,
+                  enable_adaptive_cap=True,cap_threshold=0.70,
+                  cap_factor=0.6,adaptive_cap_mode="adaptive"),
+        StageInfo(),
+    ]),
+    "Independent B/S (SD1.5)": _v2g([
+        (1.2,0.9,0.07,1.0, 0.0,0.35, 0.35,1.0),
+        (1.4,0.2,0.06,1.0, 0.0,0.35, 0.35,1.0),
+    ]),
+}
+
+all_presets: Dict[str, State] = {}
+PRESETS_PATH = pathlib.Path(__file__).parent.parent / "presets.json"
+
+def reload_presets():
+    all_presets.clear()
+    all_presets.update(default_presets)
+    all_presets.update(_load_user_presets())
+
+def _load_user_presets():
+    if not PRESETS_PATH.exists(): return {}
+    try:
+        with open(PRESETS_PATH, encoding="utf-8") as f:
+            raw = json.load(f)
+    except Exception as e:
+        print(f"[MegaFreeU] preset load error: {e}", file=sys.stderr)
+        return {}
+    result = {}
+    _state_fields = {f.name for f in dataclasses.fields(State)}
+    for k, v in raw.items():
+        try:
+            # Filter unknown keys so future/old fields don't crash State(**v)
+            known = {fk: fv for fk, fv in v.items() if fk in _state_fields}
+            result[k] = State(**known)
+        except Exception as e:
+            print(f"[MegaFreeU] skipping preset {k!r}: {e}", file=sys.stderr)
+    return result
+
+def save_presets(custom=None):
+    if custom is None: custom = get_user_presets()
+    try:
+        PRESETS_PATH.parent.mkdir(parents=True, exist_ok=True)
+        with open(PRESETS_PATH, "w", encoding="utf-8") as f:
+            json.dump({k: v.to_dict() for k,v in custom.items()}, f, indent=4)
+    except Exception as e:
+        print(f"[MegaFreeU] preset save error: {e}", file=sys.stderr)
+
+def get_user_presets():
+    return {k: v for k,v in all_presets.items() if k not in default_presets}
+
+def apply_xyz():
+    global instance
+    if pk := xyz_attrs.get("preset"):
+        if p := all_presets.get(pk):
+            instance = p.copy()
+        elif pk != "UI Settings":
+            print(f"[MegaFreeU] XYZ preset '{pk}' not found", file=sys.stderr)
+    for k, v in xyz_attrs.items():
+        if k != "preset":
+            instance.update_attr(k, v)

mega_freeu_a1111/lib_mega_freeu/unet.py

@@ -0,0 +1,559 @@
+"""
+lib_mega_freeu/unet.py — Math engine + A1111 th.cat patch
+
+BUGS FIXED vs sdwebui-freeU-extension/scripts/freeunet_hijack.py:
+  BUG 1 dtype: mask = torch.ones(..., dtype=torch.bool)
+    bool*float = NOOP, scale always 1.0
+    Fix: torch.full(..., float(scale_high))
+  BUG 2 quadrant: mask[..., crow-t:crow, ccol-t:ccol] (top-left only)
+    Fix: mask[..., crow-t:crow+t, ccol-t:ccol+t] (symmetric center)
+
+Sources:
+  sd-webui-freeu/lib_free_u/unet.py         patch(), free_u_cat_hijack(),
+    get_backbone_scale(), ratio_to_region(), filter_skip()[box],
+    get_schedule_ratio(), is_gpu_complex_supported(), lerp()
+  WAS FreeU_Advanced/nodes.py               9 blending modes, Fourier_filter() multiscale
+  ComfyUI_FreeU_V2_advanced/utils.py        Fourier_filter_gauss(), get_band_energy_stats()
+  ComfyUI_FreeU_V2_advanced/FreeU_S1S2.py  Adaptive Cap loop MAX_CAP_ITER=3
+  ComfyUI_FreeU_V2_advanced/FreeU_B1B2.py  channel_threshold, model_channels*4/2/1
+  FreeU_V2_timestepadd.py                  step-fraction timestep gating concept
+  nrs_kohaku_enhanced_v3_5.py              _freeu_b_scale_h, _freeu_fourier_filter_gaussian,
+                                            hf_boost param, on_cpu_devices dict
+"""
+import dataclasses
+import functools
+import logging
+import math
+import pathlib
+import sys
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from lib_mega_freeu import global_state
+
+# ── GPU complex support (sd-webui-freeu exact) ────────────────────────────────
+_gpu_complex_support: Optional[bool] = None
+
+def is_gpu_complex_supported(x: torch.Tensor) -> bool:
+    global _gpu_complex_support
+    if x.is_cpu:
+        return True
+    if _gpu_complex_support is not None:
+        return _gpu_complex_support
+    mps_avail = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+    try:
+        import torch_directml
+    except ImportError:
+        dml_avail = False
+    else:
+        dml_avail = torch_directml.is_available()
+    _gpu_complex_support = not (mps_avail or dml_avail)
+    if _gpu_complex_support:
+        try: torch.fft.fftn(x.float(), dim=(-2, -1))
+        except RuntimeError: _gpu_complex_support = False
+    return _gpu_complex_support
+
+_on_cpu_devices: Dict = {}
+
+# ── Blending modes (WAS nodes.py exact) ───────────────────────────────────────
+def _normalize(t):
+    mn, mx = t.min(), t.max()
+    return (t - mn) / (mx - mn + 1e-8)
+
+def _hslerp(a, b, t):
+    nc = a.size(1)
+    iv = torch.zeros(1, nc, 1, 1, device=a.device, dtype=a.dtype)
+    iv[0, 0, 0, 0] = 1.0
+    result = (1 - t) * a + t * b
+    if t < 0.5:
+        result += (torch.norm(b - a, dim=1, keepdim=True) / 6) * iv
+    else:
+        result -= (torch.norm(b - a, dim=1, keepdim=True) / 6) * iv
+    return result
+
+def _stable_slerp(a, b, t, eps=1e-6):
+    an = a / torch.linalg.norm(a, dim=1, keepdim=True).clamp_min(eps)
+    bn = b / torch.linalg.norm(b, dim=1, keepdim=True).clamp_min(eps)
+    dot = (an * bn).sum(dim=1, keepdim=True).clamp(-1.0 + eps, 1.0 - eps)
+    theta = torch.acos(dot)
+    sin_t = torch.sin(theta).clamp_min(eps)
+    s0 = torch.sin((1.0 - t) * theta) / sin_t
+    s1 = torch.sin(t * theta) / sin_t
+    slerp_out = s0 * a + s1 * b
+    lerp_out  = (1.0 - t) * a + t * b
+    use_lerp  = (theta < 1e-3).squeeze(1)
+    return torch.where(use_lerp.unsqueeze(1), lerp_out, slerp_out)
+
+BLENDING_MODES = {
+    "bislerp":       lambda a, b, t: _normalize((1 - t) * a + t * b),
+    "colorize":      lambda a, b, t: a + (b - a) * t,
+    "cosine interp": lambda a, b, t: (
+        a + b - (a - b) * torch.cos(
+            torch.tensor(math.pi, device=a.device, dtype=a.dtype) * t)) / 2,
+    "cuberp":        lambda a, b, t: a + (b - a) * (3 * t**2 - 2 * t**3),
+    "hslerp":        _hslerp,
+    "stable_slerp":  _stable_slerp,
+    "inject":        lambda a, b, t: a + b * t,
+    "lerp":          lambda a, b, t: (1 - t) * a + t * b,
+    "linear dodge":  lambda a, b, t: _normalize(a + b * t),
+}
+
+def lerp(a, b, r):
+    return (1 - r) * a + r * b
+
+# ── Backbone scaling ──────────────────────────────────────────────────────────
+def get_backbone_scale(h: torch.Tensor, backbone_factor: float, version: str):
+    if version == "1":
+        return backbone_factor
+    # V2: adaptive hidden_mean (FreeU_B1B2.py + kohaku _freeu_b_scale_h exact)
+    features_mean = h.mean(1, keepdim=True)
+    B = features_mean.shape[0]
+    hidden_max, _ = torch.max(features_mean.view(B, -1), dim=-1, keepdim=True)
+    hidden_min, _ = torch.min(features_mean.view(B, -1), dim=-1, keepdim=True)
+    denom = (hidden_max - hidden_min).clamp_min(1e-6)
+    hidden_mean = (features_mean - hidden_min.unsqueeze(2).unsqueeze(3)) \
+                  / denom.unsqueeze(2).unsqueeze(3)
+    return 1 + (backbone_factor - 1) * hidden_mean
+
+def ratio_to_region(width: float, offset: float, n: int) -> Tuple[int, int, bool]:
+    """sd-webui-freeu ratio_to_region exact."""
+    if width < 0:
+        offset += width; width = -width
+    width = min(width, 1.0)
+    if offset < 0:
+        offset = 1 + offset - int(offset)
+    offset = math.fmod(offset, 1.0)
+    if width + offset <= 1:
+        return round(offset * n), round((width + offset) * n), False
+    else:
+        return round((width + offset - 1) * n), round(offset * n), True
+
+# ── Box FFT (BUGS FIXED symmetric center + float dtype) ─────────────────────
+def filter_skip_box(x: torch.Tensor, cutoff: float,
+                    scale: float, scale_high: float = 1.0) -> torch.Tensor:
+    """
+    FreeU box filter with TWO BUGS FIXED from sdwebui-freeU-extension:
+    BUG 1 (dtype): was torch.bool mask -> scale multiplication was NOOP
+    BUG 2 (region): was [crow-t:crow, ccol-t:ccol] -> single quadrant top-left
+    Both fixed: torch.full float + symmetric [crow-t:crow+t, ccol-t:ccol+t].
+    sd-webui-freeu has these correct already, we match their implementation.
+    """
+    if scale == 1.0 and scale_high == 1.0:
+        return x
+    fft_dev = x.device if is_gpu_complex_supported(x) else torch.device("cpu")
+    x_freq = torch.fft.fftn(x.to(fft_dev, dtype=torch.float32), dim=(-2, -1))
+    x_freq = torch.fft.fftshift(x_freq, dim=(-2, -1))
+    B, C, H, W = x_freq.shape
+    mask = torch.full((B, C, H, W), float(scale_high), device=fft_dev)  # FIX: float, not bool
+    crow, ccol = H // 2, W // 2
+    tr = max(1, math.floor(crow * cutoff)) if cutoff > 0 else 1
+    tc = max(1, math.floor(ccol * cutoff)) if cutoff > 0 else 1
+    mask[..., crow - tr:crow + tr, ccol - tc:ccol + tc] = scale  # FIX: symmetric center
+    x_freq *= mask
+    x_freq = torch.fft.ifftshift(x_freq, dim=(-2, -1))
+    return torch.fft.ifftn(x_freq, dim=(-2, -1)).real.to(device=x.device, dtype=x.dtype)
+
+# ── Box + WAS multiscale overlay (WAS nodes.py Fourier_filter exact) ─────────
+def filter_skip_box_multiscale(x: torch.Tensor, cutoff: float, scale: float,
+                                scales_preset: Optional[list],
+                                strength: float = 1.0,
+                                scale_high: float = 1.0) -> torch.Tensor:
+    """
+    WAS FreeU_Advanced/nodes.py Fourier_filter(x, threshold, scale, scales, strength).
+    threshold = cutoff: float ratio [0-1] or int pixels (WAS uses int default=1).
+    scales: None, list of (radius_px, val) single-scale, or list of lists multi-scale.
+    """
+    if scale == 1.0 and scale_high == 1.0 and scales_preset is None:
+        return x
+    fft_dev = x.device if is_gpu_complex_supported(x) else torch.device("cpu")
+    x_freq = torch.fft.fftn(x.to(fft_dev, dtype=torch.float32), dim=(-2, -1))
+    x_freq = torch.fft.fftshift(x_freq, dim=(-2, -1))
+    B, C, H, W = x_freq.shape
+    crow, ccol = H // 2, W // 2
+    if isinstance(cutoff, float) and 0 < cutoff <= 1.0:
+        tr = max(1, math.floor(crow * cutoff)); tc = max(1, math.floor(ccol * cutoff))
+    else:
+        t = max(1, int(cutoff)) if cutoff > 0 else 1; tr = tc = t
+    mask = torch.ones((B, C, H, W), device=fft_dev)
+    mask[..., crow - tr:crow + tr, ccol - tc:ccol + tc] = scale
+    if scale_high != 1.0:
+        hfm = torch.full((B, C, H, W), float(scale_high), device=fft_dev)
+        hfm[..., crow - tr:crow + tr, ccol - tc:ccol + tc] = 1.0
+        mask = mask * hfm
+    if scales_preset:
+        if isinstance(scales_preset[0], tuple):
+            # WAS single-scale mode
+            for scale_threshold, scale_value in scales_preset:
+                sv = scale_value * strength
+                sm = torch.ones((B, C, H, W), device=fft_dev)
+                st = max(1, int(scale_threshold))
+                sm[..., crow - st:crow + st, ccol - st:ccol + st] = sv
+                mask = mask + (sm - mask) * strength
+        else:
+            # WAS multi-scale mode
+            for scale_params in scales_preset:
+                if isinstance(scale_params, list):
+                    for scale_threshold, scale_value in scale_params:
+                        sv = scale_value * strength
+                        sm = torch.ones((B, C, H, W), device=fft_dev)
+                        st = max(1, int(scale_threshold))
+                        sm[..., crow - st:crow + st, ccol - st:ccol + st] = sv
+                        mask = mask + (sm - mask) * strength
+    x_freq = x_freq * mask
+    x_freq = torch.fft.ifftshift(x_freq, dim=(-2, -1))
+    return torch.fft.ifftn(x_freq, dim=(-2, -1)).real.to(device=x.device, dtype=x.dtype)
+
+# ── Gaussian FFT (ComfyUI utils.py exact) ────────────────────────────────────
+def fourier_filter_gauss(x: torch.Tensor, radius_ratio: float,
+                          scale: float, hf_boost: float = 1.0) -> torch.Tensor:
+    """
+    ComfyUI_FreeU_V2_advanced/utils.py Fourier_filter_gauss() exact.
+    Also matches kohaku _freeu_fourier_filter_gaussian().
+    R = max(1, int(min(H,W)*radius_ratio))
+    sigma_f = R^2/2
+    center = exp(-dist2/sigma_f)
+    mask = scale*center + hf_boost*(1-center)
+    """
+    x_f = torch.fft.fftn(x.float(), dim=(-2, -1))
+    x_f = torch.fft.fftshift(x_f, dim=(-2, -1))
+    B, C, H, W = x_f.shape
+    R = max(1, int(min(H, W) * radius_ratio))
+    sigma_f = max(1e-6, (R * R) / 2.0)
+    yy, xx = torch.meshgrid(
+        torch.arange(H, device=x.device, dtype=torch.float32) - H // 2,
+        torch.arange(W, device=x.device, dtype=torch.float32) - W // 2,
+        indexing="ij")
+    center = torch.exp(-(yy**2 + xx**2) / sigma_f)
+    mask = (scale * center + hf_boost * (1.0 - center)).view(1, 1, H, W)
+    x_f = x_f * mask
+    x_f = torch.fft.ifftshift(x_f, dim=(-2, -1))
+    return torch.fft.ifftn(x_f, dim=(-2, -1)).real.to(x.dtype)
+
+# ── Band energy stats (ComfyUI utils.py exact) ────────────────────────────────
+def get_band_energy_stats(x: torch.Tensor, R: int) -> Tuple[float, float, float]:
+    """ComfyUI_FreeU_V2_advanced/utils.py get_band_energy_stats() exact."""
+    xf = torch.fft.fftn(x.float(), dim=(-2, -1))
+    xf = torch.fft.fftshift(xf, dim=(-2, -1))
+    B, C, H, W = xf.shape
+    yy, xx = torch.meshgrid(
+        torch.arange(H, device=x.device, dtype=torch.float32) - H // 2,
+        torch.arange(W, device=x.device, dtype=torch.float32) - W // 2,
+        indexing="ij")
+    lf_mask = (yy**2 + xx**2) <= (R * R)
+    mag2 = xf.real**2 + xf.imag**2
+    # FIX: expand_as requires same ndim; use 2D mask on last dims
+    lf_e  = mag2[:, :, lf_mask].mean().item()   if lf_mask.any()   else 0.0
+    hf_e  = mag2[:, :, ~lf_mask].mean().item()  if (~lf_mask).any() else 0.0
+    cover = lf_mask.sum().item() / (H * W) * 100.0
+    return lf_e, hf_e, cover
+
+# ── Adaptive Cap Gaussian (FreeU_S1S2.py MAX_CAP_ITER=3 exact) ───────────────
+def filter_skip_gaussian_adaptive(hsp: torch.Tensor,
+                                   si: "global_state.StageInfo",
+                                   verbose: bool = False) -> torch.Tensor:
+    """
+    ComfyUI_FreeU_V2_advanced/FreeU_S1S2.py exact algorithm:
+    1. Compute LF/HF ratio before.
+    2. Apply Gaussian filter.
+    3. If enable_adaptive_cap and drop > cap_threshold: loop up to MAX_CAP_ITER=3.
+       adaptive mode: eff_factor = cap_factor * (cap_threshold / drop)
+       fixed mode:    eff_factor = cap_factor
+       capped_s = 1 - eff_factor*(1-s_scale)  [interpolate FROM ORIGINAL]
+       capped_s = max(capped_s, current_s*(1+1e-4))
+       Re-apply from original_hsp with capped_s.
+    hf_boost combined = max(si.hf_boost, si.skip_high_end_factor)  [kohaku pattern]
+    """
+    s_scale  = si.skip_factor
+    radius_r = si.fft_radius_ratio
+    hf_boost = max(si.hf_boost, si.skip_high_end_factor)
+    orig_dev = hsp.device
+    H, W     = hsp.shape[-2:]
+    R_eff    = max(1, int(min(H, W) * radius_r))
+
+    # CRITICAL ORDER: init cpu-fallback flag and helpers BEFORE any FFT call
+    use_cpu = _on_cpu_devices.get(orig_dev, not is_gpu_complex_supported(hsp))
+    if use_cpu:
+        _on_cpu_devices[orig_dev] = True
+
+    def _tod(t):            # to FFT-safe device
+        return t.cpu() if use_cpu else t
+
+    def _fromd(t):          # back to original device
+        return t.to(orig_dev) if use_cpu else t
+
+    def _energy(t):
+        return get_band_energy_stats(_tod(t), R_eff)
+
+    def _filt(inp, sc):
+        nonlocal use_cpu
+        try:
+            out = fourier_filter_gauss(_tod(inp), radius_r, sc, hf_boost)
+            return _fromd(out)
+        except Exception:
+            if not use_cpu:
+                logging.warning(f"[MegaFreeU] {orig_dev} -> CPU fallback for FFT")
+                _on_cpu_devices[orig_dev] = True
+                use_cpu = True
+                return fourier_filter_gauss(inp.cpu(), radius_r, sc, hf_boost).to(orig_dev)
+            return inp
+
+    # Pre-filter energy (now safe on all devices)
+    lf_b, hf_b, cover = _energy(hsp)
+    ratio_b = lf_b / hf_b if hf_b > 1e-6 else float("inf")
+    if verbose:
+        logging.info(f"[MegaFreeU] Gauss {H}x{W} R={R_eff}px cov={cover:.1f}% "
+                     f"LF={lf_b:.3e} HF={hf_b:.3e} ratio_b={ratio_b:.4f}")
+
+    hsp_filt = _filt(hsp, s_scale)
+    if not si.enable_adaptive_cap:
+        return hsp_filt
+
+    MAX_CAP_ITER = 3
+    original_hsp = hsp
+    current_s    = s_scale
+    lf_a, hf_a, _ = _energy(hsp_filt)
+    ratio_a = lf_a / hf_a if hf_a > 1e-6 else float("inf")
+    drop    = 1.0 - (ratio_a / ratio_b) if ratio_b > 1e-6 else 0.0
+    orig_drop = drop
+    iters = 0
+    hsp_cur = hsp_filt
+
+    while (si.enable_adaptive_cap
+           and drop > si.cap_threshold
+           and current_s < 0.999
+           and iters < MAX_CAP_ITER):
+
+        if iters == 0:
+            logging.warning(f"[MegaFreeU] Over-attenuation: drop={drop*100:.1f}% > "
+                            f"{si.cap_threshold*100:.1f}%  s={s_scale:.4f}")
+
+        eff_f = si.cap_factor
+        if si.adaptive_cap_mode == "adaptive":
+            eff_f = si.cap_factor * (si.cap_threshold / max(drop, 1e-8))
+
+        capped_s = 1.0 - eff_f * (1.0 - s_scale)           # interpolate from ORIGINAL s
+        capped_s = max(capped_s, current_s * (1.0 + 1e-4))  # only ever relax
+        if abs(capped_s - current_s) < 1e-4:
+            if verbose: logging.info("  Cap converged.")
+            break
+
+        if verbose:
+            logging.info(f"  Cap iter {iters+1}: s {current_s:.4f}->{capped_s:.4f} eff={eff_f:.4f}")
+
+        try:
+            hsp_new = _filt(original_hsp, capped_s)
+        except Exception as e:
+            logging.error(f"[MegaFreeU] cap re-apply error: {e}")
+            hsp_cur = original_hsp   # restore to original on error (ComfyUI FreeU_S1S2.py pattern)
+            break
+
+        hsp_cur   = hsp_new
+        lf_a, hf_a, _ = _energy(hsp_cur)
+        ratio_a = lf_a / hf_a if hf_a > 1e-6 else float("inf")
+        drop    = 1.0 - (ratio_a / ratio_b) if ratio_b > 1e-6 else 0.0
+        current_s = capped_s
+        iters += 1
+
+    if iters > 0 or verbose:
+        logging.info(f"[MegaFreeU] Cap done: {orig_drop*100:.1f}%->{drop*100:.1f}% "
+                     f"({iters} iters  s_final={current_s:.4f})")
+    return hsp_cur
+
+# ── Schedule (sd-webui-freeu exact) ──────────────────────────────────────────
+def get_schedule_ratio() -> float:
+    from modules import shared
+    st    = global_state.instance
+    steps = shared.state.sampling_steps or 20
+    cur   = global_state.current_sampling_step
+    start = _to_step(st.start_ratio, steps)
+    stop  = _to_step(st.stop_ratio, steps)
+    if start == stop:
+        smooth = 0.0
+    elif cur < start:
+        smooth = min(1.0, max(0.0, cur / (start + 1e-8)))
+    else:
+        smooth = min(1.0, max(0.0, 1 + (cur - start) / (start - stop + 1e-8)))
+    flat = 1.0 if start <= cur < stop else 0.0
+    return lerp(flat, smooth, st.transition_smoothness)
+
+def get_stage_bsratio(b_start: float, b_end: float) -> float:
+    """Independent B/S timestep range gate (FreeU_V2_timestepadd concept -> step fraction)."""
+    from modules import shared
+    steps = max(shared.state.sampling_steps or 20, 1)
+    cur   = global_state.current_sampling_step
+    pct   = cur / (steps - 1) if steps > 1 else 0.0
+    return 1.0 if b_start <= pct <= b_end else 0.0
+
+def _to_step(v, steps):
+    return int(v * steps) if isinstance(v, float) else int(v)
+
+# ── Stage auto-detection (FreeU_B1B2.py + kohaku exact) ──────────────────────
+_stage_channels: Tuple[int, int, int] = (1280, 640, 320)
+
+def detect_model_channels():
+    global _stage_channels
+    try:
+        from modules import shared
+        mc = int(shared.sd_model.model.diffusion_model.model_channels)
+        _stage_channels = (mc * 4, mc * 2, mc * 1)
+    except Exception:
+        _stage_channels = (1280, 640, 320)
+
+def get_stage_index(dims: int, channel_threshold: int = 96) -> Optional[int]:
+    """FreeU_B1B2.py abs(ch - target) <= channel_threshold proximity match."""
+    for i, target in enumerate(_stage_channels):
+        if abs(dims - target) <= channel_threshold:
+            return i
+    return None
+
+# ── Override scales parser (WAS nodes.py format exact) ───────────────────────
+def parse_override_scales(text: str) -> Optional[List]:
+    if not text or not text.strip():
+        return None
+    result = []
+    for line in text.strip().splitlines():
+        line = line.strip()
+        if not line or line.startswith(("#", "!", "//")):
+            continue
+        parts = line.split(",")
+        if len(parts) == 2:
+            try:
+                result.append((int(parts[0].strip()), float(parts[1].strip())))
+            except ValueError:
+                pass
+    return result if result else None
+
+class _VerboseRef:
+    value: bool = False
+verbose_ref = _VerboseRef()
+
+# ── Core th.cat hijack (sd-webui-freeu exact + extended) ─────────────────────
+def free_u_cat_hijack(hs, *args, original_function, **kwargs):
+    """
+    Intercepts torch.cat([h, h_skip], dim=1) in UNet output_blocks.
+    Signature: kwargs=={"dim":1} and len(hs)==2 (sd-webui-freeu exact check).
+
+    Why th.cat over alternatives:
+    - sdwebui-freeU-extension CondFunc(UNetModel.forward): rewrites full forward,
+      incompatible with other extensions, plus 2 bugs in fourier mask.
+    - kohaku register_forward_hook: output already concatenated,
+      can't cleanly separate h from h_skip for independent filtering.
+    - th.cat hijack: intercepts exactly [h, h_skip] before concatenation. CORRECT.
+    """
+    st = global_state.instance
+    if not st.enable:
+        return original_function(hs, *args, **kwargs)
+
+    sched = get_schedule_ratio()
+    if sched == 0:
+        return original_function(hs, *args, **kwargs)
+
+    try:
+        h, h_skip = hs
+        if list(kwargs.keys()) != ["dim"] or kwargs.get("dim", -1) != 1:
+            return original_function(hs, *args, **kwargs)
+    except (ValueError, TypeError):
+        return original_function(hs, *args, **kwargs)
+
+    dims      = int(h.shape[1])
+    stage_idx = get_stage_index(dims, st.channel_threshold)
+    if stage_idx is None:
+        return original_function(hs, *args, **kwargs)
+
+    si      = st.stage_infos[stage_idx]
+    version = st.version
+    verbose = verbose_ref.value
+
+    # ── BACKBONE ─────────────────────────────────────────────────────────────
+    b_gate = get_stage_bsratio(si.b_start_ratio, si.b_end_ratio)
+    eff_b  = sched * b_gate
+
+    if eff_b > 0.0 and abs(si.backbone_factor - 1.0) > 1e-6:
+        try:
+            rbegin, rend, rinv = ratio_to_region(si.backbone_width, si.backbone_offset, dims)
+            ch_idx = torch.arange(dims, device=h.device)
+            mask   = (rbegin <= ch_idx) & (ch_idx <= rend)
+            if rinv: mask = ~mask
+            mask = mask.reshape(1, -1, 1, 1).to(h.dtype)
+
+            eff_factor = float(lerp(1.0, si.backbone_factor, eff_b))
+            scale      = get_backbone_scale(h, eff_factor, version)
+            # h_scaled_full: full h with mask region scaled, rest unchanged
+            # This matches original: h *= mask*scale + (1-mask)
+            h_scaled_full = h * (mask * scale + (1.0 - mask))
+
+            bmode = si.backbone_blend_mode
+            if bmode in BLENDING_MODES and abs(si.backbone_blend - 1.0) > 1e-6:
+                # Blend on FULL tensors so modes like slerp/hslerp see proper norms.
+                # Then restore unmasked channels to original h.
+                h_blended = BLENDING_MODES[bmode](h, h_scaled_full, si.backbone_blend)
+                h = h * (1.0 - mask) + h_blended * mask
+            else:
+                h = h_scaled_full
+        except Exception as e:
+            logging.warning(f"[MegaFreeU] B-scaling stage {stage_idx}: {e}")
+
+    # ── SKIP / FOURIER ────────────────────────────────────────────────────────
+    s_gate = get_stage_bsratio(si.s_start_ratio, si.s_end_ratio)
+    eff_s  = sched * s_gate
+
+    if eff_s > 0.0 and (abs(si.skip_factor - 1.0) > 1e-6
+                         or abs(si.hf_boost - 1.0) > 1e-6
+                         or abs(si.skip_high_end_factor - 1.0) > 1e-6):
+        try:
+            s_scale = float(lerp(1.0, si.skip_factor, eff_s))
+            s_high  = float(lerp(1.0, si.skip_high_end_factor, eff_s))
+
+            if si.fft_type == "gaussian":
+                hf_b = float(lerp(1.0, si.hf_boost, eff_s))
+                si_eff = dataclasses.replace(si, skip_factor=s_scale, skip_high_end_factor=s_high, hf_boost=hf_b)
+                h_skip = filter_skip_gaussian_adaptive(h_skip, si_eff, verbose)
+            else:
+                override  = parse_override_scales(st.override_scales)
+                ms_preset = override or global_state.MSCALES.get(st.multiscale_mode)
+                if ms_preset is not None:
+                    h_skip = filter_skip_box_multiscale(
+                        h_skip, si.skip_cutoff, s_scale, ms_preset,
+                        st.multiscale_strength, s_high)
+                else:
+                    h_skip = filter_skip_box(h_skip, si.skip_cutoff, s_scale, s_high)
+        except Exception as e:
+            logging.warning(f"[MegaFreeU] skip filter stage {stage_idx}: {e}")
+
+    return original_function([h, h_skip], *args, **kwargs)
+
+# ── Patch (sd-webui-freeu exact + ControlNet) ─────────────────────────────────
+_patched = False  # guard against double-patch on hot-reload
+
+def patch():
+    global _patched
+    try:
+        from modules.sd_hijack_unet import th
+    except ImportError:
+        print("[MegaFreeU] sd_hijack_unet not available", file=sys.stderr); return
+
+    if _patched or (hasattr(th.cat, "func") and getattr(th.cat.func, "__name__", "") == "free_u_cat_hijack"):
+        return  # already patched (by name; handles module reload)
+    th.cat = functools.partial(free_u_cat_hijack, original_function=th.cat)
+    _patched = True
+
+    cn_status = "enabled"
+    try:
+        from modules import scripts
+        cn_paths = [
+            str(pathlib.Path(scripts.basedir()).parent.parent / "extensions-builtin" / "sd-webui-controlnet"),
+            str(pathlib.Path(scripts.basedir()).parent / "sd-webui-controlnet"),
+        ]
+        sys.path[0:0] = cn_paths
+        try:
+            import scripts.hook as cn_hook
+            cn_hook.th.cat = functools.partial(free_u_cat_hijack, original_function=cn_hook.th.cat)
+        except ImportError:
+            cn_status = "disabled"
+        finally:
+            for p in cn_paths:
+                if p in sys.path: sys.path.remove(p)
+    except Exception:
+        cn_status = "error"
+
+    print(f"[MegaFreeU] th.cat patched  ControlNet: *{cn_status}*")

mega_freeu_a1111/lib_mega_freeu/xyz_grid.py

@@ -0,0 +1,97 @@
+"""
+lib_mega_freeu/xyz_grid.py — XYZ/XY grid axes
+
+Source: sd-webui-freeu/lib_free_u/xyz_grid.py (exact find_xyz_module check)
+Extended with all Mega FreeU per-stage params.
+"""
+import sys
+from types import ModuleType
+from typing import Optional
+from modules import scripts
+from lib_mega_freeu import global_state
+
+
+def patch():
+    xyz_module = _find_xyz_module()
+    if xyz_module is None:
+        print("[MegaFreeU] xyz_grid.py not found — XYZ disabled", file=sys.stderr)
+        return
+
+    def _apply(k, key_map=None):
+        def cb(_p, v, _vs):
+            if key_map is not None:
+                v = key_map.get(v, v)
+            global_state.xyz_attrs[k] = v
+        return cb
+
+    opts = [
+        xyz_module.AxisOption("[MegaFreeU] Enabled",        _bool,    _apply("enable"),
+                              choices=lambda: ["False","True"]),
+        xyz_module.AxisOption("[MegaFreeU] Version",         str,      _apply("version",
+                              key_map=global_state.ALL_VERSIONS),
+                              choices=lambda: list(global_state.ALL_VERSIONS.keys())),
+        xyz_module.AxisOption("[MegaFreeU] Preset",          str,      _apply("preset"),
+                              choices=_choices_preset),
+        xyz_module.AxisOption("[MegaFreeU] Start At Step",   _num,     _apply("start_ratio")),
+        xyz_module.AxisOption("[MegaFreeU] Stop At Step",    _num,     _apply("stop_ratio")),
+        xyz_module.AxisOption("[MegaFreeU] Smoothness",      float,    _apply("transition_smoothness")),
+        xyz_module.AxisOption("[MegaFreeU] Multiscale Mode", str,      _apply("multiscale_mode"),
+                              choices=lambda: list(global_state.MSCALES.keys())),
+        xyz_module.AxisOption("[MegaFreeU] Multiscale Str",  float,    _apply("multiscale_strength")),
+        xyz_module.AxisOption("[MegaFreeU] Ch Threshold",    int,      _apply("channel_threshold")),
+    ]
+
+    for i in range(global_state.STAGES_COUNT):
+        n = i + 1
+        opts += [
+            xyz_module.AxisOption(f"[MFU] S{n} B Scale",     float, _apply(f"b{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} B Offset",    float, _apply(f"o{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} B Width",     float, _apply(f"w{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} Blend Mode",  str,   _apply(f"bm{i}"),
+                                  choices=lambda: global_state.BLEND_MODE_NAMES),
+            xyz_module.AxisOption(f"[MFU] S{n} Blend Str",   float, _apply(f"bb{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} S Scale",     float, _apply(f"s{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} Cutoff",      float, _apply(f"t{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} High-End",    float, _apply(f"h{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} HF Boost",    float, _apply(f"hfb{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} B Start%",    float, _apply(f"bs{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} B End%",      float, _apply(f"be{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} S Start%",    float, _apply(f"ss{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} S End%",      float, _apply(f"se{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} FFT Type",    str,   _apply(f"ft{i}"),
+                                  choices=lambda: global_state.FFT_TYPES),
+            xyz_module.AxisOption(f"[MFU] S{n} Radius",      float, _apply(f"fr{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} Cap Enable",  _bool, _apply(f"cap{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} Cap Thresh",  float, _apply(f"ct{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} Cap Factor",  float, _apply(f"cf{i}")),
+            xyz_module.AxisOption(f"[MFU] S{n} Cap Mode",    str,   _apply(f"acm{i}"),
+                                  choices=lambda: ["adaptive", "fixed"]),
+        ]
+
+    xyz_module.axis_options.extend(opts)
+
+
+def _find_xyz_module() -> Optional[ModuleType]:
+    """Exact check from sd-webui-freeu/lib_free_u/xyz_grid.py"""
+    for data in scripts.scripts_data:
+        if data.script_class.__module__ in {"xyz_grid.py", "xy_grid.py"} and hasattr(data, "module"):
+            return data.module
+    return None
+
+
+def _choices_preset():
+    presets = list(global_state.all_presets.keys())
+    presets.insert(0, "UI Settings")
+    return presets
+
+
+def _bool(s):
+    s = str(s).lower()
+    if s in ("true","1","yes"): return True
+    if s in ("false","0","no"): return False
+    return bool(s)
+
+
+def _num(s):
+    try: return int(s)
+    except (ValueError, TypeError): return float(s)

mega_freeu_a1111/scripts/mega_freeu.py

@@ -0,0 +1,765 @@
+"""
+scripts/mega_freeu.py - Mega FreeU for A1111 / Forge
+
+Combined from 5 sources:
+  1. sd-webui-freeu         th.cat hijack, V1/V2 backbone, box filter, schedule,
+                             presets JSON, PNG metadata, XYZ, ControlNet, region masking,
+                             dict-API compat (alwayson_scripts legacy)
+  2. WAS FreeU_Advanced     9 blending modes, 13 multi-scale FFT presets, override_scales,
+                             Post-CFG Shift (WAS_PostCFGShift ported to A1111 callback)
+                             NOTE: target_block / input_block / middle_block / slice_b1/b2
+                             were not ported — th.cat hijack works on output-side skip concat.
+  3. ComfyUI_FreeU_V2_Adv   Gaussian filter, Adaptive Cap (MAX_CAP_ITER=3),
+                             independent B/S timestep ranges per-stage, channel_threshold
+  4. FreeU_V2_timestepadd   b_start/b_end%, s_start/s_end% per-stage gating
+                             NOTE: gating uses step-fraction (cur/total), not percent_to_sigma
+                             as in original ComfyUI sources. Conceptually equivalent.
+  5. nrs_kohaku_v3.5        hf_boost param, on_cpu_devices dict, gaussian standalone
+
+BUGS FIXED vs sdwebui-freeU-extension:
+  BUG 1: bool mask in Fourier filter (scale multiplication was NOOP)
+  BUG 2: single-quadrant mask instead of symmetric center
+"""
+import dataclasses
+import json
+from typing import List
+
+import gradio as gr
+from modules import script_callbacks, scripts, shared, processing
+
+from lib_mega_freeu import global_state, unet, xyz_grid
+
+_steps_comps = {"txt2img": None, "img2img": None}
+_steps_cbs   = {"txt2img": [], "img2img": []}
+
+_SF = [f.name for f in dataclasses.fields(global_state.StageInfo)]
+_SN = len(_SF)  # 19 fields per stage
+
+
+def _stage_ui(idx, si, elem_id_fn):
+    n  = idx + 1
+    ch = {0: "~1280ch (deep)", 1: "~640ch (mid)", 2: "~320ch (shallow)"}.get(idx, f"stage{n}")
+
+    with gr.Accordion(open=(idx < 2), label=f"Stage {n}  ({ch})"):
+
+        # Backbone
+        gr.HTML(f"<p style=\'margin:4px 0;font-size:.82em;color:#aaa;\'>Backbone h (B)</p>")
+        with gr.Row():
+            bf = gr.Slider(label=f"B{n} Scale",  minimum=-1, maximum=3, step=0.001,
+                           value=si.backbone_factor,
+                           info=">1 strengthens backbone features. V2: adaptive per-region.")
+            bo = gr.Slider(label=f"B{n} Offset", minimum=0,  maximum=1, step=0.001,
+                           value=si.backbone_offset, info="Channel region start [0-1].")
+            bw = gr.Slider(label=f"B{n} Width",  minimum=-1, maximum=1, step=0.001,
+                           value=si.backbone_width, info="Channel region width. Negative=invert.")
+        with gr.Row():
+            bm = gr.Dropdown(label=f"B{n} Blend Mode",
+                             choices=global_state.BLEND_MODE_NAMES,
+                             value=si.backbone_blend_mode,
+                             info="lerp=default, stable_slerp=quality, inject=additive")
+            bb = gr.Slider(label=f"B{n} Blend Str", minimum=0, maximum=2, step=0.001,
+                           value=si.backbone_blend)
+        gr.HTML("<p style=\'font-size:.75em;color:#888;margin:2px 0;\'>B timestep range (ComfyUI V2)</p>")
+        with gr.Row():
+            bsr = gr.Slider(label=f"B{n} Start%", minimum=0, maximum=1, step=0.001,
+                            value=si.b_start_ratio, info="B activates at this step fraction.")
+            ber = gr.Slider(label=f"B{n} End%",   minimum=0, maximum=1, step=0.001,
+                            value=si.b_end_ratio,   info="B stops. 0.35=structure phase only.")
+
+        # Skip / FFT
+        gr.HTML(f"<p style=\'margin:8px 0 4px;font-size:.82em;color:#aaa;\'>Skip h_skip (S) - Fourier Filter</p>")
+        with gr.Row():
+            sf  = gr.Slider(label=f"S{n} LF Scale",     minimum=-1, maximum=3, step=0.001,
+                            value=si.skip_factor,
+                            info="<1 suppresses LF components. 0.2=strong suppression.")
+            she = gr.Slider(label=f"S{n} HF (Box)",     minimum=-1, maximum=3, step=0.001,
+                            value=si.skip_high_end_factor,
+                            info="HF scale outside LF region (box filter). >1=boost HF.")
+            hfb = gr.Slider(label=f"S{n} HF Boost (Gauss)", minimum=0, maximum=3, step=0.001,
+                            value=si.hf_boost,
+                            info="Gaussian explicit HF multiplier. Combined as max(hf_boost, high_end).")
+        with gr.Row():
+            ft  = gr.Radio(label=f"S{n} FFT Type",
+                           choices=global_state.FFT_TYPES, value=si.fft_type,
+                           info="gaussian=smooth no-ringing. box=original FreeU (both bugs fixed).")
+            sco = gr.Slider(label=f"S{n} Cutoff (Box)",   minimum=0, maximum=1, step=0.001,
+                            value=si.skip_cutoff, info="Box: LF cutoff fraction. 0=1px default.")
+            srr = gr.Slider(label=f"S{n} Radius (Gauss)", minimum=0.01, maximum=0.5, step=0.001,
+                            value=si.fft_radius_ratio,
+                            info="Gaussian R=ratio*min(H,W). 0.07=moderate LF.")
+        gr.HTML("<p style=\'font-size:.75em;color:#888;margin:2px 0;\'>S timestep range (ComfyUI V2)</p>")
+        with gr.Row():
+            ssr = gr.Slider(label=f"S{n} Start%", minimum=0, maximum=1, step=0.001,
+                            value=si.s_start_ratio,
+                            info="S activates. Tip: set = B End% for clean phase separation.")
+            ser = gr.Slider(label=f"S{n} End%",   minimum=0, maximum=1, step=0.001,
+                            value=si.s_end_ratio,   info="S stops. 1.0=to last step.")
+
+        # Adaptive Cap
+        gr.HTML("<p style=\'font-size:.75em;color:#888;margin:4px 0;\'>Adaptive Cap - prevents LF over-attenuation (FreeU_S1S2.py)</p>")
+        with gr.Row():
+            eac = gr.Checkbox(label=f"S{n} Enable Cap", value=si.enable_adaptive_cap,
+                              info="Iteratively weakens Gaussian if LF/HF drop exceeds threshold.")
+            ct  = gr.Slider(label="Threshold", minimum=0, maximum=1, step=0.001,
+                            value=si.cap_threshold, info="Max allowed LF/HF ratio drop. 0.35=35%.")
+            cf  = gr.Slider(label="Factor",    minimum=0, maximum=1, step=0.001,
+                            value=si.cap_factor, info="Relaxation factor. 0.6=moderate.")
+            cm  = gr.Radio(label="Mode", choices=["adaptive", "fixed"],
+                           value=si.adaptive_cap_mode,
+                           info="adaptive: scales factor with over-attenuation. fixed: always cap_factor.")
+
+    # Return exactly in _SF field order
+    return [bf, sf, bo, bw, sco, she, bm, bb, bsr, ber, ssr, ser, ft, srr, hfb, eac, ct, cf, cm]
+
+
+class MegaFreeUScript(scripts.Script):
+
+    def title(self): return "Mega FreeU"
+    def show(self, is_img2img): return scripts.AlwaysVisible
+
+    def ui(self, is_img2img):
+        global_state.reload_presets()
+        pnames  = list(global_state.all_presets.keys())
+        def_sis = global_state.all_presets[pnames[0]].stage_infos
+
+        with gr.Accordion(open=False, label="Mega FreeU"):
+
+            # Top bar
+            with gr.Row():
+                enabled = gr.Checkbox(label="Enable Mega FreeU", value=False)
+                version = gr.Dropdown(
+                    label="Version",
+                    choices=list(global_state.ALL_VERSIONS.keys()),
+                    value="Version 2",
+                    elem_id=self.elem_id("version"),
+                    info="V2=adaptive hidden-mean backbone. V1=flat multiplier.")
+
+            with gr.Row():
+                preset_dd = gr.Dropdown(
+                    label="Preset", choices=pnames, value=pnames[0],
+                    allow_custom_value=True,
+                    elem_id=self.elem_id("preset_name"),
+                    info="Apply loads settings. Custom name enables Save. Delete auto-saves.")
+                btn_apply   = gr.Button("Apply",   size="sm", elem_classes="tool")
+                btn_save    = gr.Button("Save",    size="sm", elem_classes="tool")
+                btn_refresh = gr.Button("Refresh", size="sm", elem_classes="tool")
+                btn_delete  = gr.Button("Delete",  size="sm", elem_classes="tool")
+
+            # Global schedule
+            gr.HTML("<p style=\'font-size:.82em;color:#aaa;margin:6px 0 2px;\'>Global Schedule</p>")
+            with gr.Row():
+                start_r = gr.Slider(label="Start At", elem_id=self.elem_id("start_at_step"),
+                                    minimum=0, maximum=1, step=0.001, value=0)
+                stop_r  = gr.Slider(label="Stop At",  elem_id=self.elem_id("stop_at_step"),
+                                    minimum=0, maximum=1, step=0.001, value=1)
+                smooth  = gr.Slider(label="Transition Smoothness",
+                                    elem_id=self.elem_id("transition_smoothness"),
+                                    minimum=0, maximum=1, step=0.001, value=0,
+                                    info="0=hard on/off. 1=smooth fade.")
+
+            # Box Multi-Scale (WAS FreeU_Advanced)
+            with gr.Accordion(open=False, label="Box Multi-Scale FFT  (WAS FreeU_Advanced)"):
+                gr.HTML("<p style=\'font-size:.8em;color:#888;\'>Applied on top of Box filter. Ignored in Gaussian mode.</p>")
+                with gr.Row():
+                    ms_mode = gr.Dropdown(label="Multiscale Mode",
+                                          choices=list(global_state.MSCALES.keys()),
+                                          value="Default")
+                    ms_str  = gr.Slider(label="Strength", minimum=0, maximum=1,
+                                        step=0.001, value=1.0)
+                ov_scales = gr.Textbox(
+                    label="Override Scales  (WAS format: radius_px, scale  per line, # comments)",
+                    lines=3,
+                    placeholder="# Example custom scales:\n10, 1.5\n20, 0.8",
+                    value="")
+
+            with gr.Row():
+                ch_thresh = gr.Slider(
+                    label="Channel Match Threshold (+-)",
+                    elem_id=self.elem_id("ch_thresh"),
+                    minimum=0, maximum=256, step=1, value=96,
+                    info="Stage channel tolerance. 96=standard (FreeU_B1B2.py default).")
+
+            # Per-stage accordions
+            flat_comps: List = []
+            for i in range(global_state.STAGES_COUNT):
+                si = def_sis[i] if i < len(def_sis) else global_state.StageInfo()
+                flat_comps.extend(_stage_ui(i, si, self.elem_id))
+
+            # Post-CFG Shift (WAS_PostCFGShift -> A1111)
+            with gr.Accordion(open=False, label="Post-CFG Shift  (WAS_PostCFGShift -> A1111 callback)"):
+                gr.HTML("<p style=\'font-size:.8em;color:#888;\'>Runs after combine_denoised. Blends denoised*b into output via on_cfg_after_cfg callback.</p>")
+                with gr.Row():
+                    pcfg_en    = gr.Checkbox(label="Enable Post-CFG Shift", value=False)
+                    pcfg_steps = gr.Slider(label="Max Steps", minimum=1, maximum=200,
+                                           step=1, value=20,
+                                           info="Apply only to first N steps.")
+                with gr.Row():
+                    pcfg_mode = gr.Dropdown(label="Blend Mode",
+                                             choices=global_state.BLEND_MODE_NAMES,
+                                             value="inject")
+                    pcfg_bl   = gr.Slider(label="Blend", minimum=0, maximum=5,
+                                           step=0.001, value=1.0)
+                    pcfg_b    = gr.Slider(label="B Factor", minimum=0, maximum=5,
+                                           step=0.001, value=1.1,
+                                           info=">1 amplifies shift.")
+                with gr.Row():
+                    pcfg_fou  = gr.Checkbox(label="Apply Fourier Filter", value=False)
+                    pcfg_mmd  = gr.Dropdown(label="Fourier Multiscale",
+                                             choices=list(global_state.MSCALES.keys()),
+                                             value="Default")
+                    pcfg_mst  = gr.Slider(label="Fourier Strength", minimum=0, maximum=1,
+                                           step=0.001, value=1.0)
+                with gr.Row():
+                    pcfg_thr  = gr.Slider(label="Threshold (px)", minimum=1, maximum=20,
+                                           step=1, value=1,
+                                           info="Box filter LF radius in pixels.")
+                    pcfg_s    = gr.Slider(label="S Scale", minimum=0, maximum=3,
+                                           step=0.001, value=0.5)
+                    pcfg_gain = gr.Slider(label="Force Gain", minimum=0, maximum=5,
+                                           step=0.01, value=1.0,
+                                           info="Final output multiplier.")
+
+            verbose = gr.Checkbox(label="Verbose Logging (Adaptive Cap, energy stats)", value=False)
+
+            # Hidden PNG infotext components
+            sched_info        = gr.HTML(visible=False)
+            stages_info       = gr.HTML(visible=False)
+            version_info      = gr.HTML(visible=False)
+            ms_mode_info      = gr.HTML(visible=False)
+            ms_str_info       = gr.HTML(visible=False)
+            ov_scales_info    = gr.HTML(visible=False)
+            ch_thresh_info    = gr.HTML(visible=False)
+            postcfg_info      = gr.HTML(visible=False)
+            verbose_info      = gr.HTML(visible=False)
+            # Legacy sd-webui-freeu keys for backward compat
+            legacy_sched_info   = gr.HTML(visible=False)
+            legacy_stages_info  = gr.HTML(visible=False)
+            legacy_version_info = gr.HTML(visible=False)
+
+            # Preset buttons
+            def _btn_upd(name):
+                ex  = name in global_state.all_presets
+                usr = name not in global_state.default_presets
+                return (gr.update(interactive=ex),
+                        gr.update(interactive=usr),
+                        gr.update(interactive=usr and ex))
+
+            preset_dd.change(fn=_btn_upd, inputs=[preset_dd],
+                             outputs=[btn_apply, btn_save, btn_delete])
+
+            def _apply_p(name):
+                p = global_state.all_presets.get(name)
+                n_extras = 20  # 8 main + 11 Post-CFG + 1 verbose
+                if p is None:
+                    return [gr.skip()] * (n_extras + len(flat_comps))
+                flat = []
+                for si in p.stage_infos:
+                    for f in _SF:
+                        flat.append(getattr(si, f))
+                vlabel = global_state.REVERSED_VERSIONS.get(p.version, "Version 2")
+                return (
+                    gr.update(value=p.start_ratio),
+                    gr.update(value=p.stop_ratio),
+                    gr.update(value=p.transition_smoothness),
+                    gr.update(value=vlabel),
+                    gr.update(value=p.multiscale_mode),
+                    gr.update(value=p.multiscale_strength),
+                    gr.update(value=p.override_scales),
+                    gr.update(value=p.channel_threshold),
+                    gr.update(value=p.pcfg_enabled),
+                    gr.update(value=p.pcfg_steps),
+                    gr.update(value=p.pcfg_mode),
+                    gr.update(value=p.pcfg_blend),
+                    gr.update(value=p.pcfg_b),
+                    gr.update(value=p.pcfg_fourier),
+                    gr.update(value=p.pcfg_ms_mode),
+                    gr.update(value=p.pcfg_ms_str),
+                    gr.update(value=p.pcfg_threshold),
+                    gr.update(value=p.pcfg_s),
+                    gr.update(value=p.pcfg_gain),
+                    gr.update(value=p.verbose),
+                    *[gr.update(value=v) for v in flat],
+                )
+
+            btn_apply.click(
+                fn=_apply_p,
+                inputs=[preset_dd],
+                outputs=[
+                    start_r, stop_r, smooth, version,
+                    ms_mode, ms_str, ov_scales, ch_thresh,
+                    pcfg_en, pcfg_steps, pcfg_mode, pcfg_bl, pcfg_b,
+                    pcfg_fou, pcfg_mmd, pcfg_mst, pcfg_thr, pcfg_s, pcfg_gain,
+                    verbose,
+                    *flat_comps,
+                ]
+            )
+
+            def _save_p(
+                name, sr, sp, sm, ver, msm, mss, ovs, cht,
+                p_en, p_steps, p_mode, p_bl, p_b,
+                p_four, p_mmd, p_mst, p_thr, p_s, p_gain,
+                v_log,
+                *flat
+            ):
+                sis = _flat_to_sis(flat)
+                vc  = global_state.ALL_VERSIONS.get(ver, "1")
+                global_state.all_presets[name] = global_state.State(
+                    start_ratio=sr, stop_ratio=sp, transition_smoothness=sm,
+                    version=vc,
+                    multiscale_mode=msm,
+                    multiscale_strength=float(mss),
+                    override_scales=ovs or "",
+                    channel_threshold=int(cht),
+                    stage_infos=sis,
+                    pcfg_enabled=bool(p_en),
+                    pcfg_steps=int(p_steps),
+                    pcfg_mode=str(p_mode),
+                    pcfg_blend=float(p_bl),
+                    pcfg_b=float(p_b),
+                    pcfg_fourier=bool(p_four),
+                    pcfg_ms_mode=str(p_mmd),
+                    pcfg_ms_str=float(p_mst),
+                    pcfg_threshold=int(p_thr),
+                    pcfg_s=float(p_s),
+                    pcfg_gain=float(p_gain),
+                    verbose=bool(v_log),
+                )
+                global_state.save_presets()
+                return (
+                    gr.update(choices=list(global_state.all_presets.keys())),
+                    gr.update(interactive=True),
+                    gr.update(interactive=True),
+                )
+
+            btn_save.click(
+                fn=_save_p,
+                inputs=[
+                    preset_dd, start_r, stop_r, smooth, version,
+                    ms_mode, ms_str, ov_scales, ch_thresh,
+                    pcfg_en, pcfg_steps, pcfg_mode, pcfg_bl, pcfg_b,
+                    pcfg_fou, pcfg_mmd, pcfg_mst, pcfg_thr, pcfg_s, pcfg_gain,
+                    verbose,
+                    *flat_comps,
+                ],
+                outputs=[preset_dd, btn_apply, btn_delete]
+            )
+
+            def _refresh_p(name):
+                global_state.reload_presets()
+                ex  = name in global_state.all_presets
+                usr = name not in global_state.default_presets
+                ch  = list(global_state.all_presets.keys())
+                return (gr.update(choices=ch, value=name),
+                        gr.update(interactive=ex), gr.update(interactive=usr),
+                        gr.update(interactive=usr and ex))
+
+            btn_refresh.click(fn=_refresh_p, inputs=[preset_dd],
+                              outputs=[preset_dd, btn_apply, btn_save, btn_delete])
+
+            def _delete_p(name):
+                if name in global_state.all_presets and name not in global_state.default_presets:
+                    idx   = list(global_state.all_presets.keys()).index(name)
+                    del global_state.all_presets[name]
+                    global_state.save_presets()
+                    names = list(global_state.all_presets.keys())
+                    name  = names[min(idx, len(names) - 1)]
+                ex  = name in global_state.all_presets
+                usr = name not in global_state.default_presets
+                return (gr.update(choices=list(global_state.all_presets.keys()), value=name),
+                        gr.update(interactive=ex), gr.update(interactive=usr),
+                        gr.update(interactive=usr and ex))
+
+            btn_delete.click(fn=_delete_p, inputs=[preset_dd],
+                             outputs=[preset_dd, btn_apply, btn_save, btn_delete])
+
+            # PNG schedule restore
+            def _restore_sched(info, steps):
+                if not info: return [gr.skip()] * 4
+                try:
+                    parts = info.split(", ")
+                    sr, sp, sm = parts[0], parts[1], parts[2]
+                    total = max(int(float(steps)), 1)
+                    def _r(v):
+                        n = float(v.strip())
+                        return n / total if n > 1.0 else n
+                    return (gr.update(value=""), gr.update(value=_r(sr)),
+                            gr.update(value=_r(sp)), gr.update(value=float(sm)))
+                except Exception:
+                    return [gr.skip()] * 4
+
+            def _reg_sched_cb(steps_comp):
+                sched_info.change(fn=_restore_sched,
+                                  inputs=[sched_info, steps_comp],
+                                  outputs=[sched_info, start_r, stop_r, smooth])
+
+            mode_key = "img2img" if is_img2img else "txt2img"
+            if _steps_comps[mode_key] is None:
+                _steps_cbs[mode_key].append(_reg_sched_cb)
+            else:
+                _reg_sched_cb(_steps_comps[mode_key])
+
+            def _restore_stages(info):
+                n_out = 2 + len(flat_comps)
+                if not info: return [gr.skip()] * n_out
+                try:
+                    raw_list = json.loads(info)
+                    sis = []
+                    for d in raw_list:
+                        known = {k: v for k, v in d.items()
+                                 if k in global_state.STAGE_FIELD_NAMES}
+                        sis.append(global_state.StageInfo(**known))
+                    while len(sis) < global_state.STAGES_COUNT:
+                        sis.append(global_state.StageInfo())
+                except Exception:
+                    return [gr.skip()] * n_out
+                flat = []
+                for si in sis:
+                    for f in _SF:
+                        flat.append(getattr(si, f, getattr(global_state.StageInfo(), f)))
+                auto_en = shared.opts.data.get("mega_freeu_png_auto_enable", True)
+                return (gr.update(value=""), gr.update(value=auto_en),
+                        *[gr.update(value=v) for v in flat])
+
+            stages_info.change(fn=_restore_stages, inputs=[stages_info],
+                               outputs=[stages_info, enabled, *flat_comps])
+
+            def _restore_ver(info):
+                if not info: return [gr.skip()] * 2
+                lbl = global_state.REVERSED_VERSIONS.get(info.strip(), info.strip())
+                return gr.update(value=""), gr.update(value=lbl)
+
+            version_info.change(fn=_restore_ver, inputs=[version_info],
+                                outputs=[version_info, version])
+
+            # ── New extended PNG restore callbacks ─────────────────────────────
+            def _restore_ms_mode(info):
+                if not info: return gr.skip(), gr.skip()
+                return gr.update(value=""), gr.update(value=info.strip())
+
+            def _restore_ms_str(info):
+                if not info: return gr.skip(), gr.skip()
+                try: return gr.update(value=""), gr.update(value=float(info.strip()))
+                except Exception: return gr.skip(), gr.skip()
+
+            def _restore_ov_scales(info):
+                if info is None: return gr.skip(), gr.skip()
+                return gr.update(value=""), gr.update(value=info)
+
+            def _restore_ch_thresh(info):
+                if not info: return gr.skip(), gr.skip()
+                try: return gr.update(value=""), gr.update(value=int(float(info.strip())))
+                except Exception: return gr.skip(), gr.skip()
+
+            def _restore_verbose(info):
+                if not info: return gr.skip(), gr.skip()
+                return gr.update(value=""), gr.update(value=(info.strip().lower() == "true"))
+
+            def _restore_postcfg(info):
+                n = 12
+                if not info: return [gr.skip()] * n
+                try:
+                    d = json.loads(info)
+                    return (
+                        gr.update(value=""),
+                        gr.update(value=bool(d.get("enabled",  False))),
+                        gr.update(value=int(d.get("steps",     20))),
+                        gr.update(value=str(d.get("mode",      "inject"))),
+                        gr.update(value=float(d.get("blend",   1.0))),
+                        gr.update(value=float(d.get("b",       1.1))),
+                        gr.update(value=bool(d.get("fourier",  False))),
+                        gr.update(value=str(d.get("ms_mode",   "Default"))),
+                        gr.update(value=float(d.get("ms_str",  1.0))),
+                        gr.update(value=int(d.get("threshold", 1))),
+                        gr.update(value=float(d.get("s",       0.5))),
+                        gr.update(value=float(d.get("gain",    1.0))),
+                    )
+                except Exception:
+                    return [gr.skip()] * n
+
+            ms_mode_info.change(fn=_restore_ms_mode, inputs=[ms_mode_info],
+                                outputs=[ms_mode_info, ms_mode])
+            ms_str_info.change(fn=_restore_ms_str, inputs=[ms_str_info],
+                               outputs=[ms_str_info, ms_str])
+            ov_scales_info.change(fn=_restore_ov_scales, inputs=[ov_scales_info],
+                                  outputs=[ov_scales_info, ov_scales])
+            ch_thresh_info.change(fn=_restore_ch_thresh, inputs=[ch_thresh_info],
+                                  outputs=[ch_thresh_info, ch_thresh])
+            verbose_info.change(fn=_restore_verbose, inputs=[verbose_info],
+                                outputs=[verbose_info, verbose])
+            postcfg_info.change(fn=_restore_postcfg, inputs=[postcfg_info],
+                                outputs=[postcfg_info,
+                                         pcfg_en, pcfg_steps, pcfg_mode, pcfg_bl, pcfg_b,
+                                         pcfg_fou, pcfg_mmd, pcfg_mst, pcfg_thr, pcfg_s, pcfg_gain])
+
+            # Legacy sd-webui-freeu keys — reuse same restore logic
+            legacy_sched_info.change(fn=lambda info, steps: _restore_sched(info, steps),
+                                     inputs=[legacy_sched_info, _steps_comps.get(mode_key) or legacy_sched_info],
+                                     outputs=[legacy_sched_info, start_r, stop_r, smooth])
+            legacy_stages_info.change(fn=_restore_stages, inputs=[legacy_stages_info],
+                                      outputs=[legacy_stages_info, enabled, *flat_comps])
+            legacy_version_info.change(fn=_restore_ver, inputs=[legacy_version_info],
+                                       outputs=[legacy_version_info, version])
+
+            self.infotext_fields = [
+                (sched_info,         "MegaFreeU Schedule"),
+                (stages_info,        "MegaFreeU Stages"),
+                (version_info,       "MegaFreeU Version"),
+                (ms_mode_info,       "MegaFreeU Multiscale Mode"),
+                (ms_str_info,        "MegaFreeU Multiscale Strength"),
+                (ov_scales_info,     "MegaFreeU Override Scales"),
+                (ch_thresh_info,     "MegaFreeU Channel Threshold"),
+                (postcfg_info,       "MegaFreeU PostCFG"),
+                (verbose_info,       "MegaFreeU Verbose"),
+                # Backward compat with sd-webui-freeu generated PNGs
+                (legacy_sched_info,  "FreeU Schedule"),
+                (legacy_stages_info, "FreeU Stages"),
+                (legacy_version_info,"FreeU Version"),
+            ]
+            self.paste_field_names = [f for _, f in self.infotext_fields]
+
+        return [
+            enabled, version, preset_dd,
+            start_r, stop_r, smooth,
+            ms_mode, ms_str, ov_scales,
+            ch_thresh,
+            *flat_comps,
+            pcfg_en, pcfg_steps, pcfg_mode, pcfg_bl, pcfg_b,
+            pcfg_fou, pcfg_mmd, pcfg_mst, pcfg_thr, pcfg_s, pcfg_gain,
+            verbose,
+        ]
+
+    def process(self, p: processing.StableDiffusionProcessing, *args):
+        # ── Branch 1: old sd-webui-freeu API (dict passed as first arg) ───────
+        if args and isinstance(args[0], dict):
+            global_state.instance = global_state.State(**{
+                k: v for k, v in args[0].items()
+                if k in {f.name for f in dataclasses.fields(global_state.State)}
+            })
+            global_state.apply_xyz()
+            global_state.xyz_attrs.clear()
+            st = global_state.instance
+            unet.verbose_ref.value = bool(getattr(st, "verbose", False))
+            if getattr(st, "pcfg_enabled", False):
+                p._mega_pcfg = {
+                    "enabled":   True,
+                    "steps":     st.pcfg_steps,
+                    "mode":      st.pcfg_mode,
+                    "blend":     st.pcfg_blend,
+                    "b":         st.pcfg_b,
+                    "fourier":   st.pcfg_fourier,
+                    "ms_mode":   st.pcfg_ms_mode,
+                    "ms_str":    st.pcfg_ms_str,
+                    "threshold": st.pcfg_threshold,
+                    "s":         st.pcfg_s,
+                    "gain":      st.pcfg_gain,
+                    "step":      0,
+                }
+            else:
+                p._mega_pcfg = {"enabled": False}
+            if st.enable:
+                unet.detect_model_channels()
+                unet._on_cpu_devices.clear()
+                _write_generation_params(p, st)
+            return
+
+        # ── Branch 2: normal UI call ───────────────────────────────────────────
+        (enabled, version, preset_dd,
+         start_r, stop_r, smooth,
+         ms_mode, ms_str, ov_scales,
+         ch_thresh, *rest) = args
+
+        n_sv = _SN * global_state.STAGES_COUNT
+        flat_stage = rest[:n_sv]
+        post       = rest[n_sv:]          # 11 pcfg params + verbose
+
+        verbose = bool(post[11]) if len(post) > 11 else False
+        unet.verbose_ref.value = verbose
+
+        # Write UI values into instance BEFORE apply_xyz so XYZ can override any of them
+        inst = global_state.instance
+        inst.enable               = bool(enabled)
+        inst.start_ratio          = start_r
+        inst.stop_ratio           = stop_r
+        inst.transition_smoothness = smooth
+        inst.version              = global_state.ALL_VERSIONS.get(version, "1")
+        inst.multiscale_mode      = ms_mode
+        inst.multiscale_strength  = float(ms_str)
+        inst.override_scales      = ov_scales or ""
+        inst.channel_threshold    = int(ch_thresh)
+        inst.stage_infos          = _flat_to_sis(flat_stage)
+
+        # Sync Post-CFG into instance state so presets/PNG capture it
+        pcfg = post[:11]
+        if len(pcfg) >= 11:
+            inst.pcfg_enabled   = bool(pcfg[0])
+            inst.pcfg_steps     = int(pcfg[1])
+            inst.pcfg_mode      = str(pcfg[2])
+            inst.pcfg_blend     = float(pcfg[3])
+            inst.pcfg_b         = float(pcfg[4])
+            inst.pcfg_fourier   = bool(pcfg[5])
+            inst.pcfg_ms_mode   = str(pcfg[6])
+            inst.pcfg_ms_str    = float(pcfg[7])
+            inst.pcfg_threshold = int(pcfg[8])
+            inst.pcfg_s         = float(pcfg[9])
+            inst.pcfg_gain      = float(pcfg[10])
+        inst.verbose = verbose
+
+        # apply_xyz() may replace global_state.instance with a preset copy;
+        # take the fresh reference AFTER so PNG metadata / verbose use the final state.
+        global_state.apply_xyz()
+        global_state.xyz_attrs.clear()
+        st = global_state.instance   # ← fresh ref post-XYZ
+
+        # ── Post-CFG: set up ALWAYS (independent of main Enable) ──────────────
+        if st.pcfg_enabled:
+            p._mega_pcfg = {
+                "enabled":   True,
+                "steps":     st.pcfg_steps,
+                "mode":      st.pcfg_mode,
+                "blend":     st.pcfg_blend,
+                "b":         st.pcfg_b,
+                "fourier":   st.pcfg_fourier,
+                "ms_mode":   st.pcfg_ms_mode,
+                "ms_str":    st.pcfg_ms_str,
+                "threshold": st.pcfg_threshold,
+                "s":         st.pcfg_s,
+                "gain":      st.pcfg_gain,
+                "step":      0,
+            }
+        else:
+            p._mega_pcfg = {"enabled": False}
+
+        if not st.enable:
+            # Write partial params so PNG records the session even when disabled
+            _write_generation_params(p, st)
+            return
+
+        unet.detect_model_channels()
+        unet._on_cpu_devices.clear()
+
+        _write_generation_params(p, st)
+
+        if unet.verbose_ref.value:
+            print(f"[MegaFreeU] v{st.version}  "
+                  f"start={st.start_ratio:.3f} stop={st.stop_ratio:.3f} "
+                  f"smooth={st.transition_smoothness:.3f}  "
+                  f"ch_thresh=+-{st.channel_threshold}")
+            for i, si in enumerate(st.stage_infos):
+                ch = unet._stage_channels[i] if i < len(unet._stage_channels) else "?"
+                print(f"  Stage {i+1} ({ch}ch): "
+                      f"b={si.backbone_factor:.3f} [{si.b_start_ratio:.2f}-{si.b_end_ratio:.2f}] "
+                      f"{si.backbone_blend_mode}:{si.backbone_blend:.2f}  "
+                      f"s={si.skip_factor:.3f} [{si.s_start_ratio:.2f}-{si.s_end_ratio:.2f}] "
+                      f"fft={si.fft_type} r={si.fft_radius_ratio:.3f} "
+                      f"hfe={si.skip_high_end_factor:.2f} hfb={si.hf_boost:.2f}  "
+                      f"cap={'ON' if si.enable_adaptive_cap else 'off'} "
+                      f"({si.cap_threshold:.2f}/{si.cap_factor:.2f} {si.adaptive_cap_mode})")
+
+    def process_batch(self, p, *args, **kwargs):
+        global_state.current_sampling_step = 0
+        # FIX: reset PostCFG step counter for each image in batch
+        if hasattr(p, "_mega_pcfg"):
+            p._mega_pcfg["step"] = 0
+
+    def postprocess(self, p, processed, *args, **kwargs):
+        """Clean up per-image state after generation."""
+        if hasattr(p, "_mega_pcfg"):
+            p._mega_pcfg = {"enabled": False}
+
+
+def _write_generation_params(p, st):
+    """Write full Mega FreeU state into PNG extra_generation_params."""
+    p.extra_generation_params["MegaFreeU Schedule"] = (
+        f"{st.start_ratio}, {st.stop_ratio}, {st.transition_smoothness}")
+    p.extra_generation_params["MegaFreeU Stages"] = (
+        json.dumps([si.to_dict() for si in st.stage_infos]))
+    p.extra_generation_params["MegaFreeU Version"] = st.version
+    p.extra_generation_params["MegaFreeU Multiscale Mode"]     = st.multiscale_mode
+    p.extra_generation_params["MegaFreeU Multiscale Strength"] = str(st.multiscale_strength)
+    p.extra_generation_params["MegaFreeU Override Scales"]     = st.override_scales or ""
+    p.extra_generation_params["MegaFreeU Channel Threshold"]   = str(st.channel_threshold)
+    p.extra_generation_params["MegaFreeU Verbose"]             = str(st.verbose)
+    if st.pcfg_enabled:
+        p.extra_generation_params["MegaFreeU PostCFG"] = json.dumps({
+            "enabled":   st.pcfg_enabled,
+            "steps":     st.pcfg_steps,
+            "mode":      st.pcfg_mode,
+            "blend":     st.pcfg_blend,
+            "b":         st.pcfg_b,
+            "fourier":   st.pcfg_fourier,
+            "ms_mode":   st.pcfg_ms_mode,
+            "ms_str":    st.pcfg_ms_str,
+            "threshold": st.pcfg_threshold,
+            "s":         st.pcfg_s,
+            "gain":      st.pcfg_gain,
+        })
+
+
+def _flat_to_sis(flat) -> List[global_state.StageInfo]:
+    result = []
+    for i in range(global_state.STAGES_COUNT):
+        chunk = flat[i * _SN:(i + 1) * _SN]
+        si = global_state.StageInfo()
+        for j, fname in enumerate(_SF):
+            if j < len(chunk):
+                setattr(si, fname, chunk[j])
+        result.append(si)
+    return result
+
+
+# Callbacks
+def _on_cfg_step(*_args, **_kwargs):
+    global_state.current_sampling_step += 1
+
+def _on_cfg_post(params):
+    """WAS_PostCFGShift ported to A1111 on_cfg_after_cfg callback (exact algorithm)."""
+    p = getattr(params, "p", None)
+    if p is None:
+        p = getattr(getattr(params, "denoiser", None), "p", None)
+    if p is None: return
+    cfg = getattr(p, "_mega_pcfg", None)
+    if not cfg or not cfg.get("enabled"): return
+    cfg["step"] = cfg.get("step", 0) + 1
+    if cfg["step"] > cfg["steps"]: return
+    x = params.x
+    fn = unet.BLENDING_MODES.get(cfg["mode"], unet.BLENDING_MODES["inject"])
+    y  = fn(x, x * cfg["b"], cfg["blend"])
+    if cfg["fourier"]:
+        ms = global_state.MSCALES.get(cfg["ms_mode"])
+        y  = unet.filter_skip_box_multiscale(
+            y, cfg["threshold"], cfg["s"], ms, cfg["ms_str"])
+    if cfg["gain"] != 1.0:
+        y = y * float(cfg["gain"])
+    params.x = y
+
+try:
+    script_callbacks.on_cfg_after_cfg(_on_cfg_step)
+    script_callbacks.on_cfg_after_cfg(_on_cfg_post)
+except AttributeError:
+    # webui < 1.6.0 (sd-webui-freeu compatibility note)
+    script_callbacks.on_cfg_denoised(_on_cfg_step)
+    script_callbacks.on_cfg_denoised(_on_cfg_post)
+
+def _on_after_component(component, **kwargs):
+    eid = kwargs.get("elem_id", "")
+    for key, sid in [("txt2img", "txt2img_steps"), ("img2img", "img2img_steps")]:
+        if eid == sid:
+            _steps_comps[key] = component
+            for cb in _steps_cbs[key]: cb(component)
+            _steps_cbs[key].clear()
+
+script_callbacks.on_after_component(_on_after_component)
+
+def _on_ui_settings():
+    shared.opts.add_option(
+        "mega_freeu_png_auto_enable",
+        shared.OptionInfo(
+            default=True,
+            label="Auto-enable Mega FreeU when loading PNG info from a FreeU generation",
+            section=("mega_freeu", "Mega FreeU")))
+
+script_callbacks.on_ui_settings(_on_ui_settings)
+script_callbacks.on_before_ui(xyz_grid.patch)
+
+# Install th.cat patch at import (sd-webui-freeu pattern)
+unet.patch()

mega_freeu_a1111/tests/README.md

@@ -0,0 +1,31 @@
+# Mega FreeU — Test Suite
+
+Tests run without A1111 or PyTorch installed. A NumPy-based torch mock is used instead.
+
+## Requirements
+
+```
+pip install numpy
+```
+
+## Running
+
+```bash
+# From the extension root:
+python tests/test_core.py        # 144 tests — math, filters, blending, schedule, state
+python tests/test_fixes.py       # 36 tests  — dict-API, PNG metadata, Post-CFG, XYZ
+python tests/test_preset_pcfg.py # 32 tests  — full preset save/apply round-trip
+
+# Or run all at once:
+for f in tests/test_core.py tests/test_fixes.py tests/test_preset_pcfg.py; do
+    python "$f" && echo "--- $f PASSED ---" || echo "--- $f FAILED ---"
+done
+```
+
+## What is covered
+
+| File | Tests | Coverage |
+|------|-------|----------|
+| `test_core.py` | 144 | `ratio_to_region`, `lerp`, `get_backbone_scale` (V1+V2), `filter_skip_box`, `fourier_filter_gauss`, `get_band_energy_stats`, `filter_skip_box_multiscale`, all 9 blending modes, `parse_override_scales`, `_normalize`, `get_schedule_ratio`, `get_stage_bsratio`, `StageInfo`/`State` dataclasses, `update_attr`, `_load_user_presets`, `filter_skip_gaussian_adaptive` (no-cap/cap/fixed/aggressive), backbone blend math, `apply_xyz`, `detect_model_channels`, `_flat_to_sis`, PostCFG step counter |
+| `test_fixes.py` | 36 | `State` pcfg/verbose fields + round-trip, `_load_user_presets` with pcfg, `_write_generation_params` PNG keys, Post-CFG independent of Enable, dict-API compat |
+| `test_preset_pcfg.py` | 32 | Full preset save/apply with all 20 fields, pcfg disabled case, unknown preset, dict-API pcfg propagation |

mega_freeu_a1111/tests/mock_torch.py

@@ -0,0 +1,229 @@
+import math, types, sys
+import numpy as np
+
+class Tensor:
+    def __init__(self, data, device='cpu', dtype=None):
+        if isinstance(data, np.ndarray):
+            self._d = data
+        elif isinstance(data, Tensor):
+            self._d = data._d.copy()
+        else:
+            self._d = np.array(data, dtype=np.float32)
+        self.device = device
+        self.dtype  = dtype or 'float32'
+
+    @property
+    def shape(self): return self._d.shape
+    def size(self, dim=None): return self.shape[dim] if dim is not None else self.shape
+    @property
+    def is_cpu(self): return True
+
+    # dtype / device
+    def float(self):
+        return Tensor(self._d.astype(np.float32), self.device, 'float32')
+    def to(self, *a, **kw): return self
+    def cpu(self): return Tensor(self._d.copy(), 'cpu', self.dtype)
+    def item(self): return float(self._d.flat[0])
+
+    # shape ops
+    def view(self, *shape):
+        s = shape[0] if len(shape)==1 and isinstance(shape[0],(tuple,list)) else shape
+        return Tensor(self._d.reshape(s), self.device, self.dtype)
+    def reshape(self, *s): return self.view(*s)
+    def unsqueeze(self, dim):
+        return Tensor(np.expand_dims(self._d, dim), self.device, self.dtype)
+    def squeeze(self, dim=None):
+        d = self._d.squeeze() if dim is None else self._d.squeeze(dim)
+        return Tensor(d, self.device, self.dtype)
+    def expand_as(self, other):
+        return Tensor(np.broadcast_to(self._d, other.shape).copy(), self.device, self.dtype)
+
+    # .real / .imag — real tensors just return themselves
+    @property
+    def real(self): return Tensor(self._d.real.astype(np.float32), self.device, self.dtype)
+    @property
+    def imag(self): return Tensor(np.zeros_like(self._d.real, dtype=np.float32), self.device, self.dtype)
+
+    # reductions
+    def mean(self, dim=None, keepdim=False):
+        if dim is None: return Tensor(np.array(self._d.mean(), np.float32), self.device)
+        return Tensor(self._d.mean(axis=dim, keepdims=keepdim).astype(np.float32), self.device)
+    def sum(self, dim=None, keepdim=False):
+        if dim is None: return Tensor(np.array(self._d.sum(), np.float32), self.device)
+        return Tensor(self._d.sum(axis=dim, keepdims=keepdim).astype(np.float32), self.device)
+    def any(self): return bool(self._d.any())
+    def clamp(self, lo=None, hi=None):
+        return Tensor(np.clip(self._d, lo, hi).astype(np.float32), self.device, self.dtype)
+    def clamp_min(self, v):
+        return Tensor(np.maximum(self._d, v).astype(np.float32), self.device, self.dtype)
+    def abs(self): return Tensor(np.abs(self._d).astype(np.float32), self.device, self.dtype)
+    def min(self, dim=None, keepdim=False):
+        if dim is None: return Tensor(np.array(self._d.min(), np.float32), self.device)
+        v=self._d.min(axis=dim,keepdims=keepdim); i=self._d.argmin(axis=dim)
+        return Tensor(v.astype(np.float32),self.device), Tensor(i.astype(np.float32),self.device)
+    def max(self, dim=None, keepdim=False):
+        if dim is None: return Tensor(np.array(self._d.max(), np.float32), self.device)
+        v=self._d.max(axis=dim,keepdims=keepdim); i=self._d.argmax(axis=dim)
+        return Tensor(v.astype(np.float32),self.device), Tensor(i.astype(np.float32),self.device)
+
+    # indexing
+    def __getitem__(self, idx):
+        if isinstance(idx, Tensor):
+            idx = idx._d.astype(bool)
+        elif isinstance(idx, tuple):
+            idx = tuple(i._d.astype(bool) if isinstance(i, Tensor) else i for i in idx)
+        return Tensor(self._d[idx], self.device, self.dtype)
+    def __setitem__(self, idx, val):
+        self._d[idx] = val._d if isinstance(val, Tensor) else val
+
+    # arithmetic helpers
+    def _v(self, o): return o._d if isinstance(o, Tensor) else o
+    def __add__(self, o):  return Tensor((self._d + self._v(o)).astype(np.float32), self.device)
+    def __radd__(self, o): return self.__add__(o)
+    def __sub__(self, o):  return Tensor((self._d - self._v(o)).astype(np.float32), self.device)
+    def __rsub__(self, o): return Tensor((self._v(o) - self._d).astype(np.float32), self.device)
+    def __mul__(self, o):  return Tensor((self._d * self._v(o)).astype(np.float32), self.device)
+    def __rmul__(self, o): return self.__mul__(o)
+    def __truediv__(self, o):  return Tensor((self._d / self._v(o)).astype(np.float32), self.device)
+    def __rtruediv__(self, o): return Tensor((self._v(o) / self._d).astype(np.float32), self.device)
+    def __neg__(self): return Tensor(-self._d, self.device, self.dtype)
+    def __pow__(self, n): return Tensor(self._d**n, self.device, self.dtype)
+    def __imul__(self, o): self._d = (self._d * self._v(o)).astype(np.float32); return self
+    def __abs__(self): return self.abs()
+
+    # comparisons → float32 0/1 tensor (bool-compatible)
+    def __le__(self, o): return Tensor((self._d <= self._v(o)).astype(np.float32), self.device)
+    def __lt__(self, o): return Tensor((self._d <  self._v(o)).astype(np.float32), self.device)
+    def __ge__(self, o): return Tensor((self._d >= self._v(o)).astype(np.float32), self.device)
+    def __gt__(self, o): return Tensor((self._d >  self._v(o)).astype(np.float32), self.device)
+    def __eq__(self, o): return bool(np.array_equal(self._d, self._v(o))) if isinstance(o, (Tensor,np.ndarray)) else Tensor((self._d == o).astype(np.float32), self.device)
+    def __invert__(self):   # ~bool_tensor
+        return Tensor((~self._d.astype(bool)).astype(np.float32), self.device)
+    def __bool__(self): return bool(self._d.flat[0])
+    def __repr__(self): return f"Tensor({self._d.shape})"
+
+
+# ComplexTensor — returned by FFT ops
+class CTensor(Tensor):
+    def __init__(self, data, device='cpu', dtype=None):
+        super().__init__(data, device, dtype)
+    @property
+    def real(self):
+        return Tensor(self._d.real.astype(np.float32), self.device, self.dtype)
+    @property
+    def imag(self):
+        return Tensor(self._d.imag.astype(np.float32), self.device, self.dtype)
+    def __mul__(self, o):
+        v = o._d if isinstance(o, Tensor) else o
+        return CTensor((self._d * v), self.device)
+    def __imul__(self, o):
+        v = o._d if isinstance(o, Tensor) else o
+        self._d = self._d * v
+        return self
+
+
+# ── torch namespace ──────────────────────────────────────────────────────────
+torch = types.ModuleType('torch')
+torch.Tensor  = Tensor
+torch.float32 = 'float32'
+torch.device  = lambda s: s
+
+torch.full   = lambda shape, v, device='cpu', **kw: Tensor(np.full(shape, float(v), np.float32), device)
+torch.ones   = lambda *sh, device='cpu', **kw: Tensor(np.ones(sh[0] if len(sh)==1 else sh, np.float32), device)
+torch.zeros  = lambda *sh, device='cpu', **kw: Tensor(np.zeros(sh[0] if len(sh)==1 else sh, np.float32), device)
+torch.tensor = lambda v, device='cpu', dtype=None, **kw: Tensor(np.array(v, dtype=np.float32), device)
+torch.arange = lambda n, device='cpu', dtype=None, **kw: Tensor(np.arange(n, dtype=np.float32), device)
+
+def _meshgrid(y, x, indexing='ij'):
+    yy, xx = np.meshgrid(y._d, x._d, indexing=indexing)
+    return Tensor(yy.astype(np.float32), y.device), Tensor(xx.astype(np.float32), x.device)
+torch.meshgrid = _meshgrid
+
+torch.exp  = lambda x: Tensor(np.exp(np.clip(x._d,-500,500)).astype(np.float32), x.device)
+torch.sin  = lambda x: Tensor(np.sin(x._d).astype(np.float32), x.device)
+torch.cos  = lambda x: Tensor(np.cos(x._d).astype(np.float32), x.device)
+torch.acos = lambda x: Tensor(np.arccos(np.clip(x._d,-1+1e-7,1-1e-7)).astype(np.float32), x.device)
+torch.abs  = lambda x: Tensor(np.abs(x._d).astype(np.float32), x.device)
+torch.sqrt = lambda x: Tensor(np.sqrt(np.maximum(x._d,0)).astype(np.float32), x.device)
+torch.norm = lambda x, dim=None, keepdim=False, **kw: Tensor(
+    np.linalg.norm(x._d, axis=dim, keepdims=keepdim).astype(np.float32), x.device)
+
+def _max(x, dim=None, keepdim=False):
+    if dim is None: return float(x._d.max())
+    v = x._d.max(axis=dim, keepdims=keepdim)
+    i = x._d.argmax(axis=dim)
+    return Tensor(v.astype(np.float32), x.device), Tensor(i.astype(np.float32), x.device)
+def _min(x, dim=None, keepdim=False):
+    if dim is None: return float(x._d.min())
+    v = x._d.min(axis=dim, keepdims=keepdim)
+    i = x._d.argmin(axis=dim)
+    return Tensor(v.astype(np.float32), x.device), Tensor(i.astype(np.float32), x.device)
+torch.max = _max
+torch.min = _min
+
+def _where(c, a, b):
+    cd = c._d.astype(bool) if isinstance(c, Tensor) else np.array(c, bool)
+    ad = a._d if isinstance(a, Tensor) else a
+    bd = b._d if isinstance(b, Tensor) else b
+    return Tensor(np.where(cd, ad, bd).astype(np.float32),
+                  (a if isinstance(a,Tensor) else b).device)
+torch.where = _where
+
+class _Linalg:
+    @staticmethod
+    def norm(x, dim=None, keepdim=False, **kw):
+        return Tensor(
+            np.linalg.norm(x._d, axis=dim, keepdims=keepdim).astype(np.float32),
+            x.device)
+torch.linalg = _Linalg
+
+class _FFT:
+    @staticmethod
+    def fftn(x, dim=None):
+        ax = tuple(dim) if dim is not None else None
+        return CTensor(np.fft.fftn(x._d.astype(complex), axes=ax), x.device)
+    @staticmethod
+    def ifftn(x, dim=None):
+        ax = tuple(dim) if dim is not None else None
+        r = np.fft.ifftn(x._d, axes=ax).real.astype(np.float32)
+        # Return a plain Tensor but with .real property (already float32)
+        return Tensor(r, x.device)
+    @staticmethod
+    def fftshift(x, dim=None):
+        ax = tuple(dim) if dim is not None else None
+        d = np.fft.fftshift(x._d, axes=ax)
+        return CTensor(d, x.device) if isinstance(x, CTensor) else Tensor(d, x.device)
+    @staticmethod
+    def ifftshift(x, dim=None):
+        ax = tuple(dim) if dim is not None else None
+        d = np.fft.ifftshift(x._d, axes=ax)
+        return CTensor(d, x.device) if isinstance(x, CTensor) else Tensor(d, x.device)
+torch.fft = _FFT
+
+class _Backends:
+    class mps:
+        @staticmethod
+        def is_available(): return False
+torch.backends = _Backends
+
+sys.modules['torch'] = torch
+
+# ── modules mock ─────────────────────────────────────────────────────────────
+class _State:
+    sampling_steps = 20
+class _SharedObj:
+    state = _State()
+    class opts:
+        data = {}
+
+mmod = types.ModuleType('modules')
+mmod.shared = _SharedObj()
+sys.modules['modules'] = mmod
+for sub in ['modules.shared','modules.scripts','modules.processing',
+            'modules.script_callbacks','gradio']:
+    sys.modules[sub] = types.ModuleType(sub)
+
+import logging
+sys.modules['logging'] = logging
+# expose modules.shared in the right place
+import types as _t

mega_freeu_a1111/tests/test_core.py

@@ -0,0 +1,534 @@
+import pathlib, sys
+exec(open(str(pathlib.Path(__file__).parent / 'mock_torch.py')).read())
+
+import sys, math, types, dataclasses, json, tempfile, os, pathlib
+import numpy as np
+import pathlib; sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))
+
+# ── load ──────────────────────────────────────────────────────────────────────
+import importlib.util
+
+def load_mod(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    m = importlib.util.module_from_spec(spec)
+    sys.modules[name] = m
+    spec.loader.exec_module(m)
+    return m
+
+lib_pkg = types.ModuleType('lib_mega_freeu')
+sys.modules['lib_mega_freeu'] = lib_pkg
+
+GS = load_mod('lib_mega_freeu.global_state',
+              str(pathlib.Path(__file__).parent.parent / 'lib_mega_freeu' / 'global_state.py'))
+UN = load_mod('lib_mega_freeu.unet',
+              str(pathlib.Path(__file__).parent.parent / 'lib_mega_freeu' / 'unet.py'))
+print("Loaded OK\n")
+
+# ── test helpers ──────────────────────────────────────────────────────────────
+P=0; F=0; ERRS=[]
+
+def ok(t): global P; P+=1; print(f"  ✓  {t}")
+def ng(t,m=""): global F; F+=1; ERRS.append(f"{t}: {m}"); print(f"  ✗  {t}  {m}")
+def chk(t, cond, m=""): ok(t) if cond else ng(t, m)
+def near(t, got, want, tol=1e-4):
+    g = float(got) if hasattr(got,'item') else float(got)
+    (ok(t+f" ({g:.5g})") if abs(g-want)<=tol else ng(t, f"got {g:.5g} want {want}"))
+def shp(t, tensor, expected):
+    chk(t+f" shape={tensor.shape}", tuple(tensor.shape)==tuple(expected), f"≠{expected}")
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("═"*54)
+print("1. RATIO_TO_REGION")
+print("═"*54)
+R = UN.ratio_to_region
+s,e,inv = R(0.5, 0.0, 100)
+chk("[0,50] no-inv", s==0 and e==50 and not inv, f"s={s},e={e},inv={inv}")
+s,e,inv = R(0.5, 0.5, 100)
+chk("[50,100] no-inv", s==50 and e==100 and not inv)
+s,e,inv = R(0.7, 0.5, 100)
+chk("width+offset>1 → inverted", inv and s==20 and e==50, f"s={s},e={e},inv={inv}")
+s,e,inv = R(-0.3, 0.5, 100)    # negative width
+chk("neg width no crash", True)
+s,e,inv = R(0.0, 0.0, 100)
+chk("zero width: s==e", s==e)
+s,e,inv = R(1.0, 0.0, 100)     # full width → [0,100]
+chk("full width [0,100]", s==0 and e==100 and not inv)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n2. LERP")
+near("lerp(0,1,.5)", UN.lerp(0,1,.5), .5)
+near("lerp t=0 → a", UN.lerp(5,10,0), 5)
+near("lerp t=1 → b", UN.lerp(5,10,1), 10)
+near("lerp(2,4,.25)", UN.lerp(2,4,.25), 2.5)
+a=Tensor(np.array([[[[2.,4.]]]])); b=Tensor(np.array([[[[6.,8.]]]]))
+r=UN.lerp(a,b,0.5)
+chk("tensor lerp values", np.allclose(r._d, [[[[4.,6.]]]]))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n3. GET_BACKBONE_SCALE")
+h = Tensor(np.random.randn(2,8,4,4).astype(np.float32))
+# V1
+v1 = UN.get_backbone_scale(h, 1.3, "1")
+near("V1=1.3", v1 if isinstance(v1,(int,float)) else v1.item(), 1.3)
+# V2
+v2 = UN.get_backbone_scale(h, 1.4, "2")
+shp("V2 shape=(2,1,4,4)", v2, (2,1,4,4))
+chk("V2 vals in (0,2.5)", v2._d.min()>0 and v2._d.max()<2.5,
+    f"[{v2._d.min():.3f},{v2._d.max():.3f}]")
+# V2 factor=1.0 → all 1s
+v2id = UN.get_backbone_scale(h, 1.0, "2")
+chk("V2 factor=1→all-ones", np.allclose(v2id._d, 1.0, atol=1e-5))
+# V2 factor=2.0 → range [1,2]
+v2f2 = UN.get_backbone_scale(h, 2.0, "2")
+chk("V2 factor=2 → [1,2]",
+    v2f2._d.min()>=1.0-1e-5 and v2f2._d.max()<=2.0+1e-5,
+    f"[{v2f2._d.min():.4f},{v2f2._d.max():.4f}]")
+# zero input → no nan/crash
+hz = Tensor(np.zeros((1,8,4,4),np.float32))
+v2z = UN.get_backbone_scale(hz, 1.5, "2")
+chk("V2 zero input: no nan", not np.isnan(v2z._d).any())
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n4. FILTER_SKIP_BOX")
+x = Tensor(np.random.randn(1,4,16,16).astype(np.float32))
+# identity
+out = UN.filter_skip_box(x, 0.5, 1.0, 1.0)
+chk("box identity", np.allclose(out._d, x._d, atol=1e-4))
+# scale=0
+out0 = UN.filter_skip_box(x, 0.5, 0.0, 1.0)
+shp("box s=0 shape", out0, (1,4,16,16))
+chk("box s=0 ≠ input", not np.allclose(out0._d, x._d, atol=1e-3))
+# scale_high=0
+outh = UN.filter_skip_box(x, 0.5, 1.0, 0.0)
+chk("box h=0 ≠ input", not np.allclose(outh._d, x._d, atol=1e-3))
+# cutoff=0 edge
+outc0 = UN.filter_skip_box(x, 0.0, 0.5, 1.0)
+shp("box cutoff=0 shape", outc0, (1,4,16,16))
+# int cutoff
+outi = UN.filter_skip_box(x, 2, 0.7, 1.0)
+shp("box int cutoff", outi, (1,4,16,16))
+# batch > 1
+xb = Tensor(np.random.randn(3,8,16,16).astype(np.float32))
+outb = UN.filter_skip_box(xb, 0.3, 0.7, 1.0)
+shp("box batch=3", outb, (3,8,16,16))
+# no nan
+chk("box no nan", not np.isnan(out0._d).any())
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n5. FOURIER_FILTER_GAUSS")
+xg = Tensor(np.random.randn(1,4,16,16).astype(np.float32))
+outg = UN.fourier_filter_gauss(xg, 0.1, 0.8, 1.2)
+shp("gauss shape", outg, (1,4,16,16))
+chk("gauss s=.8 ≠ input", not np.allclose(outg._d, xg._d, atol=1e-3))
+# identity
+outgi = UN.fourier_filter_gauss(xg, 0.1, 1.0, 1.0)
+chk("gauss identity ≈ input", np.allclose(outgi._d, xg._d, atol=1e-4),
+    f"maxdiff={abs(outgi._d-xg._d).max():.2e}")
+# no nan
+chk("gauss no nan", not np.isnan(outg._d).any())
+# batch=3
+xg3 = Tensor(np.random.randn(3,8,16,16).astype(np.float32))
+outg3 = UN.fourier_filter_gauss(xg3, 0.08, 0.9, 1.0)
+shp("gauss batch=3", outg3, (3,8,16,16))
+# tiny radius (R=1)
+outgt = UN.fourier_filter_gauss(xg, 0.01, 0.5, 1.0)
+shp("gauss tiny R=1 shape", outgt, (1,4,16,16))
+# large radius (R=max)
+outgl = UN.fourier_filter_gauss(xg, 0.49, 0.5, 1.0)
+shp("gauss large radius shape", outgl, (1,4,16,16))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n6. GET_BAND_ENERGY_STATS")
+xe = Tensor(np.random.randn(2,4,16,16).astype(np.float32))
+lf,hf,cov = UN.get_band_energy_stats(xe, 3)
+chk("lf>0", lf>0)
+chk("hf>0", hf>0)
+chk("0<cov<100", 0<cov<100, f"cov={cov:.1f}")
+# zeros → energies=0
+xze = Tensor(np.zeros((1,4,8,8),np.float32))
+lf0,hf0,_ = UN.get_band_energy_stats(xze, 2)
+chk("zeros → lf=hf=0", lf0==0 and hf0==0)
+# R=1 minimal → some coverage
+lf1,hf1,cov1 = UN.get_band_energy_stats(xe, 1)
+chk("R=1 cover>0", cov1>0)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n7. FILTER_SKIP_BOX_MULTISCALE")
+xm = Tensor(np.random.randn(1,4,32,32).astype(np.float32))
+# identity
+outmi = UN.filter_skip_box_multiscale(xm, 0.3, 1.0, None, 1.0, 1.0)
+chk("ms identity", np.allclose(outmi._d, xm._d, atol=1e-4))
+# single-scale preset
+outss = UN.filter_skip_box_multiscale(xm, 0.3, 0.7, [(8,1.5)], 1.0)
+shp("ms single-scale shape", outss, (1,4,32,32))
+chk("ms single differs", not np.allclose(outss._d, xm._d, atol=1e-3))
+# multi-scale preset
+outms = UN.filter_skip_box_multiscale(xm, 0.3, 0.7, [[(5,0.0),(15,1.0)]], 1.0)
+shp("ms multi-scale shape", outms, (1,4,32,32))
+# scale_high
+outsh = UN.filter_skip_box_multiscale(xm, 0.3, 0.7, None, 1.0, 1.5)
+chk("ms scale_high=1.5 differs", not np.allclose(outsh._d, xm._d, atol=1e-3))
+# no nan
+chk("ms no nan", not np.isnan(outss._d).any())
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n8. BLENDING MODES")
+a = Tensor(np.full((1,4,4,4), 2.0, np.float32))
+b = Tensor(np.full((1,4,4,4), 4.0, np.float32))
+for mname, fn in UN.BLENDING_MODES.items():
+    try:
+        out = fn(a, b, 0.5)
+        shp(f"  {mname}", out, (1,4,4,4))
+        chk(f"  {mname} no nan", not np.isnan(out._d).any())
+    except Exception as e:
+        ng(f"  {mname} CRASHED", str(e))
+
+lerp_fn = UN.BLENDING_MODES['lerp']
+near("lerp t=0→a", lerp_fn(a,b,0)._d.mean(), 2.0)
+near("lerp t=1→b", lerp_fn(a,b,1)._d.mean(), 4.0)
+near("lerp t=.5→3", lerp_fn(a,b,.5)._d.mean(), 3.0)
+
+inj = UN.BLENDING_MODES['inject']
+near("inject a+b*.5=4", inj(a,b,.5)._d.mean(), 4.0)
+near("inject t=0→a", inj(a,b,0)._d.mean(), 2.0)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n9. PARSE_OVERRIDE_SCALES")
+P2 = UN.parse_override_scales
+chk("None→None", P2(None) is None)
+chk("empty→None", P2("") is None)
+chk("comment-only→None", P2("# x\n! y\n// z") is None)
+r = P2("10, 1.5\n20, 0.8")
+chk("2 lines → 2 entries", r is not None and len(r)==2)
+chk("entry[0]=(10,1.5)", r and r[0]==(10,1.5), str(r))
+r2 = P2("# comment\n5, 2.0\n! skip\n15, 0.5")
+chk("with comments → 2 entries", r2 and len(r2)==2)
+r3 = P2("10,1.5\nbad_line\n20,0.8")  # malformed line skipped
+chk("malformed skipped", r3 and len(r3)==2, str(r3))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n10. _NORMALIZE")
+n = UN._normalize
+chk("norm [0,3]: min=0 max=1",
+    abs(n(Tensor(np.array([[[[0.,1.,2.,3.]]]])))._d.min())<1e-5 and
+    abs(n(Tensor(np.array([[[[0.,1.,2.,3.]]]])))._d.max()-1.0)<1e-5)
+chk("norm const: no crash/nan",
+    not np.isnan(n(Tensor(np.ones((1,1,4,4),np.float32)*5))._d).any())
+chk("norm negatives: min=0 max=1",
+    abs(n(Tensor(np.array([[[[-2.,-1.,0.,1.,2.]]]])))._d.max()-1.0)<1e-5)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n11. GET_SCHEDULE_RATIO")
+mmod.shared.state.sampling_steps = 20
+GS.instance = GS.State(start_ratio=0.0, stop_ratio=1.0, transition_smoothness=0.0)
+GS.current_sampling_step = 0
+near("full range step=0 → 1", UN.get_schedule_ratio(), 1.0)
+GS.current_sampling_step = 10
+near("full range step=10 → 1", UN.get_schedule_ratio(), 1.0)
+GS.current_sampling_step = 25
+near("past stop → 0", UN.get_schedule_ratio(), 0.0)
+GS.instance = GS.State(start_ratio=0.5, stop_ratio=1.0, transition_smoothness=0.0)
+GS.current_sampling_step = 5
+near("before start → 0", UN.get_schedule_ratio(), 0.0)
+GS.current_sampling_step = 15
+near("in range → 1", UN.get_schedule_ratio(), 1.0)
+GS.instance = GS.State(start_ratio=0.5, stop_ratio=0.5, transition_smoothness=0.0)
+near("start==stop → 0", UN.get_schedule_ratio(), 0.0)
+GS.instance = GS.State(start_ratio=0.0, stop_ratio=1.0, transition_smoothness=1.0)
+GS.current_sampling_step = 0
+r = UN.get_schedule_ratio()
+chk("smoothness=1 step=0: 0≤r≤1", 0.0<=r<=1.0)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n12. GET_STAGE_BSRATIO")
+bsr = UN.get_stage_bsratio
+mmod.shared.state.sampling_steps = 20
+GS.current_sampling_step = 5   # pct≈0.26
+near("bsr [0,1]→1", bsr(0.0,1.0), 1.0)
+near("bsr [0.5,1]→0 (pct<0.5)", bsr(0.5,1.0), 0.0)
+near("bsr [0,0.5]→1 (pct<0.5)", bsr(0.0,0.5), 1.0)
+GS.current_sampling_step = 18  # pct≈0.95
+near("bsr [0,0.5]→0 (pct>0.5)", bsr(0.0,0.5), 0.0)
+near("bsr [0.5,1]→1 (pct>0.5)", bsr(0.5,1.0), 1.0)
+mmod.shared.state.sampling_steps = 1
+GS.current_sampling_step = 0
+near("steps=1 pct=0 in [0,1]→1", bsr(0.0,1.0), 1.0)
+mmod.shared.state.sampling_steps = 20
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n13. STAGEINFO + STATE")
+si = GS.StageInfo()
+chk("default bf=1", si.backbone_factor==1.0)
+chk("default fft=box", si.fft_type=="box")
+chk("default cap=off", not si.enable_adaptive_cap)
+chk("19 STAGE_FIELD_NAMES", len(GS.STAGE_FIELD_NAMES)==19)
+
+st = GS.State(version="Version 2")
+chk("version coerced '2'", st.version=="2")
+st2 = GS.State(version="1")
+chk("version '1' stays '1'", st2.version=="1")
+
+# unknown dict key in stage_infos
+st3 = GS.State(stage_infos=[{'backbone_factor':1.5,'UNKNOWN':999}])
+chk("dict unknown key ignored", st3.stage_infos[0].backbone_factor==1.5)
+
+# padding to STAGES_COUNT
+st4 = GS.State(stage_infos=[GS.StageInfo(backbone_factor=2.0)])
+chk("pads to 3 stages", len(st4.stage_infos)==3)
+chk("pad default bf=1", st4.stage_infos[1].backbone_factor==1.0)
+
+# round-trip
+st5 = GS.State(version="2", stage_infos=[GS.StageInfo(backbone_factor=1.7)])
+d = st5.to_dict()
+chk("to_dict no 'enable'", 'enable' not in d)
+chk("to_dict has stages", 'stage_infos' in d)
+fields = {f.name for f in dataclasses.fields(GS.State)}
+st6 = GS.State(**{k:v for k,v in d.items() if k in fields})
+chk("round-trip version", st6.version=="2")
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n14. UPDATE_ATTR (XYZ shorthands)")
+st = GS.State()
+st.update_attr("b0", 1.5);   chk("b0→bf stage0", st.stage_infos[0].backbone_factor==1.5)
+st.update_attr("s1", 0.3);   chk("s1→sf stage1", st.stage_infos[1].skip_factor==0.3)
+st.update_attr("ft2","gaussian"); chk("ft2→fft_type stage2", st.stage_infos[2].fft_type=="gaussian")
+st.update_attr("acm0","fixed");   chk("acm0→cap_mode stage0", st.stage_infos[0].adaptive_cap_mode=="fixed")
+st.update_attr("cap1", True);     chk("cap1→enable_adaptive_cap", st.stage_infos[1].enable_adaptive_cap==True)
+st.update_attr("ct0", 0.4);       chk("ct0→cap_threshold", st.stage_infos[0].cap_threshold==0.4)
+st.update_attr("start_ratio", 0.2); chk("start_ratio direct", st.start_ratio==0.2)
+st.update_attr("enable", True);   chk("enable direct", st.enable==True)
+# unknown key → no crash
+try:
+    st.update_attr("UNKNOWN_KEY", 99)
+    chk("unknown key: no crash", True)
+except Exception as e:
+    chk("unknown key: no crash", False, str(e))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n15. _LOAD_USER_PRESETS robustness")
+# Good + bad preset in same file
+pdata = {
+    "good": {"start_ratio":0.0,"stop_ratio":1.0,"transition_smoothness":0.0,
+             "version":"2","multiscale_mode":"Default","multiscale_strength":1.0,
+             "override_scales":"","channel_threshold":96,
+             "stage_infos":[{"backbone_factor":1.3}]},
+    "with_unknown": {"start_ratio":0.0,"stop_ratio":1.0,"transition_smoothness":0.0,
+                     "version":"1","multiscale_mode":"Default","multiscale_strength":1.0,
+                     "override_scales":"","channel_threshold":96,
+                     "FUTURE_FIELD":"ignored","stage_infos":[]}
+}
+with tempfile.NamedTemporaryFile(mode='w',suffix='.json',delete=False) as f:
+    json.dump(pdata, f); tmp = f.name
+GS.PRESETS_PATH = pathlib.Path(tmp)
+res = GS._load_user_presets()
+chk("good preset loaded", "good" in res)
+chk("good preset bf=1.3", res.get("good") and res["good"].stage_infos[0].backbone_factor==1.3)
+chk("no crash on unknown field preset", True)
+os.unlink(tmp)
+
+# Invalid JSON → {}
+with tempfile.NamedTemporaryFile(mode='w',suffix='.json',delete=False) as f:
+    f.write("{bad json!!!"); tmp2 = f.name
+GS.PRESETS_PATH = pathlib.Path(tmp2)
+chk("invalid JSON → {}", GS._load_user_presets() == {})
+os.unlink(tmp2)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n16. FILTER_SKIP_GAUSSIAN_ADAPTIVE")
+hs = Tensor(np.random.randn(1,4,16,16).astype(np.float32))
+
+# no cap
+si_nc = GS.StageInfo(skip_factor=0.8, fft_type='gaussian',
+                     fft_radius_ratio=0.1, hf_boost=1.2,
+                     skip_high_end_factor=1.1, enable_adaptive_cap=False)
+out_nc = UN.filter_skip_gaussian_adaptive(hs, si_nc)
+shp("no-cap shape", out_nc, (1,4,16,16))
+chk("no-cap differs from input", not np.allclose(out_nc._d, hs._d, atol=1e-3))
+chk("no-cap no nan", not np.isnan(out_nc._d).any())
+
+# identity (scale=1, hfb=1)
+si_id = GS.StageInfo(skip_factor=1.0, fft_type='gaussian',
+                     fft_radius_ratio=0.1, hf_boost=1.0,
+                     skip_high_end_factor=1.0, enable_adaptive_cap=False)
+out_id = UN.filter_skip_gaussian_adaptive(hs, si_id)
+chk("gauss identity ≈ input", np.allclose(out_id._d, hs._d, atol=1e-4),
+    f"maxdiff={abs(out_id._d-hs._d).max():.2e}")
+
+# with cap
+si_cap = GS.StageInfo(skip_factor=0.3, fft_type='gaussian',
+                      fft_radius_ratio=0.15, hf_boost=1.0,
+                      skip_high_end_factor=1.0, enable_adaptive_cap=True,
+                      cap_threshold=0.35, cap_factor=0.6, adaptive_cap_mode='adaptive')
+out_cap = UN.filter_skip_gaussian_adaptive(hs, si_cap)
+shp("cap shape", out_cap, (1,4,16,16))
+chk("cap no nan", not np.isnan(out_cap._d).any())
+
+# fixed cap mode
+si_fixed = GS.StageInfo(skip_factor=0.3, fft_type='gaussian',
+                        fft_radius_ratio=0.15, hf_boost=1.0,
+                        skip_high_end_factor=1.0, enable_adaptive_cap=True,
+                        cap_threshold=0.35, cap_factor=0.6, adaptive_cap_mode='fixed')
+out_fixed = UN.filter_skip_gaussian_adaptive(hs, si_fixed)
+shp("fixed-cap shape", out_fixed, (1,4,16,16))
+chk("fixed-cap no nan", not np.isnan(out_fixed._d).any())
+
+# very aggressive scale (s=0.0) with cap
+si_agg = GS.StageInfo(skip_factor=0.0, fft_type='gaussian',
+                      fft_radius_ratio=0.1, hf_boost=1.0,
+                      skip_high_end_factor=1.0, enable_adaptive_cap=True,
+                      cap_threshold=0.35, cap_factor=0.6, adaptive_cap_mode='adaptive')
+out_agg = UN.filter_skip_gaussian_adaptive(hs, si_agg)
+shp("aggressive cap shape", out_agg, (1,4,16,16))
+chk("aggressive cap no nan", not np.isnan(out_agg._d).any())
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n17. BACKBONE BLEND MATH (unit test)")
+h = Tensor(np.full((1,8,4,4), 3.0, np.float32))
+dims = 8
+rbegin, rend, rinv = UN.ratio_to_region(0.5, 0.0, dims)
+mask_np = np.zeros(dims, np.float32)
+if not rinv: mask_np[rbegin:rend] = 1.0
+else: mask_np[:rend]=1.0; mask_np[rbegin:]=1.0
+mask_t = Tensor(mask_np.reshape(1,-1,1,1))
+
+# V1 scale=2.0: masked → 6, unmasked → 3
+scale = 2.0
+h_scaled = h * (mask_t * scale + (1.0 - mask_t))
+masked = h_scaled._d[0, :rend, 0, 0]
+unmasked = h_scaled._d[0, rend:, 0, 0]
+chk("masked ch → 6.0",   np.allclose(masked, 6.0, atol=1e-4), str(masked))
+chk("unmasked ch → 3.0", np.allclose(unmasked, 3.0, atol=1e-4), str(unmasked))
+
+# blend lerp 0.5: masked → lerp(3,6,0.5)=4.5
+lerp_fn = UN.BLENDING_MODES['lerp']
+h_scaled_full = h * (mask_t * scale + (1.0 - mask_t))
+h_blended = lerp_fn(h, h_scaled_full, 0.5)
+h_out = h * (1.0 - mask_t) + h_blended * mask_t
+chk("blend lerp 0.5 masked→4.5",
+    np.allclose(h_out._d[0,:rend,0,0], 4.5, atol=1e-4),
+    str(h_out._d[0,:rend,0,0]))
+chk("blend lerp 0.5 unmasked→3.0",
+    np.allclose(h_out._d[0,rend:,0,0], 3.0, atol=1e-4))
+
+# inject blend t=0.5: h + h_scaled*0.5 = 3 + 6*0.5 = 6
+inj_fn = UN.BLENDING_MODES['inject']
+h_inj = inj_fn(h, h_scaled_full, 0.5)
+h_out2 = h * (1.0 - mask_t) + h_inj * mask_t
+chk("blend inject masked→h+h_scaled*.5",
+    np.allclose(h_out2._d[0,:rend,0,0], 3+6*0.5, atol=1e-4),
+    str(h_out2._d[0,:rend,0,0]))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n18. APPLY_XYZ")
+GS.instance = GS.State(); GS.xyz_attrs.clear()
+orig = GS.instance.stage_infos[0].backbone_factor
+GS.apply_xyz()
+chk("empty attrs: unchanged", GS.instance.stage_infos[0].backbone_factor==orig)
+
+GS.xyz_attrs['b0'] = 2.5; GS.apply_xyz()
+chk("b0=2.5 applied", GS.instance.stage_infos[0].backbone_factor==2.5)
+GS.xyz_attrs.clear()
+
+# Preset
+GS.reload_presets()
+pname = list(GS.all_presets.keys())[0]
+pbf = GS.all_presets[pname].stage_infos[0].backbone_factor
+GS.xyz_attrs['preset'] = pname; GS.apply_xyz()
+chk("preset applied", abs(GS.instance.stage_infos[0].backbone_factor-pbf)<1e-5)
+GS.xyz_attrs.clear()
+
+# Unknown preset → warning, instance unchanged
+GS.instance = GS.State()
+bf_before = GS.instance.stage_infos[0].backbone_factor
+GS.xyz_attrs['preset'] = 'NO_SUCH_PRESET_XYZ'; GS.apply_xyz()
+chk("unknown preset: no crash", True)
+GS.xyz_attrs.clear()
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n19. DETECT_MODEL_CHANNELS")
+UN.detect_model_channels()
+chk("fallback=(1280,640,320)", UN._stage_channels==(1280,640,320))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n20. _FLAT_TO_SIS (simulated)")
+_SF = [f.name for f in dataclasses.fields(GS.StageInfo)]
+_SN = len(_SF)
+def flat_to_sis(flat):
+    res=[]
+    for i in range(GS.STAGES_COUNT):
+        chunk=flat[i*_SN:(i+1)*_SN]
+        si_new=GS.StageInfo()
+        for j,fname in enumerate(_SF):
+            if j<len(chunk): setattr(si_new,fname,chunk[j])
+        res.append(si_new)
+    return res
+
+si0=GS.StageInfo(backbone_factor=1.7, skip_factor=0.3, fft_type='gaussian')
+flat=[]
+for si_x in [si0,GS.StageInfo(),GS.StageInfo()]:
+    for f in _SF: flat.append(getattr(si_x,f))
+sis=flat_to_sis(flat)
+chk("flat bf=1.7", sis[0].backbone_factor==1.7)
+chk("flat sf=0.3", sis[0].skip_factor==0.3)
+chk("flat fft=gaussian", sis[0].fft_type=='gaussian')
+chk("flat 3 stages", len(sis)==3)
+chk("flat short: no crash", len(flat_to_sis(flat[:20]))==3)
+
+# ══════════════════════════════════════════════════════════════════════════════
+print("\n21. POSTCFG STEP COUNTER LOGIC")
+# Simulate _on_cfg_post step counting
+cfg = {"enabled":True,"steps":3,"mode":"lerp","blend":0.5,"b":1.1,
+       "fourier":False,"ms_mode":"Default","ms_str":1.0,
+       "threshold":1,"s":0.5,"gain":1.0,"step":0}
+
+class FakeParams:
+    def __init__(self): self.x = Tensor(np.ones((1,4,8,8),np.float32)*2.0)
+
+class FakeP:
+    _mega_pcfg = cfg
+
+class FakeDenoiser:
+    p = FakeP()
+
+# simulate _on_cfg_post inline
+def run_post(params):
+    p = getattr(params, "p", None)
+    if p is None:
+        p = getattr(getattr(params,"denoiser",None),"p",None)
+    if p is None: return False
+    c = getattr(p,"_mega_pcfg",None)
+    if not c or not c.get("enabled"): return False
+    c["step"] = c.get("step",0)+1
+    if c["step"] > c["steps"]: return False
+    x = params.x
+    fn = UN.BLENDING_MODES.get(c["mode"], UN.BLENDING_MODES["inject"])
+    params.x = fn(x, x*c["b"], c["blend"])
+    return True
+
+# via p attribute
+fp1 = FakeParams(); fp1.p = FakeP(); fp1.p._mega_pcfg = {"enabled":True,"steps":2,"mode":"lerp","blend":0.5,"b":1.0,"fourier":False,"step":0,"gain":1.0}
+ran1 = run_post(fp1); chk("postcfg step1 ran", ran1)
+ran2 = run_post(fp1); chk("postcfg step2 ran", ran2)
+ran3 = run_post(fp1); chk("postcfg step3 → past limit, skipped", not ran3)
+
+# via denoiser.p
+fp2 = FakeParams()
+class D2:
+    class p2:
+        _mega_pcfg = {"enabled":True,"steps":1,"mode":"inject","blend":0.5,"b":1.1,"fourier":False,"step":0,"gain":1.0}
+    p = p2
+fp2.denoiser = D2
+chk("postcfg via denoiser.p: no crash", run_post(fp2))
+
+# disabled → skip
+fp3 = FakeParams(); fp3.p = type('P',(),{'_mega_pcfg':{"enabled":False}})()
+chk("postcfg disabled → skip", not run_post(fp3))
+
+# ══════════════════════════════════════════════════════════════════════════════
+print(f"\n{'═'*54}")
+print(f"TOTAL: {P} PASS   {F} FAIL")
+if ERRS:
+    print("\nFailed:")
+    for e in ERRS: print(f"  • {e}")
+else:
+    print("ALL TESTS PASSED ✓")

mega_freeu_a1111/tests/test_fixes.py

@@ -0,0 +1,166 @@
+import pathlib, sys
+exec(open(str(pathlib.Path(__file__).parent / 'mock_torch.py')).read())
+import sys, math, types, dataclasses, json
+import numpy as np
+import pathlib; sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))
+import importlib.util
+
+def load_mod(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    m = importlib.util.module_from_spec(spec); sys.modules[name]=m; spec.loader.exec_module(m); return m
+
+lib_pkg = types.ModuleType('lib_mega_freeu'); sys.modules['lib_mega_freeu'] = lib_pkg
+GS = load_mod('lib_mega_freeu.global_state', str(pathlib.Path(__file__).parent.parent / 'lib_mega_freeu' / 'global_state.py'))
+UN = load_mod('lib_mega_freeu.unet', str(pathlib.Path(__file__).parent.parent / 'lib_mega_freeu' / 'unet.py'))
+
+P=0; F=0; ERRS=[]
+def ok(t): global P; P+=1; print(f"  ✓  {t}")
+def ng(t,m=""): global F; F+=1; ERRS.append(f"{t}: {m}"); print(f"  ✗  {t}  {m}")
+def chk(t, c, m=""): ok(t) if c else ng(t, m)
+
+print("═"*52)
+print("FIX 1: State has pcfg_ and verbose fields")
+print("═"*52)
+st = GS.State()
+chk("pcfg_enabled default False", not st.pcfg_enabled)
+chk("pcfg_steps default 20", st.pcfg_steps == 20)
+chk("pcfg_mode default inject", st.pcfg_mode == "inject")
+chk("verbose default False", not st.verbose)
+
+st2 = GS.State(pcfg_enabled=True, pcfg_b=1.5, pcfg_steps=10, verbose=True)
+chk("pcfg_enabled=True", st2.pcfg_enabled)
+chk("pcfg_b=1.5", st2.pcfg_b == 1.5)
+chk("verbose=True", st2.verbose)
+
+print("\n═"*27)
+print("FIX 2: to_dict() round-trip includes pcfg fields")
+d = st2.to_dict()
+chk("to_dict has pcfg_enabled", "pcfg_enabled" in d)
+chk("to_dict has pcfg_b", "pcfg_b" in d)
+chk("to_dict has verbose", "verbose" in d)
+chk("to_dict no 'enable'", "enable" not in d)
+st3 = GS.State(**{k:v for k,v in d.items() if k in {f.name for f in dataclasses.fields(GS.State)}})
+chk("round-trip pcfg_b=1.5", st3.pcfg_b == 1.5)
+chk("round-trip verbose=True", st3.verbose)
+
+print("\n═"*27)
+print("FIX 3: _load_user_presets saves/restores pcfg")
+import json, tempfile, os, pathlib
+preset_data = {
+    "my_pcfg_preset": {
+        "start_ratio": 0.0, "stop_ratio": 1.0, "transition_smoothness": 0.0,
+        "version": "2", "multiscale_mode": "Default", "multiscale_strength": 1.0,
+        "override_scales": "", "channel_threshold": 96,
+        "pcfg_enabled": True, "pcfg_b": 1.8, "pcfg_steps": 5,
+        "pcfg_mode": "lerp", "pcfg_blend": 0.7, "verbose": True,
+        "stage_infos": []
+    }
+}
+with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+    json.dump(preset_data, f); tmp = f.name
+GS.PRESETS_PATH = pathlib.Path(tmp)
+res = GS._load_user_presets()
+chk("preset loaded", "my_pcfg_preset" in res)
+p = res.get("my_pcfg_preset")
+chk("pcfg_enabled restored", p and p.pcfg_enabled == True)
+chk("pcfg_b restored", p and p.pcfg_b == 1.8)
+chk("pcfg_steps restored", p and p.pcfg_steps == 5)
+chk("verbose restored", p and p.verbose == True)
+os.unlink(tmp)
+
+print("\n═"*27)
+print("FIX 4: _write_generation_params (simulate)")
+class FakePNG:
+    extra_generation_params = {}
+
+st4 = GS.State(
+    start_ratio=0.1, stop_ratio=0.9, transition_smoothness=0.5,
+    version="2", multiscale_mode="Multi-Bandpass", multiscale_strength=0.8,
+    override_scales="10, 1.5\n20, 0.8", channel_threshold=64,
+    pcfg_enabled=True, pcfg_b=1.2, pcfg_steps=15, pcfg_mode="inject",
+    pcfg_blend=1.0, pcfg_fourier=False, pcfg_ms_mode="Default",
+    pcfg_ms_str=1.0, pcfg_threshold=1, pcfg_s=0.5, pcfg_gain=1.0,
+    verbose=True,
+    stage_infos=[GS.StageInfo(backbone_factor=1.3)]
+)
+
+# Simulate _write_generation_params
+fp = FakePNG()
+fp.extra_generation_params["MegaFreeU Schedule"] = f"{st4.start_ratio}, {st4.stop_ratio}, {st4.transition_smoothness}"
+fp.extra_generation_params["MegaFreeU Stages"] = json.dumps([si.to_dict() for si in st4.stage_infos])
+fp.extra_generation_params["MegaFreeU Version"] = st4.version
+fp.extra_generation_params["MegaFreeU Multiscale Mode"] = st4.multiscale_mode
+fp.extra_generation_params["MegaFreeU Multiscale Strength"] = str(st4.multiscale_strength)
+fp.extra_generation_params["MegaFreeU Override Scales"] = st4.override_scales
+fp.extra_generation_params["MegaFreeU Channel Threshold"] = str(st4.channel_threshold)
+fp.extra_generation_params["MegaFreeU Verbose"] = str(st4.verbose)
+if st4.pcfg_enabled:
+    fp.extra_generation_params["MegaFreeU PostCFG"] = json.dumps({
+        "enabled": st4.pcfg_enabled, "steps": st4.pcfg_steps,
+        "mode": st4.pcfg_mode, "blend": st4.pcfg_blend,
+        "b": st4.pcfg_b, "fourier": st4.pcfg_fourier,
+        "ms_mode": st4.pcfg_ms_mode, "ms_str": st4.pcfg_ms_str,
+        "threshold": st4.pcfg_threshold, "s": st4.pcfg_s, "gain": st4.pcfg_gain,
+    })
+
+eg = fp.extra_generation_params
+chk("PNG has Schedule",           "MegaFreeU Schedule" in eg)
+chk("PNG has Stages",             "MegaFreeU Stages" in eg)
+chk("PNG has Version",            "MegaFreeU Version" in eg)
+chk("PNG has Multiscale Mode",    "MegaFreeU Multiscale Mode" in eg)
+chk("PNG has Multiscale Strength","MegaFreeU Multiscale Strength" in eg)
+chk("PNG has Override Scales",    "MegaFreeU Override Scales" in eg)
+chk("PNG has Channel Threshold",  "MegaFreeU Channel Threshold" in eg)
+chk("PNG has Verbose",            "MegaFreeU Verbose" in eg)
+chk("PNG has PostCFG",            "MegaFreeU PostCFG" in eg)
+
+# Verify PostCFG round-trip
+pcfg_d = json.loads(eg["MegaFreeU PostCFG"])
+chk("PostCFG b=1.2",   pcfg_d["b"] == 1.2)
+chk("PostCFG steps=15",pcfg_d["steps"] == 15)
+
+# Verify multiscale restore
+chk("ms_mode=Multi-Bandpass", eg["MegaFreeU Multiscale Mode"] == "Multi-Bandpass")
+chk("ch_thresh=64", eg["MegaFreeU Channel Threshold"] == "64")
+
+print("\n═"*27)
+print("FIX 5: Post-CFG independent of Enable (simulate process logic)")
+# The fix: pcfg is set BEFORE checking st.enable
+# Test: when enabled=False but pcfg_enabled=True, pcfg still created
+
+class FakeP2:
+    extra_generation_params = {}
+    _mega_pcfg = None
+
+fp2 = FakeP2()
+# Simulate new process() logic for disabled main FreeU + enabled Post-CFG
+st_disabled = GS.State(enable=False, pcfg_enabled=True, pcfg_b=1.5, pcfg_steps=10)
+# Post-CFG created regardless
+if st_disabled.pcfg_enabled:
+    fp2._mega_pcfg = {"enabled": True, "b": st_disabled.pcfg_b, "steps": st_disabled.pcfg_steps, "step": 0}
+else:
+    fp2._mega_pcfg = {"enabled": False}
+
+chk("pcfg set even when main disabled", fp2._mega_pcfg["enabled"] == True)
+chk("pcfg_b=1.5 propagated", fp2._mega_pcfg["b"] == 1.5)
+
+print("\n═"*27)
+print("FIX 6: dict-API compat (old sd-webui-freeu alwayson_scripts)")
+# Simulate the dict branch of process()
+dict_args = {
+    "enable": True, "start_ratio": 0.2, "stop_ratio": 0.8,
+    "version": "2", "multiscale_mode": "Default"
+}
+fields = {f.name for f in dataclasses.fields(GS.State)}
+GS.instance = GS.State(**{k:v for k,v in dict_args.items() if k in fields})
+chk("dict API: start_ratio=0.2", GS.instance.start_ratio == 0.2)
+chk("dict API: version coerced '2'", GS.instance.version == "2")
+chk("dict API: pcfg defaults", not GS.instance.pcfg_enabled)
+
+print(f"\n{'═'*52}")
+print(f"NEW FIXES: {P} PASS   {F} FAIL")
+if ERRS:
+    print("\nFailed:")
+    for e in ERRS: print(f"  • {e}")
+else:
+    print("ALL NEW FIX TESTS PASSED ✓")

mega_freeu_a1111/tests/test_preset_pcfg.py

@@ -0,0 +1,160 @@
+import pathlib, sys
+exec(open(str(pathlib.Path(__file__).parent / 'mock_torch.py')).read())
+import sys, types, dataclasses, json
+import numpy as np
+import pathlib; sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))
+import importlib.util
+
+def load(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    m = importlib.util.module_from_spec(spec); sys.modules[name]=m; spec.loader.exec_module(m); return m
+
+lib = types.ModuleType('lib_mega_freeu'); sys.modules['lib_mega_freeu'] = lib
+GS = load('lib_mega_freeu.global_state', str(pathlib.Path(__file__).parent.parent / 'lib_mega_freeu' / 'global_state.py'))
+UN = load('lib_mega_freeu.unet', str(pathlib.Path(__file__).parent.parent / 'lib_mega_freeu' / 'unet.py'))
+
+P=0; F=0; ERRS=[]
+def ok(t): global P; P+=1; print(f"  ✓  {t}")
+def ng(t,m=""): global F; F+=1; ERRS.append(f"{t}: {m}"); print(f"  ✗  {t}  {m}")
+def chk(t, c, m=""): ok(t) if c else ng(t, m)
+
+_SF = [f.name for f in dataclasses.fields(GS.StageInfo)]
+_SN = len(_SF)
+
+def _flat_to_sis(flat):
+    result = []
+    for i in range(GS.STAGES_COUNT):
+        chunk = flat[i*_SN:(i+1)*_SN]
+        si = GS.StageInfo()
+        for j, fname in enumerate(_SF):
+            if j < len(chunk): setattr(si, fname, chunk[j])
+        result.append(si)
+    return result
+
+# Simulate _save_p with all fields
+def _save_p(name, sr, sp, sm, ver, msm, mss, ovs, cht,
+            p_en, p_steps, p_mode, p_bl, p_b,
+            p_four, p_mmd, p_mst, p_thr, p_s, p_gain,
+            v_log, *flat):
+    sis = _flat_to_sis(flat)
+    vc  = GS.ALL_VERSIONS.get(ver, "1")
+    GS.all_presets[name] = GS.State(
+        start_ratio=sr, stop_ratio=sp, transition_smoothness=sm,
+        version=vc, multiscale_mode=msm, multiscale_strength=float(mss),
+        override_scales=ovs or "", channel_threshold=int(cht),
+        stage_infos=sis,
+        pcfg_enabled=bool(p_en), pcfg_steps=int(p_steps),
+        pcfg_mode=str(p_mode), pcfg_blend=float(p_bl), pcfg_b=float(p_b),
+        pcfg_fourier=bool(p_four), pcfg_ms_mode=str(p_mmd),
+        pcfg_ms_str=float(p_mst), pcfg_threshold=int(p_thr),
+        pcfg_s=float(p_s), pcfg_gain=float(p_gain), verbose=bool(v_log),
+    )
+
+# Simulate _apply_p
+def _apply_p(name):
+    p = GS.all_presets.get(name)
+    if p is None: return None
+    flat = []
+    for si in p.stage_infos:
+        for f in _SF: flat.append(getattr(si, f))
+    return {
+        "start_ratio": p.start_ratio, "stop_ratio": p.stop_ratio,
+        "smooth": p.transition_smoothness,
+        "version": GS.REVERSED_VERSIONS.get(p.version, "Version 2"),
+        "ms_mode": p.multiscale_mode, "ms_str": p.multiscale_strength,
+        "ov_scales": p.override_scales, "ch_thresh": p.channel_threshold,
+        "pcfg_en": p.pcfg_enabled, "pcfg_steps": p.pcfg_steps,
+        "pcfg_mode": p.pcfg_mode, "pcfg_bl": p.pcfg_blend,
+        "pcfg_b": p.pcfg_b, "pcfg_fou": p.pcfg_fourier,
+        "pcfg_mmd": p.pcfg_ms_mode, "pcfg_mst": p.pcfg_ms_str,
+        "pcfg_thr": p.pcfg_threshold, "pcfg_s": p.pcfg_s,
+        "pcfg_gain": p.pcfg_gain, "verbose": p.verbose,
+        "flat": flat,
+    }
+
+print("═"*52)
+print("TEST: Full preset save/apply round-trip")
+print("═"*52)
+
+# Build flat for 3 stages
+si0 = GS.StageInfo(backbone_factor=1.3, skip_factor=0.8)
+flat_in = []
+for si in [si0, GS.StageInfo(), GS.StageInfo()]:
+    for f in _SF: flat_in.append(getattr(si, f))
+
+# Save
+_save_p(
+    "full_preset",
+    0.1, 0.9, 0.5, "Version 2",
+    "Multi-Bandpass", 0.7, "10,1.5", 64,
+    True, 15, "lerp", 0.8, 1.4,
+    True, "Default", 0.9, 3, 0.6, 1.2,
+    True,
+    *flat_in
+)
+
+chk("preset saved", "full_preset" in GS.all_presets)
+saved = GS.all_presets["full_preset"]
+chk("saved pcfg_enabled=True", saved.pcfg_enabled == True)
+chk("saved pcfg_steps=15",     saved.pcfg_steps == 15)
+chk("saved pcfg_mode=lerp",    saved.pcfg_mode == "lerp")
+chk("saved pcfg_blend=0.8",    abs(saved.pcfg_blend - 0.8) < 1e-5)
+chk("saved pcfg_b=1.4",        abs(saved.pcfg_b - 1.4) < 1e-5)
+chk("saved pcfg_fourier=True", saved.pcfg_fourier == True)
+chk("saved pcfg_ms_mode=Default", saved.pcfg_ms_mode == "Default")
+chk("saved pcfg_ms_str=0.9",   abs(saved.pcfg_ms_str - 0.9) < 1e-5)
+chk("saved pcfg_threshold=3",  saved.pcfg_threshold == 3)
+chk("saved pcfg_s=0.6",        abs(saved.pcfg_s - 0.6) < 1e-5)
+chk("saved pcfg_gain=1.2",     abs(saved.pcfg_gain - 1.2) < 1e-5)
+chk("saved verbose=True",      saved.verbose == True)
+chk("saved ms_mode=Multi-Bandpass", saved.multiscale_mode == "Multi-Bandpass")
+chk("saved ms_str=0.7",        abs(saved.multiscale_strength - 0.7) < 1e-5)
+chk("saved ch_thresh=64",      saved.channel_threshold == 64)
+chk("saved bf=1.3",            abs(saved.stage_infos[0].backbone_factor - 1.3) < 1e-5)
+
+# Apply
+restored = _apply_p("full_preset")
+chk("apply: pcfg_en=True",    restored["pcfg_en"] == True)
+chk("apply: pcfg_steps=15",   restored["pcfg_steps"] == 15)
+chk("apply: pcfg_mode=lerp",  restored["pcfg_mode"] == "lerp")
+chk("apply: pcfg_b=1.4",      abs(restored["pcfg_b"] - 1.4) < 1e-5)
+chk("apply: pcfg_fourier=True", restored["pcfg_fou"] == True)
+chk("apply: verbose=True",    restored["verbose"] == True)
+chk("apply: ms_mode restored", restored["ms_mode"] == "Multi-Bandpass")
+chk("apply: ch_thresh=64",    restored["ch_thresh"] == 64)
+chk("apply: bf=1.3 in flat",  abs(restored["flat"][0] - 1.3) < 1e-5)
+
+# Save with pcfg disabled → defaults
+_save_p(
+    "no_pcfg",
+    0.0, 1.0, 0.0, "Version 1",
+    "Default", 1.0, "", 96,
+    False, 20, "inject", 1.0, 1.1,
+    False, "Default", 1.0, 1, 0.5, 1.0,
+    False,
+    *flat_in
+)
+r2 = _apply_p("no_pcfg")
+chk("no_pcfg: pcfg_en=False",  r2["pcfg_en"] == False)
+chk("no_pcfg: verbose=False",  r2["verbose"] == False)
+
+# Unknown preset → None
+r3 = _apply_p("does_not_exist")
+chk("unknown preset → None",   r3 is None)
+
+# dict-API branch: pcfg_enabled passed through dict
+GS.instance = GS.State()
+d_api = {"enable": True, "pcfg_enabled": True, "pcfg_b": 1.9, "verbose": True}
+fields = {f.name for f in dataclasses.fields(GS.State)}
+GS.instance = GS.State(**{k:v for k,v in d_api.items() if k in fields})
+chk("dict-API: pcfg_enabled propagated", GS.instance.pcfg_enabled == True)
+chk("dict-API: pcfg_b=1.9",  abs(GS.instance.pcfg_b - 1.9) < 1e-5)
+chk("dict-API: verbose=True", GS.instance.verbose == True)
+
+print(f"\n{'═'*52}")
+print(f"PRESET PCFG ROUND-TRIP: {P} PASS   {F} FAIL")
+if ERRS:
+    print("\nFailed:")
+    for e in ERRS: print(f"  • {e}")
+else:
+    print("ALL PASSED ✓")