diff --git a/.gitignore b/.gitignore index 8ae4ae3b23..033ba3233b 100644 --- a/.gitignore +++ b/.gitignore @@ -30,6 +30,8 @@ examples/python/hf_cache examples/csharp/ModelChat/models !test/models/gemma4/* +!/test/models/lfm2-audio/ +!test/models/lfm2-audio/* !test/models/hf-internal-testing/* !test/models/lfm2-vl/* !test/models/marian-batch/* diff --git a/README.md b/README.md index a4b7251a1f..478b565734 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ See documentation at the [ONNX Runtime website](https://onnxruntime.ai/docs/gena | Support matrix | Supported now | Under development | On the roadmap| | -------------- | ------------- | ----------------- | -------------- | -| Model architectures | AMD OLMo
ChatGLM
DeepSeek
ERNIE 4.5
Fara
Gemma
gpt-oss
Granite
Granite MoE Hybrid
HunYuan Dense V1
InternLM2
LFM2 (language + vision)
Llama
Mistral
Nemotron
Phi (language + vision)
Qwen (language + vision)
SmolLM3
Whisper | Stable diffusion | Multi-modal models | +| Model architectures | AMD OLMo
ChatGLM
DeepSeek
ERNIE 4.5
Fara
Gemma
gpt-oss
Granite
Granite MoE Hybrid
HunYuan Dense V1
InternLM2
LFM2 (language + vision + speech input)
Llama
Mistral
Nemotron
Phi (language + vision)
Qwen (language + vision)
SmolLM3
Whisper | Stable diffusion | Multi-modal models | | API | Python
C#
C/C++
Java ^ | Objective-C || | O/S | Linux
Windows
Mac
Android || iOS ||| | Architecture | x86
x64
arm64 |||| diff --git a/examples/python/common.py b/examples/python/common.py index bf23a6914e..a5f408a0a6 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -322,6 +322,11 @@ def get_user_content(model_type: str, num_images: int, num_audios: int, prompt: # per projected vision feature, and <|image_end|>, based on the resized image's patch grid. image_tags = "".join(["" for _ in range(num_images)]) content = image_tags + prompt + elif model_type == "lfm2_audio": + # LFM2-Audio: the C++ audio processor replaces each <|audio|> with one placeholder per + # encoder frame (one per 80 ms of audio) and splices the encoder output in at those positions. + audio_tags = "".join(["<|audio|>" for _ in range(num_audios)]) + content = audio_tags + prompt elif model_type == "mistral3": # Pixtral / Ministral-3 VLM: the C++ image processor expands each # [IMG] into the full token sequence based on image resolution. diff --git a/examples/python/lfm2-audio-tutorial.md b/examples/python/lfm2-audio-tutorial.md new file mode 100644 index 0000000000..9bec5a9347 --- /dev/null +++ b/examples/python/lfm2-audio-tutorial.md @@ -0,0 +1,561 @@ +# Build your LFM2-Audio / LFM2.5-Audio ONNX models for ONNX Runtime GenAI + +LFM2-Audio pairs a FastConformer speech encoder (based on +[nvidia/canary-180m-flash](https://huggingface.co/nvidia/canary-180m-flash)) with an LFM2 decoder +(the hybrid conv/attention model described in +[the LFM2 support PR](https://github.com/microsoft/onnxruntime-genai/pull/1979)). ONNX Runtime GenAI +runs the speech-to-text half of it — ASR and spoken-prompt chat — as three ONNX models: + +| Model | Inputs | Outputs | +| --- | --- | --- | +| speech encoder (`audio_encoder.onnx`) | `mel_spectrogram`, `mel_lengths` | `audio_embeddings`, `audio_lengths` | +| `embeddings.onnx` | `input_ids`, `audio_features` | `inputs_embeds` | +| `model.onnx` (decoder) | `inputs_embeds`, `attention_mask`, KV + conv cache | `logits` | + +Only the decoder is produced by the model builder in this repository; the speech encoder is +published as ONNX by LiquidAI, and the embedding model is a small graph you build once. + +The model can also *speak*. Two more graphs, both published by LiquidAI, give it a voice here: a +depthformer that turns a decoder hidden state into a frame of audio codes, and the audio embedding that +feeds each frame back to the decoder. [Speech output](#6-speech-output) adds them; without them the +model answers in text only. + +## Steps + +1. [Build the decoder](#1-build-the-decoder) +2. [Get the speech encoder and build the embedding model](#2-get-the-speech-encoder-and-build-the-embedding-model) +3. [Write `genai_config.json`](#3-write-genai_configjson) +4. [Choose the precisions](#4-choose-the-precisions) +5. [Run the model](#5-run-the-model) +6. [Speech output](#6-speech-output) +7. [Known limitations](#7-known-limitations) + +## 1. Build the decoder + +```bash +# Download the PyTorch model +$ huggingface-cli download LiquidAI/LFM2.5-Audio-1.5B --local-dir ./lfm2.5-audio/pytorch + +# Build the decoder as INT4 with FP32 inputs/outputs for CPU +$ python3 -m onnxruntime_genai.models.builder \ + -i ./lfm2.5-audio/pytorch \ + -o ./lfm2.5-audio/cpu \ + -p int4 \ + -e cpu \ + --extra_options exclude_embeds=true +``` + +`-m LiquidAI/LFM2.5-Audio-1.5B` instead of `-i` works too, and downloads the checkpoint itself. + +The three checkpoints LiquidAI publishes as PyTorch — `LFM2-Audio-1.5B`, `LFM2.5-Audio-1.5B` and +`LFM2.5-Audio-1.5B-JP` — build the same way. Their `"lfm"`, `"encoder"` and `"preprocessor"` settings +are identical, so the front-end defaults below fit all three; only the weights differ. The GGUF +repositories are for llama.cpp and are not inputs to this builder. + +The checkpoint has no `model_type` and no transformers model class: its `config.json` nests the LFM2 +decoder config under `"lfm"`, next to the speech encoder, depthformer and mel front-end settings, and +the checkpoint stores the decoder under the `lfm.` prefix. The builder recognizes +`Lfm2AudioForConditionalGeneration`, reads that nested config and loads only the decoder, whose +logits are tied to the token embeddings. + +`exclude_embeds=true` is what makes this a speech pipeline stage: the decoder then takes +`inputs_embeds` instead of `input_ids`, so the embedding model can splice the encoder output into the +token embeddings before the decoder runs. The builder writes `"type": "lfm2_audio"` into +`genai_config.json` for this case, and `"type": "lfm2_audio_text"` when the flag is omitted — that +second form is a plain text-only LFM2 model that happens to come from an audio checkpoint, and it +runs through the normal `AppendTokens` path with no speech or embedding model. + +## 2. Get the speech encoder and build the embedding model + +The speech encoder is `onnx/audio_encoder.onnx` in +[LiquidAI/LFM2.5-Audio-1.5B-ONNX](https://huggingface.co/LiquidAI/LFM2.5-Audio-1.5B-ONNX) +(`audio_encoder_fp16.onnx` and `audio_encoder_q4.onnx` are there too). Download the graph together +with every `*.onnx_data*` file next to it. That repository is the export of `LFM2.5-Audio-1.5B`, so +its encoder belongs to that checkpoint only; `LFM2-Audio-1.5B` and `LFM2.5-Audio-1.5B-JP` have their +own encoder weights and need their own export, from `conformer.*` and `audio_adapter.*` in the +checkpoint (`liquid_audio.model.conformer.ConformerEncoder` plus its adapter MLP, traced for one +clip). It has the signature ONNX Runtime GenAI expects, for a +single clip at a time: + +| Name | Shape | Type | +| --- | --- | --- | +| `mel_spectrogram` (input) | `[1, num_frames, 128]` | float | +| `mel_lengths` (input) | `[1]` | int64 | +| `audio_embeddings` (output) | `[1, ceil(num_frames / 8), hidden_size]` | float | +| `audio_lengths` (output) | `[1]` | int64 | + +The encoder subsamples the mel frames by 8 (three stride-2 convolutions) and its adapter projects the +result to the decoder's hidden size, so one output frame covers 80 ms of audio. `mel_lengths` gives +it the clip's real length; the runtime always passes the clip's own frames, so nothing is padded. +The `audio_lengths` output is not read — the runtime derives the same counts itself, before the +encoder runs, so that it can size the prompt — but the graph produces it. + +The published export only handles one clip per run: its subsampling mask is traced for a batch of +one, and a wider `mel_spectrogram` fails inside the graph. The runtime therefore runs it once per +clip. An encoder exported to take a real batch would work here too, one clip at a time. + +The embedding model looks up `input_ids` in the decoder's embedding table and scatters +`audio_features` into the positions holding the audio placeholder token. Build it from +`lfm.embed_tokens.weight` in the checkpoint: + +```python +import json +from pathlib import Path + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper +from safetensors.torch import load_file + +pytorch_dir = Path("./lfm2.5-audio/pytorch") +output_dir = Path("./lfm2.5-audio/cpu") + +config = json.loads((pytorch_dir / "config.json").read_text())["lfm"] +hidden_size, vocab_size = config["hidden_size"], config["vocab_size"] +# Pick any id the text never uses; it must match model.audio_token_id in genai_config.json. +audio_token_id = 133 # <|reserved_123|> + +embed_weight = load_file(pytorch_dir / "model.safetensors")["lfm.embed_tokens.weight"] +table = numpy_helper.from_array(embed_weight.float().numpy(), name="embed_tokens.weight") + +graph = helper.make_graph( + [ + helper.make_node("Gather", ["embed_tokens.weight", "input_ids"], ["text_embeds"], axis=0), + helper.make_node("Shape", ["text_embeds"], ["embeds_shape"]), + helper.make_node("Reshape", ["text_embeds", "flat_rows"], ["flat_embeds"]), + helper.make_node("Reshape", ["input_ids", "flat"], ["flat_ids"]), + helper.make_node("Equal", ["flat_ids", "audio_token_id"], ["is_audio"]), + helper.make_node("NonZero", ["is_audio"], ["audio_positions_t"]), + helper.make_node("Transpose", ["audio_positions_t"], ["audio_positions"], perm=[1, 0]), + helper.make_node("ScatterND", ["flat_embeds", "audio_positions", "audio_features"], ["merged"]), + helper.make_node("Reshape", ["merged", "embeds_shape"], ["inputs_embeds"]), + ], + "embedding", + [ + helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch_size", "sequence_length"]), + helper.make_tensor_value_info("audio_features", TensorProto.FLOAT, ["num_audio_tokens", hidden_size]), + ], + [helper.make_tensor_value_info("inputs_embeds", TensorProto.FLOAT, ["batch_size", "sequence_length", hidden_size])], + initializer=[ + table, + numpy_helper.from_array(np.array([-1, hidden_size], np.int64), name="flat_rows"), + numpy_helper.from_array(np.array([-1], np.int64), name="flat"), + numpy_helper.from_array(np.array(audio_token_id, np.int64), name="audio_token_id"), + ], +) +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)], ir_version=8) +onnx.checker.check_model(model) +onnx.save(model, output_dir / "embeddings.onnx", save_as_external_data=True, location="embeddings.onnx.data") +``` + +## 3. Write `genai_config.json` + +Add `embedding` and `speech` sections and an `audio_token_id` to the `genai_config.json` the builder +produced. Leave the `decoder` and `search` sections it wrote alone: + +```json +{ + "model": { + "audio_token_id": 133, + "bos_token_id": 1, + "context_length": 32768, + "decoder": { "...": "written by the model builder" }, + "embedding": { + "filename": "embeddings.onnx", + "inputs": { + "input_ids": "input_ids", + "audio_features": "audio_features" + }, + "outputs": { + "inputs_embeds": "inputs_embeds" + }, + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [] + } + }, + "eos_token_id": [7, 128, 130], + "pad_token_id": 0, + "speech": { + "filename": "audio_encoder.onnx", + "inputs": { + "audio_embeds": "mel_spectrogram", + "audio_lengths": "mel_lengths", + "audio_sizes": "audio_sizes" + }, + "outputs": { + "audio_features": "audio_embeddings" + }, + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [] + } + }, + "type": "lfm2_audio", + "vocab_size": 65536 + }, + "search": { "...": "written by the model builder" } +} +``` + +Keep the `eos_token_id` list the builder wrote. Besides `<|im_end|>` (7) it holds `<|audio_start|>` +(128) and `<|text_end|>` (130), where the model turns from text to speech; without them generation +runs on into positions meant for the audio head and returns fluent nonsense. See +[the modes](#the-models-modes). + +`filename` is whatever the encoder file is called on disk — keep the name it was published under, +for the reason in [Choose the precisions](#4-choose-the-precisions). + +`audio_sizes` is not an encoder input; the name only tells the runtime what to call the per-clip +token counts it computes, so leave it as is unless it collides with a real input of your graph. +`audio_token_id` must be the id the embedding model scatters over, and must not be an id the +tokenizer can emit for ordinary text — the reserved range of the LFM2 tokenizer is the natural home +for it. + +The audio front end needs no configuration file: it defaults to the settings every published +LFM2-Audio checkpoint shares (`"preprocessor"` in `config.json`), which are NeMo's +`AudioToMelSpectrogramPreprocessor`: + +| Setting | Value | `genai_config.json` override | +| --- | --- | --- | +| Sample rate | 16000 | `model.sample_rate` | +| Mel bins | 128 | `model.num_mels` | +| FFT size | 512 | `model.fft_size` | +| Window | 400 samples (25 ms), symmetric Hann | `model.win_length` | +| Hop | 160 samples (10 ms) | `model.hop_length` | +| Pre-emphasis | 0.97 | `model.preemph` | +| Log guard | 2⁻²⁴ | `model.log_eps` | +| Normalization epsilon | 1e-5 | `model.norm_eps` | +| Encoder subsampling | 8 | `model.subsampling_factor` | + +Set the overrides only for a fine-tune that changed them; the defaults already match all the +published models. + +## 4. Choose the precisions + +`-p int4` quantizes the decoder, and the decoder alone. The embedding model and the speech encoder +are assembled by hand, so they stay at whatever precision you built or downloaded them at — and +together they are bigger than the quantized decoder. Measured on LFM2.5-Audio-1.5B, transcribing +three clips (16 kHz, 44.1 kHz stereo, and an mp3): + +| Decoder | Embedding table | Encoder | Size | Transcripts | +| --- | --- | --- | --- | --- | +| fp32 | fp32 | fp32 | 5736 MB | baseline | +| int4 | fp32 | fp32 | 1787 MB | correct; one comma differs from fp32 | +| int4 | fp16 | fp32 | 1518 MB | identical to the row above | +| int4 | fp16 | fp16 | 1278 MB | identical to the row above | +| **int4** | **fp16** | **q4** | **1178 MB** | **identical to the row above** | + +Quantizing the other two components is free here: every int4 row produced the same tokens as every +other, so the last row is the one to build. The only difference anywhere is the int4 decoder against +fp32 — `And so, my fellow Americans` became `And so my fellow Americans` on one clip, with the words +otherwise unchanged. + +For the embedding table, `astype(np.float16)` on the weight in the snippet above is the whole change; +keep the graph's output fp32 by casting after the `Gather`, so the decoder still receives what it +declares. A 4-bit table would be smaller still — the builder's own text-only export gets one, through +`GatherBlockQuantized` over the quantized `lm_head` weights — but `exclude_embeds` is what puts the +table outside the decoder in the first place, so the pipeline cannot share that copy. + +For the encoder, take `audio_encoder_q4.onnx` instead of `audio_encoder.onnx`. **Keep the published +file names.** Each variant's graph refers to its own `*.onnx_data` by name, so renaming the pair to +`speech.onnx` breaks the lookup with `External data path validation failed`. Leave both files as they +are and point `model.speech.filename` at the name you downloaded. + +## 5. Run the model + +[`model-mm.py`](model-mm.py) drives any multi-modal model in this repository: + +```bash +$ python3 model-mm.py -m ./lfm2.5-audio/cpu -e cpu --audio_paths ./question.wav +``` + +The prompt must hold exactly one `<|audio|>` marker per clip, in clip order; the processor rejects +any other count rather than guessing, and it accepts one prompt at a time (a single-entry list is +fine, batching is not). The C++ audio processor replaces each marker with one `audio_token_id` per +encoder frame, so the number of placeholders always matches the number of features the encoder +produced. Text on either side of a marker is tokenized on its own, which is what +`liquid_audio.ChatState` does. + +LFM2-Audio has no audio token of its own and no chat-template entry for audio: `<|audio|>` is this +runtime's marker, the same role `` plays for LFM2-VL. + +``` +<|startoftext|><|im_start|>system +Perform ASR.<|im_end|> +<|im_start|>user +<|audio|><|im_end|> +<|im_start|>assistant +``` + +### The model's modes + +The system prompt picks the task, and the reference implementation pairs each with a generation +routine of its own: + +| Mode | System prompt | Reference routine | Output | Text-only build | With [speech output](#6-speech-output) | +| --- | --- | --- | --- | --- | --- | +| Chat | none | `generate_sequential` | text | works | works | +| ASR | `Perform ASR.` | `generate_sequential` | text | works | works | +| TTS | `Perform TTS. Use the UK male voice.` (also US male, US female, UK female) | `generate_sequential` | speech | stops immediately | works | +| Interleaved | `Respond with interleaved text and audio.` | `generate_interleaved` | text and speech, alternating | do not use | works, with `audio_interleaved=True` | + +**ASR** and **chat** need nothing more than the three models above. A spoken or typed question with +no system prompt gets a text answer, and both match the reference step for step: on the reference's +own `question.wav` and three typed questions the text logits agree to within 6e-4 and every token is +the same. + +`LFM2.5-Audio-1.5B-JP` takes its own prompts, `Perform ASR in japanese.` and +`Perform TTS in japanese.`; the rest of this section applies to it unchanged. + +Two tokens mark the turn passing from text to speech: `<|audio_start|>` when the rest of the answer is +spoken, and `<|text_end|>` when the text half of an interleaved answer is done. In a text-only build +the builder puts both in `eos_token_id`, because the positions after either are audio codes meant for +the depthformer, and read off the text head they decode to fluent, plausible nonsense. + +**TTS** in a text-only build produces nothing: the whole text stream the model emits is +`<|audio_start|>` followed by `<|im_end|>`, everything in between being audio. Generation stops on +that first token, and the stop token is not added to the sequence, so there is nothing after the +prompt to decode: expect an empty string. + +**Interleaved** cannot be used in a text-only build. The reference writes six text tokens, then +`interleaved_n_audio` audio frames (12, or 9 for the JP checkpoint) whose embeddings go back into the +context, then six more text tokens, and so on. Without audio frames to feed, the text head is read at +positions the model means for speech: the first six tokens are right and the rest is mostly the +non-breaking-space token. For a text answer, leave the system prompt out and use chat mode. + +Sampling: the reference generates text greedily in every mode, so `do_sample=False` is right here. +The temperatures and `top_k` values quoted for the model apply to the audio codes only; see +[Speech output](#6-speech-output) for where they go. + +## 6. Speech output + +Speech needs the decoder's hidden states, two more graphs and a few lines of `genai_config.json`. + +**Build the decoder with its hidden states.** The depthformer reads the hidden state the logits are +projected from, so add `include_hidden_states=true`: + +```bash +$ python3 -m onnxruntime_genai.models.builder \ + -i ./lfm2.5-audio/pytorch -o ./lfm2.5-audio/cpu -p int4 -e cpu \ + --extra_options exclude_embeds=true include_hidden_states=true +``` + +**Get the two graphs** from +[LiquidAI/LFM2.5-Audio-1.5B-ONNX](https://huggingface.co/LiquidAI/LFM2.5-Audio-1.5B-ONNX): +`onnx/vocoder_depthformer.onnx` and `onnx/audio_embedding.onnx` (with their `*.onnx_data` files for +the `fp16` and `q4` variants), and `onnx/audio_detokenizer.onnx` for turning the codes into sound +afterwards. Like the encoder they belong to `LFM2.5-Audio-1.5B`; the other checkpoints need their own +export, which LiquidAI's [onnx-export](https://github.com/Liquid4All/onnx-export) produces. + +**Fix the depthformer's rotary embedding.** As published, the graph's attention nodes carry rotary +tables but do not apply them (`do_rotary` is unset), and the tables are computed for `theta` 10000 +where the model uses 1,000,000. The first codebook of each frame sits at position zero and is +unaffected; the other seven come out with logits up to 2.5 away from the PyTorch model's. The speech +is intelligible either way, since the first codebook carries most of what is said, but only the +corrected graph gives the model's own codes: with both fixes every codebook agrees to 1e-5. + +```python +import numpy as np +import onnx +from onnx import helper, numpy_helper + +model = onnx.load("vocoder_depthformer.onnx") +for node in model.graph.node: + if node.op_type == "GroupQueryAttention" and not any(a.name == "do_rotary" for a in node.attribute): + node.attribute.append(helper.make_attribute("do_rotary", 1)) +for table in model.graph.initializer: + if table.name.endswith((".gqa_cos", ".gqa_sin")): + positions, half = numpy_helper.to_array(table).shape + angles = np.outer(np.arange(positions), 1.0 / 1_000_000 ** (np.arange(0, 2 * half, 2) / (2 * half))) + values = np.cos(angles) if table.name.endswith("cos") else np.sin(angles) + table.CopyFrom(numpy_helper.from_array(values.astype(np.float32), table.name)) +onnx.save(model, "vocoder_depthformer.onnx") +``` + +The same defect is in all four published depthformers (`fp16`, `q4` and `q8` too), and the same fix +applies to each. Which variants are worth taking, measured against the PyTorch model: + +| Graph | fp32 | fp16 | q8 | q4 | +| --- | --- | --- | --- | --- | +| `vocoder_depthformer` (fixed), largest logit difference | 9e-6 | 7e-3 | not measured | 2.0, and 37 of 48 codes agree | +| `audio_detokenizer`, waveform SNR on real speech | 112 dB | 51 dB | 18 dB | 3 dB | +| `audio_embedding`, against the checkpoint | exact | 6e-8 | exact | exact | + +Take `fp32` or `fp16`. The `q4` depthformer loses more to quantization than the rotary fix gains, and +the `q4` detokenizer's waveform is barely correlated with the model's, though the speech is still +understood: all four detokenizers read back word for word through ASR. The `q4` and `q8` audio +embeddings are the `fp32` table under another name, so they save nothing. LiquidAI's own `decoder*.onnx` +graphs, which this pipeline does not use, do apply their rotary embedding. + +**Add `audio_output` to `genai_config.json`, and take the two switch tokens out of `eos_token_id`.** +They end a text-only answer; here they are where speech begins, and the model is refused if they are +still stop tokens: + +```json +{ + "model": { + "eos_token_id": 7, + "audio_output": { + "depthformer": { "filename": "vocoder_depthformer.onnx" }, + "embedding": { "filename": "audio_embedding.onnx" } + } + } +} +``` + +`audio_output` also takes `num_codebooks` (8), `codebook_size` (2049, the last entry being the +end-of-audio code), `audio_start_token_id` (128), `text_end_token_id` (130), `interleaved_n_text` (6) +and `interleaved_n_audio` (12). The defaults are the published checkpoints'; set +`interleaved_n_audio` to 9 for `LFM2.5-Audio-1.5B-JP`. + +**Generate.** Three search options steer the speech, next to the usual ones: + +| Option | Default | Meaning | +| --- | --- | --- | +| `audio_interleaved` | `false` | Alternate text and speech by count, as `generate_interleaved` does. Leave it off for TTS, ASR and chat, which switch on `<|audio_start|>` alone. | +| `audio_temperature` | `1.0` | Temperature of the audio codes. `0` takes the likeliest code. The model card uses 0.8 for TTS and 1.0 for interleaved. | +| `audio_top_k` | `4` | Audio codes kept when sampling. `1` takes the likeliest code. The model card uses 64 for TTS and 4 for interleaved. | + +The text is still decoded greedily, and `random_seed` seeds the audio codes as well. + +```python +import numpy as np +import onnxruntime_genai as og + +model = og.Model("./lfm2.5-audio/cpu") +processor = model.create_multimodal_processor() +tokenizer = og.Tokenizer(model) + +prompt = ( + "<|startoftext|><|im_start|>system\nRespond with interleaved text and audio.<|im_end|>\n" + "<|im_start|>user\n<|audio|><|im_end|>\n<|im_start|>assistant\n" +) +inputs = processor(prompt, audios=og.Audios.open("question.wav")) +prompt_length = inputs["input_ids"].as_numpy().shape[1] + +params = og.GeneratorParams(model) +params.set_search_options( + do_sample=False, max_length=prompt_length + 512, audio_interleaved=True, audio_temperature=1.0, audio_top_k=4 +) +generator = og.Generator(model, params) +generator.set_inputs(inputs) +while not generator.is_done(): + generator.generate_next_token() + +answer = generator.get_sequence(0)[prompt_length:] +audio_token_id, audio_start, text_end = 133, 128, 130 +text = tokenizer.decode(np.array([t for t in answer if t not in (audio_token_id, audio_start, text_end)], np.int32)) +codes = generator.get_output("audio_codes") # [num_frames, 8] int64, 80 ms of speech per frame +``` + +Each audio frame takes one place in the sequence, held by `audio_token_id`, so `max_length` counts +frames as well as text tokens; twelve and a half frames make a second of speech. The frames themselves +come from `get_output("audio_codes")`, at any point during generation or after it, with the +end-of-audio frames left out. + +**Turn the codes into sound** with `audio_detokenizer.onnx`, which gives the log-magnitude and phase +of a short-time Fourier transform, six columns per frame, and an inverse transform with the +reference's "same" padding. This graph needs no correction: it matches the PyTorch detokenizer to +1e-5. + +```python +import onnxruntime as ort +import soundfile as sf + +n_fft, hop = 1280, 320 +detokenizer = ort.InferenceSession("audio_detokenizer.onnx") +# Only the first codebook carries the end-of-audio code on purpose; the detokenizer takes 0..2047. +features = detokenizer.run(None, {"audio_codes": np.minimum(codes, 2047).T[None]})[0][0] +spectrum = (np.exp(features[:, : n_fft // 2 + 1]) * np.exp(1j * features[:, n_fft // 2 + 1 :])).T + +window = 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(n_fft) / n_fft) # torch.hann_window: periodic +frames = np.fft.irfft(spectrum, n_fft, axis=0) * window[:, None] +wave = np.zeros((frames.shape[1] - 1) * hop + n_fft) +envelope = np.zeros_like(wave) +for t in range(frames.shape[1]): + wave[t * hop : t * hop + n_fft] += frames[:, t] + envelope[t * hop : t * hop + n_fft] += window**2 +pad = (n_fft - hop) // 2 +sf.write("answer.wav", (wave[pad:-pad] / envelope[pad:-pad]).astype(np.float32), 24000) +``` + +**How it compares with the reference.** With the audio codes taken greedily on both sides +(`audio_top_k=1`), so that the two can be compared at all, this runtime reproduces liquid-audio +exactly: every text token and every audio frame of an interleaved answer to the reference's +`question.wav` (200 items, 149 of them frames), of two interleaved answers to typed questions, and of a +TTS sentence, on Linux x64 and on macOS arm64 alike. With the model card's sampling the speech reads back through this runtime's own ASR as +the text the model wrote: *Red, blue, and yellow are the three primary colors. Would you like to hear +how they’re used in art?* comes back word for word, and *The quick brown fox jumps over the lazy dog.* +comes back exactly in TTS. + +## 7. Known limitations + +**The waveform is made outside the runtime.** Generation gives audio codes; the detokenizer and the +inverse transform that turn them into sound are the few lines of Python in +[Speech output](#6-speech-output), not part of the generation loop. + +**Speech output is for one sequence.** A batch size of 1 and no beam search, as for speech input. +After an answer that ended in speech, a new turn appended to the same generator starts in text. + +**One prompt at a time, decoded greedily.** Several clips in one prompt work, but batched prompts +and beam search (`num_beams` above 1) are refused; the reference decodes its text greedily too. The +clips do not share an encoder run: the published encoder export is traced for a single clip, so each +one is encoded on its own frames and the results are concatenated in prompt order. + +**In ASR mode, two clips in one turn give one transcript.** Asked to transcribe a turn holding two +clips, the model writes out the last one and leaves the first, whichever order they come in. That is +the model and not this runtime: the reference implementation, given the same two clips through +`ChatState.add_audio`, produces the same tokens exactly, in both orders. Transcribe one clip per +request. + +**LoRA adapters are not supported.** The builder refuses `adapter_path` for these checkpoints: an +adapter trained on them names the decoder `lfm.*`, which does not match the decoder the builder +loads. Merge the adapter into the checkpoint first. + +**More than two channels has to be downmixed first.** Mono and stereo are handled: a stereo clip is +mixed down and reaches the front end at its true length. Wider audio is not — the decoder +deinterleaves it as though it were stereo, so a six channel clip arrives three times too long and +garbled, and transcribes to nonsense. Nothing in the decoder's interface reports how many channels a +file had, so this cannot be caught and refused here. Mix down to mono or stereo before passing the +file in. + +**Audio is not chunked.** The whole clip goes through the encoder in one run, and one hour of audio +is 450k mel frames, so memory grows with the clip length. The reference implementation has the same +shape; for long-form transcription, split the audio yourself. + +**The mel front end differs from the reference by about 1e-4.** It computes the same NeMo pipeline in +float32 where the reference uses float64 intermediates in places, and the FFT is a different +implementation. The difference is far below the audio's own quantization noise. + +**Resampling is the reference's.** A clip at any other rate is decoded at its own rate and brought to +16 kHz with `torchaudio`'s resampler (Hann-windowed sinc), upwards as well as downwards, which is what +`ChatState.add_audio` does. The audio decoder's own resampler is a different filter and changes the +log-mel frames by up to about 2, enough to flip a token: on the reference's 24 kHz `asr_jp.wav` it +turned one homophone, and resampled this way the transcript matches token for token. + +**CPU is the verified path; a GPU needs every graph on it.** All three modes were checked against +the reference on CPU, x64 and arm64, and match token for token and frame for frame. On **CUDA** they +match too, but only once *every* graph runs there: build the model for CUDA, or leave the embedding +and speech entries without their own `session_options` so they inherit the decoder's provider. A +mixed configuration — the decoder on the GPU and the embedding or encoder left on CPU, which is what +a CPU export's `genai_config.json` describes — is refused with an error naming the entry to remove: +the CPU graph would be handed buffers that live on the decoder's device and write host bytes into +them ([onnxruntime-genai#2604](https://github.com/microsoft/onnxruntime-genai/issues/2604)). The encoder's convolutions need cuDNN +on the library path. Expect the audio codes to drift a little against a CPU run: on one of the test +clips two of 1192 codes differed, both in the last two codebooks, with every text token identical. + +**WebGPU runs the text-in modes.** Interleaved and TTS match the reference once the `audio_output` +graphs keep their own CPU `session_options`: ONNX Runtime's WebGPU `GroupQueryAttention` fails +whenever the present KV cache has to grow past the buffer it was given, which is every step of the +depthformer's loop ([onnxruntime#32716](https://github.com/microsoft/onnxruntime/issues/32716)). +Speech *input* does not work there yet, for a reason of this runtime's own rather than the +encoder's: the features buffer lands on the wrong device +([onnxruntime-genai#2604](https://github.com/microsoft/onnxruntime-genai/issues/2604)), so a prompt +with audio is refused while the encoder or embedding keeps its own CPU `session_options`. Note also +that the shipped `onnxruntime-ep-webgpu` 0.3.0 cannot load the encoder at the default optimisation +level — it rejects the fused `Conv` activation that a current `onnxruntime` emits — but does load +and run it at `ORT_ENABLE_BASIC`. + +**Dither is off, as in the reference's eval mode.** NeMo adds 1e-5 of white noise to the samples +during training only; the runtime never does. diff --git a/src/config.cpp b/src/config.cpp index 6c0e726538..f2023b5bc1 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -1473,6 +1473,8 @@ struct SpeechInputs_Element : JSON::Element { v_.attention_mask = JSON::Get(value); } else if (name == "audio_sizes") { v_.audio_sizes = JSON::Get(value); + } else if (name == "audio_lengths") { + v_.audio_lengths = JSON::Get(value); } else if (name == "audio_projection_mode") { v_.audio_projection_mode = JSON::Get(value); } else { @@ -1542,6 +1544,68 @@ struct Speech_Element : JSON::Element { SpeechOutputs_Element outputs_{v_.outputs}; }; +struct AudioOutputGraph_Element : JSON::Element { + explicit AudioOutputGraph_Element(Config::Model::AudioOutput::Graph& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "filename") { + v_.filename = JSON::Get(value); + } else { + throw JSON::unknown_value_error{}; + } + } + + Element& OnObject(std::string_view name) override { + if (name == "session_options") { + v_.session_options = Config::SessionOptions{}; + session_options_ = std::make_unique(*v_.session_options); + return *session_options_; + } + throw JSON::unknown_value_error{}; + } + + private: + Config::Model::AudioOutput::Graph& v_; + std::unique_ptr session_options_; +}; + +struct AudioOutput_Element : JSON::Element { + explicit AudioOutput_Element(Config::Model::AudioOutput& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "num_codebooks") { + v_.num_codebooks = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "codebook_size") { + v_.codebook_size = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "audio_start_token_id") { + v_.audio_start_token_id = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "text_end_token_id") { + v_.text_end_token_id = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "interleaved_n_text") { + v_.interleaved_n_text = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "interleaved_n_audio") { + v_.interleaved_n_audio = SafeDoubleToInt(JSON::Get(value), name); + } else { + throw JSON::unknown_value_error{}; + } + } + + Element& OnObject(std::string_view name) override { + if (name == "depthformer") { + return depthformer_; + } + if (name == "embedding") { + return embedding_; + } + throw JSON::unknown_value_error{}; + } + + private: + Config::Model::AudioOutput& v_; + AudioOutputGraph_Element depthformer_{v_.depthformer}; + AudioOutputGraph_Element embedding_{v_.embedding}; +}; + struct JoinerInputs_Element : JSON::Element { explicit JoinerInputs_Element(Config::Model::Joiner::Inputs& v) : v_{v} {} @@ -2160,6 +2224,9 @@ struct Model_Element : JSON::Element { if (name == "speech") { return speech_; } + if (name == "audio_output") { + return audio_output_; + } if (name == "joiner") { return joiner_; } @@ -2196,6 +2263,7 @@ struct Model_Element : JSON::Element { Vision_Element vision_{v_.vision}; Embedding_Element embedding_{v_.embedding}; Speech_Element speech_{v_.speech}; + AudioOutput_Element audio_output_{v_.audio_output}; Joiner_Element joiner_{v_.joiner}; VAD_Element vad_{v_.vad}; Moonshine_Element moonshine_{v_.moonshine}; @@ -2310,6 +2378,12 @@ struct Search_Element : JSON::Element { v_.early_stopping = JSON::Get(value); } else if (name == "blank_penalty") { v_.blank_penalty = static_cast(JSON::Get(value)); + } else if (name == "audio_interleaved") { + v_.audio_interleaved = JSON::Get(value); + } else if (name == "audio_temperature") { + v_.audio_temperature = static_cast(JSON::Get(value)); + } else if (name == "audio_top_k") { + v_.audio_top_k = SafeDoubleToInt(JSON::Get(value), name); } else { throw JSON::unknown_value_error{}; } diff --git a/src/config.h b/src/config.h index 64022a32e3..13208cf02e 100644 --- a/src/config.h +++ b/src/config.h @@ -48,6 +48,7 @@ struct Config { // Speech encoder names static constexpr std::string_view AudioAttentionMaskName = "audio_attention_mask"; static constexpr std::string_view AudioSizesName = "audio_sizes"; + static constexpr std::string_view AudioLengthsName = "audio_lengths"; static constexpr std::string_view AudioProjectionModeName = "audio_projection_mode"; static constexpr std::string_view AudioFeaturesName = "audio_features"; static constexpr std::string_view NumAudioTokens = "num_audio_tokens"; @@ -317,6 +318,7 @@ struct Config { std::string audio_embeds{Defaults::AudioEmbedsName}; std::string attention_mask{Defaults::AudioAttentionMaskName}; std::string audio_sizes{Defaults::AudioSizesName}; + std::string audio_lengths{Defaults::AudioLengthsName}; // per-clip valid frame count of audio_embeds (LFM2-Audio) std::string audio_projection_mode{Defaults::AudioProjectionModeName}; } inputs; @@ -325,6 +327,24 @@ struct Config { } outputs; } speech; + // Speech output of LFM2-Audio: the depthformer turns one decoder hidden state into a frame of + // audio codes, one per codebook, and the audio embedding turns that frame back into the decoder's + // next input. Empty filenames leave the model answering in text only. + struct AudioOutput { + struct Graph { + std::string filename; + std::optional session_options; + }; + Graph depthformer; + Graph embedding; + int num_codebooks{8}; + int codebook_size{2049}; // entries per codebook; the last one is the end-of-audio code + int audio_start_token_id{128}; // <|audio_start|>: the rest of a sequential answer is speech + int text_end_token_id{130}; // <|text_end|>: the text of an interleaved answer is over + int interleaved_n_text{6}; // text tokens per turn of an interleaved answer + int interleaved_n_audio{12}; // audio frames per turn of an interleaved answer + } audio_output; + struct Joiner { std::string filename; std::optional session_options; @@ -734,6 +754,9 @@ struct Config { int random_seed{-1}; // -1 = Seed with random device, otherwise use value to seed RNG std::optional chunk_size; // Chunk size for prefill chunking during context processing. If present, chunking is enabled with the chunk size > 0. float blank_penalty{}; // Penalty applied to blank token logits in CTC/RNNT decoding. Default 0 means no penalty. + bool audio_interleaved{}; // LFM2-Audio: alternate text tokens and audio frames by count (interleaved mode) rather than switching on <|audio_start|>. + float audio_temperature{1.0f}; // LFM2-Audio: temperature the audio codes are sampled with. 0 takes the most likely code. + int audio_top_k{4}; // LFM2-Audio: number of most likely audio codes kept when sampling. 1 takes the most likely code. } search; struct Speculative { diff --git a/src/generator/generators.cpp b/src/generator/generators.cpp index 0b63554382..3608babb57 100644 --- a/src/generator/generators.cpp +++ b/src/generator/generators.cpp @@ -971,7 +971,7 @@ void Generator::RewindToLength(size_t new_length) { if (ModelType::IsTransducer(model_type) || ModelType::IsStreamingEncDecASR(model_type)) throw std::runtime_error("RewindTo is not supported for streaming ASR models (" + model_type + ")."); if (model_type == "whisper" || model_type == "phi3v" || model_type == "decoder-pipeline" || - ModelType::IsLFM2(model_type) || model_type == "lfm2_vl") + ModelType::IsLFM2(model_type) || model_type == "lfm2_vl" || model_type == "lfm2_audio") throw std::runtime_error("RewindTo is currently not supported for " + model_type + "."); const size_t current_length = search_->GetSequenceLength(); if (new_length > current_length) diff --git a/src/generator/generators.h b/src/generator/generators.h index 4a865a374f..40a1ee4734 100644 --- a/src/generator/generators.h +++ b/src/generator/generators.h @@ -70,6 +70,12 @@ DeviceSpan WrapTensor(DeviceInterface& device, OrtValue& value) { DeviceSpan ByteWrapTensor(DeviceInterface& device, OrtValue& value); +// The interface that matches where a tensor's memory actually is. A model that runs on a device does +// not necessarily keep its tensors there: WebGPU without graph capture leaves them host-side, and ORT +// places some outputs on the CPU whatever the EP. Wrapping such a tensor with the device interface +// hands a host pointer to the EP's copy, which crashes inside it. +DeviceInterface& DeviceForTensor(const OrtValue& value, DeviceInterface& device); + // OgaSequences are a vector of int32 vectors using TokenSequences = std::vector>; diff --git a/src/models/lfm2_audio_output.cpp b/src/models/lfm2_audio_output.cpp new file mode 100644 index 0000000000..f3d728450c --- /dev/null +++ b/src/models/lfm2_audio_output.cpp @@ -0,0 +1,283 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "generator/generators.h" +#include "models/lfm2_audio_output.h" +#include "models/multi_modal.h" +#include "models/utils.h" + +#include +#include +#include +#include + +namespace Generators { + +namespace { + +// Names of LiquidAI's depthformer and audio embedding exports (github.com/Liquid4All/onnx-export). +constexpr const char* kDepthformerInputs[] = {"hidden_states", "depth_slices_in", "step_idx", "prev_token", + "past_keys", "past_values", "seqlens_k", "total_seq_len"}; +constexpr const char* kDepthformerOutputs[] = {"logits", "depth_slices", "new_keys", "new_values"}; +constexpr const char* kEmbeddingInputs[] = {"audio_codes"}; +constexpr const char* kEmbeddingOutputs[] = {"audio_embeds"}; + +std::vector InputShape(OrtSession& session, const char* name, const char* graph) { + const auto names = session.GetInputNames(); + for (size_t i = 0; i < names.size(); ++i) { + if (names[i] == name) { + return session.GetInputTypeInfo(i)->GetTensorTypeAndShapeInfo().GetShape(); + } + } + throw std::runtime_error(std::string("Lfm2AudioOutput: the ") + graph + " model has no input named \"" + name + "\"."); +} + +template +std::unique_ptr MakeTensor(std::span values, std::vector shape) { + auto tensor = OrtValue::CreateTensor(Ort::Allocator::GetWithDefaultOptions(), shape); + std::copy(values.begin(), values.end(), tensor->template GetTensorMutableData()); + return tensor; +} + +template +std::unique_ptr MakeScalar(T value) { + return MakeTensor(std::span{&value, 1}, {}); +} + +} // namespace + +Lfm2AudioOutput::Lfm2AudioOutput(const MultiModalLanguageModel& model, const GeneratorParams& params) + : model_{model}, + config_{model.config_->model.audio_output}, + interleaved_{params.search.audio_interleaved}, + temperature_{params.search.audio_temperature}, + top_k_{params.search.audio_top_k}, + placeholder_token_id_{model.config_->model.audio_token_id}, + rng_{CreateRandomGenerator(params.search.random_seed)}, + modality_left_{config_.interleaved_n_text} { + if (params.BatchBeamSize() != 1) { + throw std::runtime_error("Lfm2AudioOutput: speech output needs a batch size of 1 and no beam search."); + } + + const auto hidden_shape = InputShape(*model_.depthformer_session_, "hidden_states", "depthformer"); + const auto slices_shape = InputShape(*model_.depthformer_session_, "depth_slices_in", "depthformer"); + const auto keys_shape = InputShape(*model_.depthformer_session_, "past_keys", "depthformer"); + if (hidden_shape.size() != 2 || slices_shape.size() != 3 || keys_shape.size() != 5 || + config_.num_codebooks <= 0 || slices_shape[1] != config_.num_codebooks) { + throw std::runtime_error("Lfm2AudioOutput: the depthformer model does not have the expected inputs for " + + std::to_string(config_.num_codebooks) + " codebooks."); + } + hidden_size_ = hidden_shape[1]; + depth_size_ = slices_shape[2]; + num_layers_ = keys_shape[0]; + num_kv_heads_ = keys_shape[2]; + head_size_ = keys_shape[4]; + + // RunDepthformer checks it against the logits the depthformer produces. + if (config_.codebook_size <= 0) { + throw std::runtime_error("Lfm2AudioOutput: model.audio_output.codebook_size " + + std::to_string(config_.codebook_size) + " must be positive."); + } + + // The search has to pick the placeholder whatever its sampling settings are. + const int vocab_size = model_.config_->model.vocab_size; + if (placeholder_token_id_ <= 0 || placeholder_token_id_ >= vocab_size) { + throw std::runtime_error("Lfm2AudioOutput: model.audio_token_id " + std::to_string(placeholder_token_id_) + + " must be set to a token inside the vocabulary of " + std::to_string(vocab_size) + + " tokens."); + } + placeholder_logits_ = model_.p_device_->Allocate(static_cast(vocab_size)); + auto logits = placeholder_logits_.CpuSpan(); + std::fill(logits.begin(), logits.end(), std::numeric_limits::lowest()); + logits[static_cast(placeholder_token_id_)] = 0.0f; + placeholder_logits_.CopyCpuToDevice(); +} + +void Lfm2AudioOutput::BeginStep(DeviceSpan& next_tokens, bool is_prompt) { + if (is_prompt || next_tokens.size() != 1) { + // A new turn starts in text, whatever the last one ended in. + modality_ = Modality::Text; + previous_was_text_ = false; + text_done_ = false; + modality_left_ = config_.interleaved_n_text; + pending_embedding_.clear(); + } else if (previous_was_text_) { + const int32_t token = next_tokens.CopyDeviceToCpu()[0]; + if (interleaved_) { + text_done_ = text_done_ || token == config_.text_end_token_id; + if (modality_left_ == 0 || text_done_) { + modality_ = Modality::Audio; + modality_left_ = config_.interleaved_n_audio; + } + } else if (token == config_.audio_start_token_id) { + modality_ = Modality::Audio; + } + } + + --modality_left_; + previous_was_text_ = modality_ == Modality::Text; +} + +void Lfm2AudioOutput::WritePendingFrame(OrtValue& inputs_embeds) { + const auto info = inputs_embeds.GetTensorTypeAndShapeInfo(); + if (info->GetElementCount() != pending_embedding_.size()) { + throw std::runtime_error("Lfm2AudioOutput: expected the decoder to take one embedding of " + + std::to_string(pending_embedding_.size()) + " values after an audio frame, got " + + std::to_string(info->GetElementCount()) + "."); + } + + auto bytes = ByteWrapTensor(DeviceForTensor(inputs_embeds, *model_.p_device_inputs_), inputs_embeds); + auto cpu = bytes.CpuSpan(); + if (info->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + std::memcpy(cpu.data(), pending_embedding_.data(), pending_embedding_.size() * sizeof(float)); + } else if (info->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) { + auto* half = reinterpret_cast(cpu.data()); + std::transform(pending_embedding_.begin(), pending_embedding_.end(), half, FastFloat32ToFloat16); + } else { + throw std::runtime_error("Lfm2AudioOutput: the decoder's inputs_embeds must be float or float16."); + } + bytes.CopyCpuToDevice(); + pending_embedding_.clear(); +} + +DeviceSpan Lfm2AudioOutput::SampleFrame(OrtValue& hidden_states) { + const auto info = hidden_states.GetTensorTypeAndShapeInfo(); + const auto shape = info->GetShape(); + if (shape.empty() || shape.back() != hidden_size_) { + throw std::runtime_error("Lfm2AudioOutput: the decoder's hidden states do not match the depthformer's input of " + + std::to_string(hidden_size_) + " values."); + } + + // The last position of [batch, sequence, hidden_size]. + const size_t count = static_cast(hidden_size_); + const size_t first = info->GetElementCount() - count; + auto bytes = ByteWrapTensor(DeviceForTensor(hidden_states, *model_.p_device_inputs_), hidden_states).CopyDeviceToCpu(); + std::vector hidden(count); + if (info->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + std::memcpy(hidden.data(), bytes.data() + first * sizeof(float), count * sizeof(float)); + } else if (info->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) { + const auto* half = reinterpret_cast(bytes.data()) + first; + std::transform(half, half + count, hidden.begin(), FastFloat16ToFloat32); + } else { + throw std::runtime_error("Lfm2AudioOutput: the decoder's hidden states must be float or float16."); + } + + std::vector frame = RunDepthformer(hidden); + + // As the reference orders it: the count first, then the end-of-audio code, which wins. + if (interleaved_ && modality_left_ == 0 && !text_done_) { + modality_ = Modality::Text; + modality_left_ = config_.interleaved_n_text; + } + const int64_t end_of_audio = config_.codebook_size - 1; + if (frame[0] == end_of_audio) { + std::fill(frame.begin(), frame.end(), end_of_audio); + modality_ = Modality::Text; + } else { + audio_codes_.insert(audio_codes_.end(), frame.begin(), frame.end()); + } + + EmbedFrame(frame); + return placeholder_logits_; +} + +std::vector Lfm2AudioOutput::RunDepthformer(std::span hidden) { + const auto hidden_tensor = MakeTensor(hidden, {1, hidden_size_}); + const std::vector cache_shape{num_layers_, 1, num_kv_heads_, 0, head_size_}; + auto keys = OrtValue::CreateTensor(Ort::Allocator::GetWithDefaultOptions(), cache_shape); + auto values = OrtValue::CreateTensor(Ort::Allocator::GetWithDefaultOptions(), cache_shape); + const std::vector no_slices(static_cast(config_.num_codebooks * depth_size_), 0.0f); + auto slices = MakeTensor(no_slices, {1, config_.num_codebooks, depth_size_}); + + // One run per codebook: each code is sampled given the ones before it in the frame. The projection + // of the hidden state is computed on the first run and handed back in on the rest. + std::vector frame; + int64_t previous_code = 0; + for (int32_t step = 0; step < config_.num_codebooks; ++step) { + const auto step_index = MakeScalar(step); + const auto previous = MakeTensor(std::span{&previous_code, 1}, {1}); + const auto past_length = MakeTensor(std::span{&step, 1}, {1}); + const auto total_length = MakeScalar(step + 1); + const OrtValue* inputs[] = {hidden_tensor.get(), slices.get(), step_index.get(), previous.get(), + keys.get(), values.get(), past_length.get(), total_length.get()}; + auto outputs = model_.depthformer_session_->Run(nullptr, kDepthformerInputs, inputs, std::size(inputs), + kDepthformerOutputs, std::size(kDepthformerOutputs)); + + // The code is sampled from the first codebook_size logits, so there have to be at least that many. + const auto logits_info = outputs[0]->GetTensorTypeAndShapeInfo(); + if (logits_info->GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + throw std::runtime_error(std::string("Lfm2AudioOutput: the depthformer's logits must be float, got ") + + TypeToString(logits_info->GetElementType()) + "."); + } + const auto logits_shape = logits_info->GetShape(); + if (logits_shape.empty() || logits_shape.back() < config_.codebook_size) { + throw std::runtime_error("Lfm2AudioOutput: model.audio_output.codebook_size is " + + std::to_string(config_.codebook_size) + ", but the depthformer produces " + + std::to_string(logits_shape.empty() ? 0 : logits_shape.back()) + " logits."); + } + const float* logits = outputs[0]->GetTensorData(); + previous_code = SampleCode({logits, static_cast(config_.codebook_size)}); + frame.push_back(previous_code); + if (step == 0) { + slices = std::move(outputs[1]); + } + keys = std::move(outputs[2]); + values = std::move(outputs[3]); + } + return frame; +} + +int64_t Lfm2AudioOutput::SampleCode(std::span logits) { + if (temperature_ <= 0.0f || top_k_ == 1) { + return std::max_element(logits.begin(), logits.end()) - logits.begin(); + } + + std::vector order(logits.size()); + std::iota(order.begin(), order.end(), int64_t{0}); + const size_t kept = top_k_ > 0 ? std::min(static_cast(top_k_), order.size()) : order.size(); + std::partial_sort(order.begin(), order.begin() + kept, order.end(), + [&](int64_t a, int64_t b) { return logits[a] > logits[b]; }); + + std::vector weights(kept); + const float highest = logits[order[0]]; + for (size_t i = 0; i < kept; ++i) { + weights[i] = std::exp(static_cast(logits[order[i]] - highest) / temperature_); + } + std::discrete_distribution distribution(weights.begin(), weights.end()); + return order[distribution(rng_)]; +} + +void Lfm2AudioOutput::EmbedFrame(std::span frame) { + // Every codebook has its own rows of the one embedding table. + std::vector rows(frame.size()); + for (size_t codebook = 0; codebook < frame.size(); ++codebook) { + rows[codebook] = frame[codebook] + static_cast(codebook) * config_.codebook_size; + } + const auto codes = MakeTensor(rows, {1, static_cast(rows.size())}); + const OrtValue* inputs[] = {codes.get()}; + auto outputs = model_.audio_embedding_session_->Run(nullptr, kEmbeddingInputs, inputs, std::size(inputs), + kEmbeddingOutputs, std::size(kEmbeddingOutputs)); + + const auto info = outputs[0]->GetTensorTypeAndShapeInfo(); + const auto shape = info->GetShape(); // [1, num_codebooks, hidden_size] + if (info->GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT || shape.empty() || + info->GetElementCount() != frame.size() * static_cast(shape.back())) { + throw std::runtime_error("Lfm2AudioOutput: the audio embedding model must produce one float embedding per codebook."); + } + const size_t width = static_cast(shape.back()); + const float* embeddings = outputs[0]->GetTensorData(); + pending_embedding_.assign(width, 0.0f); + for (size_t codebook = 0; codebook < frame.size(); ++codebook) { + for (size_t i = 0; i < width; ++i) { + pending_embedding_[i] += embeddings[codebook * width + i]; + } + } +} + +OrtValue* Lfm2AudioOutput::GetAudioCodes() { + const int64_t num_frames = static_cast(audio_codes_.size()) / config_.num_codebooks; + audio_codes_tensor_ = MakeTensor(audio_codes_, {num_frames, config_.num_codebooks}); + return audio_codes_tensor_.get(); +} + +} // namespace Generators diff --git a/src/models/lfm2_audio_output.h b/src/models/lfm2_audio_output.h new file mode 100644 index 0000000000..f13b8538cd --- /dev/null +++ b/src/models/lfm2_audio_output.h @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include + +#include "models/model.h" + +namespace Generators { + +struct MultiModalLanguageModel; + +// Speech output of LFM2-Audio, as liquid-audio's generate_sequential and generate_interleaved +// produce it. +// +// The decoder emits one hidden state per step. While the answer is text the search samples a token +// from the logits as usual; while it is speech the depthformer turns that hidden state into one +// audio frame of num_codebooks codes instead, and the sum of the codes' embeddings is the decoder's +// next input. The generator's sequence records each frame as one model.audio_token_id placeholder, +// so lengths and stopping work unchanged, and the codes themselves are read with +// generator.get_output("audio_codes"). +// +// The turn passes to speech on <|audio_start|> and back on the end-of-audio code. With +// search.audio_interleaved it also alternates by count, interleaved_n_text tokens against +// interleaved_n_audio frames, until <|text_end|> leaves the rest of the turn to speech. +struct Lfm2AudioOutput { + Lfm2AudioOutput(const MultiModalLanguageModel& model, const GeneratorParams& params); + + // Called at the top of every pipeline run with the tokens the search just chose. + void BeginStep(DeviceSpan& next_tokens, bool is_prompt); + + // True when the decoder's next input is the last audio frame rather than a token. + bool HasPendingFrame() const { return !pending_embedding_.empty(); } + // Overwrites the decoder's [1, 1, hidden_size] input with the last frame's embedding. + void WritePendingFrame(OrtValue& inputs_embeds); + + bool InAudio() const { return modality_ == Modality::Audio; } + // Samples the next frame from the decoder's last hidden state and returns logits that make the + // search pick the placeholder token. + DeviceSpan SampleFrame(OrtValue& hidden_states); + + // Every frame so far, [num_frames, num_codebooks] int64, end-of-audio frames left out. + OrtValue* GetAudioCodes(); + + private: + enum struct Modality { Text, + Audio }; + + std::vector RunDepthformer(std::span hidden); + int64_t SampleCode(std::span logits); + void EmbedFrame(std::span frame); + + const MultiModalLanguageModel& model_; + const Config::Model::AudioOutput& config_; + const bool interleaved_; + const float temperature_; + const int top_k_; + const int32_t placeholder_token_id_; + std::mt19937 rng_; + + // Depthformer geometry, read off the graph. + int64_t hidden_size_{}; + int64_t depth_size_{}; + int64_t num_layers_{}; + int64_t num_kv_heads_{}; + int64_t head_size_{}; + + Modality modality_{Modality::Text}; + bool previous_was_text_{}; + bool text_done_{}; + int modality_left_{}; + + std::vector pending_embedding_; // the last frame, as the decoder's next input + std::vector audio_codes_; // num_frames * num_codebooks + std::unique_ptr audio_codes_tensor_; + DeviceSpan placeholder_logits_; +}; + +} // namespace Generators diff --git a/src/models/model.cpp b/src/models/model.cpp index 4c6f465719..321d8ca0ce 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -949,6 +949,8 @@ std::shared_ptr CreateModel(OrtEnv& ort_env, std::unique_ptr conf return std::make_shared(std::move(config), ort_env); if (ModelType::IsTDT(config->model.type)) return std::make_shared(std::move(config), ort_env); + if (config->model.type == "lfm2_audio") + return std::make_shared(std::move(config), ort_env, /*vision=*/false, /*speech=*/true); if (ModelType::IsALM(config->model.type)) return std::make_shared(std::move(config), ort_env); if (ModelType::IsVLM(config->model.type)) diff --git a/src/models/model_type.h b/src/models/model_type.h index 4690ba4f67..d5e9ea07be 100644 --- a/src/models/model_type.h +++ b/src/models/model_type.h @@ -17,7 +17,7 @@ namespace Generators { struct ModelType { inline static bool IsLLM(const std::string& model_type) { // Large-language model (LLM) - static constexpr std::array LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma3_vl_text", "gemma4_text", "gpt2", "gptoss", "granite", "granitemoehybrid", "hunyuandensev1", "internlm2", "lfm2", "lfm2_moe", "lfm2_vl_text", "llama", "mistral", "mistral3_text", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen2_5_vl_text", "qwen3", "qwen3_vl_text", "qwen3_5_moe_text", "qwen3_5_text", "smollm3"}; + static constexpr std::array LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma3_vl_text", "gemma4_text", "gpt2", "gptoss", "granite", "granitemoehybrid", "hunyuandensev1", "internlm2", "lfm2", "lfm2_audio_text", "lfm2_moe", "lfm2_vl_text", "llama", "mistral", "mistral3_text", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen2_5_vl_text", "qwen3", "qwen3_vl_text", "qwen3_5_moe_text", "qwen3_5_text", "smollm3"}; return std::find(LLM.begin(), LLM.end(), model_type) != LLM.end(); } @@ -28,8 +28,9 @@ struct ModelType { } inline static bool IsALM(const std::string& model_type) { - // Audio-language model (ALM) - static constexpr std::array ALM = {"whisper"}; + // Audio-language model (ALM): audio in, text out. "lfm2_audio" runs the multi-modal pipeline + // (speech encoder -> embedding -> LFM2 decoder), "whisper" its own encoder-decoder model. + static constexpr std::array ALM = {"lfm2_audio", "whisper"}; return std::find(ALM.begin(), ALM.end(), model_type) != ALM.end(); } @@ -86,9 +87,10 @@ struct ModelType { inline static bool IsLFM2(const std::string& model_type) { // Liquid Foundation Model 2: hybrid attention/conv architecture with conv state cache. // The MoE variant (LFM2-8B-A1B, LFM2.5-8B-A1B, LFM2-24B-A2B) has the same runtime shape, and - // "lfm2_vl_text" is an LFM2-VL checkpoint exported as a plain text decoder; both have the same - // conv/attention layer mix and need the same conv state cache. - return model_type == "lfm2" || model_type == "lfm2_moe" || model_type == "lfm2_vl_text"; + // "lfm2_vl_text" / "lfm2_audio_text" are LFM2-VL / LFM2-Audio checkpoints exported as plain text + // decoders; all have the same conv/attention layer mix and need the same conv state cache. + return model_type == "lfm2" || model_type == "lfm2_moe" || model_type == "lfm2_vl_text" || + model_type == "lfm2_audio_text"; } }; diff --git a/src/models/multi_modal.cpp b/src/models/multi_modal.cpp index 911f284b6f..c1982fb235 100644 --- a/src/models/multi_modal.cpp +++ b/src/models/multi_modal.cpp @@ -95,8 +95,51 @@ int64_t GetImageFeatureBatchSize(const std::vector& extra_inputs) { } // namespace +void CheckLfm2AudioSessionDevices(const Config& config, DeviceType decoder_device, DeviceType inputs_device, + bool with_audio) { + // Whether a buffer allocated on the device is device memory. OpenVINO allocates from the CPU and + // QNN from shared memory, both of which a CPU session can use. + const auto is_device_memory = [](DeviceType device) { + switch (device) { + case DeviceType::CUDA: + case DeviceType::DML: + case DeviceType::WEBGPU: + case DeviceType::NvTensorRtRtx: + case DeviceType::RyzenAI: + case DeviceType::AMDGPU: + return true; + default: + return false; + } + }; + const auto check = [](const std::optional& options, const char* graph, + const char* buffers, DeviceType device) { + // Without session_options of its own a graph follows the decoder; with them, no provider but CPU + // leaves it on CPU. + if (!options.has_value() || !std::all_of(options->providers.begin(), options->providers.end(), + [](const std::string& provider) { return provider == "CPU"; })) { + return; + } + throw std::runtime_error(std::string("lfm2_audio: model.") + graph + ".session_options run the " + graph + + " model on CPU, but " + buffers + " in " + to_string(device) + + " memory, which it cannot use. Remove model." + graph + + ".session_options so that it runs on the decoder's device."); + }; + + if (is_device_memory(inputs_device)) { + check(config.model.embedding.session_options, "embedding", "the decoder takes its inputs", inputs_device); + } + if (with_audio && is_device_memory(decoder_device)) { + check(config.model.speech.session_options, "speech", "the audio features are passed", decoder_device); + check(config.model.embedding.session_options, "embedding", "the audio features are passed", decoder_device); + } +} + MultiModalLanguageModel::MultiModalLanguageModel(std::unique_ptr config, OrtEnv& ort_env, bool vision, bool speech) : Model(std::move(config)) { + if (config_->model.type == "lfm2_audio") { + CheckLfm2AudioSessionDevices(*config_, p_device_->GetType(), p_device_inputs_->GetType(), /*with_audio=*/false); + } // The non-decoder models don't support graph capture because of control flow nodes, so disable graph capture for them if (vision) { vision_session_options_ = OrtSessionOptions::Create(); @@ -116,6 +159,28 @@ MultiModalLanguageModel::MultiModalLanguageModel(std::unique_ptr config, embedding_session_ = CreateSession(ort_env, config_->model.embedding.filename, embedding_session_options_.get()); decoder_session_ = CreateSession(ort_env, config_->model.decoder.filename, session_options_.get()); + const auto& audio_output = config_->model.audio_output; + if (!audio_output.depthformer.filename.empty() || !audio_output.embedding.filename.empty()) { + if (audio_output.depthformer.filename.empty() || audio_output.embedding.filename.empty()) { + throw std::runtime_error("model.audio_output needs both depthformer.filename and embedding.filename."); + } + // With speech output these tokens hand the turn to the depthformer; as stop tokens they would end + // it there instead, which is what a text-only export needs and this one must not have. + for (const int token : {audio_output.audio_start_token_id, audio_output.text_end_token_id}) { + if (std::find(config_->model.eos_token_id.begin(), config_->model.eos_token_id.end(), token) != + config_->model.eos_token_id.end()) { + throw std::runtime_error("model.eos_token_id holds " + std::to_string(token) + + ", which model.audio_output uses to start speech. Remove it from eos_token_id."); + } + } + depthformer_session_options_ = OrtSessionOptions::Create(); + CreateSessionOptionsFromConfig(audio_output.depthformer.session_options.has_value() ? audio_output.depthformer.session_options.value() : config_->model.decoder.session_options, *depthformer_session_options_, true, /*disable_graph_capture=*/true); + depthformer_session_ = CreateSession(ort_env, audio_output.depthformer.filename, depthformer_session_options_.get()); + audio_embedding_session_options_ = OrtSessionOptions::Create(); + CreateSessionOptionsFromConfig(audio_output.embedding.session_options.has_value() ? audio_output.embedding.session_options.value() : config_->model.decoder.session_options, *audio_embedding_session_options_, true, /*disable_graph_capture=*/true); + audio_embedding_session_ = CreateSession(ort_env, audio_output.embedding.filename, audio_embedding_session_options_.get()); + } + session_info_.Add(*decoder_session_); session_info_.Add(*embedding_session_); if (speech) { @@ -603,6 +668,154 @@ DeviceSpan SpeechState::Run(int current_length, DeviceSpan& next return {}; } +void Lfm2AudioSpeechState::SetExtraInputs(const std::vector& extra_inputs, const int64_t num_audio_tokens) { + // The feature buffer is one sequence wide, so beams would fail as a shape mismatch inside the + // encoder. The reference decodes greedily. + if (params_->search.num_beams > 1) { + throw std::runtime_error("Lfm2AudioSpeechState: beam search is not supported for lfm2_audio; got num_beams " + + std::to_string(params_->search.num_beams) + ". Set num_beams to 1."); + } + + SpeechState::SetExtraInputs(extra_inputs, num_audio_tokens); + + // audio_sizes holds the decoder tokens each clip contributes; its sum is num_audio_tokens. + tokens_per_clip_.clear(); + for (const auto& input : extra_inputs) { + if (input.name == model_.config_->model.speech.inputs.audio_sizes) { + const auto info = input.tensor->ort_tensor_->GetTensorTypeAndShapeInfo(); + const int64_t* sizes = input.tensor->ort_tensor_->GetTensorData(); + tokens_per_clip_.assign(sizes, sizes + info->GetElementCount()); + break; + } + } +} + +Lfm2AudioSpeechState::SpeechBindings Lfm2AudioSpeechState::ResolveBindings() const { + SpeechBindings bindings; + bindings.mel_index = FindInput(model_.config_->model.speech.inputs.audio_embeds); + bindings.lengths_index = FindInput(model_.config_->model.speech.inputs.audio_lengths); + bindings.features_index = FindOutput(model_.config_->model.speech.outputs.audio_features); + + const auto mel_info = inputs_[bindings.mel_index]->GetTensorTypeAndShapeInfo(); + const auto mel_shape = mel_info->GetShape(); // [num_clips, longest_clip, num_mels] + if (mel_shape.size() != 3) { + throw std::runtime_error("Lfm2AudioSpeechState: expected a 3D [num_clips, num_frames, num_mels] mel tensor, got rank " + + std::to_string(mel_shape.size()) + "."); + } + bindings.num_clips = mel_shape[0]; + bindings.longest_clip = mel_shape[1]; + bindings.num_mels = mel_shape[2]; + bindings.mel_type = mel_info->GetElementType(); + + if (bindings.num_clips != static_cast(tokens_per_clip_.size())) { + throw std::runtime_error("Lfm2AudioSpeechState: the mel tensor holds " + std::to_string(bindings.num_clips) + + " clips but audio_sizes has " + std::to_string(tokens_per_clip_.size()) + " entries."); + } + const auto lengths_info = inputs_[bindings.lengths_index]->GetTensorTypeAndShapeInfo(); + if (static_cast(lengths_info->GetElementCount()) != bindings.num_clips) { + throw std::runtime_error("Lfm2AudioSpeechState: the mel tensor holds " + std::to_string(bindings.num_clips) + + " clips but " + model_.config_->model.speech.inputs.audio_lengths + " has " + + std::to_string(lengths_info->GetElementCount()) + " entries."); + } + + const auto features_info = outputs_[bindings.features_index]->GetTensorTypeAndShapeInfo(); + const auto features_shape = features_info->GetShape(); // [1, num_audio_tokens, hidden_size] + bindings.features_type = features_info->GetElementType(); + bindings.hidden_size = features_shape.back(); + return bindings; +} + +std::unique_ptr Lfm2AudioSpeechState::RunClip(const SpeechBindings& bindings, int64_t index, + int64_t num_frames) { + const size_t mel_element_size = Ort::SizeOf(bindings.mel_type); + const size_t clip_stride = static_cast(bindings.longest_clip * bindings.num_mels) * mel_element_size; + const auto* mel_data = static_cast(inputs_[bindings.mel_index]->GetTensorRawData()); + + auto clip_mel = OrtValue::CreateTensor(Ort::Allocator::GetWithDefaultOptions(), + std::vector{1, num_frames, bindings.num_mels}, bindings.mel_type); + std::memcpy(clip_mel->GetTensorMutableRawData(), mel_data + static_cast(index) * clip_stride, + static_cast(num_frames * bindings.num_mels) * mel_element_size); + + auto clip_length = OrtValue::CreateTensor(Ort::Allocator::GetWithDefaultOptions(), std::vector{1}); + clip_length->GetTensorMutableData()[0] = num_frames; + + auto clip_features = OrtValue::CreateTensor( + model_.p_device_->GetAllocator(), + std::vector{1, tokens_per_clip_[static_cast(index)], bindings.hidden_size}, + bindings.features_type); + + OrtValue* mel_batch = inputs_[bindings.mel_index]; + OrtValue* lengths = inputs_[bindings.lengths_index]; + OrtValue* features = outputs_[bindings.features_index]; + inputs_[bindings.mel_index] = clip_mel.get(); + inputs_[bindings.lengths_index] = clip_length.get(); + outputs_[bindings.features_index] = clip_features.get(); + State::Run(*model_.speech_session_); + inputs_[bindings.mel_index] = mel_batch; + inputs_[bindings.lengths_index] = lengths; + outputs_[bindings.features_index] = features; + + return clip_features; +} + +DeviceSpan Lfm2AudioSpeechState::Run(int current_length, DeviceSpan& next_tokens, DeviceSpan next_indices) { + CheckLfm2AudioSessionDevices(*model_.config_, model_.p_device_->GetType(), model_.p_device_inputs_->GetType(), + /*with_audio=*/true); + if (model_.config_->model.speech.run_options.has_value()) { + State::SetRunOptions(model_.config_->model.speech.run_options.value()); + } + // A single clip fills the whole mel tensor and the whole feature buffer; run it as it stands. + if (tokens_per_clip_.size() <= 1) { + State::Run(*model_.speech_session_); + return {}; + } + + const SpeechBindings bindings = ResolveBindings(); + const int64_t* frames_per_clip = inputs_[bindings.lengths_index]->GetTensorData(); + auto features_bytes = ByteWrapTensor(*model_.p_device_, *outputs_[bindings.features_index]); + const size_t feature_row_bytes = static_cast(bindings.hidden_size) * Ort::SizeOf(bindings.features_type); + size_t destination = 0; + + // The published encoder export is traced for one clip (its subsampling mask cannot broadcast over + // a batch), so run it once per clip on that clip's own frames — which also keeps the padding out + // of the encoder entirely — and concatenate the results in clip order. + for (int64_t clip = 0; clip < bindings.num_clips; ++clip) { + const int64_t num_frames = frames_per_clip[clip]; + if (num_frames <= 0 || num_frames > bindings.longest_clip) { + throw std::runtime_error("Lfm2AudioSpeechState: clip " + std::to_string(clip) + " reports " + + std::to_string(num_frames) + " mel frames, outside the 1.." + + std::to_string(bindings.longest_clip) + " the mel tensor holds."); + } + + auto clip_features = RunClip(bindings, clip, num_frames); + const size_t clip_bytes = static_cast(tokens_per_clip_[static_cast(clip)]) * feature_row_bytes; + features_bytes.subspan(destination, clip_bytes).CopyFrom(ByteWrapTensor(*model_.p_device_, *clip_features)); + destination += clip_bytes; + } + return {}; +} + +size_t Lfm2AudioSpeechState::FindInput(const std::string& name) const { + for (size_t i = 0; i < input_names_.size(); ++i) { + if (name == input_names_[i]) return i; + } + throw std::runtime_error("Lfm2AudioSpeechState: speech input \"" + name + "\" is not bound."); +} + +size_t Lfm2AudioSpeechState::FindOutput(const std::string& name) const { + for (size_t i = 0; i < output_names_.size(); ++i) { + if (name == output_names_[i]) return i; + } + throw std::runtime_error("Lfm2AudioSpeechState: speech output \"" + name + "\" is not bound."); +} + +std::unique_ptr CreateSpeechState(const MultiModalLanguageModel& model, const GeneratorParams& params) { + if (model.config_->model.type == "lfm2_audio") { + return std::make_unique(model, params); + } + return std::make_unique(model, params); +} + EmbeddingState::EmbeddingState(const MultiModalLanguageModel& model, const GeneratorParams& params) : State{params, model}, model_{model} { @@ -818,10 +1031,13 @@ MultiModalPipelineState::MultiModalPipelineState(const MultiModalLanguageModel& vision_state_ = CreateVisionState(model_, params); } if (model_.speech_session_) { - speech_state_ = std::make_unique(model_, params); + speech_state_ = CreateSpeechState(model_, params); } embedding_state_ = std::make_unique(model, params); decoder_state_ = std::make_unique(model_, sequence_lengths, params); + if (model_.depthformer_session_) { + audio_output_ = std::make_unique(model_, params); + } if (vision_state_ != nullptr && model_.config_->model.vision.adapter_filename.has_value() && num_image_tokens_ > 0) { const auto lora_adapter = (model_.config_->config_path / fs::path(*model_.config_->model.vision.adapter_filename)).string(); @@ -877,6 +1093,9 @@ DeviceSpan MultiModalPipelineState::Run(int current_length, DeviceSpan |embeddings_model| -> inputs_embeds // - inputs_embeds -> |decoder_model| -> logits + if (audio_output_) { + audio_output_->BeginStep(next_tokens, is_prompt_); + } embedding_state_->UpdateInputsOutputs(next_tokens, is_prompt_); // Prefill chunking (search.chunk_size): during the prompt stage the decoder can process the @@ -930,15 +1149,36 @@ DeviceSpan MultiModalPipelineState::Run(int current_length, DeviceSpaninputs_embeds_.ReuseEmbeddingsBuffer(decoder_state_->inputs_embeds_); if (embedding_state_->per_layer_inputs_ && decoder_state_->per_layer_inputs_) { embedding_state_->per_layer_inputs_->ReuseEmbeddingsBuffer(*decoder_state_->per_layer_inputs_); } - embedding_state_->Run(current_length, next_tokens, next_indices); - return decoder_state_->Run(current_length, next_tokens, next_indices); + if (audio_output_ && audio_output_->HasPendingFrame()) { + // The token is only the placeholder of the last audio frame: the decoder takes the frame itself, + // and the embedding model would look for audio features to put in the placeholder's place. + audio_output_->WritePendingFrame(*decoder_state_->inputs_embeds_.Get()); + } else { + embedding_state_->Run(current_length, next_tokens, next_indices); + } + auto logits = decoder_state_->Run(current_length, next_tokens, next_indices); + return audio_output_ ? SampleAudioOrText(logits) : logits; +} + +DeviceSpan MultiModalPipelineState::SampleAudioOrText(DeviceSpan logits) { + if (!audio_output_->InAudio()) { + return logits; + } + const std::string& configured = model_.config_->model.decoder.outputs.hidden_states; + const std::string name = configured.empty() ? std::string(Config::Defaults::HiddenStatesName) : configured; + OrtValue* hidden_states = decoder_state_->GetOutput(name.c_str()); + if (!hidden_states) { + throw std::runtime_error("Speech output needs the decoder's \"" + name + + "\" output. Build the decoder with --extra_options include_hidden_states=true."); + } + return audio_output_->SampleFrame(*hidden_states); } OrtValue* MultiModalPipelineState::GetInput(const char* name) { @@ -978,6 +1218,10 @@ OrtValue* MultiModalPipelineState::GetInput(const char* name) { }; OrtValue* MultiModalPipelineState::GetOutput(const char* name) { + if (audio_output_ && std::strcmp(name, "audio_codes") == 0) { + return audio_output_->GetAudioCodes(); + } + if (vision_state_) { // Check if output name is in vision state's outputs for (size_t i = 0; i < vision_state_->output_names_.size(); i++) { diff --git a/src/models/multi_modal.h b/src/models/multi_modal.h index 7726f292b8..35d0554b0c 100644 --- a/src/models/multi_modal.h +++ b/src/models/multi_modal.h @@ -19,9 +19,17 @@ #include "models/io/position_inputs.h" #include "model_type.h" #include "models/io/recurrent_state.h" +#include "models/lfm2_audio_output.h" namespace Generators { +// lfm2_audio exchanges tensors with its embedding and speech sessions in buffers allocated on the +// decoder's devices. Throws if either has session_options of its own that leave it on CPU while a +// buffer it would be handed is device memory, which the session would treat as host memory and +// corrupt: the decoder's inputs for the embedding session, and with_audio, the audio features for both. +void CheckLfm2AudioSessionDevices(const Config& config, DeviceType decoder_device, DeviceType inputs_device, + bool with_audio); + struct MultiModalLanguageModel : Model { MultiModalLanguageModel(std::unique_ptr config, OrtEnv& ort_env, bool vision, bool speech); MultiModalLanguageModel(const MultiModalLanguageModel&) = delete; @@ -37,6 +45,12 @@ struct MultiModalLanguageModel : Model { std::unique_ptr vision_session_options_; std::unique_ptr speech_session_options_; std::unique_ptr embedding_session_options_; + + // LFM2-Audio speech output, present when model.audio_output names the two graphs. + std::unique_ptr depthformer_session_; // hidden_states -> one audio code per run, a frame in num_codebooks runs + std::unique_ptr audio_embedding_session_; // audio_codes -> audio_embeds, summed into the decoder's next input + std::unique_ptr depthformer_session_options_; + std::unique_ptr audio_embedding_session_options_; }; // Base VisionState: runs vision.onnx with a single State::Run() call. @@ -133,10 +147,10 @@ struct SpeechState : State { SpeechState(const SpeechState&) = delete; SpeechState& operator=(const SpeechState&) = delete; - void SetExtraInputs(const std::vector& extra_inputs, const int64_t num_audio_tokens); + virtual void SetExtraInputs(const std::vector& extra_inputs, const int64_t num_audio_tokens); DeviceSpan Run(int current_length, DeviceSpan& next_tokens, DeviceSpan next_indices = {}) override; - private: + protected: friend struct MultiModalPipelineState; const MultiModalLanguageModel& model_; @@ -145,6 +159,49 @@ struct SpeechState : State { std::unique_ptr audio_features_; }; +// Lfm2AudioSpeechState: per-clip encoder loop for LFM2-Audio. +// +// The processor stacks the clips of one prompt into a zero-padded [N, T_max, num_mels] mel tensor +// with their real frame counts alongside. The published encoder export is traced for a single clip +// (its subsampling mask cannot broadcast over a batch), so with several clips this subclass slices +// each clip's own frames out of that tensor, runs the encoder on [1, T_i, num_mels], and writes the +// results one after another into the contiguous [1, total_tokens, hidden] feature buffer the +// embedding model expects, in clip order. +struct Lfm2AudioSpeechState : SpeechState { + using SpeechState::SpeechState; // inherit constructor + + void SetExtraInputs(const std::vector& extra_inputs, const int64_t num_audio_tokens) override; + DeviceSpan Run(int current_length, DeviceSpan& next_tokens, DeviceSpan next_indices = {}) override; + + private: + // Where the processor's staged mel tensor and the encoder's feature buffer are bound. + struct SpeechBindings { + size_t mel_index{}; // the [num_clips, longest, num_mels] staging tensor + size_t lengths_index{}; // its per-clip frame counts + size_t features_index{}; // the [1, total_tokens, hidden] buffer the embedding model reads + int64_t num_clips{}; + int64_t longest_clip{}; + int64_t num_mels{}; + int64_t hidden_size{}; + ONNXTensorElementDataType mel_type{}; + ONNXTensorElementDataType features_type{}; + }; + + // Reads the bound shapes and checks them against the per-clip token counts. + SpeechBindings ResolveBindings() const; + + // Runs the encoder on clip `index`'s own frames and returns its features. + std::unique_ptr RunClip(const SpeechBindings& bindings, int64_t index, int64_t num_frames); + + size_t FindInput(const std::string& name) const; + size_t FindOutput(const std::string& name) const; + + std::vector tokens_per_clip_; +}; + +// Factory: pick the right SpeechState subclass based on model type. +std::unique_ptr CreateSpeechState(const MultiModalLanguageModel& model, const GeneratorParams& params); + struct EmbeddingState : State { EmbeddingState(const MultiModalLanguageModel& model, const GeneratorParams& params); EmbeddingState(const EmbeddingState&) = delete; @@ -221,6 +278,8 @@ struct MultiModalPipelineState : State { private: void UpdateInputsOutputs(const DeviceSpan& next_tokens, DeviceSpan next_indices, int current_length); + // The decoder's logits while the answer is text, or the next audio frame's placeholder while it is speech. + DeviceSpan SampleAudioOrText(DeviceSpan logits); const MultiModalLanguageModel& model_; int64_t num_image_tokens_{}; @@ -230,6 +289,7 @@ struct MultiModalPipelineState : State { std::unique_ptr speech_state_; std::unique_ptr embedding_state_; std::unique_ptr decoder_state_; + std::unique_ptr audio_output_; // LFM2-Audio speech output, when the model has it std::shared_ptr adapters_; bool is_prompt_{true}; diff --git a/src/models/preprocessing/genai_tokenizer.cpp b/src/models/preprocessing/genai_tokenizer.cpp index 53309dc134..783c04bf57 100644 --- a/src/models/preprocessing/genai_tokenizer.cpp +++ b/src/models/preprocessing/genai_tokenizer.cpp @@ -114,6 +114,12 @@ std::vector Tokenizer::Encode(const char* text) const { } std::string Tokenizer::Decode(std::span tokens) const { + // The detokenizer dies with SIGFPE on an empty sequence, which a generation that stops on its + // first token produces. + if (tokens.empty()) { + return {}; + } + OrtxPtr ortx_string_array; CheckResult(OrtxDetokenize1D(tokenizer_, reinterpret_cast(tokens.data()), tokens.size(), ortx_string_array.Address())); diff --git a/src/models/preprocessing/lfm2_audio_processor.cpp b/src/models/preprocessing/lfm2_audio_processor.cpp new file mode 100644 index 0000000000..11343f1f8c --- /dev/null +++ b/src/models/preprocessing/lfm2_audio_processor.cpp @@ -0,0 +1,409 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "generator/generators.h" +#include "models/model.h" +#include "models/preprocessing/genai_tokenizer.h" +#include "models/preprocessing/lfm2_audio_processor.h" +#include "nemo_mel_spectrogram.h" + +#include +#include +#include +#include + +namespace Generators { + +namespace { + +// The list overload of MultiModalProcessor::Process fills payload.prompts and leaves payload.prompt +// empty. Batching is not supported, so a single-entry list is the prompt and anything longer is an error +// rather than silently dropped text. +std::string ResolvePrompt(const Payload& payload) { + if (payload.prompts.empty()) { + return payload.prompt; + } + if (payload.prompts.size() != 1) { + throw std::runtime_error("Lfm2AudioProcessor: batched prompts are not supported; got " + + std::to_string(payload.prompts.size()) + " prompts. Pass a single prompt string."); + } + return payload.prompts[0] ? std::string(payload.prompts[0]) : std::string{}; +} + +// The encoder graph's input names come from genai_config.json; a wrong name would otherwise be +// dropped silently by ExtraInputs and only surface as ORT's "Missing Input" at run time. +ONNXTensorElementDataType SpeechInputType(const SessionInfo& session_info, const std::string& name, + const char* config_field) { + if (!session_info.HasInput(name)) { + throw std::runtime_error("Lfm2AudioProcessor: the speech model has no input named \"" + name + + "\". Point model.speech.inputs." + config_field + + " in genai_config.json at the name the speech model uses."); + } + return session_info.GetInputDataType(name); +} + +std::unique_ptr MakeInt64Tensor(const std::vector& values, std::vector shape, + Ort::Allocator& allocator) { + auto tensor = OrtValue::CreateTensor(allocator, shape); + std::copy(values.begin(), values.end(), tensor->GetTensorMutableData()); + return tensor; +} + +std::unique_ptr MakeInputIds(const std::vector& input_ids, Ort::Allocator& allocator) { + auto tensor = OrtValue::CreateTensor(allocator, std::vector{1, static_cast(input_ids.size())}); + std::copy(input_ids.begin(), input_ids.end(), tensor->GetTensorMutableData()); + return tensor; +} + +// torch.hann_window(win_length, periodic=False), the window NeMo's FilterbankFeatures registers. +std::vector SymmetricHannWindow(int win_length) { + constexpr double kPi = 3.14159265358979323846; + std::vector window(static_cast(win_length)); + for (int n = 0; n < win_length; ++n) { + window[static_cast(n)] = static_cast(0.5 * (1.0 - std::cos(2.0 * kPi * n / (win_length - 1)))); + } + return window; +} + +} // namespace + +int64_t Lfm2AudioNumMelFrames(int64_t num_samples, const Lfm2AudioMelConfig& config) { + // torch.stft(center=True): fft_size / 2 zeros on both sides, one frame per hop_length. + const int64_t padded = num_samples + 2 * (config.fft_size / 2); + if (padded < config.fft_size) { + return 0; + } + return (padded - config.fft_size) / config.hop_length + 1; +} + +int64_t Lfm2AudioNumValidMelFrames(int64_t num_samples, const Lfm2AudioMelConfig& config) { + // FilterbankFeatures.get_seq_len: (num_samples + 2 * (fft_size // 2) - fft_size) // hop_length. + const int64_t padded = num_samples + 2 * (config.fft_size / 2); + if (padded < config.fft_size) { + return 0; + } + return (padded - config.fft_size) / config.hop_length; +} + +int64_t Lfm2AudioNumTokens(int64_t num_mel_frames, int64_t subsampling_factor) { + if (subsampling_factor <= 0) { + throw std::runtime_error("Lfm2AudioProcessor: subsampling_factor must be positive."); + } + return (num_mel_frames + subsampling_factor - 1) / subsampling_factor; +} + +std::vector ResampleLfm2Audio(const float* pcm, int64_t num_samples, int64_t source_rate, int64_t target_rate) { + if (source_rate <= 0 || target_rate <= 0) { + throw std::runtime_error("Lfm2AudioProcessor: cannot resample from " + std::to_string(source_rate) + " Hz to " + + std::to_string(target_rate) + " Hz."); + } + if (source_rate == target_rate) { + return {pcm, pcm + num_samples}; + } + + constexpr double kPi = 3.14159265358979323846; + constexpr double kLowpassFilterWidth = 6.0; + constexpr double kRolloff = 0.99; + const int64_t divisor = std::gcd(source_rate, target_rate); + const int64_t orig = source_rate / divisor; + const int64_t target = target_rate / divisor; + const double base_freq = static_cast(std::min(orig, target)) * kRolloff; + const int64_t width = static_cast(std::ceil(kLowpassFilterWidth * static_cast(orig) / base_freq)); + const int64_t kernel_size = 2 * width + orig; + + // One kernel per output phase, built in double and stored in float as torchaudio does. + std::vector kernels(static_cast(target * kernel_size)); + for (int64_t phase = 0; phase < target; ++phase) { + for (int64_t k = 0; k < kernel_size; ++k) { + double t = (-static_cast(phase) / static_cast(target) + + static_cast(k - width) / static_cast(orig)) * + base_freq; + t = std::clamp(t, -kLowpassFilterWidth, kLowpassFilterWidth); + const double window = std::pow(std::cos(t * kPi / kLowpassFilterWidth / 2.0), 2.0); + t *= kPi; + const double sinc = t == 0.0 ? 1.0 : std::sin(t) / t; + kernels[static_cast(phase * kernel_size + k)] = + static_cast(sinc * window * base_freq / static_cast(orig)); + } + } + + std::vector padded(static_cast(width + num_samples + width + orig), 0.0f); + std::copy(pcm, pcm + num_samples, padded.begin() + width); + + // ceil(target * num_samples / orig) samples; the strided convolution makes a few more, which + // torchaudio trims. + const int64_t num_outputs = (target * num_samples + orig - 1) / orig; + std::vector resampled(static_cast(num_outputs)); + for (int64_t i = 0; i < num_outputs; ++i) { + const float* frame = padded.data() + (i / target) * orig; + const float* kernel = kernels.data() + (i % target) * kernel_size; + float value = 0.0f; + for (int64_t k = 0; k < kernel_size; ++k) { + value += frame[k] * kernel[k]; + } + resampled[static_cast(i)] = value; + } + return resampled; +} + +std::vector ComputeLfm2AudioMel(const float* pcm, int64_t num_samples, const Lfm2AudioMelConfig& config, + int64_t& num_frames) { + if (config.num_mels <= 0 || config.fft_size <= 0 || config.hop_length <= 0 || config.win_length < 2 || + config.win_length > config.fft_size) { + throw std::runtime_error("Lfm2AudioProcessor: invalid mel configuration (num_mels " + std::to_string(config.num_mels) + + ", fft_size " + std::to_string(config.fft_size) + ", win_length " + + std::to_string(config.win_length) + ", hop_length " + std::to_string(config.hop_length) + ")."); + } + + num_frames = Lfm2AudioNumMelFrames(num_samples, config); + const int64_t valid_frames = Lfm2AudioNumValidMelFrames(num_samples, config); + // The per-feature normalization divides by valid_frames - 1; NeMo raises on a single frame and + // produces NaNs on none, so refuse both up front with the clip length that would be needed. + if (valid_frames < 2) { + throw std::runtime_error("Lfm2AudioProcessor: audio clip is too short (" + std::to_string(num_samples) + + " samples at " + std::to_string(config.sample_rate) + " Hz); at least " + + std::to_string(2 * config.hop_length) + " samples are needed."); + } + + const auto mel_filters = nemo_mel::CreateMelFilterbank(config.num_mels, config.fft_size, config.sample_rate); + const auto window = SymmetricHannWindow(config.win_length); + + std::vector preemphasized(static_cast(num_samples)); + nemo_mel::ApplyPreemphasis(pcm, static_cast(num_samples), config.preemph, 0.0f, preemphasized.data()); + + // Constant (zero) padding on both sides: NeMo's stft uses pad_mode="constant". + const int64_t pad = config.fft_size / 2; + std::vector padded(static_cast(pad + num_samples + pad), 0.0f); + std::copy(preemphasized.begin(), preemphasized.end(), padded.begin() + pad); + + // The win_length window sits centered inside the fft_size frame; the power spectrum does not + // depend on where the windowed samples land within the FFT buffer. + const int64_t win_offset = (config.fft_size - config.win_length) / 2; + const int num_bins = config.fft_size / 2 + 1; + const auto num_mels = static_cast(config.num_mels); + + std::vector mel(static_cast(num_frames) * num_mels); + std::vector magnitudes; + for (int64_t t = 0; t < num_frames; ++t) { + const float* frame = padded.data() + t * config.hop_length + win_offset; + nemo_mel::ComputeSTFTFrame(frame, window.data(), config.win_length, config.fft_size, magnitudes); + float* row = mel.data() + static_cast(t) * num_mels; + for (size_t m = 0; m < num_mels; ++m) { + float value = 0.0f; + for (int k = 0; k < num_bins; ++k) { + value += mel_filters[m][static_cast(k)] * magnitudes[static_cast(k)]; + } + row[m] = std::log(value + config.log_eps); + } + } + + // normalize_batch(normalize_type="per_feature"): mean and unbiased standard deviation over the + // valid frames only, then the frames past the valid length are set to pad_value (0). + for (size_t m = 0; m < num_mels; ++m) { + double sum = 0.0; + for (int64_t t = 0; t < valid_frames; ++t) { + sum += mel[static_cast(t) * num_mels + m]; + } + const float mean = static_cast(sum / static_cast(valid_frames)); + double variance = 0.0; + for (int64_t t = 0; t < valid_frames; ++t) { + const double delta = mel[static_cast(t) * num_mels + m] - mean; + variance += delta * delta; + } + const float std = static_cast(std::sqrt(variance / static_cast(valid_frames - 1))) + config.norm_eps; + for (int64_t t = 0; t < valid_frames; ++t) { + float& value = mel[static_cast(t) * num_mels + m]; + value = (value - mean) / std; + } + } + std::fill(mel.begin() + static_cast(valid_frames) * num_mels, mel.end(), 0.0f); + return mel; +} + +std::vector SplitLfm2AudioPrompt(const std::string& prompt, size_t num_clips) { + const std::string marker{kLfm2AudioMarker}; + std::vector segments; + size_t position = 0; + for (size_t match = prompt.find(marker, position); match != std::string::npos; match = prompt.find(marker, position)) { + segments.push_back(prompt.substr(position, match - position)); + position = match + marker.size(); + } + segments.push_back(prompt.substr(position)); + + // Same rule as the LFM2-VL processor applies to : one marker per clip, in clip order. + // Silently inventing or dropping placeholders would misalign the encoder features with the + // decoder positions. + const size_t num_markers = segments.size() - 1; + if (num_markers != num_clips) { + throw std::runtime_error("Lfm2AudioProcessor: the prompt contains " + std::to_string(num_markers) + " " + marker + + " markers but " + std::to_string(num_clips) + " audio clips were provided. Put exactly one " + + marker + " in the prompt per clip."); + } + return segments; +} + +Lfm2AudioProcessor::Lfm2AudioProcessor(Config& config, const SessionInfo& session_info) + : mel_type_{SpeechInputType(session_info, config.model.speech.inputs.audio_embeds, "audio_embeds")}, + subsampling_factor_{config.model.subsampling_factor > 0 ? config.model.subsampling_factor : 8}, + audio_token_id_{config.model.audio_token_id} { + if (mel_type_ != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT && mel_type_ != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) { + throw std::runtime_error("Lfm2AudioProcessor: speech input \"" + config.model.speech.inputs.audio_embeds + + "\" must be float or float16, got " + TypeToString(mel_type_) + "."); + } + // The clip lengths are always emitted as int64, so fail at load rather than at the first clip if + // the graph was exported with another integer type. + const auto lengths_type = SpeechInputType(session_info, config.model.speech.inputs.audio_lengths, "audio_lengths"); + if (lengths_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) { + throw std::runtime_error("Lfm2AudioProcessor: speech input \"" + config.model.speech.inputs.audio_lengths + + "\" must be int64, got " + TypeToString(lengths_type) + "."); + } + if (audio_token_id_ <= 0) { + throw std::runtime_error( + "Lfm2AudioProcessor: model.audio_token_id must be set in genai_config.json to the id " + "the embedding model replaces with the audio features."); + } + if (audio_token_id_ >= config.model.vocab_size) { + throw std::runtime_error("Lfm2AudioProcessor: model.audio_token_id " + std::to_string(audio_token_id_) + + " is outside the vocabulary of " + std::to_string(config.model.vocab_size) + " tokens."); + } + + // The mel front end defaults to the settings of the published checkpoints; a genai_config.json can + // override any of them with the same top-level fields the Parakeet and Nemotron models use. + const auto& m = config.model; + if (m.num_mels > 0) mel_config_.num_mels = m.num_mels; + if (m.fft_size > 0) mel_config_.fft_size = m.fft_size; + if (m.win_length > 0) mel_config_.win_length = m.win_length; + if (m.hop_length > 0) mel_config_.hop_length = m.hop_length; + if (m.sample_rate > 0) mel_config_.sample_rate = m.sample_rate; + if (m.preemph > 0.0f) mel_config_.preemph = m.preemph; + if (m.log_eps > 0.0f) mel_config_.log_eps = m.log_eps; + if (m.norm_eps > 0.0f) mel_config_.norm_eps = m.norm_eps; + + config.AddMapping(std::string(Config::Defaults::InputIdsName), config.model.embedding.inputs.input_ids); + config.AddMapping(std::string(Config::Defaults::AudioEmbedsName), config.model.speech.inputs.audio_embeds); + config.AddMapping(std::string(Config::Defaults::AudioLengthsName), config.model.speech.inputs.audio_lengths); + config.AddMapping(std::string(Config::Defaults::AudioSizesName), config.model.speech.inputs.audio_sizes); +} + +Lfm2AudioProcessor::ClipMel Lfm2AudioProcessor::ComputeClipMel(const Audios& audios, size_t index) const { + // Decoded at the file's own rate and resampled here, the way the reference does it: the decoder's + // resampler is a different filter, enough to change a token, and it cannot resample upwards. + ort_extensions::OrtxObjectPtr decoded; + CheckResult(OrtxDecodeAudio(audios.audios_.get(), index, /*target_sample_rate=*/0, /*stereo_to_mono=*/1, + decoded.ToBeAssigned())); + + ort_extensions::OrtxObjectPtr pcm_tensor; + CheckResult(OrtxTensorResultGetAt(decoded.get(), 0, pcm_tensor.ToBeAssigned())); + const float* pcm{}; + const int64_t* pcm_shape{}; + size_t pcm_dims{}; + CheckResult(OrtxGetTensorData(pcm_tensor.get(), reinterpret_cast(&pcm), &pcm_shape, &pcm_dims)); + + int64_t num_samples = 0; + if (pcm_dims == 1) { + num_samples = pcm_shape[0]; + } else if (pcm_dims == 2 && pcm_shape[0] == 1) { + num_samples = pcm_shape[1]; + } else { + throw std::runtime_error("Lfm2AudioProcessor: expected mono PCM from the audio decoder for clip " + + std::to_string(index) + ", got a rank " + std::to_string(pcm_dims) + " tensor."); + } + + ort_extensions::OrtxObjectPtr rate_tensor; + CheckResult(OrtxTensorResultGetAt(decoded.get(), 1, rate_tensor.ToBeAssigned())); + const int64_t* source_rate{}; + CheckResult(OrtxGetTensorData(rate_tensor.get(), reinterpret_cast(&source_rate), nullptr, nullptr)); + + const std::vector samples = ResampleLfm2Audio(pcm, num_samples, *source_rate, mel_config_.sample_rate); + + ClipMel clip; + clip.mel = ComputeLfm2AudioMel(samples.data(), static_cast(samples.size()), mel_config_, clip.num_frames); + clip.num_tokens = Lfm2AudioNumTokens(clip.num_frames, subsampling_factor_); + return clip; +} + +std::vector Lfm2AudioProcessor::MakePromptTokens(const Tokenizer& tokenizer, + const std::vector& segments, + const std::vector& clips) const { + // Text segments are tokenized on their own, as the reference ChatState does, and each clip + // contributes one placeholder per encoder frame between them. + std::vector input_ids; + for (size_t i = 0; i < segments.size(); ++i) { + if (!segments[i].empty()) { + const std::vector ids = tokenizer.Encode(segments[i].c_str()); + input_ids.insert(input_ids.end(), ids.begin(), ids.end()); + } + if (i < clips.size()) { + input_ids.insert(input_ids.end(), static_cast(clips[i].num_tokens), audio_token_id_); + } + } + return input_ids; +} + +void Lfm2AudioProcessor::AddAudioTensors(const std::vector& clips, Ort::Allocator& allocator, + NamedTensors& named_tensors) const { + const int64_t num_clips = static_cast(clips.size()); + const int64_t num_mels = mel_config_.num_mels; + int64_t longest = 0; + std::vector mel_lengths; + std::vector tokens_per_clip; + for (const auto& clip : clips) { + longest = std::max(longest, clip.num_frames); + mel_lengths.push_back(clip.num_frames); + tokens_per_clip.push_back(clip.num_tokens); + } + + // The clips are staged in one zero-padded [num_clips, longest, num_mels] tensor with their real + // lengths alongside. It is a container, not a batched encoder run: Lfm2AudioSpeechState slices + // each clip's own frames back out and runs the encoder on them one clip at a time. + auto batch = OrtValue::CreateTensor(allocator, std::vector{num_clips, longest, num_mels}); + float* batch_data = batch->GetTensorMutableData(); + std::fill_n(batch_data, static_cast(num_clips) * static_cast(longest * num_mels), 0.0f); + for (size_t i = 0; i < clips.size(); ++i) { + std::copy(clips[i].mel.begin(), clips[i].mel.end(), batch_data + i * static_cast(longest * num_mels)); + } + + if (mel_type_ == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + named_tensors.emplace(std::string(Config::Defaults::AudioEmbedsName), std::make_shared(std::move(batch))); + } else { + std::unique_ptr converted; + Cast(*batch, converted, *GetDeviceInterface(DeviceType::CPU), mel_type_); + named_tensors.emplace(std::string(Config::Defaults::AudioEmbedsName), std::make_shared(std::move(converted))); + } + named_tensors.emplace(std::string(Config::Defaults::AudioLengthsName), + std::make_shared(MakeInt64Tensor(mel_lengths, {num_clips}, allocator))); + named_tensors.emplace(std::string(Config::Defaults::AudioSizesName), + std::make_shared(MakeInt64Tensor(tokens_per_clip, {num_clips}, allocator))); +} + +std::unique_ptr Lfm2AudioProcessor::Process(const Tokenizer& tokenizer, const Payload& payload) const { + const std::string prompt = ResolvePrompt(payload); + const Audios* audios = payload.audios; + const size_t num_clips = audios ? audios->num_audios_ : 0; + Ort::Allocator& allocator{Ort::Allocator::GetWithDefaultOptions()}; + auto named_tensors = std::make_unique(); + + // With no clips the prompt must not ask for any; the split with zero clips enforces that. + const std::vector segments = SplitLfm2AudioPrompt(prompt, num_clips); + + std::vector clips; + clips.reserve(num_clips); + for (size_t i = 0; i < num_clips; ++i) { + clips.push_back(ComputeClipMel(*audios, i)); + } + + named_tensors->emplace(std::string(Config::Defaults::InputIdsName), + std::make_shared(MakeInputIds(MakePromptTokens(tokenizer, segments, clips), allocator))); + + if (clips.empty()) { + // The pipeline reads audio_sizes to skip the speech run. + named_tensors->emplace(std::string(Config::Defaults::AudioSizesName), + std::make_shared(MakeInt64Tensor({0}, {1}, allocator))); + return named_tensors; + } + + AddAudioTensors(clips, allocator, *named_tensors); + return named_tensors; +} + +} // namespace Generators diff --git a/src/models/preprocessing/lfm2_audio_processor.h b/src/models/preprocessing/lfm2_audio_processor.h new file mode 100644 index 0000000000..e42d4d611a --- /dev/null +++ b/src/models/preprocessing/lfm2_audio_processor.h @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include + +#include "models/preprocessing/processor.h" + +namespace Generators { + +// Marker the prompt carries once per audio clip, in clip order. LFM2-Audio has no audio input token +// of its own: the reference implementation (liquid-audio's ChatState) splices the encoder output +// between the text tokens by position, so the marker is a convention of this runtime, and the +// processor replaces it with model.audio_token_id repeated once per encoder frame. +inline constexpr char kLfm2AudioMarker[] = "<|audio|>"; + +// Log-mel front end of the FastConformer encoder, NeMo's AudioToMelSpectrogramPreprocessor with +// the settings of the published checkpoints (config.json "preprocessor"). +struct Lfm2AudioMelConfig { + int num_mels{128}; + int fft_size{512}; + int win_length{400}; // window_size 0.025 s at 16 kHz + int hop_length{160}; // window_stride 0.01 s at 16 kHz + int sample_rate{16000}; + float preemph{0.97f}; + float log_eps{5.960464477539063e-08f}; // log_zero_guard_value 2^-24, added before the log + float norm_eps{1e-5f}; // added to the per-feature standard deviation +}; + +// Frame accounting of the reference preprocessor. torch.stft with center=True yields one frame per +// hop plus one, but NeMo's get_seq_len counts only num_samples / hop_length of them as valid: the +// per-feature normalization uses the valid frames and the rest are zeroed, yet every frame is still +// handed to the encoder and counted towards the decoder tokens. +int64_t Lfm2AudioNumMelFrames(int64_t num_samples, const Lfm2AudioMelConfig& config); +int64_t Lfm2AudioNumValidMelFrames(int64_t num_samples, const Lfm2AudioMelConfig& config); + +// Decoder tokens one clip contributes: the encoder subsamples the mel frames by subsampling_factor +// with ceil rounding (three stride-2 convolutions with padding 1 for the shipped models). +int64_t Lfm2AudioNumTokens(int64_t num_mel_frames, int64_t subsampling_factor); + +// torchaudio.functional.resample with its defaults (Hann-windowed sinc, lowpass_filter_width 6, +// rolloff 0.99), which is what the reference's ChatState.add_audio brings every clip to the +// encoder's rate with, upwards as well as downwards. +std::vector ResampleLfm2Audio(const float* pcm, int64_t num_samples, int64_t source_rate, int64_t target_rate); + +// Computes the normalized log-mel spectrogram of one mono clip, time-major [num_frames, num_mels], +// the layout the encoder's mel_spectrogram input takes. Throws if the clip is too short to +// normalize (fewer than two valid frames). +std::vector ComputeLfm2AudioMel(const float* pcm, int64_t num_samples, const Lfm2AudioMelConfig& config, + int64_t& num_frames); + +// Splits the prompt at each marker. Throws unless the prompt holds exactly one marker per clip, the +// rule the vision processor applies to , so the encoder features and the placeholders can +// never be misaligned or silently dropped. Returns num_clips + 1 text segments. +std::vector SplitLfm2AudioPrompt(const std::string& prompt, size_t num_clips); + +struct Lfm2AudioProcessor : Processor { + Lfm2AudioProcessor(Config& config, const SessionInfo& session_info); + + std::unique_ptr Process(const Tokenizer& tokenizer, const Payload& payload) const override; + + private: + // One clip's log-mel frames, with what the rest of the pipeline has to know about them. + struct ClipMel { + std::vector mel; // num_frames * num_mels, frame-major + int64_t num_frames{}; // what the encoder is given for this clip + int64_t num_tokens{}; // placeholders it is worth, one per encoder frame + }; + + // Decodes clip `index` to mono PCM at the encoder's sample rate and turns it into log-mel frames. + ClipMel ComputeClipMel(const Audios& audios, size_t index) const; + + // The prompt's text tokenized segment by segment, with each clip's placeholders spliced between. + std::vector MakePromptTokens(const Tokenizer& tokenizer, const std::vector& segments, + const std::vector& clips) const; + + // Stages the clips in one zero-padded [num_clips, longest, num_mels] tensor, cast to the speech + // model's input type, and adds it with the per-clip frame and placeholder counts. + void AddAudioTensors(const std::vector& clips, Ort::Allocator& allocator, NamedTensors& named_tensors) const; + + ONNXTensorElementDataType mel_type_; + Lfm2AudioMelConfig mel_config_; + int64_t subsampling_factor_{}; + int32_t audio_token_id_{}; +}; + +} // namespace Generators diff --git a/src/models/preprocessing/multi_modal_processor.cpp b/src/models/preprocessing/multi_modal_processor.cpp index 45b31646e0..a0ae8443fd 100644 --- a/src/models/preprocessing/multi_modal_processor.cpp +++ b/src/models/preprocessing/multi_modal_processor.cpp @@ -7,6 +7,7 @@ #include "models/preprocessing/gemma4_multimodal_processor.h" #include "models/preprocessing/gemma_image_processor.h" #include "models/preprocessing/genai_tokenizer.h" +#include "models/preprocessing/lfm2_audio_processor.h" #include "models/preprocessing/lfm2_vl_image_processor.h" #include "models/preprocessing/mistral3_image_processor.h" #include "models/preprocessing/parakeet_processor.h" @@ -32,6 +33,7 @@ MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& sess {"parakeet_tdt", Processor::Create}, {"phi4mm", Processor::Create}, {"gemma3", Processor::Create}, + {"lfm2_audio", Processor::Create}, {"lfm2_vl", Processor::Create}, {"gemma4", Processor::Create}, {"mistral3", Processor::Create}, diff --git a/src/models/utils.cpp b/src/models/utils.cpp index 60c1f86bbf..0a270664c3 100644 --- a/src/models/utils.cpp +++ b/src/models/utils.cpp @@ -10,6 +10,11 @@ DeviceSpan ByteWrapTensor(DeviceInterface& device, OrtValue& value) { return device.WrapMemory(std::span{value.GetTensorMutableData(), info->GetElementCount() * Ort::SizeOf(info->GetElementType())}); } +DeviceInterface& DeviceForTensor(const OrtValue& value, DeviceInterface& device) { + const bool on_cpu = value.GetTensorMemoryInfo().GetDeviceType() == OrtMemoryInfoDeviceType_CPU; + return on_cpu ? *GetDeviceInterface(DeviceType::CPU) : device; +} + const char* TypeToString(ONNXTensorElementDataType type) { switch (type) { case Ort::TypeToTensorType: diff --git a/src/python/py/models/README.md b/src/python/py/models/README.md index 1d8d081d5c..94d26904fc 100644 --- a/src/python/py/models/README.md +++ b/src/python/py/models/README.md @@ -74,7 +74,7 @@ The tool currently supports the following model architectures. - Granite MoE Hybrid - HunYuan Dense V1 - InternLM2 -- LFM2 (text and the decoder of LFM2-VL / LFM2.5-VL) +- LFM2 (text and the decoders of LFM2-VL / LFM2.5-VL and LFM2-Audio / LFM2.5-Audio) - LFM2 MoE - Llama - Mistral diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 703ebb192a..f2b5ef7c09 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -28,6 +28,7 @@ GraniteMoEHybridModel, HunyuanDenseV1Model, InternLM2Model, + LFM2AudioModel, LFM2Model, LFM2MoEModel, LlamaModel, @@ -101,7 +102,13 @@ def get_hf_details(model_name, input_path, cache_dir, extra_options): hf_token = extra_options.get("hf_token", True) hf_remote = extra_options.get("hf_remote", False) - config = AutoConfig.from_pretrained(hf_name, token=hf_token, trust_remote_code=hf_remote, **extra_kwargs) + try: + config = AutoConfig.from_pretrained(hf_name, token=hf_token, trust_remote_code=hf_remote, **extra_kwargs) + except ValueError: + # LFM2-Audio checkpoints have no model_type; their LFM2 decoder config is nested. + config = LFM2AudioModel.load_config(hf_name, token=hf_token, **extra_kwargs) + if config is None: + raise tokenizer = AutoTokenizer.from_pretrained(hf_name, token=hf_token, trust_remote_code=hf_remote, **extra_kwargs) add_special_token_ids(config, tokenizer) if extra_options.get("adapter_path", False): @@ -628,6 +635,10 @@ def create_model( # With the embedding layer excluded the decoder is one stage of the LFM2-VL vision pipeline; # otherwise it is a standalone text model that happens to come from a VLM checkpoint. onnx_model.model_type = "lfm2_vl" if onnx_model.exclude_embeds else "lfm2_vl_text" + elif config.architectures[0] == "Lfm2AudioForConditionalGeneration": + onnx_model = LFM2AudioModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) + # Same split as LFM2-VL: the decoder stage of the lfm2_audio pipeline, or a plain text model. + onnx_model.model_type = "lfm2_audio" if onnx_model.exclude_embeds else "lfm2_audio_text" elif config.architectures[0] == "LlamaForCausalLM": onnx_model = LlamaModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) elif config.architectures[0] == "MistralForCausalLM": diff --git a/src/python/py/models/builders/__init__.py b/src/python/py/models/builders/__init__.py index 0925693898..b5bf8df3a6 100644 --- a/src/python/py/models/builders/__init__.py +++ b/src/python/py/models/builders/__init__.py @@ -14,7 +14,7 @@ from .granite import GraniteModel, GraniteMoEHybridModel from .hunyuan import HunyuanDenseV1Model from .internlm import InternLM2Model -from .lfm2 import LFM2Model, LFM2MoEModel +from .lfm2 import LFM2AudioModel, LFM2Model, LFM2MoEModel from .llama import LlamaModel from .mistral import Mistral3TextModel, MistralModel from .mtp import MTPModel @@ -57,6 +57,7 @@ "GraniteModel", "HunyuanDenseV1Model", "InternLM2Model", + "LFM2AudioModel", "LFM2Model", "LFM2MoEModel", "LlamaModel", diff --git a/src/python/py/models/builders/lfm2.py b/src/python/py/models/builders/lfm2.py index c7e4dce6e0..f0e053296a 100644 --- a/src/python/py/models/builders/lfm2.py +++ b/src/python/py/models/builders/lfm2.py @@ -4,7 +4,11 @@ # license information. # Modifications Copyright(C) 2026 Advanced Micro Devices, Inc. All rights reserved. # -------------------------------------------------------------------------- +import json +import os + import onnx_ir as ir +import torch from .base import Model @@ -190,6 +194,110 @@ def update_genai_config(self, genai_config): decoder["conv_cache_size"] = self.conv_L_cache - 1 +class LFM2AudioModel(LFM2Model): + """The LFM2 decoder of an LFM2-Audio checkpoint, exported for text generation. + + The checkpoint stores the decoder as an Lfm2Model under the "lfm." prefix, with the logits tied + to the token embeddings; the audio encoder, depthformer and audio embeddings that surround it in + the checkpoint are left out. With the embedding layer excluded the export is the decoder stage of + the lfm2_audio pipeline, whose embedding model splices the audio encoder output into the token + embeddings; otherwise it is a plain text model. + """ + + @classmethod + def load_config(cls, model_name_or_path, **kwargs): + """The LFM2 decoder config of an LFM2-Audio checkpoint, or None if it is not one. + + These checkpoints have no model type and no transformers model class, and their config.json + nests the decoder config under "lfm" next to the audio encoder, depthformer and mel front-end + settings, so AutoConfig cannot read them. The builder calls this once AutoConfig has refused + the checkpoint and before it knows which model class to use, which is why it cannot be an + instance method. + """ + from transformers import Lfm2Config, PretrainedConfig + + architecture = "Lfm2AudioForConditionalGeneration" + config_dict, _ = PretrainedConfig.get_config_dict(model_name_or_path, **kwargs) + if config_dict.get("architectures") != [architecture]: + return None + + config = Lfm2Config(**config_dict["lfm"]) + config.architectures = [architecture] + config.name_or_path = model_name_or_path + return config + + def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): + super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) + self.decoder_config = config + + def update_genai_config(self, genai_config): + super().update_genai_config(genai_config) + + # These checkpoints answer in speech as well as text, and two tokens mark the turn passing + # from one to the other: <|audio_start|> when the whole answer is spoken, and <|text_end|> + # when text and speech are interleaved. Everything after either is audio codes for the + # depthformer, which this runtime does not have, and sampling those positions from the text + # head gives fluent nonsense. Stop on them instead, keeping whatever text came first. + modality_switch_token_ids = [128, 130] # <|audio_start|>, <|text_end|> + eos_token_id = genai_config["model"].get("eos_token_id") + eos_token_ids = list(eos_token_id) if isinstance(eos_token_id, list) else [eos_token_id] + eos_token_ids += [token for token in modality_switch_token_ids if token not in eos_token_ids] + genai_config["model"]["eos_token_id"] = eos_token_ids + + def load_weights(self, input_path): + # A fine-tune of these checkpoints names the decoder "lfm.*", which matches nothing in the + # Lfm2ForCausalLM loaded here: PEFT only warns, and the export would carry untrained LoRA weights. + if "adapter_path" in self.extra_options: + raise ValueError("LoRA adapters are not supported for LFM2-Audio checkpoints; merge the adapter first.") + + if input_path.endswith(".gguf") or self.quant_type is not None: + return super().load_weights(input_path) + + from transformers import Lfm2ForCausalLM + + with torch.device("meta"): + model = Lfm2ForCausalLM(self.decoder_config) + state_dict = {} + for name, tensor in self.load_lfm2_audio_checkpoint(input_path).items(): + if name.startswith("lfm."): + state_dict["model." + name[len("lfm.") :]] = tensor + missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True) + if unexpected or set(missing) - {"lm_head.weight"}: + raise ValueError( + f"The LFM2 decoder of {input_path} does not match its config: missing {sorted(missing)}, unexpected {sorted(unexpected)}." + ) + model.tie_weights() + return model.eval() + + def load_lfm2_audio_checkpoint(self, input_path): + """The checkpoint tensors, from a local directory or the Hugging Face Hub.""" + from safetensors.torch import load_file + + # `-i` gives a local directory. With `-m` it is empty and the checkpoint is the one the + # config was read from, which is the repository name. + source = input_path or self.model_name_or_path + if os.path.isdir(source): + checkpoint_dir = source + else: + from huggingface_hub import snapshot_download + + checkpoint_dir = snapshot_download( + source, cache_dir=self.cache_dir, token=self.hf_token, allow_patterns=["*.safetensors", "*.json"] + ) + + index_path = os.path.join(checkpoint_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path) as f: + shards = sorted(set(json.load(f)["weight_map"].values())) + else: + shards = ["model.safetensors"] + + tensors = {} + for shard in shards: + tensors.update(load_file(os.path.join(checkpoint_dir, shard))) + return tensors + + class LFM2MoEModel(LFM2Model): """LFM2-MoE builder (LFM2-8B-A1B, LFM2.5-8B-A1B, LFM2-24B-A2B). diff --git a/src/python/python.cpp b/src/python/python.cpp index 3f2885794a..c982cc2820 100644 --- a/src/python/python.cpp +++ b/src/python/python.cpp @@ -675,12 +675,19 @@ PYBIND11_MODULE(onnxruntime_genai, m) { pybind11::class_(m, "Images") .def_static("open", [](pybind11::args image_paths) { std::vector image_paths_string; - std::vector image_paths_vector; + image_paths_string.reserve(image_paths.size()); for (auto image_path : image_paths) { if (!pybind11::isinstance(image_path)) throw std::runtime_error("Image paths must be strings."); image_paths_string.push_back(image_path.cast()); - image_paths_vector.push_back(image_paths_string.back().c_str()); + } + + // Take the pointers only once the strings have stopped moving: a short string keeps its + // characters inside the string object, so growing the vector would dangle them. + std::vector image_paths_vector; + image_paths_vector.reserve(image_paths_string.size()); + for (const auto& image_path : image_paths_string) { + image_paths_vector.push_back(image_path.c_str()); } return OgaImages::Load(image_paths_vector); @@ -703,13 +710,19 @@ PYBIND11_MODULE(onnxruntime_genai, m) { pybind11::class_(m, "Audios") .def_static("open", [](pybind11::args audio_paths) { std::vector audio_paths_string; - std::vector audio_paths_vector; - + audio_paths_string.reserve(audio_paths.size()); for (const auto& audio_path : audio_paths) { if (!pybind11::isinstance(audio_path)) throw std::runtime_error("Audio paths must be strings."); audio_paths_string.push_back(audio_path.cast()); - audio_paths_vector.push_back(audio_paths_string.back().c_str()); + } + + // Take the pointers only once the strings have stopped moving: a short string keeps its + // characters inside the string object, so growing the vector would dangle them. + std::vector audio_paths_vector; + audio_paths_vector.reserve(audio_paths_string.size()); + for (const auto& audio_path : audio_paths_string) { + audio_paths_vector.push_back(audio_path.c_str()); } return OgaAudios::Load(audio_paths_vector); diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index f5b5c5345c..65e468535f 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -15,6 +15,7 @@ file(GLOB test_srcs CONFIGURE_DEPENDS # Keep tests with dedicated white-box executables out of the public-API unit_tests target. list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/dflash2_config_test.cpp") list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/kv_cache_tests.cpp") +list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/lfm2_audio_frames_test.cpp") list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/lfm2_vl_image_geometry_test.cpp") list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/reinit_tests.cpp") list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/search_checkpoint_tests.cpp") @@ -202,6 +203,7 @@ file(GLOB engine_unit_test_srcs CONFIGURE_DEPENDS ) list(APPEND engine_unit_test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/dflash2_config_test.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/cpp/lfm2_audio_frames_test.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/cpp/lfm2_vl_image_geometry_test.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/cpp/search_checkpoint_tests.cpp" ) diff --git a/test/cpp/lfm2_audio_frames_test.cpp b/test/cpp/lfm2_audio_frames_test.cpp new file mode 100644 index 0000000000..16672a3848 --- /dev/null +++ b/test/cpp/lfm2_audio_frames_test.cpp @@ -0,0 +1,255 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include + +#include "generator/generators.h" +#include "models/multi_modal.h" +#include "models/preprocessing/lfm2_audio_processor.h" + +// LFM2-Audio turns a clip into log-mel frames, one per 10 ms hop, and the encoder subsamples those +// by 8, so the decoder sees one token per 80 ms. These tests pin the three things that have to agree +// for the pipeline to line up: how many frames a clip is worth, how many of them the reference +// preprocessor counts as valid, and how the prompt is split around the audio markers. + +namespace Generators::test { +namespace { + +constexpr Lfm2AudioMelConfig kMelConfig{}; // settings of the published checkpoints +constexpr int64_t kSubsamplingFactor = 8; + +template +std::string CaptureThrowMessage(Fn&& fn) { + try { + fn(); + } catch (const std::exception& e) { + return e.what(); + } + return {}; +} + +} // namespace + +TEST(Lfm2AudioFramesTest, OneSecondIsOneHundredFramesPlusTheCenterPadOne) { + // torch.stft(center=True) pads by fft_size / 2 on both sides, which buys exactly one extra frame. + EXPECT_EQ(Lfm2AudioNumMelFrames(16000, kMelConfig), 101); + // NeMo's get_seq_len counts num_samples / hop_length instead, one fewer. + EXPECT_EQ(Lfm2AudioNumValidMelFrames(16000, kMelConfig), 100); +} + +TEST(Lfm2AudioFramesTest, PartialHopsAreTruncated) { + EXPECT_EQ(Lfm2AudioNumMelFrames(16159, kMelConfig), 101); + EXPECT_EQ(Lfm2AudioNumMelFrames(16160, kMelConfig), 102); + EXPECT_EQ(Lfm2AudioNumValidMelFrames(16159, kMelConfig), 100); + EXPECT_EQ(Lfm2AudioNumValidMelFrames(16160, kMelConfig), 101); +} + +TEST(Lfm2AudioFramesTest, ClipShorterThanTheWindowStillHasFrames) { + // The center padding alone is fft_size samples wide, so even one sample produces frames; they are + // just not all valid. + EXPECT_EQ(Lfm2AudioNumMelFrames(1, kMelConfig), 1); + EXPECT_EQ(Lfm2AudioNumValidMelFrames(1, kMelConfig), 0); + EXPECT_EQ(Lfm2AudioNumMelFrames(0, kMelConfig), 1); +} + +TEST(Lfm2AudioFramesTest, TokensRoundUpToWholeEncoderFrames) { + EXPECT_EQ(Lfm2AudioNumTokens(8, kSubsamplingFactor), 1); + EXPECT_EQ(Lfm2AudioNumTokens(9, kSubsamplingFactor), 2); + EXPECT_EQ(Lfm2AudioNumTokens(101, kSubsamplingFactor), 13); + EXPECT_EQ(Lfm2AudioNumTokens(0, kSubsamplingFactor), 0); +} + +TEST(Lfm2AudioFramesTest, TokensRejectANonPositiveSubsamplingFactor) { + EXPECT_NE(CaptureThrowMessage([] { Lfm2AudioNumTokens(64, 0); }).find("subsampling_factor must be positive"), + std::string::npos); +} + +TEST(Lfm2AudioMelTest, NormalizesEveryValidFrameAndZeroesTheRest) { + // Half a second of a tone: every mel bin should end up zero-mean over the valid frames, and the + // frames past the valid length are the pad_value NeMo writes. + constexpr int64_t kNumSamples = 8000; + std::vector pcm(kNumSamples); + for (int64_t i = 0; i < kNumSamples; ++i) { + pcm[static_cast(i)] = 0.25f * std::sin(2.0f * 3.14159265358979323846f * 440.0f * static_cast(i) / 16000.0f); + } + + int64_t num_frames = 0; + const auto mel = ComputeLfm2AudioMel(pcm.data(), kNumSamples, kMelConfig, num_frames); + + const int64_t valid_frames = Lfm2AudioNumValidMelFrames(kNumSamples, kMelConfig); + ASSERT_EQ(num_frames, Lfm2AudioNumMelFrames(kNumSamples, kMelConfig)); + ASSERT_EQ(num_frames, valid_frames + 1); + ASSERT_EQ(mel.size(), static_cast(num_frames) * static_cast(kMelConfig.num_mels)); + + for (int64_t m = 0; m < kMelConfig.num_mels; ++m) { + double sum = 0.0; + double square_sum = 0.0; + for (int64_t t = 0; t < valid_frames; ++t) { + const double value = mel[static_cast(t) * static_cast(kMelConfig.num_mels) + static_cast(m)]; + sum += value; + square_sum += value * value; + } + EXPECT_NEAR(sum / valid_frames, 0.0, 1e-4) << "mel bin " << m << " is not zero-mean"; + // Unbiased standard deviation of 1, up to the epsilon added to the divisor. + EXPECT_NEAR(std::sqrt(square_sum / (valid_frames - 1)), 1.0, 1e-3) << "mel bin " << m << " is not unit-variance"; + } + for (int64_t m = 0; m < kMelConfig.num_mels; ++m) { + EXPECT_EQ(mel[static_cast(valid_frames) * static_cast(kMelConfig.num_mels) + static_cast(m)], 0.0f); + } +} + +TEST(Lfm2AudioMelTest, RejectsAClipTooShortToNormalize) { + const std::vector pcm(160, 0.1f); // one valid frame; the standard deviation would divide by zero + int64_t num_frames = 0; + const auto message = CaptureThrowMessage([&] { ComputeLfm2AudioMel(pcm.data(), 160, kMelConfig, num_frames); }); + EXPECT_NE(message.find("too short"), std::string::npos) << message; + EXPECT_NE(message.find("320 samples are needed"), std::string::npos) << message; +} + +TEST(Lfm2AudioMelTest, RejectsAWindowLongerThanTheFftSize) { + Lfm2AudioMelConfig config = kMelConfig; + config.win_length = config.fft_size + 1; + int64_t num_frames = 0; + const std::vector pcm(16000, 0.0f); + EXPECT_NE(CaptureThrowMessage([&] { ComputeLfm2AudioMel(pcm.data(), 16000, config, num_frames); }) + .find("invalid mel configuration"), + std::string::npos); +} + +TEST(Lfm2AudioPromptTest, SplitsAroundEveryMarker) { + const auto segments = SplitLfm2AudioPrompt("before <|audio|> middle <|audio|> after", 2); + ASSERT_EQ(segments.size(), 3u); + EXPECT_EQ(segments[0], "before "); + EXPECT_EQ(segments[1], " middle "); + EXPECT_EQ(segments[2], " after"); +} + +TEST(Lfm2AudioPromptTest, KeepsEmptySegmentsAroundAdjacentMarkers) { + // Adjacent markers and a prompt that is nothing but a marker still need their empty text segments, + // so the clips stay in order and nothing is dropped. + const auto segments = SplitLfm2AudioPrompt("<|audio|><|audio|>", 2); + ASSERT_EQ(segments.size(), 3u); + EXPECT_EQ(segments[0], ""); + EXPECT_EQ(segments[1], ""); + EXPECT_EQ(segments[2], ""); +} + +TEST(Lfm2AudioPromptTest, TextOnlyPromptWithNoClipsIsOneSegment) { + const auto segments = SplitLfm2AudioPrompt("no audio here", 0); + ASSERT_EQ(segments.size(), 1u); + EXPECT_EQ(segments[0], "no audio here"); +} + +TEST(Lfm2AudioPromptTest, RejectsAMarkerCountThatDoesNotMatchTheClips) { + EXPECT_NE(CaptureThrowMessage([] { SplitLfm2AudioPrompt("only text", 1); }) + .find("contains 0 <|audio|> markers but 1 audio clips"), + std::string::npos); + EXPECT_NE(CaptureThrowMessage([] { SplitLfm2AudioPrompt("<|audio|> and <|audio|>", 1); }) + .find("contains 2 <|audio|> markers but 1 audio clips"), + std::string::npos); + EXPECT_NE(CaptureThrowMessage([] { SplitLfm2AudioPrompt("<|audio|>", 0); }) + .find("contains 1 <|audio|> markers but 0 audio clips"), + std::string::npos); +} + +TEST(Lfm2AudioResampleTest, SameRateIsACopy) { + const std::vector samples{0.1f, -0.2f, 0.3f}; + EXPECT_EQ(ResampleLfm2Audio(samples.data(), 3, 16000, 16000), samples); +} + +TEST(Lfm2AudioResampleTest, LengthIsTheCeilingOfTheRateRatio) { + // torchaudio trims to ceil(target * num_samples / source), up as well as down. + const std::vector samples(1001, 0.0f); + EXPECT_EQ(ResampleLfm2Audio(samples.data(), 1001, 44100, 16000).size(), 364u); + EXPECT_EQ(ResampleLfm2Audio(samples.data(), 1001, 48000, 16000).size(), 334u); + EXPECT_EQ(ResampleLfm2Audio(samples.data(), 1001, 8000, 16000).size(), 2002u); +} + +TEST(Lfm2AudioResampleTest, KeepsAConstantSignalConstantAwayFromTheEdges) { + // The kernels of every output phase sum to one, so only the zero padding at the ends shows. + const std::vector samples(4410, 0.5f); + for (const int64_t source_rate : {8000, 22050, 44100}) { + const auto resampled = ResampleLfm2Audio(samples.data(), 4410, source_rate, 16000); + for (size_t i = resampled.size() / 4; i < 3 * resampled.size() / 4; ++i) { + ASSERT_NEAR(resampled[i], 0.5f, 2e-3f) << source_rate << " Hz, sample " << i; + } + } +} + +TEST(Lfm2AudioResampleTest, RejectsARateThatIsNotPositive) { + const std::vector samples(16, 0.0f); + EXPECT_NE(CaptureThrowMessage([&] { ResampleLfm2Audio(samples.data(), 16, 0, 16000); }).find("cannot resample"), + std::string::npos); +} + +// The embedding and speech sessions are handed buffers allocated on the decoder's devices, so neither +// may be left on CPU by session_options of its own while such a buffer is device memory. +namespace { + +Config CpuSubModelsConfig() { + Config config; // as a CPU export writes it: both graphs have session_options with no provider + config.model.speech.session_options = Config::SessionOptions{}; + config.model.embedding.session_options = Config::SessionOptions{}; + return config; +} + +} // namespace + +TEST(Lfm2AudioSessionDevicesTest, ACpuEmbeddingCannotTakeTheInputsOfADecoderOnCuda) { + const auto message = CaptureThrowMessage([] { + CheckLfm2AudioSessionDevices(CpuSubModelsConfig(), DeviceType::CUDA, DeviceType::CUDA, /*with_audio=*/false); + }); + EXPECT_NE(message.find("model.embedding.session_options run the embedding model on CPU, but the decoder takes " + "its inputs in CUDA memory"), + std::string::npos) + << message; +} + +TEST(Lfm2AudioSessionDevicesTest, CpuSubModelsTakeTextOnlyPromptsWhenTheInputsAreOnTheHost) { + // WebGPU without graph capture keeps the decoder's inputs on the host, so text-only prompts run. + EXPECT_NO_THROW( + CheckLfm2AudioSessionDevices(CpuSubModelsConfig(), DeviceType::WEBGPU, DeviceType::CPU, /*with_audio=*/false)); + + // Audio features are allocated on the decoder's device, which a CPU encoder cannot write. + const auto message = CaptureThrowMessage([] { + CheckLfm2AudioSessionDevices(CpuSubModelsConfig(), DeviceType::WEBGPU, DeviceType::CPU, /*with_audio=*/true); + }); + EXPECT_NE(message.find("model.speech.session_options run the speech model on CPU, but the audio features are " + "passed in WebGPU memory"), + std::string::npos) + << message; + + // Nor can a CPU embedding read them. + auto config = CpuSubModelsConfig(); + config.model.speech.session_options.reset(); + EXPECT_NE(CaptureThrowMessage([&] { + CheckLfm2AudioSessionDevices(config, DeviceType::WEBGPU, DeviceType::CPU, /*with_audio=*/true); + }).find("model.embedding.session_options"), + std::string::npos); +} + +TEST(Lfm2AudioSessionDevicesTest, SubModelsThatCanUseTheBuffersAreAccepted) { + // Without session_options of their own both graphs follow the decoder. + EXPECT_NO_THROW(CheckLfm2AudioSessionDevices(Config{}, DeviceType::CUDA, DeviceType::CUDA, /*with_audio=*/true)); + + auto config = CpuSubModelsConfig(); + config.model.speech.session_options->providers = {"cuda"}; + config.model.embedding.session_options->providers = {"cuda"}; + EXPECT_NO_THROW(CheckLfm2AudioSessionDevices(config, DeviceType::CUDA, DeviceType::CUDA, /*with_audio=*/true)); + + // An explicit CPU provider is still CPU. + config.model.speech.session_options->providers = {"CPU"}; + EXPECT_THROW(CheckLfm2AudioSessionDevices(config, DeviceType::CUDA, DeviceType::CUDA, /*with_audio=*/true), + std::runtime_error); + + // Everything on CPU, and a decoder whose buffers a CPU session can use: OpenVINO allocates from the CPU. + for (const auto device : {DeviceType::CPU, DeviceType::OpenVINO}) { + EXPECT_NO_THROW(CheckLfm2AudioSessionDevices(CpuSubModelsConfig(), device, device, /*with_audio=*/true)); + } +} + +} // namespace Generators::test diff --git a/test/cpp/lfm2_cache_layer_types_test.cpp b/test/cpp/lfm2_cache_layer_types_test.cpp index fb146c6cb4..2491471147 100644 --- a/test/cpp/lfm2_cache_layer_types_test.cpp +++ b/test/cpp/lfm2_cache_layer_types_test.cpp @@ -183,3 +183,10 @@ TEST(LFM2CacheLayerTypesValidationTest, Lfm2VlTextTypeUsesLfm2Runtime) { SkipIfModelUnavailable(); ExpectLfm2Runtime("lfm2_vl_text", "vl_text_type"); } + +// An LFM2-Audio checkpoint exported without exclude_embeds is labeled "lfm2_audio_text": a plain +// LFM2 decoder rather than a stage of the speech pipeline. +TEST(LFM2CacheLayerTypesValidationTest, Lfm2AudioTextTypeUsesLfm2Runtime) { + SkipIfModelUnavailable(); + ExpectLfm2Runtime("lfm2_audio_text", "audio_text_type"); +} diff --git a/test/models/lfm2-audio/dummy_audio_embedding.onnx b/test/models/lfm2-audio/dummy_audio_embedding.onnx new file mode 100644 index 0000000000..cced6a01fa Binary files /dev/null and b/test/models/lfm2-audio/dummy_audio_embedding.onnx differ diff --git a/test/models/lfm2-audio/dummy_depthformer.onnx b/test/models/lfm2-audio/dummy_depthformer.onnx new file mode 100644 index 0000000000..44201df59a Binary files /dev/null and b/test/models/lfm2-audio/dummy_depthformer.onnx differ diff --git a/test/models/lfm2-audio/dummy_embedding.onnx b/test/models/lfm2-audio/dummy_embedding.onnx new file mode 100644 index 0000000000..659101cc82 Binary files /dev/null and b/test/models/lfm2-audio/dummy_embedding.onnx differ diff --git a/test/models/lfm2-audio/dummy_speech.onnx b/test/models/lfm2-audio/dummy_speech.onnx new file mode 100644 index 0000000000..1adb422c86 Binary files /dev/null and b/test/models/lfm2-audio/dummy_speech.onnx differ diff --git a/test/models/lfm2-audio/dummy_text.onnx b/test/models/lfm2-audio/dummy_text.onnx new file mode 100644 index 0000000000..299c2d07ee Binary files /dev/null and b/test/models/lfm2-audio/dummy_text.onnx differ diff --git a/test/models/lfm2-audio/genai_config.json b/test/models/lfm2-audio/genai_config.json new file mode 100644 index 0000000000..9f8d66933e --- /dev/null +++ b/test/models/lfm2-audio/genai_config.json @@ -0,0 +1,92 @@ +{ + "model": { + "audio_token_id": 133, + "bos_token_id": 1, + "context_length": 4096, + "decoder": { + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [] + }, + "filename": "dummy_text.onnx", + "head_size": 16, + "hidden_size": 32, + "inputs": { + "inputs_embeds": "inputs_embeds", + "attention_mask": "attention_mask", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + "past_conv_names": "past.%d.conv" + }, + "outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + "present_conv_names": "present.%d.conv" + }, + "num_attention_heads": 2, + "num_hidden_layers": 4, + "num_key_value_heads": 2, + "layer_types": [ + "conv", + "full_attention", + "conv", + "full_attention" + ], + "conv_cache_size": 2 + }, + "eos_token_id": [ + 7, + 128, + 130 + ], + "pad_token_id": 0, + "type": "lfm2_audio", + "vocab_size": 65536, + "embedding": { + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [] + }, + "filename": "dummy_embedding.onnx", + "inputs": { + "input_ids": "input_ids", + "audio_features": "audio_features" + }, + "outputs": { + "inputs_embeds": "inputs_embeds" + } + }, + "speech": { + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [] + }, + "filename": "dummy_speech.onnx", + "inputs": { + "audio_embeds": "mel_spectrogram", + "audio_lengths": "mel_lengths", + "audio_sizes": "audio_sizes" + }, + "outputs": { + "audio_features": "audio_embeddings" + } + } + }, + "search": { + "diversity_penalty": 0.0, + "do_sample": false, + "early_stopping": true, + "length_penalty": 1.0, + "max_length": 4096, + "min_length": 0, + "no_repeat_ngram_size": 0, + "num_beams": 1, + "num_return_sequences": 1, + "past_present_share_buffer": false, + "repetition_penalty": 1.0, + "temperature": 1.0, + "top_k": 1, + "top_p": 1.0 + } +} diff --git a/test/models/lfm2-audio/tokenizer.zip b/test/models/lfm2-audio/tokenizer.zip new file mode 100644 index 0000000000..3e36ce5ae8 Binary files /dev/null and b/test/models/lfm2-audio/tokenizer.zip differ diff --git a/test/python/builder/test_lfm2_audio.py b/test/python/builder/test_lfm2_audio.py new file mode 100644 index 0000000000..bac381a762 --- /dev/null +++ b/test/python/builder/test_lfm2_audio.py @@ -0,0 +1,343 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License + +"""Unit tests for the LFM2-Audio model builder wiring. + +LFM2-Audio checkpoints (LFM2-Audio-1.5B, LFM2.5-Audio-1.5B) have no transformers model class: their +config.json nests the LFM2 decoder config under "lfm" and the checkpoint stores the decoder weights +under the "lfm." prefix next to the audio encoder, depthformer and audio embeddings. The builder has +to read that nested config, load only the decoder with its logits tied to the token embeddings, and +label the export `lfm2_audio` (a pipeline stage) or `lfm2_audio_text` (a text-only decoder). +""" + +from __future__ import annotations + +import json +import sys +import types + +import pytest +from _builder_test_utils import load_builder_module +from test_lfm2_vl import builder_module + +base_module = load_builder_module("base") +lfm2_module = load_builder_module("lfm2") +LFM2AudioModel = lfm2_module.LFM2AudioModel + +torch = pytest.importorskip("torch") +transformers = pytest.importorskip("transformers") +save_file = pytest.importorskip("safetensors.torch").save_file +# `_builder_test_utils` substitutes a stub for transformers when it is missing, which has none of the +# LFM2 classes these tests build a checkpoint from. +pytestmark = pytest.mark.skipif( + not hasattr(transformers, "Lfm2ForCausalLM"), reason="transformers does not provide the LFM2 model classes" +) + +DECODER_CONFIG = { + "architectures": ["Lfm2ForCausalLM"], + "block_auto_adjust_ff_dim": True, + "block_ff_dim": 96, + "block_ffn_dim_multiplier": 1.0, + "block_multiple_of": 16, + "conv_L_cache": 3, + "conv_bias": False, + "hidden_size": 32, + "intermediate_size": 96, + "layer_types": ["conv", "full_attention", "conv", "full_attention"], + "model_type": "lfm2", + "norm_eps": 1e-05, + "num_attention_heads": 4, + "num_hidden_layers": 4, + "num_key_value_heads": 2, + "rope_theta": 1000000, + "vocab_size": 512, +} + + +def _write_checkpoint(path, decoder_weights=None, shards=1): + """An LFM2-Audio checkpoint directory: nested config.json plus lfm.* and audio tensors.""" + path.mkdir(parents=True, exist_ok=True) + config = { + "architectures": ["Lfm2AudioForConditionalGeneration"], + "codebooks": 8, + "preprocessor": {"sample_rate": 16000, "features": 128, "n_fft": 512}, + "encoder": {"feat_in": 128, "d_model": 512, "n_layers": 17}, + "lfm": DECODER_CONFIG, + "depthformer": {"layers": 6, "dim": 1024, "tie": True}, + } + (path / "config.json").write_text(json.dumps(config)) + + if decoder_weights is None: + torch.manual_seed(0) + model = transformers.Lfm2ForCausalLM(transformers.Lfm2Config(**DECODER_CONFIG)) + decoder_weights = { + "lfm." + name[len("model.") :]: tensor + for name, tensor in model.state_dict().items() + if name.startswith("model.") + } + tensors = dict(decoder_weights) + tensors["conformer.pre_encode.out.weight"] = torch.zeros(4, 4) + tensors["audio_adapter.model.1.weight"] = torch.zeros(32, 8) + tensors["audio_embedding.embedding.weight"] = torch.zeros(16, 32) + tensors["depth_linear.weight"] = torch.zeros(8, 32) + tensors["codebook_offsets"] = torch.arange(8) * 2049 + tensors = {name: tensor.contiguous() for name, tensor in tensors.items()} + if shards == 1: + save_file(tensors, str(path / "model.safetensors")) + else: + # What `save_pretrained` writes once a checkpoint is over its shard size: several files and + # an index naming which tensor lives in which. An fp32 fine-tune of this model reaches it. + names = sorted(tensors) + weight_map = {} + for shard in range(shards): + filename = f"model-{shard + 1:05d}-of-{shards:05d}.safetensors" + part = {name: tensors[name] for name in names[shard::shards]} + save_file(part, str(path / filename)) + weight_map.update(dict.fromkeys(part, filename)) + (path / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}, "weight_map": weight_map})) + return decoder_weights + + +def test_lfm2_audio_load_config_reads_the_nested_decoder_config(tmp_path): + _write_checkpoint(tmp_path) + + config = LFM2AudioModel.load_config(str(tmp_path)) + + assert config.architectures == ["Lfm2AudioForConditionalGeneration"] + assert config._name_or_path == str(tmp_path) + assert config.hidden_size == DECODER_CONFIG["hidden_size"] + assert config.layer_types == DECODER_CONFIG["layer_types"] + assert config.conv_L_cache == DECODER_CONFIG["conv_L_cache"] + assert config.tie_word_embeddings + + +def test_lfm2_audio_load_config_is_none_for_other_checkpoints(tmp_path): + (tmp_path / "config.json").write_text(json.dumps(DECODER_CONFIG)) + assert LFM2AudioModel.load_config(str(tmp_path)) is None + + +def _hf_details_with(monkeypatch, auto_config): + """`get_hf_details` with AutoConfig and the tokenizer replaced, and the calls it made.""" + calls = [] + + def from_pretrained(name, **kwargs): + calls.append("AutoConfig") + return auto_config(name) + + def load_config(name, **kwargs): + calls.append("load_config") + return LFM2AudioModel.load_config(name) + + monkeypatch.setattr(builder_module, "AutoConfig", types.SimpleNamespace(from_pretrained=from_pretrained)) + monkeypatch.setattr( + builder_module, "AutoTokenizer", types.SimpleNamespace(from_pretrained=lambda *args, **kwargs: None) + ) + monkeypatch.setattr(builder_module, "add_special_token_ids", lambda config, tokenizer: None) + monkeypatch.setattr(builder_module, "LFM2AudioModel", types.SimpleNamespace(load_config=load_config)) + return calls + + +def test_lfm2_audio_config_is_read_only_once_autoconfig_refuses_the_checkpoint(monkeypatch, tmp_path): + _write_checkpoint(tmp_path) + + def refuse(name): + raise ValueError(f"Unrecognized model in {name}. Should have a `model_type` key in its config.json.") + + calls = _hf_details_with(monkeypatch, refuse) + details = builder_module.get_hf_details("", str(tmp_path), str(tmp_path / "cache"), {}) + + assert calls == ["AutoConfig", "load_config"] + assert details["hf_config"].architectures == ["Lfm2AudioForConditionalGeneration"] + + +def test_lfm2_audio_config_lookup_is_skipped_for_every_other_model(monkeypatch, tmp_path): + # Every build goes through get_hf_details; other models must not pay for a second config read. + recognized = types.SimpleNamespace(architectures=["LlamaForCausalLM"]) + calls = _hf_details_with(monkeypatch, lambda name: recognized) + + details = builder_module.get_hf_details("", str(tmp_path), str(tmp_path / "cache"), {}) + + assert calls == ["AutoConfig"] + assert details["hf_config"] is recognized + + +def test_lfm2_audio_config_lookup_keeps_autoconfig_error_for_an_unknown_checkpoint(monkeypatch, tmp_path): + (tmp_path / "config.json").write_text(json.dumps({"architectures": ["SomethingElse"]})) + + def refuse(name): + raise ValueError("Unrecognized model") + + _hf_details_with(monkeypatch, refuse) + with pytest.raises(ValueError, match="Unrecognized model"): + builder_module.get_hf_details("", str(tmp_path), str(tmp_path / "cache"), {}) + + +def _audio_builder(tmp_path, extra_options=None, model_name_or_path=None): + model = LFM2AudioModel.__new__(LFM2AudioModel) + model.decoder_config = LFM2AudioModel.load_config(str(tmp_path)) + model.quant_type = None + model.cache_dir = str(tmp_path / "cache") + model.hf_token = True + model.extra_options = extra_options or {} + model.model_name_or_path = model_name_or_path or str(tmp_path) + return model + + +def test_lfm2_audio_load_weights_loads_the_decoder_with_tied_logits(tmp_path): + decoder_weights = _write_checkpoint(tmp_path) + + loaded = _audio_builder(tmp_path).load_weights(str(tmp_path)) + + assert type(loaded).__name__ == "Lfm2ForCausalLM" + torch.testing.assert_close( + loaded.model.layers[0].conv.in_proj.weight, decoder_weights["lfm.layers.0.conv.in_proj.weight"] + ) + torch.testing.assert_close(loaded.model.embedding_norm.weight, decoder_weights["lfm.embedding_norm.weight"]) + assert loaded.lm_head.weight.data_ptr() == loaded.model.embed_tokens.weight.data_ptr() + assert not any("conformer" in name or "depth" in name or "audio" in name for name, _ in loaded.named_parameters()) + assert not any(parameter.is_meta for parameter in loaded.parameters()) + + +@pytest.mark.parametrize("shards", [2, 3]) +def test_lfm2_audio_load_weights_reads_a_sharded_checkpoint(tmp_path, shards): + # These checkpoints are one file today, but `save_pretrained` shards anything over its size + # limit, which an fp32 fine-tune of this model exceeds. The tensors have to be gathered from + # every shard the index names, not just the first. + decoder_weights = _write_checkpoint(tmp_path, shards=shards) + assert not (tmp_path / "model.safetensors").exists(), "the sharded layout has no single file" + + loaded = _audio_builder(tmp_path).load_weights(str(tmp_path)) + + torch.testing.assert_close( + loaded.model.layers[0].conv.in_proj.weight, decoder_weights["lfm.layers.0.conv.in_proj.weight"] + ) + torch.testing.assert_close(loaded.model.embedding_norm.weight, decoder_weights["lfm.embedding_norm.weight"]) + assert not any(parameter.is_meta for parameter in loaded.parameters()) + + +def test_lfm2_audio_load_weights_rejects_a_checkpoint_missing_decoder_tensors(tmp_path): + torch.manual_seed(0) + model = transformers.Lfm2ForCausalLM(transformers.Lfm2Config(**DECODER_CONFIG)) + decoder_weights = { + "lfm." + name[len("model.") :]: tensor + for name, tensor in model.state_dict().items() + if name.startswith("model.") + } + del decoder_weights["lfm.layers.1.self_attn.q_proj.weight"] + _write_checkpoint(tmp_path, decoder_weights) + + with pytest.raises( + ValueError, match=r"does not match its config: missing .*'model\.layers\.1\.self_attn\.q_proj\.weight'" + ): + _audio_builder(tmp_path).load_weights(str(tmp_path)) + + +def test_lfm2_audio_load_weights_downloads_the_checkpoint_named_by_the_config(monkeypatch, tmp_path): + # `-m ` leaves the input path empty; the checkpoint to fetch is then the one the config + # was read from. Passing the empty path straight to the Hub asks it for a repository called "". + _write_checkpoint(tmp_path) + asked = [] + + def fake_snapshot_download(repo_id, **kwargs): + asked.append((repo_id, kwargs)) + return str(tmp_path) + + hub = types.ModuleType("huggingface_hub") + hub.snapshot_download = fake_snapshot_download + monkeypatch.setitem(sys.modules, "huggingface_hub", hub) + + builder = _audio_builder(tmp_path, model_name_or_path="LiquidAI/LFM2.5-Audio-1.5B") + loaded = builder.load_weights("") + + assert type(loaded).__name__ == "Lfm2ForCausalLM" + assert asked[0][0] == "LiquidAI/LFM2.5-Audio-1.5B" + assert asked[0][1]["cache_dir"] == builder.cache_dir + + +def test_lfm2_audio_load_weights_prefers_a_local_input_path(monkeypatch, tmp_path): + _write_checkpoint(tmp_path) + + def fail(*args, **kwargs): + raise AssertionError("a local checkpoint directory must not be downloaded") + + hub = types.ModuleType("huggingface_hub") + hub.snapshot_download = fail + monkeypatch.setitem(sys.modules, "huggingface_hub", hub) + + loaded = _audio_builder(tmp_path, model_name_or_path="LiquidAI/LFM2.5-Audio-1.5B").load_weights(str(tmp_path)) + + assert type(loaded).__name__ == "Lfm2ForCausalLM" + + +def test_lfm2_audio_load_weights_refuses_a_lora_adapter(tmp_path): + # An adapter trained on these checkpoints is keyed "lfm.*" and matches nothing in the decoder + # loaded here, so PEFT would attach untrained weights with only a warning. + _write_checkpoint(tmp_path) + builder = _audio_builder(tmp_path, {"adapter_path": "some/adapter"}) + + with pytest.raises(ValueError, match="LoRA adapters are not supported for LFM2-Audio"): + builder.load_weights(str(tmp_path)) + + +@pytest.mark.parametrize( + "existing,expected", + [ + (7, [7, 128, 130]), + ([7], [7, 128, 130]), + ([7, 2], [7, 2, 128, 130]), + ([7, 128, 130], [7, 128, 130]), + ([128], [128, 130]), + ], + ids=["scalar", "list", "several", "already-there", "partly-there"], +) +def test_lfm2_audio_stops_when_the_model_starts_speaking(existing, expected): + # The model answers in speech too. <|audio_start|> marks a wholly spoken answer and <|text_end|> + # the text half of an interleaved one ending; after either, the positions are audio codes meant + # for the depthformer, which this runtime does not have, so generation has to end there rather + # than read them off the text head. + builder = LFM2AudioModel.__new__(LFM2AudioModel) + builder.layer_types = ["conv", "full_attention"] + builder.conv_L_cache = 3 + genai_config = {"model": {"decoder": {}, "eos_token_id": existing}} + + builder.update_genai_config(genai_config) + + assert genai_config["model"]["eos_token_id"] == expected + # The LFM2 decoder settings the base class writes are still there. + assert genai_config["model"]["decoder"]["layer_types"] == ["conv", "full_attention"] + assert genai_config["model"]["decoder"]["conv_cache_size"] == 2 + + +@pytest.mark.parametrize("exclude_embeds,expected_type", [(True, "lfm2_audio"), (False, "lfm2_audio_text")]) +def test_lfm2_audio_architecture_labels_the_export(monkeypatch, tmp_path, exclude_embeds, expected_type): + captured = {} + + class FakeLFM2AudioModel: + model_type = "lfm2" + + def __init__(self, *args): + captured["args"] = args + self.exclude_embeds = exclude_embeds + + def make_genai_config(self, *args): + captured["model_type"] = self.model_type + + def save_processing(self, *args): + pass + + config = types.SimpleNamespace(architectures=["Lfm2AudioForConditionalGeneration"]) + monkeypatch.setattr(builder_module, "LFM2AudioModel", FakeLFM2AudioModel) + + builder_module.create_model( + "LiquidAI/LFM2.5-Audio-1.5B", + str(tmp_path / "input"), + str(tmp_path / "output"), + "int4", + "cpu", + str(tmp_path / "cache"), + config_only=True, + hf_details={"extra_kwargs": {}, "hf_name": "LiquidAI/LFM2.5-Audio-1.5B", "hf_config": config}, + ) + + assert captured["args"][0] is config + assert captured["model_type"] == expected_type diff --git a/test/python/builder/test_qmoe_weights.py b/test/python/builder/test_qmoe_weights.py index 7e09a01405..c2216934a4 100644 --- a/test/python/builder/test_qmoe_weights.py +++ b/test/python/builder/test_qmoe_weights.py @@ -131,6 +131,7 @@ def _load_builder_cli_module(monkeypatch): "GraniteModel", "HunyuanDenseV1Model", "InternLM2Model", + "LFM2AudioModel", "LFM2Model", "LFM2MoEModel", "LlamaModel", diff --git a/test/python/create/create_dummy_lfm2_audio_models.py b/test/python/create/create_dummy_lfm2_audio_models.py new file mode 100644 index 0000000000..4fc22a52f6 --- /dev/null +++ b/test/python/create/create_dummy_lfm2_audio_models.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License + +""" +Generate tiny ONNX models for LFM2-Audio pipeline testing. + +Writes the three models the lfm2_audio pipeline runs, with the input/output contract of the published +LFM2-Audio / LFM2.5-Audio encoder export (LiquidAI/LFM2.5-Audio-1.5B-ONNX), plus the genai_config.json +that wires them together: + + dummy_speech.onnx mel_spectrogram [N, T, 128], mel_lengths [N] (int64) + -> audio_embeddings [N, ceil(T / 8), hidden], audio_lengths [N] (int64) + Every 8 mel frames are averaged and projected, so the features depend on the + audio and the frame count matches the placeholders the processor writes. + dummy_embedding.onnx input_ids, audio_features -> inputs_embeds + Token lookup with audio_features scattered over the positions holding + AUDIO_TOKEN_ID, the graph the tutorial describes for the real model. + dummy_text.onnx The LFM2 decoder of the LFM2-VL fixture (see create_dummy_lfm2_vl_models.py): + each logit depends on the current token and the two before it through a causal + convolution over past.%d.conv, so generation only matches between runs when the + conv state is carried correctly. + + dummy_depthformer.onnx, dummy_audio_embedding.onnx + Speech output, with the contract of LiquidAI's vocoder_depthformer.onnx and + audio_embedding.onnx. The depthformer gives one codebook's logits per run; they + depend on the hidden state, the codebook, the code before it and the cache of + the runs before, so a frame only comes out right when all four are carried. + The decoder has the hidden_states output they read. genai_config.json leaves + them out: the tests that use them add model.audio_output themselves. + +The tokenizer is checked in separately, and zipped: tokenizer.zip holds tokenizer.json and +tokenizer_config.json from LiquidAI/LFM2.5-Audio-1.5B (LFM2-Audio-1.5B ships the same files), whose +4.8 MB compresses to under 1 MB. test_lfm2_audio_models.py unpacks it into a temporary directory +next to these models. + +Usage: + python create_dummy_lfm2_audio_models.py --output test/models/lfm2-audio +""" + +import argparse +import json +import os +import sys + +import numpy as np +from onnx import TensorProto, helper, numpy_helper + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from create_dummy_lfm2_vl_models import ( + CONV_CACHE_SIZE, + HEAD_SIZE, + HIDDEN_SIZE, + LAYER_TYPES, + LOGIT_VOCAB, + NUM_KV_HEADS, + VOCAB_SIZE, + _const, + _save, + create_dummy_decoder_model, +) + +NUM_MELS = 128 +SUBSAMPLING_FACTOR = 8 +# <|reserved_123|>: the placeholder the processor writes once per encoder frame. LFM2-Audio has no +# audio input token of its own, so any single token that never occurs in text will do. +AUDIO_TOKEN_ID = 133 + + +def create_dummy_speech_model(output_path: str, rng: np.random.Generator): + """mel_spectrogram, mel_lengths -> audio_embeddings [N, ceil(T / 8), hidden], audio_lengths""" + inputs = [ + helper.make_tensor_value_info("mel_spectrogram", TensorProto.FLOAT, ["batch_size", "time_steps", NUM_MELS]), + helper.make_tensor_value_info("mel_lengths", TensorProto.INT64, ["batch_size"]), + ] + outputs = [ + helper.make_tensor_value_info( + "audio_embeddings", TensorProto.FLOAT, ["batch_size", "reduced_time", HIDDEN_SIZE] + ), + helper.make_tensor_value_info("audio_lengths", TensorProto.INT64, ["batch_size"]), + ] + projection = numpy_helper.from_array( + (rng.standard_normal((NUM_MELS, HIDDEN_SIZE)) / np.sqrt(NUM_MELS)).astype(np.float32), name="projection" + ) + nodes = [ + _const("subsampling_minus_one", [SUBSAMPLING_FACTOR - 1]), + _const("subsampling", [SUBSAMPLING_FACTOR]), + # [N, T, mels] -> [N, mels, T] so the pool runs over time; ceil_mode keeps the partial last window + helper.make_node("Transpose", ["mel_spectrogram"], ["mel_nct"], perm=[0, 2, 1]), + helper.make_node( + "AveragePool", + ["mel_nct"], + ["pooled_nct"], + kernel_shape=[SUBSAMPLING_FACTOR], + strides=[SUBSAMPLING_FACTOR], + ceil_mode=1, + ), + helper.make_node("Transpose", ["pooled_nct"], ["pooled"], perm=[0, 2, 1]), + helper.make_node("MatMul", ["pooled", "projection"], ["audio_embeddings"]), + # audio_lengths = ceil(mel_lengths / subsampling_factor) + helper.make_node("Add", ["mel_lengths", "subsampling_minus_one"], ["lengths_rounded"]), + helper.make_node("Div", ["lengths_rounded", "subsampling"], ["audio_lengths"]), + ] + _save(helper.make_graph(nodes, "speech", inputs, outputs, initializer=[projection]), output_path) + + +def create_dummy_embedding_model(output_path: str, rng: np.random.Generator): + """input_ids, audio_features -> inputs_embeds, with the audio features scattered over the placeholders""" + inputs = [ + helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch_size", "sequence_length"]), + helper.make_tensor_value_info("audio_features", TensorProto.FLOAT, ["num_audio_tokens", HIDDEN_SIZE]), + ] + outputs = [ + helper.make_tensor_value_info( + "inputs_embeds", TensorProto.FLOAT, ["batch_size", "sequence_length", HIDDEN_SIZE] + ) + ] + table = numpy_helper.from_array( + rng.standard_normal((LOGIT_VOCAB, HIDDEN_SIZE)).astype(np.float32), name="embed_table" + ) + nodes = [ + _const("table_rows", [LOGIT_VOCAB]), + _const("audio_token_id", [AUDIO_TOKEN_ID]), + _const("flat_rows", [-1, HIDDEN_SIZE]), + _const("flat", [-1]), + helper.make_node("Mod", ["input_ids", "table_rows"], ["row_ids"]), + helper.make_node("Gather", ["embed_table", "row_ids"], ["text_embeds"], axis=0), + # Positions of the placeholders in the flattened [batch * sequence] token stream, in order, + # receive the audio features row by row. + helper.make_node("Shape", ["text_embeds"], ["embeds_shape"]), + helper.make_node("Reshape", ["text_embeds", "flat_rows"], ["flat_embeds"]), + helper.make_node("Reshape", ["input_ids", "flat"], ["flat_ids"]), + helper.make_node("Equal", ["flat_ids", "audio_token_id"], ["is_audio"]), + helper.make_node("NonZero", ["is_audio"], ["audio_positions_t"]), + helper.make_node("Transpose", ["audio_positions_t"], ["audio_positions"], perm=[1, 0]), + helper.make_node("ScatterND", ["flat_embeds", "audio_positions", "audio_features"], ["merged"]), + helper.make_node("Reshape", ["merged", "embeds_shape"], ["inputs_embeds"]), + ] + _save(helper.make_graph(nodes, "embedding", inputs, outputs, initializer=[table]), output_path) + + +NUM_CODEBOOKS = 8 +CODEBOOK_SIZE = 2049 # 2048 codes and the end-of-audio code +DEPTH_SIZE = 16 +TABLE_ROWS = 251 # the embedding tables are folded to keep the files small; prime, so codebooks do not collide + + +def create_dummy_depthformer_model(output_path: str, rng: np.random.Generator): + """One codebook's logits per run, with the inputs and outputs of LiquidAI's vocoder_depthformer.onnx""" + cache_shape = [1, "batch", 1, "past_len", DEPTH_SIZE] + inputs = [ + helper.make_tensor_value_info("hidden_states", TensorProto.FLOAT, ["batch", HIDDEN_SIZE]), + helper.make_tensor_value_info("depth_slices_in", TensorProto.FLOAT, ["batch", NUM_CODEBOOKS, DEPTH_SIZE]), + helper.make_tensor_value_info("step_idx", TensorProto.INT64, []), + helper.make_tensor_value_info("prev_token", TensorProto.INT64, ["batch"]), + helper.make_tensor_value_info("past_keys", TensorProto.FLOAT, cache_shape), + helper.make_tensor_value_info("past_values", TensorProto.FLOAT, cache_shape), + helper.make_tensor_value_info("seqlens_k", TensorProto.INT32, ["batch"]), + helper.make_tensor_value_info("total_seq_len", TensorProto.INT32, []), + ] + new_cache_shape = [1, "batch", 1, "new_len", DEPTH_SIZE] + outputs = [ + helper.make_tensor_value_info("logits", TensorProto.FLOAT, ["batch", CODEBOOK_SIZE]), + helper.make_tensor_value_info("depth_slices", TensorProto.FLOAT, ["batch", NUM_CODEBOOKS, DEPTH_SIZE]), + helper.make_tensor_value_info("new_keys", TensorProto.FLOAT, new_cache_shape), + helper.make_tensor_value_info("new_values", TensorProto.FLOAT, new_cache_shape), + ] + initializers = [ + numpy_helper.from_array( + rng.standard_normal((HIDDEN_SIZE, NUM_CODEBOOKS * DEPTH_SIZE)).astype(np.float32), name="depth_linear" + ), + numpy_helper.from_array(rng.standard_normal((TABLE_ROWS, DEPTH_SIZE)).astype(np.float32), name="code_table"), + numpy_helper.from_array(rng.standard_normal((DEPTH_SIZE, CODEBOOK_SIZE)).astype(np.float32), name="to_logits"), + ] + nodes = [ + _const("slices_shape", [-1, NUM_CODEBOOKS, DEPTH_SIZE]), + _const("cache_entry_shape", [1, -1, 1, 1, DEPTH_SIZE]), + _const("zero", [0], dims=[]), + _const("past_axis", [3]), + _const("table_rows", [TABLE_ROWS]), + # The projection of the hidden state is only computed on the first codebook's run; later runs + # are handed it back, as the real graph does to skip the projection. + helper.make_node("MatMul", ["hidden_states", "depth_linear"], ["projected"]), + helper.make_node("Reshape", ["projected", "slices_shape"], ["computed_slices"]), + helper.make_node("Equal", ["step_idx", "zero"], ["is_first"]), + helper.make_node("Where", ["is_first", "computed_slices", "depth_slices_in"], ["depth_slices"]), + helper.make_node("Gather", ["depth_slices", "step_idx"], ["slice"], axis=1), + helper.make_node("Mod", ["prev_token", "table_rows"], ["previous_row"]), + helper.make_node("Gather", ["code_table", "previous_row"], ["previous"], axis=0), + helper.make_node("Add", ["slice", "previous"], ["position"]), + # What the earlier codebooks of this frame left in the cache. + helper.make_node("ReduceSum", ["past_keys", "past_axis"], ["history_5d"], keepdims=0), + helper.make_node("Reshape", ["history_5d", _shape_name("history")], ["history"]), + helper.make_node("Add", ["position", "history"], ["mixed"]), + helper.make_node("MatMul", ["mixed", "to_logits"], ["logits"]), + helper.make_node("Reshape", ["position", "cache_entry_shape"], ["cache_entry"]), + helper.make_node("Concat", ["past_keys", "cache_entry"], ["new_keys"], axis=3), + helper.make_node("Concat", ["past_values", "cache_entry"], ["new_values"], axis=3), + ] + nodes.insert(0, _const(_shape_name("history"), [-1, DEPTH_SIZE])) + _save(helper.make_graph(nodes, "depthformer", inputs, outputs, initializer=initializers), output_path) + + +def _shape_name(name: str) -> str: + return f"{name}_shape" + + +def create_dummy_audio_embedding_model(output_path: str, rng: np.random.Generator): + """audio_codes [batch, length], each offset into its codebook's rows -> audio_embeds [batch, length, hidden]""" + inputs = [helper.make_tensor_value_info("audio_codes", TensorProto.INT64, ["batch_size", "audio_length"])] + outputs = [ + helper.make_tensor_value_info("audio_embeds", TensorProto.FLOAT, ["batch_size", "audio_length", HIDDEN_SIZE]) + ] + table = numpy_helper.from_array( + rng.standard_normal((TABLE_ROWS, HIDDEN_SIZE)).astype(np.float32), name="audio_table" + ) + nodes = [ + _const("table_rows", [TABLE_ROWS]), + helper.make_node("Mod", ["audio_codes", "table_rows"], ["rows"]), + helper.make_node("Gather", ["audio_table", "rows"], ["audio_embeds"], axis=0), + ] + _save(helper.make_graph(nodes, "audio_embedding", inputs, outputs, initializer=[table]), output_path) + + +def create_genai_config(output_path: str): + session_options = {"log_id": "onnxruntime-genai", "provider_options": []} + config = { + "model": { + "audio_token_id": AUDIO_TOKEN_ID, + "bos_token_id": 1, + "context_length": 4096, + "decoder": { + "session_options": session_options, + "filename": "dummy_text.onnx", + "head_size": HEAD_SIZE, + "hidden_size": HIDDEN_SIZE, + "inputs": { + "inputs_embeds": "inputs_embeds", + "attention_mask": "attention_mask", + "past_key_names": "past_key_values.%d.key", + "past_value_names": "past_key_values.%d.value", + "past_conv_names": "past.%d.conv", + }, + "outputs": { + "logits": "logits", + "present_key_names": "present.%d.key", + "present_value_names": "present.%d.value", + "present_conv_names": "present.%d.conv", + }, + "num_attention_heads": NUM_KV_HEADS, + "num_hidden_layers": len(LAYER_TYPES), + "num_key_value_heads": NUM_KV_HEADS, + "layer_types": LAYER_TYPES, + "conv_cache_size": CONV_CACHE_SIZE, + }, + "eos_token_id": [7, 128, 130], # <|im_end|>, <|audio_start|>, <|text_end|>: what the builder writes + "pad_token_id": 0, + "type": "lfm2_audio", + "vocab_size": VOCAB_SIZE, + "embedding": { + "session_options": session_options, + "filename": "dummy_embedding.onnx", + "inputs": {"input_ids": "input_ids", "audio_features": "audio_features"}, + "outputs": {"inputs_embeds": "inputs_embeds"}, + }, + "speech": { + "session_options": session_options, + "filename": "dummy_speech.onnx", + "inputs": { + "audio_embeds": "mel_spectrogram", + "audio_lengths": "mel_lengths", + "audio_sizes": "audio_sizes", + }, + "outputs": {"audio_features": "audio_embeddings"}, + }, + }, + "search": { + "diversity_penalty": 0.0, + "do_sample": False, + "early_stopping": True, + "length_penalty": 1.0, + "max_length": 4096, + "min_length": 0, + "no_repeat_ngram_size": 0, + "num_beams": 1, + "num_return_sequences": 1, + "past_present_share_buffer": False, + "repetition_penalty": 1.0, + "temperature": 1.0, + "top_k": 1, + "top_p": 1.0, + }, + } + with open(output_path, "w") as f: + json.dump(config, f, indent=4) + f.write("\n") + + +def main(): + parser = argparse.ArgumentParser(description="Generate tiny ONNX models for LFM2-Audio pipeline testing") + parser.add_argument("--output", type=str, default="test/models/lfm2-audio", help="Output directory") + args = parser.parse_args() + + os.makedirs(args.output, exist_ok=True) + rng = np.random.default_rng(42) + create_dummy_speech_model(os.path.join(args.output, "dummy_speech.onnx"), rng) + create_dummy_embedding_model(os.path.join(args.output, "dummy_embedding.onnx"), rng) + create_dummy_decoder_model(os.path.join(args.output, "dummy_text.onnx"), rng, hidden_states_output=True) + create_dummy_depthformer_model(os.path.join(args.output, "dummy_depthformer.onnx"), rng) + create_dummy_audio_embedding_model(os.path.join(args.output, "dummy_audio_embedding.onnx"), rng) + create_genai_config(os.path.join(args.output, "genai_config.json")) + print(f"Wrote dummy LFM2-Audio models to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/test/python/create/create_dummy_lfm2_vl_models.py b/test/python/create/create_dummy_lfm2_vl_models.py index d64d1ce977..cabb39713a 100644 --- a/test/python/create/create_dummy_lfm2_vl_models.py +++ b/test/python/create/create_dummy_lfm2_vl_models.py @@ -119,8 +119,12 @@ def create_dummy_embedding_model(output_path: str, rng: np.random.Generator): _save(helper.make_graph(nodes, "embedding", inputs, outputs, initializer=[table]), output_path) -def create_dummy_decoder_model(output_path: str, rng: np.random.Generator): - """inputs_embeds, attention_mask, KV cache, conv state -> logits, present KV cache, present conv state""" +def create_dummy_decoder_model(output_path: str, rng: np.random.Generator, hidden_states_output: bool = False): + """inputs_embeds, attention_mask, KV cache, conv state -> logits, present KV cache, present conv state + + With hidden_states_output the values the logits are projected from come out too, as a decoder built + with include_hidden_states=true gives them. + """ conv_layers = [i for i, t in enumerate(LAYER_TYPES) if t == "conv"] attention_layers = [i for i, t in enumerate(LAYER_TYPES) if t == "full_attention"] kv_shape_past = ["batch_size", NUM_KV_HEADS, "past_sequence_length", HEAD_SIZE] @@ -189,6 +193,14 @@ def create_dummy_decoder_model(output_path: str, rng: np.random.Generator): helper.make_node("Concat", [f"past_key_values.{i}.{kind}", "kv_new"], [f"present.{i}.{kind}"], axis=2) ) + if hidden_states_output: + outputs.append( + helper.make_tensor_value_info( + "hidden_states", TensorProto.FLOAT, ["batch_size", "sequence_length", HIDDEN_SIZE] + ) + ) + nodes.append(helper.make_node("Identity", ["mixed_bsh"], ["hidden_states"])) + _save(helper.make_graph(nodes, "decoder", inputs, outputs, initializer=initializers), output_path) diff --git a/test/python/models/test_lfm2_audio_frames.py b/test/python/models/test_lfm2_audio_frames.py new file mode 100644 index 0000000000..b9673c0781 --- /dev/null +++ b/test/python/models/test_lfm2_audio_frames.py @@ -0,0 +1,182 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +""" +LFM2-Audio frame and token accounting tests. + +Three pieces have to agree on how many decoder tokens a clip is worth, or the encoder output and the +audio placeholders in the prompt will not line up: + + * NeMo's `AudioToMelSpectrogramPreprocessor`, which the reference implementation wraps: its STFT + decides how many mel frames a clip has, and its `get_seq_len` decides how many of those count as + valid for the per-feature normalization, + * the FastConformer encoder, whose three stride-2 convolutions subsample those frames by 8, + * `Lfm2AudioProcessor` in C++, which writes one placeholder per encoder frame. + +These tests are pure arithmetic reimplementations of all three, so they run without a model or a +built runtime. The real C++ processor is exercised against the tiny test model in +test_lfm2_audio_models.py, and `test_matches_liquid_audio` below checks this arithmetic against the +reference implementation itself whenever it is installed. + +Run with: + python -m pytest test/python/models/test_lfm2_audio_frames.py -v +""" + +import math + +import pytest + +# Defaults from LiquidAI/LFM2.5-Audio-1.5B's config.json ("preprocessor" and "encoder"), shared +# byte for byte by LFM2-Audio-1.5B and LFM2.5-Audio-1.5B-JP. +SAMPLE_RATE = 16000 +WINDOW_SIZE = 0.025 # seconds +WINDOW_STRIDE = 0.01 # seconds +WIN_LENGTH = int(WINDOW_SIZE * SAMPLE_RATE) # 400 +HOP_LENGTH = int(WINDOW_STRIDE * SAMPLE_RATE) # 160 +FFT_SIZE = 512 +NUM_MELS = 128 +# Three stride-2 depthwise convolutions in the encoder's pre-encode stage. +SUBSAMPLING_FACTOR = 8 +# The shortest clip the per-feature normalization can handle: its standard deviation divides by +# (valid frames - 1), so one valid frame is not enough. +MIN_VALID_FRAMES = 2 + + +def num_mel_frames(num_samples: int) -> int: + """Frames out of `torch.stft(..., center=True)`: one per hop, plus one for the centre padding.""" + return num_samples // HOP_LENGTH + 1 + + +def num_valid_mel_frames(num_samples: int) -> int: + """NeMo's `get_seq_len`: the frames the per-feature normalization averages over. + + One fewer than the STFT produces. The extra frame is still handed to the encoder and still + counts towards the decoder tokens; it is just zeroed instead of normalized. + """ + return num_samples // HOP_LENGTH + + +def num_audio_tokens(num_frames: int) -> int: + """Decoder tokens for a clip: one per encoder frame, the tail frame rounded up.""" + return math.ceil(num_frames / SUBSAMPLING_FACTOR) + + +def seconds_to_tokens(seconds: float) -> int: + return num_audio_tokens(num_mel_frames(int(seconds * SAMPLE_RATE))) + + +class TestMelFrames: + @pytest.mark.parametrize( + "num_samples,frames,valid", + [ + (0, 1, 0), # centre padding alone still yields a frame + (1, 1, 0), + (159, 1, 0), + (160, 2, 1), # exactly one hop + (16000, 101, 100), # one second + (16159, 101, 100), # a partial hop is truncated + (16160, 102, 101), + ], + ) + def test_known_lengths(self, num_samples, frames, valid): + assert num_mel_frames(num_samples) == frames + assert num_valid_mel_frames(num_samples) == valid + + @pytest.mark.parametrize("num_samples", [0, 1, 160, 4000, 16000, 44100, 160000]) + def test_valid_frames_are_one_short_of_the_stft(self, num_samples): + # This off-by-one is the whole reason the two counts exist separately: the last frame is + # zeroed by the normalization but still reaches the encoder. + assert num_valid_mel_frames(num_samples) == num_mel_frames(num_samples) - 1 + + def test_one_second_is_one_hundred_hops(self): + assert num_valid_mel_frames(SAMPLE_RATE) == int(1.0 / WINDOW_STRIDE) + + @pytest.mark.parametrize("num_samples", [320, 1000, 16000, 100000]) + def test_frame_count_is_monotonic_in_the_clip_length(self, num_samples): + assert num_mel_frames(num_samples) <= num_mel_frames(num_samples + 1) + assert num_mel_frames(num_samples) < num_mel_frames(num_samples + HOP_LENGTH) + + def test_shortest_normalizable_clip(self): + # Two valid frames need two whole hops; anything shorter the processor rejects. + shortest = MIN_VALID_FRAMES * HOP_LENGTH + assert num_valid_mel_frames(shortest) == MIN_VALID_FRAMES + assert num_valid_mel_frames(shortest - 1) < MIN_VALID_FRAMES + + def test_window_fits_inside_the_fft(self): + # A window longer than the FFT size cannot be centred in it; the processor refuses that. + assert WIN_LENGTH <= FFT_SIZE + assert (FFT_SIZE - WIN_LENGTH) % 2 == 0, "the window is centred, so the padding has to split evenly" + + +class TestAudioTokens: + @pytest.mark.parametrize( + "frames,tokens", + [(0, 0), (1, 1), (8, 1), (9, 2), (16, 2), (101, 13), (126, 16), (1001, 126)], + ) + def test_known_frame_counts(self, frames, tokens): + assert num_audio_tokens(frames) == tokens + + @pytest.mark.parametrize("frames", [1, 7, 8, 9, 63, 64, 65, 1000]) + def test_tail_frames_round_up(self, frames): + # Rounding down would drop the tail of every clip whose length is not a multiple of 8, and + # the placeholders would then outnumber the features. + assert num_audio_tokens(frames) == -(-frames // SUBSAMPLING_FACTOR) + assert num_audio_tokens(frames) * SUBSAMPLING_FACTOR >= frames + + @pytest.mark.parametrize("seconds", [0.1, 0.5, 1.0, 2.5, 10.0, 60.0]) + def test_one_token_per_eighty_milliseconds(self, seconds): + # 8 frames x 10 ms: the figure the tutorial quotes for prompt budgeting. + assert seconds_to_tokens(seconds) == pytest.approx(seconds / 0.08, abs=1) + + def test_one_minute_of_audio(self): + assert seconds_to_tokens(60.0) == 751 + + +class TestPlaceholders: + @pytest.mark.parametrize("clips", [[1.0], [0.5, 1.7], [0.2, 0.2, 0.2]]) + def test_placeholders_match_the_encoder_output(self, clips): + # The processor writes this many placeholders; the speech state concatenates exactly this + # many feature rows. A mismatch is what the scatter in the embedding model would trip over. + per_clip = [seconds_to_tokens(seconds) for seconds in clips] + assert sum(per_clip) == sum(num_audio_tokens(num_mel_frames(int(s * SAMPLE_RATE))) for s in clips) + assert all(count > 0 for count in per_clip) + + def test_padding_a_batch_does_not_change_any_clip(self): + # Clips are staged in one padded tensor, but each is encoded on its own frames, so a clip's + # token count never depends on what it is batched with. + short, long = 0.5, 1.7 + assert seconds_to_tokens(short) == num_audio_tokens(num_mel_frames(int(short * SAMPLE_RATE))) + padded_frames = num_mel_frames(int(long * SAMPLE_RATE)) + assert num_audio_tokens(padded_frames) != seconds_to_tokens(short), "the test lengths must differ" + + +@pytest.mark.parametrize("seconds", [0.25, 1.0, 2.5]) +def test_matches_liquid_audio(seconds): + """The reference implementation itself, when the liquid-audio package is installed.""" + torch = pytest.importorskip("torch") + processor_module = pytest.importorskip("liquid_audio.model.conformer.processor") + + preprocessor = processor_module.AudioToMelSpectrogramPreprocessor( + sample_rate=SAMPLE_RATE, + normalize="per_feature", + window_size=WINDOW_SIZE, + window_stride=WINDOW_STRIDE, + window="hann", + features=NUM_MELS, + n_fft=FFT_SIZE, + log=True, + frame_splicing=1, + dither=1.0e-05, + pad_to=0, + pad_value=0.0, + ).eval() + + num_samples = int(seconds * SAMPLE_RATE) + with torch.no_grad(): + mel, lengths = preprocessor(torch.zeros(1, num_samples), torch.tensor([num_samples])) + + assert mel.shape[1] == NUM_MELS + assert mel.shape[2] == num_mel_frames(num_samples) + assert int(lengths[0]) == num_valid_mel_frames(num_samples) diff --git a/test/python/models/test_lfm2_audio_models.py b/test/python/models/test_lfm2_audio_models.py new file mode 100644 index 0000000000..0ca15ce9fa --- /dev/null +++ b/test/python/models/test_lfm2_audio_models.py @@ -0,0 +1,1258 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License + +""" +Tests for the LFM2-Audio pipeline against the tiny models in test/models/lfm2-audio. + +The models are generated by test/python/create/create_dummy_lfm2_audio_models.py, and the tokenizer +is checked in zipped alongside them, so `_model_path` assembles a runnable model directory in a +temporary directory once per session. The speech model +emits one feature per 8 mel frames, the embedding model scatters those features over the placeholders +the processor writes, and the decoder's logits depend on the conv state, so generation only succeeds, +and only matches between runs, when the processor, the speech/embedding hand-off and the conv state +cache agree. + +Run with: + python -m pytest test/python/models/test_lfm2_audio_models.py --test_models test/models +""" + +from __future__ import annotations + +import json +import math +import os +import shutil +import tempfile +import wave +import zipfile +from pathlib import Path + +import numpy as np +import onnxruntime_genai as og +import pytest +from _test_utils import is_webgpu_ep_available + +AUDIO_MARKER = "<|audio|>" +AUDIO_START_TOKEN_ID = 128 # <|audio_start|>: the answer is spoken from here on +TEXT_END_TOKEN_ID = 130 # <|text_end|>: the text half of an interleaved answer is over +AUDIO_TOKEN_ID = 133 # <|reserved_123|>, model.audio_token_id in the fixture's genai_config.json +NUM_MELS = 128 +SAMPLE_RATE = 16000 +FFT_SIZE = 512 +WIN_LENGTH = 400 +HOP_LENGTH = 160 +SUBSAMPLING_FACTOR = 8 +LOG_EPS = 2.0**-24 +NORM_EPS = 1e-5 +# Settings of LiquidAI/LFM2-Audio-1.5B and LFM2.5-Audio-1.5B (config.json "preprocessor"). +PREPROCESSOR_CONFIG = { + "sample_rate": SAMPLE_RATE, + "normalize": "per_feature", + "window_size": 0.025, + "window_stride": 0.01, + "window": "hann", + "features": NUM_MELS, + "n_fft": FFT_SIZE, + "log": True, + "frame_splicing": 1, + "dither": 1.0e-05, + "pad_to": 0, + "pad_value": 0.0, +} + + +# The tokenizer is checked in zipped: its 4.8 MB of vocabulary compresses to under 1 MB, and no test +# reads it as text. The files come from LiquidAI/LFM2.5-Audio-1.5B; _write_tokenizer_archive rebuilds it. +TOKENIZER_ARCHIVE = "tokenizer.zip" +TOKENIZER_FILES = ("tokenizer.json", "tokenizer_config.json") + + +def _write_tokenizer_archive(directory: Path, archive_path: Path): + """Zips the tokenizer files in `directory`, with the timestamp and mode pinned so that refreshing + the archive from the same files does not show up as a changed binary.""" + with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED, compresslevel=9) as archive: + for name in TOKENIZER_FILES: + info = zipfile.ZipInfo(name, date_time=(1980, 1, 1, 0, 0, 0)) + info.compress_type = zipfile.ZIP_DEFLATED + info.external_attr = 0o644 << 16 + archive.writestr(info, (directory / name).read_bytes(), compresslevel=9) + + +# One assembled directory per source, kept alive (and so undeleted) for the whole session. +_MODEL_DIRECTORIES: dict[str, tempfile.TemporaryDirectory] = {} + + +def _model_path(test_data_path) -> str: + """The fixture model directory, with the tokenizer unpacked, assembled once per session. + + Everything the runtime loads has to sit in one directory, so the ONNX models and the config are + copied next to the unpacked tokenizer rather than the archive being unpacked over the checkout. + """ + source = Path(test_data_path) / "lfm2-audio" + if not (source / "genai_config.json").exists(): + pytest.skip(f"LFM2-Audio test model not found at {source}") + + key = os.fspath(source) + if key not in _MODEL_DIRECTORIES: + holder = tempfile.TemporaryDirectory(prefix="lfm2-audio-") + destination = Path(holder.name) + for item in source.iterdir(): + if item.is_file() and item.name != TOKENIZER_ARCHIVE: + shutil.copy(item, destination) + with zipfile.ZipFile(source / TOKENIZER_ARCHIVE) as archive: + # Named rather than extractall: the archive holds exactly these two files, and naming + # them keeps a stray path in the archive from landing outside the directory. + for name in TOKENIZER_FILES: + (destination / name).write_bytes(archive.read(name)) + _MODEL_DIRECTORIES[key] = holder + return _MODEL_DIRECTORIES[key].name + + +def test_lfm2_audio_tokenizer_archive_can_be_rebuilt(test_data_path, tmp_path): + # Rebuilding from the unpacked files gives the same archive every time, holding what is checked in. + unpacked = Path(_model_path(test_data_path)) + _write_tokenizer_archive(unpacked, tmp_path / "first.zip") + _write_tokenizer_archive(unpacked, tmp_path / "second.zip") + + assert (tmp_path / "first.zip").read_bytes() == (tmp_path / "second.zip").read_bytes() + with ( + zipfile.ZipFile(tmp_path / "first.zip") as rebuilt, + zipfile.ZipFile(Path(test_data_path) / "lfm2-audio" / TOKENIZER_ARCHIVE) as checked_in, + ): + assert rebuilt.namelist() == checked_in.namelist() == list(TOKENIZER_FILES) + for name in TOKENIZER_FILES: + assert rebuilt.read(name) == checked_in.read(name) + assert rebuilt.getinfo(name).date_time == checked_in.getinfo(name).date_time + + +def _copy_model(test_data_path, tmp_path) -> Path: + destination = tmp_path / "lfm2-audio" + shutil.copytree(_model_path(test_data_path), destination) + return destination + + +def _edit_json(path: Path, edit): + config = json.loads(path.read_text()) + edit(config) + path.write_text(json.dumps(config)) + + +def _synthetic_signal(seconds: float, seed: int = 0) -> np.ndarray: + """A few tones plus noise, quantized to 16-bit like the WAV that carries it.""" + rng = np.random.default_rng(seed) + t = np.arange(int(seconds * SAMPLE_RATE)) / SAMPLE_RATE + signal = 0.3 * np.sin(2 * np.pi * 440 * t) + 0.1 * np.sin(2 * np.pi * 3000 * t) + 0.05 * rng.standard_normal(t.size) + return np.round(np.clip(signal, -1, 1) * 32767).astype(np.int16).astype(np.float32) / 32768.0 + + +def _write_wav(path: Path, samples: np.ndarray, sample_rate: int = SAMPLE_RATE, channels: int = 1) -> str: + pcm = np.round(samples * 32768.0).clip(-32768, 32767).astype(" 1: + pcm = np.repeat(pcm[:, None], channels, axis=1) + with wave.open(os.fspath(path), "wb") as handle: + handle.setnchannels(channels) + handle.setsampwidth(2) + handle.setframerate(sample_rate) + handle.writeframes(pcm.tobytes()) + return os.fspath(path) + + +def _num_frames(num_samples: int) -> int: + # torch.stft(center=True): one frame per hop plus one. + return num_samples // HOP_LENGTH + 1 + + +def _num_tokens(num_samples: int) -> int: + return math.ceil(_num_frames(num_samples) / SUBSAMPLING_FACTOR) + + +def _slaney_mel_filterbank() -> np.ndarray: + """librosa.filters.mel(sr, n_fft, n_mels, fmin=0, fmax=sr/2, norm="slaney") without librosa.""" + + def hz_to_mel(hz): + hz = np.asarray(hz, dtype=np.float64) + mel = hz / (200.0 / 3) + log_region = hz >= 1000.0 + mel[log_region] = 15.0 + np.log(hz[log_region] / 1000.0) / (np.log(6.4) / 27.0) + return mel + + def mel_to_hz(mel): + mel = np.asarray(mel, dtype=np.float64) + hz = mel * (200.0 / 3) + log_region = mel >= 15.0 + hz[log_region] = 1000.0 * np.exp((np.log(6.4) / 27.0) * (mel[log_region] - 15.0)) + return hz + + fft_freqs = np.linspace(0, SAMPLE_RATE / 2, FFT_SIZE // 2 + 1) + mel_points = mel_to_hz(np.linspace(hz_to_mel([0.0])[0], hz_to_mel([SAMPLE_RATE / 2])[0], NUM_MELS + 2)) + fdiff = np.diff(mel_points) + ramps = mel_points[:, None] - fft_freqs[None, :] + weights = np.maximum(0, np.minimum(-ramps[:-2] / fdiff[:-1, None], ramps[2:] / fdiff[1:, None])) + weights *= (2.0 / (mel_points[2 : NUM_MELS + 2] - mel_points[:NUM_MELS]))[:, None] + return weights.astype(np.float32) + + +def _reference_mel( + samples: np.ndarray, preemph: float = 0.97, log_eps: float = LOG_EPS, norm_eps: float = NORM_EPS +) -> np.ndarray: + """NeMo's AudioToMelSpectrogramPreprocessor in eval mode, as liquid-audio ships it, in numpy. + + Pre-emphasis, constant center padding, symmetric Hann window, power spectrum, Slaney mel filter + bank, log with the 2^-24 guard, then per-feature normalization over the num_samples // hop + frames get_seq_len counts as valid, with the frames past that zeroed. + """ + preemphasized = np.concatenate([samples[:1], samples[1:] - preemph * samples[:-1]]) + padded = np.pad(preemphasized, FFT_SIZE // 2) + num_frames = _num_frames(len(samples)) + window = np.zeros(FFT_SIZE, np.float64) + window[(FFT_SIZE - WIN_LENGTH) // 2 :][:WIN_LENGTH] = 0.5 * ( + 1 - np.cos(2 * np.pi * np.arange(WIN_LENGTH) / (WIN_LENGTH - 1)) + ) + frames = np.stack([padded[t * HOP_LENGTH : t * HOP_LENGTH + FFT_SIZE] * window for t in range(num_frames)]) + power = np.abs(np.fft.rfft(frames, axis=1)) ** 2 + log_mel = np.log(power @ _slaney_mel_filterbank().T.astype(np.float64) + log_eps) + + valid = len(samples) // HOP_LENGTH + mean = log_mel[:valid].mean(axis=0) + std = np.sqrt(((log_mel[:valid] - mean) ** 2).sum(axis=0) / (valid - 1)) + norm_eps + normalized = (log_mel - mean) / std + normalized[valid:] = 0.0 + return normalized.astype(np.float32) + + +def _reference_resample(samples: np.ndarray, source_rate: int, target_rate: int = SAMPLE_RATE) -> np.ndarray: + """torchaudio.functional.resample with its defaults, which liquid-audio's ChatState.add_audio uses. + + A Hann-windowed sinc with lowpass_filter_width 6 and rolloff 0.99: one kernel per output phase, + built in float64 and applied in float32, strided over the zero-padded input. + """ + if source_rate == target_rate: + return samples + divisor = math.gcd(source_rate, target_rate) + orig, new = source_rate // divisor, target_rate // divisor + base_freq = min(orig, new) * 0.99 + width = math.ceil(6 * orig / base_freq) + idx = np.arange(-width, width + orig, dtype=np.float64) / orig + t = (np.arange(0, -new, -1, dtype=np.float64)[:, None] / new + idx[None, :]) * base_freq + t = t.clip(-6, 6) + window = np.cos(t * np.pi / 6 / 2) ** 2 + t *= np.pi + kernels = (np.where(t == 0, 1.0, np.sin(t) / np.where(t == 0, 1.0, t)) * window * base_freq / orig).astype( + np.float32 + ) + + padded = np.pad(samples.astype(np.float32), (width, width + orig)) + steps = len(samples) // orig + 1 + frames = np.lib.stride_tricks.sliding_window_view(padded, 2 * width + orig)[::orig][:steps] + return (frames @ kernels.T).reshape(-1)[: math.ceil(new * len(samples) / orig)] + + +def _process(model_path: str, prompt: str, audios): + model = og.Model(model_path) + inputs = model.create_multimodal_processor()(prompt, audios=audios) + return model, inputs + + +def _placeholder_runs(input_ids: np.ndarray) -> list[int]: + """Lengths of the consecutive AUDIO_TOKEN_ID runs, in order.""" + runs = [] + run = 0 + for token in input_ids: + if token == AUDIO_TOKEN_ID: + run += 1 + elif run: + runs.append(run) + run = 0 + if run: + runs.append(run) + return runs + + +def test_lfm2_audio_processor_single_clip(test_data_path, tmp_path): + samples = _synthetic_signal(1.3) + clip = _write_wav(tmp_path / "clip.wav", samples) + prompt = f"<|startoftext|><|im_start|>user\n{AUDIO_MARKER}<|im_end|>\n<|im_start|>assistant\n" + model, inputs = _process(_model_path(test_data_path), prompt, og.Audios.open(clip)) + tokenizer = og.Tokenizer(model) + + # Same token stream as liquid-audio's ChatState: the text on either side of the clip is tokenized + # on its own and one placeholder per encoder frame sits between. + num_frames = _num_frames(len(samples)) + expected_ids = ( + tokenizer.encode("<|startoftext|><|im_start|>user\n").tolist() + + [AUDIO_TOKEN_ID] * _num_tokens(len(samples)) + + tokenizer.encode("<|im_end|>\n<|im_start|>assistant\n").tolist() + ) + input_ids = inputs["input_ids"].as_numpy() + assert input_ids.shape == (1, len(expected_ids)) + assert input_ids[0].tolist() == expected_ids + + mel = inputs["audio_embeds"].as_numpy() + assert mel.shape == (1, num_frames, NUM_MELS) + assert mel.dtype == np.float32 + np.testing.assert_array_equal(inputs["audio_lengths"].as_numpy(), [num_frames]) + np.testing.assert_array_equal(inputs["audio_sizes"].as_numpy(), [_num_tokens(len(samples))]) + + +def test_lfm2_audio_processor_text_only(test_data_path): + _, inputs = _process(_model_path(test_data_path), "<|startoftext|>Just text.", None) + assert AUDIO_TOKEN_ID not in inputs["input_ids"].as_numpy() + np.testing.assert_array_equal(inputs["audio_sizes"].as_numpy(), [0]) + with pytest.raises(RuntimeError, match="audio_embeds"): + _ = inputs["audio_embeds"] + + +def test_lfm2_audio_processor_two_clips_of_different_lengths(test_data_path, tmp_path): + short = _synthetic_signal(0.5, seed=1) + long = _synthetic_signal(1.7, seed=2) + clips = [_write_wav(tmp_path / "short.wav", short), _write_wav(tmp_path / "long.wav", long)] + prompt = f"First {AUDIO_MARKER} then {AUDIO_MARKER} done" + _, inputs = _process(_model_path(test_data_path), prompt, og.Audios.open(*clips)) + + # Clips are staged as one zero-padded tensor, each with its real length, in prompt order; the + # placeholders follow the same order. The encoder itself runs once per clip on those lengths. + assert _placeholder_runs(inputs["input_ids"].as_numpy()[0]) == [_num_tokens(len(short)), _num_tokens(len(long))] + mel = inputs["audio_embeds"].as_numpy() + assert mel.shape == (2, _num_frames(len(long)), NUM_MELS) + np.testing.assert_array_equal(inputs["audio_lengths"].as_numpy(), [_num_frames(len(short)), _num_frames(len(long))]) + np.testing.assert_array_equal(inputs["audio_sizes"].as_numpy(), [_num_tokens(len(short)), _num_tokens(len(long))]) + assert np.all(mel[0, _num_frames(len(short)) :] == 0) + np.testing.assert_allclose(mel[0, : _num_frames(len(short))], _reference_mel(short), atol=2e-3) + np.testing.assert_allclose(mel[1], _reference_mel(long), atol=2e-3) + + +@pytest.mark.parametrize( + "prompt,num_clips,expected", + [ + ("No marker here", 1, "contains 0 <|audio|> markers but 1 audio clips"), + (f"{AUDIO_MARKER} twice {AUDIO_MARKER}", 1, "contains 2 <|audio|> markers but 1 audio clips"), + (f"marker without a clip {AUDIO_MARKER}", 0, "contains 1 <|audio|> markers but 0 audio clips"), + ], +) +def test_lfm2_audio_rejects_marker_count_mismatch(test_data_path, tmp_path, prompt, num_clips, expected): + clips = [_write_wav(tmp_path / f"{i}.wav", _synthetic_signal(0.4, seed=i)) for i in range(num_clips)] + audios = og.Audios.open(*clips) if clips else None + with pytest.raises(RuntimeError, match=expected): + _process(_model_path(test_data_path), prompt, audios) + + +def test_lfm2_audio_open_bytes_matches_open(test_data_path, tmp_path): + clip = Path(_write_wav(tmp_path / "clip.wav", _synthetic_signal(0.9, seed=7))) + model = og.Model(_model_path(test_data_path)) + processor = model.create_multimodal_processor() + + from_path = processor(AUDIO_MARKER, audios=og.Audios.open(os.fspath(clip))) + from_bytes = processor(AUDIO_MARKER, audios=og.Audios.open_bytes(clip.read_bytes())) + + np.testing.assert_array_equal(from_bytes["input_ids"].as_numpy(), from_path["input_ids"].as_numpy()) + np.testing.assert_array_equal(from_bytes["audio_embeds"].as_numpy(), from_path["audio_embeds"].as_numpy()) + np.testing.assert_array_equal(from_bytes["audio_sizes"].as_numpy(), from_path["audio_sizes"].as_numpy()) + + +def _signal_at(source_rate: int, seconds: float, seed: int) -> np.ndarray: + """`_synthetic_signal` stretched to another rate, quantized to 16-bit like the WAV that carries it.""" + base = _synthetic_signal(seconds, seed=seed) + stretched = np.interp(np.linspace(0, 1, int(seconds * source_rate)), np.linspace(0, 1, base.size), base) + return (np.round(stretched * 32768.0).clip(-32768, 32767) / 32768.0).astype(np.float32) + + +@pytest.mark.parametrize("source_rate", [8000, 22050, 24000, 44100, 48000]) +def test_lfm2_audio_resamples_the_way_the_reference_does(test_data_path, tmp_path, source_rate): + # The clip is decoded at its own rate and brought to the encoder's 16 kHz with torchaudio's + # resampler, upwards as well as downwards, as ChatState.add_audio does. The decoder's own + # resampler is a different filter: near enough to transcribe, far enough to change a token. + samples = _signal_at(source_rate, 0.75, seed=8) + clip = _write_wav(tmp_path / "clip.wav", samples, sample_rate=source_rate) + + _, inputs = _process(_model_path(test_data_path), AUDIO_MARKER, og.Audios.open(clip)) + + resampled = _reference_resample(samples, source_rate) + assert len(resampled) == math.ceil(len(samples) * SAMPLE_RATE / source_rate) + np.testing.assert_array_equal(inputs["audio_lengths"].as_numpy(), [_num_frames(len(resampled))]) + np.testing.assert_array_equal(inputs["audio_sizes"].as_numpy(), [_num_tokens(len(resampled))]) + np.testing.assert_allclose(inputs["audio_embeds"].as_numpy()[0], _reference_mel(resampled), atol=2e-3) + + +@pytest.mark.parametrize("source_rate", [8000, 44100]) +def test_lfm2_audio_reference_resampler_is_torchaudio(source_rate): + """The numpy resampler above against the real one, when torchaudio is installed.""" + torch = pytest.importorskip("torch") + torchaudio = pytest.importorskip("torchaudio") + samples = _signal_at(source_rate, 0.5, seed=14) + + expected = torchaudio.functional.resample(torch.from_numpy(samples)[None], source_rate, SAMPLE_RATE)[0].numpy() + + # Float32 sums taken in a different order. + np.testing.assert_allclose(_reference_resample(samples, source_rate), expected, atol=1e-5) + + +@pytest.mark.parametrize( + "field,value", + [("hop_length", 320), ("win_length", 320), ("fft_size", 1024), ("num_mels", 64), ("sample_rate", 8000)], +) +def test_lfm2_audio_honors_the_mel_config_overrides(test_data_path, tmp_path, field, value): + # The front end defaults to the published checkpoints' settings, and genai_config.json can + # override each one; a fine-tune that changed them would otherwise be silently mis-framed. + model_path = _copy_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: config["model"].update({field: value})) + samples = _synthetic_signal(1.0, seed=9) + clip = _write_wav(tmp_path / "clip.wav", samples) + + _, inputs = _process(os.fspath(model_path), AUDIO_MARKER, og.Audios.open(clip)) + mel = inputs["audio_embeds"].as_numpy() + + expected_mels = value if field == "num_mels" else NUM_MELS + expected_hop = value if field == "hop_length" else HOP_LENGTH + # A different sample rate resamples the clip instead of changing the framing arithmetic. + expected_samples = len(samples) * value // SAMPLE_RATE if field == "sample_rate" else len(samples) + expected_frames = expected_samples // expected_hop + 1 + + assert mel.shape[2] == expected_mels + assert abs(mel.shape[1] - expected_frames) <= 1, f"{mel.shape[1]} frames, expected about {expected_frames}" + np.testing.assert_array_equal(inputs["audio_lengths"].as_numpy(), [mel.shape[1]]) + + +@pytest.mark.parametrize( + "field,value", + [("preemph", 0.5), ("log_eps", 1e-4), ("norm_eps", 0.5)], +) +def test_lfm2_audio_honors_the_mel_value_overrides(test_data_path, tmp_path, field, value): + # These three change the numbers rather than the shape, so a default silently ignoring the + # override would look perfectly healthy. Check them against the same arithmetic done here. + model_path = _copy_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: config["model"].update({field: value})) + samples = _synthetic_signal(1.0, seed=12) + clip = _write_wav(tmp_path / "clip.wav", samples) + + _, inputs = _process(os.fspath(model_path), AUDIO_MARKER, og.Audios.open(clip)) + mel = inputs["audio_embeds"].as_numpy()[0] + + np.testing.assert_allclose(mel, _reference_mel(samples, **{field: value}), atol=2e-3) + # And it is genuinely a different front end from the default one. + assert not np.allclose(mel, _reference_mel(samples), atol=2e-3) + + +def test_lfm2_audio_subsampling_factor_override_changes_the_placeholder_count(test_data_path, tmp_path): + model_path = _copy_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: config["model"].update({"subsampling_factor": 4})) + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(1.0, seed=10)) + + _, inputs = _process(os.fspath(model_path), AUDIO_MARKER, og.Audios.open(clip)) + + num_frames = int(inputs["audio_lengths"].as_numpy()[0]) + assert int(inputs["audio_sizes"].as_numpy()[0]) == math.ceil(num_frames / 4) + assert _placeholder_runs(inputs["input_ids"].as_numpy()[0]) == [math.ceil(num_frames / 4)] + + +def test_lfm2_audio_rejects_a_clip_too_short_to_normalize(test_data_path, tmp_path): + # Per-feature normalization needs two valid frames (num_samples // hop >= 2). + clip = _write_wav(tmp_path / "blip.wav", _synthetic_signal(0.015)) + with pytest.raises(RuntimeError, match="too short"): + _process(_model_path(test_data_path), AUDIO_MARKER, og.Audios.open(clip)) + + +def test_lfm2_audio_single_prompt_list_matches_the_string_form(test_data_path, tmp_path): + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.6)) + model = og.Model(_model_path(test_data_path)) + processor = model.create_multimodal_processor() + prompt = f"Transcribe {AUDIO_MARKER}" + from_string = processor(prompt, audios=og.Audios.open(clip)) + from_list = processor([prompt], audios=og.Audios.open(clip)) + np.testing.assert_array_equal(from_list["input_ids"].as_numpy(), from_string["input_ids"].as_numpy()) + np.testing.assert_array_equal(from_list["audio_embeds"].as_numpy(), from_string["audio_embeds"].as_numpy()) + + +def test_lfm2_audio_rejects_batched_prompts(test_data_path): + model = og.Model(_model_path(test_data_path)) + with pytest.raises(RuntimeError, match="batched prompts are not supported"): + model.create_multimodal_processor()(["one", "two"], audios=None) + + +def test_lfm2_audio_mel_matches_reference(test_data_path, tmp_path): + samples = _synthetic_signal(2.0, seed=3) + clip = _write_wav(tmp_path / "clip.wav", samples) + _, inputs = _process(_model_path(test_data_path), AUDIO_MARKER, og.Audios.open(clip)) + mel = inputs["audio_embeds"].as_numpy()[0] + + reference = _reference_mel(samples) + assert mel.shape == reference.shape + np.testing.assert_allclose(mel, reference, atol=2e-3) + # NeMo counts one frame fewer than the STFT produces; that last frame is zeroed but still fed. + assert np.all(mel[-1] == 0) + assert np.any(mel[-2] != 0) + + +def test_lfm2_audio_more_than_two_channels_is_not_downmixed(test_data_path, tmp_path): + # The decoder downmixes one and two channels correctly, but deinterleaves anything wider as + # though it were stereo, so a six channel clip arrives three times too long and garbled. Nothing + # in the decoder's API reports the channel count, so this cannot be caught here; the tutorial + # says to downmix first. This pins the behaviour so a fix upstream shows up as a failure here. + samples = _synthetic_signal(0.6, seed=13) + frames = {} + for channels in (1, 2, 6): + clip = _write_wav(tmp_path / f"c{channels}.wav", samples, channels=channels) + _, inputs = _process(_model_path(test_data_path), AUDIO_MARKER, og.Audios.open(clip)) + frames[channels] = int(inputs["audio_lengths"].as_numpy()[0]) + + assert frames[2] == frames[1], "two channels are downmixed, so the clip keeps its length" + # Compare the hops rather than the frames: every clip gets one extra frame from the centre + # padding, so only `frames - 1` scales with the length. + assert frames[6] - 1 == (frames[1] - 1) * 3, ( + f"six channels still arrive as {frames[6]} frames against {frames[1]}; if these now match, " + "the decoder has learned to downmix and the tutorial's warning can go" + ) + + +def test_lfm2_audio_downmixes_stereo_to_mono(test_data_path, tmp_path): + samples = _synthetic_signal(0.8, seed=4) + mono = _write_wav(tmp_path / "mono.wav", samples) + stereo = _write_wav(tmp_path / "stereo.wav", samples, channels=2) + model_path = _model_path(test_data_path) + _, from_mono = _process(model_path, AUDIO_MARKER, og.Audios.open(mono)) + _, from_stereo = _process(model_path, AUDIO_MARKER, og.Audios.open(stereo)) + np.testing.assert_allclose(from_stereo["audio_embeds"].as_numpy(), from_mono["audio_embeds"].as_numpy(), atol=1e-4) + + +def test_lfm2_audio_mel_matches_liquid_audio(test_data_path, tmp_path): + """The reference implementation itself, when the liquid-audio package is installed.""" + torch = pytest.importorskip("torch") + processor_module = pytest.importorskip("liquid_audio.model.conformer.processor") + reference_processor = processor_module.AudioToMelSpectrogramPreprocessor(**PREPROCESSOR_CONFIG).eval() + + samples = _synthetic_signal(1.5, seed=5) + clip = _write_wav(tmp_path / "clip.wav", samples) + _, inputs = _process(_model_path(test_data_path), AUDIO_MARKER, og.Audios.open(clip)) + mel = inputs["audio_embeds"].as_numpy()[0] + + with torch.no_grad(): + wave_tensor = torch.from_numpy(samples)[None, :] + reference, _ = reference_processor(wave_tensor, torch.tensor([wave_tensor.shape[1]])) + np.testing.assert_allclose(mel, reference[0].T.numpy(), atol=2e-3) + + +def test_lfm2_audio_casts_mel_to_the_speech_input_type(test_data_path, tmp_path): + onnx = pytest.importorskip("onnx") + model_path = _copy_model(test_data_path, tmp_path) + speech_path = model_path / "dummy_speech.onnx" + speech = onnx.load(speech_path) + # Take mel_spectrogram as float16 and cast it back to float32 for the rest of the graph. + speech.graph.input[0].type.tensor_type.elem_type = onnx.TensorProto.FLOAT16 + for node in speech.graph.node: + node.input[:] = ["mel_spectrogram_fp16" if name == "mel_spectrogram" else name for name in node.input] + speech.graph.node.insert( + 0, onnx.helper.make_node("Cast", ["mel_spectrogram"], ["mel_spectrogram_fp16"], to=onnx.TensorProto.FLOAT) + ) + onnx.save(speech, speech_path) + + samples = _synthetic_signal(0.7, seed=6) + clip = _write_wav(tmp_path / "clip.wav", samples) + _, inputs = _process(os.fspath(model_path), AUDIO_MARKER, og.Audios.open(clip)) + mel = inputs["audio_embeds"].as_numpy() + assert mel.dtype == np.float16 + np.testing.assert_allclose(mel[0].astype(np.float32), _reference_mel(samples), atol=2e-2) + + # And the fp16 features still flow through generation. + generator, _ = _generate(os.fspath(model_path), AUDIO_MARKER, og.Audios.open(clip), num_tokens=2) + assert len(generator.get_sequence(0)) == mel.shape[1] // SUBSAMPLING_FACTOR + 1 + 2 + + +@pytest.mark.parametrize("field", ["audio_embeds", "audio_lengths"]) +def test_lfm2_audio_rejects_speech_input_name_missing_from_the_graph(test_data_path, tmp_path, field): + model_path = _copy_model(test_data_path, tmp_path) + _edit_json( + model_path / "genai_config.json", lambda config: config["model"]["speech"]["inputs"].update({field: "nope"}) + ) + with pytest.raises(RuntimeError, match=f'no input named "nope".*model.speech.inputs.{field}'): + og.Model(os.fspath(model_path)).create_multimodal_processor() + + +def test_lfm2_audio_rejects_non_int64_lengths(test_data_path, tmp_path): + onnx = pytest.importorskip("onnx") + model_path = _copy_model(test_data_path, tmp_path) + speech_path = model_path / "dummy_speech.onnx" + speech = onnx.load(speech_path) + # Take mel_lengths as int32 and widen it for the rest of the graph. + speech.graph.input[1].type.tensor_type.elem_type = onnx.TensorProto.INT32 + for node in speech.graph.node: + node.input[:] = ["mel_lengths_i64" if name == "mel_lengths" else name for name in node.input] + speech.graph.node.insert( + 0, onnx.helper.make_node("Cast", ["mel_lengths"], ["mel_lengths_i64"], to=onnx.TensorProto.INT64) + ) + onnx.save(speech, speech_path) + with pytest.raises(RuntimeError, match='"mel_lengths" must be int64, got int32'): + og.Model(os.fspath(model_path)).create_multimodal_processor() + + +def test_lfm2_audio_requires_audio_token_id(test_data_path, tmp_path): + model_path = _copy_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: config["model"].pop("audio_token_id")) + with pytest.raises(RuntimeError, match="audio_token_id must be set"): + og.Model(os.fspath(model_path)).create_multimodal_processor() + + +def test_lfm2_audio_requires_audio_token_id_inside_the_vocabulary(test_data_path, tmp_path): + model_path = _copy_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: config["model"].update({"audio_token_id": 65536})) + with pytest.raises(RuntimeError, match="audio_token_id 65536 is outside the vocabulary of 65536 tokens"): + og.Model(os.fspath(model_path)).create_multimodal_processor() + + +@pytest.mark.skipif(not og.is_cuda_available(), reason="needs a CUDA decoder next to the CPU sub-models") +def test_lfm2_audio_rejects_cpu_sub_models_next_to_a_cuda_decoder(test_data_path, tmp_path): + # The fixture's speech and embedding entries have session_options of their own, as a CPU export + # writes them, so only the decoder moves to CUDA: the embedding session would be handed its inputs + # in CUDA memory and write host bytes over them. + config = og.Config(os.fspath(_copy_model(test_data_path, tmp_path))) + config.clear_providers() + config.append_provider("cuda") + with pytest.raises(RuntimeError, match="embedding.session_options run the embedding model on CPU"): + og.Model(config) + + +def _run_prompt(model, prompt, audios, num_tokens): + inputs = model.create_multimodal_processor()(prompt, audios=audios) + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=inputs["input_ids"].as_numpy().shape[1] + num_tokens) + generator = og.Generator(model, params) + generator.set_inputs(inputs) + while not generator.is_done(): + generator.generate_next_token() + return generator.get_sequence(0) + + +def _model_on(model_path, provider): + config = og.Config(os.fspath(model_path)) + config.clear_providers() + config.append_provider(provider) + return og.Model(config) + + +@pytest.mark.skipif(not og.is_cuda_available(), reason="needs a CUDA decoder next to a CPU encoder") +def test_lfm2_audio_refuses_audio_for_a_cpu_encoder_next_to_a_cuda_decoder(test_data_path, tmp_path): + # With the embedding following the decoder only the encoder is left on CPU. Text-only prompts never + # run it; a prompt with audio is refused before it writes its features into CUDA memory. + model_path = _copy_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: config["model"]["embedding"].pop("session_options")) + model = _model_on(model_path, "cuda") + + prompt = "<|startoftext|>Transcribe. " + assert len(_run_prompt(model, prompt, None, 4)) > 0 + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5, seed=72)) + with pytest.raises( + RuntimeError, + match="speech.session_options run the speech model on CPU, but the audio features are passed in CUDA memory", + ): + _run_prompt(model, prompt + AUDIO_MARKER, og.Audios.open(clip), 4) + + +@pytest.mark.skipif(not is_webgpu_ep_available(), reason="needs the WebGPU plug-in EP") +def test_lfm2_audio_keeps_text_prompts_on_webgpu_with_cpu_sub_models(test_data_path, tmp_path): + # WebGPU without graph capture keeps the decoder's inputs on the host, so a CPU export's config still + # runs text-only prompts. Only a prompt with audio, whose features are in WebGPU memory, is refused. + model = _model_on(_copy_model(test_data_path, tmp_path), "webgpu") + + prompt = "<|startoftext|>Transcribe. " + assert len(_run_prompt(model, prompt, None, 4)) > 0 + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5, seed=73)) + with pytest.raises( + RuntimeError, + match="speech.session_options run the speech model on CPU, but the audio features are passed in WebGPU memory", + ): + _run_prompt(model, prompt + AUDIO_MARKER, og.Audios.open(clip), 4) + + +def _generate(model_path: str, prompt: str, audios, num_tokens: int, chunk_size: int | None = None): + model = og.Model(model_path) + processor = model.create_multimodal_processor() + params = og.GeneratorParams(model) + options = {"do_sample": False, "max_length": 4096} + if chunk_size is not None: + options["chunk_size"] = chunk_size + params.set_search_options(**options) + + inputs = processor(prompt, audios=audios) + generator = og.Generator(model, params) + generator.set_inputs(inputs) + for _ in range(num_tokens): + generator.generate_next_token() + return generator, inputs + + +def _reference_generation(model_path: str, inputs, num_tokens: int) -> np.ndarray: + """Greedy generation replayed with plain ONNX Runtime. + + Runs the speech model once per clip on that clip's own frames — the padding never reaches the + encoder, which is what the runtime does and what the published encoder export requires — then + scatters the concatenated features through the embedding model, and finally one decoder run per + token carrying the KV cache and conv state by hand. + """ + ort = pytest.importorskip("onnxruntime") + config = json.loads((Path(model_path) / "genai_config.json").read_text())["model"] + decoder_config = config["decoder"] + speech = ort.InferenceSession(os.path.join(model_path, config["speech"]["filename"])) + embedding = ort.InferenceSession(os.path.join(model_path, config["embedding"]["filename"])) + decoder = ort.InferenceSession(os.path.join(model_path, decoder_config["filename"])) + output_names = [output.name for output in decoder.get_outputs()] + hidden_size = decoder_config["hidden_size"] + + input_ids = inputs["input_ids"].as_numpy()[0] + tokens_per_clip = inputs["audio_sizes"].as_numpy() + if tokens_per_clip.sum() > 0: + mel = inputs["audio_embeds"].as_numpy() + frames_per_clip = inputs["audio_lengths"].as_numpy() + per_clip = [ + speech.run(None, {"mel_spectrogram": mel[i : i + 1, :frames], "mel_lengths": frames_per_clip[i : i + 1]})[ + 0 + ][0] + for i, frames in enumerate(frames_per_clip) + ] + audio_features = np.concatenate( + [features[:count] for features, count in zip(per_clip, tokens_per_clip, strict=True)] + ) + else: + audio_features = np.zeros((0, hidden_size), np.float32) + + state = {} + for layer, layer_type in enumerate(decoder_config["layer_types"]): + if layer_type == "conv": + state[f"past.{layer}.conv"] = np.zeros((1, hidden_size, decoder_config["conv_cache_size"]), np.float32) + else: + for kind in ("key", "value"): + state[f"past_key_values.{layer}.{kind}"] = np.zeros( + (1, decoder_config["num_key_value_heads"], 0, decoder_config["head_size"]), np.float32 + ) + + sequence = [int(token) for token in input_ids] + step_ids = sequence + for _ in range(num_tokens): + embeds = embedding.run( + None, {"input_ids": np.array([step_ids], dtype=np.int64), "audio_features": audio_features} + )[0] + feeds = {"inputs_embeds": embeds, "attention_mask": np.ones((1, len(sequence)), np.int64), **state} + outputs = dict(zip(output_names, decoder.run(None, feeds), strict=True)) + for name in state: + state[name] = outputs[name.replace("past_key_values.", "present.").replace("past.", "present.")] + next_token = int(outputs["logits"][0, -1].argmax()) + sequence.append(next_token) + step_ids = [next_token] + audio_features = np.zeros((0, hidden_size), np.float32) + return np.array(sequence, dtype=np.int32) + + +@pytest.mark.parametrize( + "num_clips,chunk_size", + [(0, None), (0, 3), (1, None), (1, 3), (1, 10), (2, None), (2, 5)], + ids=["text", "text-chunk3", "audio", "audio-chunk3", "audio-chunk10", "two-clips", "two-clips-chunk5"], +) +def test_lfm2_audio_generation_matches_onnxruntime_reference(test_data_path, tmp_path, num_clips, chunk_size): + # The decoder's logits read the conv state over the last few embeddings, so the prompt ends with + # the audio: the first generated tokens depend on the scattered features, and the runtime only + # matches the hand-carried reference if it passes both the features and the conv state along, + # including across prefill chunk boundaries. The clips run longest first, so the last one is the + # one the mel batch pads: feeding the encoder that padding instead of the clip's own frames + # changes its final feature, and with it the first token generated. + model_path = _model_path(test_data_path) + clips = [_write_wav(tmp_path / f"{i}.wav", _synthetic_signal(1.7 - 0.7 * i, seed=10 + i)) for i in range(num_clips)] + prompt = "<|startoftext|>Transcribe the following. " + " and ".join([AUDIO_MARKER] * num_clips) + num_tokens = 6 + + generator, inputs = _generate(model_path, prompt, og.Audios.open(*clips) if clips else None, num_tokens, chunk_size) + input_ids = inputs["input_ids"].as_numpy()[0] + assert chunk_size is None or len(input_ids) > chunk_size, "the prompt must span several chunks" + + np.testing.assert_array_equal(generator.get_sequence(0), _reference_generation(model_path, inputs, num_tokens)) + + +def test_lfm2_audio_generation_depends_on_the_audio(test_data_path, tmp_path): + model_path = _model_path(test_data_path) + outputs = [] + for seed in (20, 21): + clip = _write_wav(tmp_path / f"{seed}.wav", _synthetic_signal(0.5, seed=seed)) + generator, inputs = _generate(model_path, f"<|startoftext|>{AUDIO_MARKER}", og.Audios.open(clip), num_tokens=4) + outputs.append(generator.get_sequence(0)[len(inputs["input_ids"].as_numpy()[0]) :].tolist()) + assert outputs[0] != outputs[1], "two different clips must not decode to the same tokens" + + +def test_lfm2_audio_features_follow_the_clip_order(test_data_path, tmp_path): + # Two clips of the same length, so a swapped concatenation would still have the right shape. + model_path = _model_path(test_data_path) + clips = [_write_wav(tmp_path / f"{i}.wav", _synthetic_signal(1.0, seed=30 + i)) for i in range(2)] + prompt = "<|startoftext|>Compare {} and {}" + + outputs = [] + for order in (clips, clips[::-1]): + generator, inputs = _generate( + model_path, prompt.format(AUDIO_MARKER, AUDIO_MARKER), og.Audios.open(*order), num_tokens=4 + ) + start = len(inputs["input_ids"].as_numpy()[0]) + outputs.append(generator.get_sequence(0)[start:].tolist()) + assert outputs[0] != outputs[1], "swapping the clips must change what the decoder sees" + + +def _generate_with_eos(model_path, clip, eos, num_tokens): + """Greedy generation from a copied model directory, stopping on the token ids in `eos`.""" + _edit_json(model_path / "genai_config.json", lambda config: config["model"].update({"eos_token_id": eos})) + model = og.Model(os.fspath(model_path)) + processor = model.create_multimodal_processor() + inputs = processor(AUDIO_MARKER, audios=og.Audios.open(clip)) + prompt_length = inputs["input_ids"].as_numpy().shape[1] + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=prompt_length + num_tokens) + generator = og.Generator(model, params) + generator.set_inputs(inputs) + generated = [] + while not generator.is_done(): + generator.generate_next_token() + generated.append(int(generator.get_next_tokens()[0])) + return generated + + +def test_lfm2_audio_generation_stops_when_the_model_starts_speaking(test_data_path, tmp_path): + # Everything the model emits after <|audio_start|> is audio codes for a depthformer this runtime + # does not have, so the builder puts that token in eos_token_id and generation has to end on it. + # The fixture decoder never emits 128 on its own, so this pins the mechanism on a token it does + # emit: whichever id is named as the audio stop, the run ends at its first occurrence. + model_path = _copy_model(test_data_path, tmp_path) + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5, seed=40)) + + unrestricted = _generate_with_eos(model_path, clip, [7], num_tokens=40) + assert len(unrestricted) == 40, "the fixture should run to max_length without a stop token" + + stop_token = unrestricted[5] + first = unrestricted.index(stop_token) + + stopped = _generate_with_eos(model_path, clip, [stop_token], num_tokens=40) + + assert stopped == unrestricted[: first + 1], "generation must end at the stop token, with nothing after it" + + +def test_lfm2_audio_modality_stops_leave_normal_stopping_alone(test_data_path, tmp_path): + # Adding the two modality-switch tokens alongside <|im_end|> must not change a run that stays in + # text the whole way. + model_path = _copy_model(test_data_path, tmp_path) + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5, seed=41)) + switches = [AUDIO_START_TOKEN_ID, TEXT_END_TOKEN_ID] + + without = _generate_with_eos(model_path, clip, [7], num_tokens=25) + with_switch_stops = _generate_with_eos(model_path, clip, [7, *switches], num_tokens=25) + + assert not set(switches) & set(without), "the fixture is not expected to emit the switch tokens" + assert with_switch_stops == without + + +def test_lfm2_audio_fixture_config_is_what_the_builder_writes(test_data_path): + # The tiny model stands in for a built one; its stop tokens have to match what the builder emits + # for these checkpoints, or the tests above would be exercising a different contract. + config = json.loads((Path(_model_path(test_data_path)) / "genai_config.json").read_text()) + assert config["model"]["type"] == "lfm2_audio" + assert config["model"]["audio_token_id"] == AUDIO_TOKEN_ID + assert config["model"]["eos_token_id"] == [7, AUDIO_START_TOKEN_ID, TEXT_END_TOKEN_ID] + + +# Speech output. The fixture's depthformer and audio embedding have the contract of LiquidAI's exports, +# and its decoder has the hidden_states output they read; model.audio_output switches them on. +NUM_CODEBOOKS = 8 +CODEBOOK_SIZE = 2049 +INTERLEAVED_N_TEXT = 6 +INTERLEAVED_N_AUDIO = 12 + + +def _speech_model(test_data_path, tmp_path, **audio_output) -> Path: + model_path = _copy_model(test_data_path, tmp_path) + + def edit(config): + # The modality-switch tokens must not be stop tokens once there is speech to switch to. + config["model"]["eos_token_id"] = 7 + config["model"]["audio_output"] = { + "depthformer": {"filename": "dummy_depthformer.onnx"}, + "embedding": {"filename": "dummy_audio_embedding.onnx"}, + **audio_output, + } + + _edit_json(model_path / "genai_config.json", edit) + return model_path + + +def _generate_speech(model_path, prompt, audios, num_items, **search): + model = og.Model(os.fspath(model_path)) + inputs = model.create_multimodal_processor()(prompt, audios=audios) + prompt_length = inputs["input_ids"].as_numpy().shape[1] + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=prompt_length + num_items, **search) + generator = og.Generator(model, params) + generator.set_inputs(inputs) + while not generator.is_done(): + generator.generate_next_token() + return inputs, generator.get_sequence(0)[prompt_length:].tolist(), generator.get_output("audio_codes") + + +def _reference_speech(model_path, inputs, num_items, interleaved, audio_start=128, codebook_size=CODEBOOK_SIZE): + """liquid-audio's generate_sequential / generate_interleaved over the fixture, audio taken greedily. + + Returns the generated sequence, with AUDIO_TOKEN_ID standing for each audio frame as the runtime + records it, and the frames themselves without the end-of-audio ones. + """ + ort = pytest.importorskip("onnxruntime") + config = json.loads((Path(model_path) / "genai_config.json").read_text())["model"] + decoder_config = config["decoder"] + session = lambda name: ort.InferenceSession(os.path.join(model_path, name)) # noqa: E731 + speech, embedding, decoder = (session(config[k]["filename"]) for k in ("speech", "embedding", "decoder")) + depthformer, audio_embedding = session("dummy_depthformer.onnx"), session("dummy_audio_embedding.onnx") + hidden_size = decoder_config["hidden_size"] + no_features = np.zeros((0, hidden_size), np.float32) + + features = no_features + if inputs["audio_sizes"].as_numpy().sum() > 0: + mel, frames = inputs["audio_embeds"].as_numpy(), inputs["audio_lengths"].as_numpy() + features = np.concatenate( + [ + speech.run(None, {"mel_spectrogram": mel[i : i + 1, :n], "mel_lengths": frames[i : i + 1]})[0][0] + for i, n in enumerate(frames) + ] + ) + embeds = embedding.run( + None, {"input_ids": inputs["input_ids"].as_numpy().astype(np.int64), "audio_features": features} + )[0] + + state = {} + for layer, layer_type in enumerate(decoder_config["layer_types"]): + if layer_type == "conv": + state[f"past.{layer}.conv"] = np.zeros((1, hidden_size, decoder_config["conv_cache_size"]), np.float32) + else: + for kind in ("key", "value"): + state[f"past_key_values.{layer}.{kind}"] = np.zeros( + (1, decoder_config["num_key_value_heads"], 0, decoder_config["head_size"]), np.float32 + ) + output_names = [output.name for output in decoder.get_outputs()] + + def sample_frame(hidden): + keys = np.zeros((1, 1, 1, 0, 16), np.float32) + values, slices, previous, frame = keys.copy(), np.zeros((1, NUM_CODEBOOKS, 16), np.float32), 0, [] + for step in range(NUM_CODEBOOKS): + logits, new_slices, keys, values = depthformer.run( + ["logits", "depth_slices", "new_keys", "new_values"], + { + "hidden_states": hidden[None], + "depth_slices_in": slices, + "step_idx": np.array(step, np.int64), + "prev_token": np.array([previous], np.int64), + "past_keys": keys, + "past_values": values, + "seqlens_k": np.array([step], np.int32), + "total_seq_len": np.array(step + 1, np.int32), + }, + ) + slices = new_slices if step == 0 else slices + previous = int(logits[0, :codebook_size].argmax()) + frame.append(previous) + return np.array(frame, np.int64) + + sequence, frames, total = [], [], 0 + in_audio, left, text_done = False, INTERLEAVED_N_TEXT, False + for _ in range(num_items): + left -= 1 + total += embeds.shape[1] + feeds = {"inputs_embeds": embeds, "attention_mask": np.ones((1, total), np.int64), **state} + outputs = dict(zip(output_names, decoder.run(None, feeds), strict=True)) + for name in state: + state[name] = outputs[name.replace("past_key_values.", "present.").replace("past.", "present.")] + + if not in_audio: + token = int(outputs["logits"][0, -1].argmax()) + sequence.append(token) + if token == 7: + break + if interleaved: + text_done = text_done or token == 130 + if left == 0 or text_done: + in_audio, left = True, INTERLEAVED_N_AUDIO + elif token == audio_start: + in_audio = True + embeds = embedding.run(None, {"input_ids": np.array([[token]], np.int64), "audio_features": no_features})[0] + else: + frame = sample_frame(outputs["hidden_states"][0, -1]) + if interleaved and left == 0 and not text_done: + in_audio, left = False, INTERLEAVED_N_TEXT + if frame[0] == codebook_size - 1: + frame[:] = codebook_size - 1 + in_audio = False + else: + frames.append(frame.copy()) + sequence.append(AUDIO_TOKEN_ID) + rows = (frame + np.arange(NUM_CODEBOOKS) * codebook_size)[None] + embeds = audio_embedding.run(None, {"audio_codes": rows})[0].sum(axis=1, keepdims=True) + return sequence, np.array(frames, np.int64).reshape(-1, NUM_CODEBOOKS) + + +@pytest.mark.parametrize("num_clips", [0, 1], ids=["typed", "spoken"]) +def test_lfm2_audio_interleaved_speech_matches_reference(test_data_path, tmp_path, num_clips): + # Six text tokens, twelve audio frames, and round again: each frame is eight codes from the + # depthformer, and its embedding rather than a token's is what the decoder reads next. The + # sequence keeps one placeholder per frame, and the codes come out on their own. + model_path = _speech_model(test_data_path, tmp_path) + clips = [_write_wav(tmp_path / "question.wav", _synthetic_signal(0.9, seed=70))] * num_clips + prompt = "<|startoftext|>Answer aloud. " + AUDIO_MARKER * num_clips + + inputs, sequence, codes = _generate_speech( + model_path, prompt, og.Audios.open(*clips) if clips else None, 45, audio_interleaved=True, audio_top_k=1 + ) + expected_sequence, expected_codes = _reference_speech(os.fspath(model_path), inputs, 45, interleaved=True) + + layout = "".join("A" if token == AUDIO_TOKEN_ID else "T" for token in sequence) + assert layout.startswith("T" * INTERLEAVED_N_TEXT + "A" * INTERLEAVED_N_AUDIO + "T" * INTERLEAVED_N_TEXT) + assert sequence == expected_sequence + np.testing.assert_array_equal(codes, expected_codes) + assert codes.shape == (layout.count("A"), NUM_CODEBOOKS) + + +def _text_stream(test_data_path, tmp_path, num_tokens): + """What the fixture writes when nothing turns it to speech.""" + clip = _write_wav(tmp_path / "probe.wav", _synthetic_signal(0.5, seed=71)) + return _generate_with_eos(_copy_model(test_data_path, tmp_path / "probe"), clip, [7], num_tokens) + + +def test_lfm2_audio_sequential_speech_starts_at_audio_start_and_ends_at_end_of_audio(test_data_path, tmp_path): + # TTS and the like: text until <|audio_start|>, speech until the end-of-audio code, then text + # again. The fixture emits neither on its own, so the start token is one it does write, and the + # codebook is cut down to sixteen entries so that the last of them, end-of-audio, comes up. + text = _text_stream(test_data_path, tmp_path, 12) + audio_start = text[3] + model_path = _speech_model(test_data_path, tmp_path, audio_start_token_id=audio_start, codebook_size=16) + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5, seed=71)) + + inputs, sequence, codes = _generate_speech(model_path, AUDIO_MARKER, og.Audios.open(clip), 60, audio_top_k=1) + expected_sequence, expected_codes = _reference_speech( + os.fspath(model_path), inputs, 60, interleaved=False, audio_start=audio_start, codebook_size=16 + ) + + assert sequence[: text.index(audio_start) + 1] == text[: text.index(audio_start) + 1] + assert sequence[text.index(audio_start) + 1] == AUDIO_TOKEN_ID, "speech starts right after <|audio_start|>" + assert sequence == expected_sequence + np.testing.assert_array_equal(codes, expected_codes) + # End-of-audio frames hold a place in the sequence but are not speech. + assert 0 < len(codes) < sequence.count(AUDIO_TOKEN_ID) + assert codes[:, 0].max() < 15, "only the first codebook carries the end-of-audio code" + after_speech = sequence[text.index(audio_start) + 1 :] + assert any(token != AUDIO_TOKEN_ID for token in after_speech), "the turn goes back to text after end-of-audio" + + +def test_lfm2_audio_speech_output_leaves_a_text_answer_alone(test_data_path, tmp_path): + text = _text_stream(test_data_path, tmp_path, 20) + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5, seed=71)) + + _, sequence, codes = _generate_speech( + _speech_model(test_data_path, tmp_path), AUDIO_MARKER, og.Audios.open(clip), 20 + ) + + assert sequence == text + assert codes.shape == (0, NUM_CODEBOOKS) + + +def test_lfm2_audio_speech_sampling_follows_the_seed(test_data_path, tmp_path): + model_path = _speech_model(test_data_path, tmp_path) + # The fixture's logits are far apart, so it takes a high temperature for the four likeliest to share the odds. + options = {"audio_interleaved": True, "audio_temperature": 100.0, "audio_top_k": 4} + run = lambda **search: _generate_speech(model_path, "Say something.", None, 30, **options, **search)[2] # noqa: E731 + + first, again, other = run(random_seed=3), run(random_seed=3), run(random_seed=4) + greedy = _generate_speech(model_path, "Say something.", None, 30, audio_interleaved=True, audio_top_k=1)[2] + + np.testing.assert_array_equal(first, again) + assert not np.array_equal(first, other) + assert not np.array_equal(first, greedy) + # The first frame follows text only, so it is sampled from the same logits whatever the seed: + # with top_k 4 every code in it has to be one of the four likeliest. + assert first.shape == greedy.shape and first.min() >= 0 and first.max() < CODEBOOK_SIZE + + +@pytest.mark.parametrize( + "edit,expected", + [ + (lambda m: m.update({"eos_token_id": [7, 128]}), "eos_token_id holds 128.*Remove it"), + (lambda m: m.update({"eos_token_id": [7, 130]}), "eos_token_id holds 130.*Remove it"), + (lambda m: m["audio_output"].pop("embedding"), "needs both depthformer.filename and embedding.filename"), + (lambda m: m["audio_output"].update({"num_codebooks": 4}), "expected inputs for 4 codebooks"), + (lambda m: m["audio_output"].update({"num_codebooks": -8}), "expected inputs for -8 codebooks"), + (lambda m: m["audio_output"].update({"codebook_size": -1}), "codebook_size -1 must be positive"), + (lambda m: m["audio_output"].update({"codebook_size": 0}), "codebook_size 0 must be positive"), + (lambda m: m.update({"audio_token_id": 65536}), "audio_token_id 65536 .* vocabulary of 65536"), + (lambda m: m.update({"audio_token_id": -1}), "audio_token_id -1 .* vocabulary of 65536"), + ], + ids=[ + "audio-start-is-a-stop-token", + "text-end-is-a-stop-token", + "one-graph-missing", + "codebook-count", + "negative-codebook-count", + "negative-codebook-size", + "zero-codebook-size", + "audio-token-past-the-vocabulary", + "negative-audio-token", + ], +) +def test_lfm2_audio_rejects_a_speech_output_config_that_cannot_work(test_data_path, tmp_path, edit, expected): + model_path = _speech_model(test_data_path, tmp_path) + _edit_json(model_path / "genai_config.json", lambda config: edit(config["model"])) + + with pytest.raises(RuntimeError, match=expected): + model = og.Model(os.fspath(model_path)) + og.Generator(model, og.GeneratorParams(model)) + + +def test_lfm2_audio_rejects_a_codebook_size_past_the_depthformer_logits(test_data_path, tmp_path): + # The codes are sampled from the first codebook_size logits; one more than there are would read + # past the end of them. + model_path = _speech_model(test_data_path, tmp_path, codebook_size=CODEBOOK_SIZE + 1) + with pytest.raises( + RuntimeError, + match=f"codebook_size is {CODEBOOK_SIZE + 1}, but the depthformer produces {CODEBOOK_SIZE} logits", + ): + _generate_speech(model_path, "<|startoftext|>Answer aloud. ", None, 10, audio_interleaved=True, audio_top_k=1) + + +def _replace_graph_output(onnx, path: Path, name: str, build) -> None: + """Has build(raw) -> (nodes, initializers, elem_type) produce the graph output `name` from what did.""" + model = onnx.load(path) + raw = f"{name}_raw" + for node in model.graph.node: + node.output[:] = [raw if output == name else output for output in node.output] + nodes, initializers, elem_type = build(raw) + model.graph.node.extend(nodes) + model.graph.initializer.extend(initializers) + output = next(output for output in model.graph.output if output.name == name) + output.type.tensor_type.ClearField("shape") + output.type.tensor_type.elem_type = elem_type + onnx.save(model, path) + + +def test_lfm2_audio_rejects_depthformer_logits_that_are_not_float(test_data_path, tmp_path): + onnx = pytest.importorskip("onnx") + model_path = _speech_model(test_data_path, tmp_path) + _replace_graph_output( + onnx, + model_path / "dummy_depthformer.onnx", + "logits", + lambda raw: ( + [onnx.helper.make_node("Cast", [raw], ["logits"], to=onnx.TensorProto.FLOAT16)], + [], + onnx.TensorProto.FLOAT16, + ), + ) + with pytest.raises(RuntimeError, match="depthformer's logits must be float, got float16"): + _generate_speech(model_path, "<|startoftext|>Answer aloud. ", None, 10, audio_interleaved=True, audio_top_k=1) + + +def test_lfm2_audio_rejects_an_audio_embedding_short_of_a_codebook(test_data_path, tmp_path): + # A frame's embedding is the sum of one row per codebook; fewer rows than codes would be read past. + onnx = pytest.importorskip("onnx") + model_path = _speech_model(test_data_path, tmp_path) + ints = onnx.numpy_helper.from_array + + def keep_half_the_codebooks(raw): + slice_bounds = [ + ints(np.array([0], np.int64), "keep_starts"), + ints(np.array([NUM_CODEBOOKS // 2], np.int64), "keep_ends"), + ints(np.array([1], np.int64), "keep_axes"), + ] + node = onnx.helper.make_node("Slice", [raw, "keep_starts", "keep_ends", "keep_axes"], ["audio_embeds"]) + return [node], slice_bounds, onnx.TensorProto.FLOAT + + _replace_graph_output(onnx, model_path / "dummy_audio_embedding.onnx", "audio_embeds", keep_half_the_codebooks) + with pytest.raises(RuntimeError, match="audio embedding model must produce one float embedding per codebook"): + _generate_speech(model_path, "<|startoftext|>Answer aloud. ", None, 10, audio_interleaved=True, audio_top_k=1) + + +def test_lfm2_audio_rejects_a_negative_codebook_count_the_graph_leaves_open(test_data_path, tmp_path): + # With the codebook dimension symbolic the graph reports it as -1, which a count of -1 would match. + onnx = pytest.importorskip("onnx") + model_path = _speech_model(test_data_path, tmp_path, num_codebooks=-1) + depthformer = onnx.load(model_path / "dummy_depthformer.onnx") + slices = next(value for value in depthformer.graph.input if value.name == "depth_slices_in") + slices.type.tensor_type.shape.dim[1].dim_param = "num_codebooks" + onnx.save(depthformer, model_path / "dummy_depthformer.onnx") + + with pytest.raises(RuntimeError, match="expected inputs for -1 codebooks"): + model = og.Model(os.fspath(model_path)) + og.Generator(model, og.GeneratorParams(model)) + + +def test_lfm2_audio_speech_output_needs_the_decoder_hidden_states(test_data_path, tmp_path): + onnx = pytest.importorskip("onnx") + model_path = _speech_model(test_data_path, tmp_path) + decoder_path = model_path / "dummy_text.onnx" + decoder = onnx.load(decoder_path) + del decoder.graph.output[[output.name for output in decoder.graph.output].index("hidden_states")] + onnx.save(decoder, decoder_path) + + with pytest.raises(RuntimeError, match="include_hidden_states=true"): + _generate_speech(model_path, "Say something.", None, 12, audio_interleaved=True) + + +def _two_clip_inputs(model_path, tmp_path): + clips = [_write_wav(tmp_path / f"{i}.wav", _synthetic_signal(1.5 - 0.5 * i, seed=50 + i)) for i in range(2)] + model = og.Model(os.fspath(model_path)) + inputs = model.create_multimodal_processor()(f"{AUDIO_MARKER} and {AUDIO_MARKER}", audios=og.Audios.open(*clips)) + return model, inputs + + +def _set_inputs(model, inputs): + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=4096) + og.Generator(model, params).set_inputs(inputs) + + +@pytest.mark.parametrize( + "corrupt,expected", + [ + (lambda i, lengths: i.__setitem__("audio_embeds", np.zeros((2, 5), np.float32)), "expected a 3D"), + ( + lambda i, lengths: i.__setitem__("audio_embeds", np.zeros((3, 20, NUM_MELS), np.float32)), + "holds 3 clips but audio_sizes has 2 entries", + ), + (lambda i, lengths: i.__setitem__("audio_lengths", np.array([8], np.int64)), "mel_lengths has 1 entries"), + # Past the end of the staged tensor: the clip's frames would be copied out of bounds. + ( + lambda i, lengths: i.__setitem__("audio_lengths", np.array([lengths[0], 10**6], np.int64)), + r"clip 1 reports 1000000 mel frames, outside the 1\.\.", + ), + ( + lambda i, lengths: i.__setitem__("audio_lengths", np.array([lengths[0], 0], np.int64)), + r"clip 1 reports 0 mel frames, outside the 1\.\.", + ), + ], + ids=["mel-rank", "clip-count", "lengths-count", "frames-past-the-end", "frames-zero"], +) +def test_lfm2_audio_rejects_inconsistent_speech_inputs(test_data_path, tmp_path, corrupt, expected): + # The tensors the processor hands back can be replaced before they reach the generator, so the + # speech state checks that they still agree with one another. Without these the mismatch would + # surface as a confusing shape error from inside the encoder, or as an out-of-bounds read. + model, inputs = _two_clip_inputs(_model_path(test_data_path), tmp_path) + corrupt(inputs, inputs["audio_lengths"].as_numpy().copy()) + + with pytest.raises(RuntimeError, match=expected): + _set_inputs(model, inputs) + + +@pytest.mark.parametrize("num_clips", [0, 1, 2], ids=["text", "one-clip", "two-clips"]) +def test_lfm2_audio_rejects_beam_search(test_data_path, tmp_path, num_clips): + # Without the check a single clip fails with a shape mismatch from inside the encoder. + clips = [_write_wav(tmp_path / f"{i}.wav", _synthetic_signal(1.5 - 0.5 * i, seed=60 + i)) for i in range(num_clips)] + model = og.Model(_model_path(test_data_path)) + prompt = "Hi " + " and ".join([AUDIO_MARKER] * num_clips) + inputs = model.create_multimodal_processor()(prompt, audios=og.Audios.open(*clips) if clips else None) + params = og.GeneratorParams(model) + params.set_search_options(do_sample=False, max_length=4096, num_beams=2) + + with pytest.raises(RuntimeError, match="beam search is not supported for lfm2_audio; got num_beams 2"): + og.Generator(model, params).set_inputs(inputs) + + +def test_lfm2_audio_rejects_rewind(test_data_path, tmp_path): + clip = _write_wav(tmp_path / "clip.wav", _synthetic_signal(0.5)) + generator, _ = _generate(_model_path(test_data_path), f"{AUDIO_MARKER}Hi", og.Audios.open(clip), num_tokens=1) + sequence = generator.get_sequence(0).copy() + + # The conv state cannot be rewound; the generator must refuse before touching the sequence. + with pytest.raises(RuntimeError, match="RewindTo is currently not supported for lfm2_audio"): + generator.rewind_to(0) + np.testing.assert_array_equal(generator.get_sequence(0), sequence) diff --git a/test/python/test_onnxruntime_genai_api.py b/test/python/test_onnxruntime_genai_api.py index d2dc7e8289..f7fb608e38 100644 --- a/test/python/test_onnxruntime_genai_api.py +++ b/test/python/test_onnxruntime_genai_api.py @@ -67,6 +67,15 @@ def test_tokenizer_create_from_config_and_path(test_data_path): assert tokenizer_from_path.decode(tokenizer_from_config.encode(text)) == text +def test_tokenizer_decodes_no_tokens_to_an_empty_string(test_data_path): + # What is left after the prompt when generation stops on its first token; it used to SIGFPE. + model_path = os.fspath(Path(test_data_path) / "models" / "hf-internal-testing" / "tiny-random-gpt2-fp32") + tokenizer = og.Tokenizer(model_path) + + assert tokenizer.decode(np.array([], dtype=np.int32)) == "" + assert tokenizer.decode(np.arange(4, dtype=np.int32)[4:]) == "" + + def test_telemetry_control(): og.disable_telemetry_events() og.enable_telemetry_events() @@ -979,6 +988,26 @@ def test_whisper_preprocessing_multiple_audios(test_data_path, relative_model_pa _ = processor(prompts, audios=audios) +@pytest.mark.parametrize("opener", ["images", "audios"]) +def test_open_keeps_short_paths_alive(test_data_path, tmp_path, opener): + # The paths are copied into a vector first and only then turned into char pointers: a short + # string keeps its characters inside the string object, so taking the pointer as the vector + # grows leaves the earlier entries pointing at moved-from memory. + source = Path(test_data_path) / ("images/cars.jpg" if opener == "images" else "audios/jfk.flac") + suffix = source.suffix + short_dir = tmp_path / "s" + short_dir.mkdir() + paths = [] + for index in range(4): + path = short_dir / f"{index}{suffix}" + shutil.copy(source, path) + paths.append(os.fspath(path)) + + opened = (og.Images if opener == "images" else og.Audios).open(*paths) + + assert opened is not None + + def test_streaming_asr_create(asr_speech_model_path): """Test that Generator + StreamingProcessor can be created from a streaming ASR model.""" model = og.Model(asr_speech_model_path)