diff --git a/NAM/wavenet.h b/NAM/wavenet.h deleted file mode 100644 index 4aeafe3..0000000 --- a/NAM/wavenet.h +++ /dev/null @@ -1,811 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include - -#include "activations.h" -#include "conv1d.h" -#include "dsp.h" -#include "gating_activations.h" -#include "film.h" -#include "json.hpp" - -namespace nam -{ -namespace wavenet -{ - -/// \brief Gating mode for WaveNet layers -/// -/// Determines how the layer processes the doubled bottleneck channels when gating is enabled. -enum class GatingMode -{ - NONE, ///< No gating or blending - standard activation - GATED, ///< Traditional gating (element-wise multiplication of activated channels) - BLENDED ///< Blending (weighted average between activated and pre-activated values) -}; - -/// \brief Helper function for backward compatibility with boolean gated parameter -/// \param gated Boolean indicating whether gating should be enabled -/// \return GatingMode::GATED if gated is true, GatingMode::NONE otherwise -inline GatingMode gating_mode_from_bool(bool gated) -{ - return gated ? GatingMode::GATED : GatingMode::NONE; -} - -/// \brief Parameters for head1x1 configuration -/// -/// Configures an optional 1x1 convolution that outputs directly to the head (skip connection) -/// instead of using the activation output directly. -struct Head1x1Params -{ - /// \brief Constructor - /// \param active_ Whether the head1x1 convolution is active - /// \param out_channels_ Number of output channels for the head1x1 convolution - /// \param groups_ Number of groups for grouped convolution - Head1x1Params(bool active_, int out_channels_, int groups_) - : active(active_) - , out_channels(out_channels_) - , groups(groups_) - { - } - - const bool active; ///< Whether the head1x1 convolution is active - const int out_channels; ///< Number of output channels - const int groups; ///< Number of groups for grouped convolution -}; - -/// \brief Parameters for layer1x1 configuration -/// -/// Configures an optional 1x1 convolution that processes the activation output -/// for the residual connection to the next layer. -struct Layer1x1Params -{ - /// \brief Constructor - /// \param active_ Whether the layer1x1 convolution is active - /// \param groups_ Number of groups for grouped convolution - Layer1x1Params(bool active_, int groups_) - : active(active_) - , groups(groups_) - { - } - - const bool active; ///< Whether the layer1x1 convolution is active - const int groups; ///< Number of groups for grouped convolution -}; - -/// \brief Parameters for FiLM (Feature-wise Linear Modulation) configuration -/// -/// FiLM applies per-channel scaling and optional shifting based on conditioning input. -struct _FiLMParams -{ - /// \brief Constructor - /// \param active_ Whether FiLM is active at this location - /// \param shift_ Whether to apply both scale and shift (true) or only scale (false) - /// \param groups_ Number of groups for grouped convolution in the condition-to-scale-shift submodule (default: 1) - _FiLMParams(bool active_, bool shift_, int groups_ = 1) - : active(active_) - , shift(shift_) - , groups(groups_) - { - } - const bool active; ///< Whether FiLM is active - const bool shift; ///< Whether to apply shift in addition to scale - const int groups; ///< Number of groups for grouped convolution in the condition-to-scale-shift submodule -}; - -/// \brief Parameters for constructing a single Layer -/// -/// Contains all configuration needed to construct a _Layer -struct LayerParams -{ - /// \brief Constructor - /// \param condition_size_ Size of the conditioning input - /// \param channels_ Number of input/output channels from layer to layer - /// \param bottleneck_ Internal channel count - /// \param kernel_size_ Kernel size for the dilated convolution - /// \param dilation_ Dilation factor for the convolution - /// \param activation_config_ Primary activation function configuration - /// \param gating_mode_ Gating mode (NONE, GATED, or BLENDED) - /// \param groups_input_ Number of groups for the input convolution - /// \param groups_input_mixin_ Number of groups for the input mixin convolution - /// \param layer1x1_params_ Configuration of the optional layer1x1 convolution - /// \param head1x1_params_ Configuration of the optional head1x1 convolution - /// \param secondary_activation_config_ Secondary activation (for gating/blending) - /// \param conv_pre_film_params_ FiLM parameters before the input convolution - /// \param conv_post_film_params_ FiLM parameters after the input convolution - /// \param input_mixin_pre_film_params_ FiLM parameters before the input mixin - /// \param input_mixin_post_film_params_ FiLM parameters after the input mixin - /// \param activation_pre_film_params_ FiLM parameters after the input/mixin summed output before activation - /// \param activation_post_film_params_ FiLM parameters after the activation output before the layer1x1 convolution - /// \param _layer1x1_post_film_params_ FiLM parameters after the layer1x1 convolution - /// \param head1x1_post_film_params_ FiLM parameters after the head1x1 convolution - LayerParams(const int condition_size_, const int channels_, const int bottleneck_, const int kernel_size_, - const int dilation_, const activations::ActivationConfig& activation_config_, - const GatingMode gating_mode_, const int groups_input_, const int groups_input_mixin_, - const Layer1x1Params& layer1x1_params_, const Head1x1Params& head1x1_params_, - const activations::ActivationConfig& secondary_activation_config_, - const _FiLMParams& conv_pre_film_params_, const _FiLMParams& conv_post_film_params_, - const _FiLMParams& input_mixin_pre_film_params_, const _FiLMParams& input_mixin_post_film_params_, - const _FiLMParams& activation_pre_film_params_, const _FiLMParams& activation_post_film_params_, - const _FiLMParams& _layer1x1_post_film_params_, const _FiLMParams& head1x1_post_film_params_) - : condition_size(condition_size_) - , channels(channels_) - , bottleneck(bottleneck_) - , kernel_size(kernel_size_) - , dilation(dilation_) - , activation_config(activation_config_) - , gating_mode(gating_mode_) - , groups_input(groups_input_) - , groups_input_mixin(groups_input_mixin_) - , layer1x1_params(layer1x1_params_) - , head1x1_params(head1x1_params_) - , secondary_activation_config(secondary_activation_config_) - , conv_pre_film_params(conv_pre_film_params_) - , conv_post_film_params(conv_post_film_params_) - , input_mixin_pre_film_params(input_mixin_pre_film_params_) - , input_mixin_post_film_params(input_mixin_post_film_params_) - , activation_pre_film_params(activation_pre_film_params_) - , activation_post_film_params(activation_post_film_params_) - , _layer1x1_post_film_params(_layer1x1_post_film_params_) - , head1x1_post_film_params(head1x1_post_film_params_) - { - } - - const int condition_size; ///< Size of the conditioning input - const int channels; ///< Number of input/output channels from layer to layer - const int bottleneck; ///< Internal channel count - const int kernel_size; ///< Kernel size for the dilated convolution - const int dilation; ///< Dilation factor for the convolution - const activations::ActivationConfig activation_config; ///< Primary activation function configuration - const GatingMode gating_mode; ///< Gating mode (NONE, GATED, or BLENDED) - const int groups_input; ///< Number of groups for the input convolution - const int groups_input_mixin; ///< Number of groups for the input mixin convolution - const Layer1x1Params layer1x1_params; ///< Configuration of the optional layer1x1 convolution - const Head1x1Params head1x1_params; ///< Configuration of the optional head1x1 convolution - const activations::ActivationConfig secondary_activation_config; ///< Secondary activation (for gating/blending) - const _FiLMParams conv_pre_film_params; ///< FiLM parameters before the input convolution - const _FiLMParams conv_post_film_params; ///< FiLM parameters after the input convolution - const _FiLMParams input_mixin_pre_film_params; ///< FiLM parameters before the input mixin - const _FiLMParams input_mixin_post_film_params; ///< FiLM parameters after the input mixin - const _FiLMParams activation_pre_film_params; ///< FiLM parameters before activation - const _FiLMParams activation_post_film_params; ///< FiLM parameters after activation - const _FiLMParams _layer1x1_post_film_params; ///< FiLM parameters after the layer1x1 convolution (layer1x1_post_film) - const _FiLMParams head1x1_post_film_params; ///< FiLM parameters after the head1x1 convolution -}; - -/// \brief A single WaveNet layer block -/// -/// A WaveNet layer performs the following operations: -/// 1. Dilated convolution on the input (with optional pre/post-FiLM) -/// 2. Input mixin (conditioning input processing, with optional pre/post-FiLM) -/// 3. Sum of conv and input mixin outputs -/// 4. Activation (with optional gating/blending and pre/post FiLM) -/// 5. Optional layer1x1 convolution for the next layer (with optional post-FiLM) -/// 6. Optional 1x1 convolution for the head output (with optional post-FiLM) -/// 7. Residual connection (input + layer1x1 output, or just input if layer1x1 inactive) and skip connection (to next -/// layer) -/// -/// The layer supports multiple gating modes and FiLM at various points in the computation. -/// See the walkthrough documentation for detailed step-by-step explanation. -class _Layer -{ -public: - /// \brief Constructor with LayerParams - /// \param params Parameters for constructing the layer - /// \throws std::invalid_argument If head1x1_post_film_params is active but head1x1 is not, or if layer1x1 is inactive - /// but bottleneck != channels - _Layer(const LayerParams& params) - : _conv(params.channels, (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck, - params.kernel_size, true, params.dilation, params.groups_input) - , _input_mixin(params.condition_size, - (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck, false, - params.groups_input_mixin) - , _activation(activations::Activation::get_activation(params.activation_config)) - , _gating_mode(params.gating_mode) - , _bottleneck(params.bottleneck) - { - if (params.layer1x1_params.active) - { - _layer1x1 = std::make_unique(params.bottleneck, params.channels, true, params.layer1x1_params.groups); - } - else - { - // Validation: if layer1x1 is inactive, bottleneck must equal channels - if (params.bottleneck != params.channels) - { - throw std::invalid_argument("When layer1x1.active is false, bottleneck (" + std::to_string(params.bottleneck) - + ") must equal channels (" + std::to_string(params.channels) + ")"); - } - // If there's a post-layer1x1 FiLM but no layer1x1, this is redundant--don't allow it - if (params._layer1x1_post_film_params.active) - { - throw std::invalid_argument("layer1x1_post_film cannot be active when layer1x1 is not active"); - } - } - - if (params.head1x1_params.active) - { - _head1x1 = std::make_unique( - params.bottleneck, params.head1x1_params.out_channels, true, params.head1x1_params.groups); - } - else - { - // If there's a post-head 1x1 FiLM but no head 1x1, this is redundant--don't allow it - if (params.head1x1_post_film_params.active) - { - throw std::invalid_argument("Do not use post-head 1x1 FiLM if there is no head 1x1"); - } - } - - // When no head1x1 and no gating, _output_head would be a straight copy of _z. - // Skip the copy and return _z directly from GetOutputHead(). - _skip_head_copy = !params.head1x1_params.active && params.gating_mode == GatingMode::NONE; - - // Validate & initialize gating/blending activation - if (params.gating_mode == GatingMode::GATED) - { - _gating_activation = std::make_unique( - _activation, activations::Activation::get_activation(params.secondary_activation_config), params.bottleneck); - } - else if (params.gating_mode == GatingMode::BLENDED) - { - _blending_activation = std::make_unique( - _activation, activations::Activation::get_activation(params.secondary_activation_config), params.bottleneck); - } - - // Initialize FiLM objects - if (params.conv_pre_film_params.active) - { - _conv_pre_film = std::make_unique( - params.condition_size, params.channels, params.conv_pre_film_params.shift, params.conv_pre_film_params.groups); - } - if (params.conv_post_film_params.active) - { - const int conv_out_channels = - (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck; - _conv_post_film = std::make_unique(params.condition_size, conv_out_channels, - params.conv_post_film_params.shift, params.conv_post_film_params.groups); - } - if (params.input_mixin_pre_film_params.active) - { - _input_mixin_pre_film = - std::make_unique(params.condition_size, params.condition_size, params.input_mixin_pre_film_params.shift, - params.input_mixin_pre_film_params.groups); - } - if (params.input_mixin_post_film_params.active) - { - const int input_mixin_out_channels = - (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck; - _input_mixin_post_film = - std::make_unique(params.condition_size, input_mixin_out_channels, - params.input_mixin_post_film_params.shift, params.input_mixin_post_film_params.groups); - } - if (params.activation_pre_film_params.active) - { - const int z_channels = (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck; - _activation_pre_film = - std::make_unique(params.condition_size, z_channels, params.activation_pre_film_params.shift, - params.activation_pre_film_params.groups); - } - if (params.activation_post_film_params.active) - { - _activation_post_film = - std::make_unique(params.condition_size, params.bottleneck, params.activation_post_film_params.shift, - params.activation_post_film_params.groups); - } - if (params._layer1x1_post_film_params.active && params.layer1x1_params.active) - { - _layer1x1_post_film = - std::make_unique(params.condition_size, params.channels, params._layer1x1_post_film_params.shift, - params._layer1x1_post_film_params.groups); - } - if (params.head1x1_post_film_params.active && params.head1x1_params.active) - { - _head1x1_post_film = - std::make_unique(params.condition_size, params.head1x1_params.out_channels, - params.head1x1_post_film_params.shift, params.head1x1_post_film_params.groups); - } - }; - - /// \brief Resize all arrays to be able to process maxBufferSize frames - /// \param maxBufferSize Maximum number of frames to process in a single call - void SetMaxBufferSize(const int maxBufferSize); - - /// \brief Set the parameters (weights) of this module - /// \param weights Iterator to the weights vector. Will be advanced as weights are consumed. - void set_weights_(std::vector::iterator& weights); - - /// \brief Process a block of frames - /// - /// Performs the complete layer computation: - /// 1. Input convolution (with optional pre/post-FiLM) - /// 2. Input mixin processing (with optional pre/post-FiLM) - /// 3. Sum and activation (with optional gating/blending and pre/post-FiLM) - /// 4. Optional layer1x1 convolution toward the skip connection for next layer (with optional post-FiLM) - /// 5. Optional 1x1 convolution for the head output (with optional post-FiLM) - /// 6. Store outputs for next layer and the layer array head - /// - /// \param input Input from previous layer (channels x num_frames) - /// \param condition Conditioning input (condition_size x num_frames) - /// \param num_frames Number of frames to process - /// - /// Outputs are stored internally and accessible via GetOutputNextLayer() and GetOutputHead(). - /// Only the first num_frames columns of the output buffers are valid. - void Process(const Eigen::MatrixXf& input, const Eigen::MatrixXf& condition, const int num_frames); - - /// \brief Get the number of channels expected as input/output from this layer - /// \return Number of channels - long get_channels() const { return this->_conv.get_in_channels(); }; - - /// \brief Get the dilation of the input convolution layer - /// \return Dilation factor - int get_dilation() const { return this->_conv.get_dilation(); }; - - /// \brief Get the kernel size of the input convolution layer - /// \return Kernel size - long get_kernel_size() const { return this->_conv.get_kernel_size(); }; - - /// \brief Get output to next layer (residual connection: input + layer1x1 output) - /// - /// Returns the full pre-allocated buffer; only the first num_frames columns - /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. - /// \return Reference to the output buffer (channels x maxBufferSize) - Eigen::MatrixXf& GetOutputNextLayer() { return this->_output_next_layer; } - - /// \brief Get output to next layer (const version) - /// \return Const reference to the output buffer - const Eigen::MatrixXf& GetOutputNextLayer() const { return this->_output_next_layer; } - - /// \brief Get output to head (skip connection: activated conv output) - /// - /// Returns the full pre-allocated buffer; only the first num_frames columns - /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. - /// When _skip_head_copy is true (no head1x1, no gating), returns _z directly - /// to avoid a redundant memcpy. - /// \return Reference to the head output buffer - Eigen::MatrixXf& GetOutputHead() { return _skip_head_copy ? this->_z : this->_output_head; } - - /// \brief Get output to head (const version) - /// \return Const reference to the head output buffer - const Eigen::MatrixXf& GetOutputHead() const { return _skip_head_copy ? this->_z : this->_output_head; } - - /// \brief Access Conv1D for Reset() propagation (needed for _LayerArray) - /// \return Reference to the internal Conv1D object - Conv1D& get_conv() { return _conv; } - - /// \brief Access Conv1D (const version) - /// \return Const reference to the internal Conv1D object - const Conv1D& get_conv() const { return _conv; } - -private: - // The dilated convolution at the front of the block - Conv1D _conv; - // Input mixin - Conv1x1 _input_mixin; - // The post-activation layer1x1 convolution (optional) - std::unique_ptr _layer1x1; - // The post-activation 1x1 convolution outputting to the head, optional - std::unique_ptr _head1x1; - // The internal state - Eigen::MatrixXf _z; - // Output to next layer (residual connection: input + layer1x1 output, or just input if layer1x1 inactive) - Eigen::MatrixXf _output_next_layer; - // Output to head (skip connection: activated conv output) - Eigen::MatrixXf _output_head; - - activations::Activation::Ptr _activation; - const GatingMode _gating_mode; - const int _bottleneck; // Internal channel count (not doubled when gated) - bool _skip_head_copy = false; // When true, GetOutputHead() returns _z directly (no head1x1, no gating) - - // Gating/blending activation objects - std::unique_ptr _gating_activation; - std::unique_ptr _blending_activation; - - // FiLM objects for feature-wise linear modulation - std::unique_ptr _conv_pre_film; - std::unique_ptr _conv_post_film; - std::unique_ptr _input_mixin_pre_film; - std::unique_ptr _input_mixin_post_film; - std::unique_ptr _activation_pre_film; - std::unique_ptr _activation_post_film; - std::unique_ptr _layer1x1_post_film; - std::unique_ptr _head1x1_post_film; -}; - -/// \brief Parameters for constructing a LayerArray -/// -/// Contains all configuration needed to construct a _LayerArray with multiple layers -/// sharing the same channel count and kernel size. Each layer can have its own activation configuration. -class LayerArrayParams -{ -public: - /// \brief Constructor - /// \param input_size_ Input size (number of channels) to the layer array - /// \param condition_size_ Size of the conditioning input - /// \param head_size_ Size of the head output (after head rechannel) - /// \param channels_ Number of channels in each layer - /// \param bottleneck_ Bottleneck size (internal channel count) - /// \param kernel_sizes_ Per-layer kernel sizes, one per layer - /// \param dilations_ Vector of dilation factors, one per layer - /// \param activation_configs_ Vector of primary activation configurations, one per layer - /// \param gating_modes_ Vector of gating modes, one per layer - /// \param head_kernel_size_ Kernel size of the head rechannel conv (>= 1) - /// \param head_bias_ Whether to use bias in the head rechannel - /// \param groups_input Number of groups for input convolutions - /// \param groups_input_mixin_ Number of groups for input mixin convolutions - /// \param layer1x1_params_ Parameters for optional layer1x1 convolutions - /// \param head1x1_params_ Parameters for optional head1x1 convolutions - /// \param secondary_activation_configs_ Vector of secondary activation configs for gating/blending, one per layer - /// \param conv_pre_film_params_ FiLM parameters before input convolutions - /// \param conv_post_film_params_ FiLM parameters after input convolutions - /// \param input_mixin_pre_film_params_ FiLM parameters before input mixin - /// \param input_mixin_post_film_params_ FiLM parameters after input mixin - /// \param activation_pre_film_params_ FiLM parameters before activation - /// \param activation_post_film_params_ FiLM parameters after activation - /// \param _layer1x1_post_film_params_ FiLM parameters after layer1x1 convolutions - /// \param head1x1_post_film_params_ FiLM parameters after head1x1 convolutions - /// \throws std::invalid_argument If dilations, activation_configs, gating_modes, or secondary_activation_configs - /// sizes don't match - LayerArrayParams(const int input_size_, const int condition_size_, const int head_size_, const int head_kernel_size_, - const int channels_, const int bottleneck_, const std::vector&& kernel_sizes_, - const std::vector&& dilations_, - const std::vector&& activation_configs_, - const std::vector&& gating_modes_, const bool head_bias_, const int groups_input, - const int groups_input_mixin_, const Layer1x1Params& layer1x1_params_, - const Head1x1Params& head1x1_params_, - const std::vector&& secondary_activation_configs_, - const _FiLMParams& conv_pre_film_params_, const _FiLMParams& conv_post_film_params_, - const _FiLMParams& input_mixin_pre_film_params_, const _FiLMParams& input_mixin_post_film_params_, - const _FiLMParams& activation_pre_film_params_, const _FiLMParams& activation_post_film_params_, - const _FiLMParams& _layer1x1_post_film_params_, const _FiLMParams& head1x1_post_film_params_) - : input_size(input_size_) - , condition_size(condition_size_) - , head_size(head_size_) - , head_kernel_size(head_kernel_size_) - , channels(channels_) - , bottleneck(bottleneck_) - , kernel_sizes(std::move(kernel_sizes_)) - , dilations(std::move(dilations_)) - , activation_configs(std::move(activation_configs_)) - , gating_modes(std::move(gating_modes_)) - , head_bias(head_bias_) - , groups_input(groups_input) - , groups_input_mixin(groups_input_mixin_) - , layer1x1_params(layer1x1_params_) - , head1x1_params(head1x1_params_) - , secondary_activation_configs(std::move(secondary_activation_configs_)) - , conv_pre_film_params(conv_pre_film_params_) - , conv_post_film_params(conv_post_film_params_) - , input_mixin_pre_film_params(input_mixin_pre_film_params_) - , input_mixin_post_film_params(input_mixin_post_film_params_) - , activation_pre_film_params(activation_pre_film_params_) - , activation_post_film_params(activation_post_film_params_) - , _layer1x1_post_film_params(_layer1x1_post_film_params_) - , head1x1_post_film_params(head1x1_post_film_params_) - { - if (head_kernel_size < 1) - { - throw std::invalid_argument("LayerArrayParams: head_kernel_size must be >= 1"); - } - const size_t num_layers = dilations.size(); - if (kernel_sizes.empty()) - { - throw std::invalid_argument("LayerArrayParams: kernel_sizes must not be empty"); - } - if (kernel_sizes.size() != num_layers) - { - throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) - + ") must match kernel_sizes size (" + std::to_string(kernel_sizes.size()) + ")"); - } - if (activation_configs.size() != num_layers) - { - throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) - + ") must match activation_configs size (" + std::to_string(activation_configs.size()) - + ")"); - } - if (gating_modes.size() != num_layers) - { - throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) - + ") must match gating_modes size (" + std::to_string(gating_modes.size()) + ")"); - } - if (secondary_activation_configs.size() != num_layers) - { - throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) - + ") must match secondary_activation_configs size (" - + std::to_string(secondary_activation_configs.size()) + ")"); - } - } - - const int input_size; ///< Input size (number of channels) - const int condition_size; ///< Size of conditioning input - const int head_size; ///< Size of head output (after rechannel) - const int head_kernel_size; ///< Kernel size of head rechannel convolution (>= 1) - const int channels; ///< Number of channels in each layer - const int bottleneck; ///< Bottleneck size (internal channel count) - std::vector kernel_sizes; ///< Per-layer kernel sizes, one per layer - std::vector dilations; ///< Dilation factors, one per layer - std::vector activation_configs; ///< Primary activation configurations, one per layer - std::vector gating_modes; ///< Gating modes, one per layer - const bool head_bias; ///< Whether to use bias in head rechannel - const int groups_input; ///< Number of groups for input convolutions - const int groups_input_mixin; ///< Number of groups for input mixin - const Layer1x1Params layer1x1_params; ///< Parameters for optional layer1x1 - const Head1x1Params head1x1_params; ///< Parameters for optional head1x1 - std::vector - secondary_activation_configs; ///< Secondary activation configs for gating/blending, one per layer - const _FiLMParams conv_pre_film_params; ///< FiLM params before input conv - const _FiLMParams conv_post_film_params; ///< FiLM params after input conv - const _FiLMParams input_mixin_pre_film_params; ///< FiLM params before input mixin - const _FiLMParams input_mixin_post_film_params; ///< FiLM params after input mixin - const _FiLMParams activation_pre_film_params; ///< FiLM params before activation - const _FiLMParams activation_post_film_params; ///< FiLM params after activation - const _FiLMParams _layer1x1_post_film_params; ///< FiLM params after layer1x1 conv - const _FiLMParams head1x1_post_film_params; ///< FiLM params after head1x1 conv -}; - -/// \brief An array of layers with the same channels, kernel sizes, and activations -/// -/// A LayerArray chains multiple _Layer objects together, processing them sequentially. -/// Each layer processes the output of the previous layer (residual connection). -/// All layers contribute to a shared head output (skip connection) that is accumulated -/// and then projected to the final head size. -/// -/// The LayerArray handles: -/// - Input projection to match layer channel count -/// - Processing layers in sequence with residual connections -/// - Accumulating head outputs from all layers -/// - Projecting the accumulated head output to the final head size -class _LayerArray -{ -public: - /// \brief Constructor with LayerArrayParams - /// \param params Parameters for constructing the layer array - _LayerArray(const LayerArrayParams& params); - - /// \brief Resize all arrays to be able to process maxBufferSize frames - /// \param maxBufferSize Maximum number of frames to process in a single call - void SetMaxBufferSize(const int maxBufferSize); - - /// \brief Process without a given previous head input (first layer array) - /// - /// Zeros head accumulated output before proceeding. Used for the first layer array in a WaveNet. - /// \param layer_inputs Input to the layer array (input_size x num_frames) - /// \param condition Conditioning input (condition_size x num_frames) - /// \param num_frames Number of frames to process - void Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition, const int num_frames); - - /// \brief Process with a given previous head input (subsequent layer arrays) - /// - /// Copies head input before proceeding. Used for subsequent layer arrays that accumulate - /// head outputs from previous arrays. - /// \param layer_inputs Input to the layer array (input_size x num_frames) - /// \param condition Conditioning input (condition_size x num_frames) - /// \param head_inputs Head input from previous layer array (head_input_size x num_frames) - /// \param num_frames Number of frames to process - void Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition, - const Eigen::MatrixXf& head_inputs, const int num_frames); - - /// \brief Get output from last layer (for next layer array) - /// - /// Returns the full pre-allocated buffer; only the first num_frames columns - /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. - /// \return Reference to the layer output buffer (channels x maxBufferSize) - Eigen::MatrixXf& GetLayerOutputs() { return this->_layer_outputs; } - - /// \brief Get output from last layer (const version) - /// \return Const reference to the layer output buffer - const Eigen::MatrixXf& GetLayerOutputs() const { return this->_layer_outputs; } - - /// \brief Get head outputs (post head-rechannel) - /// - /// Returns the full pre-allocated buffer; only the first num_frames columns - /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. - /// \return Reference to the head output buffer (head_size x maxBufferSize) - Eigen::MatrixXf& GetHeadOutputs(); - - /// \brief Get head outputs (const version) - /// \return Const reference to the head output buffer - const Eigen::MatrixXf& GetHeadOutputs() const; - - /// \brief Set the parameters (weights) of this module - /// \param it Iterator to the weights vector. Will be advanced as weights are consumed. - void set_weights_(std::vector::iterator& it); - - /// \brief Get the "zero-indexed" receptive field - /// - /// The receptive field is the number of input samples that affect the output. - /// A 1x1 convolution is defined to have a zero-indexed receptive field of zero. - /// \return Receptive field size - long get_receptive_field() const; - -private: - // The rechannel before the layers - Conv1x1 _rechannel; - - // The layer objects - std::vector<_Layer> _layers; - // Output from last layer (for next layer array) - Eigen::MatrixXf _layer_outputs; - // Accumulated head inputs from all layers - // Size is _head_output_size (= head1x1.out_channels if head1x1 active, else bottleneck) - Eigen::MatrixXf _head_inputs; - - // Rechannel for the head (_head_output_size -> head_size), causal Conv1D (dilation 1) - Conv1D _head_rechannel; - - // Head output size from each layer (head1x1.out_channels if active, else bottleneck) - const int _head_output_size; - - long _get_channels() const; - // Common processing logic after head inputs are set - void ProcessInner(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition, const int num_frames); -}; - -/// \brief Parameters for the optional post-stack head (matches Python ``nam.models.wavenet._head.Head``). -/// JSON export omits ``in_channels`` (implied by last layer array ``head_size``); load sets it from there. -struct WaveNetHeadParams -{ - int in_channels; - int channels; - int out_channels; - std::vector kernel_sizes; - activations::ActivationConfig activation_config; -}; - -/// \brief Post-stack head: repeated (activation → Conv1D) with dilation 1, stride 1, valid (causal streaming) conv. -class PostStackHead -{ -public: - explicit PostStackHead(const WaveNetHeadParams& params); - - void set_weights_(std::vector::iterator& weights); - void SetMaxBufferSize(int maxBufferSize); - long receptive_field() const; - int in_channels() const { return _in_channels; } - int out_channels() const { return _out_channels; } - - /// \param work Input buffer (in_channels × maxBufferSize); first in_channels×num_frames scaled by head_scale; - /// may be modified in place. - void process(Eigen::MatrixXf& work, int num_frames); - - const Eigen::MatrixXf& get_last_output() const { return _convs.back().GetOutput(); } - -private: - std::vector _convs; - std::vector _activations; - int _in_channels; - int _out_channels; -}; - -/// \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& layer_array_params, const float head_scale, - const bool with_head, std::optional head_params, std::vector weights, - std::unique_ptr 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& 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::iterator& weights); - -protected: - // Element-wise arrays: - Eigen::MatrixXf _condition_input; - Eigen::MatrixXf _condition_output; - std::unique_ptr _condition_dsp; - // Temporary buffers for condition DSP processing (to avoid allocations in _process_condition) - std::vector> _condition_dsp_input_buffers; - std::vector> _condition_dsp_output_buffers; - std::vector _condition_dsp_input_ptrs; - std::vector _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<_LayerArray> _layer_arrays; - - float _head_scale; - - std::unique_ptr _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 layer_array_params; - float head_scale; - bool with_head; - std::optional head_params; - std::unique_ptr 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 create(std::vector 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 create_config(const nlohmann::json& config, double sampleRate); -}; // namespace wavenet -}; // namespace nam diff --git a/NAM/wavenet/detail.h b/NAM/wavenet/detail.h new file mode 100644 index 0000000..1e5d10e --- /dev/null +++ b/NAM/wavenet/detail.h @@ -0,0 +1,388 @@ +#pragma once + +#include "params.h" + +#include +#include +#include +#include + +#include + +#include "../conv1d.h" +#include "../gating_activations.h" +#include "../film.h" + +namespace nam +{ +namespace wavenet +{ +namespace detail +{ + +/// \brief A single WaveNet layer block +/// +/// A WaveNet layer performs the following operations: +/// 1. Dilated convolution on the input (with optional pre/post-FiLM) +/// 2. Input mixin (conditioning input processing, with optional pre/post-FiLM) +/// 3. Sum of conv and input mixin outputs +/// 4. Activation (with optional gating/blending and pre/post FiLM) +/// 5. Optional layer1x1 convolution for the next layer (with optional post-FiLM) +/// 6. Optional 1x1 convolution for the head output (with optional post-FiLM) +/// 7. Residual connection (input + layer1x1 output, or just input if layer1x1 inactive) and skip connection (to next +/// layer) +/// +/// The layer supports multiple gating modes and FiLM at various points in the computation. +/// See the walkthrough documentation for detailed step-by-step explanation. +class Layer +{ +public: + /// \brief Constructor with LayerParams + /// \param params Parameters for constructing the layer + /// \throws std::invalid_argument If head1x1_post_film_params is active but head1x1 is not, or if layer1x1 is inactive + /// but bottleneck != channels + Layer(const LayerParams& params) + : _conv(params.channels, (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck, + params.kernel_size, true, params.dilation, params.groups_input) + , _input_mixin(params.condition_size, + (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck, false, + params.groups_input_mixin) + , _activation(activations::Activation::get_activation(params.activation_config)) + , _gating_mode(params.gating_mode) + , _bottleneck(params.bottleneck) + { + if (params.layer1x1_params.active) + { + _layer1x1 = std::make_unique(params.bottleneck, params.channels, true, params.layer1x1_params.groups); + } + else + { + // Validation: if layer1x1 is inactive, bottleneck must equal channels + if (params.bottleneck != params.channels) + { + throw std::invalid_argument("When layer1x1.active is false, bottleneck (" + std::to_string(params.bottleneck) + + ") must equal channels (" + std::to_string(params.channels) + ")"); + } + // If there's a post-layer1x1 FiLM but no layer1x1, this is redundant--don't allow it + if (params._layer1x1_post_film_params.active) + { + throw std::invalid_argument("layer1x1_post_film cannot be active when layer1x1 is not active"); + } + } + + if (params.head1x1_params.active) + { + _head1x1 = std::make_unique( + params.bottleneck, params.head1x1_params.out_channels, true, params.head1x1_params.groups); + } + else + { + // If there's a post-head 1x1 FiLM but no head 1x1, this is redundant--don't allow it + if (params.head1x1_post_film_params.active) + { + throw std::invalid_argument("Do not use post-head 1x1 FiLM if there is no head 1x1"); + } + } + + // When no head1x1 and no gating, _output_head would be a straight copy of _z. + // Skip the copy and return _z directly from GetOutputHead(). + _skip_head_copy = !params.head1x1_params.active && params.gating_mode == GatingMode::NONE; + + // Validate & initialize gating/blending activation + if (params.gating_mode == GatingMode::GATED) + { + _gating_activation = std::make_unique( + _activation, activations::Activation::get_activation(params.secondary_activation_config), params.bottleneck); + } + else if (params.gating_mode == GatingMode::BLENDED) + { + _blending_activation = std::make_unique( + _activation, activations::Activation::get_activation(params.secondary_activation_config), params.bottleneck); + } + + // Initialize FiLM objects + if (params.conv_pre_film_params.active) + { + _conv_pre_film = std::make_unique( + params.condition_size, params.channels, params.conv_pre_film_params.shift, params.conv_pre_film_params.groups); + } + if (params.conv_post_film_params.active) + { + const int conv_out_channels = + (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck; + _conv_post_film = std::make_unique(params.condition_size, conv_out_channels, + params.conv_post_film_params.shift, params.conv_post_film_params.groups); + } + if (params.input_mixin_pre_film_params.active) + { + _input_mixin_pre_film = + std::make_unique(params.condition_size, params.condition_size, params.input_mixin_pre_film_params.shift, + params.input_mixin_pre_film_params.groups); + } + if (params.input_mixin_post_film_params.active) + { + const int input_mixin_out_channels = + (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck; + _input_mixin_post_film = + std::make_unique(params.condition_size, input_mixin_out_channels, + params.input_mixin_post_film_params.shift, params.input_mixin_post_film_params.groups); + } + if (params.activation_pre_film_params.active) + { + const int z_channels = (params.gating_mode != GatingMode::NONE) ? 2 * params.bottleneck : params.bottleneck; + _activation_pre_film = + std::make_unique(params.condition_size, z_channels, params.activation_pre_film_params.shift, + params.activation_pre_film_params.groups); + } + if (params.activation_post_film_params.active) + { + _activation_post_film = + std::make_unique(params.condition_size, params.bottleneck, params.activation_post_film_params.shift, + params.activation_post_film_params.groups); + } + if (params._layer1x1_post_film_params.active && params.layer1x1_params.active) + { + _layer1x1_post_film = + std::make_unique(params.condition_size, params.channels, params._layer1x1_post_film_params.shift, + params._layer1x1_post_film_params.groups); + } + if (params.head1x1_post_film_params.active && params.head1x1_params.active) + { + _head1x1_post_film = + std::make_unique(params.condition_size, params.head1x1_params.out_channels, + params.head1x1_post_film_params.shift, params.head1x1_post_film_params.groups); + } + }; + + /// \brief Resize all arrays to be able to process maxBufferSize frames + /// \param maxBufferSize Maximum number of frames to process in a single call + void SetMaxBufferSize(const int maxBufferSize); + + /// \brief Set the parameters (weights) of this module + /// \param weights Iterator to the weights vector. Will be advanced as weights are consumed. + void set_weights_(std::vector::iterator& weights); + + /// \brief Process a block of frames + /// + /// Performs the complete layer computation: + /// 1. Input convolution (with optional pre/post-FiLM) + /// 2. Input mixin processing (with optional pre/post-FiLM) + /// 3. Sum and activation (with optional gating/blending and pre/post-FiLM) + /// 4. Optional layer1x1 convolution toward the skip connection for next layer (with optional post-FiLM) + /// 5. Optional 1x1 convolution for the head output (with optional post-FiLM) + /// 6. Store outputs for next layer and the layer array head + /// + /// \param input Input from previous layer (channels x num_frames) + /// \param condition Conditioning input (condition_size x num_frames) + /// \param num_frames Number of frames to process + /// + /// Outputs are stored internally and accessible via GetOutputNextLayer() and GetOutputHead(). + /// Only the first num_frames columns of the output buffers are valid. + void Process(const Eigen::MatrixXf& input, const Eigen::MatrixXf& condition, const int num_frames); + + /// \brief Get the number of channels expected as input/output from this layer + /// \return Number of channels + long get_channels() const { return this->_conv.get_in_channels(); }; + + /// \brief Get the dilation of the input convolution layer + /// \return Dilation factor + int get_dilation() const { return this->_conv.get_dilation(); }; + + /// \brief Get the kernel size of the input convolution layer + /// \return Kernel size + long get_kernel_size() const { return this->_conv.get_kernel_size(); }; + + /// \brief Get output to next layer (residual connection: input + layer1x1 output) + /// + /// Returns the full pre-allocated buffer; only the first num_frames columns + /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. + /// \return Reference to the output buffer (channels x maxBufferSize) + Eigen::MatrixXf& GetOutputNextLayer() { return this->_output_next_layer; } + + /// \brief Get output to next layer (const version) + /// \return Const reference to the output buffer + const Eigen::MatrixXf& GetOutputNextLayer() const { return this->_output_next_layer; } + + /// \brief Get output to head (skip connection: activated conv output) + /// + /// Returns the full pre-allocated buffer; only the first num_frames columns + /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. + /// When _skip_head_copy is true (no head1x1, no gating), returns _z directly + /// to avoid a redundant memcpy. + /// \return Reference to the head output buffer + Eigen::MatrixXf& GetOutputHead() { return _skip_head_copy ? this->_z : this->_output_head; } + + /// \brief Get output to head (const version) + /// \return Const reference to the head output buffer + const Eigen::MatrixXf& GetOutputHead() const { return _skip_head_copy ? this->_z : this->_output_head; } + + /// \brief Access Conv1D for Reset() propagation (needed for LayerArray) + /// \return Reference to the internal Conv1D object + Conv1D& get_conv() { return _conv; } + + /// \brief Access Conv1D (const version) + /// \return Const reference to the internal Conv1D object + const Conv1D& get_conv() const { return _conv; } + +private: + // The dilated convolution at the front of the block + Conv1D _conv; + // Input mixin + Conv1x1 _input_mixin; + // The post-activation layer1x1 convolution (optional) + std::unique_ptr _layer1x1; + // The post-activation 1x1 convolution outputting to the head, optional + std::unique_ptr _head1x1; + // The internal state + Eigen::MatrixXf _z; + // Output to next layer (residual connection: input + layer1x1 output, or just input if layer1x1 inactive) + Eigen::MatrixXf _output_next_layer; + // Output to head (skip connection: activated conv output) + Eigen::MatrixXf _output_head; + + activations::Activation::Ptr _activation; + const GatingMode _gating_mode; + const int _bottleneck; // Internal channel count (not doubled when gated) + bool _skip_head_copy = false; // When true, GetOutputHead() returns _z directly (no head1x1, no gating) + + // Gating/blending activation objects + std::unique_ptr _gating_activation; + std::unique_ptr _blending_activation; + + // FiLM objects for feature-wise linear modulation + std::unique_ptr _conv_pre_film; + std::unique_ptr _conv_post_film; + std::unique_ptr _input_mixin_pre_film; + std::unique_ptr _input_mixin_post_film; + std::unique_ptr _activation_pre_film; + std::unique_ptr _activation_post_film; + std::unique_ptr _layer1x1_post_film; + std::unique_ptr _head1x1_post_film; +}; + +/// \brief An array of layers with the same channels, kernel sizes, and activations +/// +/// A LayerArray chains multiple Layer objects together, processing them sequentially. +/// Each layer processes the output of the previous layer (residual connection). +/// All layers contribute to a shared head output (skip connection) that is accumulated +/// and then projected to the final head size. +/// +/// The LayerArray handles: +/// - Input projection to match layer channel count +/// - Processing layers in sequence with residual connections +/// - Accumulating head outputs from all layers +/// - Projecting the accumulated head output to the final head size +class LayerArray +{ +public: + /// \brief Constructor with LayerArrayParams + /// \param params Parameters for constructing the layer array + LayerArray(const LayerArrayParams& params); + + /// \brief Resize all arrays to be able to process maxBufferSize frames + /// \param maxBufferSize Maximum number of frames to process in a single call + void SetMaxBufferSize(const int maxBufferSize); + + /// \brief Process without a given previous head input (first layer array) + /// + /// Zeros head accumulated output before proceeding. Used for the first layer array in a WaveNet. + /// \param layer_inputs Input to the layer array (input_size x num_frames) + /// \param condition Conditioning input (condition_size x num_frames) + /// \param num_frames Number of frames to process + void Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition, const int num_frames); + + /// \brief Process with a given previous head input (subsequent layer arrays) + /// + /// Copies head input before proceeding. Used for subsequent layer arrays that accumulate + /// head outputs from previous arrays. + /// \param layer_inputs Input to the layer array (input_size x num_frames) + /// \param condition Conditioning input (condition_size x num_frames) + /// \param head_inputs Head input from previous layer array (head_input_size x num_frames) + /// \param num_frames Number of frames to process + void Process(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition, + const Eigen::MatrixXf& head_inputs, const int num_frames); + + /// \brief Get output from last layer (for next layer array) + /// + /// Returns the full pre-allocated buffer; only the first num_frames columns + /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. + /// \return Reference to the layer output buffer (channels x maxBufferSize) + Eigen::MatrixXf& GetLayerOutputs() { return this->_layer_outputs; } + + /// \brief Get output from last layer (const version) + /// \return Const reference to the layer output buffer + const Eigen::MatrixXf& GetLayerOutputs() const { return this->_layer_outputs; } + + /// \brief Get head outputs (post head-rechannel) + /// + /// Returns the full pre-allocated buffer; only the first num_frames columns + /// are valid for a given processing call. Slice with .leftCols(num_frames) as needed. + /// \return Reference to the head output buffer (head_size x maxBufferSize) + Eigen::MatrixXf& GetHeadOutputs(); + + /// \brief Get head outputs (const version) + /// \return Const reference to the head output buffer + const Eigen::MatrixXf& GetHeadOutputs() const; + + /// \brief Set the parameters (weights) of this module + /// \param it Iterator to the weights vector. Will be advanced as weights are consumed. + void set_weights_(std::vector::iterator& it); + + /// \brief Get the "zero-indexed" receptive field + /// + /// The receptive field is the number of input samples that affect the output. + /// A 1x1 convolution is defined to have a zero-indexed receptive field of zero. + /// \return Receptive field size + long get_receptive_field() const; + +private: + // The rechannel before the layers + Conv1x1 _rechannel; + + // The layer objects + std::vector _layers; + // Output from last layer (for next layer array) + Eigen::MatrixXf _layer_outputs; + // Accumulated head inputs from all layers + // Size is _head_output_size (= head1x1.out_channels if head1x1 active, else bottleneck) + Eigen::MatrixXf _head_inputs; + + // Rechannel for the head (_head_output_size -> head_size), causal Conv1D (dilation 1) + Conv1D _head_rechannel; + + // Head output size from each layer (head1x1.out_channels if active, else bottleneck) + const int _head_output_size; + + long _get_channels() const; + // Common processing logic after head inputs are set + void ProcessInner(const Eigen::MatrixXf& layer_inputs, const Eigen::MatrixXf& condition, const int num_frames); +}; + +/// \brief Post-stack head: repeated (activation → Conv1D) with dilation 1, stride 1, valid (causal streaming) conv. +class Head +{ +public: + explicit Head(const HeadParams& params); + + void set_weights_(std::vector::iterator& weights); + void SetMaxBufferSize(int maxBufferSize); + long receptive_field() const; + int in_channels() const { return _in_channels; } + int out_channels() const { return _out_channels; } + + /// \param work Input buffer (in_channels × maxBufferSize); first in_channels×num_frames scaled by head_scale; + /// may be modified in place. + void process(Eigen::MatrixXf& work, int num_frames); + + const Eigen::MatrixXf& get_last_output() const { return _convs.back().GetOutput(); } + +private: + std::vector _convs; + std::vector _activations; + int _in_channels; + int _out_channels; +}; + +} // namespace detail +} // namespace wavenet +} // namespace nam diff --git a/NAM/wavenet.cpp b/NAM/wavenet/model.cpp similarity index 94% rename from NAM/wavenet.cpp rename to NAM/wavenet/model.cpp index f45b9f4..92d5e18 100644 --- a/NAM/wavenet.cpp +++ b/NAM/wavenet/model.cpp @@ -7,19 +7,19 @@ #include -#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++) @@ -27,10 +27,10 @@ nam::wavenet::PostStackHead::PostStackHead(const WaveNetHeadParams& params) 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); @@ -39,19 +39,19 @@ nam::wavenet::PostStackHead::PostStackHead(const WaveNetHeadParams& params) } } -void nam::wavenet::PostStackHead::set_weights_(std::vector::iterator& weights) +void nam::wavenet::detail::Head::set_weights_(std::vector::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++) @@ -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++) { @@ -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); @@ -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::iterator& weights) +void nam::wavenet::detail::Layer::set_weights_(std::vector::iterator& weights) { this->_conv.set_weights_(weights); this->_input_mixin.set_weights_(weights); @@ -159,7 +159,8 @@ void nam::wavenet::_Layer::set_weights_(std::vector::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 @@ -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) @@ -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); @@ -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++) @@ -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 @@ -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); @@ -506,18 +507,18 @@ 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::iterator& weights) +void nam::wavenet::detail::LayerArray::set_weights_(std::vector::iterator& weights) { this->_rechannel.set_weights_(weights); for (size_t i = 0; i < this->_layers.size(); i++) @@ -525,7 +526,7 @@ void nam::wavenet::_LayerArray::set_weights_(std::vector::iterator& weigh 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; } @@ -533,7 +534,7 @@ long nam::wavenet::_LayerArray::_get_channels() const namespace { int wave_net_output_channels(const std::vector& layer_array_params, - const bool with_head, const std::optional& head_params) + const bool with_head, const std::optional& head_params) { if (layer_array_params.empty()) throw std::runtime_error("WaveNet requires at least one layer array"); @@ -547,9 +548,9 @@ int wave_net_output_channels(const std::vector& nam::wavenet::WaveNet::WaveNet(const int in_channels, const std::vector& layer_array_params, - const float head_scale, const bool with_head, - std::optional head_params, std::vector weights, - std::unique_ptr condition_dsp, const double expected_sample_rate) + const float head_scale, const bool with_head, std::optional head_params, + std::vector weights, std::unique_ptr 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) @@ -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(*head_params); + this->_post_stack_head = std::make_unique(*head_params); } else if (head_params.has_value()) throw std::runtime_error("WaveNet: head configuration provided but with_head is false"); @@ -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) { @@ -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()) diff --git a/NAM/wavenet/model.h b/NAM/wavenet/model.h new file mode 100644 index 0000000..968baf8 --- /dev/null +++ b/NAM/wavenet/model.h @@ -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 +#include +#include + +#include + +#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& layer_array_params, const float head_scale, + const bool with_head, std::optional head_params, std::vector weights, + std::unique_ptr 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& 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::iterator& weights); + +protected: + // Element-wise arrays: + Eigen::MatrixXf _condition_input; + Eigen::MatrixXf _condition_output; + std::unique_ptr _condition_dsp; + // Temporary buffers for condition DSP processing (to avoid allocations in _process_condition) + std::vector> _condition_dsp_input_buffers; + std::vector> _condition_dsp_output_buffers; + std::vector _condition_dsp_input_ptrs; + std::vector _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 _layer_arrays; + + float _head_scale; + + std::unique_ptr _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 layer_array_params; + float head_scale; + bool with_head; + std::optional head_params; + std::unique_ptr 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 create(std::vector 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 create_config(const nlohmann::json& config, double sampleRate); + +} // namespace wavenet +} // namespace nam diff --git a/NAM/wavenet/params.h b/NAM/wavenet/params.h new file mode 100644 index 0000000..3ac3849 --- /dev/null +++ b/NAM/wavenet/params.h @@ -0,0 +1,316 @@ +#pragma once + +#include +#include +#include + +#include "../activations.h" + +namespace nam +{ +namespace wavenet +{ + +/// \brief Gating mode for WaveNet layers +/// +/// Determines how the layer processes the doubled bottleneck channels when gating is enabled. +enum class GatingMode +{ + NONE, ///< No gating or blending - standard activation + GATED, ///< Traditional gating (element-wise multiplication of activated channels) + BLENDED ///< Blending (weighted average between activated and pre-activated values) +}; + +/// \brief Helper function for backward compatibility with boolean gated parameter +/// \param gated Boolean indicating whether gating should be enabled +/// \return GatingMode::GATED if gated is true, GatingMode::NONE otherwise +inline GatingMode gating_mode_from_bool(bool gated) +{ + return gated ? GatingMode::GATED : GatingMode::NONE; +} + +/// \brief Parameters for head1x1 configuration +/// +/// Configures an optional 1x1 convolution that outputs directly to the head (skip connection) +/// instead of using the activation output directly. +struct Head1x1Params +{ + /// \brief Constructor + /// \param active_ Whether the head1x1 convolution is active + /// \param out_channels_ Number of output channels for the head1x1 convolution + /// \param groups_ Number of groups for grouped convolution + Head1x1Params(bool active_, int out_channels_, int groups_) + : active(active_) + , out_channels(out_channels_) + , groups(groups_) + { + } + + const bool active; ///< Whether the head1x1 convolution is active + const int out_channels; ///< Number of output channels + const int groups; ///< Number of groups for grouped convolution +}; + +/// \brief Parameters for layer1x1 configuration +/// +/// Configures an optional 1x1 convolution that processes the activation output +/// for the residual connection to the next layer. +struct Layer1x1Params +{ + /// \brief Constructor + /// \param active_ Whether the layer1x1 convolution is active + /// \param groups_ Number of groups for grouped convolution + Layer1x1Params(bool active_, int groups_) + : active(active_) + , groups(groups_) + { + } + + const bool active; ///< Whether the layer1x1 convolution is active + const int groups; ///< Number of groups for grouped convolution +}; + +/// \brief Parameters for FiLM (Feature-wise Linear Modulation) configuration +/// +/// FiLM applies per-channel scaling and optional shifting based on conditioning input. +struct _FiLMParams +{ + /// \brief Constructor + /// \param active_ Whether FiLM is active at this location + /// \param shift_ Whether to apply both scale and shift (true) or only scale (false) + /// \param groups_ Number of groups for grouped convolution in the condition-to-scale-shift submodule (default: 1) + _FiLMParams(bool active_, bool shift_, int groups_ = 1) + : active(active_) + , shift(shift_) + , groups(groups_) + { + } + const bool active; ///< Whether FiLM is active + const bool shift; ///< Whether to apply shift in addition to scale + const int groups; ///< Number of groups for grouped convolution in the condition-to-scale-shift submodule +}; + +/// \brief Parameters for constructing a single Layer +/// +/// Contains all configuration needed to construct a detail::Layer +struct LayerParams +{ + /// \brief Constructor + /// \param condition_size_ Size of the conditioning input + /// \param channels_ Number of input/output channels from layer to layer + /// \param bottleneck_ Internal channel count + /// \param kernel_size_ Kernel size for the dilated convolution + /// \param dilation_ Dilation factor for the convolution + /// \param activation_config_ Primary activation function configuration + /// \param gating_mode_ Gating mode (NONE, GATED, or BLENDED) + /// \param groups_input_ Number of groups for the input convolution + /// \param groups_input_mixin_ Number of groups for the input mixin convolution + /// \param layer1x1_params_ Configuration of the optional layer1x1 convolution + /// \param head1x1_params_ Configuration of the optional head1x1 convolution + /// \param secondary_activation_config_ Secondary activation (for gating/blending) + /// \param conv_pre_film_params_ FiLM parameters before the input convolution + /// \param conv_post_film_params_ FiLM parameters after the input convolution + /// \param input_mixin_pre_film_params_ FiLM parameters before the input mixin + /// \param input_mixin_post_film_params_ FiLM parameters after the input mixin + /// \param activation_pre_film_params_ FiLM parameters after the input/mixin summed output before activation + /// \param activation_post_film_params_ FiLM parameters after the activation output before the layer1x1 convolution + /// \param _layer1x1_post_film_params_ FiLM parameters after the layer1x1 convolution + /// \param head1x1_post_film_params_ FiLM parameters after the head1x1 convolution + LayerParams(const int condition_size_, const int channels_, const int bottleneck_, const int kernel_size_, + const int dilation_, const activations::ActivationConfig& activation_config_, + const GatingMode gating_mode_, const int groups_input_, const int groups_input_mixin_, + const Layer1x1Params& layer1x1_params_, const Head1x1Params& head1x1_params_, + const activations::ActivationConfig& secondary_activation_config_, + const _FiLMParams& conv_pre_film_params_, const _FiLMParams& conv_post_film_params_, + const _FiLMParams& input_mixin_pre_film_params_, const _FiLMParams& input_mixin_post_film_params_, + const _FiLMParams& activation_pre_film_params_, const _FiLMParams& activation_post_film_params_, + const _FiLMParams& _layer1x1_post_film_params_, const _FiLMParams& head1x1_post_film_params_) + : condition_size(condition_size_) + , channels(channels_) + , bottleneck(bottleneck_) + , kernel_size(kernel_size_) + , dilation(dilation_) + , activation_config(activation_config_) + , gating_mode(gating_mode_) + , groups_input(groups_input_) + , groups_input_mixin(groups_input_mixin_) + , layer1x1_params(layer1x1_params_) + , head1x1_params(head1x1_params_) + , secondary_activation_config(secondary_activation_config_) + , conv_pre_film_params(conv_pre_film_params_) + , conv_post_film_params(conv_post_film_params_) + , input_mixin_pre_film_params(input_mixin_pre_film_params_) + , input_mixin_post_film_params(input_mixin_post_film_params_) + , activation_pre_film_params(activation_pre_film_params_) + , activation_post_film_params(activation_post_film_params_) + , _layer1x1_post_film_params(_layer1x1_post_film_params_) + , head1x1_post_film_params(head1x1_post_film_params_) + { + } + + const int condition_size; ///< Size of the conditioning input + const int channels; ///< Number of input/output channels from layer to layer + const int bottleneck; ///< Internal channel count + const int kernel_size; ///< Kernel size for the dilated convolution + const int dilation; ///< Dilation factor for the convolution + const activations::ActivationConfig activation_config; ///< Primary activation function configuration + const GatingMode gating_mode; ///< Gating mode (NONE, GATED, or BLENDED) + const int groups_input; ///< Number of groups for the input convolution + const int groups_input_mixin; ///< Number of groups for the input mixin convolution + const Layer1x1Params layer1x1_params; ///< Configuration of the optional layer1x1 convolution + const Head1x1Params head1x1_params; ///< Configuration of the optional head1x1 convolution + const activations::ActivationConfig secondary_activation_config; ///< Secondary activation (for gating/blending) + const _FiLMParams conv_pre_film_params; ///< FiLM parameters before the input convolution + const _FiLMParams conv_post_film_params; ///< FiLM parameters after the input convolution + const _FiLMParams input_mixin_pre_film_params; ///< FiLM parameters before the input mixin + const _FiLMParams input_mixin_post_film_params; ///< FiLM parameters after the input mixin + const _FiLMParams activation_pre_film_params; ///< FiLM parameters before activation + const _FiLMParams activation_post_film_params; ///< FiLM parameters after activation + const _FiLMParams _layer1x1_post_film_params; ///< FiLM parameters after the layer1x1 convolution (layer1x1_post_film) + const _FiLMParams head1x1_post_film_params; ///< FiLM parameters after the head1x1 convolution +}; + +/// \brief Parameters for constructing a LayerArray +/// +/// Contains all configuration needed to construct a detail::LayerArray with multiple layers +/// sharing the same channel count and kernel size. Each layer can have its own activation configuration. +class LayerArrayParams +{ +public: + /// \brief Constructor + /// \param input_size_ Input size (number of channels) to the layer array + /// \param condition_size_ Size of the conditioning input + /// \param head_size_ Size of the head output (after head rechannel) + /// \param channels_ Number of channels in each layer + /// \param bottleneck_ Bottleneck size (internal channel count) + /// \param kernel_sizes_ Per-layer kernel sizes, one per layer + /// \param dilations_ Vector of dilation factors, one per layer + /// \param activation_configs_ Vector of primary activation configurations, one per layer + /// \param gating_modes_ Vector of gating modes, one per layer + /// \param head_kernel_size_ Kernel size of the head rechannel conv (>= 1) + /// \param head_bias_ Whether to use bias in the head rechannel + /// \param groups_input Number of groups for input convolutions + /// \param groups_input_mixin_ Number of groups for input mixin convolutions + /// \param layer1x1_params_ Parameters for optional layer1x1 convolutions + /// \param head1x1_params_ Parameters for optional head1x1 convolutions + /// \param secondary_activation_configs_ Vector of secondary activation configs for gating/blending, one per layer + /// \param conv_pre_film_params_ FiLM parameters before input convolutions + /// \param conv_post_film_params_ FiLM parameters after input convolutions + /// \param input_mixin_pre_film_params_ FiLM parameters before input mixin + /// \param input_mixin_post_film_params_ FiLM parameters after input mixin + /// \param activation_pre_film_params_ FiLM parameters before activation + /// \param activation_post_film_params_ FiLM parameters after activation + /// \param _layer1x1_post_film_params_ FiLM parameters after layer1x1 convolutions + /// \param head1x1_post_film_params_ FiLM parameters after head1x1 convolutions + /// \throws std::invalid_argument If dilations, activation_configs, gating_modes, or secondary_activation_configs + /// sizes don't match + LayerArrayParams(const int input_size_, const int condition_size_, const int head_size_, const int head_kernel_size_, + const int channels_, const int bottleneck_, const std::vector&& kernel_sizes_, + const std::vector&& dilations_, + const std::vector&& activation_configs_, + const std::vector&& gating_modes_, const bool head_bias_, const int groups_input, + const int groups_input_mixin_, const Layer1x1Params& layer1x1_params_, + const Head1x1Params& head1x1_params_, + const std::vector&& secondary_activation_configs_, + const _FiLMParams& conv_pre_film_params_, const _FiLMParams& conv_post_film_params_, + const _FiLMParams& input_mixin_pre_film_params_, const _FiLMParams& input_mixin_post_film_params_, + const _FiLMParams& activation_pre_film_params_, const _FiLMParams& activation_post_film_params_, + const _FiLMParams& _layer1x1_post_film_params_, const _FiLMParams& head1x1_post_film_params_) + : input_size(input_size_) + , condition_size(condition_size_) + , head_size(head_size_) + , head_kernel_size(head_kernel_size_) + , channels(channels_) + , bottleneck(bottleneck_) + , kernel_sizes(std::move(kernel_sizes_)) + , dilations(std::move(dilations_)) + , activation_configs(std::move(activation_configs_)) + , gating_modes(std::move(gating_modes_)) + , head_bias(head_bias_) + , groups_input(groups_input) + , groups_input_mixin(groups_input_mixin_) + , layer1x1_params(layer1x1_params_) + , head1x1_params(head1x1_params_) + , secondary_activation_configs(std::move(secondary_activation_configs_)) + , conv_pre_film_params(conv_pre_film_params_) + , conv_post_film_params(conv_post_film_params_) + , input_mixin_pre_film_params(input_mixin_pre_film_params_) + , input_mixin_post_film_params(input_mixin_post_film_params_) + , activation_pre_film_params(activation_pre_film_params_) + , activation_post_film_params(activation_post_film_params_) + , _layer1x1_post_film_params(_layer1x1_post_film_params_) + , head1x1_post_film_params(head1x1_post_film_params_) + { + if (head_kernel_size < 1) + { + throw std::invalid_argument("LayerArrayParams: head_kernel_size must be >= 1"); + } + const size_t num_layers = dilations.size(); + if (kernel_sizes.empty()) + { + throw std::invalid_argument("LayerArrayParams: kernel_sizes must not be empty"); + } + if (kernel_sizes.size() != num_layers) + { + throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) + + ") must match kernel_sizes size (" + std::to_string(kernel_sizes.size()) + ")"); + } + if (activation_configs.size() != num_layers) + { + throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) + + ") must match activation_configs size (" + std::to_string(activation_configs.size()) + + ")"); + } + if (gating_modes.size() != num_layers) + { + throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) + + ") must match gating_modes size (" + std::to_string(gating_modes.size()) + ")"); + } + if (secondary_activation_configs.size() != num_layers) + { + throw std::invalid_argument("LayerArrayParams: dilations size (" + std::to_string(num_layers) + + ") must match secondary_activation_configs size (" + + std::to_string(secondary_activation_configs.size()) + ")"); + } + } + + const int input_size; ///< Input size (number of channels) + const int condition_size; ///< Size of conditioning input + const int head_size; ///< Size of head output (after rechannel) + const int head_kernel_size; ///< Kernel size of head rechannel convolution (>= 1) + const int channels; ///< Number of channels in each layer + const int bottleneck; ///< Bottleneck size (internal channel count) + std::vector kernel_sizes; ///< Per-layer kernel sizes, one per layer + std::vector dilations; ///< Dilation factors, one per layer + std::vector activation_configs; ///< Primary activation configurations, one per layer + std::vector gating_modes; ///< Gating modes, one per layer + const bool head_bias; ///< Whether to use bias in head rechannel + const int groups_input; ///< Number of groups for input convolutions + const int groups_input_mixin; ///< Number of groups for input mixin + const Layer1x1Params layer1x1_params; ///< Parameters for optional layer1x1 + const Head1x1Params head1x1_params; ///< Parameters for optional head1x1 + std::vector + secondary_activation_configs; ///< Secondary activation configs for gating/blending, one per layer + const _FiLMParams conv_pre_film_params; ///< FiLM params before input conv + const _FiLMParams conv_post_film_params; ///< FiLM params after input conv + const _FiLMParams input_mixin_pre_film_params; ///< FiLM params before input mixin + const _FiLMParams input_mixin_post_film_params; ///< FiLM params after input mixin + const _FiLMParams activation_pre_film_params; ///< FiLM params before activation + const _FiLMParams activation_post_film_params; ///< FiLM params after activation + const _FiLMParams _layer1x1_post_film_params; ///< FiLM params after layer1x1 conv + const _FiLMParams head1x1_post_film_params; ///< FiLM params after head1x1 conv +}; + +/// \brief Parameters for the optional post-stack head (matches Python ``nam.models.wavenet._head.Head``). +/// JSON export omits ``in_channels`` (implied by last layer array ``head_size``); load sets it from there. +struct HeadParams +{ + int in_channels; + int channels; + int out_channels; + std::vector kernel_sizes; + activations::ActivationConfig activation_config; +}; + +} // namespace wavenet +} // namespace nam diff --git a/NAM/slimmable_wavenet.cpp b/NAM/wavenet/slimmable.cpp similarity index 99% rename from NAM/slimmable_wavenet.cpp rename to NAM/wavenet/slimmable.cpp index 5be26e2..603239a 100644 --- a/NAM/slimmable_wavenet.cpp +++ b/NAM/wavenet/slimmable.cpp @@ -1,5 +1,5 @@ -#include "slimmable_wavenet.h" -#include "get_dsp.h" +#include "slimmable.h" +#include "../get_dsp.h" #include #include diff --git a/NAM/slimmable_wavenet.h b/NAM/wavenet/slimmable.h similarity index 96% rename from NAM/slimmable_wavenet.h rename to NAM/wavenet/slimmable.h index 5fb28f4..a798108 100644 --- a/NAM/slimmable_wavenet.h +++ b/NAM/wavenet/slimmable.h @@ -3,11 +3,11 @@ #include #include -#include "dsp.h" +#include "../dsp.h" #include "json.hpp" -#include "model_config.h" -#include "slimmable.h" -#include "wavenet.h" +#include "../model_config.h" +#include "../slimmable.h" +#include "model.h" namespace nam { diff --git a/docs/api/wavenet.rst b/docs/api/wavenet.rst index 571c7e4..5a973a3 100644 --- a/docs/api/wavenet.rst +++ b/docs/api/wavenet.rst @@ -9,11 +9,11 @@ WaveNet API :project: NeuralAmpModelerCore :members: -.. doxygenclass:: nam::wavenet::_LayerArray +.. doxygenclass:: nam::wavenet::detail::LayerArray :project: NeuralAmpModelerCore :members: -.. doxygenclass:: nam::wavenet::_Layer +.. doxygenclass:: nam::wavenet::detail::Layer :project: NeuralAmpModelerCore :members: diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 37164d6..9cadccd 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -1,4 +1,6 @@ -file(GLOB_RECURSE NAM_SOURCES ../NAM/*.cpp ../NAM/*.c ../NAM*.h) +file(GLOB NAM_SOURCES_TOP "${CMAKE_CURRENT_SOURCE_DIR}/../NAM/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/../NAM/*.c") +file(GLOB NAM_SOURCES_SUB "${CMAKE_CURRENT_SOURCE_DIR}/../NAM/*/*.cpp") +set(NAM_SOURCES ${NAM_SOURCES_TOP} ${NAM_SOURCES_SUB}) # TODO: add loadmodel and run_tests to TOOLS? set(TOOLS benchmodel) diff --git a/tools/test/test_wavenet/test_condition_processing.cpp b/tools/test/test_wavenet/test_condition_processing.cpp index f929483..8f38cc3 100644 --- a/tools/test/test_wavenet/test_condition_processing.cpp +++ b/tools/test/test_wavenet/test_condition_processing.cpp @@ -7,7 +7,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" #include "NAM/dsp.h" namespace test_wavenet diff --git a/tools/test/test_wavenet/test_factory.cpp b/tools/test/test_wavenet/test_factory.cpp index ed06c22..9c13cf6 100644 --- a/tools/test/test_wavenet/test_factory.cpp +++ b/tools/test/test_wavenet/test_factory.cpp @@ -8,7 +8,7 @@ #include "json.hpp" #include "NAM/get_dsp.h" -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { diff --git a/tools/test/test_wavenet/test_full.cpp b/tools/test/test_wavenet/test_full.cpp index 0c09c2e..abd9a33 100644 --- a/tools/test/test_wavenet/test_full.cpp +++ b/tools/test/test_wavenet/test_full.cpp @@ -7,7 +7,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { diff --git a/tools/test/test_wavenet/test_head1x1.cpp b/tools/test/test_wavenet/test_head1x1.cpp index 8bb31b3..6821a27 100644 --- a/tools/test/test_wavenet/test_head1x1.cpp +++ b/tools/test/test_wavenet/test_head1x1.cpp @@ -6,7 +6,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -19,13 +19,13 @@ static nam::wavenet::_FiLMParams make_default_film_params() } // Helper function to create a Layer with default FiLM parameters -static nam::wavenet::_Layer make_layer(const int condition_size, const int channels, const int bottleneck, - const int kernel_size, const int dilation, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const int groups_input, - const int groups_input_mixin, const int groups_1x1, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::Layer make_layer(const int condition_size, const int channels, const int bottleneck, + const int kernel_size, const int dilation, + const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const int groups_input, + const int groups_input_mixin, const int groups_1x1, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); // Create layer1x1_params with active=true and groups=groups_1x1 for backward compatibility @@ -34,7 +34,7 @@ static nam::wavenet::_Layer make_layer(const int condition_size, const int chann gating_mode, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, secondary_activation_config, film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_Layer(layer_params); + return nam::wavenet::detail::Layer(layer_params); } void test_head1x1_inactive() diff --git a/tools/test/test_wavenet/test_kernel_sizes.cpp b/tools/test/test_wavenet/test_kernel_sizes.cpp index e266981..a0a5918 100644 --- a/tools/test/test_wavenet/test_kernel_sizes.cpp +++ b/tools/test/test_wavenet/test_kernel_sizes.cpp @@ -5,7 +5,7 @@ #include "json.hpp" -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -55,7 +55,7 @@ void test_kernel_size_int_compat() } // Verify that receptive field computation uses the (uniform) kernel size - nam::wavenet::_LayerArray array(p); + nam::wavenet::detail::LayerArray array(p); const long receptive_field = array.get_receptive_field(); long expected_receptive_field = 0; @@ -104,7 +104,7 @@ void test_kernel_size_per_layer_array() assert(p.dilations == expected_dilations); // Verify that receptive field computation uses the per-layer kernel sizes - nam::wavenet::_LayerArray array(p); + nam::wavenet::detail::LayerArray array(p); const long receptive_field = array.get_receptive_field(); long expected_receptive_field = 0; diff --git a/tools/test/test_wavenet/test_layer.cpp b/tools/test/test_wavenet/test_layer.cpp index 4494781..b30cd93 100644 --- a/tools/test/test_wavenet/test_layer.cpp +++ b/tools/test/test_wavenet/test_layer.cpp @@ -6,7 +6,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -19,13 +19,13 @@ static nam::wavenet::_FiLMParams make_default_film_params() } // Helper function to create a Layer with default FiLM parameters -static nam::wavenet::_Layer make_layer(const int condition_size, const int channels, const int bottleneck, - const int kernel_size, const int dilation, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const int groups_input, - const int groups_input_mixin, const int groups_1x1, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::Layer make_layer(const int condition_size, const int channels, const int bottleneck, + const int kernel_size, const int dilation, + const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const int groups_input, + const int groups_input_mixin, const int groups_1x1, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); // Create layer1x1_params with active=true and groups=groups_1x1 for backward compatibility @@ -34,7 +34,7 @@ static nam::wavenet::_Layer make_layer(const int condition_size, const int chann gating_mode, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, secondary_activation_config, film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_Layer(layer_params); + return nam::wavenet::detail::Layer(layer_params); } void test_gated() { diff --git a/tools/test/test_wavenet/test_layer1x1.cpp b/tools/test/test_wavenet/test_layer1x1.cpp index 889f210..8af466c 100644 --- a/tools/test/test_wavenet/test_layer1x1.cpp +++ b/tools/test/test_wavenet/test_layer1x1.cpp @@ -7,7 +7,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -20,21 +20,21 @@ static nam::wavenet::_FiLMParams make_default_film_params() } // Helper function to create a Layer with default FiLM parameters -static nam::wavenet::_Layer make_layer(const int condition_size, const int channels, const int bottleneck, - const int kernel_size, const int dilation, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const int groups_input, - const int groups_input_mixin, - const nam::wavenet::Layer1x1Params& layer1x1_params, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::Layer make_layer(const int condition_size, const int channels, const int bottleneck, + const int kernel_size, const int dilation, + const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const int groups_input, + const int groups_input_mixin, + const nam::wavenet::Layer1x1Params& layer1x1_params, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); nam::wavenet::LayerParams layer_params(condition_size, channels, bottleneck, kernel_size, dilation, activation_config, gating_mode, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, secondary_activation_config, film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_Layer(layer_params); + return nam::wavenet::detail::Layer(layer_params); } void test_layer1x1_active() @@ -198,7 +198,7 @@ void test_layer1x1_inactive_bottleneck_mismatch() bool threw_exception = false; try { - auto layer = nam::wavenet::_Layer(layer_params); + auto layer = nam::wavenet::detail::Layer(layer_params); } catch (const std::invalid_argument& e) { @@ -236,7 +236,7 @@ void test_layer1x1_post_film_active() nam::activations::ActivationConfig{}, film_params, film_params, film_params, film_params, film_params, film_params, layer1x1_post_film_params, film_params); - auto layer = nam::wavenet::_Layer(layer_params); + auto layer = nam::wavenet::detail::Layer(layer_params); // Set weights: conv, input_mixin, layer1x1, layer1x1_post_film // With bottleneck=channels=2: @@ -314,7 +314,7 @@ void test_layer1x1_post_film_inactive_with_layer1x1_inactive() bool threw_exception = false; try { - auto layer = nam::wavenet::_Layer(layer_params); + auto layer = nam::wavenet::detail::Layer(layer_params); } catch (const std::invalid_argument& e) { diff --git a/tools/test/test_wavenet/test_layer_array.cpp b/tools/test/test_wavenet/test_layer_array.cpp index 0dc1fde..9ffc4da 100644 --- a/tools/test/test_wavenet/test_layer_array.cpp +++ b/tools/test/test_wavenet/test_layer_array.cpp @@ -6,7 +6,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -19,15 +19,13 @@ static nam::wavenet::_FiLMParams make_default_film_params() } // Helper function to create a LayerArray with default FiLM parameters -static nam::wavenet::_LayerArray make_layer_array(const int input_size, const int condition_size, const int head_size, - const int channels, const int bottleneck, const int kernel_size, - const std::vector& dilations, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const bool head_bias, - const int groups_input, const int groups_input_mixin, - const nam::wavenet::Layer1x1Params& layer1x1_params, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::LayerArray make_layer_array( + const int input_size, const int condition_size, const int head_size, const int channels, const int bottleneck, + const int kernel_size, const std::vector& dilations, const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const bool head_bias, const int groups_input, + const int groups_input_mixin, const nam::wavenet::Layer1x1Params& layer1x1_params, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); // Duplicate activation_config, gating_mode, and secondary_activation_config for each layer (based on dilations size) @@ -42,7 +40,7 @@ static nam::wavenet::_LayerArray make_layer_array(const int input_size, const in std::move(activation_configs), std::move(gating_modes), head_bias, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, std::move(secondary_activation_configs), film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_LayerArray(params); + return nam::wavenet::detail::LayerArray(params); } // Test layer array construction and basic processing void test_layer_array_basic() @@ -227,7 +225,7 @@ void test_layer_array_different_activations() std::move(activation_configs), std::move(gating_modes), head_bias, groups, groups_input_mixin, layer1x1_params, head1x1_params, std::move(secondary_activation_configs), film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - nam::wavenet::_LayerArray layer_array(params); + nam::wavenet::detail::LayerArray layer_array(params); const int numFrames = 4; layer_array.SetMaxBufferSize(numFrames); @@ -310,7 +308,7 @@ void test_layer_array_different_activations() std::move(dilations_all_relu), std::move(all_relu_configs), std::move(all_none_gating_modes), head_bias, groups, groups_input_mixin, layer1x1_params, head1x1_params, std::move(all_empty_secondary_configs), film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - nam::wavenet::_LayerArray layer_array_all_relu(params_all_relu); + nam::wavenet::detail::LayerArray layer_array_all_relu(params_all_relu); layer_array_all_relu.SetMaxBufferSize(numFrames); // Create weights for all-NONE version (simpler, no gating) diff --git a/tools/test/test_wavenet/test_layer_head_config.cpp b/tools/test/test_wavenet/test_layer_head_config.cpp index 847810a..9eac89b 100644 --- a/tools/test/test_wavenet/test_layer_head_config.cpp +++ b/tools/test/test_wavenet/test_layer_head_config.cpp @@ -5,7 +5,7 @@ #include "json.hpp" -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -60,7 +60,7 @@ void test_nested_head_with_kernel_size_three() assert(p.head_kernel_size == 3); assert(p.head_bias == true); - nam::wavenet::_LayerArray array(p); + nam::wavenet::detail::LayerArray array(p); assert(array.get_receptive_field() == 2); // one dilated layer: 0 + (3-1) head rechannel } diff --git a/tools/test/test_wavenet/test_output_head.cpp b/tools/test/test_wavenet/test_output_head.cpp index 4762df5..5998a2a 100644 --- a/tools/test/test_wavenet/test_output_head.cpp +++ b/tools/test/test_wavenet/test_output_head.cpp @@ -7,7 +7,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" namespace test_wavenet { @@ -41,13 +41,13 @@ static nam::wavenet::LayerArrayParams make_layer_array_params( void test_post_stack_head_receptive_field() { - nam::wavenet::WaveNetHeadParams p; + nam::wavenet::HeadParams p; p.in_channels = 2; p.channels = 3; p.out_channels = 1; p.kernel_sizes = {3, 5}; p.activation_config = nam::activations::ActivationConfig::simple(nam::activations::ActivationType::Tanh); - nam::wavenet::PostStackHead head(p); + nam::wavenet::detail::Head head(p); // Python: 1 + (3-1) + (5-1) = 7 assert(head.receptive_field() == 7); } @@ -78,7 +78,7 @@ void test_wavenet_with_post_stack_head_processes() std::vector layer_array_params; layer_array_params.push_back(std::move(layer_params)); - nam::wavenet::WaveNetHeadParams hp; + nam::wavenet::HeadParams hp; hp.in_channels = 1; hp.channels = 1; hp.out_channels = 1; @@ -95,7 +95,7 @@ void test_wavenet_with_post_stack_head_processes() std::unique_ptr condition_dsp = nullptr; auto wavenet = std::make_unique(input_size, layer_array_params, head_scale, with_head, - std::optional(std::move(hp)), + std::optional(std::move(hp)), std::move(weights), std::move(condition_dsp), 48000.0); const int numFrames = 8; @@ -143,7 +143,7 @@ void test_wavenet_with_two_layer_post_stack_head_applies_activation_per_layer_in std::vector layer_array_params; layer_array_params.push_back(std::move(layer_params)); - nam::wavenet::WaveNetHeadParams hp; + nam::wavenet::HeadParams hp; hp.in_channels = 1; hp.channels = 1; hp.out_channels = 1; @@ -167,7 +167,7 @@ void test_wavenet_with_two_layer_post_stack_head_applies_activation_per_layer_in std::unique_ptr condition_dsp = nullptr; auto wavenet = std::make_unique(input_size, layer_array_params, head_scale, with_head, - std::optional(std::move(hp)), + std::optional(std::move(hp)), std::move(weights), std::move(condition_dsp), 48000.0); const int numFrames = 8; diff --git a/tools/test/test_wavenet/test_real_time_safe.cpp b/tools/test/test_wavenet/test_real_time_safe.cpp index b590551..9ed29a8 100644 --- a/tools/test/test_wavenet/test_real_time_safe.cpp +++ b/tools/test/test_wavenet/test_real_time_safe.cpp @@ -9,7 +9,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" #include "NAM/conv1d.h" #include "../allocation_tracking.h" @@ -24,33 +24,31 @@ static nam::wavenet::_FiLMParams make_default_film_params() } // Helper function to create a Layer with default FiLM parameters -static nam::wavenet::_Layer make_layer(const int condition_size, const int channels, const int bottleneck, - const int kernel_size, const int dilation, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const int groups_input, - const int groups_input_mixin, - const nam::wavenet::Layer1x1Params& layer1x1_params, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::Layer make_layer(const int condition_size, const int channels, const int bottleneck, + const int kernel_size, const int dilation, + const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const int groups_input, + const int groups_input_mixin, + const nam::wavenet::Layer1x1Params& layer1x1_params, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); nam::wavenet::LayerParams layer_params(condition_size, channels, bottleneck, kernel_size, dilation, activation_config, gating_mode, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, secondary_activation_config, film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_Layer(layer_params); + return nam::wavenet::detail::Layer(layer_params); } // Helper function to create a LayerArray with default FiLM parameters -static nam::wavenet::_LayerArray make_layer_array(const int input_size, const int condition_size, const int head_size, - const int channels, const int bottleneck, const int kernel_size, - const std::vector& dilations, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const bool head_bias, - const int groups_input, const int groups_input_mixin, - const nam::wavenet::Layer1x1Params& layer1x1_params, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::LayerArray make_layer_array( + const int input_size, const int condition_size, const int head_size, const int channels, const int bottleneck, + const int kernel_size, const std::vector& dilations, const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const bool head_bias, const int groups_input, + const int groups_input_mixin, const nam::wavenet::Layer1x1Params& layer1x1_params, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); // Duplicate activation_config, gating_mode, and secondary_activation_config for each layer (based on dilations size) @@ -65,7 +63,7 @@ static nam::wavenet::_LayerArray make_layer_array(const int input_size, const in std::move(activation_configs), std::move(gating_modes), head_bias, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, std::move(secondary_activation_configs), film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_LayerArray(params); + return nam::wavenet::detail::LayerArray(params); } // Helper function to create LayerArrayParams with default FiLM parameters @@ -91,7 +89,7 @@ static nam::wavenet::LayerArrayParams make_layer_array_params( } // Helper function to create a Layer with all FiLMs active -static nam::wavenet::_Layer make_layer_all_films( +static nam::wavenet::detail::Layer make_layer_all_films( const int condition_size, const int channels, const int bottleneck, const int kernel_size, const int dilation, const nam::activations::ActivationConfig& activation_config, const nam::wavenet::GatingMode gating_mode, const int groups_input, const int groups_input_mixin, const nam::wavenet::Layer1x1Params& layer1x1_params, @@ -106,7 +104,7 @@ static nam::wavenet::_Layer make_layer_all_films( gating_mode, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, secondary_activation_config, film_params, film_params, film_params, film_params, film_params, film_params, film_params, head1x1_post_film_params); - return nam::wavenet::_Layer(layer_params); + return nam::wavenet::detail::Layer(layer_params); } // Test that pre-allocated Eigen operations with noalias() don't allocate @@ -739,7 +737,7 @@ void test_layer_post_activation_film_gated_realtime_safe() inactive_film, // _1x1_post_film inactive_film // head1x1_post_film ); - auto layer = nam::wavenet::_Layer(layer_params); + auto layer = nam::wavenet::detail::Layer(layer_params); // Set weights - Order: conv, input_mixin, 1x1, then FiLMs // NOTE: In GATED mode, conv and input_mixin output 2*bottleneck channels! @@ -846,7 +844,7 @@ void test_layer_post_activation_film_blended_realtime_safe() inactive_film, // _1x1_post_film inactive_film // head1x1_post_film ); - auto layer = nam::wavenet::_Layer(layer_params); + auto layer = nam::wavenet::detail::Layer(layer_params); // Set weights - Order: conv, input_mixin, 1x1, then FiLMs // NOTE: In BLENDED mode, conv and input_mixin output 2*bottleneck channels! @@ -1196,7 +1194,7 @@ void test_process_3in_2out_realtime_safe() } } -// WaveNet::process() with optional post-stack head (multi-layer PostStackHead) must not allocate or free. +// WaveNet::process() with optional post-stack head (multi-layer detail::Head) must not allocate or free. void test_process_with_post_stack_head_realtime_safe() { const int input_size = 1; @@ -1223,7 +1221,7 @@ void test_process_with_post_stack_head_realtime_safe() std::move(dilations), activation, gating_mode, head_bias, groups, groups_input_mixin, layer1x1_params, head1x1_params, nam::activations::ActivationConfig{})); - nam::wavenet::WaveNetHeadParams head_params; + nam::wavenet::HeadParams head_params; head_params.in_channels = 1; head_params.channels = 1; head_params.out_channels = 1; @@ -1243,7 +1241,7 @@ void test_process_with_post_stack_head_realtime_safe() std::unique_ptr condition_dsp = nullptr; auto wavenet = std::make_unique(input_size, layer_array_params, head_scale, with_head, - std::optional(std::move(head_params)), + std::optional(std::move(head_params)), std::move(weights), std::move(condition_dsp), 48000.0); const int maxBufferSize = 256; diff --git a/tools/test/test_wavenet_configurable_gating.cpp b/tools/test/test_wavenet_configurable_gating.cpp index 62d5a17..5df3ee6 100644 --- a/tools/test/test_wavenet_configurable_gating.cpp +++ b/tools/test/test_wavenet_configurable_gating.cpp @@ -4,7 +4,7 @@ #include #include -#include "NAM/wavenet.h" +#include "NAM/wavenet/model.h" #include "NAM/gating_activations.h" namespace test_wavenet_configurable_gating @@ -16,21 +16,21 @@ static nam::wavenet::_FiLMParams make_default_film_params() } // Helper function to create a Layer with default FiLM parameters -static nam::wavenet::_Layer make_layer(const int condition_size, const int channels, const int bottleneck, - const int kernel_size, const int dilation, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const int groups_input, - const int groups_input_mixin, - const nam::wavenet::Layer1x1Params& layer1x1_params, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::Layer make_layer(const int condition_size, const int channels, const int bottleneck, + const int kernel_size, const int dilation, + const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const int groups_input, + const int groups_input_mixin, + const nam::wavenet::Layer1x1Params& layer1x1_params, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); nam::wavenet::LayerParams layer_params(condition_size, channels, bottleneck, kernel_size, dilation, activation_config, gating_mode, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, secondary_activation_config, film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_Layer(layer_params); + return nam::wavenet::detail::Layer(layer_params); } // Helper function to create LayerArrayParams with default FiLM parameters @@ -56,15 +56,13 @@ static nam::wavenet::LayerArrayParams make_layer_array_params( } // Helper function to create a LayerArray with default FiLM parameters -static nam::wavenet::_LayerArray make_layer_array(const int input_size, const int condition_size, const int head_size, - const int channels, const int bottleneck, const int kernel_size, - const std::vector& dilations, - const nam::activations::ActivationConfig& activation_config, - const nam::wavenet::GatingMode gating_mode, const bool head_bias, - const int groups_input, const int groups_input_mixin, - const nam::wavenet::Layer1x1Params& layer1x1_params, - const nam::wavenet::Head1x1Params& head1x1_params, - const nam::activations::ActivationConfig& secondary_activation_config) +static nam::wavenet::detail::LayerArray make_layer_array( + const int input_size, const int condition_size, const int head_size, const int channels, const int bottleneck, + const int kernel_size, const std::vector& dilations, const nam::activations::ActivationConfig& activation_config, + const nam::wavenet::GatingMode gating_mode, const bool head_bias, const int groups_input, + const int groups_input_mixin, const nam::wavenet::Layer1x1Params& layer1x1_params, + const nam::wavenet::Head1x1Params& head1x1_params, + const nam::activations::ActivationConfig& secondary_activation_config) { auto film_params = make_default_film_params(); // Duplicate activation_config, gating_mode, and secondary_activation_config for each layer (based on dilations size) @@ -79,7 +77,7 @@ static nam::wavenet::_LayerArray make_layer_array(const int input_size, const in std::move(activation_configs), std::move(gating_modes), head_bias, groups_input, groups_input_mixin, layer1x1_params, head1x1_params, std::move(secondary_activation_configs), film_params, film_params, film_params, film_params, film_params, film_params, film_params, film_params); - return nam::wavenet::_LayerArray(params); + return nam::wavenet::detail::LayerArray(params); } class TestConfigurableGating @@ -198,7 +196,7 @@ class TestConfigurableGating static void test_layer_array_construction() { - // Test _LayerArray construction with configurable activations + // Test LayerArray construction with configurable activations const int input_size = 1; const int condition_size = 1; const int head_size = 2;