From b2ba79a4cd820e23d8caeeb3cf0798dccfa858fb Mon Sep 17 00:00:00 2001 From: Mrinaal Dogra Date: Tue, 14 Apr 2026 05:55:47 +0000 Subject: [PATCH] [test]: add Z-Image scheduler parity and reference timestep option --- .../scheduling_flow_match_euler_discrete.py | 18 ++- .../zimage/test_zimage_scheduler_parity.py | 130 ++++++++++++++++++ 2 files changed, 147 insertions(+), 1 deletion(-) create mode 100644 tests/local_tests/zimage/test_zimage_scheduler_parity.py diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 5e0d65853..2c7c696e1 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -107,6 +107,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, + use_reference_discrete_timesteps: bool = False, base_shift: float | None = 0.5, max_shift: float | None = 1.15, base_image_seq_len: int | None = 256, @@ -350,7 +351,22 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, if timesteps_array is None: t_max = self._sigma_to_t(self.sigma_max) t_min = self._sigma_to_t(self.sigma_min) - timesteps_array = np.linspace(t_max, t_min, num_inference_steps) + if self.config.use_reference_discrete_timesteps: + # Some reference schedulers (for example Z-Image) build a + # num_steps+1 linspace and drop the terminal point. + timesteps_array = np.linspace( + t_max, + t_min, + num_inference_steps + 1, + dtype=np.float32, + )[:-1] + else: + timesteps_array = np.linspace( + t_max, + t_min, + num_inference_steps, + dtype=np.float32, + ) sigmas_array = timesteps_array / self.config.num_train_timesteps else: sigmas_array = np.array(sigmas).astype(np.float32) diff --git a/tests/local_tests/zimage/test_zimage_scheduler_parity.py b/tests/local_tests/zimage/test_zimage_scheduler_parity.py new file mode 100644 index 000000000..1d1edc764 --- /dev/null +++ b/tests/local_tests/zimage/test_zimage_scheduler_parity.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Parity tests for Z-Image FlowMatchEulerDiscreteScheduler. + +These tests compare FastVideo's scheduler implementation against the +reference scheduler shipped in the local Z-Image repository. + +Usage: + pytest tests/local_tests/zimage/test_zimage_scheduler_parity.py -v +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler as FastVideoScheduler, +) + + +REPO_ROOT = Path(__file__).resolve().parents[3] +ZIMAGE_SRC = REPO_ROOT / "Z-Image" / "src" +ZIMAGE_SCHEDULER_CFG = ( + REPO_ROOT / "official_weights" / "Z-Image" / "scheduler" / "scheduler_config.json" +) + + +if str(ZIMAGE_SRC) not in sys.path: + sys.path.insert(0, str(ZIMAGE_SRC)) + +try: + from zimage.scheduler import FlowMatchEulerDiscreteScheduler as ReferenceScheduler +except Exception as exc: # pragma: no cover - handled by skip in fixture + ReferenceScheduler = None + _IMPORT_ERROR = exc +else: + _IMPORT_ERROR = None + + +@pytest.fixture(scope="module") +def scheduler_config() -> dict: + if not ZIMAGE_SCHEDULER_CFG.exists(): + pytest.skip(f"Z-Image scheduler config not found: {ZIMAGE_SCHEDULER_CFG}") + with ZIMAGE_SCHEDULER_CFG.open("r", encoding="utf-8") as f: + return json.load(f) + + +@pytest.fixture(scope="module") +def scheduler_pair(scheduler_config: dict): + if ReferenceScheduler is None: + pytest.skip(f"Cannot import Z-Image reference scheduler: {_IMPORT_ERROR}") + + kwargs = { + "num_train_timesteps": scheduler_config["num_train_timesteps"], + "shift": scheduler_config["shift"], + "use_dynamic_shifting": scheduler_config["use_dynamic_shifting"], + } + ref = ReferenceScheduler(**kwargs) + fv = FastVideoScheduler(**kwargs, use_reference_discrete_timesteps=True) + return ref, fv + + +def _run_step_loop(ref_scheduler, fv_scheduler, sample_shape=(2, 4, 32, 32)): + torch.manual_seed(123) + + ref_sample = torch.randn(sample_shape, dtype=torch.float32) + fv_sample = ref_sample.clone() + + for t in ref_scheduler.timesteps: + model_output = torch.randn_like(ref_sample) + + ref_next = ref_scheduler.step( + model_output, + t, + ref_sample, + return_dict=False, + )[0] + fv_next = fv_scheduler.step( + model_output, + t, + fv_sample, + return_dict=False, + )[0] + + assert_close(ref_next, fv_next, atol=1e-6, rtol=1e-6) + + ref_sample = ref_next + fv_sample = fv_next + + +def test_zimage_scheduler_parity_default_schedule_and_step(scheduler_pair): + ref, fv = scheduler_pair + + num_inference_steps = 8 + ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu") + fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu") + + assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6) + assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6) + + _run_step_loop(ref, fv) + + +def test_zimage_scheduler_parity_dynamic_shifting_with_mu(scheduler_config: dict): + if ReferenceScheduler is None: + pytest.skip(f"Cannot import Z-Image reference scheduler: {_IMPORT_ERROR}") + + kwargs = { + "num_train_timesteps": scheduler_config["num_train_timesteps"], + "shift": scheduler_config["shift"], + "use_dynamic_shifting": True, + } + ref = ReferenceScheduler(**kwargs) + fv = FastVideoScheduler(**kwargs, use_reference_discrete_timesteps=True) + + mu = 0.75 + num_inference_steps = 8 + ref.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu) + fv.set_timesteps(num_inference_steps=num_inference_steps, device="cpu", mu=mu) + + assert_close(ref.timesteps, fv.timesteps, atol=1e-4, rtol=1e-6) + assert_close(ref.sigmas, fv.sigmas, atol=1e-7, rtol=1e-6) + + _run_step_loop(ref, fv)