voice_clone.py

"""
Voice cloning module — wraps Qwen3-TTS for zero-shot voice cloning.

Supports two backends:
  - Base 1.7B: zero-shot voice cloning from reference audio
  - CustomVoice 0.6B: fast predefined speakers (no cloning, lower latency)

Voice profiles are:
  - Cached in-memory for fast access during a session
  - Persisted to Voice_Profile/ folder as .pt files for reuse across restarts
  - Auto-loaded on startup if saved profiles exist

Latency optimizations applied:
  - bfloat16 / float16 precision (auto-detected per GPU arch)
  - FlashAttention-2 when available
  - Streaming mode (non_streaming_mode=False)
  - Reduced sampling params (top_k=20, temperature=0.7)
  - Capped max_new_tokens=1024
  - torch.set_float32_matmul_precision('high')
  - Reference audio trimmed to 3-5s for faster embedding extraction

Usage:
    profile_id = create_voice_profile(ref_audio_path, voice_name="Mom")
    wav, sr = synthesize_cloned(text, profile_id)
"""
from __future__ import annotations

import json
import logging
import os
import uuid
import threading
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

from runtime_config import GPU_INFERENCE_LOCK, VOICE_PROFILE_DIR

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Global PyTorch optimizations
# ---------------------------------------------------------------------------
torch.set_float32_matmul_precision("high")

# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------
# Base model for zero-shot voice cloning (1.7B) — used for BOTH cloned and stock voice
BASE_MODEL_ID = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"

# Stock voice reference audio (pre-generated "vivian" sample)
VIVIAN_REF_PATH = Path(__file__).parent / "assets" / "vivian_reference.wav"

# Profile ID for the built-in stock voice
STOCK_VOICE_PROFILE_ID = "__stock_vivian__"

# Optimized generation parameters (reduced from defaults: top_k=50, temp=0.9, max=2048)
GENERATION_PARAMS = dict(
    top_k=20,
    temperature=0.7,
    subtalker_top_k=20,
    subtalker_temperature=0.7,
    max_new_tokens=1024,
)

# Reference audio limits (seconds) — 3-5s is optimal for Qwen3-TTS
REF_AUDIO_MIN_SEC = 3.0
REF_AUDIO_MAX_SEC = 10.0
REF_AUDIO_TARGET_SR = 24000

# ---------------------------------------------------------------------------
# Voice profile persistence
# ---------------------------------------------------------------------------
VOICE_PROFILE_DIR.mkdir(exist_ok=True)

# Default profile ID — used when no cloned voice exists
DEFAULT_PROFILE_ID = "__default__"

# ---------------------------------------------------------------------------
# Server-side cache: { profile_id -> VoiceClonePromptItem list }
# ---------------------------------------------------------------------------
_PROFILE_CACHE: dict[str, list] = {}
_cache_lock = threading.Lock()

_qwen_tts_model = None
_model_lock = threading.Lock()


def _select_dtype() -> torch.dtype:
    """Pick optimal dtype based on GPU architecture."""
    if not torch.cuda.is_available():
        return torch.float32
    cap = torch.cuda.get_device_capability()
    # bfloat16 requires compute capability >= 8.0 (Ampere+)
    if cap[0] >= 8:
        return torch.bfloat16
    return torch.float16


def _select_attn_impl() -> str:
    """Use FlashAttention-2 if available, else SDPA (PyTorch native)."""
    try:
        import flash_attn  # noqa: F401
        return "flash_attention_2"
    except ImportError:
        logger.info("flash-attn not installed — using SDPA attention.")
        return "sdpa"


def get_qwen_tts():
    """Lazy-load Qwen3-TTS Base 1.7B model for cloning. Thread-safe."""
    global _qwen_tts_model
    if _qwen_tts_model is None:
        with _model_lock:
            if _qwen_tts_model is None:
                from qwen_tts import Qwen3TTSModel

                model_id = BASE_MODEL_ID
                device = "cuda" if torch.cuda.is_available() else "cpu"

                logger.info("Loading %s on %s (cloning model)...", model_id, device)
                attn_impl = _select_attn_impl()
                _qwen_tts_model = Qwen3TTSModel.from_pretrained(
                    model_id,
                    device_map=device,
                    attn_implementation=attn_impl,
                )
                logger.info("Qwen3-TTS Base loaded on %s (attn=%s).", device, attn_impl)
    return _qwen_tts_model


def _try_torch_compile(wrapper):
    """Best-effort torch.compile on model submodules. Disabled for stability."""
    # torch.compile can cause CUDA asserts on some GPU architectures (T4/Turing)
    # Disable for now in favor of stability
    return


