146 lines
5.4 KiB
Python
146 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Generate talker codec tokens on PC using the full Python pipeline,
|
||
then save them for the tablet to decode (CP + VQ + Decoder).
|
||
|
||
This bypasses the broken talker ONNX/GGUF export and proves
|
||
that the rest of the tablet pipeline (CP + decoder) works.
|
||
|
||
Output:
|
||
- talker_output.bin: binary file with cb0 tokens + hidden states
|
||
Format: int32 n_tokens, then for each token:
|
||
int32 cb0, float32[1024] hidden_state
|
||
"""
|
||
import sys
|
||
sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages")
|
||
import warnings; warnings.filterwarnings("ignore")
|
||
import torch, numpy as np, struct
|
||
|
||
model_path = "/home/alf/.cache/huggingface/hub/models--Qwen--Qwen3-TTS-12Hz-0.6B-Base/snapshots/5d83992436eae1d760afd27aff78a71d676296fc"
|
||
|
||
from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration
|
||
from transformers import AutoTokenizer
|
||
|
||
print("Loading model...")
|
||
model = Qwen3TTSForConditionalGeneration.from_pretrained(model_path, dtype=torch.float32, device_map='cpu')
|
||
model.eval()
|
||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||
spk = torch.from_numpy(np.load('models_qnn/qwen3-tts-embeddings/damien_spk_embedding.npy')).float()
|
||
|
||
TEXT = "Bonjour, je m'appelle Kazeia, je suis encore en phase de développement."
|
||
|
||
# Hook to capture hidden states from talker
|
||
hidden_states = []
|
||
cb0_tokens = []
|
||
|
||
orig_forward = model.talker.model.forward
|
||
|
||
def capture_forward(*args, **kwargs):
|
||
result = orig_forward(*args, **kwargs)
|
||
# result is a BaseModelOutputWithPast — hidden states are in last_hidden_state
|
||
if hasattr(result, 'last_hidden_state'):
|
||
h = result.last_hidden_state
|
||
if h.shape[1] == 1: # decode step
|
||
hidden_states.append(h.detach().squeeze().cpu().numpy())
|
||
return result
|
||
|
||
model.talker.model.forward = capture_forward
|
||
|
||
# Also hook codec_head to capture cb0 tokens
|
||
orig_codec_head = model.talker.codec_head
|
||
|
||
class HookedCodecHead(torch.nn.Module):
|
||
def __init__(self, head):
|
||
super().__init__()
|
||
self.head = head
|
||
def forward(self, x):
|
||
logits = self.head(x)
|
||
return logits
|
||
|
||
model.talker.codec_head = HookedCodecHead(orig_codec_head)
|
||
|
||
prompt = f"<|im_start|>assistant\n{TEXT}<|endoftext|>\n<|im_start|>assistant\n"
|
||
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
|
||
|
||
print(f"Generating: '{TEXT}'")
|
||
print(f"Input tokens: {input_ids.shape[1]}")
|
||
|
||
torch.manual_seed(42)
|
||
with torch.no_grad():
|
||
result = model.generate(
|
||
input_ids=[input_ids],
|
||
languages=["french"],
|
||
speakers=[""],
|
||
voice_clone_prompt={
|
||
"x_vector_only_mode": [True],
|
||
"icl_mode": [False],
|
||
"ref_code": None,
|
||
"ref_spk_embedding": [spk],
|
||
"x_vector": [spk],
|
||
},
|
||
do_sample=True,
|
||
temperature=0.9,
|
||
top_k=50,
|
||
repetition_penalty=1.05,
|
||
max_new_tokens=200,
|
||
)
|
||
|
||
# result is a list of [audio_codes] tensors
|
||
codes = result[0][0] # [n_tokens, 16] (16 codebooks)
|
||
n_tokens = codes.shape[0]
|
||
print(f"Generated {n_tokens} tokens × {codes.shape[1]} codebooks")
|
||
print(f"Duration: {n_tokens/12:.2f}s")
|
||
print(f"First 5 cb0: {codes[:5, 0].tolist()}")
|
||
print(f"Hidden states captured: {len(hidden_states)}")
|
||
|
||
# Save ALL 16 codebooks directly — tablet doesn't need to run CP at all!
|
||
# Format: int32 n_tokens, int32 n_codebooks, then [n_tokens × n_codebooks] int32
|
||
OUT = "/opt/Kazeia/tablet_codes.bin"
|
||
with open(OUT, "wb") as f:
|
||
f.write(struct.pack("<ii", n_tokens, codes.shape[1]))
|
||
for t in range(n_tokens):
|
||
for cb in range(codes.shape[1]):
|
||
f.write(struct.pack("<i", int(codes[t, cb].item())))
|
||
|
||
print(f"\nSaved {n_tokens} × {codes.shape[1]} codes to {OUT}")
|
||
print(f"File size: {4*2 + n_tokens*codes.shape[1]*4} bytes")
|
||
|
||
# Also save as numpy for easy loading
|
||
np.save("/opt/Kazeia/tablet_codes.npy", codes.numpy())
|
||
|
||
# Decode locally for comparison WAV
|
||
print("\nDecoding locally...")
|
||
speech_tokenizer_path = "/home/alf/.cache/huggingface/hub/models--Qwen--Qwen3-TTS-Tokenizer-12Hz/snapshots/7dd38ad4e9bad454aae9cd937d0cd577604fe229"
|
||
from qwen_tts.core.tokenizer_25hz.modeling_qwen3_tts_tokenizer_v1 import Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model
|
||
import json
|
||
from safetensors.torch import load_file
|
||
|
||
with open(f"{speech_tokenizer_path}/config.json") as f:
|
||
tok_cfg = json.load(f)
|
||
tok_config = Qwen3TTSTokenizerV1Config(**tok_cfg)
|
||
tok_model = Qwen3TTSTokenizerV1Model(tok_config)
|
||
tok_state = load_file(f"{speech_tokenizer_path}/model.safetensors")
|
||
tok_model.load_state_dict(tok_state, strict=False)
|
||
tok_model.eval()
|
||
|
||
# Decode codes → audio
|
||
with torch.no_grad():
|
||
audio_codes = codes.unsqueeze(0).long() # [1, n_tokens, 16]
|
||
# The decode method needs xvectors and ref_mels, but for x_vector_only_mode
|
||
# we can try direct decode
|
||
# Actually, the tokenizer decode takes audio_codes and produces waveform
|
||
dummy_xvec = spk.unsqueeze(0)
|
||
dummy_mel = torch.zeros(1, 1, 128)
|
||
output = tok_model.decode(audio_codes[:, :, 0], dummy_xvec, dummy_mel)
|
||
|
||
print(f"Decode output type: {type(output)}")
|
||
if hasattr(output, 'audio_values') or isinstance(output, tuple):
|
||
audio = output[0] if isinstance(output, tuple) else output.audio_values
|
||
if isinstance(audio, list):
|
||
audio = audio[0]
|
||
audio_np = audio.numpy() if hasattr(audio, 'numpy') else np.array(audio)
|
||
print(f"Audio: {audio_np.shape}, range [{audio_np.min():.3f}, {audio_np.max():.3f}]")
|
||
import soundfile as sf
|
||
sf.write("/opt/Kazeia/kazeia_tablet_codes_decoded_pc.wav", audio_np, 24000)
|
||
print("Saved decoded WAV")
|