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
112 lines
4.9 KiB
C++
112 lines
4.9 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/model_spec/metadata.h"
|
|
#include "engine/framework/runtime/session_base.h"
|
|
#include "engine/framework/runtime/streaming_tts_session.h"
|
|
#include "engine/models/kokoro_tts/assets.h"
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
namespace kokoro_ggml {
|
|
class KokoroIstftNetDecoderRuntime;
|
|
class KokoroPredictorRuntime;
|
|
}
|
|
|
|
namespace engine::models::kokoro_tts {
|
|
|
|
std::shared_ptr<runtime::IVoiceModelLoader> make_kokoro_tts_loader();
|
|
|
|
/// One timing per PHONEME GROUP -- a run of tokens between the space tokens Kokoro's vocabulary
|
|
/// carries -- from the durations its duration predictor produced, appended to `out`.
|
|
///
|
|
/// A group is NOT a written word: on the text path eSpeak-ng merges function words, so
|
|
/// `on the` arrives as one group. See the note on the definition in session.cpp.
|
|
///
|
|
/// Declared here rather than kept in session.cpp's anonymous namespace so it can be tested
|
|
/// directly: it is a pure function of its arguments, and the pad/space boundaries, the
|
|
/// punctuation-only groups and the frames->samples scale are all worth pinning down without
|
|
/// standing up a session. It takes the vocabulary rather than the whole KokoroAssets for the same
|
|
/// reason -- that is all it needs.
|
|
///
|
|
/// `chunk_start_sample` offsets the spans into a buffer several chunks are being merged into.
|
|
/// Reports nothing at all rather than throwing if the token and duration counts disagree: the
|
|
/// audio is the product and an empty word list is a state every caller already handles.
|
|
void append_kokoro_word_timings(
|
|
std::vector<runtime::WordTimestamp> & out,
|
|
const std::vector<int32_t> & input_ids,
|
|
const std::vector<int32_t> & durations,
|
|
const std::unordered_map<std::string, int32_t> & vocab,
|
|
size_t chunk_samples,
|
|
int64_t chunk_start_sample);
|
|
|
|
struct KokoroSynthesisInput;
|
|
class KokoroTTSSession final
|
|
: public runtime::RuntimeSessionBase
|
|
, public runtime::IOfflineVoiceTaskSession
|
|
, public runtime::StreamingTtsSessionBase {
|
|
public:
|
|
KokoroTTSSession(
|
|
runtime::TaskSpec task,
|
|
runtime::SessionOptions options,
|
|
std::shared_ptr<const KokoroAssets> assets,
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract);
|
|
~KokoroTTSSession() override;
|
|
|
|
std::string family() const override;
|
|
runtime::VoiceTaskKind task_kind() const override;
|
|
runtime::RunMode run_mode() const override;
|
|
void prepare(const runtime::SessionPreparationRequest & request) override;
|
|
runtime::TaskResult run(const runtime::TaskRequest & request) override;
|
|
|
|
private:
|
|
runtime::TaskResult synthesize(const runtime::TaskRequest & request, bool streaming);
|
|
runtime::TaskResult generate_stream(const runtime::TaskRequest & request) override;
|
|
void reset_stream_state() override {}
|
|
|
|
struct DecoderCapacityContract {
|
|
int64_t decoder_frame_capacity = 0;
|
|
int64_t conditioning_sample_capacity = 0;
|
|
int64_t conditioning_frame_capacity = 0;
|
|
};
|
|
|
|
runtime::MappedGraphCapacityAdapter make_graph_capacity_adapter();
|
|
std::vector<int64_t> prepared_graph_capacities() const;
|
|
DecoderCapacityContract make_decoder_capacity_contract(int64_t decoder_frame_capacity) const;
|
|
void prepare_graph_capacity(int64_t capacity);
|
|
void prepare_decoder_graph_capacity(int64_t capacity);
|
|
|
|
runtime::TaskSpec task_;
|
|
std::shared_ptr<const KokoroAssets> assets_;
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract_;
|
|
std::shared_ptr<const kokoro_ggml::KokoroWeights> weights_;
|
|
runtime::GraphCapacityController graph_capacity_controller_;
|
|
// artifacts.main from kokoro_tts.exec_plan.generator when RKNN is selected.
|
|
std::string rknn_generator_path_;
|
|
// exec_plan.generator params: overlap-blend / per-window loudness correction.
|
|
int64_t rknn_generator_overlap_ = 0;
|
|
bool rknn_generator_gain_ = false;
|
|
bool neon_generator_ = false;
|
|
int64_t fixed_token_capacity_ = 0;
|
|
int64_t pre_tail_token_capacity_ = 0;
|
|
uint64_t rng_seed_ = 0;
|
|
engine::assets::TensorStorageType matmul_weight_storage_type_ = engine::assets::TensorStorageType::Native;
|
|
engine::assets::TensorStorageType conv_weight_storage_type_ = engine::assets::TensorStorageType::Native;
|
|
size_t weight_context_bytes_ = 512ull * 1024ull * 1024ull;
|
|
size_t predictor_duration_graph_bytes_ = 384ull * 1024ull * 1024ull;
|
|
size_t predictor_text_graph_bytes_ = 256ull * 1024ull * 1024ull;
|
|
size_t predictor_tail_graph_bytes_ = 640ull * 1024ull * 1024ull;
|
|
std::string cached_request_key_;
|
|
std::unique_ptr<KokoroSynthesisInput> cached_input_;
|
|
int64_t prepared_decoder_capacity_ = 0;
|
|
std::unique_ptr<kokoro_ggml::KokoroIstftNetDecoderRuntime> prepared_decoder_;
|
|
DecoderCapacityContract prepared_decoder_context_ = {};
|
|
int64_t prepared_session_capacity_ = 0;
|
|
std::unique_ptr<kokoro_ggml::KokoroPredictorRuntime> prepared_predictor_;
|
|
};
|
|
|
|
} // namespace engine::models::kokoro_tts
|