在实时语音交互上超过GPT-4o，端到端语音模型Mini-Omni部署

Mini-Omni是清华大学开源的多模态大型语言模型，具备实时语音输入和流式音频输出的能力。

Mini-Omni模型能够一边听、一边说，一边思考，类似于ChatGPT的语言对话模式。

Mini-Omni模型的主要特点是能够直接通过音频模态进行推理，并生成流式输出，而不需要依赖额外的文本到语音(TTS)系统，从而减少了延迟。

Mini-Omni模型的架构在Qwen2-0.5B基础上进行了增强，使用了Whisper-small编码器来有效处理语音输入。

Mini-Omni模型采用了并行文本-音频生成方法，通过批量并行解码生成语音和文本，确保了模型在不同模态间的推理能力不受损害。

Mini-Omni模型还引入了VoiceAssistant-400K数据集，用于对优化语音输出的模型进行微调。

github项目地址：https://github.com/gpt-omni/mini-omni。

一、环境安装

1、python环境

建议安装python版本在3.10以上。

2、pip库安装

pip install torch==2.3.1+cu118 torchvision==0.18.1+cu118 torchaudio==2.3.1 --extra-index-url https://download.pytorch.org/whl/cu118

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

3、模型下载：

git lfs install

git clone https://huggingface.co/gpt-omni/mini-omni

二、功能测试

1、运行测试：

（1）python代码调用测试

