Fabio
c36510c21e
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
108 lines
3.9 KiB
C++
108 lines
3.9 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/core/module.h"
|
|
#include "engine/framework/modules/activation_modules.h"
|
|
#include "engine/framework/modules/conv_modules.h"
|
|
#include "engine/framework/runtime/session.h"
|
|
#include "engine/models/irodori_tts/assets.h"
|
|
|
|
#include <ggml-backend.h>
|
|
|
|
#include <cstddef>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace engine::core {
|
|
class BackendWeightStore;
|
|
}
|
|
|
|
namespace engine::models::irodori_tts {
|
|
|
|
struct IrodoriCodecResidualUnitWeights {
|
|
modules::Snake1dWeights snake0;
|
|
modules::Conv1dWeights conv0;
|
|
modules::Snake1dWeights snake1;
|
|
modules::Conv1dWeights conv1;
|
|
};
|
|
|
|
struct IrodoriCodecDecoderBlockWeights {
|
|
modules::Snake1dWeights up_snake;
|
|
modules::ConvTranspose1dWeights up_conv;
|
|
IrodoriCodecResidualUnitWeights residual_0;
|
|
IrodoriCodecResidualUnitWeights residual_1;
|
|
IrodoriCodecResidualUnitWeights residual_2;
|
|
};
|
|
|
|
struct IrodoriCodecEncoderBlockWeights {
|
|
IrodoriCodecResidualUnitWeights residual_0;
|
|
IrodoriCodecResidualUnitWeights residual_1;
|
|
IrodoriCodecResidualUnitWeights residual_2;
|
|
modules::Snake1dWeights down_snake;
|
|
modules::Conv1dWeights down_conv;
|
|
};
|
|
|
|
struct IrodoriCodecWatermarkPassthroughWeights {
|
|
modules::Snake1dWeights snake;
|
|
modules::Conv1dWeights conv;
|
|
};
|
|
|
|
struct IrodoriCodecWeights {
|
|
std::shared_ptr<core::BackendWeightStore> store;
|
|
modules::Conv1dWeights encoder_input;
|
|
std::vector<IrodoriCodecEncoderBlockWeights> encoder_blocks;
|
|
modules::Snake1dWeights encoder_output_snake;
|
|
modules::Conv1dWeights encoder_output;
|
|
modules::Conv1dWeights quantizer_in_proj;
|
|
modules::Conv1dWeights quantizer_out_proj;
|
|
modules::Conv1dWeights decoder_input;
|
|
std::vector<IrodoriCodecDecoderBlockWeights> decoder_blocks;
|
|
IrodoriCodecWatermarkPassthroughWeights watermark_passthrough;
|
|
};
|
|
|
|
IrodoriCodecWeights
|
|
load_irodori_codec_weights(const IrodoriTTSAssets &assets, ggml_backend_t backend,
|
|
core::BackendType backend_type,
|
|
size_t weight_context_bytes,
|
|
assets::TensorStorageType conv_storage_type);
|
|
|
|
core::TensorValue build_irodori_codec_decode(
|
|
core::ModuleBuildContext &ctx, const core::TensorValue &latent_btd,
|
|
const IrodoriCodecWeights &weights, const IrodoriCodecConfig &config);
|
|
|
|
core::TensorValue build_irodori_codec_encode(
|
|
core::ModuleBuildContext &ctx, const core::TensorValue &waveform_bct,
|
|
const IrodoriCodecWeights &weights, const IrodoriCodecConfig &config);
|
|
|
|
class IrodoriCodec {
|
|
public:
|
|
IrodoriCodec(std::shared_ptr<const IrodoriTTSAssets> assets,
|
|
core::ExecutionContext &execution_context,
|
|
size_t graph_arena_bytes, size_t weight_context_bytes,
|
|
assets::TensorStorageType weight_storage_type);
|
|
~IrodoriCodec();
|
|
|
|
runtime::AudioBuffer decode(const std::vector<float> &latent,
|
|
int64_t latent_steps, int64_t target_samples,
|
|
const std::function<void(const runtime::AudioBuffer &)> &on_audio = {},
|
|
int64_t stream_frames = 0);
|
|
std::vector<float> encode_reference(const runtime::AudioBuffer &audio,
|
|
int64_t &latent_steps_out);
|
|
// Decode in fixed-size windows of chunk_steps latent frames plus
|
|
// overlap_steps of context on each side, so the decode graph (and its
|
|
// activation buffer) no longer grows with the output length. 0 disables.
|
|
void set_decode_chunking(int64_t chunk_steps, int64_t overlap_steps);
|
|
// The same for encoding reference audio: windows of chunk_steps latent
|
|
// frames (hop_length samples each) plus overlap_steps on each side.
|
|
void set_encode_chunking(int64_t chunk_steps, int64_t overlap_steps);
|
|
void release_graphs();
|
|
|
|
private:
|
|
class Impl;
|
|
std::unique_ptr<Impl> impl_;
|
|
};
|
|
|
|
} // namespace engine::models::irodori_tts
|