def _trim_reference_audio(audio_path: str) -> str:
    """
    Trim reference audio to REF_AUDIO_MAX_SEC seconds if longer.
    Returns path to trimmed file (or original if already short enough).
    """
    try:
        info = sf.info(audio_path)
        duration = info.duration
        if duration <= REF_AUDIO_MAX_SEC:
            return audio_path

        logger.info(
            "Reference audio %.1fs exceeds %.0fs limit — trimming.",
            duration, REF_AUDIO_MAX_SEC,
        )
        data, sr = sf.read(audio_path)
        max_samples = int(REF_AUDIO_MAX_SEC * sr)
        trimmed = data[:max_samples]

        trimmed_path = audio_path + ".trimmed.wav"
        sf.write(trimmed_path, trimmed, sr)
        return trimmed_path
    except Exception as exc:
        logger.warning("Could not trim reference audio: %s", exc)
        return audio_path


# ---------------------------------------------------------------------------
# Profile persistence (disk ↔ memory)
# ---------------------------------------------------------------------------

def save_profile_to_disk(profile_id: str, voice_name: str = "Cloned Voice") -> Path:
    """
    Save a cached voice profile to Voice_Profile/ as a .pt file + metadata JSON.
    Returns the path to the saved .pt file.
    """
    with _cache_lock:
        prompt_items = _PROFILE_CACHE.get(profile_id)
    if prompt_items is None:
        raise ValueError(f"Profile '{profile_id}' not found in cache.")

    profile_dir = VOICE_PROFILE_DIR / profile_id
    profile_dir.mkdir(exist_ok=True)

    # Serialize VoiceClonePromptItem fields
    serializable = []
    for item in prompt_items:
        serializable.append({
            "ref_code": item.ref_code.cpu() if item.ref_code is not None else None,
            "ref_spk_embedding": item.ref_spk_embedding.cpu(),
            "x_vector_only_mode": item.x_vector_only_mode,
            "icl_mode": item.icl_mode,
            "ref_text": item.ref_text,
        })

    pt_path = profile_dir / "profile.pt"
    torch.save(serializable, pt_path)

    # Save metadata
    meta = {"profile_id": profile_id, "voice_name": voice_name}
    meta_path = profile_dir / "metadata.json"
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    logger.info("Voice profile '%s' (%s) saved to %s", profile_id, voice_name, profile_dir)
    return pt_path


def load_profile_from_disk(profile_id: str) -> bool:
    """
    Load a voice profile from Voice_Profile/<profile_id>/profile.pt into memory cache.
    Returns True if loaded successfully, False otherwise.
    """
    profile_dir = VOICE_PROFILE_DIR / profile_id
    pt_path = profile_dir / "profile.pt"

    if not pt_path.exists():
        logger.warning("Profile file not found: %s", pt_path)
        return False

    try:
        from qwen_tts.inference.qwen3_tts_model import VoiceClonePromptItem

        raw_items = torch.load(pt_path, map_location="cpu", weights_only=False)
        prompt_items = []
        for item_dict in raw_items:
            prompt_items.append(VoiceClonePromptItem(
                ref_code=item_dict["ref_code"],
                ref_spk_embedding=item_dict["ref_spk_embedding"],
                x_vector_only_mode=item_dict["x_vector_only_mode"],
                icl_mode=item_dict["icl_mode"],
                ref_text=item_dict.get("ref_text"),
            ))

        with _cache_lock:
            _PROFILE_CACHE[profile_id] = prompt_items

        logger.info("Voice profile '%s' loaded from disk.", profile_id)
        return True
    except Exception as exc:
        logger.exception("Failed to load profile '%s': %s", profile_id, exc)
        return False


def list_saved_profiles() -> list[dict]:
    """
    List all saved voice profiles from Voice_Profile/ directory.
    Returns list of {profile_id, voice_name, path} dicts, newest first.
    """
    profiles = []
    if not VOICE_PROFILE_DIR.exists():
        return profiles

    for entry in VOICE_PROFILE_DIR.iterdir():
        if not entry.is_dir():
            continue
        pt_path = entry / "profile.pt"
        meta_path = entry / "metadata.json"
        if not pt_path.exists():
            continue

        voice_name = "Cloned Voice"
        if meta_path.exists():
            try:
                with open(meta_path, encoding="utf-8") as f:
                    meta = json.load(f)
                voice_name = meta.get("voice_name", voice_name)
            except Exception:
                pass

        profiles.append({
            "profile_id": entry.name,
            "voice_name": voice_name,
            "path": str(entry),
            "mtime": pt_path.stat().st_mtime,
        })

    # Newest first
    profiles.sort(key=lambda p: p["mtime"], reverse=True)
    return profiles


def load_default_profile() -> str | None:
    """
    Load the most recently saved voice profile from Voice_Profile/ into memory.
    Returns the profile_id if found, None if no saved profiles exist.
    This is called on app startup to restore the last cloned voice.
    """
    saved = list_saved_profiles()
    if not saved:
        logger.info("No saved voice profiles found — using stock voice as default.")
        return None

    newest = saved[0]
    profile_id = newest["profile_id"]
    if load_profile_from_disk(profile_id):
        logger.info(
            "Default voice profile loaded: '%s' (%s)",
            profile_id, newest["voice_name"],
        )
        return profile_id

    return None


# ---------------------------------------------------------------------------
# Core API
# ---------------------------------------------------------------------------

