diff --git a/omnivoice/cli/demo.py b/omnivoice/cli/demo.py index 8dd3f2c7..091103f5 100644 --- a/omnivoice/cli/demo.py +++ b/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,6 +25,8 @@ import argparse import logging +import os +import tempfile from typing import Any, Dict import gradio as gr @@ -32,6 +34,7 @@ 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, diff --git a/omnivoice/cli/infer.py b/omnivoice/cli/infer.py index ce4657da..8e4eba16 100644 --- a/omnivoice/cli/infer.py +++ b/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}") diff --git a/omnivoice/cli/infer_batch.py b/omnivoice/cli/infer_batch.py index b1d19d60..374e8bcb 100644 --- a/omnivoice/cli/infer_batch.py +++ b/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( ( diff --git a/omnivoice/utils/audio.py b/omnivoice/utils/audio.py index 02c40100..0ccd4d1b 100644 --- a/omnivoice/utils/audio.py +++ b/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, diff --git a/start-demo.bat b/start-demo.bat new file mode 100644 index 00000000..04fecc40 --- /dev/null +++ b/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