long_text_batch.py

"""
Long text: one continuous generation vs. sentence batch.

Two ways to voice a paragraph, measured against each other:

  A) naive   -- hand the whole paragraph to generate() as one string. The
                model keeps it in one piece as long as the estimated audio
                stays under audio_chunk_threshold (default 30 s), so a
                12-25 s paragraph lands in a bad middle zone: too long to
                stay natural, too short to trigger the built-in chunking.

  B) batched -- split into sentences, pass them as a list in a single
                generate() call so the GPU processes them together, then
                join the pieces with a short silence. Every sentence stays
                in the range the model handles well, and the shared voice
                clone prompt keeps the voice identical across them.

Both write a wav, so the difference can be judged by ear, not just by RTF.

Run from the project root (omnivoice_lab):
    uv run --project OmniVoice python scripts/long_text_batch.py
"""

import os

# Force huggingface_hub offline BEFORE anything imports it. Every file we
# need is already on disk, so a network call would mean something is wrong.
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import re
import time

import numpy as np
import soundfile as sf
import torch
from dotenv import load_dotenv
from omnivoice import OmniVoice

load_dotenv()

MODEL_DIR = os.getenv("OMNIVOICE_MODEL_DIR")
AUDIO_DIR = os.path.join("data", "audio_ref")
REFERENCE = os.getenv("OMNIVOICE_REF_NAME", "my_voice_berlin_k")

REF_AUDIO = os.path.join(AUDIO_DIR, f"{REFERENCE}.wav")
REF_TEXT_FILE = os.path.join(AUDIO_DIR, f"reference_txt_{REFERENCE}.txt")

OUTPUT_DIR = os.path.join("data", "output", "long_text")

LANGUAGE = "de"

# Silence inserted between sentences when joining. The model already pads
# each piece by pad_duration (0.1 s per side by default), so a little goes
# a long way -- raise it for a more deliberate, slower delivery.
GAP_SECONDS = 0.2

TEXT = (
    "Das Wetter ist heute deutlich besser als gestern, wir könnten spazieren gehen. "
    "Aber vorsichtig beim Laufen, da können die Beine schon mal schwer werden. "
    "Das ist dann nicht so dolle. "
    "Am besten nehmen wir den Weg am Kanal entlang, der ist flach und gut zu gehen. "
    "Wenn es später doch noch regnet, kehren wir einfach irgendwo ein."
)


def split_sentences(text):
    """Split on sentence-ending punctuation followed by whitespace.

    Deliberately simple. It does not know German abbreviations, so "z. B."
    or "Dr. Meyer" would be split in the wrong place -- check the printed
    sentence list if the output sounds chopped.
    """
    text = " ".join(text.split())
    parts = re.split(r"(?<=[.!?])\s+", text)
    return [p.strip() for p in parts if p.strip()]


def join_with_gaps(pieces, sample_rate, gap_seconds):
    """Concatenate audio chunks with a fixed silence between them."""
    gap = np.zeros(int(gap_seconds * sample_rate), dtype=pieces[0].dtype)
    joined = []
    for index, piece in enumerate(pieces):
        if index:
            joined.append(gap)
        joined.append(piece)
    return np.concatenate(joined)


