app.py

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "zipvoice"))

import json
import logging
import tempfile

import soundfile as sf
import torch
import torchaudio
from soe_vinorm import SoeNormalizer


# torchaudio >= 2.9 removed the legacy backends and now requires torchcodec,
# which is GPU-only. Patch load/save to use soundfile for CPU compatibility.
def _sf_load(
    uri,
    frame_offset=0,
    num_frames=-1,
    normalize=True,
    channels_first=True,
    format=None,
    buffer_size=4096,
    backend=None,
):
    data, sr = sf.read(str(uri), dtype="float32", always_2d=True)
    tensor = torch.from_numpy(data.T)  # (channels, frames)
    if frame_offset > 0:
        tensor = tensor[:, frame_offset:]
    if num_frames > 0:
        tensor = tensor[:, :num_frames]
    return tensor, sr


def _sf_save(uri, src, sample_rate, **kwargs):
    data = src.numpy().T  # (frames, channels)
    sf.write(str(uri), data, sample_rate)


torchaudio.load = _sf_load
torchaudio.save = _sf_save

import gradio as gr
import spaces
from huggingface_hub import hf_hub_download
from lhotse.utils import fix_random_seed
from zipvoice.bin.infer_zipvoice import generate_sentence, get_vocoder
from zipvoice.models.zipvoice_distill import ZipVoiceDistill
from zipvoice.tokenizer.tokenizer import EspeakTokenizer
from zipvoice.utils.checkpoint import load_checkpoint
from zipvoice.utils.feature import VocosFbank

logging.basicConfig(
    format="%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s",
    level=logging.INFO,
)

HF_REPO = "sleeper371/zipvoice_vi"
SAMPLING_RATE = 24000
_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
AUDIO_PROMPTS_DIR = os.path.join(_BASE_DIR, "audio_prompts")
_model = None
_vocoder = None
_tokenizer = None
_feature_extractor = None

normalizer = SoeNormalizer()


# ---------------------------------------------------------------------------
# Preset voices
# ---------------------------------------------------------------------------


def _load_preset_voices() -> dict:
    """Scan audio_prompts/ and return {display_name: (wav_path, transcript)}."""
    voices = {}
    if not os.path.isdir(AUDIO_PROMPTS_DIR):
        return voices
    for wav_file in sorted(os.listdir(AUDIO_PROMPTS_DIR)):
        if not wav_file.endswith(".wav"):
            continue
        stem = wav_file[:-4]
        txt_file = os.path.join(AUDIO_PROMPTS_DIR, stem + ".txt")
        if not os.path.isfile(txt_file):
            continue
        with open(txt_file, encoding="utf-8") as f:
            transcript = f.read().strip()
        display = stem.replace("_", " ").title()
        voices[display] = (os.path.join(AUDIO_PROMPTS_DIR, wav_file), transcript)
    return voices


PRESET_VOICES = _load_preset_voices()
PRESET_NAMES = list(PRESET_VOICES.keys())


def select_preset_voice(name: str):
    if name and name in PRESET_VOICES:
        wav_path, transcript = PRESET_VOICES[name]
        return wav_path, transcript
    return None, ""


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------


def _resolve_model_file(filename: str) -> str:
    local = os.path.join(_BASE_DIR, filename)
    if os.path.isfile(local):
        logging.info("Using local file: %s", local)
        return local
    logging.info("Downloading %s from %s ...", filename, HF_REPO)
    return hf_hub_download(HF_REPO, filename=filename)


def load_models():
    global _model, _vocoder, _tokenizer, _feature_extractor

    if _model is not None:
        return

    model_ckpt = _resolve_model_file("model.pt")
    model_config_path = _resolve_model_file("model.json")
    token_file = _resolve_model_file("tokens.txt")

    _tokenizer = EspeakTokenizer(token_file=token_file, lang="vi")
    tokenizer_config = {
        "vocab_size": _tokenizer.vocab_size,
        "pad_id": _tokenizer.pad_id,
    }

    with open(model_config_path, "r") as f:
        model_config = json.load(f)

    _model = ZipVoiceDistill(**model_config["model"], **tokenizer_config)
    load_checkpoint(filename=model_ckpt, model=_model, strict=True)
    _model.eval()

    _vocoder = get_vocoder()
    _vocoder.eval()

    _feature_extractor = VocosFbank()

    logging.info("All models loaded successfully.")


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


