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

//! DDIM (Denoising Diffusion Implicit Models) sampler implementation
//!
//! Provides deterministic and stochastic sampling with improved efficiency

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

/// Configuration for DDIM sampler
#[derive(Debug, Clone)]
pub struct DDIMConfig {
    /// Number of denoising steps
    pub num_inference_steps: usize,
    /// Stochasticity parameter (0 = deterministic, 1 = DDPM)
    pub eta: f32,
    /// Whether to clip samples to [-1, 1]
    pub clip_sample: bool,
    /// Set alpha to 1 for the final step
    pub set_alpha_to_one: bool,
    /// Skip Pseudo Runge-Kutta steps
    pub skip_prk_steps: bool,
}

impl Default for DDIMConfig {
    fn default() -> Self {
        Self {
            num_inference_steps: 50,
            eta: 0.0,
            clip_sample: true,
            set_alpha_to_one: false,
            skip_prk_steps: true,
        }
    }
}

/// DDIM sampler for efficient diffusion sampling
#[derive(Debug, Clone)]
pub struct DDIMSampler {
    config: DDIMConfig,
    num_train_timesteps: usize,
    timesteps: Vec<usize>,
    alphas_cumprod: Vec<f32>,
    device: Device,
}

impl DDIMSampler {
    /// Create a new DDIM sampler
    pub fn new(config: DDIMConfig, num_train_timesteps: usize, device: &Device) -> Result<Self> {
        if config.num_inference_steps == 0 {
            return Err(DiffusionError::Scheduler {
                message: "Number of inference steps must be greater than 0".to_string(),
            });
        }

        // Compute timesteps for inference
        let timesteps = Self::compute_timesteps(num_train_timesteps, config.num_inference_steps);

        // Compute alpha schedule
        let alphas_cumprod = Self::compute_alphas_cumprod(num_train_timesteps);

        Ok(Self {
            config,
            num_train_timesteps,
            timesteps,
            alphas_cumprod,
            device: device.clone(),
        })
    }

    /// Create sampler with rescaled betas for zero terminal SNR
    pub fn with_rescaled_betas(
        config: DDIMConfig,
        num_train_timesteps: usize,
        device: &Device,
    ) -> Result<Self> {
        let mut sampler = Self::new(config, num_train_timesteps, device)?;
        sampler.rescale_betas_for_zero_snr();
        Ok(sampler)
    }

    /// Sample from the model
    pub fn sample<F>(&self, initial_noise: &Tensor, model_fn: F) -> Result<Tensor>
    where
        F: FnMut(&Tensor, usize) -> Result<Tensor>,
    {
        self.sample_with_seed(initial_noise, model_fn, None)
    }

    /// Sample with specific random seed
    pub fn sample_with_seed<F>(
        &self,
        initial_noise: &Tensor,
        mut model_fn: F,
        seed: Option<u64>,
    ) -> Result<Tensor>
    where
        F: FnMut(&Tensor, usize) -> Result<Tensor>,
    {
        let mut sample = initial_noise.clone();
        let mut rng = seed.map(|s| StdRng::seed_from_u64(s));

        // Reverse diffusion process
        for i in 0..self.timesteps.len() {
            let t = self.timesteps[i];

            // Get model prediction
            let model_output = model_fn(&sample, t)?;

            // Compute next sample
            sample = self.step(&sample, &model_output, t, rng.as_mut())?;
        }

        Ok(sample)
    }

    /// Perform one step of the reverse diffusion process
    pub fn step(
        &self,
        sample: &Tensor,
        model_output: &Tensor,
        timestep: usize,
        rng: Option<&mut StdRng>,
    ) -> Result<Tensor> {
        // Get current and previous timestep indices
        let curr_idx = self
            .timesteps
            .iter()
            .position(|&t| t == timestep)
            .ok_or_else(|| DiffusionError::Scheduler {
                message: format!("Timestep {} not found in schedule", timestep),
            })?;

        let prev_timestep = if curr_idx + 1 < self.timesteps.len() {
            self.timesteps[curr_idx + 1]
        } else {
            0
        };

        // Get alpha values
        let alpha_prod_t = self.alphas_cumprod[timestep];
        let alpha_prod_t_prev = if prev_timestep > 0 {
            self.alphas_cumprod[prev_timestep]
        } else if self.config.set_alpha_to_one {
            1.0
        } else {
            self.alphas_cumprod[0]
        };

        let beta_prod_t = 1.0 - alpha_prod_t;
        let beta_prod_t_prev = 1.0 - alpha_prod_t_prev;

        // Compute predicted original sample
        let pred_original =
            self.compute_predicted_original(sample, model_output, alpha_prod_t, beta_prod_t)?;

        // Clip if needed
        let pred_original = if self.config.clip_sample {
            self.clip_sample(&pred_original)?
        } else {
            pred_original
        };

        // Compute variance
        let variance = self.compute_variance(timestep, prev_timestep)?;
        let std_dev = variance.sqrt();

        // Compute direction pointing to x_t
        let pred_sample_direction = self.compute_pred_sample_direction(
            &pred_original,
            model_output,
            alpha_prod_t_prev,
            beta_prod_t_prev,
        )?;

        // Compute previous sample
        let mut prev_sample = pred_original
            .mul_scalar(alpha_prod_t_prev.sqrt())?
            .add(&pred_sample_direction.mul_scalar(beta_prod_t_prev.sqrt())?)?;

        // Add noise if eta > 0 (stochastic)
        if self.config.eta > 0.0 && rng.is_some() && timestep > 0 {
            let noise = Tensor::randn(sample.shape().dims(), &self.device)?;
            prev_sample = prev_sample.add(&noise.mul_scalar(std_dev)?)?;
        }

        Ok(prev_sample)
    }

