ComfyUI-RK-Sampler

Batched Runge-Kutta Samplers for ComfyUI

Supports most practical Explicit Runge-Kutta (ERK) methods.

Tested on SD1.5, SDXL, and SD3.

Features

• Parallel ODE solvers for fast batch processing
• Explicit and Embedded Explicit Runge-Kutta methods
• PID controller for adaptive step sizing with tunable settings
• Scheduled controller for fixed step sizing (determined by the given $\sigma$ schedule)

Installation

ComfyUI-Manager

ComfyUI Manager Menu > Custom Nodes Manager > ComfyUI-RK-Sampler > Install

Manual installation

From ComfyUI/custom_nodes and ComfyUI virtual environment:

git clone https://github.com/wootwootwootwoot/ComfyUI-RK-Sampler.git
pip install torchode

Usage

Basic workflow

From Add Node:

sampling > custom_sampling > samplers > Runge-Kutta Sampler

Best defaults

• These methods support normal as well as high CFG scales, but the results may depend on the specific model.
• For SDXL, the best results I got were from CFG scales between 7-35.
• If you don't know the right step count or CFG scale:
  • Try the adaptive_pid controller with the base CFG and increment it until the results get worse.
  • Tune the CFG to the desired output.
  • Try the fixed_scheduled controller:
    • Use the Align Your Steps scheduler.
    • Set the scheduler's step count to be the same as the number of steps taken by the adaptive_pid controller.
    • Leave the CFG scale unchanged.
  • Tune the scheduler step count.
• When using adaptive_pid or adaptive_scipy controllers:
  • Set log_absolute_tolerance to -3.5:
    • Decrease it to set more strict tolerances (better results) in exchange for slower inference times.
    • Increase it to set less strict tolerances (worse results) in exchange for faster inference times.
  • Set log_relative_tolerance to be 1-2 more than log_absolute_tolerance.

Adaptive step size

method: ae_bosh3
step_size_controller: adaptive_pid
log_absolute_tolerance: -3.5
log_relative_tolerance: -2.5
pcoeff: 0
icoeff: 1
dcoeff: 0
norm: rms_norm
enable_dt_min: false
enable_dt_max: true
dt_min: -0.1
dt_max: 0
safety: 0.9
factor_min: 0.2
factor_max: 10
max_steps: 2147483647
min_sigma: 1e-5
cfg: 7-35

Fixed step size

method: fe_ralston3
step_size_controller: fixed_scheduled
scheduler: Align Your Steps
steps: 28-150
cfg: 7-35

Choose a step size controller

step_size_controller: Controller to determine the step size taken on each sampling step.

adaptive controllers: Automatically determines the step size(s). The scheduler choice and scheduler step count does not matter since they only use the start and end timesteps.

fixed controllers: Uses the step size(s) provided by the scheduler. Works like a normal non-adaptive sampler in this case.

• adaptive_pid: A proportional–integral–derivative (PID) controller. Works with a-class methods.
• fixed_scheduled: A controller that uses the $\sigma$ (timestep) schedule from the scheduler. Works with a-class and f-class methods.
• adaptive_scipy: A basic integral controller wrapped from scipy. Works with s-class methods.

Choose a method

method: Determines the solver method used.

a, f, and s classes

• a = adaptive, f = fixed, s = scipy, e = explicit
• Use a-class methods with either:
  • adaptive_pid for automatically determined step sizes/count.
  • fixed_scheduled for scheduler determined step sizes/count.
• Use f-class methods with fixed_scheduled.
• Use s-class methods with adaptive_scipy.

Recommended methods

• Try ae_bosh3, ae_dopri5, and ae_fehlberg5 with the adaptive_pid step size controller.
• Try fe_ralston3, ae_bosh3, and fe_ssprk3 with the fixed_scheduled step size controller.

Quality ranking

