diff --git a/modules/sampling/turbo/__init__.py b/modules/sampling/turbo/__init__.py index e6b7a4a0..72bc8f40 100644 --- a/modules/sampling/turbo/__init__.py +++ b/modules/sampling/turbo/__init__.py @@ -44,6 +44,13 @@ SETTINGS = { "hint": "Replaces the chosen sampler with a step rule made for few-step turbo runs. No extra model calls " "(except the optional polish step). Works with any step count and with Two stages.", }, + "keep_sampler": { + "type": "bool", "default": False, + "label": "Keep the chosen sampler", + "hint": "The chosen sampler stays in charge of the steps. Turbo adds its fresh noise and its polish step on top. " + "Second order late is not used, the sampler has its own.", + "when": {"enabled": True}, + }, "noise_start": { "type": "float", "default": 0.4, "min": 0.0, "max": 1.0, "step": 0.05, "label": "Noise at the start", "ui": "slider", @@ -60,7 +67,7 @@ SETTINGS = { "type": "float", "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "label": "Second order late", "ui": "slider", "hint": "How much the late steps lean on the step before. 0 = off. It fades in over the run. Strong values can distort the picture.", - "when": {"enabled": True}, + "when": {"enabled": True, "keep_sampler": False}, }, "polish": { "type": "float", "default": 0.25, "min": 0.0, "max": 0.6, "step": 0.01, @@ -161,6 +168,70 @@ def make(settings): return sample +def churn(x, denoised, sigma, eta, noise): + """The answer `denoised` of a model call at `sigma` on `x`, changed so that any sampler which steps from it + moves toward a point whose noise is `eta` fresh. With D the answer and e = (x - (1 - sigma) D) / sigma the + noise it implies, the new answer is D + sigma * (e - e'), e' = sqrt(1 - eta^2) e + eta * noise. For Euler that + lands exactly on `(1 - to) D + to e'`, the Turbo step; other samplers read it as a nearby picture guess.""" + sigma = torch.as_tensor(sigma, device=x.device, dtype=x.dtype).reshape(-1)[0] + implied = (x - (1 - sigma) * denoised) / sigma + keep = math.sqrt(max(0.0, 1.0 - eta * eta)) + return denoised + sigma * ((1.0 - keep) * implied - eta * noise) + + +def make_on_top(inner, settings): + """`inner` (the chosen sampler's own function) with Turbo's fresh noise on each step's answer and its polish + step after the run. The sampler is not replaced; it sees a model whose answers carry the changes.""" + n_start, n_end, tail = settings["noise_start"], settings["noise_end"], settings["polish"] + + @torch.no_grad() + def sample(model, x, sigmas, extra_args=None, callback=None, disable=None, **options): + extra_args = {} if extra_args is None else extra_args + guider = getattr(model, "inner_model", None) + stage = dit_hooks.stage_of(guider) + full = getattr(guider, dit_hooks.FULL, None) + steps = len(sigmas) - 1 + total = (len(full) - 1) if full is not None else steps + n0 = total - steps if stage == 2 else 0 + base = int(extra_args.get("seed") or 0) + gen = torch.Generator(device=_parts(x)[0].device) + gen.manual_seed(base + 1000003 * (n0 + 1)) + ones = x.new_ones([x.shape[0]]) + log.once(f"turbo_top:{stage}:{steps}:{n0}", log.INFO, TITLE, + f"{'stage ' + str(stage) if stage else 'run'}: the chosen sampler keeps its steps, fresh noise and polish on top") + + class Churned: + def __getattr__(self, name): + return getattr(model, name) + + def __call__(self, latent, sigma, **kwargs): + answer = model(latent, sigma, **kwargs) + s0 = sigma.reshape(-1)[0] + hit = torch.nonzero(torch.isclose(sigmas.to(device=s0.device, dtype=s0.dtype), s0, rtol=1e-4)) + if not len(hit): + return answer # a call between steps (a midpoint): left alone + i = int(hit[0]) + if i >= steps or float(sigmas[i + 1]) == 0.0: + return answer # the last step must land on the clean picture + eta = n_start + (n_end - n_start) * ((n0 + i) / max(1, total - 1)) + if eta <= 0.0: + return answer + lx, la = _parts(latent), _parts(answer) + la[0] = churn(lx[0], la[0], s0, eta, _noise(lx[0], gen)) # the picture only; the sound gets none + return _join(answer, la) + + result = inner(Churned(), x, sigmas, extra_args=extra_args, callback=callback, disable=disable, **options) + if tail > 0.0 and stage != 1: + noisy = [] + for p in _parts(result): + t = torch.as_tensor(tail, device=p.device, dtype=p.dtype) + noisy.append((1 - t) * p + t * _noise(p, gen)) + t = torch.as_tensor(tail, device=_parts(result)[0].device, dtype=_parts(result)[0].dtype) + result = model(_join(result, noisy), t * ones, **extra_args) + return result + return sample + + def _settings(values): return {name: float(values.get(name, spec["default"])) for name, spec in SETTINGS.items() if name != "enabled"} @@ -198,14 +269,16 @@ def install(patcher, values, key): if inner is None: log.once("turbo:shape", log.INFO, TITLE, "off this run: the sampler is not a plain sampler function") return executor(guider, sigmas, *args, **kwargs) - sampler.sampler_function = make(settings) # swapped for this call and passed ON: later wrappers still run + sampler.sampler_function = (make_on_top(inner, settings) if settings["keep_sampler"] else make(settings)) # swapped for this call and passed ON: later wrappers still run try: return executor(guider, sigmas, *args, **kwargs) finally: sampler.sampler_function = inner patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample) - return (f"fresh noise {settings['noise_start']:g} to {settings['noise_end']:g}, second order {settings['second_order']:g}" + kept = bool(settings["keep_sampler"]) + return (f"fresh noise {settings['noise_start']:g} to {settings['noise_end']:g}" + + ("" if kept else f", second order {settings['second_order']:g}") + (", on the chosen sampler" if kept else "") + (f", polish at {settings['polish']:g}" if settings["polish"] > 0 else "")) diff --git a/modules/sampling/turbo/tests/test_turbo.py b/modules/sampling/turbo/tests/test_turbo.py index fb6e1f6b..34f6e996 100644 --- a/modules/sampling/turbo/tests/test_turbo.py +++ b/modules/sampling/turbo/tests/test_turbo.py @@ -3,7 +3,7 @@ import pytest import torch -from modules.sampling.turbo import SETTINGS, install, make, step_to +from modules.sampling.turbo import SETTINGS, churn, install, make, make_on_top, step_to from modules._core import dit_hooks from modules.sampling.two_stage.tests.test_two_stage import model # noqa: F401 -- the module-scoped fixture @@ -172,3 +172,48 @@ def test_the_switch_takes_over_whatever_sampler_is_chosen(model, sampler): note, out, calls = _guided(model, values, (1.0, 0.7, 0.3, 0.0), sampler=sampler) assert "fresh noise" in note and len(calls) == 4 and calls[-1] == pytest.approx(0.05) assert torch.isfinite(out.unbind()[0]).all() + + +def test_churn_with_no_fresh_noise_changes_nothing_and_follows_dtype_and_device(): + x, d = torch.randn(1, 3, 2, 4, 4, dtype=torch.float64), torch.randn(1, 3, 2, 4, 4, dtype=torch.float64) + n = torch.randn_like(x) + sigma = torch.tensor([0.6], dtype=torch.float32) # a different dtype from the picture + assert torch.equal(churn(x, d, sigma, 0.0, n), d) + out = churn(x, d, sigma, 0.5, n) + assert out.dtype == torch.float64 and out.device == x.device and not torch.equal(out, d) + + +def test_an_euler_step_from_the_churned_answer_lands_on_the_turbo_step(): + x, d, n = torch.randn(1, 3, 2, 4, 4), torch.randn(1, 3, 2, 4, 4), torch.randn(1, 3, 2, 4, 4) + sigma, to, eta = torch.tensor(0.7), torch.tensor(0.4), 0.45 + answer = churn(x, d, sigma, eta, n) + euler = x + ((x - answer) / sigma) * (to - sigma) # what Euler does with an answer + assert torch.allclose(euler, step_to(x, d, sigma, to, eta, n), atol=1e-5) + + +def _euler(model, x, sigmas, extra_args=None, callback=None, disable=None, **_): + ones = x.new_ones([x.shape[0]]) + for i in range(len(sigmas) - 1): + d = (x - model(x, sigmas[i] * ones, **(extra_args or {}))) / sigmas[i] + x = x + d * (sigmas[i + 1] - sigmas[i]) + return x + + +def test_on_top_of_a_sampler_with_no_fresh_noise_and_no_polish_is_that_sampler(): + target, sigmas = torch.randn(1, 3, 2, 4, 4), torch.tensor([1.0, 0.7, 0.3, 0.0]) + noise = torch.randn_like(target) + plain = _euler(Oracle(target), noise.clone(), sigmas) + top = make_on_top(_euler, _settings(noise_start=0.0, noise_end=0.0, polish=0.0))(Oracle(target), noise.clone(), sigmas) + assert torch.equal(plain, top) + + +def test_on_top_of_a_sampler_adds_fresh_noise_before_the_last_step_and_one_polish_call(): + target, sigmas = torch.randn(1, 3, 2, 4, 4), torch.tensor([1.0, 0.7, 0.3, 0.0]) + noise = torch.randn_like(target) + oracle = Oracle(target) + run = make_on_top(_euler, _settings(noise_start=0.5, noise_end=0.5, polish=0.1)) + a = run(oracle, noise.clone(), sigmas, extra_args={"seed": 3}) + b = run(Oracle(target), noise.clone(), sigmas, extra_args={"seed": 3}) + assert torch.equal(a, b) # the same seed gives the same picture + assert oracle.calls == pytest.approx([1.0, 0.7, 0.3, 0.1]) # one call per step, then the polish call + assert not torch.equal(a, _euler(Oracle(target), noise.clone(), sigmas))