    /// Compute predicted original sample from model output
    fn compute_predicted_original(
        &self,
        sample: &Tensor,
        model_output: &Tensor,
        alpha_prod_t: f32,
        beta_prod_t: f32,
    ) -> Result<Tensor> {
        // x_0 = (x_t - sqrt(1 - alpha_t) * eps) / sqrt(alpha_t)
        let scaled_sample = sample.div_scalar(alpha_prod_t.sqrt())?;
        let scaled_noise = model_output.mul_scalar(beta_prod_t.sqrt() / alpha_prod_t.sqrt())?;
        scaled_sample
            .sub(&scaled_noise)
            .map_err(DiffusionError::Tensor)
    }

    /// Compute direction pointing to x_t
    fn compute_pred_sample_direction(
        &self,
        pred_original: &Tensor,
        model_output: &Tensor,
        alpha_prod_t_prev: f32,
        beta_prod_t_prev: f32,
    ) -> Result<Tensor> {
        // Direction = sqrt(1 - alpha_t-1) * eps
        model_output
            .mul_scalar(beta_prod_t_prev.sqrt())
            .map_err(DiffusionError::Tensor)
    }

    /// Compute variance for the step
    pub fn compute_variance(&self, timestep: usize, prev_timestep: usize) -> Result<f32> {
        if self.config.eta == 0.0 {
            return Ok(0.0);
        }

        let alpha_prod_t = self.alphas_cumprod[timestep];
        let alpha_prod_t_prev = if prev_timestep > 0 {
            self.alphas_cumprod[prev_timestep]
        } else {
            1.0
        };

        let beta_prod_t = 1.0 - alpha_prod_t;
        let beta_prod_t_prev = 1.0 - alpha_prod_t_prev;

        // variance = eta^2 * (1 - alpha_t-1) / (1 - alpha_t) * (1 - alpha_t / alpha_t-1)
        let variance = self.config.eta.powi(2) * beta_prod_t_prev / beta_prod_t
            * (1.0 - alpha_prod_t / alpha_prod_t_prev);

        Ok(variance)
    }

    /// Clip sample values to reasonable range
    fn clip_sample(&self, sample: &Tensor) -> Result<Tensor> {
        // Clip to [-3, 3] which covers most of the Gaussian distribution
        sample.clamp(-3.0, 3.0).map_err(DiffusionError::Tensor)
    }

    /// Compute timesteps for inference
    fn compute_timesteps(num_train_timesteps: usize, num_inference_steps: usize) -> Vec<usize> {
        let step = num_train_timesteps / num_inference_steps;
        (0..num_inference_steps)
            .map(|i| num_train_timesteps - 1 - i * step)
            .collect()
    }

    /// Compute cumulative product of alphas
    fn compute_alphas_cumprod(num_train_timesteps: usize) -> Vec<f32> {
        // Linear beta schedule
        let beta_start = 0.0001;
        let beta_end = 0.02;

        let betas: Vec<f32> = (0..num_train_timesteps)
            .map(|i| {
                beta_start + (beta_end - beta_start) * (i as f32) / (num_train_timesteps - 1) as f32
            })
            .collect();

        let alphas: Vec<f32> = betas.iter().map(|b| 1.0 - b).collect();

        // Cumulative product
        let mut alphas_cumprod = vec![1.0; num_train_timesteps];
        alphas_cumprod[0] = alphas[0];
        for i in 1..num_train_timesteps {
            alphas_cumprod[i] = alphas_cumprod[i - 1] * alphas[i];
        }

        alphas_cumprod
    }

    /// Rescale betas for zero terminal SNR
    fn rescale_betas_for_zero_snr(&mut self) {
        // Ensure terminal alpha is close to 0
        let last_idx = self.alphas_cumprod.len() - 1;
        self.alphas_cumprod[last_idx] = 0.0001;

        // Recompute to maintain monotonicity
        for i in (1..last_idx).rev() {
            if self.alphas_cumprod[i] < self.alphas_cumprod[i + 1] {
                self.alphas_cumprod[i] = self.alphas_cumprod[i + 1] * 1.001;
            }
        }
    }

    // Getter methods for testing
    pub fn num_inference_steps(&self) -> usize {
        self.config.num_inference_steps
    }

    pub fn eta(&self) -> f32 {
        self.config.eta
    }

    pub fn get_timesteps(&self) -> &[usize] {
        &self.timesteps
    }

    pub fn get_alpha_prod_t(&self, t: usize) -> Result<f32> {
        self.alphas_cumprod
            .get(t)
            .copied()
            .ok_or_else(|| DiffusionError::Scheduler {
                message: format!("Timestep {} out of range", t),
            })
    }
}