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

//! Classifier-Free Guidance (CFG) implementation for diffusion models
//!
//! Improves sample quality by amplifying the difference between conditional and unconditional predictions

use crate::error::{DiffusionError, Result};
use rtx_tensor::{Device, Tensor};

/// Configuration for Classifier-Free Guidance
#[derive(Debug, Clone)]
pub struct CFGConfig {
    /// Main guidance scale (typically 7.5 for good results)
    pub guidance_scale: f32,
    /// Scale for unconditional model (usually 1.0)
    pub unconditional_guidance_scale: f32,
    /// Rescale factor to prevent oversaturation
    pub guidance_rescale: f32,
    /// Enable dynamic thresholding
    pub dynamic_thresholding: bool,
    /// Percentile for dynamic thresholding
    pub thresholding_percentile: f32,
}

impl Default for CFGConfig {
    fn default() -> Self {
        Self {
            guidance_scale: 7.5,
            unconditional_guidance_scale: 1.0,
            guidance_rescale: 0.0,
            dynamic_thresholding: false,
            thresholding_percentile: 0.995,
        }
    }
}

/// Guidance scale schedule types
#[derive(Debug, Clone)]
pub enum GuidanceScale {
    /// Constant scale throughout sampling
    Constant(f32),
    /// Linear interpolation from start to end
    Linear { start: f32, end: f32 },
    /// Cosine schedule
    Cosine { start: f32, end: f32 },
}

/// Classifier-Free Guidance for improved sample quality
#[derive(Debug, Clone)]
pub struct ClassifierFreeGuidance {
    config: CFGConfig,
    scale_schedule: GuidanceScale,
    per_channel_scales: Option<Vec<f32>>,
    device: Device,
}

impl ClassifierFreeGuidance {
    /// Create new CFG with configuration
    pub fn new(config: CFGConfig, device: &Device) -> Result<Self> {
        if config.guidance_scale < 0.0 {
            return Err(DiffusionError::Scheduler {
                message: "Guidance scale must be non-negative".to_string(),
            });
        }

        let guidance_scale = config.guidance_scale;
        Ok(Self {
            config,
            scale_schedule: GuidanceScale::Constant(guidance_scale),
            per_channel_scales: None,
            device: device.clone(),
        })
    }

    /// Create CFG with time-varying guidance scale
    pub fn with_schedule(scale_schedule: GuidanceScale, device: &Device) -> Result<Self> {
        let config = CFGConfig::default();
        Ok(Self {
            config,
            scale_schedule,
            per_channel_scales: None,
            device: device.clone(),
        })
    }

    /// Create CFG with per-channel guidance scales
    pub fn with_per_channel_scales(scales: Vec<f32>, device: &Device) -> Result<Self> {
        if scales.is_empty() {
            return Err(DiffusionError::Scheduler {
                message: "Per-channel scales cannot be empty".to_string(),
            });
        }

        let config = CFGConfig::default();
        let guidance_scale = config.guidance_scale;
        Ok(Self {
            config,
            scale_schedule: GuidanceScale::Constant(guidance_scale),
            per_channel_scales: Some(scales),
            device: device.clone(),
        })
    }

    /// Apply classifier-free guidance
    pub fn apply(&self, conditional: &Tensor, unconditional: &Tensor) -> Result<Tensor> {
        self.apply_with_scale(conditional, unconditional, self.config.guidance_scale)
    }

    /// Apply CFG with specific timestep (for scheduled guidance)
    pub fn apply_with_timestep(
        &self,
        conditional: &Tensor,
        unconditional: &Tensor,
        timestep: f32,
    ) -> Result<Tensor> {
        let scale = self.get_scale_at_timestep(timestep);
        self.apply_with_scale(conditional, unconditional, scale)
    }

    /// Apply CFG with multiple conditional inputs
    pub fn apply_multi(
        &self,
        conditionals: &[&Tensor],
        unconditional: &Tensor,
        weights: &[f32],
    ) -> Result<Tensor> {
        if conditionals.len() != weights.len() {
            return Err(DiffusionError::Scheduler {
                message: "Number of conditionals must match number of weights".to_string(),
            });
        }

        // Weighted combination of conditionals
        let mut combined = Tensor::zeros_like(conditionals[0])?;
        for (cond, &weight) in conditionals.iter().zip(weights.iter()) {
            combined = combined.add(&cond.mul_scalar(weight)?)?;
        }

        self.apply(&combined, unconditional)
    }

    /// Apply CFG with negative prompting
    pub fn apply_with_negative(&self, positive: &Tensor, negative: &Tensor) -> Result<Tensor> {
        // Use negative as the unconditional
        self.apply(positive, negative)
    }

