src/tts/f5_tts.py

"""
F5-TTS voice cloning engine.

Generates speech in a target speaker's voice given a short reference WAV.
Falls back to None gracefully if f5-tts is not installed or the GPU is
unavailable — the caller then falls back to MMS-TTS.

Install:
    pip install f5-tts>=1.0.0

Reference:
    SWivid/F5-TTS  (HuggingFace / GitHub)
    Model: ~750 MB, downloaded on first use to HF cache.
"""
from __future__ import annotations

import logging
import threading
from pathlib import Path
from typing import Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)

_lock   = threading.Lock()
_model  = None   # F5TTS instance, loaded lazily


def _load_model():
    global _model
    if _model is not None:
        return _model
    with _lock:
        if _model is None:
            from f5_tts.api import F5TTS  # type: ignore
            _model = F5TTS(model_type="F5TTS")
            logger.info("F5-TTS model loaded.")
    return _model


def synthesize(
    text: str,
    ref_wav_path: str,
    ref_text: str = "",
    speed: float = 1.0,
    device: str = "cuda",
) -> Optional[Tuple[np.ndarray, int]]:
    """
    Generate speech for `text` using `ref_wav_path` as the speaker reference.

    Args:
        text:         Text to synthesize (Bambara, Fula, French, or English).
        ref_wav_path: Path to reference audio (WAV, 5–30 s of the target speaker).
        ref_text:     Transcript of the reference audio.  If empty the model
                      uses in-context inference (slightly lower quality but still
                      good for voice matching).
        speed:        Speaking rate multiplier.  1.0 = normal.
        device:       "cuda" or "cpu".  CPU is 30-60 s/sentence — use GPU.

    Returns:
        (waveform_float32, sample_rate) or None on failure.
    """
    if not text.strip():
        return None

    try:
        import torch
        model = _load_model()

        wav, sr, _ = model.infer(
            ref_file=ref_wav_path,
            ref_text=ref_text.strip(),
            gen_text=text.strip(),
            speed=speed,
            target_rms=0.1,
            cross_fade_duration=0.15,
            nfe_step=32,
            cfg_strength=2.0,
            show_info=False,
            progress=None,
        )

        if isinstance(wav, torch.Tensor):
            wav = wav.cpu().float().numpy()
        else:
            wav = np.asarray(wav, dtype=np.float32)

        return wav, int(sr)

    except ImportError:
        logger.warning(
            "f5-tts not installed — voice cloning disabled.  "
            "Add 'f5-tts>=1.0.0' to requirements.txt."
        )
        return None
    except Exception as exc:
        logger.error("F5-TTS synthesis failed: %s", exc)
        return None


def to_wav_24k(audio_path: str) -> str:
    """
    Resample any audio file to 24 kHz mono WAV (F5-TTS preferred sample rate).
    Returns the path to the converted file (same stem, .wav extension).
    Modifies in-place if the input is already a WAV — otherwise writes a new file.
    """
    import librosa
    import soundfile as sf

    out_path = str(Path(audio_path).with_suffix(".f5ref.wav"))
    audio, _ = librosa.load(audio_path, sr=24_000, mono=True)
    sf.write(out_path, audio, 24_000)
    return out_path