Tested on RTX3090, SDXL, 896x1152, CFG=30, batch size 1, fixed_scheduled, Align Your Steps, 28 steps. | Rank | Name | Method | Order | NFEs | Time | | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | | 1 | fe_ralston3 | Ralston | 3 | 3 | 23.08s | | 2 | ae_bosh3 | Bogacki–Shampine | 3 | 3 | 23.03s | | 3 | fe_ssprk3 | Strong Stability Preserving Runge-Kutta | 3 | 3 | 22.97s | | 4 | fe_kutta4 | Runge-Kutta | 4 | 4 | 29.87s | | 5 | fe_kutta_38th4 | Runge-Kutta (3/8-rule) | 4 | 4 | 30.10s | | 6 | ae_dopri5 | Dormand–Prince | 5 | 6 | 44.03s | | 7 | ae_fehlberg5 | Runge–Kutta–Fehlberg | 5 | 6 | 44.28s | | 8 | ae_heun_euler2 | Heun–Euler | 2 | 2 | 15.87s | | 9 | fe_kutta3 | Runge-Kutta | 3 | 3 | 22.83s | | 10 | ae_ralston2 | Ralston | 2 | 2 | 16.05s | | 11 | ae_cash_karp5 | Cash-Karp | 5 | 6 | 44.47s | | 12 | fe_ralston4 | Ralston | 4 | 4 | 29.69s | | 13 | fe_euler1 | Forward Euler | 1 | 1 | 9.00s | | 14 | fe_wray3 | van der Houwen and Wray | 3 | 3 | 22.96s | | 15 | fe_heun3 | Heun | 3 | 3 | 23.13s | | 16 | ae_tsit5 | Tsitouras | 5 | 6 | 43.90s | | 17 | ae_midpoint2 | Explicit midpoint | 2 | 2 | 15.97s | | 18 | ae_fehlberg2 | Runge–Kutta–Fehlberg | 2 | 3 | 23.23s | | 19 | ae_dopri8 | Dormand–Prince | 8 | 13 | 94.10s |

Explicit methods from scipy.integrate

These methods are wrapped implementations of explicit solvers from scipy.integrate. | Name | Method | Order | NFEs | | ----------- | ----------- | ----------- | ----------- | | se_RK23 | Runge-Kutta | 3 | 3 | | se_RK45 | Runge-Kutta | 5 | 6 | | se_DOP853 | Dormand-Prince | 8 | 12 |

Solver settings

| Option | Applies to | Description | | ----------- | ----------- | ----------- | | method | adaptive_pid, fixed_scheduled, adaptive_scipy | Solver method. | | step_size_controller | adaptive_pid, fixed_scheduled, adaptive_scipy | Step size controller. | | log_absolute_tolerance | adaptive_pid, adaptive_scipy | $log_{10}$ of the threshold below which the solver does not worry about the accuracy of the solution since it is effectively 0. log_absolute_tolerance must be less than or equal to log_relative_tolerance. | | log_relative_tolerance | adaptive_pid, adaptive_scipy | $log_{10}$ of the threshold for the relative error of a single step of the integrator. | | pcoeff | adaptive_pid | Coefficients for the proportional term of the PID controller. | | icoeff | adaptive_pid | Coefficients for the integral term of the PID controller. P/I/D of 0/1/0 corresponds to a basic integral controller. | | dcoeff | adaptive_pid | Coefficients for the derivative term of the PID controller. | | norm | adaptive_pid, adaptive_scipy | Normalization function for error control. Step sizes are chosen so that norm(error / (absolute_tolerance + relative_tolerance * y)) is approximately one. | | enable_dt_min | adaptive_pid, adaptive_scipy | Enable clamping of the minimum step size to take to dt_min. | | enable_dt_max | adaptive_pid, adaptive_scipy | Enable clamping of the maximum step size to take to dt_max. | | dt_min | adaptive_pid, adaptive_scipy | The dt_min value to clamp to. Since we are solving a reverse-time ODE, this value should be negative. dt_min must be less than or equal to dt_max. | | dt_max | adaptive_pid, adaptive_scipy | The dt_max value to clamp to. Since we are solving a reverse-time ODE, this value should be negative. Clamped to 0 by default to force a monotonic solve. | | safety | adaptive_pid, adaptive_scipy | Multiplicative safety factor. | | factor_min | adaptive_pid, adaptive_scipy | Minimum amount a step size can be decreased relative to the previous step. factor_min must be less than or equal to factor_max. | | factor_max | adaptive_pid, adaptive_scipy | Maximum amount a step size can be increased relative to the previous step. | | max_steps | adaptive_pid, adaptive_scipy | Maximum amount of steps an adaptive step size controller is allowed to take. Taking more steps than max_steps will return an error. | | min_sigma | adaptive_pid, adaptive_scipy | Lower bound for $\sigma$ to consider the IVP solve to be complete. |

