56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
import mlx.core as mx
|
|
|
|
if TYPE_CHECKING:
|
|
from mflux.config.runtime_config import RuntimeConfig
|
|
|
|
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
|
|
|
|
|
|
class LinearScheduler(BaseScheduler):
|
|
def __init__(self, runtime_config: "RuntimeConfig"):
|
|
self.runtime_config = runtime_config
|
|
self._sigmas = self._get_sigmas()
|
|
self._timesteps = self._get_timesteps()
|
|
|
|
@property
|
|
def sigmas(self) -> mx.array:
|
|
return self._sigmas
|
|
|
|
@property
|
|
def timesteps(self) -> mx.array:
|
|
return self._timesteps
|
|
|
|
def _get_sigmas(self) -> mx.array:
|
|
model_config = self.runtime_config.model_config
|
|
sigmas = mx.linspace(
|
|
1.0,
|
|
1.0 / self.runtime_config.num_inference_steps,
|
|
self.runtime_config.num_inference_steps,
|
|
)
|
|
sigmas = mx.array(sigmas).astype(mx.float32)
|
|
sigmas = mx.concatenate([sigmas, mx.zeros(1)])
|
|
if model_config.requires_sigma_shift:
|
|
y1 = 0.5
|
|
x1 = 256
|
|
m = (1.15 - y1) / (4096 - x1)
|
|
b = y1 - m * x1
|
|
mu = m * self.runtime_config.width * self.runtime_config.height / 256 + b
|
|
mu = mx.array(mu)
|
|
shifted_sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / sigmas - 1))
|
|
shifted_sigmas[-1] = 0
|
|
return shifted_sigmas
|
|
else:
|
|
return sigmas
|
|
|
|
def _get_timesteps(self) -> mx.array:
|
|
num_steps = self.runtime_config.num_inference_steps
|
|
timesteps = mx.arange(num_steps, dtype=mx.float32)
|
|
|
|
return timesteps
|
|
|
|
def step(self, model_output: mx.array, timestep: int, sample: mx.array, **kwargs) -> mx.array:
|
|
dt = self._sigmas[timestep + 1] - self._sigmas[timestep]
|
|
return sample + model_output * dt
|