335 lines
9.9 KiB
Rust
335 lines
9.9 KiB
Rust
//! Tests for Classifier-Free Guidance in diffusion models
|
|
//!
|
|
//! TDD: Tests define behavior for CFG implementation
|
|
|
|
use rtx_diffuse::{CFGConfig, ClassifierFreeGuidance, DiffusionScheduler, GuidanceScale, Result};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
// Helper function to replace approx::assert_abs_diff_eq
|
|
fn assert_abs_diff_eq(a: f32, b: f32, epsilon: f32) {
|
|
assert!(
|
|
(a - b).abs() < epsilon,
|
|
"Expected |{} - {}| < {}, but got {}",
|
|
a,
|
|
b,
|
|
epsilon,
|
|
(a - b).abs()
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_creation() {
|
|
let device = Device::cpu();
|
|
|
|
// Create CFG with standard configuration
|
|
let config = CFGConfig {
|
|
guidance_scale: 7.5,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device);
|
|
assert!(cfg.is_ok());
|
|
|
|
let cfg = cfg.unwrap();
|
|
assert_eq!(cfg.guidance_scale(), 7.5);
|
|
assert!(!cfg.uses_dynamic_thresholding());
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_conditional_unconditional_mixing() {
|
|
let device = Device::cpu();
|
|
|
|
let config = CFGConfig {
|
|
guidance_scale: 7.5,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
// Create conditional and unconditional predictions
|
|
let shape = vec![1, 3, 64, 64];
|
|
let conditional = Tensor::ones(&shape, &device).unwrap();
|
|
let unconditional = Tensor::zeros(&shape, &device).unwrap();
|
|
|
|
// Apply CFG: output = unconditional + scale * (conditional - unconditional)
|
|
let guided = cfg.apply(&conditional, &unconditional).unwrap();
|
|
|
|
// Check the formula: guided = 0 + 7.5 * (1 - 0) = 7.5
|
|
let guided_data = guided.to_vec().unwrap();
|
|
assert_abs_diff_eq(guided_data[0], 7.5, 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_with_batch() {
|
|
let device = Device::cpu();
|
|
|
|
let config = CFGConfig {
|
|
guidance_scale: 5.0,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
// Batch of 4 samples
|
|
let batch_size = 4;
|
|
let shape = vec![batch_size, 3, 32, 32];
|
|
|
|
let conditional = Tensor::randn(&shape, &device).unwrap();
|
|
let unconditional = Tensor::randn(&shape, &device).unwrap();
|
|
|
|
let guided = cfg.apply(&conditional, &unconditional).unwrap();
|
|
|
|
// Output should have same shape
|
|
assert_eq!(guided.shape().dims(), shape);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_guidance_rescale() {
|
|
let device = Device::cpu();
|
|
|
|
// Test guidance rescaling to prevent oversaturation
|
|
let config = CFGConfig {
|
|
guidance_scale: 10.0,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.7, // Rescale factor
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
let shape = vec![1, 3, 64, 64];
|
|
let conditional = Tensor::ones(&shape, &device)
|
|
.unwrap()
|
|
.mul_scalar(2.0)
|
|
.unwrap();
|
|
let unconditional = Tensor::ones(&shape, &device).unwrap();
|
|
|
|
let guided = cfg.apply(&conditional, &unconditional).unwrap();
|
|
|
|
// With rescaling, the values should be moderated
|
|
let guided_data = guided.to_vec().unwrap();
|
|
|
|
// Without rescale: 1 + 10 * (2 - 1) = 11
|
|
// With rescale, it should be less
|
|
assert!(
|
|
guided_data[0] < 11.0,
|
|
"Rescaling should reduce guidance magnitude"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_dynamic_thresholding() {
|
|
let device = Device::cpu();
|
|
|
|
let config = CFGConfig {
|
|
guidance_scale: 15.0, // High guidance
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: true,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
// Create predictions that would overflow without thresholding
|
|
let shape = vec![1, 3, 32, 32];
|
|
let mut cond_data = vec![0.0f32; 3 * 32 * 32];
|
|
let mut uncond_data = vec![0.0f32; 3 * 32 * 32];
|
|
|
|
// Set some extreme values
|
|
for i in 0..100 {
|
|
cond_data[i] = 5.0;
|
|
uncond_data[i] = -1.0;
|
|
}
|
|
|
|
let conditional = Tensor::from_data(cond_data, shape.clone(), &device).unwrap();
|
|
let unconditional = Tensor::from_data(uncond_data, shape, &device).unwrap();
|
|
|
|
let guided = cfg.apply(&conditional, &unconditional).unwrap();
|
|
let guided_data = guided.to_vec().unwrap();
|
|
|
|
// Dynamic thresholding should prevent extreme values
|
|
let max_val = guided_data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
let min_val = guided_data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
|
|
// Values should be clamped to reasonable range
|
|
assert!(
|
|
max_val < 100.0,
|
|
"Dynamic thresholding should prevent extreme positive values"
|
|
);
|
|
assert!(
|
|
min_val > -100.0,
|
|
"Dynamic thresholding should prevent extreme negative values"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_variable_guidance_scale() {
|
|
let device = Device::cpu();
|
|
|
|
// Test with time-varying guidance scale
|
|
let cfg = ClassifierFreeGuidance::with_schedule(
|
|
GuidanceScale::Linear {
|
|
start: 1.0,
|
|
end: 10.0,
|
|
},
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let shape = vec![1, 3, 32, 32];
|
|
let conditional = Tensor::ones(&shape, &device).unwrap();
|
|
let unconditional = Tensor::zeros(&shape, &device).unwrap();
|
|
|
|
// Early timestep (t=0.0) - should use start scale
|
|
let guided_early = cfg
|
|
.apply_with_timestep(&conditional, &unconditional, 0.0)
|
|
.unwrap();
|
|
let early_data = guided_early.to_vec().unwrap();
|
|
assert_abs_diff_eq(early_data[0], 1.0, 1e-5); // scale = 1.0
|
|
|
|
// Late timestep (t=1.0) - should use end scale
|
|
let guided_late = cfg
|
|
.apply_with_timestep(&conditional, &unconditional, 1.0)
|
|
.unwrap();
|
|
let late_data = guided_late.to_vec().unwrap();
|
|
assert_abs_diff_eq(late_data[0], 10.0, 1e-5); // scale = 10.0
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_with_multiple_conditions() {
|
|
let device = Device::cpu();
|
|
|
|
// Test CFG with multiple conditioning signals (e.g., text + image)
|
|
let config = CFGConfig {
|
|
guidance_scale: 7.5,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
let shape = vec![1, 3, 32, 32];
|
|
|
|
// Multiple conditional predictions (e.g., from different prompts)
|
|
let cond1 = Tensor::ones(&shape, &device).unwrap();
|
|
let cond2 = Tensor::ones(&shape, &device)
|
|
.unwrap()
|
|
.mul_scalar(0.5)
|
|
.unwrap();
|
|
let unconditional = Tensor::zeros(&shape, &device).unwrap();
|
|
|
|
// Combine conditions (average in this case)
|
|
let combined = cfg
|
|
.apply_multi(&[&cond1, &cond2], &unconditional, &[0.7, 0.3])
|
|
.unwrap();
|
|
|
|
// Check weighted combination
|
|
let combined_data = combined.to_vec().unwrap();
|
|
// Expected: 0 + 7.5 * ((0.7 * 1 + 0.3 * 0.5) - 0) = 7.5 * 0.85 = 6.375
|
|
assert_abs_diff_eq(combined_data[0], 6.375, 1e-4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_negative_prompting() {
|
|
let device = Device::cpu();
|
|
|
|
let config = CFGConfig {
|
|
guidance_scale: 7.5,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
let shape = vec![1, 3, 32, 32];
|
|
|
|
// Positive condition
|
|
let positive = Tensor::ones(&shape, &device).unwrap();
|
|
// Negative condition (what to avoid)
|
|
let negative = Tensor::ones(&shape, &device)
|
|
.unwrap()
|
|
.mul_scalar(-0.5)
|
|
.unwrap();
|
|
|
|
// Apply with negative prompting
|
|
let guided = cfg.apply_with_negative(&positive, &negative).unwrap();
|
|
|
|
// Should push away from negative
|
|
let guided_data = guided.to_vec().unwrap();
|
|
assert!(
|
|
guided_data[0] > 1.0,
|
|
"Should amplify away from negative prompt"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_gradient_scaling() {
|
|
let device = Device::cpu();
|
|
|
|
// Test gradient scaling for training stability
|
|
let config = CFGConfig {
|
|
guidance_scale: 7.5,
|
|
unconditional_guidance_scale: 1.0,
|
|
guidance_rescale: 0.0,
|
|
dynamic_thresholding: false,
|
|
thresholding_percentile: 0.995,
|
|
};
|
|
|
|
let cfg = ClassifierFreeGuidance::new(config, &device).unwrap();
|
|
|
|
// Create tensors that require gradients
|
|
let shape = vec![2, 3, 32, 32];
|
|
let mut conditional = Tensor::randn(&shape, &device).unwrap();
|
|
conditional.set_requires_grad(true);
|
|
let mut unconditional = Tensor::randn(&shape, &device).unwrap();
|
|
unconditional.set_requires_grad(true);
|
|
|
|
let guided = cfg.apply(&conditional, &unconditional).unwrap();
|
|
|
|
// Check that gradients can flow
|
|
assert!(guided.requires_grad());
|
|
}
|
|
|
|
#[test]
|
|
fn test_cfg_per_channel_guidance() {
|
|
let device = Device::cpu();
|
|
|
|
// Test different guidance scales per channel (RGB)
|
|
let cfg = ClassifierFreeGuidance::with_per_channel_scales(
|
|
vec![5.0, 7.5, 10.0], // Different scales for R, G, B
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let shape = vec![1, 3, 32, 32];
|
|
let conditional = Tensor::ones(&shape, &device).unwrap();
|
|
let unconditional = Tensor::zeros(&shape, &device).unwrap();
|
|
|
|
let guided = cfg.apply(&conditional, &unconditional).unwrap();
|
|
let guided_data = guided.to_vec().unwrap();
|
|
|
|
// Check different scales applied to different channels
|
|
let r_val = guided_data[0]; // Red channel, first pixel
|
|
let g_val = guided_data[32 * 32]; // Green channel, first pixel
|
|
let b_val = guided_data[2 * 32 * 32]; // Blue channel, first pixel
|
|
|
|
assert_abs_diff_eq(r_val, 5.0, 1e-5);
|
|
assert_abs_diff_eq(g_val, 7.5, 1e-5);
|
|
assert_abs_diff_eq(b_val, 10.0, 1e-5);
|
|
}
|