Skip to content

Commit 05c1919

Browse files
dhunstackdhunstack
committed
[feat] Add Cpp scripts for inference
Add Cpp scripts for running inference on ONNX exported Demucs Signed-off-by: Anmol Mishra <anmolmishra1997@gmail.com>
1 parent 83a2217 commit 05c1919

14 files changed

Lines changed: 1173 additions & 2 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,4 @@ Session.vim
1717
*.onnx
1818
*.ort
1919
*.config
20+
/cppscripts/build

‎.gitmodules‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
[submodule "cppscripts/dependencies/eigen"]
2+
path = cppscripts/dependencies/eigen
3+
url = https://gitlab.com/libeigen/eigen.git
4+
[submodule "cppscripts/dependencies/libnyquist"]
5+
path = cppscripts/dependencies/libnyquist
6+
url = https://github.com/ddiakopoulos/libnyquist.git

‎cppscripts/Makefile‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
default: cli
2+
3+
cli:
4+
cmake -S src_cli -B build/build-cli -DCMAKE_BUILD_TYPE=Release
5+
cmake --build build/build-cli -- -j16
6+
7+
cli-debug:
8+
cmake -S src_cli -B build/build-cli -DCMAKE_BUILD_TYPE=Debug
9+
cmake --build build/build-cli -- -j16
10+
11+
clean-all:
12+
rm -rf build
13+
14+
clean-cli:
15+
rm -rf build/build-cli

‎cppscripts/dependencies/eigen‎

Submodule eigen added at 8e60d41

‎cppscripts/dependencies/libnyquist‎

Submodule libnyquist added at 767efd9

