mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-09-20 01:31:42 +02:00
Merge branch 'upstream' into concedo_experimental
# Conflicts: # .github/workflows/release.yml # .github/workflows/server.yml # examples/model-conversion/requirements.txt # examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py # examples/speculative-simple/README.md # examples/speculative-simple/speculative-simple.cpp # ggml/src/ggml-opencl/ggml-opencl.cpp # requirements/requirements-convert_hf_to_gguf.txt # requirements/requirements-convert_lora_to_gguf.txt # scripts/hip/gcn-cdna-vgpr-check.py # tests/test-backend-ops.cpp # tests/test-chat.cpp # tools/imatrix/imatrix.cpp # tools/mtmd/CMakeLists.txt
This commit is contained in:
Concedo
@@ -59,8 +59,10 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i
|
||||
|
||||
### Checklist for porting new audio generation models to mtmd
|
||||
|
||||
1. Establish a list of reusable and missing components from the current mtmd implementation.
|
||||
2. For GGUF conversion:
|
||||
1. Make sure to consult merged PRs about adding new TTS models, especially reviewer comments
|
||||
- Example: https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+mtmd+tts+is%3Amerged
|
||||
2. Establish a list of reusable and missing components from the current mtmd implementation.
|
||||
3. For GGUF conversion:
|
||||
- Backbone model should be converted to a normal text model (loadable via `libllama`)
|
||||
- If model used hard-coded embedding row ID, append them to token embeddings and assign token name for them (see `qwen3tts.py`)
|
||||
- If model have a specific output logits head for audio codes (usually semantic code), keep the head as-is and pad the logits at inference time (see `src/models/qwen3vl.cpp`)
|
||||
@@ -70,12 +72,17 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i
|
||||
- For tensor naming:
|
||||
- Prefixed with `a.*` for tensors used by speaker encoder pipeline
|
||||
- Prefixed with `a.gen.*` for generation stages (code / mel-spectrogram / PCM generation)
|
||||
3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this:
|
||||
- For GGUF metadata:
|
||||
- Reuse as many existing keys as possible
|
||||
- In most cases, you can hard-code model configs in the model graph class, or in `clip_hparams`
|
||||
- If some values need to be exposed to the `mtmd_helper` layer, hard-code them in `mtmd_helper` and distinguish by pipeline and `mtmd_gen_audio_info::model_variant` if necessary
|
||||
- Do NOT add new GGUF metadata or new fields to `mtmd_gen_audio_info` unless you can prove that you absolutely need them
|
||||
4. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this:
|
||||
- 10-20% changes is to add new backbone (text) model and conversion
|
||||
- 60% changes inside `mtmd-helper-gen.cpp`
|
||||
- 10% changes inside `libmtmd` and `clip.cpp` systems
|
||||
- The rest downstream code (CLI, server) should have no changes at all
|
||||
4. Update usage documentation in `tools/tts/README.md`
|
||||
5. Update usage documentation in `tools/tts/README.md`
|
||||
|
||||
IMPORTANT: If your model needs changes that don't fit the existing infrastructure, **open an issue first for discussion**.
|
||||
|
||||
|
||||
+39
-1
@@ -92,7 +92,9 @@
|
||||
#define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size
|
||||
// audio generation (gen-audio)-specific
|
||||
#define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities
|
||||
#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor"
|
||||
// name of the weight variant, for settings that are not in the checkpoint
|
||||
#define KEY_GEN_AUDIO_VARIANT "clip.gen.audio.model_variant"
|
||||
#define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor"
|
||||
|
||||
//
|
||||
// tensor name constants
|
||||
@@ -246,6 +248,38 @@
|
||||
#define TN_A_GEN_WAV_DAC_POST_SNAKE "a.gen.wav.dac.post_snake.%s"
|
||||
#define TN_A_GEN_WAV_DAC_POST_CONV "a.gen.wav.dac.post_conv.%s"
|
||||
|
||||
// pocket-tts
|
||||
#define TN_A_SEANET_CONV_IN "a.seanet.conv_in.%s"
|
||||
#define TN_A_SEANET_CONV_OUT "a.seanet.conv_out.%s"
|
||||
#define TN_A_SEANET_RES_CONV1 "a.seanet.blk.%d.res_conv1.%s"
|
||||
#define TN_A_SEANET_RES_CONV2 "a.seanet.blk.%d.res_conv2.%s"
|
||||
#define TN_A_SEANET_SCALE_CONV "a.seanet.blk.%d.scale_conv.%s"
|
||||
#define TN_A_SPEAKER_PROJ "a.speaker_proj.%s"
|
||||
#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s"
|
||||
#define TN_A_GEN_FLOW_INPUT_PROJ "a.gen.flow.input_proj.%s"
|
||||
#define TN_A_GEN_FLOW_COND_EMBD "a.gen.flow.cond_embd.%s"
|
||||
#define TN_A_GEN_FLOW_TIME_FREQS "a.gen.flow.time.%d.freqs"
|
||||
#define TN_A_GEN_FLOW_TIME_UP "a.gen.flow.time.%d.up.%s"
|
||||
#define TN_A_GEN_FLOW_TIME_DOWN "a.gen.flow.time.%d.down.%s"
|
||||
#define TN_A_GEN_FLOW_TIME_NORM "a.gen.flow.time.%d.norm"
|
||||
#define TN_A_GEN_FLOW_BLK_NORM "a.gen.flow.blk.%d.norm.%s"
|
||||
#define TN_A_GEN_FLOW_BLK_UP "a.gen.flow.blk.%d.up.%s"
|
||||
#define TN_A_GEN_FLOW_BLK_DOWN "a.gen.flow.blk.%d.down.%s"
|
||||
#define TN_A_GEN_FLOW_BLK_ADA "a.gen.flow.blk.%d.ada.%s"
|
||||
#define TN_A_GEN_FLOW_FINAL_ADA "a.gen.flow.final.ada.%s"
|
||||
#define TN_A_GEN_FLOW_FINAL_PROJ "a.gen.flow.final.proj.%s"
|
||||
#define TN_A_GEN_OUT_EOS "a.gen.out_eos.%s"
|
||||
#define TN_A_GEN_INPUT_LINEAR "a.gen.input_linear.%s"
|
||||
#define TN_A_GEN_EMB_MEAN "a.gen.emb_mean"
|
||||
#define TN_A_GEN_EMB_STD "a.gen.emb_std"
|
||||
#define TN_A_GEN_WAV_QUANT_OUT "a.gen.wav.quant_out.%s"
|
||||
#define TN_A_GEN_WAV_UPSAMPLE "a.gen.wav.upsample.%s"
|
||||
#define TN_A_GEN_WAV_SEANET_CONV_IN "a.gen.wav.seanet.conv_in.%s"
|
||||
#define TN_A_GEN_WAV_SEANET_CONV_OUT "a.gen.wav.seanet.conv_out.%s"
|
||||
#define TN_A_GEN_WAV_SEANET_RES_CONV1 "a.gen.wav.seanet.blk.%d.res_conv1.%s"
|
||||
#define TN_A_GEN_WAV_SEANET_RES_CONV2 "a.gen.wav.seanet.blk.%d.res_conv2.%s"
|
||||
#define TN_A_GEN_WAV_SEANET_SCALE_CONV "a.gen.wav.seanet.blk.%d.scale_conv.%s"
|
||||
|
||||
// cogvlm
|
||||
#define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s"
|
||||
#define TN_MM_H_TO_4H "mm.up.%s"
|
||||
@@ -455,6 +489,8 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_MIMO_AUDIO,
|
||||
PROJECTOR_TYPE_QWEN3TTS_SPKENC,
|
||||
PROJECTOR_TYPE_QWEN3TTS_GEN,
|
||||
PROJECTOR_TYPE_POCKETTTS_SPKENC,
|
||||
PROJECTOR_TYPE_POCKETTTS_GEN,
|
||||
PROJECTOR_TYPE_MUSE_GLIMMER,
|
||||
PROJECTOR_TYPE_UNKNOWN,
|
||||
};
|
||||
@@ -515,6 +551,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_PARAKEET, "parakeet"},
|
||||
{ PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"},
|
||||
{ PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"},
|
||||
{ PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"},
|
||||
{ PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"},
|
||||
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
|
||||
};
|
||||
|
||||
|
||||
@@ -141,6 +141,20 @@ struct clip_hparams {
|
||||
int32_t rvq_num_quantizers = 0;
|
||||
std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
|
||||
|
||||
// threshold for the "out_eos_score" graph output
|
||||
float gen_eos_threshold = 0.0f;
|
||||
|
||||
// name of the weight variant, some pipelines tune themselves on it
|
||||
std::string gen_model_variant;
|
||||
|
||||
// pocket-tts
|
||||
static constexpr int32_t pockettts_max_spk_seconds = 30;
|
||||
int32_t seanet_n_stage = 0;
|
||||
std::vector<int32_t> seanet_ratios; // encoder order (reversed compared to the config)
|
||||
int32_t mimi_downsample = 0; // encoder frame rate / model frame rate
|
||||
int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames
|
||||
int32_t flow_n_step = 1; // lsd_decode steps
|
||||
|
||||
// qwen3tts code2wav
|
||||
int32_t wav_tfm_n_layer = 0;
|
||||
int32_t wav_tfm_n_embd = 0;
|
||||
@@ -402,6 +416,63 @@ struct qf_block {
|
||||
std::vector<clip_layer> qf_proj_layers;
|
||||
};
|
||||
|
||||
// pocket-tts SEANet stack, used in both directions:
|
||||
// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out
|
||||
// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out
|
||||
struct clip_seanet {
|
||||
// one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input
|
||||
struct stage {
|
||||
ggml_tensor * res_conv1_w = nullptr;
|
||||
ggml_tensor * res_conv1_b = nullptr;
|
||||
ggml_tensor * res_conv2_w = nullptr;
|
||||
ggml_tensor * res_conv2_b = nullptr;
|
||||
ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder)
|
||||
ggml_tensor * scale_conv_b = nullptr;
|
||||
};
|
||||
|
||||
ggml_tensor * conv_in_w = nullptr;
|
||||
ggml_tensor * conv_in_b = nullptr;
|
||||
ggml_tensor * conv_out_w = nullptr;
|
||||
ggml_tensor * conv_out_b = nullptr;
|
||||
std::vector<stage> stages;
|
||||
};
|
||||
|
||||
// pocket-tts flow-matching decoder (SimpleMLPAdaLN)
|
||||
struct clip_flow_net {
|
||||
// AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual
|
||||
struct block {
|
||||
ggml_tensor * norm_w = nullptr;
|
||||
ggml_tensor * norm_b = nullptr;
|
||||
ggml_tensor * up_w = nullptr;
|
||||
ggml_tensor * up_b = nullptr;
|
||||
ggml_tensor * down_w = nullptr;
|
||||
ggml_tensor * down_b = nullptr;
|
||||
ggml_tensor * ada_w = nullptr; // -> shift, scale, gate
|
||||
ggml_tensor * ada_b = nullptr;
|
||||
};
|
||||
|
||||
// timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm
|
||||
struct time_embd {
|
||||
ggml_tensor * freqs = nullptr;
|
||||
ggml_tensor * up_w = nullptr;
|
||||
ggml_tensor * up_b = nullptr;
|
||||
ggml_tensor * down_w = nullptr;
|
||||
ggml_tensor * down_b = nullptr;
|
||||
ggml_tensor * norm = nullptr; // RMSNorm alpha
|
||||
};
|
||||
|
||||
ggml_tensor * input_proj_w = nullptr;
|
||||
ggml_tensor * input_proj_b = nullptr;
|
||||
ggml_tensor * cond_embd_w = nullptr;
|
||||
ggml_tensor * cond_embd_b = nullptr;
|
||||
ggml_tensor * final_ada_w = nullptr; // -> shift, scale
|
||||
ggml_tensor * final_ada_b = nullptr;
|
||||
ggml_tensor * final_proj_w = nullptr;
|
||||
ggml_tensor * final_proj_b = nullptr;
|
||||
std::vector<time_embd> time;
|
||||
std::vector<block> blocks;
|
||||
};
|
||||
|
||||
// qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
struct clip_code2wav {
|
||||
// "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it
|
||||
@@ -699,6 +770,24 @@ struct clip_model {
|
||||
// qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
clip_code2wav c2w;
|
||||
|
||||
// pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path)
|
||||
clip_seanet seanet;
|
||||
|
||||
// pocket-tts: voice latent -> backbone embd (speaker path)
|
||||
ggml_tensor * spk_proj_w = nullptr;
|
||||
ggml_tensor * downsample_w = nullptr;
|
||||
|
||||
// pocket-tts: flow-matching decoder, backbone hidden state -> next latent
|
||||
clip_flow_net flow;
|
||||
ggml_tensor * gen_out_eos_w = nullptr;
|
||||
ggml_tensor * gen_out_eos_b = nullptr;
|
||||
ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd
|
||||
ggml_tensor * gen_emb_mean = nullptr;
|
||||
ggml_tensor * gen_emb_std = nullptr;
|
||||
ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim
|
||||
ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate
|
||||
std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
|
||||
|
||||
// cogvlm
|
||||
ggml_tensor * mm_post_fc_norm_w = nullptr;
|
||||
ggml_tensor * mm_post_fc_norm_b = nullptr;
|
||||
|
||||
+276
-30
@@ -64,6 +64,9 @@
|
||||
#include "models/paddleocr.cpp"
|
||||
#include "models/parakeet.cpp"
|
||||
#include "models/pixtral.cpp"
|
||||
#include "models/pockettts-gen.cpp"
|
||||
#include "models/pockettts-seanet.cpp"
|
||||
#include "models/pockettts-spkenc.cpp"
|
||||
#include "models/qwen2vl.cpp"
|
||||
#include "models/qwen3vl.cpp"
|
||||
#include "models/qwen3a.cpp"
|
||||
@@ -228,6 +231,10 @@ struct clip_ctx {
|
||||
|
||||
bool support_batch = false;
|
||||
|
||||
// for audio gen, reseeded only when the caller asks for another seed
|
||||
std::mt19937 rng{std::random_device{}()};
|
||||
uint32_t rng_seed = UINT32_MAX;
|
||||
|
||||
clip_ctx(clip_context_params & ctx_params) {
|
||||
flash_attn_type = ctx_params.flash_attn_type;
|
||||
no_alloc = ctx_params.no_alloc;
|
||||
@@ -1113,6 +1120,25 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
{
|
||||
builder = std::make_unique<clip_graph_qwen3tts_spkenc>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
{
|
||||
builder = std::make_unique<clip_graph_pockettts_spkenc>(ctx, img);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
{
|
||||
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
|
||||
const int n_step = ctx->model.hparams.flow_n_step;
|
||||
const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0];
|
||||
GGML_ASSERT(n_step > 0);
|
||||
GGML_ASSERT(n_latent > 0);
|
||||
// "inp_feats" takes the caller's buffer as-is, the graph must consume all of it
|
||||
if (params && params->feats) {
|
||||
GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0);
|
||||
GGML_ASSERT(params->feats->size() >= (size_t) n_latent);
|
||||
}
|
||||
const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1;
|
||||
builder = std::make_unique<clip_graph_pockettts_gen>(ctx, img, gen_process, n_step, n_frames);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
|
||||
@@ -1354,6 +1380,7 @@ struct clip_model_loader {
|
||||
// these are unused, but still need to be set to avoid issues
|
||||
hparams.image_size = 0;
|
||||
hparams.patch_size = 1;
|
||||
get_string(KEY_GEN_AUDIO_VARIANT, hparams.gen_model_variant, false);
|
||||
|
||||
} else {
|
||||
GGML_ASSERT(false && "unknown modality");
|
||||
@@ -1498,7 +1525,7 @@ struct clip_model_loader {
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
{
|
||||
get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor);
|
||||
get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor);
|
||||
GGML_ASSERT(hparams.subsampling_factor == 8 &&
|
||||
"subsampling_factor must match the conv strides in clip_graph_parakeet::build()");
|
||||
get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size);
|
||||
@@ -1827,6 +1854,22 @@ struct clip_model_loader {
|
||||
// matches the reference decoder's sliding_window (speech_tokenizer/config.json)
|
||||
hparams.wav_tfm_swa = 72;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
{
|
||||
// mimi front-end takes the raw waveform, no mel
|
||||
hparams.audio_sample_rate = 24000;
|
||||
// seanet ratios are [6,5,4] in the config, the encoder reverses them
|
||||
hparams.seanet_ratios = { 4, 5, 6 };
|
||||
hparams.seanet_n_stage = (int32_t) hparams.seanet_ratios.size();
|
||||
hparams.mimi_downsample = 16;
|
||||
// matches the reference transformer's "context"
|
||||
hparams.mimi_tfm_context = 250;
|
||||
hparams.rope_theta = 10000.0f;
|
||||
// flow_lm defaults, see pocket_tts/default_parameters.py
|
||||
hparams.flow_n_step = 1;
|
||||
hparams.gen_eos_threshold = -4.0f;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PADDLEOCR:
|
||||
{
|
||||
hparams.n_merge = 2;
|
||||
@@ -2029,7 +2072,9 @@ struct clip_model_loader {
|
||||
|
||||
// GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT
|
||||
// checks below do not apply to it.
|
||||
const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA;
|
||||
// pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform
|
||||
const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA &&
|
||||
model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC;
|
||||
|
||||
// Validate audio hparams loaded from GGUF metadata
|
||||
if (hparams.n_mel_bins <= 0 || (fft_based && hparams.n_mel_bins > 256)) {
|
||||
@@ -2107,6 +2152,31 @@ struct clip_model_loader {
|
||||
return cur;
|
||||
};
|
||||
|
||||
// pocket-tts: the encoder and the decoder share the same layout, only the prefix differs
|
||||
auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) {
|
||||
const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN;
|
||||
const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT;
|
||||
const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1;
|
||||
const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2;
|
||||
const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV;
|
||||
|
||||
seanet.conv_in_w = get_tensor(string_format(conv_in, "weight"));
|
||||
seanet.conv_in_b = get_tensor(string_format(conv_in, "bias"));
|
||||
seanet.conv_out_w = get_tensor(string_format(conv_out, "weight"));
|
||||
seanet.conv_out_b = get_tensor(string_format(conv_out, "bias"));
|
||||
|
||||
seanet.stages.resize(hparams.seanet_n_stage);
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
auto & stage = seanet.stages[i];
|
||||
stage.res_conv1_w = get_tensor(string_format(res1, i, "weight"));
|
||||
stage.res_conv1_b = get_tensor(string_format(res1, i, "bias"));
|
||||
stage.res_conv2_w = get_tensor(string_format(res2, i, "weight"));
|
||||
stage.res_conv2_b = get_tensor(string_format(res2, i, "bias"));
|
||||
stage.scale_conv_w = get_tensor(string_format(scale, i, "weight"));
|
||||
stage.scale_conv_b = get_tensor(string_format(scale, i, "bias"));
|
||||
}
|
||||
};
|
||||
|
||||
auto get_vector = [&](const std::string & name) {
|
||||
std::vector<float> result;
|
||||
auto it = tensor_offset.find(name);
|
||||
@@ -2168,7 +2238,8 @@ struct clip_model_loader {
|
||||
|
||||
const bool has_standard_layers = (
|
||||
model.proj_type != PROJECTOR_TYPE_GEMMA3NV &&
|
||||
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC);
|
||||
model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC &&
|
||||
model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN);
|
||||
|
||||
// layers
|
||||
const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0;
|
||||
@@ -2842,6 +2913,81 @@ struct clip_model_loader {
|
||||
model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight"));
|
||||
model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias"));
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
{
|
||||
load_seanet(model.seanet, false);
|
||||
model.downsample_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight"));
|
||||
model.spk_proj_w = get_tensor(string_format(TN_A_SPEAKER_PROJ, "weight"));
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
{
|
||||
auto & flow = model.flow;
|
||||
flow.input_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "weight"));
|
||||
flow.input_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "bias"));
|
||||
flow.cond_embd_w = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "weight"));
|
||||
flow.cond_embd_b = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "bias"));
|
||||
flow.final_ada_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "weight"));
|
||||
flow.final_ada_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "bias"));
|
||||
flow.final_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "weight"));
|
||||
flow.final_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "bias"));
|
||||
|
||||
flow.time.resize(2);
|
||||
for (size_t i = 0; i < flow.time.size(); i++) {
|
||||
auto & t = flow.time[i];
|
||||
t.freqs = get_tensor(string_format(TN_A_GEN_FLOW_TIME_FREQS, (int) i));
|
||||
t.up_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "weight"));
|
||||
t.up_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "bias"));
|
||||
t.down_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "weight"));
|
||||
t.down_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "bias"));
|
||||
t.norm = get_tensor(string_format(TN_A_GEN_FLOW_TIME_NORM, (int) i));
|
||||
}
|
||||
|
||||
// one AdaLN block per flow depth, the count is only known from the tensors
|
||||
for (int il = 0; ; il++) {
|
||||
ggml_tensor * probe = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "weight"), false);
|
||||
if (probe == nullptr) {
|
||||
break;
|
||||
}
|
||||
clip_flow_net::block blk;
|
||||
blk.norm_w = probe;
|
||||
blk.norm_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "bias"));
|
||||
blk.up_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "weight"));
|
||||
blk.up_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "bias"));
|
||||
blk.down_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "weight"));
|
||||
blk.down_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "bias"));
|
||||
blk.ada_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "weight"));
|
||||
blk.ada_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "bias"));
|
||||
flow.blocks.push_back(blk);
|
||||
}
|
||||
|
||||
model.gen_out_eos_w = get_tensor(string_format(TN_A_GEN_OUT_EOS, "weight"));
|
||||
model.gen_out_eos_b = get_tensor(string_format(TN_A_GEN_OUT_EOS, "bias"));
|
||||
model.gen_input_lin_w = get_tensor(string_format(TN_A_GEN_INPUT_LINEAR, "weight"));
|
||||
model.gen_emb_mean = get_tensor(TN_A_GEN_EMB_MEAN);
|
||||
model.gen_emb_std = get_tensor(TN_A_GEN_EMB_STD);
|
||||
|
||||
// mimi decoder
|
||||
model.gen_quant_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_OUT, "weight"));
|
||||
model.gen_upsample_w = get_tensor(string_format(TN_A_GEN_WAV_UPSAMPLE, "weight"));
|
||||
load_seanet(model.seanet, true);
|
||||
model.gen_tfm_layers.resize(hparams.n_layer);
|
||||
for (int il = 0; il < hparams.n_layer; il++) {
|
||||
auto & layer = model.gen_tfm_layers[il];
|
||||
const char * p = "a.gen.wav.tfm";
|
||||
layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight"));
|
||||
layer.ln_1_b = get_tensor(string_format(TN_LN_1, p, il, "bias"));
|
||||
layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight"));
|
||||
layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight"));
|
||||
layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight"));
|
||||
layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight"));
|
||||
layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight"));
|
||||
layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight"));
|
||||
layer.ln_2_b = get_tensor(string_format(TN_LN_2, p, il, "bias"));
|
||||
layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight"));
|
||||
layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight"));
|
||||
layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight"));
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
// code_predictor
|
||||
@@ -4147,6 +4293,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
// one hidden-state vector fed back to the talker per call
|
||||
n_patches = 1;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
{
|
||||
// one conditioning row per 12.5Hz frame
|
||||
const int hop = ctx->model.hparams.mimi_downsample * 120;
|
||||
n_patches = img->nx() / hop;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
{
|
||||
// one latent per call for GEN_CODE, GEN_WAV sizes its input from the caller
|
||||
n_patches = 1;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GRANITE4_VISION:
|
||||
{
|
||||
// Per-tile output token count: each projector block outputs
|
||||
@@ -4188,6 +4345,15 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
return clip_encode(ctx, ¶ms);
|
||||
}
|
||||
|
||||
// persisted state slots of the gen-audio decoder, per pipeline
|
||||
static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
switch (model.proj_type) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model);
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model);
|
||||
default: return {};
|
||||
}
|
||||
}
|
||||
|
||||
bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
const clip_image_f32_batch & imgs = *params->imgs;
|
||||
int n_batch_cur = imgs.entries.size();
|
||||
@@ -4203,6 +4369,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
clip_model_loader::warmup(*ctx, *params->imgs);
|
||||
}
|
||||
|
||||
if (params->seed != ctx->rng_seed) {
|
||||
ctx->rng_seed = params->seed;
|
||||
ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed);
|
||||
}
|
||||
|
||||
// build the inference graph
|
||||
ggml_backend_sched_reset(ctx->sched.get());
|
||||
ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build();
|
||||
@@ -4247,6 +4418,50 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur));
|
||||
};
|
||||
|
||||
// upload the decoder state from the previous call, or zero-fill on a cold start
|
||||
auto set_gen_state_in = [&]() {
|
||||
size_t offset = 0;
|
||||
for (const auto & slot : list_gen_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
|
||||
const size_t nb = ggml_nbytes(t);
|
||||
if (params->state_in && params->state_in->size() >= offset + nb) {
|
||||
ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
|
||||
} else {
|
||||
std::vector<uint8_t> zeros(nb, 0);
|
||||
ggml_backend_tensor_set(t, zeros.data(), 0, nb);
|
||||
}
|
||||
offset += nb;
|
||||
}
|
||||
};
|
||||
|
||||
// rope positions and attention mask of the mimi transformers (pocket-tts).
|
||||
// the mask is causal with a sliding window, see _build_attention_mask() in the reference
|
||||
auto set_pockettts_tfm_inputs = [&]() {
|
||||
const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos"));
|
||||
GGML_ASSERT(n_pos > 0);
|
||||
std::vector<int32_t> positions((size_t) n_pos);
|
||||
for (int64_t i = 0; i < n_pos; i++) {
|
||||
positions[(size_t) i] = (int32_t) i;
|
||||
}
|
||||
set_input_i32("inp_pos", positions);
|
||||
|
||||
// the preprocessor truncates the waveform to keep this mask bounded
|
||||
const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120;
|
||||
GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask");
|
||||
|
||||
const int64_t context = hparams.mimi_tfm_context;
|
||||
std::vector<float> mask((size_t) n_pos * n_pos, -INFINITY);
|
||||
for (int64_t q = 0; q < n_pos; q++) {
|
||||
for (int64_t k = 0; k < n_pos; k++) {
|
||||
const int64_t delta = q - k;
|
||||
if (delta >= 0 && delta < context) {
|
||||
mask[(size_t) q * n_pos + k] = 0.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
set_input_f32("kq_mask", mask);
|
||||
};
|
||||
|
||||
// set input pixel values
|
||||
if (!imgs.is_audio) {
|
||||
size_t nelem = 0;
|
||||
@@ -4290,8 +4505,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
}
|
||||
set_input_f32("inp_raw", inp_raw);
|
||||
|
||||
} else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_WAV)) {
|
||||
// audio input, code2wav is not here: its only input is "inp_codes", set in the switch below
|
||||
} else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
// audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below
|
||||
GGML_ASSERT(imgs.entries.size() == 1);
|
||||
|
||||
const auto & mel_inp = imgs.entries[0];
|
||||
@@ -4824,6 +5039,30 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
}
|
||||
set_input_i32("patches", patches);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
{
|
||||
set_pockettts_tfm_inputs();
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
{
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
GGML_ASSERT(params->feats != nullptr);
|
||||
set_input_f32("inp_feats", *params->feats);
|
||||
// positions and mask are derived in-graph from the persisted counter
|
||||
set_gen_state_in();
|
||||
} else {
|
||||
// flow matching starts from gaussian noise, std = sqrt(temp)
|
||||
ggml_tensor * t = get_inp_tensor("inp_noise");
|
||||
// Config.default_temperature, for a caller that does not set one
|
||||
const float temp = params->temp > 0.0f ? params->temp : 0.7f;
|
||||
std::normal_distribution<float> dist(0.0f, std::sqrt(temp));
|
||||
std::vector<float> noise(ggml_nelements(t));
|
||||
for (auto & v : noise) {
|
||||
v = dist(ctx->rng);
|
||||
}
|
||||
set_input_f32("inp_noise", noise);
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_GEMMA4V:
|
||||
case PROJECTOR_TYPE_GEMMA4UV:
|
||||
{
|
||||
@@ -4948,20 +5187,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
}
|
||||
}
|
||||
set_input_i32("inp_codes", codes);
|
||||
|
||||
// upload the state from the previous call, or zero-fill on a cold start
|
||||
size_t offset = 0;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
|
||||
const size_t nb = ggml_nbytes(t);
|
||||
if (params->state_in && params->state_in->size() >= offset + nb) {
|
||||
ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
|
||||
} else {
|
||||
std::vector<uint8_t> zeros(nb, 0);
|
||||
ggml_backend_tensor_set(t, zeros.data(), 0, nb);
|
||||
}
|
||||
offset += nb;
|
||||
}
|
||||
set_gen_state_in();
|
||||
} else {
|
||||
// code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it
|
||||
const int64_t vocab0 = model.gen_code_out_embd_w->ne[1];
|
||||
@@ -4973,11 +5199,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
set_input_i32("inp_code0", code0);
|
||||
|
||||
// one uniform(0,1) draw per codebook, used by do_sampling()
|
||||
static std::mt19937 rng{ std::random_device{}() };
|
||||
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
|
||||
for (int64_t g = 0; g < n_acoustic; g++) {
|
||||
std::vector<float> r = { dist(rng) };
|
||||
std::vector<float> r = { dist(ctx->rng) };
|
||||
set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
|
||||
}
|
||||
}
|
||||
@@ -5430,14 +5655,31 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
// for audio gen models
|
||||
//
|
||||
|
||||
// optional outputs: a pipeline yields codes or feats, and not all have an eos head
|
||||
if (params->out_codes != nullptr) {
|
||||
ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes");
|
||||
if (codes == nullptr) {
|
||||
GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor");
|
||||
if (codes != nullptr) {
|
||||
auto & out_codes = *params->out_codes;
|
||||
out_codes.resize(ggml_nelements(codes));
|
||||
ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes));
|
||||
}
|
||||
}
|
||||
if (params->out_feats != nullptr) {
|
||||
ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats");
|
||||
if (feats != nullptr) {
|
||||
auto & out_feats = *params->out_feats;
|
||||
out_feats.resize(ggml_nelements(feats));
|
||||
ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats));
|
||||
}
|
||||
}
|
||||
if (params->out_is_eos != nullptr) {
|
||||
ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score");
|
||||
if (eos != nullptr) {
|
||||
GGML_ASSERT(ggml_nelements(eos) == 1);
|
||||
float score = 0.0f;
|
||||
ggml_backend_tensor_get(eos, &score, 0, sizeof(float));
|
||||
*params->out_is_eos = score > hparams.gen_eos_threshold;
|
||||
}
|
||||
auto & out_codes = *params->out_codes;
|
||||
out_codes.resize(ggml_nelements(codes));
|
||||
ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes));
|
||||
}
|
||||
if (params->out_audio != nullptr) {
|
||||
ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio");
|
||||
@@ -5449,9 +5691,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio));
|
||||
|
||||
// drop the tail audio that comes from the code-0 rear padding
|
||||
const int64_t n_codes = model.gen_code_head_w->ne[2] + 1;
|
||||
const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0;
|
||||
const int64_t n_frames_w = hparams.wav_tfm_swa;
|
||||
const int64_t n_frames = (int64_t) params->codes->size() / n_codes;
|
||||
const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w;
|
||||
if (n_frames < n_frames_w) {
|
||||
const size_t hop = out_audio.size() / n_frames_w;
|
||||
out_audio.resize((size_t) n_frames * hop);
|
||||
@@ -5460,12 +5702,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
if (params->state_out != nullptr) {
|
||||
auto & state_out = *params->state_out;
|
||||
size_t total = 0;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
for (const auto & slot : list_gen_state_slots(hparams, model)) {
|
||||
total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float);
|
||||
}
|
||||
state_out.resize(total);
|
||||
size_t offset = 0;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
for (const auto & slot : list_gen_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str());
|
||||
if (t == nullptr) {
|
||||
GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str());
|
||||
@@ -5613,6 +5855,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
return ctx->model.mm_fc_w->ne[2];
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
return ctx->model.gen_code_out_embd_w->ne[0];
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
return ctx->model.spk_proj_w->ne[1];
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
return ctx->model.gen_input_lin_w->ne[1];
|
||||
case PROJECTOR_TYPE_PARAKEET:
|
||||
return ctx->model.mm_1_w->ne[1];
|
||||
default:
|
||||
|
||||
@@ -104,9 +104,14 @@ struct clip_encode_params {
|
||||
int32_t top_k = 50;
|
||||
float top_p = 1.0f;
|
||||
std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
|
||||
std::vector<float> * out_feats = nullptr; // continuous counterpart of out_codes
|
||||
uint32_t seed = UINT32_MAX; // UINT32_MAX for random
|
||||
float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders
|
||||
bool * out_is_eos = nullptr;
|
||||
|
||||
// GEN_WAV
|
||||
const std::vector<int32_t> * codes = nullptr; // this frame's 16 RVQ codes
|
||||
const std::vector<float> * feats = nullptr; // continuous counterpart of codes
|
||||
std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32
|
||||
const std::vector<uint8_t> * state_in = nullptr; // state from previous call, null or wrong size means cold start
|
||||
std::vector<uint8_t> * state_out = nullptr; // state for the next call
|
||||
|
||||
@@ -318,6 +318,59 @@ struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
};
|
||||
};
|
||||
|
||||
//
|
||||
// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder.
|
||||
// stateless unless state_in is populated: convs then pad instead of carrying left-context.
|
||||
//
|
||||
struct clip_graph_pockettts_seanet : clip_graph {
|
||||
clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {}
|
||||
ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); }
|
||||
|
||||
// per-call streaming state, keyed by slot name (see list_pockettts_state_slots)
|
||||
std::map<std::string, ggml_tensor *> state_in;
|
||||
mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
|
||||
|
||||
ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
|
||||
bool pad_replicate = false, const std::string & state_name = "") const;
|
||||
ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
|
||||
const std::string & state_name = "") const;
|
||||
ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
|
||||
const std::string & state_prefix = "") const;
|
||||
|
||||
// x: [T, C] -> [T / hop, dim]
|
||||
ggml_tensor * encode(ggml_tensor * x) const;
|
||||
// x: [T, dim] -> [T * hop, 1], streams when state_in is populated
|
||||
ggml_tensor * decode(ggml_tensor * x) const;
|
||||
};
|
||||
|
||||
// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows
|
||||
struct clip_graph_pockettts_spkenc : clip_graph {
|
||||
clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const;
|
||||
};
|
||||
|
||||
//
|
||||
// pocket-tts generation:
|
||||
// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call
|
||||
// GEN_WAV = mimi decoder, a window of latents -> PCM
|
||||
//
|
||||
struct clip_graph_pockettts_gen : clip_graph {
|
||||
clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames)
|
||||
: clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
clip_gen_process_type gen_process;
|
||||
int n_step; // lsd_decode steps, fixed at graph-build time
|
||||
int n_frames; // GEN_WAV only: number of latents to decode
|
||||
|
||||
// AdaLN modulation: x * (1 + scale) + shift
|
||||
ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const;
|
||||
ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const;
|
||||
ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const;
|
||||
};
|
||||
|
||||
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
|
||||
struct c2w_state_slot {
|
||||
std::string name;
|
||||
@@ -326,6 +379,9 @@ struct c2w_state_slot {
|
||||
};
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
// same, for the streaming mimi decoder (pocket-tts GEN_WAV)
|
||||
std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
struct clip_graph_kimik25 : clip_graph {
|
||||
clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -0,0 +1,291 @@
|
||||
#include "models.h"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
// pocket-tts generation stages
|
||||
//
|
||||
// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score
|
||||
// GEN_WAV : a window of latents -> PCM, through the mimi decoder
|
||||
//
|
||||
// there is no codebook anywhere, "codes" in the mtmd API are continuous features here
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const {
|
||||
ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f));
|
||||
return ggml_add(ctx0, cur, shift);
|
||||
}
|
||||
|
||||
// see TimestepEmbedder in the reference
|
||||
ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const {
|
||||
// t is a graph-build constant, so the cos/sin table can be folded into a scaled copy
|
||||
ggml_tensor * args = ggml_scale(ctx0, te.freqs, t);
|
||||
ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0);
|
||||
|
||||
ggml_tensor * cur = build_mm(te.up_w, emb);
|
||||
cur = ggml_add(ctx0, cur, te.up_b);
|
||||
cur = ggml_silu(ctx0, cur);
|
||||
cur = build_mm(te.down_w, cur);
|
||||
cur = ggml_add(ctx0, cur, te.down_b);
|
||||
|
||||
// this "RMSNorm" divides by the unbiased variance, not the mean square
|
||||
// it also rescales the input, not the centered value, see _rms_norm() in mlp.py
|
||||
{
|
||||
const int64_t n = cur->ne[0];
|
||||
ggml_tensor * mean = ggml_mean(ctx0, cur);
|
||||
ggml_tensor * dev = ggml_sub(ctx0, cur, mean);
|
||||
ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev));
|
||||
var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f);
|
||||
cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var));
|
||||
cur = ggml_mul(ctx0, cur, te.norm);
|
||||
}
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
// one velocity evaluation: v(cond, s, t, x)
|
||||
ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const {
|
||||
const auto & flow = model.flow;
|
||||
|
||||
ggml_tensor * cur = build_mm(flow.input_proj_w, x);
|
||||
cur = ggml_add(ctx0, cur, flow.input_proj_b);
|
||||
|
||||
// the two time conditions are averaged, then added to the projected backbone state
|
||||
ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t));
|
||||
ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size());
|
||||
|
||||
ggml_tensor * c = build_mm(flow.cond_embd_w, cond);
|
||||
c = ggml_add(ctx0, c, flow.cond_embd_b);
|
||||
|
||||
ggml_tensor * y = ggml_add(ctx0, ts, c);
|
||||
cb(y, "flow_cond", -1);
|
||||
|
||||
const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0];
|
||||
|
||||
for (size_t il = 0; il < flow.blocks.size(); il++) {
|
||||
const auto & blk = flow.blocks[il];
|
||||
|
||||
ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y));
|
||||
mod = ggml_add(ctx0, mod, blk.ada_b);
|
||||
|
||||
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
|
||||
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
|
||||
ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]);
|
||||
|
||||
ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il);
|
||||
h = modulate(h, shift, scale);
|
||||
h = build_mm(blk.up_w, h);
|
||||
h = ggml_add(ctx0, h, blk.up_b);
|
||||
h = ggml_silu(ctx0, h);
|
||||
h = build_mm(blk.down_w, h);
|
||||
h = ggml_add(ctx0, h, blk.down_b);
|
||||
|
||||
cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h));
|
||||
cb(cur, "flow_blk", (int) il);
|
||||
}
|
||||
|
||||
// final layer: the norm has no weights, only the AdaLN modulation
|
||||
ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y));
|
||||
mod = ggml_add(ctx0, mod, flow.final_ada_b);
|
||||
|
||||
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
|
||||
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
|
||||
|
||||
cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1);
|
||||
cur = modulate(cur, shift, scale);
|
||||
cur = build_mm(flow.final_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, flow.final_proj_b);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context
|
||||
// and the transposed-conv overlap tails
|
||||
std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
std::vector<c2w_state_slot> slots;
|
||||
if (model.gen_upsample_w == nullptr) {
|
||||
return slots; // not a pocket-tts decoder
|
||||
}
|
||||
const auto & seanet = model.seanet;
|
||||
|
||||
// the slots below are sized from these
|
||||
GGML_ASSERT(!model.gen_tfm_layers.empty());
|
||||
GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage);
|
||||
GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage);
|
||||
GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0);
|
||||
|
||||
slots.push_back({"tfm_pos", 1, 1});
|
||||
|
||||
const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1];
|
||||
const int64_t prefix = hparams.mimi_tfm_context - 1;
|
||||
for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) {
|
||||
slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix});
|
||||
slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix});
|
||||
}
|
||||
|
||||
// upsample is depthwise, its output channel count is the input one
|
||||
slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]});
|
||||
|
||||
slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]});
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
const auto & stage = seanet.stages[i];
|
||||
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
|
||||
slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]});
|
||||
slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]});
|
||||
}
|
||||
slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]});
|
||||
|
||||
return slots;
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_pockettts_gen::build() {
|
||||
if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) {
|
||||
// the backbone hidden state arrives as the single batch entry
|
||||
ggml_tensor * h_state = build_inp_raw(1);
|
||||
h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1);
|
||||
|
||||
// end-of-speech probe, thresholded on the host side
|
||||
ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state);
|
||||
eos = ggml_add(ctx0, eos, model.gen_out_eos_b);
|
||||
ggml_set_name(eos, "out_eos_score");
|
||||
ggml_set_output(eos);
|
||||
ggml_build_forward_expand(gf, eos);
|
||||
|
||||
const int64_t n_latent = model.gen_input_lin_w->ne[0];
|
||||
|
||||
ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1);
|
||||
ggml_set_name(noise, "inp_noise");
|
||||
ggml_set_input(noise);
|
||||
|
||||
// lsd_decode: integrate the velocity field from the noise sample
|
||||
ggml_tensor * cur = noise;
|
||||
for (int i = 0; i < n_step; i++) {
|
||||
const float s = (float) i / (float) n_step;
|
||||
const float t = (float) (i + 1) / (float) n_step;
|
||||
ggml_tensor * v = flow_forward(h_state, cur, s, t);
|
||||
cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step));
|
||||
}
|
||||
cb(cur, "flow_latent", -1);
|
||||
|
||||
ggml_set_name(cur, "out_feats");
|
||||
ggml_set_output(cur);
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
// the same latent, projected into the backbone's input space for the next step
|
||||
ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur);
|
||||
cb(embd, "gen_embd", -1);
|
||||
ggml_build_forward_expand(gf, embd);
|
||||
|
||||
return gf;
|
||||
}
|
||||
|
||||
// GEN_WAV: [32, n_frames] latents -> PCM
|
||||
ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32,
|
||||
model.gen_input_lin_w->ne[0], n_frames);
|
||||
ggml_set_name(feats, "inp_feats");
|
||||
ggml_set_input(feats);
|
||||
|
||||
// denormalize, then the DummyQuantizer up-projection
|
||||
ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean);
|
||||
cur = build_mm(model.gen_quant_out_w, cur);
|
||||
cb(cur, "quant_out", -1);
|
||||
|
||||
clip_graph_pockettts_seanet seanet(*this);
|
||||
for (const auto & slot : list_pockettts_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
|
||||
ggml_set_name(t, ("state_in_" + slot.name).c_str());
|
||||
ggml_set_input(t);
|
||||
seanet.state_in[slot.name] = t;
|
||||
}
|
||||
|
||||
// model frame rate -> encoder frame rate, depthwise transposed conv
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up");
|
||||
cb(cur, "mimi_upsample", -1);
|
||||
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
|
||||
// positions continue across calls, the counter lives in the state
|
||||
const int64_t n_pos = cur->ne[1];
|
||||
const int64_t prefix = hparams.mimi_tfm_context - 1;
|
||||
const int64_t n_kv = prefix + n_pos;
|
||||
|
||||
ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1);
|
||||
ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base),
|
||||
GGML_TYPE_I32);
|
||||
seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)});
|
||||
|
||||
// banded causal mask over [cached prefix | this chunk]
|
||||
// the last factor masks out cache rows that hold no real frame yet
|
||||
ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1);
|
||||
ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos);
|
||||
ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k);
|
||||
|
||||
ggml_tensor * keep = ggml_mul(ctx0,
|
||||
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0
|
||||
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context
|
||||
keep = ggml_mul(ctx0, keep,
|
||||
ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix)));
|
||||
ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1);
|
||||
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
const auto & layer = model.gen_tfm_layers[il];
|
||||
ggml_tensor * inp = cur;
|
||||
|
||||
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
|
||||
|
||||
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
|
||||
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
|
||||
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
|
||||
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
|
||||
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// prepend the cached window, then keep this chunk's tail for the next call
|
||||
const std::string k_name = "tfm_k_" + std::to_string(il);
|
||||
const std::string v_name = "tfm_v_" + std::to_string(il);
|
||||
ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name),
|
||||
ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1);
|
||||
ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1);
|
||||
seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix,
|
||||
k_full->nb[1], (size_t) n_pos * k_full->nb[1]))});
|
||||
seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix,
|
||||
v_full->nb[1], (size_t) n_pos * v_full->nb[1]))});
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1);
|
||||
|
||||
cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
|
||||
inp = cur;
|
||||
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
|
||||
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
}
|
||||
cb(cur, "mimi_dec_tfm", -1);
|
||||
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = seanet.decode(cur);
|
||||
|
||||
for (const auto & s : seanet.state_out) {
|
||||
ggml_set_name(s.second, ("state_out_" + s.first).c_str());
|
||||
ggml_set_output(s.second);
|
||||
ggml_build_forward_expand(gf, s.second);
|
||||
}
|
||||
|
||||
// [n_samples, 1] -> [n_samples], clamped like the reference output
|
||||
cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
|
||||
cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f);
|
||||
ggml_set_name(cur, "out_audio");
|
||||
ggml_set_output(cur);
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
return gf;
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
#include "models.h"
|
||||
|
||||
// SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py
|
||||
//
|
||||
// tensors are T-first here: [T, C]
|
||||
// the convs are causal: left context comes from a state slot, or from padding on a cold start
|
||||
|
||||
static int64_t div_ceil(int64_t a, int64_t b) {
|
||||
return a / b + (a % b ? 1 : 0);
|
||||
}
|
||||
|
||||
// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC]
|
||||
// the convs are causal, so the whole K - stride padding goes on the left
|
||||
ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
|
||||
bool pad_replicate, const std::string & state_name) const {
|
||||
const int64_t k_size = (w->ne[0] - 1) * dilation + 1;
|
||||
const int64_t p_total = k_size - stride;
|
||||
|
||||
// trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py
|
||||
const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride);
|
||||
const int64_t ideal_len = n_frames * stride + k_size - p_total;
|
||||
const int64_t p_extra = ideal_len - x->ne[0];
|
||||
|
||||
if (!state_name.empty() && p_total > 0) {
|
||||
// streaming: the left context is the tail of the previous call
|
||||
ggml_tensor * left = state_in.at(state_name); // [p_total, IC]
|
||||
x = ggml_concat(ctx0, left, x, 0);
|
||||
state_out.push_back({state_name,
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1],
|
||||
(size_t) (x->ne[0] - p_total) * x->nb[0]))});
|
||||
} else if (pad_replicate && p_total > 0) {
|
||||
// the resamplers repeat the first frame instead of zero-padding
|
||||
ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0);
|
||||
ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1);
|
||||
x = ggml_concat(ctx0, left, x, 0);
|
||||
x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0);
|
||||
} else {
|
||||
x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0);
|
||||
}
|
||||
|
||||
ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation);
|
||||
y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]);
|
||||
if (b) {
|
||||
y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC]
|
||||
// the K - stride overlap tail belongs to the next call: added to its head when streaming, else dropped
|
||||
ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
|
||||
const std::string & state_name) const {
|
||||
const int64_t K = w->ne[0];
|
||||
const int64_t T = x->ne[0];
|
||||
const int64_t p_total = K - stride;
|
||||
const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1;
|
||||
const int64_t OC = depthwise ? w->ne[2] : w->ne[1];
|
||||
const int64_t emit_len = T * stride;
|
||||
|
||||
// one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride
|
||||
ggml_tensor * col;
|
||||
if (depthwise) {
|
||||
// one group per channel: a batched matmul over the channels scales the kernel by each step
|
||||
ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC]
|
||||
ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC]
|
||||
col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC]
|
||||
col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T]
|
||||
col = ggml_reshape_2d(ctx0, col, K * OC, T);
|
||||
} else {
|
||||
ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]);
|
||||
w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC]
|
||||
ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T]
|
||||
col = ggml_mul_mat(ctx0, w2, xt);
|
||||
}
|
||||
ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [emit_len + p_total, OC]
|
||||
|
||||
ggml_tensor * out;
|
||||
if (state_name.empty() || p_total == 0) {
|
||||
out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0));
|
||||
} else {
|
||||
// overlap-add the tail the previous call held back
|
||||
ggml_tensor * prev = state_in.at(state_name); // [p_total, OC]
|
||||
ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev);
|
||||
if (emit_len > p_total) {
|
||||
ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1],
|
||||
(size_t) p_total * full->nb[0]);
|
||||
out = ggml_concat(ctx0, head, rest, 0);
|
||||
} else {
|
||||
out = head;
|
||||
}
|
||||
state_out.push_back({state_name,
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1],
|
||||
(size_t) emit_len * full->nb[0]))});
|
||||
}
|
||||
|
||||
if (b) {
|
||||
out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
|
||||
const std::string & state_prefix) const {
|
||||
ggml_tensor * h = ggml_elu(ctx0, x);
|
||||
h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix);
|
||||
h = ggml_elu(ctx0, h);
|
||||
// the second conv is pointwise, it needs no left context
|
||||
h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1);
|
||||
return ggml_add(ctx0, x, h);
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const {
|
||||
const auto & seanet = model.seanet;
|
||||
|
||||
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1);
|
||||
cb(cur, "seanet_enc_in", -1);
|
||||
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
const auto & stage = seanet.stages[i];
|
||||
const int stride = hparams.seanet_ratios[i];
|
||||
|
||||
cur = res_unit(cur, stage, 1);
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1);
|
||||
cb(cur, "seanet_enc_stage", i);
|
||||
}
|
||||
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1);
|
||||
cb(cur, "seanet_enc_out", -1);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const {
|
||||
const auto & seanet = model.seanet;
|
||||
const bool stream = !state_in.empty();
|
||||
|
||||
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false,
|
||||
stream ? "dec_in" : "");
|
||||
cb(cur, "seanet_dec_in", -1);
|
||||
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
const auto & stage = seanet.stages[i];
|
||||
// the decoder mirrors the encoder, so the ratios are walked backwards
|
||||
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
|
||||
const std::string id = std::to_string(i);
|
||||
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride,
|
||||
stream ? "dec_up_" + id : "");
|
||||
cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : "");
|
||||
cb(cur, "seanet_dec_stage", i);
|
||||
}
|
||||
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false,
|
||||
stream ? "dec_out" : "");
|
||||
cb(cur, "seanet_dec_out", -1);
|
||||
|
||||
return cur;
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
#include "models.h"
|
||||
|
||||
// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame
|
||||
// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight
|
||||
|
||||
// pre-norm block with layer scale on both residual paths, see mimi_transformer.py
|
||||
ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const {
|
||||
ggml_tensor * inp = cur;
|
||||
|
||||
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
|
||||
|
||||
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
|
||||
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
|
||||
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
|
||||
|
||||
const int64_t n_pos = cur->ne[1];
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
|
||||
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
|
||||
inp = cur;
|
||||
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
|
||||
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_pockettts_spkenc::build() {
|
||||
// the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1]
|
||||
ggml_tensor * inp_raw = build_inp_raw(1);
|
||||
ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]);
|
||||
|
||||
clip_graph_pockettts_seanet seanet(*this);
|
||||
cur = seanet.encode(cur);
|
||||
cb(cur, "mimi_enc", -1);
|
||||
|
||||
// [T, 512] -> transformer works on [512, T]
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
|
||||
ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]);
|
||||
ggml_set_name(inp_pos, "inp_pos");
|
||||
ggml_set_input(inp_pos);
|
||||
|
||||
// the mimi transformer is causal with a sliding window, see _build_attention_mask()
|
||||
ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]);
|
||||
ggml_set_name(kq_mask, "kq_mask");
|
||||
ggml_set_input(kq_mask);
|
||||
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il);
|
||||
}
|
||||
cb(cur, "mimi_enc_tfm", -1);
|
||||
|
||||
// downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32]
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true);
|
||||
cb(cur, "mimi_downsample", -1);
|
||||
|
||||
// voice latent -> backbone embd
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = build_mm(model.spk_proj_w, cur);
|
||||
cb(cur, "spk_proj", -1);
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
return gf;
|
||||
}
|
||||
@@ -610,6 +610,10 @@ std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, c
|
||||
const auto & c2w = model.c2w;
|
||||
std::vector<c2w_state_slot> slots;
|
||||
|
||||
if (c2w.pre_conv_w == nullptr) {
|
||||
return slots; // not a code2wav model, it keeps no state between calls
|
||||
}
|
||||
|
||||
slots.push_back({"tfm_pos", 1, 1});
|
||||
|
||||
// prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward)
|
||||
|
||||
@@ -1423,3 +1423,41 @@ std::vector<float> mtmd_audio_streaming_istft::flush() {
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
//
|
||||
// mtmd_audio_preprocessor_pockettts
|
||||
//
|
||||
// mimi takes the raw 24kHz waveform, there is no mel front-end
|
||||
// the samples are handed over as a single-row "mel", to reuse the normal chunk path
|
||||
//
|
||||
|
||||
bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples,
|
||||
size_t n_samples,
|
||||
std::vector<mtmd_audio_mel> & output) {
|
||||
// the encoder needs whole frames, see pad_for_conv1d() in the reference
|
||||
const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120;
|
||||
if (n_samples == 0 || frame_size <= 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// the mimi transformer mask is dense, so cost is quadratic in the reference length
|
||||
const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate;
|
||||
if ((int64_t) n_samples > max_samples) {
|
||||
LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__,
|
||||
(double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds);
|
||||
n_samples = (size_t) max_samples;
|
||||
}
|
||||
|
||||
const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size;
|
||||
const int64_t n_padded = n_frames * frame_size;
|
||||
|
||||
mtmd_audio_mel out;
|
||||
out.n_mel = 1;
|
||||
out.n_len = n_padded;
|
||||
out.n_len_org = (int64_t) n_samples;
|
||||
out.data.assign((size_t) n_padded, 0.0f);
|
||||
std::copy(samples, samples + n_samples, out.data.begin());
|
||||
|
||||
output.push_back(std::move(out));
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -129,6 +129,13 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor {
|
||||
mtmd_audio_cache cache;
|
||||
};
|
||||
|
||||
// mimi convolves the waveform directly, so this only pads it to a whole number of frames
|
||||
struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
|
||||
mtmd_audio_preprocessor_pockettts(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
|
||||
void initialize() override {}
|
||||
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
|
||||
};
|
||||
|
||||
struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
|
||||
mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { }
|
||||
void initialize() override;
|
||||
|
||||
+573
-10
@@ -5,6 +5,8 @@
|
||||
#include "../src/llama-ext.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
@@ -87,7 +89,8 @@ public:
|
||||
virtual int32_t step_prompt(int32_t n_batch) = 0;
|
||||
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token,
|
||||
// those read what they need from h_state_in instead
|
||||
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0;
|
||||
// set out_stop on end-of-speech, h_state_out must be null if no frame is generated
|
||||
virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 0;
|
||||
virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0;
|
||||
|
||||
protected:
|
||||
@@ -200,8 +203,10 @@ public:
|
||||
prompt_pos = 0;
|
||||
|
||||
pos = 0;
|
||||
top_k = inp->top_k > 0 ? inp->top_k : 50;
|
||||
top_p = inp->top_p > 0 ? inp->top_p : 1.0f;
|
||||
const mtmd_gen_inp def = mtmd_gen_inp_default(mctx);
|
||||
top_k = inp->top_k > 0 ? inp->top_k : def.top_k;
|
||||
top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
|
||||
seed = inp->seed;
|
||||
out_type = inp->out_type;
|
||||
|
||||
// the prompt above holds the whole text stream up to tts_eos, so every generated
|
||||
@@ -241,13 +246,26 @@ public:
|
||||
return n_prompt - prompt_pos;
|
||||
}
|
||||
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override {
|
||||
mtmd_gen_inp inp{};
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
|
||||
if (sampled == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
// backbone signals end-of-speech with a token, no frame for this step
|
||||
if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) {
|
||||
*out_stop = true;
|
||||
*h_state_out = nullptr;
|
||||
return 0;
|
||||
}
|
||||
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
inp.code0 = sampled - codec_0;
|
||||
inp.embd = const_cast<float *>(h_state_in);
|
||||
inp.top_k = top_k;
|
||||
inp.top_p = top_p;
|
||||
inp.seed = seed;
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
|
||||
@@ -384,10 +402,11 @@ private:
|
||||
if (codes_buf.empty()) {
|
||||
return true;
|
||||
}
|
||||
mtmd_gen_inp inp{};
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
|
||||
inp.codes = codes_buf.data();
|
||||
inp.n_codes = codes_buf.size();
|
||||
inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation
|
||||
inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
|
||||
inp.state_size = c2w_state.size();
|
||||
mtmd_gen_out out{};
|
||||
@@ -427,8 +446,9 @@ private:
|
||||
std::unique_ptr<decode_embd_batch> prompt_batch;
|
||||
int n_prompt = 0;
|
||||
int prompt_pos = 0;
|
||||
int32_t top_k = 50;
|
||||
float top_p = 1.0f;
|
||||
int32_t top_k = 50;
|
||||
float top_p = 1.0f;
|
||||
uint32_t seed = UINT32_MAX;
|
||||
std::vector<int32_t> codes_buf;
|
||||
std::vector<uint8_t> c2w_state;
|
||||
std::vector<float> audio_pcm;
|
||||
@@ -438,10 +458,547 @@ private:
|
||||
std::vector<char> out_buf;
|
||||
};
|
||||
|
||||
// settings that only live in the reference's per-pack yaml, not in the checkpoint
|
||||
// the english packs share the same shapes and tokenizer, but disagree on these
|
||||
// all three are 0 / false when the pack does not tune them, the model default is then used
|
||||
struct pockettts_pack_settings {
|
||||
float temp = 0.0f;
|
||||
int frames_after_eos = 0;
|
||||
bool pad_short_text = false;
|
||||
};
|
||||
|
||||
static pockettts_pack_settings pockettts_pack(const char * variant) {
|
||||
static const std::unordered_map<std::string, pockettts_pack_settings> packs = {
|
||||
{ "english", { 0.3f, 0, false } },
|
||||
{ "english_2026-01", { 0.7f, 0, true } },
|
||||
{ "english_2026-04", { 0.3f, 0, false } },
|
||||
{ "french_24l", { 0.7f, 8, false } },
|
||||
};
|
||||
auto it = packs.find(variant ? variant : "");
|
||||
if (it == packs.end()) {
|
||||
LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n",
|
||||
variant ? variant : "");
|
||||
return {};
|
||||
}
|
||||
return it->second;
|
||||
}
|
||||
|
||||
// pocket-tts: the backbone emits no token, the flow net turns each hidden state into a latent
|
||||
// the end-of-speech head also lives in the mmproj
|
||||
class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
|
||||
public:
|
||||
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
|
||||
|
||||
void reset() override {
|
||||
seq_id = 0;
|
||||
pos = 0;
|
||||
feats_buf.clear();
|
||||
dec_state.clear();
|
||||
audio_pcm.clear();
|
||||
h_state_buf.clear();
|
||||
out_buf.clear();
|
||||
prompt_embd_buf.clear();
|
||||
prompt_batch.reset();
|
||||
n_prompt = 0;
|
||||
prompt_pos = 0;
|
||||
step_idx = 0;
|
||||
eos_step = -1;
|
||||
chunks.clear();
|
||||
chunk_idx = 0;
|
||||
n_voice_pos = 0;
|
||||
chunk_budget = 0;
|
||||
}
|
||||
|
||||
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
|
||||
reset();
|
||||
seq_id = inp->seq_id;
|
||||
|
||||
if (!ensure_cache()) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::vector<float> voice;
|
||||
if (inp->speaker_ref) {
|
||||
if (!encode_speaker(inp->speaker_ref, voice)) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
pack = pockettts_pack(info.model_variant);
|
||||
|
||||
const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len),
|
||||
pack.pad_short_text);
|
||||
if (text.empty()) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: empty prompt\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::vector<llama_token> ids(text.size() + 16);
|
||||
int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(),
|
||||
(int32_t) ids.size(), false, false);
|
||||
if (n_ids <= 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
|
||||
return 1;
|
||||
}
|
||||
ids.resize((size_t) n_ids);
|
||||
|
||||
// long inputs degrade badly, so each chunk restarts from the voice conditioning
|
||||
// see split_into_best_sentences() in the reference
|
||||
chunks = split_chunks(ids);
|
||||
chunk_idx = 0;
|
||||
if (chunks.size() > 1) {
|
||||
LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size());
|
||||
}
|
||||
|
||||
const int n_e = n_embd;
|
||||
|
||||
// sequence order is voice, then text, then the audio BOS that starts generation
|
||||
if (!voice.empty()) {
|
||||
GGML_ASSERT(voice.size() % (size_t) n_e == 0);
|
||||
if (bos_before_voice != LLAMA_TOKEN_NULL) {
|
||||
push_embd_row(prompt_embd_buf, bos_before_voice);
|
||||
}
|
||||
prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end());
|
||||
}
|
||||
// every later chunk rewinds to here and re-prompts, so the voice stays primed
|
||||
n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e);
|
||||
|
||||
for (llama_token t : chunks[0]) {
|
||||
push_embd_row(prompt_embd_buf, t);
|
||||
}
|
||||
push_embd_row(prompt_embd_buf, audio_bos);
|
||||
arm_chunk_budget(0);
|
||||
|
||||
n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
|
||||
prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e));
|
||||
prompt_batch->set_position_normal(0, seq_id);
|
||||
prompt_pos = 0;
|
||||
|
||||
seed = inp->seed;
|
||||
out_type = inp->out_type;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t step_prompt(int32_t n_batch) override {
|
||||
GGML_ASSERT(n_batch > 0);
|
||||
if (prompt_pos >= n_prompt) {
|
||||
return 0;
|
||||
}
|
||||
const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos);
|
||||
llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch);
|
||||
|
||||
if ((prompt_pos + n_tokens_batch) == n_prompt) {
|
||||
batch_view.logits[n_tokens_batch - 1] = 1;
|
||||
}
|
||||
|
||||
if (llama_decode(lctx, batch_view) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
pos += n_tokens_batch;
|
||||
prompt_pos += n_tokens_batch;
|
||||
|
||||
if (prompt_pos >= n_prompt) {
|
||||
prompt_batch.reset();
|
||||
prompt_embd_buf.clear();
|
||||
return 0;
|
||||
}
|
||||
return n_prompt - prompt_pos;
|
||||
}
|
||||
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
|
||||
(void) sampled; // the backbone output is continuous, there is no token to consume
|
||||
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
inp.embd = const_cast<float *>(h_state_in);
|
||||
// clip only reseeds when the seed changes, so pass the same one on every step
|
||||
inp.seed = seed;
|
||||
if (pack.temp > 0.0f) {
|
||||
inp.temp = pack.temp;
|
||||
}
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
|
||||
return 1;
|
||||
}
|
||||
if (out.is_eos && eos_step < 0) {
|
||||
eos_step = step_idx;
|
||||
}
|
||||
// the frame of the stopping step is discarded, matching _autoregressive_generation().
|
||||
// the budget is the reference's fallback for a chunk whose eos head never fires
|
||||
const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) ||
|
||||
step_idx >= chunk_budget;
|
||||
if (chunk_done) {
|
||||
if (eos_step < 0) {
|
||||
LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx);
|
||||
}
|
||||
return finish_chunk(h_state_out, out_stop);
|
||||
}
|
||||
|
||||
feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats);
|
||||
step_idx++;
|
||||
if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) {
|
||||
if (!flush_gen_wav()) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
decode_embd_batch batch_embd(const_cast<float *>(out.embd), 1, 1, n_embd);
|
||||
batch_embd.set_position_normal(pos, seq_id);
|
||||
batch_embd.batch.logits[0] = 1;
|
||||
pos++;
|
||||
|
||||
if (llama_decode(lctx, batch_embd.batch) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
const float * he = llama_get_embeddings_ith(lctx, -1);
|
||||
h_state_buf.assign(he, he + n_embd);
|
||||
*h_state_out = h_state_buf.data();
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
|
||||
if (!flush_gen_wav()) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
*out_sample_rate = info.sample_rate;
|
||||
if (out_n_samples) {
|
||||
*out_n_samples = (int64_t) audio_pcm.size();
|
||||
}
|
||||
|
||||
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
|
||||
*out_data = (const char *) audio_pcm.data();
|
||||
*out_data_len = audio_pcm.size() * sizeof(float);
|
||||
return 0;
|
||||
}
|
||||
|
||||
out_buf.clear();
|
||||
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
|
||||
return 1;
|
||||
}
|
||||
*out_data = out_buf.data();
|
||||
*out_data_len = out_buf.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
private:
|
||||
bool ensure_cache() {
|
||||
if (specials_ok) {
|
||||
return true;
|
||||
}
|
||||
// bos_before_voice is optional, some packs do not insert it
|
||||
bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>");
|
||||
audio_bos = find_special_token(vocab, "<|audio_bos|>");
|
||||
if (audio_bos == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> in vocab\n");
|
||||
return false;
|
||||
}
|
||||
const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
|
||||
if (n_tok_embd == 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n");
|
||||
return false;
|
||||
}
|
||||
tok_embd.resize(n_tok_embd);
|
||||
if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n");
|
||||
return false;
|
||||
}
|
||||
GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0);
|
||||
specials_ok = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
// the table can be shorter than the vocab, so bound the row lookup
|
||||
void push_embd_row(std::vector<float> & dst, llama_token t) const {
|
||||
const size_t n_rows = tok_embd.size() / (size_t) n_embd;
|
||||
GGML_ASSERT(t >= 0 && (size_t) t < n_rows);
|
||||
dst.insert(dst.end(),
|
||||
tok_embd.begin() + (size_t) t * n_embd,
|
||||
tok_embd.begin() + (size_t) (t + 1) * n_embd);
|
||||
}
|
||||
|
||||
// token ids of the pieces the reference splits on, see split_into_best_sentences().
|
||||
// the leading token is dropped, it is the tokenizer's dummy prefix
|
||||
std::vector<llama_token> punct_ids(const char * s) const {
|
||||
std::vector<llama_token> ids(16);
|
||||
const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false);
|
||||
if (n <= 1) {
|
||||
return {};
|
||||
}
|
||||
return std::vector<llama_token>(ids.begin() + 1, ids.begin() + n);
|
||||
}
|
||||
|
||||
// cut after runs of boundary tokens, so punctuation stays with the sentence it ends
|
||||
static std::vector<std::vector<llama_token>> split_on(const std::vector<llama_token> & ids,
|
||||
const std::vector<llama_token> & boundary) {
|
||||
std::vector<std::vector<llama_token>> out;
|
||||
size_t start = 0;
|
||||
bool prev_was_boundary = false;
|
||||
for (size_t i = 0; i < ids.size(); i++) {
|
||||
const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end();
|
||||
if (!is_boundary && prev_was_boundary) {
|
||||
out.emplace_back(ids.begin() + start, ids.begin() + i);
|
||||
start = i;
|
||||
}
|
||||
prev_was_boundary = is_boundary;
|
||||
}
|
||||
out.emplace_back(ids.begin() + start, ids.end());
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<std::vector<llama_token>> split_chunks(const std::vector<llama_token> & ids) const {
|
||||
if ((int) ids.size() <= max_chunk_tokens) {
|
||||
return { ids };
|
||||
}
|
||||
const std::vector<llama_token> eos_punct = punct_ids(".!...?");
|
||||
const std::vector<llama_token> mid_punct = punct_ids(",;:");
|
||||
|
||||
// oversized sentences are split again on weaker punctuation, else words get skipped
|
||||
std::vector<std::vector<llama_token>> segments;
|
||||
for (auto & seg : split_on(ids, eos_punct)) {
|
||||
if ((int) seg.size() <= max_chunk_tokens) {
|
||||
segments.push_back(std::move(seg));
|
||||
continue;
|
||||
}
|
||||
auto sub = split_on(seg, mid_punct);
|
||||
if (sub.size() > 1) {
|
||||
for (auto & s : sub) {
|
||||
segments.push_back(std::move(s));
|
||||
}
|
||||
} else {
|
||||
segments.push_back(std::move(seg));
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<llama_token>> out;
|
||||
for (auto & seg : segments) {
|
||||
if (seg.empty()) {
|
||||
continue;
|
||||
}
|
||||
if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) {
|
||||
out.back().insert(out.back().end(), seg.begin(), seg.end());
|
||||
} else {
|
||||
out.push_back(std::move(seg));
|
||||
}
|
||||
}
|
||||
if (out.empty()) {
|
||||
out.push_back(ids);
|
||||
}
|
||||
for (const auto & c : out) {
|
||||
if ((int) c.size() > max_chunk_tokens) {
|
||||
LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, "
|
||||
"generation may skip words\n", c.size(), max_chunk_tokens);
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// _estimate_max_gen_len() plus the per-chunk tail guess, both in frames
|
||||
void arm_chunk_budget(size_t idx) {
|
||||
const int n_tok = (int) chunks[idx].size();
|
||||
chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate);
|
||||
// the pack may pin the tail, else the reference guesses it from the word count
|
||||
frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3);
|
||||
step_idx = 0;
|
||||
eos_step = -1;
|
||||
}
|
||||
|
||||
// ends the current chunk and, if there is another, re-prompts it on top of the voice
|
||||
int32_t finish_chunk(const float ** h_state_out, bool * out_stop) {
|
||||
if (!flush_gen_wav()) {
|
||||
return 1;
|
||||
}
|
||||
// the decoder restarts too, the next chunk's audio is not continuous with this one
|
||||
dec_state.clear();
|
||||
|
||||
if (chunk_idx + 1 >= chunks.size()) {
|
||||
*out_stop = true;
|
||||
*h_state_out = nullptr;
|
||||
return 0;
|
||||
}
|
||||
chunk_idx++;
|
||||
|
||||
// drop this chunk's text and audio, keep the voice conditioning
|
||||
llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1);
|
||||
pos = n_voice_pos;
|
||||
|
||||
const int n_e = n_embd;
|
||||
prompt_embd_buf.clear();
|
||||
for (llama_token t : chunks[chunk_idx]) {
|
||||
push_embd_row(prompt_embd_buf, t);
|
||||
}
|
||||
push_embd_row(prompt_embd_buf, audio_bos);
|
||||
arm_chunk_budget(chunk_idx);
|
||||
|
||||
const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e);
|
||||
GGML_ASSERT(n_rows > 0);
|
||||
decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e);
|
||||
batch.set_position_normal(pos, seq_id);
|
||||
batch.batch.logits[n_rows - 1] = 1;
|
||||
if (llama_decode(lctx, batch.batch) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n");
|
||||
return 1;
|
||||
}
|
||||
pos += n_rows;
|
||||
prompt_embd_buf.clear();
|
||||
|
||||
const float * he = llama_get_embeddings_ith(lctx, -1);
|
||||
h_state_buf.assign(he, he + n_embd);
|
||||
*h_state_out = h_state_buf.data();
|
||||
*out_stop = false;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// same normalization as prepare_text_prompt() in the reference, it affects quality
|
||||
static std::string prepare_text(const std::string & in, bool pad_short) {
|
||||
std::string s;
|
||||
s.reserve(in.size() + 1);
|
||||
for (char c : in) {
|
||||
if (c == '\n' || c == '\r') {
|
||||
s += ' ';
|
||||
} else if (c == ';') {
|
||||
s += ',';
|
||||
} else {
|
||||
s += c;
|
||||
}
|
||||
}
|
||||
const size_t b = s.find_first_not_of(' ');
|
||||
const size_t e = s.find_last_not_of(' ');
|
||||
if (b == std::string::npos) {
|
||||
return "";
|
||||
}
|
||||
s = s.substr(b, e - b + 1);
|
||||
if (s[0] >= 'a' && s[0] <= 'z') {
|
||||
s[0] = (char) (s[0] - 'a' + 'A');
|
||||
}
|
||||
const unsigned char last = (unsigned char) s.back();
|
||||
if (std::isalnum(last)) {
|
||||
s += '.';
|
||||
}
|
||||
if (pad_short && count_words(s) < 5) {
|
||||
s = std::string(8, ' ') + s;
|
||||
}
|
||||
return s;
|
||||
}
|
||||
|
||||
static int count_words(const std::string & s) {
|
||||
int n = 0;
|
||||
bool in_word = false;
|
||||
for (char c : s) {
|
||||
if (c == ' ') {
|
||||
in_word = false;
|
||||
} else if (!in_word) {
|
||||
in_word = true;
|
||||
n++;
|
||||
}
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
// runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
|
||||
bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
|
||||
if (!mtmd_support_audio(mctx)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
|
||||
return false;
|
||||
}
|
||||
const std::string marker = mtmd_default_marker();
|
||||
mtmd_input_text text{ marker.c_str(), marker.size(), false, true };
|
||||
mtmd_input_chunks * chunks = mtmd_input_chunks_init();
|
||||
const mtmd_bitmap * bptr = bitmap;
|
||||
bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
|
||||
if (ok) {
|
||||
ok = false;
|
||||
for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
|
||||
const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
|
||||
if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
|
||||
continue;
|
||||
}
|
||||
if (mtmd_encode_chunk(mctx, chunk) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n");
|
||||
break;
|
||||
}
|
||||
const float * embd = mtmd_get_output_embd(mctx);
|
||||
const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
|
||||
out.assign(embd, embd + n);
|
||||
ok = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
mtmd_input_chunks_free(chunks);
|
||||
return ok;
|
||||
}
|
||||
|
||||
// decodes the buffered latents, the mimi decoder state carries over between calls
|
||||
bool flush_gen_wav() {
|
||||
if (feats_buf.empty()) {
|
||||
return true;
|
||||
}
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
|
||||
inp.feats = feats_buf.data();
|
||||
inp.n_feats = feats_buf.size();
|
||||
inp.seed = seed;
|
||||
inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
|
||||
inp.state_size = dec_state.size();
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
|
||||
return false;
|
||||
}
|
||||
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
|
||||
dec_state.assign(out.state_data, out.state_data + out.state_size);
|
||||
feats_buf.clear();
|
||||
return true;
|
||||
}
|
||||
|
||||
pockettts_pack_settings pack;
|
||||
bool specials_ok = false;
|
||||
llama_token bos_before_voice = LLAMA_TOKEN_NULL;
|
||||
llama_token audio_bos = LLAMA_TOKEN_NULL;
|
||||
std::vector<float> tok_embd;
|
||||
|
||||
llama_seq_id seq_id = 0;
|
||||
int pos = 0;
|
||||
std::vector<float> prompt_embd_buf;
|
||||
std::unique_ptr<decode_embd_batch> prompt_batch;
|
||||
int n_prompt = 0;
|
||||
int prompt_pos = 0;
|
||||
uint32_t seed = UINT32_MAX;
|
||||
// end-of-speech is latched, then a few more frames are generated as tail padding
|
||||
int step_idx = 0;
|
||||
int eos_step = -1;
|
||||
int frames_after_eos = 3;
|
||||
static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference
|
||||
static constexpr double frame_rate = 12.5;
|
||||
std::vector<std::vector<llama_token>> chunks;
|
||||
size_t chunk_idx = 0;
|
||||
int n_voice_pos = 0; // KV positions held by the voice conditioning
|
||||
int chunk_budget = 0;
|
||||
|
||||
// latents are decoded a window at a time, the decoder state bridges the windows
|
||||
size_t window_frames = 8;
|
||||
std::vector<float> feats_buf;
|
||||
std::vector<uint8_t> dec_state;
|
||||
std::vector<float> audio_pcm;
|
||||
std::vector<float> h_state_buf;
|
||||
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
|
||||
std::vector<char> out_buf;
|
||||
};
|
||||
|
||||
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
|
||||
switch (mtmd_gen_audio_get_info(mctx).type) {
|
||||
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new qwen3tts_gen_audio_pipeline(lctx, mctx));
|
||||
case MTMD_GEN_AUDIO_TYPE_POCKETTTS:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new pockettts_gen_audio_pipeline(lctx, mctx));
|
||||
default:
|
||||
return nullptr;
|
||||
}
|
||||
@@ -483,11 +1040,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled,
|
||||
const float * h_state_in, const float ** h_state_out) {
|
||||
const float * h_state_in, const float ** h_state_out,
|
||||
bool * out_stop) {
|
||||
if (!ctx->pipeline) {
|
||||
return 1;
|
||||
}
|
||||
return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out);
|
||||
bool stop = false;
|
||||
const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop);
|
||||
if (out_stop) {
|
||||
*out_stop = stop;
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate,
|
||||
|
||||
@@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp {
|
||||
mtmd_bitmap * speaker_ref; // optional, can be NULL
|
||||
const char * lang; // optional, can be NULL
|
||||
|
||||
int32_t top_k;
|
||||
float top_p;
|
||||
int32_t top_k;
|
||||
float top_p;
|
||||
uint32_t seed; // UINT32_MAX for random (default: random)
|
||||
|
||||
enum mtmd_helper_gen_audio_outtype out_type;
|
||||
};
|
||||
@@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt(
|
||||
int32_t n_batch);
|
||||
|
||||
// generates one frame; must only be called after step_prompt() has returned 0
|
||||
// h_state_out is valid until next step_gen() or reset() call
|
||||
// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token
|
||||
// out_stop (optional) is set on end-of-speech, the caller must then stop the loop
|
||||
// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated
|
||||
MTMD_API int32_t mtmd_helper_gen_audio_step_gen(
|
||||
mtmd_helper_gen_audio * ctx,
|
||||
llama_token sampled,
|
||||
const float * h_state_in,
|
||||
const float ** h_state_out);
|
||||
const float ** h_state_out,
|
||||
bool * out_stop);
|
||||
|
||||
// out_data valid until next get_output() or reset() call
|
||||
// out_n_samples (optional, can be NULL) receives the number of generated PCM samples
|
||||
@@ -261,8 +265,8 @@ struct gen_audio {
|
||||
int32_t step_prompt(int32_t n_batch) {
|
||||
return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch);
|
||||
}
|
||||
int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) {
|
||||
return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out);
|
||||
int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) {
|
||||
return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop);
|
||||
}
|
||||
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) {
|
||||
return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples);
|
||||
|
||||
+77
-16
@@ -477,6 +477,7 @@ struct mtmd_context {
|
||||
// generation context
|
||||
struct clip_ctx * ctx_gen_a; // audio
|
||||
std::vector<int32_t> gen_out_codes; // this frame's 16 sampled codes (GEN_CODE)
|
||||
std::vector<float> gen_out_feats; // this frame's continuous features, if any (GEN_CODE)
|
||||
std::vector<float> gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE)
|
||||
std::vector<float> gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV)
|
||||
std::vector<uint8_t> gen_out_state; // state to feed into the next GEN_WAV call
|
||||
@@ -979,6 +980,10 @@ struct mtmd_context {
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_qwen3tts_spk>(ctx_a);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
{
|
||||
audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
|
||||
} break;
|
||||
default:
|
||||
throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
|
||||
}
|
||||
@@ -1798,16 +1803,22 @@ float * mtmd_get_output_embd(mtmd_context * ctx) {
|
||||
//
|
||||
|
||||
mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
mtmd_gen_audio_info info;
|
||||
mtmd_gen_audio_info info{};
|
||||
info.model_variant = "";
|
||||
if (!ctx->ctx_gen_a) {
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
return info;
|
||||
}
|
||||
info.model_variant = clip_get_hparams(ctx->ctx_gen_a)->gen_model_variant.c_str();
|
||||
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS;
|
||||
info.sample_rate = 24000;
|
||||
break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS;
|
||||
info.sample_rate = 24000;
|
||||
break;
|
||||
default:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
break;
|
||||
@@ -1815,6 +1826,33 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
return info;
|
||||
}
|
||||
|
||||
mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
|
||||
mtmd_gen_inp inp{};
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
inp.seed = UINT32_MAX;
|
||||
if (!ctx->ctx_gen_a) {
|
||||
return inp;
|
||||
}
|
||||
|
||||
switch (clip_get_projector_type(ctx->ctx_gen_a)) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
// https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base/blob/main/generation_config.json
|
||||
inp.top_k = 50;
|
||||
inp.top_p = 1.0f;
|
||||
inp.temp = 0.9f; // TODO: handle this on graph
|
||||
break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
// https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py
|
||||
inp.top_k = 50;
|
||||
inp.top_p = 1.0f;
|
||||
inp.temp = 0.7f;
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
return inp;
|
||||
}
|
||||
|
||||
static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
|
||||
clip_ctx * ctx_clip = ctx->ctx_gen_a;
|
||||
if (!ctx_clip) {
|
||||
@@ -1822,6 +1860,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
return 1;
|
||||
}
|
||||
|
||||
*out = {};
|
||||
|
||||
if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) {
|
||||
const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip);
|
||||
|
||||
@@ -1835,16 +1875,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
|
||||
std::vector<float> out_embd(n_embd);
|
||||
std::vector<int32_t> out_codes;
|
||||
std::vector<float> out_feats;
|
||||
bool is_eos = false;
|
||||
|
||||
clip_encode_params params;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
|
||||
params.out_embd = &out_embd;
|
||||
params.out_codes = &out_codes;
|
||||
params.code0 = inp->code0;
|
||||
params.top_k = inp->top_k;
|
||||
params.top_p = inp->top_p;
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_CODE;
|
||||
params.out_embd = &out_embd;
|
||||
params.out_codes = &out_codes;
|
||||
params.out_feats = &out_feats;
|
||||
params.code0 = inp->code0;
|
||||
params.top_k = inp->top_k;
|
||||
params.top_p = inp->top_p;
|
||||
params.seed = inp->seed;
|
||||
params.temp = inp->temp;
|
||||
params.out_is_eos = &is_eos;
|
||||
|
||||
if (!clip_encode(ctx_clip, ¶ms)) {
|
||||
LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
|
||||
@@ -1853,19 +1899,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
|
||||
ctx->gen_out_embd = std::move(out_embd);
|
||||
ctx->gen_out_codes = std::move(out_codes);
|
||||
ctx->gen_out_feats = std::move(out_feats);
|
||||
|
||||
out->embd = ctx->gen_out_embd.data();
|
||||
out->codes = ctx->gen_out_codes.data();
|
||||
out->n_codes = ctx->gen_out_codes.size();
|
||||
out->embd = ctx->gen_out_embd.data();
|
||||
out->codes = ctx->gen_out_codes.data();
|
||||
out->n_codes = ctx->gen_out_codes.size();
|
||||
out->feats = ctx->gen_out_feats.data();
|
||||
out->n_feats = ctx->gen_out_feats.size();
|
||||
out->is_eos = is_eos;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
if (!inp->codes || inp->n_codes == 0) {
|
||||
LOG_ERR("%s: codes required for gen_wav\n", __func__);
|
||||
const bool has_codes = inp->codes && inp->n_codes > 0;
|
||||
const bool has_feats = inp->feats && inp->n_feats > 0;
|
||||
if (has_codes == has_feats) {
|
||||
LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
|
||||
std::vector<int32_t> in_codes;
|
||||
std::vector<float> in_feats;
|
||||
if (has_codes) {
|
||||
in_codes.assign(inp->codes, inp->codes + inp->n_codes);
|
||||
} else {
|
||||
in_feats.assign(inp->feats, inp->feats + inp->n_feats);
|
||||
}
|
||||
std::vector<uint8_t> in_state;
|
||||
if (inp->state_data) {
|
||||
in_state.assign(inp->state_data, inp->state_data + inp->state_size);
|
||||
@@ -1885,7 +1943,10 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
|
||||
params.imgs = &batch;
|
||||
params.n_threads = ctx->n_threads;
|
||||
params.gen_process = CLIP_GEN_PROCESS_GEN_WAV;
|
||||
params.codes = &in_codes;
|
||||
// gen_wav draws no randomness, but keep the seed so it does not reseed mid-generation
|
||||
params.seed = inp->seed;
|
||||
params.codes = has_codes ? &in_codes : nullptr;
|
||||
params.feats = has_feats ? &in_feats : nullptr;
|
||||
params.out_audio = &ctx->gen_out_audio;
|
||||
params.state_in = inp->state_data ? &in_state : nullptr;
|
||||
params.state_out = &ctx->gen_out_state;
|
||||
|
||||
+21
-1
@@ -344,18 +344,25 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname);
|
||||
enum mtmd_gen_audio_type {
|
||||
MTMD_GEN_AUDIO_TYPE_NONE, // not supported
|
||||
MTMD_GEN_AUDIO_TYPE_QWEN3TTS,
|
||||
MTMD_GEN_AUDIO_TYPE_POCKETTTS,
|
||||
};
|
||||
|
||||
struct mtmd_gen_audio_info {
|
||||
enum mtmd_gen_audio_type type;
|
||||
int32_t sample_rate; // in Hz, for example 24000 for qwen3tts
|
||||
const char * model_variant; // name of the weight variant, can be nullptr if not applicable
|
||||
};
|
||||
|
||||
MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx);
|
||||
|
||||
|
||||
enum mtmd_gen_process_type {
|
||||
MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to semantic (codes, mel-spectrogram, etc.)
|
||||
MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio
|
||||
// for qwen3tts, this is code2wav
|
||||
// for pocket-tts, this is mimi decoder
|
||||
};
|
||||
|
||||
struct mtmd_gen_inp {
|
||||
enum mtmd_gen_process_type type;
|
||||
|
||||
@@ -364,21 +371,30 @@ struct mtmd_gen_inp {
|
||||
float * embd; // the hidden state from backbone, must have n_text_embd elements
|
||||
int32_t top_k;
|
||||
float top_p;
|
||||
uint32_t seed; // UINT32_MAX for random
|
||||
float temp; // sampling temperature, or noise scale for flow-matching decoders
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
// pass either codes (discrete) or feats (continuous), depending on the pipeline
|
||||
int32_t * codes;
|
||||
size_t n_codes;
|
||||
const float * feats;
|
||||
size_t n_feats;
|
||||
const char * state_data;
|
||||
size_t state_size;
|
||||
};
|
||||
|
||||
struct mtmd_gen_out {
|
||||
// note: output memory is allocated by the context, valid until next process() call
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_CODE
|
||||
const int32_t * codes;
|
||||
size_t n_codes;
|
||||
size_t n_codes;
|
||||
const float * feats; // continuous counterpart of codes
|
||||
size_t n_feats;
|
||||
const float * embd; // the generated hidden state, to be fed back to backbone
|
||||
// it must have n_text_embd elements
|
||||
bool is_eos; // only set by pipelines having the EOS head inside mmproj
|
||||
|
||||
// for MTMD_GEN_PROCESS_TYPE_GEN_WAV
|
||||
const float * audio;
|
||||
@@ -386,6 +402,10 @@ struct mtmd_gen_out {
|
||||
const char * state_data;
|
||||
size_t state_size;
|
||||
};
|
||||
|
||||
// defaults tuned for the loaded pipeline, callers override only what they care about
|
||||
MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
|
||||
|
||||
// note: this API is stateless, caller must handle state management and audio frame accumulation
|
||||
MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
|
||||
const struct mtmd_gen_inp * inp,
|
||||
|
||||
@@ -2,11 +2,5 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu
|
||||
pillow~=11.3.0
|
||||
|
||||
## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility
|
||||
torch==2.11.0; platform_machine != "s390x" # check_requirements: ignore "=="
|
||||
torch==2.11.0 # check_requirements: ignore "=="
|
||||
torchvision==0.26.0; platform_machine != "s390x" # check_requirements: ignore "=="
|
||||
|
||||
# torch s390x packages can only be found from nightly builds
|
||||
--extra-index-url https://download.pytorch.org/whl/nightly
|
||||
torch>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "=="
|
||||
torchvision>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "=="
|
||||
|
||||
@@ -397,12 +397,7 @@ struct server_slot {
|
||||
|
||||
bool need_embd() const {
|
||||
GGML_ASSERT(task);
|
||||
return task->need_embd() || (spec && common_speculative_need_embd(spec));
|
||||
}
|
||||
|
||||
bool need_embd_nextn() const {
|
||||
GGML_ASSERT(task);
|
||||
return spec && common_speculative_need_embd_nextn(spec);
|
||||
return task->need_embd();
|
||||
}
|
||||
|
||||
// if the context does not have a memory module then all embeddings have to be computed within a single ubatch
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#include "server-tools.h"
|
||||
|
||||
#include "subproc.h"
|
||||
#include "base64.hpp"
|
||||
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
@@ -864,6 +865,7 @@ static bool path_glob_match(const std::string & pattern, const std::string & rel
|
||||
//
|
||||
|
||||
static constexpr size_t SERVER_TOOL_READ_FILE_MAX_SIZE = 16 * 1024; // 16 KB
|
||||
static constexpr size_t SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64 = 32 * 1024 * 1024; // 32 MB
|
||||
|
||||
struct server_tool_read_file : server_tool {
|
||||
server_tool_read_file() {
|
||||
@@ -899,6 +901,8 @@ struct server_tool_read_file : server_tool {
|
||||
int start_line = json_value(params, "start_line", 1);
|
||||
int end_line = json_value(params, "end_line", -1); // -1 = no limit
|
||||
bool append_loc = json_value(params, "append_loc", false);
|
||||
// comes from the x-resp-type header, the model cannot ask for it
|
||||
bool as_base64 = json_value(params, "resp_type", std::string()) == "base64";
|
||||
|
||||
auto io = make_tools_io(params);
|
||||
|
||||
@@ -906,6 +910,23 @@ struct server_tool_read_file : server_tool {
|
||||
if (!io->file_size(path, file_size)) {
|
||||
return {{"error", "cannot stat file: " + path}};
|
||||
}
|
||||
|
||||
if (as_base64) {
|
||||
if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64) {
|
||||
return {{"error", string_format(
|
||||
"file too large (%zu bytes, max %zu)",
|
||||
(size_t)file_size, SERVER_TOOL_READ_FILE_MAX_SIZE_BASE64)}};
|
||||
}
|
||||
std::string content;
|
||||
if (!io->read_file(path, content)) {
|
||||
return {{"error", "failed to open file: " + path}};
|
||||
}
|
||||
return {
|
||||
{"base64", base64::encode(content.data(), content.size())},
|
||||
{"size_bytes", (size_t) content.size()},
|
||||
};
|
||||
}
|
||||
|
||||
if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE && end_line == -1) {
|
||||
return {{"error", string_format(
|
||||
"file too large (%zu bytes, max %zu). Use start_line/end_line to read a portion.",
|
||||
@@ -2135,6 +2156,15 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools,
|
||||
params["runtime"] = runtime->spec();
|
||||
}
|
||||
|
||||
// x-resp-type header is only used by read_file for now
|
||||
if (params.contains("resp_type")) {
|
||||
params.erase("resp_type");
|
||||
}
|
||||
auto resp_type = get_header(req.headers, "x-resp-type");
|
||||
if (!resp_type.empty()) {
|
||||
params["resp_type"] = resp_type;
|
||||
}
|
||||
|
||||
server_tool & tool = find_tool(tools, tool_name, stream);
|
||||
|
||||
if (stream) {
|
||||
|
||||
@@ -27,8 +27,8 @@ def test_with_and_without_draft():
|
||||
global server
|
||||
request = {
|
||||
"prompt": "I believe the meaning of life is",
|
||||
"temperature": 0.8,
|
||||
"top_k": 40,
|
||||
"temperature": 0.2,
|
||||
"top_k": 5,
|
||||
"seed": 4242,
|
||||
"n_predict": 16,
|
||||
"return_tokens": True,
|
||||
@@ -36,7 +36,6 @@ def test_with_and_without_draft():
|
||||
|
||||
server.model_draft = None # disable draft model
|
||||
server.spec_type = None
|
||||
server.backend_sampling = True
|
||||
server.start()
|
||||
res = server.make_request("POST", "/completion", data=request)
|
||||
assert res.status_code == 200
|
||||
@@ -45,7 +44,6 @@ def test_with_and_without_draft():
|
||||
|
||||
# create new server with draft model
|
||||
create_server()
|
||||
server.backend_sampling = True
|
||||
server.start()
|
||||
res = server.make_request("POST", "/completion", data=request)
|
||||
assert res.status_code == 200
|
||||
|
||||
@@ -32,3 +32,28 @@ llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \
|
||||
--tts-speaker-file speaker.mp3 \
|
||||
--output out.wav
|
||||
```
|
||||
|
||||
## Pocket TTS
|
||||
|
||||
Available params:
|
||||
- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it
|
||||
- Note: `lang` is not used, the language is a property of the weights
|
||||
|
||||
Example usage:
|
||||
|
||||
```sh
|
||||
llama-tts -m pocket-tts.gguf \
|
||||
-mm mmproj-pocket-tts.gguf \
|
||||
-p "Hello world" \
|
||||
--tts-speaker-file speaker.mp3 \
|
||||
--output out.wav
|
||||
```
|
||||
|
||||
**Note for GGUF conversion:**
|
||||
|
||||
The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/<name>` directories, **not** the root directory:
|
||||
|
||||
```sh
|
||||
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf
|
||||
python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf
|
||||
```
|
||||
|
||||
+10
-5
@@ -119,6 +119,7 @@ int main(int argc, char ** argv) {
|
||||
inp.lang = params.tts_lang.c_str();
|
||||
inp.top_k = params.sampling.top_k;
|
||||
inp.top_p = params.sampling.top_p;
|
||||
inp.seed = params.sampling.seed;
|
||||
inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
|
||||
|
||||
//
|
||||
@@ -143,8 +144,7 @@ int main(int argc, char ** argv) {
|
||||
}
|
||||
}
|
||||
|
||||
const llama_vocab * vocab = llama_model_get_vocab(model);
|
||||
|
||||
// note: some pipelines ignore this token and use the hidden state instead
|
||||
auto sample_semantic_code = [&]() -> llama_token {
|
||||
llama_token t = common_sampler_sample(smpl, lctx, -1);
|
||||
common_sampler_accept(smpl, t, true);
|
||||
@@ -159,19 +159,24 @@ int main(int argc, char ** argv) {
|
||||
tts_timings timings;
|
||||
const int64_t t_gen_start_us = ggml_time_us();
|
||||
|
||||
for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) {
|
||||
bool stop = false;
|
||||
while (!stop && n_frames < max_new) {
|
||||
const float * h_next = nullptr;
|
||||
|
||||
// stage 2+3: semantic --> acoustic details --> audio waveform
|
||||
// step_gen() runs both stages and returns new h_state for next step
|
||||
if (gen.step_gen(sampled, h_state, &h_next) != 0) {
|
||||
if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) {
|
||||
LOG_ERR("step_gen failed at frame %d\n", n_frames);
|
||||
return 1;
|
||||
}
|
||||
if (!h_next) {
|
||||
break; // stopped without generating a frame
|
||||
}
|
||||
|
||||
n_frames++;
|
||||
h_state = h_next;
|
||||
sampled = sample_semantic_code();
|
||||
timings.report(n_frames + 1);
|
||||
timings.report(n_frames);
|
||||
}
|
||||
const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6;
|
||||
|
||||
|
||||
+5
-4
@@ -3,6 +3,7 @@
|
||||
import { ChevronDown } from '@lucide/svelte';
|
||||
import * as Collapsible from '$lib/components/ui/collapsible';
|
||||
import { STATS_UNITS } from '$lib/constants';
|
||||
import { gaugePopup } from '$lib/stores/context-gauge-popup.svelte';
|
||||
|
||||
interface Props {
|
||||
currentRead: number;
|
||||
@@ -30,19 +31,19 @@
|
||||
transientDetails
|
||||
}: Props = $props();
|
||||
|
||||
let open = $state(false);
|
||||
|
||||
const hasCumulative = $derived(cumulativeRead > 0 || cumulativeOutput > 0);
|
||||
const hasCurrent = $derived(currentRead > 0 || currentOutput > 0);
|
||||
</script>
|
||||
|
||||
<Collapsible.Root bind:open class="mt-3 border-t border-border/50 pt-4">
|
||||
<Collapsible.Root bind:open={gaugePopup.detailsOpen} class="mt-3 border-t border-border/50 pt-4">
|
||||
<Collapsible.Trigger
|
||||
class="flex w-full cursor-pointer items-center gap-1 text-xs text-muted-foreground hover:text-foreground"
|
||||
>
|
||||
<span>Token usage details</span>
|
||||
|
||||
<ChevronDown class={'ml-auto h-3 w-3 transition-transform' + (open ? ' rotate-180' : '')} />
|
||||
<ChevronDown
|
||||
class={'ml-auto h-3 w-3 transition-transform' + (gaugePopup.detailsOpen ? ' rotate-180' : '')}
|
||||
/>
|
||||
</Collapsible.Trigger>
|
||||
|
||||
<Collapsible.Content class="flex flex-col gap-4 text-xs pt-4">
|
||||
|
||||
+3
@@ -7,6 +7,7 @@
|
||||
import ChatMessageToolCallBlockGetInfo from './ChatMessageToolCallBlockGetInfo.svelte';
|
||||
import ChatMessageToolCallBlockGrepSearch from './ChatMessageToolCallBlockGrepSearch.svelte';
|
||||
import ChatMessageToolCallBlockReadFile from './ChatMessageToolCallBlockReadFile.svelte';
|
||||
import ChatMessageToolCallBlockReadMedia from './ChatMessageToolCallBlockReadMedia.svelte';
|
||||
import ChatMessageToolCallBlockRunJavascript from './ChatMessageToolCallBlockRunJavascript.svelte';
|
||||
import ChatMessageToolCallBlockSearchResults from './ChatMessageToolCallBlockSearchResults.svelte';
|
||||
import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte';
|
||||
@@ -45,6 +46,8 @@
|
||||
<ChatMessageToolCallBlockGetInfo {section} {isStreaming} />
|
||||
{:else if section.toolName === BuiltInTool.READ_FILE}
|
||||
<ChatMessageToolCallBlockReadFile {section} {open} {isStreaming} {onToggle} />
|
||||
{:else if section.toolName === BuiltInTool.READ_MEDIA}
|
||||
<ChatMessageToolCallBlockReadMedia {section} {open} {isStreaming} {onToggle} />
|
||||
{:else if section.toolName === BuiltInTool.EDIT_FILE}
|
||||
<ChatMessageToolCallBlockEditFile {section} {open} {isStreaming} {onToggle} />
|
||||
{:else if section.toolName === BuiltInTool.WRITE_FILE}
|
||||
|
||||
+26
-11
@@ -8,14 +8,16 @@
|
||||
import { MarkdownContent, SyntaxHighlightedCode } from '$lib/components/app';
|
||||
import { MAX_HEIGHT_CODE_BLOCK } from '$lib/constants';
|
||||
import { getBuiltinToolUi } from '$lib/constants/built-in-tools';
|
||||
import { FileTypeText, ToolResultKind } from '$lib/enums';
|
||||
import { AttachmentType, FileTypeText, MimeTypeAudio, ToolResultKind } from '$lib/enums';
|
||||
import type { DatabaseMessageExtra } from '$lib/types';
|
||||
import {
|
||||
type AgenticSection,
|
||||
classifyToolResult,
|
||||
formatJsonPretty,
|
||||
parseToolResultWithImages
|
||||
parseToolResultWithMedia,
|
||||
type ToolResultLine
|
||||
} from '$lib/utils';
|
||||
import { createBase64DataUrl } from '$lib/utils/data-url';
|
||||
|
||||
interface Props {
|
||||
section: AgenticSection;
|
||||
@@ -29,8 +31,8 @@
|
||||
|
||||
const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? '');
|
||||
const outputKind = $derived(classifyToolResult(section.toolResult));
|
||||
const parsedLines = $derived(
|
||||
section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
|
||||
const parsedLines: ToolResultLine[] = $derived(
|
||||
section.toolResult ? parseToolResultWithMedia(section.toolResult, attachments) : []
|
||||
);
|
||||
</script>
|
||||
|
||||
@@ -103,13 +105,26 @@
|
||||
<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">
|
||||
{line.text}
|
||||
</div>
|
||||
{#if line.image}
|
||||
<img
|
||||
src={line.image.base64Url}
|
||||
alt={line.image.name}
|
||||
class="mt-2 mb-2 h-auto max-w-full rounded-lg"
|
||||
loading="lazy"
|
||||
/>
|
||||
{#if line.media}
|
||||
{#if line.media.type === AttachmentType.AUDIO}
|
||||
{@const audioMimeType = line.media.mimeType ?? MimeTypeAudio.MP3_MPEG}
|
||||
<div class="mt-2 mb-2">
|
||||
<audio controls class="w-full rounded-lg">
|
||||
<source
|
||||
src={createBase64DataUrl(audioMimeType, line.media.base64Data)}
|
||||
type={audioMimeType}
|
||||
/>
|
||||
Your browser does not support the audio element.
|
||||
</audio>
|
||||
</div>
|
||||
{:else}
|
||||
<img
|
||||
src={line.media.base64Url}
|
||||
alt={line.media.name}
|
||||
class="mt-2 mb-2 h-auto max-w-full rounded-lg"
|
||||
loading="lazy"
|
||||
/>
|
||||
{/if}
|
||||
{/if}
|
||||
{/each}
|
||||
</div>
|
||||
|
||||
+5
-5
@@ -23,7 +23,7 @@
|
||||
isExitCodeSummaryLine,
|
||||
parseExecShellCommandError,
|
||||
parseExecShellCommandExitStatus,
|
||||
parseToolResultWithImages,
|
||||
parseToolResultWithMedia,
|
||||
type ToolResultLine
|
||||
} from '$lib/utils';
|
||||
|
||||
@@ -53,7 +53,7 @@
|
||||
);
|
||||
|
||||
const parsedLines: ToolResultLine[] = $derived(
|
||||
section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
|
||||
section.toolResult ? parseToolResultWithMedia(section.toolResult, attachments) : []
|
||||
);
|
||||
|
||||
// Drop the trailing "[exit code: N]" line - rendered as a colored
|
||||
@@ -223,10 +223,10 @@
|
||||
>
|
||||
{#each outputLines as line, i (i)}
|
||||
<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">{line.text}</div>
|
||||
{#if line.image}
|
||||
{#if line.media}
|
||||
<img
|
||||
src={line.image.base64Url}
|
||||
alt={line.image.name}
|
||||
src={line.media.base64Url}
|
||||
alt={line.media.name}
|
||||
class="mt-2 mb-2 h-auto max-w-full rounded-lg"
|
||||
loading="lazy"
|
||||
/>
|
||||
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
<script lang="ts">
|
||||
import { parseReadMediaMeta } from './parsers/read-media';
|
||||
import ToolCallBlock from './ToolCallBlock.svelte';
|
||||
import { ATTACHMENT_SAVED_REGEX } from '$lib/constants/agentic';
|
||||
import { AttachmentType, MimeTypeAudio } from '$lib/enums';
|
||||
import type { DatabaseMessageExtraAudioFile, DatabaseMessageExtraImageFile } from '$lib/types';
|
||||
import { type AgenticSection } from '$lib/utils';
|
||||
import { createBase64DataUrl } from '$lib/utils/data-url';
|
||||
|
||||
interface Props {
|
||||
section: AgenticSection;
|
||||
open: boolean;
|
||||
isStreaming: boolean;
|
||||
onToggle?: () => void;
|
||||
}
|
||||
|
||||
let { isStreaming, onToggle, open, section }: Props = $props();
|
||||
|
||||
const readMediaMeta = $derived(parseReadMediaMeta(section));
|
||||
|
||||
// extractBase64Attachments swapped the data URI line for [Attachment saved: name]
|
||||
// and moved the bytes to the message extras, so the name is the only link back
|
||||
const mediaAttachment = $derived.by(() => {
|
||||
const extras = section.toolResultExtras;
|
||||
|
||||
if (!extras || extras.length === 0) return null;
|
||||
|
||||
const match = section.toolResult?.match(ATTACHMENT_SAVED_REGEX);
|
||||
|
||||
if (!match) return null;
|
||||
|
||||
const attachmentName = match[1];
|
||||
|
||||
return (
|
||||
extras.find(
|
||||
(e): e is DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile =>
|
||||
(e.type === AttachmentType.IMAGE || e.type === AttachmentType.AUDIO) &&
|
||||
e.name === attachmentName
|
||||
) ?? null
|
||||
);
|
||||
});
|
||||
|
||||
const audioMimeType = $derived(readMediaMeta?.mimeType ?? MimeTypeAudio.MP3_MPEG);
|
||||
</script>
|
||||
|
||||
<ToolCallBlock {section} {open} {isStreaming} meta={readMediaMeta} {onToggle}>
|
||||
{#snippet titleSnippet()}
|
||||
<span class="text-muted-foreground">Read media </span>
|
||||
<span class="font-mono">{readMediaMeta?.fileName}</span>
|
||||
{/snippet}
|
||||
|
||||
{#snippet children(_meta, _ctx)}
|
||||
{#if section.toolResult}
|
||||
{#if !mediaAttachment}
|
||||
<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
|
||||
Media attachment not found in message extras
|
||||
</div>
|
||||
{:else if mediaAttachment.type === AttachmentType.AUDIO}
|
||||
<div class="mt-2">
|
||||
<audio controls class="w-full rounded-lg">
|
||||
<source
|
||||
src={createBase64DataUrl(audioMimeType, mediaAttachment.base64Data)}
|
||||
type={audioMimeType}
|
||||
/>
|
||||
Your browser does not support the audio element.
|
||||
</audio>
|
||||
</div>
|
||||
{:else}
|
||||
<div class="mt-2">
|
||||
<img
|
||||
src={mediaAttachment.base64Url}
|
||||
alt={readMediaMeta?.fileName ?? 'media'}
|
||||
class="max-h-[60vh] max-w-full rounded-lg object-contain shadow-lg"
|
||||
loading="lazy"
|
||||
/>
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
{#if readMediaMeta?.sizeBytes || readMediaMeta?.mimeType}
|
||||
<div class="mt-2 flex gap-4 text-xs text-muted-foreground">
|
||||
{#if readMediaMeta?.sizeBytes}
|
||||
<span>Size: {readMediaMeta.sizeBytes} bytes</span>
|
||||
{/if}
|
||||
{#if readMediaMeta?.mimeType}
|
||||
<span>MIME: {readMediaMeta.mimeType}</span>
|
||||
{/if}
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
{#if readMediaMeta?.path}
|
||||
<div class="mt-1 font-mono text-xs text-muted-foreground/60">{readMediaMeta.path}</div>
|
||||
{/if}
|
||||
{:else}
|
||||
<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
|
||||
Waiting for media data...
|
||||
</div>
|
||||
{/if}
|
||||
{/snippet}
|
||||
</ToolCallBlock>
|
||||
+56
@@ -0,0 +1,56 @@
|
||||
import { FILE_PATH_SEPARATOR_REGEX, NEWLINE } from '$lib/constants/code';
|
||||
import {
|
||||
PREFIX_FILE,
|
||||
PREFIX_MIME,
|
||||
PREFIX_SIZE,
|
||||
READ_MEDIA_SIZE_REGEX
|
||||
} from '$lib/constants/read-media';
|
||||
import type { AgenticSection } from '$lib/utils';
|
||||
|
||||
export interface ReadMediaMeta {
|
||||
fileName: string;
|
||||
path: string;
|
||||
sizeBytes?: number;
|
||||
mimeType?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse read_media tool result to extract metadata.
|
||||
* Expected format (after extractBase64Attachments processing):
|
||||
* File: /path/to/file.png
|
||||
* Size: 12345 bytes
|
||||
* MIME: image/png
|
||||
* [Attachment saved: mcp-attachment-xxx.png]
|
||||
*
|
||||
* The data URI line is replaced by the attachment marker by
|
||||
* agenticStore.extractBase64Attachments before storage.
|
||||
*/
|
||||
export function parseReadMediaMeta(section: AgenticSection): ReadMediaMeta | null {
|
||||
if (!section.toolResult) return null;
|
||||
|
||||
const lines = section.toolResult.split(NEWLINE);
|
||||
|
||||
let fileName = '';
|
||||
let path = '';
|
||||
let sizeBytes: number | undefined;
|
||||
let mimeType: string | undefined;
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
|
||||
if (trimmed.startsWith(PREFIX_FILE)) {
|
||||
path = trimmed.slice(PREFIX_FILE.length).trim();
|
||||
fileName = path.split(FILE_PATH_SEPARATOR_REGEX).pop() ?? path;
|
||||
} else if (trimmed.startsWith(PREFIX_SIZE)) {
|
||||
const match = trimmed.match(READ_MEDIA_SIZE_REGEX);
|
||||
|
||||
if (match) sizeBytes = Number(match[1]);
|
||||
} else if (trimmed.startsWith(PREFIX_MIME)) {
|
||||
mimeType = trimmed.slice(PREFIX_MIME.length).trim();
|
||||
}
|
||||
}
|
||||
|
||||
if (!path) return null;
|
||||
|
||||
return { fileName, mimeType, path, sizeBytes };
|
||||
}
|
||||
@@ -10,6 +10,7 @@
|
||||
import {
|
||||
Braces,
|
||||
Clock,
|
||||
Eye,
|
||||
FilePen,
|
||||
FilePlus,
|
||||
FileSearch,
|
||||
@@ -47,6 +48,7 @@ export const BUILTIN_TOOL_UI: Readonly<Record<BuiltInTool, BuiltinToolUiEntry>>
|
||||
source: ToolSource.BUILTIN
|
||||
},
|
||||
[BuiltInTool.READ_FILE]: { icon: FileText, label: 'Read file', source: ToolSource.BUILTIN },
|
||||
[BuiltInTool.READ_MEDIA]: { icon: Eye, label: 'Read media', source: ToolSource.FRONTEND },
|
||||
[BuiltInTool.RUN_JAVASCRIPT]: {
|
||||
icon: Braces,
|
||||
label: 'Run JavaScript',
|
||||
|
||||
@@ -18,6 +18,9 @@ export const TRIM_TRAILING_PADDING_REGEX = /(?:\n[ \t]*)+$/;
|
||||
// `C:\foo\bar.txt`. Used wherever a parameter accepts a user-supplied path.
|
||||
export const FILE_PATH_SEPARATOR_REGEX = /[\\/]/;
|
||||
|
||||
// Separates a file name from its extension, e.g. the '.' in `cover.png`.
|
||||
export const FILE_EXTENSION_SEPARATOR = '.';
|
||||
|
||||
// Matches the `text:` prefix that file-type identifiers use to denote a
|
||||
// plain-text language (e.g. `text:typescript`). Used by tool-call renderers
|
||||
// to recover the underlying highlight.js language.
|
||||
|
||||
@@ -49,6 +49,7 @@ export * from './sse';
|
||||
export * from './precision';
|
||||
export * from './processing-info';
|
||||
export * from './pwa';
|
||||
export * from './read-media';
|
||||
export * from './routes';
|
||||
export * from './sandbox';
|
||||
export * from './settings-keys';
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { MimeTypeImage } from '$lib/enums';
|
||||
import { MimeTypeAudio, MimeTypeImage } from '$lib/enums';
|
||||
|
||||
// File extension patterns for resource type detection
|
||||
export const IMAGE_FILE_EXTENSION_REGEX = /\.(png|jpg|jpeg|gif|svg|webp)$/i;
|
||||
@@ -27,6 +27,9 @@ export const MCP_RESOURCE_ATTACHMENT_ID_PREFIX = 'res';
|
||||
// Default file extension for unknown image types
|
||||
export const DEFAULT_IMAGE_EXTENSION = 'img';
|
||||
|
||||
// Default file extension for unknown audio types
|
||||
export const DEFAULT_AUDIO_EXTENSION = 'mp3';
|
||||
|
||||
// Default filename for resource content downloads
|
||||
export const DEFAULT_RESOURCE_FILENAME = 'resource.txt';
|
||||
|
||||
@@ -53,3 +56,18 @@ export const IMAGE_MIME_TO_EXTENSION: Record<string, string> = {
|
||||
[MimeTypeImage.PNG]: 'png',
|
||||
[MimeTypeImage.WEBP]: 'webp'
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Mapping from audio MIME types to file extensions.
|
||||
* Used for generating attachment filenames from MIME types.
|
||||
*/
|
||||
export const AUDIO_MIME_TO_EXTENSION: Record<string, string> = {
|
||||
[MimeTypeAudio.MP3]: 'mp3',
|
||||
[MimeTypeAudio.MP3_MPEG]: 'mp3',
|
||||
[MimeTypeAudio.VND_WAVE]: 'wav',
|
||||
[MimeTypeAudio.WAV]: 'wav',
|
||||
[MimeTypeAudio.WAVE]: 'wav',
|
||||
[MimeTypeAudio.X_PN_WAV]: 'wav',
|
||||
[MimeTypeAudio.X_WAV]: 'wav',
|
||||
[MimeTypeAudio.X_WAVE]: 'wav'
|
||||
} as const;
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
import {
|
||||
BuiltInTool,
|
||||
JsonSchemaType,
|
||||
MimeTypeAudio,
|
||||
MimeTypeImage,
|
||||
ToolCallType
|
||||
} from '$lib/enums';
|
||||
import type { OpenAIToolDefinition } from '$lib/types';
|
||||
|
||||
export const READ_MEDIA_TOOL_NAME = BuiltInTool.READ_MEDIA;
|
||||
|
||||
// header lines of the tool result, parsed back by the read_media renderer
|
||||
export const PREFIX_FILE = 'File: ';
|
||||
export const PREFIX_SIZE = 'Size: ';
|
||||
export const PREFIX_MIME = 'MIME: ';
|
||||
|
||||
/** Byte count of the `Size: ` header line, e.g. `Size: 12345 bytes` -> capture group 1 is `12345`. */
|
||||
export const READ_MEDIA_SIZE_REGEX = new RegExp(`^${PREFIX_SIZE}\\s*(\\d+)\\s*bytes`);
|
||||
|
||||
/** Image extensions the tool accepts. The server decodes images with stb_image, which has no webp or tiff. */
|
||||
export const READ_MEDIA_IMAGE_MIME: Record<string, string> = {
|
||||
gif: MimeTypeImage.GIF,
|
||||
jpeg: MimeTypeImage.JPEG,
|
||||
jpg: MimeTypeImage.JPEG,
|
||||
png: MimeTypeImage.PNG
|
||||
} as const;
|
||||
|
||||
/** Audio extensions the tool accepts. The `input_audio` API only takes wav and mp3. */
|
||||
export const READ_MEDIA_AUDIO_MIME: Record<string, string> = {
|
||||
mp3: MimeTypeAudio.MP3_MPEG,
|
||||
wav: MimeTypeAudio.WAV
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Build the read_media tool definition for the modalities the active model has.
|
||||
* At least one of the two flags must be true, otherwise the tool is not offered
|
||||
* at all - a model that cannot see or hear has nothing to do with the bytes.
|
||||
*/
|
||||
export function buildReadMediaToolDefinition(
|
||||
supportsVision: boolean,
|
||||
supportsAudio: boolean
|
||||
): OpenAIToolDefinition {
|
||||
const kinds: string[] = [];
|
||||
|
||||
if (supportsVision) kinds.push(`images (${Object.keys(READ_MEDIA_IMAGE_MIME).join(', ')})`);
|
||||
|
||||
if (supportsAudio) kinds.push(`audio (${Object.keys(READ_MEDIA_AUDIO_MIME).join(', ')})`);
|
||||
|
||||
return {
|
||||
function: {
|
||||
description: `Read a media file and attach it to the conversation so it can be perceived directly. Supports ${kinds.join(' and ')}.`,
|
||||
name: READ_MEDIA_TOOL_NAME,
|
||||
parameters: {
|
||||
properties: {
|
||||
path: {
|
||||
description: 'Path to the media file',
|
||||
type: JsonSchemaType.STRING
|
||||
}
|
||||
},
|
||||
required: ['path'],
|
||||
type: JsonSchemaType.OBJECT
|
||||
}
|
||||
},
|
||||
type: ToolCallType.FUNCTION
|
||||
};
|
||||
}
|
||||
@@ -3,6 +3,12 @@ import { ToolSource } from '$lib/enums/tools.enums';
|
||||
/** HTTP header carrying the working directory a tool call runs in. The server resolves relative paths against it; the model cannot override it. */
|
||||
export const X_TOOL_CWD_HEADER = 'x-tool-cwd';
|
||||
|
||||
/** HTTP header asking the server to encode a tool's output differently, e.g. read_file returning base64. Not a tool parameter, so it stays out of the definition the model sees. */
|
||||
export const X_RESP_TYPE_HEADER = 'x-resp-type';
|
||||
|
||||
/** `X_RESP_TYPE_HEADER` value that makes read_file return the raw bytes as base64 instead of text. */
|
||||
export const RESP_TYPE_BASE64 = 'base64';
|
||||
|
||||
export const TOOL_GROUP_LABELS = {
|
||||
[ToolSource.BUILTIN]: 'Built-in',
|
||||
[ToolSource.CUSTOM]: 'JSON Schema',
|
||||
|
||||
@@ -163,6 +163,7 @@ export enum FileExtensionText {
|
||||
// MIME type prefixes and includes for content detection
|
||||
export enum MimeTypePrefix {
|
||||
IMAGE = 'image/',
|
||||
AUDIO = 'audio/',
|
||||
TEXT = 'text'
|
||||
}
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ export enum GlobSearchType {
|
||||
*/
|
||||
export enum BuiltInTool {
|
||||
READ_FILE = 'read_file',
|
||||
READ_MEDIA = 'read_media',
|
||||
EDIT_FILE = 'edit_file',
|
||||
WRITE_FILE = 'write_file',
|
||||
GET_DATETIME = 'get_datetime',
|
||||
|
||||
@@ -112,7 +112,7 @@ export function useContextGauge(): UseContextGaugeReturn {
|
||||
return chatStore.getConversationModel(activeMessages() as DatabaseMessage[]);
|
||||
});
|
||||
const isActiveModelLoaded = $derived(
|
||||
activeModelId !== null && modelsStore.isModelLoaded(activeModelId)
|
||||
activeModelId !== null && (!isRouterMode() || modelsStore.isModelLoaded(activeModelId))
|
||||
);
|
||||
const isActiveModelLoading = $derived(
|
||||
activeModelId !== null && modelsStore.isModelOperationInProgress(activeModelId)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { settingsStore } from '../stores/settings.svelte';
|
||||
import { getAudioInputFormat } from '../utils/audio-format';
|
||||
import { capImageDataURLSize } from '../utils/cap-img-size';
|
||||
import {
|
||||
API_CHAT,
|
||||
@@ -20,18 +21,12 @@ import {
|
||||
import {
|
||||
AttachmentType,
|
||||
ContentPartType,
|
||||
FileTypeAudio,
|
||||
MessageRole,
|
||||
MimeTypeAudio,
|
||||
ReasoningFormat,
|
||||
StreamConnectionState
|
||||
} from '$lib/enums';
|
||||
import { modelsStore } from '$lib/stores/models.svelte';
|
||||
import type {
|
||||
AudioInputFormat,
|
||||
DatabaseMessageExtraMcpPrompt,
|
||||
DatabaseMessageExtraMcpResource
|
||||
} from '$lib/types';
|
||||
import type { DatabaseMessageExtraMcpPrompt, DatabaseMessageExtraMcpResource } from '$lib/types';
|
||||
import type {
|
||||
ApiChatCompletionToolCall,
|
||||
ApiChatMessageContentPart,
|
||||
@@ -43,23 +38,6 @@ import { getAuthHeaders, getJsonHeaders } from '$lib/utils/api-headers';
|
||||
import { formatAttachmentText } from '$lib/utils/formatters';
|
||||
import { streamIdentity } from '$lib/utils/stream-identity';
|
||||
|
||||
function getAudioInputFormat(mimeType: string): AudioInputFormat {
|
||||
const normalizedMimeType = mimeType.trim().toLowerCase();
|
||||
|
||||
if (
|
||||
normalizedMimeType === MimeTypeAudio.WAV ||
|
||||
normalizedMimeType === MimeTypeAudio.WAVE ||
|
||||
normalizedMimeType === MimeTypeAudio.X_WAV ||
|
||||
normalizedMimeType === MimeTypeAudio.X_WAVE ||
|
||||
normalizedMimeType === MimeTypeAudio.VND_WAVE ||
|
||||
normalizedMimeType === MimeTypeAudio.X_PN_WAV
|
||||
) {
|
||||
return FileTypeAudio.WAV;
|
||||
}
|
||||
|
||||
return FileTypeAudio.MP3;
|
||||
}
|
||||
|
||||
interface ResumableStreamState {
|
||||
bytesReceived: number;
|
||||
updatedAt: number;
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
import { ToolsService } from './tools.service';
|
||||
import {
|
||||
FILE_EXTENSION_SEPARATOR,
|
||||
FILE_PATH_SEPARATOR_REGEX,
|
||||
NEWLINE,
|
||||
PREFIX_FILE,
|
||||
PREFIX_MIME,
|
||||
PREFIX_SIZE,
|
||||
READ_MEDIA_AUDIO_MIME,
|
||||
READ_MEDIA_IMAGE_MIME,
|
||||
RESP_TYPE_BASE64
|
||||
} from '$lib/constants';
|
||||
import { BuiltInTool, ToolResponseField } from '$lib/enums';
|
||||
import type { ToolExecutionResult } from '$lib/types';
|
||||
|
||||
/** Modalities of the model the tool call runs for. */
|
||||
export interface ReadMediaCapabilities {
|
||||
audio: boolean;
|
||||
vision: boolean;
|
||||
}
|
||||
|
||||
/** Lowercase extension of a path, without the dot. Empty when the file name has none. */
|
||||
function fileExtension(path: string): string {
|
||||
const name = path.split(FILE_PATH_SEPARATOR_REGEX).pop() ?? '';
|
||||
const dot = name.lastIndexOf(FILE_EXTENSION_SEPARATOR);
|
||||
|
||||
return dot > 0 ? name.slice(dot + 1).toLowerCase() : '';
|
||||
}
|
||||
|
||||
/**
|
||||
* **ReadMediaService** - frontend executor for the `read_media` tool
|
||||
*
|
||||
* The tool is synthetic: no such tool exists on the server. It reads the file
|
||||
* through the built-in `read_file` tool with the `base64` response type, then
|
||||
* turns the bytes into a data URI line. The agentic store lifts that line into
|
||||
* an image or audio attachment on the tool result message, which is what makes
|
||||
* the model perceive the file instead of reading a wall of base64.
|
||||
*
|
||||
* Living in the frontend is what lets it exist only for models that can
|
||||
* actually use the result - the server has no idea which model is selected.
|
||||
*
|
||||
* @see buildReadMediaToolDefinition in constants/read-media.ts - tool schema sent to the LLM
|
||||
* @see agenticStore in stores/agentic.svelte.ts - tool dispatch and attachment extraction
|
||||
*/
|
||||
export class ReadMediaService {
|
||||
static async executeTool(
|
||||
params: Record<string, unknown>,
|
||||
capabilities: ReadMediaCapabilities,
|
||||
signal?: AbortSignal,
|
||||
cwd?: string
|
||||
): Promise<ToolExecutionResult> {
|
||||
const path = typeof params.path === 'string' ? params.path : '';
|
||||
|
||||
if (!path) {
|
||||
return { content: 'Error: missing "path" argument.', isError: true };
|
||||
}
|
||||
|
||||
const extension = fileExtension(path);
|
||||
const imageMime = READ_MEDIA_IMAGE_MIME[extension];
|
||||
const audioMime = READ_MEDIA_AUDIO_MIME[extension];
|
||||
|
||||
let resolvedMime: string | undefined;
|
||||
|
||||
if (imageMime && capabilities.vision) resolvedMime = imageMime;
|
||||
else if (audioMime && capabilities.audio) resolvedMime = audioMime;
|
||||
|
||||
if (!resolvedMime) {
|
||||
const supported = [
|
||||
...(capabilities.vision ? Object.keys(READ_MEDIA_IMAGE_MIME) : []),
|
||||
...(capabilities.audio ? Object.keys(READ_MEDIA_AUDIO_MIME) : [])
|
||||
];
|
||||
// an unreadable-by-this-model file is a dead end, so say why instead of failing silently
|
||||
const reason =
|
||||
imageMime || audioMime
|
||||
? `the current model cannot perceive ".${extension}" files`
|
||||
: `".${extension}" is not a supported media type`;
|
||||
|
||||
return {
|
||||
content: `Error: ${reason}. Supported: ${supported.join(', ')}.`,
|
||||
isError: true
|
||||
};
|
||||
}
|
||||
|
||||
const raw = await ToolsService.executeToolRaw(
|
||||
BuiltInTool.READ_FILE,
|
||||
{ path },
|
||||
signal,
|
||||
cwd,
|
||||
RESP_TYPE_BASE64
|
||||
);
|
||||
|
||||
if (ToolResponseField.ERROR in raw) {
|
||||
return { content: String(raw[ToolResponseField.ERROR]), isError: true };
|
||||
}
|
||||
|
||||
const base64 = typeof raw.base64 === 'string' ? raw.base64 : '';
|
||||
|
||||
if (!base64) {
|
||||
return { content: `Error: no data returned for ${path}.`, isError: true };
|
||||
}
|
||||
|
||||
const sizeBytes = typeof raw.size_bytes === 'number' ? raw.size_bytes : 0;
|
||||
const content = [
|
||||
`${PREFIX_FILE}${path}`,
|
||||
`${PREFIX_SIZE}${sizeBytes} bytes`,
|
||||
`${PREFIX_MIME}${resolvedMime}`,
|
||||
`data:${resolvedMime};base64,${base64}`
|
||||
].join(NEWLINE);
|
||||
|
||||
return { content, isError: false };
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,5 @@
|
||||
import { base } from '$app/paths';
|
||||
import { API_TOOLS, X_TOOL_CWD_HEADER } from '$lib/constants';
|
||||
import { API_TOOLS, X_RESP_TYPE_HEADER, X_TOOL_CWD_HEADER } from '$lib/constants';
|
||||
import { ToolResponseField } from '$lib/enums';
|
||||
import type { ServerBuiltinToolInfo, ToolExecutionResult } from '$lib/types';
|
||||
import { apiFetch } from '$lib/utils';
|
||||
@@ -51,16 +51,26 @@ export class ToolsService {
|
||||
* Execute a built-in tool and return the raw JSON response. Unlike
|
||||
* executeTool, this preserves structured fields (e.g. file_glob_search's
|
||||
* `entries` and `base`) that the flattened ToolExecutionResult drops.
|
||||
*
|
||||
* @param respType - sent as the x-resp-type request header. Only read_file
|
||||
* honors it, with `base64` to get the raw bytes instead of decoded text.
|
||||
*/
|
||||
static async executeToolRaw(
|
||||
toolName: string,
|
||||
params: Record<string, unknown>,
|
||||
signal?: AbortSignal,
|
||||
cwd?: string
|
||||
cwd?: string,
|
||||
respType?: string
|
||||
): Promise<Record<string, unknown>> {
|
||||
const headers: Record<string, string> = {};
|
||||
|
||||
if (cwd) headers[X_TOOL_CWD_HEADER] = cwd;
|
||||
|
||||
if (respType) headers[X_RESP_TYPE_HEADER] = respType;
|
||||
|
||||
return apiFetch<Record<string, unknown>>(API_TOOLS.EXECUTE, {
|
||||
body: JSON.stringify({ params, tool: toolName }),
|
||||
headers: cwd ? { [X_TOOL_CWD_HEADER]: cwd } : undefined,
|
||||
headers: Object.keys(headers).length > 0 ? headers : undefined,
|
||||
method: 'POST',
|
||||
signal
|
||||
});
|
||||
|
||||
@@ -22,7 +22,9 @@
|
||||
|
||||
import { DEFAULT_AGENTIC_CONFIG, NEWLINE } from '$lib/constants';
|
||||
import {
|
||||
AUDIO_MIME_TO_EXTENSION,
|
||||
DATA_URI_BASE64_REGEX,
|
||||
DEFAULT_AUDIO_EXTENSION,
|
||||
DEFAULT_IMAGE_EXTENSION,
|
||||
IMAGE_MIME_TO_EXTENSION,
|
||||
MCP_ATTACHMENT_NAME_PREFIX
|
||||
@@ -36,6 +38,7 @@ import {
|
||||
ToolCallType
|
||||
} from '$lib/enums';
|
||||
import { ChatService } from '$lib/services';
|
||||
import { ReadMediaService } from '$lib/services/read-media.service';
|
||||
import { SandboxService } from '$lib/services/sandbox.service';
|
||||
import { ToolsService } from '$lib/services/tools.service';
|
||||
import { conversationsStore } from '$lib/stores/conversations.svelte';
|
||||
@@ -75,9 +78,10 @@ import type {
|
||||
import type {
|
||||
DatabaseMessage,
|
||||
DatabaseMessageExtra,
|
||||
DatabaseMessageExtraAudioFile,
|
||||
DatabaseMessageExtraImageFile
|
||||
} from '$lib/types/database';
|
||||
import { isAbortError } from '$lib/utils';
|
||||
import { getAudioInputFormat, isAbortError } from '$lib/utils';
|
||||
import { SvelteMap } from 'svelte/reactivity';
|
||||
|
||||
function createDefaultSession(): AgenticSession {
|
||||
@@ -900,7 +904,18 @@ class AgenticStore {
|
||||
if (executionResult.isError) toolSuccess = false;
|
||||
} else if (toolSource === ToolSource.FRONTEND) {
|
||||
const args = this.parseToolArguments(toolCall.function.arguments);
|
||||
const executionResult = await SandboxService.executeTool(toolName, args, signal);
|
||||
const executionResult =
|
||||
toolName === BuiltInTool.READ_MEDIA
|
||||
? await ReadMediaService.executeTool(
|
||||
args,
|
||||
{
|
||||
audio: modelsStore.modelSupportsAudio(effectiveModel),
|
||||
vision: modelsStore.modelSupportsVision(effectiveModel)
|
||||
},
|
||||
signal,
|
||||
conversationsStore.activeConversation?.cwd
|
||||
)
|
||||
: await SandboxService.executeTool(toolName, args, signal);
|
||||
|
||||
result = executionResult.content;
|
||||
|
||||
@@ -990,7 +1005,19 @@ class AgenticStore {
|
||||
];
|
||||
|
||||
for (const attachment of attachments) {
|
||||
if (attachment.type === AttachmentType.IMAGE) {
|
||||
if (attachment.type === AttachmentType.AUDIO) {
|
||||
if (modelsStore.modelSupportsAudio(effectiveModel)) {
|
||||
contentParts.push({
|
||||
input_audio: {
|
||||
data: (attachment as DatabaseMessageExtraAudioFile).base64Data,
|
||||
format: getAudioInputFormat(
|
||||
(attachment as DatabaseMessageExtraAudioFile).mimeType
|
||||
)
|
||||
},
|
||||
type: ContentPartType.INPUT_AUDIO
|
||||
});
|
||||
}
|
||||
} else if (attachment.type === AttachmentType.IMAGE) {
|
||||
if (modelsStore.modelSupportsVision(effectiveModel)) {
|
||||
contentParts.push({
|
||||
image_url: {
|
||||
@@ -1101,6 +1128,18 @@ class AgenticStore {
|
||||
return `[Attachment saved: ${name}]`;
|
||||
}
|
||||
|
||||
if (mimeType.startsWith(MimeTypePrefix.AUDIO)) {
|
||||
// audio extras hold the bare base64, the input_audio part has no room for a data URI
|
||||
attachments.push({
|
||||
base64Data,
|
||||
mimeType,
|
||||
name,
|
||||
type: AttachmentType.AUDIO
|
||||
});
|
||||
|
||||
return `[Attachment saved: ${name}]`;
|
||||
}
|
||||
|
||||
return line;
|
||||
});
|
||||
|
||||
@@ -1108,7 +1147,9 @@ class AgenticStore {
|
||||
}
|
||||
|
||||
private buildAttachmentName(mimeType: string, index: number): string {
|
||||
const extension = IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION;
|
||||
const extension = mimeType.startsWith(MimeTypePrefix.AUDIO)
|
||||
? (AUDIO_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_AUDIO_EXTENSION)
|
||||
: (IMAGE_MIME_TO_EXTENSION[mimeType] ?? DEFAULT_IMAGE_EXTENSION);
|
||||
|
||||
return `${MCP_ATTACHMENT_NAME_PREFIX}-${Date.now()}-${index}.${extension}`;
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ import {
|
||||
let closeTimer: ReturnType<typeof setTimeout> | undefined;
|
||||
let lastPointerType = '';
|
||||
|
||||
export const gaugePopup = $state({ bottom: 0, centerX: 0, open: false });
|
||||
export const gaugePopup = $state({ bottom: 0, centerX: 0, detailsOpen: false, open: false });
|
||||
|
||||
function openFrom(trigger: HTMLElement): void {
|
||||
clearTimeout(closeTimer);
|
||||
|
||||
@@ -18,7 +18,10 @@ import { ModelsService } from '$lib/services/models.service';
|
||||
import { PropsService } from '$lib/services/props.service';
|
||||
import { conversationsStore } from '$lib/stores/conversations.svelte';
|
||||
import { isRouterMode, serverStore } from '$lib/stores/server.svelte';
|
||||
import { getAuthHeaders, TTLCache } from '$lib/utils';
|
||||
// deep imports, not the '$lib/utils' barrel: it re-exports modules that reach back
|
||||
// into the stores, and going through it here would read a half-built module
|
||||
import { getAuthHeaders } from '$lib/utils/api-headers';
|
||||
import { TTLCache } from '$lib/utils/cache-ttl';
|
||||
import {
|
||||
detectThinkingSupport,
|
||||
detectThinkingSupportWithReason
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import {
|
||||
buildReadMediaToolDefinition,
|
||||
buildSandboxToolDefinition,
|
||||
DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY,
|
||||
HOME_TILDE,
|
||||
@@ -15,6 +16,7 @@ import {
|
||||
} from '$lib/enums';
|
||||
import { ToolsService } from '$lib/services/tools.service';
|
||||
import { mcpStore } from '$lib/stores/mcp.svelte';
|
||||
import { modelsStore, selectedModelName } from '$lib/stores/models.svelte';
|
||||
import { config } from '$lib/stores/settings.svelte';
|
||||
import type { OpenAIToolDefinition, ToolEntry, ToolGroup } from '$lib/types';
|
||||
import { SvelteMap, SvelteSet } from 'svelte/reactivity';
|
||||
@@ -168,9 +170,42 @@ class ToolsStore {
|
||||
}
|
||||
|
||||
get frontendTools(): OpenAIToolDefinition[] {
|
||||
return config().jsSandboxEnabled
|
||||
? [buildSandboxToolDefinition(!!config().symbolicMathEnabled)]
|
||||
: [];
|
||||
const tools: OpenAIToolDefinition[] = [];
|
||||
|
||||
if (config().jsSandboxEnabled) {
|
||||
tools.push(buildSandboxToolDefinition(!!config().symbolicMathEnabled));
|
||||
}
|
||||
|
||||
const readMedia = this.readMediaTool();
|
||||
|
||||
if (readMedia) tools.push(readMedia);
|
||||
|
||||
return tools;
|
||||
}
|
||||
|
||||
/**
|
||||
* `read_media` runs in the frontend on top of the server's `read_file`, so it
|
||||
* exists only when that tool is served and the active model can perceive the
|
||||
* bytes. The server cannot make this call - it does not know which model the
|
||||
* conversation uses.
|
||||
*/
|
||||
private readMediaTool(): OpenAIToolDefinition | null {
|
||||
const hasReadFile = this._builtinTools.some(
|
||||
(def) => def.function.name === BuiltInTool.READ_FILE
|
||||
);
|
||||
|
||||
if (!hasReadFile) return null;
|
||||
|
||||
const model = selectedModelName() ?? modelsStore.models[0]?.model ?? '';
|
||||
|
||||
if (!model) return null;
|
||||
|
||||
const vision = modelsStore.modelSupportsVision(model);
|
||||
const audio = modelsStore.modelSupportsAudio(model);
|
||||
|
||||
if (!vision && !audio) return null;
|
||||
|
||||
return buildReadMediaToolDefinition(vision, audio);
|
||||
}
|
||||
|
||||
get customTools(): OpenAIToolDefinition[] {
|
||||
|
||||
@@ -50,11 +50,11 @@ export interface AgenticSection {
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents a tool result line that may reference an image attachment
|
||||
* Represents a tool result line that may reference a media attachment (image or audio)
|
||||
*/
|
||||
export type ToolResultLine = {
|
||||
text: string;
|
||||
image?: DatabaseMessageExtraImageFile;
|
||||
media?: DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile;
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -301,16 +301,16 @@ export function splitSearchSummaryList(
|
||||
return { lines };
|
||||
}
|
||||
|
||||
/** Bounded cache for parseToolResultWithImages results. */
|
||||
/** Bounded cache for parseToolResultWithMedia results. */
|
||||
const TOOL_RESULT_LINES_CACHE_MAX_SIZE = 32;
|
||||
const toolResultLinesCache = new Map<string, ToolResultLine[]>();
|
||||
|
||||
/**
|
||||
* Parse tool result text into lines, matching image attachments by name.
|
||||
* Parse tool result text into lines, matching media attachments (images and audio) by name.
|
||||
* Memoized: called per render during streaming on unchanged tool result
|
||||
* strings with unchanged extras.
|
||||
*/
|
||||
export function parseToolResultWithImages(
|
||||
export function parseToolResultWithMedia(
|
||||
toolResult: string,
|
||||
extras?: DatabaseMessageExtra[]
|
||||
): ToolResultLine[] {
|
||||
@@ -332,12 +332,13 @@ export function parseToolResultWithImages(
|
||||
if (!match || !extras) return { text: line };
|
||||
|
||||
const attachmentName = match[1];
|
||||
const image = extras.find(
|
||||
(e): e is DatabaseMessageExtraImageFile =>
|
||||
e.type === AttachmentType.IMAGE && e.name === attachmentName
|
||||
const media = extras.find(
|
||||
(e): e is DatabaseMessageExtraImageFile | DatabaseMessageExtraAudioFile =>
|
||||
(e.type === AttachmentType.IMAGE || e.type === AttachmentType.AUDIO) &&
|
||||
e.name === attachmentName
|
||||
);
|
||||
|
||||
return { image, text: line };
|
||||
return { media, text: line };
|
||||
});
|
||||
|
||||
if (toolResultLinesCache.size >= TOOL_RESULT_LINES_CACHE_MAX_SIZE) {
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
import { FileTypeAudio, MimeTypeAudio } from '$lib/enums';
|
||||
import type { AudioInputFormat } from '$lib/types/api';
|
||||
|
||||
/**
|
||||
* Map a MIME type to the AudioInputFormat expected by the API.
|
||||
*/
|
||||
export function getAudioInputFormat(mimeType: string): AudioInputFormat {
|
||||
const normalizedMimeType = mimeType.trim().toLowerCase();
|
||||
|
||||
if (
|
||||
normalizedMimeType === MimeTypeAudio.WAV ||
|
||||
normalizedMimeType === MimeTypeAudio.WAVE ||
|
||||
normalizedMimeType === MimeTypeAudio.X_WAV ||
|
||||
normalizedMimeType === MimeTypeAudio.X_WAVE ||
|
||||
normalizedMimeType === MimeTypeAudio.VND_WAVE ||
|
||||
normalizedMimeType === MimeTypeAudio.X_PN_WAV
|
||||
) {
|
||||
return FileTypeAudio.WAV;
|
||||
}
|
||||
|
||||
return FileTypeAudio.MP3;
|
||||
}
|
||||
@@ -248,7 +248,7 @@ export {
|
||||
export {
|
||||
deriveAgenticSections,
|
||||
buildAssistantRawOutput,
|
||||
parseToolResultWithImages,
|
||||
parseToolResultWithMedia,
|
||||
splitSearchSummaryList,
|
||||
hasAgenticContent,
|
||||
classifyToolResult,
|
||||
@@ -325,3 +325,6 @@ export { uuid } from './uuid';
|
||||
|
||||
// CSS utilities
|
||||
export { remToPx } from './css';
|
||||
|
||||
// Audio format helper (used by agentic store and chat service)
|
||||
export { getAudioInputFormat } from './audio-format';
|
||||
|
||||
@@ -33,12 +33,12 @@ npx vitest --project=client --run tests/client/agentic-stream.perf.svelte.test.t
|
||||
|
||||
The point of the harness is the _scaling curve_, not any single number.
|
||||
|
||||
| Knob | Reads on |
|
||||
| --------------------------- | ---------------------------------------------------------------------------------------------------- |
|
||||
| `priorToolCalls` (0/1/5/20) | the reactive fan-out. Flat => no fan-out. Linear => confirmed. |
|
||||
| `toolResultBytes` | whole-blob string scans (`extractSearchResults`, `parseToolResultWithImages`, `classifyToolResult`). |
|
||||
| `editFileEdits` | `computeLineDiff`, the O(m\*n) LCS. |
|
||||
| `openCodeFence` | `hljs.highlightAuto` on partial code. |
|
||||
| Knob | Reads on |
|
||||
| --------------------------- | --------------------------------------------------------------------------------------------------- |
|
||||
| `priorToolCalls` (0/1/5/20) | the reactive fan-out. Flat => no fan-out. Linear => confirmed. |
|
||||
| `toolResultBytes` | whole-blob string scans (`extractSearchResults`, `parseToolResultWithMedia`, `classifyToolResult`). |
|
||||
| `editFileEdits` | `computeLineDiff`, the O(m\*n) LCS. |
|
||||
| `openCodeFence` | `hljs.highlightAuto` on partial code. |
|
||||
|
||||
Deliberately no hard assertions: CI timing is noisy and the value here is the
|
||||
before/after delta, not a gate.
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
//
|
||||
// Run: npx vitest bench --project=unit tests/unit/agentic-hotpath.bench.ts
|
||||
|
||||
import { classifyToolResult, parseToolResultWithImages } from '$lib/utils/agentic';
|
||||
import { classifyToolResult, parseToolResultWithMedia } from '$lib/utils/agentic';
|
||||
import { detectIncompleteCodeBlock, highlightCode } from '$lib/utils/code';
|
||||
import { computeLineDiff } from '$lib/utils/compute-line-diff';
|
||||
import { preprocessLaTeX } from '$lib/utils/latex-protection';
|
||||
@@ -200,17 +200,17 @@ describe('exit-code regex', () => {
|
||||
|
||||
// --- per-line result parsers ----------------------------------------------
|
||||
|
||||
describe('parseToolResultWithImages', () => {
|
||||
describe('parseToolResultWithMedia', () => {
|
||||
bench('1KB', () => {
|
||||
parseToolResultWithImages(SHELL_OUTPUT_1KB, []);
|
||||
parseToolResultWithMedia(SHELL_OUTPUT_1KB, []);
|
||||
});
|
||||
|
||||
bench('200KB', () => {
|
||||
parseToolResultWithImages(SHELL_OUTPUT_200KB, []);
|
||||
parseToolResultWithMedia(SHELL_OUTPUT_200KB, []);
|
||||
});
|
||||
|
||||
bench('2MB', () => {
|
||||
parseToolResultWithImages(SHELL_OUTPUT_2MB, []);
|
||||
parseToolResultWithMedia(SHELL_OUTPUT_2MB, []);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user