[Bug] model.config.stream_input=False 时，音频 chunk embedding 与 audio_bounds 逐组对齐导致 shape mismatch

环境

Model: openbmb/MiniCPM-o-4_5
Transformers: 4.51.0
Inference API: model.chat()
Audio sampling rate: 16000 Hz

故障现象

在普通音频推理中设置：

model.config.stream_input = False

调用 model.chat() 时，可能在 get_omni_embedding() 中报错：

ValueError: Shape mismatch: Trying to assign embeddings of shape
torch.Size([300, 4096]) to input indices of length 334

保持模型默认配置：

model.config.stream_input = True

同一输入可以正常推理。

最小复现

import librosa
import torch
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "openbmb/MiniCPM-o-4_5",
    trust_remote_code=True,
    attn_implementation="sdpa",
    torch_dtype=torch.bfloat16,
    init_vision=False,
    init_audio=True,
    init_tts=False,
).eval().cuda()

model.config.stream_input = False

audio, _ = librosa.load(
    "example.wav",
    sr=16000,
    mono=True,
)

msgs = [
    {
        "role": "user",
        "content": [
            audio,
            "Describe the audio.",
        ],
    }
]

result = model.chat(
    msgs=msgs,
    max_new_tokens=32,
    do_sample=False,
    generate_audio=False,
)

print(result)

相关调用链

该问题发生在音频已经完成编码、准备写入语言模型输入 embedding 的阶段：

model.chat()
    ↓
processor
    ├─ 生成 audio_features
    └─ 根据 prompt 中的音频 placeholder 生成 audio_bounds
    ↓
get_audio_embedding()
    └─ 生成 audio_embeddings
    ↓
get_omni_embedding()
    ↓
根据 model.config.stream_input
选择 audio_embeddings 的回填方式

这里的两个数据结构分别表示：

audio_embeddings
    音频编码器产生的 embedding 列表。
    一条音频可能因为内部切块而对应多个 tensor。

audio_bounds
    prompt 中需要被音频 embedding 替换的 token 区间。
    一条完整音频可能只对应一个 bound。

例如，一条音频可能得到：

audio_embeddings[0] = [
    Tensor[300, hidden_size],
    Tensor[34, hidden_size],
]

但 prompt 中只有一个完整音频占位区间：

audio_bounds[0] = [
    [start, start + 334],
]

两边的总长度相同：

300 + 34 = 334

但分组方式不同：

audio_embeddings = [300, 34]
audio_bounds      = [334]

当前源码行为

model.config.stream_input=True

当前 True 分支先拼接同一个 sample 的所有音频 embedding：

audio_embs = torch.cat(
    audio_embeddings[i],
    dim=0,
)

然后按照 audio_bounds 的长度依次切片回填：

audio_start_pos = 0

for bound in audio_bounds[i]:
    audio_len = bound[1] - bound[0]

    input_embeddings[i, bound[0]:bound[1]] = audio_embs[
        audio_start_pos : audio_start_pos + audio_len
    ]

    audio_start_pos += audio_len

对于上面的例子：

[300, 34]
    ↓ torch.cat
334
    ↓
写入长度为 334 的 audio bound

因此不会报错。

model.config.stream_input=False

当前 False 分支直接逐组配对：

audio_embs = audio_embeddings[i]
bounds = audio_bounds[i]

for embs, bound in zip(audio_embs, bounds):
    audio_indices = torch.arange(
        bound[0],
        bound[1],
        dtype=torch.long,
    ).to(input_embeddings.device)

    if embs.shape[0] != len(audio_indices):
        raise ValueError(...)

对于：

audio_embeddings = [300, 34]
audio_bounds      = [334]

zip() 第一次配对的是：

300 ↔ 334

因此立即触发：

300 != 334

并抛出 shape mismatch。

剩余的 34 个 embedding 还没有参与回填。

根本原因

False 分支错误地假设：

audio_embeddings[i][0] 必须对应 audio_bounds[i][0]
audio_embeddings[i][1] 必须对应 audio_bounds[i][1]
……

但两边的分组来源不同：

audio_embeddings 的分组
    由音频内部处理和切块方式决定。

audio_bounds 的分组
    由 prompt 中的音频 placeholder 决定。

因此不能保证两边的 group 数量和 group 边界完全一致。

模型真正需要保证的是：

所有 audio embedding 的总长度
==
所有 audio bound 的总长度

而不是要求每一组都一一对应。

该问题可以概括为：

model.config.stream_input=False 分支使用 zip()，
错误地把音频内部 chunk 分组
当成了 prompt 中 audio bound 的分组。

期望行为

当 model.config.stream_input=False 时：

1. 如果 audio_embeddings 与 audio_bounds 本来可以逐组匹配，保持现有行为。
2. 如果逐组无法匹配，但总长度相同，应先拼接 embedding，再按照 audio_bounds 顺序回填。
3. 只有总长度也不相同时，才抛出 shape mismatch。

已合并，感谢您的关注和修复。