Kazeia-central/kazeia_central/voice/enroll.py

"""Enrôlement vocal CosyVoice : WAV de référence → artefact `.cvps`.

⚠️ DOIT s'exécuter sous l'environnement encodeur (`cv_venv` : torch + cosyvoice +
gguf + onnxruntime), PAS sous l'interpréteur de l'API Kazeia-central (Python 3.14,
sans torch). C'est pourquoi ce module est invoqué par un **worker subprocess**
(`voice/worker.py`) tournant sous `cv_venv`, pas importé en-process par l'API.

Réimplémente la logique des 25 lignes de `/opt/Kazeia/cv_make_prompt.py` (cf.
VOICE_ENROLLMENT_SPEC §2), avec deux différences voulues :
- chemins repo/teacher **configurables** (env `CV_REPO`, `CV_TEACHER`), pas codés en dur ;
- teacher chargé **une seule fois** (singleton) pour enchaîner les voix à chaud.

Contrat de sortie (VOICE_ENROLLMENT_SPEC §3) — conteneur GGUF, arch
`cosyvoice-prompt-speech`, 4 tenseurs aux noms/dtypes EXACTS (le loader C++
`nativeLoadVoice` ne tolère aucune dérive) :
  feat      mel prompt      numpy [T, 80]  float32   (GGUF: ne0=80, ne1=T)
  embedding x-vector        numpy [1, dim] float32   (GGUF: ne0=dim, ne1=1)
  tokens    speech tokens   numpy [n]      int32
  text      transcription   numpy [octets] int8      (UTF-8 brut)
"""

from __future__ import annotations

import logging
import os
import sys

import numpy as np

# La stack CosyVoice/librosa/numba crache du DEBUG sur stdout/stderr. Le worker (à
# venir) utilisera stdout pour un protocole JSON → on muselle ces loggers ici pour
# garder stdout propre. Indispensable avant de brancher l'IPC.
for _noisy in ("numba", "matplotlib", "jieba", "modelscope", "cosyvoice"):
    logging.getLogger(_noisy).setLevel(logging.WARNING)
os.environ.setdefault("NUMBA_DISABLE_JIT", "0")

CV_REPO = os.environ.get("CV_REPO", "/opt/Kazeia/cosyvoice-repo")
CV_TEACHER = os.environ.get("CV_TEACHER", "/opt/Kazeia/_models_dl/cosyvoice3-0.5b")

GGUF_ARCH = "cosyvoice-prompt-speech"

# Préparation du clip. Le tokenizer s3 plante DUR au-delà de 30 s ("do not support
# extract speech token for audio longer than 30s") et la similarité chute sous ~8 s
# (§4). Les captures device ne sont pas calibrées (ex. 111 s stéréo 44 kHz observé) →
# on normalise systématiquement : mono, 16 kHz, silences rognés, fenêtre cible ~15 s.
# 22 s (≤ _HARD_MAX_S) : la phrase de consentement (~13 s) + la fin phonétiquement
# riche (nasales, semi-voyelles, ch/j/r, + question/exclamation) entrent ENTIÈREMENT
# dans le prompt enrôlé ; sinon la tête légale remplit seule les 15 s et la fin
# (l'info utile au clonage) est tronquée. Prompt plus long ⇒ meilleure similarité.
CV_TARGET_S = float(os.environ.get("CV_TARGET_S", "22.0"))
_PREP_SR = 16000
_HARD_MAX_S = 28.0   # marge sous la limite dure 30 s du tokenizer
_MIN_WARN_S = 6.0

_teacher_singleton = None


class EnrollmentError(RuntimeError):
    pass


def prepare_wav(in_path: str, out_path: str, target_s: float = CV_TARGET_S) -> dict:
    """Normalise un WAV de capture en segment exploitable : mono, 16 kHz, silences de
    bord rognés, tronqué à `target_s` (borné dur sous 30 s). Renvoie un résumé."""
    import librosa
    import soundfile as sf

    y, _ = librosa.load(in_path, sr=_PREP_SR, mono=True)   # downmix + resample
    orig_s = len(y) / _PREP_SR
    yt, _ = librosa.effects.trim(y, top_db=30)             # rogne silences de bord
    if yt.size == 0:
        yt = y
    keep_s = min(target_s, _HARD_MAX_S)
    yt = yt[: int(keep_s * _PREP_SR)]
    final_s = len(yt) / _PREP_SR
    sf.write(out_path, yt, _PREP_SR, subtype="PCM_16")
    return {"orig_s": round(orig_s, 1), "final_s": round(final_s, 1),
            "too_short": final_s < _MIN_WARN_S}