Comparison

Related projects

This extension improves upon ComfyUI-ODE by adding support for parallel processing, more controllability, high-quality live previews, a PID controller, and support for more fixed and adaptive step size solvers.

Speed

Tested on RTX3090, SDXL, 896x1152, CFG=30, adaptive_pid 0/1/0, ae_bosh3 (ComfyUI-RK-Sampler) / bosh3 (ComfyUI-ODE), log_absolute_tolerance=-3.5, log_relative_tolerance=-2.5. | Batch size | ComfyUI-RK-Sampler | ComfyUI-ODE | | ----------- | ----------- | ----------- | | 1 | 1m2s | 1m6s | | 2 | 2m3s | 2m7s | | 4 | 4m4s | 4m27s | | 8 | 7m39s | 9m3s | | 16 | 15m29s | 17m22s | | 32 | 30m18s | 36m34s |

Explanation

From Wikipedia:

In numerical analysis, the Runge–Kutta methods are a family of implicit and explicit iterative methods, which include the Euler method, used in temporal discretization for the approximate solutions of simultaneous nonlinear equations. These methods were developed around 1900 by the German mathematicians Carl Runge and Wilhelm Kutta.

The Runge-Kutta methods are a family of methods used for solving approximate solutions of ODEs by iterative discretization (or, if in diffusion terms, by sampling).

Runge-Kutta methods generally have less discretization error than standard diffusion sampling methods, allowing for the use of high CFG scales (within practical limits) to create high-quality results without artifacts.

Conditions for new solvers and methods

• No implicit solvers:
  • Implicit solvers require root-finding.
  • This means a backward pass through the model has to be ran to compute the Jacobian.
  • The Jacobian is then used during LBFGS optimization.
  • This is also very costly in terms of compute and memory usage.
• No non-RK methods:
  • This would cut off linear multistep methods like Adams-Bashforth or Adams predictor-corrector.
  • From my testing, RK methods perform better, even for explicit RK vs. predictor-corrector linear multistep.
• No duplicate methods:
  • If two methods have different coefficients, they aren't duplicates.
  • SciPy methods have different coefficients and implementations.

Changelog

24/07/24

• Added solver settings for adaptive_scipy.
• Progress bar improvements:
  • Adaptive solvers now show the number of steps taken.
  • Accurate $\sigma$ info is now displayed.
• Bugfixes and small refactors.

23/07/24

• Installation from ComfyUI-Manager.
• Bugfixes and small refactors.

22/07/24

• Added wrappers for explicit solvers from scipy.integrate:
  • se_RK23
  • se_RK45
  • se_DOP853
• Some notes about the wrapped scipy solvers:
  • To use the new solvers, select adaptive_scipy as the step size controller.
  • No fixed step size controllers are available for wrapped scipy methods.
  • These methods do not support parallel IVP solve, meaning the batch elements are processed sequentially.
  • Implicit solvers from scipy.integrate were skipped as the root finding step takes too long.
• The refactors from this update can cause the results to differ from the previous version. This is due to changes in floating point precision casting and operation orders.