‎cppscripts/src/demucs.hpp‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
#ifndef MODEL_HPP
2+
#define MODEL_HPP
3+
4+
#include "dsp.hpp"
5+
#include "tensor.hpp"
6+
#include <Eigen/Dense>
7+
#include <array>
8+
#include <functional>
9+
#include <iostream>
10+
#include <string>
11+
#include <vector>
12+
#include <onnxruntime/onnxruntime_cxx_api.h>
13+
14+
namespace demucsonnx
15+
{
16+
extern Ort::AllocatorWithDefaultOptions allocator;
17+
extern Ort::RunOptions run_options;
18+
19+
// Define a type for your callback function
20+
using ProgressCallback = std::function<void(float, const std::string &)>;
21+
22+
const int FREQ_BRANCH_LEN = 336;
23+
const int TIME_BRANCH_LEN_IN = 343980;
24+
25+
struct demucs_model {
26+
std::unique_ptr<Ort::Session> sess; // Smart pointer to allow "empty" state
27+
int nb_sources = 0;
28+
Ort::Env env{ORT_LOGGING_LEVEL_ERROR, "demucs_onnx"}; // Persistent environment
29+
std::vector<std::string> input_names; // Persistent input names
30+
std::vector<std::string> output_names; // Persistent output names
31+
32+
std::vector<const char*> input_names_ptrs;
33+
std::vector<const char*> output_names_ptrs;
34+
35+
// Constructor (optionally initialize here if needed)
36+
demucs_model() = default;
37+
};
38+
39+
bool load_model(const char *model_data,
40+
int n_bytes,
41+
struct demucs_model &model,
42+
Ort::SessionOptions &session_options);
43+
44+
bool load_model(const std::vector<char> &model_data,
45+
struct demucs_model &model,
46+
Ort::SessionOptions &session_options);
47+
48+
struct demucs_segment_buffers
49+
{
50+
int segment_samples;
51+
int le;
52+
int pad;
53+
int pad_end;
54+
int padded_segment_samples;
55+
int nb_stft_frames;
56+
int nb_stft_bins;
57+
58+
Eigen::Tensor3dXf targets_out;
59+
Eigen::MatrixXf padded_mix;
60+
Eigen::Tensor3dXcf z;
61+
62+
std::vector<int64_t> x_onnx_in_shape;
63+
std::vector<int64_t> xt_onnx_in_shape;
64+
65+
std::vector<int64_t> x_onnx_out_shape;
66+
std::vector<int64_t> xt_onnx_out_shape;
67+
68+
std::vector<Ort::Value> input_tensors;
69+
std::vector<Ort::Value> output_tensors;
70+
71+
// constructor for demucs_segment_buffers that takes int parameters
72+
73+
// let's do pesky precomputing of the signal repadding to 1/4 hop
74+
// for time and frequency alignment
75+
demucs_segment_buffers(int nb_channels, int segment_samples, int nb_sources)
76+
: segment_samples(segment_samples),
77+
le(int(std::ceil((float)segment_samples / (float)FFT_HOP_SIZE))),
78+
pad(std::floor((float)FFT_HOP_SIZE / 2.0f) * 3),
79+
pad_end(pad + le * FFT_HOP_SIZE - segment_samples),
80+
padded_segment_samples(segment_samples + pad + pad_end),
81+
nb_stft_frames(segment_samples / demucsonnx::FFT_HOP_SIZE + 1),
82+
nb_stft_bins(demucsonnx::FFT_WINDOW_SIZE / 2 + 1),
83+
targets_out(nb_sources, nb_channels, segment_samples),
84+
padded_mix(nb_channels, padded_segment_samples),
85+
z(nb_channels, nb_stft_bins, nb_stft_frames+4),
86+
// complex-as-channels implies 2*nb_channels for real+imag
87+
x_onnx_in_shape({1, 2 * nb_channels, nb_stft_bins - 1, nb_stft_frames}),
88+
xt_onnx_in_shape({1, nb_channels, segment_samples}),
89+
x_onnx_out_shape({1, nb_sources, 2 * nb_channels, nb_stft_bins - 1, nb_stft_frames}),
90+
xt_onnx_out_shape({1, nb_sources, nb_channels, segment_samples})
91+
{
92+
// precompute the input tensors
93+
// inputs in form (xt, x)
94+
input_tensors.push_back(Ort::Value::CreateTensor<float>(
95+
demucsonnx::allocator,
96+
xt_onnx_in_shape.data(),
97+
xt_onnx_in_shape.size()));
98+
99+
// input_tensors.push_back(Ort::Value::CreateTensor<float>(
100+
// demucsonnx::allocator,
101+
// x_onnx_in_shape.data(),
102+
// x_onnx_in_shape.size()));
103+
104+
// precompute the output tensors
105+
// outputs in form (x_out, xt_out)
106+
// output_tensors.push_back(Ort::Value::CreateTensor<float>(
107+
// demucsonnx::allocator,
108+
// x_onnx_out_shape.data(),
109+
// x_onnx_out_shape.size()));
110+
111+
output_tensors.push_back(Ort::Value::CreateTensor<float>(
112+
demucsonnx::allocator,
113+
xt_onnx_out_shape.data(),
114+
xt_onnx_out_shape.size()));
115+
};
116+
};
117+
118+
const float SEGMENT_LEN_SECS = 7.8; // 8 seconds, the demucs chunk size
119+
const float SEGMENT_OVERLAP_SECS = 0.25; // 0.25 overlap
120+
const float MAX_SHIFT_SECS = 0.5; // max shift
121+
const float OVERLAP = 0.25; // overlap between segments
122+
const float TRANSITION_POWER = 1.0; // transition between segments
123+
124+
Eigen::Tensor3dXf demucs_inference(struct demucs_model &model,
125+
const Eigen::MatrixXf &audio,
126+
ProgressCallback cb);
127+
128+
void model_inference(struct demucs_model &model,
129+
struct demucsonnx::demucs_segment_buffers &buffers);
130+
} // namespace demucsonnx
131+
132+
#endif // MODEL_HPP

‎cppscripts/src/dsp.cpp‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
#include "dsp.hpp"
2+
#include <algorithm>
3+
#include <cmath>
4+
#include <cstdlib>
5+
#include <iostream>
6+
#include <memory>
7+
#include <string>
8+
#include <unsupported/Eigen/FFT>
9+
#include <vector>
10+
11+
// forward declaration of inner stft
12+
void stft_inner(struct demucsonnx::stft_buffers &stft_buf,
13+
Eigen::FFT<float> &cfg);
14+
15+
void istft_inner(struct demucsonnx::stft_buffers &stft_buf,
16+
Eigen::FFT<float> &cfg);
17+
18+
// reflect padding
19+
void pad_signal(struct demucsonnx::stft_buffers &stft_buf)
20+
{
21+
// copy from stft_buf.padded_waveform_mono_in+pad into stft_buf.pad_start,
22+
// stft_buf.pad_end
23+
std::copy_n(stft_buf.padded_waveform_mono_in.begin() + stft_buf.pad,
24+
stft_buf.pad, stft_buf.pad_start.begin());
25+
std::copy_n(stft_buf.padded_waveform_mono_in.end() - 2 * stft_buf.pad,
26+
stft_buf.pad, stft_buf.pad_end.begin());
27+
28+
std::reverse(stft_buf.pad_start.begin(), stft_buf.pad_start.end());
29+
std::reverse(stft_buf.pad_end.begin(), stft_buf.pad_end.end());
30+
31+
// copy stft_buf.pad_start into stft_buf.padded_waveform_mono_in
32+
std::copy_n(stft_buf.pad_start.begin(), stft_buf.pad,
33+
stft_buf.padded_waveform_mono_in.begin());
34+
35+
// copy stft_buf.pad_end into stft_buf.padded_waveform_mono_in
36+
std::copy_n(stft_buf.pad_end.begin(), stft_buf.pad,
37+
stft_buf.padded_waveform_mono_in.end() - stft_buf.pad);
38+
}
39+
40+
Eigen::FFT<float> get_fft_cfg()
41+
{
42+
Eigen::FFT<float> cfg;
43+
44+
cfg.SetFlag(Eigen::FFT<float>::Speedy);
45+
// cfg.SetFlag(Eigen::FFT<float>::HalfSpectrum);
46+
// cfg.SetFlag(Eigen::FFT<float>::Unscaled);
47+
48+
return cfg;
49+
}
50+
51+
void demucsonnx::stft(
52+
struct stft_buffers &stft_buf,
53+
const Eigen::MatrixXf &waveform,
54+
Eigen::Tensor3dXcf &spec)
55+
{
56+
// get the fft config
57+
Eigen::FFT<float> cfg = get_fft_cfg();
58+
59+
/*****************************************/
60+
/* operate on each channel sequentially */
61+
/*****************************************/
62+
63+
for (int channel = 0; channel < 2; ++channel)
64+
{
65+
Eigen::VectorXf row_vec = waveform.row(channel);
66+
67+
std::copy_n(row_vec.data(), row_vec.size(),
68+
stft_buf.padded_waveform_mono_in.begin() + stft_buf.pad);
69+
70+
// apply padding equivalent to center padding with center=True
71+
// in torch.stft:
72+
// https://pytorch.org/docs/stable/generated/torch.stft.html
73+
74+
// reflect pads stft_buf.padded_waveform_mono in-place
75+
pad_signal(stft_buf);
76+
77+
// does forward fft on stft_buf.padded_waveform_mono, stores spectrum in
78+
// complex_spec_mono
79+
stft_inner(stft_buf, cfg);
80+
81+
for (int i = 0; i < stft_buf.nb_bins; ++i)
82+
{
83+
for (int j = 0; j < stft_buf.nb_frames; ++j)
84+
{
85+
spec(channel, i, j) = stft_buf.complex_spec_mono[j][i];
86+
}
87+
}
88+
}
89+
}
90+
91+
void demucsonnx::istft(
92+
struct stft_buffers &stft_buf,
93+
const Eigen::Tensor3dXcf &spec,
94+
Eigen::MatrixXf &waveform)
95+
{
96+
// get the fft config
97+
Eigen::FFT<float> cfg = get_fft_cfg();
98+
99+
/*****************************************/
100+
/* operate on each channel sequentially */
101+
/*****************************************/
102+
103+
for (int channel = 0; channel < 2; ++channel)
104+
{
105+
// Populate the nested vectors
106+
for (int i = 0; i < stft_buf.nb_bins; ++i)
107+
{
108+
for (int j = 0; j < stft_buf.nb_frames; ++j)
109+
{
110+
stft_buf.complex_spec_mono[j][i] = spec(channel, i, j);
111+
}
112+
}
113+
114+
// does inverse fft on stft_buf.complex_spec_mono, stores waveform in
115+
// padded_waveform_mono
116+
istft_inner(stft_buf, cfg);
117+
118+
// copies waveform_mono into stft_buf.waveform past first pad samples
119+
waveform.row(channel) = Eigen::Map<Eigen::MatrixXf>(
120+
stft_buf.padded_waveform_mono_out.data() + stft_buf.pad, 1,
121+
stft_buf.padded_waveform_mono_out.size() - FFT_WINDOW_SIZE);
122+
}
123+
}
124+
125+
void stft_inner(struct demucsonnx::stft_buffers &stft_buf,
126+
Eigen::FFT<float> &cfg)
127+
{
128+
int frame_idx = 0;
129+
130+
// Loop over the waveform with a stride of hop_size
131+
for (std::size_t start = 0;
132+
start <=
133+
stft_buf.padded_waveform_mono_in.size() - demucsonnx::FFT_WINDOW_SIZE;
134+
start += demucsonnx::FFT_HOP_SIZE)
135+
{
136+
// Apply window and run FFT
137+
for (int i = 0; i < demucsonnx::FFT_WINDOW_SIZE; ++i)
138+
{
139+
stft_buf.windowed_waveform_mono[i] =
140+
stft_buf.padded_waveform_mono_in[start + i] *
141+
stft_buf.window[i];
142+
}
143+
cfg.fwd(stft_buf.complex_spec_mono[frame_idx],
144+
stft_buf.windowed_waveform_mono);
145+
// now scale stft_buf.complex_spec_mono[frame_idx] by 1.0f /
146+
// sqrt(float(FFT_WINDOW_SIZE)))
147+
148+
for (int i = 0; i < demucsonnx::FFT_WINDOW_SIZE / 2 + 1; ++i)
149+
{
150+
stft_buf.complex_spec_mono[frame_idx][i] *=
151+
1.0f / sqrt(float(demucsonnx::FFT_WINDOW_SIZE));
152+
}
153+
frame_idx++;
154+
}
155+
}
156+
157+
void istft_inner(struct demucsonnx::stft_buffers &stft_buf,
158+
Eigen::FFT<float> &cfg)
159+
{
160+
// clear padded_waveform_mono
161+
std::fill(stft_buf.padded_waveform_mono_out.begin(),
162+
stft_buf.padded_waveform_mono_out.end(), 0.0f);
163+
164+
// Loop over the input with a stride of (hop_size)
165+
for (int start = 0; start < stft_buf.nb_frames * demucsonnx::FFT_HOP_SIZE;
166+
start += demucsonnx::FFT_HOP_SIZE)
167+
{
168+
int frame_idx = start / demucsonnx::FFT_HOP_SIZE;
169+
// undo sqrt(nfft) scaling
170+
for (int i = 0; i < demucsonnx::FFT_WINDOW_SIZE / 2 + 1; ++i)
171+
{
172+
stft_buf.complex_spec_mono[frame_idx][i] *=
173+
sqrt(float(demucsonnx::FFT_WINDOW_SIZE));
174+
}
175+
// Run iFFT
176+
cfg.inv(stft_buf.windowed_waveform_mono,
177+
stft_buf.complex_spec_mono[frame_idx]);
178+
179+
// Apply window and add to output
180+
for (int i = 0; i < demucsonnx::FFT_WINDOW_SIZE; ++i)
181+
{
182+
// x[start+i] is the sum of squared window values
183+
// https://github.com/librosa/librosa/blob/main/librosa/core/spectrum.py#L613
184+
// 1e-8f is a small number to avoid division by zero
185+
stft_buf.padded_waveform_mono_out[start + i] +=
186+
stft_buf.windowed_waveform_mono[i] * stft_buf.window[i] * 1.0f /
187+
float(demucsonnx::FFT_WINDOW_SIZE) /
188+
(stft_buf.normalized_window[start + i] + 1e-8f);
189+
}
190+
}
191+
}

0 commit comments

Comments
 (0)