907 lines
45 KiB
Python
Executable File
907 lines
45 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# convert.py: Qwen3-TTS HF checkpoint -> GGUF.
|
|
#
|
|
# Reads a HF safetensors checkpoint and writes a GGUF F32 file with a
|
|
# flat tensor naming scheme and arch-specific KV metadata. Two converter
|
|
# paths live here :
|
|
# tokenizer Mimi-style 12Hz audio codec (encoder + decoder + RVQ) with
|
|
# precomputed codebook embeddings (embedding_sum / clamp(
|
|
# cluster_usage, min=eps)) for single-shot F.embedding lookup.
|
|
# talker Qwen3-style autoregressive LM (talker + 5-layer code
|
|
# predictor MTP head) plus the optional ECAPA-TDNN speaker
|
|
# encoder shipped with Base checkpoints. BPE vocab and
|
|
# merges of the Qwen2 tokenizer are embedded in the GGUF.
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from safetensors import safe_open
|
|
|
|
import gguf
|
|
|
|
# RVQ codebook epsilon clamp, matches Qwen3TTS EuclideanCodebook.epsilon
|
|
RVQ_EPS = 1e-5
|
|
|
|
# Tokenizer 12Hz fixed shape: 4 DAC blocks (one per upsample stride) and
|
|
# 3 residual units per block (dilations 1, 3, 9).
|
|
DAC_NUM_BLOCKS = 4
|
|
DAC_RU_PER_BLOCK = 3
|
|
|
|
|
|
def rename_pre_transformer(name: str) -> str:
|
|
# decoder.pre_transformer.layers.{i}.input_layernorm.weight -> tok_dec.pre_tfm.blk.{i}.attn_norm.weight
|
|
# decoder.pre_transformer.layers.{i}.post_attention_layernorm.weight -> tok_dec.pre_tfm.blk.{i}.ffn_norm.weight
|
|
# decoder.pre_transformer.layers.{i}.self_attn.q_proj.weight -> tok_dec.pre_tfm.blk.{i}.attn_q.weight
|
|
# decoder.pre_transformer.layers.{i}.self_attn.{k,v,o}_proj.weight -> tok_dec.pre_tfm.blk.{i}.attn_{k,v,output}.weight
|
|
# decoder.pre_transformer.layers.{i}.self_attn_layer_scale.scale -> tok_dec.pre_tfm.blk.{i}.attn_scale
|
|
# decoder.pre_transformer.layers.{i}.mlp.{gate,up,down}_proj.weight -> tok_dec.pre_tfm.blk.{i}.ffn_{gate,up,down}.weight
|
|
# decoder.pre_transformer.layers.{i}.mlp_layer_scale.scale -> tok_dec.pre_tfm.blk.{i}.ffn_scale
|
|
assert name.startswith("decoder.pre_transformer.layers.")
|
|
parts = name.split(".")
|
|
idx = parts[3]
|
|
rest = parts[4:]
|
|
if rest == ["input_layernorm", "weight"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.attn_norm.weight"
|
|
if rest == ["post_attention_layernorm", "weight"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.ffn_norm.weight"
|
|
if rest[:2] == ["self_attn", "q_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.attn_q.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "k_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.attn_k.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "v_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.attn_v.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "o_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.attn_output.{rest[-1]}"
|
|
if rest == ["self_attn_layer_scale", "scale"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.attn_scale"
|
|
if rest[:2] == ["mlp", "gate_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.ffn_gate.{rest[-1]}"
|
|
if rest[:2] == ["mlp", "up_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.ffn_up.{rest[-1]}"
|
|
if rest[:2] == ["mlp", "down_proj"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.ffn_down.{rest[-1]}"
|
|
if rest == ["mlp_layer_scale", "scale"]:
|
|
return f"tok_dec.pre_tfm.blk.{idx}.ffn_scale"
|
|
raise ValueError(f"Unhandled pre_transformer tensor : {name}")
|
|
|
|
|
|
def rename_pre_conv(name: str) -> str:
|
|
# decoder.pre_conv.conv.{weight,bias} -> tok_dec.pre_conv.{weight,bias}
|
|
suffix = name.rsplit(".", 1)[-1]
|
|
return f"tok_dec.pre_conv.{suffix}"
|
|
|
|
|
|
def rename_upsample(name: str) -> str:
|
|
# Two ModuleList per stage: index 0 is the CausalTransConv, index 1
|
|
# is the ConvNeXt block.
|
|
# decoder.upsample.{i}.0.conv.{weight,bias} -> tok_dec.upsample.{i}.conv.{weight,bias}
|
|
# decoder.upsample.{i}.1.dwconv.conv.{weight,bias} -> tok_dec.upsample.{i}.dwconv.{weight,bias}
|
|
# decoder.upsample.{i}.1.norm.{weight,bias} -> tok_dec.upsample.{i}.norm.{weight,bias}
|
|
# decoder.upsample.{i}.1.pwconv{1,2}.{weight,bias} -> tok_dec.upsample.{i}.pwconv{1,2}.{weight,bias}
|
|
# decoder.upsample.{i}.1.gamma -> tok_dec.upsample.{i}.gamma
|
|
parts = name.split(".")
|
|
assert parts[0] == "decoder" and parts[1] == "upsample"
|
|
block_idx = parts[2]
|
|
sub = parts[3]
|
|
if sub == "0":
|
|
suffix = parts[-1]
|
|
return f"tok_dec.upsample.{block_idx}.conv.{suffix}"
|
|
if sub == "1":
|
|
if parts[4] == "dwconv":
|
|
suffix = parts[-1]
|
|
return f"tok_dec.upsample.{block_idx}.dwconv.{suffix}"
|
|
if parts[4] in ("norm", "pwconv1", "pwconv2"):
|
|
suffix = parts[-1]
|
|
return f"tok_dec.upsample.{block_idx}.{parts[4]}.{suffix}"
|
|
if parts[4] == "gamma":
|
|
return f"tok_dec.upsample.{block_idx}.gamma"
|
|
raise ValueError(f"Unhandled upsample tensor : {name}")
|
|
|
|
|
|
def precompute_codebook(embedding_sum: np.ndarray, cluster_usage: np.ndarray) -> np.ndarray:
|
|
# Qwen3TTS EuclideanCodebook.decode :
|
|
# embedding = embedding_sum / cluster_usage.clamp(min=epsilon)[:, None]
|
|
# Stored in F32 to match runtime precision of the codebook lookup. We
|
|
# pre-divide at convert time so the runtime can read a ready to use
|
|
# F.embedding table straight from the GGUF.
|
|
usage = np.clip(cluster_usage, RVQ_EPS, None).astype(np.float32)
|
|
sums = embedding_sum.astype(np.float32)
|
|
return sums / usage[:, None]
|
|
|
|
|
|
def rename_decoder_chain(name: str) -> str:
|
|
# Direct {i} preservation per the koboldcpp tok_dec convention :
|
|
# decoder.decoder.{i}.block.0.alpha -> tok_dec.dec.{i}.snake.alpha
|
|
# decoder.decoder.{i}.block.0.beta -> tok_dec.dec.{i}.snake.beta
|
|
# decoder.decoder.{i}.block.1.conv.{weight,bias} -> tok_dec.dec.{i}.conv_t.{weight,bias}
|
|
# decoder.decoder.{i}.block.{j}.act{1,2}.{alpha,beta} -> tok_dec.dec.{i}.res.{j-2}.act{1,2}.{alpha,beta}
|
|
# decoder.decoder.{i}.block.{j}.conv{1,2}.conv.{w,b} -> tok_dec.dec.{i}.res.{j-2}.conv{1,2}.{w,b}
|
|
# decoder.decoder.0.conv.{weight,bias} -> tok_dec.dec.0.conv.{weight,bias}
|
|
# decoder.decoder.5.{alpha,beta} -> tok_dec.dec.5.snake.{alpha,beta}
|
|
# decoder.decoder.6.conv.{weight,bias} -> tok_dec.dec.6.conv.{weight,bias}
|
|
parts = name.split(".")
|
|
assert parts[0] == "decoder" and parts[1] == "decoder"
|
|
idx = int(parts[2])
|
|
if idx == 0:
|
|
return f"tok_dec.dec.0.conv.{parts[-1]}"
|
|
if idx == 5:
|
|
return f"tok_dec.dec.5.snake.{parts[-1]}"
|
|
if idx == 6:
|
|
return f"tok_dec.dec.6.conv.{parts[-1]}"
|
|
sub = int(parts[4])
|
|
if sub == 0:
|
|
return f"tok_dec.dec.{idx}.snake.{parts[-1]}"
|
|
if sub == 1:
|
|
return f"tok_dec.dec.{idx}.conv_t.{parts[-1]}"
|
|
if sub in (2, 3, 4):
|
|
ru = sub - 2
|
|
rest = parts[5]
|
|
if rest in ("act1", "act2"):
|
|
return f"tok_dec.dec.{idx}.res.{ru}.{rest}.{parts[-1]}"
|
|
if rest in ("conv1", "conv2"):
|
|
return f"tok_dec.dec.{idx}.res.{ru}.{rest}.{parts[-1]}"
|
|
raise ValueError(f"Unhandled decoder chain tensor : {name}")
|
|
|
|
|
|
def rename_seanet(name: str) -> str:
|
|
# encoder.encoder.layers.{idx}.conv.{weight,bias} -> tok_enc.conv.{idx}.{weight,bias}
|
|
# encoder.encoder.layers.{idx}.block.{j}.conv.{weight,bias} -> tok_enc.res.{idx}.blk.{j}.{weight,bias}
|
|
# We pass the raw Python ModuleList index through so the loader
|
|
# reconstructs the SEANet topology from the upsampling_ratios array.
|
|
parts = name.split(".")
|
|
assert parts[0] == "encoder" and parts[1] == "encoder" and parts[2] == "layers"
|
|
idx = parts[3]
|
|
suffix = parts[-1]
|
|
if len(parts) == 6 and parts[4] == "conv":
|
|
return f"tok_enc.conv.{idx}.{suffix}"
|
|
if len(parts) == 8 and parts[4] == "block" and parts[6] == "conv":
|
|
sub = parts[5]
|
|
return f"tok_enc.res.{idx}.blk.{sub}.{suffix}"
|
|
raise ValueError(f"Unhandled SEANet tensor : {name}")
|
|
|
|
|
|
def rename_encoder_transformer(name: str) -> str:
|
|
# encoder.encoder_transformer.layers.{i}.input_layernorm.{weight,bias} -> tok_enc.blk.{i}.attn_norm.{weight,bias}
|
|
# encoder.encoder_transformer.layers.{i}.post_attention_layernorm.{weight,bias} -> tok_enc.blk.{i}.ffn_norm.{weight,bias}
|
|
# encoder.encoder_transformer.layers.{i}.self_attn.{q,k,v,o}_proj.weight -> tok_enc.blk.{i}.attn_{q,k,v,output}.weight
|
|
# encoder.encoder_transformer.layers.{i}.self_attn_layer_scale.scale -> tok_enc.blk.{i}.attn_scale
|
|
# encoder.encoder_transformer.layers.{i}.mlp.{fc1,fc2}.weight -> tok_enc.blk.{i}.{ffn_up,ffn_down}.weight
|
|
# encoder.encoder_transformer.layers.{i}.mlp_layer_scale.scale -> tok_enc.blk.{i}.ffn_scale
|
|
assert name.startswith("encoder.encoder_transformer.layers.")
|
|
parts = name.split(".")
|
|
idx = parts[3]
|
|
rest = parts[4:]
|
|
if rest[:1] == ["input_layernorm"]:
|
|
return f"tok_enc.blk.{idx}.attn_norm.{rest[-1]}"
|
|
if rest[:1] == ["post_attention_layernorm"]:
|
|
return f"tok_enc.blk.{idx}.ffn_norm.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "q_proj"]:
|
|
return f"tok_enc.blk.{idx}.attn_q.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "k_proj"]:
|
|
return f"tok_enc.blk.{idx}.attn_k.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "v_proj"]:
|
|
return f"tok_enc.blk.{idx}.attn_v.{rest[-1]}"
|
|
if rest[:2] == ["self_attn", "o_proj"]:
|
|
return f"tok_enc.blk.{idx}.attn_output.{rest[-1]}"
|
|
if rest == ["self_attn_layer_scale", "scale"]:
|
|
return f"tok_enc.blk.{idx}.attn_scale"
|
|
if rest[:2] == ["mlp", "fc1"]:
|
|
return f"tok_enc.blk.{idx}.ffn_up.{rest[-1]}"
|
|
if rest[:2] == ["mlp", "fc2"]:
|
|
return f"tok_enc.blk.{idx}.ffn_down.{rest[-1]}"
|
|
if rest == ["mlp_layer_scale", "scale"]:
|
|
return f"tok_enc.blk.{idx}.ffn_scale"
|
|
raise ValueError(f"Unhandled encoder transformer tensor : {name}")
|
|
|
|
|
|
def rename_encoder_downsample(name: str) -> str:
|
|
# encoder.downsample.conv.{weight,bias} -> tok_enc.downsample.{weight,bias}
|
|
suffix = name.rsplit(".", 1)[-1]
|
|
return f"tok_enc.downsample.{suffix}"
|
|
|
|
|
|
def rename_encoder_quantizer_proj(name: str) -> str:
|
|
# encoder.quantizer.{semantic|acoustic}_residual_vector_quantizer.{input_proj,output_proj}.weight ->
|
|
# tok_enc.vq_{semantic|acoustic}.{input_proj,output_proj}.weight
|
|
if "input_proj" in name:
|
|
proj = "input_proj"
|
|
elif "output_proj" in name:
|
|
proj = "output_proj"
|
|
else:
|
|
raise ValueError(f"Unknown encoder quantizer proj : {name}")
|
|
if "semantic_residual_vector_quantizer" in name:
|
|
return f"tok_enc.vq_semantic.{proj}.weight"
|
|
if "acoustic_residual_vector_quantizer" in name:
|
|
return f"tok_enc.vq_acoustic.{proj}.weight"
|
|
raise ValueError(f"Unhandled encoder quantizer proj : {name}")
|
|
|
|
|
|
def convert_tokenizer_12hz(checkpoint_dir: Path, out_path: Path) -> int:
|
|
cfg_path = checkpoint_dir / "config.json"
|
|
st_path = checkpoint_dir / "model.safetensors"
|
|
if not cfg_path.is_file() or not st_path.is_file():
|
|
print(f"[Convert] FATAL: missing checkpoint files in {checkpoint_dir}")
|
|
return 1
|
|
|
|
cfg = json.loads(cfg_path.read_text())
|
|
dec = cfg["decoder_config"]
|
|
enc = cfg["encoder_config"]
|
|
|
|
arch = "qwen3-tts-tokenizer"
|
|
writer = gguf.GGUFWriter(str(out_path), arch)
|
|
|
|
writer.add_string("general.name", "Qwen3-TTS-Tokenizer-12Hz")
|
|
|
|
# Tokenizer-level metadata
|
|
writer.add_uint32("qwen3-tts-tokenizer.input_sample_rate", cfg["input_sample_rate"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.output_sample_rate", cfg["output_sample_rate"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decode_upsample_rate", cfg["decode_upsample_rate"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encode_downsample_rate", cfg["encode_downsample_rate"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder_valid_num_quantizers", cfg["encoder_valid_num_quantizers"])
|
|
|
|
# Decoder-level metadata
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.latent_dim", dec["latent_dim"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.codebook_dim", dec["codebook_dim"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.codebook_size", dec["codebook_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.decoder_dim", dec["decoder_dim"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.hidden_size", dec["hidden_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.intermediate_size", dec["intermediate_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.head_dim", dec["head_dim"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.num_attention_heads", dec["num_attention_heads"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.num_key_value_heads", dec["num_key_value_heads"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.num_hidden_layers", dec["num_hidden_layers"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.num_quantizers", dec["num_quantizers"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.num_semantic_quantizers", dec["num_semantic_quantizers"])
|
|
writer.add_float32("qwen3-tts-tokenizer.decoder.rms_norm_eps", dec["rms_norm_eps"])
|
|
writer.add_float32("qwen3-tts-tokenizer.decoder.rope_theta", float(dec["rope_theta"]))
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.sliding_window", dec["sliding_window"])
|
|
writer.add_float32("qwen3-tts-tokenizer.decoder.layer_scale_initial_scale", dec["layer_scale_initial_scale"])
|
|
writer.add_array("qwen3-tts-tokenizer.decoder.upsample_rates", dec["upsample_rates"])
|
|
writer.add_array("qwen3-tts-tokenizer.decoder.upsampling_ratios", dec["upsampling_ratios"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.vector_quantization_hidden_dim", dec["vector_quantization_hidden_dimension"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.decoder.codebook_dim_internal", 256) # the actual codebook vector dim before output_proj
|
|
|
|
# Encoder-level metadata. The encoder is a Mimi-style stack: SEANet conv
|
|
# downsampler -> 8-layer Mimi transformer -> 1 conv downsample -> RVQ.
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_filters", enc["num_filters"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.kernel_size", enc["kernel_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.last_kernel_size", enc["last_kernel_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.residual_kernel_size", enc["residual_kernel_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_residual_layers", enc["num_residual_layers"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.dilation_growth_rate", enc["dilation_growth_rate"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.compress", enc["compress"])
|
|
writer.add_array("qwen3-tts-tokenizer.encoder.upsampling_ratios", enc["upsampling_ratios"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.hidden_size", enc["hidden_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.intermediate_size", enc["intermediate_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.head_dim", enc["head_dim"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_attention_heads", enc["num_attention_heads"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_key_value_heads", enc["num_key_value_heads"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_hidden_layers", enc["num_hidden_layers"])
|
|
writer.add_float32("qwen3-tts-tokenizer.encoder.norm_eps", enc["norm_eps"])
|
|
writer.add_float32("qwen3-tts-tokenizer.encoder.rope_theta", float(enc["rope_theta"]))
|
|
writer.add_float32("qwen3-tts-tokenizer.encoder.layer_scale_initial_scale", enc["layer_scale_initial_scale"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.codebook_dim", enc["codebook_dim"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.codebook_size", enc["codebook_size"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_quantizers", enc["num_quantizers"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.num_semantic_quantizers", enc["num_semantic_quantizers"])
|
|
writer.add_uint32("qwen3-tts-tokenizer.encoder.vector_quantization_hidden_dim", enc["vector_quantization_hidden_dimension"])
|
|
|
|
# Walk safetensors
|
|
n_added = 0
|
|
n_skipped_encoder_extra_acoustic = 0
|
|
# Truncation policy: the encoder ships 32 acoustic codebooks (semantic 1
|
|
# + acoustic 31), but encoder_valid_num_quantizers=16 means only the first
|
|
# 16 (1 semantic + 15 acoustic) are consumed at encode time and matched
|
|
# by the decoder. We drop the unused acoustic 15..30 to stay aligned with
|
|
# the decoder side and shrink the GGUF.
|
|
encoder_valid = cfg["encoder_valid_num_quantizers"]
|
|
encoder_acoustic_kept = encoder_valid - cfg["encoder_config"]["num_semantic_quantizers"]
|
|
|
|
# Pair embed_sum and cluster_usage entries by (origin, side, layer)
|
|
# for fusion at the end of the walk.
|
|
rvq_buffers = {}
|
|
|
|
with safe_open(str(st_path), framework="pt") as f:
|
|
all_keys = list(f.keys())
|
|
for k in all_keys:
|
|
t = f.get_tensor(k)
|
|
arr = t.numpy().astype(np.float32)
|
|
|
|
# Encoder SEANet conv stack
|
|
if k.startswith("encoder.encoder.layers."):
|
|
writer.add_tensor(rename_seanet(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Encoder Mimi-style transformer
|
|
if k.startswith("encoder.encoder_transformer."):
|
|
writer.add_tensor(rename_encoder_transformer(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Encoder downsample (final conv k=4 stride=2 between transformer
|
|
# output and the RVQ)
|
|
if k.startswith("encoder.downsample."):
|
|
writer.add_tensor(rename_encoder_downsample(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Encoder quantizer
|
|
if k.startswith("encoder.quantizer."):
|
|
# input_proj projects 512 -> 256 before the codebook lookup,
|
|
# output_proj projects 256 -> 512 to reconstruct the residual
|
|
# during the RVQ encode loop. Both are needed at encode time.
|
|
if k.endswith(".input_proj.weight") or k.endswith(".output_proj.weight"):
|
|
writer.add_tensor(rename_encoder_quantizer_proj(k), arr)
|
|
n_added += 1
|
|
continue
|
|
# Codebook tensors : pair embed_sum and cluster_usage and
|
|
# emit a single pre-fused codebook tensor (koboldcpp name,
|
|
# our pre-fusion semantics so the runtime can read F.embedding
|
|
# straight from disk with no extra division).
|
|
if "._codebook." in k or ".codebook." in k:
|
|
parts = k.split(".")
|
|
side_full = parts[2]
|
|
if side_full.startswith("semantic"):
|
|
side = "semantic"
|
|
elif side_full.startswith("acoustic"):
|
|
side = "acoustic"
|
|
else:
|
|
raise ValueError(f"Unknown encoder quantizer side : {k}")
|
|
layer_idx = int(parts[4])
|
|
field_raw = parts[-1]
|
|
|
|
# Apply truncation : only keep the first encoder_acoustic_kept
|
|
# acoustic layers ; semantic always has 1 layer.
|
|
if side == "acoustic" and layer_idx >= encoder_acoustic_kept:
|
|
n_skipped_encoder_extra_acoustic += 1
|
|
continue
|
|
|
|
if field_raw == "initialized":
|
|
# Boolean flag, not used at runtime
|
|
continue
|
|
field = "embedding_sum" if field_raw == "embed_sum" else field_raw
|
|
rvq_buffers.setdefault(("encoder", side, layer_idx), {})[field] = arr
|
|
continue
|
|
raise ValueError(f"Unhandled encoder quantizer tensor : {k}")
|
|
|
|
# Decoder pre-conv
|
|
if k.startswith("decoder.pre_conv."):
|
|
writer.add_tensor(rename_pre_conv(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Decoder pre-transformer
|
|
if k.startswith("decoder.pre_transformer."):
|
|
# Top level pre transformer projections and norm.
|
|
if k == "decoder.pre_transformer.input_proj.weight":
|
|
writer.add_tensor("tok_dec.pre_tfm.input_proj.weight", arr)
|
|
n_added += 1
|
|
continue
|
|
if k == "decoder.pre_transformer.input_proj.bias":
|
|
writer.add_tensor("tok_dec.pre_tfm.input_proj.bias", arr)
|
|
n_added += 1
|
|
continue
|
|
if k == "decoder.pre_transformer.output_proj.weight":
|
|
writer.add_tensor("tok_dec.pre_tfm.output_proj.weight", arr)
|
|
n_added += 1
|
|
continue
|
|
if k == "decoder.pre_transformer.output_proj.bias":
|
|
writer.add_tensor("tok_dec.pre_tfm.output_proj.bias", arr)
|
|
n_added += 1
|
|
continue
|
|
if k == "decoder.pre_transformer.norm.weight":
|
|
writer.add_tensor("tok_dec.pre_tfm.norm.weight", arr)
|
|
n_added += 1
|
|
continue
|
|
# Per layer transformer block.
|
|
if k.startswith("decoder.pre_transformer.layers."):
|
|
writer.add_tensor(rename_pre_transformer(k), arr)
|
|
n_added += 1
|
|
continue
|
|
raise ValueError(f"Unhandled decoder.pre_transformer tensor : {k}")
|
|
|
|
# Decoder upsample stage
|
|
if k.startswith("decoder.upsample."):
|
|
writer.add_tensor(rename_upsample(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Decoder DAC chain (decoder.decoder.{0..6}.*)
|
|
if k.startswith("decoder.decoder."):
|
|
writer.add_tensor(rename_decoder_chain(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Decoder quantizer side
|
|
if k.startswith("decoder.quantizer."):
|
|
# output_proj is needed at decode time
|
|
if k.endswith(".output_proj.weight"):
|
|
if "rvq_first" in k:
|
|
writer.add_tensor("tok_dec.vq_first.output_proj.weight", arr)
|
|
elif "rvq_rest" in k:
|
|
writer.add_tensor("tok_dec.vq_rest.output_proj.weight", arr)
|
|
else:
|
|
raise ValueError(f"Unknown quantizer output_proj : {k}")
|
|
n_added += 1
|
|
continue
|
|
# input_proj on the decoder side has the same value as the
|
|
# encoder side and is unused at decode time : skip rather
|
|
# than carry a duplicate.
|
|
if k.endswith(".input_proj.weight"):
|
|
continue
|
|
# Codebook tensors : pair embed_sum and cluster_usage and
|
|
# emit a single pre-fused codebook tensor under the
|
|
# koboldcpp tok_dec.vq_{first,rest} namespace.
|
|
if "._codebook." in k:
|
|
parts = k.split(".")
|
|
side = parts[2]
|
|
layer_idx = int(parts[5])
|
|
field = parts[-1]
|
|
rvq_buffers.setdefault(("decoder", side, layer_idx), {})[field] = arr
|
|
continue
|
|
raise ValueError(f"Unhandled quantizer tensor : {k}")
|
|
|
|
raise ValueError(f"Unhandled top-level tensor : {k}")
|
|
|
|
# Fuse paired embed_sum and cluster_usage into a single pre-divided
|
|
# codebook tensor per layer per side. Output names follow koboldcpp.
|
|
n_codebooks_emitted = 0
|
|
for (origin, side, layer_idx), buf in sorted(rvq_buffers.items()):
|
|
if "cluster_usage" not in buf or "embedding_sum" not in buf:
|
|
print(f"[Convert] WARNING: incomplete codebook ({origin}, {side}, {layer_idx}): {list(buf.keys())}")
|
|
continue
|
|
emb = precompute_codebook(buf["embedding_sum"], buf["cluster_usage"])
|
|
if origin == "decoder":
|
|
# HF source uses rvq_first / rvq_rest, koboldcpp emits vq_first
|
|
# / vq_rest. Strip the leading "r" so the output matches the
|
|
# tok_dec.vq_{first,rest}.{layer}.codebook convention used by
|
|
# the runtime loaders and the koboldcpp HF release.
|
|
assert side in ("rvq_first", "rvq_rest"), f"Unexpected decoder side : {side}"
|
|
out_side = side[1:]
|
|
writer.add_tensor(f"tok_dec.{out_side}.{layer_idx}.codebook", emb)
|
|
else:
|
|
writer.add_tensor(f"tok_enc.vq_{side}.{layer_idx}.codebook", emb)
|
|
n_codebooks_emitted += 1
|
|
|
|
print(f"[Convert] Tensors: {n_added} written, {n_codebooks_emitted} codebooks fused")
|
|
print(f"[Convert] Truncate: {n_skipped_encoder_extra_acoustic} encoder acoustic codebook tensors dropped (kept {encoder_acoustic_kept}/{cfg['encoder_config']['num_quantizers'] - cfg['encoder_config']['num_semantic_quantizers']})")
|
|
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file()
|
|
writer.close()
|
|
print(f"[Convert] Wrote {out_path}")
|
|
return 0
|
|
|
|
|
|
def rename_talker_layer(name: str) -> str:
|
|
# talker.model.layers.{i}.input_layernorm.weight -> talker.blk.{i}.attn_norm.weight
|
|
# talker.model.layers.{i}.post_attention_layernorm.weight -> talker.blk.{i}.ffn_norm.weight
|
|
# talker.model.layers.{i}.self_attn.{q,k,v,o}_proj.weight -> talker.blk.{i}.attn_{q,k,v,output}.weight
|
|
# talker.model.layers.{i}.self_attn.{q,k}_norm.weight -> talker.blk.{i}.attn_{q,k}_norm.weight
|
|
# talker.model.layers.{i}.mlp.{gate,up,down}_proj.weight -> talker.blk.{i}.ffn_{gate,up,down}.weight
|
|
assert name.startswith("talker.model.layers.")
|
|
return _xlate_lm_layer(name, prefix="talker.model.layers.", out_prefix="talker.blk")
|
|
|
|
|
|
def rename_code_predictor_layer(name: str) -> str:
|
|
# talker.code_predictor.model.layers.{i}.<role> -> code_pred.blk.{i}.<role>
|
|
assert name.startswith("talker.code_predictor.model.layers.")
|
|
return _xlate_lm_layer(name, prefix="talker.code_predictor.model.layers.", out_prefix="code_pred.blk")
|
|
|
|
|
|
def rename_text_projection(name: str) -> str:
|
|
# talker.text_projection.linear_fc{1,2}.{weight,bias} -> talker.text_proj.fc{1,2}.{weight,bias}
|
|
assert name.startswith("talker.text_projection.")
|
|
return name.replace("talker.text_projection.linear_fc", "talker.text_proj.fc")
|
|
|
|
|
|
# Common HF -> llama.cpp suffix table for the Qwen3 backbone, shared between
|
|
# the Talker and the Code Predictor since both follow the same architecture.
|
|
_LM_LAYER_SUFFIX = {
|
|
"input_layernorm.weight": "attn_norm.weight",
|
|
"post_attention_layernorm.weight": "ffn_norm.weight",
|
|
"self_attn.q_proj.weight": "attn_q.weight",
|
|
"self_attn.k_proj.weight": "attn_k.weight",
|
|
"self_attn.v_proj.weight": "attn_v.weight",
|
|
"self_attn.o_proj.weight": "attn_output.weight",
|
|
"self_attn.q_norm.weight": "attn_q_norm.weight",
|
|
"self_attn.k_norm.weight": "attn_k_norm.weight",
|
|
"mlp.gate_proj.weight": "ffn_gate.weight",
|
|
"mlp.up_proj.weight": "ffn_up.weight",
|
|
"mlp.down_proj.weight": "ffn_down.weight",
|
|
}
|
|
|
|
|
|
def _xlate_lm_layer(name: str, prefix: str, out_prefix: str) -> str:
|
|
rest = name[len(prefix):]
|
|
dot = rest.find(".")
|
|
layer_idx = rest[:dot]
|
|
suffix = rest[dot + 1:]
|
|
new_suffix = _LM_LAYER_SUFFIX.get(suffix)
|
|
if new_suffix is None:
|
|
raise ValueError(f"Unhandled LM layer suffix : {suffix} (full name : {name})")
|
|
return f"{out_prefix}.{layer_idx}.{new_suffix}"
|
|
|
|
|
|
def load_bpe_vocab(checkpoint_dir: Path):
|
|
# Load Qwen2 BPE vocab + merges from a HF checkpoint directory and
|
|
# produce the (tokens, token_types, merges) triple expected by the
|
|
# GGUF tokenizer convention. Special tokens listed in
|
|
# tokenizer_config.json:added_tokens_decoder are tagged as user-defined
|
|
# (token_type 4) so the runtime can recognise them as verbatim chunks.
|
|
vocab = json.loads((checkpoint_dir / "vocab.json").read_text())
|
|
tok_cfg = json.loads((checkpoint_dir / "tokenizer_config.json").read_text())
|
|
added = tok_cfg.get("added_tokens_decoder", {})
|
|
|
|
max_id = max(vocab.values())
|
|
for sid in added.keys():
|
|
max_id = max(max_id, int(sid))
|
|
|
|
tokens = [None] * (max_id + 1)
|
|
token_types = [1] * (max_id + 1) # 1 = normal
|
|
|
|
for tok, tid in vocab.items():
|
|
tokens[tid] = tok
|
|
|
|
for sid_str, info in added.items():
|
|
sid = int(sid_str)
|
|
tokens[sid] = info["content"]
|
|
token_types[sid] = 4 # 4 = user-defined / special
|
|
|
|
# Empty slots (gaps in id space) get a placeholder so the GGUF array
|
|
# has no None entries.
|
|
for i, tok in enumerate(tokens):
|
|
if tok is None:
|
|
tokens[i] = f"<|unused-{i}|>"
|
|
token_types[i] = 5 # 5 = unused
|
|
|
|
# merges.txt : first line may be a "#version" comment, skip it.
|
|
merges_lines = (checkpoint_dir / "merges.txt").read_text().splitlines()
|
|
merges = [ln for ln in merges_lines if ln and not ln.startswith("#")]
|
|
|
|
return tokens, token_types, merges
|
|
|
|
|
|
def convert_talker_base(checkpoint_dir: Path, out_path: Path, model_size: str) -> int:
|
|
cfg_path = checkpoint_dir / "config.json"
|
|
st_path = checkpoint_dir / "model.safetensors"
|
|
gen_path = checkpoint_dir / "generation_config.json"
|
|
if not cfg_path.is_file() or not st_path.is_file():
|
|
print(f"[Convert] FATAL: missing checkpoint files in {checkpoint_dir}")
|
|
return 1
|
|
|
|
cfg = json.loads(cfg_path.read_text())
|
|
talker_cfg = cfg["talker_config"]
|
|
cp_cfg = talker_cfg["code_predictor_config"]
|
|
spk_cfg = cfg.get("speaker_encoder_config")
|
|
gen_cfg = json.loads(gen_path.read_text()) if gen_path.is_file() else {}
|
|
|
|
arch = "qwen3-tts"
|
|
writer = gguf.GGUFWriter(str(out_path), arch)
|
|
|
|
writer.add_string("general.name", f"Qwen3-TTS-12Hz-{model_size}-{cfg['tts_model_type']}")
|
|
|
|
# Top-level TTS metadata
|
|
writer.add_string("qwen3-tts.tokenizer_type", cfg["tokenizer_type"])
|
|
writer.add_string("qwen3-tts.model_size", cfg["tts_model_size"])
|
|
writer.add_string("qwen3-tts.model_type", cfg["tts_model_type"])
|
|
writer.add_uint32("qwen3-tts.num_code_groups", talker_cfg["num_code_groups"])
|
|
|
|
# Talker LM hyperparameters
|
|
writer.add_uint32("qwen3-tts.talker.embedding_length", talker_cfg["hidden_size"])
|
|
writer.add_uint32("qwen3-tts.talker.feed_forward_length", talker_cfg["intermediate_size"])
|
|
writer.add_uint32("qwen3-tts.talker.block_count", talker_cfg["num_hidden_layers"])
|
|
writer.add_uint32("qwen3-tts.talker.attention.head_count", talker_cfg["num_attention_heads"])
|
|
writer.add_uint32("qwen3-tts.talker.attention.head_count_kv", talker_cfg["num_key_value_heads"])
|
|
writer.add_uint32("qwen3-tts.talker.attention.key_length", talker_cfg["head_dim"])
|
|
writer.add_uint32("qwen3-tts.talker.vocab_size", talker_cfg["vocab_size"])
|
|
writer.add_uint32("qwen3-tts.talker.text_vocab_size", talker_cfg["text_vocab_size"])
|
|
writer.add_uint32("qwen3-tts.talker.text_hidden_size", talker_cfg["text_hidden_size"])
|
|
writer.add_uint32("qwen3-tts.talker.context_length", talker_cfg["max_position_embeddings"])
|
|
writer.add_float32("qwen3-tts.talker.rope.freq_base", float(talker_cfg["rope_theta"]))
|
|
writer.add_float32("qwen3-tts.talker.attention.layer_norm_rms_epsilon", float(talker_cfg["rms_norm_eps"]))
|
|
writer.add_uint32("qwen3-tts.talker.position_id_per_seconds", talker_cfg["position_id_per_seconds"])
|
|
rope_scaling = talker_cfg.get("rope_scaling") or {}
|
|
if "mrope_section" in rope_scaling:
|
|
writer.add_array("qwen3-tts.talker.rope.mrope_section", rope_scaling["mrope_section"])
|
|
writer.add_bool("qwen3-tts.talker.mrope_interleaved", bool(rope_scaling.get("interleaved", False)))
|
|
|
|
# Code predictor (subtalker) hyperparameters
|
|
writer.add_uint32("qwen3-tts.code_pred.embedding_length", cp_cfg["hidden_size"])
|
|
writer.add_uint32("qwen3-tts.code_pred.feed_forward_length", cp_cfg["intermediate_size"])
|
|
writer.add_uint32("qwen3-tts.code_pred.block_count", cp_cfg["num_hidden_layers"])
|
|
writer.add_uint32("qwen3-tts.code_pred.attention.head_count", cp_cfg["num_attention_heads"])
|
|
writer.add_uint32("qwen3-tts.code_pred.attention.head_count_kv", cp_cfg["num_key_value_heads"])
|
|
writer.add_uint32("qwen3-tts.code_pred.attention.key_length", cp_cfg["head_dim"])
|
|
writer.add_uint32("qwen3-tts.code_pred.vocab_size", cp_cfg["vocab_size"])
|
|
writer.add_uint32("qwen3-tts.code_pred.context_length", cp_cfg["max_position_embeddings"])
|
|
writer.add_float32("qwen3-tts.code_pred.rope.freq_base", float(cp_cfg["rope_theta"]))
|
|
writer.add_float32("qwen3-tts.code_pred.attention.layer_norm_rms_epsilon", float(cp_cfg["rms_norm_eps"]))
|
|
|
|
# Speaker encoder hyperparameters (Base checkpoints only). CustomVoice
|
|
# and VoiceDesign carry no speaker encoder so the keys are skipped
|
|
# entirely, the runtime detects the absence via tensor lookup.
|
|
if spk_cfg is not None:
|
|
writer.add_uint32("qwen3-tts.spk_enc.embedding_length", spk_cfg["enc_dim"])
|
|
writer.add_uint32("qwen3-tts.spk_enc.sample_rate", spk_cfg["sample_rate"])
|
|
|
|
# Codec stream special tokens
|
|
writer.add_uint32("qwen3-tts.codec.pad_id", talker_cfg["codec_pad_id"])
|
|
writer.add_uint32("qwen3-tts.codec.bos_id", talker_cfg["codec_bos_id"])
|
|
writer.add_uint32("qwen3-tts.codec.eos_id", talker_cfg["codec_eos_token_id"])
|
|
writer.add_uint32("qwen3-tts.codec.think_id", talker_cfg["codec_think_id"])
|
|
writer.add_uint32("qwen3-tts.codec.nothink_id", talker_cfg["codec_nothink_id"])
|
|
writer.add_uint32("qwen3-tts.codec.think_bos_id", talker_cfg["codec_think_bos_id"])
|
|
writer.add_uint32("qwen3-tts.codec.think_eos_id", talker_cfg["codec_think_eos_id"])
|
|
|
|
# Language id table flattened to two parallel arrays. Names stay as in
|
|
# the upstream config so the runtime can pass --lang chinese verbatim.
|
|
lang_map = talker_cfg.get("codec_language_id") or {}
|
|
lang_names = list(lang_map.keys())
|
|
lang_ids = [int(lang_map[k]) for k in lang_names]
|
|
writer.add_array("qwen3-tts.codec.language_names", lang_names)
|
|
writer.add_array("qwen3-tts.codec.language_ids", lang_ids)
|
|
|
|
# Speaker table for CustomVoice variants. Three parallel arrays indexed
|
|
# by speaker position : name, codec embedding id, and optional dialect
|
|
# name pulled from codec_language_id. Empty dialect string means the
|
|
# speaker keeps the user supplied language. Skipped entirely for Base
|
|
# and VoiceDesign which have no spk_id map.
|
|
spk_map = talker_cfg.get("spk_id") or {}
|
|
if spk_map:
|
|
dialect_map = talker_cfg.get("spk_is_dialect") or {}
|
|
spk_names = list(spk_map.keys())
|
|
spk_ids = [int(spk_map[k]) for k in spk_names]
|
|
spk_dialects = [dialect_map.get(k) or "" for k in spk_names]
|
|
spk_dialects = [d if isinstance(d, str) else "" for d in spk_dialects]
|
|
writer.add_array("qwen3-tts.codec.speaker_names", spk_names)
|
|
writer.add_array("qwen3-tts.codec.speaker_ids", spk_ids)
|
|
writer.add_array("qwen3-tts.codec.speaker_dialects", spk_dialects)
|
|
|
|
# Text-side special tokens (Qwen2 BPE)
|
|
writer.add_uint32("qwen3-tts.text.im_start_id", cfg["im_start_token_id"])
|
|
writer.add_uint32("qwen3-tts.text.im_end_id", cfg["im_end_token_id"])
|
|
writer.add_uint32("qwen3-tts.text.tts_pad_id", cfg["tts_pad_token_id"])
|
|
writer.add_uint32("qwen3-tts.text.tts_bos_id", cfg["tts_bos_token_id"])
|
|
writer.add_uint32("qwen3-tts.text.tts_eos_id", cfg["tts_eos_token_id"])
|
|
|
|
# Default sampling parameters from generation_config.json
|
|
if gen_cfg:
|
|
if "do_sample" in gen_cfg:
|
|
writer.add_bool("generation.do_sample", bool(gen_cfg["do_sample"]))
|
|
if "top_k" in gen_cfg:
|
|
writer.add_uint32("generation.top_k", int(gen_cfg["top_k"]))
|
|
if "top_p" in gen_cfg:
|
|
writer.add_float32("generation.top_p", float(gen_cfg["top_p"]))
|
|
if "temperature" in gen_cfg:
|
|
writer.add_float32("generation.temperature", float(gen_cfg["temperature"]))
|
|
if "repetition_penalty" in gen_cfg:
|
|
writer.add_float32("generation.repetition_penalty", float(gen_cfg["repetition_penalty"]))
|
|
if "subtalker_dosample" in gen_cfg:
|
|
writer.add_bool("generation.subtalker_do_sample", bool(gen_cfg["subtalker_dosample"]))
|
|
if "subtalker_top_k" in gen_cfg:
|
|
writer.add_uint32("generation.subtalker_top_k", int(gen_cfg["subtalker_top_k"]))
|
|
if "subtalker_top_p" in gen_cfg:
|
|
writer.add_float32("generation.subtalker_top_p", float(gen_cfg["subtalker_top_p"]))
|
|
if "subtalker_temperature" in gen_cfg:
|
|
writer.add_float32("generation.subtalker_temperature", float(gen_cfg["subtalker_temperature"]))
|
|
if "max_new_tokens" in gen_cfg:
|
|
writer.add_uint32("generation.max_new_tokens", int(gen_cfg["max_new_tokens"]))
|
|
|
|
# BPE tokenizer payload
|
|
bpe_tokens, bpe_token_types, bpe_merges = load_bpe_vocab(checkpoint_dir)
|
|
writer.add_string("tokenizer.ggml.model", "gpt2")
|
|
writer.add_array("tokenizer.ggml.tokens", bpe_tokens)
|
|
writer.add_array("tokenizer.ggml.token_type", bpe_token_types)
|
|
writer.add_array("tokenizer.ggml.merges", bpe_merges)
|
|
writer.add_uint32("tokenizer.ggml.eos_token_id", 151643) # <|endoftext|>
|
|
|
|
# Top level talker tensors. Renames mirror the koboldcpp TENSOR_MAP.
|
|
TALKER_TOP = {
|
|
"talker.model.codec_embedding.weight": "talker.codec_embd.weight",
|
|
"talker.model.text_embedding.weight": "talker.text_embd.weight",
|
|
"talker.model.norm.weight": "talker.output_norm.weight",
|
|
"talker.codec_head.weight": "talker.codec_head.weight",
|
|
"talker.text_projection.linear_fc1.weight": "talker.text_proj.fc1.weight",
|
|
"talker.text_projection.linear_fc1.bias": "talker.text_proj.fc1.bias",
|
|
"talker.text_projection.linear_fc2.weight": "talker.text_proj.fc2.weight",
|
|
"talker.text_projection.linear_fc2.bias": "talker.text_proj.fc2.bias",
|
|
}
|
|
|
|
# Top level code predictor tensors.
|
|
CP_TOP = {
|
|
"talker.code_predictor.model.norm.weight": "code_pred.output_norm.weight",
|
|
"talker.code_predictor.small_to_mtp_projection.weight": "code_pred.mtp_proj.weight",
|
|
"talker.code_predictor.small_to_mtp_projection.bias": "code_pred.mtp_proj.bias",
|
|
}
|
|
|
|
# Speaker encoder rename table, koboldcpp SPEAKER_ENCODER_PATTERNS plus
|
|
# standalone tensors. block 0 is the entry conv, blocks 1 to 3 hold the
|
|
# Res2Net + SE + TDNN stack.
|
|
SPK_TOP = {
|
|
"speaker_encoder.blocks.0.conv.weight": "spk_enc.conv0.weight",
|
|
"speaker_encoder.blocks.0.conv.bias": "spk_enc.conv0.bias",
|
|
"speaker_encoder.asp.conv.weight": "spk_enc.asp.conv.weight",
|
|
"speaker_encoder.asp.conv.bias": "spk_enc.asp.conv.bias",
|
|
"speaker_encoder.asp.tdnn.conv.weight": "spk_enc.asp.tdnn.weight",
|
|
"speaker_encoder.asp.tdnn.conv.bias": "spk_enc.asp.tdnn.bias",
|
|
"speaker_encoder.mfa.conv.weight": "spk_enc.mfa.weight",
|
|
"speaker_encoder.mfa.conv.bias": "spk_enc.mfa.bias",
|
|
"speaker_encoder.fc.weight": "spk_enc.fc.weight",
|
|
"speaker_encoder.fc.bias": "spk_enc.fc.bias",
|
|
}
|
|
|
|
def rename_speaker_encoder_block(k: str) -> str:
|
|
# speaker_encoder.blocks.{i}.res2net_block.blocks.{j}.conv.{weight,bias} -> spk_enc.blk.{i}.res2net.{j}.{weight,bias}
|
|
# speaker_encoder.blocks.{i}.se_block.conv{1,2}.{weight,bias} -> spk_enc.blk.{i}.se.conv{1,2}.{weight,bias}
|
|
# speaker_encoder.blocks.{i}.tdnn{1,2}.conv.{weight,bias} -> spk_enc.blk.{i}.tdnn{1,2}.{weight,bias}
|
|
parts = k.split(".")
|
|
idx = parts[2]
|
|
if parts[3] == "res2net_block":
|
|
sub = parts[5]
|
|
suffix = parts[-1]
|
|
return f"spk_enc.blk.{idx}.res2net.{sub}.{suffix}"
|
|
if parts[3] == "se_block":
|
|
which = parts[4] # conv1 or conv2
|
|
suffix = parts[-1]
|
|
return f"spk_enc.blk.{idx}.se.{which}.{suffix}"
|
|
if parts[3] in ("tdnn1", "tdnn2"):
|
|
return f"spk_enc.blk.{idx}.{parts[3]}.{parts[-1]}"
|
|
raise ValueError(f"Unhandled speaker encoder block tensor : {k}")
|
|
|
|
# Walk safetensors
|
|
n_added = 0
|
|
n_unhandled = 0
|
|
with safe_open(str(st_path), framework="pt") as f:
|
|
all_keys = sorted(list(f.keys()))
|
|
for k in all_keys:
|
|
t = f.get_tensor(k)
|
|
arr = t.float().numpy()
|
|
|
|
# Talker transformer layers
|
|
if k.startswith("talker.model.layers."):
|
|
writer.add_tensor(rename_talker_layer(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Talker top level (codec_embedding, text_embedding, norm,
|
|
# codec_head, text_projection)
|
|
if k in TALKER_TOP:
|
|
writer.add_tensor(TALKER_TOP[k], arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Code predictor transformer layers
|
|
if k.startswith("talker.code_predictor.model.layers."):
|
|
writer.add_tensor(rename_code_predictor_layer(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Code predictor codec embeddings (one per acoustic codebook)
|
|
if k.startswith("talker.code_predictor.model.codec_embedding."):
|
|
idx = k.split(".")[4]
|
|
writer.add_tensor(f"code_pred.codec_embd.{idx}.weight", arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Code predictor lm heads (one per acoustic codebook)
|
|
if k.startswith("talker.code_predictor.lm_head."):
|
|
idx = k.split(".")[3]
|
|
writer.add_tensor(f"code_pred.lm_head.{idx}.weight", arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Code predictor top level (final norm, MTP projection)
|
|
if k in CP_TOP:
|
|
writer.add_tensor(CP_TOP[k], arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
# Speaker encoder, only present in Base checkpoints
|
|
if k in SPK_TOP:
|
|
writer.add_tensor(SPK_TOP[k], arr)
|
|
n_added += 1
|
|
continue
|
|
if k.startswith("speaker_encoder.blocks.") and not k.startswith("speaker_encoder.blocks.0."):
|
|
writer.add_tensor(rename_speaker_encoder_block(k), arr)
|
|
n_added += 1
|
|
continue
|
|
|
|
print(f"[Convert] WARNING: unhandled tensor : {k} shape={tuple(t.shape)}")
|
|
n_unhandled += 1
|
|
|
|
print(f"[Convert] Tensors: {n_added} written, {n_unhandled} unhandled")
|
|
print(f"[Convert] BPE: {len(bpe_tokens)} tokens, {len(bpe_merges)} merges")
|
|
|
|
writer.write_header_to_file()
|
|
writer.write_kv_data_to_file()
|
|
writer.write_tensors_to_file()
|
|
writer.close()
|
|
print(f"[Convert] Wrote {out_path}")
|
|
return 0
|
|
|
|
|
|
CHECKPOINT_DIR = "checkpoints"
|
|
OUTPUT_DIR = "models"
|
|
|
|
# Talker checkpoints follow the upstream pattern Qwen3-TTS-12Hz-{size}-{kind}
|
|
# where size is 0.6B or 1.7B and kind is Base, CustomVoice or VoiceDesign.
|
|
# The compiled regex captures both groups for the GGUF filename suffix and
|
|
# the model_size argument fed to convert_talker_base.
|
|
TALKER_RE = re.compile(r"^Qwen3-TTS-12Hz-([0-9.]+B)-(\w+)$")
|
|
|
|
|
|
def classify(dir_name: str):
|
|
"""Return (kind, model_size) for a known checkpoint directory or None.
|
|
kind is 'tokenizer' or 'talker'. model_size is '0.6B' / '1.7B' for
|
|
talker, None for tokenizer."""
|
|
if "Tokenizer" in dir_name:
|
|
return ("tokenizer", None)
|
|
m = TALKER_RE.match(dir_name)
|
|
if m:
|
|
return ("talker", m.group(1))
|
|
return None
|
|
|
|
|
|
def output_path_for(out_dir: Path, kind: str, dir_name: str) -> Path:
|
|
"""Map a checkpoint directory name to its F32 GGUF output path."""
|
|
if kind == "tokenizer":
|
|
return out_dir / "qwen-tokenizer-12hz-F32.gguf"
|
|
m = TALKER_RE.match(dir_name)
|
|
short = f"{m.group(1).lower()}-{m.group(2).lower()}"
|
|
return out_dir / f"qwen-talker-{short}-F32.gguf"
|
|
|
|
|
|
def main() -> int:
|
|
ckpt_root = Path(CHECKPOINT_DIR)
|
|
out_dir = Path(OUTPUT_DIR)
|
|
|
|
if not ckpt_root.is_dir():
|
|
print(f"[Convert] FATAL: {ckpt_root}/ not found")
|
|
return 1
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
converted = 0
|
|
skipped_unknown: list[str] = []
|
|
rc = 0
|
|
|
|
for name in sorted(os.listdir(ckpt_root)):
|
|
ckpt = ckpt_root / name
|
|
if not ckpt.is_dir():
|
|
continue
|
|
|
|
info = classify(name)
|
|
if info is None:
|
|
skipped_unknown.append(name)
|
|
continue
|
|
|
|
kind, model_size = info
|
|
out = output_path_for(out_dir, kind, name)
|
|
|
|
if out.exists():
|
|
print(f"[Convert] skip {out.name}: exists")
|
|
converted += 1
|
|
continue
|
|
|
|
if kind == "tokenizer":
|
|
rc |= convert_tokenizer_12hz(ckpt, out)
|
|
else:
|
|
rc |= convert_talker_base(ckpt, out, model_size)
|
|
converted += 1
|
|
|
|
if skipped_unknown:
|
|
print(f"[Convert] skipped (unknown): {', '.join(skipped_unknown)}")
|
|
print(f"[Convert] done : {converted} model(s) in {out_dir}")
|
|
return rc
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|