cpu-variant/app.py

"""
OmniVoice on CPU — Hugging Face Space
=====================================

A CPU-only Gradio front end for k2-fsa/OmniVoice (0.6B zero-shot TTS).

Design notes for CPU:
  - float32 everywhere (fp16 CPU kernels are slow or unsupported).
  - Model is loaded once at startup and reused across requests.
  - Torch thread count is pinned to the vCPUs the Space actually has.
  - Diffusion steps default to 16 instead of 32 (roughly halves compute).
  - Text length is capped so a single request cannot hog the worker.
"""

import os
import time
import tempfile
import threading

# Must be set before torch is imported to take effect on some builds.
_CPU_COUNT = os.cpu_count() or 2
os.environ.setdefault("OMP_NUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("MKL_NUM_THREADS", str(_CPU_COUNT))
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import gradio as gr
import numpy as np
import soundfile as sf
import torch

from omnivoice import OmniVoice

# --------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------

MODEL_ID = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
SAMPLE_RATE = 24000
MAX_CHARS = int(os.environ.get("MAX_CHARS", "300"))
MAX_REF_SECONDS = 12.0

torch.set_num_threads(_CPU_COUNT)
torch.set_grad_enabled(False)

_model = None
_model_lock = threading.Lock()


def get_model():
    """Load the model once, on first use, and reuse it afterwards."""
    global _model
    if _model is None:
        with _model_lock:
            if _model is None:
                print(f"Loading {MODEL_ID} on CPU with {_CPU_COUNT} threads...")
                t0 = time.time()
                _model = OmniVoice.from_pretrained(
                    MODEL_ID,
                    device_map="cpu",
                    dtype=torch.float32,
                )
                print(f"Model ready in {time.time() - t0:.1f}s")
    return _model


# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------


def _write_wav(audio: np.ndarray) -> str:
    path = tempfile.mktemp(suffix=".wav")
    sf.write(path, audio, SAMPLE_RATE)
    return path


def _check_text(text: str) -> str:
    text = (text or "").strip()
    if not text:
        raise gr.Error("Enter some text to speak.")
    if len(text) > MAX_CHARS:
        raise gr.Error(
            f"Text is {len(text)} characters. Keep it under {MAX_CHARS} "
            "so CPU generation finishes in a reasonable time."
        )
    return text


def _check_ref_audio(path: str) -> str:
    if not path:
        raise gr.Error("Upload or record a reference clip first.")
    info = sf.info(path)
    if info.duration > MAX_REF_SECONDS:
        raise gr.Error(
            f"Reference clip is {info.duration:.1f}s. Trim it to "
            f"{MAX_REF_SECONDS:.0f}s or less — long references slow CPU "
            "inference and degrade cloning quality."
        )
    return path


def _report(elapsed: float, audio: np.ndarray, steps: int) -> str:
    duration = len(audio) / SAMPLE_RATE
    rtf = elapsed / duration if duration else 0.0
    return (
        f"Generated {duration:.1f}s of audio in {elapsed:.1f}s "
        f"(RTF {rtf:.2f}) at {steps} diffusion steps, {_CPU_COUNT} CPU threads."
    )


# --------------------------------------------------------------------------
# Generation modes
# --------------------------------------------------------------------------


def clone_voice(text, ref_audio, ref_text, steps, speed, progress=gr.Progress()):
    text = _check_text(text)
    ref_audio = _check_ref_audio(ref_audio)

    progress(0.1, desc="Loading model")
    model = get_model()

    if not (ref_text or "").strip():
        progress(0.2, desc="Transcribing reference with Whisper (slow on CPU)")

    progress(0.35, desc="Generating speech")
    t0 = time.time()
    audio = model.generate(
        text=text,
        ref_audio=ref_audio,
        ref_text=(ref_text or "").strip() or None,
        num_step=int(steps),
        speed=float(speed),
    )[0]
    elapsed = time.time() - t0

    return _write_wav(audio), _report(elapsed, audio, int(steps))


def design_voice(text, gender, age, pitch, accent, extra, steps, speed,
                 progress=gr.Progress()):
    text = _check_text(text)

    attributes = [a for a in (gender, age, pitch, accent) if a and a != "any"]
    if (extra or "").strip():
        attributes.append(extra.strip())
    if not attributes:
        raise gr.Error("Pick at least one voice attribute, or use the Auto voice tab.")
    instruct = ", ".join(attributes)

    progress(0.1, desc="Loading model")
    model = get_model()

    progress(0.35, desc="Generating speech")
    t0 = time.time()
    audio = model.generate(
        text=text,
        instruct=instruct,
        num_step=int(steps),
        speed=float(speed),
    )[0]
    elapsed = time.time() - t0

    return _write_wav(audio), f"Voice: {instruct}\n" + _report(elapsed, audio, int(steps))


def auto_voice(text, steps, speed, progress=gr.Progress()):
    text = _check_text(text)

    progress(0.1, desc="Loading model")
    model = get_model()

    progress(0.35, desc="Generating speech")
    t0 = time.time()
    audio = model.generate(
        text=text,
        num_step=int(steps),
        speed=float(speed),
    )[0]
    elapsed = time.time() - t0

    return _write_wav(audio), _report(elapsed, audio, int(steps))


# --------------------------------------------------------------------------
# Interface
# --------------------------------------------------------------------------

CSS = """
.gradio-container { max-width: 1040px !important; }
#header h1 { margin-bottom: 0.15rem; font-weight: 650; letter-spacing: -0.01em; }
#header p { margin-top: 0; opacity: 0.72; }
#speed-note { font-size: 0.85rem; opacity: 0.7; }
footer { visibility: hidden; }
"""


def settings_row():
    """Shared generation controls. Returns (steps, speed)."""
    with gr.Row():
        steps = gr.Slider(
            8, 32, value=16, step=8,
            label="Diffusion steps",
            info="Fewer steps = faster, slightly rougher. 16 is a good CPU default.",
        )
        speed = gr.Slider(
            0.7, 1.4, value=1.0, step=0.05,
            label="Speaking rate",
            info="Above 1.0 speaks faster and shortens the audio.",
        )
    return steps, speed


with gr.Blocks(title="OmniVoice on CPU", css=CSS, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# OmniVoice on CPU\n"
        "Zero-shot text to speech in 600+ languages, running on a CPU Space. "
        "Generation takes roughly as long as the clip itself — a 10 second "
        "line lands in about 10 to 25 seconds. The first request also pays a "
        "one-time model load.",
        elem_id="header",
    )

    with gr.Tabs():
        # ---------------- Voice cloning ----------------
        with gr.Tab("Clone a voice"):
            with gr.Row():
                with gr.Column(scale=3):
                    clone_text = gr.Textbox(
                        label="Text to speak",
                        placeholder="Type what the cloned voice should say.",
                        lines=3,
                        max_lines=6,
                    )
                    clone_ref_audio = gr.Audio(
                        label="Reference clip (3–10 seconds)",
                        type="filepath",
                        sources=["upload", "microphone"],
                    )
                    clone_ref_text = gr.Textbox(
                        label="What the reference clip says",
                        placeholder="Leave blank to transcribe it automatically — "
                                    "much slower on CPU.",
                        lines=2,
                    )
                    clone_steps, clone_speed = settings_row()
                    clone_btn = gr.Button("Generate speech", variant="primary")
                with gr.Column(scale=2):
                    clone_out = gr.Audio(label="Result", type="filepath")
                    clone_status = gr.Textbox(label="Timing", lines=3, interactive=False)

            clone_btn.click(
                clone_voice,
                inputs=[clone_text, clone_ref_audio, clone_ref_text,
                        clone_steps, clone_speed],
                outputs=[clone_out, clone_status],
            )

        # ---------------- Voice design ----------------
        with gr.Tab("Design a voice"):
            with gr.Row():
                with gr.Column(scale=3):
                    design_text = gr.Textbox(
                        label="Text to speak",
                        placeholder="Type what the designed voice should say.",
                        lines=3,
                        max_lines=6,
                    )
                    with gr.Row():
                        d_gender = gr.Dropdown(
                            ["any", "male", "female"], value="female", label="Gender")
                        d_age = gr.Dropdown(
                            ["any", "child", "young", "middle-aged", "elderly"],
                            value="any", label="Age")
                    with gr.Row():
                        d_pitch = gr.Dropdown(
                            ["any", "very low pitch", "low pitch", "medium pitch",
                             "high pitch", "very high pitch"],
                            value="any", label="Pitch")
                        d_accent = gr.Dropdown(
                            ["any", "american accent", "british accent",
                             "australian accent", "indian accent"],
                            value="any", label="Accent")
                    d_extra = gr.Textbox(
                        label="Other attributes",
                        placeholder="e.g. whisper, 四川话",
                    )
                    design_steps, design_speed = settings_row()
                    design_btn = gr.Button("Generate speech", variant="primary")
                with gr.Column(scale=2):
                    design_out = gr.Audio(label="Result", type="filepath")
                    design_status = gr.Textbox(label="Timing", lines=3, interactive=False)

            gr.Markdown(
                "Voice design was trained on Chinese and English only. It "
                "generalizes to other languages but results get unpredictable "
                "for low-resource ones."
            )

            design_btn.click(
                design_voice,
                inputs=[design_text, d_gender, d_age, d_pitch, d_accent,
                        d_extra, design_steps, design_speed],
                outputs=[design_out, design_status],
            )

        # ---------------- Auto voice ----------------
        with gr.Tab("Auto voice"):
            with gr.Row():
                with gr.Column(scale=3):
                    auto_text = gr.Textbox(
                        label="Text to speak",
                        placeholder="The model picks a voice for you.",
                        lines=3,
                        max_lines=6,
                    )
                    auto_steps, auto_speed = settings_row()
                    auto_btn = gr.Button("Generate speech", variant="primary")
                with gr.Column(scale=2):
                    auto_out = gr.Audio(label="Result", type="filepath")
                    auto_status = gr.Textbox(label="Timing", lines=3, interactive=False)

            auto_btn.click(
                auto_voice,
                inputs=[auto_text, auto_steps, auto_speed],
                outputs=[auto_out, auto_status],
            )

    gr.Markdown(
        "Inline controls you can use in any text box: `[laughter]`, `[sigh]`, "
        "`[question-en]`, `[surprise-ah]` and other non-verbal tags; CMU "
        "phonemes in brackets for English (`[B EY1 S]`); pinyin with tone "
        "numbers for Chinese (`打ZHE2`).\n\n"
        "Do not clone anyone's voice without their permission."
    )


if __name__ == "__main__":
    # Queue with a single worker: CPU inference is not something to run
    # concurrently on 2 vCPUs.
    demo.queue(max_size=12, default_concurrency_limit=1).launch()