rtx-nn / rtx-multimodal: cargo fmt reformatting

Pure formatting changes across rtx-nn (conv_transpose1d, conv/mod, rnn/lstm)
and rtx-multimodal (audio/generation, audio/source_separation): multi-line
braces, trailing commas, import ordering. No logic changes.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-04-27 18:37:07 -07:00
co-authored by Claude Opus 4.7
parent ed4e5e4b85
commit 6d59251c51
8 changed files with 163 additions and 81 deletions
@@ -74,7 +74,9 @@ impl ConvTranspose1dConfig {
return Err(NNError::InvalidParameter("channels must be > 0".into())); return Err(NNError::InvalidParameter("channels must be > 0".into()));
} }
if self.kernel_size == 0 || self.stride == 0 || self.dilation == 0 { if self.kernel_size == 0 || self.stride == 0 || self.dilation == 0 {
return Err(NNError::InvalidParameter("kernel/stride/dilation must be > 0".into())); return Err(NNError::InvalidParameter(
"kernel/stride/dilation must be > 0".into(),
));
} }
if self.groups == 0 { if self.groups == 0 {
return Err(NNError::InvalidParameter("groups must be > 0".into())); return Err(NNError::InvalidParameter("groups must be > 0".into()));
@@ -101,8 +103,7 @@ impl ConvTranspose1dConfig {
/// Calculate output length for a given input length. /// Calculate output length for a given input length.
pub fn output_length(&self, input_length: usize) -> usize { pub fn output_length(&self, input_length: usize) -> usize {
(input_length - 1) * self.stride (input_length - 1) * self.stride - 2 * self.padding
- 2 * self.padding
+ self.dilation * (self.kernel_size - 1) + self.dilation * (self.kernel_size - 1)
+ self.output_padding + self.output_padding
+ 1 + 1
@@ -147,7 +148,11 @@ impl ConvTranspose1d {
let out_channels_per_group = config.out_channels / config.groups; let out_channels_per_group = config.out_channels / config.groups;
// Weight shape: [in_channels, out_channels/groups, kernel_size] // Weight shape: [in_channels, out_channels/groups, kernel_size]
let weight_shape = vec![config.in_channels, out_channels_per_group, config.kernel_size]; let weight_shape = vec![
config.in_channels,
out_channels_per_group,
config.kernel_size,
];
let mut weight = Tensor::zeros(weight_shape, device)?; let mut weight = Tensor::zeros(weight_shape, device)?;
let fan_in = config.in_channels * config.kernel_size / config.groups; let fan_in = config.in_channels * config.kernel_size / config.groups;
@@ -171,12 +176,24 @@ impl ConvTranspose1d {
}) })
} }
pub fn in_channels(&self) -> usize { self.config.in_channels } pub fn in_channels(&self) -> usize {
pub fn out_channels(&self) -> usize { self.config.out_channels } self.config.in_channels
pub fn kernel_size(&self) -> usize { self.config.kernel_size } }
pub fn stride(&self) -> usize { self.config.stride } pub fn out_channels(&self) -> usize {
pub fn weight(&self) -> &Tensor { &self.weight } self.config.out_channels
pub fn bias(&self) -> Option<&Tensor> { self.bias.as_ref() } }
pub fn kernel_size(&self) -> usize {
self.config.kernel_size
}
pub fn stride(&self) -> usize {
self.config.stride
}
pub fn weight(&self) -> &Tensor {
&self.weight
}
pub fn bias(&self) -> Option<&Tensor> {
self.bias.as_ref()
}
/// Perform the transposed convolution. /// Perform the transposed convolution.
/// ///
@@ -256,7 +273,8 @@ impl crate::layers::Module for ConvTranspose1d {
if shape.dims()[1] != self.config.in_channels { if shape.dims()[1] != self.config.in_channels {
return Err(NNError::InvalidParameter(format!( return Err(NNError::InvalidParameter(format!(
"Expected {} input channels, got {}", "Expected {} input channels, got {}",
self.config.in_channels, shape.dims()[1] self.config.in_channels,
shape.dims()[1]
))); )));
} }
self.conv_transpose1d_forward(input) self.conv_transpose1d_forward(input)
@@ -287,7 +305,10 @@ impl crate::layers::Module for ConvTranspose1d {
} }
fn to_device(&mut self, device: &Device) -> Result<()> { fn to_device(&mut self, device: &Device) -> Result<()> {
self.weight = self.weight.to_device(device).map_err(|e| NNError::Tensor(e))?; self.weight = self
.weight
.to_device(device)
.map_err(|e| NNError::Tensor(e))?;
if let Some(ref b) = self.bias { if let Some(ref b) = self.bias {
self.bias = Some(b.to_device(device).map_err(|e| NNError::Tensor(e))?); self.bias = Some(b.to_device(device).map_err(|e| NNError::Tensor(e))?);
} }
@@ -321,7 +342,9 @@ mod tests {
#[test] #[test]
fn config_output_length_with_padding() { fn config_output_length_with_padding() {
let cfg = ConvTranspose1dConfig::new(1, 1, 3).with_stride(2).with_padding(1); let cfg = ConvTranspose1dConfig::new(1, 1, 3)
.with_stride(2)
.with_padding(1);
// (10-1)*2 - 2 + 1*(3-1) + 0 + 1 = 19 // (10-1)*2 - 2 + 1*(3-1) + 0 + 1 = 19
assert_eq!(cfg.output_length(10), 19); assert_eq!(cfg.output_length(10), 19);
} }
@@ -333,15 +356,17 @@ mod tests {
let input = Tensor::randn(&[2, 4, 10], &device).unwrap(); let input = Tensor::randn(&[2, 4, 10], &device).unwrap();
let output = layer.forward(&input).unwrap(); let output = layer.forward(&input).unwrap();
let dims = output.shape().dims(); let dims = output.shape().dims();
assert_eq!(dims[0], 2); // batch assert_eq!(dims[0], 2); // batch
assert_eq!(dims[1], 8); // out_channels assert_eq!(dims[1], 8); // out_channels
assert_eq!(dims[2], 12); // (10-1)*1 + 3-1 + 1 = 12 assert_eq!(dims[2], 12); // (10-1)*1 + 3-1 + 1 = 12
} }
#[test] #[test]
fn forward_shape_stride2() { fn forward_shape_stride2() {
let device = Device::Cpu; let device = Device::Cpu;
let cfg = ConvTranspose1dConfig::new(4, 8, 4).with_stride(2).with_padding(1); let cfg = ConvTranspose1dConfig::new(4, 8, 4)
.with_stride(2)
.with_padding(1);
let layer = ConvTranspose1d::from_config(cfg, &device).unwrap(); let layer = ConvTranspose1d::from_config(cfg, &device).unwrap();
let input = Tensor::randn(&[1, 4, 16], &device).unwrap(); let input = Tensor::randn(&[1, 4, 16], &device).unwrap();
let output = layer.forward(&input).unwrap(); let output = layer.forward(&input).unwrap();
@@ -369,14 +394,13 @@ mod tests {
#[test] #[test]
fn no_bias_has_fewer_params() { fn no_bias_has_fewer_params() {
let device = Device::Cpu; let device = Device::Cpu;
let with_bias = ConvTranspose1d::from_config( let with_bias =
ConvTranspose1dConfig::new(4, 8, 3), ConvTranspose1d::from_config(ConvTranspose1dConfig::new(4, 8, 3), &device).unwrap();
&device,
).unwrap();
let no_bias = ConvTranspose1d::from_config( let no_bias = ConvTranspose1d::from_config(
ConvTranspose1dConfig::new(4, 8, 3).with_bias(false), ConvTranspose1dConfig::new(4, 8, 3).with_bias(false),
&device, &device,
).unwrap(); )
.unwrap();
assert_eq!(with_bias.parameters().len(), 2); assert_eq!(with_bias.parameters().len(), 2);
assert_eq!(no_bias.parameters().len(), 1); assert_eq!(no_bias.parameters().len(), 1);
} }
@@ -397,7 +421,11 @@ mod tests {
let output_data = output.to_cpu().unwrap(); let output_data = output.to_cpu().unwrap();
for (i, &v) in input_data.iter().enumerate() { for (i, &v) in input_data.iter().enumerate() {
assert!((output_data[i] - v).abs() < 1e-6, "mismatch at {i}: {v} vs {}", output_data[i]); assert!(
(output_data[i] - v).abs() < 1e-6,
"mismatch at {i}: {v} vs {}",
output_data[i]
);
} }
} }
} }
+1 -1
View File
@@ -13,9 +13,9 @@ pub mod conv_transpose1d;
// Re-export main types // Re-export main types
pub use common::{PaddingMode, im2col, padding}; pub use common::{PaddingMode, im2col, padding};
pub use conv_transpose1d::{ConvTranspose1d, ConvTranspose1dConfig};
pub use conv1d::{Conv1d, Conv1dConfig, Conv1dPadding}; pub use conv1d::{Conv1d, Conv1dConfig, Conv1dPadding};
pub use conv2d::{Conv2d, Conv2dConfig}; pub use conv2d::{Conv2d, Conv2dConfig};
pub use conv_transpose1d::{ConvTranspose1d, ConvTranspose1dConfig};
// Re-export for backward compatibility // Re-export for backward compatibility
pub use common::PaddingMode as ConvPaddingMode; pub use common::PaddingMode as ConvPaddingMode;
+21 -22
View File
@@ -162,12 +162,7 @@ impl LSTM {
/// ///
/// Input: `[batch, seq_len, input_size]` as flat f32 vec. /// Input: `[batch, seq_len, input_size]` as flat f32 vec.
/// Returns `LSTMOutput` with output and final states. /// Returns `LSTMOutput` with output and final states.
pub fn forward_cpu( pub fn forward_cpu(&self, input: &[f32], batch: usize, seq_len: usize) -> Result<LSTMOutput> {
&self,
input: &[f32],
batch: usize,
seq_len: usize,
) -> Result<LSTMOutput> {
let h = self.config.hidden_size; let h = self.config.hidden_size;
let dirs = self.config.num_directions(); let dirs = self.config.num_directions();
let num_layers = self.config.num_layers; let num_layers = self.config.num_layers;
@@ -264,15 +259,17 @@ impl LSTM {
pub fn forward_tensor(&self, input: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { pub fn forward_tensor(&self, input: &Tensor) -> Result<(Tensor, Tensor, Tensor)> {
let dims = input.shape().dims(); let dims = input.shape().dims();
if dims.len() != 3 { if dims.len() != 3 {
return Err(NNError::InvalidParameter( return Err(NNError::InvalidParameter(format!(
format!("LSTM expects 3D input [batch, seq, features], got {}D", dims.len()) "LSTM expects 3D input [batch, seq, features], got {}D",
)); dims.len()
)));
} }
let (batch, seq_len, feat) = (dims[0], dims[1], dims[2]); let (batch, seq_len, feat) = (dims[0], dims[1], dims[2]);
if feat != self.config.input_size { if feat != self.config.input_size {
return Err(NNError::InvalidParameter( return Err(NNError::InvalidParameter(format!(
format!("Expected input_size {}, got {}", self.config.input_size, feat) "Expected input_size {}, got {}",
)); self.config.input_size, feat
)));
} }
let input_data = input.to_cpu().map_err(|e| NNError::Tensor(e))?; let input_data = input.to_cpu().map_err(|e| NNError::Tensor(e))?;
@@ -288,7 +285,9 @@ impl LSTM {
Ok((output, h_n, c_n)) Ok((output, h_n, c_n))
} }
pub fn config(&self) -> &LSTMConfig { &self.config } pub fn config(&self) -> &LSTMConfig {
&self.config
}
} }
#[inline] #[inline]
@@ -343,7 +342,7 @@ mod tests {
let result = lstm.forward_cpu(&input, batch, seq_len).unwrap(); let result = lstm.forward_cpu(&input, batch, seq_len).unwrap();
assert_eq!(result.output_shape, [2, 5, 32]); // 16 * 2 directions assert_eq!(result.output_shape, [2, 5, 32]); // 16 * 2 directions
assert_eq!(result.state_shape, [2, 2, 16]); // 1 layer * 2 directions assert_eq!(result.state_shape, [2, 2, 16]); // 1 layer * 2 directions
} }
#[test] #[test]
@@ -362,7 +361,9 @@ mod tests {
#[test] #[test]
fn forward_multilayer_bidirectional_shape() { fn forward_multilayer_bidirectional_shape() {
let cfg = LSTMConfig::new(8, 16).with_num_layers(2).with_bidirectional(true); let cfg = LSTMConfig::new(8, 16)
.with_num_layers(2)
.with_bidirectional(true);
let lstm = LSTM::new(cfg, &Device::Cpu).unwrap(); let lstm = LSTM::new(cfg, &Device::Cpu).unwrap();
let batch = 3; let batch = 3;
@@ -371,7 +372,7 @@ mod tests {
let result = lstm.forward_cpu(&input, batch, seq_len).unwrap(); let result = lstm.forward_cpu(&input, batch, seq_len).unwrap();
assert_eq!(result.output_shape, [3, 7, 32]); // 16 * 2 assert_eq!(result.output_shape, [3, 7, 32]); // 16 * 2
assert_eq!(result.state_shape, [4, 3, 16]); // 2 layers * 2 directions assert_eq!(result.state_shape, [4, 3, 16]); // 2 layers * 2 directions
} }
#[test] #[test]
@@ -392,11 +393,8 @@ mod tests {
let cfg = LSTMConfig::new(4, 8); let cfg = LSTMConfig::new(4, 8);
let lstm = LSTM::new(cfg, &Device::Cpu).unwrap(); let lstm = LSTM::new(cfg, &Device::Cpu).unwrap();
let input = Tensor::from_data( let input =
vec![0.1f32; 2 * 5 * 4], Tensor::from_data(vec![0.1f32; 2 * 5 * 4], vec![2, 5, 4], &Device::Cpu).unwrap();
vec![2, 5, 4],
&Device::Cpu,
).unwrap();
let (output, h_n, c_n) = lstm.forward_tensor(&input).unwrap(); let (output, h_n, c_n) = lstm.forward_tensor(&input).unwrap();
assert_eq!(output.shape().dims(), &[2, 5, 8]); assert_eq!(output.shape().dims(), &[2, 5, 8]);
@@ -420,7 +418,8 @@ mod tests {
vec![0.0f32; 2 * 5 * 3], // 3 != 4 vec![0.0f32; 2 * 5 * 3], // 3 != 4
vec![2, 5, 3], vec![2, 5, 3],
&Device::Cpu, &Device::Cpu,
).unwrap(); )
.unwrap();
assert!(lstm.forward_tensor(&input).is_err()); assert!(lstm.forward_tensor(&input).is_err());
} }
} }
@@ -7,7 +7,7 @@
pub mod stable_audio; pub mod stable_audio;
#[cfg(feature = "generation")] #[cfg(feature = "generation")]
pub use stable_audio::{StableAudioConfig, StableAudioModel, StableAudioError}; pub use stable_audio::{StableAudioConfig, StableAudioError, StableAudioModel};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
@@ -67,8 +67,9 @@ impl StableAudioModel {
"Loading Stable Audio Open model" "Loading Stable Audio Open model"
); );
let text_encoder = OnnxSession::from_file(&config.text_encoder_path, onnx_config.clone()) let text_encoder =
.map_err(|e| StableAudioError::ModelLoad(format!("text encoder: {e}")))?; OnnxSession::from_file(&config.text_encoder_path, onnx_config.clone())
.map_err(|e| StableAudioError::ModelLoad(format!("text encoder: {e}")))?;
let diffusion = OnnxSession::from_file(&config.diffusion_path, onnx_config.clone()) let diffusion = OnnxSession::from_file(&config.diffusion_path, onnx_config.clone())
.map_err(|e| StableAudioError::ModelLoad(format!("diffusion: {e}")))?; .map_err(|e| StableAudioError::ModelLoad(format!("diffusion: {e}")))?;
@@ -78,11 +79,19 @@ impl StableAudioModel {
info!("Stable Audio Open loaded successfully"); info!("Stable Audio Open loaded successfully");
Ok(Self { text_encoder, diffusion, vae_decoder, config }) Ok(Self {
text_encoder,
diffusion,
vae_decoder,
config,
})
} }
/// Generate audio from a text prompt. /// Generate audio from a text prompt.
pub fn generate(&mut self, params: &GenerationParams) -> Result<GenerationOutput, StableAudioError> { pub fn generate(
&mut self,
params: &GenerationParams,
) -> Result<GenerationOutput, StableAudioError> {
info!(prompt = %params.prompt, duration = params.duration_secs, steps = params.steps, "Generating audio"); info!(prompt = %params.prompt, duration = params.duration_secs, steps = params.steps, "Generating audio");
// Step 1: Encode text prompt // Step 1: Encode text prompt
@@ -90,7 +99,8 @@ impl StableAudioModel {
debug!("Text encoded"); debug!("Text encoded");
// Step 2: Initialize latent noise // Step 2: Initialize latent noise
let latent_frames = (params.duration_secs * self.config.sample_rate as f64 / 512.0) as usize; let latent_frames =
(params.duration_secs * self.config.sample_rate as f64 / 512.0) as usize;
let latent = self.initialize_latent(latent_frames, params.seed)?; let latent = self.initialize_latent(latent_frames, params.seed)?;
debug!(latent_frames, "Latent initialized"); debug!(latent_frames, "Latent initialized");
@@ -119,7 +129,11 @@ impl StableAudioModel {
Ok(dummy) Ok(dummy)
} }
fn initialize_latent(&self, frames: usize, seed: Option<u64>) -> Result<Tensor, StableAudioError> { fn initialize_latent(
&self,
frames: usize,
seed: Option<u64>,
) -> Result<Tensor, StableAudioError> {
// Initialize with random noise (or seeded noise for reproducibility) // Initialize with random noise (or seeded noise for reproducibility)
let latent = Tensor::randn(&[1, self.config.latent_dim, frames], &Device::Cpu) let latent = Tensor::randn(&[1, self.config.latent_dim, frames], &Device::Cpu)
.map_err(|e| StableAudioError::Inference(e.to_string()))?; .map_err(|e| StableAudioError::Inference(e.to_string()))?;
@@ -113,7 +113,11 @@ impl DemucsModel {
/// For audio longer than `segment_length`, the input is split into overlapping /// For audio longer than `segment_length`, the input is split into overlapping
/// segments, each processed independently, then recombined via overlap-add with /// segments, each processed independently, then recombined via overlap-add with
/// a triangular cross-fade window. /// a triangular cross-fade window.
pub fn separate(&mut self, waveform: &[f32], channels: usize) -> Result<Vec<StemOutput>, DemucsError> { pub fn separate(
&mut self,
waveform: &[f32],
channels: usize,
) -> Result<Vec<StemOutput>, DemucsError> {
if waveform.is_empty() { if waveform.is_empty() {
return Err(DemucsError::EmptyInput); return Err(DemucsError::EmptyInput);
} }
@@ -138,7 +142,8 @@ impl DemucsModel {
// 4. Run inference on each segment // 4. Run inference on each segment
let stem_types = StemType::stems_for(self.config.num_stems); let stem_types = StemType::stems_for(self.config.num_stems);
let mut stem_accumulators: Vec<Vec<f32>> = vec![vec![0.0; total_frames * channels]; self.config.num_stems]; let mut stem_accumulators: Vec<Vec<f32>> =
vec![vec![0.0; total_frames * channels]; self.config.num_stems];
let mut weight_accumulator: Vec<f32> = vec![0.0; total_frames]; let mut weight_accumulator: Vec<f32> = vec![0.0; total_frames];
for (start_frame, chunk) in &segments { for (start_frame, chunk) in &segments {
@@ -208,7 +213,12 @@ impl DemucsModel {
} }
/// Segment audio into overlapping chunks. /// Segment audio into overlapping chunks.
fn segment(&self, channel_first: &[f32], channels: usize, total_frames: usize) -> Vec<(usize, Vec<f32>)> { fn segment(
&self,
channel_first: &[f32],
channels: usize,
total_frames: usize,
) -> Vec<(usize, Vec<f32>)> {
let seg_len = self.config.segment_length; let seg_len = self.config.segment_length;
let hop = ((1.0 - self.config.overlap) * seg_len as f32) as usize; let hop = ((1.0 - self.config.overlap) * seg_len as f32) as usize;
let hop = hop.max(1); let hop = hop.max(1);
@@ -272,14 +282,19 @@ impl DemucsModel {
let mut inputs = HashMap::new(); let mut inputs = HashMap::new();
inputs.insert("mix".to_string(), &input_tensor); inputs.insert("mix".to_string(), &input_tensor);
let outputs = self.session.run(inputs) let outputs = self
.session
.run(inputs)
.map_err(|e| DemucsError::Inference(e.to_string()))?; .map_err(|e| DemucsError::Inference(e.to_string()))?;
// Extract the output tensor (first output, whatever its name) // Extract the output tensor (first output, whatever its name)
let output_tensor = outputs.into_values().next() let output_tensor = outputs
.into_values()
.next()
.ok_or_else(|| DemucsError::Inference("no output tensor from ONNX model".into()))?; .ok_or_else(|| DemucsError::Inference("no output tensor from ONNX model".into()))?;
let output_data = output_tensor.to_vec_f32() let output_data = output_tensor
.to_vec_f32()
.map_err(|e| DemucsError::Inference(e.to_string()))?; .map_err(|e| DemucsError::Inference(e.to_string()))?;
Ok(output_data) Ok(output_data)
@@ -292,7 +307,11 @@ fn normalize(samples: &[f32]) -> (Vec<f32>, f32) {
return (vec![], 1.0); return (vec![], 1.0);
} }
let mean_sq: f64 = samples.iter().map(|&s| (s as f64) * (s as f64)).sum::<f64>() / samples.len() as f64; let mean_sq: f64 = samples
.iter()
.map(|&s| (s as f64) * (s as f64))
.sum::<f64>()
/ samples.len() as f64;
let rms = mean_sq.sqrt() as f32; let rms = mean_sq.sqrt() as f32;
let scale = rms.max(1e-8); let scale = rms.max(1e-8);
@@ -19,8 +19,8 @@
use rtx_nn::layers::{ use rtx_nn::layers::{
Module, Module,
conv::conv1d::{Conv1d, Conv1dConfig, Conv1dPadding},
conv::conv_transpose1d::{ConvTranspose1d, ConvTranspose1dConfig}, conv::conv_transpose1d::{ConvTranspose1d, ConvTranspose1dConfig},
conv::conv1d::{Conv1d, Conv1dConfig, Conv1dPadding},
rnn::lstm::{LSTM, LSTMConfig}, rnn::lstm::{LSTM, LSTMConfig},
}; };
use rtx_tensor::{Device, Tensor}; use rtx_tensor::{Device, Tensor};
@@ -152,8 +152,7 @@ impl HtDemucsNative {
let lstm_cfg = LSTMConfig::new(bottleneck_ch, bottleneck_ch) let lstm_cfg = LSTMConfig::new(bottleneck_ch, bottleneck_ch)
.with_num_layers(config.lstm_layers) .with_num_layers(config.lstm_layers)
.with_bidirectional(true); .with_bidirectional(true);
let lstm = LSTM::new(lstm_cfg, device) let lstm = LSTM::new(lstm_cfg, device).map_err(|e| HtDemucsError::Build(e.to_string()))?;
.map_err(|e| HtDemucsError::Build(e.to_string()))?;
// Linear projection after BiLSTM (2*hidden → hidden) // Linear projection after BiLSTM (2*hidden → hidden)
let proj_cfg = Conv1dConfig::new(bottleneck_ch * 2, bottleneck_ch, 1).with_bias(true); let proj_cfg = Conv1dConfig::new(bottleneck_ch * 2, bottleneck_ch, 1).with_bias(true);
@@ -195,10 +194,13 @@ impl HtDemucsNative {
for (d, enc) in self.encoder_convs.iter().enumerate() { for (d, enc) in self.encoder_convs.iter().enumerate() {
skips.push(x.clone()); skips.push(x.clone());
x = enc.forward(&x) x = enc
.forward(&x)
.map_err(|e| HtDemucsError::Forward(format!("encoder[{d}]: {e}")))?; .map_err(|e| HtDemucsError::Forward(format!("encoder[{d}]: {e}")))?;
// ReLU activation // ReLU activation
let data = x.to_cpu().map_err(|e| HtDemucsError::Forward(e.to_string()))?; let data = x
.to_cpu()
.map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let relu_data: Vec<f32> = data.iter().map(|&v| v.max(0.0)).collect(); let relu_data: Vec<f32> = data.iter().map(|&v| v.max(0.0)).collect();
x = Tensor::from_data(relu_data, x.shape().dims().to_vec(), &self.device) x = Tensor::from_data(relu_data, x.shape().dims().to_vec(), &self.device)
.map_err(|e| HtDemucsError::Forward(e.to_string()))?; .map_err(|e| HtDemucsError::Forward(e.to_string()))?;
@@ -209,7 +211,9 @@ impl HtDemucsNative {
let enc_dims = x.shape().dims(); let enc_dims = x.shape().dims();
let (b, c, t) = (enc_dims[0], enc_dims[1], enc_dims[2]); let (b, c, t) = (enc_dims[0], enc_dims[1], enc_dims[2]);
let x_data = x.to_cpu().map_err(|e| HtDemucsError::Forward(e.to_string()))?; let x_data = x
.to_cpu()
.map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let mut transposed = vec![0.0f32; b * t * c]; let mut transposed = vec![0.0f32; b * t * c];
for bi in 0..b { for bi in 0..b {
for ci in 0..c { for ci in 0..c {
@@ -222,19 +226,24 @@ impl HtDemucsNative {
let lstm_input = Tensor::from_data(transposed, vec![b, t, c], &self.device) let lstm_input = Tensor::from_data(transposed, vec![b, t, c], &self.device)
.map_err(|e| HtDemucsError::Forward(e.to_string()))?; .map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let (lstm_out, _, _) = self.lstm.forward_tensor(&lstm_input) let (lstm_out, _, _) = self
.lstm
.forward_tensor(&lstm_input)
.map_err(|e| HtDemucsError::Forward(format!("lstm: {e}")))?; .map_err(|e| HtDemucsError::Forward(format!("lstm: {e}")))?;
// Project BiLSTM output (2*hidden → hidden) and transpose back to [batch, channels, time] // Project BiLSTM output (2*hidden → hidden) and transpose back to [batch, channels, time]
let lo_dims = lstm_out.shape().dims(); let lo_dims = lstm_out.shape().dims();
let lo_data = lstm_out.to_cpu().map_err(|e| HtDemucsError::Forward(e.to_string()))?; let lo_data = lstm_out
.to_cpu()
.map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let lstm_ch = lo_dims[2]; // 2 * hidden_size let lstm_ch = lo_dims[2]; // 2 * hidden_size
let mut back = vec![0.0f32; b * lstm_ch * t]; let mut back = vec![0.0f32; b * lstm_ch * t];
for bi in 0..b { for bi in 0..b {
for ti in 0..t { for ti in 0..t {
for ci in 0..lstm_ch { for ci in 0..lstm_ch {
back[bi * lstm_ch * t + ci * t + ti] = lo_data[bi * t * lstm_ch + ti * lstm_ch + ci]; back[bi * lstm_ch * t + ci * t + ti] =
lo_data[bi * t * lstm_ch + ti * lstm_ch + ci];
} }
} }
} }
@@ -243,7 +252,9 @@ impl HtDemucsNative {
.map_err(|e| HtDemucsError::Forward(e.to_string()))?; .map_err(|e| HtDemucsError::Forward(e.to_string()))?;
// 1x1 projection: [batch, 2*hidden, time] → [batch, hidden, time] // 1x1 projection: [batch, 2*hidden, time] → [batch, hidden, time]
x = self.output_conv.forward(&x) x = self
.output_conv
.forward(&x)
.map_err(|e| HtDemucsError::Forward(format!("projection: {e}")))?; .map_err(|e| HtDemucsError::Forward(format!("projection: {e}")))?;
// Decoder pass with skip connections // Decoder pass with skip connections
@@ -252,8 +263,12 @@ impl HtDemucsNative {
let skip = &skips[skip_idx]; let skip = &skips[skip_idx];
// Concatenate skip connection along channel dimension // Concatenate skip connection along channel dimension
let x_data = x.to_cpu().map_err(|e| HtDemucsError::Forward(e.to_string()))?; let x_data = x
let s_data = skip.to_cpu().map_err(|e| HtDemucsError::Forward(e.to_string()))?; .to_cpu()
.map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let s_data = skip
.to_cpu()
.map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let x_dims = x.shape().dims(); let x_dims = x.shape().dims();
let s_dims = skip.shape().dims(); let s_dims = skip.shape().dims();
@@ -285,13 +300,16 @@ impl HtDemucsNative {
x = Tensor::from_data(cat_data, vec![xb, cat_c, min_t], &self.device) x = Tensor::from_data(cat_data, vec![xb, cat_c, min_t], &self.device)
.map_err(|e| HtDemucsError::Forward(e.to_string()))?; .map_err(|e| HtDemucsError::Forward(e.to_string()))?;
x = dec.forward(&x) x = dec
.forward(&x)
.map_err(|e| HtDemucsError::Forward(format!("decoder[{d}]: {e}")))?; .map_err(|e| HtDemucsError::Forward(format!("decoder[{d}]: {e}")))?;
} }
// Reshape output: [batch, num_sources * audio_channels, samples] // Reshape output: [batch, num_sources * audio_channels, samples]
// → split into [batch, audio_channels, samples] per source // → split into [batch, audio_channels, samples] per source
let out_data = x.to_cpu().map_err(|e| HtDemucsError::Forward(e.to_string()))?; let out_data = x
.to_cpu()
.map_err(|e| HtDemucsError::Forward(e.to_string()))?;
let out_dims = x.shape().dims(); let out_dims = x.shape().dims();
let out_len = out_dims[2].min(samples); // Trim to original length let out_len = out_dims[2].min(samples); // Trim to original length
let num_sources = self.config.num_sources; let num_sources = self.config.num_sources;
@@ -382,11 +400,8 @@ mod tests {
let model = HtDemucsNative::new(cfg, &Device::Cpu).unwrap(); let model = HtDemucsNative::new(cfg, &Device::Cpu).unwrap();
// Input: batch=1, stereo, 4096 samples // Input: batch=1, stereo, 4096 samples
let input = Tensor::from_data( let input =
vec![0.01f32; 1 * 2 * 4096], Tensor::from_data(vec![0.01f32; 1 * 2 * 4096], vec![1, 2, 4096], &Device::Cpu).unwrap();
vec![1, 2, 4096],
&Device::Cpu,
).unwrap();
let stems = model.forward(&input).unwrap(); let stems = model.forward(&input).unwrap();
assert_eq!(stems.len(), 4); assert_eq!(stems.len(), 4);
@@ -31,11 +31,18 @@ impl StemType {
pub fn stems_for(num_stems: usize) -> Vec<StemType> { pub fn stems_for(num_stems: usize) -> Vec<StemType> {
match num_stems { match num_stems {
6 => vec![ 6 => vec![
StemType::Drums, StemType::Bass, StemType::Other, StemType::Drums,
StemType::Vocals, StemType::Guitar, StemType::Piano, StemType::Bass,
StemType::Other,
StemType::Vocals,
StemType::Guitar,
StemType::Piano,
], ],
_ => vec![ _ => vec![
StemType::Drums, StemType::Bass, StemType::Other, StemType::Vocals, StemType::Drums,
StemType::Bass,
StemType::Other,
StemType::Vocals,
], ],
} }
} }