# Enrolment offline : wav de reference -> artefact prompt_speech GGUF pour cosyvoice.cpp.
# 4 tenseurs : feat (mel prompt [80,T] f32), embedding (campplus [dim,1] f32),
# tokens (s3tok i32), text (transcription brute i8). Pas de crc32 -> pas de check au load.
import sys, numpy as np, torch
sys.path.insert(0, "/opt/Kazeia/cosyvoice-repo")
sys.path.insert(0, "/opt/Kazeia/cosyvoice-repo/third_party/Matcha-TTS")
from cosyvoice.cli.cosyvoice import CosyVoice3
import gguf

REF = sys.argv[1]
REF_TXT = sys.argv[2]
OUT = sys.argv[3]

cv = CosyVoice3("/opt/Kazeia/_models_dl/cosyvoice3-0.5b", load_trt=False, fp16=False)
fe = cv.frontend
text = open(REF_TXT).read().strip() if REF_TXT.endswith(".txt") else REF_TXT

feat, _ = fe._extract_speech_feat(REF)          # [1, T, 80]
emb = fe._extract_spk_embedding(REF)            # [1, dim]
tok, _ = fe._extract_speech_token(REF)          # [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)     # [bytes]
print(f"feat {feat.shape}, emb {emb.shape}, tokens {tok.shape}, text {len(txt)}B = {text[:50]!r}")

w = gguf.GGUFWriter(OUT, "cosyvoice-prompt-speech")
w.add_tensor("feat", feat)         # numpy [T,80] -> ne0=80, ne1=T
w.add_tensor("embedding", emb)     # [1,dim] -> ne0=dim, ne1=1
w.add_tensor("tokens", tok)        # [n] i32
w.add_tensor("text", txt)          # [bytes] i8
w.write_header_to_file()
w.write_kv_data_to_file()
w.write_tensors_to_file()
w.close()
print("artefact ecrit:", OUT)