feat(cli): add MP3 output support alongside WAV

omnivoice/cli/demo.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python3
 # Copyright    2026  Xiaomi Corp.        (authors:  Han Zhu)
 #
 # See ../../LICENSE for clarification regarding multiple authors
@@ -25,13 +25,16 @@
 
 import argparse
 import logging
+import os
+import tempfile
 from typing import Any, Dict
 
 import gradio as gr
 import numpy as np
 import torch
 
 from omnivoice import OmniVoice, OmniVoiceGenerationConfig
+from omnivoice.utils.audio import save_audio
 from omnivoice.utils.common import get_best_device
 from omnivoice.utils.lang_map import LANG_NAMES, lang_display_name
 
@@ -170,6 +173,7 @@ def _gen_core(
         duration,
         preprocess_prompt,
         postprocess_output,
+        output_format,
         mode,
         ref_text=None,
     ):
@@ -211,8 +215,26 @@ def _gen_core(
         except Exception as e:
             return None, f"Error: {type(e).__name__}: {e}"
 
-        waveform = (audio[0] * 32767).astype(np.int16)
-        return (sampling_rate, waveform), "Done."
+        if output_format == "mp3":
+            # Save to a system temp file and return the filepath.
+            # Gradio serves files from the system temp directory.
+            tmp = tempfile.NamedTemporaryFile(
+                suffix=".mp3", delete=False, dir=tempfile.gettempdir()
+            )
+            tmp.close()
+            try:
+                save_audio(audio[0], tmp.name, sampling_rate)
+            except Exception:
+                try:
+                    os.unlink(tmp.name)
+                except OSError:
+                    pass
+                return None, "Failed to encode MP3."
+            return tmp.name, "Done. (MP3)"
+        else:
+            # WAV: return numpy tuple (Gradio's default, plays in browser).
+            waveform = (audio[0] * 32767).astype(np.int16)
+            return (sampling_rate, waveform), "Done."
 
     # Allow external wrappers (e.g. spaces.GPU for ZeroGPU Spaces)
     _gen = generate_fn if generate_fn is not None else _gen_core
@@ -295,7 +317,7 @@ def _gen_settings():
             )
         return ns, gs, dn, sp, du, pp, po
 
