[update] minimind-3o

pull/3/head
jingyaogong 2026-05-06 13:53:28 +08:00
parent 30b9d1c96b
commit c009567d06
77 changed files with 36293 additions and 0 deletions

8
.gitignore vendored Normal file
View File

@ -0,0 +1,8 @@
__pycache__/
*.pyc
.DS_Store
.vscode/
.cursor/
.claude/
.specstory/
.env

673
README.md Normal file
View File

@ -0,0 +1,673 @@
<div align="center">
![logo](./images/logo.png)
</div>
<div align="center">
![visitors](https://visitor-badge.laobi.icu/badge?page_id=jingyaogong/minimind-o)
[![GitHub Repo stars](https://img.shields.io/github/stars/jingyaogong/minimind-o?style=social)](https://github.com/jingyaogong/minimind-o/stargazers)
[![GitHub Code License](https://img.shields.io/github/license/jingyaogong/minimind-o?v=1)](LICENSE)
[![GitHub last commit](https://img.shields.io/github/last-commit/jingyaogong/minimind-o)](https://github.com/jingyaogong/minimind-o/commits/master)
[![GitHub pull request](https://img.shields.io/badge/PRs-welcome-blue)](https://github.com/jingyaogong/minimind-o/pulls)
[![Collection](https://img.shields.io/badge/🤗-MiniMind--O%20%20Collection-blue)](https://huggingface.co/collections/jingyaogong/minimind-o)
[![Technical Report](https://img.shields.io/badge/Technical%20Report-arXiv-red)](http://arxiv.org/abs/2605.03937)
</div>
<div align="center">
<h3>"大道至简"</h3>
</div>
<div align="center">
中文 | [English](./README_en.md)
</div>
* 此开源项目旨在从 0 完整实现一个小规模的端到端 Omni 模型,单一权重同时支持文 / 音 / 图三模态输入与文本 / 流式语音输出。
* 其中 `minimind-3o` 仅 ~0.1B,普通个人 GPU 即可完成训练、CPU即可快速推理,是当前公开模型中规模最小的完整 Omni 实现(或之一)。
* 开源 mini 与 full 两套训练数据:mini 单卡 3090 上约 2 小时跑通完整链路,便于入门;full 与发布权重对应。
* 开源 Omni 模型的完整代码与技术报告,覆盖 Thinker–Talker 双路径、流式语音生成、实时打断、近似双工交互、音色克隆与电话模式 WebUI。
* 所有核心算法代码均从 0 使用 PyTorch 原生实现,不依赖三方框架提供的高层抽象。
* MiniMind-O 进一步延续了 [MiniMind](https://github.com/jingyaogong/minimind)(语言)与 [MiniMind-V](https://github.com/jingyaogong/minimind-v)(视觉多模态)的设计范式。
> 注:"约 2 小时" 指 mini 数据集在单张 NVIDIA RTX 3090 上跑完 SFT 的实测耗时。
---
<div align="center">
[📄 MiniMind-O Technical Report](http://arxiv.org/abs/2605.03937)
https://github.com/user-attachments/assets/10cbcc5f-4e70-45cf-bdc5-d6361e40bb86
[🔗 在线体验](https://modelscope.cn/studios/gongjy/MiniMind-O) &nbsp;|&nbsp; [🔗 视频介绍](https://www.bilibili.com/video/BV121421h71e)
</div>
---
# 📌 项目介绍
继 [MiniMind](https://github.com/jingyaogong/minimind)(LLM)和 [MiniMind-V](https://github.com/jingyaogong/minimind-v)(VLM)之后,MiniMind-O 是这个系列的第三站。所谓 Omni,就是让一个模型同时具备听、看、说的多模态交互能力:接收文本、语音和视觉信号,输出文本与流式语音。
或许 GPT-4o 让人第一次感受到足够自然的流式语音交互形态,随后 Mini-Omni2、Moshi、GLM-4-Voice、Qwen3-Omni 等开源工作陆续出现。但如果目标不是直接调用这些参数庞大的现成权重,而是从 0 读懂、训练、改动一个完整 Omni 模型,开源社区仍然急缺足够轻量、链路完整的起点。要把语音真正纳入 Omni 模型,一种做法是把 ASR、LLM、TTS 串成级联链路:语音先转文字,LLM 处理后再合成语音。这条路工程上直接,但中间多了一次文本转写,延迟、语气和情绪信息都会受到影响。
MiniMind-O 尝试补上已知的空位:让语音和文本在 hidden state 层面直接连通,在主 backbone 仅 0.1B 的规模下保留端到端 Omni 链路。Talker 侧采用 MTP(Multi-Token Prediction)一次预测多层 Mimi codes,再配合 VAD 支持实时打断与近似双工交互,这是足够实用的工程路线之一。本项目的代码、模型权重、训练数据和技术报告全部完整开源,单张 RTX 3090 上约 2 小时即可跑通 mini 数据集训练。目标依旧:让每个人都能从第一行代码读起,自己动手,从 0 训练一个能听、能看、能思考、能说的模型:
![](images/omni_io_flow.png)
😊 一起感受创造的乐趣吧!
---
#### 🎉 项目包含以下内容
- 提供完整的 MiniMind-O 结构代码:Thinker、独立 Talker、audio / vision projector、Mimi codebook 接口以及 MTP audio head。
- 提供 SFT 全链路训练流程,覆盖 T2A、I2T、A2A 三类数据,支持全参数训练、音频投影层训练、视觉投影层训练与 DDP 多卡训练。
- 提供 mini 与 full 两套训练数据:mini 便于快速入门,单卡 3090 上约 2 小时可跑通;full 与发布权重对应,覆盖中文语音与图像任务。
- 提供多种内置音色、unseen 音色与任意参考音频的音色克隆能力,便于复现音色控制实验。
- 提供完整的推理与 Demo 工具,支持 CLI 推理、Web UI、流式播放、barge-in 打断和电话模式。
- 关键模块均从 0 用 PyTorch 原生实现,不依赖三方高层封装;同时兼容 `transformers` Tokenizer 与原生权重格式。
- 配套技术报告覆盖架构、训练曲线、CER / WER 评估、音色克隆相似度与跨模型对比,链接见顶部 Tech Report 区。
#### 🎉 已发布模型列表
| 模型 | 参数(主干) | Release |
|---|---|---|
| minimind-3o | ~0.1B | 2026.05.05 |
| minimind-3o-moe | ~0.3B-A0.1B | 2026.05.05 |
---
#### 👉 更新日志
<details close>
<summary> <b>🔥 2026-05-05</b> </summary>
- MiniMind-O 首次开源,发布 `minimind-3o`(115M)与 `minimind-3o-moe`(312M-A115M)
- Thinker–Talker 双路径架构,Talker 采用 MTP 预测多层 Mimi codes,支持 24 kHz 流式语音生成与 barge-in 打断
- 音频编解码器采用 Mimi(8 层 codebook,12.5 Hz,24 kHz),Talker 在 codebook 接口上使用共享主体与轻量 adapter
- 语音 / 视觉特征分别由冻结的 SenseVoice-Small 与 SigLIP2 编码,再通过两层 MLP projector 注入 MiniMind 隐空间
- 同步发布 mini 与 full 两套训练数据,mini 单卡 3090 ~2h 即可跑通整条 Thinker–Talker 链路
- 内置 5 个 voice prompt + 7 个 unseen voice prompt,提供音色克隆与电话模式 WebUI
</details>
# 📌 快速开始
<details style="color:rgb(128,128,128)">
<summary>分享本人的软硬件配置(仅供参考)</summary>
* CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
* RAM: 128 GB
* GPU: NVIDIA GeForce RTX 3090(24GB) * 8
* Ubuntu==20.04
* CUDA==12.2
* Python==3.10
* [requirements.txt](./requirements.txt)
</details>
## 第0步(必须)
### 1' 环境准备
```bash
# 克隆仓库代码
git clone --depth 1 https://github.com/jingyaogong/minimind-o
# 安装必要依赖
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
```
### 2' 下载资源
```bash
# 下载 SenseVoice-Small 语音编码器到 ./model/SenseVoiceSmall
modelscope download --model gongjy/SenseVoiceSmall --local_dir ./model/SenseVoiceSmall
# 下载 SigLIP2 视觉编码器到 ./model/siglip2-base-p32-256-ve
modelscope download --model gongjy/siglip2-base-p32-256-ve --local_dir ./model/siglip2-base-p32-256-ve
# 下载 Mimi 音频编解码器到 ./model/mimi
modelscope download --model gongjy/mimi --local_dir ./model/mimi
# 下载 CAMPPlus 说话人编码器到 ./model/campplus
modelscope download --model gongjy/campplus --local_dir ./model/campplus
# 下载 MiniMind 语言模型权重到 ./out 目录下(作为训练 Omni 的基座语言模型)
modelscope download --model gongjy/minimind-3o-pytorch llm_768.pth --local_dir ./out
```
注:也可从 [ModelScope Collection](https://modelscope.cn/collections/gongjy/MiniMind-O) 或 [HuggingFace Collection](https://huggingface.co/collections/jingyaogong/minimind-o) 选择对应内容 `git clone`(需LFS)下载,此处不再赘述。
完成后,结构应如下:
```text
minimind-o/
├── model/
│ ├── SenseVoiceSmall/
│ ├── siglip2-base-p32-256-ve/
│ ├── mimi/
│ ├── campplus/
│ └── ...
├── out/
│ └── llm_768.pth
└── ...
```
## Ⅰ 🚀 模型推理
### 1' 下载发布权重
```bash
# 下载发布权重到 ./out 目录下
modelscope download --model gongjy/minimind-3o-pytorch --local_dir ./out
```
### 2' 命令行问答
```bash
python eval_omni.py --load_from model --weight sft_omni
```
如果使用 transformers 格式模型,可先下载模型目录:
```bash
git clone https://huggingface.co/jingyaogong/minimind-3o
python eval_omni.py --load_from minimind-3o
```
### 3' 启动 WebUI(可选)
```bash
# ⚠️ 须先将 transformers 格式模型文件夹复制到 ./scripts/ 目录下,web_demo_omni 脚本会自动扫描该目录下包含权重文件的子文件夹,如不存在则报错
cp -r minimind-3o ./scripts/minimind-3o
cd scripts && python web_demo_omni.py
```
## Ⅱ 🛠️ 模型训练
<details style="color:rgb(128,128,128)">
<summary>注:提前测试Torch是否可用cuda</summary>
```python
import torch
print(torch.cuda.is_available())
```
如果不可用,请自行去[torch_stable](https://download.pytorch.org/whl/torch_stable.html)下载whl文件安装。
</details>
### 1' 下载数据
快速开始时,推荐从[数据集链接](https://huggingface.co/datasets/jingyaogong/minimind-o_dataset)只下载 `_mini` 数据集,并放到 `./dataset` 下。
### 2' 开始训练
推荐 mini 训练管线如下,默认在 `trainer/` 目录下执行,可直接 `cd trainer && bash train.sh`:
```bash
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_t2a_mini.parquet --epochs 1 --batch_size 40 --use_compile 1 --from_weight llm --save_weight sft_zero --max_seq_len 512 --use_wandb --use_moe 0
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_a2a_mini.parquet --epochs 1 --batch_size 40 --use_compile 0 --from_weight sft_zero --save_weight sft_zero --max_seq_len 640 --mode audio_proj --use_wandb --use_moe 0
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 2e-5 --data_path ../dataset/sft_a2a_mini.parquet --epochs 1 --batch_size 16 --use_compile 0 --from_weight sft_zero --save_weight sft_zero --max_seq_len 768 --use_wandb --use_moe 0
```
### 3' 测试已训练模型(可选)
确保需要测试的模型 `*.pth` 文件已保存于 `./out/` 目录下。
```bash
python eval_omni.py --weight sft_omni
```
# 📌 模型细节
MiniMind-O 的基座语言模型来自孪生项目 [MiniMind](https://github.com/jingyaogong/minimind),LLM 的结构与训练细节可移步该项目查阅。即使不了解 LLM 细节,也可直接参照上方"快速开始"流程训练一个 MiniMind-O。
## Ⅰ 架构总览
![](./images/architecture.jpg)
MiniMind-O 的主体由 Thinker 和 Talker 两条路径组成。Thinker 负责理解文本、语音和图像输入,并生成语义层面的文本回复;Talker 则在 Thinker 给出的语义条件上,通过 MTP 同步预测多层 Mimi audio codes,最后由音频解码器还原成流式语音。这样做的目的不是把 ASR、LLM、TTS 简单串起来,而是在一个统一序列里同时保留文本推理、语音输出和流式交互能力。
文本输入直接进入语言主干;语音和图像分别经过 Audio Encoder 与 Vision Encoder 提取特征,再映射到 MiniMind 的隐空间中。音色信息由 Speaker Encoder 或参考音频 codes 提供,推理时可以配合 VAD 实现边听边答、实时打断和近似双工交互。更细的 projector 结构、序列排布和训练目标在后文展开,代码层面的实现细节可直接参考 `model/model_omni.py` 与[技术报告](http://arxiv.org/abs/2605.03937)。
![](./images/input_token_layout.jpg)
图中展示了文本 token、语音特征、图像特征和音色条件在输入序列中的布局方式。
## Ⅱ Thinker 侧多模态理解
Thinker 负责统一接收文本、语音和图像信息,并生成语义层面的文本回复。文本 token 直接进入语言主干,语音和图像特征则通过对应 projector 注入到占位符位置,使不同模态最终落到同一条序列中建模。
## Ⅲ 中间层 Bridge
Thinker 向 Talker 传递的表征取自中间层,而不是 embedding 层或最后一层。embedding 层语义信息不足,最后一层又更贴近 next-token prediction 目标;中间层通常已经融合了上下文和跨模态信息,同时还没有被 LM head 过度塑形,更适合作为语音生成的条件。默认 `bridge_layer = num_hidden_layers // 2 - 1`,不同规模下也可以通过配置调整。
## Ⅳ Talker 侧语音生成
Talker 负责把 Thinker 给出的语义状态转成 8 层 Mimi codebook 序列。这里采用 MTP 形式同时预测多个 audio codebook,而不是把每层 codebook 拆成独立的长链路;为了控制 0.1B 模型中的额外参数量,音频 embedding 和输出 head 采用共享主体加轻量 codebook adapter 的形式。这样既保留不同 codebook 的分布差异,也避免为每一层 codebook 复制一整套参数。
## Ⅴ 序列格式与流式解码
![](./images/sequence_format.jpg)
MiniMind-O 将文本 token 与 8 路 audio-code stream 放在同一个训练样本中:Thinker 负责文本序列,Talker 负责音频 code 序列,语音、图像和音色条件都通过占位符或 reference codes 注入。回复开始之后才计算目标文本和目标音频的损失,因此 reference 与 conditioning 区域只提供条件,不作为重构目标。
流式生成时,模型一边产生文本 token,一边通过 MTP 和延迟调度补齐 8 层 Mimi codes。Mimi 解码器可以增量恢复 24 kHz 波形,因此语音播放不必等待完整回答结束。
## Ⅵ 音色控制
音色控制采用 in-context voice cloning 的方式完成:参考音频先被编码成 voice prompt,作为上下文条件喂给 Talker,而不是通过微调权重或改写文本 prompt 来指定音色。模型也可以同时使用 speaker embedding 提供更稳定的说话人约束;推理时更换音色只需要替换这些条件信息,Thinker prompt 与 Talker 权重保持不变。
默认 release 带有 5 个内置 voice prompt(dylan、eric、serena、uncle_fu、vivian),另保留 7 个 unseen prompt 用于评估(arthur、chelsie、cherry、ethan、jennifer、momo、moon)。
## Ⅶ 模块与参数规模
MiniMind-O 所说的 0.1B,指 Thinker、Talker 和两路 projector 组成的可训练主体;落到具体发布版本上,`minimind-3o` 约 113M,`minimind-3o-moe` 约 315M。Audio Encoder、Vision Encoder 和 Speech Codec 属于冻结的外部旁路模型,负责特征提取或音频编解码,合计约 425M 参数,不计入 active MiniMind-O 参数。
下表按发布模型统计主要模块参数,Trainable 参数按 PyTorch 模块统计,tied embedding 去重计入。
| 统计口径 | minimind-3o | minimind-3o-moe |
|---|---:|---:|
| 可训练主体 | 113.13M | 314.89M |
| 冻结外部模块 | 424.70M | 424.70M |
| 运行时总加载 | 537.83M | 739.59M |
| 模块 | 具体实现 | 关键配置 | 状态 / 参数 (~3o / ~3o-moe) |
|---|---|---|---|
| Thinker | MiniMind Transformer | 8 layers, hidden 768 | trainable, 63.91M / 198.42M |
| Talker | 独立 MiniMind blocks | 4 layers, 8 codebook heads | trainable, 47.05M / 114.30M |
| Audio projector | `MMAudioProjector` | 512 → 768 | trainable, 0.99M |
| Vision projector | `MMVisionProjector` | 768 → 768 | trainable, 1.18M |
| Audio encoder | SenseVoice-Small | 16 kHz speech features | frozen, 234.00M |
| Vision encoder | SigLIP2 base-p32-256 | 256×256 image, 64 tokens | frozen, 94.55M |
| Speech codec | Mimi | 8 codebooks, 12.5 Hz, 24 kHz | frozen, 96.15M |
| Speaker condition | CAM++ embedding | 192-d speaker vector | precomputed |
# 📌 实验
## Ⅰ 数据集
数据集下载:[ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind-o_dataset) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind-o_dataset)
所有语音数据都统一转成 Mimi codes 存储,8 层 codebook,帧率 12.5 Hz;图像统一 resize 到 256×256,由 SigLIP2 P32 编码为 64 个 patch token。训练数据主要来自公开 omni / speech instruction 数据,包括 [VoiceAssistant-400K](https://huggingface.co/datasets/gpt-omni/VoiceAssistant-400K)、[UltraChat-300K-SLAM-Omni](https://huggingface.co/datasets/worstchan/UltraChat-300K-SLAM-Omni) 等;同时基于 Qwen3-TTS 进行了大量多说话人音频合成,并用 CAM++ 提取 speaker embedding 作为音色条件。I2T 数据与 [MiniMind-V](https://github.com/jingyaogong/minimind-v) 使用的视觉指令数据来源一致,原始组成和引用可参考该项目说明。
仓库提供 **mini** 与 **full** 两套训练数据。mini 从 full 中按"英文 + 无视觉"筛出,配 `train_sft_omni.py` 的默认 `--data_path` 即可使用;它的目标是用较低成本跑通 Thinker–Talker、Mimi 编解码、序列布局和音色注入链路,而不是复现发布模型的中文语音能力。中文 Talker 要同时处理更复杂的字音映射、韵律停顿和多说话人稳定性,明显比英文更难,不能依赖单卡 3090 约 2 小时的 mini 训练完成。
full 数据集与发布的 `minimind-3o` / `minimind-3o-moe` 权重对应,覆盖中英文 T2A / A2A 与图像 I2T。规模与中英文比例见下表,是论文中 CER / 音色相似度等指标的实际训练源。
其中 T2A 表示 Text-to-Audio,A2A 表示 Audio-to-Audio,I2T 表示 Image-to-Text。
| 数据集 | 子集 | 输入语音 | 输出语音 | 备注 |
|---|---|---|---|---|
| `sft_t2a_mini` | 英文 T2A | — | 约 470.14 h | mini 入门用 |
| `sft_a2a_mini` | 英文 A2A | 约 74.64 h | 约 56.60 h | mini 入门用 |
| `sft_t2a` | 中英 T2A | — | 约 1636.01 h | full 训练 |
| `sft_a2a` | 中英 A2A | 约 1711.97 h | 约 423.40 h | full 训练 |
| `sft_i2t` | 图像 I2T | — | — | full 训练 |
`sft_t2a` 中中文、英文、混合样本占比分别为 45.7%、46.5%、7.8%;`sft_a2a` 中三者分别为 70.8%、21.2%、8.0%。这个分布会直接反映到行为上:短中文和短英文回答通常较稳定,较长英文语音更容易出现读音漂移和漏词。mini 子集只保留英文,因此即便参数量和数据量都收得很紧,单语种内部的 CER 表现仍能维持在可用范围。
## Ⅱ 训练
训练入口是 `train_sft_omni.py`,推荐流程可直接参考 `trainer/train.sh`。当前 full 训练不拆复杂的多阶段预训练,而是按数据流逐步接入能力:
![](./images/training_pipeline.jpg)
- `sft_t2a`:先对齐文本到语音输出,让 Talker 学会在 Thinker 语义条件下生成 Mimi codes;
- `sft_a2a`:再接入语音输入,使模型从 speech instruction 进入同一套 Thinker–Talker 回复链路;
- `sft_i2t`:最后对齐视觉路径,其中 `vision_proj` 模式只更新视觉投影层,避免图像数据过度改写语言和语音能力。
训练模式里,`all` 会更新 MiniMind / Talker / projector,`audio_proj` 和 `vision_proj` 只用于单独对齐对应投影层;SenseVoice-Small、SigLIP2 和 Mimi 始终冻结。Dense 与 MoE 版本沿用同一套数据顺序。mini 命令只用于快速跑通链路,默认单卡 3090 约 2 小时完成;发布权重对应 full 数据训练。
下面给出 full 训练过程中的 T2A 与 A2A loss 曲线(仅供参考):
![](./images/t2a_training_curves.jpg)
> `sft_t2a`:文本到语音输出链路
![](./images/a2a_training_curves.jpg)
> `sft_a2a`:接入语音输入后的 loss
T2A 曲线已去掉早期不兼容权重 resume 造成的异常尖峰;MoE 总参数更多但 active 参数与 dense 接近,更适合作为容量分配实验参考。
## Ⅲ 模型权重
| 模型格式 | ModelScope | HuggingFace |
|---|---|---|
| PyTorch (`*.pth`) | [minimind-3o-pytorch](https://www.modelscope.cn/models/gongjy/minimind-3o-pytorch) | [minimind-3o-pytorch](https://huggingface.co/jingyaogong/minimind-3o-pytorch) |
| Transformers | [minimind-o collection](https://modelscope.cn/collections/gongjy/MiniMind-O) | [minimind-o collection](https://huggingface.co/collections/jingyaogong/minimind-o) |
> Transformers 版本包含 `minimind-3o` 与 `minimind-3o-moe`,适合直接用于 `eval_omni.py` 和 WebUI 推理;原生 PyTorch 权重主要用于训练、复现实验和继续微调。
# 📌 评估
Omni 模型目前可能还没有统一的评估口径,不同工作的 LLM 主干、音频合成器和系统目标都不一样:有的看重 LLM 本身的知识和推理,会报告 MMLU、HumanEval 等指标;有的看重流式速度和音质,有的强调语音一致性指标,也有的更关注自然交互或更大范围的 Omni 生成。这些工作大都基于 SOTA 开源 LLM 续训,而 MiniMind 的 0.06B 主干在复杂知识问答、数学推理、代码生成或开放式长回答上显然不可能形成竞争力,Talker 的自然度、韵律和稳定性也弱于成规模的系统。
所以这里无法追求综合榜单,而是落实到几项更可复现的局部评估和 use cases:Talker hidden size 消融、音色克隆相似度、相同问题和相同 ASR 流程下的 CER / WER 对比,以及 A2A、I2A 和实时交互样例。CER / WER 主要用来观察文本一致性,音质、自然度和人类偏好则留给定性样例和实际试听判断。
## Ⅰ Talker Hidden Size 消融
如果只看语音生成,Talker 做到 1024 / 2048 维、或者继续加深层数一定会更稳。但 MiniMind-O 要把完整 Omni 链路压在 0.1B 左右,不能把大部分参数都交给声学端。Thinker / Talker 解耦后,语言理解和跨模态融合主要由 Thinker 承担,Talker 只在语义条件上渲染 Mimi codes,这让小 Talker 成为可能。这里的渲染不是只预测语义 token 再交给外部声学生成器,而是由 Talker 直接生成可解码的 Mimi acoustic codes;真正的瓶颈也就在输出端:Talker 面对的是 8 层 Mimi codebook,而不是单一路径的 next-token prediction。
384 维最诱人,dense 版本可以压到 88M 左右;512 维也更轻。但表中结果说明,小并不自动等于划算:短句还能维持,中长句更容易出现漏词、重复和发音漂移。768 维最后留下来,是因为它和 MiniMind 主干维度一致,可以用 Thinker 后 4 层初始化;参数仍在 0.1B 左右,训练成本没有明显增加,一致性却稳定得多。
| Variant | Talker hidden | Params | Avg CER ↓ | Short ↓ | Mid / Long ↓ |
|---|---|---|---|---|---|
| Dense | 768 | 115.29M | **0.0897** | 0.1528 | 0.0874 / 0.0675 |
| Dense | 512 | 96.13M | 0.1745 | 0.2709 | 0.2455 / 0.0976 |
| Dense | 384 | 88.72M | 0.2767 | 0.3904 | 0.1865 / 0.4046 |
| MoE | 768 | 317.05M-A115.33M | **0.0900** | 0.2075 | 0.0533 / 0.0271 |
| MoE | 512 | 261.32M-A96.17M | 0.1265 | 0.0711 | 0.1490 / 0.1464 |
| MoE | 384 | 240.04M-A88.75M | 0.3280 | 0.3757 | 0.2777 / 0.4313 |
Dense 和 MoE 的 CER 不宜直接横向比较:同一问题下,两个 Thinker 生成的内容和长度可能不同,Talker 面对的合成难度也不同。更有意义的是看同一架构内部的趋势,768 都明显优于 512 和 384。
## Ⅱ 音色克隆相似度
音色克隆是当前版本里比较 Beta 的能力。没说错的话,多数开源 Omni 模型只支持固定输出音色,而 minimind-3o 尝试把多音色生成塞到同一套 Talker 里完成。这个目标比“能说话”更难,因为模型不仅要把内容说对,还要在生成 Mimi codes 时保留说话人的音色线索。
目前效果还谈不上高保真克隆,同一个参考音色在不同问题上并不总能保持一致,长句里也容易被发音和节奏问题带偏。但基本的男女声差异、语调倾向和一部分韵律特征是能区分出来的。
下面的 CAM++ speaker embedding 余弦相似度只作为自动化参考:Seen 组来自 `voices.pt` 中 5 个内置音色,Unseen 组来自 `voices_unseen.pt` 中 7 个训练时未见过的音色;每个音色使用同一组文本问题,只替换音色条件。
逐音色细分如下:
| Split | Speaker | Dense ↑ | MoE ↑ |
|---|---|---|---|
| Seen | dylan | 0.6997 | 0.6837 |
| Seen | eric | 0.5289 | 0.4232 |
| Seen | serena | 0.7092 | 0.7041 |
| Seen | uncle_fu | 0.7241 | 0.7337 |
| Seen | vivian | 0.5744 | 0.5888 |
| Unseen | arthur | 0.7171 | 0.6750 |
| Unseen | chelsie | 0.6437 | 0.6240 |
| Unseen | cherry | 0.5689 | 0.5678 |
| Unseen | ethan | 0.4783 | 0.4847 |
| Unseen | jennifer | 0.4749 | 0.4003 |
| Unseen | momo | 0.6470 | 0.5720 |
| Unseen | moon | 0.4282 | 0.6673 |
总体上,minimind-3o 与 minimind-3o-moe 的平均结果接近,也都略高于早期 baseline;这说明音色保持不主要由 inactive expert 容量决定,更直接的影响来自 reference 片段质量、CAM++ embedding 的可分性,以及 Talker 生成音频本身是否稳定。单个音色里,uncle_fu、serena、arthur 这类声音更容易保持住,至少一个版本能超过 0.70;eric、moon 等 outlier 则更容易受生成质量影响。换句话说,这个能力已经能区分一部分音色特征,但距离“给一段参考音频就稳定复刻”的产品级体验还有距离。
### 音色克隆消融实验(音频播放)
为了让试听更直观,这里固定 seed=42、temperature=0.7,对每个音色展示 1 个生成样例,唯一变化的是参考音频 codes 和 speaker embedding。作为对照,下面先给出不施加任何参考音色条件时的 default 输出(朗读文本对所有样例一致):
https://github.com/user-attachments/assets/b31fd8f2-e3af-4fed-ba19-65424b59bec6
#### Seen 音色
Seen 表示训练数据中出现过的音色,用来观察模型对熟悉说话人的保持情况。
<table>
<tr><th width="100">说话人</th><th width="380">参考音色</th><th width="380">输出结果</th><th width="80">平均</th></tr>
<tr><td>dylan</td><td>
https://github.com/user-attachments/assets/070ea3ab-0e8e-4aa0-84b5-af8d3c4e2725
</td><td>
https://github.com/user-attachments/assets/eb2da7ed-173c-47e9-9431-7bdb5a9b7385
</td><td>0.6712</td></tr>
<tr><td>eric</td><td>
https://github.com/user-attachments/assets/c74aa5dc-1edd-44c1-9546-6e57194c2f60
</td><td>
https://github.com/user-attachments/assets/f3fa8906-4e14-4610-a9d9-c16c915ca1b3
</td><td>0.4430</td></tr>
<tr><td>serena</td><td>
https://github.com/user-attachments/assets/0eeeac87-fa70-4025-b66e-1f0197f2b434
</td><td>
https://github.com/user-attachments/assets/c5901dca-4b2a-47f5-9b30-c89de54f908e
</td><td>0.6600</td></tr>
<tr><td>uncle_fu</td><td>
https://github.com/user-attachments/assets/fdd1bb28-6648-44bf-8bcb-4509e709e347
</td><td>
https://github.com/user-attachments/assets/95b480f1-f015-4712-8d7c-17db465f6584
</td><td>0.6632</td></tr>
<tr><td>vivian</td><td>
https://github.com/user-attachments/assets/f64731c4-67a3-4e18-b7d7-61bf44ef4bdd
</td><td>
https://github.com/user-attachments/assets/3f1cc9bb-16d2-4ce0-a473-40676cf4523e
</td><td>0.5320</td></tr>
</table>
#### Unseen 音色
Unseen 表示训练时没有见过的音色,用来观察模型能否把新的参考音色0样本直接迁移到生成语音里。
<table>
<tr><th width="100">说话人</th><th width="380">参考音色</th><th width="380">输出结果</th><th width="80">平均</th></tr>
<tr><td>arthur</td><td>
https://github.com/user-attachments/assets/46e720ba-c413-4731-8275-a9a2cb483690
</td><td>
https://github.com/user-attachments/assets/f489b463-c7d7-4ad6-9e4d-53333172c66c
</td><td>0.6479</td></tr>
<tr><td>chelsie</td><td>
https://github.com/user-attachments/assets/3430ecdb-6de8-4fb0-a6a7-ad82bdce01a1
</td><td>
https://github.com/user-attachments/assets/e598dbc2-ba28-4c38-b52d-6fa6c2349a5b
</td><td>0.5975</td></tr>
<tr><td>cherry</td><td>
https://github.com/user-attachments/assets/f9166af6-3a98-42f3-9cf8-ad105eea87d6
</td><td>
https://github.com/user-attachments/assets/eccca693-4708-409a-88f7-85eb25f66fe6
</td><td>0.5418</td></tr>
<tr><td>ethan</td><td>
https://github.com/user-attachments/assets/9c992505-2046-483e-a7cf-50ec18a5e329
</td><td>
https://github.com/user-attachments/assets/98013c5e-f5b5-4e1a-bc0e-a0f0be5d3240
</td><td>0.4323</td></tr>
<tr><td>jennifer</td><td>
https://github.com/user-attachments/assets/924b035d-5c7c-45a5-a8f8-5dbdc18f71db
</td><td>
https://github.com/user-attachments/assets/853d1370-0065-4567-9a71-dc88a6a34d56
</td><td>0.4052</td></tr>
<tr><td>momo</td><td>
https://github.com/user-attachments/assets/7e97f524-da6d-4a2f-9095-e7f99262f4a5
</td><td>
https://github.com/user-attachments/assets/4c193c8f-8750-4424-acba-2bd13089a634
</td><td>0.5968</td></tr>
<tr><td>moon</td><td>
https://github.com/user-attachments/assets/527df88a-adc0-48d3-9a6a-827ca1ba7fb0
</td><td>
https://github.com/user-attachments/assets/3f533e26-1ad8-4ab3-baf1-21267734d3ee
</td><td>0.5874</td></tr>
</table>
## Ⅲ 跨模型英文 T2A 对比
这里选了 20 个英文问题,并统一加上 `Answer briefly in one short sentence` 约束。这样做不是为了考察开放式英文能力,而是尽量把回答长度压到同一范围内;三套模型生成音频后,再统一用 Qwen3-ASR 转写,并与目标文本计算 CER / WER,用来比较 Talker 的文本一致性。
| 长度桶 | [Mini-Omni](https://huggingface.co/gpt-omni/mini-omni) CER/WER | [Mini-Omni2](https://huggingface.co/gpt-omni/mini-omni2) CER/WER | minimind-3o CER/WER |
|---|---|---|---|
| short (≤15w) | 0.0195 / 0.0384 (n=8) | 0.0503 / 0.0584 (n=14) | 0.0531 / 0.0417 (n=8) |
| mid (16–30w) | 0.0038 / 0.0052 (n=12) | 0.0062 / 0.0076 (n=6) | 0.1327 / 0.1420 (n=11) |
| long (31–60w) | — | — | 0.0431 / 0.0508 (n=1) |
≤15 词的短回复里,minimind-3o 已经接近 Mini-Omni2;真正拉开差距的是 16–30 词段。这个长度已经不是简单短语,Talker 需要在一个完整短句里同时维持发音、节奏和词面一致性,也是当前 0.1B Talker 最容易暴露不稳定性的区间。
## Ⅳ 跨模型视觉语言对比
[Mini-Omni](https://huggingface.co/gpt-omni/mini-omni) 不支持 VL 路径,因此这里只比较 [Mini-Omni2](https://huggingface.co/gpt-omni/mini-omni2)(0.5B)和 minimind-3o(0.1B)。9 张合成图像上,两个模型分别生成英文回答,再统一转写并计算 CER / WER,作为视觉到语音链路的一致性参考。
| Model | Params | Avg CER ↓ | Avg WER ↓ |
|---|---|---|---|
| [Mini-Omni2](https://huggingface.co/gpt-omni/mini-omni2) | 0.5B | 0.7609 | 0.9756 |
| minimind-3o | 0.1B | 0.8241 | 1.0293 |
这个数值不能当作开放式图像描述的绝对正确率。视觉描述存在大量等价表达,同义改写和描述顺序都会影响 CER / WER,数值整体偏高是预期现象。在同一自动流程下,minimind-3o 落后于 Mini-Omni2,但仍处在同一数量级,同时参数约为后者的 1/5。
## Ⅴ 样例
![](./images/qual_a2a.jpg)
语音到语音样例直接以真实语音作为输入,由 Thinker 组织语义,再由 Talker 渲染成语音。短回答仍然是当前更稳的区间,中文解释型问题通常能生成较连贯的回答,英文的发音和节奏相对更稳定。
<table>
<tr>
<td>
https://github.com/user-attachments/assets/c85809b2-4787-4656-9c7e-55b693798494
</td>
<td>
https://github.com/user-attachments/assets/354a5eec-c147-4d18-8c7a-942bd2a0b4b0
</td>
</tr>
</table>
![](./images/image2audio_qualitative.jpg)
图像问答样例把视觉编码、文本生成和语音渲染串在同一条链路里。当前模型通常能抓住主体物体和大致场景,但细粒度空间关系、数量和属性仍容易出错,因此更适合作为小模型 omni pipeline 的可复现基线。
<table>
<tr>
<td>
https://github.com/user-attachments/assets/244e08b0-5b12-449e-a7a2-2a2139c5d62d
</td>
<td>
https://github.com/user-attachments/assets/3e8d0a76-282d-4a9d-9726-a954cf80198a
</td>
</tr>
</table>
## Ⅵ 实时交互
![](./images/realtime_interaction.jpg)
最后是实时交互界面。用户停止说话后,Thinker 先完成语义侧的 prefill,Talker 随后开始逐步产生音频 code,Mimi decoder 则边接收 code 边写出 24 kHz 波形。Barge-in 的例子展示了另一条更接近真实对话的路径:当用户在模型说话过程中再次开口,系统会中断当前生成,重新进入 prefill–reply 流程。这里的中断检测仍然只是简单 VAD 阈值,还谈不上语义级打断;但从工程闭环看,系统已经能从 speaking 状态退回 listening 状态,并处理下一轮输入。
### 🧩 未来值得改进的方面
当前模型和大规模 Omni 系统在各方面仍有差距,这一点不需要回避。长语音自然度、复杂视觉推理、开放式英文中长回答和音色稳定性,都还不是它擅长的区间。视觉路径更接近紧凑的 vision-to-speech 链路,MoE 版本也更像一次容量分配实验,而不是同算力最优解。
这些限制也给出了后续方向:更长的 ICL 上下文、更细的 prosody 监督、更强的视觉编码器、更稳定的音色条件,以及对 Bridge layer 和 MTP codebook interface 的系统扫描,都值得继续做。
话说回来,MiniMind-O 的价值也正在这里。它把一个完整 Omni 闭环压到 0.1B 量级,并把代码、权重和主要训练数据放在同一个可检查对象里;这意味着它不只是一个 demo,而是一个足够小、足够透明、可以从头复现和继续改造的基线。对于想理解 Thinker–Talker 解耦、MTP codebook interface、in-context voice cloning 和 middle hidden bridge 这些细节的人来说,它提供的是一套可以真正动手验证的设计经验。
# 📌 致谢
> [!TIP]
> 如果您觉得 `MiniMind-O` 对您有所帮助,可以在 GitHub 上加一个⭐<br/>
> 水平有限难免存在未知的纰漏,欢迎所有人在 Issues 交流指正或提交 PR 改进项目<br/>
> 您的支持就是持续改进项目的动力,谢谢!
## 🤝贡献者
<a href="https://github.com/jingyaogong/minimind-o/graphs/contributors">
<img width="200" src="https://contrib.rocks/image?repo=jingyaogong/minimind-o" />
</a>
## 😊鸣谢
- [MiniMind](https://github.com/jingyaogong/minimind) / [MiniMind-V](https://github.com/jingyaogong/minimind-v)(基座、数据)
- [Qwen2.5-Omni / Qwen3-Omni](https://github.com/QwenLM/Qwen2.5-Omni)(灵感、数据)
- [Mini-Omni / Mini-Omni2](https://github.com/gpt-omni/mini-omni)(灵感、数据)
- [SLAM-Omni](https://aclanthology.org/2025.findings-acl.115/)(数据)
- [SenseVoice](https://arxiv.org/abs/2407.04051)(组件)
- [Mimi / Moshi](https://arxiv.org/abs/2410.00037)(组件)
## 🫶支持者
<a href="https://github.com/jingyaogong/minimind-o/stargazers">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://bytecrank.com/nastyox/reporoster/php/stargazersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
<source media="(prefers-color-scheme: light)" srcset="https://bytecrank.com/nastyox/reporoster/php/stargazersSVG.php?user=jingyaogong&repo=minimind-o"/>
<img alt="github contribution grid snake animation" src="https://bytecrank.com/nastyox/reporoster/php/stargazersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
</picture>
</a>
<a href="https://github.com/jingyaogong/minimind-o/network/members">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://bytecrank.com/nastyox/reporoster/php/forkersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
<source media="(prefers-color-scheme: light)" srcset="https://bytecrank.com/nastyox/reporoster/php/forkersSVG.php?user=jingyaogong&repo=minimind-o"/>
<img alt="github contribution grid snake animation" src="https://bytecrank.com/nastyox/reporoster/php/forkersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
</picture>
</a>
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/svg?repos=jingyaogong/minimind-o&type=Date&theme=dark"/>
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/svg?repos=jingyaogong/minimind-o&type=Date"/>
<img alt="Star History Chart" src="https://api.star-history.com/svg?repos=jingyaogong/minimind-o&type=Date&theme=dark"/>
</picture>
# 🎓 引用
如果您觉得 MiniMind-O 对您的研究或工作有所帮助,请引用:
```bibtex
% 引用 MiniMind-O 技术报告:用于讨论模型架构、训练方法与实验结论。
@article{minimind-o-report,
title = {MiniMind-O Technical Report: An Open Small-Scale Speech-Native Omni Model},
author = {Jingyao Gong},
journal = {arXiv preprint arXiv:2605.03937},
year = {2026}
}
% 引用 MiniMind-O GitHub 仓库:用于指代开源代码与发布权重。
@misc{minimind-o,
title = {MiniMind-O: Train a Tiny Omni Model from Scratch},
author = {Jingyao Gong},
year = {2026},
url = {https://github.com/jingyaogong/minimind-o},
note = {GitHub repository, accessed 2026}
}
```
# 📜 许可协议
本仓库遵循 [Apache-2.0 License](LICENSE) 开源协议。

675
README_en.md Normal file
View File

@ -0,0 +1,675 @@
<div align="center">
![logo](./images/logo.png)
</div>
<div align="center">
![visitors](https://visitor-badge.laobi.icu/badge?page_id=jingyaogong/minimind-o)
[![GitHub Repo stars](https://img.shields.io/github/stars/jingyaogong/minimind-o?style=social)](https://github.com/jingyaogong/minimind-o/stargazers)
[![GitHub Code License](https://img.shields.io/github/license/jingyaogong/minimind-o?v=1)](LICENSE)
[![GitHub last commit](https://img.shields.io/github/last-commit/jingyaogong/minimind-o)](https://github.com/jingyaogong/minimind-o/commits/master)
[![GitHub pull request](https://img.shields.io/badge/PRs-welcome-blue)](https://github.com/jingyaogong/minimind-o/pulls)
[![Collection](https://img.shields.io/badge/🤗-MiniMind--O%20%20Collection-blue)](https://huggingface.co/collections/jingyaogong/minimind-o)
[![Technical Report](https://img.shields.io/badge/Technical%20Report-arXiv-red)](http://arxiv.org/abs/2605.03937)
</div>
<div align="center">
<h3>"Less is More"</h3>
</div>
<div align="center">
[中文](./README.md) | English
</div>
* This project implements a small end-to-end Omni model from scratch, where a single set of weights jointly handles text / audio / image inputs and produces text / streaming-speech outputs.
* `minimind-3o` has only ~0.1B parameters: it can be trained on a consumer GPU and runs quickly on CPU, making it one of the smallest fully-functional Omni implementations publicly available.
* Two training datasets are released, `mini` and `full`. `mini` runs the full pipeline in about 2 hours on a single RTX 3090 and is intended for getting started; `full` corresponds to the released weights.
* The full codebase and technical report are released, covering the Thinker–Talker dual-path architecture, streaming speech generation, real-time barge-in, near-duplex interaction, voice cloning and a phone-mode WebUI.
* All core algorithmic components are implemented from scratch in native PyTorch and do not rely on high-level abstractions from third-party frameworks.
* MiniMind-O continues the design philosophy of [MiniMind](https://github.com/jingyaogong/minimind) (language) and [MiniMind-V](https://github.com/jingyaogong/minimind-v) (vision-language).
> Note: "about 2 hours" refers to the measured time of running SFT on the mini dataset using a single NVIDIA RTX 3090.
---
<div align="center">
[📄 MiniMind-O Technical Report](http://arxiv.org/abs/2605.03937)
https://github.com/user-attachments/assets/10cbcc5f-4e70-45cf-bdc5-d6361e40bb86
[🔗 Online Demo](https://modelscope.cn/studios/gongjy/MiniMind-O) &nbsp;|&nbsp; [🔗 Video Intro](https://www.bilibili.com/video/BV121421h71e)
</div>
---
# 📌 Project Introduction
After [MiniMind](https://github.com/jingyaogong/minimind) (LLM) and [MiniMind-V](https://github.com/jingyaogong/minimind-v) (VLM), MiniMind-O is the third stop in this series. By "Omni" we mean a model that can listen, see and speak at the same time: it takes text, speech and visual signals as inputs, and produces text together with streaming speech.
GPT-4o was probably the first system that made natural streaming voice interaction feel real. Since then, open-source projects such as Mini-Omni2, Moshi, GLM-4-Voice and Qwen3-Omni have gradually appeared. However, if the goal is not just to call ready-made checkpoints with billions of parameters, but to fully understand, train and modify a complete Omni model from scratch, the open-source community still lacks a sufficiently lightweight starting point with an end-to-end pipeline. A common way to bring speech into an Omni model is to chain ASR, LLM and TTS into a cascade: speech is first transcribed to text, the LLM processes it, and the answer is then synthesized back to speech. This is straightforward from an engineering perspective, but it adds an extra transcription step and noticeably hurts latency, prosody and emotional cues.
MiniMind-O attempts to fill this gap: speech and text are connected directly at the hidden-state level, while the trainable backbone remains only ~0.1B parameters and the end-to-end Omni pipeline is preserved. The Talker side adopts MTP (Multi-Token Prediction) to predict multiple Mimi codebook layers at once, and combines it with VAD to support real-time barge-in and near-duplex interaction—a practical engineering route for a tiny Omni model. The code, model weights, training data and technical report are all open-sourced. A single RTX 3090 can finish training on the mini dataset in about 2 hours. The goal remains the same: let everyone read the project from the first line of code, and train, from scratch, a model that can listen, see, think and speak:
![](images/omni_io_flow.png)
😊 Enjoy building.
---
#### 🎉 What this project provides
- A complete MiniMind-O architecture: Thinker, an independent Talker, audio / vision projectors, the Mimi codebook interface and the MTP audio head.
- A full SFT pipeline that covers T2A, I2T and A2A data, supporting full-parameter training, audio-projector-only training, vision-projector-only training, and DDP multi-GPU training.
- Two training datasets, `mini` and `full`. `mini` is meant for quick onboarding and runs the pipeline in ~2 hours on a single RTX 3090; `full` matches the released weights and covers Chinese speech and image tasks.
- Multiple built-in voice prompts, unseen voice prompts and voice cloning from arbitrary reference audio, making voice-control experiments easy to reproduce.
- A complete inference and demo toolkit: CLI, Web UI, streaming playback, barge-in interruption and a phone-mode demo.
- Key modules are written from scratch in native PyTorch without high-level third-party wrappers, while remaining compatible with `transformers` tokenizers and native weight formats.
- A companion technical report covers architecture, training curves, CER / WER evaluation, voice-cloning similarity and cross-model comparisons. See the Tech Report badge at the top.
#### 🎉 Released models
| Model | Backbone params | Release |
|---|---|---|
| minimind-3o | ~0.1B | 2026.05.05 |
| minimind-3o-moe | ~0.3B-A0.1B | 2026.05.05 |
---
#### 👉 Update Log
<details close>
<summary> <b>🔥 2026-05-05</b> </summary>
- First release of MiniMind-O: `minimind-3o` (115M) and `minimind-3o-moe` (312M-A115M).
- Thinker–Talker dual-path architecture. Talker uses MTP to predict multi-codebook Mimi codes and supports 24 kHz streaming speech generation and barge-in.
- Audio codec is Mimi (8 codebooks, 12.5 Hz, 24 kHz). Talker uses a shared backbone plus lightweight adapters at the codebook interface.
- Speech and visual features are extracted by frozen SenseVoice-Small and SigLIP2 respectively, and injected into the MiniMind hidden space through two-layer MLP projectors.
- Mini and full training datasets are released alongside; mini runs the full Thinker–Talker pipeline in ~2h on a single RTX 3090.
- 5 built-in voice prompts and 7 unseen voice prompts, with voice cloning and a phone-mode WebUI included.
</details>
# 📌 Quick Start
<details style="color:rgb(128,128,128)">
<summary>Reference hardware / software setup</summary>
* CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
* RAM: 128 GB
* GPU: NVIDIA GeForce RTX 3090 (24GB) * 8
* Ubuntu==20.04
* CUDA==12.2
* Python==3.10
* [requirements.txt](./requirements.txt)
</details>
## Step 0 (required)
### 1' Environment
```bash
# Clone the repository
git clone --depth 1 https://github.com/jingyaogong/minimind-o
# Install dependencies
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
```
### 2' Download resources
```bash
# Download SenseVoice-Small audio encoder to ./model/SenseVoiceSmall
modelscope download --model gongjy/SenseVoiceSmall --local_dir ./model/SenseVoiceSmall
# Download SigLIP2 vision encoder to ./model/siglip2-base-p32-256-ve
modelscope download --model gongjy/siglip2-base-p32-256-ve --local_dir ./model/siglip2-base-p32-256-ve
# Download Mimi audio codec to ./model/mimi
modelscope download --model gongjy/mimi --local_dir ./model/mimi
# Download CAM++ speaker encoder to ./model/campplus
modelscope download --model gongjy/campplus --local_dir ./model/campplus
# Download MiniMind LLM weights to ./out (used as the language backbone for training Omni)
modelscope download --model gongjy/minimind-3o-pytorch llm_768.pth --local_dir ./out
```
You can also `git clone` the corresponding repos from the [ModelScope Collection](https://modelscope.cn/collections/gongjy/MiniMind-O) or [HuggingFace Collection](https://huggingface.co/collections/jingyaogong/minimind-o) (LFS required); details omitted here.
After downloading, the directory should look like:
```text
minimind-o/
├── model/
│ ├── SenseVoiceSmall/
│ ├── siglip2-base-p32-256-ve/
│ ├── mimi/
│ ├── campplus/
│ └── ...
├── out/
│ └── llm_768.pth
└── ...
```
## Ⅰ 🚀 Inference
### 1' Download released weights
```bash
# Download released weights to ./out
modelscope download --model gongjy/minimind-3o-pytorch --local_dir ./out
```
### 2' Command-line chat
```bash
python eval_omni.py --load_from model --weight sft_omni
```
To use the Transformers-format model, download the model directory first:
```bash
git clone https://huggingface.co/jingyaogong/minimind-3o
python eval_omni.py --load_from minimind-3o
```
### 3' Launch WebUI (optional)
```bash
# ⚠️ Copy the Transformers-format model folder into ./scripts/. The web_demo_omni script
# automatically scans this directory for sub-folders that contain weight files; it
# raises an error if none is found.
cp -r minimind-3o ./scripts/minimind-3o
cd scripts && python web_demo_omni.py
```
## Ⅱ 🛠️ Training
<details style="color:rgb(128,128,128)">
<summary>Verify Torch is using CUDA</summary>
```python
import torch
print(torch.cuda.is_available())
```
If unavailable, please download the matching `.whl` from [torch_stable](https://download.pytorch.org/whl/torch_stable.html) and install it manually.
</details>
### 1' Download data
For a quick start, downloading only the `_mini` parquet files from the [dataset link](https://huggingface.co/datasets/jingyaogong/minimind-o_dataset) and placing them under `./dataset` is enough.
### 2' Train
The recommended mini training pipeline is shown below. It is meant to be run from the `trainer/` directory; equivalently, run `cd trainer && bash train.sh`:
```bash
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_t2a_mini.parquet --epochs 1 --batch_size 40 --use_compile 1 --from_weight llm --save_weight sft_zero --max_seq_len 512 --use_wandb --use_moe 0
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_a2a_mini.parquet --epochs 1 --batch_size 40 --use_compile 0 --from_weight sft_zero --save_weight sft_zero --max_seq_len 640 --mode audio_proj --use_wandb --use_moe 0
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 2e-5 --data_path ../dataset/sft_a2a_mini.parquet --epochs 1 --batch_size 16 --use_compile 0 --from_weight sft_zero --save_weight sft_zero --max_seq_len 768 --use_wandb --use_moe 0
```
### 3' Test the trained model (optional)
Make sure the model `*.pth` to be tested is placed under `./out/`.
```bash
python eval_omni.py --weight sft_omni
```
# 📌 Model Details
The language backbone of MiniMind-O comes from the sister project [MiniMind](https://github.com/jingyaogong/minimind). For LLM architecture and training details, please refer to that repository. Even without going into the LLM internals, you can still follow the Quick Start section above to train MiniMind-O end-to-end.
## Ⅰ Architecture overview
![](./images/architecture.jpg)
MiniMind-O consists of two paths: Thinker and Talker. Thinker is responsible for understanding text, speech and image inputs and producing a semantic-level text reply. Talker takes the semantic conditions from Thinker and uses MTP to jointly predict multi-codebook Mimi audio codes, which the audio decoder finally restores into streaming speech. The point is not to chain ASR, LLM and TTS together, but to keep text reasoning, speech generation and streaming interaction inside a single unified sequence.
Text inputs go directly into the language backbone; speech and images are first encoded by the Audio Encoder and Vision Encoder respectively, and then projected into the MiniMind hidden space. Voice information is provided either by a Speaker Encoder or by reference-audio codes; combined with VAD at inference time, this enables listen-while-speaking, real-time barge-in and near-duplex interaction. Later sections describe the projectors, sequence layout and training objectives in more detail; for code-level details, please refer to `model/model_omni.py` and the [technical report](http://arxiv.org/abs/2605.03937).
![](./images/input_token_layout.jpg)
The figure above shows how text tokens, speech features, image features and voice conditions are laid out in the input sequence.
## Ⅱ Multimodal understanding on the Thinker side
Thinker receives text, speech and image information uniformly and produces a semantic-level text reply. Text tokens enter the language backbone directly, while speech and image features are injected into placeholder positions through their respective projectors, so that all modalities are eventually modeled within the same sequence.
## Ⅲ Middle-layer Bridge
The representation passed from Thinker to Talker is taken from a middle layer rather than the embedding layer or the final layer. Embedding layers carry too little semantic information, while the final layer is overly shaped towards next-token prediction. A middle layer typically already fuses contextual and cross-modal information without being over-tuned by the LM head, which makes it a better conditioning source for speech generation. By default `bridge_layer = num_hidden_layers // 2 - 1`, and it can be adjusted through configuration at different scales.
## Ⅳ Speech generation on the Talker side
Talker turns the semantic states from Thinker into 8 streams of Mimi codebook codes. It uses MTP to predict multiple audio codebooks simultaneously, instead of running each codebook through a separate long path. To control the additional parameter count inside a 0.1B model, the audio embedding and output head share a common backbone with lightweight per-codebook adapters. This preserves the distributional differences between codebooks while avoiding a full parameter copy for each codebook layer.
## Ⅴ Sequence format and streaming decoding
![](./images/sequence_format.jpg)
MiniMind-O places text tokens and 8 audio-code streams in the same training sample: Thinker handles the text sequence, Talker handles the audio-code sequence, and speech / image / voice conditions are injected through placeholders or reference codes. Loss on target text and target audio is computed only after the reply starts; reference and conditioning regions serve only as conditions and are not part of the reconstruction target.
For streaming generation, the model emits text tokens while simultaneously filling in 8 layers of Mimi codes via MTP and a delay schedule. The Mimi decoder can incrementally reconstruct the 24 kHz waveform, so playback does not have to wait for the full reply to finish.
## Ⅵ Voice control
Voice control is realized through in-context voice cloning: reference audio is first encoded into a voice prompt, and then fed to Talker as a contextual condition, instead of fine-tuning weights or rewriting the text prompt to specify a voice. The model can additionally use a speaker embedding to provide a more stable speaker constraint; switching the voice at inference time only requires changing these conditioning inputs, while the Thinker prompt and Talker weights remain unchanged.
The default release ships with 5 built-in voice prompts (dylan, eric, serena, uncle_fu, vivian), and reserves 7 unseen prompts for evaluation (arthur, chelsie, cherry, ethan, jennifer, momo, moon).
## Ⅶ Modules and parameter scale
The "0.1B" referenced for MiniMind-O denotes the trainable backbone composed of Thinker, Talker and the two projectors. For the released checkpoints, `minimind-3o` is about 113M and `minimind-3o-moe` is about 315M. The Audio Encoder, Vision Encoder and Speech Codec are frozen external side modules used only for feature extraction or audio (de)coding; together they contain about 425M parameters and are not counted as active MiniMind-O parameters.
The table below counts the main module sizes per released model. Trainable counts are based on PyTorch modules, with tied embeddings deduplicated.
| Counting scope | minimind-3o | minimind-3o-moe |
|---|---:|---:|
| Trainable backbone | 113.13M | 314.89M |
| Frozen external modules | 424.70M | 424.70M |
| Total loaded at runtime | 537.83M | 739.59M |
| Module | Implementation | Key configuration | Status / params (~3o / ~3o-moe) |
|---|---|---|---|
| Thinker | MiniMind Transformer | 8 layers, hidden 768 | trainable, 63.91M / 198.42M |
| Talker | Standalone MiniMind blocks | 4 layers, 8 codebook heads | trainable, 47.05M / 114.30M |
| Audio projector | `MMAudioProjector` | 512 → 768 | trainable, 0.99M |
| Vision projector | `MMVisionProjector` | 768 → 768 | trainable, 1.18M |
| Audio encoder | SenseVoice-Small | 16 kHz speech features | frozen, 234.00M |
| Vision encoder | SigLIP2 base-p32-256 | 256×256 image, 64 tokens | frozen, 94.55M |
| Speech codec | Mimi | 8 codebooks, 12.5 Hz, 24 kHz | frozen, 96.15M |
| Speaker condition | CAM++ embedding | 192-d speaker vector | precomputed |
# 📌 Experiments
## Ⅰ Datasets
Dataset download: [ModelScope](https://www.modelscope.cn/datasets/gongjy/minimind-o_dataset) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind-o_dataset)
All speech data is stored uniformly as Mimi codes (8 codebooks, 12.5 Hz frame rate). Images are resized uniformly to 256×256 and encoded by SigLIP2 P32 into 64 patch tokens. The training data mainly comes from public Omni / speech-instruction corpora, including [VoiceAssistant-400K](https://huggingface.co/datasets/gpt-omni/VoiceAssistant-400K), [UltraChat-300K-SLAM-Omni](https://huggingface.co/datasets/worstchan/UltraChat-300K-SLAM-Omni) and others. A large amount of multi-speaker audio is additionally synthesized with Qwen3-TTS, and CAM++ is used to extract speaker embeddings as voice conditions. The I2T data follows the same source as the visual instruction data used in [MiniMind-V](https://github.com/jingyaogong/minimind-v); please refer to that project for the original composition and citations.
The repository ships two training sets, **mini** and **full**. The mini set is filtered from full using the "English + no-vision" criteria and works with `train_sft_omni.py` using the default `--data_path`. Its goal is to verify the Thinker–Talker pipeline, Mimi (de)coding, sequence layout and voice-injection path at low cost, rather than to reproduce the Chinese speech ability of the released models. A Chinese Talker has to handle more complex grapheme-to-phoneme mapping, prosodic pauses and multi-speaker stability, which is clearly harder than English and cannot be expected to converge within ~2 hours on a single RTX 3090.
The full set corresponds to the released `minimind-3o` / `minimind-3o-moe` checkpoints and covers Chinese-English T2A / A2A as well as image-to-text. Sizes and language ratios are listed below; this is the actual training source behind the CER / voice-similarity numbers reported in the paper.
T2A means Text-to-Audio, A2A means Audio-to-Audio, and I2T means Image-to-Text.
| Dataset | Subset | Input speech | Output speech | Note |
|---|---|---|---|---|
| `sft_t2a_mini` | English T2A | — | ~470.14 h | mini onboarding |
| `sft_a2a_mini` | English A2A | ~74.64 h | ~56.60 h | mini onboarding |
| `sft_t2a` | zh+en T2A | — | ~1636.01 h | full training |
| `sft_a2a` | zh+en A2A | ~1711.97 h | ~423.40 h | full training |
| `sft_i2t` | Image I2T | — | — | full training |
In `sft_t2a`, Chinese / English / mixed samples account for 45.7% / 46.5% / 7.8% respectively; in `sft_a2a` the ratios are 70.8% / 21.2% / 8.0%. This distribution is directly reflected in behavior: short Chinese and short English replies are usually stable, while longer English speech is more prone to mispronunciation and word omissions. The mini subset keeps only English, so even with a tight budget on parameters and data, the within-language CER stays in a usable range.
## Ⅱ Training
The training entry point is `train_sft_omni.py`, and the recommended pipeline can be found in `trainer/train.sh`. Full training is not split into multiple complex pretraining stages; instead, capabilities are introduced incrementally along the data flow:
![](./images/training_pipeline.jpg)
- `sft_t2a`: align text with speech output first, so that Talker learns to generate Mimi codes under Thinker's semantic conditions;
- `sft_a2a`: bring in speech inputs, so that the model can enter the same Thinker–Talker reply path from speech instructions;
- `sft_i2t`: align the visual path last; the `vision_proj` mode updates only the vision projector to avoid image data overwriting language and speech abilities.
Among training modes, `all` updates MiniMind / Talker / projectors, while `audio_proj` and `vision_proj` are used solely to align the corresponding projector. SenseVoice-Small, SigLIP2 and Mimi are kept frozen throughout. The Dense and MoE variants share the same data ordering. The mini commands are meant only to make the pipeline runnable end-to-end and finish in ~2 hours on a single RTX 3090 by default; the released weights correspond to full training.
T2A and A2A loss curves during full training are shown below for reference:
![](./images/t2a_training_curves.jpg)
> `sft_t2a`: text-to-speech-output path
![](./images/a2a_training_curves.jpg)
> `sft_a2a`: loss after speech inputs are added
Early spikes caused by an incompatible weight resume have been removed from the T2A curve. The MoE variant has more total parameters but a similar number of active parameters compared with Dense, which makes it more useful as a capacity-allocation reference.
## Ⅲ Model weights
| Format | ModelScope | HuggingFace |
|---|---|---|
| PyTorch (`*.pth`) | [minimind-3o-pytorch](https://www.modelscope.cn/models/gongjy/minimind-3o-pytorch) | [minimind-3o-pytorch](https://huggingface.co/jingyaogong/minimind-3o-pytorch) |
| Transformers | [minimind-o collection](https://modelscope.cn/collections/gongjy/MiniMind-O) | [minimind-o collection](https://huggingface.co/collections/jingyaogong/minimind-o) |
> The Transformers version contains both `minimind-3o` and `minimind-3o-moe` and is suitable for direct use with `eval_omni.py` and the WebUI. The native PyTorch weights are mainly intended for training, reproducing experiments and continued fine-tuning.
# 📌 Evaluation
There is currently no unified evaluation protocol for Omni models: different works differ in the LLM backbone, the audio synthesizer and the system goal. Some focus on the LLM's own knowledge and reasoning and report MMLU, HumanEval and related benchmarks; some emphasize streaming latency and audio quality; some highlight speech-consistency metrics; and others focus on natural interaction or broader Omni generation. Most of these systems are continually trained from a state-of-the-art open-source LLM, while MiniMind's 0.06B backbone is clearly not competitive on complex knowledge QA, math reasoning, code generation or long open-ended replies, and the Talker's naturalness, prosody and stability are also weaker than those of full-scale systems.
Therefore, the goal here is not to chase a comprehensive leaderboard, but to focus on a few more reproducible local evaluations and use cases: a Talker hidden-size ablation, voice-cloning similarity, CER / WER comparisons under identical questions and identical ASR pipelines, and qualitative A2A, I2A and real-time interaction examples. CER / WER are mainly used to inspect text consistency, while audio quality, naturalness and human preference are left to qualitative samples and actual listening tests.
## Ⅰ Talker hidden-size ablation
If only speech generation is considered, scaling Talker to 1024 / 2048 hidden size or stacking more layers would obviously be more stable. But MiniMind-O has to fit the entire Omni pipeline within ~0.1B parameters, and cannot afford to allocate most of the budget to the acoustic side. Once Thinker and Talker are decoupled, language understanding and cross-modal fusion are mainly carried by Thinker, while Talker only renders Mimi codes from semantic conditions; this makes a small Talker possible. The rendering here is not "predict semantic tokens and hand them off to an external acoustic model"—Talker directly produces decodable Mimi acoustic codes, so the real bottleneck is at the output side: it has to handle 8 Mimi codebooks rather than a single next-token-prediction stream.
384-d is tempting, since the dense version compresses to ~88M; 512-d is also lighter. But the table below shows that smaller does not automatically mean better allocated: short utterances remain acceptable, but medium-to-long ones are more prone to word drops, repetitions and pronunciation drift. 768-d was kept in the end because it matches the MiniMind backbone width and can be initialized from the last 4 layers of Thinker; the parameter count remains around 0.1B, training cost does not increase noticeably, and consistency is clearly more stable.
| Variant | Talker hidden | Params | Avg CER ↓ | Short ↓ | Mid / Long ↓ |
|---|---|---|---|---|---|
| Dense | 768 | 115.29M | **0.0897** | 0.1528 | 0.0874 / 0.0675 |
| Dense | 512 | 96.13M | 0.1745 | 0.2709 | 0.2455 / 0.0976 |
| Dense | 384 | 88.72M | 0.2767 | 0.3904 | 0.1865 / 0.4046 |
| MoE | 768 | 317.05M-A115.33M | **0.0900** | 0.2075 | 0.0533 / 0.0271 |
| MoE | 512 | 261.32M-A96.17M | 0.1265 | 0.0711 | 0.1490 / 0.1464 |
| MoE | 384 | 240.04M-A88.75M | 0.3280 | 0.3757 | 0.2777 / 0.4313 |
Dense and MoE CERs should not be compared directly across architectures: under the same question, the two Thinkers may produce different content with different lengths, leading to different synthesis difficulty for Talker. What matters more is the within-architecture trend: 768 clearly outperforms 512 and 384 in both cases.
## Ⅱ Voice-cloning similarity
Voice cloning is one of the more beta-quality features in this release. To our knowledge, most open-source Omni models support only fixed output voices, while minimind-3o tries to fit multi-voice generation into a single Talker. This goal is harder than simply "being able to talk", because the model needs not only to say the right content, but also to preserve speaker cues while generating Mimi codes.
Quality has not yet reached high-fidelity cloning: the same reference voice does not always stay consistent across questions, and longer utterances can drift because of pronunciation and rhythm issues. But basic male / female differences, intonation tendencies and parts of the prosody are distinguishable.
The CAM++ speaker-embedding cosine similarity below is only an automatic reference. Seen comes from the 5 built-in voices in `voices.pt`; Unseen comes from 7 voices in `voices_unseen.pt` that were never seen during training. Each voice uses the same set of text questions and only the voice condition is swapped.
Per-speaker breakdown:
| Split | Speaker | Dense ↑ | MoE ↑ |
|---|---|---|---|
| Seen | dylan | 0.6997 | 0.6837 |
| Seen | eric | 0.5289 | 0.4232 |
| Seen | serena | 0.7092 | 0.7041 |
| Seen | uncle_fu | 0.7241 | 0.7337 |
| Seen | vivian | 0.5744 | 0.5888 |
| Unseen | arthur | 0.7171 | 0.6750 |
| Unseen | chelsie | 0.6437 | 0.6240 |
| Unseen | cherry | 0.5689 | 0.5678 |
| Unseen | ethan | 0.4783 | 0.4847 |
| Unseen | jennifer | 0.4749 | 0.4003 |
| Unseen | momo | 0.6470 | 0.5720 |
| Unseen | moon | 0.4282 | 0.6673 |
Overall, minimind-3o and minimind-3o-moe land at similar averages, both slightly above the early baseline. This suggests that voice retention is not primarily determined by inactive expert capacity; the more direct factors are reference-clip quality, the separability of CAM++ embeddings, and the stability of Talker generation itself. Per-speaker, voices like uncle_fu, serena and arthur are easier to preserve, with at least one variant exceeding 0.70; outliers like eric and moon are more sensitive to generation quality. In other words, this capability already separates some speaker characteristics, but is still some distance away from a product-level "given a reference clip, faithfully reproduce its timbre" experience.
### Voice-cloning ablation samples (audio playback)
For a more direct listening test, seed=42 and temperature=0.7 are fixed, and one generated sample is shown per voice. The only variables are the reference audio codes and speaker embedding. As a control, the default output without any reference voice condition is shown first (the spoken text is identical across all samples):
https://github.com/user-attachments/assets/b31fd8f2-e3af-4fed-ba19-65424b59bec6
#### Seen voices
"Seen" means voices that appeared in training data, used to inspect how well the model preserves familiar speakers.
<table>
<tr><th width="100">Speaker</th><th width="380">Reference</th><th width="380">Output</th><th width="80">Avg</th></tr>
<tr><td>dylan</td><td>
https://github.com/user-attachments/assets/070ea3ab-0e8e-4aa0-84b5-af8d3c4e2725
</td><td>
https://github.com/user-attachments/assets/eb2da7ed-173c-47e9-9431-7bdb5a9b7385
</td><td>0.6712</td></tr>
<tr><td>eric</td><td>
https://github.com/user-attachments/assets/c74aa5dc-1edd-44c1-9546-6e57194c2f60
</td><td>
https://github.com/user-attachments/assets/f3fa8906-4e14-4610-a9d9-c16c915ca1b3
</td><td>0.4430</td></tr>
<tr><td>serena</td><td>
https://github.com/user-attachments/assets/0eeeac87-fa70-4025-b66e-1f0197f2b434
</td><td>
https://github.com/user-attachments/assets/c5901dca-4b2a-47f5-9b30-c89de54f908e
</td><td>0.6600</td></tr>
<tr><td>uncle_fu</td><td>
https://github.com/user-attachments/assets/fdd1bb28-6648-44bf-8bcb-4509e709e347
</td><td>
https://github.com/user-attachments/assets/95b480f1-f015-4712-8d7c-17db465f6584
</td><td>0.6632</td></tr>
<tr><td>vivian</td><td>
https://github.com/user-attachments/assets/f64731c4-67a3-4e18-b7d7-61bf44ef4bdd
</td><td>
https://github.com/user-attachments/assets/3f1cc9bb-16d2-4ce0-a473-40676cf4523e
</td><td>0.5320</td></tr>
</table>
#### Unseen voices
"Unseen" means voices not seen during training, used to inspect zero-shot transfer of a new reference voice into generated speech.
<table>
<tr><th width="100">Speaker</th><th width="380">Reference</th><th width="380">Output</th><th width="80">Avg</th></tr>
<tr><td>arthur</td><td>
https://github.com/user-attachments/assets/46e720ba-c413-4731-8275-a9a2cb483690
</td><td>
https://github.com/user-attachments/assets/f489b463-c7d7-4ad6-9e4d-53333172c66c
</td><td>0.6479</td></tr>
<tr><td>chelsie</td><td>
https://github.com/user-attachments/assets/3430ecdb-6de8-4fb0-a6a7-ad82bdce01a1
</td><td>
https://github.com/user-attachments/assets/e598dbc2-ba28-4c38-b52d-6fa6c2349a5b
</td><td>0.5975</td></tr>
<tr><td>cherry</td><td>
https://github.com/user-attachments/assets/f9166af6-3a98-42f3-9cf8-ad105eea87d6
</td><td>
https://github.com/user-attachments/assets/eccca693-4708-409a-88f7-85eb25f66fe6
</td><td>0.5418</td></tr>
<tr><td>ethan</td><td>
https://github.com/user-attachments/assets/9c992505-2046-483e-a7cf-50ec18a5e329
</td><td>
https://github.com/user-attachments/assets/98013c5e-f5b5-4e1a-bc0e-a0f0be5d3240
</td><td>0.4323</td></tr>
<tr><td>jennifer</td><td>
https://github.com/user-attachments/assets/924b035d-5c7c-45a5-a8f8-5dbdc18f71db
</td><td>
https://github.com/user-attachments/assets/853d1370-0065-4567-9a71-dc88a6a34d56
</td><td>0.4052</td></tr>
<tr><td>momo</td><td>
https://github.com/user-attachments/assets/7e97f524-da6d-4a2f-9095-e7f99262f4a5
</td><td>
https://github.com/user-attachments/assets/4c193c8f-8750-4424-acba-2bd13089a634
</td><td>0.5968</td></tr>
<tr><td>moon</td><td>
https://github.com/user-attachments/assets/527df88a-adc0-48d3-9a6a-827ca1ba7fb0
</td><td>
https://github.com/user-attachments/assets/3f533e26-1ad8-4ab3-baf1-21267734d3ee
</td><td>0.5874</td></tr>
</table>
## Ⅲ Cross-model English T2A comparison
We selected 20 English questions, all constrained by `Answer briefly in one short sentence`. The intent is not to evaluate open-ended English ability, but to keep response lengths within a similar range. The three models then synthesize audio, which is uniformly transcribed by Qwen3-ASR; CER / WER between transcription and target text are used to compare Talker-side textual consistency.
| Length bucket | [Mini-Omni](https://huggingface.co/gpt-omni/mini-omni) CER/WER | [Mini-Omni2](https://huggingface.co/gpt-omni/mini-omni2) CER/WER | minimind-3o CER/WER |
|---|---|---|---|
| short (≤15w) | 0.0195 / 0.0384 (n=8) | 0.0503 / 0.0584 (n=14) | 0.0531 / 0.0417 (n=8) |
| mid (16–30w) | 0.0038 / 0.0052 (n=12) | 0.0062 / 0.0076 (n=6) | 0.1327 / 0.1420 (n=11) |
| long (31–60w) | — | — | 0.0431 / 0.0508 (n=1) |
For replies of ≤15 words, minimind-3o is already close to Mini-Omni2; the gap really opens up at 16–30 words. This length is no longer a simple phrase, and Talker must keep pronunciation, rhythm and surface form consistent in a complete short sentence simultaneously. This is also the regime where the current 0.1B Talker most easily exposes its instability.
## Ⅳ Cross-model vision-language comparison
[Mini-Omni](https://huggingface.co/gpt-omni/mini-omni) does not support a VL path, so the comparison is between [Mini-Omni2](https://huggingface.co/gpt-omni/mini-omni2) (0.5B) and minimind-3o (0.1B). On 9 synthetic images, both models generate English answers, which are then uniformly transcribed and used to compute CER / WER as a vision-to-speech consistency reference.
| Model | Params | Avg CER ↓ | Avg WER ↓ |
|---|---|---|---|
| [Mini-Omni2](https://huggingface.co/gpt-omni/mini-omni2) | 0.5B | 0.7609 | 0.9756 |
| minimind-3o | 0.1B | 0.8241 | 1.0293 |
These numbers should not be read as the absolute correctness of open-ended image description. Image captioning has many equivalent expressions, and synonym choices and word order both affect CER / WER, so high absolute values are expected. Under the same automatic pipeline, minimind-3o trails behind Mini-Omni2 but stays in the same order of magnitude, with roughly 1/5 the parameters.
## Ⅴ Qualitative samples
![](./images/qual_a2a.jpg)
In speech-to-speech samples, the input is real speech, Thinker organizes the semantics, and Talker renders speech. Short replies are again the more stable regime; Chinese explanatory questions usually produce coherent answers, while English pronunciation and rhythm are relatively more stable.
<table>
<tr>
<td>
https://github.com/user-attachments/assets/c85809b2-4787-4656-9c7e-55b693798494
</td>
<td>
https://github.com/user-attachments/assets/354a5eec-c147-4d18-8c7a-942bd2a0b4b0
</td>
</tr>
</table>
![](./images/image2audio_qualitative.jpg)
The image-QA samples chain visual encoding, text generation and speech rendering inside the same path. The current model usually captures the main object and the rough scene, but fine-grained spatial relations, counts and attributes are still often wrong, which makes it more suitable as a reproducible baseline for tiny-model Omni pipelines.
<table>
<tr>
<td>
https://github.com/user-attachments/assets/244e08b0-5b12-449e-a7a2-2a2139c5d62d
</td>
<td>
https://github.com/user-attachments/assets/3e8d0a76-282d-4a9d-9726-a954cf80198a
</td>
</tr>
</table>
## Ⅵ Real-time interaction
![](./images/realtime_interaction.jpg)
This is the real-time interaction interface. Once the user stops speaking, Thinker first finishes the semantic-side prefill, Talker then starts to emit audio codes incrementally, and the Mimi decoder writes the 24 kHz waveform as it receives codes. The barge-in example shows another path that is closer to a real conversation: when the user starts speaking again while the model is talking, the system interrupts the current generation and re-enters the prefill–reply flow. Interruption detection here is still based on a simple VAD threshold, not yet semantic-level barge-in; but from an engineering loop-closure perspective, the system can already fall back from speaking to listening and process the next turn.
### 🧩 Possible future improvements
The current model still has clear gaps compared with large-scale Omni systems, and there is no need to gloss over them. Long-form speech naturalness, complex visual reasoning, open-ended English mid/long replies and voice stability are not its strong areas. The visual path is closer to a compact vision-to-speech link, and the MoE variant looks more like a capacity-allocation experiment than a same-FLOP optimum.
These limitations also point to several follow-ups: longer ICL contexts, finer prosody supervision, stronger vision encoders, more stable voice conditions, and systematic sweeps over the Bridge layer and the MTP codebook interface—all of which are worth continuing.
That said, the value of MiniMind-O lies exactly here. It compresses an entire Omni loop into the 0.1B regime, and ships code, weights and the main training data inside the same inspectable artifact. This means it is not just a demo, but a baseline small enough, transparent enough, and reproducible enough to rebuild from scratch and modify further. For people who want to understand Thinker–Talker decoupling, the MTP codebook interface, in-context voice cloning, and the middle-hidden bridge, it offers a set of design choices that can actually be verified by hand.
# 📌 Acknowledgements
> [!TIP]
> If you find `MiniMind-O` helpful, consider giving us a ⭐ on GitHub.<br/>
> Given limited bandwidth there will inevitably be unknown bugs. Discussions, corrections and PRs in Issues are welcome.<br/>
> Your support is what keeps the project moving—thank you!
## 🤝 Contributors
<a href="https://github.com/jingyaogong/minimind-o/graphs/contributors">
<img width="200" src="https://contrib.rocks/image?repo=jingyaogong/minimind-o" />
</a>
## 😊 Credits
- [MiniMind](https://github.com/jingyaogong/minimind) / [MiniMind-V](https://github.com/jingyaogong/minimind-v) (backbone, data)
- [Qwen2.5-Omni / Qwen3-Omni](https://github.com/QwenLM/Qwen2.5-Omni) (inspiration, data)
- [Mini-Omni / Mini-Omni2](https://github.com/gpt-omni/mini-omni) (inspiration, data)
- [SLAM-Omni](https://aclanthology.org/2025.findings-acl.115/) (data)
- [SenseVoice](https://arxiv.org/abs/2407.04051) (component)
- [Mimi / Moshi](https://arxiv.org/abs/2410.00037) (component)
## 🫶 Supporters
<a href="https://github.com/jingyaogong/minimind-o/stargazers">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://bytecrank.com/nastyox/reporoster/php/stargazersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
<source media="(prefers-color-scheme: light)" srcset="https://bytecrank.com/nastyox/reporoster/php/stargazersSVG.php?user=jingyaogong&repo=minimind-o"/>
<img alt="github contribution grid snake animation" src="https://bytecrank.com/nastyox/reporoster/php/stargazersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
</picture>
</a>
<a href="https://github.com/jingyaogong/minimind-o/network/members">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://bytecrank.com/nastyox/reporoster/php/forkersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
<source media="(prefers-color-scheme: light)" srcset="https://bytecrank.com/nastyox/reporoster/php/forkersSVG.php?user=jingyaogong&repo=minimind-o"/>
<img alt="github contribution grid snake animation" src="https://bytecrank.com/nastyox/reporoster/php/forkersSVG.php?user=jingyaogong&repo=minimind-o&theme=dark"/>
</picture>
</a>
<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/svg?repos=jingyaogong/minimind-o&type=Date&theme=dark"/>
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/svg?repos=jingyaogong/minimind-o&type=Date"/>
<img alt="Star History Chart" src="https://api.star-history.com/svg?repos=jingyaogong/minimind-o&type=Date&theme=dark"/>
</picture>
# 🎓 Citation
If MiniMind-O helps your research or work, please cite:
```bibtex
% Cite the technical report when referencing the model design or experimental results.
@article{minimind-o-report,
title = {MiniMind-O Technical Report: An Open Small-Scale Speech-Native Omni Model},
author = {Jingyao Gong},
journal = {arXiv preprint arXiv:2605.03937},
year = {2026}
}
% Cite the GitHub repo when referencing the open-source codebase or released weights.
@misc{minimind-o,
title = {MiniMind-O: Train a Tiny Omni Model from Scratch},
author = {Jingyao Gong},
year = {2026},
url = {https://github.com/jingyaogong/minimind-o},
note = {GitHub repository, accessed 2026}
}
```
# 📜 License
This repository is released under the [Apache-2.0 License](LICENSE).

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

After

Width:  |  Height:  |  Size: 17 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 14 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 25 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 30 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 14 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 19 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 13 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 17 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 20 KiB

Binary file not shown.

Binary file not shown.

Binary file not shown.

344
dataset/omni_dataset.py Normal file
View File

@ -0,0 +1,344 @@
import torch
import os
import math
import random
import numpy as np
import soundfile as sf
import librosa
import json
import io
from PIL import Image
from scipy.signal import resample
from torch.utils.data import Dataset
import pyarrow as pa
import pyarrow.parquet as pq
os.environ["TOKENIZERS_PARALLELISM"] = "false"
def pre_processing_chat(conversations, add_system_ratio=0.2):
if any(conv.get('tools') for conv in conversations): return conversations
SYSTEM_PROMPTS = [
"你是一个知识丰富的AI,尽力为用户提供准确的信息。",
"你是minimind,一个小巧但有用的语言模型。",
"你是一个专业的AI助手,请提供有价值的回答。",
"你是minimind,请尽力帮助用户解决问题。",
"你是一个可靠的AI,请给出准确的回答。",
"You are a helpful AI assistant.",
"You are minimind, a lightweight intelligent assistant.",
"You are a friendly chatbot. Please answer the user's questions carefully.",
"You are a knowledgeable AI. Try your best to provide accurate information.",
"You are minimind, a small but useful language model."
]
if conversations[0].get('role') != 'system':
if random.random() < add_system_ratio:
return [{'role': 'system', 'content': random.choice(SYSTEM_PROMPTS)}] + conversations
return conversations
def post_processing_chat(prompt_content, empty_think_ratio=0.2):
if '<think>\n\n</think>\n\n' in prompt_content and random.random() > empty_think_ratio:
prompt_content = prompt_content.replace('<think>\n\n</think>\n\n', '')
return prompt_content
class OmniDataset(Dataset):
def __init__(self, data_path, tokenizer, audio_processor=None, vision_processor=None,
max_length=1200, audio_special_token='<|audio_pad|>', image_special_token='<|image_pad|>',
audio_stop_token=2050, # <|audio_stop|>
audio_pad_token=2049, # <|audio_pad|>
audio_spk_token=2051, # <|audio_spk|>
audio_vocab_size=2112, # 2048 mimi codes + 64 special tokens
scheduled_sampling=0.05,
image_token_len=64):
super().__init__()
tables = [pa.Table.from_batches(pq.ParquetFile(p.strip()).iter_batches()) for p in data_path.split(',')]
tables = [t.cast(pa.schema([f.with_type(pa.large_string()) if pa.types.is_string(f.type) else f for f in t.schema])) for t in tables]
self.table = pa.concat_tables(tables, promote_options='default')
self.tokenizer = tokenizer
self.audio_processor = audio_processor
self.vision_processor = vision_processor
self.max_length = max_length
self.audio_token = audio_special_token
self.image_token_len = image_token_len
self.image_token = image_special_token * image_token_len
self.audio_stop_token = audio_stop_token
self.audio_pad_token = audio_pad_token
self.audio_spk_token = audio_spk_token
self.audio_vocab_size = audio_vocab_size
self.scheduled_sampling_prob = scheduled_sampling
self.text_vocab_size = len(tokenizer)
self.image_token_id = tokenizer.encode(image_special_token, add_special_tokens=False)[0]
self.audio_token_id = tokenizer.encode(audio_special_token, add_special_tokens=False)[0]
self.think_end_ids = tokenizer.encode('</think>\n\n', add_special_tokens=False)
self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
def __len__(self):
return len(self.table)
@staticmethod
def process_audio(audio_path, audio_processor):
"""加载音频并预处理成fbank,返回 (fbank (T,560), valid_len=encoder输出帧数)"""
wav, sr = sf.read(audio_path)
if wav.ndim > 1: wav = wav.mean(axis=1)
if sr != 16000: wav = librosa.resample(wav.astype(float), orig_sr=sr, target_sr=16000)
inputs = audio_processor(wav.astype(np.float32), sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
valid_len = inputs.attention_mask.sum().item()
return inputs.input_features.squeeze(0), valid_len
def augment_wav(self, wav, sr=16000):
# 随机变速(0.7~1.6x):改变音频时长和音调,覆盖快/慢语速
if random.random() < 0.5:
speed = random.uniform(0.7, 1.6)
wav = resample(wav, int(len(wav) / speed)).astype(np.float32)
# 随机加噪:叠加轻微高斯白噪声,模拟录音环境差异
if random.random() < 0.3:
noise = np.random.randn(len(wav)).astype(np.float32) * random.uniform(0.001, 0.01)
wav = wav + noise
# 随机音量:缩放振幅0.8~1.2倍,模拟说话音量变化
if random.random() < 0.3:
wav = wav * random.uniform(0.8, 1.2)
# 随机时间遮蔽:将0.25秒片段置零,模拟短暂静音/丢包
if random.random() < 0.2 and len(wav) > sr:
start = random.randint(0, len(wav) - sr // 4)
wav[start:start + sr // 4] = 0
# 随机低通滤波:移动平均模糊高频,模拟电话/低质量麦克风
if random.random() < 0.2:
k = random.choice([3, 5, 7])
wav = np.convolve(wav, np.ones(k) / k, mode='same').astype(np.float32)
# 随机混响:合成指数衰减脉冲响应卷积,模拟房间反射/回声
if random.random() < 0.3:
ir_len = int(sr * random.uniform(0.05, 0.2))
ir = np.random.randn(ir_len).astype(np.float32) * np.exp(-np.linspace(0, 10, ir_len))
ir[0] = 1.0
ir /= np.sqrt(np.sum(ir ** 2) + 1e-6)
wav = np.convolve(wav, ir, mode='same').astype(np.float32)
# 随机粉红噪声:1/f噪声模拟房间环境底噪(空调/远处人声)
if random.random() < 0.2:
pink = np.cumsum(np.random.randn(len(wav))).astype(np.float32)
pink /= np.max(np.abs(pink)) + 1e-6
wav = wav + pink * random.uniform(0.003, 0.015)
return np.clip(wav, -1.0, 1.0).astype(np.float32)
def augment_mel(self, fbank):
# fbank: (T, 560) — SenseVoice LFR 后的特征(时间维在前,频率维在后)
T, D = fbank.shape
# SpecAugment频率遮蔽:随机抹掉1~64维,防止对特定频段过拟合
if random.random() < 0.5:
f = random.randint(1, 64)
f0 = random.randint(0, D - f)
fbank[:, f0:f0 + f] = 0
# SpecAugment时间遮蔽:随机抹掉1~min(10,T)帧,提升对不完整输入的容错
if random.random() < 0.5 and T > 1:
t = random.randint(1, min(10, T))
t0 = random.randint(0, T - t)
fbank[t0:t0 + t, :] = 0
return fbank
def load_audio_inputs(self, audio_bytes):
if not audio_bytes: return None, 0
wav, sr = sf.read(io.BytesIO(audio_bytes))
if wav.ndim > 1: wav = wav.mean(axis=1)
if sr != 16000: wav = librosa.resample(wav.astype(float), orig_sr=sr, target_sr=16000)
wav = self.augment_wav(wav.astype(np.float32))
inputs = self.audio_processor(wav, sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
valid_len = inputs.attention_mask.sum().item()
return self.augment_mel(inputs.input_features.squeeze(0)), valid_len
def load_image_inputs(self, image_bytes):
if not image_bytes or self.vision_processor is None: return None
image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
inputs = self.vision_processor(images=image, return_tensors="pt")
if hasattr(inputs, 'keys'): return {k: v for k, v in inputs.items()}
return inputs.pixel_values
def create_chat_prompt(self, conversations, audio_features_length=0):
conversations = pre_processing_chat(conversations)
messages = []
is_last_user = lambda i: i == max(j for j, t in enumerate(conversations) if t['role'] == 'user')
for idx, turn in enumerate(conversations):
role, content = turn['role'], turn['content']
if role == 'user' and is_last_user(idx) and audio_features_length > 0:
ap = self.audio_token * audio_features_length
r = random.random()
if r < 0.4: content = ap
elif r < 0.6: content = content
elif r < 0.8: content = ap + '\n\n' + content
else: content = content + '\n\n' + ap
if '<image>' in content:
r = random.random()
if r < 0.2: content = '<image>\n' + content.replace('<image>', '').strip()
elif r < 0.4: content = '<image>\n\n' + content.replace('<image>', '').strip()
elif r < 0.6: content = content.replace('<image>', '').strip() + '\n' + '<image>'
else: content = content.replace('<image>', '').strip() + '\n\n' + '<image>'
messages.append({"role": role, "content": content})
prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
return post_processing_chat(prompt)
def generate_text_labels(self, input_ids):
labels = [-100] * len(input_ids)
ranges = []
i = 0
while i < len(input_ids):
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
start = i + len(self.bos_id)
end = start
while end < len(input_ids):
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
break
end += 1
ranges.append((start, end))
for j in range(start, min(end + len(self.eos_id), self.max_length)):
labels[j] = input_ids[j]
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
else:
i += 1
return labels, ranges
def apply_scheduled_sampling(self, input_ids, audio_labels, text_labels):
"""Scheduled Sampling: 用随机值替代部分GT,让模型学会从错误历史中恢复"""
if self.scheduled_sampling_prob <= 0:
return input_ids
audio_mask = (audio_labels != -100).any(dim=0) & (torch.rand(input_ids.size(1)) < self.scheduled_sampling_prob)
for i in range(8):
input_ids[i] = torch.where(audio_mask, torch.randint(0, self.audio_vocab_size, input_ids[i].shape), input_ids[i])
# 保护 image token 的连续性
text_mask = (text_labels != -100) & (input_ids[8] != self.image_token_id) & (torch.rand(input_ids.size(1)) < self.scheduled_sampling_prob)
input_ids[8] = torch.where(text_mask, torch.randint(0, self.text_vocab_size, input_ids[8].shape), input_ids[8])
return input_ids
def __getitem__(self, index: int):
conversations = json.loads(self.table['conversations'][index].as_py())
question_audios = self.table['question_audios'][index].as_py() if 'question_audios' in self.table.column_names else []
answer_audios = self.table['answer_audios'][index].as_py() if 'answer_audios' in self.table.column_names else []
image_bytes = self.table['image_bytes'][index].as_py() if 'image_bytes' in self.table.column_names else []
if image_bytes and not isinstance(image_bytes, list): image_bytes = [image_bytes]
ref_audios = self.table['ref_audios'][index].as_py() if 'ref_audios' in self.table.column_names else []
spk_emb_raw = self.table['spk_emb'][index].as_py() if 'spk_emb' in self.table.column_names else []
# 随机截断到某一轮(每轮=user+assistant)
asst_indices = [i for i, t in enumerate(conversations) if t['role'] == 'assistant']
if len(asst_indices) > 1:
rand_idx = random.randint(0, len(asst_indices) - 1)
# 从随机轮次开始,向前回退直到长度安全
for i in range(rand_idx, -1, -1):
conversations = conversations[:asst_indices[i] + 1]
test_prompt = self.create_chat_prompt(conversations, 0)
if len(self.tokenizer(test_prompt).input_ids) + 100 < self.max_length:
break
# 加载最后一个user的图像(按user轮次索引访问,与audio一致)
pixel_values = None
user_count = sum(1 for t in conversations if t['role'] == 'user')
if image_bytes and len(image_bytes) > 0 and self.vision_processor:
pixel_values = self.load_image_inputs(image_bytes[0])
# 只加载最后一个user的audio(按user轮次索引访问)
audio_inputs, audio_len, audio_features_length = None, 0, 0
user_count = sum(1 for t in conversations if t['role'] == 'user')
if question_audios and user_count > 0 and user_count <= len(question_audios) and self.audio_processor:
audio_bytes = question_audios[user_count - 1]
if audio_bytes:
mel, valid_len = self.load_audio_inputs(audio_bytes)
if mel is not None:
audio_inputs = mel.unsqueeze(0)
audio_len = valid_len
audio_features_length = valid_len or 1
# 混合训练时,无音频样本返回dummy tensor保持batch索引尽可能对齐 (SenseVoice: T x 560)
if audio_inputs is None and self.audio_processor:
audio_inputs = torch.zeros(1, 1, 560)
audio_len = 0
if pixel_values is None and self.vision_processor:
pixel_values = {'pixel_values': torch.zeros(1, 3, 256, 256)}
# 从answer_audios获取最后一个assistant的音频codes
last_audio_codes = None
asst_count = sum(1 for t in conversations if t['role'] == 'assistant')
if answer_audios and asst_count > 0 and asst_count <= len(answer_audios):
tokens = answer_audios[asst_count - 1]
if tokens:
audio_codes_8layers = [[] for _ in range(8)]
for i in range(0, len(tokens) - 7, 8):
for j in range(8): audio_codes_8layers[j].append(tokens[i + j])
for layer in audio_codes_8layers: layer.append(self.audio_stop_token)
last_audio_codes = audio_codes_8layers
# 生成prompt (text input_ids)
prompt = self.create_chat_prompt(conversations, audio_features_length)
if pixel_values is not None: prompt = prompt.replace('<image>', self.image_token)
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
# PAD input_ids到max_length
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
# 生成labels(只训练最后一个assistant)
text_labels, assistant_ranges = self.generate_text_labels(input_ids)
for start, end in assistant_ranges[:-1]:
text_labels[start:end] = [-100] * (end - start)
# 生成7层audio targets(只填充最后一个assistant)
Y_audio_layers = [[self.audio_pad_token] * self.max_length for _ in range(8)]
audio_labels = [[-100] * self.max_length for _ in range(8)]
if assistant_ranges and last_audio_codes:
assistant_start, assistant_end = assistant_ranges[-1]
for pos in range(assistant_start, min(assistant_end, assistant_start + 50)):
if input_ids[pos:pos + len(self.think_end_ids)] == self.think_end_ids:
assistant_start = pos + len(self.think_end_ids)
break
# spk_emb 占位 + ref_codes 右对齐(50% 概率 drop ref_codes,只保留 spk)
has_spk = bool(spk_emb_raw)
has_ref = bool(ref_audios) and random.random() > 0.5
spk_reserve = 1 if has_spk else 0
if has_ref:
ref_codes = [[] for _ in range(8)]
for i in range(0, len(ref_audios) - 7, 8):
for j in range(8): ref_codes[j].append(ref_audios[i + j])
ref_len = len(ref_codes[0])
ref_start = max(spk_reserve, assistant_start - ref_len)
for layer_idx in range(8):
codes = ref_codes[layer_idx][-(assistant_start - ref_start):] if ref_len > (assistant_start - ref_start) else ref_codes[layer_idx]
for i, code in enumerate(codes):
Y_audio_layers[layer_idx][ref_start + i] = code
else:
ref_start = assistant_start
if has_spk and ref_start > 0:
spk_pos = ref_start - 1
for layer_idx in range(8):
Y_audio_layers[layer_idx][spk_pos] = self.audio_spk_token
# target codes 填充到 assistant_start 之后(参与 loss)
for layer_idx in range(8):
codes = last_audio_codes[layer_idx]
start_pos = assistant_start + layer_idx + 1
for i, code in enumerate(codes):
if start_pos + i < self.max_length:
Y_audio_layers[layer_idx][start_pos + i] = code
audio_labels[layer_idx][start_pos + i] = code
# 构造9路输入:input_ids = (9, T) = 8路audio + 1路text
X_audio = torch.tensor([layer[:-1] for layer in Y_audio_layers], dtype=torch.long) # (8, T-1)
X_text = torch.tensor(input_ids[:-1], dtype=torch.long) # (T-1,)
input_ids = torch.cat((X_audio, X_text.unsqueeze(0)), dim=0) # (9, T-1)
text_labels = torch.tensor(text_labels[1:], dtype=torch.long) # (T-1,)
audio_labels = torch.tensor([layer[1:] for layer in audio_labels], dtype=torch.long) # (8, T-1)
input_ids = self.apply_scheduled_sampling(input_ids, audio_labels, text_labels)
spk_emb = torch.tensor(spk_emb_raw, dtype=torch.float32) if spk_emb_raw else torch.zeros(192)
return input_ids, text_labels, audio_labels, audio_inputs, audio_len, pixel_values, spk_emb
# 测试parquet数据读取
if __name__ == '__main__':
for path in ['sft_a2a.parquet']:
if not os.path.exists(path): continue
t = pa.Table.from_batches(pq.ParquetFile(path).iter_batches())
conversations = json.loads(t['conversations'][0].as_py())
answer_audios = t['answer_audios'][0].as_py() if 'answer_audios' in t.column_names else []
user_msg = conversations[0]
asst_msg = conversations[1] if len(conversations) > 1 else {}
print(f'{path}: {len(t)}条, 列{t.column_names}')
print(f' User: {user_msg["content"][:50]}...')
print(f' Asst: {asst_msg.get("content", "")[:50]}...')
if answer_audios:
print(f' answer_audios: {len(answer_audios)}轮, 首轮{len(answer_audios[0])}tokens')

244
eval_omni.py Normal file
View File

@ -0,0 +1,244 @@
import argparse
import os
import random
import time
import warnings
import torch
import soundfile as sf
from PIL import Image
from pydub import AudioSegment
from transformers import AutoTokenizer, AutoModelForCausalLM, MimiModel
from model.model_omni import MiniMindOmni, OmniConfig
from dataset.omni_dataset import OmniDataset
from trainer.trainer_utils import setup_seed, log_model_params
warnings.filterwarnings('ignore')
def init_model(args):
tokenizer = AutoTokenizer.from_pretrained(args.load_from)
if 'model' in args.load_from:
moe_suffix = '_moe' if args.use_moe else ''
ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
model = MiniMindOmni(
OmniConfig(
hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers,
use_moe=bool(args.use_moe)
),
audio_encoder_path="./model/SenseVoiceSmall",
vision_model_path="./model/siglip2-base-p32-256-ve"
)
model.load_state_dict(torch.load(ckp, map_location=args.device), strict=False)
else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
model.audio_encoder, model.audio_processor = MiniMindOmni.load_sensevoice("./model/SenseVoiceSmall")
model.vision_encoder, model.vision_processor = MiniMindOmni.load_vision("./model/siglip2-base-p32-256-ve")
log_model_params(model)
if model.audio_encoder is not None: model.audio_encoder.to(args.device)
if model.vision_encoder is not None: model.vision_encoder.to(args.device)
model.mimi_model = MimiModel.from_pretrained("./model/mimi").eval()
return model.half().eval().to(args.device), tokenizer
def eval_sample(model, tokenizer, args, idx, prompt, audio_inputs, output_name, pixel_values=None, history=None, audio_lens=None, ref_codes=None, spk_emb=None):
messages = (history or []) + [{"role": "user", "content": prompt}]
inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
x = torch.tensor(tokenizer(inputs_text).data['input_ids'], dtype=torch.long, device=args.device)[None, ...]
audio_frames = []
with torch.no_grad():
res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=args.max_new_tokens,
temperature=args.temperature, top_p=args.top_p, stream=True,
return_audio_codes=True, open_thinking=bool(args.open_thinking),
audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values,
ref_codes=ref_codes, spk_emb=spk_emb)
print('📒 [Thinker]: ', end='', flush=True)
history_idx = 0
for y, audio_frame in res_y:
if y is not None:
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
if answer and answer[-1] != '<EFBFBD>':
print(answer[history_idx:], end='', flush=True)
history_idx = len(answer)
if audio_frame:
audio_frames.append(audio_frame)
print()
if audio_frames:
print(f'🎹 [Talker]: {len(audio_frames)} frames', end=" ")
if args.decode_audio:
try:
codes = [f for f in audio_frames if f and len(f) == 8]
if not codes:
print('⚠️ 生成的Mimi codes为空,跳过保存。')
return
mimi_codes = torch.tensor(codes, dtype=torch.long).T.unsqueeze(0).to(args.device)
filtered = torch.where(mimi_codes >= 2049, torch.zeros_like(mimi_codes), mimi_codes)
audio = model.mimi_model.decode(filtered).audio_values
output_path = os.path.join(args.output_dir, output_name)
wav_path = output_path.rsplit('.', 1)[0] + '.wav'
sf.write(wav_path, audio.squeeze().float().cpu().numpy(), 24000)
AudioSegment.from_wav(wav_path).export(output_path, format='mp3', bitrate='64k')
os.remove(wav_path)
print(f'| Audio decoded to: {output_path}')
except Exception as e:
print(f'⚠️ 保存音频失败: {str(e)}')
else:
print("(decode_audio=off)\n")
def main():
parser = argparse.ArgumentParser(description="MiniMind-O Chat")
parser.add_argument('--load_from', default='model', type=str, help="模型加载路径(model=原生torch权重)")
parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
parser.add_argument('--weight', default='sft_omni', type=str, help="权重名称前缀")
parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构")
parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
parser.add_argument('--temperature', default=0.7, type=float, help="Thinker生成温度")
parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值")
parser.add_argument('--output_dir', default='./output_audio/', type=str, help="输出音频保存目录")
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
parser.add_argument('--audio_dir', default='./dataset/eval_omni/', type=str, help="测试音频目录")
parser.add_argument('--image_dir', default='./dataset/eval_omni/', type=str, help="测试图像目录")
parser.add_argument('--open_thinking', default=0, type=int, help="是否开启思考模式(0=否,1=是)(思考模式下禁用audio输出)")
parser.add_argument('--decode_audio', default=1, type=int, help="是否解码音频输出(0=否,1=是)")
parser.add_argument('--mode', default='0', type=str, help="评估模式:-1=all 0=text 1=multi 2=audio 3=clone 4=image 5=mix(逗号组合,如 2,5)")
parser.add_argument('--prompt_lang', default=0, type=int, choices=[0, 1, 2], help="问题语言:0=英文 1=中文 2=英文+中文")
args = parser.parse_args()
modes = set(args.mode.replace(',', '').replace('-1', '012345'))
os.makedirs(args.output_dir, exist_ok=True)
model, tokenizer = init_model(args)
setup_seed(int(time.time()) % 31415926)
if '0' in modes:
print('\n\n==================== text -> {text, audio} ====================')
test_prompts_en = [
"Tell me an interesting fact about space.", "How do I make a cup of coffee?", "What's the weather like today?",
"Will it rain tomorrow?", "Tell me a joke.", "Can you sing a song for me?", "Please introduce yourself."
]
test_prompts_zh = [
"告诉我一个关于太空的有趣事实。", "如何制作一杯咖啡?", "今天的天气怎么样?",
"明天会下雨吗?", "给我讲个笑话吧", "你能为我唱首歌吗?", "介绍一下你自己"
]
test_prompts = [test_prompts_en, test_prompts_zh, test_prompts_en + test_prompts_zh][args.prompt_lang]
for idx, prompt in enumerate(test_prompts):
print(f'\n📝 [text-{idx+1}]: {prompt}')
eval_sample(model, tokenizer, args, idx, prompt, None, f"text-{idx:02d}.mp3")
if '1' in modes:
print('\n\n==================== multi-turn -> {text, audio} ====================')
multi_turn_tests_zh = [
{
"history": [
{"role": "user", "content": "你好"},
{"role": "assistant", "content": "你好!有什么可以帮你的吗?"}
],
"prompt": "我想找点事做,你有什么建议吗?"
},
{
"history": [
{"role": "user", "content": "你好"},
{"role": "assistant", "content": "你好!有什么可以帮你的吗?"},
{"role": "user", "content": "我想找点事做,你有什么建议吗?"},
{"role": "assistant", "content": "可以听听音乐或者看看书,放松一下心情。"}
],
"prompt": "好的,那我去照做了,谢谢你"
}
]
multi_turn_tests_en = [
{
"history": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hello! How can I help you?"}
],
"prompt": "I want to find something to do. Do you have any suggestions?"
},
{
"history": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hello! How can I help you?"},
{"role": "user", "content": "I want to find something to do. Do you have any suggestions?"},
{"role": "assistant", "content": "You can listen to music or read a book to relax a little."}
],
"prompt": "Okay, I will try that. Thank you."
}
]
multi_turn_tests = [multi_turn_tests_en, multi_turn_tests_zh, multi_turn_tests_en + multi_turn_tests_zh][args.prompt_lang]
for idx, test in enumerate(multi_turn_tests):
print(f'\n💬 [multi-{idx+1}]')
for msg in test["history"]: print(f' {msg["role"]}: {msg["content"]}')
print(f' user: {test["prompt"]}')
eval_sample(model, tokenizer, args, idx, test["prompt"], None, f"multi-{idx:02d}.mp3", history=test["history"])
if '2' in modes:
print('\n\n==================== audio -> {text, audio} ====================')
audio_files_en = sorted([f for f in os.listdir(args.audio_dir) if f.startswith('audio-en-') and f.lower().endswith(('.mp3', '.wav'))])
audio_files_zh = sorted([f for f in os.listdir(args.audio_dir) if f.startswith('audio-zh-') and f.lower().endswith(('.mp3', '.wav'))])
audio_files = [audio_files_en, audio_files_zh, audio_files_en + audio_files_zh][args.prompt_lang]
for idx, audio_file in enumerate(audio_files):
print(f'\n🎤 [audio-{idx+1}]: {audio_file}')
mel, valid_len = OmniDataset.process_audio(os.path.join(args.audio_dir, audio_file), model.audio_processor)
audio_inputs = mel.unsqueeze(0).to(args.device)
audio_lens = torch.tensor([valid_len], device=args.device)
audio_token_len = valid_len or 1
prompt = model.config.audio_special_token * audio_token_len
eval_sample(model, tokenizer, args, idx, prompt, audio_inputs, f"audio-{idx:02d}-{os.path.splitext(audio_file)[0]}.mp3", audio_lens=audio_lens)
if '3' in modes:
print('\n\n==================== clone voice -> {text, audio} ====================')
clone_prompts_en = ["Hello, please introduce yourself.", "What's the weather like today?", "Tell me a joke."]
clone_prompts_zh = ["你好,请介绍一下你自己。", "今天天气怎么样?", "给我讲个笑话吧。"]
clone_prompts = [clone_prompts_en, clone_prompts_zh, clone_prompts_en + clone_prompts_zh][args.prompt_lang]
voices_pt = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'model', 'speaker', 'voices_unseen.pt')
voices = [('default', None, None)]
if os.path.exists(voices_pt):
voice_data = torch.load(voices_pt, map_location=args.device)
for speaker, v in sorted(voice_data.items()):
rc = v['ref_codes'].unsqueeze(0).to(args.device)
se = v['spk_emb'].half().unsqueeze(0).to(args.device) if 'spk_emb' in v else None
voices.append((speaker, rc, se))
for speaker, rc, se in voices:
info = f'ref_codes: {rc.shape[2]} frames, spk_emb: {"+" if se is not None else "-"}' if rc is not None else ('spk_emb only' if se is not None else 'default')
print(f'\n🎵 [clone: {speaker}] {info}')
for idx, prompt in enumerate(clone_prompts):
print(f' 📝 [text-{idx+1}]: {prompt}')
history = [{"role": "system", "content": "你是一个专业的语音助手,请用给定的音色风格来回答用户的问题。请尽量详细地回答,给出有价值的信息。"}]
eval_sample(model, tokenizer, args, idx, prompt, None, f"clone-{speaker}-{idx:02d}.mp3", ref_codes=rc, history=history, spk_emb=se)
if '4' in modes:
print('\n\n==================== image -> {text, audio} ====================')
image_files = sorted([f for f in os.listdir(args.image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
for idx, image_file in enumerate(image_files):
print(f'\n🖼️ [image-{idx+1}]: {image_file}')
image = Image.open(os.path.join(args.image_dir, image_file)).convert('RGB')
pixel_values = {k: v.to(args.device) for k, v in model.vision_processor(images=image, return_tensors="pt").items()}
prompts = [["Please describe this image."], ["请描述这张图片"], ["Please describe this image.", "请描述这张图片"]][args.prompt_lang]
for lang_idx, prompt_text in enumerate(prompts):
prompt = prompt_text + "\n\n" + model.config.image_special_token * model.config.image_token_len
eval_sample(model, tokenizer, args, idx, prompt, None, f"image-{idx:02d}-{lang_idx}-{os.path.splitext(image_file)[0]}.mp3", pixel_values=pixel_values)
if '5' in modes:
print('\n\n==================== text+audio+image -> {text, audio} ====================')
img_audio_files = sorted([f for f in os.listdir(args.audio_dir) if f.startswith('img-') and f.lower().endswith(('.mp3', '.wav'))])
image_files = sorted([f for f in os.listdir(args.image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
text_hints = [["Please answer me: "], ["请回答我:"], ["Please answer me: ", "请回答我:"]][args.prompt_lang]
for idx, image_file in enumerate(image_files):
audio_file = random.choice(img_audio_files)
image = Image.open(os.path.join(args.image_dir, image_file)).convert('RGB')
pixel_values = {k: v.to(args.device) for k, v in model.vision_processor(images=image, return_tensors="pt").items()}
for lang_idx, text_hint in enumerate(text_hints):
print(f'\n🌀 [mix-{idx+1}-{lang_idx}]: {text_hint} | {audio_file} | {image_file}')
mel, valid_len = OmniDataset.process_audio(os.path.join(args.audio_dir, audio_file), model.audio_processor)
audio_inputs = mel.unsqueeze(0).to(args.device)
audio_lens = torch.tensor([valid_len], device=args.device)
audio_token_len = valid_len or 1
prompt = text_hint + model.config.audio_special_token * audio_token_len + "\n\n" + model.config.image_special_token * model.config.image_token_len
eval_sample(model, tokenizer, args, idx, prompt, audio_inputs, f"mix-{idx:02d}-{lang_idx}-{os.path.splitext(image_file)[0]}.mp3", pixel_values=pixel_values, audio_lens=audio_lens)
if __name__ == "__main__":
main()

Binary file not shown.

After

Width:  |  Height:  |  Size: 131 KiB

BIN
images/architecture.jpg Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 279 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 228 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 30 KiB

BIN
images/logo.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 692 KiB

BIN
images/omni_io_flow.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 217 KiB

BIN
images/qual_a2a.jpg Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 121 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 90 KiB

BIN
images/sequence_format.jpg Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 102 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 104 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 218 KiB

0
model/__init__.py Normal file
View File

287
model/model_minimind.py Executable file
View File

@ -0,0 +1,287 @@
import math, torch, torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
# MiniMind Config
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
class MiniMindConfig(PretrainedConfig):
model_type = "minimind"
def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
super().__init__(**kwargs)
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.use_moe = use_moe
self.dropout = kwargs.get("dropout", 0.0)
self.vocab_size = kwargs.get("vocab_size", 6400)
self.bos_token_id = kwargs.get("bos_token_id", 1)
self.eos_token_id = kwargs.get("eos_token_id", 2)
self.flash_attn = kwargs.get("flash_attn", True)
self.num_attention_heads = kwargs.get("num_attention_heads", 8)
self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
self.hidden_act = kwargs.get("hidden_act", 'silu')
self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
self.rope_theta = kwargs.get("rope_theta", 1e6)
self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
self.rope_scaling = {
"beta_fast": 32,
"beta_slow": 1,
"factor": 16,
"original_max_position_embeddings": 2048,
"attention_factor": 1.0,
"type": "yarn"
} if self.inference_rope_scaling else None
### MoE specific configs (ignored if use_moe = False)
self.num_experts = kwargs.get("num_experts", 4)
self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
# MiniMind Model
# 🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏🌎🌍🌏
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-5):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
return (self.weight * self.norm(x.float())).type_as(x)
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
if rope_scaling is not None: # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp
orig_max, factor, beta_fast, beta_slow, attn_factor = (
rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
)
if end / orig_max > 1.0:
inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
freqs = freqs * (1 - ramp + ramp / factor)
t = torch.arange(end, device=freqs.device)
freqs = torch.outer(t, freqs).float()
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
return freqs_cos, freqs_sin
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
return q_embed, k_embed
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
bs, slen, num_key_value_heads, head_dim = x.shape
if n_rep == 1: return x
return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))
class Attention(nn.Module):
def __init__(self, config: MiniMindConfig):
super().__init__()
self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
self.n_local_heads = config.num_attention_heads
self.n_local_kv_heads = self.num_key_value_heads
self.n_rep = self.n_local_heads // self.n_local_kv_heads
self.head_dim = config.head_dim
self.is_causal = True
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)
self.dropout = config.dropout
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
bsz, seq_len, _ = x.shape
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xq, xk = self.q_norm(xq), self.k_norm(xk)
cos, sin = position_embeddings
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
if past_key_value is not None:
xk = torch.cat([past_key_value[0], xk], dim=1)
xv = torch.cat([past_key_value[1], xv], dim=1)
past_kv = (xk, xv) if use_cache else None
xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
else:
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output))
return output, past_kv
class FeedForward(nn.Module):
def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
super().__init__()
intermediate_size = intermediate_size or config.intermediate_size
self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
class MOEFeedForward(nn.Module):
def __init__(self, config: MiniMindConfig):
super().__init__()
self.config = config
self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)])
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x):
batch_size, seq_len, hidden_dim = x.shape
x_flat = x.view(-1, hidden_dim)
scores = F.softmax(self.gate(x_flat), dim=-1)
topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
y = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts):
mask = (topk_idx == i)
if mask.any():
token_idx = mask.any(dim=-1).nonzero().flatten()
weight = topk_weight[mask].view(-1, 1)
y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
elif self.training:
y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
if self.training and self.config.router_aux_loss_coef > 0:
load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
else:
self.aux_loss = scores.new_zeros(1).squeeze()
return y.view(batch_size, seq_len, hidden_dim)
class MiniMindBlock(nn.Module):
def __init__(self, layer_id: int, config: MiniMindConfig):
super().__init__()
self.self_attn = Attention(config)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
residual = hidden_states
hidden_states, present_key_value = self.self_attn(
self.input_layernorm(hidden_states), position_embeddings,
past_key_value, use_cache, attention_mask
)
hidden_states += residual
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
return hidden_states, present_key_value
class MiniMindModel(nn.Module):
def __init__(self, config: MiniMindConfig):
super().__init__()
self.config = config
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
self.dropout = nn.Dropout(config.dropout)
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
batch_size, seq_length = input_ids.shape
if hasattr(past_key_values, 'layers'): past_key_values = None
past_key_values = past_key_values or [None] * len(self.layers)
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
hidden_states = self.dropout(self.embed_tokens(input_ids))
# Recompute RoPE buffers lost during meta-device init (transformers>=5.x)
if self.freqs_cos[0, 0] == 0:
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
presents = []
for layer, past_key_value in zip(self.layers, past_key_values):
hidden_states, present = layer(
hidden_states,
position_embeddings,
past_key_value=past_key_value,
use_cache=use_cache,
attention_mask=attention_mask
)
presents.append(present)
hidden_states = self.norm(hidden_states)
aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
return hidden_states, presents, aux_loss
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
config_class = MiniMindConfig
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
def __init__(self, config: MiniMindConfig = None):
self.config = config or MiniMindConfig()
super().__init__(self.config)
self.model = MiniMindModel(self.config)
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
if self.config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight
self.post_init()
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
loss = None
if labels is not None:
x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
# https://github.com/jingyaogong/minimind/discussions/611
@torch.inference_mode()
def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
past_key_values = kwargs.pop("past_key_values", None)
finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
if streamer: streamer.put(input_ids.cpu())
for _ in range(max_new_tokens):
past_len = past_key_values[0][0].shape[1] if past_key_values else 0
outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
logits = outputs.logits[:, -1, :] / temperature
if repetition_penalty != 1.0:
for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty
if top_k > 0:
logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
if top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
input_ids = torch.cat([input_ids, next_token], dim=-1)
past_key_values = outputs.past_key_values if use_cache else None
if streamer: streamer.put(next_token.cpu())
if eos_token_id is not None:
finished |= next_token.squeeze(-1).eq(eos_token_id)
if finished.all(): break
if streamer: streamer.end()
if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
return input_ids

461
model/model_omni.py Normal file
View File

@ -0,0 +1,461 @@
import os, math, torch, soundfile as sf, librosa, warnings, numpy as np, onnxruntime as ort, logging, contextlib, io
from types import SimpleNamespace
from torch import nn
from torch.nn import functional as F
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
from transformers import SiglipImageProcessor, SiglipVisionModel, logging as hf_logging
from .model_minimind import *
class OmniConfig(MiniMindConfig):
model_type = "minimind-o"
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.num_talker_hidden_layers = kwargs.get("num_talker_hidden_layers", 4)
self.talker_hidden_size = kwargs.get("talker_hidden_size", 768)
self.audio_ids = kwargs.get("audio_ids", [16]) # "<|audio_pad|>" token id
self.audio_special_token = kwargs.get("audio_special_token", "<|audio_pad|>")
self.audio_hidden_size = kwargs.get("audio_hidden_size", 512)
self.audio_vocab_size = kwargs.get("audio_vocab_size", 2112)
self.audio_pad_token = kwargs.get("audio_pad_token", 2049)
self.audio_stop_token = kwargs.get("audio_stop_token", 2050)
self.audio_spk_token = kwargs.get("audio_spk_token", 2051)
self.spk_emb_size = kwargs.get("spk_emb_size", 192)
self.think_end_ids = kwargs.get("think_end_ids", [26, 234, 234]) # </think>\n\n
self.image_ids = kwargs.get("image_ids", [12]) # "<|image_pad|>" token id
self.image_special_token = kwargs.get("image_special_token", "<|image_pad|>")
self.image_hidden_size = kwargs.get("image_hidden_size", 768)
self.image_token_len = kwargs.get("image_token_len", 64)
self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1)
class MMAudioProjector(nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
self.mlp = nn.Sequential(
nn.LayerNorm(in_dim),
nn.Linear(in_dim, out_dim),
nn.GELU(),
nn.Linear(out_dim, out_dim),
)
def forward(self, x):
return self.mlp(x)
class MMVisionProjector(nn.Module):
def __init__(self, in_dim, out_dim, source_tokens=64, target_tokens=64):
super().__init__()
self.mlp = nn.Sequential(
nn.LayerNorm(in_dim),
nn.Linear(in_dim, out_dim),
nn.GELU(),
nn.Linear(out_dim, out_dim),
)
def forward(self, x):
return self.mlp(x)
class TalkerHead(nn.Module):
def __init__(self, in_features, out_features, num_layers=8, rank=256):
super().__init__()
self.num_layers = num_layers
self.base = nn.Linear(in_features, out_features, bias=False)
self.adapters = nn.ModuleList([nn.Sequential(nn.Linear(in_features, rank, bias=False), nn.GELU(), nn.Linear(rank, out_features, bias=False)) for _ in range(num_layers)])
def forward(self, x):
base_out = self.base(x)
return [base_out + adapter(x) for adapter in self.adapters]
class TalkerEmbedding(nn.Module):
def __init__(self, num_embeddings, embedding_dim, num_layers=8, rank=256):
super().__init__()
self.num_layers = num_layers
self.base = nn.Embedding(num_embeddings, embedding_dim)
self.adapters = nn.ModuleList([nn.Sequential(nn.Embedding(num_embeddings, rank), nn.GELU(), nn.Linear(rank, embedding_dim, bias=False)) for _ in range(num_layers)])
def forward(self, x):
base_out = self.base(x)
return sum(base_out[:, i, :] + self.adapters[i](x[:, i, :]) for i in range(len(self.adapters))) / self.num_layers
class SenseVoiceAudioProcessor:
def __init__(self, frontend): self.frontend = frontend
def __call__(self, wav, sampling_rate=16000, return_tensors="pt", return_attention_mask=True, **kwargs):
if isinstance(wav, np.ndarray): wav = torch.from_numpy(wav).float()
if wav.dim() == 1: wav = wav.unsqueeze(0)
with torch.no_grad():
fbank, flen = self.frontend(wav, torch.tensor([wav.size(1)]))
return SimpleNamespace(input_features=fbank, attention_mask=(torch.arange(fbank.size(1)) < flen[0]).long().unsqueeze(0))
class TalkerModule(nn.Module):
def __init__(self, config):
super().__init__()
self.talker_config = MiniMindConfig(hidden_size=config.talker_hidden_size, use_moe=config.use_moe)
self.layers = nn.ModuleList([MiniMindBlock(l, self.talker_config) for l in range(config.num_talker_hidden_layers)])
self.norm = RMSNorm(config.talker_hidden_size, eps=config.rms_norm_eps)
self.lm_head = TalkerHead(config.talker_hidden_size, config.audio_vocab_size)
self.embed_tokens = TalkerEmbedding(config.audio_vocab_size, config.talker_hidden_size)
self.codec_proj = nn.Sequential(nn.Linear(config.talker_hidden_size, config.talker_hidden_size), nn.GELU(), nn.Linear(config.talker_hidden_size, config.talker_hidden_size), RMSNorm(config.talker_hidden_size, eps=config.rms_norm_eps))
self.embed_proj = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(), nn.Linear(config.hidden_size, config.talker_hidden_size), RMSNorm(config.talker_hidden_size, eps=config.rms_norm_eps))
self.text_scale, self.audio_scale = nn.Parameter(torch.tensor(3.0)), nn.Parameter(torch.tensor(1.0))
self.spk_proj = nn.Linear(config.spk_emb_size, config.talker_hidden_size, bias=False)
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.talker_config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
class MiniMindOmni(MiniMindForCausalLM):
config_class = OmniConfig
def __init__(self, config: OmniConfig = None, audio_encoder_path="./model/SenseVoiceSmall", vision_model_path="./model/siglip2-base-p32-256-ve"):
config = config or OmniConfig()
super().__init__(config)
object.__setattr__(self, 'thinker', self.model) # alias: self.thinker == self.model
object.__setattr__(self.model, 'lm_head', self.lm_head) # alias: self.thinker.lm_head == self.lm_head
self.talker = TalkerModule(config)
self.audio_proj = MMAudioProjector(config.audio_hidden_size, config.hidden_size)
self.vision_proj = MMVisionProjector(config.image_hidden_size, config.hidden_size, target_tokens=config.image_token_len)
self.audio_pad_token, self.audio_stop_token, self.audio_spk_token = config.audio_pad_token, config.audio_stop_token, config.audio_spk_token
audio_encoder, audio_processor = self.load_sensevoice(audio_encoder_path)
object.__setattr__(self, 'audio_encoder', audio_encoder)
object.__setattr__(self, 'audio_processor', audio_processor)
vision_encoder, vision_processor = self.load_vision(vision_model_path)
object.__setattr__(self, 'vision_encoder', vision_encoder)
object.__setattr__(self, 'vision_processor', vision_processor)
@staticmethod
def load_sensevoice(path):
if not os.path.exists(path):
warnings.warn(f"[MiniMindOmni] SenseVoice path not found: {path}")
return None, None
logging.getLogger().setLevel(logging.ERROR)
hf_logging.set_verbosity_error()
with contextlib.redirect_stdout(io.StringIO()):
from funasr import AutoModel
m = AutoModel(model=path, trust_remote_code=True, disable_update=True, device="cpu")
encoder, frontend = m.model.encoder, m.kwargs["frontend"]
for p in encoder.parameters(): p.requires_grad = False
return encoder.eval().float(), SenseVoiceAudioProcessor(frontend.eval())
@torch.compiler.disable
def encode_audio_inputs(self, audio_inputs, audio_lens=None):
if (audio_inputs is None) or (self.audio_encoder is None) or (not audio_inputs.any()): return None
batch_mask = audio_inputs.flatten(1).any(1)
enc_dtype = next(self.audio_encoder.parameters()).dtype
valid_fbank = audio_inputs[batch_mask].to(dtype=enc_dtype)
if audio_lens is not None:
valid_lens = audio_lens[batch_mask].to(valid_fbank.device)
else:
valid_lens = torch.tensor([valid_fbank.size(1)] * valid_fbank.size(0), device=valid_fbank.device)
with torch.no_grad():
emb, _ = self.audio_encoder(valid_fbank, valid_lens)
proj_dtype = next(self.audio_proj.parameters()).dtype
emb_list = [self.audio_proj(emb[i, :max(1, min(valid_lens[i].item(), emb.size(1)))].unsqueeze(0).to(proj_dtype)).squeeze(0) for i in range(emb.size(0))]
if batch_mask.all(): return emb_list
out = [None] * audio_inputs.size(0)
j = 0
for i in range(audio_inputs.size(0)):
if batch_mask[i]:
out[i] = emb_list[j]
j += 1
return out
@torch.compiler.disable
def inject_audio_features(self, tokens, h, audio_feats, seqlen):
if audio_feats is None or not self.config.audio_ids:
return h
marker = self.config.audio_ids[0]
out = []
for b in range(h.size(0)):
hb, seq, i = h[b], tokens[b].tolist(), 0
af = audio_feats[b] if audio_feats[b] is not None else None
while i < len(seq):
if seq[i] == marker:
start = i
while i < len(seq) and seq[i] == marker:
i += 1
if af is not None:
inject_len = min(af.size(0), i - start)
hb = torch.cat((hb[:start], af[:inject_len], hb[start + inject_len:]), dim=0)
af = None
else:
i += 1
out.append(hb)
return torch.stack(out)
@staticmethod
def load_vision(path):
if path is None or not os.path.exists(path):
warnings.warn(f"[MiniMindOmni] Vision model path not found: {path}. vision_encoder will be None!")
return None, None
hf_logging.set_verbosity_error()
try:
model = SiglipVisionModel.from_pretrained(path)
except (RuntimeError, ValueError):
return None, None
processor = SiglipImageProcessor.from_pretrained(path)
for p in model.parameters():
p.requires_grad = False
return model.eval(), processor
@torch.compiler.disable
def get_image_embeddings(self, image_inputs):
if hasattr(image_inputs, 'keys'):
image_inputs = {k: v.squeeze(1) if v.ndim > 2 and v.shape[1] == 1 else v for k, v in image_inputs.items()}
pixel_attention_mask = image_inputs.get('pixel_attention_mask')
if pixel_attention_mask is not None and not pixel_attention_mask.any():
pv = image_inputs['pixel_values']
return pv.new_zeros(pv.size(0), pv.size(1), self.config.image_hidden_size)
with torch.no_grad():
outputs = self.vision_encoder(**image_inputs)
return outputs.last_hidden_state
@torch.compiler.disable
def encode_image_inputs(self, pixel_values):
if pixel_values is None or self.vision_encoder is None: return None
mask = pixel_values.flatten(1).any(1)
if not mask.any(): return pixel_values.new_zeros(pixel_values.size(0), self.config.image_token_len, self.config.hidden_size)
with torch.no_grad(): emb = self.vision_encoder(pixel_values=pixel_values[mask]).last_hidden_state
if emb.dim() == 2: emb = emb.unsqueeze(0)
emb = self.vision_proj(emb)
if mask.all(): return emb
idx = mask.nonzero().view(-1, 1, 1).expand_as(emb)
return emb.new_zeros(pixel_values.size(0), *emb.shape[1:]).scatter(0, idx, emb)
@torch.compiler.disable
def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
if vision_tensors is None or not self.config.image_ids:
return h
marker, vf = self.config.image_ids[0], vision_tensors
if vf.dim() == 3:
vf = vf.unsqueeze(1)
out = []
for b in range(h.size(0)):
hb, seq, k, i = h[b], tokens[b].tolist(), 0, 0
while i < len(seq):
if seq[i] == marker:
start = i
while i < len(seq) and seq[i] == marker:
i += 1
if k < vf.size(1):
hb = torch.cat((hb[:start], vf[b][k][:i - start], hb[i:]), dim=0)[:seqlen]
k += 1
else:
i += 1
out.append(hb)
return torch.stack(out)
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, audio_inputs=None, audio_lens=None, pixel_values=None, **args):
if len(input_ids.shape) == 2:
batch_size, seq_length = input_ids.shape
text_ids = input_ids
audio_ids = torch.full((batch_size, 8, seq_length), self.audio_pad_token, dtype=torch.long, device=input_ids.device)
else:
batch_size, _, seq_length = input_ids.shape
text_ids, audio_ids = input_ids[:, 8, :], input_ids[:, :8, :]
if hasattr(past_key_values, 'layers'): past_key_values = None
n_thinker, n_talker = len(self.thinker.layers), len(self.talker.layers)
past_key_values = past_key_values or ([None] * (n_thinker + n_talker))
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
# Recompute RoPE buffers lost during meta-device init (transformers>=5.x)
if self.thinker.freqs_cos[0, 0] == 0:
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
self.thinker.freqs_cos, self.thinker.freqs_sin = freqs_cos.to(input_ids.device), freqs_sin.to(input_ids.device)
if self.talker.freqs_cos[0, 0] == 0:
freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.talker.talker_config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
self.talker.freqs_cos, self.talker.freqs_sin = freqs_cos.to(input_ids.device), freqs_sin.to(input_ids.device)
presents = []
# ======= Thinker: text-only input, output text logits =======
hidden_states = self.thinker.dropout(self.thinker.embed_tokens(text_ids))
position_embeddings = (self.thinker.freqs_cos[start_pos:start_pos + seq_length], self.thinker.freqs_sin[start_pos:start_pos + seq_length])
if audio_inputs is not None and start_pos == 0:
audio_features = self.encode_audio_inputs(audio_inputs, audio_lens)
hidden_states = self.inject_audio_features(text_ids, hidden_states, audio_features, seq_length)
if pixel_values is not None and start_pos == 0:
if hasattr(pixel_values, 'keys'):
img_emb = self.get_image_embeddings(pixel_values).to(hidden_states.dtype)
vision_tensors = self.vision_proj(img_emb)
else:
if len(pixel_values.shape) == 6:
pixel_values = pixel_values.squeeze(2)
if len(pixel_values.shape) == 4:
pixel_values = pixel_values.unsqueeze(1)
bs, num, c, im_h, im_w = pixel_values.shape
stack_dim = 1 if bs > 1 else 0
vision_tensors = torch.stack([
self.encode_image_inputs(pixel_values[:, i, :, :, :])
for i in range(num)
], dim=stack_dim)
hidden_states = self.count_vision_proj(tokens=text_ids, h=hidden_states, vision_tensors=vision_tensors, seqlen=seq_length)
bridge_states = hidden_states
for i, (layer, past_key_value) in enumerate(zip(self.thinker.layers, past_key_values[:n_thinker])):
hidden_states, present = layer(hidden_states, position_embeddings, past_key_value=past_key_value, use_cache=use_cache, attention_mask=attention_mask)
presents.append(present)
if i == self.config.bridge_layer: bridge_states = hidden_states
h_thinker = self.thinker.norm(hidden_states)
# ======= Talker: thinker hidden + audio codes, output audio logits =======
talker_emb = self.talker.embed_tokens(audio_ids)
spk_emb = args.get('spk_emb', None)
if spk_emb is not None:
spk_mask = (audio_ids[:, 0, :] == self.audio_spk_token).unsqueeze(-1)
talker_emb = torch.where(spk_mask, self.talker.spk_proj(spk_emb).unsqueeze(1), talker_emb)
hidden_states = self.talker.embed_proj(bridge_states) * self.talker.text_scale + self.talker.codec_proj(talker_emb) * self.talker.audio_scale
talker_pos_emb = (self.talker.freqs_cos[start_pos:start_pos + seq_length], self.talker.freqs_sin[start_pos:start_pos + seq_length])
for layer, past_key_value in zip(self.talker.layers, past_key_values[n_thinker:]):
hidden_states, present = layer(hidden_states, talker_pos_emb, past_key_value=past_key_value, use_cache=use_cache, attention_mask=attention_mask)
presents.append(present)
h_talker = self.talker.norm(hidden_states)
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
aux_loss = sum(l.mlp.aux_loss for l in list(self.thinker.layers) + list(self.talker.layers) if isinstance(l.mlp, MOEFeedForward))
aux_loss += sum(p.sum() for p in self.audio_proj.parameters()) * 0 + sum(p.sum() for p in self.vision_proj.parameters()) * 0 + sum(p.sum() for p in self.talker.lm_head.adapters.parameters()) * 0 + sum(p.sum() for p in self.talker.spk_proj.parameters()) * 0 # dummy gradient
text_logits = self.thinker.lm_head(h_thinker[:, slice_indices, :])
audio_logits = self.talker.lm_head(h_talker[:, slice_indices, :])
out = MoeCausalLMOutputWithPast(aux_loss=aux_loss, logits=text_logits, past_key_values=presents)
out.audio_logits = audio_logits
return out
@torch.inference_mode()
def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
stream=False, rp=1., use_cache=True, return_audio_codes=False, **args):
if stream:
return self.stream_generate(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, return_audio_codes, **args)
tokens = list(self.stream_generate(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, return_audio_codes, **args))
return tokens[-1] if tokens else input_ids
def stream_generate(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, return_audio_codes=False, **args):
start_pos, past_kvs, text_finished, first_finished = input_ids.shape[1], None, False, True
audio_codes = [[] for _ in range(8)]
audio_stop_pos = [None] * 8
audio_buffer = torch.full((1, 8, start_pos), self.audio_pad_token, dtype=torch.long, device=input_ids.device)
spk_emb = args.get('spk_emb', None)
ref_codes = args.get('ref_codes', None)
ref_len = ref_codes.shape[2] if ref_codes is not None else 0
spk_reserve = 1 if spk_emb is not None else 0
fill_end = start_pos
fill_start = max(spk_reserve, start_pos - ref_len)
if ref_codes is not None and fill_start < fill_end:
audio_buffer[:, :, fill_start:fill_end] = ref_codes[:, :, -(fill_end - fill_start):]
if spk_emb is not None and fill_start > 0:
audio_buffer[:, :, fill_start - 1] = self.audio_spk_token
think_end_step, generated_tokens = None, ([] if args.get('open_thinking', False) else None)
while input_ids.shape[1] < start_pos + max_new_tokens:
if past_kvs is None or not use_cache:
out = self.forward(torch.cat((audio_buffer, input_ids.unsqueeze(1)), dim=1), past_key_values=past_kvs, use_cache=use_cache, **args)
else:
out = self.forward(torch.cat((audio_buffer[:, :, -1:], input_ids[:, -1:].unsqueeze(1)), dim=1), past_key_values=past_kvs, use_cache=use_cache, **args)
past_kvs = out.past_key_values
logits = out.logits[0, -1, :].clone() / (temperature + 1e-9)
logits[list(set(input_ids[0].tolist()))] /= rp
if top_p and top_p < 1.0:
sorted_l, sorted_i = torch.sort(logits, descending=True)
mask = torch.cumsum(F.softmax(sorted_l, dim=-1), dim=-1) > top_p
mask[1:], mask[0] = mask[:-1].clone(), False
logits[sorted_i[mask]] = -float('Inf')
text_token = torch.multinomial(F.softmax(logits, dim=-1), 1).item()
if text_finished:
text_token = args.get('enter_token_id', 201) if first_finished else args.get('pad_token_id', 0)
first_finished = False
step = input_ids.shape[1] - start_pos # 已生成token数(0=首次,此时模型处理prompt末尾token)
audio_step = step - 1 # 延迟1步:输出第1个text时无audio,输出第2个text时layer0开始
if generated_tokens is not None:
generated_tokens.append(text_token)
if not think_end_step and generated_tokens[-len(self.config.think_end_ids):] == list(self.config.think_end_ids): think_end_step = step + 2
audio_step = (step - think_end_step) if think_end_step else -1
for i, al in enumerate(out.audio_logits):
if audio_step < i:
audio_codes[i].append(self.audio_pad_token)
else:
logits_i = al[0, -1, :].clone() / 0.2
for prev_code in audio_codes[i][-3:]: logits_i[prev_code] /= 1.05
top_val, top_idx = logits_i.topk(50)
code = top_idx[torch.multinomial(F.softmax(top_val, dim=-1), 1)].item()
audio_codes[i].append(code)
if audio_stop_pos[i] is None and code >= 2048: audio_stop_pos[i] = len(audio_codes[i]) - 1
if text_finished and audio_codes[7][-1] == self.audio_stop_token: break
input_ids = torch.cat((input_ids, torch.tensor([[text_token]], device=input_ids.device)), dim=1)
audio_buffer = torch.cat((audio_buffer, torch.full((1, 8, 1), self.audio_pad_token, dtype=torch.long, device=input_ids.device)), dim=2)
for i in range(min(audio_step + 1, 8)): audio_buffer[0, i, -1] = audio_codes[i][-1]
audio_frame = None
if return_audio_codes and audio_step >= 7:
frame = [audio_codes[i][step - 7 + i] for i in range(8)]
active_layers = sum(1 for i in range(8) if audio_stop_pos[i] is None or step - 7 + i < audio_stop_pos[i])
if active_layers >= 8: audio_frame = frame
if not text_finished:
yield input_ids[:, start_pos:], audio_frame
if text_token == eos_token_id: text_finished = True
else:
yield None, audio_frame
# ==== Realtime VAD (与模型本体零耦合,纯工程层) ====
class SileroVAD:
def __init__(self, path):
opts = ort.SessionOptions()
opts.inter_op_num_threads = opts.intra_op_num_threads = 1
opts.log_severity_level = 4
self.session = ort.InferenceSession(path, providers=["CPUExecutionProvider"], sess_options=opts)
self.h, self.c = np.zeros((2, 1, 64), dtype=np.float32), np.zeros((2, 1, 64), dtype=np.float32)
def reset(self):
self.h[:], self.c[:] = 0, 0
def __call__(self, chunk, sr=16000):
out, self.h, self.c = self.session.run(None, {"input": chunk.reshape(1, -1).astype(np.float32), "h": self.h, "c": self.c, "sr": np.array(sr, dtype="int64")})
return float(out[0][0])
class RealtimeSession:
def __init__(self, vad_path, sr=16000, threshold=0.8, min_speech_ms=128, min_silence_ms=800):
self.vad, self.sr, self.threshold = SileroVAD(vad_path), sr, threshold
self.min_speech, self.min_silence = int(sr * min_speech_ms / 1000), int(sr * min_silence_ms / 1000)
self.reset()
def reset(self):
self.vad.reset()
self.buffer, self.ring, self.speaking, self.generating, self.interrupt = [], [], False, False, False
self.speech_samples = self.silence_samples = self.tail_silence = 0
def push_chunk(self, chunk, W=1024):
for i in range(0, max(len(chunk), 1), W):
w = chunk[i:i + W]
if len(w) < W:
w = np.pad(w, (0, W - len(w)))
prob = self.vad(w, self.sr)
if prob > self.threshold:
self.silence_samples = self.tail_silence = 0
self.speech_samples += len(w)
self.buffer.append(w)
if self.speech_samples >= self.min_speech and not self.speaking:
self.speaking = True
self.buffer = self.ring + self.buffer
self.ring = []
if self.generating and self.speaking:
self.interrupt = True
return 'interrupt'
elif self.speaking:
self.silence_samples += len(w)
self.tail_silence += 1
self.buffer.append(w)
if self.silence_samples >= self.min_silence:
if self.tail_silence > 1:
del self.buffer[-(self.tail_silence - 1):]
self.speaking, self.speech_samples, self.silence_samples, self.tail_silence = False, 0, 0, 0
return 'speech_end'
else:
if self.speech_samples > 0:
self.buffer.clear()
self.speech_samples = 0
self.ring = [w]
return 'listening'
def get_audio(self):
audio = np.concatenate(self.buffer) if self.buffer else np.array([], dtype=np.float32)
self.buffer.clear()
return audio

Binary file not shown.

BIN
model/speaker/voices.pt Normal file

Binary file not shown.

Binary file not shown.

31191
model/tokenizer.json Normal file

File diff suppressed because it is too large Load Diff

335
model/tokenizer_config.json Normal file
View File

@ -0,0 +1,335 @@
{
"add_bos_token": false,
"add_eos_token": false,
"add_prefix_space": false,
"added_tokens_decoder": {
"0": {
"content": "<|endoftext|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"1": {
"content": "<|im_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"2": {
"content": "<|im_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"3": {
"content": "<|object_ref_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"4": {
"content": "<|object_ref_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"5": {
"content": "<|box_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"6": {
"content": "<|box_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"7": {
"content": "<|quad_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"8": {
"content": "<|quad_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"9": {
"content": "<|vision_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"10": {
"content": "<|vision_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"11": {
"content": "<|vision_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"12": {
"content": "<|image_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"13": {
"content": "<|video_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"14": {
"content": "<|audio_start|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"15": {
"content": "<|audio_end|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"16": {
"content": "<|audio_pad|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"17": {
"content": "<tts_pad>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"18": {
"content": "<tts_text_bos>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"19": {
"content": "<tts_text_eod>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"20": {
"content": "<tts_text_bos_single>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": true
},
"21": {
"content": "<tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"22": {
"content": "</tool_call>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"23": {
"content": "<tool_response>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"24": {
"content": "</tool_response>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"25": {
"content": "<think>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"26": {
"content": "</think>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"27": {
"content": "<|buffer1|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"28": {
"content": "<|buffer2|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"29": {
"content": "<|buffer3|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"30": {
"content": "<|buffer4|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"31": {
"content": "<|buffer5|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"32": {
"content": "<|buffer6|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"33": {
"content": "<|buffer7|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"34": {
"content": "<|buffer8|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
},
"35": {
"content": "<|buffer9|>",
"lstrip": false,
"normalized": false,
"rstrip": false,
"single_word": false,
"special": false
}
},
"additional_special_tokens": [
"<|im_start|>",
"<|im_end|>",
"<|object_ref_start|>",
"<|object_ref_end|>",
"<|box_start|>",
"<|box_end|>",
"<|quad_start|>",
"<|quad_end|>",
"<|vision_start|>",
"<|vision_end|>",
"<|vision_pad|>",
"<|image_pad|>",
"<|video_pad|>",
"<|audio_start|>",
"<|audio_end|>",
"<|audio_pad|>",
"<tts_pad>",
"<tts_text_bos>",
"<tts_text_eod>",
"<tts_text_bos_single>"
],
"bos_token": "<|im_start|>",
"clean_up_tokenization_spaces": false,
"eos_token": "<|im_end|>",
"legacy": true,
"model_max_length": 131072,
"pad_token": "<|endoftext|>",
"sp_model_kwargs": {},
"spaces_between_special_tokens": false,
"unk_token": "<|endoftext|>",
"image_token": "<|image_pad|>",
"audio_token": "<|audio_pad|>",
"video_token": "<|video_pad|>",
"vision_bos_token": "<|vision_start|>",
"vision_eos_token": "<|vision_end|>",
"audio_bos_token": "<|audio_start|>",
"audio_eos_token": "<|audio_end|>",
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {{- messages[0].content + '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n {%- if message.content is string %}\n {%- set content = message.content %}\n {%- else %}\n {%- set content = '' %}\n {%- endif %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is string %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '</think>' in content %}\n {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- if true %}\n {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if open_thinking is defined and open_thinking is true %}\n {{- '<think>\\n' }}\n {%- else %}\n {{- '<think>\\n\\n</think>\\n\\n' }}\n {%- endif %}\n{%- endif %}",
"tokenizer_class": "PreTrainedTokenizerFast"
}

BIN
model/vad/silero_vad.onnx Normal file

Binary file not shown.

13
requirements.txt Normal file
View File

@ -0,0 +1,13 @@
soundfile==0.13.1
librosa==0.11.0
pyarrow==23.0.1
snac==1.2.1
transformers==4.57.6
speechbrain==1.1.0
onnxruntime==1.23.2
funasr==1.3.1
flask-sock==0.7.0
swanlab==0.7.11
# torch==2.6
# torchaudio==2.6
# torchvision==0.21.0

51
scripts/convert_omni.py Normal file
View File

@ -0,0 +1,51 @@
import os
import sys
__package__ = "scripts"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import json
import torch
import transformers
import warnings
from transformers import AutoTokenizer, AutoModelForCausalLM
from model.model_omni import MiniMindOmni, OmniConfig
warnings.filterwarnings('ignore', category=UserWarning)
def convert_torch2transformers(torch_path, transformers_path, lm_config, dtype=torch.bfloat16):
OmniConfig.register_for_auto_class()
MiniMindOmni.register_for_auto_class("AutoModelForCausalLM")
model = MiniMindOmni(lm_config, audio_encoder_path="../model/SenseVoiceSmall", vision_model_path="../model/siglip2-base-p32-256-ve")
state_dict = torch.load(torch_path, map_location='cpu')
model.load_state_dict(state_dict, strict=False)
model = model.to(dtype)
params = sum(p.numel() for p in model.parameters()) / 1e6
print(f'模型参数: {params:.2f}M')
del model.audio_encoder
del model.vision_encoder
model.save_pretrained(transformers_path, safe_serialization=False)
tokenizer = AutoTokenizer.from_pretrained('../model/')
tokenizer.save_pretrained(transformers_path)
config_path = os.path.join(transformers_path, "config.json")
config = json.load(open(config_path, 'r', encoding='utf-8'))
config['tie_word_embeddings'] = True
if int(transformers.__version__.split('.')[0]) >= 5:
tokenizer_config_path = os.path.join(transformers_path, "tokenizer_config.json")
json.dump({**json.load(open(tokenizer_config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(tokenizer_config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
config['rope_theta'] = lm_config.rope_theta; config['rope_scaling'] = None; config.pop('rope_parameters', None)
json.dump(config, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
print(f"已保存为 Transformers 格式: {transformers_path}")
def convert_transformers2torch(transformers_path, torch_path):
model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True)
torch.save(model.state_dict(), torch_path)
print(f"已保存为 PyTorch 格式: {torch_path}")
if __name__ == '__main__':
lm_config = OmniConfig(hidden_size=768, num_hidden_layers=8, use_moe=False)
torch_path = f"../out/sft_omni_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
transformers_path = '../minimind-3o'
convert_torch2transformers(torch_path, transformers_path, lm_config)

294
scripts/web_demo_omni.py Normal file
View File

@ -0,0 +1,294 @@
import os
import sys
__package__ = "scripts"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import argparse
import base64
import io
import tempfile
import warnings
import logging
import contextlib
import librosa
import soundfile as sf
import numpy as np
import torch
import gradio as gr
from queue import Queue
from threading import Thread, Lock
from PIL import Image
from pydub import AudioSegment
from transformers import AutoTokenizer, AutoModelForCausalLM, MimiModel, TextStreamer
from model.model_omni import MiniMindOmni, OmniConfig
from dataset.omni_dataset import OmniDataset
from trainer.trainer_utils import setup_seed, log_model_params
logging.getLogger().setLevel(logging.ERROR)
with contextlib.redirect_stdout(io.StringIO()):
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
warnings.filterwarnings('ignore')
model_lock = Lock()
model, tokenizer, device, mimi_model, asr_model = None, None, None, None, None
voices_data = {}
builtin_voices, clone_voices = set(), set()
def frames_to_mimi(frames):
codes = [f for f in frames if f and len(f) == 8]
if not codes:
return None
return torch.tensor(codes, dtype=torch.long).T.unsqueeze(0)
def decode_mimi(mimi_codes):
if mimi_codes is None or mimi_codes.numel() == 0:
return None
filtered = torch.where(mimi_codes >= 2049, torch.zeros_like(mimi_codes), mimi_codes).to(next(mimi_model.parameters()).device)
with torch.no_grad():
audio = mimi_model.decode(filtered).audio_values
return audio.squeeze().float().cpu().numpy()
def scan_hf_models(base_dir):
models = {}
base_dir = os.path.abspath(base_dir)
for d in sorted(os.listdir(base_dir), reverse=True):
full_path = os.path.join(base_dir, d)
if not os.path.isdir(full_path) or d.startswith('.') or d.startswith('_'):
continue
files = set(os.listdir(full_path))
has_model = bool(files & {'pytorch_model.bin', 'model.safetensors', 'pytorch_model.bin.index.json', 'model.safetensors.index.json'})
if has_model:
models[d] = full_path
return models
def load_hf_model(model_path):
global model, tokenizer
with model_lock:
[sys.modules.pop(k) for k in list(sys.modules) if 'transformers_modules' in k]
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
vision_encoder, vision_processor = MiniMindOmni.load_vision(args.vision_model)
audio_encoder, audio_processor = MiniMindOmni.load_sensevoice(args.audio_encoder)
object.__setattr__(model, 'vision_encoder', vision_encoder)
object.__setattr__(model, 'vision_processor', vision_processor)
object.__setattr__(model, 'audio_encoder', audio_encoder)
object.__setattr__(model, 'audio_processor', audio_processor)
model = model.half().eval().to(device)
if model.audio_encoder:
model.audio_encoder.to(device)
if model.vision_encoder:
model.vision_encoder.to(device)
name = os.path.basename(model_path)
params = sum(p.numel() for p in model.parameters()) / 1e6
print(f'已加载 {name},参数量:{params:.2f}M')
return f"已加载: {name} ({params:.2f}M)"
def load_voices():
global voices_data, builtin_voices, clone_voices
voices_data = {}
builtin_voices, clone_voices = set(), set()
for name, is_builtin in [('voices.pt', True), ('voices_unseen.pt', False)]:
path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'speaker', name)
if os.path.exists(path):
data = torch.load(path, map_location=device)
for speaker, v in data.items():
if speaker not in voices_data:
voices_data[speaker] = v
(builtin_voices if is_builtin else clone_voices).add(speaker)
def chat_stream(prompt, audio_input=None, image_input=None, voice_name="default", history=None, temperature=0.85, max_tokens=512):
audio_inputs, audio_lens, pixel_values, ref_codes, spk_emb = None, None, None, None, None
asr_result = [None]
asr_thread = None
if audio_input is not None:
sr, samples = audio_input
if len(samples.shape) > 1:
samples = samples.mean(axis=1)
if samples.dtype != np.float32:
samples = samples.astype(np.float32) / max(np.abs(samples).max(), 1)
if sr != 16000:
samples = librosa.resample(samples.astype(float), orig_sr=sr, target_sr=16000).astype(np.float32)
inputs = model.audio_processor(samples, sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
mel = inputs.input_features.squeeze(0)
valid_len = inputs.attention_mask.sum().item()
audio_inputs = mel.unsqueeze(0).to(device)
audio_lens = torch.tensor([valid_len], device=device)
audio_token_len = valid_len or 1
prompt = model.config.audio_special_token * audio_token_len
samples_for_asr = samples.copy()
def _do_asr():
r = asr_model.generate(input=samples_for_asr, cache={}, language='auto', use_itn=True)
asr_result[0] = rich_transcription_postprocess(r[0]['text']).strip() if r else ''
print(f'[ASR] {asr_result[0]}')
asr_thread = Thread(target=_do_asr)
asr_thread.start()
if image_input is not None:
image = Image.open(image_input).convert('RGB') if isinstance(image_input, str) else image_input.convert('RGB')
pixel_values = {k: v.to(device) for k, v in model.vision_processor(images=image, return_tensors="pt").items()}
prompt = (prompt + "\n\n" if prompt else "") + model.config.image_special_token * model.config.image_token_len
if voice_name != "default" and voice_name in voices_data:
v = voices_data[voice_name]
ref_codes = v['ref_codes'].unsqueeze(0).to(device)
spk_emb = v['spk_emb'].half().unsqueeze(0).to(device) if 'spk_emb' in v else None
messages = (history or []) + [{"role": "user", "content": prompt}]
open_thinking = bool(args.open_thinking)
inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=open_thinking)
x = torch.tensor(tokenizer(inputs_text).data['input_ids'], dtype=torch.long, device=device)[None, ...]
audio_frames = []
with torch.no_grad():
with model_lock:
res = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_tokens,
temperature=temperature, top_p=args.top_p, stream=True,
return_audio_codes=True, open_thinking=open_thinking,
audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values,
ref_codes=ref_codes, spk_emb=spk_emb)
history_idx = 0
for y, audio_frame in res:
text_chunk = None
if y is not None:
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
if answer and answer[-1] != '<EFBFBD>' and len(answer) > history_idx:
text_chunk = answer[history_idx:]
history_idx = len(answer)
if audio_frame:
audio_frames.append(audio_frame)
yield text_chunk, None, None
if asr_thread:
asr_thread.join()
if audio_frames and mimi_model:
yield None, "loading_audio", None
mimi_codes = frames_to_mimi(audio_frames)
audio_np = decode_mimi(mimi_codes)
if audio_np is not None:
yield None, (24000, audio_np), asr_result[0]
return
yield None, None, asr_result[0]
def launch_gradio(server_name="0.0.0.0", server_port=8888):
voice_choices = [("default", "default")]
for s in sorted(builtin_voices): voice_choices.append((f"[内置] {s}", s))
for s in sorted(clone_voices): voice_choices.append((f"[克隆] {s}", s))
def respond(message, audio, voice, chat_history, model_history, max_turns):
text = message.get("text", "") if isinstance(message, dict) else (message or "")
files = message.get("files", []) if isinstance(message, dict) else []
img_path = next((f for f in files if any(f.lower().endswith(e) for e in ('.png','.jpg','.jpeg','.gif','.bmp','.webp'))), None)
if not text and audio is None and img_path is None:
yield chat_history + [{"role": "assistant", "content": "请输入文本、上传图片或录制音频"}], gr.update(), gr.update(), model_history, ""
return
if audio is not None:
sr, samples = audio
display_samples = samples.mean(axis=1) if len(samples.shape) > 1 else samples
wav_path = tempfile.NamedTemporaryFile(suffix='.wav', delete=False).name
sf.write(wav_path, display_samples, sr)
chat_history = chat_history + [{"role": "user", "content": {"path": wav_path}}]
elif img_path:
chat_history = chat_history + [{"role": "user", "content": {"path": img_path}}, {"role": "user", "content": text or "请描述这张图片"}]
else:
chat_history = chat_history + [{"role": "user", "content": text}]
response_text = ""
final_audio = None
asr_text = None
chat_history = chat_history + [{"role": "assistant", "content": ""}]
yield chat_history, gr.update(value=None), None, model_history, ""
hist = model_history[-(max_turns * 2):] if max_turns > 0 else []
for text_chunk, audio_data, asr in chat_stream(
text or "", audio_input=audio, image_input=img_path,
voice_name=voice, history=hist, temperature=0.7, max_tokens=512
):
if text_chunk:
response_text += text_chunk
chat_history[-1]["content"] = response_text
yield chat_history, gr.update(), None, model_history, gr.update()
if audio_data == "loading_audio":
yield chat_history, gr.update(), None, model_history, '<div style="text-align:center;padding:10px 0;color:#aaa;font-size:13px;animation:pulse 1.5s ease-in-out infinite">正在组装语音(受限于 Gradio 特性,此处为非流式解码)...</div>'
elif audio_data:
final_audio = audio_data
if asr is not None:
asr_text = asr
user_text = asr_text if asr_text else (text or "")
if user_text:
model_history = model_history + [{"role": "user", "content": user_text}]
if response_text:
model_history = model_history + [{"role": "assistant", "content": response_text}]
yield chat_history, final_audio if final_audio else gr.update(), None, model_history, ""
with gr.Blocks(title="MiniMind-O", js="()=>{new MutationObserver(()=>{const m=document.getElementById('mic-box');if(!m)return;const h=!!m.querySelector('audio');document.body.classList.toggle('has-audio',h);const t=document.querySelector('textarea');if(t){t.placeholder=h?'已加载语音,点击发送':'输入文本';t.disabled=h}}).observe(document.body,{childList:true,subtree:true})}", css=".app{padding-top:6px!important} #component-0{gap:6px!important} #component-1{padding:2px 0!important;margin:0!important;min-height:0!important;border:none!important} #component-1 .padding{padding:0!important} #chatbox img{max-width:120px!important;max-height:120px!important;border-radius:8px} textarea{overflow-y:hidden!important;height:auto!important;min-height:30px!important;max-height:60px!important} #mic-box{max-height:150px!important;overflow:hidden!important} #mic-box .wrap span.or,#mic-box .wrap span:first-child{display:none!important} #mic-box .wrap{font-size:0!important;min-height:40px!important;padding:8px!important} #mic-box .wrap::after{content:'上传/录音';font-size:14px!important} #mic-box .mic-select{display:none!important} .has-audio textarea{opacity:0.4!important;pointer-events:none!important} .has-audio .upload-button,.has-audio [data-testid='upload-button']{display:none!important} @keyframes pulse{0%,100%{opacity:1}50%{opacity:0.4}}") as demo:
gr.HTML('<div style="text-align:center;margin:2px 0"><span style="font-size:1.2rem;font-weight:bold;font-style:italic">MiniMind-O</span> <span style="color:#999;font-size:0.8rem">text / image / audio → text + audio</span></div>')
chatbot = gr.Chatbot(label="", height=380, elem_id="chatbox", type="messages")
model_history = gr.State([])
audio_status = gr.HTML("", elem_id="audio-status")
audio_out = gr.Audio(label="语音回复", autoplay=True, elem_id="audio-out")
with gr.Row(equal_height=True):
with gr.Column(scale=0, min_width=160):
aud = gr.Audio(sources=["upload", "microphone"], type="numpy", show_label=False, elem_id="mic-box")
with gr.Column(scale=4):
msg = gr.MultimodalTextbox(placeholder="输入文本", show_label=False, submit_btn="发送")
with gr.Row():
voice_dd = gr.Dropdown(choices=voice_choices, value="default", label="音色选择", scale=0, min_width=140)
turns_dd = gr.Dropdown(choices=[0, 2, 4, 6, 8], value=0, label="多轮记忆", scale=0, min_width=120)
if model_dict:
model_dd = gr.Dropdown(choices=list(model_dict.keys()), value=current_model_name, label="模型选择", scale=1, min_width=180)
status = gr.Textbox(value=f"已加载: {current_model_name}", label="状态", interactive=False, scale=2)
model_dd.change(lambda n: load_hf_model(model_dict[n]), [model_dd], [status])
msg.submit(respond, [msg, aud, voice_dd, chatbot, model_history, turns_dd], [chatbot, audio_out, aud, model_history, audio_status])
demo.queue().launch(server_name=server_name, server_port=server_port)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="MiniMind-O Gradio Demo")
parser.add_argument('--load_from', default='./', type=str, help="transformers模型扫描目录")
parser.add_argument('--audio_encoder', default='../model/SenseVoiceSmall', type=str)
parser.add_argument('--vision_model', default='../model/siglip2-base-p32-256-ve', type=str)
parser.add_argument('--mimi_path', default='../model/mimi', type=str)
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
parser.add_argument('--open_thinking', default=0, type=int, choices=[0, 1])
parser.add_argument('--top_p', default=0.85, type=float)
parser.add_argument('--port', default=8888, type=int)
args = parser.parse_args()
device = args.device
model_dict = scan_hf_models(args.load_from)
if not model_dict:
print(f"未在 {os.path.abspath(args.load_from)} 找到 transformers 模型")
exit(1)
current_model_name = list(model_dict.keys())[0]
load_hf_model(model_dict[current_model_name])
try:
mimi_model = MimiModel.from_pretrained(args.mimi_path).eval()
print('Mimi model loaded')
except Exception:
mimi_model = None
print('Mimi model not found, audio output disabled')
with contextlib.redirect_stdout(io.StringIO()):
asr_model = AutoModel(model=args.audio_encoder, trust_remote_code=True, device=device, disable_update=True)
load_voices()
print(f'Voices loaded: {list(voices_data.keys()) or "none"}')
launch_gradio(server_port=args.port)

25
trainer/train.sh Normal file
View File

@ -0,0 +1,25 @@
## ==================== Full dataset training pipeline ====================
## Suggested full training pipeline (Dense, 4x GPU)
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_t2a.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --use_wandb --use_moe 0 # ~45min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_a2a.parquet --epochs 1 --batch_size 32 --use_compile 0 --from_weight sft_omni --max_seq_len 1024 --mode audio_proj --use_wandb --use_moe 0 # ~25min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-5 --data_path ../dataset/sft_a2a.parquet --epochs 3 --batch_size 32 --use_compile 0 --from_weight sft_omni --max_seq_len 1024 --use_wandb --use_moe 0 # 3 * ~25min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-5 --data_path ../dataset/sft_i2t.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --max_seq_len 768 --mode vision_proj --use_wandb --use_moe 0 # ~45min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_i2t.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --max_seq_len 768 --use_wandb --use_moe 0 # ~45min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_a2a.parquet --epochs 1 --batch_size 32 --use_compile 0 --from_weight sft_omni --max_seq_len 1024 --use_wandb --use_moe 0 # ~25min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_i2t.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --max_seq_len 768 --mode vision_proj --use_wandb --use_moe 0 # ~45min
## Suggested full training pipeline (MoE, 4x GPU)
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_t2a.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --use_wandb --use_moe 1 # ~45min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_a2a.parquet --epochs 1 --batch_size 32 --use_compile 0 --from_weight sft_omni --max_seq_len 1024 --mode audio_proj --use_wandb --use_moe 1 # ~25min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-5 --data_path ../dataset/sft_a2a.parquet --epochs 3 --batch_size 32 --use_compile 0 --from_weight sft_omni --max_seq_len 1024 --use_wandb --use_moe 1 # 3 * ~25min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-5 --data_path ../dataset/sft_i2t.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --max_seq_len 768 --mode vision_proj --use_wandb --use_moe 1 # ~45min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_i2t.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --max_seq_len 768 --use_wandb --use_moe 1 # ~45min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_a2a.parquet --epochs 1 --batch_size 32 --use_compile 0 --from_weight sft_omni --max_seq_len 1024 --use_wandb --use_moe 1 # ~25min
# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port 29560 --nproc_per_node 4 train_sft_omni.py --learning_rate 5e-6 --data_path ../dataset/sft_i2t.parquet --epochs 1 --batch_size 32 --use_compile 1 --from_weight sft_omni --max_seq_len 768 --mode vision_proj --use_wandb --use_moe 1 # ~45min
# ==================== (Recommend) Mini dataset training pipeline ====================
# Time estimates below are based on 1x RTX 3090, for reference only and may not be entirely accurate.
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_t2a_mini.parquet --epochs 1 --batch_size 40 --use_compile 1 --from_weight llm --save_weight sft_zero --max_seq_len 512 --use_wandb --use_moe 0 # epochs * 60min
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 5e-4 --data_path ../dataset/sft_a2a_mini.parquet --epochs 1 --batch_size 40 --use_compile 0 --from_weight sft_zero --save_weight sft_zero --max_seq_len 640 --mode audio_proj --use_wandb --use_moe 0 # epochs * 15min
CUDA_VISIBLE_DEVICES=0 torchrun --master_port 29560 --nproc_per_node 1 train_sft_omni.py --learning_rate 2e-5 --data_path ../dataset/sft_a2a_mini.parquet --epochs 1 --batch_size 16 --use_compile 0 --from_weight sft_zero --save_weight sft_zero --max_seq_len 768 --use_wandb --use_moe 0 # 2 * ~15min

263
trainer/train_sft_omni.py Normal file
View File

@ -0,0 +1,263 @@
import os
import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import argparse
import time
import warnings
import torch
import torch.nn as nn
import torch.distributed as dist
from contextlib import nullcontext
from torch import optim
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
from model.model_omni import OmniConfig
from dataset.omni_dataset import OmniDataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, init_distributed_mode, setup_seed, init_omni_model, omni_checkpoint, SkipBatchSampler, log_model_params
warnings.filterwarnings('ignore')
def omni_collate_fn(batch):
"""自定义collate函数,处理变长audio_inputs和pixel_values"""
input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb = zip(*batch)
input_ids = torch.stack(input_ids)
labels = torch.stack(labels)
audio_labels = torch.stack(audio_labels)
audio_lens = torch.tensor(audio_lens, dtype=torch.long)
valid_audios = [a for a in audio_inputs if a is not None]
if valid_audios:
max_t = max(a.size(1) for a in valid_audios)
padded = [a if a.size(1) == max_t else torch.nn.functional.pad(a, (0, 0, 0, max_t - a.size(1))) for a in valid_audios]
audio_inputs = torch.cat(padded, dim=0)
else:
audio_inputs = None
valid_images = [p for p in pixel_values if p is not None]
if valid_images:
if hasattr(valid_images[0], 'keys'):
keys = set.intersection(*[set(d.keys()) for d in valid_images])
pixel_values = {k: torch.cat([d[k] for d in valid_images], dim=0) for k in keys}
else:
pixel_values = torch.cat(valid_images, dim=0)
else:
pixel_values = None
spk_emb = torch.stack(spk_emb)
return input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
start_time = time.time()
last_step = start_step
for step, (input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb) in enumerate(loader, start=start_step + 1):
input_ids = input_ids.to(args.device)
labels = labels.to(args.device)
audio_labels = audio_labels.to(args.device)
audio_lens = audio_lens.to(args.device)
if audio_inputs is not None:
audio_inputs = audio_inputs.to(args.device)
if pixel_values is not None:
if hasattr(pixel_values, 'keys'):
pixel_values = {k: v.to(args.device) for k, v in pixel_values.items()}
else:
pixel_values = pixel_values.to(args.device)
spk_emb = spk_emb.to(args.device)
last_step = step
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
with autocast_ctx:
res = model(input_ids, audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values, spk_emb=spk_emb)
loss_fct = nn.CrossEntropyLoss(reduction='none')
# Text loss
text_loss_raw = loss_fct(res.logits.view(-1, res.logits.size(-1)), labels.view(-1))
text_mask = (labels.view(-1) != -100).float()
text_loss = (text_loss_raw * text_mask).sum() / (text_mask.sum() + 1e-9)
# Audio loss
audio_loss = res.audio_logits[0].sum() * 0
for i, al in enumerate(res.audio_logits):
al_flat = al.view(-1, al.size(-1))
target_flat = audio_labels[:, i, :].reshape(-1)
layer_loss = loss_fct(al_flat, target_flat)
valid_mask = (target_flat != -100).float()
stop_mask = (target_flat == 2050).float()
weighted_loss = layer_loss * valid_mask * (1 + stop_mask * 9)
msum = valid_mask.sum()
if msum > 0:
audio_loss = audio_loss + weighted_loss.sum() / (msum + 1e-9)
audio_loss = audio_loss / 8
loss = (text_loss + audio_loss + res.aux_loss) / args.accumulation_steps
scaler.scale(loss).backward()
if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
if step % args.log_interval == 0 or step == iters:
spend_time = time.time() - start_time
current_loss = loss.item() * args.accumulation_steps
text_loss_val = text_loss.item() if isinstance(text_loss, torch.Tensor) else 0
audio_loss_val = audio_loss.item() if isinstance(audio_loss, torch.Tensor) else 0
current_lr = optimizer.param_groups[-1]['lr']
eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, text: {text_loss_val:.4f}, audio: {audio_loss_val:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
if wandb:
wandb.log({"loss": current_loss, "text_loss": text_loss_val,
"audio_loss": audio_loss_val, "lr": current_lr, "epoch_time": eta_min})
if (step % args.save_interval == 0 or step == iters) and is_main_process():
model.eval()
moe_suffix = '_moe' if omni_config.use_moe else ''
ckp = f'{args.save_dir}/{args.save_weight}_{omni_config.hidden_size}{moe_suffix}.pth'
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
raw_model = getattr(raw_model, '_orig_mod', raw_model)
clean_state_dict = {k: v for k, v in raw_model.state_dict().items() if not k.startswith('audio_encoder.')}
torch.save({k: v.half().cpu() for k, v in clean_state_dict.items()}, ckp)
omni_checkpoint(omni_config, weight=args.save_weight, model=model, optimizer=optimizer,
epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
model.train()
del input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb, res, loss
if last_step > start_step and last_step % args.accumulation_steps != 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="MiniMind-O SFT")
parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
parser.add_argument('--save_weight', default='sft_omni', type=str, help="保存权重的前缀名")
parser.add_argument("--epochs", type=int, default=15, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--num_workers", type=int, default=4, help="数据加载线程数")
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
parser.add_argument('--max_seq_len', default=512, type=int, help="训练的最大截断长度")
parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构")
parser.add_argument("--data_path", type=str, default="../dataset/train_t2a_mini.parquet", help="训练数据路径(parquet格式)")
parser.add_argument("--audio_encoder_dir", type=str, default="../model/SenseVoiceSmall", help="音频encoder路径(SenseVoice)")
parser.add_argument("--vision_dir", type=str, default="../model/siglip2-base-p32-256-ve", help="CLIP视觉模型路径")
parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练,为none则不基于任何权重训练")
parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训(0=否,1=是)")
parser.add_argument('--freeze_backbone', default='none', type=str, choices=['none', 'all', 'last1'], help="冻结主干模型: none=全量训练, all=只训练audio层, last1=只训练最后1层+audio层")
parser.add_argument('--mode', default='all', type=str, choices=['all', 'audio_proj', 'vision_proj'], help="训练模式: all=全量训练, audio_proj=只训练audio_proj, vision_proj=只训练vision_proj")
parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
parser.add_argument("--wandb_project", type=str, default="MiniMind-O-SFT", help="wandb项目名")
parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速(0=否,1=是)")
args = parser.parse_args()
# ========== 1. 初始化环境和随机种子 ==========
local_rank = init_distributed_mode()
if dist.is_initialized():
args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
# ========== 2. 配置目录、模型参数、检查ckp ==========
os.makedirs(args.save_dir, exist_ok=True)
omni_config = OmniConfig(
hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers,
use_moe=bool(args.use_moe)
)
ckp_data = omni_checkpoint(omni_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu"
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
# ========== 4. 配wandb ==========
wandb = None
if args.use_wandb and is_main_process():
import swanlab as wandb
wandb_id = ckp_data.get('wandb_id') if ckp_data else None
resume = 'must' if wandb_id else None
wandb_run_name = f"MiniMind-O-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LR-{args.learning_rate}"
wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
# ========== 5. 定义模型、数据、优化器 ==========
model, tokenizer = init_omni_model(omni_config, from_weight=args.from_weight,
audio_encoder_path=args.audio_encoder_dir,
vision_model_path=args.vision_dir,
save_dir=args.save_dir, device=args.device,
freeze_backbone=args.freeze_backbone, from_resume=args.from_resume)
if args.use_compile == 1:
model = torch.compile(model)
if model.audio_encoder is not None: model.audio_encoder.to(args.device)
if model.vision_encoder is not None: model.vision_encoder.to(args.device)
if args.mode == 'audio_proj':
for p in model.parameters(): p.requires_grad = False
for p in model.audio_proj.parameters(): p.requires_grad = True
elif args.mode == 'vision_proj':
for p in model.parameters(): p.requires_grad = False
for p in model.vision_proj.parameters(): p.requires_grad = True
log_model_params(model)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6
Logger(f'Trainable: {trainable:.2f}M | Mode: {args.mode} | Freeze: {args.freeze_backbone} | Compile: {"on" if args.use_compile else "off"}')
# scheduled_sampling 现在会自动保护 image/audio token 的连续性
train_ds = OmniDataset(
args.data_path,
tokenizer,
audio_processor=model.audio_processor,
vision_processor=model.vision_processor,
max_length=args.max_seq_len,
image_token_len=model.config.image_token_len
)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ==========
start_epoch, start_step = 0, 0
if ckp_data:
model.load_state_dict(ckp_data['model'], strict=False)
optimizer.load_state_dict(ckp_data['optimizer'])
scaler.load_state_dict(ckp_data['scaler'])
start_epoch = ckp_data['epoch']
start_step = ckp_data.get('step', 0)
# ========== 7. DDP包模型 ==========
if dist.is_initialized():
model = DistributedDataParallel(model, device_ids=[local_rank])
# ========== 8. 开始训练 ==========
for epoch in range(start_epoch, args.epochs):
train_sampler and train_sampler.set_epoch(epoch)
setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
skip = start_step if (epoch == start_epoch and start_step > 0) else 0
batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
loader = DataLoader(train_ds, batch_sampler=batch_sampler, collate_fn=omni_collate_fn, num_workers=args.num_workers, pin_memory=True)
if skip > 0:
Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step,从step {start_step + 1}开始')
train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
else:
train_epoch(epoch, loader, len(loader), 0, wandb)
# ========== 9. 清理分布进程 ==========
if dist.is_initialized(): dist.destroy_process_group()

200
trainer/trainer_utils.py Normal file
View File

@ -0,0 +1,200 @@
"""
训练工具函数集合
"""
import os
import random
import math
import numpy as np
import torch
import torch.distributed as dist
from torch.utils.data import Sampler
from transformers import AutoTokenizer
from model.model_omni import MiniMindOmni
def is_main_process():
return not dist.is_initialized() or dist.get_rank() == 0
def Logger(content):
if is_main_process():
print(content)
def get_lr(current_step, total_steps, lr):
# 与 mmv 保持一致:初始 lr=1.0*base_lr,最终 lr=0.1*base_lr
return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * current_step / total_steps)))
def init_distributed_mode():
if int(os.environ.get("RANK", -1)) == -1:
return 0 # 非DDP模式
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
def setup_seed(seed: int):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def log_model_params(model, ignore_patterns=['audio_encoder', 'vision_encoder']):
def should_count(n): return not any(p in n for p in ignore_patterns)
total = sum(p.numel() for n, p in model.named_parameters() if should_count(n)) / 1e6
cfg = model.config
n_routed = getattr(cfg, 'n_routed_experts', getattr(cfg, 'num_experts', 0))
n_active = getattr(cfg, 'num_experts_per_tok', 0)
n_shared = getattr(cfg, 'n_shared_experts', 0)
expert = sum(p.numel() for n, p in model.named_parameters() if 'mlp.experts.0.' in n and should_count(n)) / 1e6
shared_expert = sum(p.numel() for n, p in model.named_parameters() if 'mlp.shared_experts.0.' in n and should_count(n)) / 1e6
base = total - (expert * n_routed) - (shared_expert * n_shared)
active = base + (expert * n_active) + (shared_expert * n_shared)
if active < total: Logger(f'Model Params: {total:.2f}M-A{active:.2f}M')
else: Logger(f'Model Params: {total:.2f}M')
def init_omni_model(omni_config, from_weight='full_sft', tokenizer_path='../model', audio_encoder_path='../model/SenseVoiceSmall', vision_model_path='../model/siglip2-base-p32-256-ve', save_dir='../out', device='cuda', freeze_backbone='none', from_resume=0):
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
model = MiniMindOmni(omni_config, audio_encoder_path=audio_encoder_path, vision_model_path=vision_model_path)
if from_weight != 'none':
moe_suffix = '_moe' if omni_config.use_moe else ''
weight_path = f'{save_dir}/{from_weight}_{omni_config.hidden_size}{moe_suffix}.pth'
if os.path.exists(weight_path):
weights = torch.load(weight_path, map_location=device)
param_shapes = {k: v.shape for k, v in model.named_parameters()}
incompatible = {k for k, v in weights.items() if k in param_shapes and v.shape != param_shapes[k]}
if incompatible:
Logger(f'跳过shape不匹配的权重: {incompatible}')
weights = {k: v for k, v in weights.items() if k not in incompatible}
model.load_state_dict(weights, strict=False)
Logger(f'已加载权重: {weight_path}')
if from_resume == 0 and omni_config.talker_hidden_size == omni_config.hidden_size:
n_talker = omni_config.num_talker_hidden_layers
n_thinker = len(model.thinker.layers)
has_talker = any(k.startswith('talker.layers.') for k in weights)
if not has_talker and n_talker > 0:
for i in range(n_talker):
src = n_thinker - n_talker + i
model.talker.layers[i].load_state_dict(model.thinker.layers[src].state_dict())
Logger(f'Talker层初始化: 复制thinker layers[{n_thinker-n_talker}:{n_thinker}] → talker layers[0:{n_talker}]')
# 冻结策略
if freeze_backbone == 'all':
# 冻结整个主干模型
for param in model.model.parameters():
param.requires_grad = False
elif freeze_backbone == 'last1':
# 冻结除了最后1层之外的所有层
for param in model.model.parameters():
param.requires_grad = False
# 打开最后1层
if hasattr(model.model, 'layers') and len(model.model.layers) > 0:
for param in model.model.layers[-1].parameters():
param.requires_grad = True
return model.to(device), tokenizer
def omni_checkpoint(omni_config, weight='pretrain_omni', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
os.makedirs(save_dir, exist_ok=True)
moe_path = '_moe' if omni_config.use_moe else ''
ckp_path = f'{save_dir}/{weight}_{omni_config.hidden_size}{moe_path}.pth'
resume_path = f'{save_dir}/{weight}_{omni_config.hidden_size}{moe_path}_resume.pth'
if model is not None:
from torch.nn.parallel import DistributedDataParallel
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
raw_model = getattr(raw_model, '_orig_mod', raw_model)
# 移除冻结的 audio_encoder / vision_encoder 参数(不需要保存,从预训练路径重新加载)
clean_state_dict = {k: v for k, v in raw_model.state_dict().items() if not k.startswith('audio_encoder.') and not k.startswith('vision_encoder.')}
state_dict = {k: v.half().cpu() for k, v in clean_state_dict.items()}
ckp_tmp = ckp_path + '.tmp'
torch.save(state_dict, ckp_tmp)
os.replace(ckp_tmp, ckp_path)
wandb_id = None
if wandb:
if hasattr(wandb, 'get_run'):
run = wandb.get_run()
wandb_id = getattr(run, 'id', None) if run else None
else:
wandb_id = getattr(wandb, 'id', None)
resume_data = {
'model': state_dict,
'optimizer': optimizer.state_dict(),
'epoch': epoch,
'step': step,
'world_size': dist.get_world_size() if dist.is_initialized() else 1,
'wandb_id': wandb_id
}
for key, value in kwargs.items():
if value is not None:
if hasattr(value, 'state_dict'):
if isinstance(value, DistributedDataParallel):
resume_data[key] = value.module.state_dict()
else:
resume_data[key] = value.state_dict()
else:
resume_data[key] = value
resume_tmp = resume_path + '.tmp'
torch.save(resume_data, resume_tmp)
os.replace(resume_tmp, resume_path)
else: # 加载模式
if os.path.exists(resume_path):
ckp_data = torch.load(resume_path, map_location='cpu')
saved_ws = ckp_data.get('world_size', 1)
current_ws = dist.get_world_size() if dist.is_initialized() else 1
if saved_ws != current_ws:
ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws
Logger(f'GPU数量变化({saved_ws}→{current_ws}),step已自动转换为{ckp_data["step"]}')
return ckp_data
return None
def vlm_collate_fn(batch):
input_ids = torch.stack([b[0] for b in batch])
labels = torch.stack([b[1] for b in batch])
pixel_data = [b[2] for b in batch]
if hasattr(pixel_data[0], 'keys'):
pixel_values = {k: torch.stack([d[k] for d in pixel_data]) for k in pixel_data[0].keys()}
else:
pixel_values = torch.stack(pixel_data)
return input_ids, labels, pixel_values
class SkipBatchSampler(Sampler):
def __init__(self, sampler, batch_size, skip_batches=0):
self.sampler = sampler
self.batch_size = batch_size
self.skip_batches = skip_batches
def __iter__(self):
batch = []
skipped = 0
for idx in self.sampler:
batch.append(idx)
if len(batch) == self.batch_size:
if skipped < self.skip_batches:
skipped += 1
batch = []
continue
yield batch
batch = []
if len(batch) > 0 and skipped >= self.skip_batches:
yield batch
def __len__(self):
total_batches = (len(self.sampler) + self.batch_size - 1) // self.batch_size
return max(0, total_batches - self.skip_batches)

727
webui/web_demo.html Normal file
View File

@ -0,0 +1,727 @@
<!DOCTYPE html>
<html lang="zh">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>MiniMind-O</title>
<link rel="icon" type="image/png" href="https://www.modelscope.cn/studio/gongjy/MiniMind/resolve/master/images/logo2.png">
<style>
:root {
--bg:#1e1e20; --p1:#2a2a2e; --p2:#35353b; --p3:#40404a;
--bd:#55555e; --bd2:#7a7a85;
--fg:#ececef; --mu:#9a9aa0; --dim:#8a8a92;
--u:#7fc5e0; --a:#e5b57a; --danger:#d84545; --cam:#7ad88a;
--serif:'Georgia','Times New Roman',serif; --mono:'SF Mono',Consolas,monospace;
}
*{margin:0;padding:0;box-sizing:border-box}
html,body{height:100%}
body{height:100vh;background:var(--bg);color:var(--fg);font-family:var(--serif);display:flex;align-items:center;justify-content:center;overflow:hidden}
.app{max-width:860px;width:100%;height:calc(100vh - 40px);display:flex;flex-direction:column;padding:20px 24px;background:var(--p1);border:1px solid var(--bd);overflow:hidden}
.topbar{display:flex;align-items:center;justify-content:space-between;padding:4px 0 12px;flex-shrink:0;gap:10px}
.brand{display:flex;align-items:center;gap:10px;min-width:0;flex:1}
.brand>div{min-width:0}
.brand img{height:34px}
.brand h1{font-size:1.15rem;font-weight:700;font-style:italic;color:#fff}
.brand .sub{font-size:.72rem;color:var(--mu);font-style:italic}
.topright{display:flex;gap:8px;align-items:center}
.btn{padding:5px 10px;font-size:.75rem;cursor:pointer;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:var(--serif);font-style:italic;display:inline-flex;align-items:center;gap:5px}
.btn:hover:not(:disabled){background:var(--p3);border-color:var(--bd2)}
.btn:disabled{opacity:.4;cursor:not-allowed}
.btn.danger{color:var(--danger);border-color:#6a3a3a}
.btn.danger:hover:not(:disabled){background:var(--danger);border-color:var(--danger);color:#fff}
.btn svg{width:12px;height:12px;fill:currentColor}
.model-sel{max-width:150px;padding:5px 8px;font-size:.72rem;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:var(--serif);font-style:italic;outline:none}
.brand .model-sel{max-width:260px;padding:0 34px 0 0;background:transparent;border:none;font-size:1.15rem;font-weight:700;color:#fff}
.model-sel:focus{border-color:var(--bd2)}
.mini-sel{max-width:130px;padding:5px 8px;font-size:.72rem;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:var(--serif);font-style:italic;outline:none}
.mini-sel:focus{border-color:var(--bd2)}
.voice-sel{display:none}
.more{position:relative}
.more-pop{top:100%;right:0;left:auto;margin-top:6px;min-width:130px}
.vgroup{padding:4px 10px;font-size:.6rem;color:var(--mu);font-style:italic}
.vitem{padding:5px 10px;display:flex;justify-content:space-between;align-items:center;cursor:pointer;font-size:.72rem;font-family:var(--serif);font-style:italic;color:var(--fg);gap:8px}
.vitem:hover{background:var(--p3)}
.vitem.on{color:var(--u)}
.vx{color:var(--mu);cursor:pointer;font-style:normal;font-size:.8rem;padding:0 2px;line-height:1}
.vx:hover{color:var(--danger)}
.vadd{padding:5px 10px;font-size:.72rem;color:var(--mu);cursor:pointer;font-family:var(--serif);font-style:italic;border-top:1px solid var(--bd)}
.vadd:hover{background:var(--p3);color:var(--fg)}
.status{display:inline-flex;align-items:center;gap:10px;padding:5px 10px;font-size:.65rem;font-family:var(--mono);background:var(--p2);border:1px solid var(--bd);transition:opacity .3s}
.status.hidden{opacity:0;pointer-events:none}
.status .row{display:inline-flex;gap:5px;white-space:nowrap}
.status .lab{color:var(--mu);font-style:italic}
.chat{flex:1;overflow-y:auto;padding:12px 0;min-height:0;display:flex;flex-direction:column;gap:18px;scrollbar-width:none;-ms-overflow-style:none}
.chat::-webkit-scrollbar{display:none}
.msg{display:flex;flex-direction:column;gap:6px;animation:fi .25s ease}
.msg.user{align-items:flex-end}
.msg.user .body{align-items:flex-end}
@keyframes fi{from{opacity:0;transform:translateY(3px)}to{opacity:1;transform:translateY(0)}}
.meta{font-size:.72rem;font-style:italic}
.msg.user .meta{color:var(--u)}
.msg.assistant .meta{color:var(--a)}
.body{display:flex;flex-direction:column;gap:8px}
.text{line-height:1.65;white-space:pre-wrap;word-break:break-word}
.img{display:block;max-width:160px;max-height:160px}
.msg.user .text:not(.voice){color:var(--u)}
.voice{font-family:var(--mono);font-size:1.05rem;line-height:1.2;color:var(--bd)}
.msg.user .voice{color:rgba(127,197,224,.5)}
.prog{text-decoration:underline dashed;text-decoration-color:var(--bd);text-underline-offset:4px;text-decoration-thickness:1px}
.prog .tk{cursor:pointer;padding:1px 0;transition:color .08s,background .08s}
.msg.assistant .prog .tk.played{color:var(--a)}
.msg.assistant .prog .tk:hover{background:rgba(229,181,122,.2)}
.msg.assistant .prog.on .tk.cur{background:rgba(229,181,122,.45)}
.msg.user .prog .tk.played{color:var(--u)}
.msg.user .prog .tk:hover{background:rgba(127,197,224,.18)}
.msg.user .prog.on .tk.cur{background:rgba(127,197,224,.35)}
.player{display:inline-flex;align-items:center;gap:10px;padding:6px 12px;width:320px;max-width:100%;background:var(--p2);border:1px solid var(--bd)}
.pp{background:transparent;border:none;color:inherit;cursor:pointer;font-family:var(--mono);font-size:.9rem;width:20px;text-align:left}
.pp:hover{opacity:.75}
.bar-wrap{flex:1;height:3px;cursor:pointer;background:var(--p1)}
.bar{height:100%;width:0;transition:width .08s linear}
.msg.user .bar,.msg.user .pp{background:var(--u);color:var(--u)}
.msg.assistant .bar,.msg.assistant .pp{background:var(--a);color:var(--a)}
.msg.user .bar-wrap{background:var(--p1)}
.time{font-size:.7rem;font-family:var(--mono);min-width:76px;text-align:right;color:var(--mu)}
.input-area{flex-shrink:0;padding:10px 0 4px}
.compose{display:flex;gap:6px}
.compose button{padding:0 12px;min-width:42px;cursor:pointer;font-size:.85rem;background:var(--p2);border:1px solid var(--bd);color:var(--fg);display:flex;align-items:center;justify-content:center;font-family:inherit}
.compose button svg{width:16px;height:16px;fill:currentColor}
.compose button:hover:not(:disabled){background:var(--p3);border-color:var(--bd2)}
.compose button:disabled{opacity:.4;cursor:not-allowed}
.compose button.primary{background:var(--fg);color:#1a1a1e;border-color:var(--fg)}
.compose button.primary:hover:not(:disabled){background:#fff}
.compose button.rec{animation:pulse 1s infinite;background:var(--danger);color:#fff;border-color:var(--danger)}
@keyframes pulse{0%,100%{opacity:1}50%{opacity:.5}}
.text-wrap{flex:1;display:flex;align-items:center;gap:6px;padding:0 6px 0 10px;min-width:0;background:var(--p2);border:1px solid var(--bd)}
.text-wrap:focus-within{border-color:var(--fg)}
.text-wrap input{flex:1;min-width:0;padding:10px 4px;font-size:.9rem;outline:none;background:transparent;border:none;color:inherit;font-family:inherit}
.text-wrap input::placeholder{color:var(--dim)}
.attach{display:flex;gap:4px;flex-wrap:nowrap;flex-shrink:0;max-width:45%;overflow-x:auto;scrollbar-width:none}
.attach::-webkit-scrollbar{display:none}
.attach:empty{display:none}
.chip{display:inline-flex;align-items:center;gap:6px;padding:3px 7px;font-size:.72rem;background:var(--p1);border:1px solid var(--bd);flex-shrink:0}
.chip .x{cursor:pointer;opacity:.55;font-size:.85rem}
.chip .x:hover{opacity:1}
.chip .play{cursor:pointer}
.chip img{width:22px;height:22px;object-fit:cover}
.pop{position:absolute;background:var(--p2);border:1px solid var(--bd);display:none;flex-direction:column;z-index:50;min-width:110px;font-family:var(--serif);font-style:italic;white-space:nowrap}
.pop.on{display:flex}
.pop button{padding:5px 10px;background:transparent;border:none;color:var(--fg);cursor:pointer;text-align:left;font-size:.72rem;font-family:inherit;font-style:inherit;line-height:1.2}
.pop button:hover{background:var(--p3)}
.snap{position:fixed;inset:0;z-index:200;background:rgba(0,0,0,.85);display:none;align-items:center;justify-content:center;flex-direction:column;gap:14px}
.snap.on{display:flex}
.snap video{max-width:80vw;max-height:70vh;background:#000;border:1px solid var(--bd)}
.snap .sr{display:flex;gap:10px}
.snap .sr button{padding:8px 18px;background:var(--p2);border:1px solid var(--bd);color:var(--fg);cursor:pointer;font-family:var(--serif);font-style:italic}
.snap .sr button.primary{background:var(--fg);color:#1a1a1e;border-color:var(--fg)}
.clone-card{width:min(92vw,520px);padding:22px;background:var(--p1);border:1px solid var(--bd);display:flex;flex-direction:column;gap:14px}
.clone-title{font-size:1rem;font-style:italic;color:#fff}
.clone-hint{font-size:.76rem;color:var(--mu);font-style:italic;line-height:1.6}
.clone-name{width:100%;padding:10px 12px;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:inherit;font-size:.85rem;outline:none}
.clone-name:focus{border-color:var(--fg)}
.clone-script{padding:12px 14px;background:var(--p2);border:1px solid var(--bd);font-size:1rem;line-height:1.85}
.clone-stat{min-height:18px;font-size:.72rem;color:var(--mu);font-family:var(--mono)}
/* ===== Call mode ===== */
body>.app,#call{will-change:transform,opacity,filter;transform-origin:center center}
body>.app{transition:opacity .55s cubic-bezier(.22,1,.36,1),transform .7s cubic-bezier(.22,1,.36,1),filter .55s cubic-bezier(.22,1,.36,1)}
body.call>.app{opacity:0;transform:scale(.92);filter:blur(8px);pointer-events:none}
#call{position:fixed;inset:0;z-index:100;background:var(--bg);overflow:hidden;opacity:0;transform:scale(1.08);filter:blur(10px);pointer-events:none;visibility:hidden;transition:opacity .55s cubic-bezier(.22,1,.36,1),transform .7s cubic-bezier(.22,1,.36,1),filter .55s cubic-bezier(.22,1,.36,1),visibility 0s .7s}
body.call #call{opacity:1;transform:scale(1);filter:blur(0);pointer-events:auto;visibility:visible;transition:opacity .55s cubic-bezier(.22,1,.36,1),transform .7s cubic-bezier(.22,1,.36,1),filter .55s cubic-bezier(.22,1,.36,1),visibility 0s 0s}
#call .app{max-width:560px;margin:20px auto;padding:18px 22px;position:relative}
#call .topbar{border-bottom:1px dashed var(--bd)}
#call .brand h1{font-size:1.05rem}
#call .brand .model-sel{max-width:230px;padding-right:38px;font-size:1.05rem}
#call .brand .sub{font-size:.68rem}
.pill{display:inline-flex;align-items:center;gap:6px;padding:5px 10px;font-size:.7rem;font-family:var(--mono);background:var(--p2);border:1px solid var(--bd);color:var(--mu);font-style:italic;line-height:1.2}
.pill .dot{width:6px;height:6px;background:var(--mu);border-radius:50%}
.pill.live{color:var(--u);border-color:#3c5966}
.pill.live .dot{background:var(--u);box-shadow:0 0 6px var(--u);animation:cp 1.4s ease-in-out infinite}
@keyframes cp{0%,100%{opacity:1}50%{opacity:.4}}
#call .btn.active{background:var(--u);color:#1a1a1e;border-color:var(--u)}
.hero{flex:1;display:flex;align-items:center;justify-content:center;min-height:0;padding:14px 0;position:relative}
.orb{position:relative;width:280px;height:280px;max-width:60vh;max-height:60vh;display:flex;align-items:center;justify-content:center;will-change:transform}
.orb svg{width:100%;height:100%;overflow:visible}
.blob{fill:none;stroke-linejoin:round;transition:stroke .35s,opacity .45s}
.blob.l0{stroke:#5a5a64;stroke-width:1.2;opacity:.55}
.blob.l1{stroke:var(--bd2);stroke-width:1.4;opacity:.85}
.blob.l2{stroke:var(--fg);stroke-width:1.6}
.orb.camon .blob.l0,.orb.camon .blob.l1{opacity:0}
.orb.camon .blob.l2{stroke:var(--cam);filter:drop-shadow(0 0 4px rgba(122,216,138,.45))}
svg.listen .blob.l2{stroke:var(--u);filter:drop-shadow(0 0 4px rgba(127,197,224,.45))}
svg.listen .blob.l1{stroke:rgba(127,197,224,.6)}
svg.speak .blob.l2{stroke:var(--a);filter:drop-shadow(0 0 5px rgba(229,181,122,.5))}
svg.speak .blob.l1{stroke:rgba(229,181,122,.6)}
svg.gen .blob.l2{stroke:var(--fg);filter:drop-shadow(0 0 4px rgba(236,236,239,.35))}
.core{display:none}
svg.listen .core{fill:var(--u)}
svg.speak .core{fill:var(--a)}
.ripple{fill:none;stroke:var(--mu);stroke-width:1;transform-origin:center;animation:rip 1.6s ease-out forwards}
.ripple.user{stroke:var(--u)}
.ripple.assistant{stroke:var(--a)}
@keyframes rip{0%{transform:scale(.55);opacity:.55}100%{transform:scale(1.35);opacity:0}}
.lbl{position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);text-align:center;pointer-events:none;display:flex;flex-direction:column;align-items:center;gap:6px}
.ic{width:34px;height:34px;color:var(--mu);transition:color .3s;display:flex;align-items:center;justify-content:center}
.ic svg{width:100%;height:100%;stroke:currentColor;fill:none;stroke-width:2;stroke-linecap:round;stroke-linejoin:round}
.lbl.listen .ic{color:var(--u)}
.lbl.speak .ic{color:var(--a)}
.lbl.gen .ic{color:var(--fg)}
.lbl .st{display:none}
.ic svg *{transform-box:view-box;transform-origin:12px 12px}
.lbl.idle .ic .dot{animation:breathe 2.2s ease-in-out infinite;transform-origin:center;transform-box:view-box}
@keyframes breathe{0%,100%{opacity:.4;transform:scale(.8)}50%{opacity:1;transform:scale(1.15)}}
.lbl.listen .ic .w{animation:sigin 1.6s ease-out infinite}
.lbl.listen .ic .w2{animation-delay:.3s}
.lbl.listen .ic .w3{animation-delay:.6s}
@keyframes sigin{0%{opacity:0;transform:scale(1.2)}60%{opacity:1;transform:scale(1)}100%{opacity:0;transform:scale(.85)}}
.lbl.gen .ic svg{animation:spin 1.2s cubic-bezier(.5,.1,.5,.9) infinite}
@keyframes spin{to{transform:rotate(360deg)}}
.lbl.speak .ic .a1{animation:radout 1.3s ease-out infinite}
.lbl.speak .ic .a2{animation:radout 1.3s ease-out .25s infinite}
.lbl.speak .ic .a3{animation:radout 1.3s ease-out .5s infinite}
@keyframes radout{0%{opacity:0;transform:scale(.7)}40%{opacity:1}100%{opacity:0;transform:scale(1.15)}}
.up{display:none}
.cam{position:absolute;top:50%;left:50%;width:180px;height:180px;object-fit:cover;background:var(--bg);border:1px solid var(--bd);transform:translate(-50%,-50%) scale(.3);opacity:0;transition:transform .5s cubic-bezier(.4,0,.2,1),opacity .3s,border-color .35s,box-shadow .35s;pointer-events:none;will-change:transform,opacity}
.orb.camon .cam{transform:translate(-50%,-50%) scale(1);opacity:1;border-color:var(--cam);box-shadow:0 0 14px rgba(122,216,138,.35)}
.orb.camon .core{opacity:0}
.orb.camon .lbl{top:auto;bottom:-8px;transform:translate(-50%,100%)}
.orb.camon .lbl .st{font-size:.9rem}
svg.listen ~ .cam{border-color:var(--u);box-shadow:0 0 14px rgba(127,197,224,.45)}
svg.speak ~ .cam{border-color:var(--a);box-shadow:0 0 14px rgba(229,181,122,.45)}
.latest{width:100%;padding:10px 14px;background:var(--p2);border:1px solid var(--bd);display:flex;flex-direction:column;gap:5px;flex-shrink:0;max-height:120px;overflow-y:auto;scrollbar-width:none}
.latest::-webkit-scrollbar{display:none}
.latest .lm{font-family:var(--mono);font-size:.6rem;color:var(--mu);font-style:italic}
.latest .lt{font-size:.82rem;line-height:1.55;font-style:italic;opacity:.9;white-space:pre-wrap;word-break:break-word}
.latest.user .lm{color:var(--u)}
.latest.assistant .lm{color:var(--a)}
.latest.empty .lt{color:var(--bd)}
.ctrls{display:flex;justify-content:center;gap:10px;padding:14px 0 2px;margin-top:12px;border-top:1px dashed var(--bd);flex-shrink:0}
.cbtn{width:72px;height:58px;background:var(--p2);border:1px solid var(--bd);color:var(--fg);cursor:pointer;display:flex;flex-direction:column;align-items:center;justify-content:center;gap:3px;font-family:var(--mono);font-size:.58rem;letter-spacing:1px;text-transform:uppercase;transition:all .18s}
.cbtn svg{width:18px;height:18px;fill:currentColor}
.cbtn:hover{background:var(--p3);border-color:var(--bd2)}
.cbtn.active{background:var(--u);color:#1a1a1e;border-color:var(--u)}
.cbtn.danger{color:var(--danger);border-color:#6a3a3a}
.cbtn.danger:hover{background:var(--danger);color:#fff}
.trans{position:absolute;inset:0;background:var(--p1);z-index:20;display:flex;flex-direction:column;padding:18px 22px;transform:translateY(100%);transition:transform .3s ease;will-change:transform;contain:layout paint}
.trans.on{transform:translateY(0)}
.trans .th{display:flex;justify-content:space-between;padding-bottom:10px;border-bottom:1px dashed var(--bd)}
.trans .tt{font-style:italic;font-size:.95rem;display:flex;gap:10px;align-items:center}
.trans .cnt{font-family:var(--mono);font-size:.68rem;color:var(--mu)}
#call .chat:empty::before{content:'no turns yet · say something';display:block;text-align:center;color:var(--bd);font-style:italic;padding:40px 0;font-size:.82rem}
@media(max-width:540px){
body{display:block}
.app{height:100vh;height:100dvh;max-width:none;border:none;padding:14px 16px}
.brand img{height:28px}
.brand h1{font-size:1rem}
.brand .sub{display:none}
.topbar{gap:6px}
.status{display:none!important}
.player{width:100%}
.vitem,.vadd{padding:8px 12px}
.compose button{min-width:38px;padding:0 10px}
.text-wrap input{padding:8px 4px;font-size:.85rem}
.clone-card{padding:16px}
.clone-hint{font-size:.7rem}
#call .app{margin:0;height:100vh;height:100dvh;border:none;padding:14px 16px}
.orb{width:240px;height:240px}
.cbtn{width:58px;height:52px}
}
</style>
</head>
<body>
<div class="app">
<div class="topbar">
<div class="brand">
<img src="https://www.modelscope.cn/studio/gongjy/MiniMind/resolve/master/images/logo2.png" alt="logo">
<div><select class="model-sel" id="modelSel"></select><div class="sub">multimodal inference · text / image / audio → text + audio</div></div>
</div>
<div class="topright">
<div class="status hidden" id="stat"><div class="row"><span class="lab">text_ttft</span><span id="tt">—</span></div><div class="row"><span class="lab">audio_ttft</span><span id="ta">—</span></div><div class="row"><span class="lab">status</span><span id="ts">idle</span></div></div>
<select class="voice-sel" id="voiceSel"><option value="default">default</option></select>
<div class="more">
<button class="btn" id="moreBtn">default</button>
<div class="pop more-pop" id="morePop"></div>
</div>
<button class="btn" id="callBtn"><svg viewBox="0 0 24 24"><path d="M6.62 10.79a15.05 15.05 0 0 0 6.59 6.59l2.2-2.2a1 1 0 0 1 1.02-.24c1.12.37 2.32.57 3.57.57a1 1 0 0 1 1 1V20a1 1 0 0 1-1 1A17 17 0 0 1 3 4a1 1 0 0 1 1-1h3.5a1 1 0 0 1 1 1c0 1.25.2 2.45.57 3.57.1.35 0 .74-.25 1.02l-2.2 2.2z"/></svg>Call</button>
</div>
</div>
<div class="chat" id="chat"></div>
<div class="input-area"><div class="compose">
<div style="position:relative;display:flex">
<button id="imgBtn" title="image"><svg viewBox="0 0 24 24"><path d="M21 19V5a2 2 0 0 0-2-2H5a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2zm-2 0H5l3.5-4.5 2.5 3.01L14.5 13l4.5 6zM8.5 8A1.5 1.5 0 1 1 7 9.5 1.5 1.5 0 0 1 8.5 8z"/></svg></button>
<div class="pop" id="imgPop" style="bottom:100%;left:0;margin-bottom:6px"><button id="fileBtn">from file</button><button id="shotBtn">take photo</button></div>
</div>
<input type="file" id="imgIn" accept="image/*" hidden>
<button id="recBtn" title="录音"><svg viewBox="0 0 24 24"><path d="M12 14c1.66 0 3-1.34 3-3V5c0-1.66-1.34-3-3-3S9 3.34 9 5v6c0 1.66 1.34 3 3 3zm5.91-3c-.49 0-.9.36-.98.85C16.52 14.2 14.47 16 12 16s-4.52-1.8-4.93-4.15c-.08-.49-.49-.85-.98-.85-.61 0-1.09.54-1 1.14.49 3 2.89 5.35 5.91 5.78V20c0 .55.45 1 1 1s1-.45 1-1v-2.08c3.02-.43 5.42-2.78 5.91-5.78.1-.6-.39-1.14-1-1.14z"/></svg></button>
<div class="text-wrap"><div class="attach" id="att"></div><input type="text" id="txt" placeholder="Message MiniMind-O..." autocomplete="off"></div>
<button id="sendBtn" class="primary"><svg viewBox="0 0 24 24"><path d="M2.01 21L23 12 2.01 3 2 10l15 2-15 2z"/></svg></button>
</div></div>
</div>
<div class="snap" id="snap"><video id="snapV" autoplay playsinline muted></video><div class="sr"><button id="snapCancel">cancel</button><button id="snapOk" class="primary">capture</button></div></div>
<div class="snap" id="cloneModal"><div class="clone-card"><div class="clone-title">Clone Voice</div><div class="clone-hint">请自然朗读下面这句话。建议 3 到 6 秒,环境尽量安静。录完后会自动做 1.5 倍速、提取 `CAM++` speaker embedding 与 `Mimi` ref codes,然后立即加入音色列表。</div><input class="clone-name" id="cloneName" placeholder="音色名,例如:小王" maxlength="24"><div class="clone-script" id="cloneText">今天阳光很好,欢迎来到我的语音克隆小实验。</div><div class="clone-stat" id="cloneStat">点击 start 开始录音</div><div class="sr"><button id="cloneCancel">cancel</button><button id="cloneRec" class="primary">start</button></div></div></div>
<div id="call"><div class="app">
<div class="topbar">
<div class="brand"><img src="https://www.modelscope.cn/studio/gongjy/MiniMind/resolve/master/images/logo2.png" alt=""><div><select class="model-sel" id="callModelSel"></select><div class="sub">live voice · bidirectional streaming</div></div></div>
<div class="topright">
<select class="mini-sel" id="callVoiceSel"><option value="default">default</option></select>
<button class="btn" id="logBtn"><svg viewBox="0 0 24 24"><path d="M14 2H6c-1.1 0-1.99.9-1.99 2L4 20c0 1.1.89 2 1.99 2H18c1.1 0 2-.9 2-2V8l-6-6zM6 20V4h7v5h5v11H6zm2-3h8v-2H8v2zm0-4h8v-2H8v2z"/></svg>log</button>
<span class="pill" id="pill"><span class="dot"></span><span id="connT">connecting</span></span>
</div>
</div>
<div class="hero"><div class="orb" id="orb">
<svg viewBox="-160 -160 320 320" id="osvg"><g id="rip"></g><path class="blob l0" id="b0"/><path class="blob l1" id="b1"/><path class="blob l2" id="b2"/><circle class="core" id="core" cx="0" cy="0" r="4"/></svg>
<video class="cam" id="camV" autoplay playsinline muted></video>
<div class="lbl" id="lbl"><div class="ic" id="stIc"></div><div class="st" id="stT">init</div><div class="up" id="upT">00:00:00</div></div>
</div></div>
<div class="latest empty" id="lat"><span class="lm" id="lmE">latest · awaiting first turn</span><span class="lt" id="ltE">speak naturally — model responds with voice and text</span></div>
<div class="ctrls">
<button class="cbtn" id="camBtn"><svg viewBox="0 0 24 24"><path d="M17 10.5V7c0-.55-.45-1-1-1H4c-.55 0-1 .45-1 1v10c0 .55.45 1 1 1h12c.55 0 1-.45 1-1v-3.5l4 4v-11l-4 4z"/></svg><span>cam</span></button>
<button class="cbtn danger" id="endBtn"><svg viewBox="0 0 24 24"><path d="M12 9c-1.6 0-3.15.25-4.6.72v3.1c0 .39-.23.74-.56.9-.98.49-1.87 1.12-2.66 1.85-.18.18-.43.28-.7.28-.28 0-.53-.11-.71-.29L.29 13.08c-.18-.17-.29-.42-.29-.7 0-.28.11-.53.29-.71C3.34 8.78 7.46 7 12 7s8.66 1.78 11.71 4.67c.18.18.29.43.29.71 0 .28-.11.53-.29.71l-2.48 2.48c-.18.18-.43.29-.71.29-.27 0-.52-.11-.7-.28-.79-.74-1.69-1.36-2.67-1.85-.33-.16-.56-.5-.56-.9v-3.1C15.15 9.25 13.6 9 12 9z"/></svg><span>end</span></button>
</div>
<div class="trans" id="trans"><div class="th"><div class="tt">session transcript <span class="cnt" id="tcnt">0 turns</span></div><button class="btn" id="closeBtn">× close</button></div><div class="chat" id="cchat"></div></div>
</div></div>
<script>
const $=id=>document.getElementById(id);
const chat=$('chat'),txt=$('txt'),sendBtn=$('sendBtn'),recBtn=$('recBtn'),imgBtn=$('imgBtn'),imgIn=$('imgIn'),callBtn=$('callBtn'),att=$('att');
const stat=$('stat'),tt=$('tt'),ta=$('ta'),ts=$('ts'),voiceSel=$('voiceSel'),modelSel=$('modelSel'),callModelSel=$('callModelSel'),callVoiceSel=$('callVoiceSel');
const moreBtn=$('moreBtn'),morePop=$('morePop');
const CLONE_PROMPT='今天阳光很好,欢迎来到我的语音克隆小实验。';
let voiceGroups={builtin:[],unseen:[],manual:[]};
const closeMore=()=>morePop.classList.remove('on');
const modelLabel=m=>m;
async function loadVoices(sel){
const keep=sel||voiceSel.value||'default';
const keepCall=callVoiceSel.value||'default';
voiceSel.innerHTML='<option value="default">default</option>';
const d=await fetch('/voices').then(r=>r.json());
voiceGroups={builtin:d.builtin||[],unseen:d.unseen||[],manual:d.manual||[]};
for(const [label,names] of [['builtin',voiceGroups.builtin],['unseen',voiceGroups.unseen],['manual',voiceGroups.manual]]){
if(!names.length)continue;
const g=document.createElement('optgroup');
g.label=label;
names.forEach(v=>{const o=document.createElement('option');o.value=v;o.textContent=v;g.appendChild(o)});
voiceSel.appendChild(g);
}
if([...voiceSel.options].some(o=>o.value===keep))voiceSel.value=keep;
callVoiceSel.innerHTML=voiceSel.innerHTML;
if([...callVoiceSel.options].some(o=>o.value===keepCall))callVoiceSel.value=keepCall;
renderVoiceList();
}
function renderVoiceList(){
let h='';const sel=voiceSel.value;
moreBtn.textContent=sel;
h+=`<div class="vitem${sel==='default'?' on':''}" data-voice="default">default</div>`;
for(const [g,ns] of [['builtin',voiceGroups.builtin],['unseen',voiceGroups.unseen],['manual',voiceGroups.manual]]){
if(!ns.length)continue;
h+=`<div class="vgroup">${g}</div>`;
for(const v of ns)h+=`<div class="vitem${v===sel?' on':''}" data-voice="${v}"><span>${v}</span>${g==='manual'?`<span class="vx" data-del="${v}">×</span>`:''}</div>`;
}
h+='<div class="vadd">+ clone</div>';
morePop.innerHTML=h;
}
loadVoices().catch(()=>{});
async function loadModels(){
const d=await fetch('/models').then(r=>r.json());
modelSel.innerHTML='';callModelSel.innerHTML='';
for(const m of d.models||[]){const o=document.createElement('option');o.value=m;o.textContent=modelLabel(m);modelSel.appendChild(o);callModelSel.appendChild(o.cloneNode(true))}
if(d.current){modelSel.value=d.current;callModelSel.value=d.current}
}
async function switchModelSelect(sel, oldSel){
const old=sel.dataset.current||sel.value,name=sel.value;
const wasCall=document.body.classList.contains('call');
stopChat();if(callStop)callStop();stat.classList.remove('hidden');ts.textContent='switching';sel.disabled=true;
try{
const r=await fetch('/switch_model',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({name})});
const d=await r.json();
if(!r.ok||!d.ok)throw new Error(d.error||'switch failed');
modelSel.value=name;callModelSel.value=name;modelSel.dataset.current=name;callModelSel.dataset.current=name;ts.textContent=`loaded ${name}`;if(wasCall)openCall();
}catch(e){sel.value=old;if(oldSel)oldSel.value=old;ts.textContent='switch failed';alert(e.message)}
sel.disabled=false;
}
modelSel.onchange=()=>switchModelSelect(modelSel,callModelSel);
callModelSel.onchange=()=>switchModelSelect(callModelSel,modelSel);
callVoiceSel.onchange=()=>{if(window.callSendCtx)window.callSendCtx()};
loadModels().then(()=>{modelSel.dataset.current=modelSel.value;callModelSel.dataset.current=callModelSel.value}).catch(()=>{});
/* ------------- shared helpers ------------- */
function decPcm(b64){
const bin=atob(b64),i16=new Int16Array(bin.length/2);
for(let i=0;i<i16.length;i++)i16[i]=bin.charCodeAt(i*2)|(bin.charCodeAt(i*2+1)<<8);
const f32=new Float32Array(i16.length),fade=240;
for(let i=0;i<i16.length;i++)f32[i]=i16[i]/32768;
for(let i=0;i<fade&&i<f32.length;i++){f32[i]*=i/fade;f32[f32.length-1-i]*=i/fade;}
return {i16,f32};
}
function buildWav(chunks,sr){
let n=0;for(const c of chunks)n+=c.length;
const buf=new ArrayBuffer(44+n*2),v=new DataView(buf),w=(o,s)=>{for(let i=0;i<s.length;i++)v.setUint8(o+i,s.charCodeAt(i))};
w(0,'RIFF');v.setUint32(4,36+n*2,true);w(8,'WAVE');w(12,'fmt ');v.setUint32(16,16,true);v.setUint16(20,1,true);v.setUint16(22,1,true);
v.setUint32(24,sr,true);v.setUint32(28,sr*2,true);v.setUint16(32,2,true);v.setUint16(34,16,true);w(36,'data');v.setUint32(40,n*2,true);
let o=44;for(const c of chunks)for(let i=0;i<c.length;i++){v.setInt16(o,c[i],true);o+=2;}
return new Blob([buf],{type:'audio/wav'});
}
const blobToB64=blob=>new Promise(r=>{const fr=new FileReader();fr.onloadend=()=>r(fr.result.split(',')[1]);fr.readAsDataURL(blob)});
const fmtT=s=>{if(!isFinite(s)||s<0)s=0;s=Math.floor(s);return `${Math.floor(s/60)}:${String(s%60).padStart(2,'0')}`};
const fmtDur=s=>{s=Math.max(0,Math.floor(s));const h=s/3600|0,m=(s%3600/60)|0,x=s%60;return [h,m,x].map(n=>String(n).padStart(2,'0')).join(':')};
const waveChars=n=>{const c='▁▂▃▄▅▆▇';let s='';for(let i=0;i<n;i++){const v=(Math.sin(i*.7)+Math.sin(i*.31+1.3)+Math.sin(i*1.9+.4))/3;s+=c[Math.max(0,Math.min(c.length-1,Math.floor((v+1)/2*c.length)))];}return s};
function attachProg(audio,el,durHint,initPlayed){
const raw=el.textContent;if(!raw)return;
el.innerHTML='';el.classList.add('prog');
const chars=[...raw],wts=chars.map(c=>/\s/.test(c)?.2:/[,、;;::,]/.test(c)?.5:/[.。!!??…—]/.test(c)?1:1);
const cum=[0];for(const w of wts)cum.push(cum[cum.length-1]+w);
const tot=cum[cum.length-1]||1;
const spans=chars.map(c=>{const s=document.createElement('span');s.className='tk';s.textContent=c;el.appendChild(s);return s});
if(initPlayed)spans.forEach(s=>s.classList.add('played'));
const dur=()=>isFinite(audio.duration)&&audio.duration>0?audio.duration:(durHint||0);
const tOf=i=>dur()*cum[i]/tot;
const findCur=t=>{const d=dur();if(d<=0)return -1;const tg=t/d*tot;let lo=0,hi=spans.length-1;while(lo<hi){const m=(lo+hi+1)>>1;cum[m]<=tg?lo=m:hi=m-1}return lo};
spans.forEach((s,i)=>s.onclick=()=>{if(!audio.paused&&s.classList.contains('cur')){audio.pause();return}audio.currentTime=tOf(i);if(audio.paused)audio.play()});
const upd=()=>{const t=audio.currentTime,ci=!audio.paused?findCur(t):-1;spans.forEach((s,i)=>{s.classList.toggle('played',tOf(i)<=t);s.classList.toggle('cur',i===ci)});el.classList.toggle('on',!audio.paused)};
['timeupdate','seeked','play','pause'].forEach(e=>audio.addEventListener(e,upd));
}
function mkPlayer(src){
const w=document.createElement('div');w.className='player';
w.innerHTML='<button class="pp">▶</button><div class="bar-wrap"><div class="bar"></div></div><span class="time">0:00 / 0:00</span>';
const a=new Audio(src),pp=w.querySelector('.pp'),bw=w.querySelector('.bar-wrap'),b=w.querySelector('.bar'),tm=w.querySelector('.time');
pp.onclick=()=>a.paused?a.play():a.pause();
a.onplay=()=>pp.textContent='❚❚';a.onpause=()=>pp.textContent='▶';
a.ontimeupdate=()=>{const d=a.duration||0;b.style.width=d?(a.currentTime/d*100)+'%':'0';tm.textContent=`${fmtT(a.currentTime)} / ${fmtT(d)}`};
a.onended=()=>b.style.width='0';a.onloadedmetadata=()=>tm.textContent=`0:00 / ${fmtT(a.duration)}`;
bw.onclick=e=>{const r=bw.getBoundingClientRect();if(a.duration)a.currentTime=(e.clientX-r.left)/r.width*a.duration};
return w;
}
let msgN=0;
function addMsg(root,role,opts={}){
const d=document.createElement('div');d.className='msg '+role;msgN++;
const ts=new Date().toLocaleTimeString('en-US',{hour12:false});
d.innerHTML=`<div class="meta">[${msgN}] ${role} @ ${ts}</div><div class="body"></div>`;
const body=d.querySelector('.body');
if(opts.image){const i=document.createElement('img');i.className='img';i.src=opts.image;body.appendChild(i)}
if(opts.audioUrl){
if(role==='user'&&opts.audioDuration){
const n=Math.max(14,Math.min(80,Math.round(opts.audioDuration/.15)));
const v=document.createElement('div');v.className='text voice';v.textContent=waveChars(n);body.appendChild(v);
attachProg(new Audio(opts.audioUrl),v,opts.audioDuration);
} else body.appendChild(mkPlayer(opts.audioUrl));
}
if(opts.text!==undefined){const t=document.createElement('div');t.className='text';t.textContent=opts.text;body.appendChild(t)}
root.appendChild(d);root.scrollTop=root.scrollHeight;
return d;
}
/* ------------- Chat mode ------------- */
let chatHist=[],pendImg=null,pendAud=null,rec=null,recChunks=[],recording=false,gen=false,actx=null,nextT=0,aborter=null,chatPlaying=[],chatEpoch=0;
function stopChatAudio(){for(const s of chatPlaying)try{s.stop()}catch{}chatPlaying=[];if(actx){actx.close().catch(()=>{});actx=null;nextT=0}}
function stopChat(){chatEpoch++;if(aborter)try{aborter.abort()}catch{}stopChatAudio();gen=false;sendBtn.disabled=false;ts.textContent='idle'}
function renderChips(){
att.innerHTML='';txt.disabled=!!pendAud;txt.placeholder=pendAud?'(voice attached — text disabled)':'Message MiniMind-O...';
if(pendImg){const c=document.createElement('span');c.className='chip';c.innerHTML=`<img src="${pendImg.url}"><span>image</span><span class="x">×</span>`;c.querySelector('.x').onclick=()=>{URL.revokeObjectURL(pendImg.url);pendImg=null;renderChips()};att.appendChild(c)}
if(pendAud){const c=document.createElement('span');c.className='chip';c.innerHTML=`<span class="play">▶ voice · ${pendAud.duration.toFixed(1)}s</span><span class="x">×</span>`;c.querySelector('.play').onclick=()=>new Audio(pendAud.url).play();c.querySelector('.x').onclick=()=>{URL.revokeObjectURL(pendAud.url);pendAud=null;renderChips()};att.appendChild(c)}
}
async function send(){
if(gen)return;
stopChatAudio();
const t=txt.value.trim();if(!t&&!pendImg&&!pendAud)return;
const o={};if(pendImg)o.image=pendImg.url;if(pendAud){o.audioUrl=pendAud.url;o.audioDuration=pendAud.duration}if(t)o.text=t;
addMsg(chat,'user',o);
const payload={text:t||'',audio:pendAud?.b64||null,image:pendImg?.b64||null,temperature:.7,max_tokens:512,history:chatHist,voice:voiceSel.value};
txt.value='';pendImg=null;pendAud=null;renderChips();
sendBtn.disabled=true;gen=true;stat.classList.remove('hidden');tt.textContent='—';ta.textContent='—';ts.textContent='sending';
const ad=addMsg(chat,'assistant',{text:'…'}),ab=ad.querySelector('.body'),at=ab.querySelector('.text');
let full='',segs=[],acc='';
try{
const myEpoch=++chatEpoch;
aborter=new AbortController();
const res=await fetch('/chat',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload),signal:aborter.signal});
const rd=res.body.getReader(),dec=new TextDecoder();ts.textContent='streaming';
while(true){
if(myEpoch!==chatEpoch)throw new DOMException('chat stopped','AbortError');
const {done,value}=await rd.read();if(done)break;
for(const ln of dec.decode(value).split('\n')){
if(myEpoch!==chatEpoch)throw new DOMException('chat stopped','AbortError');
if(!ln.startsWith('data: '))continue;
let d;try{d=JSON.parse(ln.slice(6))}catch{continue}
if(d.type==='user_prompt')chatHist.push({role:'user',content:d.content});
else if(d.type==='ttft'){if(d.text_ttft!==undefined)tt.textContent=d.text_ttft+' ms';if(d.audio_ttft!==undefined)ta.textContent=d.audio_ttft+' ms'}
else if(d.type==='text'){full+=d.content;at.textContent=full;chat.scrollTop=chat.scrollHeight}
else if(d.type==='pcm'){
acc+=d.c;
if(d.d){
if(!actx){actx=new AudioContext({sampleRate:24000});nextT=actx.currentTime+.1}
const {i16,f32}=decPcm(acc);
const buf=actx.createBuffer(1,f32.length,24000);buf.getChannelData(0).set(f32);
const s=actx.createBufferSource();s.buffer=buf;s.connect(actx.destination);
s.onended=()=>{chatPlaying=chatPlaying.filter(x=>x!==s)};
chatPlaying.push(s);
if(nextT<actx.currentTime)nextT=actx.currentTime;
s.start(nextT);nextT+=buf.duration-.01;segs.push(i16);acc='';
}
} else if(d.type==='error')at.textContent='⚠ '+d.content;
}
}
if(segs.length){const url=URL.createObjectURL(buildWav(segs,24000));attachProg(new Audio(url),at,null,true)}
if(full)chatHist.push({role:'assistant',content:full});
ts.textContent='done';
} catch(e){if(e.name!=='AbortError'){at.textContent='⚠ '+e.message;ts.textContent='error'}}
gen=false;sendBtn.disabled=false;aborter=null;
}
sendBtn.onclick=send;
txt.addEventListener('keydown',e=>{if(e.key==='Enter'&&!e.shiftKey){e.preventDefault();send()}});
const imgPop=$('imgPop'),fileBtn=$('fileBtn'),shotBtn=$('shotBtn'),snap=$('snap'),snapV=$('snapV'),snapOk=$('snapOk'),snapCancel=$('snapCancel');
const cloneModal=$('cloneModal'),cloneText=$('cloneText'),cloneStat=$('cloneStat'),cloneRecBtn=$('cloneRec'),cloneCancelBtn=$('cloneCancel'),cloneName=$('cloneName');
cloneText.textContent=CLONE_PROMPT;
moreBtn.onclick=e=>{e.stopPropagation();morePop.classList.toggle('on')};
morePop.onclick=e=>{
e.stopPropagation();
const del=e.target.closest('.vx');
if(del){const name=del.dataset.del;if(!confirm('确认删除 '+name+'?'))return;fetch('/delete_voice',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({name})}).then(r=>r.json()).then(d=>{if(!d.ok)throw new Error(d.error);loadVoices(voiceSel.value===name?'default':voiceSel.value)}).catch(e=>alert(e.message));return}
const add=e.target.closest('.vadd');
if(add){closeMore();stopChat();cloneModal.classList.add('on');if(!cloneName.value.trim())cloneName.value='voice_clone';cloneStat.textContent='请自然读完,不用刻意加速';return}
const item=e.target.closest('.vitem');
if(item&&item.dataset.voice){voiceSel.value=item.dataset.voice;renderVoiceList();closeMore()}
};
imgBtn.onclick=e=>{e.stopPropagation();imgPop.classList.toggle('on')};
document.addEventListener('click',e=>{if(!imgPop.contains(e.target)&&e.target!==imgBtn)imgPop.classList.remove('on');if(!morePop.contains(e.target)&&e.target!==moreBtn)closeMore()});
fileBtn.onclick=()=>{imgPop.classList.remove('on');imgIn.click()};
const setImg=(b64,url)=>{if(pendImg)URL.revokeObjectURL(pendImg.url);pendImg={b64,url};renderChips()};
imgIn.onchange=()=>{const f=imgIn.files?.[0];if(!f)return;const r=new FileReader();r.onloadend=()=>setImg(r.result.split(',')[1],URL.createObjectURL(f));r.readAsDataURL(f);imgIn.value=''};
let snapStream=null;
const closeSnap=()=>{snap.classList.remove('on');if(snapStream){snapStream.getTracks().forEach(t=>t.stop());snapStream=null}snapV.srcObject=null};
shotBtn.onclick=async()=>{
imgPop.classList.remove('on');
try{snapStream=await navigator.mediaDevices.getUserMedia({video:{facingMode:'user',width:1280,height:720}});snapV.srcObject=snapStream;snap.classList.add('on')}
catch(e){alert('camera error: '+e.message)}
};
snapCancel.onclick=closeSnap;
snapOk.onclick=()=>{
if(!snapV.videoWidth)return;
const cv=document.createElement('canvas');cv.width=snapV.videoWidth;cv.height=snapV.videoHeight;cv.getContext('2d').drawImage(snapV,0,0);
const dataUrl=cv.toDataURL('image/jpeg',.9),b64=dataUrl.split(',')[1];
cv.toBlob(blob=>setImg(b64,URL.createObjectURL(blob)),'image/jpeg',.9);
closeSnap();
};
const CLONE_MAX_SECS=6;
let cloneRec=null,cloneChunks=[],cloneRecording=false,cloneBusy=false,cloneCanceled=false,cloneStream=null,cloneDeadline=0,cloneTimer=null;
const stopCloneTimer=()=>{if(cloneTimer){clearInterval(cloneTimer);cloneTimer=null}};
const resetClone=()=>{stopCloneTimer();if(cloneStream){cloneStream.getTracks().forEach(t=>t.stop());cloneStream=null}cloneRec=null;cloneChunks=[];cloneRecording=false;cloneBusy=false;cloneCanceled=false;cloneDeadline=0;cloneRecBtn.classList.remove('rec');cloneRecBtn.textContent='start';cloneRecBtn.disabled=false;cloneCancelBtn.disabled=false};
const hideClone=()=>{cloneModal.classList.remove('on');resetClone();cloneStat.textContent='点击 start 开始录音'};
cloneCancelBtn.onclick=()=>{if(cloneBusy)return;if(cloneRecording&&cloneRec){cloneCanceled=true;cloneStat.textContent='已取消';cloneRec.stop()}else hideClone()};
cloneRecBtn.onclick=async()=>{
if(cloneBusy)return;
if(!cloneRecording){
try{
cloneStream=await navigator.mediaDevices.getUserMedia({audio:true});
cloneRec=new MediaRecorder(cloneStream);cloneChunks=[];cloneCanceled=false;
cloneRec.ondataavailable=e=>cloneChunks.push(e.data);
cloneRec.onstop=async()=>{
const canceled=cloneCanceled,blob=new Blob(cloneChunks,{type:(cloneRec&&cloneRec.mimeType)||'audio/webm'});
resetClone();
if(canceled){cloneModal.classList.remove('on');cloneStat.textContent='点击 start 开始录音';return}
cloneBusy=true;cloneRecBtn.disabled=true;cloneCancelBtn.disabled=true;cloneStat.textContent='正在提取音色…';
try{
const audio=await blobToB64(blob);
const res=await fetch('/clone_voice',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({audio,name:cloneName.value.trim()})});
const d=await res.json();
if(!res.ok||!d.ok)throw new Error(d.error||'clone failed');
await loadVoices(d.voice);
voiceSel.value=d.voice;
cloneStat.textContent='音色已生成,已自动切换到 '+d.voice;
setTimeout(()=>hideClone(),600);
}catch(e){cloneBusy=false;cloneRecBtn.disabled=false;cloneCancelBtn.disabled=false;cloneStat.textContent='失败:'+e.message}
};
cloneRec.start();cloneRecording=true;cloneRecBtn.classList.add('rec');cloneRecBtn.textContent='stop';cloneDeadline=Date.now()+CLONE_MAX_SECS*1000;
const tick=()=>{
const left=Math.max(0,Math.ceil((cloneDeadline-Date.now())/1000));
cloneStat.textContent=`正在录音,请读完这句话 · ${left}s`;
if(left<=0&&cloneRecording&&cloneRec){cloneBusy=true;cloneRecording=false;cloneRecBtn.disabled=true;cloneRecBtn.classList.remove('rec');cloneRecBtn.textContent='start';cloneStat.textContent='录音结束,正在上传…';cloneRec.stop()}
};
tick();
cloneTimer=setInterval(tick,200);
}catch(e){alert('mic error: '+e.message)}
}else{
stopCloneTimer();cloneBusy=true;cloneRecording=false;cloneRecBtn.disabled=true;cloneRecBtn.classList.remove('rec');cloneRecBtn.textContent='start';cloneStat.textContent='录音结束,正在上传…';cloneRec.stop();
}
};
recBtn.onclick=async()=>{
if(!recording){
try{
const st=await navigator.mediaDevices.getUserMedia({audio:true});
rec=new MediaRecorder(st);recChunks=[];const t0=Date.now();
rec.ondataavailable=e=>recChunks.push(e.data);
rec.onstop=()=>{
st.getTracks().forEach(t=>t.stop());
const blob=new Blob(recChunks,{type:'audio/webm'}),dur=(Date.now()-t0)/1000,r=new FileReader();
r.onloadend=()=>{const b64=r.result.split(',')[1],url=URL.createObjectURL(blob);if(pendAud)URL.revokeObjectURL(pendAud.url);pendAud={b64,url,duration:dur};renderChips()};
r.readAsDataURL(blob);
};
rec.start();recording=true;recBtn.classList.add('rec');
} catch(e){alert('mic error: '+e.message)}
} else {rec.stop();recording=false;recBtn.classList.remove('rec')}
};
/* ------------- Call mode ------------- */
let callInit=false,callStop=null;
callBtn.onclick=openCall;
function openCall(){stopChat();document.body.classList.add('call');history.pushState({call:1},'','/call');if(!callInit){callInit=true;initCall()}}
function closeCall(){document.body.classList.remove('call');if(callStop)callStop();if(location.pathname==='/call')history.pushState({},'','/')}
window.addEventListener('popstate',()=>location.pathname==='/call'?openCall():closeCall());
if(location.pathname==='/call')openCall();
function initCall(){
const stT=$('stT'),upT=$('upT'),stIc=$('stIc'),lbl=$('lbl'),osvg=$('osvg'),core=$('core'),ripG=$('rip');
const ICONS={
idle:'<svg viewBox="0 0 24 24"><circle class="dot" cx="12" cy="12" r="3" fill="currentColor" stroke="none"/></svg>',
listen:'<svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="2.2" fill="currentColor" stroke="none"/><path class="w w3" d="M7 8.5a6 6 0 0 0 0 7"/><path class="w w3" d="M17 8.5a6 6 0 0 1 0 7"/><path class="w w2" d="M4 6a9.5 9.5 0 0 0 0 12"/><path class="w w2" d="M20 6a9.5 9.5 0 0 1 0 12"/></svg>',
gen:'<svg viewBox="0 0 24 24"><circle cx="12" cy="4" r="1.8" fill="currentColor" stroke="none"/><circle cx="18.93" cy="15.5" r="1.4" fill="currentColor" stroke="none" opacity=".65"/><circle cx="5.07" cy="15.5" r="1" fill="currentColor" stroke="none" opacity=".35"/></svg>',
speak:'<svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="2.2" fill="currentColor" stroke="none"/><path class="a1" d="M8 9a5 5 0 0 1 0 6"/><path class="a1" d="M16 9a5 5 0 0 0 0 6"/><path class="a2" d="M5 6.5a9 9 0 0 1 0 11"/><path class="a2" d="M19 6.5a9 9 0 0 0 0 11"/><path class="a3" d="M2 4a13 13 0 0 1 0 16"/><path class="a3" d="M22 4a13 13 0 0 0 0 16"/></svg>',
err:'<svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="9"/><line x1="8.5" y1="8.5" x2="15.5" y2="15.5"/><line x1="15.5" y1="8.5" x2="8.5" y2="15.5"/></svg>',
};
const cchat=$('cchat'),lat=$('lat'),lmE=$('lmE'),ltE=$('ltE');
const camBtn=$('camBtn'),logBtn=$('logBtn'),endBtn=$('endBtn'),orb=$('orb'),camV=$('camV');
const pill=$('pill'),connT=$('connT'),tcnt=$('tcnt'),trans=$('trans'),closeBtn=$('closeBtn');
const blobs=[$('b0'),$('b1'),$('b2')];
let ac,ana,mic,camStream=null,camOn=false,hist=[],nextT=0;
let ws=null,wsActive=false,proc=null,rtc=null,rtSrc=null,playing=[],turns=0,mn=0;
let curMsg=null,curText='',segs=[],start=0,ripT=0,sVol=0,asstPlay=false,raf=0;
let spk=false,userBuf=[];
const flushUser=()=>{
if(!userBuf.length)return;
const dur=userBuf.reduce((a,c)=>a+c.length,0)/16000;
const url=URL.createObjectURL(buildWav(userBuf,16000));
userBuf=[];
const d=addCM('user'),body=d.querySelector('.body'),n=Math.max(14,Math.min(80,Math.round(dur/.15)));
const v=document.createElement('div');v.className='text voice';v.textContent=waveChars(n);body.appendChild(v);
attachProg(new Audio(url),v,dur);
setLatest('user',`[voice · ${dur.toFixed(1)}s]`);
};
const N=32,NC=2,NR=5;let shT=0;
const LAY=[{c:56,r:112,a:1.8,sp:.0005,ph:.6,vc:34,vr:18},{c:80,r:124,a:2.6,sp:.0008,ph:1.3,vc:52,vr:26},{c:104,r:138,a:3.6,sp:.0011,ph:2.1,vc:74,vr:34}];
const catmull=pts=>{let d=`M ${pts[0][0].toFixed(2)} ${pts[0][1].toFixed(2)} `;const n=pts.length;for(let i=0;i<n;i++){const p0=pts[(i-1+n)%n],p1=pts[i],p2=pts[(i+1)%n],p3=pts[(i+2)%n];const c1x=p1[0]+(p2[0]-p0[0])/6,c1y=p1[1]+(p2[1]-p0[1])/6,c2x=p2[0]-(p3[0]-p1[0])/6,c2y=p2[1]-(p3[1]-p1[1])/6;d+=`C ${c1x.toFixed(2)} ${c1y.toFixed(2)}, ${c2x.toFixed(2)} ${c2y.toFixed(2)}, ${p2[0].toFixed(2)} ${p2[1].toFixed(2)} `}return d+'Z'};
const setSt=(cls,t)=>{osvg.setAttribute('class',cls||'');lbl.className='lbl '+(cls||'');stT.textContent=t;const key=cls||(/error/i.test(t)?'err':'idle');stIc.innerHTML=ICONS[key]||ICONS.idle};
const setConn=(live,t)=>{pill.className='pill'+(live?' live':'');connT.textContent=t};
const ripple=k=>{const now=performance.now();if(now-ripT<550)return;ripT=now;const c=document.createElementNS('http://www.w3.org/2000/svg','circle');c.setAttribute('cx',0);c.setAttribute('cy',0);c.setAttribute('r',85);c.setAttribute('class','ripple '+k);ripG.appendChild(c);c.addEventListener('animationend',()=>c.remove())};
const upTimer=setInterval(()=>{if(start)upT.textContent=fmtDur((Date.now()-start)/1000)},1000);
const setLatest=(role,t)=>{if(!t)return;lat.classList.remove('empty','user','assistant');lat.classList.add(role);const ts=new Date().toLocaleTimeString('en-US',{hour12:false});lmE.textContent=`latest · ${role} @ ${ts}`;ltE.textContent=t;lat.scrollTop=0};
const addCM=role=>{const d=document.createElement('div');d.className='msg '+role;mn++;const ts=new Date().toLocaleTimeString('en-US',{hour12:false});d.innerHTML=`<div class="meta">[${mn}] ${role} @ ${ts}</div><div class="body"></div>`;cchat.appendChild(d);cchat.scrollTop=cchat.scrollHeight;return d};
const playPcm=b64=>{const {i16,f32}=decPcm(b64);segs.push(i16);const buf=ac.createBuffer(1,f32.length,24000);buf.getChannelData(0).set(f32);const s=ac.createBufferSource();s.buffer=buf;s.connect(ac.destination);s.onended=()=>{playing=playing.filter(x=>x!==s);if(!playing.length){asstPlay=false;if(!curMsg)setSt('','idle')}};playing.push(s);asstPlay=true;if(nextT<ac.currentTime)nextT=ac.currentTime;s.start(nextT);nextT+=buf.duration-.01};
const stopPcm=()=>{for(const s of playing)try{s.stop()}catch{}playing=[];asstPlay=false;if(ac)nextT=ac.currentTime};
const capCam=()=>{if(!camOn||!camV.videoWidth)return null;const cv=document.createElement('canvas');cv.width=camV.videoWidth;cv.height=camV.videoHeight;cv.getContext('2d').drawImage(camV,0,0);return cv.toDataURL('image/jpeg',.8).split(',')[1]};
const sendCtx=()=>{if(!ws||ws.readyState!==1)return;ws.send(JSON.stringify({type:'context',history:hist,image:camOn?capCam():null,voice:callVoiceSel.value}))};
window.callSendCtx=sendCtx;
function render(){
const now=performance.now();let vol=0,fq=null;
if(ana){const d=new Uint8Array(ana.frequencyBinCount);ana.getByteFrequencyData(d);fq=d;let s=0;for(let i=0;i<d.length;i++)s+=d[i];vol=s/d.length/255}
if(asstPlay)vol=Math.max(vol,.6+.25*Math.sin(now*.013));
sVol=sVol*.7+vol*.3;
const tgt=camOn?1:0;shT+=(tgt-shT)*.08;
const shN=NC+(NR-NC)*shT,inv=-1/shN;
for(let li=0;li<LAY.length;li++){
const L=LAY[li],hw=L.c+(L.r-L.c)*shT,hh=hw,vs=L.vc+(L.vr-L.vc)*shT;
const pts=new Array(N);
for(let i=0;i<N;i++){
const ang=(i/N)*Math.PI*2,cA=Math.cos(ang),sA=Math.sin(ang);
const n=Math.sin(now*L.sp+i*L.ph)+.55*Math.sin(now*L.sp*1.7+i*L.ph*.7+li);
let fv=0;if(fq){const idx=Math.floor((i/N)*Math.min(64,fq.length))+li*4;fv=(fq[idx%fq.length]||0)/255}
if(asstPlay)fv=Math.max(fv,.5+.45*Math.abs(Math.sin(now*.011+i*.55+li*.9)));
const g=Math.min(1,fv*.65+sVol*.55);
const tx=Math.pow(Math.abs(cA)/hw,shN),ty=Math.pow(Math.abs(sA)/hh,shN);
const r=Math.pow(tx+ty,inv)+n*L.a+g*vs;
pts[i]=[cA*r,sA*r];
}
blobs[li].setAttribute('d',catmull(pts));
}
core.setAttribute('r',(3+sVol*10).toFixed(1));
raf=requestAnimationFrame(render);
}
raf=requestAnimationFrame(render);
async function initMic(){
try{mic=await navigator.mediaDevices.getUserMedia({audio:true});ac=new AudioContext({sampleRate:24000});const s=ac.createMediaStreamSource(mic);ana=ac.createAnalyser();ana.fftSize=256;s.connect(ana);setSt('','connecting');setConn(false,'connecting')}
catch(e){setSt('','mic error');setConn(false,'mic error')}
}
camBtn.onclick=async()=>{
if(camOn){if(camStream){camStream.getTracks().forEach(t=>t.stop());camStream=null}camV.srcObject=null;orb.classList.remove('camon');camBtn.classList.remove('active');camOn=false}
else{try{camStream=await navigator.mediaDevices.getUserMedia({video:{facingMode:'user',width:640,height:480}});camV.srcObject=camStream;orb.classList.add('camon');camBtn.classList.add('active');camOn=true}catch{setSt('','cam error')}}
if(wsActive)sendCtx();
};
logBtn.onclick=()=>{trans.classList.toggle('on');logBtn.classList.toggle('active',trans.classList.contains('on'))};
closeBtn.onclick=()=>{trans.classList.remove('on');logBtn.classList.remove('active')};
endBtn.onclick=closeCall;
async function startRT(){
wsActive=true;if(!mic)await initMic();
rtc=new AudioContext({sampleRate:16000});rtSrc=rtc.createMediaStreamSource(mic);proc=rtc.createScriptProcessor(4096,1,1);
const pr=location.protocol==='https:'?'wss':'ws';
ws=new WebSocket(`${pr}://${location.host}/ws/realtime`);ws.binaryType='arraybuffer';
ws.onmessage=e=>{
const m=JSON.parse(e.data);
if(m.type==='vad'){
if(m.speaking&&!spk){sendCtx();stopPcm();ripple('user');userBuf=[]}
if(!m.speaking&&spk)flushUser();
spk=m.speaking;
if(m.speaking)setSt('listen','listening');else if(!asstPlay&&!curMsg)setSt('','idle');
} else if(m.type==='generating'){
setSt('gen','thinking');if(spk){flushUser();spk=false}
turns++;tcnt.textContent=turns+(turns===1?' turn':' turns');
curText='';segs=[];
const d=addCM('assistant'),body=d.querySelector('.body'),t=document.createElement('div');t.className='text';body.appendChild(t);
curMsg={div:d,body,textEl:t};stopPcm();
} else if(m.type==='text'){
curText+=m.content;if(curMsg)curMsg.textEl.textContent=curText;
setLatest('assistant',curText);cchat.scrollTop=cchat.scrollHeight;
} else if(m.type==='pcm'){setSt('speak','speaking');ripple('assistant');playPcm(m.data)}
else if(m.type==='done'){
if(curText)hist.push({role:'assistant',content:curText});
if(m.interrupted)stopPcm();
if(curMsg&&segs.length){const url=URL.createObjectURL(buildWav(segs,24000));attachProg(new Audio(url),curMsg.textEl,null,true)}
curMsg=null;if(!asstPlay)setSt('','idle');
}
};
ws.onclose=()=>{if(wsActive)stopRT()};
ws.onopen=()=>{
start=Date.now();setConn(true,'live');setSt('','idle');sendCtx();
proc.onaudioprocess=e=>{
if(ws.readyState!==1)return;
const f=e.inputBuffer.getChannelData(0),i16=new Int16Array(f.length);
for(let i=0;i<f.length;i++)i16[i]=Math.max(-32768,Math.min(32767,f[i]*32768));
ws.send(i16.buffer);
if(spk)userBuf.push(new Int16Array(i16));
};
rtSrc.connect(proc);proc.connect(rtc.destination);
};
}
function stopRT(){wsActive=false;if(ws){ws.close();ws=null}if(proc){proc.disconnect();proc.onaudioprocess=null;proc=null}if(rtSrc){rtSrc.disconnect();rtSrc=null}if(rtc){rtc.close();rtc=null}stopPcm();setSt('','disconnected');setConn(false,'offline')}
callStop=()=>{
stopRT();
window.callSendCtx=null;
if(camStream){camStream.getTracks().forEach(t=>t.stop());camStream=null;camOn=false;orb.classList.remove('camon');camV.srcObject=null;camBtn.classList.remove('active')}
if(mic){mic.getTracks().forEach(t=>t.stop());mic=null}
if(ac){try{ac.close()}catch{}ac=null}
if(raf)cancelAnimationFrame(raf);
clearInterval(upTimer);callInit=false;callStop=null;
};
initMic().then(startRT);
}
</script>
</body>
</html>

502
webui/web_demo.py Normal file
View File

@ -0,0 +1,502 @@
import argparse, os, sys, json, time, math, torch, threading, queue, base64, io, logging, contextlib
import numpy as np
import torchaudio
from flask import Flask, request, Response, send_from_directory
from flask_cors import CORS
from flask_sock import Sock
from PIL import Image
from pydub import AudioSegment
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from transformers import AutoTokenizer, AutoModelForCausalLM
from model.model_omni import MiniMindOmni, RealtimeSession
from trainer.trainer_utils import log_model_params
logging.getLogger().setLevel(logging.ERROR)
with contextlib.redirect_stdout(io.StringIO()):
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
app = Flask(__name__, static_folder='.')
CORS(app)
sock = Sock(app)
M = {} # model / tokenizer / device / mimi / asr / cfg
V = {} # voice_name -> {ref_codes, spk_emb}
V_builtin, V_unseen, V_manual = [], [], []
MODEL_LOCK = threading.Lock()
SAMPLES_PER_FRAME = 1920
REF_FRAMES = 300
CLONE_VOICE = 'voice_clone'
CLONE_FILE = 'voice_clone.pt'
# -------- helpers --------
def sse(d): return f"data: {json.dumps(d)}\n\n"
def scan_hf_models(base_dir):
models = {}
base_dir = os.path.abspath(base_dir)
for d in sorted(os.listdir(base_dir), reverse=True):
full_path = os.path.join(base_dir, d)
if not os.path.isdir(full_path) or d.startswith('.') or d.startswith('_'):
continue
files = set(os.listdir(full_path))
has_model = bool(files & {'pytorch_model.bin', 'model.safetensors', 'pytorch_model.bin.index.json', 'model.safetensors.index.json'})
if has_model:
models[d] = full_path
return models
def asr_run(samples):
r = M['asr'].generate(input=samples, cache={}, language='auto', use_itn=True)
return rich_transcription_postprocess(r[0]['text']).strip() if r else ''
def prep_audio(samples):
m = M['model']
proc = m.audio_processor(samples, sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
mel = proc.input_features.squeeze(0).unsqueeze(0).to(M['device'])
vlen = proc.attention_mask.sum().item()
prompt = m.config.audio_special_token * (vlen or 1)
return mel, torch.tensor([vlen], device=M['device']), prompt
def prep_image(b64):
img = Image.open(io.BytesIO(base64.b64decode(b64))).convert('RGB')
return {k: v.to(M['device']) for k, v in M['model'].vision_processor(images=img, return_tensors="pt").items()}
def build_ids(prompt, history):
tok, dev, n = M['tokenizer'], M['device'], M['cfg'].max_history_turns
hist = history[-n:] if n > 0 else []
msgs = hist + [{"role": "user", "content": prompt}]
t = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
return torch.tensor(tok(t).data['input_ids'], dtype=torch.long, device=dev)[None, ...]
def _mimi_decode(frames):
codes = [f for f in frames if f and len(f) == 8]
if not codes or not M['mimi']: return None
mc = torch.tensor(codes, dtype=torch.long).T.unsqueeze(0)
mc = torch.where(mc >= 2049, torch.zeros_like(mc), mc)
with torch.no_grad():
au = M['mimi'].decode(mc).audio_values.squeeze().cpu().numpy()
return au, mc.shape[-1]
def pcm_bytes(frames, ov):
r = _mimi_decode(frames)
if r is None: return None
au, T = r
if ov > 0: au = au[int(ov * len(au) / T):]
return (au * 32767).astype('int16').tobytes()
def stream_pcm(frames, flush=False):
"""yield (pcm_bytes,) on chunk boundaries or on final flush."""
if not M['mimi']: return
cf, ov_max, n = M['cfg'].audio_chunk_frames, M['cfg'].audio_overlap, len(frames)
if not flush and n >= cf and n % cf == 0:
ov = min(ov_max, n - cf)
p = pcm_bytes(frames[-(cf + ov):], ov)
if p: yield p
elif flush:
rem = n % cf
if rem:
ov = min(ov_max, n - rem)
p = pcm_bytes(frames[-(rem + ov):], ov)
if p: yield p
def voice_args(name):
if name and name != 'default' and name in V:
v = V[name]
dev = M['device']
rc = v['ref_codes'].unsqueeze(0).to(dev)
se = v['spk_emb'].half().unsqueeze(0).to(dev) if 'spk_emb' in v else None
return {'ref_codes': rc, 'spk_emb': se}
return {}
def register_voice(name, value, group='manual'):
V[name] = value
groups = {'builtin': V_builtin, 'unseen': V_unseen, 'manual': V_manual}
dst = groups[group]
if name not in dst:
dst.append(name)
for k, lst in groups.items():
if k != group and name in lst:
lst.remove(name)
def clone_voice_path():
return os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'speaker', CLONE_FILE)
def delete_manual_voice(name):
if name not in V_manual:
raise RuntimeError('只能删除手动克隆的音色')
out_path = clone_voice_path()
saved = torch.load(out_path, map_location='cpu') if os.path.exists(out_path) else {}
if name in saved:
saved.pop(name)
torch.save(saved, out_path)
V.pop(name, None)
if name in V_manual:
V_manual.remove(name)
def normalize_voice_name(name):
name = ' '.join(str(name or '').split())
if not name:
name = CLONE_VOICE
if len(name) > 24:
raise RuntimeError('音色名太长,建议控制在 24 个字以内')
if name.lower() == 'default':
raise RuntimeError('default 是保留名称,请换一个')
if name in V_builtin or name in V_unseen:
raise RuntimeError('该名称已被现有音色占用,请换一个')
return name
def validate_clone_audio(w16):
if w16.numel() < int(16000 * 1.8):
raise RuntimeError('录音太短,请把整句话读完')
peak = w16.abs().max().item()
frame, hop = 800, 400
if w16.numel() >= frame:
rms = w16.unfold(0, frame, hop).pow(2).mean(dim=1).sqrt().cpu().numpy()
else:
rms = np.array([w16.pow(2).mean().sqrt().item()])
hi = float(np.quantile(rms, 0.95))
lo = float(np.quantile(rms, 0.2))
if hi < 0.008:
raise RuntimeError('录音太轻,请靠近麦克风一点')
if hi > 0 and lo / hi > 0.45:
raise RuntimeError('环境噪声太大,请换安静一点的环境')
if peak > 0.995:
raise RuntimeError('录音有爆音,请离麦克风远一点')
def build_clone_voice(audio_b64):
if M.get('mimi') is None or M.get('campplus') is None or M.get('mel_fn') is None:
raise RuntimeError('Mimi 或 CAM++ 未加载')
seg = AudioSegment.from_file(io.BytesIO(base64.b64decode(audio_b64))).set_channels(1).set_sample_width(2)
if len(seg) < 1000:
raise RuntimeError('录音太短,至少读 1 秒')
try:
seg = seg.speedup(playback_speed=1.5, chunk_size=150, crossfade=25)
except Exception:
seg = seg.speedup(playback_speed=1.5)
seg24 = seg.set_frame_rate(24000)
seg16 = seg.set_frame_rate(16000)
w24 = torch.tensor(np.frombuffer(seg24.raw_data, dtype=np.int16).astype(np.float32) / 32768.0)
w16 = torch.tensor(np.frombuffer(seg16.raw_data, dtype=np.int16).astype(np.float32) / 32768.0)
validate_clone_audio(w16)
mimi_dev = next(M['mimi'].parameters()).device
mimi_dtype = torch.float16 if mimi_dev.type != 'cpu' else torch.float32
with torch.inference_mode():
t = w24.unsqueeze(0).unsqueeze(0).to(device=mimi_dev, dtype=mimi_dtype)
codes = M['mimi'].encode(t).audio_codes
nf = math.ceil(w24.shape[-1] / SAMPLES_PER_FRAME)
ref_codes = codes[0, :8, :nf].cpu()[:, :min(nf, REF_FRAMES)]
with torch.no_grad():
mel = M['mel_fn'](w16.unsqueeze(0).to(M['device']))
feat = mel.clamp(min=1e-10).log().transpose(1, 2)
feat = feat - feat.mean(dim=1, keepdim=True)
spk_emb = M['campplus'](feat).squeeze(0).cpu()
return {'ref_codes': ref_codes, 'spk_emb': spk_emb}
def run_generate(x, audio_inputs, audio_lens, pixel_values, **kw):
with MODEL_LOCK, torch.no_grad():
yield from M['model'].generate(
x, M['tokenizer'].eos_token_id, stream=True, return_audio_codes=True,
audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values, **kw)
def load_main_model(model_path, model_name):
with MODEL_LOCK:
[sys.modules.pop(k) for k in list(sys.modules) if 'transformers_modules' in k]
M.pop('model', None); M.pop('tokenizer', None)
if torch.cuda.is_available(): torch.cuda.empty_cache()
tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
m = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
vision_encoder, vision_processor = MiniMindOmni.load_vision('../model/siglip2-base-p32-256-ve')
audio_encoder, audio_processor = MiniMindOmni.load_sensevoice('../model/SenseVoiceSmall')
object.__setattr__(m, 'vision_encoder', vision_encoder)
object.__setattr__(m, 'vision_processor', vision_processor)
object.__setattr__(m, 'audio_encoder', audio_encoder)
object.__setattr__(m, 'audio_processor', audio_processor)
m = m.half().eval().to(M['device'])
if m.audio_encoder: m.audio_encoder.to(M['device'])
if m.vision_encoder: m.vision_encoder.to(M['device'])
M['tokenizer'], M['model'], M['model_name'] = tok, m, model_name
params = sum(p.numel() for p in m.parameters()) / 1e6
print(f'Loaded model: {model_name} ({params:.2f}M)')
return round(params, 2)
def prepare_turn(text, samples, image_b64, do_asr_for_image):
"""返回 (audio_inputs, audio_lens, pixel_values, prompt_for_model, user_text_for_history, asr_thread, asr_result)"""
audio_inputs = audio_lens = pixel_values = None
prompt = text or ''
user_text = text or ''
asr_thread, asr_result = None, [None]
if samples is not None:
if image_b64 and do_asr_for_image:
user_text = asr_run(samples)
prompt = user_text
else:
audio_inputs, audio_lens, prompt = prep_audio(samples)
if M['cfg'].max_history_turns > 0:
sa = samples.copy()
def _a(): asr_result[0] = asr_run(sa)
asr_thread = threading.Thread(target=_a); asr_thread.start()
if image_b64:
pixel_values = prep_image(image_b64)
m = M['model']
prompt = (prompt + "\n\n" if prompt else "") + "请描述这张图片\n\n" + m.config.image_special_token * m.config.image_token_len
return audio_inputs, audio_lens, pixel_values, prompt, user_text, asr_thread, asr_result
# -------- routes --------
@app.route('/')
def index(): return send_from_directory('.', 'web_demo.html')
@app.route('/call')
def call_page(): return send_from_directory('.', 'web_demo.html')
@app.route('/voices')
def get_voices():
return json.dumps({'builtin': sorted(V_builtin), 'unseen': sorted(V_unseen), 'manual': sorted(V_manual)})
@app.route('/models')
def get_models():
return json.dumps({'models': list(M.get('models', {}).keys()), 'current': M.get('model_name')})
@app.route('/switch_model', methods=['POST'])
def switch_model():
name = (request.json or {}).get('name')
if name not in M.get('models', {}):
return Response(json.dumps({'ok': False, 'error': 'unknown model'}), status=400, mimetype='application/json')
try:
params = load_main_model(M['models'][name], name)
return Response(json.dumps({'ok': True, 'model': name, 'params': params}), mimetype='application/json')
except Exception as e:
return Response(json.dumps({'ok': False, 'error': str(e)}), status=500, mimetype='application/json')
@app.route('/clone_voice', methods=['POST'])
def clone_voice():
d = request.json or {}
if not d.get('audio'):
return Response(json.dumps({'ok': False, 'error': 'missing audio'}), status=400, mimetype='application/json')
try:
name = normalize_voice_name(d.get('name'))
value = build_clone_voice(d['audio'])
out_path = clone_voice_path()
saved = torch.load(out_path, map_location='cpu') if os.path.exists(out_path) else {}
saved[name] = value
torch.save(saved, out_path)
register_voice(name, value, group='manual')
return Response(json.dumps({'ok': True, 'voice': name, 'path': './model/speaker/' + CLONE_FILE}), mimetype='application/json')
except Exception as e:
return Response(json.dumps({'ok': False, 'error': str(e)}), status=500, mimetype='application/json')
@app.route('/delete_voice', methods=['POST'])
def delete_voice():
d = request.json or {}
name = ' '.join(str(d.get('name') or '').split())
if not name:
return Response(json.dumps({'ok': False, 'error': 'missing name'}), status=400, mimetype='application/json')
try:
delete_manual_voice(name)
return Response(json.dumps({'ok': True, 'voice': name}), mimetype='application/json')
except Exception as e:
return Response(json.dumps({'ok': False, 'error': str(e)}), status=500, mimetype='application/json')
@app.route('/chat', methods=['POST'])
def chat():
d = request.json
history = d.get('history', [])
samples = None
if d.get('audio'):
seg = AudioSegment.from_file(io.BytesIO(base64.b64decode(d['audio']))).set_frame_rate(16000).set_channels(1).set_sample_width(2)
samples = np.frombuffer(seg.raw_data, dtype=np.int16).astype(np.float32) / 32768.0
va = voice_args(d.get('voice', 'default'))
def gen():
audio_inputs, audio_lens, pixel_values, prompt, user_text, asr_th, asr_res = prepare_turn(
d.get('text', ''), samples, d.get('image'), do_asr_for_image=True)
x = build_ids(prompt, history)
asr_sent = False
if user_text and samples is not None and d.get('image'):
yield sse({'type': 'user_prompt', 'content': user_text}); asr_sent = True
frames, text_ttft, audio_ttft = [], None, None
t0 = time.time(); hi = 0
for y, af in run_generate(x, audio_inputs, audio_lens, pixel_values,
max_new_tokens=d.get('max_tokens', 512),
temperature=d.get('temperature', 1), top_p=0.85, **va):
if not asr_sent and asr_th and not asr_th.is_alive():
asr_th.join()
if asr_res[0]: yield sse({'type': 'user_prompt', 'content': asr_res[0]})
asr_sent = True
if y is not None:
if text_ttft is None:
text_ttft = (time.time() - t0) * 1000
yield sse({'type': 'ttft', 'text_ttft': round(text_ttft, 1)})
ans = M['tokenizer'].decode(y[0].tolist(), skip_special_tokens=True)
if ans and ans[-1] != '\ufffd' and len(ans) > hi:
yield sse({'type': 'text', 'content': ans[hi:]}); hi = len(ans)
if af:
if audio_ttft is None:
audio_ttft = (time.time() - t0) * 1000
yield sse({'type': 'ttft', 'audio_ttft': round(audio_ttft, 1)})
frames.append(af)
for pcm in stream_pcm(frames):
b64 = base64.b64encode(pcm).decode()
for i in range(0, len(b64), 2000):
yield sse({'type': 'pcm', 'c': b64[i:i+2000], 'd': i+2000 >= len(b64)})
for pcm in stream_pcm(frames, flush=True):
b64 = base64.b64encode(pcm).decode()
for i in range(0, len(b64), 2000):
yield sse({'type': 'pcm', 'c': b64[i:i+2000], 'd': i+2000 >= len(b64)})
if not asr_sent:
if asr_th:
asr_th.join()
if asr_res[0]: yield sse({'type': 'user_prompt', 'content': asr_res[0]})
else:
yield sse({'type': 'user_prompt', 'content': prompt})
yield sse({'type': 'done'})
return Response(gen(), mimetype='text/event-stream', headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'})
@sock.route('/ws/realtime')
def realtime(ws):
session = RealtimeSession(M['vad_path'])
q = queue.Queue(); alive = [True]; state = {'history': [], 'image': None}
n_hist = M['cfg'].max_history_turns
def push_audio(data):
return session.push_chunk(np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0)
def set_ctx(msg):
h = msg.get('history') or []
state['history'] = h[-n_hist:] if n_hist > 0 else []
if 'image' in msg: state['image'] = msg.get('image')
if 'voice' in msg: state['voice'] = msg.get('voice', 'default')
def poll_interrupt():
while True:
try: data = q.get_nowait()
except queue.Empty: return False
if isinstance(data, bytes):
if push_audio(data) == 'interrupt': return True
ws.send(json.dumps({'type': 'vad', 'speaking': session.speaking}))
else:
m = json.loads(data)
if m.get('type') == 'context': set_ctx(m)
elif m.get('type') in ('stop', 'end'):
if m['type'] == 'end': alive[0] = False
session.interrupt = True; return True
def recv_loop():
while alive[0]:
try:
data = ws.receive(timeout=1)
if data is None: alive[0] = False; break
q.put(data)
except: alive[0] = False; break
threading.Thread(target=recv_loop, daemon=True).start()
try:
while alive[0]:
try: data = q.get(timeout=0.05)
except queue.Empty: continue
if isinstance(data, str):
m = json.loads(data)
if m.get('type') == 'context': set_ctx(m)
elif m.get('type') == 'stop': session.interrupt = True
elif m.get('type') == 'end': break
continue
if session.generating:
push_audio(data); ws.send(json.dumps({'type': 'vad', 'speaking': session.speaking})); continue
status = push_audio(data)
ws.send(json.dumps({'type': 'vad', 'speaking': session.speaking}))
if status != 'speech_end': continue
session.generating = True
audio = session.get_audio()
ws.send(json.dumps({'type': 'generating'}))
audio_inputs, audio_lens, pixel_values, prompt, user_text, asr_th, asr_res = prepare_turn(
'', audio, state['image'], do_asr_for_image=True)
if state['image']: state['image'] = None
x = build_ids(prompt, state['history'])
va_rt = voice_args(state.get('voice', 'default'))
frames, full_text, interrupted = [], '', False
for y, af in run_generate(x, audio_inputs, audio_lens, pixel_values,
max_new_tokens=512, temperature=0.7, **va_rt):
if poll_interrupt() or session.interrupt: interrupted = True; break
if y is not None:
ans = M['tokenizer'].decode(y[0].tolist(), skip_special_tokens=True)
if ans and ans[-1] != '\ufffd' and len(ans) > len(full_text):
ws.send(json.dumps({'type': 'text', 'content': ans[len(full_text):]})); full_text = ans
if af:
frames.append(af)
for pcm in stream_pcm(frames):
ws.send(json.dumps({'type': 'pcm', 'data': base64.b64encode(pcm).decode()}))
if not interrupted:
for pcm in stream_pcm(frames, flush=True):
ws.send(json.dumps({'type': 'pcm', 'data': base64.b64encode(pcm).decode()}))
if asr_th:
asr_th.join(); user_text = asr_res[0] or user_text
if n_hist > 0:
if user_text: state['history'].append({'role': 'user', 'content': user_text})
if full_text: state['history'].append({'role': 'assistant', 'content': full_text})
state['history'] = state['history'][-n_hist:]
ws.send(json.dumps({'type': 'done', 'interrupted': interrupted or session.interrupt}))
session.generating = False; session.interrupt = False
finally:
alive[0] = False
def init_model(args):
M['cfg'] = args; M['device'] = args.device
with contextlib.redirect_stdout(io.StringIO()):
M['asr'] = AutoModel(model='../model/SenseVoiceSmall', trust_remote_code=True, device=args.device, disable_update=True)
M['models'] = scan_hf_models(args.load_from)
if not M['models']:
raise RuntimeError(f"未在 {os.path.abspath(args.load_from)} 找到 transformers 模型")
model_name = next(iter(M['models']))
load_main_model(M['models'][model_name], model_name)
try:
from transformers import MimiModel
M['mimi'] = MimiModel.from_pretrained('../model/mimi').eval()
print('Mimi model loaded')
except: M['mimi'] = None
try:
from modelscope.models.audio.sv.DTDNN import CAMPPlus
M['campplus'] = CAMPPlus(feat_dim=80, embedding_size=192, growth_rate=32, bn_size=4,
init_channels=128, config_str='batchnorm-relu', memory_efficient=True)
sd = torch.load('../model/campplus/campplus_cn_common.pt', map_location='cpu')
M['campplus'].load_state_dict({k: v.float() for k, v in sd.items()})
M['campplus'] = M['campplus'].eval().to(args.device)
M['mel_fn'] = torchaudio.transforms.MelSpectrogram(
sample_rate=16000, n_fft=512, win_length=400, hop_length=160,
n_mels=80, f_min=20, f_max=7600, norm='slaney', mel_scale='slaney',
).to(args.device)
print('CAM++ loaded')
except Exception as e:
M['campplus'], M['mel_fn'] = None, None
print(f'CAM++ load failed: {e}')
M['vad_path'] = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'vad', 'silero_vad.onnx')
spk_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'speaker')
for fn, group in [('voices.pt', 'builtin'), ('voices_unseen.pt', 'unseen'), (CLONE_FILE, 'manual')]:
fp = os.path.join(spk_dir, fn)
if os.path.exists(fp):
for speaker, v in torch.load(fp, map_location=args.device).items():
if speaker not in V or fn == CLONE_FILE:
register_voice(speaker, v, group=group)
if V: print(f'Loaded {len(V)} voices: builtin={sorted(V_builtin)}, unseen={sorted(V_unseen)}, manual={sorted(V_manual)}')
log_model_params(M['model'])
print('Warmup...')
with torch.no_grad():
ids = torch.tensor([[1, 2, 3]], device=args.device)
au = torch.full((1, 8, 3), 2049, dtype=torch.long, device=args.device)
M['model'].forward(torch.cat((au, ids.unsqueeze(1)), dim=1))
if M['model'].audio_encoder: M['model'].audio_encoder(torch.zeros(1, 100, 560, device=args.device), torch.tensor([100], device=args.device))
if M['mimi']: M['mimi'].decode(torch.zeros(1, 8, 1, dtype=torch.long))
print('Warmup done!')
if __name__ == '__main__':
p = argparse.ArgumentParser()
p.add_argument('--load_from', default='./')
p.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
p.add_argument('--port', default=7860, type=int)
p.add_argument('--audio_chunk_frames', default=4, type=int)
p.add_argument('--audio_overlap', default=2, type=int)
p.add_argument('--max_history_turns', default=0, type=int)
args = p.parse_args()
init_model(args)
app.run(host='0.0.0.0', port=args.port, threaded=True)