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
111 lines
5.1 KiB
C++
111 lines
5.1 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/runtime/model.h"
|
|
#include "engine/framework/runtime/session_base.h"
|
|
#include "engine/models/omnivoice/assets.h"
|
|
#include "engine/models/omnivoice/audio_tokenizer.h"
|
|
#include "engine/models/omnivoice/generator.h"
|
|
#include "engine/models/omnivoice/postprocess.h"
|
|
#include "engine/models/omnivoice/prompt_builder.h"
|
|
#include "engine/models/omnivoice/tokenizer_text.h"
|
|
|
|
#include <cstddef>
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
namespace engine::models::omnivoice {
|
|
|
|
class OmniVoiceSession final
|
|
: public runtime::RuntimeSessionBase
|
|
, public runtime::IOfflineVoiceTaskSession
|
|
, public runtime::IStreamingVoiceTaskSession {
|
|
public:
|
|
OmniVoiceSession(
|
|
runtime::TaskSpec task,
|
|
runtime::SessionOptions options,
|
|
std::shared_ptr<const OmniVoiceAssets> assets);
|
|
|
|
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 SessionDefaults {
|
|
std::optional<runtime::Transcript> text = std::nullopt;
|
|
std::optional<runtime::AudioBuffer> reference_audio = std::nullopt;
|
|
std::optional<std::string> reference_text = std::nullopt;
|
|
std::optional<std::string> instruct = std::nullopt;
|
|
std::unordered_map<std::string, std::string> options;
|
|
};
|
|
|
|
struct ReferencePromptCacheEntry {
|
|
bool preprocess_prompt = true;
|
|
bool reference_text_provided = false;
|
|
int sample_rate = 0;
|
|
int channels = 0;
|
|
uint64_t sample_count = 0;
|
|
uint64_t sample_hash = 0;
|
|
OmniVoiceAudioTokens tokens;
|
|
};
|
|
|
|
OmniVoiceRequest make_request(const runtime::TaskRequest & request) const;
|
|
std::optional<runtime::AudioBuffer> resolve_reference_audio(const runtime::TaskRequest & request) const;
|
|
std::optional<std::string> resolve_reference_text(const runtime::TaskRequest & request) const;
|
|
std::optional<std::string> resolve_instruct(const runtime::TaskRequest & request) const;
|
|
std::unordered_map<std::string, std::string> merged_request_options(const runtime::TaskRequest & request) const;
|
|
OmniVoiceAudioTokens resolve_reference_audio_tokens(
|
|
const runtime::AudioBuffer & audio,
|
|
bool preprocess_prompt,
|
|
bool reference_text_provided);
|
|
std::vector<std::string> plan_text_chunks(const OmniVoiceRequest & request, const OmniVoicePrompt & prompt) const;
|
|
void initialize_streaming_request(const runtime::TaskRequest & request);
|
|
runtime::AudioBuffer synthesize_stream_chunk(size_t chunk_index);
|
|
|
|
runtime::TaskSpec task_;
|
|
std::shared_ptr<const OmniVoiceAssets> assets_;
|
|
size_t audio_tokenizer_graph_arena_bytes_ = 128ull * 1024ull * 1024ull;
|
|
size_t generator_prefill_graph_arena_bytes_ = 256ull * 1024ull * 1024ull;
|
|
size_t generator_decode_graph_arena_bytes_ = 256ull * 1024ull * 1024ull;
|
|
size_t audio_tokenizer_weight_context_bytes_ = 128ull * 1024ull * 1024ull;
|
|
size_t generator_weight_context_bytes_ = 256ull * 1024ull * 1024ull;
|
|
engine::assets::TensorStorageType audio_tokenizer_weight_storage_type_ = engine::assets::TensorStorageType::Native;
|
|
engine::assets::TensorStorageType audio_tokenizer_codebook_weight_storage_type_ = engine::assets::TensorStorageType::F32;
|
|
engine::assets::TensorStorageType generator_weight_storage_type_ = engine::assets::TensorStorageType::Native;
|
|
bool mem_saver_ = false;
|
|
OmniVoiceGeneratorPerfMode generator_perf_mode_ = OmniVoiceGeneratorPerfMode::Standard;
|
|
OmniVoiceTextTokenizer tokenizer_;
|
|
OmniVoiceHiggsAudioV2TokenizerRuntime audio_tokenizer_;
|
|
OmniVoicePromptBuilder prompt_builder_;
|
|
OmniVoiceQwen3DiffusionRuntime generator_;
|
|
OmniVoicePostprocessor postprocessor_;
|
|
SessionDefaults session_defaults_;
|
|
std::optional<ReferencePromptCacheEntry> reference_prompt_cache_;
|
|
std::optional<OmniVoiceRequest> stream_request_;
|
|
std::vector<std::string> stream_text_chunks_;
|
|
std::optional<OmniVoiceAudioTokens> stream_first_chunk_reference_;
|
|
std::string stream_first_chunk_text_;
|
|
runtime::AudioBuffer stream_merged_audio_;
|
|
size_t stream_chunk_index_ = 0;
|
|
int64_t stream_chunk_codebooks_ = 0;
|
|
bool stream_started_ = false;
|
|
bool stream_has_reference_audio_ = false;
|
|
std::chrono::steady_clock::time_point stream_wall_start_;
|
|
};
|
|
|
|
} // namespace engine::models::omnivoice
|