rustytorch/crates/models/rtx-csm/src/cfg_schedule.rs

//! Per-frame CFG scale schedules (Selective CFG, arXiv 2509.19668).
//!
//! Standard CFG uses a fixed `cfg_scale` for every frame. Zheng & Maleki
//! show that for zero-shot TTS, **frame-dependent** CFG schedules
//! preserve text adherence late in synthesis while keeping speaker
//! similarity high early on:
//!
//!   - Early frames (set the speaker character) → full CFG scale.
//!   - Later frames (need to commit to the text) → lower scale, since
//!     too much CFG late in generation causes over-conditioning on the
//!     speaker prompt and degrades intelligibility.
//!
//! ## Variants
//!
//! - [`CfgSchedule::Constant`] — back-compat with the old fixed-`f64`
//!   API; `cfg_scale.unwrap_or(1.0)` translates to this.
//! - [`CfgSchedule::Step`] — `scale_early` until frame `transition`,
//!   then jump to `scale_late`. Matches the paper's two-phase recipe.
//! - [`CfgSchedule::LinearRamp`] — smooth interpolation from `start` at
//!   frame 0 to `end` at `ramp_frames`, then holds at `end`. Sometimes
//!   sounds less abrupt than the step variant.

use crate::error::{CsmError, Result};

#[derive(Debug, Clone, Copy)]
pub enum CfgSchedule {
    Constant(f64),
    Step {
        early: f64,
        late: f64,
        transition: usize,
    },
    LinearRamp {
        start: f64,
        end: f64,
        ramp_frames: usize,
    },
}

impl CfgSchedule {
    pub fn scale_at(&self, frame_idx: usize) -> f64 {
        match *self {
            Self::Constant(s) => s,
            Self::Step {
                early,
                late,
                transition,
            } => {
                if frame_idx < transition {
                    early
                } else {
                    late
                }
            }
            Self::LinearRamp {
                start,
                end,
                ramp_frames,
            } => {
                if ramp_frames == 0 || frame_idx >= ramp_frames {
                    end
                } else {
                    let t = frame_idx as f64 / ramp_frames as f64;
                    start * (1.0 - t) + end * t
                }
            }
        }
    }

    /// True if any frame would invoke the CFG path (any scale > 1.0).
    pub fn is_active(&self) -> bool {
        match *self {
            Self::Constant(s) => s > 1.0,
            Self::Step { early, late, .. } => early > 1.0 || late > 1.0,
            Self::LinearRamp { start, end, .. } => start > 1.0 || end > 1.0,
        }
    }

    /// Parse the `--cfg-schedule` CLI value. Forms:
    ///   - `step:<early>:<late>:<transition>`
    ///   - `linear:<start>:<end>:<ramp_frames>`
    ///   - `const:<scale>`
    pub fn parse(s: &str) -> Result<Self> {
        let err = || {
            CsmError::Config(format!(
                "invalid --cfg-schedule '{s}'; expected step:E:L:T, linear:S:E:R, or const:X"
            ))
        };
        let mut parts = s.split(':');
        let kind = parts.next().ok_or_else(err)?;
        match kind {
            "const" => {
                let v: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                Ok(Self::Constant(v))
            }
            "step" => {
                let e: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                let l: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                let t: usize = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                Ok(Self::Step {
                    early: e,
                    late: l,
                    transition: t,
                })
            }
            "linear" => {
                let s0: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                let e: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                let r: usize = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
                Ok(Self::LinearRamp {
                    start: s0,
                    end: e,
                    ramp_frames: r,
                })
            }
            _ => Err(err()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn constant_returns_same_scale() {
        let s = CfgSchedule::Constant(2.0);
        assert_eq!(s.scale_at(0), 2.0);
        assert_eq!(s.scale_at(50), 2.0);
        assert!(s.is_active());
    }

    #[test]
    fn step_transitions_at_threshold() {
        let s = CfgSchedule::Step {
            early: 3.0,
            late: 1.5,
            transition: 10,
        };
        assert_eq!(s.scale_at(0), 3.0);
        assert_eq!(s.scale_at(9), 3.0);
        assert_eq!(s.scale_at(10), 1.5);
        assert_eq!(s.scale_at(100), 1.5);
        assert!(s.is_active());
    }

    #[test]
    fn linear_ramp_interpolates() {
        let s = CfgSchedule::LinearRamp {
            start: 3.0,
            end: 1.0,
            ramp_frames: 10,
        };
        assert_eq!(s.scale_at(0), 3.0);
        assert!((s.scale_at(5) - 2.0).abs() < 1e-9);
        assert_eq!(s.scale_at(10), 1.0);
        assert_eq!(s.scale_at(100), 1.0); // holds at end
    }

    #[test]
    fn parse_step() {
        let s = CfgSchedule::parse("step:3.0:1.5:12").unwrap();
        match s {
            CfgSchedule::Step {
                early,
                late,
                transition,
            } => {
                assert_eq!(early, 3.0);
                assert_eq!(late, 1.5);
                assert_eq!(transition, 12);
            }
            _ => panic!("wrong variant"),
        }
    }

    #[test]
    fn parse_linear() {
        let s = CfgSchedule::parse("linear:2.5:1.0:20").unwrap();
        match s {
            CfgSchedule::LinearRamp {
                start,
                end,
                ramp_frames,
            } => {
                assert_eq!(start, 2.5);
                assert_eq!(end, 1.0);
                assert_eq!(ramp_frames, 20);
            }
            _ => panic!("wrong variant"),
        }
    }

    #[test]
    fn parse_invalid() {
        assert!(CfgSchedule::parse("garbage").is_err());
        assert!(CfgSchedule::parse("step:1.5:1.0").is_err()); // missing transition
    }
}