From 0ee9a8600b3eb63e1979837ce9ef0294aae3b78a Mon Sep 17 00:00:00 2001 From: xmarre Date: Thu, 23 Apr 2026 20:14:00 +0200 Subject: [PATCH] Allow Flux2JiTSamplerImpl solver-step tracking --- comfyui_spectrum/flux.py | 1 + tests/smoke_runtime.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/comfyui_spectrum/flux.py b/comfyui_spectrum/flux.py index 6abe696..8a8c5ec 100644 --- a/comfyui_spectrum/flux.py +++ b/comfyui_spectrum/flux.py @@ -18,6 +18,7 @@ _SUPPORTED_SINGLE_EVAL_SAMPLERS = frozenset( "sample_dpmpp_2m_sde", "sample_dpmpp_3m_sde", "euler_flow", + "Flux2JiTSamplerImpl", } ) diff --git a/tests/smoke_runtime.py b/tests/smoke_runtime.py index 3b735db..2b06a94 100644 --- a/tests/smoke_runtime.py +++ b/tests/smoke_runtime.py @@ -45,6 +45,7 @@ def test_supported_single_eval_sampler_names_include_supported_variants() -> Non assert "sample_dpmpp_2m_sde" in _SUPPORTED_SINGLE_EVAL_SAMPLERS assert "sample_dpmpp_3m_sde" in _SUPPORTED_SINGLE_EVAL_SAMPLERS assert "euler_flow" in _SUPPORTED_SINGLE_EVAL_SAMPLERS + assert "Flux2JiTSamplerImpl" in _SUPPORTED_SINGLE_EVAL_SAMPLERS class _FunctionSampler: def __init__(self, fn): @@ -66,6 +67,7 @@ def test_supported_single_eval_sampler_names_include_supported_variants() -> Non raise NotImplementedError EulerFlowSampler = type("euler_flow", (), {}) + Flux2JiTSamplerImpl = type("Flux2JiTSamplerImpl", (), {}) assert _supports_solver_step_tracking(_FunctionSampler(sample_euler_ancestral)) is True assert _supports_solver_step_tracking(_FunctionSampler(sample_euler_flow)) is True @@ -73,6 +75,7 @@ def test_supported_single_eval_sampler_names_include_supported_variants() -> Non assert _supports_solver_step_tracking(_FunctionSampler(sample_dpmpp_2m_sde)) is True assert _supports_solver_step_tracking(_FunctionSampler(sample_dpmpp_3m_sde)) is True assert _supports_solver_step_tracking(EulerFlowSampler()) is True + assert _supports_solver_step_tracking(Flux2JiTSamplerImpl()) is True def test_forecaster_recomputes_coeff_on_update_not_predict() -> None: @@ -642,6 +645,7 @@ def test_flux_sampler_contract_allows_supported_single_eval_variants() -> None: "sample_dpmpp_2m_sde", "sample_dpmpp_3m_sde", "euler_flow", + "Flux2JiTSamplerImpl", } )