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
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

    # --- 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_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 
24
25
26
27
28
29
30
31
32
33
34
35
36
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))