342 lines
10 KiB
Rust
342 lines
10 KiB
Rust
//! Diffusion UNet architecture.
|
|
//!
|
|
//! This module implements a simplified UNet architecture suitable for
|
|
//! diffusion models in the context of PDE solving.
|
|
|
|
use rtx_backend::Backend;
|
|
use rtx_nn::{GenericLinear, GenericModule};
|
|
use rtx_tensor::GenericTensor;
|
|
|
|
/// Configuration for the diffusion UNet.
|
|
#[derive(Debug, Clone)]
|
|
pub struct UNetConfig {
|
|
/// Number of input channels
|
|
pub in_channels: usize,
|
|
/// Number of output channels (usually same as input)
|
|
pub out_channels: usize,
|
|
/// Base channel dimension
|
|
pub base_channels: usize,
|
|
/// Hidden dimension multiplier
|
|
pub hidden_mult: usize,
|
|
/// Dimension of time embedding
|
|
pub time_embed_dim: usize,
|
|
}
|
|
|
|
impl Default for UNetConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
in_channels: 1,
|
|
out_channels: 1,
|
|
base_channels: 64,
|
|
hidden_mult: 2,
|
|
time_embed_dim: 128,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl UNetConfig {
|
|
/// Create config for small model (for testing/debugging).
|
|
pub fn small() -> Self {
|
|
Self {
|
|
in_channels: 1,
|
|
out_channels: 1,
|
|
base_channels: 32,
|
|
hidden_mult: 2,
|
|
time_embed_dim: 64,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Sinusoidal time embedding module.
|
|
///
|
|
/// Maps scalar timestep to a fixed-dimension embedding vector using
|
|
/// sinusoidal positional encoding.
|
|
pub struct TimeEmbedding<B: Backend<FloatElem = f32>> {
|
|
/// Linear projection from embedding to hidden
|
|
fc1: GenericLinear<B>,
|
|
/// Linear projection from hidden to output
|
|
fc2: GenericLinear<B>,
|
|
/// Embedding dimension
|
|
embed_dim: usize,
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> TimeEmbedding<B> {
|
|
/// Create a new time embedding module.
|
|
pub fn new(embed_dim: usize, hidden_dim: usize, device: &B::Device) -> Self {
|
|
let fc1 = GenericLinear::new(embed_dim, hidden_dim, true, device);
|
|
let fc2 = GenericLinear::new(hidden_dim, hidden_dim, true, device);
|
|
|
|
Self {
|
|
fc1,
|
|
fc2,
|
|
embed_dim,
|
|
}
|
|
}
|
|
|
|
/// Get sinusoidal embedding for timestep.
|
|
pub fn get_embedding(&self, t: usize, device: &B::Device) -> GenericTensor<B, 2> {
|
|
let half_dim = self.embed_dim / 2;
|
|
let emb_scale = -(10000.0f32.ln()) / (half_dim as f32 - 1.0);
|
|
|
|
let mut embedding = vec![0.0f32; self.embed_dim];
|
|
|
|
for i in 0..half_dim {
|
|
let freq = (i as f32 * emb_scale).exp();
|
|
let angle = t as f32 * freq;
|
|
embedding[i] = angle.sin();
|
|
embedding[i + half_dim] = angle.cos();
|
|
}
|
|
|
|
GenericTensor::<B, 2>::from_slice(&embedding, [1, self.embed_dim], device)
|
|
}
|
|
|
|
/// Forward pass: embed timestep and project through MLP.
|
|
pub fn forward(&self, t: usize, device: &B::Device) -> GenericTensor<B, 2> {
|
|
let emb = self.get_embedding(t, device);
|
|
|
|
// MLP with SiLU activation
|
|
let h = GenericModule::forward(&self.fc1, &emb);
|
|
let h = silu(&h);
|
|
GenericModule::forward(&self.fc2, &h)
|
|
}
|
|
}
|
|
|
|
/// SiLU (Swish) activation: x * sigmoid(x)
|
|
fn silu<B: Backend<FloatElem = f32>, const D: usize>(
|
|
x: &GenericTensor<B, D>,
|
|
) -> GenericTensor<B, D> {
|
|
// x * sigmoid(x) = x / (1 + exp(-x))
|
|
let neg_x = x.mul_scalar(-1.0);
|
|
let exp_neg_x = neg_x.exp();
|
|
let one_plus_exp = exp_neg_x.add_scalar(1.0);
|
|
x.div(&one_plus_exp)
|
|
}
|
|
|
|
/// Simplified Diffusion UNet for 2D fields.
|
|
///
|
|
/// This is a simplified architecture that treats the 2D field as a flattened
|
|
/// vector and processes it through fully-connected layers with time conditioning.
|
|
///
|
|
/// Architecture:
|
|
/// - Input: [batch, channels, height, width] -> flatten to [batch, h*w*c]
|
|
/// - Encoder: fc -> hidden -> fc
|
|
/// - Time conditioning added at hidden layer
|
|
/// - Decoder: fc -> output
|
|
/// - Output: reshape back to [batch, channels, height, width]
|
|
pub struct DiffusionUNet<B: Backend<FloatElem = f32>> {
|
|
/// Configuration
|
|
config: UNetConfig,
|
|
/// Time embedding module
|
|
time_embedding: TimeEmbedding<B>,
|
|
/// Input projection
|
|
encoder1: GenericLinear<B>,
|
|
/// Hidden layer
|
|
encoder2: GenericLinear<B>,
|
|
/// Time projection to hidden
|
|
time_proj: GenericLinear<B>,
|
|
/// Decoder layer 1
|
|
decoder1: GenericLinear<B>,
|
|
/// Output projection
|
|
decoder2: GenericLinear<B>,
|
|
/// Expected spatial size (height * width)
|
|
spatial_size: usize,
|
|
}
|
|
|
|
impl<B: Backend<FloatElem = f32>> DiffusionUNet<B> {
|
|
/// Create a new diffusion UNet.
|
|
///
|
|
/// # Arguments
|
|
/// * `config` - Model configuration
|
|
/// * `spatial_size` - Expected height * width of input
|
|
/// * `device` - Device to create on
|
|
pub fn new(config: UNetConfig, spatial_size: usize, device: &B::Device) -> Self {
|
|
let input_dim = config.in_channels * spatial_size;
|
|
let hidden_dim = config.base_channels * config.hidden_mult;
|
|
let output_dim = config.out_channels * spatial_size;
|
|
|
|
// Time embedding
|
|
let time_embedding =
|
|
TimeEmbedding::new(config.time_embed_dim, config.time_embed_dim, device);
|
|
|
|
// Encoder
|
|
let encoder1 = GenericLinear::new(input_dim, config.base_channels, true, device);
|
|
let encoder2 = GenericLinear::new(config.base_channels, hidden_dim, true, device);
|
|
|
|
// Time projection to hidden dimension
|
|
let time_proj = GenericLinear::new(config.time_embed_dim, hidden_dim, true, device);
|
|
|
|
// Decoder
|
|
let decoder1 = GenericLinear::new(hidden_dim, config.base_channels, true, device);
|
|
let decoder2 = GenericLinear::new(config.base_channels, output_dim, true, device);
|
|
|
|
Self {
|
|
config,
|
|
time_embedding,
|
|
encoder1,
|
|
encoder2,
|
|
time_proj,
|
|
decoder1,
|
|
decoder2,
|
|
spatial_size,
|
|
}
|
|
}
|
|
|
|
/// Forward pass for noise prediction.
|
|
///
|
|
/// # Arguments
|
|
/// * `x` - Noisy input [batch, channels, height, width]
|
|
/// * `t` - Timestep (scalar)
|
|
/// * `device` - Device
|
|
///
|
|
/// # Returns
|
|
/// Predicted noise with same shape as input
|
|
pub fn forward(
|
|
&self,
|
|
x: &GenericTensor<B, 4>,
|
|
t: usize,
|
|
device: &B::Device,
|
|
) -> GenericTensor<B, 4> {
|
|
let shape = x.shape();
|
|
let batch = shape[0];
|
|
let channels = shape[1];
|
|
let height = shape[2];
|
|
let width = shape[3];
|
|
|
|
// Flatten: [batch, c, h, w] -> [batch, c*h*w]
|
|
let x_flat = x.reshape([batch, channels * height * width]);
|
|
|
|
// Get time embedding [1, time_embed_dim]
|
|
let time_emb = self.time_embedding.forward(t, device);
|
|
|
|
// Expand time embedding to batch size [batch, time_embed_dim]
|
|
let time_emb = expand_batch(&time_emb, batch);
|
|
|
|
// Encoder
|
|
let h = GenericModule::forward(&self.encoder1, &x_flat);
|
|
let h = silu(&h);
|
|
let h = GenericModule::forward(&self.encoder2, &h);
|
|
let h = silu(&h);
|
|
|
|
// Add time conditioning
|
|
let time_proj = GenericModule::forward(&self.time_proj, &time_emb);
|
|
let time_proj = silu(&time_proj);
|
|
let h = h.add(&time_proj);
|
|
|
|
// Decoder
|
|
let h = GenericModule::forward(&self.decoder1, &h);
|
|
let h = silu(&h);
|
|
let output = GenericModule::forward(&self.decoder2, &h);
|
|
|
|
// Reshape back: [batch, c*h*w] -> [batch, c, h, w]
|
|
output.reshape([batch, channels, height, width])
|
|
}
|
|
}
|
|
|
|
/// Expand a [1, dim] tensor to [batch, dim] by repeating.
|
|
fn expand_batch<B: Backend<FloatElem = f32>>(
|
|
x: &GenericTensor<B, 2>,
|
|
batch: usize,
|
|
) -> GenericTensor<B, 2> {
|
|
let shape = x.shape();
|
|
let dim = shape[1];
|
|
|
|
let x_vec = x.to_vec();
|
|
let mut expanded = Vec::with_capacity(batch * dim);
|
|
for _ in 0..batch {
|
|
expanded.extend_from_slice(&x_vec);
|
|
}
|
|
|
|
GenericTensor::<B, 2>::from_slice(&expanded, [batch, dim], &x.device())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rtx_backend_cpu::{CpuBackend, CpuDevice};
|
|
|
|
#[test]
|
|
fn test_time_embedding() {
|
|
let device = CpuDevice::default();
|
|
let time_emb = TimeEmbedding::<CpuBackend>::new(64, 128, &device);
|
|
|
|
let emb = time_emb.forward(100, &device);
|
|
let shape = emb.shape();
|
|
|
|
assert_eq!(shape[0], 1);
|
|
assert_eq!(shape[1], 128);
|
|
}
|
|
|
|
#[test]
|
|
fn test_silu_activation() {
|
|
let device = CpuDevice::default();
|
|
let x = GenericTensor::<CpuBackend, 2>::from_slice(&[0.0, 1.0, -1.0, 2.0], [2, 2], &device);
|
|
|
|
let y = silu(&x);
|
|
let y_vec = y.to_vec();
|
|
|
|
// silu(0) = 0
|
|
assert!((y_vec[0] - 0.0).abs() < 1e-5);
|
|
// silu(1) ≈ 0.731
|
|
assert!((y_vec[1] - 0.731).abs() < 0.01);
|
|
// silu(-1) ≈ -0.269
|
|
assert!((y_vec[2] + 0.269).abs() < 0.01);
|
|
}
|
|
|
|
#[test]
|
|
fn test_unet_creation() {
|
|
let device = CpuDevice::default();
|
|
let config = UNetConfig::small();
|
|
let spatial_size = 64; // 8x8
|
|
let _unet = DiffusionUNet::<CpuBackend>::new(config, spatial_size, &device);
|
|
}
|
|
|
|
#[test]
|
|
fn test_unet_forward() {
|
|
let device = CpuDevice::default();
|
|
let config = UNetConfig {
|
|
in_channels: 1,
|
|
out_channels: 1,
|
|
base_channels: 16,
|
|
hidden_mult: 2,
|
|
time_embed_dim: 32,
|
|
};
|
|
|
|
let height = 8;
|
|
let width = 8;
|
|
let spatial_size = height * width;
|
|
|
|
let unet = DiffusionUNet::<CpuBackend>::new(config, spatial_size, &device);
|
|
|
|
// Create test input [1, 1, 8, 8]
|
|
let input = GenericTensor::<CpuBackend, 4>::rand([1, 1, height, width], &device);
|
|
|
|
let output = unet.forward(&input, 50, &device);
|
|
let output_shape = output.shape();
|
|
|
|
assert_eq!(output_shape[0], 1); // batch
|
|
assert_eq!(output_shape[1], 1); // channels
|
|
assert_eq!(output_shape[2], height); // height
|
|
assert_eq!(output_shape[3], width); // width
|
|
}
|
|
|
|
#[test]
|
|
fn test_expand_batch() {
|
|
let device = CpuDevice::default();
|
|
let x = GenericTensor::<CpuBackend, 2>::from_slice(&[1.0, 2.0, 3.0], [1, 3], &device);
|
|
|
|
let expanded = expand_batch(&x, 4);
|
|
let shape = expanded.shape();
|
|
|
|
assert_eq!(shape[0], 4);
|
|
assert_eq!(shape[1], 3);
|
|
|
|
let data = expanded.to_vec();
|
|
// All rows should have same values
|
|
for i in 0..4 {
|
|
assert_eq!(data[i * 3], 1.0);
|
|
assert_eq!(data[i * 3 + 1], 2.0);
|
|
assert_eq!(data[i * 3 + 2], 3.0);
|
|
}
|
|
}
|
|
}
|