#!/usr/bin/env python3 """OmniVoice Word-Control โ€” ZeroGPU Space. Voice cloning with explicit word-level acoustic control (WordVoice-5A style: duration / boundary / energy / pitch / tone), on an OmniVoice checkpoint fine-tuned with inline control tokens. UI + inline-tag syntax adapted from hugging-apps/wordvoice-tts; model loading and generation flow adapted from the official k2-fsa/OmniVoice Space. """ import json import logging import os import re logging.basicConfig(level=logging.INFO) import numpy as np import spaces import torch import gradio as gr from omnivoice import OmniVoice, OmniVoiceGenerationConfig from omnivoice.data.word_control import ( BND_CLASSES, TONE_CLASSES, dur_bin, eng_bin, pit_bin, ) CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "multimodalart/omnivoice-word-control") print(f"Loading model from {CHECKPOINT} ...") model = OmniVoice.from_pretrained( CHECKPOINT, device_map="cuda", dtype=torch.float16, load_asr=True, # auto-transcribes the reference clip when no transcript given ) sampling_rate = model.sampling_rate print("Model loaded successfully!") # --------------------------------------------------------------------------- # Inline tag parsing: word[pit:0.8][dur:400] -> <|pit_18|><|dur_9|>word # --------------------------------------------------------------------------- TAG_RE = re.compile(r"\[(dur|bnd|eng|pit|ton)\s*:\s*([-+]?\d+(?:\.\d+)?|[A-Za-z]\w*)\]") def _tags_to_tokens(tags): """Map user-facing tag values to the checkpoint's control tokens (fixed training order: dur, bnd, eng, pit, ton).""" toks = [] if "dur" in tags: # milliseconds toks.append(f"<|dur_{dur_bin(float(tags['dur']) / 1000.0)}|>") if "bnd" in tags and tags["bnd"] in BND_CLASSES: toks.append(f"<|bnd_{BND_CLASSES.index(tags['bnd'])}|>") if "eng" in tags: # 0..1 toks.append(f"<|eng_{eng_bin(float(tags['eng']))}|>") if "pit" in tags: # -1..1 toks.append(f"<|pit_{pit_bin(float(tags['pit']))}|>") if "ton" in tags and tags["ton"] in TONE_CLASSES: toks.append(f"<|ton_{tags['ton']}|>") return "".join(toks) def parse_control_text(text): """Convert `word[tag:val]...` syntax into control-token-annotated text. Returns (model_text, n_tagged_words, table_rows). """ out, rows, n_tagged = [], [], 0 for token in text.split(): found = {m.group(1): m.group(2) for m in TAG_RE.finditer(token)} word = TAG_RE.sub("", token) if found: n_tagged += 1 prefix = _tags_to_tokens(found) out.append(prefix + word) rows.append((word, found, prefix)) else: out.append(word) return " ".join(out), n_tagged, rows def _fmt_plan(rows): if not rows: return "No control tags โ€” the model plans all prosody freely." md = "| word | requested | control tokens |\n|---|---|---|\n" for word, found, prefix in rows: req = ", ".join(f"{k}={v}" for k, v in found.items()) # pipes split GFM table cells even inside code spans tok = prefix.replace("|", "\\|") md += f"| {word.replace('|', '')} | {req} | `{tok}` |\n" return md # --------------------------------------------------------------------------- # Generation # --------------------------------------------------------------------------- def _gen_core( ref_audio, ref_text, text, num_step, guidance_scale, speed, duration, ): if not text or not text.strip(): return None, "Please enter the text to synthesize." if not ref_audio: return None, "Please upload or record a reference audio clip." model_text, n_tagged, rows = parse_control_text(text.strip()) gen_config = OmniVoiceGenerationConfig( num_step=int(num_step or 32), guidance_scale=float(guidance_scale) if guidance_scale is not None else 2.0, ) kw = dict(text=model_text, language="en", generation_config=gen_config) if speed is not None and float(speed) != 1.0: kw["speed"] = float(speed) if duration is not None and float(duration) > 0: kw["duration"] = float(duration) kw["voice_clone_prompt"] = model.create_voice_clone_prompt( ref_audio=ref_audio, ref_text=ref_text.strip() if ref_text and ref_text.strip() else None, ) try: audio = model.generate(**kw) except Exception as e: # noqa: BLE001 return None, f"Error: {type(e).__name__}: {e}" waveform = (np.clip(audio[0], -1.0, 1.0) * 32767).astype(np.int16) info = f"Done โ€” {len(waveform) / sampling_rate:.1f}s generated, {n_tagged} word(s) controlled.\n\n" return (sampling_rate, waveform), info + _fmt_plan(rows) @spaces.GPU(duration=90) def generate_fn(*args): return _gen_core(*args) # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- CSS = "#col-container { max-width: 1100px; margin: 0 auto; }" DESCRIPTION = """# ๐ŸŽ›๏ธ OmniVoice Word-Control โ€” voice cloning with word-level prosody control Zero-shot **voice cloning** with *explicit, decoupled word-level control* over five acoustic dimensions, ร  la [WordVoice](https://huggingface.co/papers/2607.06461) โ€” but built on [OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) (masked-diffusion LM TTS), fine-tuned on [WordVoice-5A](https://huggingface.co/datasets/XXH333/WordVoice-5A) with inline control tokens. Upload a reference clip (+ optional transcript โ€” auto-transcribed if empty), write your text, and attach **inline tags** to any word to steer its prosody. Untagged words are planned freely. [Model](https://huggingface.co/multimodalart/omnivoice-word-control) ยท fine-tune of k2-fsa/OmniVoice """ CONTROL_HELP = """ ### Inline control tags Attach one or more tags immediately after a word โ€” e.g. `crazy[eng:0.9][dur:400]` or `never[pit:0][ton:ffall]`. | Tag | Meaning | Range | |-----|---------|-------| | `[dur:N]` | word duration | milliseconds (40ms steps, 40โ€“2560) | | `[eng:x]` | energy / loudness | `0`โ€“`1` | | `[pit:x]` | pitch (core F0) | `-1`โ€“`1` | | `[bnd:b]` | pause after word | `b0` (none) โ€ฆ `b4` (long) | | `[ton:t]` | pitch contour | `flat, rise, rrise, fall, ffall, peak, valley` | Tags are converted to the checkpoint's control tokens (shown in the output table). Anything you leave untagged is planned by the model. Tip: for strict overall timing, also set the *total duration* slider. """ REF_META = json.load(open(os.path.join(os.path.dirname(__file__), "demo", "ref_en.json"))) REF_WAV = os.path.join(os.path.dirname(__file__), "demo", "ref_en.mp3") EXAMPLES = [ [REF_WAV, REF_META["text"], "I will never[pit:-0.6][ton:ffall] agree to this[bnd:b4], are[pit:0.4] you crazy[eng:0.9][dur:520]?"], [REF_WAV, REF_META["text"], "The quiet[eng:0.2] river drifted slowly[dur:700][ton:fall] under the old stone bridge."], [REF_WAV, REF_META["text"], "This is a zero shot text to speech tool with explicit word level control."], ] with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(): ref_audio = gr.Audio( sources=["upload", "microphone"], type="filepath", label="Reference audio (โ‰ค ~20s)", value=REF_WAV, ) ref_text = gr.Textbox( label="Reference transcript (optional โ€” auto-transcribed if empty)", value=REF_META["text"], lines=2, ) text = gr.Textbox( label="Text to synthesize (attach inline control tags โ€” see reference below)", value=EXAMPLES[0][2], lines=3, ) run = gr.Button("Synthesize", variant="primary") with gr.Accordion("Advanced", open=False): num_step = gr.Slider(4, 64, value=32, step=1, label="Diffusion steps") guidance_scale = gr.Slider(1.0, 6.0, value=2.0, step=0.1, label="Guidance scale") speed = gr.Slider(0.5, 2.0, value=1.0, step=0.05, label="Speed") duration = gr.Slider( 0, 30, value=0, step=0.5, label="Total duration (s) โ€” 0 = auto (set for strict timing with dur tags)", ) with gr.Column(): audio_out = gr.Audio(label="Synthesized audio", autoplay=True) info_out = gr.Markdown(label="Control plan") with gr.Accordion("Control-tag reference", open=False): gr.Markdown(CONTROL_HELP) gr.Examples( examples=EXAMPLES, inputs=[ref_audio, ref_text, text], ) run.click( generate_fn, inputs=[ref_audio, ref_text, text, num_step, guidance_scale, speed, duration], outputs=[audio_out, info_out], api_name="synthesize", ) if __name__ == "__main__": demo.queue(default_concurrency_limit=2).launch(css=CSS, mcp_server=True)