Two stages: noise at the hand-over defaults to 0.5, and the schedule the model sees holds the backed-up start so step windows still find their step

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjU9jGrByQHYf6UYMVV

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()

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 "")
 

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,

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