nodes/xmarre-ComfyUI-Spectrum-SDXL-Proper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Separate model-time and sigma coordinates in SDXL runtime

Browse Source
This commit is contained in:
xmarre
2026-03-20 12:20:21 +01:00
parent 2c5d925815
commit c670b6d9ee
3 changed files with 131 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_spectrum_sdxl/runtime.py
+36 -6
View File
@@ -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} "
comfyui_spectrum_sdxl/sdxl.py
+14 -3
View File
@@ -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="
tests/smoke_runtime.py
+81 -1
View File
@@ -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: