anima_mlx/runtime/sampling.py

"""Sampling glue for Anima MLX flow models."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Sequence

from .scheduler import FlowSchedulerConfig, simple_sigmas, time_snr_shift


@dataclass(frozen=True)
class FlowStepResult:
    model_output: Any
    denoised: Any
    next_latent: Any
    sigma: Any
    next_sigma: Any
    timestep: Any


@dataclass(frozen=True)
class FlowLoopResult:
    latent: Any
    steps: tuple[FlowStepResult, ...]


def flow_timestep(sigma: Any, config: FlowSchedulerConfig = FlowSchedulerConfig()) -> Any:
    return sigma * config.multiplier


def flow_denoised(latent: Any, model_output: Any, sigma: Any) -> Any:
    return latent - model_output * _reshape_sigma(sigma, latent)


def flow_euler_step(latent: Any, model_output: Any, sigma: Any, next_sigma: Any) -> Any:
    return latent + model_output * _reshape_sigma(next_sigma - sigma, latent)


def flow_one_step(
    latent: Any,
    model_output: Any,
    sigma: Any,
    next_sigma: Any,
    config: FlowSchedulerConfig = FlowSchedulerConfig(),
) -> FlowStepResult:
    timestep = flow_timestep(sigma, config)
    denoised = flow_denoised(latent, model_output, sigma)
    next_latent = flow_euler_step(latent, model_output, sigma, next_sigma)
    return FlowStepResult(
        model_output=model_output,
        denoised=denoised,
        next_latent=next_latent,
        sigma=sigma,
        next_sigma=next_sigma,
        timestep=timestep,
    )


def flow_euler_loop(
    latent: Any,
    sigmas: Sequence[Any],
    predict_fn: Callable[[Any, Any], Any],
    config: FlowSchedulerConfig = FlowSchedulerConfig(),
    *,
    store_steps: bool = True,
) -> FlowLoopResult:
    import mlx.core as mx

    if len(sigmas) < 2:
        raise ValueError("sigmas must contain at least two values")
    current = latent
    results: list[FlowStepResult] = []
    for index in range(len(sigmas) - 1):
        sigma = _as_mx_scalar(sigmas[index])
        next_sigma = _as_mx_scalar(sigmas[index + 1])
        model_output = predict_fn(current, sigma)
        if store_steps:
            result = flow_one_step(current, model_output, sigma, next_sigma, config=config)
            mx.eval(result.next_latent)
            results.append(result)
            current = result.next_latent
            continue
        current = flow_euler_step(current, model_output, sigma, next_sigma)
        mx.eval(current)
    return FlowLoopResult(latent=current, steps=tuple(results))


def er_sde_loop(
    latent: Any,
    sigmas: Sequence[Any],
    predict_fn: Callable[[Any, Any], Any],
    config: FlowSchedulerConfig = FlowSchedulerConfig(),
    *,
    seed: int | None = None,
    s_noise: float = 1.0,
    max_stage: int = 3,
    num_integration_points: int = 200,
    store_steps: bool = True,
) -> FlowLoopResult:
    """Run ComfyUI-style ER-SDE sampling for flow/CONST prediction models.

    `predict_fn` returns the raw flow model output. The ER-SDE update itself uses
    ComfyUI's denoised/x0 wrapper semantics, so this function converts raw model
    output with `flow_denoised()` before applying the solver update.
    """

    import mlx.core as mx

    if len(sigmas) < 2:
        raise ValueError("sigmas must contain at least two values")
    if max_stage <= 0:
        raise ValueError("max_stage must be positive")
    if num_integration_points <= 0:
        raise ValueError("num_integration_points must be positive")

    sigma_values = _offset_first_sigma_for_flow_snr([_as_float(value) for value in sigmas], config)
    er_lambdas = [_er_lambda_from_sigma_value(value) for value in sigma_values]
    point_indices = mx.arange(num_integration_points).astype(mx.float32)
    noise_sampler = _normal_noise_sampler(latent, seed)

    current = latent
    old_denoised = None
    old_denoised_d = None
    results: list[FlowStepResult] = []

    for index in range(len(sigma_values) - 1):
        sigma = _as_mx_scalar(sigma_values[index])
        next_sigma = _as_mx_scalar(sigma_values[index + 1])
        model_output = predict_fn(current, sigma)
        denoised = flow_denoised(current, model_output, sigma)

        if sigma_values[index + 1] == 0.0:
            next_latent = denoised
        else:
            er_lambda_s = _as_mx_scalar(er_lambdas[index])
            er_lambda_t = _as_mx_scalar(er_lambdas[index + 1])
            alpha_s = sigma / er_lambda_s
            alpha_t = next_sigma / er_lambda_t
            r_alpha = alpha_t / alpha_s
            r = _er_sde_noise_scaler(er_lambda_t) / _er_sde_noise_scaler(er_lambda_s)

            next_latent = r_alpha * r * current + alpha_t * (1.0 - r) * denoised

            stage_used = min(max_stage, index + 1)
            if stage_used >= 2:
                dt = er_lambda_t - er_lambda_s
                lambda_step_size = -dt / float(num_integration_points)
                lambda_pos = er_lambda_t + point_indices * lambda_step_size
                scaled_pos = _er_sde_noise_scaler(lambda_pos)

                s = mx.sum(1.0 / scaled_pos) * lambda_step_size
                previous_lambda = _as_mx_scalar(er_lambdas[index - 1])
                denoised_d = (denoised - old_denoised) / (er_lambda_s - previous_lambda)
                next_latent = next_latent + alpha_t * (dt + s * _er_sde_noise_scaler(er_lambda_t)) * denoised_d

                if stage_used >= 3:
                    previous_previous_lambda = _as_mx_scalar(er_lambdas[index - 2])
                    s_u = mx.sum((lambda_pos - er_lambda_s) / scaled_pos) * lambda_step_size
                    denoised_u = (denoised_d - old_denoised_d) / ((er_lambda_s - previous_previous_lambda) / 2.0)
                    next_latent = next_latent + alpha_t * (
                        (dt * dt) / 2.0 + s_u * _er_sde_noise_scaler(er_lambda_t)
                    ) * denoised_u
                old_denoised_d = denoised_d

            if s_noise > 0:
                noise_scale = _sqrt_non_negative((er_lambda_t * er_lambda_t) - (er_lambda_s * er_lambda_s * r * r))
                next_latent = next_latent + alpha_t * noise_sampler(sigma, next_sigma) * float(s_noise) * noise_scale

        if store_steps:
            result = FlowStepResult(
                model_output=model_output,
                denoised=denoised,
                next_latent=next_latent,
                sigma=sigma,
                next_sigma=next_sigma,
                timestep=flow_timestep(sigma, config),
            )
            mx.eval(result.next_latent)
            results.append(result)
            current = result.next_latent
        else:
            current = next_latent
            mx.eval(current)
        old_denoised = denoised

    return FlowLoopResult(latent=current, steps=tuple(results))


def classifier_free_guidance(negative: Any, positive: Any, scale: float) -> Any:
    return negative + (positive - negative) * float(scale)


def sigmas_for_steps(steps: int, config: FlowSchedulerConfig = FlowSchedulerConfig()) -> list[float]:
    return simple_sigmas(steps, config)


def _reshape_sigma(sigma: Any, reference: Any) -> Any:
    if hasattr(sigma, "shape") and len(sigma.shape) > 0:
        return sigma.reshape(sigma.shape[:1] + (1,) * (len(reference.shape) - 1))
    return sigma


def _as_mx_scalar(value: Any) -> Any:
    import mlx.core as mx

    if hasattr(value, "shape"):
        return value
    return mx.array(float(value), dtype=mx.float32)


def _as_float(value: Any) -> float:
    if hasattr(value, "item"):
        return float(value.item())
    return float(value)


def _offset_first_sigma_for_flow_snr(sigmas: list[float], config: FlowSchedulerConfig, percent_offset: float = 1e-4) -> list[float]:
    if len(sigmas) > 1 and sigmas[0] >= 1.0:
        sigmas = sigmas.copy()
        sigmas[0] = time_snr_shift(config.shift, 1.0 - percent_offset)
    return sigmas


def _er_lambda_from_sigma_value(sigma: float) -> float:
    if sigma >= 1.0:
        raise ValueError("ER-SDE requires sigma < 1 after first-sigma SNR offset")
    if sigma <= 0.0:
        return 0.0
    return sigma / (1.0 - sigma)


def _er_sde_noise_scaler(value: Any) -> Any:
    import mlx.core as mx

    return value * (mx.exp(value**0.3) + 10.0)


def _sqrt_non_negative(value: Any) -> Any:
    import mlx.core as mx

    return mx.sqrt(mx.maximum(value, mx.zeros_like(value)))


def _normal_noise_sampler(reference: Any, seed: int | None) -> Callable[[Any, Any], Any]:
    import numpy as np
    import mlx.core as mx

    rng = np.random.default_rng(None if seed is None else int(seed))

    def sample(_sigma: Any, _next_sigma: Any) -> Any:
        noise = rng.standard_normal(reference.shape).astype(np.float32)
        return mx.array(noise, dtype=reference.dtype)

    return sample