@spaces.GPU
@torch.inference_mode()
def synthesize(prompt_wav, prompt_text, text, speed, num_step, seed):
    if not prompt_wav:
        return None, "Please upload a reference audio file or select a preset voice."
    if not prompt_text.strip():
        return None, "Please enter the transcription of the reference audio."
    if not text.strip():
        return None, "Please enter text to synthesize."

    text = normalizer.normalize(text)

    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        load_models()
        _model.to(device)
        _vocoder.to(device)
        fix_random_seed(int(seed))

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            output_path = f.name

        generate_sentence(
            save_path=output_path,
            prompt_text=prompt_text,
            prompt_wav=prompt_wav,
            text=text,
            model=_model,
            vocoder=_vocoder,
            tokenizer=_tokenizer,
            feature_extractor=_feature_extractor,
            device=device,
            num_step=int(num_step),
            guidance_scale=3.0,
            speed=float(speed),
            t_shift=0.5,
            target_rms=0.1,
            feat_scale=0.1,
            sampling_rate=SAMPLING_RATE,
            max_duration=30,
            remove_long_sil=False,
        )

        return output_path, "Generation complete."
    except Exception as e:
        logging.exception("Error during synthesis")
        return None, f"Error: {e}"


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

_default_voice = PRESET_NAMES[0] if PRESET_NAMES else None
_default_wav, _default_text = select_preset_voice(_default_voice)

with gr.Blocks(title="ZipVoice Vietnamese TTS") as demo:
    gr.Markdown(
        """
# ZipVoice Vietnamese TTS

Voice-cloning text-to-speech for Vietnamese, powered by a distilled
[ZipVoice](https://github.com/k2-fsa/ZipVoice) model fine-tuned on Vietnamese data.

**How to use:**
1. Pick a preset voice **or** upload your own reference audio and type its transcription.
2. Enter the Vietnamese text you want synthesized.
3. Click **Generate**.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            if PRESET_NAMES:
                preset_dropdown = gr.Dropdown(
                    choices=PRESET_NAMES,
                    value=_default_voice,
                    label="Preset Voice",
                    info="Select a built-in voice to pre-fill the reference audio and transcription.",
                )

            prompt_wav = gr.Audio(
                label="Reference Audio (1–5 seconds)",
                type="filepath",
                value=_default_wav,
            )
            prompt_text = gr.Textbox(
                label="Reference Audio Transcription",
                placeholder="What is said in the reference audio...",
                value=_default_text,
            )
            text = gr.Textbox(
                label="Text to Synthesize",
                placeholder="Vietnamese text to convert to speech...",
                lines=4,
            )
            with gr.Accordion("Advanced Options", open=False):
                speed = gr.Slider(
                    minimum=0.5,
                    maximum=2.0,
                    value=1.0,
                    step=0.05,
                    label="Speed (1.0 = normal)",
                )
                num_step = gr.Slider(
                    minimum=1,
                    maximum=32,
                    value=4,
                    step=1,
                    label="Sampling Steps (higher = slower but potentially better)",
                )
                seed = gr.Number(value=666, label="Random Seed", precision=0)

            generate_btn = gr.Button("Generate", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_audio = gr.Audio(label="Generated Speech", type="filepath")
            status_box = gr.Textbox(label="Status", interactive=False)

    if PRESET_NAMES:
        preset_dropdown.change(
            fn=select_preset_voice,
            inputs=[preset_dropdown],
            outputs=[prompt_wav, prompt_text],
        )

    generate_btn.click(
        fn=lambda: (None, "Generating… please wait."),
        inputs=[],
        outputs=[output_audio, status_box],
        queue=False,
    ).then(
        fn=synthesize,
        inputs=[prompt_wav, prompt_text, text, speed, num_step, seed],
        outputs=[output_audio, status_box],
    )

if __name__ == "__main__":
    demo.queue().launch()