367 lines
15 KiB
Python
367 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Export Qwen3-TTS speech decoder V2 to full ONNX for GPU QNN backend.
|
||
Exact architecture from state dict analysis:
|
||
pre_conv: Conv1d(512, 1024, kernel=3, pad=1)
|
||
pre_transformer: 8-layer transformer (dim=512, attn_dim=1024, heads=16, head_dim=64)
|
||
upsample: 2× ConvTranspose1d(1024,1024) + ConvNeXtV2
|
||
BigVGAN: 7 blocks with Snake activation, upsample_rates=[8,5,4,3]
|
||
"""
|
||
import sys, os, math, warnings
|
||
sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages")
|
||
warnings.filterwarnings("ignore")
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import numpy as np
|
||
|
||
state = torch.load("/opt/Kazeia/models_qnn/qwen3-tts-native/speech_decoder_weights.pt", map_location="cpu")
|
||
print(f"Loaded {len(state)} tensors")
|
||
|
||
OUT = "/opt/Kazeia/models_qnn/qwen3-tts-decoder-full-onnx"
|
||
for d in ["pre_conv", "preprocessor", "conv_decoder"]:
|
||
os.makedirs(f"{OUT}/{d}", exist_ok=True)
|
||
|
||
SEQ = 60
|
||
|
||
# ===== RMSNorm =====
|
||
class RMSNorm(nn.Module):
|
||
def __init__(self, dim, eps=1e-5):
|
||
super().__init__()
|
||
self.weight = nn.Parameter(torch.ones(dim))
|
||
self.eps = eps
|
||
def forward(self, x):
|
||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
|
||
|
||
# ===== DiT Layer =====
|
||
class DiTLayer(nn.Module):
|
||
def __init__(self, dim=512, attn_dim=1024, n_heads=16, head_dim=64, ffn_dim=1024, rope_theta=10000.0):
|
||
super().__init__()
|
||
self.dim = dim
|
||
self.n_heads = n_heads
|
||
self.head_dim = head_dim
|
||
self.input_layernorm = RMSNorm(dim)
|
||
self.post_attention_layernorm = RMSNorm(dim)
|
||
self.self_attn = nn.Module()
|
||
self.self_attn.q_proj = nn.Linear(dim, attn_dim, bias=False)
|
||
self.self_attn.k_proj = nn.Linear(dim, attn_dim, bias=False)
|
||
self.self_attn.v_proj = nn.Linear(dim, attn_dim, bias=False)
|
||
self.self_attn.o_proj = nn.Linear(attn_dim, dim, bias=False)
|
||
self.self_attn_layer_scale = nn.Module()
|
||
self.self_attn_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01)
|
||
self.mlp = nn.Module()
|
||
self.mlp.gate_proj = nn.Linear(dim, ffn_dim, bias=False)
|
||
self.mlp.up_proj = nn.Linear(dim, ffn_dim, bias=False)
|
||
self.mlp.down_proj = nn.Linear(ffn_dim, dim, bias=False)
|
||
self.mlp_layer_scale = nn.Module()
|
||
self.mlp_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01)
|
||
# Pre-compute RoPE
|
||
freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
|
||
t = torch.arange(512).float()
|
||
emb = torch.outer(t, freqs)
|
||
self.register_buffer("rope_cos", emb.cos())
|
||
self.register_buffer("rope_sin", emb.sin())
|
||
|
||
def apply_rope(self, x, seq_len):
|
||
cos = self.rope_cos[:seq_len].unsqueeze(0).unsqueeze(0) # [1,1,T,D/2]
|
||
sin = self.rope_sin[:seq_len].unsqueeze(0).unsqueeze(0)
|
||
x1, x2 = x[..., ::2], x[..., 1::2]
|
||
out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
|
||
return out.flatten(-2)
|
||
|
||
def forward(self, x):
|
||
B, T, _ = x.shape
|
||
h = self.input_layernorm(x)
|
||
q = self.self_attn.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
|
||
k = self.self_attn.k_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
|
||
v = self.self_attn.v_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
|
||
q = self.apply_rope(q, T)
|
||
k = self.apply_rope(k, T)
|
||
attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||
mask = torch.triu(torch.full((T, T), float('-inf'), device=x.device), diagonal=1)
|
||
attn = attn + mask.unsqueeze(0).unsqueeze(0)
|
||
attn = F.softmax(attn, dim=-1)
|
||
out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(B, T, -1)
|
||
x = x + self.self_attn_layer_scale.scale * self.self_attn.o_proj(out)
|
||
h2 = self.post_attention_layernorm(x)
|
||
ffn = self.mlp.down_proj(F.silu(self.mlp.gate_proj(h2)) * self.mlp.up_proj(h2))
|
||
x = x + self.mlp_layer_scale.scale * ffn
|
||
return x
|
||
|
||
# ===== Pre-transformer =====
|
||
class PreTransformer(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.input_proj = nn.Linear(1024, 512)
|
||
self.layers = nn.ModuleList([DiTLayer() for _ in range(8)])
|
||
self.norm = RMSNorm(512)
|
||
self.output_proj = nn.Linear(512, 1024)
|
||
def forward(self, x):
|
||
x = x.transpose(1, 2) # [B,C,T] → [B,T,C]
|
||
x = self.input_proj(x)
|
||
for layer in self.layers:
|
||
x = layer(x)
|
||
x = self.norm(x)
|
||
x = self.output_proj(x)
|
||
return x.transpose(1, 2) # [B,T,C] → [B,C,T]
|
||
|
||
# ===== ConvNeXt-V2 =====
|
||
class ConvNeXtV2(nn.Module):
|
||
def __init__(self, dim=1024, int_dim=4096):
|
||
super().__init__()
|
||
self.dwconv = nn.Module()
|
||
self.dwconv.conv = nn.Conv1d(dim, dim, 7, padding=3, groups=dim)
|
||
self.norm = nn.LayerNorm(dim)
|
||
self.pwconv1 = nn.Linear(dim, int_dim)
|
||
self.pwconv2 = nn.Linear(int_dim, dim)
|
||
self.gamma = nn.Parameter(torch.ones(dim) * 0.01)
|
||
def forward(self, x):
|
||
r = x
|
||
x = self.dwconv.conv(x)
|
||
x = x.transpose(1, 2)
|
||
x = self.norm(x)
|
||
x = F.gelu(self.pwconv1(x))
|
||
x = self.pwconv2(x)
|
||
x = (self.gamma * x).transpose(1, 2)
|
||
return r + x
|
||
|
||
# ===== Upsample =====
|
||
class Upsample(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
# upsampling_ratios=[2,2] from config
|
||
# ConvTranspose1d(1024,1024,kernel=2*ratio,stride=ratio,padding=ratio//2)
|
||
w0 = state["upsample.0.0.conv.weight"]
|
||
w1 = state["upsample.1.0.conv.weight"]
|
||
# For upsampling ratio r: kernel=2*r, stride=r, pad=r//2
|
||
# But kernel=2 could mean ratio=1 or ratio=2 with pad adjustment
|
||
# From config: upsampling_ratios=[2,2], so stride=2 for both
|
||
# kernel=2, stride=2, padding=0 gives exact 2× upsample
|
||
self.conv0 = nn.ConvTranspose1d(w0.shape[0], w0.shape[1], w0.shape[2], stride=2, padding=0)
|
||
self.cnx0 = ConvNeXtV2()
|
||
self.conv1 = nn.ConvTranspose1d(w1.shape[0], w1.shape[1], w1.shape[2], stride=2, padding=0)
|
||
self.cnx1 = ConvNeXtV2()
|
||
print(f" Upsample 0: {w0.shape} stride=2")
|
||
print(f" Upsample 1: {w1.shape} stride=2")
|
||
def forward(self, x):
|
||
x = self.conv0(x)
|
||
x = self.cnx0(x)
|
||
x = self.conv1(x)
|
||
x = self.cnx1(x)
|
||
return x
|
||
|
||
# ===== SnakeBeta =====
|
||
class SnakeBeta(nn.Module):
|
||
def __init__(self, channels):
|
||
super().__init__()
|
||
# Stored as [channels] in checkpoint, reshaped to [1,channels,1] in forward
|
||
self.alpha = nn.Parameter(torch.ones(channels))
|
||
self.beta = nn.Parameter(torch.ones(channels))
|
||
def forward(self, x):
|
||
a = self.alpha.view(1, -1, 1)
|
||
b = self.beta.view(1, -1, 1)
|
||
return x + (1.0 / (b + 1e-9)) * (torch.sin(a * x) ** 2)
|
||
|
||
# ===== AMPBlock (BigVGAN resblock) =====
|
||
class AMPBlock(nn.Module):
|
||
def __init__(self, channels, kernel_size=7):
|
||
super().__init__()
|
||
self.act1 = SnakeBeta(channels)
|
||
self.act2 = SnakeBeta(channels)
|
||
self.conv1 = nn.Module()
|
||
self.conv1.conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size//2)
|
||
self.conv2 = nn.Module()
|
||
self.conv2.conv = nn.Conv1d(channels, channels, kernel_size, padding=1)
|
||
def forward(self, x):
|
||
h = self.act1(x)
|
||
h = self.conv1.conv(h)
|
||
h = self.act2(h)
|
||
h = self.conv2.conv(h)
|
||
return x + h
|
||
|
||
# ===== UpsampleBigVGANBlock =====
|
||
class UpsampleBigVGANBlock(nn.Module):
|
||
def __init__(self, ch_in, ch_out, kernel, rate, pad, block_idx):
|
||
super().__init__()
|
||
self.block = nn.ModuleList()
|
||
self.block.append(SnakeBeta(ch_in))
|
||
self.block.append(nn.ConvTranspose1d(ch_in, ch_out, kernel, stride=rate, padding=pad))
|
||
for j in range(3):
|
||
amp_w1 = state[f"decoder.{block_idx}.block.{j+2}.conv1.conv.weight"]
|
||
amp_w2 = state[f"decoder.{block_idx}.block.{j+2}.conv2.conv.weight"]
|
||
amp = AMPBlock(amp_w1.shape[0], amp_w1.shape[2])
|
||
amp.conv2.conv = nn.Conv1d(amp_w2.shape[1], amp_w2.shape[0], amp_w2.shape[2], padding=amp_w2.shape[2]//2)
|
||
self.block.append(amp)
|
||
def forward(self, x):
|
||
# block[0] = SnakeBeta, block[1] = ConvTranspose1d, block[2-4] = AMPBlocks
|
||
x = self.block[0](x)
|
||
x = self.block[1](x)
|
||
for i in range(2, 5):
|
||
x = self.block[i](x)
|
||
return x
|
||
|
||
# ===== BigVGAN =====
|
||
class BigVGAN(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
# decoder.0: initial conv [1536, 1024, 7]
|
||
w0 = state["decoder.0.conv.weight"]
|
||
ch = w0.shape[0] # 1536
|
||
self.initial_conv = nn.Module()
|
||
self.initial_conv.conv = nn.Conv1d(w0.shape[1], ch, w0.shape[2], padding=w0.shape[2]//2)
|
||
|
||
# decoder.1-4: upsample blocks
|
||
# Each: SnakeBeta(ch_in) + ConvTranspose1d(ch_in, ch_out) + 3×AMPBlock(ch_out)
|
||
self.blocks = nn.ModuleList()
|
||
for i in range(4):
|
||
block_idx = i + 1
|
||
conv_w = state[f"decoder.{block_idx}.block.1.conv.weight"]
|
||
ch_in = conv_w.shape[0] # input channels for this ConvTranspose1d (= ch from snake)
|
||
ch_out = conv_w.shape[1] # output channels (Note: ConvTranspose1d weight is [ch_in, ch_out, K])
|
||
kernel = conv_w.shape[2]
|
||
# Determine stride and padding from kernel and upsample rate
|
||
# The upsample rates are [8,5,4,3], kernel = rate*2
|
||
rate = [8, 5, 4, 3][i]
|
||
pad = (kernel - rate) // 2
|
||
|
||
blk = UpsampleBigVGANBlock(ch_in, ch_out, kernel, rate, pad, block_idx)
|
||
self.blocks.append(blk)
|
||
ch = ch_out
|
||
|
||
# decoder.5: final SnakeBeta
|
||
self.final_act = SnakeBeta(ch)
|
||
# decoder.6: final conv to audio
|
||
w6 = state["decoder.6.conv.weight"]
|
||
self.final_conv = nn.Module()
|
||
self.final_conv.conv = nn.Conv1d(w6.shape[1], w6.shape[0], w6.shape[2], padding=w6.shape[2]//2)
|
||
|
||
def forward(self, x):
|
||
x = self.initial_conv.conv(x)
|
||
for blk in self.blocks:
|
||
for m in blk.block:
|
||
x = m(x)
|
||
x = self.final_act(x)
|
||
x = self.final_conv.conv(x)
|
||
x = torch.tanh(x)
|
||
return x
|
||
|
||
# ===== Build & Load =====
|
||
print("\n--- Building pre_conv ---")
|
||
pre_conv = nn.Conv1d(512, 1024, 3, padding=1)
|
||
pre_conv.weight.data = state["pre_conv.conv.weight"]
|
||
pre_conv.bias.data = state["pre_conv.conv.bias"]
|
||
pre_conv.eval()
|
||
dummy = torch.randn(1, 512, SEQ)
|
||
with torch.no_grad():
|
||
pc = pre_conv(dummy)
|
||
print(f"pre_conv: [1,512,{SEQ}] → {list(pc.shape)}")
|
||
|
||
print("\n--- Building pre_transformer ---")
|
||
pt = PreTransformer()
|
||
pt_state = {k[len("pre_transformer."):]: v for k, v in state.items() if k.startswith("pre_transformer.")}
|
||
m, u = pt.load_state_dict(pt_state, strict=False)
|
||
print(f"pre_transformer load: missing={len(m)}, unexpected={len(u)}")
|
||
if m: print(f" Missing: {m[:3]}")
|
||
pt.eval()
|
||
with torch.no_grad():
|
||
pt_out = pt(pc)
|
||
print(f"pre_transformer: {list(pc.shape)} → {list(pt_out.shape)}")
|
||
|
||
print("\n--- Building upsample ---")
|
||
up = Upsample()
|
||
up_map = {}
|
||
for k, v in state.items():
|
||
if k.startswith("upsample."):
|
||
nk = k[len("upsample."):]
|
||
# nk = "0.0.conv.weight" → conv0.weight, "0.1.dwconv.conv.weight" → cnx0.dwconv.conv.weight
|
||
parts = nk.split(".", 2)
|
||
idx = parts[0] # 0 or 1
|
||
sub = parts[1] # 0 (conv) or 1 (convnext)
|
||
rest = parts[2] if len(parts) > 2 else ""
|
||
if sub == "0":
|
||
# conv transpose: upsample.0.0.conv.weight → conv0.weight (remove ".conv" layer)
|
||
if rest.startswith("conv."):
|
||
up_map[f"conv{idx}.{rest[5:]}"] = v
|
||
else:
|
||
up_map[f"conv{idx}.{rest}"] = v
|
||
else:
|
||
up_map[f"cnx{idx}.{rest}"] = v
|
||
m, u = up.load_state_dict(up_map, strict=False)
|
||
print(f"upsample load: missing={len(m)}, unexpected={len(u)}")
|
||
if m: print(f" Missing: {m[:3]}")
|
||
up.eval()
|
||
with torch.no_grad():
|
||
up_out = up(pt_out)
|
||
print(f"upsample: {list(pt_out.shape)} → {list(up_out.shape)}")
|
||
|
||
print("\n--- Building BigVGAN ---")
|
||
bv = BigVGAN()
|
||
bv_map = {}
|
||
for k, v in state.items():
|
||
if k.startswith("decoder."):
|
||
nk = k[len("decoder."):]
|
||
parts = nk.split(".", 1)
|
||
idx = int(parts[0])
|
||
rest = parts[1] if len(parts) > 1 else ""
|
||
if idx == 0:
|
||
# decoder.0.conv.weight → initial_conv.conv.weight
|
||
bv_map[f"initial_conv.{rest}"] = v
|
||
elif 1 <= idx <= 4:
|
||
# decoder.{i}.block.1.conv.weight → blocks.{i-1}.block.1.weight (ConvTranspose1d, no ".conv")
|
||
new_rest = rest
|
||
if rest.startswith("block.1.conv."):
|
||
# ConvTranspose1d: remove extra ".conv" layer
|
||
new_rest = "block.1." + rest[len("block.1.conv."):]
|
||
bv_map[f"blocks.{idx-1}.{new_rest}"] = v
|
||
elif idx == 5:
|
||
bv_map[f"final_act.{rest}"] = v
|
||
elif idx == 6:
|
||
bv_map[f"final_conv.{rest}"] = v
|
||
m, u = bv.load_state_dict(bv_map, strict=False)
|
||
print(f"BigVGAN load: missing={len(m)}, unexpected={len(u)}")
|
||
if m: print(f" Missing: {m[:5]}")
|
||
if u: print(f" Unexpected: {u[:5]}")
|
||
bv.eval()
|
||
with torch.no_grad():
|
||
audio = bv(up_out)
|
||
print(f"BigVGAN: {list(up_out.shape)} → {list(audio.shape)}")
|
||
|
||
# ===== Export =====
|
||
print("\n=== Exporting pre_conv ===")
|
||
torch.onnx.export(pre_conv, dummy, f"{OUT}/pre_conv/model.onnx",
|
||
input_names=["x"], output_names=["pre_conv_out"], opset_version=17)
|
||
print(f" {os.path.getsize(f'{OUT}/pre_conv/model.onnx')/1024:.0f} KB")
|
||
|
||
class Preprocessor(nn.Module):
|
||
def __init__(self, pt, up):
|
||
super().__init__(); self.pt = pt; self.up = up
|
||
def forward(self, x):
|
||
return self.up(self.pt(x))
|
||
|
||
print("\n=== Exporting preprocessor ===")
|
||
prep = Preprocessor(pt, up).eval()
|
||
with torch.no_grad():
|
||
prep_out = prep(pc)
|
||
torch.onnx.export(prep, pc, f"{OUT}/preprocessor/model.onnx",
|
||
input_names=["pre_conv_out"], output_names=["hidden"], opset_version=17)
|
||
print(f" {os.path.getsize(f'{OUT}/preprocessor/model.onnx')/1024/1024:.1f} MB")
|
||
|
||
print("\n=== Exporting conv_decoder ===")
|
||
torch.onnx.export(bv, up_out, f"{OUT}/conv_decoder/model.onnx",
|
||
input_names=["hidden"], output_names=["audio"], opset_version=17)
|
||
print(f" {os.path.getsize(f'{OUT}/conv_decoder/model.onnx')/1024/1024:.1f} MB")
|
||
|
||
# ===== Validate =====
|
||
print("\n=== Validation ===")
|
||
import onnxruntime as ort
|
||
for name, inp, ref in [("pre_conv", dummy, pc), ("preprocessor", pc, prep_out), ("conv_decoder", up_out, audio)]:
|
||
sess = ort.InferenceSession(f"{OUT}/{name}/model.onnx")
|
||
o = sess.run(None, {sess.get_inputs()[0].name: inp.detach().numpy()})[0]
|
||
d = np.max(np.abs(o - ref.detach().numpy()))
|
||
print(f" {name}: max_diff={d:.6f} shape={o.shape}")
|
||
|
||
print("\nTotal:")
|
||
for n in ["pre_conv", "preprocessor", "conv_decoder"]:
|
||
print(f" {n}: {os.path.getsize(f'{OUT}/{n}/model.onnx')/1024/1024:.1f} MB")
|