Upload scripts/analyze_model.py with huggingface_hub

scripts/analyze_model.py

@@ -0,0 +1,408 @@
+#!/usr/bin/env python3
+"""
+Phase 1: Deep Architecture Analysis of Qwen3-TTS for ExecuTorch Export
+======================================================================
+Loads the model, maps all modules with parameter counts, traces a real
+voice-clone inference to capture shapes, and identifies export blockers.
+"""
+
+import sys
+import os
+import time
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+
+# ── paths ────────────────────────────────────────────────────────────
+MODEL_PATH = os.path.expanduser("~/Documents/Qwen3-TTS/models/1.7B-Base")
+VENV_SITE = os.path.expanduser("~/Documents/Qwen3-TTS/.venv/lib/python3.10/site-packages")
+QWEN_TTS_SRC = os.path.expanduser("~/Documents/Qwen3-TTS")
+
+# Ensure the venv's site-packages is on the path so qwen_tts can be imported
+if VENV_SITE not in sys.path:
+    sys.path.insert(0, VENV_SITE)
+if QWEN_TTS_SRC not in sys.path:
+    sys.path.insert(0, QWEN_TTS_SRC)
+
+# ── helpers ──────────────────────────────────────────────────────────
+
+def count_params(module: nn.Module) -> int:
+    return sum(p.numel() for p in module.parameters())
+
+def fmt(n: int) -> str:
+    if n >= 1e9:
+        return f"{n / 1e9:.1f}B"
+    if n >= 1e6:
+        return f"{n / 1e6:.1f}M"
+    if n >= 1e3:
+        return f"{n / 1e3:.1f}K"
+    return str(n)
+
+def param_table(module: nn.Module, prefix: str = "", depth: int = 0, max_depth: int = 3):
+    """Print a hierarchical parameter table."""
+    total = count_params(module)
+    indent = "  " * depth
+    name = prefix or module.__class__.__name__
+    print(f"{indent}{name}: {fmt(total)} params")
+    if depth < max_depth:
+        for child_name, child in module.named_children():
+            child_prefix = f"{prefix}.{child_name}" if prefix else child_name
+            param_table(child, child_prefix, depth + 1, max_depth)
+
+
+# ── 1. Load Model ───────────────────────────────────────────────────
+
+print("=" * 70)
+print("PHASE 1: Deep Architecture Analysis — Qwen3-TTS 1.7B-Base")
+print("=" * 70)
+
+print("\n[1/5] Loading model from", MODEL_PATH)
+t0 = time.time()
+
+from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig
+from qwen_tts.core.models.modeling_qwen3_tts import (
+    Qwen3TTSForConditionalGeneration,
+    mel_spectrogram,
+)
+
+config = Qwen3TTSConfig.from_pretrained(MODEL_PATH)
+# Force SDPA attention for exportability
+model = Qwen3TTSForConditionalGeneration.from_pretrained(
+    MODEL_PATH,
+    config=config,
+    torch_dtype=torch.float32,
+    attn_implementation="sdpa",
+    device_map="cpu",
+)
+model.eval()
+print(f"  Loaded in {time.time() - t0:.1f}s")
+
+# ── 2. Parameter Map ────────────────────────────────────────────────
+
+print("\n[2/5] Parameter Map (hierarchical)")
+print("-" * 60)
+
+param_table(model, "Qwen3TTSForConditionalGeneration", max_depth=4)
+
+print("\n--- Top-level component sizes ---")
+components = {
+    "speaker_encoder": model.speaker_encoder,
+    "talker": model.talker,
+    "talker.model": model.talker.model,
+    "talker.text_projection": model.talker.text_projection,
+    "talker.codec_head": model.talker.codec_head,
+    "talker.code_predictor": model.talker.code_predictor,
+}
+for name, mod in components.items():
+    print(f"  {name:40s}: {fmt(count_params(mod)):>8s} params")
+
+if model.speech_tokenizer is not None and hasattr(model.speech_tokenizer, 'model'):
+    st = model.speech_tokenizer.model  # Qwen3TTSTokenizerV2Model (nn.Module)
+    print(f"  {'speech_tokenizer.model':40s}: {fmt(count_params(st)):>8s} params")
+    if hasattr(st, 'encoder'):
+        print(f"  {'speech_tokenizer.model.encoder':40s}: {fmt(count_params(st.encoder)):>8s} params")
+    if hasattr(st, 'decoder'):
+        print(f"  {'speech_tokenizer.model.decoder':40s}: {fmt(count_params(st.decoder)):>8s} params")
+
+# ── 3. Config Summary ───────────────────────────────────────────────
+
+print("\n[3/5] Key Config Values")
+print("-" * 60)
+
+tc = config.talker_config
+cpc = tc.code_predictor_config
+sec = config.speaker_encoder_config
+
+info = {
+    "Speaker Encoder": {
+        "mel_dim": sec.mel_dim,
+        "enc_dim (output)": sec.enc_dim,
+        "enc_channels": sec.enc_channels,
+        "sample_rate": sec.sample_rate,
+    },
+    "Talker (Main LM)": {
+        "hidden_size": tc.hidden_size,
+        "num_hidden_layers": tc.num_hidden_layers,
+        "num_attention_heads": tc.num_attention_heads,
+        "num_key_value_heads": tc.num_key_value_heads,
+        "head_dim": tc.head_dim,
+        "intermediate_size": tc.intermediate_size,
+        "text_vocab_size": tc.text_vocab_size,
+        "codec_vocab_size": tc.vocab_size,
+        "num_code_groups": tc.num_code_groups,
+        "max_position_embeddings": tc.max_position_embeddings,
+        "rope_scaling": tc.rope_scaling,
+    },
+    "Code Predictor": {
+        "hidden_size": cpc.hidden_size,
+        "num_hidden_layers": cpc.num_hidden_layers,
+        "num_attention_heads": cpc.num_attention_heads,
+        "num_key_value_heads": cpc.num_key_value_heads,
+        "num_code_groups": cpc.num_code_groups,
+        "vocab_size": cpc.vocab_size,
+    },
+}
+
+for section, kvs in info.items():
+    print(f"\n  {section}:")
+    for k, v in kvs.items():
+        print(f"    {k:35s}: {v}")
+
+# ── 4. Trace Real Inference ─────────────────────────────────────────
+
+print("\n[4/5] Tracing Real Voice-Clone Inference")
+print("-" * 60)
+
+# Create synthetic reference audio: 3 seconds of white noise at 24kHz
+ref_sr = 24000
+ref_duration = 3.0
+ref_audio = np.random.randn(int(ref_sr * ref_duration)).astype(np.float32) * 0.1
+
+# --- 4a. Speaker Encoder ---
+print("\n  === Speaker Encoder ===")
+mels = mel_spectrogram(
+    torch.from_numpy(ref_audio).unsqueeze(0),
+    n_fft=1024,
+    num_mels=128,
+    sampling_rate=24000,
+    hop_size=256,
+    win_size=1024,
+    fmin=0,
+    fmax=12000,
+).transpose(1, 2)
+print(f"  Mel input shape:        {list(mels.shape)}")  # [1, T, 128]
+
+with torch.no_grad():
+    spk_embed = model.speaker_encoder(mels)
+print(f"  Speaker embedding shape: {list(spk_embed.shape)}")  # [1, enc_dim]
+x_vector = spk_embed[0]
+print(f"  X-vector (per sample):   {list(x_vector.shape)}")  # [enc_dim]
+
+# --- 4b. Speech Tokenizer Encode (ref audio -> codes) ---
+print("\n  === Speech Tokenizer Encode ===")
+if model.speech_tokenizer is not None:
+    st_model = model.speech_tokenizer.model
+    ref_wav_tensor = torch.from_numpy(ref_audio).unsqueeze(0).float()  # [1, samples]
+    padding_mask = torch.ones_like(ref_wav_tensor, dtype=torch.long)
+    with torch.no_grad():
+        enc_out = st_model.encode(ref_wav_tensor, padding_mask=padding_mask, return_dict=True)
+    ref_codes = enc_out.audio_codes
+    print(f"  Ref audio samples:      {ref_wav_tensor.shape[1]}")
+    print(f"  Number of code tensors: {len(ref_codes)}")
+    for i, c in enumerate(ref_codes):
+        print(f"  ref_codes[{i}] shape:     {list(c.shape)}")  # [T, num_quantizers]
+else:
+    print("  Speech tokenizer not loaded (will skip encode)")
+    ref_codes = None
+
+# --- 4c. Talker Prefill Input Construction ---
+print("\n  === Talker Input Construction ===")
+
+# Simulate tokenized text: "<|im_start|>assistant\nHello world<|im_end|>\n<|im_start|>assistant\n"
+# Using config token IDs
+from transformers import AutoTokenizer
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+    text = "Hello world."
+    chat_text = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+    input_ids = tokenizer(chat_text, return_tensors="pt", add_special_tokens=False).input_ids
+    print(f"  Text input_ids shape:   {list(input_ids.shape)}")
+    print(f"  Text input_ids:         {input_ids[0].tolist()[:20]}...")
+except Exception as e:
+    print(f"  Tokenizer load failed: {e}")
+    # Fallback: synthetic token IDs
+    input_ids = torch.tensor([[config.im_start_token_id, 77091, 198, 9707, 1879, 13,
+                               config.im_end_token_id, 198,
+                               config.im_start_token_id, 77091, 198]])
+    print(f"  Fallback input_ids shape: {list(input_ids.shape)}")
+
+# --- 4d. Talker Key Shapes ---
+print("\n  === Talker Architecture Key Shapes ===")
+
+talker = model.talker
+
+# Text embedding
+text_emb = talker.get_text_embeddings()
+print(f"  text_embedding:         {text_emb.weight.shape}")  # [text_vocab, hidden]
+
+# Codec embedding
+codec_emb = talker.get_input_embeddings()
+print(f"  codec_embedding:        {codec_emb.weight.shape}")  # [codec_vocab, hidden]
+
+# text_projection (ResizeMLP)
+print(f"  text_projection type:   {type(talker.text_projection).__name__}")
+with torch.no_grad():
+    sample_text_hidden = text_emb(torch.tensor([[0]]))
+    proj_out = talker.text_projection(sample_text_hidden)
+print(f"  text_projection in/out: {list(sample_text_hidden.shape)} -> {list(proj_out.shape)}")
+
+# codec_head
+print(f"  codec_head:             Linear({talker.codec_head.in_features} -> {talker.codec_head.out_features})")
+
+# KV cache dimensions
+num_layers = tc.num_hidden_layers
+num_kv_heads = tc.num_key_value_heads
+head_dim = tc.head_dim
+print(f"\n  Static KV cache per layer: 2 x [B, {num_kv_heads}, max_seq_len, {head_dim}]")
+print(f"  Total KV layers:        {num_layers}")
+print(f"  Total KV cache (fp32, B=1, seq=2048): "
+      f"{2 * num_layers * num_kv_heads * 2048 * head_dim * 4 / 1e6:.1f} MB")
+
+# --- 4e. Code Predictor Key Shapes ---
+print("\n  === Code Predictor Key Shapes ===")
+cp = talker.code_predictor
+
+print(f"  small_to_mtp_projection: {type(cp.small_to_mtp_projection).__name__}")
+if hasattr(cp.small_to_mtp_projection, 'weight'):
+    print(f"    weight shape:         {list(cp.small_to_mtp_projection.weight.shape)}")
+
+print(f"  lm_heads:               {len(cp.lm_head)} heads")
+for i, head in enumerate(cp.lm_head):
+    print(f"    lm_head[{i}]:           Linear({head.in_features} -> {head.out_features})")
+
+print(f"  codec_embeddings:       {len(cp.model.codec_embedding)} embeddings")
+for i, emb in enumerate(cp.model.codec_embedding):
+    print(f"    codec_embedding[{i}]:   {emb.weight.shape}")
+
+cp_layers = cpc.num_hidden_layers
+cp_kv_heads = cpc.num_key_value_heads
+cp_head_dim = cpc.head_dim
+print(f"\n  Static KV cache per layer: 2 x [B, {cp_kv_heads}, max_seq_len, {cp_head_dim}]")
+print(f"  Total KV layers:        {cp_layers}")
+
+# --- 4f. Speech Tokenizer Decoder Key Shapes ---
+print("\n  === Speech Tokenizer Decoder Key Shapes ===")
+if model.speech_tokenizer is not None:
+    st_dec = model.speech_tokenizer.model.decoder
+    print(f"  Decoder type:           {type(st_dec).__name__}")
+    print(f"  Total params:           {fmt(count_params(st_dec))}")
+
+    # Test decode with synthetic codes
+    # codes shape: [batch, num_quantizers, seq_len]
+    test_codes = torch.randint(0, 2048, (1, 16, 10))
+    with torch.no_grad():
+        test_wav = st_dec(test_codes)
+    print(f"  Test input codes:       {list(test_codes.shape)}")
+    print(f"  Test output wav:        {list(test_wav.shape)}")
+    upsample_factor = test_wav.shape[-1] // test_codes.shape[-1]
+    print(f"  Upsample factor:        {upsample_factor}x")
+
+# ── 5. Export Blocker Analysis ───────────────────────────────────────
+
+print("\n[5/5] Export Blocker Analysis")
+print("-" * 60)
+
+blockers = []
+
+# Check speaker encoder
+print("\n  === Speaker Encoder Export Blockers ===")
+se_issues = []
+# Conv1d with padding="same" and padding_mode="reflect"
+for name, mod in model.speaker_encoder.named_modules():
+    if isinstance(mod, nn.Conv1d):
+        if hasattr(mod, 'padding') and mod.padding == 'same':
+            se_issues.append(f"Conv1d '{name}' uses padding='same' (dynamic pad calc)")
+        if hasattr(mod, 'padding_mode') and mod.padding_mode == 'reflect':
+            se_issues.append(f"Conv1d '{name}' uses padding_mode='reflect'")
+
+# AttentiveStatisticsPooling dynamic masking
+se_issues.append("AttentiveStatisticsPooling: dynamic _length_to_mask(), .repeat(), masked_fill_")
+se_issues.append("Res2NetBlock: torch.chunk + for loop (but fixed scale=8, should be OK)")
+
+for issue in se_issues:
+    print(f"  [!] {issue}")
+blockers.extend([("speaker_encoder", i) for i in se_issues])
+
+# Check talker
+print("\n  === Talker Export Blockers ===")
+t_issues = []
+t_issues.append("MROPE: 3D rotary embedding with sections [24,20,20] — need custom handling")
+t_issues.append("DynamicCache: must replace with static KV cache tensors")
+t_issues.append("create_causal_mask/create_sliding_window_causal_mask from transformers")
+t_issues.append("Two embedding tables (text + codec) with interleaving logic")
+t_issues.append("code_predictor.generate() called inside forward() — autoregressive sub-loop")
+t_issues.append("trailing_text_hidden conditional addition in decode step")
+t_issues.append("@can_return_tuple decorator")
+t_issues.append("@use_kernel_forward_from_hub on RMSNorm")
+
+for issue in t_issues:
+    print(f"  [!] {issue}")
+blockers.extend([("talker", i) for i in t_issues])
+
+# Check code predictor
+print("\n  === Code Predictor Export Blockers ===")
+cp_issues = []
+cp_issues.append("Uses GenerationMixin.generate() — full autoregressive loop")
+cp_issues.append("generation_steps counter used to index into lm_head ModuleList")
+cp_issues.append("DynamicCache")
+cp_issues.append("get_input_embeddings() returns ModuleList (indexed by generation step)")
+
+for issue in cp_issues:
+    print(f"  [!] {issue}")
+blockers.extend([("code_predictor", i) for i in cp_issues])
+
+# Check speech tokenizer
+print("\n  === Speech Tokenizer Export Blockers ===")
+st_issues = []
+if model.speech_tokenizer is not None:
+    st_issues.append("chunked_decode: while loop with dynamic chunk boundaries")
+    st_issues.append("ConvTranspose1d with dynamic slicing (right_pad removal)")
+    st_issues.append("CausalConv1d: dynamic padding calculation")
+    st_issues.append("SnakeBeta: custom activation (should be OK)")
+    st_issues.append("SplitResidualVectorQuantizer: F.embedding based (OK)")
+    st_issues.append("Transformer decoder with @dynamic_rope_update and torch.autocast")
+    st_issues.append("Sliding window attention (window=72)")
+
+for issue in st_issues:
+    print(f"  [!] {issue}")
+blockers.extend([("speech_tokenizer", i) for i in st_issues])
+
+# ── Summary ───────────────────────────────────────────────────────��──
+
+print("\n" + "=" * 70)
+print("SUMMARY")
+print("=" * 70)
+
+print(f"""
+Model: Qwen3TTSForConditionalGeneration (1.7B-Base)
+Total params: {fmt(count_params(model))}
+
+Export Targets (4 modules):
+  1. Speaker Encoder       ({fmt(count_params(model.speaker_encoder))} params) — ECAPA-TDNN
+  2. Talker (Main LM)      ({fmt(count_params(model.talker.model))} + heads) — Qwen3 28L
+  3. Code Predictor         ({fmt(count_params(model.talker.code_predictor))} params) — 5L transformer
+  4. Speech Tokenizer Dec   ({fmt(count_params(model.speech_tokenizer.model.decoder)) if model.speech_tokenizer else 'N/A'} params) — Transformer + ConvTranspose
+
+Voice Clone Pipeline:
+  ref_audio (24kHz)
+    -> mel_spectrogram -> [B, T, 128]
+    -> speaker_encoder -> x_vector [B, {sec.enc_dim}]
+
+  ref_audio -> speech_tokenizer.encode -> ref_codes [T, 16]
+
+  text -> tokenizer -> input_ids
+
+  [x_vector, ref_codes, input_ids]
+    -> talker.generate() -> codec_tokens [T', 16]
+    (internally calls code_predictor.generate() per step)
+
+  codec_tokens -> speech_tokenizer.decode -> PCM waveform
+
+Key Dimensions:
+  Talker: hidden=2048, layers=28, heads=16, kv_heads=8, head_dim=128
+  Code Predictor: hidden=1024, layers=5, heads=16, kv_heads=8
+  Codec: vocab=3072 (talker), 2048 (code_predictor), 16 code groups
+  Speaker: enc_dim={sec.enc_dim}
+
+Export Strategy:
+  Phase 2: Speaker encoder — fixed mel length, handle Conv1d padding
+  Phase 3: Talker — static KV cache, unrolled MROPE, separate prefill/decode
+  Phase 4: Code predictor — static KV, unroll 15-step generation
+  Phase 5: Vocoder (decoder only) — fixed code length, handle ConvTranspose1d
+  Phase 6: INT8 via torchao int8_weight_only (instant, no calibration)
+
+Total export blockers found: {len(blockers)}
+""")
+
+print("Phase 1 analysis complete!")