    /// Internal: Apply CFG with specific scale
    fn apply_with_scale(
        &self,
        conditional: &Tensor,
        unconditional: &Tensor,
        scale: f32,
    ) -> Result<Tensor> {
        // CFG formula: output = unconditional + scale * (conditional - unconditional)
        let mut guided = if let Some(ref channel_scales) = self.per_channel_scales {
            self.apply_per_channel(conditional, unconditional, channel_scales)?
        } else {
            let diff = conditional.sub(unconditional)?;
            let scaled_diff = diff.mul_scalar(scale)?;
            unconditional.add(&scaled_diff)?
        };

        // Apply guidance rescaling if enabled
        if self.config.guidance_rescale > 0.0 {
            guided = self.rescale_guidance(guided, conditional, unconditional)?;
        }

        // Apply dynamic thresholding if enabled
        if self.config.dynamic_thresholding {
            guided = self.apply_dynamic_thresholding(guided)?;
        }

        Ok(guided)
    }

    /// Apply per-channel guidance scales
    fn apply_per_channel(
        &self,
        conditional: &Tensor,
        unconditional: &Tensor,
        scales: &[f32],
    ) -> Result<Tensor> {
        let shape = conditional.shape().dims();
        let channels = shape[1];

        if scales.len() != channels {
            return Err(DiffusionError::Scheduler {
                message: format!("Expected {} scales, got {}", channels, scales.len()),
            });
        }

        let mut guided_data = vec![0.0f32; conditional.numel()];
        let cond_data = conditional.to_vec()?;
        let uncond_data = unconditional.to_vec()?;

        // Apply different scales to different channels
        let spatial_size = shape[2..].iter().product::<usize>();

        for batch in 0..shape[0] {
            for c in 0..channels {
                let scale = scales[c];
                let channel_offset = batch * channels * spatial_size + c * spatial_size;

                for i in 0..spatial_size {
                    let idx = channel_offset + i;
                    guided_data[idx] =
                        uncond_data[idx] + scale * (cond_data[idx] - uncond_data[idx]);
                }
            }
        }

        Tensor::from_data(guided_data, shape.to_vec(), &self.device).map_err(DiffusionError::Tensor)
    }

    /// Rescale guidance to prevent oversaturation
    fn rescale_guidance(
        &self,
        guided: Tensor,
        conditional: &Tensor,
        unconditional: &Tensor,
    ) -> Result<Tensor> {
        // Compute standard deviations
        let guided_std = self.compute_std(&guided)?;
        let cond_std = self.compute_std(conditional)?;

        // Rescale factor
        let rescale_factor = self.config.guidance_rescale;
        let target_std = cond_std * rescale_factor + guided_std * (1.0 - rescale_factor);

        // Rescale guided to match target std
        let scale = target_std / (guided_std + 1e-8);
        guided.mul_scalar(scale).map_err(DiffusionError::Tensor)
    }

    /// Apply dynamic thresholding to prevent saturation
    fn apply_dynamic_thresholding(&self, guided: Tensor) -> Result<Tensor> {
        let percentile = self.config.thresholding_percentile;

        // Compute absolute values
        let abs_guided = guided.abs()?;

        // Find threshold at percentile
        let threshold = self.compute_percentile(&abs_guided, percentile)?;

        if threshold > 1.0 {
            // Clamp and rescale
            let clamped = guided.clamp(-threshold, threshold)?;
            clamped
                .div_scalar(threshold)
                .map_err(DiffusionError::Tensor)
        } else {
            Ok(guided)
        }
    }

    /// Get guidance scale at specific timestep
    fn get_scale_at_timestep(&self, timestep: f32) -> f32 {
        match &self.scale_schedule {
            GuidanceScale::Constant(scale) => *scale,
            GuidanceScale::Linear { start, end } => start + (end - start) * timestep,
            GuidanceScale::Cosine { start, end } => {
                let cos_val = ((timestep * std::f32::consts::PI).cos() + 1.0) / 2.0;
                start + (end - start) * (1.0 - cos_val)
            }
        }
    }

    /// Compute standard deviation of tensor
    fn compute_std(&self, tensor: &Tensor) -> Result<f32> {
        let data = tensor.to_vec()?;
        let mean = data.iter().sum::<f32>() / data.len() as f32;
        let variance = data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;
        Ok(variance.sqrt())
    }

    /// Compute percentile of tensor values
    fn compute_percentile(&self, tensor: &Tensor, percentile: f32) -> Result<f32> {
        let mut data = tensor.to_vec()?;
        data.sort_by(|a, b| a.total_cmp(b));

        let idx = ((data.len() - 1) as f32 * percentile) as usize;
        Ok(data[idx])
    }

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

    pub fn uses_dynamic_thresholding(&self) -> bool {
        self.config.dynamic_thresholding
    }
}