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
71 lines
3.1 KiB
C++
71 lines
3.1 KiB
C++
#pragma once
|
|
|
|
#include "engine/community_models/kitten_tts/assets.h"
|
|
#include "engine/framework/text/espeak_phonemizer.h"
|
|
#include "engine/framework/model_spec/metadata.h"
|
|
#include "engine/framework/runtime/session_base.h"
|
|
|
|
#include <cstddef>
|
|
#include <memory>
|
|
|
|
namespace engine::models::kitten_tts {
|
|
class KittenStyleTTS2DecoderRuntime;
|
|
class KittenPredictorRuntime;
|
|
} // namespace engine::models::kitten_tts
|
|
|
|
namespace engine::models::kitten_tts {
|
|
|
|
std::shared_ptr<runtime::IVoiceModelLoader> make_kitten_tts_loader();
|
|
|
|
struct KittenSynthesisInput;
|
|
class KittenTTSSession final : public runtime::RuntimeSessionBase, public runtime::IOfflineVoiceTaskSession {
|
|
public:
|
|
KittenTTSSession(runtime::TaskSpec task, runtime::SessionOptions options,
|
|
std::shared_ptr<const KittenAssets> assets,
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract);
|
|
~KittenTTSSession() 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:
|
|
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 KittenAssets> assets_;
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract_;
|
|
std::unique_ptr<engine::audio::EspeakPhonemizer> phonemizer_;
|
|
std::shared_ptr<const KittenWeights> weights_;
|
|
runtime::GraphCapacityController graph_capacity_controller_;
|
|
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<KittenSynthesisInput> cached_input_;
|
|
int64_t prepared_decoder_capacity_ = 0;
|
|
std::unique_ptr<KittenStyleTTS2DecoderRuntime> prepared_decoder_;
|
|
DecoderCapacityContract prepared_decoder_context_ = {};
|
|
int64_t prepared_session_capacity_ = 0;
|
|
std::unique_ptr<KittenPredictorRuntime> prepared_predictor_;
|
|
};
|
|
|
|
} // namespace engine::models::kitten_tts
|