rustytorch/crates/models/rtx-diffuse/src/noise.rs

use crate::error::{DiffusionError, Result};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rtx_tensor::Tensor;

/// Noise scheduling strategies for diffusion models
#[derive(Debug, Clone)]
pub enum NoiseSchedule {
    /// Linear beta schedule from beta_start to beta_end
    Linear { beta_start: f32, beta_end: f32 },
    /// Cosine schedule for improved quality
    Cosine { s: f32 },
    /// Scaled linear schedule (used in Stable Diffusion)
    ScaledLinear { beta_start: f32, beta_end: f32 },
}

/// Manages noise generation and scheduling for diffusion processes
#[derive(Debug, Clone)]
pub struct NoiseGenerator {
    schedule: NoiseSchedule,
    num_timesteps: u32,
    betas: Vec<f32>,
    alphas: Vec<f32>,
    alpha_cumprod: Vec<f32>,
    alpha_cumprod_prev: Vec<f32>,
    sqrt_alpha_cumprod: Vec<f32>,
    sqrt_one_minus_alpha_cumprod: Vec<f32>,
    rng: StdRng,
}

impl NoiseGenerator {
    /// Create a new noise generator with the specified schedule
    pub fn new(schedule: NoiseSchedule, num_timesteps: u32, seed: Option<u64>) -> Result<Self> {
        if num_timesteps == 0 {
            return Err(DiffusionError::InvalidNoiseSchedule {
                reason: "Number of timesteps must be greater than 0".to_string(),
            });
        }

        let betas = Self::compute_betas(&schedule, num_timesteps)?;
        let alphas: Vec<f32> = betas.iter().map(|beta| 1.0 - beta).collect();

        let mut alpha_cumprod = Vec::with_capacity(num_timesteps as usize);
        let mut cumulative = 1.0;
        for alpha in &alphas {
            cumulative *= alpha;
            alpha_cumprod.push(cumulative);
        }

        let mut alpha_cumprod_prev = vec![1.0];
        alpha_cumprod_prev.extend_from_slice(&alpha_cumprod[..alpha_cumprod.len() - 1]);

        let sqrt_alpha_cumprod: Vec<f32> = alpha_cumprod.iter().map(|x| x.sqrt()).collect();
        let sqrt_one_minus_alpha_cumprod: Vec<f32> =
            alpha_cumprod.iter().map(|x| (1.0 - x).sqrt()).collect();

        let rng = match seed {
            Some(s) => StdRng::seed_from_u64(s),
            None => StdRng::from_entropy(),
        };

        Ok(Self {
            schedule,
            num_timesteps,
            betas,
            alphas,
            alpha_cumprod,
            alpha_cumprod_prev,
            sqrt_alpha_cumprod,
            sqrt_one_minus_alpha_cumprod,
            rng,
        })
    }

    /// Generate random noise with the same shape as the input tensor
    pub fn generate_noise(&mut self, shape: &[usize]) -> Result<Tensor> {
        let total_size = shape.iter().product::<usize>();
        let mut data = Vec::with_capacity(total_size);

        // Generate standard normal distribution
        for _ in 0..total_size {
            data.push(self.rng.sample(rand_distr::StandardNormal));
        }

        Tensor::new(data, shape.to_vec()).map_err(DiffusionError::Tensor)
    }

    /// Add noise to clean data according to the forward diffusion process
    pub fn add_noise(&self, x0: &Tensor, noise: &Tensor, timestep: u32) -> Result<Tensor> {
        if timestep >= self.num_timesteps {
            return Err(DiffusionError::InvalidTimestep {
                timestep: timestep as f32,
                max_timesteps: self.num_timesteps,
            });
        }

        let t = timestep as usize;
        let sqrt_alpha_cumprod = self.sqrt_alpha_cumprod[t];
        let sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alpha_cumprod[t];

        // x_t = sqrt(alpha_cumprod) * x_0 + sqrt(1 - alpha_cumprod) * noise
        let scaled_x0 = x0.scalar_mul(sqrt_alpha_cumprod)?;
        let scaled_noise = noise.scalar_mul(sqrt_one_minus_alpha_cumprod)?;

        scaled_x0.add(&scaled_noise).map_err(DiffusionError::Tensor)
    }

