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
100 lines
4.3 KiB
C++
100 lines
4.3 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/model_spec/metadata.h"
|
|
#include "engine/framework/runtime/cache_slots.h"
|
|
#include "engine/framework/runtime/session_base.h"
|
|
#include "engine/models/confucius4_tts/assets.h"
|
|
#include "engine/models/confucius4_tts/audio_features.h"
|
|
#include "engine/models/confucius4_tts/request.h"
|
|
#include "engine/models/confucius4_tts/s2a.h"
|
|
#include "engine/models/confucius4_tts/semantic_encoder.h"
|
|
#include "engine/models/confucius4_tts/style_encoder.h"
|
|
#include "engine/models/confucius4_tts/t2s.h"
|
|
#include "engine/models/confucius4_tts/tokenizer_text.h"
|
|
#include "engine/models/confucius4_tts/vocoder.h"
|
|
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::models::confucius4_tts {
|
|
|
|
std::shared_ptr<runtime::IVoiceModelLoader> make_confucius4_tts_loader();
|
|
|
|
struct ConfuciusReferenceIdentity {
|
|
std::string id;
|
|
};
|
|
|
|
struct ConfuciusPreparedVoice {
|
|
ConfuciusPreparedReferenceAudio audio;
|
|
ConfuciusSemanticEmbedding semantic;
|
|
ConfuciusStyleEmbedding style;
|
|
};
|
|
|
|
class ConfuciusSession final
|
|
: public runtime::RuntimeSessionBase,
|
|
public runtime::IOfflineVoiceTaskSession,
|
|
public runtime::IStreamingVoiceTaskSession {
|
|
public:
|
|
ConfuciusSession(
|
|
runtime::TaskSpec task,
|
|
runtime::SessionOptions options,
|
|
std::shared_ptr<const ConfuciusAssets> assets,
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract);
|
|
~ConfuciusSession() 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;
|
|
runtime::StreamingPolicy streaming_policy() const override;
|
|
void start_stream(const runtime::TaskRequest & request) override;
|
|
std::optional<runtime::StreamEvent> next_stream_event() override;
|
|
void set_stream_event_sink(runtime::StreamEventCallback sink) override;
|
|
runtime::TaskResult finish_stream() override;
|
|
void reset() override;
|
|
runtime::StreamEvent process_audio_chunk(const runtime::AudioChunk & chunk) override;
|
|
runtime::TaskResult finalize() override;
|
|
|
|
private:
|
|
struct ReferenceIdentityEqual {
|
|
bool operator()(const ConfuciusReferenceIdentity & lhs, const ConfuciusReferenceIdentity & rhs) const noexcept;
|
|
};
|
|
ConfuciusPreparedVoice & resolve_voice(const ConfuciusRequest & request);
|
|
ConfuciusPreparedVoice prepare_voice(const ConfuciusRequest & request);
|
|
runtime::AudioBuffer synthesize_segment(
|
|
const ConfuciusTextSegment & segment,
|
|
const ConfuciusPreparedVoice & voice,
|
|
const ConfuciusGenerationOptions & options,
|
|
uint64_t & rng_offset_blocks);
|
|
runtime::AudioBuffer synthesize(const ConfuciusRequest & request, const ConfuciusPreparedVoice & voice);
|
|
|
|
runtime::TaskSpec task_;
|
|
std::shared_ptr<const ConfuciusAssets> assets_;
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract_;
|
|
ConfuciusTextTokenizer tokenizer_;
|
|
std::unique_ptr<ConfuciusWav2Vec2BertRuntime> semantic_encoder_;
|
|
std::unique_ptr<ConfuciusCampplusStyleEncoder> style_encoder_;
|
|
std::unique_ptr<ConfuciusT2SRuntime> t2s_;
|
|
std::unique_ptr<ConfuciusS2ARuntime> s2a_;
|
|
std::unique_ptr<ConfuciusBigVganVocoder> vocoder_;
|
|
runtime::CacheSlots<ConfuciusReferenceIdentity, ConfuciusPreparedVoice, ReferenceIdentityEqual> reference_cache_;
|
|
std::optional<ConfuciusPreparedVoice> uncached_reference_;
|
|
std::optional<ConfuciusRequest> prepared_defaults_;
|
|
std::optional<ConfuciusRequest> streaming_request_;
|
|
const ConfuciusPreparedVoice * streaming_voice_ = nullptr;
|
|
std::vector<ConfuciusTextSegment> streaming_segments_;
|
|
std::vector<runtime::AudioBuffer> streaming_chunks_;
|
|
runtime::StreamEventCallback stream_sink_;
|
|
size_t streaming_segment_index_ = 0;
|
|
uint64_t streaming_rng_offset_blocks_ = 0;
|
|
size_t graph_arena_bytes_ = 512ull * 1024ull * 1024ull;
|
|
size_t weight_context_bytes_ = 1024ull * 1024ull * 1024ull;
|
|
engine::assets::TensorStorageType matmul_weight_storage_type_ = engine::assets::TensorStorageType::Native;
|
|
engine::assets::TensorStorageType conv_weight_storage_type_ = engine::assets::TensorStorageType::Native;
|
|
bool mem_saver_ = false;
|
|
};
|
|
|
|
} // namespace engine::models::confucius4_tts
|