-    with gr.Blocks(theme=theme, css=css, title="OmniVoice Demo") as demo:
+    with gr.Blocks(title="OmniVoice Demo") as demo:
         gr.Markdown(
             """
 # OmniVoice Demo
@@ -350,16 +372,21 @@ def _gen_settings():
                             vc_pp,
                             vc_po,
                         ) = _gen_settings()
+                        vc_format = gr.Dropdown(
+                            label="Output format / 输出格式",
+                            choices=["wav", "mp3"],
+                            value="wav",
+                            info="WAV (lossless) or MP3 128kbps (~3x smaller).",
+                        )
                         vc_btn = gr.Button("Generate / 生成", variant="primary")
                     with gr.Column(scale=1):
                         vc_audio = gr.Audio(
                             label="Output Audio / 合成结果",
-                            type="numpy",
                         )
                         vc_status = gr.Textbox(label="Status / 状态", lines=2)
 
                 def _clone_fn(
-                    text, lang, ref_aud, ref_text, instruct, ns, gs, dn, sp, du, pp, po
+                    text, lang, ref_aud, ref_text, instruct, ns, gs, dn, sp, du, pp, po, fmt
                 ):
                     return _gen(
                         text,
@@ -373,6 +400,7 @@ def _clone_fn(
                         du,
                         pp,
                         po,
+                        fmt,
                         mode="clone",
                         ref_text=ref_text or None,
                     )
@@ -392,6 +420,7 @@ def _clone_fn(
                         vc_du,
                         vc_pp,
                         vc_po,
+                        vc_format,
                     ],
                     outputs=[vc_audio, vc_status],
                 )
@@ -430,11 +459,16 @@ def _clone_fn(
                             vd_pp,
                             vd_po,
                         ) = _gen_settings()
+                        vd_format = gr.Dropdown(
+                            label="Output format / 输出格式",
+                            choices=["wav", "mp3"],
+                            value="wav",
+                            info="WAV (lossless) or MP3 128kbps (~3x smaller).",
+                        )
                         vd_btn = gr.Button("Generate / 生成", variant="primary")
                     with gr.Column(scale=1):
                         vd_audio = gr.Audio(
                             label="Output Audio / 合成结果",
-                            type="numpy",
                         )
                         vd_status = gr.Textbox(label="Status / 状态", lines=2)
 
@@ -460,7 +494,7 @@ def _build_instruct(groups):
                             parts.append(v)
                     return ", ".join(parts)
 
-                def _design_fn(text, lang, ns, gs, dn, sp, du, pp, po, *groups):
+                def _design_fn(text, lang, ns, gs, dn, sp, du, pp, po, fmt, *groups):
                     return _gen(
                         text,
                         lang,
@@ -473,6 +507,7 @@ def _design_fn(text, lang, ns, gs, dn, sp, du, pp, po, *groups):
                         du,
                         pp,
                         po,
+                        fmt,
                         mode="design",
                     )
 
@@ -488,12 +523,13 @@ def _design_fn(text, lang, ns, gs, dn, sp, du, pp, po, *groups):
                         vd_du,
                         vd_pp,
                         vd_po,
+                        vd_format,
                     ]
                     + vd_groups,
                     outputs=[vd_audio, vd_status],
                 )
 
-    return demo
+    return demo, theme, css
 
 
 # ---------------------------------------------------------------------------
@@ -525,9 +561,11 @@ def main(argv=None) -> int:
     )
     print("Model loaded.")
 
-    demo = build_demo(model, checkpoint)
+    demo, theme, css = build_demo(model, checkpoint)
 
     demo.queue().launch(
+        theme=theme,
+        css=css,
         server_name=args.ip,
         server_port=args.port,
         share=args.share,

omnivoice/cli/infer.py

@@ -4,29 +4,30 @@
 voice design, or auto voice.
 
 Usage:
-    # Voice cloning
+    # Voice cloning (WAV output)
     omnivoice-infer --model k2-fsa/OmniVoice \
         --text "Hello, this is a text for text-to-speech." \
         --ref_audio ref.wav --ref_text "Reference transcript." --output out.wav
 
-    # Voice design
+    # Voice design (MP3 output)
     omnivoice-infer --model k2-fsa/OmniVoice \
         --text "Hello, this is a text for text-to-speech." \
-        --instruct "male, British accent" --output out.wav
+        --instruct "male, British accent" --output out.mp3
 
-    # Auto voice
+    # Auto voice (MP3, explicit format)
     omnivoice-infer --model k2-fsa/OmniVoice \
-        --text "Hello, this is a text for text-to-speech." --output out.wav
+        --text "Hello, this is a text for text-to-speech." \
+        --output out.mp3 --format mp3
 """
 
 import argparse
 import logging
+import os
 
 import torch
 
-import soundfile as sf
-
 from omnivoice.models.omnivoice import OmniVoice
+from omnivoice.utils.audio import save_audio
 from omnivoice.utils.common import get_best_device, str2bool
 
 
@@ -51,7 +52,14 @@ def get_parser() -> argparse.ArgumentParser:
         "--output",
         type=str,
         required=True,
-        help="Output WAV file path.",
+        help="Output file path. Extension determines format: .wav (WAV) or .mp3 (MP3 128kbps).",
+    )
+    parser.add_argument(
+        "--format",
+        type=str,
+        default=None,
+        choices=["wav", "mp3"],
+        help="Output format (overrides extension heuristic). Default: inferred from extension.",
     )
     # Voice cloning
     parser.add_argument(
@@ -141,7 +149,14 @@ def main():
         class_temperature=args.class_temperature,
     )
 
-    sf.write(args.output, audios[0], model.sampling_rate)
+    # Honour --format flag by ensuring the output filename has the right extension
+    if args.format:
+        base = os.path.splitext(args.output)[0]
+        ext = f".{args.format}"
+        if args.output != base + ext:
+            args.output = base + ext
+
+    save_audio(audios[0], args.output, model.sampling_rate)
     logging.info(f"Saved to {args.output}")
 
 

