helena 0.1.0

Core types and component interfaces for helena, a latent data-to-waveform generation platform.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use std::f32::consts::FRAC_PI_2;

use serde::{Deserialize, Serialize};

use crate::{Error, Result};

use super::{level::NoiseLevel, ops::canonical_zero};

/// The cosine noise schedule from Nichol and Dhariwal.
///
/// `gamma(t)` falls from `1` at `t = 0` to effectively `0` at `t = 1`, with a
/// validated offset to avoid degenerate paths.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "CosineScheduleRepr")]
pub struct CosineSchedule {
    offset: f32,
}

/// Deserialization mirror routed through [`CosineSchedule::new`].
#[derive(Deserialize)]
struct CosineScheduleRepr {
    offset: f32,
}

impl TryFrom<CosineScheduleRepr> for CosineSchedule {
    type Error = Error;

    fn try_from(repr: CosineScheduleRepr) -> Result<Self> {
        Self::new(repr.offset)
    }
}

impl Default for CosineSchedule {
    fn default() -> Self {
        Self {
            offset: Self::DEFAULT_OFFSET,
        }
    }
}

impl CosineSchedule {
    /// Offset from the original cosine schedule paper.
    pub const DEFAULT_OFFSET: f32 = 0.008;

    /// Upper bound that keeps the raised cosine well conditioned in `f32`.
    pub const MAX_OFFSET: f32 = 1.0;

    /// Build a schedule with the given `offset`.
    ///
    /// Returns [`Error::Validation`] unless `offset` is finite and in
    /// `[0, MAX_OFFSET]`.
    pub fn new(offset: f32) -> Result<Self> {
        if !offset.is_finite() || !(0.0..=Self::MAX_OFFSET).contains(&offset) {
            return Err(Error::validation(format!(
                "cosine schedule offset must be finite and in [0, {}], got {offset}",
                Self::MAX_OFFSET
            )));
        }
        Ok(Self {
            offset: canonical_zero(offset),
        })
    }

    /// The signal variance `gamma(t)` at normalized time `t`.
    ///
    /// Returns [`Error::Validation`] unless `t` is finite and in `[0, 1]`.
    pub fn signal_variance(&self, t: f32) -> Result<f32> {
        if !t.is_finite() || !(0.0..=1.0).contains(&t) {
            return Err(Error::validation(format!(
                "schedule time must be finite and in [0, 1], got {t}"
            )));
        }
        Ok((self.raised_cos(t) / self.raised_cos(0.0)).clamp(0.0, 1.0))
    }

    /// The [`NoiseLevel`] at normalized time `t`.
    pub fn level(&self, t: f32) -> Result<NoiseLevel> {
        NoiseLevel::from_signal_variance(self.signal_variance(t)?)
    }

    fn raised_cos(&self, t: f32) -> f32 {
        (((t + self.offset) / (1.0 + self.offset)) * FRAC_PI_2)
            .cos()
            .powi(2)
    }
}