265 lines
8.8 KiB
Rust
265 lines
8.8 KiB
Rust
//! Diffusion model inference adapter for ensemble prediction
|
|
//!
|
|
//! Wraps a `Denoiser` and runs a single forward pass at t=1
|
|
//! (minimal noise level). The mean of the denoiser output is
|
|
//! mapped via sigmoid to direction + confidence.
|
|
|
|
use std::sync::Mutex;
|
|
|
|
use candle_core::{Device, Tensor};
|
|
use candle_nn::{VarBuilder, VarMap};
|
|
|
|
use crate::diffusion::config::DiffusionConfig;
|
|
use crate::diffusion::denoiser::Denoiser;
|
|
use crate::dqn::mixed_precision::training_dtype;
|
|
use crate::ensemble::inference_adapter::{
|
|
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction,
|
|
};
|
|
use crate::gpu::DeviceConfig;
|
|
use crate::{MLError, MLResult};
|
|
|
|
/// Inference adapter for the Diffusion (DDPM/DDIM) denoising model.
|
|
///
|
|
/// Uses the `Denoiser` network with a minimal timestep (t=1) to process
|
|
/// features. The mean of the denoiser output provides the raw signal
|
|
/// which is mapped via sigmoid to direction [-1, 1] and confidence [0, 1].
|
|
#[allow(missing_debug_implementations)]
|
|
pub struct DiffusionInferenceAdapter {
|
|
model: Mutex<Denoiser>,
|
|
varmap: VarMap,
|
|
device: Device,
|
|
data_dim: usize,
|
|
}
|
|
|
|
// SAFETY: Denoiser uses candle tensors which are Send+Sync.
|
|
// The Mutex provides exclusive access for inference calls.
|
|
#[allow(unsafe_code)]
|
|
unsafe impl Send for DiffusionInferenceAdapter {}
|
|
#[allow(unsafe_code)]
|
|
unsafe impl Sync for DiffusionInferenceAdapter {}
|
|
|
|
impl DiffusionInferenceAdapter {
|
|
/// Create a new Diffusion inference adapter.
|
|
///
|
|
/// Uses `config.seq_len * config.feature_dim` as the data dimension.
|
|
/// For 51-dim feature input: set `feature_dim=51, seq_len=1`.
|
|
pub fn new(config: DiffusionConfig) -> MLResult<Self> {
|
|
let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
|
|
let data_dim = config.seq_len * config.feature_dim;
|
|
let varmap = VarMap::new();
|
|
let vb = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
|
|
let denoiser = Denoiser::new(
|
|
data_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
config.time_embed_dim,
|
|
vb.pp("denoiser"),
|
|
&device,
|
|
)?;
|
|
|
|
Ok(Self {
|
|
model: Mutex::new(denoiser),
|
|
varmap,
|
|
device,
|
|
data_dim,
|
|
})
|
|
}
|
|
|
|
/// Load from a safetensors checkpoint.
|
|
pub fn from_checkpoint(config: DiffusionConfig, path: &str) -> MLResult<Self> {
|
|
let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
|
|
let data_dim = config.seq_len * config.feature_dim;
|
|
let mut varmap = VarMap::new();
|
|
let vb = VarBuilder::from_varmap(&varmap, training_dtype(&device), &device);
|
|
let denoiser = Denoiser::new(
|
|
data_dim,
|
|
config.hidden_dim,
|
|
config.num_layers,
|
|
config.time_embed_dim,
|
|
vb.pp("denoiser"),
|
|
&device,
|
|
)?;
|
|
|
|
varmap
|
|
.load(path)
|
|
.map_err(|e| MLError::ModelError(format!("Failed to load Diffusion checkpoint: {e}")))?;
|
|
|
|
Ok(Self {
|
|
model: Mutex::new(denoiser),
|
|
varmap,
|
|
device,
|
|
data_dim,
|
|
})
|
|
}
|
|
|
|
/// Pad or truncate feature values to match data_dim.
|
|
fn pad_features(&self, values: &[f64]) -> Vec<f32> {
|
|
let mut padded = vec![0.0_f32; self.data_dim];
|
|
let copy_len = values.len().min(self.data_dim);
|
|
for i in 0..copy_len {
|
|
if let Some(v) = values.get(i) {
|
|
if let Some(slot) = padded.get_mut(i) {
|
|
*slot = *v as f32;
|
|
}
|
|
}
|
|
}
|
|
padded
|
|
}
|
|
}
|
|
|
|
impl ModelInferenceAdapter for DiffusionInferenceAdapter {
|
|
fn model_name(&self) -> &str {
|
|
"Diffusion"
|
|
}
|
|
|
|
fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
|
let start = std::time::Instant::now();
|
|
|
|
let padded = self.pad_features(&features.values);
|
|
let input = Tensor::from_vec(padded, (1, self.data_dim), &self.device)
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion input tensor: {e}")))?;
|
|
|
|
// Timestep t=1 (minimal noise level for feature processing)
|
|
let t = Tensor::from_vec(vec![1_u32], (1,), &self.device)
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion timestep tensor: {e}")))?;
|
|
|
|
let model = self
|
|
.model
|
|
.lock()
|
|
.map_err(|e| MLError::LockError(format!("Diffusion lock poisoned: {e}")))?;
|
|
let output = model.forward(&input, &t)?;
|
|
|
|
// Mean of denoiser output -> sigmoid -> direction + confidence
|
|
let mean_val: f32 = output
|
|
.mean_all()
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion mean: {e}")))?
|
|
.to_scalar()
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion scalar: {e}")))?;
|
|
|
|
let raw_val = mean_val as f64;
|
|
let prob = 1.0 / (1.0 + (-raw_val).exp());
|
|
let direction = (2.0 * prob - 1.0).clamp(-1.0, 1.0);
|
|
let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);
|
|
|
|
let latency_us = start.elapsed().as_micros() as u64;
|
|
|
|
Ok(EnsemblePrediction {
|
|
model_name: "Diffusion".to_owned(),
|
|
direction,
|
|
confidence,
|
|
metadata: PredictionMeta {
|
|
latency_us,
|
|
quantiles: None,
|
|
attention_weights: None,
|
|
q_values: None,
|
|
},
|
|
})
|
|
}
|
|
|
|
fn predict_raw(&self, features: &FeatureVector) -> MLResult<RawPrediction> {
|
|
let padded = self.pad_features(&features.values);
|
|
let input = Tensor::from_vec(padded, (1, self.data_dim), &self.device)
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion input tensor: {e}")))?;
|
|
|
|
// Timestep t=1 (minimal noise level for feature processing)
|
|
let t = Tensor::from_vec(vec![1_u32], (1,), &self.device)
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion timestep tensor: {e}")))?;
|
|
|
|
let model = self
|
|
.model
|
|
.lock()
|
|
.map_err(|e| MLError::LockError(format!("Diffusion lock poisoned: {e}")))?;
|
|
let output = model.forward(&input, &t)?;
|
|
|
|
// Mean of denoiser output -> scalar tensor (stays on GPU)
|
|
let mean_tensor = output
|
|
.mean_all()
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion mean: {e}")))?;
|
|
|
|
// One sync for confidence (acceptable)
|
|
let raw_f32: f32 = mean_tensor
|
|
.to_scalar()
|
|
.map_err(|e| MLError::ModelError(format!("Diffusion confidence calc: {e}")))?;
|
|
let prob = 1.0 / (1.0 + (-raw_f32 as f64).exp());
|
|
let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);
|
|
|
|
Ok(RawPrediction {
|
|
direction_scalar: 0.0,
|
|
confidence,
|
|
tensor: Some(mean_tensor),
|
|
})
|
|
}
|
|
|
|
fn is_ready(&self) -> bool {
|
|
self.model.lock().is_ok()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::diffusion::config::NoiseSchedule;
|
|
|
|
fn test_config() -> DiffusionConfig {
|
|
DiffusionConfig {
|
|
feature_dim: 51,
|
|
seq_len: 1,
|
|
hidden_dim: 64,
|
|
num_layers: 2,
|
|
time_embed_dim: 32,
|
|
num_timesteps: 100,
|
|
sampling_steps: 5,
|
|
schedule: NoiseSchedule::Cosine,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_diffusion_adapter_creation() {
|
|
let adapter = DiffusionInferenceAdapter::new(test_config());
|
|
assert!(
|
|
adapter.is_ok(),
|
|
"Diffusion creation failed: {:?}",
|
|
adapter.err()
|
|
);
|
|
let adapter = adapter.unwrap();
|
|
assert_eq!(adapter.model_name(), "Diffusion");
|
|
assert!(adapter.is_ready());
|
|
}
|
|
|
|
#[test]
|
|
fn test_diffusion_adapter_predict_direction_range() {
|
|
let adapter = DiffusionInferenceAdapter::new(test_config()).unwrap();
|
|
let fv = FeatureVector {
|
|
values: vec![0.1; 51],
|
|
timestamp: 1700000000_000_000,
|
|
};
|
|
let pred = adapter.predict(&fv).unwrap();
|
|
assert!(
|
|
pred.direction >= -1.0 && pred.direction <= 1.0,
|
|
"direction {} out of [-1,1]",
|
|
pred.direction
|
|
);
|
|
assert!(
|
|
pred.confidence >= 0.0 && pred.confidence <= 1.0,
|
|
"confidence {} out of [0,1]",
|
|
pred.confidence
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_diffusion_adapter_deterministic() {
|
|
let adapter = DiffusionInferenceAdapter::new(test_config()).unwrap();
|
|
let fv = FeatureVector {
|
|
values: vec![0.1; 51],
|
|
timestamp: 1700000000_000_000,
|
|
};
|
|
let pred1 = adapter.predict(&fv).unwrap();
|
|
let pred2 = adapter.predict(&fv).unwrap();
|
|
assert_eq!(
|
|
pred1.direction, pred2.direction,
|
|
"Deterministic predictions should have same direction"
|
|
);
|
|
}
|
|
}
|