import os
import torch
import time
import lightning as L
import soundfile as sf
import whisper
from snac import SNAC
from litgpt import Tokenizer
from tqdm import tqdm
from huggingface_hub import snapshot_download
from lightning.fabric.utilities.load import _lazy_load as lazy_load
from utils.snac_utils import layershift, reconscruct_snac, reconstruct_tensors, get_time_str, get_snac, generate_audio_data
from litgpt.utils import num_parameters
from litgpt.generate.base import generate_AA, generate_ASR, generate_TA, generate_TT, generate_AT, generate_TA_BATCH, next_token_batch
from litgpt.model import GPT, Configtorch.set_printoptions(sci_mode=False)# Constants Definitions
text_vocabsize = 151936
text_specialtokens = 64
audio_vocabsize = 4096
audio_specialtokens = 64padded_text_vocabsize = text_vocabsize + text_specialtokens
padded_audio_vocabsize = audio_vocabsize + audio_specialtokens_eot = text_vocabsize
_pad_t = text_vocabsize + 1
_input_t = text_vocabsize + 2
_answer_t = text_vocabsize + 3
_asr = text_vocabsize + 4_eoa = audio_vocabsize
_pad_a = audio_vocabsize + 1
_input_a = audio_vocabsize + 2
_answer_a = audio_vocabsize + 3
_split = audio_vocabsize + 4# Utility Functions
def get_input_ids_TA(text, text_tokenizer):input_ids_item = [[] for _ in range(8)]text_tokens = text_tokenizer.encode(text)for i in range(7):input_ids_item[i] = [layershift(_pad_a, i)] * (len(text_tokens) + 2) + [layershift(_answer_a, i)]input_ids_item[i] = torch.tensor(input_ids_item[i]).unsqueeze(0)input_ids_item[-1] = [_input_t] + text_tokens.tolist() + [_eot] + [_answer_t]input_ids_item[-1] = torch.tensor(input_ids_item[-1]).unsqueeze(0)return input_ids_itemdef get_input_ids_TT(text, text_tokenizer):input_ids_item = [[] for i in range(8)]text_tokens = text_tokenizer.encode(text).tolist()for i in range(7):input_ids_item[i] = torch.tensor([layershift(_pad_a, i)] * (len(text_tokens) + 3)).unsqueeze(0)input_ids_item[-1] = [_input_t] + text_tokens + [_eot] + [_answer_t]input_ids_item[-1] = torch.tensor(input_ids_item[-1]).unsqueeze(0)return input_ids_itemdef get_input_ids_whisper(mel, leng, whispermodel, device, special_token_a=_answer_a, special_token_t=_answer_t):with torch.no_grad():mel = mel.unsqueeze(0).to(device)audio_feature = whispermodel.embed_audio(mel)[0][:leng]T = audio_feature.size(0)input_ids = []for i in range(7):input_ids_item = []input_ids_item.append(layershift(_input_a, i))input_ids_item += [layershift(_pad_a, i)] * Tinput_ids_item += [(layershift(_eoa, i)), layershift(special_token_a, i)]input_ids.append(torch.tensor(input_ids_item).unsqueeze(0))input_id_T = torch.tensor([_input_t] + [_pad_t] * T + [_eot, special_token_t])input_ids.append(input_id_T.unsqueeze(0))return audio_feature.unsqueeze(0), input_idsdef get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device):with torch.no_grad():mel = mel.unsqueeze(0).to(device)audio_feature = whispermodel.embed_audio(mel)[0][:leng]T = audio_feature.size(0)input_ids_AA, input_ids_AT = [], []for i in range(7):lang_shift = layershift(_pad_a, i)input_ids_item_AA = [layershift(_input_a, i)] + [lang_shift] * T + [(layershift(_eoa, i)), layershift(_answer_a, i)]input_ids_item_AT = [layershift(_input_a, i)] + [lang_shift] * T + [(layershift(_eoa, i)), lang_shift]input_ids_AA.append(torch.tensor(input_ids_item_AA))input_ids_AT.append(torch.tensor(input_ids_item_AT))input_id_T = torch.tensor([_input_t] + [_pad_t] * T + [_eot, _answer_t])input_ids_AA.append(input_id_T)input_ids_AT.append(input_id_T)return torch.stack([audio_feature, audio_feature]), [input_ids_AA, input_ids_AT]def load_audio(path):audio = whisper.load_audio(path)duration_ms = (len(audio) / 16000) * 1000audio = whisper.pad_or_trim(audio)mel = whisper.log_mel_spectrogram(audio)return mel, int(duration_ms / 20) + 1def model_inference(fabric, model, snacmodel, tokenizer, func, step, *args, **kwargs):with fabric.init_tensor():model.set_kv_cache(batch_size=1)output = func(fabric, *args, **kwargs)model.clear_kv_cache()return outputdef load_model(ckpt_dir, device):snacmodel = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(device)whispermodel = whisper.load_model("small").to(device)text_tokenizer = Tokenizer(ckpt_dir)fabric = L.Fabric(devices=1, strategy="auto")config = Config.from_file(ckpt_dir + "/model_config.yaml")config.post_adapter = Falsewith fabric.init_module(empty_init=False):model = GPT(config)model = fabric.setup(model)state_dict = lazy_load(ckpt_dir + "/lit_model.pth")model.load_state_dict(state_dict, strict=True)model.to(device).eval()return fabric, model, text_tokenizer, snacmodel, whispermodeldef download_model(ckpt_dir):repo_id = "gpt-omni/mini-omni"snapshot_download(repo_id, local_dir=ckpt_dir, revision="main")class OmniInference:def __init__(self, ckpt_dir='checkpoint', device='cuda'):self.device = deviceif not os.path.exists(ckpt_dir):print(f"Checkpoint directory {ckpt_dir} not found, downloading from huggingface")download_model(ckpt_dir)self.fabric, self.model, self.text_tokenizer, self.snacmodel, self.whispermodel = load_model(ckpt_dir, device)def warm_up(self, sample='./data/samples/output1.wav'):for _ in self.run_AT_batch_stream(sample):pass@torch.inference_mode()def run_AT_batch_stream(self, audio_path, stream_stride=4, max_returned_tokens=2048, temperature=0.9, top_k=1, top_p=1.0, eos_id_a=_eoa, eos_id_t=_eot):assert os.path.exists(audio_path), f"Audio file {audio_path} not found"with self.fabric.init_tensor():self.model.set_kv_cache(batch_size=2)mel, leng = load_audio(audio_path)audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, self.whispermodel, self.device)T = input_ids[0].size(1)device = input_ids[0].deviceassert max_returned_tokens > T, f"Max returned tokens {max_returned_tokens} should be greater than audio length {T}"if self.model.max_seq_length < max_returned_tokens - 1:raise NotImplementedError(f"max_seq_length {self.model.max_seq_length} needs to be >= {max_returned_tokens - 1}")input_pos = torch.tensor([T], device=device)list_output = [[] for _ in range(8)]tokens_A, token_T = next_token_batch(self.model,audio_feature.to(torch.float32).to(self.model.device),input_ids,[T - 3, T - 3],["A1T2", "A1T2"],input_pos=torch.arange(0, T, device=device),temperature=temperature,top_k=top_k,top_p=top_p,)for i in range(7):list_output[i].append(tokens_A[i].tolist()[0])list_output[7].append(token_T.tolist()[0])model_input_ids = [[] for _ in range(8)]for i in range(7):tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize + i * padded_audio_vocabsizemodel_input_ids[i].append(tokens_A[i].clone().to(device).to(torch.int32))model_input_ids[i].append(torch.tensor([layershift(4097, i)], device=device))model_input_ids[i] = torch.stack(model_input_ids[i])model_input_ids[-1].append(token_T.clone().to(torch.int32))model_input_ids[-1].append(token_T.clone().to(torch.int32))model_input_ids[-1] = torch.stack(model_input_ids[-1])text_end = Falseindex = 1nums_generate = stream_stridebegin_generate = Falsecurrent_index = 0for _ in tqdm(range(2, max_returned_tokens - T + 1)):tokens_A, token_T = next_token_batch(self.model, None, model_input_ids, None, None, input_pos=input_pos, temperature=temperature, top_k=top_k, top_p=top_p)if text_end:token_T = torch.tensor([_pad_t], device=device)if tokens_A[-1] == eos_id_a:breakif token_T == eos_id_t:text_end = Truefor i in range(7):list_output[i].append(tokens_A[i].tolist()[0])list_output[7].append(token_T.tolist()[0])model_input_ids = [[] for _ in range(8)]for i in range(7):tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize + i * padded_audio_vocabsizemodel_input_ids[i].append(tokens_A[i].clone().to(device).to(torch.int32))model_input_ids[i].append(torch.tensor([layershift(4097, i)], device=device))model_input_ids[i] = torch.stack(model_input_ids[i])model_input_ids[-1].append(token_T.clone().to(torch.int32))model_input_ids[-1].append(token_T.clone().to(torch.int32))model_input_ids[-1] = torch.stack(model_input_ids[-1])if index == 7:begin_generate = Trueif begin_generate:current_index += 1if current_index == nums_generate:current_index = 0snac = get_snac(list_output, index, nums_generate)audio_stream = generate_audio_data(snac, self.snacmodel)yield audio_streaminput_pos = input_pos.add_(1)index += 1text = self.text_tokenizer.decode(torch.tensor(list_output[-1]))print(f"Text output: {text}")self.model.clear_kv_cache()return list_outputdef test_infer():device = "cuda:0"out_dir = f"./output/{get_time_str()}"ckpt_dir = f"./checkpoint"if not os.path.exists(ckpt_dir):print(f"Checkpoint directory {ckpt_dir} not found, downloading from huggingface")download_model(ckpt_dir)fabric, model, text_tokenizer, snacmodel, whispermodel = load_model(ckpt_dir, device)task = ['A1A2', 'asr', "T1A2", "AA-BATCH", 'T1T2', 'AT']# Prepare test datatest_audio_list = sorted(os.listdir('./data/samples'))test_audio_list = [os.path.join('./data/samples', path) for path in test_audio_list]test_audio_transcripts = ["What is your name?","What are your hobbies?","Do you like Beijing?","How are you feeling today?","What is the weather like today?"]test_text_list = ["What is your name?","How are you feeling today?","Can you describe your surroundings?","What did you do yesterday?","What is your favorite book and why?","How do you make a cup of tea?","What is the weather like today?","Can you explain the concept of time?","Can you tell me a joke?"]with torch.no_grad():for task_name in task:if "A1A2" in task_name:print("===============================================================")print("                       Testing A1A2")print("===============================================================")for i, path in enumerate(test_audio_list):try:mel, leng = load_audio(path)audio_feature, input_ids = get_input_ids_whisper(mel, leng, whispermodel, device)text = model_inference(fabric, model, snacmodel, text_tokenizer, A1_A2, i,fabric, audio_feature, input_ids, leng, model, text_tokenizer, i, snacmodel, out_dir)print(f"Input: {test_audio_transcripts[i]}")print(f"Output: {text}")print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")except Exception as e:print(f"[Error] Failed to process {path}: {e}")print("===============================================================")if 'asr' in task_name:print("===============================================================")print("                       Testing ASR")print("===============================================================")for i, path in enumerate(test_audio_list):mel, leng = load_audio(path)audio_feature, input_ids = get_input_ids_whisper(mel, leng, whispermodel, device, special_token_a=_pad_a, special_token_t=_asr)output = model_inference(fabric, model, snacmodel, text_tokenizer, A1_T1, i,fabric, audio_feature, input_ids, leng, model, text_tokenizer, i).lower().replace(',', '').replace('.', '').replace('?', '')print(f"Audio path: {path}")print(f"Audio transcript: {test_audio_transcripts[i]}")print(f"ASR output: {output}")print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")print("===============================================================")if "T1A2" in task_name:print("===============================================================")print("                       Testing T1A2")print("===============================================================")for i, text in enumerate(test_text_list):input_ids = get_input_ids_TA(text, text_tokenizer)text_output = model_inference(fabric, model, snacmodel, text_tokenizer, T1_A2, i,fabric, input_ids, model, text_tokenizer, i, snacmodel, out_dir)print(f"Input: {text}")print(f"Output: {text_output}")print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")print("===============================================================")if "T1T2" in task_name:print("===============================================================")print("                       Testing T1T2")print("===============================================================")for i, text in enumerate(test_text_list):input_ids = get_input_ids_TT(text, text_tokenizer)text_output = model_inference(fabric, model, snacmodel, text_tokenizer, T1_T2, i,fabric, input_ids, model, text_tokenizer, i)print(f"Input: {text}")print(f"Output: {text_output}")print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")print("===============================================================")if "AT" in task_name:print("===============================================================")print("                       Testing A1T2")print("===============================================================")for i, path in enumerate(test_audio_list):mel, leng = load_audio(path)audio_feature, input_ids = get_input_ids_whisper(mel, leng, whispermodel, device, special_token_a=_pad_a, special_token_t=_answer_t)text = model_inference(fabric, model, snacmodel, text_tokenizer, A1_T2, i,fabric, audio_feature, input_ids, leng, model, text_tokenizer, i)print(f"Input: {test_audio_transcripts[i]}")print(f"Output: {text}")print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")print("===============================================================")if "AA-BATCH" in task_name:print("===============================================================")print("                       Testing A1A2-BATCH")print("===============================================================")for i, path in enumerate(test_audio_list):mel, leng = load_audio(path)audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device)text = model_inference(fabric, model, snacmodel, text_tokenizer, A1_A2_batch, i,fabric, audio_feature, input_ids, leng, model, text_tokenizer, i, snacmodel, out_dir)print(f"Input: {test_audio_transcripts[i]}")print(f"Output: {text}")print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")print("===============================================================")print("========================= Test End ============================")if __name__ == "__main__":test_infer()

未完......

更多详细的欢迎关注：杰哥新技术