21 lines
519 B
Python
21 lines
519 B
Python
import mlx.core as mx
|
|
|
|
from flux_1_schnell.config.config import Config
|
|
|
|
|
|
class FlowMatchEulerDiscreteNoiseScheduler:
|
|
|
|
@staticmethod
|
|
def denoise(
|
|
t: int,
|
|
noise: mx.array,
|
|
latent: mx.array,
|
|
config: Config,
|
|
) -> mx.array:
|
|
sigma = config.sigmas[t]
|
|
denoised = latent - noise * sigma
|
|
derivative = (latent - denoised) / sigma
|
|
dt = config.sigmas[t + 1] - sigma
|
|
prev_sample = latent + derivative * dt
|
|
return prev_sample
|