HIT-TMG/UniMoE-Audio-Preview

UniMoE-Audio

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

---

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

Open-source Plan

• Model Checkpoint
  • UniMoE-Audio-preview
  • UniMoE-Audio
• Training and Inference Code: HITsz-TMG/UniMoE-Audio
• Technical Report: UniMoE-Audio: Unified Speech and Music Generation with Dynamic-Capacity MoE

Evaluation

Speech Synthesis

Text to Music Generation

Video-Text to Music Generation

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

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:

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:

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:

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]