Merge pull request #13 from xmarre/codex/align-forecaster-with-sigma-axis

Fix forecaster to use sigma axis for runtime validation

comfyui_spectrum_sdxl/runtime.py

@@ -172,21 +172,21 @@ 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."""
-        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
+        """Return bounds for the active forecast axis within this stream."""
+        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))
@@ -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()} "

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,96 @@ 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_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:
@@ -706,6 +795,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()
@@ -1274,7 +1382,8 @@ 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_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()
@@ -1291,6 +1400,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()