def main():
    if not MODEL_DIR or not os.path.isdir(MODEL_DIR):
        raise SystemExit(
            f"Model directory not found. Please set OMNIVOICE_MODEL_DIR in your environment or .env (current value: {MODEL_DIR})"
        )

    if not os.path.exists(REF_AUDIO):
        raise SystemExit(
            f"\n[ERROR] Reference audio file missing!\n"
            f"Expected file: {REF_AUDIO}\n\n"
            f"How to fix:\n"
            f"1. Record a 3-10s mono WAV (24 kHz recommended) and place it in '{AUDIO_DIR}'.\n"
            f"2. Set OMNIVOICE_REF_NAME='{REFERENCE}' in .env (or adjust REFERENCE in the script).\n"
            f"See {os.path.join(AUDIO_DIR, 'README.md')} for details."
        )

    # Check for transcript under reference_txt_<name>.txt or <name>.txt
    transcript_path = REF_TEXT_FILE
    if not os.path.exists(transcript_path):
        alt_path = os.path.join(AUDIO_DIR, f"{REFERENCE}.txt")
        if os.path.exists(alt_path):
            transcript_path = alt_path
        else:
            raise SystemExit(
                f"\n[ERROR] Reference transcript file missing!\n"
                f"Expected file: {REF_TEXT_FILE} (or {alt_path})\n\n"
                f"How to fix:\n"
                f"1. Create a text file containing the exact spoken text of '{REFERENCE}.wav'.\n"
                f"2. Save it as '{REF_TEXT_FILE}'.\n"
                f"Providing the transcript is required to prevent unwanted Whisper model downloads."
            )

    with open(transcript_path, encoding="utf-8") as handle:
        ref_text = handle.read().strip()

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print(f"Loading model from {MODEL_DIR} ...")
    started = time.perf_counter()
    model = OmniVoice.from_pretrained(
        MODEL_DIR,
        device_map="cuda:0",
        dtype=torch.float16,
    )
    print(f"Model loaded in {time.perf_counter() - started:.1f} s")

    prompt = model.create_voice_clone_prompt(ref_audio=REF_AUDIO, ref_text=ref_text)
    sample_rate = getattr(model, "sampling_rate", None) or 24000

    sentences = split_sentences(TEXT)
    print(f"\nSplit into {len(sentences)} sentences:")
    for index, sentence in enumerate(sentences, start=1):
        print(f"  {index}. {sentence}")

    # Throwaway call so neither method pays the CUDA warm-up. Without this
    # whichever runs first looks slower for reasons that have nothing to do
    # with the approach being compared.
    print("\nWarming up ...")
    model.generate(text="Kurzer Aufwärmsatz.", language=LANGUAGE, voice_clone_prompt=prompt)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    # --- A: whole paragraph in one call -------------------------------
    started = time.perf_counter()
    naive_audio = model.generate(
        text=" ".join(TEXT.split()),
        language=LANGUAGE,
        voice_clone_prompt=prompt,
    )
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    naive_time = time.perf_counter() - started

    naive_wav = naive_audio[0]
    naive_duration = len(naive_wav) / sample_rate
    naive_path = os.path.join(OUTPUT_DIR, "naive.wav")
    sf.write(naive_path, naive_wav, sample_rate)

    # --- B: sentences as one batch ------------------------------------
    started = time.perf_counter()
    batched_audio = model.generate(
        text=sentences,
        language=LANGUAGE,
        voice_clone_prompt=prompt,
    )
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    batched_time = time.perf_counter() - started

    batched_wav = join_with_gaps(batched_audio, sample_rate, GAP_SECONDS)
    batched_duration = len(batched_wav) / sample_rate
    batched_path = os.path.join(OUTPUT_DIR, "batched.wav")
    sf.write(batched_path, batched_wav, sample_rate)

    # Keep the individual pieces too, to locate a bad sentence by ear.
    for index, piece in enumerate(batched_audio, start=1):
        sf.write(os.path.join(OUTPUT_DIR, f"sentence_{index:02d}.wav"), piece, sample_rate)

    print("\n--- results ---")
    print(f"A naive    : audio {naive_duration:5.2f} s  compute {naive_time:5.2f} s  "
          f"RTF {naive_time / naive_duration:.3f}  -> {naive_path}")
    print(f"B batched  : audio {batched_duration:5.2f} s  compute {batched_time:5.2f} s  "
          f"RTF {batched_time / batched_duration:.3f}  -> {batched_path}")
    print(f"\nbatch is {naive_time / batched_time:.1f}x faster in compute "
          f"({len(sentences)} sentences in one call)")
    print(f"individual sentences written to {OUTPUT_DIR}")
    print("\nNow listen to both -- speed is only half the question.")


if __name__ == "__main__":
    main()