File size: 8,179 Bytes

---
license: apache-2.0
language:
- en
- zh
tags:
- MoE
- Unified Generation
- Speech and Music
- Multi-modal
---


<h1 align="center">UniMoE-Audio</h1>

**UniMoE-Audio**  is a unified framework that seamlessly combines speech and music generation. Powered by a novel Dynamic-Capacity Mixture-of-Experts architecture. 

<div align="center" style="display: flex; justify-content: center; margin-top: 10px;">
  <a href="https://mukioxun.github.io/Uni-MoE-site/home.html"><img src="https://img.shields.io/badge/📰 -Website-228B22" style="margin-right: 5px;"></a>
  <a href="https://arxiv.org/abs/2510.13344"><img src="https://img.shields.io/badge/📄-Paper-8A2BE2" style="margin-right: 5px;"></a>
</div>

---

**If you enjoy our work or want timely updates, please give us a like and follow us.**

## Open-source Plan
- [x] Model Checkpoint 
    - [x] [UniMoE-Audio-preview](https://huggingface.co/foggyforest/UniMoE-Audio-preview)
    - [ ] [UniMoE-Audio]()
- [x] Training and Inference Code: [HITsz-TMG/UniMoE-Audio](https://github.com/HITsz-TMG/UMOE-Scaling-Unified-Multimodal-LLMs/tree/master/UniMoE-Audio)
- [x] Technical Report: [UniMoE-Audio: Unified Speech and Music Generation with Dynamic-Capacity MoE](https://arxiv.org/abs/2510.13344)

## Evaluation

### Speech Synthesis
![Speech Synthesis](./imgs/Speech_Generation.png)

### Text to Music Generation
![Text to Music Generation](./imgs/T2M.png)

### Video-Text to Music Generation
![Video-Text to Music Generation](./imgs/VT2M.png)

## Requirements

Since we have used the Qwen2.5VL model, we advise you to install transformers>=4.53.1, or you might encounter the following error:
```
KeyError: 'qwen2_vl'
```
## Quickstart

We use `qwen-vl-utils` to handle various types of visual input. You can install it using the following command:
```
pip install qwen-vl-utils
```


We use the Descript Audio Codec (DAC) for audio compression.  You can install it using the following command:
```
pip install descript-audio-codec
```
The model weight will be automatically downloaded on first run.


## Usage

Here is a code snippet to show you how to use UniMoE-Audio with `transformers`

```python
import torch
import deepspeed_utils # This line is important, do not delete it
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor

# Import from utils modules
from utils import (
    Dac,
    preprocess_codec,
    DecoderOutput,
    tts_preprocess,
    t2m_preprocess,
    v2m_preprocess,
    prepare_audio_prompt,
    generate_output
)

model_path = "/path/to/your/model"

dac = Dac()

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float32,
    attn_implementation='sdpa',
    trust_remote_code=True,
).eval()
model = model.to('cuda')

processor = AutoProcessor.from_pretrained(model_path)

```

### TTS Example:

```python
transcription = [
    "The nature reserve covers only a small part of the marsh area.",
    "我们基于动态容量混合专家框架，构建了一个统一语音和音乐生成模型。"
]
prompt_wav = "/path/to/your/voice/prompt"
prompt_transcription = "content of your voice prompt"
        
prompt_codec = preprocess_codec(model, dac.encode(prompt_wav))
text_input, tts_generation_kwargs = tts_preprocess(transcription, prompt_codec, prompt_transcription, model.device)
source_input = processor.tokenizer(text_input, add_special_tokens=False, return_tensors="pt", padding=True).to(model.device)

prefill, prefill_steps = prepare_audio_prompt(model, audio_prompts=[None] * len(transcription))
dec_output = DecoderOutput(prefill, prefill_steps, model.device)
        
with torch.no_grad():
    generated_codes, lengths_Bx = model.generate(
        input_ids=source_input.input_ids,
        attention_mask=source_input.attention_mask,
        dec_output=dec_output,
        max_tokens=10 * 50, # maximum duration of the generated audio is 10 seconds
        min_tokens=1 * 50, # minimum duration of the generated audio is 1 seconds
        temperature=1.0,
        top_p=1.0,
        cfg_filter_top_k=45,
        do_sample=True,
        use_cache=True,
        **tts_generation_kwargs
    )
        
audios = generate_output(model, generated_codes, lengths_Bx)
for i in range(len(audios)):
    output_path = os.path.join(f"./generated_speech_{i}.wav")
    dac.decode(audios[i].transpose(0, 1).unsqueeze(0), save_path=output_path, min_duration=1)

```

### T2M Example:

```python
caption = [
    "A retro-inspired synthwave track with a driving beat and nostalgic melodies. Perfect for cruising or late-night drives.",
    "A mid-tempo electronic track with a driving beat and atmospheric synth textures. Ideal for background listening or a chill dance set."
]

text_input, t2m_generation_kwargs = t2m_preprocess(caption)

source_input = processor.tokenizer(text_input, add_special_tokens=False, return_tensors="pt", padding=True).to(model.device)

prefill, prefill_steps = prepare_audio_prompt(model, audio_prompts=[None] * len(caption))
dec_output = DecoderOutput(prefill, prefill_steps, model.device)
        
with torch.no_grad():
    generated_codes, lengths_Bx = model.generate(
        input_ids=source_input.input_ids,
        attention_mask=source_input.attention_mask,
        dec_output=dec_output,
        max_tokens=20 * 50, # maximum duration of the generated audio is 20 seconds
        min_tokens=8 * 50, # minimum duration of the generated audio is 8 seconds
        temperature=1.0,
        top_p=1.0,
        cfg_filter_top_k=45,
        do_sample=True,
        use_cache=True,
        **t2m_generation_kwargs
    )
        
audios = generate_output(model, generated_codes, lengths_Bx)
for i in range(len(audios)):
    output_path = os.path.join(f"./generated_music_{i}.wav")
    dac.decode(audios[i].transpose(0, 1).unsqueeze(0), save_path=output_path, min_duration=1)


```

### V2M Example:
```python

caption = [
    "A relaxing instrumental piece featuring a simple melody played on a synth flute. The track creates a calm and peaceful atmosphere.",
]
video = [
    "/path/to/your/video/path.mp4",
]

text_input,  video_inputs, fps_inputs, v2m_generation_kwargs = v2m_preprocess(caption, video)

source_input = processor(text=text_input, images=None, videos=video_inputs, fps=fps_inputs, padding=True, return_tensors="pt", do_resize=False)
source_input = source_input.to(model.device)

prefill, prefill_steps = prepare_audio_prompt(model, audio_prompts=[None] * len(caption))
dec_output = DecoderOutput(prefill, prefill_steps, model.device)
        
with torch.no_grad():
    generated_codes, lengths_Bx = model.generate(
        input_ids=source_input.input_ids,
        pixel_values_videos=source_input.pixel_values_videos,
        video_grid_thw=source_input.video_grid_thw,
        second_per_grid_ts=source_input.second_per_grid_ts,
        attention_mask=source_input.attention_mask,
        dec_output=dec_output,
        max_tokens=20 * 50, # maximum duration of the generated audio is 20 seconds
        min_tokens=8 * 50, # minimum duration of the generated audio is 8 seconds
        temperature=1.0,
        top_p=1.0,
        cfg_filter_top_k=45,
        do_sample=True,
        use_cache=True,
        **v2m_generation_kwargs
    )
        
audios = generate_output(model, generated_codes, lengths_Bx)
for i in range(len(audios)):
    output_path = os.path.join(f"./generated_video_music_{i}.wav")
    dac.decode(audios[i].transpose(0, 1).unsqueeze(0), save_path=output_path, min_duration=1)
```




# Citation

Please cite the repo if you use the model or code in this repo.

```
@article{liu2025unimoeaudiounifiedspeechmusic,
      title={UniMoE-Audio: Unified Speech and Music Generation with Dynamic-Capacity MoE}, 
      author={Zhenyu Liu and Yunxin Li and Xuanyu Zhang and Qixun Teng and Shenyuan Jiang and Xinyu Chen and Haoyuan Shi and Jinchao Li and Qi Wang and Haolan Chen and Fanbo Meng and Mingjun Zhao and Yu Xu and Yancheng He and Baotian Hu and Min Zhang},
      year={2025},
      journal={arXiv preprint arXiv:2510.13344},
      url={https://arxiv.org/abs/2510.13344}, 
}
```

# Contract

If you encounter any issue, feel free to contact us via the email: [email protected]