nonverbal.py

"""
Non-verbal symbols of the base OmniVoice model.

The base model has no emotion tags. What it has are non-verbal symbols:
markers that insert a *sound* at that position, rather than colouring the
whole utterance. [laughter] makes the speaker laugh and then talk; it does
not make the sentence sound happy. That is what the singing finetune adds
on top -- see scripts/singing.py.

A cloned voice is used on purpose. With auto voice the speaker would change
on every call, and any difference heard could just as well come from the
new voice instead of from the tag.

Tags whose name ends in -en are English specific and are paired with an
English sentence here. The -ah / -oh / -wa / -yo / -hnn endings describe
vocalisation shapes that most likely come from the Chinese training data,
so on German text they are the least predictable of the set.

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

import os
import time

# 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 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", "nonverbal")

# (name, text, language)
CASES = [
    # Control: the same sentence without any tag, to hear what the tag adds.
    ("00_control_de", "Du hast mich wirklich erwischt, damit hatte ich nicht gerechnet.", "de"),

    # The two that should be most robust.
    ("01_laughter_de", "[laughter] Du hast mich wirklich erwischt, damit hatte ich nicht gerechnet.", "de"),
    ("02_sigh_de", "[sigh] Na gut, dann fangen wir eben noch einmal von vorne an.", "de"),

    # English-specific tags, paired with English text as the name suggests.
    ("03_confirmation_en", "[confirmation-en] Yes, that is exactly what I meant.", "en"),
    ("04_question_en", "[question-en] Are you really sure about this?", "en"),

    # Same two on German text, to see whether the suffix matters in practice.
    ("05_confirmation_de", "[confirmation-en] Ja, genau das habe ich gemeint.", "de"),
    ("06_question_de", "[question-en] Bist du dir da wirklich sicher?", "de"),

    # Vocalisation shapes, probably Chinese in origin.
    ("07_question_ah_de", "[question-ah] Wie meinst du das denn jetzt?", "de"),
    ("08_question_oh_de", "[question-oh] Und das soll ich dir einfach glauben?", "de"),
    ("09_question_ei_de", "[question-ei] Was machst du denn hier?", "de"),
    ("10_question_yi_de", "[question-yi] Hast du das etwa vergessen?", "de"),

    ("11_surprise_ah_de", "[surprise-ah] Das hätte ich jetzt nicht erwartet!", "de"),
    ("12_surprise_oh_de", "[surprise-oh] Schau dir das mal an!", "de"),
    ("13_surprise_wa_de", "[surprise-wa] Das ist ja unglaublich!", "de"),
    ("14_surprise_yo_de", "[surprise-yo] Da bist du ja endlich!", "de"),

    ("15_dissatisfaction_de", "[dissatisfaction-hnn] Also so hatte ich mir das nicht vorgestellt.", "de"),
]


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)
    print(f"Reference: {REF_AUDIO}\n")

    sample_rate = getattr(model, "sampling_rate", None) or 24000

    # Throwaway call so the first real case does not carry the CUDA warm-up.
    model.generate(text="Aufwärmen.", language="de", voice_clone_prompt=prompt)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    for name, text, language in CASES:
        started = time.perf_counter()
        audio = model.generate(
            text=text,
            language=language,
            voice_clone_prompt=prompt,
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - started

        wav = audio[0]
        duration = len(wav) / sample_rate
        sf.write(os.path.join(OUTPUT_DIR, f"{name}.wav"), wav, sample_rate)

        print(f"{name:<24} audio {duration:5.2f} s  compute {elapsed:5.2f} s")

    print(f"\nWrote {len(CASES)} files to {OUTPUT_DIR}")
    print("Start with 00_control vs 01_laughter: if the tag does nothing there,")
    print("it will not work anywhere else either.")


if __name__ == "__main__":
    main()