192 lines
7.5 KiB
Rust
192 lines
7.5 KiB
Rust
#!/usr/bin/env rust-script
|
|
//! Comprehensive demo of LoRA for Diffusion Models
|
|
//!
|
|
//! This demonstrates the complete LoRA implementation with all features:
|
|
//! - Different adapter configurations for various use cases
|
|
//! - Multiple adapter support with blending
|
|
//! - Memory-efficient parameter management
|
|
//! - Layer targeting and selective adaptation
|
|
//! - Advanced adapter management utilities
|
|
|
|
use rtx_diffuse::{
|
|
LoRAConfig, LoRAAdapter, AttentionLoRALayer, LayerTargeting,
|
|
LoRAUNet, LoRAManager, UNetConfig, DiffusionError, Result
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
fn main() -> Result<()> {
|
|
println!("🎨 LoRA for Diffusion Models - Comprehensive Demo");
|
|
println!("================================================\n");
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// 1. Different LoRA configurations for various use cases
|
|
println!("1. LoRA Configuration Presets:");
|
|
|
|
let style_config = LoRAConfig::for_style_transfer(16);
|
|
println!(" Style Transfer: rank={}, alpha={}, scaling={:.2}",
|
|
style_config.rank, style_config.alpha, style_config.scaling());
|
|
|
|
let concept_config = LoRAConfig::for_concept_learning(8);
|
|
println!(" Concept Learning: rank={}, alpha={}, scaling={:.2}",
|
|
concept_config.rank, concept_config.alpha, concept_config.scaling());
|
|
|
|
let full_config = LoRAConfig::for_full_fine_tuning(32);
|
|
println!(" Full Fine-tuning: rank={}, alpha={}, scaling={:.2}",
|
|
full_config.rank, full_config.alpha, full_config.scaling());
|
|
|
|
// 2. Create adapters for different art styles
|
|
println!("\n2. Creating Style Adapters:");
|
|
|
|
let anime_adapter = LoRAAdapter::new("anime_style", style_config.clone(), &device)?;
|
|
println!(" ✓ Anime Style: {} parameters", anime_adapter.total_parameters());
|
|
|
|
let realistic_adapter = LoRAAdapter::new("realistic_portrait", concept_config, &device)?;
|
|
println!(" ✓ Realistic Portrait: {} parameters", realistic_adapter.total_parameters());
|
|
|
|
let abstract_adapter = LoRAAdapter::new("abstract_art", full_config, &device)?;
|
|
println!(" ✓ Abstract Art: {} parameters", abstract_adapter.total_parameters());
|
|
|
|
// 3. Attention layer targeting
|
|
println!("\n3. Layer Targeting:");
|
|
|
|
let targeting = LayerTargeting::new(vec![
|
|
"self_attn".to_string(),
|
|
"cross_attn".to_string(),
|
|
"to_q".to_string(),
|
|
"to_k".to_string(),
|
|
"to_v".to_string()
|
|
]);
|
|
|
|
let test_layers = [
|
|
"encoder.0.self_attn.to_q",
|
|
"encoder.0.cross_attn.to_k",
|
|
"decoder.mlp.fc1",
|
|
"encoder.1.self_attn.to_v"
|
|
];
|
|
|
|
for layer in &test_layers {
|
|
let should_apply = targeting.should_apply_lora(layer);
|
|
println!(" {} -> {}", layer, if should_apply { "✓ Apply" } else { "✗ Skip" });
|
|
}
|
|
|
|
// 4. UNet with multiple LoRA adapters
|
|
println!("\n4. UNet LoRA Integration:");
|
|
|
|
let unet_config = UNetConfig::default();
|
|
let mut lora_unet = LoRAUNet::new(unet_config, &device)?;
|
|
|
|
// Apply adapters
|
|
lora_unet.apply_lora_adapter(anime_adapter)?;
|
|
lora_unet.apply_lora_adapter(realistic_adapter)?;
|
|
lora_unet.apply_lora_adapter(abstract_adapter)?;
|
|
|
|
println!(" ✓ Applied {} adapters", lora_unet.num_lora_adapters());
|
|
println!(" ✓ Total adapter parameters: {}", lora_unet.total_adapter_parameters());
|
|
println!(" ✓ Total adapter memory: {} MB",
|
|
lora_unet.total_adapter_memory() / (1024 * 1024));
|
|
|
|
// 5. Adapter blending and control
|
|
println!("\n5. Adapter Blending:");
|
|
|
|
// Blend adapters: 50% anime, 30% realistic, 20% abstract
|
|
lora_unet.set_adapter_weights(vec![0.5, 0.3, 0.2])?;
|
|
println!(" ✓ Set blend weights: [0.5, 0.3, 0.2]");
|
|
|
|
// Dynamically enable/disable adapters
|
|
lora_unet.set_adapter_enabled("abstract_art", false)?;
|
|
println!(" ✓ Disabled abstract art adapter");
|
|
|
|
println!(" Active adapters: {:?}", lora_unet.adapter_names());
|
|
|
|
// 6. Advanced adapter management
|
|
println!("\n6. Advanced LoRA Management:");
|
|
|
|
let mut manager = LoRAManager::new(&device);
|
|
|
|
// Create adapter library
|
|
let styles = [
|
|
("cyberpunk", 16),
|
|
("watercolor", 8),
|
|
("oil_painting", 24),
|
|
("sketch", 4)
|
|
];
|
|
|
|
manager.create_adapter_library(&styles)?;
|
|
println!(" ✓ Created adapter library with {} styles", styles.len());
|
|
|
|
// Optimal rank suggestion
|
|
let optimal_rank = manager.suggest_optimal_rank(100.0, 512); // 100MB budget, 512 dim
|
|
println!(" ✓ Suggested optimal rank for 100MB budget: {}", optimal_rank);
|
|
|
|
// 7. Memory efficiency analysis
|
|
println!("\n7. Memory Efficiency Analysis:");
|
|
|
|
// Create sample attention layer for analysis
|
|
let layer = AttentionLoRALayer::new(
|
|
"sample_attention",
|
|
512,
|
|
512,
|
|
LoRAConfig::default(),
|
|
&device
|
|
)?;
|
|
|
|
println!(" Layer: {} -> {}", layer.in_features(), layer.out_features());
|
|
println!(" LoRA parameters: {} (rank {})", layer.num_parameters(), 8);
|
|
println!(" Compression ratio: {:.1}x", layer.compression_ratio());
|
|
println!(" Memory savings: {:.1}%",
|
|
(1.0 - (1.0 / layer.compression_ratio())) * 100.0);
|
|
|
|
// 8. Parameter shapes and initialization
|
|
println!("\n8. Parameter Analysis:");
|
|
|
|
let (a_shape, b_shape) = layer.parameter_shapes();
|
|
println!(" Matrix A shape: {:?}", a_shape);
|
|
println!(" Matrix B shape: {:?}", b_shape);
|
|
println!(" Dropout rate: {}", layer.dropout_rate());
|
|
println!(" Spectral norm: {}", layer.has_spectral_norm());
|
|
println!(" Gradient checkpointing: {}", layer.has_gradient_checkpointing());
|
|
|
|
// 9. Forward pass demonstration
|
|
println!("\n9. Forward Pass Test:");
|
|
|
|
// Create test tensors
|
|
let input = create_test_tensor(&[2, 64, 512], &device)?; // [batch, seq, dim]
|
|
let weight = create_test_tensor(&[512, 512], &device)?; // [out, in]
|
|
|
|
let output = layer.forward(&input, &weight)?;
|
|
println!(" ✓ Forward pass: {:?} -> {:?}",
|
|
input.shape().dims(), output.shape().dims());
|
|
|
|
// 10. Weight merging demonstration
|
|
println!("\n10. Weight Merging:");
|
|
|
|
let merged = layer.merge_weights(&weight)?;
|
|
println!(" ✓ Merged weights shape: {:?}", merged.shape().dims());
|
|
|
|
let unmerged = layer.unmerge_weights(&merged)?;
|
|
println!(" ✓ Unmerged weights shape: {:?}", unmerged.shape().dims());
|
|
|
|
println!("\n🎉 LoRA for Diffusion Demo Complete!");
|
|
println!("\n📊 Summary of Features Demonstrated:");
|
|
println!(" • LoRA adapter creation and management");
|
|
println!(" • Multiple adapter support with weighted blending");
|
|
println!(" • Selective layer targeting for efficient fine-tuning");
|
|
println!(" • Memory-efficient low-rank adaptation (up to 100x compression)");
|
|
println!(" • Different configurations for style transfer, concept learning, etc.");
|
|
println!(" • Advanced adapter fusion and management utilities");
|
|
println!(" • Proper initialization and forward pass implementation");
|
|
println!(" • Weight merging for inference optimization");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// Helper function to create test tensors
|
|
fn create_test_tensor(shape: &[usize], device: &Device) -> Result<Tensor> {
|
|
let size = shape.iter().product::<usize>();
|
|
let data = vec![1.0f32; size];
|
|
|
|
Tensor::new(data, shape.to_vec()).map_err(|e| DiffusionError::TensorOperation {
|
|
details: format!("Failed to create test tensor: {:?}", e),
|
|
})
|
|
} |