| """Flow Matching Euler Discrete Scheduler for Nucleus-Image.""" | |
| import mlx.core as mx | |
| class FlowMatchEulerScheduler: | |
| def __init__(self, shift: float = 1.0, num_train_timesteps: int = 1000): | |
| self.shift = shift | |
| self.num_train_timesteps = num_train_timesteps | |
| self.sigmas = None | |
| self.timesteps = None | |
| def set_timesteps(self, num_inference_steps: int): | |
| sigmas = mx.linspace(1.0, 0.0, num_inference_steps + 1) | |
| if self.shift != 1.0: | |
| sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) | |
| self.sigmas = sigmas | |
| self.timesteps = sigmas[:-1] * self.num_train_timesteps | |
| def step(self, model_output, timestep_idx: int, sample): | |
| sigma = self.sigmas[timestep_idx] | |
| sigma_next = self.sigmas[timestep_idx + 1] | |
| dt = sigma_next - sigma | |
| return sample + dt * model_output | |