119 lines
5.1 KiB
Python
119 lines
5.1 KiB
Python
"""Enrôlement vocal OmniVoice : WAV de référence → clone-prompt `.ovsp`.
|
|
|
|
⚠️ Sous l'env encodeur `ov_venv` (torch + omnivoice + transformers + librosa/soundfile),
|
|
invoqué par le worker subprocess (pas en-process API py3.14).
|
|
Réimplémente /opt/Kazeia-engine/dist/omnivoice/scripts/make_ovsp.py.
|
|
Format .ovsp : int32 'OVSP' | int32 T_ref | int32 text_len | int32 tokens[8*T_ref] | char ref_text[].
|
|
"""
|
|
from __future__ import annotations
|
|
import logging, os, struct, sys
|
|
import numpy as np
|
|
|
|
for _noisy in ("numba", "matplotlib", "transformers", "omnivoice", "huggingface_hub"):
|
|
logging.getLogger(_noisy).setLevel(logging.WARNING)
|
|
|
|
OV_MODEL = os.environ.get("OV_MODEL", "k2-fsa/OmniVoice")
|
|
OVSP_MAGIC = 0x4F565350
|
|
# Cap ~16 s : T_ref ~390 -> bucket 640 (rapide). Au-delà = bucket 768/1024 (lent) +
|
|
# risque de saturer le plus grand bucket. prepare_wav partagé avec la transcription.
|
|
OV_TARGET_S = float(os.environ.get("OV_TARGET_S", "16.0"))
|
|
_PREP_SR = 24000
|
|
_HARD_MAX_S = 20.0
|
|
_MIN_WARN_S = 6.0
|
|
_model_singleton = None
|
|
|
|
class EnrollmentError(RuntimeError):
|
|
pass
|
|
|
|
def prepare_wav(in_path: str, out_path: str, target_s: float = OV_TARGET_S) -> dict:
|
|
import librosa, soundfile as sf
|
|
cap = min(target_s, _HARD_MAX_S)
|
|
y, _ = librosa.load(in_path, sr=_PREP_SR, mono=True)
|
|
orig_s = len(y) / _PREP_SR
|
|
yt, _ = librosa.effects.trim(y, top_db=30)
|
|
if yt.size == 0:
|
|
yt = y
|
|
# §7 FIX : couper sur une FRONTIÈRE DE SILENCE, jamais à `cap` brut (= en plein mot,
|
|
# ce qui faisait recracher le dernier mot tronqué en tête de chaque génération).
|
|
# Garder les segments non-silencieux jusqu'à ~cap, finir à la fin du dernier segment.
|
|
end = 0
|
|
for _s, e in librosa.effects.split(yt, top_db=30):
|
|
if e / _PREP_SR <= cap:
|
|
end = e
|
|
else:
|
|
break
|
|
if end == 0: # un seul long segment → repli au cap dur
|
|
end = min(len(yt), int(cap * _PREP_SR))
|
|
yt = yt[:end]
|
|
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_model():
|
|
global _model_singleton
|
|
if _model_singleton is not None:
|
|
return _model_singleton
|
|
try:
|
|
import torch
|
|
from omnivoice.models.omnivoice import OmniVoice
|
|
except Exception as e:
|
|
raise EnrollmentError(f"import OmniVoice échoué ({e}). Exécuter sous ov_venv.") from e
|
|
_model_singleton = OmniVoice.from_pretrained(OV_MODEL, torch_dtype=torch.float32).to("cpu").eval()
|
|
return _model_singleton
|
|
|
|
def _clone_prompt(prep_wav: str, ref_text):
|
|
import torch
|
|
m = load_model()
|
|
with torch.no_grad():
|
|
return m.create_voice_clone_prompt(ref_audio=prep_wav, ref_text=ref_text or None)
|
|
|
|
def transcribe_ref(wav_path: str) -> dict:
|
|
"""ASR du segment préparé (16 s) via OmniVoice, pour relecture opérateur."""
|
|
import tempfile
|
|
if not os.path.isfile(wav_path):
|
|
raise EnrollmentError(f"WAV introuvable: {wav_path}")
|
|
with tempfile.TemporaryDirectory() as td:
|
|
prep = os.path.join(td, "prep.wav")
|
|
prep_info = prepare_wav(wav_path, prep)
|
|
vcp = _clone_prompt(prep, None)
|
|
return {"text": (getattr(vcp, "ref_text", "") or "").strip(),
|
|
"language": None, "model": "omnivoice-asr", "prep": prep_info}
|
|
|
|
def enroll(wav_path: str, text: str, out_path: str) -> dict:
|
|
import tempfile
|
|
if not os.path.isfile(wav_path):
|
|
raise EnrollmentError(f"WAV introuvable: {wav_path}")
|
|
text = (text or "").strip()
|
|
with tempfile.TemporaryDirectory() as td:
|
|
prep = os.path.join(td, "prep.wav")
|
|
prep_info = prepare_wav(wav_path, prep)
|
|
vcp = _clone_prompt(prep, text if text else None)
|
|
toks = vcp.ref_audio_tokens.cpu().numpy().astype(np.int32) # [8, T_ref]
|
|
C, T = toks.shape
|
|
if C != 8:
|
|
raise EnrollmentError(f"attendu 8 codebooks, eu {C}")
|
|
ref_text = (getattr(vcp, "ref_text", "") or text or "").strip()
|
|
# §7 FIX (ceinture+bretelles) : tronquer à la dernière phrase complète (. ! ?), pour
|
|
# qu'un éventuel mot tronqué de fin de réf ne soit jamais recraché en tête de synthèse.
|
|
import re
|
|
_ends = list(re.finditer(r"[.!?]", ref_text))
|
|
if _ends:
|
|
ref_text = ref_text[: _ends[-1].end()].strip()
|
|
txt = ref_text.encode("utf-8")
|
|
os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
|
|
with open(out_path, "wb") as f:
|
|
f.write(struct.pack("<i", OVSP_MAGIC))
|
|
f.write(struct.pack("<i", T))
|
|
f.write(struct.pack("<i", len(txt)))
|
|
f.write(toks.tobytes())
|
|
f.write(txt)
|
|
return {"out_path": out_path, "bytes": os.path.getsize(out_path), "t_ref": int(T),
|
|
"ref_seconds": round(T / 25.0, 1), "text_bytes": len(txt),
|
|
"ref_text": ref_text, "prep": prep_info}
|
|
|
|
if __name__ == "__main__":
|
|
import json
|
|
ref, ref_txt, out = sys.argv[1], sys.argv[2], sys.argv[3]
|
|
text = "" if ref_txt == "-" else (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))
|