426 lines
15 KiB
Rust
426 lines
15 KiB
Rust
//! ControlNet TDD Tests
|
|
//!
|
|
//! Test-driven development for ControlNet implementation following:
|
|
//! "Adding Conditional Control to Text-to-Image Diffusion Models" (Zhang et al. 2023)
|
|
#![cfg(feature = "disabled_tests")]
|
|
|
|
use crate::controlnet::*;
|
|
use crate::*;
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::collections::HashMap;
|
|
|
|
/// Test configuration for ControlNet (re-export)
|
|
pub use crate::controlnet::ControlNetConfig as ControlNetTestConfig;
|
|
|
|
/// Test helper for creating test tensors
|
|
fn create_test_tensor(shape: &[usize], device: &Device) -> Result<Tensor> {
|
|
Tensor::randn(shape, device)
|
|
.map_err(|e| DiffusionError::TensorError(format!("Failed to create test tensor: {}", e)))
|
|
}
|
|
|
|
/// Test helper for creating control condition tensors
|
|
fn create_control_condition(
|
|
batch_size: usize,
|
|
channels: usize,
|
|
height: usize,
|
|
width: usize,
|
|
device: &Device,
|
|
) -> Result<Tensor> {
|
|
create_test_tensor(&[batch_size, channels, height, width], device)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod controlnet_zero_conv_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_zero_conv_initialization() {
|
|
// RED: Test fails - ZeroConv not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let in_channels = 320;
|
|
let out_channels = 320;
|
|
|
|
let zero_conv = ZeroConv::new(in_channels, out_channels).unwrap();
|
|
|
|
// Zero convolution should initialize weights to zero
|
|
let input = create_test_tensor(&[1, in_channels, 8, 8], &device).unwrap();
|
|
let output = zero_conv.forward(&input).unwrap();
|
|
|
|
// Output should be zero for zero-initialized weights
|
|
let output_data = output.to_vec().unwrap();
|
|
assert!(
|
|
output_data.iter().all(|&x| x.abs() < 1e-8),
|
|
"Zero conv output should be zero"
|
|
);
|
|
|
|
// Test that we can make it trainable
|
|
assert!(zero_conv.is_trainable(), "Zero conv should be trainable");
|
|
}
|
|
|
|
#[test]
|
|
fn test_zero_conv_gradual_learning() {
|
|
// RED: Test fails - ZeroConv gradual learning not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let mut zero_conv = ZeroConv::new(320, 320).unwrap();
|
|
|
|
// Simulate training by manually setting small weights
|
|
zero_conv.set_learning_rate(0.001);
|
|
|
|
let input = create_test_tensor(&[1, 320, 8, 8], &device).unwrap();
|
|
|
|
// Initial forward should be zero
|
|
let output1 = zero_conv.forward(&input).unwrap();
|
|
let norm1 = output1
|
|
.abs()
|
|
.unwrap()
|
|
.sum(None)
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
assert!(norm1 < 1e-6, "Initial output should be near zero");
|
|
|
|
// After simulated gradient update
|
|
zero_conv.simulate_gradient_update(&input).unwrap();
|
|
let output2 = zero_conv.forward(&input).unwrap();
|
|
let norm2 = output2
|
|
.abs()
|
|
.unwrap()
|
|
.sum(None)
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
|
|
assert!(
|
|
norm2 > norm1,
|
|
"Output should increase after gradient update"
|
|
);
|
|
assert!(
|
|
norm2 < 0.1,
|
|
"Output should still be small after initial updates"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod controlnet_conditioning_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_control_injection_at_multiple_scales() {
|
|
// RED: Test fails - ControlNet not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
let controlnet = ControlNet::new(config.clone()).unwrap();
|
|
|
|
// Create input tensors
|
|
let batch_size = 2;
|
|
let height = 64;
|
|
let width = 64;
|
|
|
|
let x =
|
|
create_test_tensor(&[batch_size, config.in_channels, height, width], &device).unwrap();
|
|
let timesteps = Tensor::from_slice(&[100.0f32, 200.0f32], &[batch_size], &device).unwrap();
|
|
let control_condition =
|
|
create_control_condition(batch_size, 3, height, width, &device).unwrap(); // Default RGB
|
|
|
|
let context = ControlContext::new()
|
|
.with_condition(control_condition)
|
|
.with_strength(0.8);
|
|
|
|
let (output, control_residuals) = controlnet.forward(&x, ×teps, &context).unwrap();
|
|
|
|
// Output should have same shape as input
|
|
assert_eq!(output.dims(), x.dims());
|
|
|
|
// Should have control residuals at multiple scales
|
|
assert!(
|
|
control_residuals.len() > 1,
|
|
"Should have multiple control residuals"
|
|
);
|
|
|
|
// Each residual should have appropriate dimensions
|
|
for (i, residual) in control_residuals.iter().enumerate() {
|
|
let dims = residual.dims();
|
|
assert_eq!(
|
|
dims[0], batch_size,
|
|
"Batch size should match for residual {}",
|
|
i
|
|
);
|
|
assert!(
|
|
dims[2] <= height,
|
|
"Height should be <= input height for residual {}",
|
|
i
|
|
);
|
|
assert!(
|
|
dims[3] <= width,
|
|
"Width should be <= input width for residual {}",
|
|
i
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_control_strength_scheduling() {
|
|
// RED: Test fails - Control strength scheduling not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
let controlnet = ControlNet::new(config.clone()).unwrap();
|
|
|
|
let batch_size = 1;
|
|
let x = create_test_tensor(&[batch_size, config.in_channels, 64, 64], &device).unwrap();
|
|
let timesteps = Tensor::from_slice(&[500.0f32], &[batch_size], &device).unwrap();
|
|
let control_condition = create_control_condition(batch_size, 3, 64, 64, &device).unwrap(); // Default RGB
|
|
|
|
// Test different control strengths
|
|
let strengths = [0.0, 0.5, 1.0];
|
|
let mut outputs = Vec::new();
|
|
|
|
for &strength in &strengths {
|
|
let context = ControlContext::new()
|
|
.with_condition(control_condition.clone())
|
|
.with_strength(strength);
|
|
|
|
let (output, _) = controlnet.forward(&x, ×teps, &context).unwrap();
|
|
outputs.push(output);
|
|
}
|
|
|
|
// With strength=0.0, should behave like no control
|
|
let no_control_context = ControlContext::new().with_strength(0.0);
|
|
let (no_control_output, _) = controlnet
|
|
.forward(&x, ×teps, &no_control_context)
|
|
.unwrap();
|
|
|
|
let diff_zero = (&outputs[0] - &no_control_output)
|
|
.unwrap()
|
|
.abs()
|
|
.unwrap()
|
|
.max()
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
assert!(
|
|
diff_zero < 1e-4,
|
|
"Zero strength should produce similar output to no control"
|
|
);
|
|
|
|
// Higher strength should produce more different results
|
|
let diff_half = (&outputs[1] - &outputs[0])
|
|
.unwrap()
|
|
.abs()
|
|
.unwrap()
|
|
.max()
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
let diff_full = (&outputs[2] - &outputs[0])
|
|
.unwrap()
|
|
.abs()
|
|
.unwrap()
|
|
.max()
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
|
|
assert!(
|
|
diff_full > diff_half,
|
|
"Full strength should differ more from zero than half strength"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_various_control_types() {
|
|
// RED: Test fails - Multiple control types not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
|
|
// Test different control types
|
|
let control_types = vec![
|
|
ControlType::Edge,
|
|
ControlType::Pose,
|
|
ControlType::Depth,
|
|
ControlType::Normal,
|
|
ControlType::Segmentation,
|
|
];
|
|
|
|
for control_type in control_types {
|
|
let controlnet = ControlNet::with_control_type(config.clone(), control_type).unwrap();
|
|
|
|
let batch_size = 1;
|
|
let x = create_test_tensor(&[batch_size, config.in_channels, 64, 64], &device).unwrap();
|
|
let timesteps = Tensor::from_slice(&[250.0f32], &[batch_size], &device).unwrap();
|
|
let control_condition = create_control_condition(
|
|
batch_size,
|
|
control_type.input_channels(),
|
|
64,
|
|
64,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let context = ControlContext::new()
|
|
.with_condition(control_condition)
|
|
.with_control_type(control_type)
|
|
.with_strength(1.0);
|
|
|
|
let result = controlnet.forward(&x, ×teps, &context);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Forward pass should work for control type {:?}",
|
|
control_type
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_trainable_copy_of_encoder_blocks() {
|
|
// RED: Test fails - Trainable encoder copy not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
let controlnet = ControlNet::new(config.clone()).unwrap();
|
|
|
|
// ControlNet should have trainable copies of UNet encoder blocks
|
|
let encoder_blocks = controlnet.get_encoder_blocks();
|
|
assert!(!encoder_blocks.is_empty(), "Should have encoder blocks");
|
|
|
|
for (i, block) in encoder_blocks.iter().enumerate() {
|
|
assert!(
|
|
block.is_trainable(),
|
|
"Encoder block {} should be trainable",
|
|
i
|
|
);
|
|
|
|
// Verify block has same architecture as UNet but different parameters
|
|
let unet = UNet::new(UNetConfig::from(config.clone())).unwrap();
|
|
let unet_encoder_blocks = unet.get_encoder_blocks();
|
|
|
|
assert_eq!(
|
|
block.parameter_count(),
|
|
unet_encoder_blocks[i].parameter_count(),
|
|
"Parameter count should match UNet encoder block {}",
|
|
i
|
|
);
|
|
|
|
// Parameters should be different (not shared)
|
|
let controlnet_params = block.get_parameters();
|
|
let unet_params = unet_encoder_blocks[i].get_parameters();
|
|
|
|
for (cp, up) in controlnet_params.iter().zip(unet_params.iter()) {
|
|
let diff = (cp - up)
|
|
.unwrap()
|
|
.abs()
|
|
.unwrap()
|
|
.max()
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
assert!(diff > 1e-6, "ControlNet and UNet parameters should differ");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod controlnet_architecture_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_controlnet_backbone_compatibility() {
|
|
// RED: Test fails - ControlNet backbone not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
let controlnet = ControlNet::new(config.clone()).unwrap();
|
|
let unet = UNet::new(UNetConfig::from(config.clone())).unwrap();
|
|
|
|
let batch_size = 1;
|
|
let x = create_test_tensor(&[batch_size, config.in_channels, 64, 64], &device).unwrap();
|
|
let timesteps = Tensor::from_slice(&[400.0f32], &[batch_size], &device).unwrap();
|
|
let control_condition = create_control_condition(batch_size, 3, 64, 64, &device).unwrap(); // Default RGB
|
|
|
|
// Test without control (should behave similarly to UNet)
|
|
let no_control_context = ControlContext::new().with_strength(0.0);
|
|
let (controlnet_output, _) = controlnet
|
|
.forward(&x, ×teps, &no_control_context)
|
|
.unwrap();
|
|
|
|
let unet_output = unet.forward(&x, ×teps, &None).unwrap();
|
|
|
|
// Outputs should be similar when no control is applied
|
|
let diff = (&controlnet_output - &unet_output)
|
|
.unwrap()
|
|
.abs()
|
|
.unwrap()
|
|
.max()
|
|
.unwrap()
|
|
.to_scalar::<f32>()
|
|
.unwrap();
|
|
assert!(
|
|
diff < 0.1,
|
|
"ControlNet without control should behave similarly to UNet"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_controlnet_memory_efficiency() {
|
|
// RED: Test fails - Memory efficiency not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
let controlnet = ControlNet::new(config.clone()).unwrap();
|
|
|
|
// Test gradient checkpointing
|
|
controlnet.enable_gradient_checkpointing(true);
|
|
assert!(controlnet.is_gradient_checkpointing_enabled());
|
|
|
|
// Test memory usage estimation
|
|
let memory_usage = controlnet.estimate_memory_usage(2, 64, 64).unwrap();
|
|
assert!(
|
|
memory_usage.parameters > 0,
|
|
"Should report parameter memory"
|
|
);
|
|
assert!(
|
|
memory_usage.activations > 0,
|
|
"Should report activation memory"
|
|
);
|
|
assert!(memory_usage.gradients > 0, "Should report gradient memory");
|
|
|
|
// With gradient checkpointing, activation memory should be lower
|
|
controlnet.enable_gradient_checkpointing(false);
|
|
let memory_usage_no_checkpoint = controlnet.estimate_memory_usage(2, 64, 64).unwrap();
|
|
|
|
assert!(
|
|
memory_usage.activations < memory_usage_no_checkpoint.activations,
|
|
"Gradient checkpointing should reduce activation memory"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_controlnet_conditioning_fusion() {
|
|
// RED: Test fails - Conditioning fusion not implemented
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = ControlNetTestConfig::default();
|
|
let controlnet = ControlNet::new(config.clone()).unwrap();
|
|
|
|
let batch_size = 2;
|
|
let x = create_test_tensor(&[batch_size, config.in_channels, 64, 64], &device).unwrap();
|
|
let timesteps = Tensor::from_slice(&[300.0f32, 400.0f32], &[batch_size], &device).unwrap();
|
|
|
|
// Test multiple conditioning inputs
|
|
let edge_condition = create_control_condition(batch_size, 1, 64, 64, &device).unwrap(); // Edge map
|
|
let depth_condition = create_control_condition(batch_size, 1, 64, 64, &device).unwrap(); // Depth map
|
|
|
|
let context = ControlContext::new()
|
|
.with_condition(edge_condition)
|
|
.with_secondary_condition("depth", depth_condition)
|
|
.with_strength(0.8)
|
|
.with_fusion_mode(FusionMode::Additive);
|
|
|
|
let result = controlnet.forward(&x, ×teps, &context);
|
|
assert!(result.is_ok(), "Multi-condition forward should work");
|
|
|
|
let (output, residuals) = result.unwrap();
|
|
assert_eq!(output.dims(), x.dims(), "Output shape should match input");
|
|
assert!(!residuals.is_empty(), "Should have control residuals");
|
|
}
|
|
}
|
|
|
|
// Note: All structs are now defined in the main controlnet module
|