anima_mlx/runtime/scheduler.py

"""Small ComfyUI-compatible scheduler helpers.

The MLX pipeline should not depend on ComfyUI at runtime, but its first
verification target must match the ComfyUI source contract. These functions
mirror the flow-scheduler math used by ComfyUI's ModelSamplingDiscreteFlow and
the `simple` KSampler scheduler.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class FlowSchedulerConfig:
    shift: float = 3.0
    multiplier: float = 1.0
    timesteps: int = 1000


def time_snr_shift(alpha: float, t: float) -> float:
    if alpha == 1.0:
        return t
    return alpha * t / (1.0 + (alpha - 1.0) * t)


def sigma_from_timestep(timestep: float, config: FlowSchedulerConfig = FlowSchedulerConfig()) -> float:
    return time_snr_shift(config.shift, timestep / config.multiplier)


def flow_sigmas(config: FlowSchedulerConfig = FlowSchedulerConfig()) -> list[float]:
    return [
        sigma_from_timestep(((index + 1) / config.timesteps) * config.multiplier, config)
        for index in range(config.timesteps)
    ]


def simple_sigmas(steps: int, config: FlowSchedulerConfig = FlowSchedulerConfig()) -> list[float]:
    if steps <= 0:
        raise ValueError("steps must be positive")
    sigmas = flow_sigmas(config)
    stride = len(sigmas) / steps
    sampled = [sigmas[-(1 + int(index * stride))] for index in range(steps)]
    sampled.append(0.0)
    return sampled