blepping
38fe739a7e
Add pingpong group merge method. Remove model call caching stuff. Allow defining pre_cfg and post_cfg filters. Handle filters changing the latent shape better. Allow disabling detecting out of order sigmas as restart sampling.
30 lines
930 B
Python
30 lines
930 B
Python
from .base import SingleStepSampler, registry
|
|
|
|
|
|
# Referenced from https://github.com/ace-step/ACE-Step/
|
|
class PingPongStep(SingleStepSampler):
|
|
name = "pingpong"
|
|
default_eta = 0.0
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
pingpong_options = self.options.pop("pingpong", {})
|
|
self.pingpong_start_step = pingpong_options.get("start_step", 0)
|
|
self.pingpong_end_step = pingpong_options.get("end_step", 0)
|
|
|
|
def step(self, x):
|
|
ss = self.ss
|
|
use_pingpong = self.pingpong_start_step <= ss.step <= self.pingpong_end_step
|
|
if not use_pingpong:
|
|
return (yield from self.euler_step(x, eta=0.0))
|
|
sn = ss.sigma_next
|
|
denoised = (
|
|
ss.denoised * (1.0 - sn) if ss.model.is_rectified_flow else ss.denoised
|
|
)
|
|
yield from self.result(denoised, sn, sigma_down=sn)
|
|
|
|
|
|
registry.add(
|
|
PingPongStep,
|
|
)
|