Upload adept_scheduler_pack.py

adept-sampler-v6/scripts/adept_scheduler_pack.py

@@ -0,0 +1,150 @@
+"""
+Adept Scheduler Pack for Automatic1111 WebUI
+Registers Adept Sampler's scheduler algorithms into A1111's native
+"Schedule type" dropdown (the one next to "Sampling method"), so they
+become usable with ANY sampler -- not only when the Adept Sampler script
+itself is active via its own "Scheduler Type" dropdown.
+
+Companion file to adept_sampler_v5.py: it must be installed in the same
+scripts/ folder. This file reads Adept's scheduler functions directly off
+the loaded adept_sampler_v5 module (no duplicated formulas), so any future
+tweak to those functions is picked up automatically -- nothing here needs
+updating when Adept's own schedulers change.
+
+If adept_sampler_v5.py isn't found (not installed, or failed to load for
+some other reason), this file logs a message and does nothing further; it
+never raises, so a missing companion file can't break WebUI startup.
+"""
+
+import os
+import torch
+from modules import scripts, script_callbacks
+import modules.sd_schedulers as sd_schedulers
+
+_ADEPT_SCRIPT_FILENAME = "adept_sampler_v5.py"
+
+# (label shown in the Schedule-type dropdown, internal name, function name on
+# the adept_sampler_v5 module). Mirrors apply_custom_scheduler's own
+# scheduler_map in adept_sampler_v5.py exactly -- kept as literal strings
+# (not imported) so this file has no hard import-time dependency on that
+# module actually existing.
+_ADEPT_SCHEDULER_SPECS = [
+    ("Adept: AOS-V",               "adept_aos_v",               "create_aos_v_sigmas"),
+    ("Adept: AOS-Epsilon",         "adept_aos_epsilon",         "create_aos_e_sigmas"),
+    ("Adept: AkashicAOS",          "adept_akashic_aos",         "create_aos_akashic_sigmas"),
+    ("Adept: Entropic",            "adept_entropic",            "create_entropic_sigmas"),
+    ("Adept: SNR-Optimized",       "adept_snr_optimized",       "create_snr_optimized_sigmas"),
+    ("Adept: Constant-Rate",       "adept_constant_rate",       "create_constant_rate_sigmas"),
+    ("Adept: Adaptive-Optimized",  "adept_adaptive_optimized",  "create_adaptive_optimized_sigmas"),
+    ("Adept: Cosine-Annealed",     "adept_cosine_annealed",     "create_cosine_sigmas"),
+    ("Adept: LogSNR-Uniform",      "adept_logsnr_uniform",      "create_logsnr_uniform_sigmas"),
+    ("Adept: Tanh Mid-Boost",      "adept_tanh_midboost",       "create_tanh_midboost_sigmas"),
+    ("Adept: Exponential Tail",    "adept_exponential_tail",    "create_exponential_tail_sigmas"),
+    ("Adept: Jittered-Karras",     "adept_jittered_karras",     "create_jittered_karras_sigmas"),
+    ("Adept: Stochastic",          "adept_stochastic",          "create_stochastic_sigmas"),
+    ("Adept: JYS (Dynamic)",       "adept_jys",                 "create_jys_sigmas"),
+    ("Adept: Hybrid JYS-Karras",   "adept_hybrid_jys_karras",   "create_hybrid_jys_karras_sigmas"),
+    ("Adept: AYS-SDXL",            "adept_ays_sdxl",            "create_ays_sdxl_sigmas"),
+    ("Adept: AkashicAOS Alt",      "adept_akashic_aos_alt",     "create_aos_akashic_alt_sigmas"),
+    ("Adept: AkashicEQFlow",       "adept_akashic_eqflow",      "create_akashic_eqflow_sigmas"),
+]
+
+
+def _find_adept_module():
+    """Locate the already-loaded adept_sampler_v5 module via A1111's own
+    script registry (scripts.scripts_data), the same technique Adept itself
+    uses to find xyz_grid.py. Deferred to on_before_ui time so load order
+    between the two files never matters -- every script's top-level code
+    has already run by the time on_before_ui callbacks fire."""
+    for data in scripts.scripts_data:
+        if os.path.basename(data.path) == _ADEPT_SCRIPT_FILENAME:
+            return data.module
+    return None
+
+
+def _make_adapter(adept_fn, label):
+    """
+    Wrap one of Adept's create_*_sigmas(sigma_max, sigma_min, num_steps,
+    device, ...) functions to match A1111's native scheduler signature:
+    function(n, sigma_min, sigma_max, device, **extra). A1111 calls this
+    with keyword arguments (n=steps, sigma_min=..., sigma_max=...,
+    device=...), never positionally, so only the parameter *names* need to
+    line up -- the underlying function's extra optional kwargs (Entropic's
+    `power`, Stochastic's `noise_type`, etc.) keep their own defaults since
+    A1111 has no UI for them and won't pass them.
+    """
+    def adapter(n, sigma_min, sigma_max, device='cpu', **_ignored):
+        try:
+            # A1111's native scheduler API passes sigma_min/sigma_max as
+            # plain Python floats (via .item()); Adept's own internal call
+            # path (apply_custom_scheduler) always passes tensor elements
+            # instead. A few of Adept's scheduler functions call torch.log()
+            # directly on these arguments, which requires a tensor -- so we
+            # convert here, at the compatibility boundary, rather than
+            # touching the already-shipped, tested adept_sampler_v5.py.
+            sigma_min_t = torch.as_tensor(sigma_min, dtype=torch.float32, device=device)
+            sigma_max_t = torch.as_tensor(sigma_max, dtype=torch.float32, device=device)
+            result = adept_fn(sigma_max_t, sigma_min_t, n, device=device)
+            if result is None or len(result) != n + 1:
+                raise ValueError(f"unexpected shape from {label}")
+            if torch.isnan(result).any() or torch.isinf(result).any():
+                raise ValueError(f"NaN/Inf from {label}")
+            return result
+        except Exception as e:
+            print(f"⚠️ Adept Scheduler Pack: '{label}' failed ({e}), falling back to Karras")
+            import k_diffusion.sampling as k_diff
+            return k_diff.get_sigmas_karras(n, sigma_min, sigma_max, device=device)
+    return adapter
+
+
+def register_adept_schedulers():
+    adept_mod = _find_adept_module()
+    if adept_mod is None:
+        print("⚠️ Adept Scheduler Pack: adept_sampler_v5.py not found -- "
+              "install it alongside this file to enable Adept schedulers "
+              "in the native Schedule type dropdown.")
+        return
+
+    # Dedup guard: safe to call more than once (e.g. UI reload).
+    if any(getattr(s, "name", "").startswith("adept_") for s in sd_schedulers.schedulers):
+        return
+
+    registered = 0
+    for label, name, fn_name in _ADEPT_SCHEDULER_SPECS:
+        adept_fn = getattr(adept_mod, fn_name, None)
+        if adept_fn is None:
+            print(f"⚠️ Adept Scheduler Pack: {fn_name} not found on adept_sampler_v5 "
+                  f"(version mismatch?), skipping '{label}'")
+            continue
+
+        sched = sd_schedulers.Scheduler(
+            name=name,
+            label=label,
+            function=_make_adapter(adept_fn, label),
+            default_rho=-1,
+            need_inner_model=False,
+        )
+        sd_schedulers.schedulers.append(sched)
+        # schedulers_map is built once (as a dict comprehension) right after
+        # the schedulers list at sd_schedulers.py's own import time -- it
+        # does NOT update itself when the list is appended to afterwards.
+        # Both the name and label must be added, matching how the map is
+        # originally built, or lookups by either key at generation time
+        # would fail to find these newly-registered entries.
+        sd_schedulers.schedulers_map[sched.name] = sched
+        sd_schedulers.schedulers_map[sched.label] = sched
+        registered += 1
+
+    print(f"✅ Adept Scheduler Pack: registered {registered}/{len(_ADEPT_SCHEDULER_SPECS)} "
+          f"schedulers into the native Schedule type dropdown")
+
+
+def on_before_ui():
+    try:
+        register_adept_schedulers()
+    except Exception:
+        import traceback
+        print(f"⚠️ Adept Scheduler Pack: registration failed:\n{traceback.format_exc()}")
+
+
+script_callbacks.on_before_ui(on_before_ui)