singing.py

"""
Singing and emotion tags (ModelsLab/omnivoice-singing).

A finetune of k2-fsa/OmniVoice with the same architecture and API, so the
only change needed is the model directory. It adds a [singing] tag and
seven emotion tags.

Every tag was trained on ENGLISH data only (GTSinger English, CREMA-D,
RAVDESS, Expresso). Applying them to German is extrapolation, so each
German case below has an English counterpart: if German singing sounds
wrong, the control tells you whether the model sings badly in general or
only outside its training language.

Licence note: the model card claims Apache 2.0, but this is a finetune of
CC-BY-NC weights and several training sets are NC as well. Treat it as
non-commercial.

Run from the project root (omnivoice_lab):
    uv run --project OmniVoice python scripts/singing.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_SINGING_MODEL_DIR")
OUTPUT_DIR = os.path.join("data", "output", "singing")

# The model card recommends 3.0 (up from the default 2.0) so the tags take
# effect more clearly.
GUIDANCE_SCALE = 3.0

# Tag strengths per the model card:
#   strong : [singing] [happy] [sad] [angry] [nervous] [whisper]
#   weak   : [calm] [excited]   -- only ~190 training clips each
#
# (name, text, language)
CASES = [
    # Baseline: no tag at all. Confirms plain speech still works.
    ("00_plain_de", "Das ist ganz normale Sprache, ohne jeden Tag.", "de"),
    ("01_plain_en", "This is plain speech without any tag at all.", "en"),

    # Singing: the main event, German against an English control.
    ("02_singing_de", "[singing] Hoch auf dem gelben Wagen sitz ich beim Schwager vorn.", "de"),
    ("03_singing_en", "[singing] Twinkle twinkle little star, how I wonder what you are.", "en"),

    # Strong emotion tags, paired DE and EN for direct comparison.
    ("04_happy_de", "[happy] Ich habe gerade die beste Nachricht des ganzen Jahres bekommen!", "de"),
    ("05_happy_en", "[happy] I just got the best news of my entire year!", "en"),
    ("06_sad_de", "[sad] Es ist still geworden, seit du fort bist.", "de"),
    ("07_sad_en", "[sad] It has become so quiet since you have been gone.", "en"),
    ("08_angry_de", "[angry] Das reicht mir jetzt endgültig, so geht das nicht weiter!", "de"),
    ("09_angry_en", "[angry] That is the absolute limit, I cannot take this anymore!", "en"),
    ("10_nervous_de", "[nervous] Ich weiß nicht, ob das wirklich eine gute Idee ist.", "de"),
    ("11_nervous_en", "[nervous] I am really not sure if this is a good idea at all.", "en"),
    ("12_whisper_de", "[whisper] Sei ganz leise, sonst hört uns noch jemand.", "de"),
    ("13_whisper_en", "[whisper] Be very quiet, otherwise someone might hear us.", "en"),

    # Weak tags, included so the difference in reliability is audible.
    ("14_calm_de", "[calm] Atme ruhig weiter, es ist alles in Ordnung.", "de"),
    ("15_excited_de", "[excited] Das hätte ich jetzt wirklich nicht erwartet!", "de"),

    # Combined tags, as shown on the model card.
    ("16_singing_sad_de", "[singing] [sad] Leiser Regen fällt auf kalten Stein, Tage die vergangen sind.", "de"),
    ("17_singing_sad_en", "[singing] [sad] Quiet rain falls on the stone, memories of days now gone.", "en"),
]


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

    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")
    print(f"guidance_scale = {GUIDANCE_SCALE}\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")
    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,
            guidance_scale=GUIDANCE_SCALE,
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - started

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

        print(f"{name:<20} audio {duration:5.2f} s  compute {elapsed:5.2f} s  "
              f"RTF {elapsed / duration:.3f}")

    print(f"\nWrote {len(CASES)} files to {OUTPUT_DIR}")
    print("Compare the _de and _en pairs: if English works and German does not,")
    print("the limit is the training language, not the model.")


if __name__ == "__main__":
    main()