def create_voice_profile(ref_audio_path: str, voice_name: str = "Cloned Voice", profile_id_override: str | None = None) -> str:
    """
    Extract speaker embedding from reference audio, cache it, and save to disk.
    Returns a profile_id string for later synthesis.

    Reference audio is trimmed to 3-10s for optimal latency.
    """
    trimmed_path = _trim_reference_audio(ref_audio_path)
    model = get_qwen_tts()

    logger.info("Creating voice profile from %s...", ref_audio_path)
    with GPU_INFERENCE_LOCK:
        prompt_items = model.create_voice_clone_prompt(
            ref_audio=trimmed_path,
            x_vector_only_mode=True,
        )

    profile_id = profile_id_override or uuid.uuid4().hex[:12]
    with _cache_lock:
        _PROFILE_CACHE[profile_id] = prompt_items

    # Persist to disk
    try:
        save_profile_to_disk(profile_id, voice_name=voice_name)
    except Exception as exc:
        logger.warning("Failed to save profile to disk: %s", exc)

    logger.info("Voice profile %s created and saved.", profile_id)
    return profile_id


def synthesize_cloned(text: str, profile_id: str) -> tuple[np.ndarray, int]:
    """
    Synthesize text using a cached voice profile.
    Returns (wav_array, sample_rate).
    """
    with _cache_lock:
        prompt_items = _PROFILE_CACHE.get(profile_id)
    if prompt_items is None:
        # Try loading from disk if not in memory
        if load_profile_from_disk(profile_id):
            with _cache_lock:
                prompt_items = _PROFILE_CACHE.get(profile_id)
    if prompt_items is None:
        raise ValueError(f"Voice profile '{profile_id}' not found. Record voice first.")

    model = get_qwen_tts()

    with GPU_INFERENCE_LOCK:
        audio_list, sample_rate = model.generate_voice_clone(
            text=text,
            language="english",
            voice_clone_prompt=prompt_items,
            non_streaming_mode=False,
            **GENERATION_PARAMS,
        )

    wav = np.concatenate(audio_list) if audio_list else np.zeros(0, dtype=np.float32)
    return wav, sample_rate


def _ensure_stock_voice_profile():
    """Ensure the stock vivian voice profile is loaded (uses Base 1.7B with reference audio)."""
    with _cache_lock:
        if STOCK_VOICE_PROFILE_ID in _PROFILE_CACHE:
            return
    # Create profile from vivian reference audio
    if not VIVIAN_REF_PATH.exists():
        logger.error("Stock voice reference not found at %s", VIVIAN_REF_PATH)
        raise FileNotFoundError(f"Stock voice reference not found: {VIVIAN_REF_PATH}")
    create_voice_profile(str(VIVIAN_REF_PATH), voice_name="Vivian (Stock)", profile_id_override=STOCK_VOICE_PROFILE_ID)
    logger.info("Stock vivian voice profile created from reference audio.")


def synthesize_custom_voice_streaming(
    text: str, speaker: str = "vivian", language: str = "english"
):
    """
    Synthesize text with the Base 1.7B model using stock vivian reference.
    Yields (wav_segment, sample_rate) tuples.
    """
    _ensure_stock_voice_profile()
    with _cache_lock:
        prompt_items = _PROFILE_CACHE.get(STOCK_VOICE_PROFILE_ID)
    if prompt_items is None:
        raise RuntimeError(
            "Stock voice profile could not be created. "
            "Check that assets/vivian_reference.wav exists and the TTS model loaded correctly."
        )

    model = get_qwen_tts()
    with GPU_INFERENCE_LOCK:
        audio_list, sample_rate = model.generate_voice_clone(
            text=text,
            language=language,
            voice_clone_prompt=prompt_items,
            non_streaming_mode=False,
            **GENERATION_PARAMS,
        )

    for segment in audio_list:
        if segment is not None and len(segment) > 0:
            yield segment, sample_rate


def synthesize_custom_voice(
    text: str, speaker: str = "vivian", language: str = "english"
) -> tuple[np.ndarray, int]:
    """
    Synthesize text with the Base 1.7B model using stock vivian reference.
    Returns (wav_array, sample_rate).
    """
    segments = list(synthesize_custom_voice_streaming(text, speaker, language))
    if segments:
        wav = np.concatenate([s for s, _ in segments])
        sr = segments[0][1]
        return wav, sr
    # Fallback: empty audio
    return np.zeros(0, dtype=np.float32), 24000


def synthesize_cloned_preview(profile_id: str) -> tuple[np.ndarray, int]:
    """Short preview sentence to verify the clone sounds right."""
    return synthesize_cloned(
        "Hello! I'm ready to read a bedtime story for you tonight.",
        profile_id,
    )


def has_profile(profile_id: str | None) -> bool:
    """Check if a voice profile exists in cache or on disk."""
    if not profile_id:
        return False
    with _cache_lock:
        if profile_id in _PROFILE_CACHE:
            return True
    # Check disk
    pt_path = VOICE_PROFILE_DIR / profile_id / "profile.pt"
    return pt_path.exists()