omnivoice/cli/infer_batch.py

@@ -29,6 +29,8 @@
     Voice cloning:   "ref_audio", "ref_text"
     Voice design:    "instruct"
     Optional:        "language_id", "duration", "speed"
+
+Output format defaults to WAV; use --format mp3 for MP3 (128 kbps).
 """
 
 import argparse
@@ -45,9 +47,8 @@
 from tqdm import tqdm
 
 from omnivoice.models.omnivoice import OmniVoice
-import soundfile as sf
 
-from omnivoice.utils.audio import load_audio
+from omnivoice.utils.audio import load_audio, save_audio
 from omnivoice.utils.common import get_best_device_with_count, str2bool
 from omnivoice.utils.data_utils import read_test_list
 from omnivoice.utils.duration import RuleDurationEstimator
@@ -88,6 +89,13 @@ def get_parser():
         required=True,
         help="Directory to save the generated audio files.",
     )
+    parser.add_argument(
+        "--format",
+        type=str,
+        default="wav",
+        choices=["wav", "mp3"],
+        help="Output format for generated audio files (default: wav).",
+    )
     parser.add_argument(
         "--num_step",
         type=int,
@@ -387,9 +395,10 @@ def run_inference_batch(
     batch_synth_time = time.time() - start_time
 
     results = []
+    ext = f".{format}"
     for save_name, audio in zip(save_names, audios):
-        save_path = os.path.join(res_dir, save_name + ".wav")
-        sf.write(save_path, audio, worker_model.sampling_rate)
+        save_path = os.path.join(res_dir, save_name + ext)
+        save_audio(audio, save_path, worker_model.sampling_rate)
         audio_duration = audio.shape[-1] / worker_model.sampling_rate
         results.append(
             (

omnivoice/utils/audio.py

@@ -26,6 +26,7 @@
 
 import io
 import logging
+from pathlib import Path
 
 import numpy as np
 import soundfile as sf
@@ -298,6 +299,45 @@ def trim_long_audio(
     return audiosegment_to_numpy(trimmed)
 
 
+def save_audio(audio: np.ndarray, save_path: str, sample_rate: int) -> None:
+    """Save a waveform to disk as WAV or MP3.
+
+    Parameters:
+        audio: numpy array of shape (C, T) with dtype float32 in [-1, 1].
+        save_path: destination file path.  Extension determines the format:
+            ``.wav`` → WAV (16-bit PCM), ``.mp3`` → MP3 (128 kbps).
+        sample_rate: sample rate in Hz (e.g. 24000).
+
+    Raises:
+        ValueError: if the file extension is not recognised.
+    """
+    ext = Path(save_path).suffix.lower()
+
+    if ext == ".wav":
+        # soundfile natively supports WAV — fast, no ffmpeg dependency.
+        sf.write(save_path, audio, sample_rate)
+
+    elif ext == ".mp3":
+        # pydub (wraps ffmpeg) handles MP3 encoding with bitrate control.
+        from pydub import AudioSegment
+
+        audio_int = (audio * 32768.0).clip(-32768, 32767).astype(np.int16)
+        if audio_int.shape[0] > 1:
+            audio_int = audio_int.T.flatten()  # interleave channels
+        segment = AudioSegment(
+            data=audio_int.tobytes(),
+            sample_width=2,
+            frame_rate=sample_rate,
+            channels=audio_int.shape[0] if audio_int.ndim > 1 else 1,
+        )
+        segment.export(save_path, format="mp3", bitrate="128k")
+
+    else:
+        raise ValueError(
+            f"Unsupported output format '{ext}'. Use '.wav' or '.mp3'."
+        )
+
+
 def cross_fade_chunks(
     chunks: list[np.ndarray],
     sample_rate: int,

start-demo.bat

@@ -0,0 +1,4 @@
+@echo off
+cd /d "%~dp0"
+.venv\Scripts\python -m omnivoice.cli.demo --ip 127.0.0.1 --port 8001
+pause