voice_clone.py

"""
Voice Cloning: synthesize new text in the voice of a reference recording.

This is the mode OmniVoice is primarily trained on, so it is the most
stable one. The transcript is read from the .txt next to the audio and
passed in explicitly -- leaving it out would pull in Whisper to transcribe
the reference on every run.

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

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_PATH = os.path.join("data", "output", f"clone_{REFERENCE}.wav")

LANGUAGE = "de"
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.""")


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(os.path.dirname(OUTPUT_PATH), 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"Reference: {REF_AUDIO}")
    print(f"  ref_text: \"{ref_text}\"")

    started = time.perf_counter()
    audio = model.generate(
        text=TEXT,
        language=LANGUAGE,
        ref_audio=REF_AUDIO,
        ref_text=ref_text,
    )
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - started

    sample_rate = getattr(model, "sampling_rate", None) or 24000
    wav = audio[0]
    duration = len(wav) / sample_rate

    sf.write(OUTPUT_PATH, wav, sample_rate)

    print(f"Wrote {OUTPUT_PATH}")
    print(f"  duration : {duration:.2f} s @ {sample_rate} Hz")
    print(f"  generated in {elapsed:.2f} s  (RTF {elapsed / duration:.3f})")


if __name__ == "__main__":
    main()