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
334 lines
15 KiB
C++
334 lines
15 KiB
C++
#include "engine/models/firered_audio/session.h"
|
|
|
|
#include "engine/framework/debug/profiler.h"
|
|
#include "engine/framework/debug/trace.h"
|
|
#include "engine/framework/runtime/options.h"
|
|
#include "engine/framework/runtime/spec_backed_model.h"
|
|
#include "engine/models/firered_audio/pipeline.h"
|
|
#include "engine/models/firered_audio/redae.h"
|
|
#include "engine/models/firered_audio/understanding.h"
|
|
|
|
#include <stdexcept>
|
|
#include <utility>
|
|
|
|
namespace engine::models::firered_audio {
|
|
namespace {
|
|
|
|
constexpr const char * kFamily = "firered_audio";
|
|
constexpr const char * kModelName = "FireRedAudio";
|
|
constexpr size_t kDefaultReferenceCacheSlots = 2;
|
|
|
|
std::shared_ptr<const FireRedAudioAssets> require_assets(std::shared_ptr<const FireRedAudioAssets> assets) {
|
|
if (assets == nullptr) {
|
|
throw std::runtime_error("FireRedAudio session requires assets");
|
|
}
|
|
return assets;
|
|
}
|
|
|
|
std::shared_ptr<const engine::model_spec::ModelContract> require_contract(
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract) {
|
|
if (contract == nullptr) {
|
|
throw std::runtime_error("FireRedAudio session requires a model contract");
|
|
}
|
|
return contract;
|
|
}
|
|
|
|
std::string request_language(const runtime::TaskRequest & request) {
|
|
if (const auto language = runtime::find_option(request.options, {"language"})) {
|
|
return *language;
|
|
}
|
|
if (request.text_input.has_value() && !request.text_input->language.empty()) {
|
|
return request.text_input->language;
|
|
}
|
|
return "zh";
|
|
}
|
|
|
|
FireRedAudioGenerationRequest make_generation_request(
|
|
const FireRedAudioAssets & assets,
|
|
const FireRedAudioTokenizer & tokenizer,
|
|
const runtime::TaskSpec & task,
|
|
const runtime::TaskRequest & request) {
|
|
FireRedAudioGenerationRequest out;
|
|
const std::string mode = runtime::find_option(request.options, {"template_name"}).value_or(
|
|
task.task == runtime::VoiceTaskKind::VoiceDesign ? "voice_design" :
|
|
(request.audio_input.has_value() ? "acoustic_edit" : "tts_clone"));
|
|
if (mode == "tts_clone") {
|
|
if (!request.text_input.has_value() || request.text_input->text.empty()) {
|
|
throw std::runtime_error("FireRedAudio clone requires text input");
|
|
}
|
|
if (!request.voice.has_value() ||
|
|
!request.voice->speaker.has_value() ||
|
|
!request.voice->speaker->audio.has_value()) {
|
|
throw std::runtime_error("FireRedAudio clone requires reference audio");
|
|
}
|
|
out.prompt_audio = *request.voice->speaker->audio;
|
|
out.prompt_audio_is_assistant = true;
|
|
const auto audio_24k = prepare_firered_prompt_audio_24k(*out.prompt_audio, assets.redae, assets.patch_encoder.patch_size);
|
|
const int64_t placeholder_count =
|
|
static_cast<int64_t>(audio_24k.size()) /
|
|
(assets.redae.audio_patch_size * assets.redae.enc_extra_downsample_rate * assets.patch_encoder.patch_size);
|
|
const std::string language = request_language(request);
|
|
const auto reference_text = runtime::find_option(request.options, {"reference_text"}).value_or("");
|
|
auto encoded = tokenizer.encode_tts_clone_prompt(
|
|
reference_text,
|
|
request.text_input->text,
|
|
language,
|
|
placeholder_count);
|
|
out.token_ids = std::move(encoded.token_ids);
|
|
out.audio_no_latent_mask = std::move(encoded.audio_no_latent_mask);
|
|
out.generated_text_language = language;
|
|
} else if (mode == "voice_design") {
|
|
if (!request.text_input.has_value() || request.text_input->text.empty()) {
|
|
throw std::runtime_error("FireRedAudio voice design requires text input");
|
|
}
|
|
const auto instruction = runtime::find_option(request.options, {"instruct", "instruction"}).value_or("");
|
|
if (instruction.empty()) {
|
|
throw std::runtime_error("FireRedAudio voice design requires instruction");
|
|
}
|
|
auto encoded = tokenizer.encode_voice_design_prompt(instruction, request.text_input->text);
|
|
out.token_ids = std::move(encoded.token_ids);
|
|
out.audio_no_latent_mask = std::move(encoded.audio_no_latent_mask);
|
|
out.generated_text_language = request_language(request);
|
|
} else if (mode == "semantic_edit" || mode == "acoustic_edit") {
|
|
if (!request.audio_input.has_value()) {
|
|
throw std::runtime_error("FireRedAudio edit requires input audio");
|
|
}
|
|
const auto instruction = runtime::find_option(request.options, {"instruct", "instruction"}).value_or("");
|
|
if (instruction.empty()) {
|
|
throw std::runtime_error("FireRedAudio edit requires instruction");
|
|
}
|
|
out.prompt_audio = *request.audio_input;
|
|
out.prompt_audio_is_assistant = false;
|
|
const auto audio_24k = prepare_firered_prompt_audio_24k(*out.prompt_audio, assets.redae, assets.patch_encoder.patch_size);
|
|
const int64_t placeholder_count =
|
|
static_cast<int64_t>(audio_24k.size()) /
|
|
(assets.redae.audio_patch_size * assets.redae.enc_extra_downsample_rate * assets.patch_encoder.patch_size);
|
|
auto encoded = tokenizer.encode_edit_prompt(
|
|
instruction,
|
|
mode == "semantic_edit" ? "semantic" : "acoustic",
|
|
placeholder_count);
|
|
out.token_ids = std::move(encoded.token_ids);
|
|
out.audio_no_latent_mask = std::move(encoded.audio_no_latent_mask);
|
|
out.generated_text_language = request_language(request);
|
|
} else {
|
|
throw std::runtime_error("unknown FireRedAudio template_name: " + mode);
|
|
}
|
|
out.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(out.seed);
|
|
out.num_inference_steps = runtime::parse_positive_i64_option(
|
|
request.options,
|
|
{"num_inference_steps"},
|
|
out.num_inference_steps);
|
|
if (const auto cfg = runtime::parse_positive_finite_float_option(request.options, {"guidance_scale"})) {
|
|
out.guidance_scale = *cfg;
|
|
}
|
|
out.max_new_audio_steps = runtime::parse_positive_i64_option(
|
|
request.options,
|
|
{"max_audio_steps", "max_new_audio_steps"},
|
|
out.max_new_audio_steps);
|
|
out.min_new_audio_steps = runtime::parse_i64_option(request.options, {"min_audio_steps", "min_new_audio_steps"}).value_or(out.min_new_audio_steps);
|
|
if (out.min_new_audio_steps < 0) {
|
|
throw std::runtime_error("FireRedAudio min_audio_steps must be non-negative");
|
|
}
|
|
out.max_new_text_tokens = runtime::parse_positive_i64_option(
|
|
request.options,
|
|
{"max_text_tokens", "max_new_text_tokens"},
|
|
out.max_new_text_tokens);
|
|
return out;
|
|
}
|
|
|
|
FireRedAudioUnderstandingRequest make_understanding_request(
|
|
const runtime::TaskRequest & request) {
|
|
if (!request.audio_input.has_value()) {
|
|
throw std::runtime_error("FireRedAudio ASR/understanding requires input audio");
|
|
}
|
|
FireRedAudioUnderstandingRequest out;
|
|
out.audio = *request.audio_input;
|
|
out.prompt = request.text_input.has_value() && !request.text_input->text.empty()
|
|
? request.text_input->text
|
|
: "Transcribe speech to text.";
|
|
out.language = request_language(request);
|
|
out.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(out.seed);
|
|
out.max_new_tokens = runtime::parse_positive_i64_option(
|
|
request.options,
|
|
{"max_tokens", "max_new_tokens"},
|
|
out.max_new_tokens);
|
|
out.enable_thinking = runtime::parse_bool_option(
|
|
runtime::find_option(request.options, {"enable_thinking"}).value_or("false"),
|
|
"enable_thinking");
|
|
const std::string mode = runtime::find_option(request.options, {"template_name"}).value_or("asr");
|
|
if (mode == "understand") {
|
|
out.do_sample = true;
|
|
out.repetition_penalty = 1.0F;
|
|
if (out.enable_thinking && !runtime::find_option(request.options, {"max_tokens", "max_new_tokens"}).has_value()) {
|
|
out.max_new_tokens = 1024;
|
|
}
|
|
} else if (mode == "asr") {
|
|
out.do_sample = false;
|
|
out.repetition_penalty = 1.1F;
|
|
if (out.enable_thinking) {
|
|
throw std::runtime_error("FireRedAudio ASR does not support enable_thinking");
|
|
}
|
|
} else {
|
|
throw std::runtime_error("unknown FireRedAudio ASR template_name: " + mode);
|
|
}
|
|
if (const auto top_k = runtime::parse_i64_option(request.options, {"top_k"})) {
|
|
out.top_k = *top_k;
|
|
}
|
|
if (const auto top_p = runtime::parse_finite_float_option(request.options, {"top_p"})) {
|
|
out.top_p = *top_p;
|
|
}
|
|
if (const auto temperature = runtime::parse_finite_float_option(request.options, {"temperature"})) {
|
|
out.temperature = *temperature;
|
|
}
|
|
return out;
|
|
}
|
|
|
|
std::unique_ptr<runtime::IVoiceTaskSession> create_firered_audio_session(
|
|
const runtime::TaskSpec & task,
|
|
const runtime::SessionOptions & options,
|
|
std::shared_ptr<const FireRedAudioAssets> assets,
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract) {
|
|
return std::make_unique<FireRedAudioSession>(
|
|
task,
|
|
options,
|
|
std::move(assets),
|
|
std::move(contract));
|
|
}
|
|
|
|
} // namespace
|
|
|
|
FireRedAudioSession::FireRedAudioSession(
|
|
runtime::TaskSpec task,
|
|
runtime::SessionOptions options,
|
|
std::shared_ptr<const FireRedAudioAssets> assets,
|
|
std::shared_ptr<const engine::model_spec::ModelContract> contract)
|
|
: runtime::RuntimeSessionBase(options),
|
|
task_(task),
|
|
assets_(require_assets(std::move(assets))),
|
|
contract_(require_contract(std::move(contract))),
|
|
tokenizer_(std::make_unique<FireRedAudioTokenizer>(assets_)) {
|
|
auto validation_options = options;
|
|
// Retain the previously advertised option, which never affected allocation.
|
|
validation_options.options.erase("firered_audio.helper_graph_arena_mb");
|
|
runtime::validate_spec_backed_session_options(validation_options, *contract_, kFamily, kModelName);
|
|
if (task_.mode != runtime::RunMode::Offline) {
|
|
throw std::runtime_error("FireRedAudio supports offline sessions");
|
|
}
|
|
if (task_.task != runtime::VoiceTaskKind::Tts &&
|
|
task_.task != runtime::VoiceTaskKind::VoiceCloning &&
|
|
task_.task != runtime::VoiceTaskKind::VoiceDesign &&
|
|
task_.task != runtime::VoiceTaskKind::AudioGeneration &&
|
|
task_.task != runtime::VoiceTaskKind::Asr) {
|
|
throw std::runtime_error("FireRedAudio supports tts, clone, voice design, gen, and asr tasks");
|
|
}
|
|
using T = engine::assets::TensorStorageType;
|
|
const auto storage_type = runtime::parse_tensor_storage_option(
|
|
options.options,
|
|
"firered_audio.weight_type",
|
|
T::Native,
|
|
{T::Native, T::F32, T::F16, T::BF16, T::Q8_0});
|
|
const auto graph_arena_bytes = runtime::parse_size_mb_option(
|
|
options.options,
|
|
{"firered_audio.graph_arena_mb"},
|
|
1024ull * 1024ull * 1024ull);
|
|
const auto weight_context_bytes = runtime::parse_size_mb_option(
|
|
options.options,
|
|
{"firered_audio.weight_context_mb"},
|
|
1024ull * 1024ull * 1024ull);
|
|
if (const auto mem_saver = runtime::find_option(options.options, {"firered_audio.mem_saver"})) {
|
|
mem_saver_ = runtime::parse_bool_option(*mem_saver, "firered_audio.mem_saver");
|
|
}
|
|
const int64_t reference_cache_slots = runtime::parse_i64_option(
|
|
options.options,
|
|
{"firered_audio.reference_cache_slots"})
|
|
.value_or(static_cast<int64_t>(kDefaultReferenceCacheSlots));
|
|
if (reference_cache_slots < 0) {
|
|
throw std::runtime_error("firered_audio.reference_cache_slots must be non-negative");
|
|
}
|
|
if (task_.task == runtime::VoiceTaskKind::Asr) {
|
|
understanding_runtime_ = std::make_unique<FireRedAudioUnderstandingRuntime>(
|
|
assets_,
|
|
execution_context(),
|
|
graph_arena_bytes,
|
|
weight_context_bytes,
|
|
storage_type,
|
|
mem_saver_);
|
|
} else {
|
|
runtime_ = std::make_unique<FireRedAudioGenerationRuntime>(
|
|
assets_,
|
|
execution_context(),
|
|
graph_arena_bytes,
|
|
weight_context_bytes,
|
|
storage_type,
|
|
static_cast<size_t>(reference_cache_slots),
|
|
mem_saver_);
|
|
}
|
|
}
|
|
|
|
FireRedAudioSession::~FireRedAudioSession() = default;
|
|
|
|
std::string FireRedAudioSession::family() const {
|
|
return kFamily;
|
|
}
|
|
|
|
runtime::VoiceTaskKind FireRedAudioSession::task_kind() const {
|
|
return task_.task;
|
|
}
|
|
|
|
runtime::RunMode FireRedAudioSession::run_mode() const {
|
|
return task_.mode;
|
|
}
|
|
|
|
void FireRedAudioSession::prepare(const runtime::SessionPreparationRequest & request) {
|
|
auto validation_options = request.options;
|
|
(void)runtime::find_option(request.options, {"instruct", "instruction"});
|
|
for (const auto * key : {"instruct", "instruction", "max_tokens", "max_new_tokens", "max_text_tokens", "max_new_text_tokens",
|
|
"max_audio_steps", "max_new_audio_steps", "min_audio_steps", "min_new_audio_steps"}) {
|
|
validation_options.erase(key);
|
|
}
|
|
runtime::validate_spec_backed_request_options(
|
|
validation_options,
|
|
*contract_,
|
|
kModelName);
|
|
mark_prepared();
|
|
}
|
|
|
|
runtime::TaskResult FireRedAudioSession::run(const runtime::TaskRequest & request) {
|
|
const auto wall_start = std::chrono::steady_clock::now();
|
|
auto validation_options = request.options;
|
|
(void)runtime::find_option(request.options, {"instruct", "instruction"});
|
|
for (const auto * key : {"instruct", "instruction", "max_tokens", "max_new_tokens", "max_text_tokens", "max_new_text_tokens",
|
|
"max_audio_steps", "max_new_audio_steps", "min_audio_steps", "min_new_audio_steps"}) {
|
|
validation_options.erase(key);
|
|
}
|
|
runtime::validate_spec_backed_request_options(
|
|
validation_options,
|
|
*contract_,
|
|
kModelName);
|
|
require_prepared("FireRedAudio run");
|
|
runtime::TaskResult result;
|
|
if (task_.task == runtime::VoiceTaskKind::Asr) {
|
|
auto parsed = make_understanding_request(request);
|
|
auto generated = understanding_runtime_->generate(parsed);
|
|
result.text_output = runtime::Transcript{std::move(generated.text), generated.language};
|
|
} else {
|
|
auto parsed = make_generation_request(*assets_, *tokenizer_, task_, request);
|
|
auto generated = runtime_->generate(parsed);
|
|
result.audio_output = std::move(generated.audio);
|
|
if (!generated.generated_text.empty()) {
|
|
result.text_output = runtime::Transcript{std::move(generated.generated_text), parsed.generated_text_language};
|
|
}
|
|
}
|
|
engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start));
|
|
return result;
|
|
}
|
|
|
|
std::shared_ptr<runtime::IVoiceModelLoader> make_firered_audio_loader() {
|
|
runtime::SpecBackedVoiceModelConfig<FireRedAudioAssets> config;
|
|
config.family = kFamily;
|
|
config.load_assets = load_firered_audio_assets;
|
|
config.create_session = create_firered_audio_session;
|
|
return runtime::make_spec_backed_voice_loader(std::move(config));
|
|
}
|
|
|
|
} // namespace engine::models::firered_audio
|