app.py

#!/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())
        md += f"| {word} | {req} | `{prefix}` |\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)