app.py

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import tempfile
import uuid
from pathlib import Path

import gradio as gr
import torch
import soundfile as sf
import torchaudio
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Channels, SampleRate, Subcategory

from data.tokenizer import AudioTokenizer, TextTokenizer
from inference_speech_editing_scale import inference_one_sample
from models.voicecraft import VoiceCraft


model_card = ModelCard(
    name="VoiceCraft",
    description=(
        "Edit English speech while retaining the surrounding recording. "
        "The model is fixed-rate and mono, at 16 kHz, so the result is resampled back to "
        "the input's rate and copied across its channels: the output has the "
        "format it came in with, but no stereo image. "
        "Repository: https://github.com/jasonppy/VoiceCraft Paper: "
        "https://arxiv.org/abs/2403.16973"
    ),
    author="Puyuan Peng, Po-Yao Huang, Abdelrahman Mohamed, David Harwath",
    tags=[Subcategory.EDITING, SampleRate(16000), Channels(1), 'speech-editing', 'domain:speech'],
)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CODEC_RATE = 50
OUTPUT_DIR = Path(tempfile.gettempdir()) / "voicecraft_outputs"
OUTPUT_DIR.mkdir(exist_ok=True)

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model and tokenizers are loaded there alone
if __name__ != "__main__":
    model = VoiceCraft.from_pretrained(
        "pyp1/VoiceCraft_830M_TTSEnhanced", revision="15ff81999431f57897a945a0d261a561a02ea342"
    ).to(DEVICE).eval()
    codec_path = hf_hub_download("pyp1/VoiceCraft", "encodec_4cb2048_giga.th",
                                 revision="52b0f909d652ea016966804767320f13ede39abd")
    audio_tokenizer = AudioTokenizer(device=DEVICE, signature=codec_path)
    text_tokenizer = TextTokenizer(backend="espeak")


# The whole recording is the model's context, and the replacement is generated a token at a
# time, so the cost grows with both, and faster than the text: a 500-character transcript
# measured at about 17 s, while a 1000-character one outran the 60 s reservation below
MAX_INPUT_SECONDS = 60
MAX_TRANSCRIPT_CHARACTERS = 500


def process_fn(audio_path: str, target_text: str, start_time: float, end_time: float) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not audio_path:
        raise gr.Error("Upload a speech recording.")
    target_text = target_text.strip()
    if not target_text:
        raise gr.Error("Enter the full transcript with the replacement text.")
    if len(target_text) > MAX_TRANSCRIPT_CHARACTERS:
        raise gr.Error(
            f"The transcript is {len(target_text)} characters long. Shorten it to "
            f"{MAX_TRANSCRIPT_CHARACTERS} or fewer."
        )

    # soundfile rather than torchaudio, whose ffmpeg backend fails outside a GPU call under
    # ZeroGPU's emulation
    try:
        duration = sf.info(audio_path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration == 0:
        raise gr.Error("The audio file is empty.")
    if duration > MAX_INPUT_SECONDS:
        raise gr.Error(
            f"The recording is {duration:.1f} s long. Trim it to {MAX_INPUT_SECONDS} seconds or "
            "less."
        )
    if not 0 <= start_time < end_time <= duration:
        raise gr.Error(f"Select an edit interval within the {duration:.2f} second recording.")

    return regenerate(audio_path, target_text, start_time, end_time, duration)


# The replacement is generated token by token until the model decides it has ended, and it
# sometimes runs on, stretching a two-second edit into ten seconds of speech. A result is
# accepted when the recording has grown by no more than this many seconds, and otherwise the
# edit is sampled again, up to the limit below, keeping the shortest result.
MAX_GROWTH_SECONDS = 2.0
MAX_ATTEMPTS = 3


@spaces.GPU(duration=60)
def regenerate(audio_path: str, target_text: str, start_time: float, end_time: float,
               duration: float) -> str:
    start = max(start_time - 0.08, 1 / CODEC_RATE)
    end = min(end_time + 0.08, duration)
    mask = torch.tensor([[round(start * CODEC_RATE), round(end * CODEC_RATE)]])
    decode_config = {
        "top_k": 0,
        "top_p": 0.9,
        "temperature": 1.0,
        "stop_repetition": 3,
        "kvcache": 1,
        "codec_sr": CODEC_RATE,
        "silence_tokens": [1388, 1898, 131],
    }
    best = None
    for attempt in range(MAX_ATTEMPTS):
        generated = inference_one_sample(
            model,
            model.args,
            model.args.phn2num,
            text_tokenizer,
            audio_tokenizer,
            audio_path,
            target_text,
            mask,
            DEVICE,
            decode_config,
        )
        candidate = generated[0].detach().cpu().float()
        seconds = candidate.shape[-1] / 16000
        print(f"VoiceCraft: attempt {attempt + 1} gives {seconds:.1f} s for a {duration:.1f} s "
              "recording", flush=True)
        if best is None or candidate.shape[-1] < best.shape[-1]:
            best = candidate
        if seconds <= duration + MAX_GROWTH_SECONDS:
            break
        # A sample that ran on is more likely to stop with silence penalized harder
        decode_config["stop_repetition"] = max(1, decode_config["stop_repetition"] - 1)

    # The model is mono at 16 kHz, so the result is resampled back to the input's rate and copied
    # across its channels
    input_info = sf.info(audio_path)
    audio = best
    if input_info.samplerate != 16000:
        audio = torchaudio.functional.resample(audio, 16000, input_info.samplerate)
    audio = audio[:1].T.repeat(1, input_info.channels)

    input_ext = Path(audio_path).suffix.lower()
    output_path = OUTPUT_DIR / f"{uuid.uuid4().hex}{input_ext if input_ext in LOSSLESS else '.wav'}"
    sf.write(str(output_path), audio.numpy(), input_info.samplerate)
    return str(output_path)


if __name__ == "__main__":
    with gr.Blocks(title="VoiceCraft") as demo:
        build_endpoint(
            model_card=model_card,
            input_components=[
                gr.Audio(type="filepath", label="Speech Recording")
                .harp_required(True)
                .set_info("English speech to edit, up to 60 seconds."),
                gr.Textbox(label="Target Transcript", lines=4)
                .harp_required(True)
                .set_info("The whole transcript of the recording, as it should read after the "
                          "edit, up to 500 characters. Only the words inside the edit interval "
                          "are regenerated, so change those and leave the rest as spoken."),
                gr.Number(label="Edit Start (seconds)", value=0.5, precision=2)
                .set_info("Where the words to replace begin, in seconds from the start of the "
                          "recording. Read the time off the ruler of the input track."),
                gr.Number(label="Edit End (seconds)", value=1.5, precision=2)
                .set_info("Where the words to replace end, in seconds. The interval should cover "
                          "whole words, and a little silence on either side helps."),
            ],
            output_components=[
                gr.Audio(type="filepath", label="Edited Speech")
                .set_info("The edited recording, at its own sample rate and channel count."),
            ],
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)