From 4ddfc343c51240f8fc5cb93255c87b0da8766772 Mon Sep 17 00:00:00 2001 From: xmarre Date: Fri, 20 Mar 2026 15:13:23 +0100 Subject: [PATCH 1/3] Use sigma for runtime forecasting --- comfyui_spectrum_sdxl/runtime.py | 29 +++++++++++++++++++++-------- tests/smoke_runtime.py | 13 +++++++------ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/comfyui_spectrum_sdxl/runtime.py b/comfyui_spectrum_sdxl/runtime.py index 7ad1fca..9b2c3e3 100644 --- a/comfyui_spectrum_sdxl/runtime.py +++ b/comfyui_spectrum_sdxl/runtime.py @@ -172,12 +172,12 @@ class SpectrumSDXLRuntime: return None return float(min(values)), float(max(values)) - def _model_time_coord_bounds( + def _observed_coord_bounds( self, state: _StreamState, current_time_coord: float, ) -> Optional[Tuple[float, float]]: - """Return bounds for the active model-time axis within this stream.""" + """Return bounds for the active forecast axis within this stream.""" values = [float(current_time_coord)] for decision in state.decisions_by_solver_step.values(): coord = decision.get("time_coord", None) @@ -196,7 +196,8 @@ class SpectrumSDXLRuntime: run_id = transformer_options.get(_RUN_ID_KEY, None) solver_step_id = transformer_options.get(_SOLVER_STEP_ID_KEY, None) sigma_coord = transformer_options.get(_TIME_COORD_KEY, None) - time_coord = transformer_options.get(_MODEL_TIME_COORD_KEY, sigma_coord) + model_time_coord = transformer_options.get(_MODEL_TIME_COORD_KEY, None) + time_coord = sigma_coord actual_forward = transformer_options.get(_ACTUAL_FORWARD_KEY, None) if run_id is None and solver_step_id is None and time_coord is None and sigma_coord is None and actual_forward is None: @@ -219,6 +220,15 @@ class SpectrumSDXLRuntime: if not math.isfinite(time_coord) or not math.isfinite(sigma_coord): return None, "invalid_solver_step_context" + if model_time_coord is not None: + try: + model_time_coord = float(model_time_coord) + except (TypeError, ValueError): + model_time_coord = None + else: + if not math.isfinite(model_time_coord): + model_time_coord = None + if actual_forward is not None and not isinstance(actual_forward, bool): return None, "invalid_solver_step_context" @@ -226,6 +236,7 @@ class SpectrumSDXLRuntime: "run_id": run_id, "solver_step_id": solver_step_id, "time_coord": time_coord, + "model_time_coord": model_time_coord, "sigma_coord": sigma_coord, "actual_forward": actual_forward, "total_steps": self._extract_total_steps(transformer_options), @@ -324,6 +335,7 @@ class SpectrumSDXLRuntime: "forecast_safe": False, "finalized": False, "time_coord": None, + "model_time_coord": None, "sigma_coord": None, "total_steps": self.num_steps(), } @@ -354,15 +366,15 @@ class SpectrumSDXLRuntime: "forecast_safe": False, "finalized": False, "time_coord": ctx["time_coord"], + "model_time_coord": ctx.get("model_time_coord", None), "sigma_coord": ctx["sigma_coord"], "total_steps": ctx["total_steps"], } state = self._ensure_stream_state(stream_key, ctx["run_id"]) - if _MODEL_TIME_COORD_KEY in transformer_options: - schedule_coord_bounds = self._model_time_coord_bounds(state, ctx["time_coord"]) - else: - schedule_coord_bounds = self._schedule_coord_bounds(transformer_options) + schedule_coord_bounds = self._schedule_coord_bounds(transformer_options) + if schedule_coord_bounds is None: + schedule_coord_bounds = self._observed_coord_bounds(state, ctx["time_coord"]) if schedule_coord_bounds is not None: state.forecaster.set_coord_bounds(*schedule_coord_bounds) existing = state.decisions_by_solver_step.get(ctx["solver_step_id"]) @@ -410,6 +422,7 @@ class SpectrumSDXLRuntime: "forecast_safe": forecast_safe, "finalized": False, "time_coord": ctx["time_coord"], + "model_time_coord": ctx.get("model_time_coord", None), "recent_validation_rel_l2": recent_validation_rel_l2, "sigma_coord": ctx["sigma_coord"], "total_steps": ctx["total_steps"], @@ -420,7 +433,7 @@ class SpectrumSDXLRuntime: print( f"[Spectrum SDXL] begin " f"run={ctx['run_id']} step={ctx['solver_step_id']} " - f"sigma={ctx['sigma_coord']} model_time={ctx['time_coord']} " + f"sigma={ctx['sigma_coord']} model_time={ctx.get('model_time_coord', None)} " f"val_rel_l2={recent_validation_rel_l2} " f"local={local_step_count} actual={actual_forward} " f"forecast_safe={forecast_safe} ready={state.forecaster.ready()} " diff --git a/tests/smoke_runtime.py b/tests/smoke_runtime.py index f187803..1a7767e 100644 --- a/tests/smoke_runtime.py +++ b/tests/smoke_runtime.py @@ -161,8 +161,8 @@ def test_explicit_solver_step_context_allows_forecast() -> None: assert runtime.last_info["forecasted_passes"] == 1 -def test_runtime_prefers_model_time_coord_for_forecasting() -> None: - """The forecast axis should follow model-time while sigma still validates the schedule.""" +def test_runtime_prefers_sigma_coord_for_forecasting_even_when_model_time_is_present() -> None: + """The forecast axis should stay on raw sigma even when model-time metadata is present.""" runtime = SpectrumSDXLRuntime(_make_relaxed_cfg()) sample_sigmas = torch.tensor([14.0, 7.0, 3.0, 1.5, 0.75, 0.25, 0.0], dtype=torch.float32) @@ -237,7 +237,8 @@ def test_runtime_prefers_model_time_coord_for_forecasting() -> None: (2, 8, 4, 4), ) - assert forecast["time_coord"] == 400.0 + assert forecast["time_coord"] == 1.5 + assert forecast["model_time_coord"] == 400.0 assert forecast["sigma_coord"] == 1.5 assert forecast["actual_forward"] is False assert forecast["forecast_safe"] is True @@ -249,8 +250,8 @@ def test_runtime_prefers_model_time_coord_for_forecasting() -> None: state = runtime.stream_states[forecast["stream_key"]] fit_cache = state.forecaster._fit_cache assert fit_cache is not None - assert torch.allclose(torch.tensor(fit_cache.coord_min), torch.tensor(100.0), atol=1e-6) - assert torch.allclose(torch.tensor(fit_cache.coord_max), torch.tensor(400.0), atol=1e-6) + assert torch.allclose(torch.tensor(fit_cache.coord_min), torch.tensor(0.25), atol=1e-6) + assert torch.allclose(torch.tensor(fit_cache.coord_max), torch.tensor(14.0), atol=1e-6) def test_explicit_solver_step_context_without_decision_still_schedules() -> None: @@ -1274,7 +1275,7 @@ def main() -> None: test_invalid_solver_step_context_fails_open() test_missing_stream_identity_fails_open() test_explicit_solver_step_context_allows_forecast() - test_runtime_prefers_model_time_coord_for_forecasting() + test_runtime_prefers_sigma_coord_for_forecasting_even_when_model_time_is_present() test_explicit_solver_step_context_without_decision_still_schedules() test_outer_step_controller_injects_context_and_resets_runs() test_outer_step_controller_same_object_restart_resets_run_state() From ee59d8ab2ec26017a0b953382500cfcaf5cd3401 Mon Sep 17 00:00:00 2001 From: xmarre Date: Fri, 20 Mar 2026 15:28:20 +0100 Subject: [PATCH 2/3] Use observed history for fallback forecast bounds --- comfyui_spectrum_sdxl/runtime.py | 12 ++++++------ tests/smoke_runtime.py | 20 ++++++++++++++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/comfyui_spectrum_sdxl/runtime.py b/comfyui_spectrum_sdxl/runtime.py index 9b2c3e3..d0d54a0 100644 --- a/comfyui_spectrum_sdxl/runtime.py +++ b/comfyui_spectrum_sdxl/runtime.py @@ -178,15 +178,15 @@ class SpectrumSDXLRuntime: current_time_coord: float, ) -> Optional[Tuple[float, float]]: """Return bounds for the active forecast axis within this stream.""" - values = [float(current_time_coord)] - for decision in state.decisions_by_solver_step.values(): - coord = decision.get("time_coord", None) - if coord is None: - continue + del current_time_coord + values = [] + for coord, _ in state.forecaster.history: try: - values.append(float(coord)) + coord = float(coord) except (TypeError, ValueError): continue + if math.isfinite(coord): + values.append(coord) if not values: return None return float(min(values)), float(max(values)) diff --git a/tests/smoke_runtime.py b/tests/smoke_runtime.py index 1a7767e..716b9af 100644 --- a/tests/smoke_runtime.py +++ b/tests/smoke_runtime.py @@ -707,6 +707,25 @@ def test_forecaster_chebyshev_fit_falls_back_to_observed_sigma_range() -> None: assert torch.isfinite(pred).all() +def test_runtime_fallback_coord_bounds_ignore_unobserved_candidate_step() -> None: + """Fallback bounds must come only from observed feature history, not pending steps.""" + runtime = SpectrumSDXLRuntime(_make_relaxed_cfg()) + + first = runtime.begin_step(_step_options("run-a", 0, 10.0, True), torch.tensor([10.0]), (1, 1, 1, 1)) + runtime.observe_actual_feature( + first["stream_key"], + first["solver_step_id"], + torch.full((1, 1, 1, 1), 10.0, dtype=torch.float16), + ) + + second = runtime.begin_step(_step_options("run-a", 1, 1.0, True), torch.tensor([1.0]), (1, 1, 1, 1)) + state = runtime.stream_states[second["stream_key"]] + + assert state.forecaster._coord_bounds is not None + assert torch.allclose(torch.tensor(state.forecaster._coord_bounds[0]), torch.tensor(10.0), atol=1e-6) + assert torch.allclose(torch.tensor(state.forecaster._coord_bounds[1]), torch.tensor(10.0), atol=1e-6) + + def test_runtime_uses_schedule_wide_sigma_bounds_for_forecast_fit() -> None: """When sample_sigmas are available, the fit must use the full active schedule span.""" cfg = _make_relaxed_cfg() @@ -1292,6 +1311,7 @@ def main() -> None: test_forecast_fallback_commits_actual_bookkeeping() test_run_id_switch_resets_stream_state() test_forecaster_chebyshev_fit_falls_back_to_observed_sigma_range() + test_runtime_fallback_coord_bounds_ignore_unobserved_candidate_step() test_runtime_uses_schedule_wide_sigma_bounds_for_forecast_fit() test_runtime_holds_forecasting_until_conservative_min_fit_points() test_runtime_blocks_forecast_when_recent_validation_error_is_bad() From ae29b59ee5df64f8237762cfbcc8e8e7f94b0dcc Mon Sep 17 00:00:00 2001 From: xmarre Date: Fri, 20 Mar 2026 15:34:43 +0100 Subject: [PATCH 3/3] Add fallback sigma forecast regression --- tests/smoke_runtime.py | 89 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/tests/smoke_runtime.py b/tests/smoke_runtime.py index 716b9af..3edb6b3 100644 --- a/tests/smoke_runtime.py +++ b/tests/smoke_runtime.py @@ -254,6 +254,94 @@ def test_runtime_prefers_sigma_coord_for_forecasting_even_when_model_time_is_pre assert torch.allclose(torch.tensor(fit_cache.coord_max), torch.tensor(14.0), atol=1e-6) +def test_runtime_prefers_sigma_coord_for_forecasting_without_sample_sigmas() -> None: + """Fallback bounds should stay on observed sigma history when schedule bounds are absent.""" + runtime = SpectrumSDXLRuntime(_make_relaxed_cfg()) + + first = runtime.begin_step( + _step_options( + "run-a", + 0, + 14.0, + True, + total_steps=8, + model_time_coord=100.0, + ), + torch.tensor([100.0]), + (2, 8, 4, 4), + ) + runtime.observe_actual_feature( + first["stream_key"], + first["solver_step_id"], + torch.full((2, 8, 4, 4), 1.0, dtype=torch.float16), + ) + + second = runtime.begin_step( + _step_options( + "run-a", + 1, + 7.0, + True, + total_steps=8, + model_time_coord=200.0, + ), + torch.tensor([200.0]), + (2, 8, 4, 4), + ) + runtime.observe_actual_feature( + second["stream_key"], + second["solver_step_id"], + torch.full((2, 8, 4, 4), 2.0, dtype=torch.float16), + ) + + third = runtime.begin_step( + _step_options( + "run-a", + 2, + 3.0, + True, + total_steps=8, + model_time_coord=300.0, + ), + torch.tensor([300.0]), + (2, 8, 4, 4), + ) + runtime.observe_actual_feature( + third["stream_key"], + third["solver_step_id"], + torch.full((2, 8, 4, 4), 3.0, dtype=torch.float16), + ) + + forecast = runtime.begin_step( + _step_options( + "run-a", + 3, + 1.5, + False, + total_steps=8, + model_time_coord=400.0, + ), + torch.tensor([400.0]), + (2, 8, 4, 4), + ) + + assert forecast["time_coord"] == 1.5 + assert forecast["model_time_coord"] == 400.0 + assert forecast["sigma_coord"] == 1.5 + assert forecast["actual_forward"] is False + assert forecast["forecast_safe"] is True + + pred = runtime.predict_feature(forecast["stream_key"], forecast["solver_step_id"]) + assert pred.shape == (2, 8, 4, 4) + assert torch.isfinite(pred).all() + + state = runtime.stream_states[forecast["stream_key"]] + fit_cache = state.forecaster._fit_cache + assert fit_cache is not None + assert torch.allclose(torch.tensor(fit_cache.coord_min), torch.tensor(3.0), atol=1e-6) + assert torch.allclose(torch.tensor(fit_cache.coord_max), torch.tensor(14.0), atol=1e-6) + + def test_explicit_solver_step_context_without_decision_still_schedules() -> None: """Outer-step context alone should be enough once the controller owns step ids.""" runtime = SpectrumSDXLRuntime(_make_relaxed_cfg()) @@ -1295,6 +1383,7 @@ def main() -> None: test_missing_stream_identity_fails_open() test_explicit_solver_step_context_allows_forecast() test_runtime_prefers_sigma_coord_for_forecasting_even_when_model_time_is_present() + test_runtime_prefers_sigma_coord_for_forecasting_without_sample_sigmas() test_explicit_solver_step_context_without_decision_still_schedules() test_outer_step_controller_injects_context_and_resets_runs() test_outer_step_controller_same_object_restart_resets_run_state()