Fix heterogeneous multimodal batch inference and per-sample TTS template handling

Summary

This PR fixes multiple issues in offline heterogeneous multimodal batch inference through model.chat() when text, image, and audio requests are mixed in the same batch.

The main issues were:

1. The TTS template state leaked from an audio sample into later text or image samples.
2. Samples without audio used an incompatible Python list for audio_feature_lens, causing heterogeneous batches to fail.
3. Audio embeddings were written back only for the final sample in the batch because the writeback logic was outside the batch loop.
4. model.chat() generated the complete batch but returned only res[0].
5. Repeated terminator tokens such as <|im_end|> leaked into decoded batch outputs.

The changes preserve single-sample behavior and the existing chunk-to-bound fallback introduced for offline audio inference.

Environment

Model: openbmb/MiniCPM-o-4_5
Transformers: 4.51.0
Inference API: model.chat()
Audio sampling rate: 16000 Hz
Batch mode: model.config.stream_input = False
Speech output: generate_audio = False

"""复现 chat() 中 use_tts_template 被 batch 顺序污染的问题。"""
import types

import librosa
import numpy as np
import torch
from PIL import Image
from transformers import AutoModel

MODEL = "/root/autodl-tmp/MiniCPM-o-4_5"
AUDIO = "/root/autodl-tmp/MMAU/mmau_0.wav"
IMAGE = "/root/autodl-tmp/heterogeneous_batch_cat.jpg"

audio, _ = librosa.load(AUDIO, sr=16000, mono=True)
audio = np.asarray(audio, dtype=np.float32)
image = Image.open(IMAGE).convert("RGB")

audio_sample = [{"role": "user", "content": [audio, "What is in the audio?"]}]
text_sample_1 = [{"role": "user", "content": "What is 2 + 2?"}]
image_sample = [{"role": "user", "content": [image, "What is in the image?"]}]
text_sample_2 = [{"role": "user", "content": "Say hello."}]

msgs_audio_pos_0 = [audio_sample, text_sample_1, image_sample, text_sample_2]
msgs_audio_pos_1 = [text_sample_1, audio_sample, image_sample, text_sample_2]
msgs_audio_pos_2 = [text_sample_1, image_sample, audio_sample, text_sample_2]
msgs_audio_pos_3 = [text_sample_1, image_sample, text_sample_2, audio_sample]

model = AutoModel.from_pretrained(
    MODEL,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
    init_tts=False,
).eval().cuda()
model.prepare_processor()

records = []
generated_batch = []
original = model.processor.tokenizer.apply_chat_template
original_generate = model.generate
model.config.stream_input = False


def traced_apply_chat_template(self, *args, **kwargs):
    prompt = original(*args, **kwargs)
    records.append((kwargs.get("use_tts_template"), prompt))
    return prompt


model.processor.tokenizer.apply_chat_template = types.MethodType(
    traced_apply_chat_template, model.processor.tokenizer
)


def traced_generate(self, *args, **kwargs):
    result, outputs = original_generate(*args, **kwargs)
    generated_batch[:] = result
    return result, outputs


model.generate = types.MethodType(traced_generate, model)


def test(name, batch):
    records.clear()
    generated_batch.clear()
    chat_return = None
    error = None
    try:
        chat_return = model.chat(
            msgs=batch,
            use_tts_template=False,
            generate_audio=False,
            enable_thinking=False,
            do_sample=False,
            max_new_tokens=64,
        )
    except Exception as exc:
        error = f"{type(exc).__name__}: {exc}"
    print(f"\n========== {name} ==========")
    print("use_tts_template:", [x[0] for x in records])
    for i, (_, prompt) in enumerate(records):
        print(f"\n[{i}]\n{prompt}")
    print(f"\n========== {name} inference results ==========")
    if error:
        print("ERROR:", error)
        return
    print("model.chat() return:", chat_return)
    for i, result in enumerate(generated_batch):
        print(f"batch[{i}]: {result}")


test("audio at batch[0]", msgs_audio_pos_0)  # [True,  True,  True,  True]
test("audio at batch[1]", msgs_audio_pos_1)  # [False, True,  True,  True]
test("audio at batch[2]", msgs_audio_pos_2)  # [False, False, True,  True]
test("audio at batch[3]", msgs_audio_pos_3)  # [False, False, False, True]

Minimal heterogeneous batch setup

The following four requests are used in every test batch:

audio_sample = [
    {
        "role": "user",
        "content": [audio, "What is in the audio?"],
    }
]

text_sample_1 = [
    {
        "role": "user",
        "content": "What is 2 + 2?",
    }
]

image_sample = [
    {
        "role": "user",
        "content": [image, "What is in the image?"],
    }
]

text_sample_2 = [
    {
        "role": "user",
        "content": "Say hello.",
    }
]

Four batches contain exactly the same requests. Only the position of the audio request changes:

