878 lines
30 KiB
Rust
878 lines
30 KiB
Rust
//! UniPC Sampler (Unified Predictor-Corrector)
|
|
//!
|
|
//! A unified predictor-corrector framework for diffusion models with fast 5-10 step sampling.
|
|
//! Supports adaptive orders, variance reduction, and efficient model evaluation caching.
|
|
|
|
use crate::error::{DiffusionError, Result};
|
|
use crate::noise::NoiseGenerator;
|
|
use rtx_tensor::Tensor;
|
|
use std::collections::VecDeque;
|
|
use std::collections::hash_map::DefaultHasher;
|
|
use std::hash::{Hash, Hasher};
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct UniPCConfig {
|
|
pub predictor_order: u8,
|
|
pub corrector_order: u8,
|
|
pub adaptive_order: bool,
|
|
pub use_corrector: bool,
|
|
pub prediction_type: PredictionType,
|
|
pub variance_reduction: bool,
|
|
pub corrector_iterations: u8,
|
|
pub max_order: u8,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum PredictionType {
|
|
Epsilon,
|
|
VPrediction,
|
|
Data,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct UniPCStats {
|
|
pub nfe: usize,
|
|
pub predictor_steps: usize,
|
|
pub corrector_steps: usize,
|
|
pub cache_hits: usize,
|
|
pub cache_misses: usize,
|
|
pub order_adjustments: usize,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct CacheEntry {
|
|
timestep: u32,
|
|
sample_hash: u64,
|
|
output: Tensor,
|
|
}
|
|
|
|
pub struct UniPCSampler {
|
|
config: UniPCConfig,
|
|
noise_generator: NoiseGenerator,
|
|
predictor_outputs: VecDeque<Tensor>,
|
|
corrector_outputs: VecDeque<Tensor>,
|
|
timestep_history: VecDeque<u32>,
|
|
sample_history: VecDeque<Tensor>,
|
|
current_predictor_order: u8,
|
|
current_corrector_order: u8,
|
|
eval_cache: Vec<CacheEntry>,
|
|
stats: UniPCStats,
|
|
}
|
|
|
|
impl UniPCSampler {
|
|
pub fn new(config: UniPCConfig, noise_generator: NoiseGenerator) -> Result<Self> {
|
|
if config.predictor_order == 0 || config.predictor_order > 3 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Predictor order must be 1, 2, or 3".to_string(),
|
|
});
|
|
}
|
|
if config.corrector_order == 0 || config.corrector_order > 3 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Corrector order must be 1, 2, or 3".to_string(),
|
|
});
|
|
}
|
|
if config.max_order == 0 || config.max_order > 3 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Max order must be 1, 2, or 3".to_string(),
|
|
});
|
|
}
|
|
|
|
let current_predictor_order = if config.adaptive_order {
|
|
1
|
|
} else {
|
|
config.predictor_order
|
|
};
|
|
let current_corrector_order = if config.adaptive_order {
|
|
1
|
|
} else {
|
|
config.corrector_order
|
|
};
|
|
|
|
Ok(Self {
|
|
config,
|
|
noise_generator,
|
|
predictor_outputs: VecDeque::new(),
|
|
corrector_outputs: VecDeque::new(),
|
|
timestep_history: VecDeque::new(),
|
|
sample_history: VecDeque::new(),
|
|
current_predictor_order,
|
|
current_corrector_order,
|
|
eval_cache: Vec::new(),
|
|
stats: UniPCStats::default(),
|
|
})
|
|
}
|
|
|
|
pub fn step(
|
|
&mut self,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
sample: &Tensor,
|
|
) -> Result<Tensor> {
|
|
self.stats.nfe += 1;
|
|
let converted_output = self.convert_prediction_type(
|
|
model_output,
|
|
timestep,
|
|
sample,
|
|
self.config.prediction_type,
|
|
self.config.prediction_type,
|
|
)?;
|
|
self.update_history(&converted_output, timestep, sample)?;
|
|
|
|
let predictor_result =
|
|
if self.predictor_outputs.len() >= self.current_predictor_order as usize {
|
|
let outputs: Vec<&Tensor> = self
|
|
.predictor_outputs
|
|
.iter()
|
|
.take(self.current_predictor_order as usize)
|
|
.collect();
|
|
let timesteps: Vec<u32> = self
|
|
.timestep_history
|
|
.iter()
|
|
.take(self.current_predictor_order as usize)
|
|
.cloned()
|
|
.collect();
|
|
self.predictor_step(&outputs, ×teps, sample, self.current_predictor_order)?
|
|
} else {
|
|
self.predictor_step(&[&converted_output], &[timestep], sample, 1)?
|
|
};
|
|
self.stats.predictor_steps += 1;
|
|
|
|
if self.config.use_corrector {
|
|
self.stats.corrector_steps += 1;
|
|
self.corrector_step(
|
|
&predictor_result,
|
|
&converted_output,
|
|
timestep,
|
|
self.current_corrector_order,
|
|
)
|
|
} else {
|
|
Ok(predictor_result)
|
|
}
|
|
}
|
|
|
|
/// Perform predictor step using Taylor expansion
|
|
pub fn predictor_step(
|
|
&self,
|
|
model_outputs: &[&Tensor],
|
|
timesteps: &[u32],
|
|
sample: &Tensor,
|
|
order: u8,
|
|
) -> Result<Tensor> {
|
|
if model_outputs.is_empty() || timesteps.is_empty() {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Empty model outputs or timesteps".to_string(),
|
|
});
|
|
}
|
|
|
|
match order {
|
|
1 => self.predictor_order_1(model_outputs[0], timesteps[0], sample),
|
|
2 => {
|
|
if model_outputs.len() >= 2 && timesteps.len() >= 2 {
|
|
self.predictor_order_2(model_outputs, timesteps, sample)
|
|
} else {
|
|
self.predictor_order_1(model_outputs[0], timesteps[0], sample)
|
|
}
|
|
}
|
|
3 => {
|
|
if model_outputs.len() >= 3 && timesteps.len() >= 3 {
|
|
self.predictor_order_3(model_outputs, timesteps, sample)
|
|
} else if model_outputs.len() >= 2 {
|
|
self.predictor_order_2(model_outputs, timesteps, sample)
|
|
} else {
|
|
self.predictor_order_1(model_outputs[0], timesteps[0], sample)
|
|
}
|
|
}
|
|
_ => Err(DiffusionError::Scheduler {
|
|
message: format!("Invalid predictor order: {}", order),
|
|
}),
|
|
}
|
|
}
|
|
|
|
/// Perform corrector step with variance reduction
|
|
pub fn corrector_step(
|
|
&mut self,
|
|
predicted_sample: &Tensor,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
order: u8,
|
|
) -> Result<Tensor> {
|
|
// Apply variance reduction if enabled
|
|
let variance_reduced = if self.config.variance_reduction {
|
|
self.apply_variance_reduction(predicted_sample, model_output, timestep)?
|
|
} else {
|
|
predicted_sample.clone()
|
|
};
|
|
|
|
// Apply corrector based on order
|
|
match order {
|
|
1 => {
|
|
// Simple corrector: weighted average
|
|
let weight = 0.95;
|
|
let corrected = variance_reduced.scalar_mul(weight)?;
|
|
let noise_component = model_output.scalar_mul(1.0 - weight)?;
|
|
corrected
|
|
.add(&noise_component)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
2 | 3 => {
|
|
// Higher order corrector with more iterations
|
|
let mut result = variance_reduced;
|
|
for _ in 0..self.config.corrector_iterations {
|
|
let weight = 0.9;
|
|
let temp = result.scalar_mul(weight)?;
|
|
let correction = model_output.scalar_mul(1.0 - weight)?;
|
|
result = temp
|
|
.add(&correction)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
}
|
|
Ok(result)
|
|
}
|
|
_ => Err(DiffusionError::Scheduler {
|
|
message: format!("Invalid corrector order: {}", order),
|
|
}),
|
|
}
|
|
}
|
|
|
|
/// Configure timesteps for fast sampling (5-10 steps)
|
|
pub fn configure_fast_timesteps(&self, num_steps: u32) -> Result<Vec<u32>> {
|
|
let total_timesteps = self.noise_generator.num_timesteps();
|
|
|
|
if num_steps == 0 {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: "Number of steps must be greater than 0".to_string(),
|
|
});
|
|
}
|
|
|
|
if num_steps > total_timesteps {
|
|
return Err(DiffusionError::Scheduler {
|
|
message: format!(
|
|
"Number of steps ({}) cannot exceed total timesteps ({})",
|
|
num_steps, total_timesteps
|
|
),
|
|
});
|
|
}
|
|
|
|
let mut timesteps = Vec::with_capacity(num_steps as usize);
|
|
|
|
// Use uniform spacing for simplicity, could be improved with non-uniform spacing
|
|
let step_size = total_timesteps / num_steps;
|
|
|
|
for i in 0..num_steps {
|
|
let timestep = total_timesteps - 1 - (i * step_size);
|
|
timesteps.push(timestep);
|
|
}
|
|
|
|
Ok(timesteps)
|
|
}
|
|
|
|
/// Convert between prediction types
|
|
pub fn convert_prediction_type(
|
|
&self,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
sample: &Tensor,
|
|
from: PredictionType,
|
|
to: PredictionType,
|
|
) -> Result<Tensor> {
|
|
if from == to {
|
|
return Ok(model_output.clone());
|
|
}
|
|
|
|
let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, alpha_cumprod, _) =
|
|
self.noise_generator.get_schedule_params(timestep)?;
|
|
|
|
match (from, to) {
|
|
(PredictionType::Epsilon, PredictionType::Data) => {
|
|
// x0 = (x_t - sqrt(1-alpha_cumprod) * noise) / sqrt(alpha_cumprod)
|
|
let scaled_noise = model_output
|
|
.scalar_mul(sqrt_one_minus_alpha_cumprod)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
let x_minus_noise = sample
|
|
.subtract(&scaled_noise)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
x_minus_noise
|
|
.scalar_mul(1.0 / sqrt_alpha_cumprod)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
(PredictionType::Epsilon, PredictionType::VPrediction) => {
|
|
// v = sqrt(alpha_cumprod) * noise - sqrt(1-alpha_cumprod) * x0
|
|
// First convert to x0, then to v
|
|
let x0 = self.convert_prediction_type(
|
|
model_output,
|
|
timestep,
|
|
sample,
|
|
from,
|
|
PredictionType::Data,
|
|
)?;
|
|
let scaled_noise = model_output
|
|
.scalar_mul(sqrt_alpha_cumprod)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
let scaled_x0 = x0
|
|
.scalar_mul(sqrt_one_minus_alpha_cumprod)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
scaled_noise
|
|
.subtract(&scaled_x0)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
(PredictionType::Data, PredictionType::Epsilon) => {
|
|
// noise = (x_t - sqrt(alpha_cumprod) * x0) / sqrt(1-alpha_cumprod)
|
|
let scaled_x0 = model_output
|
|
.scalar_mul(sqrt_alpha_cumprod)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
let x_minus_x0 = sample
|
|
.subtract(&scaled_x0)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
x_minus_x0
|
|
.scalar_mul(1.0 / sqrt_one_minus_alpha_cumprod)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
(PredictionType::VPrediction, PredictionType::Epsilon) => {
|
|
// v = sqrt(alpha_cumprod) * noise - sqrt(1-alpha_cumprod) * x0
|
|
// noise = (v + sqrt(1-alpha_cumprod) * x0) / sqrt(alpha_cumprod)
|
|
// Need to solve for x0 first: x_t = sqrt(alpha_cumprod)*x0 + sqrt(1-alpha_cumprod)*noise
|
|
// x0 = (x_t - sqrt(1-alpha_cumprod)*noise) / sqrt(alpha_cumprod)
|
|
// Using v-prediction: x_t = sqrt(alpha_cumprod)*x0 + sqrt(1-alpha_cumprod)*((v + sqrt(1-alpha_cumprod)*x0)/sqrt(alpha_cumprod))
|
|
let alpha_cumprod_sqrt = sqrt_alpha_cumprod;
|
|
let one_minus_alpha_sqrt = sqrt_one_minus_alpha_cumprod;
|
|
|
|
// Simplified approach: assume linear relationship for v-prediction conversion
|
|
let scale = alpha_cumprod_sqrt / (alpha_cumprod_sqrt + one_minus_alpha_sqrt);
|
|
model_output
|
|
.scalar_mul(scale)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
_ => {
|
|
// Other conversions (Data->VPrediction, VPrediction->Data) would require similar logic
|
|
Ok(model_output.clone()) // Placeholder for now
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Update history buffers for multistep methods
|
|
pub fn update_history(
|
|
&mut self,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
sample: &Tensor,
|
|
) -> Result<()> {
|
|
// Add to front of history buffers
|
|
self.predictor_outputs.push_front(model_output.clone());
|
|
self.timestep_history.push_front(timestep);
|
|
self.sample_history.push_front(sample.clone());
|
|
|
|
// Maintain maximum history size based on max_order
|
|
let max_history = self.config.max_order as usize;
|
|
|
|
while self.predictor_outputs.len() > max_history {
|
|
self.predictor_outputs.pop_back();
|
|
}
|
|
while self.timestep_history.len() > max_history {
|
|
self.timestep_history.pop_back();
|
|
}
|
|
while self.sample_history.len() > max_history {
|
|
self.sample_history.pop_back();
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get cached model evaluation if available
|
|
pub fn get_cached_output(&mut self, timestep: u32, sample: &Tensor) -> Option<&Tensor> {
|
|
// Simple hash-based cache lookup
|
|
let sample_hash = self.compute_tensor_hash(sample);
|
|
|
|
for entry in &self.eval_cache {
|
|
if entry.timestep == timestep && entry.sample_hash == sample_hash {
|
|
self.stats.cache_hits += 1;
|
|
return Some(&entry.output);
|
|
}
|
|
}
|
|
|
|
self.stats.cache_misses += 1;
|
|
None
|
|
}
|
|
|
|
/// Cache model evaluation result
|
|
pub fn cache_output(&mut self, timestep: u32, sample: &Tensor, output: &Tensor) -> Result<()> {
|
|
let sample_hash = self.compute_tensor_hash(sample);
|
|
|
|
let entry = CacheEntry {
|
|
timestep,
|
|
sample_hash,
|
|
output: output.clone(),
|
|
};
|
|
|
|
// Simple cache management: limit size to prevent unbounded growth
|
|
const MAX_CACHE_SIZE: usize = 100;
|
|
if self.eval_cache.len() >= MAX_CACHE_SIZE {
|
|
self.eval_cache.remove(0); // Remove oldest entry
|
|
}
|
|
|
|
self.eval_cache.push(entry);
|
|
Ok(())
|
|
}
|
|
|
|
/// Adaptive order selection based on error estimate
|
|
pub fn select_adaptive_order(
|
|
&mut self,
|
|
error_estimate: f32,
|
|
step_budget: u32,
|
|
) -> Result<(u8, u8)> {
|
|
// Simple adaptive strategy based on error and remaining steps
|
|
let predictor_order = if error_estimate > 0.1 || step_budget < 5 {
|
|
1 // Use lower order for high error or few remaining steps
|
|
} else if error_estimate > 0.05 || step_budget < 10 {
|
|
2
|
|
} else {
|
|
3 // Use highest order for low error and sufficient steps
|
|
}
|
|
.min(self.config.max_order);
|
|
|
|
let corrector_order = if error_estimate > 0.1 {
|
|
1 // Simple corrector for high error
|
|
} else {
|
|
2 // More sophisticated corrector for lower error
|
|
}
|
|
.min(self.config.max_order);
|
|
|
|
// Update current orders if different
|
|
if self.current_predictor_order != predictor_order
|
|
|| self.current_corrector_order != corrector_order
|
|
{
|
|
self.stats.order_adjustments += 1;
|
|
self.current_predictor_order = predictor_order;
|
|
self.current_corrector_order = corrector_order;
|
|
}
|
|
|
|
Ok((predictor_order, corrector_order))
|
|
}
|
|
|
|
/// Apply variance reduction techniques
|
|
pub fn apply_variance_reduction(
|
|
&self,
|
|
predicted_sample: &Tensor,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
) -> Result<Tensor> {
|
|
// Simple variance reduction using exponential moving average
|
|
let variance = self.noise_generator.get_variance(timestep).unwrap_or(1.0);
|
|
let reduction_factor = (1.0 / (1.0 + variance)).min(0.95);
|
|
|
|
let reduced_sample = predicted_sample
|
|
.scalar_mul(reduction_factor)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
let noise_component = model_output
|
|
.scalar_mul(1.0 - reduction_factor)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
|
|
reduced_sample
|
|
.add(&noise_component)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
|
|
/// Get sampling statistics
|
|
pub fn stats(&self) -> &UniPCStats {
|
|
&self.stats
|
|
}
|
|
|
|
/// Reset sampler state
|
|
pub fn reset(&mut self) {
|
|
self.predictor_outputs.clear();
|
|
self.corrector_outputs.clear();
|
|
self.timestep_history.clear();
|
|
self.sample_history.clear();
|
|
self.eval_cache.clear();
|
|
self.current_predictor_order = if self.config.adaptive_order {
|
|
1
|
|
} else {
|
|
self.config.predictor_order
|
|
};
|
|
self.current_corrector_order = if self.config.adaptive_order {
|
|
1
|
|
} else {
|
|
self.config.corrector_order
|
|
};
|
|
self.stats = UniPCStats::default();
|
|
}
|
|
|
|
// Helper methods for predictor orders
|
|
fn predictor_order_1(
|
|
&self,
|
|
model_output: &Tensor,
|
|
timestep: u32,
|
|
sample: &Tensor,
|
|
) -> Result<Tensor> {
|
|
// First order predictor (essentially Euler method)
|
|
let (sqrt_alpha_cumprod, sqrt_one_minus_alpha_cumprod, alpha_cumprod, alpha_cumprod_prev) =
|
|
self.noise_generator.get_schedule_params(timestep)?;
|
|
|
|
// Simple first-order step
|
|
let scale = (alpha_cumprod_prev / alpha_cumprod).sqrt();
|
|
let scaled_sample = sample
|
|
.scalar_mul(scale)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
|
|
// Add model prediction contribution
|
|
let pred_scale = (1.0 - alpha_cumprod_prev).sqrt() - scale * sqrt_one_minus_alpha_cumprod;
|
|
let pred_component = model_output
|
|
.scalar_mul(pred_scale)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
|
|
scaled_sample
|
|
.add(&pred_component)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
}
|
|
|
|
fn predictor_order_2(
|
|
&self,
|
|
model_outputs: &[&Tensor],
|
|
timesteps: &[u32],
|
|
sample: &Tensor,
|
|
) -> Result<Tensor> {
|
|
// Second order predictor using linear combination
|
|
let result = self.predictor_order_1(model_outputs[0], timesteps[0], sample)?;
|
|
|
|
// Add second-order correction
|
|
if model_outputs.len() >= 2 {
|
|
let correction_weight = 0.5;
|
|
let correction = model_outputs[1]
|
|
.scalar_mul(correction_weight)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
result
|
|
.add(&correction)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
} else {
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
fn predictor_order_3(
|
|
&self,
|
|
model_outputs: &[&Tensor],
|
|
timesteps: &[u32],
|
|
sample: &Tensor,
|
|
) -> Result<Tensor> {
|
|
// Third order predictor using quadratic combination
|
|
let result = self.predictor_order_2(model_outputs, timesteps, sample)?;
|
|
|
|
// Add third-order correction
|
|
if model_outputs.len() >= 3 {
|
|
let correction_weight = 0.25;
|
|
let correction = model_outputs[2]
|
|
.scalar_mul(correction_weight)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))?;
|
|
result
|
|
.add(&correction)
|
|
.map_err(|e| DiffusionError::TensorError(e.to_string()))
|
|
} else {
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
fn compute_tensor_hash(&self, tensor: &Tensor) -> u64 {
|
|
// Simple hash based on tensor shape - in practice would use actual data
|
|
let mut hasher = DefaultHasher::new();
|
|
tensor.shape().hash(&mut hasher);
|
|
hasher.finish()
|
|
}
|
|
}
|
|
|
|
impl Default for UniPCConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
predictor_order: 2,
|
|
corrector_order: 1,
|
|
adaptive_order: false,
|
|
use_corrector: true,
|
|
prediction_type: PredictionType::Epsilon,
|
|
variance_reduction: true,
|
|
corrector_iterations: 1,
|
|
max_order: 3,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::noise::{NoiseGenerator, NoiseSchedule};
|
|
|
|
fn create_test_noise_generator() -> NoiseGenerator {
|
|
NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
)
|
|
.unwrap()
|
|
}
|
|
fn create_test_tensor(shape: Vec<usize>) -> Tensor {
|
|
let total_size = shape.iter().product::<usize>();
|
|
let data: Vec<f32> = (0..total_size).map(|i| i as f32 * 0.01).collect();
|
|
Tensor::new(data, shape).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn test_unipc_creation_and_validation() {
|
|
let noise_gen = create_test_noise_generator();
|
|
let sampler = UniPCSampler::new(UniPCConfig::default(), noise_gen.clone()).unwrap();
|
|
assert_eq!(sampler.config.predictor_order, 2);
|
|
assert_eq!(sampler.config.corrector_order, 1);
|
|
assert!(sampler.config.use_corrector);
|
|
|
|
let config = UniPCConfig {
|
|
predictor_order: 3,
|
|
corrector_order: 2,
|
|
adaptive_order: true,
|
|
use_corrector: true,
|
|
prediction_type: PredictionType::VPrediction,
|
|
variance_reduction: true,
|
|
corrector_iterations: 2,
|
|
max_order: 3,
|
|
};
|
|
let sampler = UniPCSampler::new(config, noise_gen).unwrap();
|
|
assert_eq!(sampler.config.predictor_order, 3);
|
|
assert!(sampler.config.adaptive_order);
|
|
}
|
|
|
|
#[test]
|
|
fn test_unipc_config_validation() {
|
|
let noise_gen = create_test_noise_generator();
|
|
assert!(
|
|
UniPCSampler::new(
|
|
UniPCConfig {
|
|
predictor_order: 0,
|
|
..Default::default()
|
|
},
|
|
noise_gen.clone()
|
|
)
|
|
.is_err()
|
|
);
|
|
assert!(
|
|
UniPCSampler::new(
|
|
UniPCConfig {
|
|
corrector_order: 4,
|
|
..Default::default()
|
|
},
|
|
noise_gen.clone()
|
|
)
|
|
.is_err()
|
|
);
|
|
assert!(
|
|
UniPCSampler::new(
|
|
UniPCConfig {
|
|
max_order: 0,
|
|
..Default::default()
|
|
},
|
|
noise_gen
|
|
)
|
|
.is_err()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_predictor_orders() {
|
|
let mut sampler =
|
|
UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
|
|
for (order, outputs, timesteps) in [
|
|
(1, vec![&model_output], vec![500]),
|
|
(2, vec![&model_output, &model_output], vec![500, 400]),
|
|
(
|
|
3,
|
|
vec![&model_output, &model_output, &model_output],
|
|
vec![500, 400, 300],
|
|
),
|
|
] {
|
|
let result = sampler.predictor_step(&outputs, ×teps, &sample, order);
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap().shape(), sample.shape());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_corrector_functionality() {
|
|
let mut sampler = UniPCSampler::new(
|
|
UniPCConfig {
|
|
use_corrector: true,
|
|
corrector_order: 2,
|
|
corrector_iterations: 2,
|
|
..Default::default()
|
|
},
|
|
create_test_noise_generator(),
|
|
)
|
|
.unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
|
|
for order in [1, 2, 3] {
|
|
let result = sampler.corrector_step(&sample, &model_output, 500, order);
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap().shape(), sample.shape());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_full_sampling_step() {
|
|
let mut sampler =
|
|
UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
|
|
let result = sampler.step(&model_output, 500, &sample).unwrap();
|
|
assert_eq!(result.shape(), sample.shape());
|
|
assert_eq!(sampler.stats().nfe, 1);
|
|
assert!(sampler.stats().predictor_steps > 0);
|
|
|
|
let mut current_sample = sample;
|
|
let timesteps = vec![500u32, 400, 300, 200, 100];
|
|
for timestep in timesteps {
|
|
current_sample = sampler
|
|
.step(&model_output, timestep, ¤t_sample)
|
|
.unwrap();
|
|
}
|
|
assert!(sampler.stats().nfe >= 5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_prediction_type_conversion() {
|
|
let sampler =
|
|
UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let timestep = 500;
|
|
|
|
let conversions = [
|
|
(PredictionType::Epsilon, PredictionType::VPrediction),
|
|
(PredictionType::Epsilon, PredictionType::Data),
|
|
(PredictionType::Epsilon, PredictionType::Epsilon),
|
|
];
|
|
for (from, to) in conversions {
|
|
let result =
|
|
sampler.convert_prediction_type(&model_output, timestep, &sample, from, to);
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap().shape(), model_output.shape());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_cache_functionality() {
|
|
let mut sampler =
|
|
UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
|
|
assert!(sampler.get_cached_output(500, &sample).is_none());
|
|
assert_eq!(sampler.stats().cache_misses, 1);
|
|
assert!(sampler.cache_output(500, &sample, &model_output).is_ok());
|
|
|
|
// Get cached output and verify shape before checking stats
|
|
let cached_shape = {
|
|
let cached = sampler.get_cached_output(500, &sample);
|
|
assert!(cached.is_some());
|
|
cached.unwrap().shape().to_vec()
|
|
};
|
|
|
|
assert_eq!(sampler.stats().cache_hits, 1);
|
|
assert_eq!(cached_shape, model_output.shape().dims());
|
|
}
|
|
|
|
#[test]
|
|
fn test_fast_timesteps_configuration() {
|
|
let sampler =
|
|
UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
|
|
for num_steps in [5, 7, 10, 15] {
|
|
let ts = sampler.configure_fast_timesteps(num_steps).unwrap();
|
|
assert_eq!(ts.len(), num_steps as usize);
|
|
for i in 1..ts.len() {
|
|
assert!(ts[i - 1] > ts[i]);
|
|
}
|
|
}
|
|
assert!(sampler.configure_fast_timesteps(0).is_err());
|
|
assert!(sampler.configure_fast_timesteps(2000).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_adaptive_order_selection() {
|
|
let mut sampler = UniPCSampler::new(
|
|
UniPCConfig {
|
|
adaptive_order: true,
|
|
max_order: 3,
|
|
..Default::default()
|
|
},
|
|
create_test_noise_generator(),
|
|
)
|
|
.unwrap();
|
|
|
|
for (error, steps) in [(0.01, 10), (0.1, 5)] {
|
|
let (pred_order, corr_order) = sampler.select_adaptive_order(error, steps).unwrap();
|
|
assert!(pred_order >= 1 && pred_order <= 3 && corr_order >= 1 && corr_order <= 3);
|
|
}
|
|
let (pred_low, _) = sampler.select_adaptive_order(0.5, 5).unwrap();
|
|
let (pred_high, _) = sampler.select_adaptive_order(0.01, 10).unwrap();
|
|
assert!(pred_low <= pred_high);
|
|
}
|
|
|
|
#[test]
|
|
fn test_variance_reduction() {
|
|
let sampler = UniPCSampler::new(
|
|
UniPCConfig {
|
|
variance_reduction: true,
|
|
..Default::default()
|
|
},
|
|
create_test_noise_generator(),
|
|
)
|
|
.unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let reduced = sampler
|
|
.apply_variance_reduction(&sample, &model_output, 500)
|
|
.unwrap();
|
|
assert_eq!(reduced.shape(), sample.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_history_management() {
|
|
let mut sampler = UniPCSampler::new(
|
|
UniPCConfig {
|
|
predictor_order: 3,
|
|
..Default::default()
|
|
},
|
|
create_test_noise_generator(),
|
|
)
|
|
.unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
|
|
for timestep in (100..=500).step_by(50) {
|
|
assert!(
|
|
sampler
|
|
.update_history(&model_output, timestep, &sample)
|
|
.is_ok()
|
|
);
|
|
}
|
|
assert!(sampler.timestep_history.len() <= sampler.config.max_order as usize);
|
|
assert!(sampler.predictor_outputs.len() <= sampler.config.max_order as usize);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reset_functionality() {
|
|
let mut sampler =
|
|
UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();
|
|
let sample = create_test_tensor(vec![1, 3, 32, 32]);
|
|
let model_output = create_test_tensor(vec![1, 3, 32, 32]);
|
|
|
|
let _result = sampler.step(&model_output, 500, &sample).unwrap();
|
|
let _cache = sampler.cache_output(500, &sample, &model_output);
|
|
assert!(sampler.stats().nfe > 0 && !sampler.eval_cache.is_empty());
|
|
|
|
sampler.reset();
|
|
assert_eq!(sampler.stats().nfe, 0);
|
|
assert!(sampler.eval_cache.is_empty() && sampler.timestep_history.is_empty());
|
|
}
|
|
}
|