Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
73 lines
2.3 KiB
C++
73 lines
2.3 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/resource_bundle.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/runtime/model.h"
|
|
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <string_view>
|
|
#include <vector>
|
|
|
|
namespace engine::models::roformer {
|
|
|
|
inline constexpr std::string_view kMelBandRoformerFamily = "mel_band_roformer";
|
|
inline constexpr std::string_view kBsRoformerFamily = "bs_roformer";
|
|
|
|
struct RoformerArchitectureConfig {
|
|
std::string family;
|
|
int sample_rate = 0;
|
|
int channels = 0;
|
|
int chunk_size = 0;
|
|
int inference_batch_size = 0;
|
|
int inference_num_overlap = 0;
|
|
bool inference_normalize = false;
|
|
int dim = 0;
|
|
int depth = 0;
|
|
int num_bands = 0;
|
|
bool stereo = false;
|
|
int num_stems = 0;
|
|
int time_transformer_depth = 1;
|
|
int freq_transformer_depth = 1;
|
|
int linear_transformer_depth = 0;
|
|
int dim_head = 0;
|
|
int heads = 0;
|
|
int n_fft = 0;
|
|
int hop_length = 0;
|
|
int win_length = 0;
|
|
bool stft_normalized = false;
|
|
int mask_estimator_depth = 0;
|
|
int mask_estimator_linear_layers = 0;
|
|
int mlp_expansion_factor = 4;
|
|
bool skip_connection = false;
|
|
bool has_final_norm = false;
|
|
bool fused_qkv = false;
|
|
bool transformer_output_norm = true;
|
|
int stft_freq_bins = 0;
|
|
int chunk_frames = 0;
|
|
int total_band_input_dim = 0;
|
|
float rope_theta = 10000.0f;
|
|
std::vector<int64_t> band_input_dims;
|
|
std::vector<int64_t> merged_freq_indices;
|
|
std::vector<int64_t> merged_band_counts;
|
|
std::vector<std::string> instruments;
|
|
std::optional<std::string> target_instrument = std::nullopt;
|
|
};
|
|
|
|
struct RoformerAssets {
|
|
assets::ResourceBundle resources;
|
|
std::shared_ptr<const assets::TensorSource> tensor_source;
|
|
RoformerArchitectureConfig config;
|
|
};
|
|
|
|
void validate_roformer_weight_storage_type(assets::TensorStorageType storage_type);
|
|
std::shared_ptr<const RoformerAssets> load_mel_band_roformer_assets(
|
|
const runtime::ModelLoadRequest & request);
|
|
std::shared_ptr<const RoformerAssets> load_bs_roformer_assets(
|
|
const runtime::ModelLoadRequest & request);
|
|
std::shared_ptr<const RoformerAssets> load_roformer_assets(
|
|
const runtime::ModelLoadRequest & request,
|
|
std::string_view family);
|
|
|
|
} // namespace engine::models::roformer
|