Upload 9 files

.gitignore

@@ -0,0 +1,171 @@
+# Byte-compiled / optimized / DLL files
+*pyc
+*pth
+checkpoint*/
+output/
+.DS_Store
+*.wav
+*.webp
+*.jpg
+
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

Dockerfile

@@ -0,0 +1,19 @@
+# Base image set karein jo Python 3.10 par hai
+FROM python:3.10
+
+# Working directory set karein
+WORKDIR /app
+
+# Local files ko container mein copy karein
+COPY . .
+
+# Linux packages install karein
+RUN apt-get update && \
+    apt-get install -y ffmpeg
+
+# Python dependencies install karein
+RUN pip install --no-cache-dir -r requirements.txt && \
+    pip install PyAudio==0.2.14
+
+# Server aur Streamlit app run karne ke liye commands
+CMD ["bash", "-c", "python3 server.py --ip '0.0.0.0' --port 60808 & API_URL=http://0.0.0.0:60808/chat streamlit run webui/omni_streamlit.py --server.port=8501"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 gpt-omni
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,153 @@
----
-title: Bo2
-emoji: 💬
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 5.0.1
-app_file: app.py
-pinned: false
-license: llama3
----
-
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
+
+# Mini-Omni2
+
+<p align="center">
+    <img src="./data/figures/title_new.png" width="90%"/>
+</p>
+
+
+<p align="center">
+🤗 <a href="https://huggingface.co/gpt-omni/mini-omni2">Hugging Face</a>   | 📖 <a href="https://github.com/gpt-omni/mini-omni2">Github</a> 
+|     📑 <a href="https://arxiv.org/abs/2410.11190">Technical report</a> 
+</p>
+
+<h5 align="center"> If you like little Omni2, please give us a star⭐ and cite our <a href="(https://arxiv.org/abs/2410.11190)">paper</a>!</h2>
+
+## Introduction
+Mini-Omni2 is an **omni-interactive** model. It can **understand image, audio and text inputs and has end-to-end voice conversations with users**. Featuring **real-time voice output**, **omni-capable multimodal understanding** and flexible interaction **ability with interruption mechanism while speaking**.
+
+<p align="center">
+    <img src="./data/figures/framework.jpeg" width="100%"/>
+</p>
+
+
+## Updates
+
+- **2024.10:** Release the model, technical report, inference and chat demo code.
+
+## Features
+✅ **Multimodal interaction**: with the ability to understand images, speech and text, just like GPT-4o.
+
+✅ **Real-time speech-to-speech** conversational capabilities. No extra ASR or TTS models required, just like [Mini-Omni](https://github.com/gpt-omni/mini-omni).
+
+<!-- ✅ **Streaming audio output**: with first-chunk latency of audio stream less than 0.3s. -->
+
+<!-- ✅ **Duplex interaction**: hearing while speaking, it can be interrupted by key words like "stop omni". -->
+
+
+## Demo
+
+NOTE: need to unmute first.
+
+https://github.com/user-attachments/assets/ad97ca7f-f8b4-40c3-a7e8-fa54b4edf155
+
+
+## ToDo
+- [ ] update interruption mechanism
+
+
+## Install
+
+Create a new conda environment and install the required packages:
+
+```sh
+conda create -n omni python=3.10
+conda activate omni
+
+git clone https://github.com/gpt-omni/mini-omni2.git
+cd mini-omni2
+pip install -r requirements.txt
+```
+
+## Quick start
+
+**Interactive demo**
+
+- start server
+
+NOTE: you need to start the server before running the streamlit or gradio demo with API_URL set to the server address.
+
+```sh
+sudo apt-get install ffmpeg
+conda activate omni
+cd mini-omni2
+python3 server.py --ip '0.0.0.0' --port 60808
+```
+
+
+- run streamlit demo
+
+NOTE: you need to run streamlit **locally** with PyAudio installed. 
+
+```sh
+pip install PyAudio==0.2.14
+API_URL=http://0.0.0.0:60808/chat streamlit run webui/omni_streamlit.py
+```
+
+
+**Local test**
+
+```sh
+conda activate omni
+cd mini-omni2
+# test run the preset audio samples and questions
+python inference_vision.py
+```
+
+## Mini-Omni2 Overview
+
+**1. Multimodal Modeling**:
+We use multiple sequences as the input and output of the model. In the input part, we will concatenate image, audio and text features to perform a series of comprehensive tasks, as shown in the following figures. In the output part, we use text-guided delayed parallel output to generate real-time speech responses.
+<p align="center">
+    <img src="./data/figures/inputids.png" width="100%"/>
+</p>
+
+**2. Multi-stage Training**:
+We propose an efficient alignment training method and conduct encoder adaptation, modal alignment, and multimodal fine-tuning respectively in the three-stage training.
+<p align="center">
+    <img src="./data/figures/training.jpeg" width="100%"/>
+</p>
+
+<!-- **3. Cases**:
+Here are more cases of Mini-Omni2:
+<p align="center">
+    <img src="./data/figures/samples.png" width="100%"/>
+</p> -->
+
+## FAQ
+
+**1. Does the model support other languages?**
+
+No, the model is only trained on English. However, as we use whisper as the audio encoder, the model can understand other languages which is supported by whisper (like chinese), but the output is only in English.
+
+**2. Error: can not run streamlit in local browser, with remote streamlit server**
+    
+You need start streamlit **locally** with PyAudio installed.
+
+
+## Acknowledgements 
+
+- [Qwen2](https://github.com/QwenLM/Qwen2/) as the LLM backbone.
+- [litGPT](https://github.com/Lightning-AI/litgpt/) for training and inference.
+- [whisper](https://github.com/openai/whisper/)  for audio encoding.
+- [clip](https://github.com/openai/CLIP)  for image encoding.
+- [snac](https://github.com/hubertsiuzdak/snac/)  for audio decoding.
+- [CosyVoice](https://github.com/FunAudioLLM/CosyVoice) for generating synthetic speech.
+- [OpenOrca](https://huggingface.co/datasets/Open-Orca/OpenOrca) and [MOSS](https://github.com/OpenMOSS/MOSS/tree/main) for alignment.
+
+## Citation
+
+```bibtex
+@article{xie2024miniomni2opensourcegpt4ovision,
+      title={Mini-Omni2: Towards Open-source GPT-4o with Vision, Speech and Duplex Capabilities}, 
+      author={Zhifei Xie and Changqiao Wu},
+      year={2024},
+      eprint={2410.11190},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS},
+      journal={ArXiv},
+      volume={abs/2410.11190},
+}
+```
+## Star History
+[![Star History Chart](https://api.star-history.com/svg?repos=gpt-omni/mini-omni2&type=Date)](https://star-history.com/#gpt-omni/mini-omni2&Date)

__init__.py

@@ -0,0 +1,4 @@
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))

inference.py

@@ -0,0 +1,704 @@
+import os
+import lightning as L
+import torch
+import glob
+import time
+from snac import SNAC
+from litgpt import Tokenizer
+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_image_batch
+)
+import soundfile as sf
+from litgpt.model import GPT, Config
+from lightning.fabric.utilities.load import _lazy_load as lazy_load
+from utils.snac_utils import layershift, reconscruct_snac, reconstruct_tensors, get_time_str
+from utils.snac_utils import get_snac, generate_audio_data
+import whisper
+from tqdm import tqdm
+from huggingface_hub import snapshot_download
+
+
+torch.set_printoptions(sci_mode=False)
+
+
+# TODO
+text_vocabsize = 151936
+text_specialtokens = 64
+audio_vocabsize = 4096
+audio_specialtokens = 64
+
+padded_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
+_image = audio_vocabsize + 5
+_eoimage = audio_vocabsize + 6
+
+
+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_item
+
+
+def 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_item
+
+
+def 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 = whisper.decode(whispermodel,mel, options).audio_features
+        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)] * T
+        input_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_ids
+
+
+def get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device):
+    with torch.no_grad():
+        mel = mel.unsqueeze(0).to(device)
+        # audio_feature = whisper.decode(whispermodel,mel, options).audio_features
+        audio_feature = whispermodel.embed_audio(mel)[0][:leng]
+    T = audio_feature.size(0)
+    input_ids_AA = []
+    for i in range(7):
+        input_ids_item = []
+        input_ids_item.append(layershift(_input_a, i))
+        input_ids_item += [layershift(_pad_a, i)] * T
+        input_ids_item += [(layershift(_eoa, i)), layershift(_answer_a, i)]
+        input_ids_AA.append(torch.tensor(input_ids_item))
+    input_id_T = torch.tensor([_input_t] + [_pad_t] * T + [_eot, _answer_t])
+    input_ids_AA.append(input_id_T)
+
+    input_ids_AT = []
+    for i in range(7):
+        input_ids_item = []
+        input_ids_item.append(layershift(_input_a, i))
+        input_ids_item += [layershift(_pad_a, i)] * T
+        input_ids_item += [(layershift(_eoa, i)), layershift(_pad_a, i)]
+        input_ids_AT.append(torch.tensor(input_ids_item))
+    input_id_T = torch.tensor([_input_t] + [_pad_t] * T + [_eot, _answer_t])
+    input_ids_AT.append(input_id_T)
+
+    input_ids = [input_ids_AA, input_ids_AT]
+    stacked_inputids = [[] for _ in range(8)]
+    for i in range(2):
+        for j in range(8):
+            stacked_inputids[j].append(input_ids[i][j])
+    stacked_inputids = [torch.stack(tensors) for tensors in stacked_inputids]
+    return torch.stack([audio_feature, audio_feature]), stacked_inputids
+
+
+def load_audio(path):
+    audio = whisper.load_audio(path)
+    duration_ms = (len(audio) / 16000) * 1000
+    audio = whisper.pad_or_trim(audio)
+    mel = whisper.log_mel_spectrogram(audio)
+    return mel, int(duration_ms / 20) + 1
+
+
+def A1_A2_batch(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
+                snacmodel, out_dir=None):
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=2)
+    tokenlist = generate_TA_BATCH(
+        model,
+        audio_feature,
+        input_ids,
+        [leng, leng],
+        ["A1A2", "A1T2"],
+        max_returned_tokens=2048,
+        temperature=0.9,
+        top_k=1,
+        eos_id_a=_eoa,
+        eos_id_t=_eot,
+        pad_id_t=_pad_t,
+        shift=padded_text_vocabsize,
+        include_prompt=True,
+        generate_text=True,
+    )
+    text_tokenlist = tokenlist[-1]
+    if text_vocabsize in text_tokenlist:
+        text_tokenlist = text_tokenlist[: text_tokenlist.index(text_vocabsize)]
+    text = text_tokenizer.decode(torch.tensor(text_tokenlist)).strip()
+
+    audio_tokenlist = tokenlist[:-1]
+    audiolist = reconscruct_snac(audio_tokenlist)
+    audio = reconstruct_tensors(audiolist)
+    if out_dir is None:
+        out_dir = "./output/default/A1-A2-batch"
+    else:
+        out_dir = out_dir + "/A1-A2-batch"
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir)
+    with torch.inference_mode():
+        audio_hat = snacmodel.decode(audio)
+    sf.write(
+        f"{out_dir}/{step:02d}.wav",
+        audio_hat.squeeze().cpu().numpy(),
+        24000,
+    )
+    model.clear_kv_cache()
+    return text
+
+
+def A1_T2(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step):
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
+    tokenlist = generate_AT(
+        model,
+        audio_feature,
+        input_ids,
+        [leng],
+        ["AT"],
+        max_returned_tokens=2048,
+        temperature=0.9,
+        top_k=1,
+        eos_id_a=_eoa,
+        eos_id_t=_eot,
+        pad_id_t=_pad_t,
+        shift=padded_text_vocabsize,
+        include_prompt=True,
+        generate_text=True,
+    )
+    return text_tokenizer.decode(torch.tensor(tokenlist)).strip()
+
+
+def A1_A2(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
+          snacmodel, out_dir=None):
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
+    tokenlist = generate_AA(
+        model,
+        audio_feature,
+        input_ids,
+        [leng],
+        ["A1T2"],
+        max_returned_tokens=2048,
+        temperature=0.9,
+        top_k=1,
+        eos_id_a=_eoa,
+        eos_id_t=_eot,
+        pad_id_t=_pad_t,
+        shift=padded_text_vocabsize,
+        include_prompt=True,
+        generate_text=True,
+    )
+    audiolist = reconscruct_snac(tokenlist)
+    tokenlist = tokenlist[-1]
+    if text_vocabsize in tokenlist:
+        tokenlist = tokenlist[: tokenlist.index(text_vocabsize)]
+    if out_dir is None:
+        out_dir = "./output/default/A1-A2"
+    else:
+        out_dir = out_dir + "/A1-A2"
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir)
+        
+    audio = reconstruct_tensors(audiolist)
+    with torch.inference_mode():
+        audio_hat = snacmodel.decode(audio)
+    sf.write(
+        f"{out_dir}/{step:02d}.wav",
+        audio_hat.squeeze().cpu().numpy(),
+        24000,
+    )
+    model.clear_kv_cache()
+    return text_tokenizer.decode(torch.tensor(tokenlist)).strip()
+
+
+def A1_T1(fabric, audio_feature, input_ids, leng, model, text_tokenizer, step):
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
+    tokenlist = generate_ASR(
+        model,
+        audio_feature,
+        input_ids,
+        [leng],
+        ["A1T1"],
+        max_returned_tokens=2048,
+        temperature=0.9,
+        top_k=1,
+        eos_id_a=_eoa,
+        eos_id_t=_eot,
+        pad_id_t=_pad_t,
+        shift=padded_text_vocabsize,
+        include_prompt=True,
+        generate_text=True,
+    )
+    model.clear_kv_cache()
+    return text_tokenizer.decode(torch.tensor(tokenlist)).strip()
+
+
+def T1_A2(fabric, input_ids, model, text_tokenizer, step,
+          snacmodel, out_dir=None):
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
+    tokenlist = generate_TA(
+        model,
+        None,
+        input_ids,
+        None,
+        ["T1A2"],
+        max_returned_tokens=2048,
+        temperature=0.9,
+        top_k=1,
+        eos_id_a=_eoa,
+        eos_id_t=_eot,
+        pad_id_t=_pad_t,
+        shift=padded_text_vocabsize,
+        include_prompt=True,
+        generate_text=True,
+    )
+
+    audiolist = reconscruct_snac(tokenlist)
+    tokenlist = tokenlist[-1]
+
+    if text_vocabsize in tokenlist:
+        tokenlist = tokenlist[: tokenlist.index(text_vocabsize)]
+    audio = reconstruct_tensors(audiolist)
+    if out_dir is None:
+        out_dir = "./output/default/T1-A2"
+    else:
+        out_dir = out_dir + "/T1-A2"
+    if not os.path.exists(out_dir):
+        os.makedirs(out_dir)
+
+    with torch.inference_mode():
+        audio_hat = snacmodel.decode(audio)
+    sf.write(
+        f"{out_dir}/{step:02d}.wav",
+        audio_hat.squeeze().cpu().numpy(),
+        24000,
+    )
+    model.clear_kv_cache()
+    return text_tokenizer.decode(torch.tensor(tokenlist)).strip()
+
+
+def T1_T2(fabric, input_ids, model, text_tokenizer, step):
+
+    with fabric.init_tensor():
+        model.set_kv_cache(batch_size=1)
+    tokenlist = generate_TT(
+        model,
+        None,
+        input_ids,
+        None,
+        ["T1T2"],
+        max_returned_tokens=2048,
+        temperature=0.9,
+        top_k=1,
+        eos_id_a=_eoa,
+        eos_id_t=_eot,
+        pad_id_t=_pad_t,
+        shift=padded_text_vocabsize,
+        include_prompt=True,
+        generate_text=True,
+    )
+    model.clear_kv_cache()
+    return text_tokenizer.decode(torch.tensor(tokenlist)).strip()
+
+    
+def load_model(ckpt_dir, device):
+    snacmodel = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(device)
+    whisper_model_path = ckpt_dir + "/small.pt"
+    if not os.path.exists(whisper_model_path):
+        whisper_model_path = "small"
+    whispermodel = whisper.load_model(whisper_model_path).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 = False
+
+    with 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, whispermodel
+
+    
+def download_model(ckpt_dir):
+    repo_id = "gpt-omni/mini-omni2"
+    snapshot_download(repo_id, local_dir=ckpt_dir, revision="main")
+
+    
+def get_text_stream(list_output, index, text_tokenizer):
+    text_tokens = list_output[-1][index:]
+    index += len(text_tokens)
+    is_text_end = False
+    if text_vocabsize in text_tokens:
+        text_tokens = text_tokens[:text_tokens.index(text_vocabsize)]
+        is_text_end = True
+    if len(text_tokens) == 0:
+        return "", index, is_text_end
+    res_text = text_tokenizer.decode(torch.tensor(text_tokens))
+    return res_text, index, is_text_end
+
+    
+class OmniInference:
+
+    def __init__(self, ckpt_dir='./checkpoint', device='cuda:0'):
+        self.device = device
+        if 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,
+                            save_path=None
+        ):
+
+        assert os.path.exists(audio_path), f"audio file {audio_path} not found"
+        model = self.model
+
+        with self.fabric.init_tensor():
+            model.set_kv_cache(batch_size=2,device=self.device)
+
+        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].device
+
+        assert max_returned_tokens > T, f"max_returned_tokens {max_returned_tokens} should be greater than audio length {T}"
+
+        if model.max_seq_length < max_returned_tokens - 1:
+            raise NotImplementedError(
+                f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}"
+            )
+
+        input_pos = torch.tensor([T], device=device)
+        list_output = [[] for i in range(8)]
+        tokens_A, token_T = next_token_image_batch(
+            model,
+            audio_feature.to(torch.float32).to(model.device),
+            None,
+            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 i in range(8)]
+        for i in range(7):
+            tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize + i * padded_audio_vocabsize
+            model_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 = False
+        index = 1
+        nums_generate = stream_stride
+        begin_generate = False
+        current_index = 0
+
+        text_index = 0
+        is_text_end = False
+
+        for _ in tqdm(range(2, max_returned_tokens - T + 1)):
+            tokens_A, token_T = next_token_image_batch(
+                model,
+                None,
+                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:
+                break
+
+            if token_T == eos_id_t:
+                text_end = True
+
+            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 i in range(8)]
+            for i in range(7):
+                tokens_A[i] = tokens_A[i].clone() +padded_text_vocabsize + i * padded_audio_vocabsize
+                model_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 = True
+
+            if begin_generate:
+                current_index += 1
+                if current_index == nums_generate:
+                    current_index = 0
+                    snac = get_snac(list_output, index, nums_generate)
+                    audio_stream = generate_audio_data(snac, self.snacmodel, self.device)
+                    if is_text_end:
+                        text_stream = ""
+                    else:
+                        text_stream, text_index, is_text_end = get_text_stream(list_output, text_index, self.text_tokenizer)
+
+                    yield (audio_stream, text_stream)
+
+            input_pos = input_pos.add_(1)
+            index += 1
+        text = self.text_tokenizer.decode(torch.tensor(list_output[-1]))
+        print(f"text output: {text}")
+
+        if save_path is not None:
+            audiolist = reconscruct_snac(list_output)
+            audio = reconstruct_tensors(audiolist)
+            with torch.inference_mode():
+                audio_hat = self.snacmodel.decode(audio)
+                sf.write(save_path, audio_hat.squeeze().cpu().numpy(), 24000)
+        
+        model.clear_kv_cache()
+        return list_output
+
+
+def 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 data
+    # TODO
+    test_audio_list = sorted(glob.glob('./data/samples/output*.wav'))
+    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?",
+    ]
+
+    # LOAD MODEL
+    with torch.no_grad():
+        if "A1A2" in task:
+            print("===============================================================")
+            print("                       testing A1A2")
+            print("===============================================================")
+            step = 0
+            for path in test_audio_list:
+                try:
+                    mel, leng = load_audio(path)
+                    audio_feature, input_ids = get_input_ids_whisper(mel, leng, whispermodel, device)
+                    text = A1_A2(
+                        fabric,
+                        audio_feature,
+                        input_ids,
+                        leng,
+                        model,
+                        text_tokenizer,
+                        step,
+                        snacmodel,
+                        out_dir=out_dir,
+                    )
+                    print(f"input: {test_audio_transcripts[step]}")
+                    print(f"output: {text}")
+                    step += 1
+                    print(
+                        "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"
+                    )
+                except:
+                    print(f"[error] failed to process {path}")
+            print("===============================================================")
+
+        if 'asr' in task:
+            print("===============================================================")
+            print("                       testing asr")
+            print("===============================================================")
+
+            index = 0
+            step = 0
+            for path in 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 = A1_T1(fabric, audio_feature, input_ids ,leng, model, text_tokenizer, index).lower().replace(',','').replace('.','').replace('?','')
+                print(f"audio_path: {path}")
+                print(f"audio transcript: {test_audio_transcripts[index]}")
+                print(f"asr output: {output}")
+                print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
+                index += 1
+
+        if "T1A2" in task:
+            step = 0
+            print("\n")
+            print("===============================================================")
+            print("                       testing T1A2")
+            print("===============================================================")
+            for text in test_text_list:
+                input_ids = get_input_ids_TA(text, text_tokenizer)
+                text_output = T1_A2(fabric, input_ids, model, text_tokenizer, step,
+                                    snacmodel, out_dir=out_dir)
+                print(f"input: {text}")
+                print(f"output: {text_output}")
+                print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
+                step += 1
+            print("===============================================================")
+
+        if "T1T2" in task:
+            step = 0
+            print("\n")
+            print("===============================================================")
+            print("                       testing T1T2")
+            print("===============================================================")
+
+            for text in test_text_list:
+                input_ids = get_input_ids_TT(text, text_tokenizer)
+                text_output = T1_T2(fabric, input_ids, model, text_tokenizer, step)
+                print(f" Input: {text}")
+                print(f"Output: {text_output}")
+                print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
+            print("===============================================================")
+
+        if "AT" in task:
+            print("===============================================================")
+            print("                       testing A1T2")
+            print("===============================================================")
+            step = 0
+            for path in 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 = A1_T2(
+                    fabric, audio_feature, input_ids, leng, model, text_tokenizer, step
+                )
+                print(f"input: {test_audio_transcripts[step]}")
+                print(f"output: {text}")
+                step += 1
+                print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
+            print("===============================================================")
+
+        if "AA-BATCH" in task:
+            print("===============================================================")
+            print("                       testing A1A2-BATCH")
+            print("===============================================================")
+            step = 0
+            for path in test_audio_list:
+                mel, leng = load_audio(path)
+                audio_feature, input_ids = get_input_ids_whisper_ATBatch(mel, leng, whispermodel, device)
+                text = A1_A2_batch(
+                    fabric, audio_feature, input_ids, leng, model, text_tokenizer, step,
+                    snacmodel, out_dir=out_dir
+                )
+                print(f"input: {test_audio_transcripts[step]}")
+                print(f"output: {text}")
+                step += 1
+                print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
+            print("===============================================================")
+
+        print("*********************** test end *****************************")
+
+
+
+if __name__ == "__main__":
+    test_infer()

inference_vision.py

@@ -0,0 +1,259 @@
+import os
+import torch
+from litgpt.generate.base import next_token_image_batch
+import soundfile as sf
+from utils.snac_utils import layershift, reconscruct_snac, reconstruct_tensors, get_time_str
+from utils.snac_utils import get_snac, generate_audio_data
+import clip
+import inference
+from tqdm import tqdm
+from inference import OmniInference, load_model, load_audio, download_model
+from inference import text_vocabsize, padded_text_vocabsize, get_text_stream
+from PIL import Image
+
+
+torch.set_printoptions(sci_mode=False)
+
+_image = inference._image
+_eoimage = inference._eoimage
+_pad_t = inference._pad_t
+_input_t = inference._input_t
+_answer_t = inference._answer_t
+_eot = inference._eot
+_eoa = inference._eoa
+_pad_a = inference._pad_a
+_input_a = inference._input_a
+_answer_a = inference._answer_a
+
+
+def get_input_ids_ImageQA_ATBatch(mel, leng, whispermodel, device):
+    
+    with torch.no_grad():
+        mel = mel.unsqueeze(0).to(device)
+        audio_feature = whispermodel.embed_audio(mel)[0][:leng]
+    
+    audio_len = audio_feature.size(0)
+    
+    input_ids = []
+    input_ids_item = [[] for i in range(8)]
+    for i in range(7):
+        input_ids_item[i] =  [layershift(_image,i)] + [layershift(_pad_a,i)] * 50 + [layershift(_eoimage,i)] 
+        input_ids_item[i] += [layershift(_input_a,i)]+[layershift(_pad_a,i)]*(audio_len)+[layershift(_eoa,i)]
+        input_ids_item[i] += [layershift(_answer_a,i)]
+
+    input_ids_item[-1] = [_pad_t]* (52 + 2 + audio_len) + [_answer_t] 
+    input_ids_item = [torch.tensor(item) for item in input_ids_item]
+
+    input_ids.append(input_ids_item)
+
+    input_ids_item = [[] for i in range(8)]
+    for i in range(7):
+        input_ids_item[i] =  [layershift(_image,i)] + [layershift(_pad_a,i)] * 50 + [layershift(_eoimage,i)] 
+        input_ids_item[i] += [layershift(_input_a,i)]+[layershift(_pad_a,i)]*(audio_len)+[layershift(_eoa,i)] + [layershift(_pad_a,i)]
+
+    input_ids_item[-1] = [_pad_t]* (52 + 2 + audio_len) + [_answer_t] 
+
+    input_ids_item = [torch.tensor(item) for item in input_ids_item]
+    input_ids.append(input_ids_item)
+
+    stacked_inputids = [[] for _ in range(8)]
+    for i in range(2):
+        for j in range(8):
+            stacked_inputids[j].append(input_ids[i][j])
+    stacked_inputids = [torch.stack(tensors) for tensors in stacked_inputids]
+
+    return torch.stack([audio_feature,audio_feature]), stacked_inputids
+
+    
+def load_clip_model(ckpt_dir, device):
+    clip_model_path = ckpt_dir + "/ViT-B-32.pt"
+    if not os.path.exists(clip_model_path):
+        clip_model_path = "ViT-B/32"
+    clipmodel, clippreprocess = clip.load(clip_model_path, device=device)
+    return clipmodel, clippreprocess
+
+    
+class OmniVisionInference(OmniInference):
+
+    def __init__(self, ckpt_dir='./checkpoint', device='cuda:0'):
+        self.device = device
+        if 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)
+        self.clipmodel, self.clippreprocess = load_clip_model(ckpt_dir, device)
+
+    def warm_up(self, 
+                audio_sample='./data/samples/vision_qa_audio.wav',
+                image_sample='./data/samples/vision_qa_image.jpg'
+        ):
+        for _ in self.run_vision_AA_batch_stream(audio_sample, image_sample, 
+                                                 save_path="./data/samples/vision_qa_output.wav",
+                                                 warm_up=True):
+            pass
+
+    @torch.inference_mode()
+    def run_vision_AA_batch_stream(self, audio_path, image_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, 
+                                pad_id=_pad_t,
+                                save_path=None,
+                                warm_up=False
+        ):
+        with self.fabric.init_tensor():
+            self.model.set_kv_cache(batch_size=2)
+
+        model = self.model
+
+        mel, leng = load_audio(audio_path)
+        img = Image.open(image_path)
+
+        audio_feature, input_ids = get_input_ids_ImageQA_ATBatch(mel, leng, self.whispermodel, self.device)
+        ima = self.clippreprocess(img).unsqueeze(0).to(self.device)
+        ima_feature = self.clipmodel.encode_image(ima).squeeze(0).to(self.device)
+        
+        ima_feature = torch.stack([ima_feature.clone(),ima_feature.clone()]).to(self.device)
+        leng = [leng,leng]
+        task = ['ImageQA_A','ImageQA_AT']
+
+        T = input_ids[0].size(1)  
+        assert max_returned_tokens > T, f"max_returned_tokens {max_returned_tokens} should be greater than audio length {T}"
+
+        if model.max_seq_length < max_returned_tokens - 1:
+            raise NotImplementedError(
+                f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}"
+            )
+
+        list_output = [[] for i in range(8)]
+
+        tokens_A , token_T = next_token_image_batch(
+            model, 
+            audio_feature.to(torch.float32).to(self.device),
+            ima_feature.to(torch.float32).to(self.device) , 
+            input_ids , 
+            whisper_lens = leng , 
+            task = task, 
+            input_pos = torch.arange(0, T, device=self.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])
+
+        text_end = False
+        index = 1
+        nums_generate = stream_stride
+        begin_generate = False
+        current_index = 0
+        input_pos = torch.tensor([T], device=self.device)
+
+        model_input_ids = [[] for i in range(8)]
+        for i in range(7):
+            tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize+ i * 4160
+            model_input_ids[i].append(tokens_A[i].clone().to(self.device).to(torch.int32))
+            model_input_ids[i].append(torch.tensor([layershift(4097,i)],device=self.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_index = 0
+        is_text_end = False
+
+        for _ in tqdm(range(2, max_returned_tokens - T + 1)):
+            
+            tokens_A , token_T = next_token_image_batch(model, None , None , 
+                                                        input_ids = model_input_ids, 
+                                                        whisper_lens= None, 
+                                                        task = 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=self.device)
+
+            if tokens_A[-1] == eos_id_a:
+                break
+            if token_T == eos_id_t:
+                text_end = True
+
+            for i in range(7): list_output[i].append(tokens_A[i].tolist()[0])
+            list_output[7].append(token_T.tolist()[0])
+            
+
+            if index == 7:
+                begin_generate = True
+            
+            if begin_generate:
+                current_index += 1
+                if current_index == nums_generate:
+                    current_index = 0
+                    snac = get_snac(list_output,index,nums_generate)
+                    audio_stream = generate_audio_data(snac, self.snacmodel, self.device)
+                    if is_text_end:
+                        text_stream = ""
+                    else:
+                        text_stream, text_index, is_text_end = get_text_stream(list_output, text_index, self.text_tokenizer)
+
+                    yield (audio_stream, text_stream)
+
+                    if warm_up:
+                        break
+
+            input_pos = input_pos.add_(1)
+            model_input_ids = [[] for i in range(8)]
+            for i in range(7):
+                tokens_A[i] = tokens_A[i].clone() + padded_text_vocabsize+ i * 4160
+                model_input_ids[i].append(tokens_A[i].clone().to(self.device).to(torch.int32))
+                model_input_ids[i].append(torch.tensor([layershift(4097,i)],device=self.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])
+
+            index += 1    
+            
+        text_tokens = list_output[-1]
+        if text_vocabsize in text_tokens:
+            text_tokens = text_tokens[:text_tokens.index(text_vocabsize)]
+        res_text = self.text_tokenizer.decode(torch.tensor(text_tokens))
+        print(f"text output: {res_text}")
+
+        if save_path is not None:
+            audiolist = reconscruct_snac(list_output)
+            audio = reconstruct_tensors(audiolist)
+            with torch.inference_mode():
+                audio_hat = self.snacmodel.decode(audio)
+                sf.write(save_path, audio_hat.squeeze().cpu().numpy(), 24000)
+
+        model.clear_kv_cache()
+
+        
+def test_vision_infer():
+    client = OmniVisionInference()
+    client.warm_up()
+    input_audio_path = './data/samples/vision_qa_audio.wav'
+    input_image_path = './data/samples/vision_qa_image.jpg'
+
+    res_text = ""
+    for audio_stream, text_stream in client.run_vision_AA_batch_stream(
+        input_audio_path, 
+        input_image_path,
+        save_path="./vision_qa_output.wav"
+    ):
+        res_text += text_stream
+    print(f"text_output: {res_text}")
+
+
+if __name__ == "__main__":
+    test_vision_infer()

requirements.txt

@@ -1,4 +1,20 @@
-huggingface_hub==0.25.2
-transformers
-scipy
-accelerate
+torch==2.3.1
+torchvision==0.18.1
+torchaudio==2.3.1
+litgpt==0.4.3
+snac==1.2.0
+soundfile==0.12.1
+openai-whisper
+tokenizers==0.19.1 
+streamlit==1.37.1
+streamlit-webrtc
+# PyAudio==0.2.14
+pydub==0.25.1
+onnxruntime==1.19.0
+# numpy==1.26.3
+librosa==0.10.2.post1
+flask==3.0.3
+fire
+git+https://github.com/mini-omni/CLIP.git
+gradio_webrtc[vad]==0.0.11
+twilio

server.py

@@ -0,0 +1,96 @@
+import sys
+import os
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+
+import flask
+import base64
+import tempfile
+import traceback
+from flask import Flask, Response, stream_with_context
+from inference_vision import OmniVisionInference
+
+
+class OmniChatServer(object):
+    def __init__(self, ip='0.0.0.0', port=60808, run_app=True,
+                 ckpt_dir='./checkpoint', device='cuda:0') -> None:
+        server = Flask(__name__)
+        # CORS(server, resources=r"/*")
+        # server.config["JSON_AS_ASCII"] = False
+
+        self.client = OmniVisionInference(ckpt_dir, device)
+        self.client.warm_up()
+
+        server.route("/chat", methods=["POST"])(self.chat)
+
+        if run_app:
+            server.run(host=ip, port=port, threaded=False)
+        else:
+            self.server = server
+
+    def chat(self) -> Response:
+
+        req_data = flask.request.get_json()
+        try:
+            audio_data_buf = req_data["audio"].encode("utf-8")
+            audio_data_buf = base64.b64decode(audio_data_buf)
+            stream_stride = req_data.get("stream_stride", 4)
+            max_tokens = req_data.get("max_tokens", 2048)
+
+            image_data_buf = req_data.get("image", None)
+            if image_data_buf:
+                image_data_buf = image_data_buf.encode("utf-8")
+                image_data_buf = base64.b64decode(image_data_buf)
+
+            audio_path, img_path = None, None
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as audio_f, \
+                 tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as img_f:
+                audio_f.write(audio_data_buf)
+                audio_path = audio_f.name
+
+                if image_data_buf:
+                    img_f.write(image_data_buf)
+                    img_path = img_f.name
+                else:
+                    img_path = None
+
+                if img_path is not None:
+                    resp_generator = self.client.run_vision_AA_batch_stream(audio_f.name, img_f.name,
+                                                                             stream_stride, max_tokens,
+                                                                             save_path='./vision_qa_out_cache.wav')
+                else:
+                    resp_generator = self.client.run_AT_batch_stream(audio_f.name, stream_stride,
+                                                                      max_tokens,
+                                                                      save_path='./audio_qa_out_cache.wav')
+                return Response(stream_with_context(self.generator(resp_generator)),
+                                mimetype='multipart/x-mixed-replace; boundary=frame')
+        except Exception as e:
+            print(traceback.format_exc())
+            return Response("An error occurred", status=500)
+
+    def generator(self, resp_generator):
+        for audio_stream, text_stream in resp_generator:
+            yield b'\r\n--frame\r\n'
+            yield b'Content-Type: audio/wav\r\n\r\n'
+            yield audio_stream
+            yield b'\r\n--frame\r\n'
+            yield b'Content-Type: text/plain\r\n\r\n'
+            yield text_stream.encode()
+
+
+# CUDA_VISIBLE_DEVICES=1 gunicorn -w 2 -b 0.0.0.0:60808 'server:create_app()'
+def create_app():
+    server = OmniChatServer(run_app=False)
+    return server.server
+
+
+def serve(ip='0.0.0.0', port=60808, device='cuda:0'):
+
+    OmniChatServer(ip, port=port,run_app=True, device=device)
+
+
+if __name__ == "__main__":
+    import fire
+    fire.Fire(serve)
+