helena 0.1.0

use serde::{Deserialize, Serialize};

use crate::{Error, Result, Tensor};

use super::ops::{canonical_zero, combine};

/// A variance-preserving noise level.
///
/// The stored value is the signal variance `gamma = signal^2` in `[0, 1]`.
/// Signal and noise coefficients are derived from it, so invalid coefficient
/// pairs cannot be represented.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "NoiseLevelRepr")]
pub struct NoiseLevel {
    signal_variance: f32,
}

/// Deserialization mirror routed through [`NoiseLevel::from_signal_variance`].
#[derive(Deserialize)]
struct NoiseLevelRepr {
    signal_variance: f32,
}

impl TryFrom<NoiseLevelRepr> for NoiseLevel {
    type Error = Error;

    fn try_from(repr: NoiseLevelRepr) -> Result<Self> {
        Self::from_signal_variance(repr.signal_variance)
    }
}

impl NoiseLevel {
    /// Build a level from its signal variance `gamma = signal^2`.
    ///
    /// Returns [`Error::Validation`] unless `gamma` is finite and in `[0, 1]`.
    pub fn from_signal_variance(signal_variance: f32) -> Result<Self> {
        if !signal_variance.is_finite() || !(0.0..=1.0).contains(&signal_variance) {
            return Err(Error::validation(format!(
                "signal variance must be finite and in [0, 1], got {signal_variance}"
            )));
        }
        Ok(Self {
            signal_variance: canonical_zero(signal_variance),
        })
    }

    /// The `gamma = 1` endpoint: all signal, no noise.
    pub fn pure_signal() -> Self {
        Self {
            signal_variance: 1.0,
        }
    }

    /// The `gamma = 0` endpoint: all noise, no signal.
    pub fn pure_noise() -> Self {
        Self {
            signal_variance: 0.0,
        }
    }

    /// The signal variance `gamma = signal^2`.
    pub fn signal_variance(&self) -> f32 {
        self.signal_variance
    }

    /// The signal coefficient `sqrt(gamma)`.
    pub fn signal(&self) -> f32 {
        self.signal_variance.sqrt()
    }

    /// The noise coefficient `sqrt(1 - gamma)`.
    pub fn noise(&self) -> f32 {
        (1.0 - self.signal_variance).sqrt()
    }

    /// Apply the forward process: `x_t = signal * x0 + noise * eps`.
    ///
    /// `x0` and `eps` must share a shape.
    pub fn diffuse(&self, x0: &Tensor, eps: &Tensor) -> Result<Tensor> {
        combine(x0, eps, self.signal(), self.noise())
    }

    /// The v-prediction training target: `v = signal * eps - noise * x0`.
    ///
    /// `x0` and `eps` must share a shape.
    pub fn velocity_target(&self, x0: &Tensor, eps: &Tensor) -> Result<Tensor> {
        combine(eps, x0, self.signal(), -self.noise())
    }

    /// Recover the clean signal from a noisy sample and a velocity.
    ///
    /// Uses `x0 = signal * x_t - noise * v`. `x_t` and `v` must share a shape.
    pub fn signal_from_velocity(&self, x_t: &Tensor, v: &Tensor) -> Result<Tensor> {
        combine(x_t, v, self.signal(), -self.noise())
    }

    /// Recover the noise from a noisy sample and a velocity.
    ///
    /// Uses `eps = noise * x_t + signal * v`. `x_t` and `v` must share a shape.
    pub fn noise_from_velocity(&self, x_t: &Tensor, v: &Tensor) -> Result<Tensor> {
        combine(x_t, v, self.noise(), self.signal())
    }

    /// The Min-SNR-gamma loss weight for the velocity MSE at this level.
    ///
    /// A plain MSE on the velocity target implies an effective clean-target
    /// (`x0`) weight of `SNR + 1`, which diverges as `gamma -> 1` and so
    /// over-spends optimization on near-clean, high-SNR steps. Scaling the
    /// velocity MSE by this weight makes the effective clean-target weight
    /// exactly `min(SNR, clip)` (Hang et al. 2023, Min-SNR-gamma), the
    /// v-corrected form that divides the clip by `SNR + 1`.
    ///
    /// In the stored signal variance `gamma = SNR / (SNR + 1)` this is
    /// `min(gamma, clip * (1 - gamma))`: total and division-free on
    /// `gamma in [0, 1]`, zero at both endpoints, and peaking in the mid-SNR
    /// band a few-step sampler traverses coarsely.
    pub fn velocity_min_snr_weight(&self, clip: MinSnrGamma) -> f32 {
        let gamma = self.signal_variance;
        gamma.min(clip.get() * (1.0 - gamma))
    }
}

/// A validated Min-SNR-gamma clip for v-prediction loss weighting.
///
/// The clip caps the effective clean-target weight at high SNR; smaller values
/// clip more aggressively, larger values approach the unclipped noise-prediction
/// weighting. It must be finite and positive (a non-positive clip would zero or
/// invert the gradient on the high-SNR branch). [`Self::DEFAULT`] is the value
/// the originating literature recommends as a first try.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "f32")]
pub struct MinSnrGamma(f32);

impl MinSnrGamma {
    /// The Min-SNR paper's recommended starting clip.
    pub const DEFAULT: f32 = 5.0;

    /// Build a clip, requiring `clip` finite and strictly positive.
    pub fn new(clip: f32) -> Result<Self> {
        if !clip.is_finite() || clip <= 0.0 {
            return Err(Error::validation(format!(
                "min-SNR gamma clip must be finite and positive, got {clip}"
            )));
        }
        Ok(Self(clip))
    }

    /// The clip value.
    pub fn get(self) -> f32 {
        self.0
    }
}

impl TryFrom<f32> for MinSnrGamma {
    type Error = Error;

    fn try_from(clip: f32) -> Result<Self> {
        Self::new(clip)
    }
}