msgs_audio_pos_0 = [
    audio_sample,
    text_sample_1,
    image_sample,
    text_sample_2,
]

msgs_audio_pos_1 = [
    text_sample_1,
    audio_sample,
    image_sample,
    text_sample_2,
]

msgs_audio_pos_2 = [
    text_sample_1,
    image_sample,
    audio_sample,
    text_sample_2,
]

msgs_audio_pos_3 = [
    text_sample_1,
    image_sample,
    text_sample_2,
    audio_sample,
]

The model is configured for offline batch inference:

model.config.stream_input = False

Each batch is passed directly to model.chat():

result = model.chat(
    msgs=msgs,
    use_tts_template=False,
    generate_audio=False,
    enable_thinking=False,
    do_sample=False,
    max_new_tokens=64,
)

Problem 1: use_tts_template leaks across batch samples

Previous behavior

use_tts_template is a parameter of the entire chat() call. However, the previous implementation modified it inside the per-sample loop:

elif isinstance(c, np.ndarray):
    audios.append(c)
    audio_parts.append(i)
    cur_msgs.append("<audio>./</audio>")
    use_tts_template = True

Once an audio sample set the shared variable to True, it was never restored before the next sample was processed.

This made prompt construction depend on batch ordering.

Reproduction before the fix

audio at batch[0]: [True,  True,  True,  True]
audio at batch[1]: [False, True,  True,  True]
audio at batch[2]: [False, False, True,  True]
audio at batch[3]: [False, False, False, True]

For example, when the audio request was first, the following pure text request incorrectly received a TTS prompt:

<|im_start|>user
What is 2 + 2?<|im_end|>
<|im_start|>assistant
<think>

</think>

<|tts_bos|>

The request itself did not change. Its prompt changed only because an audio request appeared earlier in the batch.

Fix

Each sample now receives an independent template state:

use_tts_template_list = []

for image, msgs in zip(images_list, msgs_list):
    sample_use_tts_template = use_tts_template

Audio detection modifies only the current sample:

elif isinstance(c, np.ndarray):
    audios.append(c)
    audio_parts.append(i)
    cur_msgs.append("<audio>./</audio>")
    sample_use_tts_template = True

The sample-local value is passed to the tokenizer:

self.processor.tokenizer.apply_chat_template(
    copy_msgs,
    tokenize=False,
    add_generation_prompt=False if teacher_forcing else True,
    use_tts_template=sample_use_tts_template,
    enable_thinking=enable_thinking,
)

The aggregate state is retained for the existing downstream speech-generation condition:

use_tts_template_list.append(sample_use_tts_template)
use_tts_template = any(use_tts_template_list)

Behavior after the fix

audio at batch[0]: [True,  False, False, False]
audio at batch[1]: [False, True,  False, False]
audio at batch[2]: [False, False, True,  False]
audio at batch[3]: [False, False, False, True]

Only the request containing audio receives <|tts_bos|>.

Problem 2: samples without audio break heterogeneous batches

Previous behavior

For a sample without audio, audio_feature_lens_list contained a Python list:

audio_feature_lens_list.append([])

Audio samples contained tensors, so a heterogeneous batch produced a structure similar to:

[
    tensor([1000]),
    [],
    [],
    [],
]

Later, get_audio_embedding() calls:

audio_feature_lens = torch.hstack(audio_feature_lens_raw)

Mixing tensors and Python lists caused:

TypeError: expected Tensor as element N in argument 0, but got list

The exact element index changed with the audio request's batch position.

Fix

Samples without audio now use an empty tensor with the expected integer dtype:

audio_feature_lens_list.append(
    torch.empty(0, dtype=torch.long)
)

This preserves the per-sample batch structure and allows torch.hstack() to process mixed audio and non-audio samples.

No real audio feature or audio length is changed by this fix.

Problem 3: audio embedding writeback only handles the final batch sample

Previous behavior

In the offline branch, only variable assignment was inside the batch loop:

for i in range(bs):
    audio_embs = audio_embeddings[i]
    bounds = audio_bounds[i]

one_to_one_match = ...

if one_to_one_match:
    input_embeddings[i, ...] = ...

After the loop finishes, i, audio_embs, and bounds refer to the last batch sample. Therefore, embedding matching and writeback are performed only for that sample.

If the audio request is not the final batch item, its audio embedding is never inserted into input_embeddings.

Fix

The existing matching and writeback logic is moved inside the batch loop:

for i in range(bs):
    audio_embs = audio_embeddings[i]
    bounds = audio_bounds[i]

    one_to_one_match = len(audio_embs) == len(bounds) and all(
        embs.shape[0] == int(bound[1] - bound[0])
        for embs, bound in zip(audio_embs, bounds)
    )

    if one_to_one_match:
        for embs, bound in zip(audio_embs, bounds):
            input_embeddings[i, bound[0] : bound[1]] = embs.to(
                device=input_embeddings.device,
                dtype=input_embeddings.dtype,
            )
    else:
        ...