def load_teacher():
    """Charge le teacher CosyVoice3-0.5B (~9,1 Go) UNE fois. Coûteux (dizaines de
    secondes) → garder le process chaud (worker). Lève EnrollmentError si l'env
    encodeur n'est pas provisionné (torch/cosyvoice/teacher absents)."""
    global _teacher_singleton
    if _teacher_singleton is not None:
        return _teacher_singleton
    if not os.path.isdir(CV_TEACHER):
        raise EnrollmentError(
            f"teacher CosyVoice introuvable: {CV_TEACHER} "
            f"(provisionner le modèle 0.5B ~9,1 Go, ou définir CV_TEACHER)")
    if not os.path.isdir(CV_REPO):
        raise EnrollmentError(f"repo CosyVoice introuvable: {CV_REPO} (définir CV_REPO)")
    for p in (CV_REPO, os.path.join(CV_REPO, "third_party/Matcha-TTS")):
        if p not in sys.path:
            sys.path.insert(0, p)
    try:
        from cosyvoice.cli.cosyvoice import CosyVoice3
    except Exception as e:  # torch/cosyvoice absents = mauvais interpréteur
        raise EnrollmentError(
            f"import CosyVoice échoué ({e}). Exécuter sous l'env encodeur "
            f"(cv_venv : torch + cosyvoice + gguf).") from e
    _teacher_singleton = CosyVoice3(CV_TEACHER, load_trt=False, fp16=False)
    return _teacher_singleton


def enroll(wav_path: str, text: str, out_path: str) -> dict:
    """WAV + transcription → `.cvps` à `out_path`. Renvoie un résumé (shapes/taille)
    pour validation. `text` doit correspondre au contenu parlé du WAV (§4)."""
    import gguf

    if not os.path.isfile(wav_path):
        raise EnrollmentError(f"WAV introuvable: {wav_path}")
    text = (text or "").strip()
    if not text:
        raise EnrollmentError("transcription vide : CosyVoice exige le texte de référence")

    import tempfile

    cv = load_teacher()
    fe = cv.frontend

    # Normalise la capture (mono/16k/≤~15s) dans un fichier temporaire avant extraction.
    with tempfile.TemporaryDirectory() as td:
        prep = os.path.join(td, "prep.wav")
        prep_info = prepare_wav(wav_path, prep)
        feat, _ = fe._extract_speech_feat(prep)      # [1, T, 80]
        emb = fe._extract_spk_embedding(prep)        # [1, dim]
        tok, _ = fe._extract_speech_token(prep)      # [1, n]

    feat = feat.squeeze(0).cpu().numpy().astype(np.float32)       # [T, 80]
    emb = emb.cpu().numpy().astype(np.float32)                    # [1, dim]
    tok = tok.squeeze(0).cpu().numpy().astype(np.int32)           # [n]
    txt = np.frombuffer(text.encode("utf-8"), dtype=np.int8)      # [octets]

    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
    w = gguf.GGUFWriter(out_path, GGUF_ARCH)
    w.add_tensor("feat", feat)
    w.add_tensor("embedding", emb)
    w.add_tensor("tokens", tok)
    w.add_tensor("text", txt)
    w.write_header_to_file()
    w.write_kv_data_to_file()
    w.write_tensors_to_file()
    w.close()

    return {
        "out_path": out_path,
        "bytes": os.path.getsize(out_path),
        "feat_shape": list(feat.shape),
        "embedding_dim": int(emb.shape[1]),
        "n_tokens": int(tok.shape[0]),
        "text_bytes": int(txt.shape[0]),
        "prep": prep_info,
    }


if __name__ == "__main__":
    # Usage : cv_venv/bin/python -m kazeia_central.voice.enroll <wav> <text|fichier.txt> <out.cvps>
    import json

    ref, ref_txt, out = sys.argv[1], sys.argv[2], sys.argv[3]
    text = open(ref_txt).read() if ref_txt.endswith(".txt") and os.path.isfile(ref_txt) else ref_txt
    print(json.dumps(enroll(ref, text, out), ensure_ascii=False))