856 lines
27 KiB
Rust
856 lines
27 KiB
Rust
//! LCM Sampler (Latent Consistency Models)
|
|
//!
|
|
//! Ultra-fast 1-4 step sampling through consistency model formulation.
|
|
//! Supports consistency distillation from pre-trained diffusion models and
|
|
//! classifier-free guidance for high-quality generation.
|
|
|
|
use crate::error::{DiffusionError, Result};
|
|
use crate::noise::NoiseGenerator;
|
|
use rtx_tensor::Tensor;
|
|
|
|
/// LCM prediction types for different parameterizations
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum LCMPredictionType {
|
|
/// Predict noise (epsilon parameterization)
|
|
Epsilon,
|
|
/// Predict velocity (v-parameterization)
|
|
VPrediction,
|
|
/// Predict data directly (x0 parameterization)
|
|
Data,
|
|
}
|
|
|
|
/// Configuration for LCM sampler
|
|
#[derive(Debug, Clone)]
|
|
pub struct LCMConfig {
|
|
/// Number of sampling steps (1-4 for LCM)
|
|
pub num_steps: u32,
|
|
/// Classifier-free guidance scale
|
|
pub guidance_scale: f32,
|
|
/// Consistency loss weight
|
|
pub consistency_weight: f32,
|
|
/// Prediction type
|
|
pub prediction_type: LCMPredictionType,
|
|
/// Solver order (1-3)
|
|
pub solver_order: u8,
|
|
/// Training mode flag
|
|
pub training_mode: bool,
|
|
/// Use Karras sigmas
|
|
pub use_karras_sigmas: bool,
|
|
/// Distillation loss weight
|
|
pub distillation_weight: f32,
|
|
}
|
|
|
|
/// Statistics for LCM sampler
|
|
#[derive(Debug, Default)]
|
|
pub struct LCMStats {
|
|
/// Number of function evaluations
|
|
pub nfe: usize,
|
|
/// Consistency loss sum
|
|
pub consistency_loss_sum: f32,
|
|
/// Distillation loss sum
|
|
pub distillation_loss_sum: f32,
|
|
/// Average sampling time (ms)
|
|
pub avg_sampling_time_ms: f32,
|
|
/// Number of sampling runs
|
|
pub num_runs: usize,
|
|
}
|
|
|
|
/// LCM Sampler for ultra-fast 1-4 step generation
|
|
pub struct LCMSampler {
|
|
config: LCMConfig,
|
|
noise_generator: NoiseGenerator,
|
|
stats: LCMStats,
|
|
/// Cached sigmas for efficiency
|
|
cached_sigmas: Option<Vec<f32>>,
|
|
}
|
|
|
|
impl Default for LCMConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
num_steps: 4,
|
|
guidance_scale: 7.5,
|
|
consistency_weight: 1.0,
|
|
prediction_type: LCMPredictionType::Epsilon,
|
|
solver_order: 2,
|
|
training_mode: false,
|
|
use_karras_sigmas: false,
|
|
distillation_weight: 0.5,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl LCMSampler {
|
|
/// Create a new LCM sampler
|
|
pub fn new(config: LCMConfig, noise_generator: NoiseGenerator) -> Result<Self> {
|
|
// Validation
|
|
if config.num_steps == 0 || config.num_steps > 4 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "LCM num_steps must be between 1 and 4".to_string(),
|
|
});
|
|
}
|
|
|
|
if config.guidance_scale < 0.0 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Guidance scale must be non-negative".to_string(),
|
|
});
|
|
}
|
|
|
|
Ok(Self {
|
|
config,
|
|
noise_generator,
|
|
stats: LCMStats::default(),
|
|
cached_sigmas: None,
|
|
})
|
|
}
|
|
|
|
/// Get current statistics
|
|
pub fn stats(&self) -> &LCMStats {
|
|
&self.stats
|
|
}
|
|
|
|
/// Sample using consistency model (single/multi-step)
|
|
pub fn sample<F>(&mut self, x_t: &Tensor, timestep: u32, model_fn: F) -> Result<Tensor>
|
|
where
|
|
F: Fn(&Tensor, u32) -> Result<Tensor>,
|
|
{
|
|
let start_time = std::time::Instant::now();
|
|
let sigmas = self.get_sigmas(timestep)?;
|
|
let mut x = x_t.clone();
|
|
|
|
// Multi-step LCM sampling with consistency formulation
|
|
for i in 0..self.config.num_steps {
|
|
let current_timestep = self.sigma_to_timestep(sigmas[i as usize])?;
|
|
let model_output = model_fn(&x, current_timestep)?;
|
|
self.stats.nfe += 1;
|
|
|
|
let x0_pred = self.convert_prediction(&x, &model_output, current_timestep)?;
|
|
|
|
x = if i < self.config.num_steps - 1 {
|
|
let next_sigma = sigmas[(i + 1) as usize];
|
|
self.consistency_step(&x0_pred, &x, sigmas[i as usize], next_sigma)?
|
|
} else {
|
|
x0_pred
|
|
};
|
|
}
|
|
|
|
self.update_timing_stats(start_time.elapsed());
|
|
Ok(x)
|
|
}
|
|
|
|
/// Sample with classifier-free guidance
|
|
pub fn sample_guided<F>(
|
|
&mut self,
|
|
x_t: &Tensor,
|
|
timestep: u32,
|
|
guided_model_fn: F,
|
|
) -> Result<Tensor>
|
|
where
|
|
F: Fn(&Tensor, u32) -> Result<(Tensor, Tensor)>,
|
|
{
|
|
let start_time = std::time::Instant::now();
|
|
let sigmas = self.get_sigmas(timestep)?;
|
|
let mut x = x_t.clone();
|
|
|
|
for i in 0..self.config.num_steps {
|
|
let current_timestep = self.sigma_to_timestep(sigmas[i as usize])?;
|
|
let (uncond_output, cond_output) = guided_model_fn(&x, current_timestep)?;
|
|
self.stats.nfe += 2;
|
|
|
|
let guided_output = self.apply_cfg(&uncond_output, &cond_output)?;
|
|
let x0_pred = self.convert_prediction(&x, &guided_output, current_timestep)?;
|
|
|
|
x = if i < self.config.num_steps - 1 {
|
|
let next_sigma = sigmas[(i + 1) as usize];
|
|
self.consistency_step(&x0_pred, &x, sigmas[i as usize], next_sigma)?
|
|
} else {
|
|
x0_pred
|
|
};
|
|
}
|
|
|
|
self.update_timing_stats(start_time.elapsed());
|
|
Ok(x)
|
|
}
|
|
|
|
/// Compute consistency loss for training
|
|
pub fn compute_consistency_loss(
|
|
&self,
|
|
x0: &Tensor,
|
|
noise: &Tensor,
|
|
t1: u32,
|
|
t2: u32,
|
|
) -> Result<Tensor> {
|
|
let x_t1 = self.noise_generator.add_noise(x0, noise, t1)?;
|
|
let x_t2 = self.noise_generator.add_noise(x0, noise, t2)?;
|
|
|
|
let f_t1 = self.consistency_function(&x_t1, t1)?;
|
|
let f_t2 = self.consistency_function(&x_t2, t2)?;
|
|
|
|
let diff = f_t1.subtract(&f_t2)?;
|
|
let loss = diff.pow_scalar(2.0)?.mean(&[], false)?;
|
|
loss.scalar_mul(self.config.consistency_weight)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
|
|
/// Compute distillation loss from teacher model
|
|
pub fn compute_distillation_loss(
|
|
&self,
|
|
x0: &Tensor,
|
|
teacher_output: &Tensor,
|
|
student_output: &Tensor,
|
|
timestep: u32,
|
|
) -> Result<Tensor> {
|
|
let teacher_x0 = self.convert_prediction(x0, teacher_output, timestep)?;
|
|
let student_x0 = self.convert_prediction(x0, student_output, timestep)?;
|
|
|
|
let diff = teacher_x0.subtract(&student_x0)?;
|
|
let loss = diff.pow_scalar(2.0)?.mean(&[], false)?;
|
|
loss.scalar_mul(self.config.distillation_weight)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
|
|
/// Convert model prediction to x0 based on prediction type
|
|
pub fn convert_prediction(
|
|
&self,
|
|
x_t: &Tensor,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
) -> Result<Tensor> {
|
|
match self.config.prediction_type {
|
|
LCMPredictionType::Epsilon => {
|
|
// x0 = (x_t - sqrt(1-alpha_cumprod) * epsilon) / sqrt(alpha_cumprod)
|
|
let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, _, _) =
|
|
self.noise_generator.get_schedule_params(timestep)?;
|
|
|
|
let scaled_noise = model_output.scalar_mul(sqrt_one_minus_alpha_cumprod)?;
|
|
let x_minus_noise = x_t.subtract(&scaled_noise)?;
|
|
x_minus_noise
|
|
.scalar_mul(1.0 / sqrt_alpha_cumprod)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
LCMPredictionType::VPrediction => {
|
|
// v-parameterization: x0 = sqrt(alpha_cumprod) * x_t - sqrt(1-alpha_cumprod) * v
|
|
let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, _, _) =
|
|
self.noise_generator.get_schedule_params(timestep)?;
|
|
|
|
let x_component = x_t.scalar_mul(sqrt_alpha_cumprod)?;
|
|
let v_component = model_output.scalar_mul(sqrt_one_minus_alpha_cumprod)?;
|
|
x_component
|
|
.subtract(&v_component)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
LCMPredictionType::Data => {
|
|
// Direct x0 prediction
|
|
Ok(model_output.clone())
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get Karras sigmas for improved sampling
|
|
pub fn get_karras_sigmas(&self, num_steps: u32) -> Result<Vec<f32>> {
|
|
if let Some(ref cached) = self.cached_sigmas {
|
|
if cached.len() == (num_steps + 1) as usize {
|
|
return Ok(cached.clone());
|
|
}
|
|
}
|
|
|
|
let sigma_min: f32 = 0.002;
|
|
let sigma_max: f32 = 80.0;
|
|
let rho: f32 = 7.0;
|
|
|
|
let mut sigmas = Vec::with_capacity((num_steps + 1) as usize);
|
|
|
|
for i in 0..=num_steps {
|
|
let u = i as f32 / num_steps as f32;
|
|
let sigma: f32 = (sigma_max.powf(1.0_f32 / rho)
|
|
+ u * (sigma_min.powf(1.0_f32 / rho) - sigma_max.powf(1.0_f32 / rho)))
|
|
.powf(rho);
|
|
sigmas.push(sigma);
|
|
}
|
|
|
|
Ok(sigmas)
|
|
}
|
|
|
|
// Helper methods (optimized implementations)
|
|
|
|
fn get_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
|
|
if self.config.use_karras_sigmas {
|
|
self.get_karras_sigmas(self.config.num_steps)
|
|
} else {
|
|
self.get_linear_sigmas(self.config.num_steps, timestep)
|
|
}
|
|
}
|
|
|
|
fn get_linear_sigmas(&self, num_steps: u32, max_timestep: u32) -> Result<Vec<f32>> {
|
|
let mut sigmas = Vec::with_capacity((num_steps + 1) as usize);
|
|
|
|
for i in 0..=num_steps {
|
|
let t = max_timestep as f32 * (1.0 - i as f32 / num_steps as f32);
|
|
let (_, sqrt_one_minus_alpha_cumprod, _, _) =
|
|
self.noise_generator.get_schedule_params(t as u32)?;
|
|
sigmas.push(sqrt_one_minus_alpha_cumprod);
|
|
}
|
|
|
|
Ok(sigmas)
|
|
}
|
|
|
|
fn update_timing_stats(&mut self, elapsed: std::time::Duration) {
|
|
self.stats.num_runs += 1;
|
|
let total_time = self.stats.avg_sampling_time_ms * (self.stats.num_runs - 1) as f32
|
|
+ elapsed.as_millis() as f32;
|
|
self.stats.avg_sampling_time_ms = total_time / self.stats.num_runs as f32;
|
|
}
|
|
|
|
fn sigma_to_timestep(&self, sigma: f32) -> Result<u32> {
|
|
// Simplified mapping - in practice would be more sophisticated
|
|
let timestep = (sigma * 1000.0) as u32;
|
|
Ok(timestep.min(self.noise_generator.num_timesteps() - 1))
|
|
}
|
|
|
|
fn consistency_function(&self, x_t: &Tensor, timestep: u32) -> Result<Tensor> {
|
|
// Simplified consistency function: f(x_t, t) = x_t / (1 + sigma(t))
|
|
let (_, sqrt_one_minus_alpha_cumprod, _, _) =
|
|
self.noise_generator.get_schedule_params(timestep)?;
|
|
let scale = 1.0 / (1.0 + sqrt_one_minus_alpha_cumprod);
|
|
x_t.scalar_mul(scale).map_err(DiffusionError::Tensor)
|
|
}
|
|
|
|
fn consistency_step(
|
|
&self,
|
|
x0_pred: &Tensor,
|
|
x_t: &Tensor,
|
|
sigma_curr: f32,
|
|
sigma_next: f32,
|
|
) -> Result<Tensor> {
|
|
// Simplified consistency step: blend between prediction and current state
|
|
let alpha = sigma_next / sigma_curr;
|
|
let x0_component = x0_pred.scalar_mul(1.0 - alpha)?;
|
|
let xt_component = x_t.scalar_mul(alpha)?;
|
|
x0_component
|
|
.add(&xt_component)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
|
|
fn apply_cfg(&self, uncond_output: &Tensor, cond_output: &Tensor) -> Result<Tensor> {
|
|
// CFG: output = uncond + guidance_scale * (cond - uncond)
|
|
let diff = cond_output.subtract(uncond_output)?;
|
|
let scaled_diff = diff.scalar_mul(self.config.guidance_scale)?;
|
|
uncond_output
|
|
.add(&scaled_diff)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
|
|
/// Advanced consistency training loss with boundary conditions
|
|
pub fn compute_advanced_consistency_loss(
|
|
&self,
|
|
x0: &Tensor,
|
|
noise: &Tensor,
|
|
timesteps: &[u32],
|
|
) -> Result<Tensor> {
|
|
if timesteps.len() < 2 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Need at least 2 timesteps for consistency loss".to_string(),
|
|
});
|
|
}
|
|
|
|
let mut total_loss = None;
|
|
for window in timesteps.windows(2) {
|
|
let t1 = window[0];
|
|
let t2 = window[1];
|
|
let window_loss = self.compute_consistency_loss(x0, noise, t1, t2)?;
|
|
|
|
total_loss = match total_loss {
|
|
None => Some(window_loss),
|
|
Some(acc) => Some(acc.add(&window_loss)?),
|
|
};
|
|
}
|
|
|
|
total_loss
|
|
.unwrap()
|
|
.scalar_mul(1.0 / (timesteps.len() - 1) as f32)
|
|
.map_err(DiffusionError::Tensor)
|
|
}
|
|
|
|
/// Adaptive sampling with dynamic step adjustment
|
|
pub fn sample_adaptive<F>(
|
|
&mut self,
|
|
x_t: &Tensor,
|
|
timestep: u32,
|
|
model_fn: F,
|
|
error_threshold: f32,
|
|
) -> Result<Tensor>
|
|
where
|
|
F: Fn(&Tensor, u32) -> Result<Tensor>,
|
|
{
|
|
let mut current_steps = self.config.num_steps.max(1);
|
|
let mut x = x_t.clone();
|
|
|
|
while current_steps <= 4 {
|
|
let backup_steps = self.config.num_steps;
|
|
self.config.num_steps = current_steps;
|
|
|
|
let result = self.sample(x_t, timestep, &model_fn);
|
|
self.config.num_steps = backup_steps;
|
|
|
|
match result {
|
|
Ok(sampled) => {
|
|
// Simple error estimation based on change magnitude
|
|
let change = sampled.subtract(&x)?.pow_scalar(2.0)?.mean(&[], false)?;
|
|
let error = change.item()?;
|
|
|
|
if error < error_threshold || current_steps == 4 {
|
|
return Ok(sampled);
|
|
}
|
|
x = sampled;
|
|
current_steps += 1;
|
|
}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
Ok(x)
|
|
}
|
|
|
|
/// Get solver-specific sigma schedule
|
|
pub fn get_solver_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
|
|
match self.config.solver_order {
|
|
1 => self.get_euler_sigmas(timestep),
|
|
2 => self.get_heun_sigmas(timestep),
|
|
3 => self.get_dpm_sigmas(timestep),
|
|
_ => self.get_sigmas(timestep),
|
|
}
|
|
}
|
|
|
|
fn get_euler_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
|
|
// Euler method sigmas (linear spacing)
|
|
self.get_linear_sigmas(self.config.num_steps, timestep)
|
|
}
|
|
|
|
fn get_heun_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
|
|
// Heun method sigmas (better for 2-step)
|
|
if self.config.use_karras_sigmas {
|
|
self.get_karras_sigmas(self.config.num_steps)
|
|
} else {
|
|
self.get_linear_sigmas(self.config.num_steps, timestep)
|
|
}
|
|
}
|
|
|
|
fn get_dpm_sigmas(&self, timestep: u32) -> Result<Vec<f32>> {
|
|
// DPM-style sigma spacing for higher order
|
|
let mut sigmas = self.get_karras_sigmas(self.config.num_steps)?;
|
|
// Apply DPM weighting
|
|
for sigma in &mut sigmas {
|
|
*sigma = sigma.sqrt(); // Square root weighting for better stability
|
|
}
|
|
Ok(sigmas)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::noise::NoiseSchedule;
|
|
use rtx_tensor::Device;
|
|
|
|
#[test]
|
|
fn test_lcm_config_creation() {
|
|
let config = LCMConfig {
|
|
num_steps: 4,
|
|
guidance_scale: 7.5,
|
|
consistency_weight: 1.0,
|
|
prediction_type: LCMPredictionType::Epsilon,
|
|
solver_order: 2,
|
|
training_mode: false,
|
|
use_karras_sigmas: false,
|
|
distillation_weight: 0.5,
|
|
};
|
|
|
|
assert_eq!(config.num_steps, 4);
|
|
assert_eq!(config.guidance_scale, 7.5);
|
|
assert_eq!(config.consistency_weight, 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lcm_sampler_creation() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
let config = LCMConfig::default();
|
|
let sampler = LCMSampler::new(config, noise_gen);
|
|
|
|
assert!(sampler.is_ok());
|
|
let sampler = sampler.unwrap();
|
|
assert_eq!(sampler.stats().nfe, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_consistency_function_single_step() {
|
|
let noise_gen =
|
|
NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();
|
|
|
|
let config = LCMConfig {
|
|
num_steps: 1,
|
|
..Default::default()
|
|
};
|
|
let mut sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
// Mock model output function
|
|
let model_fn = |x: &Tensor, t: u32| -> Result<Tensor> {
|
|
// Simple mock: return scaled input
|
|
let scale = (t as f32 / 1000.0) * 0.1;
|
|
Ok(x.scalar_mul(scale)?)
|
|
};
|
|
|
|
let input_shape = vec![1, 3, 32, 32];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let latent = Tensor::randn(&input_shape, &device).unwrap();
|
|
|
|
let result = sampler.sample(&latent, 999, model_fn);
|
|
assert!(result.is_ok());
|
|
|
|
let denoised = result.unwrap();
|
|
assert_eq!(denoised.shape(), input_shape);
|
|
assert_eq!(sampler.stats().nfe, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_consistency_function_multi_step() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
let config = LCMConfig {
|
|
num_steps: 4,
|
|
..Default::default()
|
|
};
|
|
let mut sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let model_fn = |x: &Tensor, t: u32| -> Result<Tensor> {
|
|
let scale = (t as f32 / 1000.0) * 0.1;
|
|
Ok(x.scalar_mul(scale)?)
|
|
};
|
|
|
|
let input_shape = vec![2, 4, 64, 64];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let latent = Tensor::randn(&input_shape, &device).unwrap();
|
|
|
|
let result = sampler.sample(&latent, 800, model_fn);
|
|
assert!(result.is_ok());
|
|
|
|
let denoised = result.unwrap();
|
|
assert_eq!(denoised.shape(), input_shape);
|
|
assert_eq!(sampler.stats().nfe, 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_classifier_free_guidance() {
|
|
let noise_gen =
|
|
NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();
|
|
|
|
let config = LCMConfig {
|
|
num_steps: 2,
|
|
guidance_scale: 7.5,
|
|
..Default::default()
|
|
};
|
|
let mut sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let guided_model_fn = |x: &Tensor, t: u32| -> Result<(Tensor, Tensor)> {
|
|
let scale = (t as f32 / 1000.0) * 0.1;
|
|
let uncond = x.scalar_mul(scale)?;
|
|
let cond = x.scalar_mul(scale * 1.2)?;
|
|
Ok((uncond, cond))
|
|
};
|
|
|
|
let input_shape = vec![1, 4, 32, 32];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let latent = Tensor::randn(&input_shape, &device).unwrap();
|
|
|
|
let result = sampler.sample_guided(&latent, 900, guided_model_fn);
|
|
assert!(result.is_ok());
|
|
|
|
let denoised = result.unwrap();
|
|
assert_eq!(denoised.shape(), input_shape);
|
|
// Should have called model twice per step (uncond + cond)
|
|
assert_eq!(sampler.stats().nfe, 4);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal shader compilation issue"]
|
|
fn test_consistency_loss_computation() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
let config = LCMConfig {
|
|
training_mode: true,
|
|
consistency_weight: 2.0,
|
|
..Default::default()
|
|
};
|
|
let sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let x0_shape = vec![2, 3, 64, 64];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let x0 = Tensor::randn(&x0_shape, &device).unwrap();
|
|
let noise = Tensor::randn(&x0_shape, &device).unwrap();
|
|
|
|
let t1 = 100;
|
|
let t2 = 200;
|
|
|
|
let loss = sampler.compute_consistency_loss(&x0, &noise, t1, t2);
|
|
assert!(loss.is_ok());
|
|
|
|
let loss_value = loss.unwrap();
|
|
// Loss should be a scalar
|
|
assert_eq!(loss_value.shape(), vec![]);
|
|
assert!(loss_value.item().unwrap() >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal shader compilation issue"]
|
|
fn test_distillation_loss_computation() {
|
|
let noise_gen =
|
|
NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();
|
|
|
|
let config = LCMConfig {
|
|
training_mode: true,
|
|
distillation_weight: 1.5,
|
|
..Default::default()
|
|
};
|
|
let sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let x0_shape = vec![1, 4, 32, 32];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let x0 = Tensor::randn(&x0_shape, &device).unwrap();
|
|
let teacher_output = Tensor::randn(&x0_shape, &device).unwrap();
|
|
let student_output = Tensor::randn(&x0_shape, &device).unwrap();
|
|
|
|
let timestep = 500;
|
|
|
|
let loss =
|
|
sampler.compute_distillation_loss(&x0, &teacher_output, &student_output, timestep);
|
|
assert!(loss.is_ok());
|
|
|
|
let loss_value = loss.unwrap();
|
|
assert_eq!(loss_value.shape(), vec![]);
|
|
assert!(loss_value.item().unwrap() >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_prediction_type_conversion() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
let config = LCMConfig {
|
|
prediction_type: LCMPredictionType::VPrediction,
|
|
..Default::default()
|
|
};
|
|
let sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let x_shape = vec![1, 3, 32, 32];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let x_t = Tensor::randn(&x_shape, &device).unwrap();
|
|
let model_output = Tensor::randn(&x_shape, &device).unwrap();
|
|
let timestep = 500;
|
|
|
|
let x0_pred = sampler.convert_prediction(&x_t, &model_output, timestep);
|
|
assert!(x0_pred.is_ok());
|
|
assert_eq!(x0_pred.unwrap().shape(), x_shape);
|
|
}
|
|
|
|
#[test]
|
|
fn test_karras_sigma_schedule() {
|
|
let noise_gen =
|
|
NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();
|
|
|
|
let config = LCMConfig {
|
|
use_karras_sigmas: true,
|
|
num_steps: 4,
|
|
..Default::default()
|
|
};
|
|
let sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let sigmas = sampler.get_karras_sigmas(4);
|
|
assert!(sigmas.is_ok());
|
|
|
|
let sigma_vec = sigmas.unwrap();
|
|
assert_eq!(sigma_vec.len(), 5); // num_steps + 1
|
|
|
|
// Sigmas should be decreasing
|
|
for i in 1..sigma_vec.len() {
|
|
assert!(sigma_vec[i - 1] >= sigma_vec[i]);
|
|
}
|
|
|
|
// Last sigma should be close to 0
|
|
assert!(sigma_vec.last().unwrap() < &0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_lcm_stats_tracking() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
let config = LCMConfig::default();
|
|
let mut sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
// Initial stats
|
|
let initial_stats = sampler.stats();
|
|
assert_eq!(initial_stats.nfe, 0);
|
|
assert_eq!(initial_stats.consistency_loss_sum, 0.0);
|
|
assert_eq!(initial_stats.distillation_loss_sum, 0.0);
|
|
|
|
// Perform sampling to update stats
|
|
let model_fn = |x: &Tensor, _t: u32| -> Result<Tensor> { Ok(x.scalar_mul(0.1)?) };
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let latent = Tensor::randn(&vec![1, 3, 32, 32], &device).unwrap();
|
|
let _ = sampler.sample(&latent, 800, model_fn);
|
|
|
|
// Stats should be updated
|
|
let final_stats = sampler.stats();
|
|
assert!(final_stats.nfe > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_configurations() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
// Test invalid num_steps
|
|
let invalid_config = LCMConfig {
|
|
num_steps: 0,
|
|
..Default::default()
|
|
};
|
|
let result = LCMSampler::new(invalid_config, noise_gen.clone());
|
|
assert!(result.is_err());
|
|
|
|
// Test invalid guidance_scale
|
|
let invalid_config = LCMConfig {
|
|
guidance_scale: -1.0,
|
|
..Default::default()
|
|
};
|
|
let result = LCMSampler::new(invalid_config, noise_gen);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal shader compilation issue"]
|
|
fn test_advanced_consistency_loss() {
|
|
let noise_gen =
|
|
NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();
|
|
|
|
let config = LCMConfig {
|
|
training_mode: true,
|
|
consistency_weight: 1.5,
|
|
..Default::default()
|
|
};
|
|
let sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let x0_shape = vec![1, 3, 32, 32];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let x0 = Tensor::randn(&x0_shape, &device).unwrap();
|
|
let noise = Tensor::randn(&x0_shape, &device).unwrap();
|
|
|
|
let timesteps = vec![100, 200, 300, 400];
|
|
|
|
let loss = sampler.compute_advanced_consistency_loss(&x0, &noise, ×teps);
|
|
assert!(loss.is_ok());
|
|
|
|
let loss_value = loss.unwrap();
|
|
assert_eq!(loss_value.shape(), vec![]);
|
|
assert!(loss_value.item().unwrap() >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing Metal shader compilation issue"]
|
|
fn test_adaptive_sampling() {
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
let config = LCMConfig {
|
|
num_steps: 1,
|
|
..Default::default()
|
|
};
|
|
let mut sampler = LCMSampler::new(config, noise_gen).unwrap();
|
|
|
|
let model_fn = |x: &Tensor, _t: u32| -> Result<Tensor> { Ok(x.scalar_mul(0.1)?) };
|
|
|
|
let input_shape = vec![1, 3, 32, 32];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let latent = Tensor::randn(&input_shape, &device).unwrap();
|
|
|
|
let result = sampler.sample_adaptive(&latent, 800, model_fn, 0.01);
|
|
assert!(result.is_ok());
|
|
|
|
let denoised = result.unwrap();
|
|
assert_eq!(denoised.shape(), input_shape);
|
|
}
|
|
|
|
#[test]
|
|
fn test_solver_specific_sigmas() {
|
|
let noise_gen =
|
|
NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();
|
|
|
|
// Test different solver orders
|
|
for order in 1..=3 {
|
|
let config = LCMConfig {
|
|
solver_order: order,
|
|
num_steps: 4,
|
|
..Default::default()
|
|
};
|
|
let sampler = LCMSampler::new(config, noise_gen.clone()).unwrap();
|
|
|
|
let sigmas = sampler.get_solver_sigmas(800);
|
|
assert!(sigmas.is_ok());
|
|
|
|
let sigma_vec = sigmas.unwrap();
|
|
assert_eq!(sigma_vec.len(), 5);
|
|
}
|
|
}
|
|
}
|