Files

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())