diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 22b6d187e1a..3745f3e7b6c 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -908,6 +908,120 @@ bool mtmd_audio_preprocessor_granite_speech::preprocess(const float * // mtmd_audio_preprocessor_gemma4a // +bool mtmd_audio_compute_gemma4_features(const float * samples, + size_t n_samples, + int sample_rate, + int n_mel, + int n_fft, + int window_len, + int hop_len, + std::vector & features, + int & n_frames_out) { + features.clear(); + n_frames_out = 0; + if (samples == nullptr || n_samples == 0 || sample_rate <= 0 || n_mel <= 0 || n_fft <= 0 || + window_len <= 0 || hop_len <= 0) { + return false; + } + + // Gemma4 audio frontend mirrors the original Python pipeline: + // truncate to 30s, right-pad waveform to a multiple of 128, left-pad by + // half a window, unfold 321 samples, and FFT only the first 320 samples. + const size_t max_length = 480000; + const size_t n_valid = std::min(n_samples, max_length); + const size_t pad_right = (128 - (n_valid % 128)) % 128; + + std::vector waveform(n_valid + pad_right, 0.0f); + std::copy(samples, samples + n_valid, waveform.data()); + + std::vector attention_mask(waveform.size(), 0); + std::fill(attention_mask.begin(), attention_mask.begin() + n_valid, 1); + + const int pad_left = window_len / 2; + const int frame_size_for_unfold = window_len + 1; + + std::vector padded_samples((size_t) pad_left + waveform.size(), 0.0f); + std::copy(waveform.begin(), waveform.end(), padded_samples.begin() + pad_left); + + std::vector padded_mask((size_t) pad_left + attention_mask.size(), 0); + std::copy(attention_mask.begin(), attention_mask.end(), padded_mask.begin() + pad_left); + + if (padded_samples.size() < (size_t) frame_size_for_unfold) { + return false; + } + + const int n_frames = (int) ((padded_samples.size() - (size_t) frame_size_for_unfold) / (size_t) hop_len) + 1; + if (n_frames <= 0) { + return false; + } + + mtmd_audio_cache cache; + cache.fill_sin_cos_table(n_fft); + cache.hann_window.assign(window_len, 0.0f); + for (uint32_t i = 0; i < (uint32_t) window_len; i++) { + cache.hann_window[i] = 0.5f - 0.5f * cosf((2.0f * (float) M_PI * i) / window_len); + } + cache.fill_mel_filterbank_matrix( + n_mel, n_fft, sample_rate, + 0.0f, sample_rate / 2.0f, + /*slaney_area_norm=*/ false, + /*scale=*/ 1.0f, + /*use_htk=*/ true); + + const int n_fft_bins = 1 + (n_fft / 2); + features.assign((size_t) n_frames * (size_t) n_mel, 0.0f); + + const int n_threads = n_frames >= 128 ? std::min(4, n_frames) : 1; + auto worker = [&](int ith) { + std::vector fft_in((size_t) n_fft * 2, 0.0f); + std::vector fft_out((size_t) n_fft * 2 * 2 * 2, 0.0f); + std::vector magnitudes((size_t) n_fft_bins, 0.0f); + + for (int frame = ith; frame < n_frames; frame += n_threads) { + const int frame_start = frame * hop_len; + const int frame_end_mask = frame_start + frame_size_for_unfold - 1; + if (frame_end_mask < 0 || frame_end_mask >= (int) padded_mask.size() || padded_mask[frame_end_mask] == 0) { + continue; + } + + std::fill(fft_in.begin(), fft_in.end(), 0.0f); + for (int i = 0; i < window_len; ++i) { + fft_in[i] = padded_samples[(size_t) frame_start + (size_t) i] * cache.hann_window[(size_t) i]; + } + + fft(cache, fft_in.data(), n_fft, fft_out.data()); + + for (int i = 0; i < n_fft_bins; ++i) { + const float re = fft_out[(size_t) i * 2 + 0]; + const float im = fft_out[(size_t) i * 2 + 1]; + magnitudes[(size_t) i] = sqrtf(re * re + im * im); + } + + for (int mel = 0; mel < n_mel; ++mel) { + double sum = 0.0; + for (int i = 0; i < n_fft_bins; ++i) { + sum += (double) magnitudes[(size_t) i] * + (double) cache.filters.data[(size_t) mel * (size_t) n_fft_bins + (size_t) i]; + } + features[(size_t) frame * (size_t) n_mel + (size_t) mel] = (float) log(sum + 0.001); + } + } + }; + + std::vector workers; + workers.reserve((size_t) std::max(0, n_threads - 1)); + for (int ith = 1; ith < n_threads; ++ith) { + workers.emplace_back(worker, ith); + } + worker(0); + for (auto & thread : workers) { + thread.join(); + } + + n_frames_out = n_frames; + return true; +} + void mtmd_audio_preprocessor_gemma4a::initialize() { cache.fill_sin_cos_table(hparams.audio_n_fft); diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index e116e10c121..862e2fbe9fe 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -111,6 +111,17 @@ bool mtmd_audio_compute_log_mel_spectrogram(const float * samples, bool use_natural_log, bool norm_per_feature, mtmd_audio_mel & out); + +// Gemma4 audio frontend features. Output layout: [n_frames, n_mel]. +bool mtmd_audio_compute_gemma4_features(const float * samples, + size_t n_samples, + int sample_rate, + int n_mel, + int n_fft, + int window_len, + int hop_len, + std::vector & features, + int & n_frames_out); struct mtmd_audio_preprocessor_qwen3a : mtmd_audio_preprocessor { mtmd_audio_preprocessor_qwen3a(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} void initialize() override; diff --git a/tools/mtmd/mtmd-cli-smt.cpp b/tools/mtmd/mtmd-cli-smt.cpp index 783ba7ed750..4fae9ddf44a 100644 --- a/tools/mtmd/mtmd-cli-smt.cpp +++ b/tools/mtmd/mtmd-cli-smt.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -166,6 +167,10 @@ static bool arch_is_qwen3asr(const std::string & arch_name) { return contains_icase(arch_name, "qwen3asr"); } +static bool arch_is_gemma4_audio(const std::string & arch_name) { + return arch_name == "Gemma4Audio"; +} + static std::pair infer_image_grid_xy(int n_tokens) { if (n_tokens <= 0) { return {0, 0}; @@ -378,13 +383,19 @@ resolve_image_boundary_tokens(llama_context * lctx, static std::pair, std::vector> resolve_audio_boundary_tokens(llama_context * lctx, const std::string & arch_name) { - if (!arch_is_qwen3asr(arch_name)) { - return {}; + if (arch_is_qwen3asr(arch_name)) { + return { + tokenize_exact_special(lctx, "<|audio_start|>"), + tokenize_exact_special(lctx, "<|audio_end|>") + }; } - return { - tokenize_exact_special(lctx, "<|audio_start|>"), - tokenize_exact_special(lctx, "<|audio_end|>") - }; + if (arch_is_gemma4_audio(arch_name)) { + return { + tokenize_exact_special(lctx, "<|audio>"), + tokenize_exact_special(lctx, "") + }; + } + return {}; } static void replace_all(std::string & s, const std::string & from, const std::string & to) { @@ -832,7 +843,7 @@ struct mtmd_cli_smt_context { if (try_audio) { try { - smt_audio_ctx = smt_audio_context::create(smt_config_dir); + smt_audio_ctx = smt_audio_context::create(smt_config_dir, params.warmup); if (hidden_size == 0) { hidden_size = smt_audio_ctx->hidden_size(); } else if (hidden_size != smt_audio_ctx->hidden_size()) { @@ -1068,6 +1079,23 @@ static std::string format_qwen3asr_audio_prompt(const mtmd_cli_smt_context & ctx return prompt; } +static std::string format_gemma4_audio_prompt(const mtmd_cli_smt_context & ctx, const common_chat_msg & msg) { + std::string prompt; + prompt.reserve(msg.content.size() + ctx.pending_media.size() * 16 + 128); + prompt += "<|turn>user\n"; + for (const auto & media : ctx.pending_media) { + GGML_ASSERT(media.type == smt_chunk_type::audio); + prompt += k_media_marker; + } + const std::string user_text = strip_media_markers_from_prompt(msg.content); + if (!user_text.empty()) { + prompt += user_text; + } + prompt += "\n"; + prompt += "<|turn>model\n"; + return prompt; +} + // ============================================================ // Eval message - core multimodal processing // ============================================================ @@ -1100,10 +1128,19 @@ static int eval_message_smt(mtmd_cli_smt_context & ctx, common_chat_msg & msg) { ctx.smt_audio_ctx && arch_is_qwen3asr(ctx.smt_audio_ctx->architecture()) && ctx.has_pending_audio_only(); + const bool use_gemma4_audio_prompt = + msg.role == "user" && + ctx.smt_audio_ctx && + arch_is_gemma4_audio(ctx.smt_audio_ctx->architecture()) && + ctx.has_pending_audio_only(); if (use_qwen3asr_prompt) { formatted_chat = format_qwen3asr_audio_prompt(ctx, msg); ctx.chat_history.push_back(msg); add_bos = false; + } else if (use_gemma4_audio_prompt) { + formatted_chat = format_gemma4_audio_prompt(ctx, msg); + ctx.chat_history.push_back(msg); + add_bos = ctx.chat_history.size() == 1; } else { formatted_chat = chat_add_and_format(ctx, msg); } @@ -1267,11 +1304,11 @@ int mtmd_cli_smt_run(int argc, char ** argv, common_params params) { return 1; } - mtmd_cli_smt_context ctx(params, params.smt_config_dir); - bool is_single_turn = !params.prompt.empty() && !params.image.empty(); int n_predict = params.n_predict < 0 ? INT_MAX : params.n_predict; + mtmd_cli_smt_context ctx(params, params.smt_config_dir); + // Ctrl+C handling { #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) diff --git a/tools/mtmd/smt-audio-wrapper.cpp b/tools/mtmd/smt-audio-wrapper.cpp index e9ea0e96dc4..7335b5d766c 100644 --- a/tools/mtmd/smt-audio-wrapper.cpp +++ b/tools/mtmd/smt-audio-wrapper.cpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include #include @@ -37,8 +39,12 @@ extern const OrtApi * g_ort; namespace { +constexpr int32_t k_gemma4_default_warmup_feature_frames = 135; +constexpr const char * k_gemma4_audio_architecture = "Gemma4Audio"; + struct smt_audio_config { std::vector architectures; + std::string encoder_model_path; std::string frontend_model_path; std::string backend_model_path; std::unordered_map ep_config; @@ -53,8 +59,39 @@ struct smt_audio_config { int32_t inter_thread_num = 1; int32_t lfr_m = 0; int32_t lfr_n = 0; + int32_t feature_frames = 0; +}; + +struct gemma4_encoder_input_shape { + int32_t feature_frames = 0; + int32_t num_mel_bins = 0; + + bool dynamic_feature_frames() const { + return feature_frames <= 0; + } +}; + +struct gemma4_encoder_session { + explicit gemma4_encoder_session(Ort::Session && session_in) : + session(std::move(session_in)) {} + + std::string model_path; + Ort::Session session{ nullptr }; + std::vector input_names; + std::vector output_names; + std::vector input_names_raw; + std::vector output_names_raw; + gemma4_encoder_input_shape input_shape; }; +static bool arch_is_gemma4_audio(const std::string & arch_name) { + return arch_name == k_gemma4_audio_architecture; +} + +static bool uses_gemma4_single_encoder(const smt_audio_config & config, const std::string & arch_name) { + return !config.encoder_model_path.empty() && arch_is_gemma4_audio(arch_name); +} + static std::string read_file_to_string(const std::string & path) { std::ifstream file(path); if (!file.is_open()) { @@ -88,6 +125,16 @@ static std::string trim_ascii(std::string value) { return value; } +static std::string join_path(const std::string & dir, const std::string & name) { + if (dir.empty() || dir == ".") { + return name; + } + if (dir.back() == '/' || dir.back() == '\\') { + return dir + name; + } + return dir + "/" + name; +} + static std::string extract_string_value(const std::string & text, const std::string & key) { const std::string marker = "\"" + key + "\""; const size_t key_pos = text.find(marker); @@ -327,6 +374,13 @@ static bool parse_audio_config_block(const std::string & config_dir, const std::string audio_block = content.substr(audio_block_start, audio_block_end - audio_block_start + 1); warn_legacy_spacemit_ep_config_if_needed(audio_block, "audio_model"); + config.encoder_model_path = normalize_path(config_dir, extract_string_value(audio_block, "encoder_model_path")); + if (config.encoder_model_path.empty()) { + config.encoder_model_path = normalize_path(config_dir, extract_string_value(audio_block, "encoder_path")); + } + if (config.encoder_model_path.empty()) { + config.encoder_model_path = normalize_path(config_dir, extract_string_value(audio_block, "onnx_model_path")); + } config.frontend_model_path = normalize_path(config_dir, extract_string_value(audio_block, "frontend_model_path")); if (config.frontend_model_path.empty()) { config.frontend_model_path = normalize_path(config_dir, extract_string_value(audio_block, "frontend_path")); @@ -348,6 +402,7 @@ static bool parse_audio_config_block(const std::string & config_dir, config.hop_len = (int32_t) extract_int64_value(audio_block, "hop_len", config.hop_len); config.lfr_m = (int32_t) extract_int64_value(audio_block, "lfr_m", config.lfr_m); config.lfr_n = (int32_t) extract_int64_value(audio_block, "lfr_n", config.lfr_n); + config.feature_frames = (int32_t) extract_int64_value(audio_block, "feature_frames", config.feature_frames); config.ep_config = extract_string_map(audio_block, "ep_config"); apply_legacy_spacemit_ep_config(audio_block, config); config.architectures = extract_string_array(content, "architectures"); @@ -527,6 +582,7 @@ static void append_optional_spacemit_ep(Ort::SessionOptions & session_options const char * session_name, const smt_audio_config & config) { std::unordered_map provider_options = config.ep_config; + // Add defaults if not specified in ep_config if (provider_options.find("SPACEMIT_EP_INTRA_THREAD_NUM") == provider_options.end()) { provider_options["SPACEMIT_EP_INTRA_THREAD_NUM"] = std::to_string(config.intra_thread_num); @@ -573,12 +629,36 @@ static std::vector get_io_names(Ort::Session & session, bool inputs return names; } +static gemma4_encoder_input_shape get_gemma4_encoder_input_shape(const Ort::Session & session) { + for (size_t i = 0; i < session.GetInputCount(); ++i) { + const auto input_info = session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo(); + const auto shape = input_info.GetShape(); + if (shape.size() != 3) { + continue; + } + + gemma4_encoder_input_shape result; + result.feature_frames = shape[1] > 0 ? (int32_t) shape[1] : 0; + result.num_mel_bins = shape[2] > 0 ? (int32_t) shape[2] : 0; + return result; + } + + return {}; +} + static Ort::Value make_tensor_f32(const std::vector & shape, std::vector & data) { Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); return Ort::Value::CreateTensor(memory_info, data.data(), data.size(), shape.data(), shape.size()); } +static Ort::Value make_tensor_bool(const std::vector & shape, std::vector & data) { + Ort::MemoryInfo memory_info = + Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + return Ort::Value::CreateTensor(memory_info, data.data(), data.size() * sizeof(uint8_t), shape.data(), shape.size(), + ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL); +} + } // namespace struct smt_audio_context::impl { @@ -589,6 +669,11 @@ struct smt_audio_context::impl { Ort::Session frontend_session{ nullptr }; Ort::Session backend_session{ nullptr }; + bool warmup_encoder_sessions = false; + std::mutex encoder_session_mutex; + std::string encoder_session_model_path; + std::unique_ptr encoder_session; + std::vector frontend_input_names; std::vector frontend_output_names; std::vector frontend_input_names_raw; @@ -600,6 +685,76 @@ struct smt_audio_context::impl { std::vector backend_output_names_raw; std::string arch_name; + + void warmup_gemma4_encoder_session(gemma4_encoder_session & encoder) const { + std::cerr << "[SMT][audio] warmup encoder ONNX session"; + if (!arch_name.empty()) { + std::cerr << " for " << arch_name; + } + std::cerr << ": " << encoder.model_path << "\n"; + + const int32_t warmup_frames = + encoder.input_shape.dynamic_feature_frames() ? + (config.feature_frames > 0 ? config.feature_frames : k_gemma4_default_warmup_feature_frames) : + encoder.input_shape.feature_frames; + std::vector feature_data((size_t) warmup_frames * config.num_mel_bins, 0.0f); + std::vector feature_mask((size_t) warmup_frames, 1); + const std::vector feature_shape = { 1, warmup_frames, config.num_mel_bins }; + const std::vector mask_shape = { 1, warmup_frames }; + auto feature_tensor = make_tensor_f32(feature_shape, feature_data); + auto mask_tensor = make_tensor_bool(mask_shape, feature_mask); + std::array inputs = { std::move(feature_tensor), std::move(mask_tensor) }; + (void) encoder.session.Run(Ort::RunOptions{ nullptr }, encoder.input_names_raw.data(), inputs.data(), + inputs.size(), encoder.output_names_raw.data(), encoder.output_names_raw.size()); + } + + gemma4_encoder_session & get_gemma4_encoder_session(const std::string & model_path) { + if (encoder_session && encoder_session_model_path == model_path) { + return *encoder_session; + } + encoder_session.reset(); + encoder_session_model_path.clear(); + + Ort::SessionOptions encoder_options; + encoder_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); + const bool ep_affinity_is_configured = has_spacemit_ep_affinity(config); + if (!ep_affinity_is_configured) { + encoder_options.SetIntraOpNumThreads(config.intra_thread_num); + encoder_options.SetInterOpNumThreads(config.inter_thread_num); + } + + std::cerr << "[SMT][audio] using ORT CPU for Gemma4 encoder: " << model_path << "\n"; + auto encoder = std::make_unique(Ort::Session(env, model_path.c_str(), encoder_options)); + encoder->model_path = model_path; + encoder->input_names = get_io_names(encoder->session, true); + encoder->output_names = get_io_names(encoder->session, false); + encoder->input_names_raw = make_name_ptrs(encoder->input_names); + encoder->output_names_raw = make_name_ptrs(encoder->output_names); + encoder->input_shape = get_gemma4_encoder_input_shape(encoder->session); + + if (encoder->input_names_raw.size() != 2 || encoder->output_names_raw.size() != 2) { + throw std::runtime_error("Unexpected SMT audio single-encoder ONNX IO signature"); + } + if (encoder->input_shape.num_mel_bins <= 0) { + encoder->input_shape.num_mel_bins = config.num_mel_bins; + } + if (encoder->input_shape.feature_frames <= 0 && config.feature_frames > 0) { + encoder->input_shape.feature_frames = config.feature_frames; + } + if (encoder->input_shape.num_mel_bins > 0 && encoder->input_shape.num_mel_bins != config.num_mel_bins) { + throw std::runtime_error("Gemma4 audio encoder num_mel_bins mismatch: config " + + std::to_string(config.num_mel_bins) + ", ONNX " + + std::to_string(encoder->input_shape.num_mel_bins)); + } + + if (warmup_encoder_sessions) { + warmup_gemma4_encoder_session(*encoder); + } + + encoder_session_model_path = model_path; + encoder_session = std::move(encoder); + return *encoder_session; + } }; smt_audio_context::~smt_audio_context() = default; @@ -619,11 +774,22 @@ std::unique_ptr smt_audio_context::create(const std::string & d.arch_name = "Qwen3ASRForConditionalGeneration"; } - if (d.config.frontend_model_path.empty() || d.config.backend_model_path.empty()) { - throw std::runtime_error("Missing SMT audio frontend/backend model path"); + const bool gemma4_single_encoder = uses_gemma4_single_encoder(d.config, d.arch_name); + if (!d.config.encoder_model_path.empty() && !gemma4_single_encoder) { + throw std::runtime_error("SMT audio encoder_model_path is currently supported only for Gemma4 audio models"); } - if (d.config.d_model <= 0 || d.config.hidden_size <= 0) { - throw std::runtime_error("Invalid SMT audio model dimensions"); + + if (gemma4_single_encoder) { + if (d.config.hidden_size <= 0) { + throw std::runtime_error("Invalid Gemma4 audio single-encoder hidden_size"); + } + } else { + if (d.config.frontend_model_path.empty() || d.config.backend_model_path.empty()) { + throw std::runtime_error("Missing SMT audio frontend/backend model path"); + } + if (d.config.d_model <= 0 || d.config.hidden_size <= 0) { + throw std::runtime_error("Invalid SMT audio model dimensions"); + } } onnxruntime::g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION); @@ -642,105 +808,109 @@ std::unique_ptr smt_audio_context::create(const std::string & << " to avoid conflicting with EP-managed affinity\n"; } - append_optional_spacemit_ep(d.frontend_options, "frontend", d.config); - append_optional_spacemit_ep(d.backend_options, "backend", d.config); - - d.frontend_session = Ort::Session(d.env, d.config.frontend_model_path.c_str(), d.frontend_options); - d.backend_session = Ort::Session(d.env, d.config.backend_model_path.c_str(), d.backend_options); - - d.frontend_input_names = get_io_names(d.frontend_session, true); - d.frontend_output_names = get_io_names(d.frontend_session, false); - d.frontend_input_names_raw = make_name_ptrs(d.frontend_input_names); - d.frontend_output_names_raw = make_name_ptrs(d.frontend_output_names); - - d.backend_input_names = get_io_names(d.backend_session, true); - d.backend_output_names = get_io_names(d.backend_session, false); - d.backend_input_names_raw = make_name_ptrs(d.backend_input_names); - d.backend_output_names_raw = make_name_ptrs(d.backend_output_names); - - if (d.frontend_input_names_raw.size() != 1 || d.frontend_output_names_raw.size() != 1 || - (d.backend_input_names_raw.size() != 1 && d.backend_input_names_raw.size() != 2) || - d.backend_output_names_raw.size() != 1) { - throw std::runtime_error("Unexpected SMT audio ONNX IO signature"); - } - - if (warmup) { - std::cerr << "[SMT][audio] warmup ONNX sessions"; - if (!d.arch_name.empty()) { - std::cerr << " for " << d.arch_name; + if (gemma4_single_encoder) { + d.warmup_encoder_sessions = warmup; + } else { + append_optional_spacemit_ep(d.frontend_options, "frontend", d.config); + append_optional_spacemit_ep(d.backend_options, "backend", d.config); + + d.frontend_session = Ort::Session(d.env, d.config.frontend_model_path.c_str(), d.frontend_options); + d.backend_session = Ort::Session(d.env, d.config.backend_model_path.c_str(), d.backend_options); + + d.frontend_input_names = get_io_names(d.frontend_session, true); + d.frontend_output_names = get_io_names(d.frontend_session, false); + d.frontend_input_names_raw = make_name_ptrs(d.frontend_input_names); + d.frontend_output_names_raw = make_name_ptrs(d.frontend_output_names); + + d.backend_input_names = get_io_names(d.backend_session, true); + d.backend_output_names = get_io_names(d.backend_session, false); + d.backend_input_names_raw = make_name_ptrs(d.backend_input_names); + d.backend_output_names_raw = make_name_ptrs(d.backend_output_names); + + if (d.frontend_input_names_raw.size() != 1 || d.frontend_output_names_raw.size() != 1 || + (d.backend_input_names_raw.size() != 1 && d.backend_input_names_raw.size() != 2) || + d.backend_output_names_raw.size() != 1) { + throw std::runtime_error("Unexpected SMT audio ONNX IO signature"); } - std::cerr << "\n"; - int warmup_t_out; - std::vector warmup_hidden; + if (warmup) { + std::cerr << "[SMT][audio] warmup ONNX sessions"; + if (!d.arch_name.empty()) { + std::cerr << " for " << d.arch_name; + } + std::cerr << "\n"; - if (d.config.lfr_m > 0) { - const int warmup_frames = 10; - const int feat_dim = d.config.num_mel_bins * d.config.lfr_m; + int warmup_t_out; + std::vector warmup_hidden; - std::vector frontend_input_data((size_t) warmup_frames * feat_dim, 0.0f); - const std::vector frontend_input_shape = { 1, warmup_frames, (int64_t) feat_dim }; - auto frontend_input = make_tensor_f32(frontend_input_shape, frontend_input_data); + if (d.config.lfr_m > 0) { + const int warmup_frames = 10; + const int feat_dim = d.config.num_mel_bins * d.config.lfr_m; - std::cerr << "[SMT][audio] warmup frontend ONNX session (FunASR): " << d.config.frontend_model_path << "\n"; - auto frontend_outputs = d.frontend_session.Run(Ort::RunOptions{ nullptr }, d.frontend_input_names_raw.data(), - &frontend_input, 1, d.frontend_output_names_raw.data(), 1); - if (frontend_outputs.empty()) { - throw std::runtime_error("SMT audio warmup frontend returned no outputs"); - } - const auto frontend_output_info = frontend_outputs[0].GetTensorTypeAndShapeInfo(); - if (frontend_output_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { - throw std::runtime_error("SMT audio warmup frontend output must be float32"); - } - auto shape = frontend_output_info.GetShape(); - warmup_t_out = (int) shape[1]; - warmup_hidden.resize((size_t) warmup_t_out * (size_t) d.config.d_model, 0.0f); - const float * frontend_output = frontend_outputs[0].GetTensorData(); - std::memcpy(warmup_hidden.data(), frontend_output, warmup_hidden.size() * sizeof(float)); - } else { - const int chunk_frames = 100; - const int chunk_tokens = 13; - warmup_t_out = chunk_tokens; - - std::vector frontend_input_data((size_t) d.config.num_mel_bins * chunk_frames, 0.0f); - const std::vector frontend_input_shape = { 1, d.config.num_mel_bins, chunk_frames }; - auto frontend_input = make_tensor_f32(frontend_input_shape, frontend_input_data); - - std::cerr << "[SMT][audio] warmup frontend ONNX session: " << d.config.frontend_model_path << "\n"; - auto frontend_outputs = d.frontend_session.Run(Ort::RunOptions{ nullptr }, d.frontend_input_names_raw.data(), - &frontend_input, 1, d.frontend_output_names_raw.data(), 1); - - warmup_hidden.resize((size_t) warmup_t_out * (size_t) d.config.d_model, 0.0f); - if (frontend_outputs.empty()) { - throw std::runtime_error("SMT audio warmup frontend returned no outputs"); - } - const auto frontend_output_info = frontend_outputs[0].GetTensorTypeAndShapeInfo(); - if (frontend_output_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { - throw std::runtime_error("SMT audio warmup frontend output must be float32"); - } - const int64_t frontend_output_elems = frontend_output_info.GetElementCount(); - if (frontend_output_elems < 0 || (size_t) frontend_output_elems < warmup_hidden.size()) { - throw std::runtime_error("SMT audio warmup frontend output is smaller than expected"); - } - const float * frontend_output = frontend_outputs[0].GetTensorData(); - std::memcpy(warmup_hidden.data(), frontend_output, warmup_hidden.size() * sizeof(float)); - } + std::vector frontend_input_data((size_t) warmup_frames * feat_dim, 0.0f); + const std::vector frontend_input_shape = { 1, warmup_frames, (int64_t) feat_dim }; + auto frontend_input = make_tensor_f32(frontend_input_shape, frontend_input_data); - const std::vector backend_hidden_shape = { 1, warmup_t_out, d.config.d_model }; - auto hidden_tensor = make_tensor_f32(backend_hidden_shape, warmup_hidden); + std::cerr << "[SMT][audio] warmup frontend ONNX session (FunASR): " << d.config.frontend_model_path << "\n"; + auto frontend_outputs = d.frontend_session.Run(Ort::RunOptions{ nullptr }, d.frontend_input_names_raw.data(), + &frontend_input, 1, d.frontend_output_names_raw.data(), 1); + if (frontend_outputs.empty()) { + throw std::runtime_error("SMT audio warmup frontend returned no outputs"); + } + const auto frontend_output_info = frontend_outputs[0].GetTensorTypeAndShapeInfo(); + if (frontend_output_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + throw std::runtime_error("SMT audio warmup frontend output must be float32"); + } + auto shape = frontend_output_info.GetShape(); + warmup_t_out = (int) shape[1]; + warmup_hidden.resize((size_t) warmup_t_out * (size_t) d.config.d_model, 0.0f); + const float * frontend_output = frontend_outputs[0].GetTensorData(); + std::memcpy(warmup_hidden.data(), frontend_output, warmup_hidden.size() * sizeof(float)); + } else { + const int chunk_frames = 100; + const int chunk_tokens = 13; + warmup_t_out = chunk_tokens; + + std::vector frontend_input_data((size_t) d.config.num_mel_bins * chunk_frames, 0.0f); + const std::vector frontend_input_shape = { 1, d.config.num_mel_bins, chunk_frames }; + auto frontend_input = make_tensor_f32(frontend_input_shape, frontend_input_data); + + std::cerr << "[SMT][audio] warmup frontend ONNX session: " << d.config.frontend_model_path << "\n"; + auto frontend_outputs = d.frontend_session.Run(Ort::RunOptions{ nullptr }, d.frontend_input_names_raw.data(), + &frontend_input, 1, d.frontend_output_names_raw.data(), 1); + + warmup_hidden.resize((size_t) warmup_t_out * (size_t) d.config.d_model, 0.0f); + if (frontend_outputs.empty()) { + throw std::runtime_error("SMT audio warmup frontend returned no outputs"); + } + const auto frontend_output_info = frontend_outputs[0].GetTensorTypeAndShapeInfo(); + if (frontend_output_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + throw std::runtime_error("SMT audio warmup frontend output must be float32"); + } + const int64_t frontend_output_elems = frontend_output_info.GetElementCount(); + if (frontend_output_elems < 0 || (size_t) frontend_output_elems < warmup_hidden.size()) { + throw std::runtime_error("SMT audio warmup frontend output is smaller than expected"); + } + const float * frontend_output = frontend_outputs[0].GetTensorData(); + std::memcpy(warmup_hidden.data(), frontend_output, warmup_hidden.size() * sizeof(float)); + } - std::cerr << "[SMT][audio] warmup backend ONNX session: " << d.config.backend_model_path << "\n"; - if (d.backend_input_names_raw.size() == 2) { - std::vector attention_mask((size_t) warmup_t_out * (size_t) warmup_t_out, 0.0f); - const std::vector backend_mask_shape = { 1, 1, warmup_t_out, warmup_t_out }; - auto mask_tensor = make_tensor_f32(backend_mask_shape, attention_mask); - std::array backend_inputs = { std::move(hidden_tensor), std::move(mask_tensor) }; - (void) d.backend_session.Run(Ort::RunOptions{ nullptr }, d.backend_input_names_raw.data(), - backend_inputs.data(), backend_inputs.size(), - d.backend_output_names_raw.data(), 1); - } else { - (void) d.backend_session.Run(Ort::RunOptions{ nullptr }, d.backend_input_names_raw.data(), - &hidden_tensor, 1, d.backend_output_names_raw.data(), 1); + const std::vector backend_hidden_shape = { 1, warmup_t_out, d.config.d_model }; + auto hidden_tensor = make_tensor_f32(backend_hidden_shape, warmup_hidden); + + std::cerr << "[SMT][audio] warmup backend ONNX session: " << d.config.backend_model_path << "\n"; + if (d.backend_input_names_raw.size() == 2) { + std::vector attention_mask((size_t) warmup_t_out * (size_t) warmup_t_out, 0.0f); + const std::vector backend_mask_shape = { 1, 1, warmup_t_out, warmup_t_out }; + auto mask_tensor = make_tensor_f32(backend_mask_shape, attention_mask); + std::array backend_inputs = { std::move(hidden_tensor), std::move(mask_tensor) }; + (void) d.backend_session.Run(Ort::RunOptions{ nullptr }, d.backend_input_names_raw.data(), + backend_inputs.data(), backend_inputs.size(), + d.backend_output_names_raw.data(), 1); + } else { + (void) d.backend_session.Run(Ort::RunOptions{ nullptr }, d.backend_input_names_raw.data(), + &hidden_tensor, 1, d.backend_output_names_raw.data(), 1); + } } } @@ -762,6 +932,112 @@ std::vector smt_audio_context::encode_audio(const std::string & audio_pat } ggml_trace_log_end("decode_audio_file", "Audio", NULL); + if (uses_gemma4_single_encoder(d.config, d.arch_name)) { + std::vector features; + int n_feature_frames = 0; + ggml_trace_log_begin("compute_gemma4_features", "Audio", NULL); + if (!mtmd_audio_compute_gemma4_features(samples.data(), samples.size(), d.config.sample_rate, + d.config.num_mel_bins, d.config.n_fft, d.config.window_len, + d.config.hop_len, features, n_feature_frames)) { + ggml_trace_log_end("compute_gemma4_features", "Audio", NULL); + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("failed to compute Gemma4 audio features"); + } + ggml_trace_log_end("compute_gemma4_features", "Audio", NULL); + + const std::string & encoder_model_path = d.config.encoder_model_path; + if (encoder_model_path.empty()) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("Gemma4 audio encoder_model_path is empty"); + } + + std::lock_guard encoder_lock(d.encoder_session_mutex); + auto & encoder = d.get_gemma4_encoder_session(encoder_model_path); + + const int32_t encoder_feature_frames = + encoder.input_shape.dynamic_feature_frames() ? n_feature_frames : encoder.input_shape.feature_frames; + if (encoder_feature_frames <= 0) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("Gemma4 audio feature length must be positive"); + } + + if (!encoder.input_shape.dynamic_feature_frames() && n_feature_frames > encoder_feature_frames) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("Gemma4 audio feature length " + std::to_string(n_feature_frames) + + " exceeds ONNX static feature_frames " + + std::to_string(encoder_feature_frames) + " for " + encoder.model_path); + } + + std::vector feature_data((size_t) encoder_feature_frames * d.config.num_mel_bins, 0.0f); + for (int frame = 0; frame < n_feature_frames; ++frame) { + std::memcpy(feature_data.data() + (size_t) frame * d.config.num_mel_bins, + features.data() + (size_t) frame * d.config.num_mel_bins, + (size_t) d.config.num_mel_bins * sizeof(float)); + } + std::vector feature_mask((size_t) encoder_feature_frames, 0); + std::fill(feature_mask.begin(), feature_mask.begin() + n_feature_frames, (uint8_t) 1); + + const std::vector feature_shape = { 1, encoder_feature_frames, d.config.num_mel_bins }; + const std::vector mask_shape = { 1, encoder_feature_frames }; + auto feature_tensor = make_tensor_f32(feature_shape, feature_data); + auto mask_tensor = make_tensor_bool(mask_shape, feature_mask); + std::array inputs = { std::move(feature_tensor), std::move(mask_tensor) }; + + ggml_trace_log_begin("encoder_session_run", "Audio", NULL); + auto outputs = encoder.session.Run(Ort::RunOptions{ nullptr }, encoder.input_names_raw.data(), + inputs.data(), inputs.size(), encoder.output_names_raw.data(), + encoder.output_names_raw.size()); + ggml_trace_log_end("encoder_session_run", "Audio", NULL); + + if (outputs.size() < 2) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("Gemma4 audio encoder returned fewer than 2 outputs"); + } + + const auto embd_info = outputs[0].GetTensorTypeAndShapeInfo(); + if (embd_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("Gemma4 audio_embeddings output must be float32"); + } + const int64_t embd_elems = embd_info.GetElementCount(); + if (embd_elems <= 0 || embd_elems % d.config.hidden_size != 0) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("invalid Gemma4 audio_embeddings shape"); + } + const int n_audio_tokens = (int) (embd_elems / d.config.hidden_size); + + const auto mask_info = outputs[1].GetTensorTypeAndShapeInfo(); + if (mask_info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL || + mask_info.GetElementCount() < (size_t) n_audio_tokens) { + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + throw std::runtime_error("invalid Gemma4 audio_token_mask shape"); + } + + const float * embd = outputs[0].GetTensorData(); + const bool * mask = outputs[1].GetTensorData(); + std::vector audio_embd; + audio_embd.reserve((size_t) n_audio_tokens * (size_t) d.config.hidden_size); + for (int token = 0; token < n_audio_tokens; ++token) { + if (!mask[token]) { + continue; + } + const float * src = embd + (size_t) token * (size_t) d.config.hidden_size; + audio_embd.insert(audio_embd.end(), src, src + d.config.hidden_size); + } + + ggml_trace_log_end("encode_audio", "Audio", NULL); + ggml_profile_flush_tls(); + return audio_embd; + } + int t_out; std::vector hidden_states; diff --git a/tools/mtmd/smt-audio-wrapper.h b/tools/mtmd/smt-audio-wrapper.h index bea3aad52fa..c5fca5ddaa7 100644 --- a/tools/mtmd/smt-audio-wrapper.h +++ b/tools/mtmd/smt-audio-wrapper.h @@ -15,7 +15,7 @@ struct smt_audio_context { // 2. a split-encoder metadata json generated by the export script. static std::unique_ptr create(const std::string & config_dir, bool warmup = true); - // Encode an audio file using the split frontend/backend ONNX pipeline. + // Encode an audio file using the configured SMT audio ONNX pipeline. // Returns audio embedding vector (n_tokens * hidden_size floats). std::vector encode_audio(const std::string & audio_path); diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 7c5420a5658..6d436269abb 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1087,6 +1087,55 @@ static common_chat_params build_qwen3asr_audio_chat_params(const common_chat_tem return params; } +static common_chat_templates_inputs remove_media_from_chat_inputs(const common_chat_templates_inputs & inputs) { + common_chat_templates_inputs out = inputs; + for (auto & msg : out.messages) { + msg.content = collect_message_text_without_media(msg); + msg.content_parts.clear(); + } + return out; +} + +static common_chat_params build_gemma4_audio_chat_params(const server_chat_params & opt, + const common_chat_templates_inputs & inputs) { + common_chat_params params = common_chat_templates_apply(opt.tmpls.get(), remove_media_from_chat_inputs(inputs)); + + size_t media_marker_count = 0; + std::string user_text; + + for (const auto & msg : inputs.messages) { + bool has_media = false; + std::string text = collect_message_text_without_media(msg, &has_media); + media_marker_count += count_message_media_markers(msg); + if ((msg.role == "system" || msg.role == "developer") && !text.empty()) { + if (!user_text.empty()) { + user_text += "\n"; + } + user_text += text; + } else if (msg.role == "user" && has_media) { + if (!user_text.empty() && !text.empty()) { + user_text += "\n"; + } + user_text += std::move(text); + } + } + + params.prompt = "<|turn>user\n"; + for (size_t i = 0; i < media_marker_count; ++i) { + params.prompt += "<__media__>"; + } + if (!user_text.empty()) { + params.prompt += user_text; + } + params.prompt += "\n"; + if (inputs.add_generation_prompt) { + params.prompt += params.generation_prompt.empty() ? "<|turn>model\n" : params.generation_prompt; + } + params.message_spans.clear(); + + return params; +} + static bool should_use_qwen3asr_audio_prompt(const server_chat_params & opt, const std::vector & out_files) { if (out_files.empty() || !opt.allow_audio || opt.media_backend != "smt" || opt.tmpls == nullptr) { @@ -1098,6 +1147,21 @@ static bool should_use_qwen3asr_audio_prompt(const server_chat_params & opt tmpl_src.find("<|im_start|>assistant") != std::string::npos && tmpl_src.find("media_marker") == std::string::npos; } + +static bool should_use_gemma4_audio_prompt(const server_chat_params & opt, + bool has_audio_media, + bool has_image_media, + const std::vector & out_files) { + const bool has_audio_input = has_audio_media || (!has_image_media && !out_files.empty()); + if (!has_audio_input || has_image_media || !opt.allow_audio || opt.media_backend != "smt" || opt.tmpls == nullptr) { + return false; + } + + const std::string tmpl_src = common_chat_templates_source(opt.tmpls.get(), ""); + return tmpl_src.find("<|turn>model") != std::string::npos && + tmpl_src.find("<|channel>thought") != std::string::npos && + tmpl_src.find("media_marker") == std::string::npos; +} #endif server_tokens process_mtmd_prompt(mtmd_context * mctx, std::string prompt, std::vector files) { @@ -1519,6 +1583,8 @@ json oaicompat_chat_params_parse(json & body, /* openai api } #if defined(LLAMA_SERVER_SMT_VISION) apply_vision_history_mode(messages, vision_history); + bool has_image_media = false; + bool has_audio_media = false; #endif for (auto & msg : messages) { @@ -1554,6 +1620,9 @@ json oaicompat_chat_params_parse(json & body, /* openai api json image_url = json_value(p, "image_url", json::object()); handle_media(out_files, image_url, opt.media_path, opt.image_bin_only); +#if defined(LLAMA_SERVER_SMT_VISION) + has_image_media = true; +#endif p["type"] = "media_marker"; p["text"] = get_media_marker(); @@ -1573,6 +1642,9 @@ json oaicompat_chat_params_parse(json & body, /* openai api } auto decoded_data = base64_decode(data); // expected to be base64 encoded out_files.push_back(decoded_data); +#if defined(LLAMA_SERVER_SMT_VISION) + has_audio_media = true; +#endif // TODO: add audio_url support by reusing handle_media() @@ -1651,6 +1723,8 @@ json oaicompat_chat_params_parse(json & body, /* openai api chat_params = build_paddleocr_chat_params(inputs); } else if (should_use_qwen3asr_audio_prompt(opt, out_files)) { chat_params = build_qwen3asr_audio_chat_params(inputs); + } else if (should_use_gemma4_audio_prompt(opt, has_audio_media, has_image_media, out_files)) { + chat_params = build_gemma4_audio_chat_params(opt, inputs); } else #endif { diff --git a/tools/server/server-smt-vision.cpp b/tools/server/server-smt-vision.cpp index 0cfbcce9dc7..7b0e60ce96a 100644 --- a/tools/server/server-smt-vision.cpp +++ b/tools/server/server-smt-vision.cpp @@ -553,7 +553,7 @@ static lingbot_map_postprocess_result lingbot_postprocess_reconstruction( return result; } -bool server_smt_vision_config_is_lingbot_map(const std::string & config_dir) { +static bool server_smt_vision_config_has_architecture(const std::string & config_dir, const std::string & target) { if (config_dir.empty()) { return false; } @@ -572,12 +572,12 @@ bool server_smt_vision_config_is_lingbot_map(const std::string & config_dir) { const auto & arch = config.at("architectures"); if (arch.is_array()) { for (const auto & value : arch) { - if (value.is_string() && value.get() == "LingBotMapFor3DReconstruction") { + if (value.is_string() && value.get() == target) { return true; } } } else if (arch.is_string()) { - return arch.get() == "LingBotMapFor3DReconstruction"; + return arch.get() == target; } } catch (...) { return false; @@ -585,6 +585,15 @@ bool server_smt_vision_config_is_lingbot_map(const std::string & config_dir) { return false; } +bool server_smt_vision_config_is_lingbot_map(const std::string & config_dir) { + return server_smt_vision_config_has_architecture(config_dir, "LingBotMapFor3DReconstruction"); +} + +static bool server_smt_vision_config_is_audio_only(const std::string & config_dir) { + return server_smt_vision_config_has_architecture(config_dir, "Gemma4Audio") || + server_smt_vision_config_has_architecture(config_dir, "Qwen3ASRForConditionalGeneration"); +} + static std::string fnv_hash(const uint8_t * data, size_t len) { const uint64_t fnv_prime = 0x100000001b3ULL; uint64_t hash = 0xcbf29ce484222325ULL; @@ -614,6 +623,10 @@ static bool arch_is_qwen3asr(const std::string & arch_name) { return contains_icase(arch_name, "qwen3asr"); } +static bool arch_is_gemma4_audio(const std::string & arch_name) { + return arch_name == "Gemma4Audio"; +} + static bool arch_is_funasr(const std::string & arch_name) { return contains_icase(arch_name, "funasr"); } @@ -743,10 +756,13 @@ static std::pair, std::vector> resolve_ima static std::pair, std::vector> resolve_audio_boundary_tokens( llama_context * lctx, const std::string & arch_name) { - if (!arch_is_qwen3asr(arch_name)) { - return {}; + if (arch_is_qwen3asr(arch_name)) { + return { tokenize_exact_special(lctx, "<|audio_start|>"), tokenize_exact_special(lctx, "<|audio_end|>") }; + } + if (arch_is_gemma4_audio(arch_name)) { + return { tokenize_exact_special(lctx, "<|audio>"), tokenize_exact_special(lctx, "") }; } - return { tokenize_exact_special(lctx, "<|audio_start|>"), tokenize_exact_special(lctx, "<|audio_end|>") }; + return {}; } static bool looks_like_audio_file(const std::vector & data) { @@ -968,15 +984,17 @@ server_smt_vision_context * server_smt_vision_init(llama_context * lctx, const s return ctx.release(); } - try { - ctx->smt_vision = smt_vision_context::create(config_dir, warmup); - ctx->hidden_size = (int32_t) ctx->smt_vision->hidden_size(); - primary_architecture = ctx->smt_vision->architecture(); - auto boundaries = resolve_image_boundary_tokens(lctx, primary_architecture); - ctx->tok_img_beg = std::move(boundaries.first); - ctx->tok_img_end = std::move(boundaries.second); - } catch (const std::exception & e) { - LOG_WRN("[server-smt] failed to initialize SMT vision backend from '%s': %s\n", config_dir.c_str(), e.what()); + if (!server_smt_vision_config_is_audio_only(config_dir)) { + try { + ctx->smt_vision = smt_vision_context::create(config_dir, warmup); + ctx->hidden_size = (int32_t) ctx->smt_vision->hidden_size(); + primary_architecture = ctx->smt_vision->architecture(); + auto boundaries = resolve_image_boundary_tokens(lctx, primary_architecture); + ctx->tok_img_beg = std::move(boundaries.first); + ctx->tok_img_end = std::move(boundaries.second); + } catch (const std::exception & e) { + LOG_WRN("[server-smt] failed to initialize SMT vision backend from '%s': %s\n", config_dir.c_str(), e.what()); + } } try {