add sampler api endpoints

Signed-off-by: vladmandic <mandic00@live.com>
2026-09-19 01:04:32 +02:00 · 2026-02-04 13:08:31 +01:00
parent d7ca4f63a7
commit d9a2a21c8c
14 changed files with 160 additions and 32 deletions
@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+
+"""
+get list of all samplers and details of current sampler
+"""
+
+import sys
+import logging
+import urllib3
+import requests
+
+
+url = "http://127.0.0.1:7860"
+user = ""
+password = ""
+
+log_format = '%(asctime)s %(levelname)s: %(message)s'
+logging.basicConfig(level = logging.INFO, format = log_format)
+log = logging.getLogger("sd")
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+log.info('available samplers')
+auth = requests.auth.HTTPBasicAuth(user, password) if len(user) > 0 and len(password) > 0 else None
+req = requests.get(f'{url}/sdapi/v1/samplers', verify=False, auth=auth, timeout=60)
+if req.status_code != 200:
+    log.error({ 'url': req.url, 'request': req.status_code, 'reason': req.reason })
+    exit(1)
+res = req.json()
+for item in res:
+    log.info(item)
+
+log.info('current sampler')
+req = requests.get(f'{url}/sdapi/v1/sampler', verify=False, auth=auth, timeout=60)
+res = req.json()
+log.info(res)
@@ -0,0 +1,262 @@
+import os
+import sys
+import time
+import numpy as np
+import torch
+
+# Ensure we can import modules
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../")))
+
+from modules.errors import log
+from modules.res4lyf import (
+    BASE, SIMPLE, VARIANTS,
+    RESUnifiedScheduler, RESMultistepScheduler, RESDEISMultistepScheduler,
+    ETDRKScheduler, LawsonScheduler, ABNorsettScheduler, PECScheduler,
+    RiemannianFlowScheduler, RESSinglestepScheduler, RESSinglestepSDEScheduler,
+    RESMultistepSDEScheduler, SimpleExponentialScheduler, LinearRKScheduler,
+    LobattoScheduler, GaussLegendreScheduler, RungeKutta44Scheduler,
+    RungeKutta57Scheduler, RungeKutta67Scheduler, SpecializedRKScheduler,
+    BongTangentScheduler, CommonSigmaScheduler, RadauIIAScheduler,
+    LangevinDynamicsScheduler
+)
+from modules.schedulers.scheduler_vdm import VDMScheduler
+from modules.schedulers.scheduler_unipc_flowmatch import FlowUniPCMultistepScheduler
+from modules.schedulers.scheduler_ufogen import UFOGenScheduler
+from modules.schedulers.scheduler_tdd import TDDScheduler
+from modules.schedulers.scheduler_tcd import TCDScheduler
+from modules.schedulers.scheduler_flashflow import FlashFlowMatchEulerDiscreteScheduler
+from modules.schedulers.scheduler_dpm_flowmatch import FlowMatchDPMSolverMultistepScheduler
+from modules.schedulers.scheduler_dc import DCSolverMultistepScheduler
+from modules.schedulers.scheduler_bdia import BDIA_DDIMScheduler
+
+def test_scheduler(name, scheduler_class, config):
+    try:
+        scheduler = scheduler_class(**config)
+    except Exception as e:
+        log.error(f'scheduler="{name}" cls={scheduler_class} config={config} error="Init failed: {e}"')
+        return False
+
+    num_steps = 20
+    scheduler.set_timesteps(num_steps)
+
+    sample = torch.randn((1, 4, 64, 64))
+    has_changed = False
+    t0 = time.time()
+    messages = []
+
+    try:
+        for i, t in enumerate(scheduler.timesteps):
+            # Simulate model output (noise or x0 or v), Using random noise for stability check
+            model_output = torch.randn_like(sample)
+
+            # Scaling Check
+            step_idx = scheduler.step_index if hasattr(scheduler, "step_index") and scheduler.step_index is not None else i
+            # Clamp index
+            if hasattr(scheduler, 'sigmas'):
+                step_idx = min(step_idx, len(scheduler.sigmas) - 1)
+                sigma = scheduler.sigmas[step_idx]
+            else:
+                sigma = torch.tensor(1.0) # Dummy for non-sigma schedulers
+
+            # Re-introduce scaling calculation first
+            scaled_sample = scheduler.scale_model_input(sample, t)
+
+            if config.get("prediction_type") == "flow_prediction" or name in ["UFOGenScheduler", "TDDScheduler", "TCDScheduler", "BDIA_DDIMScheduler", "DCSolverMultistepScheduler"]:
+                # Some new schedulers don't use K-diffusion scaling
+                expected_scale = 1.0
+            else:
+                expected_scale = 1.0 / ((sigma**2 + 1) ** 0.5)
+
+            # Simple check with loose tolerance due to float precision
+            expected_scaled_sample = sample * expected_scale
+            if not torch.allclose(scaled_sample, expected_scaled_sample, atol=1e-4):
+                # If failed, double check if it's just 'sample' (no scaling)
+                if torch.allclose(scaled_sample, sample, atol=1e-4):
+                    messages.append('warning="scaling is identity"')
+                else:
+                    log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} expected={expected_scale} error="scaling mismatch"')
+                    return False
+
+            if torch.isnan(scaled_sample).any():
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="NaN in scaled_sample"')
+                return False
+
+            if torch.isinf(scaled_sample).any():
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Inf in scaled_sample"')
+                return False
+
+            output = scheduler.step(model_output, t, sample)
+
+            # Shape and Dtype check
+            if output.prev_sample.shape != sample.shape:
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Shape mismatch: {output.prev_sample.shape} vs {sample.shape}"')
+                return False
+            if output.prev_sample.dtype != sample.dtype:
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Dtype mismatch: {output.prev_sample.dtype} vs {sample.dtype}"')
+                return False
+
+            # Update check: Did the sample change?
+            if not torch.equal(sample, output.prev_sample):
+                has_changed = True
+
+            # Sample Evolution Check
+            step_diff = (sample - output.prev_sample).abs().mean().item()
+            if step_diff < 1e-6:
+                messages.append(f'warning="minimal sample change: {step_diff}"')
+
+            sample = output.prev_sample
+
+            if torch.isnan(sample).any():
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="NaN in sample"')
+                return False
+
+            if torch.isinf(sample).any():
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="Inf in sample"')
+                return False
+
+            # Divergence check
+            if sample.abs().max() > 1e10:
+                log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} step={i} error="divergence detected"')
+                return False
+
+        # External check for Sigma Monotonicity
+        if hasattr(scheduler, 'sigmas'):
+            sigmas = scheduler.sigmas.cpu().numpy()
+            if len(sigmas) > 1:
+                diffs = np.diff(sigmas) # Check if potentially monotonic decreasing (standard) OR increasing (some flow/inverse setups). We allow flat sections (diff=0) hence 1e-6 slack
+                is_monotonic_decreasing = np.all(diffs <= 1e-6)
+                is_monotonic_increasing = np.all(diffs >= -1e-6)
+                if not (is_monotonic_decreasing or is_monotonic_increasing):
+                    messages.append('warning="sigmas are not monotonic"')
+
+    except Exception as e:
+        log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} exception: {e}')
+        import traceback
+        traceback.print_exc()
+        return False
+
+    if not has_changed:
+        log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} error="sample never changed"')
+        return False
+
+    final_std = sample.std().item()
+    if final_std > 50.0 or final_std < 0.1:
+        log.error(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} std={final_std} error="variance drift"')
+
+    t1 = time.time()
+    messages = list(set(messages))
+    log.info(f'scheduler="{name}" cls={scheduler.__class__.__name__} config={config} time={t1-t0} messages={messages}')
+    return True
+
+def run_tests():
+    prediction_types = ["epsilon", "v_prediction", "sample"] # flow_prediction is special, usually requires flow sigmas or specific setup, checking standard ones first
+
+    # Test BASE schedulers with their specific parameters
+    log.warning('type="base"')
+    for name, cls in BASE:
+        configs = []
+
+        # prediction_types
+        for pt in prediction_types:
+            configs.append({"prediction_type": pt})
+
+        # Specific params for specific classes
+        if cls == RESUnifiedScheduler:
+            rk_types = ["res_2m", "res_3m", "res_2s", "res_3s", "res_5s", "res_6s", "deis_1s", "deis_2m", "deis_3m"]
+            for rk in rk_types:
+                for pt in prediction_types:
+                    configs.append({"rk_type": rk, "prediction_type": pt})
+
+        elif cls == RESMultistepScheduler:
+            variants = ["res_2m", "res_3m", "deis_2m", "deis_3m"]
+            for v in variants:
+                for pt in prediction_types:
+                    configs.append({"variant": v, "prediction_type": pt})
+
+        elif cls == RESDEISMultistepScheduler:
+            for order in range(1, 6):
+                for pt in prediction_types:
+                    configs.append({"solver_order": order, "prediction_type": pt})
+
+        elif cls == ETDRKScheduler:
+            variants = ["etdrk2_2s", "etdrk3_a_3s", "etdrk3_b_3s", "etdrk4_4s", "etdrk4_4s_alt"]
+            for v in variants:
+                for pt in prediction_types:
+                    configs.append({"variant": v, "prediction_type": pt})
+
+        elif cls == LawsonScheduler:
+            variants = ["lawson2a_2s", "lawson2b_2s", "lawson4_4s"]
+            for v in variants:
+                for pt in prediction_types:
+                    configs.append({"variant": v, "prediction_type": pt})
+
+        elif cls == ABNorsettScheduler:
+            variants = ["abnorsett_2m", "abnorsett_3m", "abnorsett_4m"]
+            for v in variants:
+                for pt in prediction_types:
+                    configs.append({"variant": v, "prediction_type": pt})
+
+        elif cls == PECScheduler:
+            variants = ["pec423_2h2s", "pec433_2h3s"]
+            for v in variants:
+                for pt in prediction_types:
+                    configs.append({"variant": v, "prediction_type": pt})
+
+        elif cls == RiemannianFlowScheduler:
+            metrics = ["euclidean", "hyperbolic", "spherical", "lorentzian"]
+            for m in metrics:
+                configs.append({"metric_type": m, "prediction_type": "epsilon"}) # Flow usually uses v or raw, but epsilon check matches others
+
+        if not configs:
+            for pt in prediction_types:
+                configs.append({"prediction_type": pt})
+
+        for conf in configs:
+            test_scheduler(name, cls, conf)
+
+    log.warning('type="simple"')
+    for name, cls in SIMPLE:
+        for pt in prediction_types:
+            test_scheduler(name, cls, {"prediction_type": pt})
+
+    log.warning('type="variants"')
+    for name, cls in VARIANTS:
+        # these classes preset their variants/rk_types in __init__ so we just test prediction types
+        for pt in prediction_types:
+            test_scheduler(name, cls, {"prediction_type": pt})
+
+    # Extra robustness check: Flow Prediction Type
+    log.warning('type="flow"')
+    flow_schedulers = [
+        # res4lyf schedulers
+        RESUnifiedScheduler, RESMultistepScheduler, ABNorsettScheduler,
+        RESSinglestepScheduler, RESSinglestepSDEScheduler, RESDEISMultistepScheduler,
+        RESMultistepSDEScheduler, ETDRKScheduler, LawsonScheduler, PECScheduler,
+        SimpleExponentialScheduler, LinearRKScheduler, LobattoScheduler,
+        GaussLegendreScheduler, RungeKutta44Scheduler, RungeKutta57Scheduler,
+        RungeKutta67Scheduler, SpecializedRKScheduler, BongTangentScheduler,
+        CommonSigmaScheduler, RadauIIAScheduler, LangevinDynamicsScheduler,
+        RiemannianFlowScheduler,
+        # sdnext schedulers
+        FlowUniPCMultistepScheduler, FlashFlowMatchEulerDiscreteScheduler, FlowMatchDPMSolverMultistepScheduler,
+    ]
+    for cls in flow_schedulers:
+        test_scheduler(cls.__name__, cls, {"prediction_type": "flow_prediction", "use_flow_sigmas": True})
+
+    log.warning('type="sdnext"')
+    extended_schedulers = [
+        VDMScheduler,
+        UFOGenScheduler,
+        TDDScheduler,
+        TCDScheduler,
+        DCSolverMultistepScheduler,
+        BDIA_DDIMScheduler
+    ]
+    for cls in extended_schedulers:
+        # Most of these support standard prediction types, try epsilon as default safest bet
+        # Some might be flow matching specific, we can try robust default list
+        # For now, just test default init
+        test_scheduler(cls.__name__, cls, {"prediction_type": "epsilon"})
+
+if __name__ == "__main__":
+    run_tests()