#!/usr/bin/env python3
"""Convertit speaker_encoder_weights.pt (Qwen3-TTS ECAPA-TDNN) en GGUF f16.

Input  : /opt/Kazeia/to_delete/models_qnn/qwen3-tts-native/speaker_encoder_weights.pt
Output : <out>.gguf (~17 MB f16) lisible par ggml C++ via gguf_init_from_file().

Architecture cible :
  blocks.0          : Conv1d(128 -> 512, k=5)
  blocks.1,2,3      : SE-Res2Net (tdnn1 1x1 + 7x Conv k=3 + tdnn2 1x1 + se_block)
  mfa               : Conv1d(1536 -> 1536, k=1)   ← concat 3 blocks outputs
  asp.tdnn          : Conv1d(4608 -> 128, k=1)    ← Attentive Statistics Pooling
  asp.conv          : Conv1d(128 -> 1536, k=1)
  fc                : Linear(3072 -> 1024)        ← x_vector final

Mel input attendu : [B=1, n_mels=128, T] avec SR=24kHz, N_FFT=1024, HOP=256.
"""
import sys, argparse, os
import numpy as np
import torch

# On utilise le gguf package du venv qnn (déjà installé pour les outils llama.cpp).
sys.path.insert(0, '/opt/Kazeia/qnn_venv/lib/python3.10/site-packages')
import gguf  # noqa: E402

PT_PATH = '/opt/Kazeia/to_delete/models_qnn/qwen3-tts-native/speaker_encoder_weights.pt'
ARCH    = 'qwen3-tts-speaker-encoder'

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', required=True, help='Output GGUF path')
    ap.add_argument('--dtype', choices=['f32', 'f16'], default='f16')
    args = ap.parse_args()

    sd = torch.load(PT_PATH, map_location='cpu', weights_only=False)
    if isinstance(sd, dict) and 'state_dict' in sd:
        sd = sd['state_dict']
    print(f"Loaded {len(sd)} tensors from {PT_PATH}")

    writer = gguf.GGUFWriter(args.out, ARCH)

    # Metadata utilitaire (lue par le C++ via gguf_find_key())
    writer.add_string(f'{ARCH}.architecture', 'ecapa_tdnn')
    writer.add_uint32(f'{ARCH}.n_mels',       128)
    writer.add_uint32(f'{ARCH}.sample_rate',  24000)
    writer.add_uint32(f'{ARCH}.n_fft',        1024)
    writer.add_uint32(f'{ARCH}.hop_size',     256)
    writer.add_uint32(f'{ARCH}.win_size',     1024)
    writer.add_float32(f'{ARCH}.fmin',        0.0)
    writer.add_float32(f'{ARCH}.fmax',        12000.0)
    writer.add_uint32(f'{ARCH}.embedding_dim', 1024)
    writer.add_uint32(f'{ARCH}.n_blocks',     3)  # SE-Res2Net blocks (blocks.1..3)
    writer.add_uint32(f'{ARCH}.channels',     512)

    total_bytes = 0
    for name, tensor in sd.items():
        # PyTorch tensors -> numpy. Conv weights restent en F32 -> on caste en f16 si demandé.
        arr = tensor.detach().cpu().numpy()
        # PyTorch Conv1d weights shape : [C_out, C_in, K] -> on garde
        # (ggml ggml_conv_1d s'attend à w[K, C_in, C_out] mais il est aussi possible
        # de tout passer en mul_mat custom — on garde le layout PyTorch et on adapte
        # côté C++ via reshape).
        if args.dtype == 'f16' and arr.dtype == np.float32:
            # Si c'est un bias ou un tensor < 1024 elements, on garde en f32 (peu de gain
            # f16 dessus, et les biais sont sensibles au précision).
            if arr.size < 1024:
                pass  # garde f32
            else:
                arr = arr.astype(np.float16)
        writer.add_tensor(name, arr)
        total_bytes += arr.nbytes

    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_tensors_to_file()
    writer.close()
    print(f"\nWrote {args.out}")
    print(f"  total weights : {total_bytes/1024/1024:.1f} MB ({args.dtype})")
    print(f"  file size     : {os.path.getsize(args.out)/1024/1024:.1f} MB")

if __name__ == '__main__':
    main()