debug_singing_duration.py

"""
Diagnostic: does voice cloning starve singing of time?

Hypothesis
----------
With a voice clone prompt, _estimate_target_tokens() derives the output
length from the reference: reference audio length divided by reference
transcript length, i.e. from the SPEAKING rate. Without a prompt it falls
back to a fixed heuristic ("Nice to meet you.", 25 tokens).

Singing needs several times more audio per syllable than speech. If the
budget comes from a speaking reference, it is far too short, the model
cannot sing and falls back to speaking -- and the time budgeted for the
literal tag characters gets filled by reading them out loud.

Design
------
The singing cases use ENGLISH. German singing already fails without a
clone prompt (the tags were trained on English only), so testing German
here would confound two separate causes. English singing is known to work
without a prompt, which leaves the clone prompt as the only variable.

Before generating, the estimated duration is printed for each case. That
number alone may confirm or kill the hypothesis before a single wav is
written.

Run from the project root (omnivoice_lab):
    uv run --project OmniVoice python scripts/debug_singing_duration.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")
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", "debug_duration")

GUIDANCE_SCALE = 3.0

SING_EN = "[singing] Twinkle twinkle little star, how I wonder what you are."
ANGRY_DE = "[angry] Ich habe heute Morgen deine Nachricht gelesen."

# (name, text, language, use_clone, extra kwargs for generate)
CASES = [
    # --- Group A: English Singing WITHOUT Clone (Auto-Voice) ---
    ("a1_noclone_sing_estimated", SING_EN, "en", False, {}),
    ("a2_noclone_sing_duration12", SING_EN, "en", False, {"duration": 12.0}),
    ("a3_noclone_sing_speed05", SING_EN, "en", False, {"speed": 0.5}),

    # --- Group B: English Singing WITH Clone (Reference Speaker) ---
    ("b1_clone_sing_estimated", SING_EN, "en", True, {}),
    ("b2_clone_sing_duration12", SING_EN, "en", True, {"duration": 12.0}),
    ("b3_clone_sing_speed05", SING_EN, "en", True, {"speed": 0.5}),

    # --- Group C: German Emotion WITH Clone ([angry] Tag Readout Check) ---
    ("c1_clone_angry_estimated", ANGRY_DE, "de", True, {}),
    ("c2_clone_angry_duration6", ANGRY_DE, "de", True, {"duration": 6.0}),
]


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})"
        )

    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} ...")
    model = OmniVoice.from_pretrained(
        MODEL_DIR,
        device_map="cuda:0",
        dtype=torch.float16,
    )

    prompt = model.create_voice_clone_prompt(ref_audio=REF_AUDIO, ref_text=ref_text)
    num_ref_tokens = prompt.ref_audio_tokens.size(-1)
    frame_rate = model.audio_tokenizer.config.frame_rate
    sample_rate = getattr(model, "sampling_rate", None) or 24000

    print(f"\nreference : {ref_text!r}")
    print(f"            {len(ref_text)} chars, {num_ref_tokens} audio tokens, "
          f"{num_ref_tokens / frame_rate:.2f} s at {frame_rate} Hz frame rate")

    # --- instrumentation: what length does the model plan? -------------
    print("\n--- estimated target duration before generation ---")
    for name, text, _language, use_clone, extra in CASES:
        if "duration" in extra:
            print(f"{name:<26} fixed by caller: {extra['duration']:.2f} s")
            continue
        speed = extra.get("speed", 1.0)
        if use_clone:
            tokens = model._estimate_target_tokens(text, ref_text, num_ref_tokens, speed=speed)
        else:
            tokens = model._estimate_target_tokens(text, None, None, speed=speed)
        print(f"{name:<26} {tokens / frame_rate:6.2f} s  ({tokens} tokens)")

    # Throwaway call so the first real case does not carry the CUDA warm-up.
    model.generate(text="Warm up.", language="en")
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    print("\n--- generating ---")
    for name, text, language, use_clone, extra in CASES:
        started = time.perf_counter()
        audio = model.generate(
            text=text,
            language=language,
            guidance_scale=GUIDANCE_SCALE,
            voice_clone_prompt=prompt if use_clone else None,
            **extra,
        )
        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)

        clone_flag = "clone" if use_clone else "no clone"
        print(f"{name:<26} {clone_flag:<9} audio {duration:5.2f} s  "
              f"compute {elapsed:5.2f} s")

    print(f"\nWrote {len(CASES)} files to {OUTPUT_DIR}")
    print("\nWhat to listen for:")
    print("  Group A (a1 vs a2/a3) : Does unprompted English singing improve with 12s / speed 0.5?")
    print("  Group A vs B (a vs b) : Does English singing work without clone, but collapse with clone?")
    print("  Group B (b1 vs b2/b3) : Does extra duration (12s) or speed 0.5 rescue English cloned singing?")
    print("  Group C (c1 vs c2)    : Does extra duration stop the literal reading of '[angry]' in German?")


if __name__ == "__main__":
    main()