diff --git a/comfyui_spectrum_sdxl/runtime.py b/comfyui_spectrum_sdxl/runtime.py index 9d04fc5..c46d4eb 100644 --- a/comfyui_spectrum_sdxl/runtime.py +++ b/comfyui_spectrum_sdxl/runtime.py @@ -17,6 +17,7 @@ StreamKey = Tuple[Tuple[str, ...], Tuple[int, ...], Tuple[int, ...]] _RUN_ID_KEY = "spectrum_run_id" _SOLVER_STEP_ID_KEY = "spectrum_solver_step_id" _TIME_COORD_KEY = "spectrum_time_coord" +_MODEL_TIME_COORD_KEY = "spectrum_model_time_coord" _ACTUAL_FORWARD_KEY = "spectrum_actual_forward" _TOTAL_STEPS_KEY = "spectrum_total_steps" @@ -169,14 +170,34 @@ class SpectrumSDXLRuntime: return None return float(min(values)), float(max(values)) + def _model_time_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 + try: + values.append(float(coord)) + except (TypeError, ValueError): + continue + if not values: + return None + return float(min(values)), float(max(values)) + def _solver_step_context(self, transformer_options: Dict[str, Any]) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: """Validate and extract an explicit outer-step Spectrum context.""" run_id = transformer_options.get(_RUN_ID_KEY, None) solver_step_id = transformer_options.get(_SOLVER_STEP_ID_KEY, None) - time_coord = transformer_options.get(_TIME_COORD_KEY, None) + sigma_coord = transformer_options.get(_TIME_COORD_KEY, None) + time_coord = transformer_options.get(_MODEL_TIME_COORD_KEY, 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 actual_forward is 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: return None, "missing_solver_step_context" if run_id is None or solver_step_id is None or time_coord is None: return None, "invalid_solver_step_context" @@ -190,9 +211,10 @@ class SpectrumSDXLRuntime: try: time_coord = float(time_coord) + sigma_coord = float(sigma_coord) except (TypeError, ValueError): return None, "invalid_solver_step_context" - if not math.isfinite(time_coord): + if not math.isfinite(time_coord) or not math.isfinite(sigma_coord): return None, "invalid_solver_step_context" if actual_forward is not None and not isinstance(actual_forward, bool): @@ -202,6 +224,7 @@ class SpectrumSDXLRuntime: "run_id": run_id, "solver_step_id": solver_step_id, "time_coord": time_coord, + "sigma_coord": sigma_coord, "actual_forward": actual_forward, "total_steps": self._extract_total_steps(transformer_options), }, None @@ -299,6 +322,7 @@ class SpectrumSDXLRuntime: "forecast_safe": False, "finalized": False, "time_coord": None, + "sigma_coord": None, "total_steps": self.num_steps(), } @@ -306,9 +330,9 @@ class SpectrumSDXLRuntime: self.last_info["num_steps"] = ctx["total_steps"] expected_time_coord = self._expected_time_coord(transformer_options, ctx["solver_step_id"]) if expected_time_coord is not None and not math.isclose( - float(ctx["time_coord"]), float(expected_time_coord), rel_tol=1e-6, abs_tol=1e-6 + float(ctx["sigma_coord"]), float(expected_time_coord), rel_tol=1e-6, abs_tol=1e-6 ): - self._disable_forecasting("solver-step time_coord did not match the active schedule") + self._disable_forecasting("solver-step sigma_coord did not match the active schedule") if stream_key is None: self._disable_forecasting("missing_stream_identity") @@ -328,11 +352,15 @@ class SpectrumSDXLRuntime: "forecast_safe": False, "finalized": False, "time_coord": ctx["time_coord"], + "sigma_coord": ctx["sigma_coord"], "total_steps": ctx["total_steps"], } state = self._ensure_stream_state(stream_key, ctx["run_id"]) - schedule_coord_bounds = self._schedule_coord_bounds(transformer_options) + 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) 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"]) @@ -375,6 +403,7 @@ class SpectrumSDXLRuntime: "forecast_safe": forecast_safe, "finalized": False, "time_coord": ctx["time_coord"], + "sigma_coord": ctx["sigma_coord"], "total_steps": ctx["total_steps"], } state.decisions_by_solver_step[ctx["solver_step_id"]] = decision @@ -383,6 +412,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"local={local_step_count} actual={actual_forward} " f"forecast_safe={forecast_safe} ready={state.forecaster.ready()} " f"ws={state.curr_ws:.3f} cached={state.num_consecutive_cached_steps} " diff --git a/comfyui_spectrum_sdxl/sdxl.py b/comfyui_spectrum_sdxl/sdxl.py index 727fe40..b518e8b 100644 --- a/comfyui_spectrum_sdxl/sdxl.py +++ b/comfyui_spectrum_sdxl/sdxl.py @@ -20,6 +20,7 @@ _MODEL_FUNCTION_WRAPPER_KEY = "spectrum_sdxl_model_function_wrapper" _RUN_ID_KEY = "spectrum_run_id" _SOLVER_STEP_ID_KEY = "spectrum_solver_step_id" _TIME_COORD_KEY = "spectrum_time_coord" +_MODEL_TIME_COORD_KEY = "spectrum_model_time_coord" _TOTAL_STEPS_KEY = "spectrum_total_steps" @@ -226,7 +227,7 @@ class _SpectrumOuterStepController: "[Spectrum SDXL] controller " f"run={transformer_options.get(_RUN_ID_KEY)} " f"step={transformer_options.get(_SOLVER_STEP_ID_KEY)} " - f"time={transformer_options.get(_TIME_COORD_KEY)} " + f"sigma={transformer_options.get(_TIME_COORD_KEY)} " f"total={transformer_options.get(_TOTAL_STEPS_KEY)} " f"model_options_id={id(model_options)} " f"transformer_options_id={id(transformer_options)}" @@ -274,7 +275,7 @@ class _SpectrumModelFunctionWrapper: "[Spectrum SDXL] model_wrapper " f"run={transformer_options.get(_RUN_ID_KEY)} " f"step={transformer_options.get(_SOLVER_STEP_ID_KEY)} " - f"time={transformer_options.get(_TIME_COORD_KEY)} " + f"sigma={transformer_options.get(_TIME_COORD_KEY)} " f"total={transformer_options.get(_TOTAL_STEPS_KEY)} " f"forced_actual={transformer_options.get('spectrum_actual_forward', None)}" ) @@ -367,12 +368,22 @@ def _wrap_sdxl_unet_forward(inner: Any) -> None: if transformer_options is None: transformer_options = {} + model_time_coord = None + if timesteps is not None: + try: + model_time_coord = float(timesteps.detach().flatten()[0].item()) + except Exception: + model_time_coord = None + if model_time_coord is not None: + transformer_options[_MODEL_TIME_COORD_KEY] = model_time_coord + if runtime.cfg.debug: print( "[Spectrum SDXL] unet " f"run={transformer_options.get(_RUN_ID_KEY)} " f"step={transformer_options.get(_SOLVER_STEP_ID_KEY)} " - f"time={transformer_options.get(_TIME_COORD_KEY)} " + f"sigma={transformer_options.get(_TIME_COORD_KEY)} " + f"model_time={transformer_options.get(_MODEL_TIME_COORD_KEY)} " f"total={transformer_options.get(_TOTAL_STEPS_KEY)} " f"transformer_options_id={id(transformer_options)} " f"stream_key_present=" diff --git a/tests/smoke_runtime.py b/tests/smoke_runtime.py index 4e18b9b..20d7ee3 100644 --- a/tests/smoke_runtime.py +++ b/tests/smoke_runtime.py @@ -12,6 +12,7 @@ from comfyui_spectrum_sdxl.forecast import ChebyshevFeatureForecaster from comfyui_spectrum_sdxl.runtime import SpectrumSDXLRuntime from comfyui_spectrum_sdxl.sdxl import ( _RUNTIME_KEY, + _MODEL_TIME_COORD_KEY, _SpectrumModelFunctionWrapper, _SpectrumOuterStepController, _wrap_sdxl_unet_forward, @@ -38,6 +39,8 @@ def _step_options( total_steps: int = 6, uuid: str = "stream-a", cond: int = 0, + model_time_coord: float | None = None, + sample_sigmas: torch.Tensor | None = None, ): """Create explicit outer-step context for one logical stream.""" options = { @@ -48,6 +51,10 @@ def _step_options( "uuids": [uuid], "cond_or_uncond": [cond], } + if model_time_coord is not None: + options[_MODEL_TIME_COORD_KEY] = model_time_coord + if sample_sigmas is not None: + options["sample_sigmas"] = sample_sigmas if actual_forward is not None: options["spectrum_actual_forward"] = actual_forward return options @@ -132,6 +139,79 @@ 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.""" + runtime = SpectrumSDXLRuntime(_make_cfg()) + sample_sigmas = torch.tensor([14.0, 7.0, 3.0, 1.5, 0.75, 0.25, 0.0], dtype=torch.float32) + + first = runtime.begin_step( + _step_options( + "run-a", + 0, + 14.0, + True, + total_steps=6, + model_time_coord=100.0, + sample_sigmas=sample_sigmas, + ), + 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=6, + model_time_coord=200.0, + sample_sigmas=sample_sigmas, + ), + 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), + ) + + forecast = runtime.begin_step( + _step_options( + "run-a", + 2, + 3.0, + False, + total_steps=6, + model_time_coord=300.0, + sample_sigmas=sample_sigmas, + ), + torch.tensor([300.0]), + (2, 8, 4, 4), + ) + + assert forecast["time_coord"] == 300.0 + assert forecast["sigma_coord"] == 3.0 + 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(100.0), atol=1e-6) + assert torch.allclose(torch.tensor(fit_cache.coord_max), torch.tensor(300.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_cfg()) @@ -248,7 +328,7 @@ def test_runtime_disables_forecasting_when_time_coord_does_not_match_schedule() ) assert decision["actual_forward"] is True assert decision["forecast_safe"] is False - assert runtime.last_info["forecast_disable_reason"] == "solver-step time_coord did not match the active schedule" + assert runtime.last_info["forecast_disable_reason"] == "solver-step sigma_coord did not match the active schedule" def test_model_function_wrapper_injects_context_for_bypassed_guider_path() -> None: