Skip to content

[BREAKING] Improve organization of WaveNet code & internalized functionality #250

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sdatkinson merged 2 commits into from
Apr 16, 2026
Merged

[BREAKING] Improve organization of WaveNet code & internalized functionality #250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8193b51
Refactor WaveNet into NAM/wavenet/ (params, detail, model)
sdatkinson Apr 16, 2026
735c1da
Move slimmable WaveNet sources to wavenet/slimmable.{h,cpp}
sdatkinson Apr 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
811 changes: 0 additions & 811 deletions NAM/wavenet.h
View file
Open in desktop

This file was deleted.

388 changes: 388 additions & 0 deletions NAM/wavenet/detail.h
View file
Open in desktop

Large diffs are not rendered by default.

75 changes: 38 additions & 37 deletions NAM/wavenet.cpp → NAM/wavenet/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,30 +7,30 @@

#include <Eigen/Dense>

#include "get_dsp.h"
#include "registry.h"
#include "slimmable_wavenet.h"
#include "wavenet.h"
#include "../get_dsp.h"
#include "../registry.h"
#include "slimmable.h"
#include "model.h"

// PostStackHead (WaveNet post-stack head) =====================================
// detail::Head (WaveNet post-stack head) =====================================

nam::wavenet::PostStackHead::PostStackHead(const WaveNetHeadParams& params)
nam::wavenet::detail::Head::Head(const HeadParams& params)
: _in_channels(params.in_channels)
, _out_channels(params.out_channels)
{
if (params.kernel_sizes.empty())
throw std::runtime_error("PostStackHead: kernel_sizes must be non-empty");
throw std::runtime_error("WaveNet Head: kernel_sizes must be non-empty");
const size_t n = params.kernel_sizes.size();
int cin = params.in_channels;
for (size_t i = 0; i < n; i++)
{
const int cout = (i + 1 == n) ? params.out_channels : params.channels;
const int k = params.kernel_sizes[i];
if (k < 1)
throw std::runtime_error("PostStackHead: kernel_sizes entries must be >= 1");
throw std::runtime_error("WaveNet Head: kernel_sizes entries must be >= 1");
nam::activations::Activation::Ptr act = nam::activations::Activation::get_activation(params.activation_config);
if (act == nullptr)
throw std::runtime_error("PostStackHead: unsupported activation for post-stack head");
throw std::runtime_error("WaveNet Head: unsupported activation for post-stack head");
_activations.push_back(std::move(act));
nam::Conv1D conv;
conv.set_size_(cin, cout, k, true, 1, 1);
Expand All @@ -39,19 +39,19 @@ nam::wavenet::PostStackHead::PostStackHead(const WaveNetHeadParams& params)
}
}

void nam::wavenet::PostStackHead::set_weights_(std::vector<float>::iterator& weights)
void nam::wavenet::detail::Head::set_weights_(std::vector<float>::iterator& weights)
{
for (size_t i = 0; i < _convs.size(); i++)
_convs[i].set_weights_(weights);
}

void nam::wavenet::PostStackHead::SetMaxBufferSize(const int maxBufferSize)
void nam::wavenet::detail::Head::SetMaxBufferSize(const int maxBufferSize)
{
for (size_t i = 0; i < _convs.size(); i++)
_convs[i].SetMaxBufferSize(maxBufferSize);
}

long nam::wavenet::PostStackHead::receptive_field() const
long nam::wavenet::detail::Head::receptive_field() const
{
long rf = 1;
for (size_t i = 0; i < _convs.size(); i++)
Expand All @@ -62,7 +62,7 @@ long nam::wavenet::PostStackHead::receptive_field() const
return rf;
}

void nam::wavenet::PostStackHead::process(Eigen::MatrixXf& work, const int num_frames)
void nam::wavenet::detail::Head::process(Eigen::MatrixXf& work, const int num_frames)
{
for (size_t i = 0; i < _convs.size(); i++)
{
Expand All @@ -83,7 +83,7 @@ void nam::wavenet::PostStackHead::process(Eigen::MatrixXf& work, const int num_f

// Layer ======================================================================

void nam::wavenet::_Layer::SetMaxBufferSize(const int maxBufferSize)
void nam::wavenet::detail::Layer::SetMaxBufferSize(const int maxBufferSize)
{
_conv.SetMaxBufferSize(maxBufferSize);
_input_mixin.SetMaxBufferSize(maxBufferSize);
Expand Down Expand Up @@ -128,7 +128,7 @@ void nam::wavenet::_Layer::SetMaxBufferSize(const int maxBufferSize)
this->_head1x1_post_film->SetMaxBufferSize(maxBufferSize);
}

void nam::wavenet::_Layer::set_weights_(std::vector<float>::iterator& weights)
void nam::wavenet::detail::Layer::set_weights_(std::vector<float>::iterator& weights)
{
this->_conv.set_weights_(weights);
this->_input_mixin.set_weights_(weights);
Expand Down Expand Up @@ -159,7 +159,8 @@ void nam::wavenet::_Layer::set_weights_(std::vector<float>::iterator& weights)
this->_head1x1_post_film->set_weights_(weights);
}

void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::MatrixXf& condition, const int num_frames)
void nam::wavenet::detail::Layer::Process(const Eigen::MatrixXf& input, const Eigen::MatrixXf& condition,
const int num_frames)
{
const long bottleneck = this->_bottleneck; // Use the actual bottleneck value, not the doubled output channels

Expand Down Expand Up @@ -372,7 +373,7 @@ void nam::wavenet::_Layer::Process(const Eigen::MatrixXf& input, const Eigen::Ma

// LayerArray =================================================================

nam::wavenet::_LayerArray::_LayerArray(const LayerArrayParams& params)
nam::wavenet::detail::LayerArray::LayerArray(const LayerArrayParams& params)
: _rechannel(params.input_size, params.channels, false)
, _head_rechannel(params.head1x1_params.active ? params.head1x1_params.out_channels : params.bottleneck,
params.head_size, params.head_kernel_size, params.head_bias ? 1 : 0, 1, 1)
Expand All @@ -389,11 +390,11 @@ nam::wavenet::_LayerArray::_LayerArray(const LayerArrayParams& params)
params.conv_pre_film_params, params.conv_post_film_params, params.input_mixin_pre_film_params,
params.input_mixin_post_film_params, params.activation_pre_film_params, params.activation_post_film_params,
params._layer1x1_post_film_params, params.head1x1_post_film_params);
this->_layers.push_back(_Layer(layer_params));
this->_layers.push_back(Layer(layer_params));
}
}

void nam::wavenet::_LayerArray::SetMaxBufferSize(const int maxBufferSize)
void nam::wavenet::detail::LayerArray::SetMaxBufferSize(const int maxBufferSize)
{
_rechannel.SetMaxBufferSize(maxBufferSize);
_head_rechannel.SetMaxBufferSize(maxBufferSize);
Expand All @@ -409,7 +410,7 @@ void nam::wavenet::_LayerArray::SetMaxBufferSize(const int maxBufferSize)
}


long nam::wavenet::_LayerArray::get_receptive_field() const
long nam::wavenet::detail::LayerArray::get_receptive_field() const
{
long result = 0;
for (size_t i = 0; i < this->_layers.size(); i++)
Expand All @@ -419,16 +420,16 @@ long nam::wavenet::_LayerArray::get_receptive_field() const
}


void nam::wavenet::_LayerArray::Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
const int num_frames)
void nam::wavenet::detail::LayerArray::Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
const int num_frames)
{
// Zero head inputs accumulator (first layer array)
this->_head_inputs.setZero();
ProcessInner(layer_inputs, condition, num_frames);
}

void nam::wavenet::_LayerArray::Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
const Eigen::MatrixXf& head_inputs, const int num_frames)
void nam::wavenet::detail::LayerArray::Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
const Eigen::MatrixXf& head_inputs, const int num_frames)
{
// Copy head inputs from previous layer array - use memcpy for pure copy
#ifdef NAM_USE_INLINE_GEMM
Expand All @@ -442,8 +443,8 @@ void nam::wavenet::_LayerArray::Process(const Eigen::MatrixXf& layer_inputs, con
ProcessInner(layer_inputs, condition, num_frames);
}

void nam::wavenet::_LayerArray::ProcessInner(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition,
const int num_frames)
void nam::wavenet::detail::LayerArray::ProcessInner(const Eigen::MatrixXf& layer_inputs,
const Eigen::MatrixXf& condition, const int num_frames)
{
// Process rechannel and get output
this->_rechannel.process_(layer_inputs, num_frames);
Expand Down Expand Up @@ -506,34 +507,34 @@ void nam::wavenet::_LayerArray::ProcessInner(const Eigen::MatrixXf& layer_inputs
}


Eigen::MatrixXf& nam::wavenet::_LayerArray::GetHeadOutputs()
Eigen::MatrixXf& nam::wavenet::detail::LayerArray::GetHeadOutputs()
{
return this->_head_rechannel.GetOutput();
}

const Eigen::MatrixXf& nam::wavenet::_LayerArray::GetHeadOutputs() const
const Eigen::MatrixXf& nam::wavenet::detail::LayerArray::GetHeadOutputs() const
{
return this->_head_rechannel.GetOutput();
}


void nam::wavenet::_LayerArray::set_weights_(std::vector<float>::iterator& weights)
void nam::wavenet::detail::LayerArray::set_weights_(std::vector<float>::iterator& weights)
{
this->_rechannel.set_weights_(weights);
for (size_t i = 0; i < this->_layers.size(); i++)
this->_layers[i].set_weights_(weights);
this->_head_rechannel.set_weights_(weights);
}

long nam::wavenet::_LayerArray::_get_channels() const
long nam::wavenet::detail::LayerArray::_get_channels() const
{
return this->_layers.size() > 0 ? this->_layers[0].get_channels() : 0;
}

namespace
{
int wave_net_output_channels(const std::vector<nam::wavenet::LayerArrayParams>& layer_array_params,
const bool with_head, const std::optional<nam::wavenet::WaveNetHeadParams>& head_params)
const bool with_head, const std::optional<nam::wavenet::HeadParams>& head_params)
{
if (layer_array_params.empty())
throw std::runtime_error("WaveNet requires at least one layer array");
Expand All @@ -547,9 +548,9 @@ int wave_net_output_channels(const std::vector<nam::wavenet::LayerArrayParams>&

nam::wavenet::WaveNet::WaveNet(const int in_channels,
const std::vector<nam::wavenet::LayerArrayParams>& layer_array_params,
const float head_scale, const bool with_head,
std::optional<WaveNetHeadParams> head_params, std::vector<float> weights,
std::unique_ptr<DSP> condition_dsp, const double expected_sample_rate)
const float head_scale, const bool with_head, std::optional<HeadParams> head_params,
std::vector<float> weights, std::unique_ptr<DSP> condition_dsp,
const double expected_sample_rate)
: DSP(in_channels, wave_net_output_channels(layer_array_params, with_head, head_params), expected_sample_rate)
, _condition_dsp(std::move(condition_dsp))
, _head_scale(head_scale)
Expand All @@ -576,7 +577,7 @@ nam::wavenet::WaveNet::WaveNet(const int in_channels,
<< layer_array_params.back().head_size << ")";
throw std::runtime_error(ss.str());
}
this->_post_stack_head = std::make_unique<PostStackHead>(*head_params);
this->_post_stack_head = std::make_unique<detail::Head>(*head_params);
}
else if (head_params.has_value())
throw std::runtime_error("WaveNet: head configuration provided but with_head is false");
Expand All @@ -595,7 +596,7 @@ nam::wavenet::WaveNet::WaveNet(const int in_channels,
throw std::runtime_error(ss.str().c_str());
}
}
this->_layer_arrays.push_back(nam::wavenet::_LayerArray(layer_array_params[i]));
this->_layer_arrays.push_back(nam::wavenet::detail::LayerArray(layer_array_params[i]));
if (i > 0)
if (layer_array_params[i].channels != layer_array_params[i - 1].head_size)
{
Expand Down Expand Up @@ -1149,7 +1150,7 @@ nam::wavenet::WaveNetConfig nam::wavenet::parse_config_json(const nlohmann::json
if (wc.with_head)
{
const nlohmann::json& hj = config["head"];
WaveNetHeadParams hp;
HeadParams hp;
const int implied_in = wc.layer_array_params.back().head_size;
// New trainer export omits in_channels (single source: last layer head_size). Legacy .nam may include it.
if (hj.find("in_channels") != hj.end() && !hj["in_channels"].is_null())
Expand Down
147 changes: 147 additions & 0 deletions NAM/wavenet/model.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,147 @@
#pragma once

// This header defines the WaveNet end-user model: ``WaveNet`` (DSP), ``WaveNetConfig``, and JSON helpers
// ``parse_config_json`` / ``create_config``. Lower-level building blocks live in ``params.h`` and ``detail.h``.

#include <memory>
#include <optional>
#include <vector>

#include <Eigen/Dense>

#include "../dsp.h"
#include "json.hpp"

#include "detail.h"

namespace nam
{
namespace wavenet
{

/// \brief The main WaveNet model
///
/// WaveNet is a dilated convolutional neural network architecture for audio processing.
/// It consists of multiple LayerArrays, each containing multiple layers with increasing
/// dilation factors. The model processes audio through:
///
/// 1. Condition DSP (optional) - processes input to generate conditioning signal
/// 2. LayerArrays - sequential processing with residual and skip connections
/// 3. Head scaling - final output scaling
///
/// The model supports real-time audio processing with pre-allocated buffers.
class WaveNet : public DSP
{
public:
/// \brief Constructor
/// \param in_channels Number of input channels
/// \param layer_array_params Parameters for each layer array
/// \param head_scale Scaling factor applied to the final head output
/// \param with_head Whether to apply the optional post-stack head (Conv1D stack after layer arrays)
/// \param head_params Configuration for the post-stack head when ``with_head`` is true
/// \param weights Model weights (will be consumed during construction)
/// \param condition_dsp Optional DSP module for processing the conditioning input
/// \param expected_sample_rate Expected sample rate in Hz (-1.0 if unknown)
WaveNet(const int in_channels, const std::vector<LayerArrayParams>& layer_array_params, const float head_scale,
const bool with_head, std::optional<HeadParams> head_params, std::vector<float> weights,
std::unique_ptr<DSP> condition_dsp, const double expected_sample_rate = -1.0);

/// \brief Destructor
~WaveNet() = default;

/// \brief Process audio frames
///
/// Implements the DSP::process() interface. Processes input audio through the
/// complete WaveNet pipeline and writes to output.
/// \param input Input audio buffers (in_channels x frames)
/// \param output Output audio buffers (out_channels x frames)
/// \param num_frames Number of frames to process
void process(NAM_SAMPLE** input, NAM_SAMPLE** output, const int num_frames) override;

/// \brief Set model weights from a vector
/// \param weights Vector containing all model weights
void set_weights_(std::vector<float>& weights);

/// \brief Set model weights from an iterator
/// \param weights Iterator to the weights vector. Will be advanced as weights are consumed.
void set_weights_(std::vector<float>::iterator& weights);

protected:
// Element-wise arrays:
Eigen::MatrixXf _condition_input;
Eigen::MatrixXf _condition_output;
std::unique_ptr<DSP> _condition_dsp;
// Temporary buffers for condition DSP processing (to avoid allocations in _process_condition)
std::vector<std::vector<NAM_SAMPLE>> _condition_dsp_input_buffers;
std::vector<std::vector<NAM_SAMPLE>> _condition_dsp_output_buffers;
std::vector<NAM_SAMPLE*> _condition_dsp_input_ptrs;
std::vector<NAM_SAMPLE*> _condition_dsp_output_ptrs;

/// \brief Resize all buffers to handle maxBufferSize frames
/// \param maxBufferSize Maximum number of frames to process in a single call
void SetMaxBufferSize(const int maxBufferSize) override;

/// \brief Compute the conditioning array to be given to the layer arrays
///
/// Processes the condition input through the condition DSP (if present) or
/// passes it through directly.
/// \param num_frames Number of frames to process
virtual void _process_condition(const int num_frames);

/// \brief Fill in the "condition" array that's fed into the various parts of the net
///
/// Copies input audio into the condition buffer for processing.
/// \param input Input audio buffers
/// \param num_frames Number of frames to process
virtual void _set_condition_array(NAM_SAMPLE** input, const int num_frames);

/// \brief Get the number of conditioning inputs
///
/// For standard WaveNet, this is just the audio input (same as input channels).
/// \return Number of conditioning input channels
virtual int _get_condition_dim() const { return NumInputChannels(); };

private:
std::vector<detail::LayerArray> _layer_arrays;

float _head_scale;

std::unique_ptr<detail::Head> _post_stack_head;
/// Scratch (in_channels × maxBufferSize) for scaled head input when ``_post_stack_head`` is used
Eigen::MatrixXf _scaled_head_scratch;

int mPrewarmSamples = 0; // Pre-compute during initialization
int PrewarmSamples() override { return mPrewarmSamples; };
};

/// \brief Configuration for a WaveNet model
struct WaveNetConfig : public ModelConfig
{
int in_channels;
std::vector<LayerArrayParams> layer_array_params;
float head_scale;
bool with_head;
std::optional<HeadParams> head_params;
std::unique_ptr<DSP> condition_dsp;

// Move-only due to unique_ptr
WaveNetConfig() = default;
WaveNetConfig(WaveNetConfig&&) = default;
WaveNetConfig& operator=(WaveNetConfig&&) = default;
WaveNetConfig(const WaveNetConfig&) = delete;
WaveNetConfig& operator=(const WaveNetConfig&) = delete;

std::unique_ptr<DSP> create(std::vector<float> weights, double sampleRate) override;
};

/// \brief Parse WaveNet configuration from JSON
/// \param config JSON configuration object
/// \param expectedSampleRate Expected sample rate in Hz (-1.0 if unknown)
/// \return WaveNetConfig
WaveNetConfig parse_config_json(const nlohmann::json& config, const double expectedSampleRate);

/// \brief Config parser for ConfigParserRegistry
std::unique_ptr<ModelConfig> create_config(const nlohmann::json& config, double sampleRate);

} // namespace wavenet
} // namespace nam
Loading
Loading