diff --git a/modules/sampling/turbo/tests/test_turbo.py b/modules/sampling/turbo/tests/test_turbo.py index 8bb57663..2483d2b9 100644 --- a/modules/sampling/turbo/tests/test_turbo.py +++ b/modules/sampling/turbo/tests/test_turbo.py @@ -126,7 +126,7 @@ def test_off_installs_nothing(model): ((1.0, 0.6, 0.3, 0.0), 1)]) def test_runs_through_two_stages_with_any_split(model, sigmas, cut): values = {"enabled": True, **_settings(polish=0.05)} - note, out, calls = _guided(model, values, sigmas, {"enabled": True, "size": 0.5, "steps": cut}) + note, out, calls = _guided(model, values, sigmas, {"enabled": True, "size": 0.5, "steps": cut, "back_up": 0.0}) assert "fresh noise" in note v, a = out.unbind() assert torch.isfinite(v).all() and torch.isfinite(a).all() diff --git a/modules/sampling/two_stage/__init__.py b/modules/sampling/two_stage/__init__.py index 8d7303ac..9cdbb3b2 100644 --- a/modules/sampling/two_stage/__init__.py +++ b/modules/sampling/two_stage/__init__.py @@ -58,7 +58,7 @@ SETTINGS = { "when": {"enabled": True}, }, "back_up": { - "type": "float", "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "ui": "slider", + "type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, "ui": "slider", "label": "Noise added at the hand-over", "hint": "How far back the full-size stage starts. 0 continues exactly where the first stage stopped. Higher values add " "fresh noise to the enlarged picture, so the full-size stage can redraw soft or doubled outlines.", @@ -96,7 +96,7 @@ def install(patcher, values, key): size, steps = float(values.get("size", 0.7)), int(values.get("steps", 2)) resize = values.get("upscaler") if values.get("resize") == "upscaler" and values.get("upscaler") else None method = values.get("resize") if values.get("resize") in METHODS else "bicubic" - patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, make_wrapper(size, steps, resize, method, float(values.get("back_up", 0.0)))) + patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, make_wrapper(size, steps, resize, method, float(values.get("back_up", 0.5)))) patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, f"{key}.schedule", keep_whole_schedule) return f"first {steps} step(s) at {size:.0%} of the size" + (f", enlarged by {resize}" if resize else f", enlarged by {method}" if method != "bicubic" else "") diff --git a/modules/sampling/two_stage/stages.py b/modules/sampling/two_stage/stages.py index 258df8df..5d365e5b 100644 --- a/modules/sampling/two_stage/stages.py +++ b/modules/sampling/two_stage/stages.py @@ -232,7 +232,9 @@ def _staged(original, guider, size, steps, resize, method, back_up, *args, **kwa except Exception as e: # noqa: BLE001 -- a missing file must not lose the run log.warning(TITLE, f"latent upscaler not used, bicubic instead: {e}") log.info(TITLE, f"first {steps} step(s) at latent {h1}x{w1} instead of {h}x{w}, then {len(second) - 1} at the full size") - setattr(guider, FULL, sigmas) + whole = sigmas.clone() + whole[steps] = second[0] # the step index is found by sigma, so the schedule the model sees holds the one it is called at + setattr(guider, FULL, whole) setattr(guider, dit_hooks.STAGE, 1) try: run(small_noise, small_image, sampler, first, None, first_stage_callback, disable_pbar, seed, diff --git a/modules/sampling/two_stage/tests/test_two_stage.py b/modules/sampling/two_stage/tests/test_two_stage.py index 874bca8d..0fde8785 100644 --- a/modules/sampling/two_stage/tests/test_two_stage.py +++ b/modules/sampling/two_stage/tests/test_two_stage.py @@ -74,10 +74,23 @@ def test_a_run_that_cannot_be_cut_is_left_alone(model): def test_every_stage_sees_the_whole_schedule(model): seen = [] sigmas = (1.0, 0.9, 0.7, 0.3, 0.0) - _run(model, {"enabled": True, "size": 0.5, "steps": 2}, sigmas=sigmas, seen=seen) + _run(model, {"enabled": True, "size": 0.5, "steps": 2, "back_up": 0.0}, sigmas=sigmas, seen=seen) assert len(seen) == 4 and all(s == pytest.approx(list(sigmas)) for s in seen) +def test_backing_up_still_lets_every_call_find_its_step(model): + """Step windows find the step by its sigma in the schedule the model sees: the backed-up start must be in it.""" + from core import dit_hooks + found = [] + import comfy.samplers + seen = [] + _run(model, {"enabled": True, "size": 0.5, "steps": 2, "back_up": 0.5}, sigmas=(1.0, 0.9, 0.7, 0.3, 0.0), seen=seen) + assert len(seen) == 4 and all(len(s) == 5 and abs(s[2] - 0.8) < 1e-3 for s in seen) + for i, sig in enumerate([1.0, 0.9, 0.8, 0.3]): + where = dit_hooks.current_step({"sample_sigmas": torch.tensor(seen[0]), "sigmas": torch.tensor([sig])}) + assert where == (i, 4) + + def test_the_whole_schedule_lives_where_the_stage_schedule_does(): from ..stages import FULL, keep_whole_schedule