Skip to content

Schedulers

mlx_diffuser.schedulers.Scheduler

Abstract base. Subclasses implement the four core methods below.

Source code in src/mlx_diffuser/schedulers/base.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
class Scheduler:
    """Abstract base. Subclasses implement the four core methods below."""

    config_class: type[SchedulerConfig] = SchedulerConfig

    def __init__(self, config: SchedulerConfig):
        self.config = config
        self.timesteps: mx.array | None = None
        self.num_inference_steps: int | None = None
        self._step_index = 0

    # --- training ---------------------------------------------------------
    def sample_timesteps(self, batch_size: int, key: mx.array) -> mx.array:
        """Draw a batch of training timesteps in this scheduler's convention."""
        raise NotImplementedError

    def add_noise(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        """Forward process: corrupt ``x0`` with ``noise`` at timestep ``t``."""
        raise NotImplementedError

    def get_target(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        """The quantity the network is trained to predict (per ``prediction_type``)."""
        raise NotImplementedError

    # --- sampling ---------------------------------------------------------
    def set_timesteps(self, num_inference_steps: int) -> None:
        """Configure the inference timestep grid (descending) and reset state."""
        raise NotImplementedError

    def step(
        self, model_output: mx.array, t: mx.array, sample: mx.array, key: mx.array | None = None
    ) -> mx.array:
        """Take one reverse step, returning the previous (less-noisy) sample."""
        raise NotImplementedError

    def scale_model_input(self, sample: mx.array, t: mx.array) -> mx.array:
        """Optional pre-network input scaling (identity unless overridden)."""
        return sample

    def set_begin_index(self, begin_index: int) -> None:
        """Start sampling partway through the configured schedule.

        Image/video-to-X pipelines use this after adding noise to encoded input
        latents. Keeping the step cursor on the scheduler avoids pipeline code
        reaching into scheduler internals.
        """
        if self.timesteps is None:
            raise RuntimeError("call set_timesteps() before set_begin_index().")
        if not 0 <= begin_index < len(self.timesteps):
            raise ValueError(
                f"begin_index must be in [0, {len(self.timesteps) - 1}], got {begin_index}."
            )
        self._step_index = begin_index

    # --- persistence ------------------------------------------------------
    def save_pretrained(self, save_directory) -> None:
        self.config.save(save_directory)

    @classmethod
    def from_config(cls, config: SchedulerConfig) -> Scheduler:
        return cls(config)

add_noise(x0, noise, t)

Forward process: corrupt x0 with noise at timestep t.

Source code in src/mlx_diffuser/schedulers/base.py 
88
89
90
def add_noise(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
    """Forward process: corrupt ``x0`` with ``noise`` at timestep ``t``."""
    raise NotImplementedError

get_target(x0, noise, t)

The quantity the network is trained to predict (per prediction_type).

Source code in src/mlx_diffuser/schedulers/base.py 
92
93
94
def get_target(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
    """The quantity the network is trained to predict (per ``prediction_type``)."""
    raise NotImplementedError

sample_timesteps(batch_size, key)

Draw a batch of training timesteps in this scheduler's convention.

Source code in src/mlx_diffuser/schedulers/base.py 
84
85
86
def sample_timesteps(self, batch_size: int, key: mx.array) -> mx.array:
    """Draw a batch of training timesteps in this scheduler's convention."""
    raise NotImplementedError

scale_model_input(sample, t)

Optional pre-network input scaling (identity unless overridden).

Source code in src/mlx_diffuser/schedulers/base.py 
107
108
109
def scale_model_input(self, sample: mx.array, t: mx.array) -> mx.array:
    """Optional pre-network input scaling (identity unless overridden)."""
    return sample

set_begin_index(begin_index)

Start sampling partway through the configured schedule.

Image/video-to-X pipelines use this after adding noise to encoded input latents. Keeping the step cursor on the scheduler avoids pipeline code reaching into scheduler internals.

Source code in src/mlx_diffuser/schedulers/base.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def set_begin_index(self, begin_index: int) -> None:
    """Start sampling partway through the configured schedule.

    Image/video-to-X pipelines use this after adding noise to encoded input
    latents. Keeping the step cursor on the scheduler avoids pipeline code
    reaching into scheduler internals.
    """
    if self.timesteps is None:
        raise RuntimeError("call set_timesteps() before set_begin_index().")
    if not 0 <= begin_index < len(self.timesteps):
        raise ValueError(
            f"begin_index must be in [0, {len(self.timesteps) - 1}], got {begin_index}."
        )
    self._step_index = begin_index

set_timesteps(num_inference_steps)

Configure the inference timestep grid (descending) and reset state.

Source code in src/mlx_diffuser/schedulers/base.py 
97
98
99
def set_timesteps(self, num_inference_steps: int) -> None:
    """Configure the inference timestep grid (descending) and reset state."""
    raise NotImplementedError

step(model_output, t, sample, key=None)

Take one reverse step, returning the previous (less-noisy) sample.

Source code in src/mlx_diffuser/schedulers/base.py 
101
102
103
104
105
def step(
    self, model_output: mx.array, t: mx.array, sample: mx.array, key: mx.array | None = None
) -> mx.array:
    """Take one reverse step, returning the previous (less-noisy) sample."""
    raise NotImplementedError

mlx_diffuser.schedulers.DDPMScheduler

Bases: Scheduler

Source code in src/mlx_diffuser/schedulers/ddpm.py 
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class DDPMScheduler(Scheduler):
    config_class = DDPMConfig
    config: DDPMConfig

    def __init__(self, config: DDPMConfig | None = None):
        super().__init__(config or DDPMConfig())
        betas = make_betas(
            self.config.beta_schedule,
            self.config.num_train_timesteps,
            self.config.beta_start,
            self.config.beta_end,
        )
        self.betas = betas
        self.alphas = 1.0 - betas
        self.alphas_cumprod = mx.cumprod(self.alphas)
        self._stride = 1

    # --- coefficient helpers ---------------------------------------------
    def _acp(self, t: mx.array) -> mx.array:
        return self.alphas_cumprod[t]

    def _sqrt_terms(self, t: mx.array, ndim: int) -> tuple[mx.array, mx.array]:
        acp = self._acp(t)
        return expand_to(mx.sqrt(acp), ndim), expand_to(mx.sqrt(1.0 - acp), ndim)

    # --- training ---------------------------------------------------------
    def sample_timesteps(self, batch_size: int, key: mx.array) -> mx.array:
        return mx.random.randint(0, self.config.num_train_timesteps, (batch_size,), key=key)

    def add_noise(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        sqrt_acp, sqrt_omacp = self._sqrt_terms(t, x0.ndim)
        return sqrt_acp * x0 + sqrt_omacp * noise

    def get_target(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        pt = self.config.prediction_type
        if pt == "epsilon":
            return noise
        if pt == "sample":
            return x0
        if pt == "v_prediction":
            sqrt_acp, sqrt_omacp = self._sqrt_terms(t, x0.ndim)
            return sqrt_acp * noise - sqrt_omacp * x0
        raise ValueError(f"{type(self).__name__} does not support prediction_type={pt!r}.")

    def predict_x0(self, model_output: mx.array, t: mx.array, sample: mx.array) -> mx.array:
        """Recover predicted clean sample x0 from the network output."""
        sqrt_acp, sqrt_omacp = self._sqrt_terms(t, sample.ndim)
        pt = self.config.prediction_type
        if pt == "epsilon":
            x0 = (sample - sqrt_omacp * model_output) / sqrt_acp
        elif pt == "sample":
            x0 = model_output
        elif pt == "v_prediction":
            x0 = sqrt_acp * sample - sqrt_omacp * model_output
        else:  # pragma: no cover - guarded in get_target
            raise ValueError(pt)
        if self.config.clip_sample:
            r = self.config.clip_sample_range
            x0 = mx.clip(x0, -r, r)
        return x0

    # --- sampling ---------------------------------------------------------
    def set_timesteps(self, num_inference_steps: int) -> None:
        T = self.config.num_train_timesteps
        if num_inference_steps > T:
            raise ValueError(
                f"num_inference_steps ({num_inference_steps}) > num_train_timesteps ({T})."
            )
        self.num_inference_steps = num_inference_steps
        self._stride = T // num_inference_steps
        self.timesteps = (mx.arange(num_inference_steps) * self._stride)[::-1]
        self._step_index = 0

    def step(
        self, model_output: mx.array, t: mx.array, sample: mx.array, key: mx.array | None = None
    ) -> mx.array:
        ti = int(t.item()) if isinstance(t, mx.array) else int(t)
        prev = ti - self._stride
        acp_t = self.alphas_cumprod[ti]
        acp_prev = self.alphas_cumprod[prev] if prev >= 0 else mx.array(1.0)

        x0 = self.predict_x0(model_output, mx.array(ti), sample)

        beta_t = 1.0 - acp_t
        beta_prev = 1.0 - acp_prev
        current_alpha = acp_t / acp_prev
        current_beta = 1.0 - current_alpha

        x0_coef = mx.sqrt(acp_prev) * current_beta / beta_t
        sample_coef = mx.sqrt(current_alpha) * beta_prev / beta_t
        mean = x0_coef * x0 + sample_coef * sample

        if prev < 0:
            return mean
        var = current_beta * beta_prev / beta_t
        key = key if key is not None else mx.random.key(ti)
        noise = mx.random.normal(sample.shape, key=key)
        return mean + mx.sqrt(mx.maximum(var, 1e-20)) * noise

predict_x0(model_output, t, sample)

Recover predicted clean sample x0 from the network output.

Source code in src/mlx_diffuser/schedulers/ddpm.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def predict_x0(self, model_output: mx.array, t: mx.array, sample: mx.array) -> mx.array:
    """Recover predicted clean sample x0 from the network output."""
    sqrt_acp, sqrt_omacp = self._sqrt_terms(t, sample.ndim)
    pt = self.config.prediction_type
    if pt == "epsilon":
        x0 = (sample - sqrt_omacp * model_output) / sqrt_acp
    elif pt == "sample":
        x0 = model_output
    elif pt == "v_prediction":
        x0 = sqrt_acp * sample - sqrt_omacp * model_output
    else:  # pragma: no cover - guarded in get_target
        raise ValueError(pt)
    if self.config.clip_sample:
        r = self.config.clip_sample_range
        x0 = mx.clip(x0, -r, r)
    return x0

mlx_diffuser.schedulers.DDIMScheduler

Bases: DDPMScheduler

Source code in src/mlx_diffuser/schedulers/ddim.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class DDIMScheduler(DDPMScheduler):
    config_class = DDIMConfig
    config: DDIMConfig

    def __init__(self, config: DDIMConfig | None = None):
        super().__init__(config or DDIMConfig())

    def step(
        self, model_output: mx.array, t: mx.array, sample: mx.array, key: mx.array | None = None
    ) -> mx.array:
        ti = int(t.item()) if isinstance(t, mx.array) else int(t)
        prev = ti - self._stride
        acp_t = self.alphas_cumprod[ti]
        acp_prev = self.alphas_cumprod[prev] if prev >= 0 else mx.array(1.0)

        x0 = self.predict_x0(model_output, mx.array(ti), sample)
        pred_eps = (sample - mx.sqrt(acp_t) * x0) / mx.sqrt(1.0 - acp_t)

        eta = self.config.eta
        sigma = (
            (eta * mx.sqrt((1.0 - acp_prev) / (1.0 - acp_t)) * mx.sqrt(1.0 - acp_t / acp_prev))
            if prev >= 0
            else mx.array(0.0)
        )

        dir_xt = mx.sqrt(mx.maximum(1.0 - acp_prev - sigma**2, 0.0)) * pred_eps
        prev_sample = mx.sqrt(acp_prev) * x0 + dir_xt

        if eta > 0 and prev >= 0:
            key = key if key is not None else mx.random.key(ti)
            prev_sample = prev_sample + sigma * mx.random.normal(sample.shape, key=key)
        return prev_sample

mlx_diffuser.schedulers.EulerDiscreteScheduler

Bases: DDPMScheduler

Source code in src/mlx_diffuser/schedulers/euler.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class EulerDiscreteScheduler(DDPMScheduler):
    config_class = EulerConfig
    config: EulerConfig

    def __init__(self, config: EulerConfig | None = None):
        super().__init__(config or EulerConfig())
        # Full-resolution sigma table over the training schedule.
        self._train_sigmas = mx.sqrt((1.0 - self.alphas_cumprod) / self.alphas_cumprod)
        self.sigmas: mx.array | None = None

    def set_timesteps(self, num_inference_steps: int) -> None:
        T = self.config.num_train_timesteps
        self.num_inference_steps = num_inference_steps
        train_sigmas = np.array(self._train_sigmas)  # ascending in index
        if self.config.timestep_spacing == "leading":
            step_ratio = T // num_inference_steps
            ts = (np.arange(num_inference_steps) * step_ratio).round()[::-1].astype(np.float64)
            ts = ts + self.config.steps_offset
        else:  # "linspace": fractional timesteps high->low
            ts = np.linspace(0, T - 1, num_inference_steps)[::-1].copy()
        sigmas = np.interp(ts, np.arange(T), train_sigmas)
        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
        self.sigmas = mx.array(sigmas)
        self.timesteps = mx.array(ts.astype(np.float32))
        self._step_index = 0

    @property
    def init_noise_sigma(self) -> float:
        """Std-dev for the initial latent noise (matches diffusers EulerDiscrete)."""
        assert self.sigmas is not None, "call set_timesteps() before sampling"
        sigma_max = float(self.sigmas[0])
        if self.config.timestep_spacing in ("linspace", "trailing"):
            return sigma_max
        return (sigma_max**2 + 1.0) ** 0.5

    def scale_model_input(self, sample: mx.array, t: mx.array) -> mx.array:
        assert self.sigmas is not None, "call set_timesteps() before sampling"
        sigma = self.sigmas[self._step_index]
        return sample / mx.sqrt(sigma**2 + 1.0)

    def _pred_original(self, model_output: mx.array, sigma: mx.array, sample: mx.array) -> mx.array:
        pt = self.config.prediction_type
        if pt == "epsilon":
            return sample - sigma * model_output
        if pt == "sample":
            return model_output
        if pt == "v_prediction":
            return model_output * (-sigma / mx.sqrt(sigma**2 + 1.0)) + sample / (sigma**2 + 1.0)
        raise ValueError(f"EulerDiscreteScheduler does not support prediction_type={pt!r}.")

    def step(
        self, model_output: mx.array, t: mx.array, sample: mx.array, key: mx.array | None = None
    ) -> mx.array:
        assert self.sigmas is not None, "call set_timesteps() before sampling"
        i = self._step_index
        sigma = self.sigmas[i]
        sigma_next = self.sigmas[i + 1]
        pred_original = self._pred_original(model_output, sigma, sample)
        derivative = (sample - pred_original) / sigma
        dt = sigma_next - sigma
        self._step_index += 1
        return sample + derivative * dt

    def add_noise_sigma(self, x0: mx.array, noise: mx.array, sigma: mx.array) -> mx.array:
        """VE-style corruption used by img2img-style sampling starts."""
        return x0 + expand_to(sigma, x0.ndim) * noise

init_noise_sigma property

Std-dev for the initial latent noise (matches diffusers EulerDiscrete).

add_noise_sigma(x0, noise, sigma)

VE-style corruption used by img2img-style sampling starts.

Source code in src/mlx_diffuser/schedulers/euler.py 
88
89
90
def add_noise_sigma(self, x0: mx.array, noise: mx.array, sigma: mx.array) -> mx.array:
    """VE-style corruption used by img2img-style sampling starts."""
    return x0 + expand_to(sigma, x0.ndim) * noise

mlx_diffuser.schedulers.FlowMatchEulerScheduler

Bases: Scheduler

Source code in src/mlx_diffuser/schedulers/flow_match_euler.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class FlowMatchEulerScheduler(Scheduler):
    config_class = FlowMatchConfig
    config: FlowMatchConfig

    def __init__(self, config: FlowMatchConfig | None = None):
        super().__init__(config or FlowMatchConfig())
        self.sigmas: mx.array | None = None

    def _shift(self, sigma: mx.array) -> mx.array:
        s = self.config.shift
        if s == 1.0:
            return sigma
        return s * sigma / (1.0 + (s - 1.0) * sigma)

    # --- training ---------------------------------------------------------
    def sample_timesteps(self, batch_size: int, key: mx.array) -> mx.array:
        u = mx.random.uniform(shape=(batch_size,), key=key)
        return self._shift(u)

    def add_noise(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        sigma = expand_to(t, x0.ndim)
        return (1.0 - sigma) * x0 + sigma * noise

    def get_target(self, x0: mx.array, noise: mx.array, t: mx.array) -> mx.array:
        if self.config.prediction_type != "velocity":
            raise ValueError("FlowMatchEulerScheduler only supports prediction_type='velocity'.")
        return noise - x0

    # --- sampling ---------------------------------------------------------
    def set_timesteps(self, num_inference_steps: int) -> None:
        self.num_inference_steps = num_inference_steps
        # Match the SD3/Flux/WAN flow convention: space sigmas over
        # ``[1, 1/num_train_timesteps]`` (so the last model evaluation is at very
        # low noise), apply the resolution shift, then append a final sigma of 0.
        # Sampling at a coarse final sigma instead leaves the sample under-denoised.
        sigma_min = 1.0 / self.config.num_train_timesteps
        sigmas = mx.linspace(1.0, sigma_min, num_inference_steps)
        sigmas = self._shift(sigmas)
        self.sigmas = mx.concatenate([sigmas, mx.zeros((1,))])
        self.timesteps = sigmas
        self._step_index = 0

    def set_sigmas(self, sigmas: mx.array) -> None:
        """Drive the integrator from an externally-computed sigma schedule (FLUX / SD3).

        ``sigmas`` is the descending list of flow times in ``[0, 1]`` (one per step, high
        noise first); a terminal ``0`` is appended automatically. The model is conditioned
        on each ``sigma`` directly. Used by pipelines that compute a resolution-dependent
        (``mu``-shifted) schedule themselves rather than the static-shift default.
        """
        self.num_inference_steps = int(sigmas.shape[0])
        self.sigmas = mx.concatenate([sigmas, mx.zeros((1,))])
        self.timesteps = sigmas
        self._step_index = 0

    def step(
        self, model_output: mx.array, t: mx.array, sample: mx.array, key: mx.array | None = None
    ) -> mx.array:
        assert self.sigmas is not None, "call set_timesteps() before sampling"
        i = self._step_index
        sigma = self.sigmas[i]
        sigma_next = self.sigmas[i + 1]
        self._step_index += 1
        return sample + (sigma_next - sigma) * model_output

set_sigmas(sigmas)

Drive the integrator from an externally-computed sigma schedule (FLUX / SD3).

sigmas is the descending list of flow times in [0, 1] (one per step, high noise first); a terminal 0 is appended automatically. The model is conditioned on each sigma directly. Used by pipelines that compute a resolution-dependent (mu-shifted) schedule themselves rather than the static-shift default.

Source code in src/mlx_diffuser/schedulers/flow_match_euler.py 
68
69
70
71
72
73
74
75
76
77
78
79
def set_sigmas(self, sigmas: mx.array) -> None:
    """Drive the integrator from an externally-computed sigma schedule (FLUX / SD3).

    ``sigmas`` is the descending list of flow times in ``[0, 1]`` (one per step, high
    noise first); a terminal ``0`` is appended automatically. The model is conditioned
    on each ``sigma`` directly. Used by pipelines that compute a resolution-dependent
    (``mu``-shifted) schedule themselves rather than the static-shift default.
    """
    self.num_inference_steps = int(sigmas.shape[0])
    self.sigmas = mx.concatenate([sigmas, mx.zeros((1,))])
    self.timesteps = sigmas
    self._step_index = 0

mlx_diffuser.schedulers.load_scheduler(path)

Load a scheduler from a directory containing config.json.

The concrete class is selected from the config's _class_name tag written at save time, falling back to DDPM when absent.

Source code in src/mlx_diffuser/schedulers/__init__.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
def load_scheduler(path: str | Path) -> Scheduler:
    """Load a scheduler from a directory containing ``config.json``.

    The concrete class is selected from the config's ``_class_name`` tag written
    at save time, falling back to DDPM when absent.
    """
    path = Path(path)
    config_path = path / CONFIG_NAME if path.is_dir() else path
    data = json.loads(config_path.read_text())
    cls = SCHEDULERS.get(data.get("_class_name", "DDPMConfig"))
    if cls is None:
        raise ValueError(f"Unknown scheduler config {data.get('_class_name')!r}.")
    return cls(cls.config_class.from_dict(data))