124 lines
5.1 KiB
Python
124 lines
5.1 KiB
Python
#!/usr/bin/env python3
|
||
"""Extract embedding tables from Qwen3-TTS for the Android app."""
|
||
import torch
|
||
import numpy as np
|
||
import os
|
||
|
||
OUT = "/opt/Kazeia/tts_npy_exports"
|
||
os.makedirs(OUT, exist_ok=True)
|
||
NATIVE = "/opt/Kazeia/models_qnn/qwen3-tts-native"
|
||
TALKER_HF = "/opt/Kazeia/models_qnn/talker_hf"
|
||
|
||
# 1. Text embeddings + projection
|
||
print("=== text_components.pt ===")
|
||
tc = torch.load(f"{NATIVE}/text_components.pt", map_location="cpu", weights_only=False)
|
||
text_emb_dict = tc['text_embedding'] # OrderedDict
|
||
text_proj_dict = tc['text_projection'] # OrderedDict
|
||
|
||
print("text_embedding keys:", list(text_emb_dict.keys()))
|
||
print("text_projection keys:", list(text_proj_dict.keys()))
|
||
|
||
# text_embedding.weight: [151936, 2048] — full LLM vocab, but we only need first 1050 tokens
|
||
text_emb_w = text_emb_dict['weight'] # [151936, 2048]
|
||
print(f"text_embedding.weight: {text_emb_w.shape}")
|
||
|
||
# Only keep the first 1050 text tokens (CODEC_OFFSET in the engine)
|
||
text_emb_w = text_emb_w[:1050] # [1050, 2048]
|
||
print(f"text_embedding (trimmed): {text_emb_w.shape}")
|
||
|
||
# text_projection is a 2-layer MLP: fc1(2048→1024) + GELU + fc2(1024→1024)
|
||
fc1_w = text_proj_dict['linear_fc1.weight'] # [hidden, 2048]
|
||
fc1_b = text_proj_dict['linear_fc1.bias']
|
||
fc2_w = text_proj_dict['linear_fc2.weight'] # [1024, hidden]
|
||
fc2_b = text_proj_dict['linear_fc2.bias']
|
||
print(f"text_projection: fc1={fc1_w.shape}, fc2={fc2_w.shape}")
|
||
|
||
# Apply: x = fc2(silu(fc1(text_emb))) — hidden_act = "silu" from config.json
|
||
x = text_emb_w.float() @ fc1_w.float().T + fc1_b.float()
|
||
x = torch.nn.functional.silu(x)
|
||
text_projected = x @ fc2_w.float().T + fc2_b.float()
|
||
print(f"text_embeds_projected: {text_projected.shape}")
|
||
np.save(f"{OUT}/text_embeds_projected.npy", text_projected.numpy())
|
||
|
||
# 2. Codec embedding from talker HF safetensors
|
||
print("\n=== talker safetensors ===")
|
||
from safetensors.torch import load_file
|
||
talker_st = load_file(f"{TALKER_HF}/model.safetensors")
|
||
# Find the embed_tokens / codec_embedding
|
||
for k in sorted(talker_st.keys()):
|
||
if 'embed' in k.lower():
|
||
print(f" {k}: {talker_st[k].shape}")
|
||
|
||
# The talker's token embedding is the codec embedding
|
||
codec_key = None
|
||
for k in talker_st.keys():
|
||
if 'embed_tokens' in k or 'codec_embedding' in k:
|
||
codec_key = k
|
||
break
|
||
if codec_key:
|
||
codec_emb = talker_st[codec_key].float().numpy()
|
||
print(f"codec_embedding ({codec_key}): {codec_emb.shape}")
|
||
np.save(f"{OUT}/codec_embedding.npy", codec_emb)
|
||
else:
|
||
print("WARNING: codec_embedding not found in talker safetensors")
|
||
# Fallback: try lm_head transposed
|
||
for k in talker_st.keys():
|
||
print(f" {k}: {talker_st[k].shape}")
|
||
|
||
# 3. TTS special embeddings [3, 1024] = bos(2149), eos(2150), pad(2148)
|
||
# These are TEXT token IDs that go through text_embedding + text_projection (SiLU MLP)
|
||
# NOT codec_embedding! The IDs are in the LLM text vocab space (151936)
|
||
text_emb_full = text_emb_dict['weight'] # [151936, 2048] — full vocab
|
||
special_ids = [2149, 2150, 2148] # bos, eos, pad
|
||
special_text = text_emb_full[special_ids].float() # [3, 2048]
|
||
# Apply text_projection: fc1 + SiLU + fc2
|
||
h = special_text @ fc1_w.float().T + fc1_b.float()
|
||
h = torch.nn.functional.silu(h)
|
||
special_projected = (h @ fc2_w.float().T + fc2_b.float()).numpy()
|
||
np.save(f"{OUT}/tts_special_embeds.npy", special_projected)
|
||
print(f"tts_special_embeds: {special_projected.shape} (projected via text_projection SiLU)")
|
||
|
||
# 4. Code predictor embeddings — 15 codebook embedding tables [15, 2048, 1024]
|
||
print("\n=== code_predictor_weights.pt ===")
|
||
cp = torch.load(f"{NATIVE}/code_predictor_weights.pt", map_location="cpu", weights_only=False)
|
||
cp_embs = []
|
||
for i in range(15):
|
||
key = f"model.codec_embedding.{i}.weight"
|
||
if key in cp:
|
||
cp_embs.append(cp[key].float().numpy())
|
||
if cp_embs:
|
||
cp_all = np.stack(cp_embs) # [15, 2048, 1024]
|
||
np.save(f"{OUT}/code_predictor_embeddings.npy", cp_all)
|
||
print(f"code_predictor_embeddings: {cp_all.shape}")
|
||
|
||
# 5. VQ codebooks — extract actual codebook vectors from embedding_sum / cluster_usage
|
||
print("\n=== speech_decoder_weights.pt (VQ codebooks) ===")
|
||
sd = torch.load(f"{NATIVE}/speech_decoder_weights.pt", map_location="cpu", weights_only=False)
|
||
|
||
# First codebook (rvq_first) + 15 rest codebooks (rvq_rest)
|
||
codebooks = []
|
||
# rvq_first
|
||
usage = sd['quantizer.rvq_first.vq.layers.0._codebook.cluster_usage']
|
||
emb_sum = sd['quantizer.rvq_first.vq.layers.0._codebook.embedding_sum']
|
||
# actual codebook = embedding_sum / cluster_usage (EMA-style)
|
||
cb = (emb_sum / usage.unsqueeze(1).clamp(min=1)).float().numpy()
|
||
codebooks.append(cb)
|
||
print(f" rvq_first codebook: {cb.shape}")
|
||
|
||
# rvq_rest has 15 layers (indices 0-14)
|
||
for i in range(15):
|
||
usage = sd[f'quantizer.rvq_rest.vq.layers.{i}._codebook.cluster_usage']
|
||
emb_sum = sd[f'quantizer.rvq_rest.vq.layers.{i}._codebook.embedding_sum']
|
||
cb = (emb_sum / usage.unsqueeze(1).clamp(min=1)).float().numpy()
|
||
codebooks.append(cb)
|
||
|
||
vq_all = np.stack(codebooks) # [16, 2048, 256]
|
||
np.save(f"{OUT}/vq_codebooks.npy", vq_all)
|
||
print(f"vq_codebooks: {vq_all.shape} (16 codebooks × 2048 entries × 256 dim)")
|
||
|
||
# Summary
|
||
print(f"\n=== Output: {OUT} ===")
|
||
for f in sorted(os.listdir(OUT)):
|
||
sz = os.path.getsize(f"{OUT}/{f}")
|
||
print(f" {f}: {sz/1024/1024:.1f} MB")
|