    /// Get the variance at a specific timestep
    pub fn get_variance(&self, timestep: u32) -> Result<f32> {
        if timestep >= self.num_timesteps {
            return Err(DiffusionError::InvalidTimestep {
                timestep: timestep as f32,
                max_timesteps: self.num_timesteps,
            });
        }

        let t = timestep as usize;
        if t == 0 {
            return Ok(0.0);
        }

        let beta_t = self.betas[t];
        let alpha_cumprod_t = self.alpha_cumprod[t];
        let alpha_cumprod_prev = self.alpha_cumprod_prev[t];

        // Variance = beta_t * (1 - alpha_cumprod_prev) / (1 - alpha_cumprod_t)
        let variance = beta_t * (1.0 - alpha_cumprod_prev) / (1.0 - alpha_cumprod_t);
        Ok(variance)
    }

    /// Get noise schedule parameters
    pub fn get_schedule_params(&self, timestep: u32) -> Result<(f32, f32, f32, f32)> {
        if timestep >= self.num_timesteps {
            return Err(DiffusionError::InvalidTimestep {
                timestep: timestep as f32,
                max_timesteps: self.num_timesteps,
            });
        }

        let t = timestep as usize;
        Ok((
            self.sqrt_alpha_cumprod[t],
            self.sqrt_one_minus_alpha_cumprod[t],
            self.alpha_cumprod[t],
            self.alpha_cumprod_prev[t],
        ))
    }

    fn compute_betas(schedule: &NoiseSchedule, num_timesteps: u32) -> Result<Vec<f32>> {
        match schedule {
            NoiseSchedule::Linear {
                beta_start,
                beta_end,
            } => {
                if *beta_start <= 0.0 || *beta_end <= 0.0 || beta_start >= beta_end {
                    return Err(DiffusionError::InvalidNoiseSchedule {
                        reason: "Linear schedule requires 0 < beta_start < beta_end".to_string(),
                    });
                }

                let mut betas = Vec::with_capacity(num_timesteps as usize);
                for i in 0..num_timesteps {
                    let t = i as f32 / (num_timesteps - 1) as f32;
                    let beta = beta_start + t * (beta_end - beta_start);
                    betas.push(beta);
                }
                Ok(betas)
            }

            NoiseSchedule::Cosine { s } => {
                if *s <= 0.0 {
                    return Err(DiffusionError::InvalidNoiseSchedule {
                        reason: "Cosine schedule requires s > 0".to_string(),
                    });
                }

                let mut betas = Vec::with_capacity(num_timesteps as usize);
                let f = |t: f32| -> f32 {
                    let angle = std::f32::consts::PI / 2.0 * ((t + s) / (1.0 + s));
                    angle.cos().powi(2)
                };

                for i in 0..num_timesteps {
                    let t1 = i as f32 / num_timesteps as f32;
                    let t2 = (i + 1) as f32 / num_timesteps as f32;
                    let beta = 1.0 - f(t2) / f(t1);
                    betas.push(beta.min(0.999)); // Clamp to avoid numerical issues
                }
                Ok(betas)
            }

            NoiseSchedule::ScaledLinear {
                beta_start,
                beta_end,
            } => {
                if *beta_start <= 0.0 || *beta_end <= 0.0 || beta_start >= beta_end {
                    return Err(DiffusionError::InvalidNoiseSchedule {
                        reason: "Scaled linear schedule requires 0 < beta_start < beta_end"
                            .to_string(),
                    });
                }

                let mut betas = Vec::with_capacity(num_timesteps as usize);
                for i in 0..num_timesteps {
                    let t = i as f32 / (num_timesteps - 1) as f32;
                    let scaled_beta = beta_start.sqrt() + t * (beta_end.sqrt() - beta_start.sqrt());
                    betas.push(scaled_beta * scaled_beta);
                }
                Ok(betas)
            }
        }
    }

    pub fn num_timesteps(&self) -> u32 {
        self.num_timesteps
    }
}