Only indentation changes here. The existing matching, validation, concatenation, and writeback behavior is preserved.

Compatibility with PR #18

PR #18 fixed the case where one audio input produces multiple internal embedding chunks while the prompt contains a differently grouped set of audio bounds.

For example:

audio embedding chunks: [300, 34]
audio bounds:            [334]

The existing logic first attempts one-to-one matching:

one_to_one_match = len(audio_embs) == len(bounds) and all(
    embs.shape[0] == int(bound[1] - bound[0])
    for embs, bound in zip(audio_embs, bounds)
)

If grouping differs, it concatenates the chunks and splits them according to the bounds:

flat_audio_embs = torch.cat(audio_embs, dim=0)
total_bound_len = sum(int(bound[1] - bound[0]) for bound in bounds)

This PR does not change that behavior. It makes the PR #18 logic run independently for every batch sample.

Problem 4: chat() returns only the first batch result

Previous behavior

The generation layer returned all batch results, but chat() selected only the first one:

answer = res[0]

For batch_size=4, callers therefore received one string instead of four results.

Fix

Every generated answer is cleaned, then the return type is selected according to whether the input is batched:

answers = [
    answer.split("<|tts_eos|>")[0]
    if answer is not None
    else None
    for answer in res
]

answer = answers if batched else answers[0]

This preserves the existing single-sample API:

single sample -> str
batched input -> list[str]

return_prompt=True follows the same rule:

return answer, prompts_lists if batched else prompts_lists[0]

Problem 5: repeated terminator tokens leak into decoded output

Previous behavior

Some requests in a heterogeneous batch finish earlier than others. Their generated sequences can contain repeated <|im_end|> tokens while the rest of the batch continues.

The previous decoder removed only one terminator and only when it was the final token:

if result[-1] in terminators:
    result = result[:-1]

This produced outputs such as:

4<|im_end|><|im_end|><|im_end|>...

Fix

The decoder now truncates at the first generated terminator:

for idx, token_id in enumerate(result):
    if token_id in terminators:
        result = result[:idx]
        break

It also safely handles an empty result before checking for the BOS token:

if len(result) > 0 and result[0] == tokenizer.bos_id:
    result = result[1:]

This modifies only decoded text. It does not alter:

• outputs.sequences
• generation hidden states
• tts_bound
• TTS conditioning embeddings

The existing <|tts_eos|> cleanup in chat() remains unchanged.

Validation results

The same four requests were tested with the audio request at each possible batch position.

TTS template state

audio at batch[0]: [True,  False, False, False]
audio at batch[1]: [False, True,  False, False]
audio at batch[2]: [False, False, True,  False]
audio at batch[3]: [False, False, False, True]

Inference results

Audio at batch[0]:

[
    "Labels: Firing muskets",
    "4",
    "The image shows a close-up of an orange tabby cat...",
    "Hello! How can I help you today?",
]

Audio at batch[1]:

[
    "4",
    "Labels: Firing muskets",
    "The image shows a close-up of an orange tabby cat...",
    "Hello! How can I help you today?",
]

Audio at batch[2]:

[
    "4",
    "The image shows a close-up of an orange tabby cat...",
    "Labels: Firing muskets",
    "Hello! How can I help you today?",
]

Audio at batch[3]:

[
    "4",
    "The image shows a close-up of an orange tabby cat...",
    "Hello! How can I help you today?",
    "Labels: Firing muskets",
]

All four batches completed inference and returned four clean results. No repeated <|im_end|> tokens remained in the public chat() output.

Backward compatibility

Preserved behavior

• Non-batched calls still return a single string.
• Batched calls return a list whose length matches the batch size.
• Explicit use_tts_template=True still enables the TTS template.
• A single audio sample still automatically enables its TTS template.
• The PR #18 chunk-concatenation fallback is preserved.
• Streaming input behavior is unchanged.
• Raw generation sequences and hidden states used by TTS are unchanged.

Intended behavior change

Text and image samples after an audio sample no longer receive a TTS template unless use_tts_template=True was explicitly requested for the entire call.

Scope and limitations

This PR targets offline heterogeneous batch inference:

model.config.stream_input = False
generate_audio = False

stream_input=True remains limited to batch_size=1 by the existing implementation:

assert bs == 1, "audio stream_input mode only support batch size 1"

Batched speech waveform generation with generate_audio=True is not addressed. The existing non-streaming speech-generation path still operates on index 0:

inputs["input_ids"][0]
outputs.sequences[0]
input_audios_list[0]

The changes in this PR preserve single-sample TTS behavior but do not add multi-sample TTS waveform output.

Files changed

modeling_minicpmo.py
processing_minicpmo.py