UmbraAudio/scripts/extract_embeddings.py

#!/usr/bin/env python3
# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Extract Qwen TTS speaker embeddings from existing voice reference WAVs.

Usage:
    python scripts/extract_embeddings.py
    python scripts/extract_embeddings.py --speaker narrator
    python scripts/extract_embeddings.py --force  # re-extract all
    python scripts/extract_embeddings.py --design "A deep male voice" --slug new_speaker

Loads the Base model once, iterates over voice ref WAVs in .umbra/voices/,
extracts the 2048-dim speaker embedding via Qwen TTS's internal encoder,
and saves each as a .npy file in .umbra/voice_embeddings/.
"""
from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

project_root = Path(__file__).resolve().parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

from umbraaudio.config import load_config  # noqa: E402
from umbraaudio.voices.qwen_embeddings import QwenEmbeddingCache  # noqa: E402

BASE_MODEL_PATH = os.getenv(
    "QWEN_BASE_MODEL",
    "/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-Base",
)
VOICEDESIGN_MODEL_PATH = os.getenv(
    "QWEN_VOICEDESIGN_MODEL",
    "/Volumes/1TB_APL_Storage/UmbraAudio/models/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
)


def extract_all_embeddings(
    project_name: str,
    speaker_filter: str | None = None,
    force: bool = False,
) -> None:
    config = load_config(project_name)
    voices_dir = config.data_dir / "voices"
    embeddings_dir = config.data_dir / "voice_embeddings"

    if not voices_dir.is_dir():
        print(f"ERROR: No voices directory at {voices_dir}")
        return

    cache = QwenEmbeddingCache(embeddings_dir)
    wav_files = sorted(voices_dir.glob("*_ref.wav"))
    print(f"Found {len(wav_files)} voice ref WAVs in {voices_dir}")

    if speaker_filter:
        wav_files = [w for w in wav_files if speaker_filter.lower() in w.stem.lower()]
        print(f"Filtered to {len(wav_files)} WAVs matching '{speaker_filter}'")

    to_extract = []
    for wav_path in wav_files:
        slug = wav_path.stem.replace("_ref", "")
        if force or not cache.has(slug):
            to_extract.append((slug, wav_path))

    if not to_extract:
        print("All embeddings already cached. Use --force to re-extract.")
        return

    print(f"Will extract {len(to_extract)} embeddings")

    print("Loading Base model...")
    from qwen_tts import Qwen3TTSModel

    model = Qwen3TTSModel.from_pretrained(
        BASE_MODEL_PATH,
        device_map="mps",
        dtype=torch.bfloat16,
    )
    print("Base model loaded")

    extracted = 0
    failed = 0
    t0 = time.time()

    for slug, wav_path in to_extract:
        try:
            audio, sr = sf.read(str(wav_path), dtype="float32")
            if sr != 24000:
                print(f"  SKIP {slug}: sample rate {sr} != 24000")
                failed += 1
                continue

            emb = cache.extract_and_save(
                slug, audio, sr, model.model
            )
            print(
                f"  OK {slug}: shape={emb.shape}, "
                f"norm={np.linalg.norm(emb):.2f}"
            )
            extracted += 1
        except Exception as e:
            print(f"  FAIL {slug}: {e}")
            failed += 1

    elapsed = time.time() - t0
    print(f"\nDone: {extracted} extracted, {failed} failed, {elapsed:.1f}s")
    print(f"Cached embeddings: {cache.count()} total in {embeddings_dir}")


def create_embedding_from_design(
    project_name: str,
    instruct: str,
    slug: str,
    seed: int = 42,
    force: bool = False,
) -> None:
    """Generate a voice via VoiceDesign, extract its embedding, and cache it."""
    config = load_config(project_name)
    embeddings_dir = config.data_dir / "voice_embeddings"
    cache = QwenEmbeddingCache(embeddings_dir)

    if cache.has(slug) and not force:
        print(f"Embedding for '{slug}' already exists. Use --force to regenerate.")
        return

    print("Loading VoiceDesign model...")
    from qwen_tts import Qwen3TTSModel

    vd_model = Qwen3TTSModel.from_pretrained(
        VOICEDESIGN_MODEL_PATH,
        device_map="mps",
        dtype=torch.bfloat16,
    )
    print("VoiceDesign model loaded")

    test_text = "When the sunlight strikes raindrops in the air, they act as a prism."
    print(f"Generating voice with instruct: {instruct[:80]}...")
    wavs, sr = vd_model.generate_voice_design(
        text=test_text,
        instruct=instruct,
        language="english",
        temperature=0.5,
        do_sample=True,
        seed=seed,
    )
    audio = wavs[0]
    if isinstance(audio, list):
        audio = audio[0]
    audio_np = np.array(audio)
    print(f"Generated {len(audio_np)/sr:.2f}s of audio")

    del vd_model
    torch.mps.empty_cache()

    print("Loading Base model for embedding extraction...")
    base_model = Qwen3TTSModel.from_pretrained(
        BASE_MODEL_PATH,
        device_map="mps",
        dtype=torch.bfloat16,
    )
    print("Base model loaded")

    emb = cache.extract_and_save(slug, audio_np, sr, base_model.model)
    print(f"Saved embedding for '{slug}': shape={emb.shape}, norm={np.linalg.norm(emb):.2f}")

    wav_path = embeddings_dir / f"{slug}_design_ref.wav"
    sf.write(str(wav_path), audio_np, sr)
    print(f"Saved reference audio to {wav_path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract Qwen TTS speaker embeddings from voice refs"
    )
    parser.add_argument(
        "--project",
        default="DealHunterShip_4.0e_whole-to-audio",
        help="Project name (default: DealHunterShip_4.0e_whole-to-audio)",
    )
    parser.add_argument("--speaker", help="Extract only this speaker (substring match)")
    parser.add_argument("--force", action="store_true", help="Re-extract even if cached")
    parser.add_argument("--design", help="VoiceDesign instruct text (create new voice)")
    parser.add_argument("--slug", help="Speaker slug for --design mode")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for --design")
    args = parser.parse_args()

    if args.design:
        if not args.slug:
            parser.error("--slug is required with --design")
        create_embedding_from_design(args.project, args.design, args.slug, args.seed, args.force)
    else:
        extract_all_embeddings(args.project, args.speaker, args.force)


if __name__ == "__main__":
    main()