Skip to content

feat(mtmd): add Gemma4 audio SMT support #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
muggle-stack wants to merge 1 commit into
base: spacemit-mtmd
Choose a base branch
from
+651 −121
Open

feat(mtmd): add Gemma4 audio SMT support #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 114 additions & 0 deletions tools/mtmd/mtmd-audio.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<float> & 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<float> waveform(n_valid + pad_right, 0.0f);
std::copy(samples, samples + n_valid, waveform.data());

std::vector<uint8_t> 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<float> padded_samples((size_t) pad_left + waveform.size(), 0.0f);
std::copy(waveform.begin(), waveform.end(), padded_samples.begin() + pad_left);

std::vector<uint8_t> 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<float> fft_in((size_t) n_fft * 2, 0.0f);
std::vector<float> fft_out((size_t) n_fft * 2 * 2 * 2, 0.0f);
std::vector<float> 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<std::thread> 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);

Expand Down
11 changes: 11 additions & 0 deletions tools/mtmd/mtmd-audio.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<float> & 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;
Expand Down
55 changes: 46 additions & 9 deletions tools/mtmd/mtmd-cli-smt.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#include <cctype>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <stdexcept>
#include <utility>
Expand Down Expand Up @@ -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<int, int> infer_image_grid_xy(int n_tokens) {
if (n_tokens <= 0) {
return {0, 0};
Expand Down Expand Up @@ -378,13 +383,19 @@ resolve_image_boundary_tokens(llama_context * lctx,

static std::pair<std::vector<llama_token>, std::vector<llama_token>>
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, "<audio|>")
};
}
return {};
}

static void replace_all(std::string & s, const std::string & from, const std::string & to) {
Expand Down Expand Up @@ -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()) {
Expand Down Expand Up @@ -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 += "<turn|>\n";
prompt += "<|turn>model\n";
return prompt;
}

// ============================================================
// Eval message - core multimodal processing
// ============================================================
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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__))
Expand Down
Loading