commit add3f940a0d229dfac6cb5f35afde51bf99afff0 Author: Pascal Date: Sun May 10 15:55:49 2026 +0200 Initial release diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..58bda9c --- /dev/null +++ b/.clang-format @@ -0,0 +1,170 @@ +--- +Language: Cpp +AlignAfterOpenBracket: Align +AlignArrayOfStructures: Left +AlignConsecutiveAssignments: AcrossComments +AlignConsecutiveBitFields: AcrossComments +AlignConsecutiveDeclarations: AcrossComments +AlignConsecutiveMacros: AcrossComments +# AlignConsecutiveShortCaseStatements: AcrossComments +AlignEscapedNewlines: Left # LeftWithLastLine +AlignOperands: Align +AlignTrailingComments: + Kind: Always + OverEmptyLines: 1 +AllowAllArgumentsOnNextLine: true +AllowAllParametersOfDeclarationOnNextLine: false +# AllowBreakBeforeNoexceptSpecifier: OnlyWithParen +AllowShortBlocksOnASingleLine: Never +AllowShortCaseLabelsOnASingleLine: false +AllowShortFunctionsOnASingleLine: Inline +AllowShortIfStatementsOnASingleLine: Never +AllowShortLambdasOnASingleLine: Inline +AllowShortLoopsOnASingleLine: false +AlwaysBreakBeforeMultilineStrings: true +# Treat CUDA keywords/attributes as "attribute macros" and avoid breaking lines inside them +AttributeMacros: + - __host__ + - __device__ + - __global__ + - __forceinline__ + - __launch_bounds__ +BinPackArguments: true +BinPackParameters: false # OnePerLine +BitFieldColonSpacing: Both +BreakBeforeBraces: Custom # Attach +BraceWrapping: + AfterCaseLabel: true + AfterClass: false + AfterControlStatement: false + AfterEnum: false + AfterFunction: false + AfterNamespace: false + AfterObjCDeclaration: false + AfterStruct: false + AfterUnion: false + AfterExternBlock: false + BeforeCatch: false + BeforeElse: false + BeforeLambdaBody: false + BeforeWhile: false + IndentBraces: false + SplitEmptyFunction: false + SplitEmptyRecord: false + SplitEmptyNamespace: false +# BreakAdjacentStringLiterals: true +BreakAfterAttributes: Never +BreakBeforeBinaryOperators: None +BreakBeforeInlineASMColon: OnlyMultiline +BreakBeforeTernaryOperators: false +# BreakBinaryOperations: Never +BreakConstructorInitializers: AfterColon +# BreakFunctionDefinitionParameters: false +BreakInheritanceList: AfterComma +BreakStringLiterals: true +# BreakTemplateDeclarations: Yes +ColumnLimit: 120 +CommentPragmas: '^ IWYU pragma:' +CompactNamespaces: false +ConstructorInitializerIndentWidth: 4 +ContinuationIndentWidth: 4 +Cpp11BracedListStyle: false +DerivePointerAlignment: false +DisableFormat: false +EmptyLineBeforeAccessModifier: Leave +EmptyLineAfterAccessModifier: Never +ExperimentalAutoDetectBinPacking: false +FixNamespaceComments: true +IncludeBlocks: Regroup +IncludeCategories: + - Regex: '".*"' + Priority: 1 + SortPriority: 0 + - Regex: '^<.*\.h>' + Priority: 2 + SortPriority: 0 + - Regex: '^<.*' + Priority: 3 + SortPriority: 0 + - Regex: '.*' + Priority: 4 + SortPriority: 0 +IncludeIsMainRegex: '([-_](test|unittest))?$' +IncludeIsMainSourceRegex: '' +IndentAccessModifiers: false +IndentCaseBlocks: true +IndentCaseLabels: true +IndentExternBlock: NoIndent +IndentGotoLabels: false +IndentPPDirectives: AfterHash +IndentWidth: 4 +IndentWrappedFunctionNames: false +InsertBraces: true # NOTE: may lead to incorrect formatting +InsertNewlineAtEOF: true +JavaScriptQuotes: Leave +JavaScriptWrapImports: true +KeepEmptyLinesAtTheStartOfBlocks: false +LambdaBodyIndentation: Signature +LineEnding: LF +MacroBlockBegin: '' +MacroBlockEnd: '' +MaxEmptyLinesToKeep: 1 +NamespaceIndentation: None +ObjCBinPackProtocolList: Auto +ObjCBlockIndentWidth: 4 +ObjCSpaceAfterProperty: true +ObjCSpaceBeforeProtocolList: true +PPIndentWidth: -1 +PackConstructorInitializers: CurrentLine +PenaltyBreakAssignment: 2 +PenaltyBreakBeforeFirstCallParameter: 1 +PenaltyBreakComment: 300 +PenaltyBreakFirstLessLess: 120 +PenaltyBreakString: 1000 +PenaltyBreakTemplateDeclaration: 10 +PenaltyExcessCharacter: 1000000 +PenaltyReturnTypeOnItsOwnLine: 200 +PointerAlignment: Middle +QualifierAlignment: Left +#QualifierOrder: ['static', 'inline', 'friend', 'constexpr', 'const', 'volatile', 'type', 'restrict'] +RawStringFormats: + - Language: Cpp + Delimiters: + - cc + - CC + - cpp + - Cpp + - CPP + - 'c++' + - 'C++' + CanonicalDelimiter: '' +ReferenceAlignment: Middle +ReflowComments: false # IndentOnly +SeparateDefinitionBlocks: Always +SortIncludes: CaseInsensitive +SortUsingDeclarations: LexicographicNumeric +SpaceAfterCStyleCast: true +SpaceAfterLogicalNot: false +SpaceAfterTemplateKeyword: true +SpaceBeforeAssignmentOperators: true +SpaceBeforeCpp11BracedList: false +SpaceBeforeCtorInitializerColon: true +SpaceBeforeInheritanceColon: true +SpaceBeforeParens: ControlStatements +SpaceBeforeRangeBasedForLoopColon: true +SpaceInEmptyBlock: false +SpaceInEmptyParentheses: false +SpacesBeforeTrailingComments: 2 +SpacesInAngles: Never +SpacesInContainerLiterals: true +SpacesInLineCommentPrefix: + Minimum: 1 + Maximum: -1 +SpacesInParentheses: false +SpacesInSquareBrackets: false +SpaceBeforeSquareBrackets: false +Standard: c++17 +TabWidth: 4 +UseTab: Never +WhitespaceSensitiveMacros: ['STRINGIZE'] +... diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..6313b56 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +* text=auto eol=lf diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..28aab30 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +build/ +*.wav + +tests/*/ + +checkpoints/ +models/*.gguf +__pycache__/ diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..e6a087d --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "ggml"] + path = ggml + url = https://github.com/ServeurpersoCom/ggml.git diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..7e884db --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,109 @@ +cmake_minimum_required(VERSION 3.14) +project(qwentts-ggml LANGUAGES C CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# version.h: embed git commit hash into all binaries. +# runs on every build, only rewrites if the hash changed. +set(VERSION_OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/version.h") +add_custom_target(version ALL + COMMAND "${CMAKE_COMMAND}" "-DSRC_DIR=${CMAKE_CURRENT_SOURCE_DIR}" "-DOUTPUT=${VERSION_OUTPUT}" + -P "${CMAKE_CURRENT_SOURCE_DIR}/tools/version.cmake" + BYPRODUCTS "${VERSION_OUTPUT}" + COMMENT "Checking git version" +) + +# pthread: required explicitly on older glibc (< 2.34) where libpthread +# is not merged into libc. Modern distros link it implicitly but aarch64 +# and older x86_64 toolchains need the explicit dependency. +find_package(Threads REQUIRED) + +# Suppress MSVC fopen/sprintf deprecation warnings +if(MSVC) + add_compile_definitions(_CRT_SECURE_NO_WARNINGS) +endif() + +# Put executables and backend .so in the same directory (build root). +# Without this, ggml defaults to bin/ for .so but executables stay in root, +# and ggml_backend_load_all() can't find the backends at runtime. +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) +set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) + +# Audio tokenizer tensor names can exceed default GGML_MAX_NAME of 64 +add_compile_definitions(GGML_MAX_NAME=128) + +# Harden: mark fread/fwrite/etc with warn_unused_result on all platforms +if(NOT MSVC) + add_compile_definitions(_FORTIFY_SOURCE=2) +endif() + +# CUDA architectures: cover Turing to Blackwell for distributed binaries. +# Users can override with -DCMAKE_CUDA_ARCHITECTURES=native for local builds. +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + find_package(CUDAToolkit QUIET) + if(CUDAToolkit_FOUND AND CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.8") + set(CMAKE_CUDA_ARCHITECTURES "75-virtual;80-virtual;86-real;89-real;120a-real;121a-real") + else() + set(CMAKE_CUDA_ARCHITECTURES "75-virtual;80-virtual;86-real;89-real") + endif() +endif() + +# ggml as subdirectory, inherits GGML_CUDA, GGML_METAL, etc. from cmake flags +add_subdirectory(ggml) + +# Shared compile options and ggml linkage +macro(link_ggml_backends target) + target_include_directories(${target} PRIVATE + ${CMAKE_SOURCE_DIR}/src + ${CMAKE_SOURCE_DIR} + ${CMAKE_BINARY_DIR} + ) + target_include_directories(${target} SYSTEM PRIVATE + ${CMAKE_SOURCE_DIR}/ggml/include + ) + if(MSVC) + target_compile_options(${target} PRIVATE /W4 /wd4100 /wd4505) + else() + target_compile_options(${target} PRIVATE -Wall -Wextra -Wshadow -Wconversion + -Wno-unused-parameter -Wno-unused-function -Wno-sign-conversion) + endif() + target_link_libraries(${target} PRIVATE ggml Threads::Threads) + if(TARGET ggml-base) + target_link_libraries(${target} PRIVATE ggml-base) + endif() + foreach(backend cpu blas cuda metal vulkan) + if(TARGET ggml-${backend}) + get_target_property(CURRENT_BACKEND_TYPE ggml-${backend} TYPE) + if (CURRENT_BACKEND_TYPE STREQUAL "MODULE_LIBRARY") + # DL mode: backend is loaded at runtime via dlopen, + # skip all link-time deps. + continue() + endif() + target_link_libraries(${target} PRIVATE ggml-${backend}) + endif() + endforeach() + add_dependencies(${target} version) +endmacro() + +# Core library shared between binaries. Holds the shared infrastructure +# (error/log routing, future common helpers) that any binary linking the +# pipeline needs. STATIC because we have actual sources now. +add_library(qwen-core STATIC src/qt-error.cpp) +target_include_directories(qwen-core PUBLIC src) +target_link_libraries(qwen-core PUBLIC ggml) + +# quantize: GGUF requantizer (BF16 -> K-quants), shared policy with +# omnivoice.cpp / acestep.cpp. +add_executable(quantize tools/quantize.cpp) +link_ggml_backends(quantize) + +# qwen-codec : standalone codec CLI (codes <-> WAV via 12Hz tokenizer) +add_executable(qwen-codec tools/qwen-codec.cpp src/pipeline-codec.cpp) +target_link_libraries(qwen-codec PRIVATE qwen-core) +link_ggml_backends(qwen-codec) + +# qwen-tts : full TTS pipeline (Talker LM + 12Hz tokenizer decoder). +add_executable(qwen-tts tools/qwen-tts.cpp src/pipeline-tts.cpp src/pipeline-codec.cpp src/prompt-builder.cpp src/talker-forward.cpp src/code-predictor-forward.cpp) +target_link_libraries(qwen-tts PRIVATE qwen-core) +link_ggml_backends(qwen-tts) diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..c90faa8 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023-2026 The omnivoice.cpp authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/buildall.cmd b/buildall.cmd new file mode 100644 index 0000000..092979c --- /dev/null +++ b/buildall.cmd @@ -0,0 +1,12 @@ +@echo off + +call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" + +rem rd /s /q build 2>nul +mkdir build 2>nul +cd build + +cmake .. -DGGML_CPU_ALL_VARIANTS=ON -DGGML_CUDA=ON -DGGML_VULKAN=ON -DGGML_BACKEND_DL=ON +cmake --build . --config Release -j %NUMBER_OF_PROCESSORS% + +cd .. diff --git a/buildall.sh b/buildall.sh new file mode 100755 index 0000000..8ce58a6 --- /dev/null +++ b/buildall.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +rm -rf build +mkdir build +cd build + +export PATH=/usr/local/cuda/bin:$PATH + +cmake .. -DGGML_CPU_ALL_VARIANTS=ON -DGGML_CUDA=ON -DGGML_VULKAN=ON -DGGML_BACKEND_DL=ON +cmake --build . --config Release -j "$(nproc)" diff --git a/buildcpu.sh b/buildcpu.sh new file mode 100755 index 0000000..bfdf14e --- /dev/null +++ b/buildcpu.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +rm -rf build +mkdir build +cd build + +cmake .. -DGGML_BLAS=ON +cmake --build . --config Release -j "$(nproc)" diff --git a/buildcuda.cmd b/buildcuda.cmd new file mode 100644 index 0000000..e64cbac --- /dev/null +++ b/buildcuda.cmd @@ -0,0 +1,12 @@ +@echo off + +call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" + +rem rd /s /q build 2>nul +mkdir build 2>nul +cd build + +cmake .. -DGGML_CUDA=ON +cmake --build . --config Release -j %NUMBER_OF_PROCESSORS% + +cd .. diff --git a/buildcuda.sh b/buildcuda.sh new file mode 100755 index 0000000..67f711f --- /dev/null +++ b/buildcuda.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +rm -rf build +mkdir build +cd build + +cmake .. -DGGML_CUDA=ON -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc +cmake --build . --config Release -j "$(nproc)" diff --git a/buildtermux.sh b/buildtermux.sh new file mode 100755 index 0000000..711d9e7 --- /dev/null +++ b/buildtermux.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +rm -rf build +mkdir build +cd build + +cmake .. -DGGML_BLAS=ON -DBLAS_INCLUDE_DIRS=$PREFIX/include/openblas +cmake --build . --config Release -j "$(nproc)" diff --git a/buildvulkan.cmd b/buildvulkan.cmd new file mode 100644 index 0000000..354d000 --- /dev/null +++ b/buildvulkan.cmd @@ -0,0 +1,12 @@ +@echo off + +call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" + +rem rd /s /q build 2>nul +mkdir build 2>nul +cd build + +cmake .. -DGGML_VULKAN=ON +cmake --build . --config Release -j %NUMBER_OF_PROCESSORS% + +cd .. diff --git a/buildvulkan.sh b/buildvulkan.sh new file mode 100755 index 0000000..307e912 --- /dev/null +++ b/buildvulkan.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +rm -rf build +mkdir build +cd build + +cmake .. -DGGML_VULKAN=ON +cmake --build . --config Release -j "$(nproc)" diff --git a/checkpoints.sh b/checkpoints.sh new file mode 100755 index 0000000..f2f5b9c --- /dev/null +++ b/checkpoints.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# Download Qwen3-TTS checkpoints from HuggingFace. +# Usage: ./checkpoints.sh [variant] +# variant : tokenizer | 0.6b-base | 0.6b-customvoice | 1.7b-base +# 1.7b-customvoice | 1.7b-voicedesign | all (default) + +set -eu + +DIR="checkpoints" +mkdir -p "$DIR" + +HF="hf download --quiet" + +dl_repo() { + local name="$1" repo="$2" + local target="$DIR/$name" + if [ -d "$target" ] && [ "$(ls "$target"/*.safetensors 2>/dev/null | wc -l)" -gt 0 ]; then + echo "[OK] $name" + return + fi + echo "[Download] $name <- $repo" + $HF "$repo" --local-dir "$target" +} + +variant="${1:-all}" + +case "$variant" in + tokenizer|all) + dl_repo "Qwen3-TTS-Tokenizer-12Hz" "Qwen/Qwen3-TTS-Tokenizer-12Hz" + ;;& + 0.6b-base|all) + dl_repo "Qwen3-TTS-12Hz-0.6B-Base" "Qwen/Qwen3-TTS-12Hz-0.6B-Base" + ;;& + 0.6b-customvoice|all) + dl_repo "Qwen3-TTS-12Hz-0.6B-CustomVoice" "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" + ;;& + 1.7b-base|all) + dl_repo "Qwen3-TTS-12Hz-1.7B-Base" "Qwen/Qwen3-TTS-12Hz-1.7B-Base" + ;;& + 1.7b-customvoice|all) + dl_repo "Qwen3-TTS-12Hz-1.7B-CustomVoice" "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + ;;& + 1.7b-voicedesign|all) + dl_repo "Qwen3-TTS-12Hz-1.7B-VoiceDesign" "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign" + ;; + *) + echo "Unknown variant: $variant" + echo "Valid: tokenizer | 0.6b-base | 0.6b-customvoice | 1.7b-base | 1.7b-customvoice | 1.7b-voicedesign | all" + exit 1 + ;; +esac diff --git a/convert.py b/convert.py new file mode 100755 index 0000000..b959448 --- /dev/null +++ b/convert.py @@ -0,0 +1,907 @@ +#!/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_uint32("qwen3-tts-tokenizer.encoder.sliding_window", enc["sliding_window"]) + 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}. -> code_pred.blk.{i}. + 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()) diff --git a/examples/base.cmd b/examples/base.cmd new file mode 100644 index 0000000..c09282b --- /dev/null +++ b/examples/base.cmd @@ -0,0 +1,11 @@ +@echo off + +set PATH=%~dp0..\build\Release;%PATH% + +qwen-tts.exe ^ + --model ..\models\qwen-talker-1.7b-base-Q8_0.gguf ^ + --codec ..\models\qwen-tokenizer-12hz-Q8_0.gguf ^ + --lang English ^ + -o base.wav < prompt.txt + +pause diff --git a/examples/base.sh b/examples/base.sh new file mode 100755 index 0000000..c0fef56 --- /dev/null +++ b/examples/base.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +set -eu + +../build/qwen-tts \ + --model ../models/qwen-talker-1.7b-base-Q8_0.gguf \ + --codec ../models/qwen-tokenizer-12hz-Q8_0.gguf \ + --lang English \ + -o base.wav < prompt.txt diff --git a/examples/clone.cmd b/examples/clone.cmd new file mode 100644 index 0000000..b21ab21 --- /dev/null +++ b/examples/clone.cmd @@ -0,0 +1,15 @@ +@echo off + +set PATH=%~dp0..\build\Release;%PATH% + +set /p REFTEXT= +#include +#include +#include +#include +#include + +// wav.h: WAV reader (returns interleaved, we deinterleave below) +#include "wav.h" + +// audio-resample.h: sample rate conversion +#include "audio-resample.h" + +// case-insensitive extension check +static bool audio_io_ends_with(const char * str, const char * suffix) { + int slen = (int) strlen(str); + int xlen = (int) strlen(suffix); + if (slen < xlen) { + return false; + } + for (int i = 0; i < xlen; i++) { + char a = str[slen - xlen + i]; + char b = suffix[i]; + if (a >= 'A' && a <= 'Z') { + a += 32; + } + if (b >= 'A' && b <= 'Z') { + b += 32; + } + if (a != b) { + return false; + } + } + return true; +} + +// Load entire file into memory. Caller frees the returned pointer. +static uint8_t * audio_io_load_file(const char * path, size_t * size_out) { + *size_out = 0; + FILE * fp = fopen(path, "rb"); + if (!fp) { + fprintf(stderr, "[Audio] Cannot open %s\n", path); + return NULL; + } + fseek(fp, 0, SEEK_END); + long fsize = ftell(fp); + fseek(fp, 0, SEEK_SET); + + uint8_t * buf = (uint8_t *) malloc((size_t) fsize); + if (!buf) { + fclose(fp); + return NULL; + } + size_t nr = fread(buf, 1, (size_t) fsize, fp); + fclose(fp); + if (nr != (size_t) fsize) { + free(buf); + return NULL; + } + + *size_out = (size_t) fsize; + return buf; +} + +// Decode WAV from memory buffer. Returns planar stereo float [L:T][R:T]. +static float * audio_io_read_wav_buf(const uint8_t * data, size_t size, int * T_out, int * sr_out) { + *T_out = 0; + *sr_out = 0; + + int T = 0, sr = 0; + float * interleaved = read_wav_buf(data, size, &T, &sr); + if (!interleaved) { + return NULL; + } + + float * planar = (float *) malloc((size_t) T * 2 * sizeof(float)); + if (!planar) { + free(interleaved); + return NULL; + } + for (int t = 0; t < T; t++) { + planar[t] = interleaved[t * 2 + 0]; + planar[T + t] = interleaved[t * 2 + 1]; + } + free(interleaved); + + *T_out = T; + *sr_out = sr; + return planar; +} + +// Read WAV file. Returns planar stereo float [L: T][R: T]. Caller frees. +static float * audio_read(const char * path, int * T_out, int * sr_out) { + size_t size = 0; + uint8_t * buf = audio_io_load_file(path, &size); + if (!buf) { + *T_out = 0; + *sr_out = 0; + return NULL; + } + float * result = audio_io_read_wav_buf(buf, size, T_out, sr_out); + free(buf); + return result; +} + +// Read WAV, resample to target_sr, downmix to mono. +// Returns a flat buffer of T floats at target_sr mono. Caller frees. +static float * audio_read_mono(const char * path, int target_sr, int * T_out) { + int T = 0; + int sr = 0; + float * raw = audio_read(path, &T, &sr); + if (!raw) { + *T_out = 0; + return NULL; + } + + // Resample planar stereo to target_sr first to keep both channels + // coherent when the source rate differs. + float * stereo_rs = raw; + int T_rs = T; + if (sr != target_sr) { + fprintf(stderr, "[Audio-Resample] %d Hz -> %d Hz, %d samples...\n", sr, target_sr, T); + int T_new = 0; + float * resampled = audio_resample(raw, T, sr, target_sr, 2, &T_new); + free(raw); + if (!resampled) { + fprintf(stderr, "[Audio-Resample] Resample failed\n"); + *T_out = 0; + return NULL; + } + fprintf(stderr, "[Audio-Resample] Done: %d -> %d samples\n", T, T_new); + stereo_rs = resampled; + T_rs = T_new; + } + + // Downmix planar [L:T][R:T] to mono = 0.5 * (L + R). + float * mono = (float *) malloc((size_t) T_rs * sizeof(float)); + if (!mono) { + free(stereo_rs); + *T_out = 0; + return NULL; + } + const float * left = stereo_rs; + const float * right = stereo_rs + (size_t) T_rs; + for (int i = 0; i < T_rs; i++) { + mono[i] = 0.5f * (left[i] + right[i]); + } + free(stereo_rs); + + *T_out = T_rs; + return mono; +} + +// WAV output format +enum WavFormat { + WAV_S16, // 16-bit signed integer PCM (classic RIFF, default) + WAV_S24, // 24-bit signed integer PCM (classic RIFF) + WAV_F32, // 32-bit IEEE 754 float (classic RIFF, fmt_tag=3) +}; + +// Parse a CLI format string into a WavFormat. Accepts: wav16, wav24, wav32. +// Returns false on unknown format. +static bool audio_parse_format(const char * s, WavFormat & wav_fmt) { + if (!s) { + return false; + } + if (!strcmp(s, "wav16")) { + wav_fmt = WAV_S16; + return true; + } + if (!strcmp(s, "wav24")) { + wav_fmt = WAV_S24; + return true; + } + if (!strcmp(s, "wav32")) { + wav_fmt = WAV_F32; + return true; + } + return false; +} + +// Byte-level write helpers (endian-safe) + +static void wav_write_u16le(char *& p, uint16_t x) { + *p++ = (char) (x & 0xff); + *p++ = (char) ((x >> 8) & 0xff); +} + +static void wav_write_u24le(char *& p, uint32_t x) { + *p++ = (char) (x & 0xff); + *p++ = (char) ((x >> 8) & 0xff); + *p++ = (char) ((x >> 16) & 0xff); +} + +static void wav_write_u32le(char *& p, uint32_t x) { + *p++ = (char) (x & 0xff); + *p++ = (char) ((x >> 8) & 0xff); + *p++ = (char) ((x >> 16) & 0xff); + *p++ = (char) ((x >> 24) & 0xff); +} + +static float wav_clamp1(float x) { + return x < -1.0f ? -1.0f : (x > 1.0f ? 1.0f : x); +} + +static float wav_sanitize(float x) { + return std::isfinite(x) ? x : 0.0f; +} + +// Classic RIFF header: fmt_tag 1 (PCM int) or 3 (IEEE float), 16-byte fmt chunk +static void wav_write_header_basic(char *& p, int T_audio, int sr, int n_channels, int bits, uint16_t fmt_tag) { + uint32_t bytes_per_sample = (uint32_t) bits / 8; + uint32_t byte_rate = (uint32_t) sr * (uint32_t) n_channels * bytes_per_sample; + uint16_t block_align = (uint16_t) (n_channels * (int) bytes_per_sample); + uint32_t data_size = (uint32_t) T_audio * (uint32_t) n_channels * bytes_per_sample; + uint32_t file_size = 36 + data_size; + + memcpy(p, "RIFF", 4); + p += 4; + wav_write_u32le(p, file_size); + memcpy(p, "WAVE", 4); + p += 4; + + memcpy(p, "fmt ", 4); + p += 4; + wav_write_u32le(p, 16); + wav_write_u16le(p, fmt_tag); + wav_write_u16le(p, (uint16_t) n_channels); + wav_write_u32le(p, (uint32_t) sr); + wav_write_u32le(p, byte_rate); + wav_write_u16le(p, block_align); + wav_write_u16le(p, (uint16_t) bits); + + memcpy(p, "data", 4); + p += 4; + wav_write_u32le(p, data_size); +} + +// Encode mono float to WAV 16-bit signed integer PCM in memory. +// 44-byte classic RIFF header (fmt_tag=1) + int16 samples. +// Clamps to [-1, +1], coerces NaN/Inf to zero. +static std::string audio_encode_wav_s16(const float * audio, int T_audio, int sr) { + int n_channels = 1; + int data_size = T_audio * n_channels * 2; + + std::string out; + out.resize(44 + (size_t) data_size); + char * p = &out[0]; + + wav_write_header_basic(p, T_audio, sr, n_channels, 16, 1); + + for (int t = 0; t < T_audio; t++) { + int16_t s = (int16_t) (wav_clamp1(wav_sanitize(audio[t])) * 32767.0f); + wav_write_u16le(p, (uint16_t) s); + } + + return out; +} + +// Encode mono float to WAV 24-bit signed integer PCM in memory. +// 44-byte classic RIFF header (fmt_tag=1) + int24 samples. +// Clamps to [-1, +1], coerces NaN/Inf to zero. +static std::string audio_encode_wav_s24(const float * audio, int T_audio, int sr) { + int n_channels = 1; + int data_size = T_audio * n_channels * 3; + + std::string out; + out.resize(44 + (size_t) data_size); + char * p = &out[0]; + + wav_write_header_basic(p, T_audio, sr, n_channels, 24, 1); + + for (int t = 0; t < T_audio; t++) { + int32_t s = (int32_t) (wav_clamp1(wav_sanitize(audio[t])) * 8388607.0f); + wav_write_u24le(p, (uint32_t) s); + } + + return out; +} + +// Encode mono float to WAV 32-bit IEEE 754 float in memory. +// 44-byte classic RIFF header (fmt_tag=3) + float32 samples. +// Coerces NaN/Inf to zero. No clamping: output may exceed [-1, +1]. +static std::string audio_encode_wav_f32(const float * audio, int T_audio, int sr) { + int n_channels = 1; + int data_size = T_audio * n_channels * 4; + + std::string out; + out.resize(44 + (size_t) data_size); + char * p = &out[0]; + + wav_write_header_basic(p, T_audio, sr, n_channels, 32, 3); + + for (int t = 0; t < T_audio; t++) { + float f = wav_sanitize(audio[t]); + uint32_t u; + memcpy(&u, &f, 4); + wav_write_u32le(p, u); + } + + return out; +} + +// Encode mono float to WAV in memory in the requested format. +// audio is flat mono [T], pre-normalized by caller. +// NaN and Inf are coerced to zero. S16/S24 clamp to [-1, +1]. +static std::string audio_encode_wav(const float * audio, int T_audio, int sr, WavFormat fmt = WAV_S16) { + switch (fmt) { + case WAV_S16: + return audio_encode_wav_s16(audio, T_audio, sr); + case WAV_S24: + return audio_encode_wav_s24(audio, T_audio, sr); + case WAV_F32: + return audio_encode_wav_f32(audio, T_audio, sr); + } + return audio_encode_wav_s16(audio, T_audio, sr); +} + +// Write mono float audio to WAV file in the requested format. path "-" +// streams the encoded WAV to stdout (pipe friendly). S16/S24 hard clip +// to [-1, +1], F32 preserves the full range. +static bool audio_write_wav(const char * path, const float * audio, int T_audio, int sr, WavFormat fmt = WAV_S16) { + std::string wav = audio_encode_wav(audio, T_audio, sr, fmt); + if (wav.empty()) { + return false; + } + + const bool to_stdout = (path[0] == '-' && path[1] == '\0'); + FILE * fp = to_stdout ? stdout : fopen(path, "wb"); + if (!fp) { + fprintf(stderr, "[WAV] Cannot open %s for writing\n", path); + return false; + } + if (fwrite(wav.data(), 1, wav.size(), fp) != wav.size()) { + fprintf(stderr, "[WAV] Failed to write %s\n", path); + if (!to_stdout) { + fclose(fp); + } + return false; + } + if (to_stdout) { + fflush(fp); + } else { + fclose(fp); + } + + const char * fmt_name = (fmt == WAV_S16) ? "S16" : (fmt == WAV_S24) ? "S24" : "F32"; + fprintf(stderr, "[WAV] Wrote %s: %d samples, %d Hz, mono %s\n", to_stdout ? "" : path, T_audio, sr, + fmt_name); + return true; +} diff --git a/src/audio-mel.h b/src/audio-mel.h new file mode 100644 index 0000000..77dfda2 --- /dev/null +++ b/src/audio-mel.h @@ -0,0 +1,222 @@ +#pragma once +// audio-mel.h : log mel spectrogram extractor matching Qwen3TTS upstream. +// +// Pipeline mirrored from qwen_tts/core/models/modeling_qwen3_tts.py +// mel_spectrogram() at lines 399 to 464 : +// +// pad reflect by (n_fft - hop) / 2 +// torch.stft(n_fft, hop, win=n_fft, hann_periodic, center=False) +// mag = sqrt(real^2 + imag^2 + 1e-9) +// mel = librosa.filters.mel(sr, n_fft, n_mels, fmin, fmax) @ mag +// log_mel = log(max(mel, 1e-5)) +// +// Spec for the speaker encoder path : +// sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 +// +// GGML strategy : no native FFT op, so the DFT is folded into two +// real matmuls. We precompute on CPU two F32 matrices : +// dft_real [n_freq, n_fft] with cos(2 pi k n / n_fft) +// dft_imag [n_freq, n_fft] with -sin(2 pi k n / n_fft) +// where n_freq = n_fft / 2 + 1. The framing uses ggml_im2col on the +// padded signal viewed as a 1D conv input. See audio_mel_build_graph +// for the graph topology, and audio_mel_compute_constants for the +// cos/sin and mel basis baking. + +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include + +struct AudioMelConfig { + int sample_rate; + int n_fft; + int hop; + int n_mels; + float fmin; + float fmax; +}; + +// CPU side constants : Hann window, DFT real/imag matrices, mel filter. +// Allocated once per AudioMelConfig and uploaded to the backend as +// regular ggml tensors during graph build. +struct AudioMelConstants { + AudioMelConfig cfg; + int n_freq; + std::vector hann; // [n_fft] + std::vector dft_real; // [n_freq, n_fft] row major ne=(n_fft, n_freq) + std::vector dft_imag; // [n_freq, n_fft] row major + std::vector mel_basis; // [n_mels, n_freq] row major ne=(n_freq, n_mels) +}; + +// Slaney mel scale, the default of librosa.filters.mel. +static inline float audio_mel_hz_to_mel(float hz) { + // Slaney : linear below 1000 Hz, log above. + const float f_min = 0.0f; + const float f_sp = 200.0f / 3.0f; + const float min_log_hz = 1000.0f; + const float min_log_mel = (min_log_hz - f_min) / f_sp; + const float logstep = std::log(6.4f) / 27.0f; + if (hz < min_log_hz) { + return (hz - f_min) / f_sp; + } + return min_log_mel + std::log(hz / min_log_hz) / logstep; +} + +static inline float audio_mel_mel_to_hz(float mel) { + const float f_min = 0.0f; + const float f_sp = 200.0f / 3.0f; + const float min_log_hz = 1000.0f; + const float min_log_mel = (min_log_hz - f_min) / f_sp; + const float logstep = std::log(6.4f) / 27.0f; + if (mel < min_log_mel) { + return f_min + f_sp * mel; + } + return min_log_hz * std::exp(logstep * (mel - min_log_mel)); +} + +// Bake CPU constants once. Reproduces librosa.filters.mel(slaney) and +// torch.hann_window(periodic=True) bit-for-bit on F32, with the cos/sin +// DFT matrix evaluated at double precision then cast to float. +static void audio_mel_compute_constants(const AudioMelConfig & cfg, AudioMelConstants & c) { + c.cfg = cfg; + c.n_freq = cfg.n_fft / 2 + 1; + + // Hann periodic : 0.5 * (1 - cos(2 pi i / N)) for i in [0, N). + c.hann.assign(cfg.n_fft, 0.0f); + for (int i = 0; i < cfg.n_fft; i++) { + c.hann[i] = 0.5f * (1.0f - (float) std::cos(2.0 * M_PI * (double) i / (double) cfg.n_fft)); + } + + // DFT matrices, real and imag part, F32. Computed in F64 to keep + // the trig roundoff below the F32 ULP threshold. + c.dft_real.assign((size_t) c.n_freq * (size_t) cfg.n_fft, 0.0f); + c.dft_imag.assign((size_t) c.n_freq * (size_t) cfg.n_fft, 0.0f); + for (int k = 0; k < c.n_freq; k++) { + for (int n = 0; n < cfg.n_fft; n++) { + double th = 2.0 * M_PI * (double) k * (double) n / (double) cfg.n_fft; + c.dft_real[(size_t) k * (size_t) cfg.n_fft + (size_t) n] = (float) std::cos(th); + c.dft_imag[(size_t) k * (size_t) cfg.n_fft + (size_t) n] = (float) (-std::sin(th)); + } + } + + // Slaney mel filterbank : n_mels triangular filters between fmin and + // fmax, normalized by 2 / (mel_freqs[i+2] - mel_freqs[i]). Matches + // librosa.filters.mel(htk=False, norm='slaney') byte for byte. + const float fmin = cfg.fmin; + const float fmax = (cfg.fmax <= 0.0f) ? (float) cfg.sample_rate * 0.5f : cfg.fmax; + const float mmin = audio_mel_hz_to_mel(fmin); + const float mmax = audio_mel_hz_to_mel(fmax); + + std::vector mel_pts((size_t) cfg.n_mels + 2); + for (int i = 0; i < cfg.n_mels + 2; i++) { + mel_pts[(size_t) i] = mmin + (mmax - mmin) * (float) i / (float) (cfg.n_mels + 1); + } + std::vector hz_pts((size_t) cfg.n_mels + 2); + for (int i = 0; i < cfg.n_mels + 2; i++) { + hz_pts[(size_t) i] = audio_mel_mel_to_hz(mel_pts[(size_t) i]); + } + std::vector fft_freqs((size_t) c.n_freq); + for (int k = 0; k < c.n_freq; k++) { + fft_freqs[(size_t) k] = (float) k * (float) cfg.sample_rate / (float) cfg.n_fft; + } + + c.mel_basis.assign((size_t) cfg.n_mels * (size_t) c.n_freq, 0.0f); + for (int m = 0; m < cfg.n_mels; m++) { + const float lo = hz_pts[(size_t) m]; + const float md = hz_pts[(size_t) m + 1]; + const float hi = hz_pts[(size_t) m + 2]; + for (int k = 0; k < c.n_freq; k++) { + const float f = fft_freqs[(size_t) k]; + float up = (f - lo) / (md - lo); + float down = (hi - f) / (hi - md); + float w = std::fmin(up, down); + if (w < 0.0f) { + w = 0.0f; + } + c.mel_basis[(size_t) m * (size_t) c.n_freq + (size_t) k] = w; + } + // Slaney area normalization : 2 / (hi - lo). + const float enorm = 2.0f / (hi - lo); + for (int k = 0; k < c.n_freq; k++) { + c.mel_basis[(size_t) m * (size_t) c.n_freq + (size_t) k] *= enorm; + } + } +} + +// Build the GGML graph that turns a [T_in] f32 audio waveform into +// a [n_mels, T_frames] f32 log mel spectrogram. The signal is reflect +// padded by (n_fft - hop) / 2 on each side before framing, mirroring +// torch.nn.functional.pad(mode="reflect") used upstream. +// +// Inputs : +// audio [T_padded] f32, already reflect padded by the caller +// hann [n_fft] f32, host constant +// dft_real [n_fft, n_freq] f32, host constant +// dft_imag [n_fft, n_freq] f32, host constant +// mel_basis [n_freq, n_mels] f32, host constant +// +// Output : [n_mels, T_frames] f32 log mel. +// +// The im2col path produces frames [n_fft, T_frames] T-fastest, which is +// the layout ggml_mul_mat expects on the right operand (ne[0] = K = n_fft, +// ne[1] = M = T_frames). +static struct ggml_tensor * audio_mel_build_graph(struct ggml_context * ctx, + struct ggml_tensor * audio_padded, + struct ggml_tensor * hann, + struct ggml_tensor * dft_real, + struct ggml_tensor * dft_imag, + struct ggml_tensor * mel_basis, + const AudioMelConfig & cfg) { + const int n_fft = cfg.n_fft; + const int hop = cfg.hop; + + // Shape audio as [T_padded, 1, 1, 1] so im2col reads it as a 1D + // conv input with C_in = 1. ggml_im2col expects [IW, IH=1, IC, N=1] + // and returns [K_w, IC * K_h, OW, N]. With IH = 1 and K_h = 1, the + // output collapses to [n_fft, 1, T_frames, 1] which is what we need. + struct ggml_tensor * a4d = ggml_reshape_4d(ctx, audio_padded, audio_padded->ne[0], 1, 1, 1); + struct ggml_tensor * dummy = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_fft, 1, 1, 1); + ggml_set_name(dummy, "mel.im2col_dummy_kernel"); + + // im2col(s0=hop, s1=1, p0=0, p1=0, d0=1, d1=1, is_2D=false). + // Output dtype F32 to match audio dtype. Output shape with is_2D=false + // is [K * IC, OW, IC_outer, 1] which collapses to [n_fft, T_frames, 1, 1] + // since our IC = 1 and IC_outer = 1. + struct ggml_tensor * frames = ggml_im2col(ctx, dummy, a4d, hop, 1, 0, 0, 1, 1, false, GGML_TYPE_F32); + // frames has ne=(n_fft, T_frames, 1, 1). Reshape to [n_fft, T_frames]. + frames = ggml_reshape_2d(ctx, frames, n_fft, frames->ne[1]); + + // Multiply by hann window, broadcast over T_frames. ggml_mul + // broadcasts ne[0]-equal operands when one has ne[1]=1. hann is + // already [n_fft, 1]. + struct ggml_tensor * hann_2d = ggml_reshape_2d(ctx, hann, n_fft, 1); + frames = ggml_mul(ctx, frames, hann_2d); + + // STFT real and imag parts via two F32 matmuls. dft_* have layout + // [n_fft, n_freq] in ggml notation, so mul_mat returns [n_freq, T_frames]. + struct ggml_tensor * spec_re = ggml_mul_mat(ctx, dft_real, frames); + struct ggml_tensor * spec_im = ggml_mul_mat(ctx, dft_imag, frames); + ggml_mul_mat_set_prec(spec_re, GGML_PREC_F32); + ggml_mul_mat_set_prec(spec_im, GGML_PREC_F32); + + // Magnitude with the same eps as torch upstream (1e-9 added to the + // power, then sqrt). ggml_scale_bias does s*a + b so we add the eps + // without conjuring a backend dependent constant tensor. + struct ggml_tensor * mag2 = ggml_add(ctx, ggml_sqr(ctx, spec_re), ggml_sqr(ctx, spec_im)); + mag2 = ggml_scale_bias(ctx, mag2, 1.0f, 1e-9f); + struct ggml_tensor * mag = ggml_sqrt(ctx, mag2); + + // mel_basis [n_freq, n_mels] @ mag [n_freq, T_frames] -> [n_mels, T_frames]. + struct ggml_tensor * mel = ggml_mul_mat(ctx, mel_basis, mag); + ggml_mul_mat_set_prec(mel, GGML_PREC_F32); + + // log(max(mel, 1e-5)). ggml has clamp + log primitives. + mel = ggml_clamp(ctx, mel, 1e-5f, 1e30f); + mel = ggml_log(ctx, mel); + ggml_set_name(mel, "mel.log_mel"); + + return mel; +} diff --git a/src/audio-postproc.h b/src/audio-postproc.h new file mode 100644 index 0000000..a4c6bbc --- /dev/null +++ b/src/audio-postproc.h @@ -0,0 +1,387 @@ +#pragma once +// audio-postproc.h: TTS waveform post-processing +// +// Generic post-processing for neural TTS output : silence trimming +// (pydub-strict bit-for-bit), fade-in/out, padding. Public functions take +// and return float32 mono PCM in [-1, 1] at the pipeline sample rate. +// Internal silence detection runs on int16 samples to match pydub. +// Math reference : omnivoice/utils/audio.py (1:1 port). + +#include +#include +#include +#include +#include + +// RMS of an int16 slice [start, start + n) clamped to s16.size(). A slice that +// extends past the end shrinks accordingly. Empty slices return 0.0, matching +// pydub's AudioSegment.rms on empty segments. +static double postproc_slice_rms_s16(const std::vector & s16, size_t start, size_t n) { + size_t end = start + n; + if (end > s16.size()) { + end = s16.size(); + } + + if (start >= end) { + return 0.0; + } + + int64_t ssq = 0; + for (size_t i = start; i < end; i++) { + int32_t s = s16[i]; + ssq += (int64_t) s * (int64_t) s; + } + + size_t cnt = end - start; + return std::sqrt((double) ssq / (double) cnt); +} + +// Converts float32 [-1, 1] to int16 with the exact pydub recipe: +// (audio * 32768.0).clip(-32768, 32767).astype(int16). Truncation toward 0, +// matching numpy's astype(int16). +static std::vector postproc_f32_to_s16(const std::vector & a) { + std::vector out(a.size()); + for (size_t i = 0; i < a.size(); i++) { + double v = (double) a[i] * 32768.0; + if (v > 32767.0) { + v = 32767.0; + } + + if (v < -32768.0) { + v = -32768.0; + } + + out[i] = (int16_t) v; + } + + return out; +} + +// Inverse of postproc_f32_to_s16: int16 -> float32 via division by 32768.0. +static std::vector postproc_s16_to_f32(const std::vector & s16) { + std::vector out(s16.size()); + for (size_t i = 0; i < s16.size(); i++) { + out[i] = (float) ((double) s16[i] / 32768.0); + } + + return out; +} + +// pydub.silence.detect_silence ported to int16 samples. seek_step and +// min_silence_len are in samples. Returns inclusive ranges [start, end] in +// samples where end = start + min_silence_len of the last hit, exactly as +// pydub builds them. +static std::vector> postproc_detect_silence(const std::vector & s16, + int min_silence_len, + double thresh_lin, + int seek_step) { + std::vector> ranges; + int seg_len = (int) s16.size(); + + if (seg_len < min_silence_len) { + return ranges; + } + + int last_slice_start = seg_len - min_silence_len; + + std::vector starts; + for (int i = 0; i <= last_slice_start; i += seek_step) { + starts.push_back(i); + } + + if ((last_slice_start % seek_step) != 0) { + starts.push_back(last_slice_start); + } + + std::vector silence_starts; + for (int i : starts) { + double r = postproc_slice_rms_s16(s16, (size_t) i, (size_t) min_silence_len); + if (r <= thresh_lin) { + silence_starts.push_back(i); + } + } + + if (silence_starts.empty()) { + return ranges; + } + + int prev_i = silence_starts[0]; + int range_start = prev_i; + + for (size_t k = 1; k < silence_starts.size(); k++) { + int si = silence_starts[k]; + bool continuous = (si == prev_i + seek_step); + bool has_gap = (si > prev_i + min_silence_len); + + if (!continuous && has_gap) { + ranges.push_back({ range_start, prev_i + min_silence_len }); + range_start = si; + } + + prev_i = si; + } + + ranges.push_back({ range_start, prev_i + min_silence_len }); + return ranges; +} + +// pydub.silence.detect_nonsilent: invert detect_silence over [0, seg_len]. +static std::vector> postproc_detect_nonsilent(const std::vector & s16, + int min_silence_len, + double thresh_lin, + int seek_step) { + std::vector> nonsilent; + int seg_len = (int) s16.size(); + auto silent = postproc_detect_silence(s16, min_silence_len, thresh_lin, seek_step); + + if (silent.empty()) { + nonsilent.push_back({ 0, seg_len }); + return nonsilent; + } + + if (silent.front().first == 0 && silent.front().second == seg_len) { + return nonsilent; + } + + int prev_end = 0; + int last_end = 0; + + for (const auto & r : silent) { + nonsilent.push_back({ prev_end, r.first }); + prev_end = r.second; + last_end = r.second; + } + + if (last_end != seg_len) { + nonsilent.push_back({ prev_end, seg_len }); + } + + if (!nonsilent.empty() && nonsilent.front().first == 0 && nonsilent.front().second == 0) { + nonsilent.erase(nonsilent.begin()); + } + + return nonsilent; +} + +// pydub.silence.detect_leading_silence ported to int16. chunk_n is in samples. +// Returns the sample index where the leading silence ends (clamped to len). +static int postproc_detect_leading_silence(const std::vector & s16, double thresh_lin, int chunk_n) { + int trim = 0; + int seg_len = (int) s16.size(); + + while (trim < seg_len) { + int slice_end = std::min(trim + chunk_n, seg_len); + int n = slice_end - trim; + double r = postproc_slice_rms_s16(s16, (size_t) trim, (size_t) n); + + // pydub compares dBFS < threshold; in linear amplitude that is + // r < thresh_lin (strict), since dBFS is monotonic in r and r=0 + // gives -inf which is always below any finite threshold. + if (r >= thresh_lin) { + break; + } + + trim += chunk_n; + } + + if (trim > seg_len) { + trim = seg_len; + } + + return trim; +} + +// remove_silence: strict 1:1 port of omnivoice/utils/audio.py:remove_silence. +// Removes mid silences longer than mid_sil_ms (kept down to mid_sil_ms via +// pydub split_on_silence with keep_silence == mid_sil_ms), then trims the +// leading and trailing silences leaving lead_sil_ms / trail_sil_ms intact. +// thresh_db is the dBFS threshold (default -50 dBFS in upstream). +static void remove_silence(std::vector & a, + int sr, + int mid_sil_ms, + int lead_sil_ms, + int trail_sil_ms, + double thresh_db) { + if (a.empty()) { + return; + } + + std::vector s16 = postproc_f32_to_s16(a); + double thresh_lin = 32768.0 * std::pow(10.0, thresh_db / 20.0); + int seek_step = sr / 100; // 10 ms + + // Mid silence removal via split_on_silence + concat. + if (mid_sil_ms > 0) { + int min_sil_n = sr * mid_sil_ms / 1000; + int keep_n = min_sil_n; + + auto nonsilent = postproc_detect_nonsilent(s16, min_sil_n, thresh_lin, seek_step); + + std::vector> output_ranges; + output_ranges.reserve(nonsilent.size()); + for (const auto & r : nonsilent) { + output_ranges.push_back({ r.first - keep_n, r.second + keep_n }); + } + + // pydub pairwise overlap dedup: split overlap at the midpoint. + for (size_t i = 0; i + 1 < output_ranges.size(); i++) { + int last_end = output_ranges[i].second; + int next_start = output_ranges[i + 1].first; + if (next_start < last_end) { + int mid = (last_end + next_start) / 2; + output_ranges[i].second = mid; + output_ranges[i + 1].first = mid; + } + } + + // Concat clipped slices. Empty slices contribute nothing, matching + // AudioSegment.silent(0) += seg semantics. + std::vector out; + out.reserve(s16.size()); + int seg_len = (int) s16.size(); + + for (const auto & r : output_ranges) { + int cs = std::max(0, r.first); + int ce = std::min(seg_len, r.second); + if (cs < ce) { + out.insert(out.end(), s16.begin() + cs, s16.begin() + ce); + } + } + + s16 = std::move(out); + } + + // Edge trimming: leading then trailing via reverse trick. + int chunk_n = sr / 100; // 10 ms + + int trim_lead = postproc_detect_leading_silence(s16, thresh_lin, chunk_n); + trim_lead = std::max(0, trim_lead - sr * lead_sil_ms / 1000); + if (trim_lead > 0) { + s16.erase(s16.begin(), s16.begin() + std::min(trim_lead, (int) s16.size())); + } + + std::reverse(s16.begin(), s16.end()); + + int trim_trail = postproc_detect_leading_silence(s16, thresh_lin, chunk_n); + trim_trail = std::max(0, trim_trail - sr * trail_sil_ms / 1000); + if (trim_trail > 0) { + s16.erase(s16.begin(), s16.begin() + std::min(trim_trail, (int) s16.size())); + } + + std::reverse(s16.begin(), s16.end()); + + a = postproc_s16_to_f32(s16); +} + +// peak_normalize_half: rescale so peak amplitude becomes 0.5 (-6 dBFS). +// Mirrors the no-ref branch of _post_process_audio in omnivoice.py. +static void peak_normalize_half(std::vector & a) { + if (a.empty()) { + return; + } + + float peak = 0.0f; + for (float s : a) { + float v = std::fabs(s); + if (v > peak) { + peak = v; + } + } + + if (peak > 1e-6f) { + float k = 0.5f / peak; + for (float & s : a) { + s *= k; + } + } +} + +// fade_and_pad: linear fade-in / fade-out on the first and last fade_dur +// seconds, then pad pad_dur seconds of silence on each side. 1:1 port of +// fade_and_pad_audio in omnivoice/utils/audio.py. +static void fade_and_pad(std::vector & a, int sr, double fade_dur, double pad_dur) { + if (a.empty()) { + return; + } + + int fade_n = (int) (fade_dur * (double) sr); + int pad_n = (int) (pad_dur * (double) sr); + + if (fade_n > 0) { + int k = std::min(fade_n, (int) a.size() / 2); + if (k > 0) { + int denom = std::max(k - 1, 1); + + for (int i = 0; i < k; i++) { + float w = (float) i / (float) denom; + a[(size_t) i] *= w; + } + + for (int i = 0; i < k; i++) { + float w = 1.0f - (float) i / (float) denom; + a[a.size() - (size_t) k + (size_t) i] *= w; + } + } + } + + if (pad_n > 0) { + std::vector padded((size_t) pad_n + a.size() + (size_t) pad_n, 0.0f); + std::copy(a.begin(), a.end(), padded.begin() + pad_n); + a = std::move(padded); + } +} + +// cross_fade_chunks: concatenate audio chunks with a silence_dur gap split +// into fade_out, pure silence, fade_in. 1:1 port of cross_fade_chunks in +// omnivoice/utils/audio.py. +static std::vector cross_fade_chunks(const std::vector> & chunks, + int sr, + double silence_dur) { + if (chunks.empty()) { + return std::vector(); + } + + if (chunks.size() == 1) { + return chunks[0]; + } + + int total_n = (int) (silence_dur * (double) sr); + int fade_n = total_n / 3; + int silence_n = fade_n; + + std::vector merged = chunks[0]; + + for (size_t i = 1; i < chunks.size(); i++) { + const auto & chunk = chunks[i]; + + // Fade-out tail of merged. + int fout_n = std::min(fade_n, (int) merged.size()); + if (fout_n > 0) { + int denom = std::max(fout_n - 1, 1); + for (int j = 0; j < fout_n; j++) { + float w = 1.0f - (float) j / (float) denom; + merged[merged.size() - (size_t) fout_n + (size_t) j] *= w; + } + } + + // Silence gap. + if (silence_n > 0) { + merged.insert(merged.end(), (size_t) silence_n, 0.0f); + } + + // Fade-in head of next chunk (worked on a copy to keep input const). + std::vector head = chunk; + int fin_n = std::min(fade_n, (int) head.size()); + if (fin_n > 0) { + int denom = std::max(fin_n - 1, 1); + for (int j = 0; j < fin_n; j++) { + float w = (float) j / (float) denom; + head[(size_t) j] *= w; + } + } + + merged.insert(merged.end(), head.begin(), head.end()); + } + + return merged; +} diff --git a/src/audio-resample.h b/src/audio-resample.h new file mode 100644 index 0000000..770c4ca --- /dev/null +++ b/src/audio-resample.h @@ -0,0 +1,177 @@ +#pragma once +// audio-resample.h: torchaudio.functional.resample compatible reimplementation. +// Hann-windowed sinc interpolation with rolloff=0.99 and lowpass_filter_width=6, +// matching torchaudio defaults bit for bit. +// +// Reference: torchaudio/functional/functional.py, _get_sinc_resample_kernel +// and _apply_sinc_resample_kernel. +// +// Algorithm: +// gcd = gcd(sr_in, sr_out) +// orig = sr_in / gcd +// new = sr_out / gcd +// base = min(orig, new) * rolloff +// width = ceil(lpfw * orig / base) +// kernel_size = 2 * width + orig +// kernel[j, k] = sinc(t * pi) * hann(t)^2 * (base / orig) +// with t = clamp(((k - width) / orig - j / new) * base, +// -lpfw, lpfw) +// target_length = ceil(sr_out * n_in / sr_in) +// +// Apply: pad (width, width + orig), strided conv1d, transpose, truncate. + +#include +#include +#include +#include +#include + +#ifndef M_PI +# define M_PI 3.14159265358979323846 +#endif + +#define AUDIO_RESAMPLE_LPFW 6 +#define AUDIO_RESAMPLE_ROLLOFF 0.99 + +static int audio_resample_gcd(int a, int b) { + while (b != 0) { + int t = b; + b = a % b; + a = t; + } + return a; +} + +// Build the Hann-sinc polyphase kernel [new_freq_red, kernel_size] in row major. +// new_freq_red and orig_freq_red are sr_out/gcd and sr_in/gcd respectively. +static std::vector audio_resample_build_kernel(int orig, int newf, int * out_width, int * out_kernel_size) { + int base_int = (orig < newf) ? orig : newf; + double base = (double) base_int * AUDIO_RESAMPLE_ROLLOFF; + int width = (int) std::ceil((double) AUDIO_RESAMPLE_LPFW * (double) orig / base); + int K = 2 * width + orig; + + std::vector ker((size_t) newf * (size_t) K); + + double scale = base / (double) orig; + double inv_o = 1.0 / (double) orig; + double inv_n = 1.0 / (double) newf; + double pi = M_PI; + + for (int j = 0; j < newf; j++) { + double t_off = (double) (-j) * inv_n; + for (int k = 0; k < K; k++) { + double idx_k = (double) (k - width) * inv_o; + double t = (t_off + idx_k) * base; + if (t < -AUDIO_RESAMPLE_LPFW) { + t = -AUDIO_RESAMPLE_LPFW; + } + if (t > AUDIO_RESAMPLE_LPFW) { + t = AUDIO_RESAMPLE_LPFW; + } + + double w = std::cos(t * pi / (double) AUDIO_RESAMPLE_LPFW / 2.0); + w = w * w; + + double tp = t * pi; + double sinc = (tp == 0.0) ? 1.0 : std::sin(tp) / tp; + + ker[(size_t) j * (size_t) K + (size_t) k] = (float) (sinc * w * scale); + } + } + + *out_width = width; + *out_kernel_size = K; + return ker; +} + +// Resample one mono channel from sr_in to sr_out. Returns ceil(sr_out*n_in/sr_in) +// samples. Caller passes the kernel + width built once via audio_resample_build_kernel. +static void audio_resample_apply_mono(const float * in, + int n_in, + int orig, + int newf, + int width, + int kernel_size, + const float * kernel, + float * out, + long long target_length) { + int K = kernel_size; + int Np = n_in + 2 * width + orig; + std::vector padded((size_t) Np, 0.0f); + std::memcpy(padded.data() + width, in, (size_t) n_in * sizeof(float)); + + int n_per_chan = (Np - K) / orig + 1; + long long total = (long long) n_per_chan * (long long) newf; + long long out_len = (target_length < total) ? target_length : total; + + for (long long t_out = 0; t_out < out_len; t_out++) { + int chan = (int) (t_out % (long long) newf); + int pos = (int) (t_out / (long long) newf); + const float * w = kernel + (size_t) chan * (size_t) K; + const float * x = padded.data() + (size_t) pos * (size_t) orig; + float sum = 0.0f; + for (int k = 0; k < K; k++) { + sum += x[k] * w[k]; + } + out[(size_t) t_out] = sum; + } +} + +// Public API: resample a planar (or mono) f32 buffer from sr_in to sr_out. +// in: float buffer with channels stored planar [ch0: n_in][ch1: n_in][...]. +// n_in: per-channel input sample count. +// nch: number of channels. +// n_out: receives the per-channel output sample count. +// +// Returns a malloc'd planar buffer [ch0: *n_out][ch1: *n_out][...]. +// Caller must free() the result. NULL on error. +static float * audio_resample(const float * in, int n_in, int sr_in, int sr_out, int nch, int * n_out) { + if (!in || n_in <= 0 || sr_in <= 0 || sr_out <= 0 || nch <= 0) { + *n_out = 0; + return NULL; + } + + // Passthrough when source and target rates match. + if (sr_in == sr_out) { + size_t sz = (size_t) n_in * (size_t) nch * sizeof(float); + float * out = (float *) malloc(sz); + if (!out) { + fprintf(stderr, "[Audio-Resample] OOM passthrough buffer (%zu bytes)\n", sz); + *n_out = 0; + return NULL; + } + *n_out = n_in; + memcpy(out, in, sz); + return out; + } + + int g = audio_resample_gcd(sr_in, sr_out); + int orig = sr_in / g; + int newf = sr_out / g; + + int width = 0, kernel_size = 0; + std::vector kernel = audio_resample_build_kernel(orig, newf, &width, &kernel_size); + + long long target = (long long) std::ceil((double) sr_out * (double) n_in / (double) sr_in); + if (target <= 0) { + *n_out = 0; + return NULL; + } + + *n_out = (int) target; + + float * out = (float *) malloc((size_t) target * (size_t) nch * sizeof(float)); + if (!out) { + fprintf(stderr, "[Audio-Resample] OOM output buffer\n"); + *n_out = 0; + return NULL; + } + + for (int ch = 0; ch < nch; ch++) { + const float * src = in + (size_t) ch * (size_t) n_in; + float * dst = out + (size_t) ch * (size_t) target; + audio_resample_apply_mono(src, n_in, orig, newf, width, kernel_size, kernel.data(), dst, target); + } + + return out; +} diff --git a/src/backend.h b/src/backend.h new file mode 100644 index 0000000..5a687ac --- /dev/null +++ b/src/backend.h @@ -0,0 +1,158 @@ +#pragma once +// backend.h: shared GGML backend initialization +// +// All modules use the same pattern: load all backends, pick best GPU, +// keep CPU as fallback. Single shared backend across modules in the +// same binary, refcounted. + +#include "ggml-backend.h" +#include "qt-error.h" + +#include +#include +#include +#include + +struct BackendPair { + ggml_backend_t backend; + ggml_backend_t cpu_backend; + bool has_gpu; +}; + +// Cached backend state (shared across all modules in the same binary) +static BackendPair g_backend_cache = {}; +static int g_backend_refs = 0; + +// Physical core count heuristic (logical / 2 for HT/SMT). +// Used for GGML CPU thread count: GEMM shares SIMD units across hyperthreads, +// so one thread per physical core is optimal. +static int backend_cpu_n_threads(void) { + int n = (int) std::thread::hardware_concurrency() / 2; + return n > 0 ? n : 1; +} + +// Standalone CPU backend via Registry API (DL-safe, no ggml-cpu.h needed). +// Sets thread count via proc address since ggml_backend_cpu_device_init_backend +// ignores its params string and always defaults to GGML_DEFAULT_N_THREADS (4). +// Returns NULL on failure. +static ggml_backend_t cpu_backend_new(int n_threads) { + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + ggml_backend_t cpu = NULL; + if (cpu_dev) { + cpu = ggml_backend_dev_init(cpu_dev, NULL); + } + if (!cpu) { + cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, NULL); + } + if (!cpu) { + return NULL; + } + + ggml_backend_dev_t dev = ggml_backend_get_device(cpu); + ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : NULL; + if (reg) { + auto set_fn = + (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads"); + if (set_fn) { + set_fn(cpu, n_threads); + } + } + return cpu; +} + +// Initialize backends: load all available (CUDA, Metal, Vulkan...), +// pick the best one, keep CPU as fallback. +// label: log prefix, e.g. "DiT", "VAE", "LM" +// Subsequent calls reuse the same backend (single VMM pool). +static BackendPair backend_init(const char * label) { + if (g_backend_refs > 0) { + g_backend_refs++; + qt_log(QT_LOG_INFO, "[Load] %s backend: %s (shared)", label, ggml_backend_name(g_backend_cache.backend)); + return g_backend_cache; + } + + ggml_backend_load_all(); + BackendPair bp = {}; + + // GGML_BACKEND env var: force a specific device instead of auto-best. + // Device names: CUDA0, Vulkan0, CPU, BLAS (see ggml_backend_dev_name). + const char * force_backend = std::getenv("GGML_BACKEND"); + if (force_backend) { + bp.backend = ggml_backend_init_by_name(force_backend, nullptr); + if (!bp.backend) { + std::string avail; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + if (i > 0) { + avail += " "; + } + avail += ggml_backend_dev_name(ggml_backend_dev_get(i)); + } + qt_throw("GGML_BACKEND=%s not found. Available: %s", force_backend, avail.c_str()); + } + } else { + bp.backend = ggml_backend_init_best(); + } + if (!bp.backend) { + qt_throw("no backend available"); + } + bool best_is_cpu = (strcmp(ggml_backend_name(bp.backend), "CPU") == 0); + int n_threads = backend_cpu_n_threads(); + if (best_is_cpu) { + ggml_backend_free(bp.backend); + bp.backend = cpu_backend_new(n_threads); + bp.cpu_backend = bp.backend; + } else { + bp.cpu_backend = cpu_backend_new(n_threads); + } + if (!bp.cpu_backend) { + qt_throw("failed to init CPU backend"); + } + bp.has_gpu = !best_is_cpu; + qt_log(QT_LOG_INFO, "[Load] %s backend: %s (CPU threads: %d)", label, ggml_backend_name(bp.backend), n_threads); + + g_backend_cache = bp; + g_backend_refs = 1; + return bp; +} + +// Release a backend reference. Frees GPU + CPU backends when refcount hits 0. +static void backend_release(ggml_backend_t backend, ggml_backend_t cpu_backend) { + if (g_backend_refs <= 0) { + return; + } + g_backend_refs--; + if (g_backend_refs == 0) { + if (backend && backend != cpu_backend) { + ggml_backend_free(backend); + } + if (cpu_backend) { + ggml_backend_free(cpu_backend); + } + g_backend_cache = {}; + } +} + +// Create a scheduler from a backend pair. +// max_nodes: graph size hint (4096 for small models, 8192 for large) +// When a GPU is present, use its host buffer type for the CPU backend. +// Pinned memory lets the scheduler keep more ops on GPU instead of +// falling back to CPU with plain malloc. +static ggml_backend_sched_t backend_sched_new(BackendPair bp, int max_nodes) { + ggml_backend_t backends[2] = { bp.backend, bp.cpu_backend }; + ggml_backend_buffer_type_t bufts[2] = { NULL, NULL }; + int n = (bp.backend == bp.cpu_backend) ? 1 : 2; + + bufts[0] = ggml_backend_get_default_buffer_type(bp.backend); + if (n == 2) { + ggml_backend_dev_t gpu_dev = ggml_backend_get_device(bp.backend); + ggml_backend_buffer_type_t host_buft = gpu_dev ? ggml_backend_dev_host_buffer_type(gpu_dev) : NULL; + bufts[1] = host_buft ? host_buft : ggml_backend_get_default_buffer_type(bp.cpu_backend); + } + + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, n, max_nodes, false, true); + if (!sched) { + qt_log(QT_LOG_ERROR, "[Load] failed to create scheduler"); + return nullptr; + } + return sched; +} diff --git a/src/bpe.h b/src/bpe.h new file mode 100644 index 0000000..ab9148a --- /dev/null +++ b/src/bpe.h @@ -0,0 +1,633 @@ +#pragma once +// bpe.h, Qwen3/GPT-2 byte-level BPE tokenizer (CPU-only, no dependencies) +// +// Parses the tokenizer.json fields stored in a GGUF model and produces a +// byte-level BPE encoder/decoder. Arch-specific special tokens (text +// markers, language tags, audio sentinels) are loaded through +// bpe_load_specials_from_keys with a caller-provided list of GGUF KV keys. +// Loads vocab + merges from a GGUF tokenizer payload. Handles byte-level +// encoding, GPT-2 regex pre-tokenizer, BPE merges, and a registry of +// verbatim special tokens (endoftext plus any caller-registered tokens +// such as TTS style markers and language tags). + +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include + +// GPT-2 byte-level encoding table +// Maps byte [0..255] -> Unicode char for BPE vocab keys. +// Printable ASCII stays as-is, control/space bytes get remapped. +static void build_byte_encoder(std::string byte2str[256]) { + // Standard GPT-2 byte encoder + int bs[256], cs[256], n = 0, total = 0; + // Printable ranges that map to themselves + for (int b = '!'; b <= '~'; b++) { + bs[total] = b; + cs[total] = b; + total++; + } + for (int b = 0xA1; b <= 0xAC; b++) { + bs[total] = b; + cs[total] = b; + total++; + } + for (int b = 0xAE; b <= 0xFF; b++) { + bs[total] = b; + cs[total] = b; + total++; + } + // Remaining bytes get mapped to 256+ + bool used[256] = {}; + for (int i = 0; i < total; i++) { + used[bs[i]] = true; + } + for (int b = 0; b < 256; b++) { + if (!used[b]) { + bs[total] = b; + cs[total] = 256 + n; + n++; + total++; + } + } + assert(total == 256); + // Convert codepoints to UTF-8 strings + for (int i = 0; i < 256; i++) { + int cp = cs[i]; + char buf[4]; + int len; + if (cp < 0x80) { + buf[0] = (char) cp; + len = 1; + } else if (cp < 0x800) { + buf[0] = (char) (0xC0 | (cp >> 6)); + buf[1] = (char) (0x80 | (cp & 0x3F)); + len = 2; + } else { + buf[0] = (char) (0xE0 | (cp >> 12)); + buf[1] = (char) (0x80 | ((cp >> 6) & 0x3F)); + buf[2] = (char) (0x80 | (cp & 0x3F)); + len = 3; + } + byte2str[bs[i]] = std::string(buf, len); + } +} + +// UTF-8 helpers +static int utf8_codepoint(const char * s, int * advance) { + unsigned char c = s[0]; + if (c < 0x80) { + *advance = 1; + return c; + } + if ((c & 0xE0) == 0xC0) { + *advance = 2; + return ((c & 0x1F) << 6) | (s[1] & 0x3F); + } + if ((c & 0xF0) == 0xE0) { + *advance = 3; + return ((c & 0x0F) << 12) | ((s[1] & 0x3F) << 6) | (s[2] & 0x3F); + } + if ((c & 0xF8) == 0xF0) { + *advance = 4; + return ((c & 0x07) << 18) | ((s[1] & 0x3F) << 12) | ((s[2] & 0x3F) << 6) | (s[3] & 0x3F); + } + // invalid lead byte : advance one and return the raw byte to avoid an + // infinite loop. The Python tokenizer never reaches this path since Python + // str guarantees valid UTF 8 ; in C++ the std::string input has no such + // guarantee, so this branch handles malformed input defensively. + *advance = 1; + return c; +} + +// Unicode category checks (simplified but covers Latin + common scripts) +static bool is_letter(int cp) { + if ((cp >= 'A' && cp <= 'Z') || (cp >= 'a' && cp <= 'z')) { + return true; + } + if (cp < 0x80) { + return false; + } + // Latin Extended: U+00C0-U+00D6, U+00D8-U+00F6, U+00F8-U+01BF + Latin Extended-A/B + if (cp >= 0xC0 && cp <= 0x024F && cp != 0xD7 && cp != 0xF7) { + return true; + } + // Common CJK, Cyrillic, Greek, Arabic, etc., treat as letters + if (cp >= 0x0370 && cp <= 0x1FFF) { + return true; // Greek, Cyrillic, Armenian, etc. + } + if (cp >= 0x2C00 && cp <= 0x2DFF) { + return true; // Georgian, etc. + } + if (cp >= 0x3040 && cp <= 0x9FFF) { + return true; // CJK + } + if (cp >= 0xAC00 && cp <= 0xD7AF) { + return true; // Korean + } + if (cp >= 0xF900 && cp <= 0xFAFF) { + return true; // CJK compatibility + } + if (cp >= 0x10000) { + return true; // SMP, mostly letters/symbols + } + return false; +} + +static bool is_digit(int cp) { + return cp >= '0' && cp <= '9'; +} + +static bool is_whitespace(int cp) { + return cp == ' ' || cp == '\t' || cp == '\n' || cp == '\r' || cp == 0x0B || cp == 0x0C || cp == 0xA0 || + cp == 0x2000 || cp == 0x2001 || cp == 0x2002 || cp == 0x200B; +} + +static bool is_newline(int cp) { + return cp == '\n' || cp == '\r'; +} + +// GPT-2 pre-tokenizer regex (manual implementation) +// Pattern: (?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| +// \s?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+ +// Splits text into non-overlapping chunks (on original text, not byte-encoded). +static std::vector gpt2_pre_tokenize(const std::string & text) { + std::vector chunks; + const char * s = text.c_str(); + int len = (int) text.size(); + int i = 0; + + while (i < len) { + int adv; + int cp = utf8_codepoint(s + i, &adv); + + // Rule 1: Contractions 's 't 're 've 'm 'll 'd + if ((cp == '\'' || cp == 0x2019) && i + adv < len) { + const char * rest = s + i + adv; + int rlen = len - i - adv; + auto try_match = [&](const char * suffix, int slen) -> bool { + if (rlen >= slen) { + // case-insensitive compare + for (int k = 0; k < slen; k++) { + char c1 = rest[k], c2 = suffix[k]; + if (c1 >= 'A' && c1 <= 'Z') { + c1 = (char) (c1 + 32); + } + if (c1 != c2) { + return false; + } + } + // next char should NOT be a letter + if (rlen > slen) { + int a2; + int cp2 = utf8_codepoint(rest + slen, &a2); + if (is_letter(cp2)) { + return false; + } + } + chunks.push_back(std::string(s + i, adv + slen)); + i += adv + slen; + return true; + } + return false; + }; + if (try_match("ll", 2)) { + continue; + } + if (try_match("re", 2)) { + continue; + } + if (try_match("ve", 2)) { + continue; + } + if (try_match("s", 1)) { + continue; + } + if (try_match("t", 1)) { + continue; + } + if (try_match("m", 1)) { + continue; + } + if (try_match("d", 1)) { + continue; + } + } + + // Rule 2: [^\r\n\p{L}\p{N}]?\p{L}+ + if (is_letter(cp)) { + int start = i; + i += adv; + while (i < len) { + int a2; + int cp2 = utf8_codepoint(s + i, &a2); + if (!is_letter(cp2)) { + break; + } + i += a2; + } + chunks.push_back(std::string(s + start, i - start)); + continue; + } + if (!is_newline(cp) && !is_letter(cp) && !is_digit(cp) && !is_whitespace(cp)) { + // Non-letter/number/space, check if followed by letters + int start = i; + int after = i + adv; + if (after < len) { + int a2; + int cp2 = utf8_codepoint(s + after, &a2); + if (is_letter(cp2)) { + i = after + a2; + while (i < len) { + int a3; + int cp3 = utf8_codepoint(s + i, &a3); + if (!is_letter(cp3)) { + break; + } + i += a3; + } + chunks.push_back(std::string(s + start, i - start)); + continue; + } + } + } + + // Rule 3: \p{N}+ (digits, consume consecutively) + if (is_digit(cp)) { + int start = i; + while (i < len && is_digit((unsigned char) s[i])) { + i++; + } + // GPT-2 regex matches single \p{N}, let's match one at a time + // to be safe, but in practice consecutive digits usually merge anyway. + // The regex is \p{N} (single digit), so split each digit: + for (int j = start; j < i; j++) { + chunks.push_back(std::string(s + j, 1)); + } + continue; + } + + // Rule 5: \s*[\r\n]+ (newlines with optional leading whitespace) + if (is_newline(cp)) { + int start = i; + while (i < len && is_newline((unsigned char) s[i])) { + i++; + } + chunks.push_back(std::string(s + start, i - start)); + continue; + } + + // Rule 6: whitespace handling + // Regex order: \s+(?!\S) first (trailing whitespace), then \s+ as fallback + // \s+(?!\S) backtracks: consumes whitespace NOT followed by non-whitespace + // This peels off leading spaces, leaving the last space to combine with the next word + if (is_whitespace(cp)) { + int start = i; + // Find end of whitespace run + int ws_end = i + adv; + while (ws_end < len && is_whitespace((unsigned char) s[ws_end]) && !is_newline((unsigned char) s[ws_end])) { + ws_end++; + } + // Check what follows the whitespace run + bool followed_by_non_ws = + (ws_end < len && !is_whitespace((unsigned char) s[ws_end]) && !is_newline((unsigned char) s[ws_end])); + if (followed_by_non_ws && ws_end - start > 1) { + // \s+(?!\S) matches all but the last space + // Leave one space for the next iteration to combine with word + int trailing = ws_end - 1; + chunks.push_back(std::string(s + start, trailing - start)); + i = trailing; + continue; + } + // Single space followed by word: combine space + word as one chunk + i = start + adv; + if (i < len) { + int a2; + int cp2 = utf8_codepoint(s + i, &a2); + if (is_letter(cp2)) { + i += a2; + while (i < len) { + int a3; + int cp3 = utf8_codepoint(s + i, &a3); + if (!is_letter(cp3)) { + break; + } + i += a3; + } + chunks.push_back(std::string(s + start, i - start)); + continue; + } + if (is_digit(cp2)) { + chunks.push_back(std::string(s + start, i - start)); + continue; + } + if (!is_whitespace(cp2) && !is_newline(cp2)) { + int pstart = start; + while (i < len) { + int a3; + int cp3 = utf8_codepoint(s + i, &a3); + if (is_whitespace(cp3) || is_letter(cp3) || is_digit(cp3)) { + break; + } + i += a3; + } + while (i < len && is_newline((unsigned char) s[i])) { + i++; + } + chunks.push_back(std::string(s + pstart, i - pstart)); + continue; + } + } + // Trailing whitespace (end of string or before newline), consume all + i = ws_end; + while (i < len) { + int a2; + int cp2 = utf8_codepoint(s + i, &a2); + if (!is_whitespace(cp2)) { + break; + } + i += a2; + } + chunks.push_back(std::string(s + start, i - start)); + continue; + } + + // Rule 4: [^\s\p{L}\p{N}]+[\r\n]* (punctuation/symbols) + { + int start = i; + i += adv; + while (i < len) { + int a2; + int cp2 = utf8_codepoint(s + i, &a2); + if (is_whitespace(cp2) || is_letter(cp2) || is_digit(cp2) || is_newline(cp2)) { + break; + } + i += a2; + } + // trailing newlines + while (i < len && is_newline((unsigned char) s[i])) { + i++; + } + chunks.push_back(std::string(s + start, i - start)); + } + } + return chunks; +} + +// BPE tokenizer struct +struct BPETokenizer { + std::unordered_map vocab; // token_str -> id + std::unordered_map merges; // "a b" -> rank + std::string byte2str[256]; // byte -> GPT-2 UTF-8 string + int eos_id; // <|endoftext|> + int n_vocab; + std::vector id_to_str; // id -> token_str (reverse vocab) + + // Registered special tokens. Each (str, id) pair is matched verbatim in + // bpe_encode and emitted as a single id, bypassing the BPE merge passes. + std::vector> specials; +}; + +// Register a special token. Strings already registered are skipped (no dup). +static void bpe_add_special(BPETokenizer * tok, const std::string & str, int id) { + for (const auto & sp : tok->specials) { + if (sp.first == str) { + return; + } + } + tok->specials.emplace_back(str, id); +} + +// Load tokenizer from GGUF KV (tokenizer.ggml.tokens + tokenizer.ggml.merges) +static bool load_bpe_from_gguf(BPETokenizer * tok, const char * gguf_path) { + build_byte_encoder(tok->byte2str); + + struct gguf_init_params gp = { true, NULL }; + struct gguf_context * ctx = gguf_init_from_file(gguf_path, gp); + if (!ctx) { + fprintf(stderr, "[BPE] Failed to open %s\n", gguf_path); + return false; + } + + int64_t tok_key = gguf_find_key(ctx, "tokenizer.ggml.tokens"); + int64_t mrg_key = gguf_find_key(ctx, "tokenizer.ggml.merges"); + if (tok_key < 0 || mrg_key < 0) { + fprintf(stderr, "[BPE] Tokenizer not found in %s\n", gguf_path); + gguf_free(ctx); + return false; + } + + int n_tokens = (int) gguf_get_arr_n(ctx, tok_key); + int n_merges = (int) gguf_get_arr_n(ctx, mrg_key); + + for (int i = 0; i < n_tokens; i++) { + const char * s = gguf_get_arr_str(ctx, tok_key, (size_t) i); + tok->vocab[std::string(s)] = i; + } + + for (int i = 0; i < n_merges; i++) { + const char * s = gguf_get_arr_str(ctx, mrg_key, (size_t) i); + tok->merges[std::string(s)] = i; + } + + gguf_free(ctx); + + tok->n_vocab = (int) tok->vocab.size(); + + tok->id_to_str.resize(tok->n_vocab); + for (auto & kv : tok->vocab) { + if (kv.second >= 0 && kv.second < tok->n_vocab) { + tok->id_to_str[kv.second] = kv.first; + } + } + + // Resolve eos_id from the vocab itself rather than hard-coding 151643. + // Falls back to -1 if the standard sentinel is absent. + auto eos_it = tok->vocab.find("<|endoftext|>"); + tok->eos_id = (eos_it != tok->vocab.end()) ? eos_it->second : -1; + if (tok->eos_id >= 0) { + bpe_add_special(tok, "<|endoftext|>", tok->eos_id); + } + + fprintf(stderr, "[BPE] Loaded from GGUF: %d vocab, %d merges, eos_id=%d\n", tok->n_vocab, n_merges, tok->eos_id); + return true; +} + +// Read arch-specific special tokens from a caller-provided list of GGUF KV +// keys. Each key holds a u32 vocab id, mapped back to its vocab string and +// registered through bpe_add_special. The endoftext sentinel is already +// registered by load_bpe_from_gguf, so callers should not list it here. +static bool bpe_load_specials_from_keys(BPETokenizer * tok, + const char * gguf_path, + const char * const * keys, + int n_keys) { + struct gguf_init_params gp = { true, NULL }; + struct gguf_context * ctx = gguf_init_from_file(gguf_path, gp); + if (!ctx) { + fprintf(stderr, "[BPE] Failed to open %s for specials\n", gguf_path); + return false; + } + + int n_added = 0; + for (int i = 0; i < n_keys; i++) { + int64_t k = gguf_find_key(ctx, keys[i]); + if (k < 0) { + fprintf(stderr, "[BPE] WARNING: missing %s in GGUF\n", keys[i]); + continue; + } + int id = (int) gguf_get_val_u32(ctx, k); + if (id < 0 || id >= tok->n_vocab) { + fprintf(stderr, "[BPE] WARNING: %s id=%d out of vocab range\n", keys[i], id); + continue; + } + const std::string & s = tok->id_to_str[id]; + if (s.empty()) { + fprintf(stderr, "[BPE] WARNING: %s id=%d has empty vocab string\n", keys[i], id); + continue; + } + bpe_add_special(tok, s, id); + n_added++; + } + + gguf_free(ctx); + fprintf(stderr, "[BPE] Registered %d arch special tokens (total specials=%zu)\n", n_added, tok->specials.size()); + return true; +} + +// Byte-level encode: raw text bytes -> GPT-2 BPE string +static std::string byte_level_encode(const BPETokenizer * tok, const std::string & text) { + std::string out; + for (unsigned char c : text) { + out += tok->byte2str[c]; + } + return out; +} + +// BPE merge algorithm +// Input: list of symbols (strings). Merges pairs by priority. +static std::vector bpe_merge(const std::unordered_map & merge_rank, + const std::vector & symbols) { + if (symbols.size() <= 1) { + return symbols; + } + + std::vector work = symbols; + + while (work.size() > 1) { + // Find the pair with lowest rank (highest priority) + int best_rank = INT_MAX; + int best_pos = -1; + for (int i = 0; i < (int) work.size() - 1; i++) { + std::string key = work[i] + " " + work[i + 1]; + auto it = merge_rank.find(key); + if (it != merge_rank.end() && it->second < best_rank) { + best_rank = it->second; + best_pos = i; + } + } + if (best_pos < 0) { + break; // no more merges + } + + // Merge the pair + std::string merged = work[best_pos] + work[best_pos + 1]; + work[best_pos] = merged; + work.erase(work.begin() + best_pos + 1); + } + return work; +} + +// Encode a single pre-tokenized chunk -> token ids +static void encode_chunk(const BPETokenizer * tok, const std::string & chunk, std::vector & ids) { + // Byte-level encode + std::string encoded = byte_level_encode(tok, chunk); + + // Split into individual UTF-8 characters (each is a BPE symbol) + std::vector symbols; + const char * s = encoded.c_str(); + int len = (int) encoded.size(); + int i = 0; + while (i < len) { + int adv; + utf8_codepoint(s + i, &adv); + symbols.push_back(std::string(s + i, adv)); + i += adv; + } + + // Apply BPE merges + std::vector merged = bpe_merge(tok->merges, symbols); + + // Look up in vocab + for (const auto & piece : merged) { + auto it = tok->vocab.find(piece); + if (it != tok->vocab.end()) { + ids.push_back(it->second); + } else { + // Fallback: encode each byte individually (should not happen with byte-level BPE) + fprintf(stderr, "[BPE] WARNING: unknown token '%s'\n", piece.c_str()); + for (unsigned char c : piece) { + auto it2 = tok->vocab.find(std::string(1, c)); + if (it2 != tok->vocab.end()) { + ids.push_back(it2->second); + } + } + } + } +} + +// Full encode : text -> token ids. +// Walks the text from left to right, matching any registered special token +// verbatim. For each segment between specials, runs the GPT-2 byte-level +// pre-tokenizer + BPE merges. The endoftext sentinel is auto-registered as +// a special by load_bpe_from_gguf, so existing call sites that embed +// "<|endoftext|>" in the input text keep working. +// add_eos = true appends the eos_id at the end (post-processor behavior). +static std::vector bpe_encode(const BPETokenizer * tok, const std::string & text, bool add_eos = true) { + std::vector ids; + + auto encode_segment = [&](const std::string & seg) { + if (seg.empty()) { + return; + } + auto chunks = gpt2_pre_tokenize(seg); + for (const auto & chunk : chunks) { + encode_chunk(tok, chunk, ids); + } + }; + + size_t pos = 0; + while (pos < text.size()) { + // Find the leftmost occurrence of any registered special token. + size_t best_pos = std::string::npos; + int best_idx = -1; + for (size_t i = 0; i < tok->specials.size(); i++) { + size_t p = text.find(tok->specials[i].first, pos); + if (p != std::string::npos && p < best_pos) { + best_pos = p; + best_idx = (int) i; + } + } + + if (best_idx < 0) { + encode_segment(text.substr(pos)); + break; + } + + if (best_pos > pos) { + encode_segment(text.substr(pos, best_pos - pos)); + } + const auto & sp = tok->specials[(size_t) best_idx]; + ids.push_back(sp.second); + pos = best_pos + sp.first.size(); + } + + if (add_eos && tok->eos_id >= 0) { + ids.push_back(tok->eos_id); + } + return ids; +} diff --git a/src/causal-trans-conv.h b/src/causal-trans-conv.h new file mode 100644 index 0000000..53c6afd --- /dev/null +++ b/src/causal-trans-conv.h @@ -0,0 +1,171 @@ +#pragma once +// causal-trans-conv.h: Causal ConvTranspose1d primitive for the +// Qwen3-TTS 12Hz tokenizer decoder. +// +// PyTorch reference (Qwen3TTSTokenizerV2CausalTransConvNet): +// y = ConvTranspose1d(x, k, stride) # raw length (T-1)*stride + K +// y = y[..., : y.shape[-1] - (K - stride)] # right-trim K-stride frames +// final length: T * stride +// +// GGML implementation: the weight is pre-permuted at load time from the +// PyTorch (IC, OC, K) layout to a [IC, K*OC] layout with k varying +// faster than oc inside K*OC. The forward graph multiplies this weight +// against a channels-first input via ggml_mul_mat to produce a column +// matrix [K*OC, T_in], scatters it into [T_raw, OC] via ggml_col2im_1d +// with padding=0, right-trims to [T_in*stride, OC], transposes to +// channels-first [OC, T_in*stride], and adds the bias. + +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include + +// Load a ConvTranspose1d weight stored on disk in PyTorch layout +// (IC, OC, K) and pre-permute it to ggml [IC, K*OC] with k fastest in +// K*OC. Source dtype must be F32. +// +// src flat[ic*OC*K + oc*K + k] = w[ic][oc][k] PyTorch row-major +// dst flat[(oc*K + k)*IC + ic] = w[ic][oc][k] ggml row-major, ne=(IC, K*OC) +static struct ggml_tensor * qwen_load_ctw_f32(WeightCtx * wctx, const GGUFModel & gf, const std::string & name) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str()); + if (!src) { + fprintf(stderr, "[CausalTransConv] FATAL: tensor '%s' not found\n", name.c_str()); + exit(1); + } + // Source dtype is F32 in the F32 master, F16 in the quantized variants + // since 3D conv weights cannot be Q8_0 / Q4_K_M and ggml falls back to + // F16 in the quantizer. Both paths cast to F32 here ; the K*OC*IC + // permutation always lands in a freshly allocated F32 buffer anyway. + if (src->type != GGML_TYPE_F32 && src->type != GGML_TYPE_F16) { + fprintf(stderr, "[CausalTransConv] FATAL: '%s' expected F32 or F16, got type %d\n", name.c_str(), + (int) src->type); + exit(1); + } + int K = (int) src->ne[0]; + int OC = (int) src->ne[1]; + int IC = (int) src->ne[2]; + + struct ggml_tensor * dst = ggml_new_tensor_2d(wctx->ctx, GGML_TYPE_F32, IC, K * OC); + ggml_set_name(dst, name.c_str()); + + const void * raw = gf_get_data(gf, name.c_str()); + auto buf = std::make_unique((size_t) IC * (size_t) K * (size_t) OC); + float * dstp = buf.get(); + + auto load_src = [&](size_t idx) -> float { + if (src->type == GGML_TYPE_F32) { + return ((const float *) raw)[idx]; + } + return ggml_fp16_to_fp32(((const ggml_fp16_t *) raw)[idx]); + }; + + for (int ic = 0; ic < IC; ic++) { + for (int oc = 0; oc < OC; oc++) { + for (int k = 0; k < K; k++) { + dstp[(size_t) (oc * K + k) * IC + ic] = load_src((size_t) ic * OC * K + oc * K + k); + } + } + } + + wctx->pending.push_back({ dst, dstp, (size_t) IC * (size_t) K * (size_t) OC * sizeof(float), 0 }); + wctx->staging.push_back(std::move(buf)); + return dst; +} + +// Causal ConvTranspose1d forward graph. +// w_perm: [IC, K*OC] f32, pre-permuted by qwen_load_ctw_f32 +// b: [OC] f32 or NULL +// x: [T_in, IC] f32, T-first +// stride: upsample factor +// kernel: kernel size +// oc: output channels (must match the K*OC factorization of w_perm) +// Returns [T_in*stride, OC] f32, T-first. +static struct ggml_tensor * qwen_causal_trans_conv1d(struct ggml_context * ctx, + struct ggml_tensor * w_perm, + struct ggml_tensor * b, + struct ggml_tensor * x, + int stride, + int kernel, + int oc) { + int trim = kernel - stride; + + // Transpose x to channels-first [IC, T_in] for the mul_mat contraction + struct ggml_tensor * xt = ggml_cont(ctx, ggml_transpose(ctx, x)); + + // mul_mat contracts over IC: col [K*OC, T_in] + struct ggml_tensor * col = ggml_mul_mat(ctx, w_perm, xt); + + // col2im_1d with padding=0: [T_raw, OC] T-first, T_raw = (T_in-1)*stride + K + struct ggml_tensor * y = ggml_col2im_1d(ctx, col, stride, oc, 0); + + // Right-trim K-stride frames -> [T_in*stride, OC] T-first + if (trim > 0) { + int64_t T_keep = y->ne[0] - trim; + y = ggml_view_2d(ctx, y, T_keep, y->ne[1], y->nb[1], 0); + } + + if (b) { + // bias [OC] broadcasts as (1, OC) onto (T, OC) via ne[0]=1 + struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, b->ne[0]); + y = ggml_add(ctx, y, b2d); + } + return y; +} + +// Causal Conv1d with optional stride. Left pad with (kernel_eff - stride), +// add an extra right pad to align with stride boundaries, then run a +// standard ggml_conv_1d. Matches MimiConv1d.causal forward exactly: +// kernel_eff = (k - 1) * d + 1 +// padding_total = kernel_eff - stride +// extra_pad = ceil((T + padding_total - kernel_eff) / stride) * stride +// + kernel_eff - padding_total - T +// = (T - 1) % stride for the common case +// The output length is (T + padding_total + extra_pad - kernel_eff) / stride + 1 +// = ceil(T / stride). Stride defaults to 1 to preserve the Qwen3 causal +// path used by pre_conv and the DAC decoder. +// w: [k, IC, OC] f32, source layout (K, IC, OC) maps to ggml ne directly +// b: [OC] f32 or NULL +// x: [T, IC] f32 T-first +// Returns [ceil(T / stride), OC] f32 T-first. +static struct ggml_tensor * qwen_causal_conv1d(struct ggml_context * ctx, + struct ggml_tensor * w, + struct ggml_tensor * b, + struct ggml_tensor * x, + int k, + int d, + int s = 1) { + int OC = (int) w->ne[2]; + int kernel_eff = (k - 1) * d + 1; + int padding_tot = kernel_eff - s; + + // Mimi extra padding: ensures the causal conv lands on a stride boundary + // by adding zeros on the right before the convolution. + int T = (int) x->ne[0]; + int n_frames = (T + padding_tot - kernel_eff + s - 1) / s + 1; + int ideal_len = (n_frames - 1) * s + kernel_eff - padding_tot; + int extra_pad = ideal_len - T; + if (extra_pad < 0) { + extra_pad = 0; + } + + struct ggml_tensor * y = x; + if (padding_tot > 0 || extra_pad > 0) { + y = ggml_pad_ext(ctx, y, padding_tot, extra_pad, 0, 0, 0, 0, 0, 0); + } + + // ggml_conv_1d expects 3D input [T, IC, N], add the batch dim + y = ggml_reshape_3d(ctx, y, y->ne[0], y->ne[1], 1); + y = ggml_conv_1d(ctx, w, y, s, 0, d); + // squeeze batch back to 2D + y = ggml_reshape_2d(ctx, y, y->ne[0], y->ne[1]); + + if (b) { + struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, OC); + y = ggml_add(ctx, y, b2d); + } + return y; +} diff --git a/src/code-predictor-forward.cpp b/src/code-predictor-forward.cpp new file mode 100644 index 0000000..942d271 --- /dev/null +++ b/src/code-predictor-forward.cpp @@ -0,0 +1,301 @@ +// code-predictor-forward.cpp : eager full-recompute graph for the +// Qwen3-TTS code predictor (5-layer Qwen3 stack with plain 1D NEOX +// RoPE, GQA attention with QK-norm, SwiGLU MLP, head-per-codebook +// output projection). +// +// The predictor architecture mirrors the Talker block, the only +// differences are : +// - 5 layers instead of 28 +// - plain 1D RoPE (no multimodal sections) +// - one private embedding table and one private linear head per +// acoustic codebook (1..15) +// +// The single-frame loop here recomputes the full graph at every step g +// (0..14) over a sequence of length g+2. With 5 layers and at most 16 +// tokens per recompute this is sub-millisecond on modern GPUs. + +#include "code-predictor-forward.h" + +#include "debug.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include +#include +#include + +// One Qwen3 decoder block, identical structural pattern to the Talker +// layer, just with plain 1D RoPE on the position axis. Returns the +// layer output [hidden, T]. +static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context * ctx, + const CodePredictorWeights * cw, + const TalkerLayer & layer, + struct ggml_tensor * x, + struct ggml_tensor * positions, + struct ggml_tensor * mask, + int T) { + const int n_q_heads = cw->num_attention_heads; + const int n_kv = cw->num_key_value_heads; + const int hd = cw->head_dim; + const float eps = cw->rms_norm_eps; + + struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps); + h = ggml_mul(ctx, h, layer.input_norm_w); + + struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h); + struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h); + struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h); + + q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T); + k = ggml_reshape_3d(ctx, k, hd, n_kv, T); + v = ggml_reshape_3d(ctx, v, hd, n_kv, T); + + q = ggml_rms_norm(ctx, q, eps); + q = ggml_mul(ctx, q, layer.attn.q_norm_w); + k = ggml_rms_norm(ctx, k, eps); + k = ggml_mul(ctx, k, layer.attn.k_norm_w); + + q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, cw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + + struct ggml_tensor * q_p = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + struct ggml_tensor * k_p = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + struct ggml_tensor * v_p = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); + + struct ggml_tensor * scores = ggml_mul_mat(ctx, k_p, q_p); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + + float scale = 1.0f / sqrtf((float) hd); + scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); + + struct ggml_tensor * attn = ggml_mul_mat(ctx, v_p, scores); + attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3)); + attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T); + + struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn); + x = ggml_add(ctx, x, o); + + struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps); + h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w); + + struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2); + struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2); + gate = ggml_silu(ctx, gate); + struct ggml_tensor * gu = ggml_mul(ctx, gate, up); + struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu); + + x = ggml_add(ctx, x, mlp); + return x; +} + +// Run one recompute pass for predictor step g, with sub_input of length +// T = g + 2 vectors flattened row-major as [T, hidden]. Fills logits_out +// with the output of lm_head[g_head] applied to the final-norm of the +// last position. Returns false on failure. +static bool code_predictor_recompute(const CodePredictorWeights * cw, + ggml_backend_sched_t sched, + const float * sub_input, + int T, + int talker_hidden, + int g_head, + std::vector * logits_out) { + const int vocab = cw->vocab_size; + const int n_layers = cw->num_hidden_layers; + + const int max_nodes = 32 * n_layers + 64; + const size_t arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false); + + struct ggml_init_params gp = { arena_bytes, NULL, true }; + struct ggml_context * gctx = ggml_init(gp); + if (!gctx) { + fprintf(stderr, "[CodePredictor] FATAL: ggml_init failed\n"); + return false; + } + + // Input lives in talker_hidden because every entry of sub_input is a + // talker-side hidden : either the talker last hidden state, or one of + // the codec_embedding rows that are also talker-sized in the upstream + // checkpoint. The mtp_proj node below brings them down to predictor + // hidden when the variant has a non-Identity projection. + struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, T); + struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T); + struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, T, T); + ggml_set_name(x_in, "sub_input"); + ggml_set_name(pos_in, "positions"); + ggml_set_name(mask_in, "causal_mask"); + + struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false); + + // small_to_mtp projection : Linear(talker_hidden -> hidden) with bias. + // When mtp_proj is absent (Identity case) the input is already at the + // predictor hidden dimension and we feed x_in straight to the layers. + struct ggml_tensor * h = x_in; + if (cw->mtp_proj_w) { + h = ggml_mul_mat(gctx, cw->mtp_proj_w, h); + if (cw->mtp_proj_b) { + h = ggml_add(gctx, h, cw->mtp_proj_b); + } + ggml_set_name(h, "mtp_proj_out"); + } + + for (int l = 0; l < n_layers; l++) { + h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, T); + } + + struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps); + h_final = ggml_mul(gctx, h_final, cw->norm_w); + + struct ggml_tensor * logits = ggml_mul_mat(gctx, cw->lm_head[(size_t) g_head], h_final); + ggml_set_name(logits, "logits"); + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + + if (!ggml_backend_sched_alloc_graph(sched, gf)) { + fprintf(stderr, "[CodePredictor] FATAL: graph allocation failed\n"); + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return false; + } + + ggml_backend_tensor_set(x_in, sub_input, 0, (size_t) T * (size_t) talker_hidden * sizeof(float)); + + { + std::vector pos((size_t) T); + for (int i = 0; i < T; i++) { + pos[(size_t) i] = i; + } + ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t)); + } + + { + std::vector mask((size_t) T * (size_t) T, -INFINITY); + for (int q = 0; q < T; q++) { + for (int k = 0; k <= q; k++) { + mask[(size_t) q * (size_t) T + (size_t) k] = 0.0f; + } + } + ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(float)); + } + + if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "[CodePredictor] FATAL: graph compute failed\n"); + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return false; + } + + logits_out->resize((size_t) vocab); + size_t row_bytes = (size_t) vocab * sizeof(float); + ggml_backend_tensor_get(logits, logits_out->data(), (size_t) (T - 1) * row_bytes, row_bytes); + + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return true; +} + +// Read one row of an embedding table to f32. Reads from the backend +// (the predictor weights live there) via ggml_backend_tensor_get, +// dispatched through ggml_get_type_traits so quants are accepted. +static void embed_row_from_backend(struct ggml_tensor * t, int row_id, int dim, float * dst) { + if (t->ne[0] != dim) { + fprintf(stderr, "[CodePredictor] FATAL: embed dim mismatch %lld vs %d\n", (long long) t->ne[0], dim); + std::exit(1); + } + if (row_id < 0 || row_id >= (int) t->ne[1]) { + fprintf(stderr, "[CodePredictor] FATAL: row %d out of range (vocab=%lld)\n", row_id, (long long) t->ne[1]); + std::exit(1); + } + const size_t row_bytes = ggml_row_size(t->type, dim); + if (t->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(t, dst, (size_t) row_id * row_bytes, row_bytes); + return; + } + const struct ggml_type_traits * tt = ggml_get_type_traits(t->type); + if (!tt || !tt->to_float) { + fprintf(stderr, "[CodePredictor] FATAL: unsupported embed dtype %d\n", (int) t->type); + std::exit(1); + } + std::vector tmp(row_bytes); + ggml_backend_tensor_get(t, tmp.data(), (size_t) row_id * row_bytes, row_bytes); + tt->to_float(tmp.data(), dst, dim); +} + +bool code_predictor_step(const TalkerWeights * tw, + const CodePredictorWeights * cw, + ggml_backend_sched_t sched, + const float * talker_hidden_last, + int c0, + float temperature, + int top_k, + float top_p, + int64_t seed, + int64_t subseq_base, + const char * dump_dir, + CodePredictorOutput * out) { + // sub_input lives at the talker hidden dimension because both the + // talker last hidden and the codec_embedding rows feeding the sub + // network are talker-sized in the upstream checkpoint. The recompute + // graph projects them down to predictor hidden when mtp_proj is + // present, or feeds them straight when the projection is identity. + const int talker_hidden = tw->hidden_size; + const int n_acoustic = cw->num_acoustic_codebooks; + + out->codes.assign((size_t) (n_acoustic + 1), 0); + out->codes[0] = c0; + + std::vector sub_input((size_t) (n_acoustic + 1) * (size_t) talker_hidden, 0.0f); + int T = 0; + + std::memcpy(sub_input.data() + (size_t) T * (size_t) talker_hidden, talker_hidden_last, + (size_t) talker_hidden * sizeof(float)); + T++; + + embed_row_from_backend(tw->codec_embedding, c0, talker_hidden, + sub_input.data() + (size_t) T * (size_t) talker_hidden); + T++; + + // Acoustic codebooks have no repetition penalty in upstream defaults + // so we pass a null history. + for (int g = 0; g < n_acoustic; g++) { + std::vector logits; + if (!code_predictor_recompute(cw, sched, sub_input.data(), T, talker_hidden, g, &logits)) { + return false; + } + float u_g = 0.0f; + int cg = sample_top_k_p(logits.data(), (int) logits.size(), temperature, top_k, top_p, 1.0f, nullptr, 0, seed, + subseq_base + 1 + g, &u_g); + // Trace the first 32 samples unconditionally : that window + // matches the Python harness seq < 32 trace, so [Sample-CP] and + // [Sample-PY] align by subseq id when the cossim test runs. + if (subseq_base + 1 + g < 32) { + fprintf(stderr, "[Sample-CP] g=%d c=%d u=%.10f subseq=%lld\n", g, cg, (double) u_g, + (long long) (subseq_base + 1 + g)); + } + if (cg < 0) { + fprintf(stderr, "[CodePredictor] FATAL: sample returned no candidate\n"); + return false; + } + out->codes[(size_t) (g + 1)] = cg; + + if (g + 1 < n_acoustic) { + embed_row_from_backend(cw->codec_embedding[(size_t) g], cg, talker_hidden, + sub_input.data() + (size_t) T * (size_t) talker_hidden); + T++; + } + } + + if (dump_dir) { + DebugDumper d; + debug_init(&d, dump_dir); + std::vector codes32(out->codes.begin(), out->codes.end()); + int n = (int) codes32.size(); + debug_dump_i32_as_f32(&d, "codes-step0", codes32.data(), &n, 1); + } + + return true; +} diff --git a/src/code-predictor-forward.h b/src/code-predictor-forward.h new file mode 100644 index 0000000..7116b3a --- /dev/null +++ b/src/code-predictor-forward.h @@ -0,0 +1,51 @@ +#pragma once +// code-predictor-forward.h : run the 5-layer Qwen3 code predictor over a +// growing context to produce the 15 acoustic codes of one audio frame. +// +// Input : +// talker_hidden_last [hidden] f32 -- last position hidden state from +// the Talker forward (post final norm) +// c0 -- semantic code sampled from the +// Talker codec_head (codebook 0) +// Output : +// codes[16] = [c0, c1, ..., c15] -- the full set of codes for one +// frame, ready for decode through +// the codec +// +// Phase 4.5 runs the predictor without a KV cache : every step rebuilds +// the full graph over a context of length g+2 (g being the predictor +// step, 0..14). With 5 layers and at most 16 tokens this is well below +// the threshold where caching would matter. A KV-cached variant lands +// in the generation loop phase. + +#include "code-predictor-weights.h" +#include "ggml-backend.h" +#include "sampling.h" +#include "talker-weights.h" + +#include +#include + +struct CodePredictorOutput { + // Sixteen codes : c0 from the talker plus c1..c15 from the predictor. + std::vector codes; +}; + +// Run the predictor for one audio frame. Caller passes the talker hidden +// state for the current frame and the already-sampled c0. Sampling +// parameters control greedy (temperature <= 0) vs stochastic. subseq_base +// is the Philox subsequence of the c0 sample for this step ; the 15 +// acoustic samples consume subseq_base + 1 .. subseq_base + 15. +// Returns the full vector of 16 codes. dump_dir may be NULL. +bool code_predictor_step(const TalkerWeights * tw, + const CodePredictorWeights * cw, + ggml_backend_sched_t sched, + const float * talker_hidden_last, + int c0, + float temperature, + int top_k, + float top_p, + int64_t seed, + int64_t subseq_base, + const char * dump_dir, + CodePredictorOutput * out); diff --git a/src/code-predictor-weights.h b/src/code-predictor-weights.h new file mode 100644 index 0000000..1e7e230 --- /dev/null +++ b/src/code-predictor-weights.h @@ -0,0 +1,183 @@ +#pragma once +// code-predictor-weights.h : 5-layer Qwen3 stack that predicts the +// acoustic codebooks 1..15 of every audio frame conditioned on the +// Talker hidden state and the codebook 0 token just sampled. +// +// Architecture mirrors the Talker block (pre-norm, GQA attention with +// QK-norm, SwiGLU MLP) with one important difference : RoPE is plain +// 1D (half-split, neox-style in GGUF terms) at freq base 1e6, not the +// multimodal interleaved variant the Talker uses. +// +// The MTP head carries fifteen private embedding tables and fifteen +// private linear heads, one pair per acoustic codebook. The talker +// codebook 0 stays handled by talker.codec_embedding and talker.codec_head. +// +// Tensor naming (convert.py output) : +// code_predictor.norm.weight [hidden] +// code_predictor.codec_embedding.{0..14}.weight [vocab, hidden] +// code_predictor.lm_head.{0..14}.weight [vocab, hidden] +// code_predictor.layers.{0..N-1}.input_layernorm.weight [hidden] +// code_predictor.layers.{0..N-1}.post_attention_layernorm.weight [hidden] +// code_predictor.layers.{0..N-1}.attn.{q,k,v,o}_proj.weight +// code_predictor.layers.{0..N-1}.attn.{q,k}_norm.weight [head_dim] +// code_predictor.layers.{0..N-1}.mlp.{gate,up,down}_proj.weight + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "talker-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include + +struct CodePredictorWeights { + int hidden_size; + int intermediate_size; + int num_hidden_layers; + int num_attention_heads; + int num_key_value_heads; + int head_dim; + int vocab_size; + int max_position_embeddings; + int num_acoustic_codebooks; // num_code_groups - 1 + float rope_theta; + float rms_norm_eps; + + struct ggml_tensor * norm_w; + std::vector codec_embedding; // size num_acoustic_codebooks + std::vector lm_head; // size num_acoustic_codebooks + + // Optional small_to_mtp projection that brings the talker hidden + // dimension down to the predictor hidden dimension when the two + // differ (1.7B-base case : 2048 -> 1024). Both tensors are NULL when + // the upstream sets nn.Identity() i.e. talker_hidden == predictor_hidden + // (0.6B case). Loaded with gf_try_load_tensor so absence is silent. + struct ggml_tensor * mtp_proj_w; + struct ggml_tensor * mtp_proj_b; + + // Layers reuse the same TalkerLayer struct since the per-layer + // tensor set is identical. Only the model-level wiring differs. + std::vector layers; + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +static bool code_predictor_weights_load(CodePredictorWeights * cw, const GGUFModel & gf, ggml_backend_t backend) { + cw->hidden_size = (int) gf_get_u32(gf, "qwen3-tts.code_pred.embedding_length"); + cw->intermediate_size = (int) gf_get_u32(gf, "qwen3-tts.code_pred.feed_forward_length"); + cw->num_hidden_layers = (int) gf_get_u32(gf, "qwen3-tts.code_pred.block_count"); + cw->num_attention_heads = (int) gf_get_u32(gf, "qwen3-tts.code_pred.attention.head_count"); + cw->num_key_value_heads = (int) gf_get_u32(gf, "qwen3-tts.code_pred.attention.head_count_kv"); + cw->head_dim = (int) gf_get_u32(gf, "qwen3-tts.code_pred.attention.key_length"); + cw->vocab_size = (int) gf_get_u32(gf, "qwen3-tts.code_pred.vocab_size"); + cw->max_position_embeddings = (int) gf_get_u32(gf, "qwen3-tts.code_pred.context_length"); + cw->rope_theta = gf_get_f32(gf, "qwen3-tts.code_pred.rope.freq_base"); + cw->rms_norm_eps = gf_get_f32(gf, "qwen3-tts.code_pred.attention.layer_norm_rms_epsilon"); + + int num_code_groups = (int) gf_get_u32(gf, "qwen3-tts.num_code_groups"); + if (num_code_groups <= 1) { + fprintf(stderr, "[CodePredictor] FATAL: invalid num_code_groups=%d\n", num_code_groups); + return false; + } + cw->num_acoustic_codebooks = num_code_groups - 1; + + if (cw->num_hidden_layers <= 0 || cw->hidden_size <= 0) { + fprintf(stderr, "[CodePredictor] FATAL: invalid hyperparameters (layers=%d hidden=%d)\n", cw->num_hidden_layers, + cw->hidden_size); + return false; + } + + cw->layers.resize((size_t) cw->num_hidden_layers); + cw->codec_embedding.resize((size_t) cw->num_acoustic_codebooks); + cw->lm_head.resize((size_t) cw->num_acoustic_codebooks); + + int n_tensors = 1 // final norm + + 2 // mtp_proj weight + bias (when present) + + 2 * cw->num_acoustic_codebooks // 15 embeds + 15 heads + + cw->num_hidden_layers * 11 // 11 per layer + + 8; // headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + cw->norm_w = gf_load_tensor(&wctx, gf, "code_pred.output_norm.weight"); + + // Optional projection talker_hidden -> predictor_hidden. Absent in + // checkpoints where talker.hidden_size == code_pred.hidden_size + // because upstream uses nn.Identity in that case. + cw->mtp_proj_w = gf_try_load_tensor(&wctx, gf, "code_pred.mtp_proj.weight"); + cw->mtp_proj_b = gf_try_load_tensor(&wctx, gf, "code_pred.mtp_proj.bias"); + + for (int g = 0; g < cw->num_acoustic_codebooks; g++) { + char name[160]; + snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g); + cw->codec_embedding[(size_t) g] = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.lm_head.%d.weight", g); + cw->lm_head[(size_t) g] = gf_load_tensor(&wctx, gf, name); + } + + for (int l = 0; l < cw->num_hidden_layers; l++) { + TalkerLayer & layer = cw->layers[(size_t) l]; + char name[160]; + + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_norm.weight", l); + layer.input_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_norm.weight", l); + layer.post_attn_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_q.weight", l); + layer.attn.q_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_k.weight", l); + layer.attn.k_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_v.weight", l); + layer.attn.v_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_output.weight", l); + layer.attn.o_proj_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_q_norm.weight", l); + layer.attn.q_norm_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.blk.%d.attn_k_norm.weight", l); + layer.attn.k_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_gate.weight", l); + layer.mlp.gate_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_up.weight", l); + layer.mlp.up_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_down.weight", l); + layer.mlp.down_proj_w = gf_load_tensor(&wctx, gf, name); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[CodePredictor] FATAL: backend allocation failed\n"); + return false; + } + cw->weight_ctx = wctx.ctx; + cw->weight_buf = wctx.buffer; + + fprintf(stderr, + "[CodePredictor] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, " + "FFN %d, RoPE theta %.0f, %d acoustic codebooks (vocab %d each), mtp_proj %s\n", + cw->num_hidden_layers, cw->hidden_size, cw->num_attention_heads, cw->num_key_value_heads, cw->head_dim, + cw->intermediate_size, (double) cw->rope_theta, cw->num_acoustic_codebooks, cw->vocab_size, + cw->mtp_proj_w ? "linear" : "identity"); + return true; +} + +static void code_predictor_weights_free(CodePredictorWeights * cw) { + if (cw->weight_buf) { + ggml_backend_buffer_free(cw->weight_buf); + cw->weight_buf = NULL; + } + if (cw->weight_ctx) { + ggml_free(cw->weight_ctx); + cw->weight_ctx = NULL; + } + cw->layers.clear(); + cw->codec_embedding.clear(); + cw->lm_head.clear(); +} diff --git a/src/convnext-block.h b/src/convnext-block.h new file mode 100644 index 0000000..a135de3 --- /dev/null +++ b/src/convnext-block.h @@ -0,0 +1,191 @@ +#pragma once +// convnext-block.h : 2-block upsample stage for the Qwen3-TTS 12Hz +// tokenizer decoder. +// +// Each block is a CausalTransConv1d (kernel 2, stride 2) followed by a +// ConvNeXt block. The ConvNeXt block is : +// x = x + gamma * pwconv2(gelu(pwconv1(layernorm(dwconv(x))))) +// where dwconv is a depthwise causal Conv1d (kernel 7, dilation 1), +// pwconv1 / pwconv2 are pointwise Linears (1024 -> 4096 -> 1024), and +// gamma is a per-channel LayerScale parameter. The two blocks together +// upsample the temporal axis by 4x while keeping the channel count at +// latent_dim (1024). + +#include "causal-trans-conv.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include + +#define QWEN_UPSAMPLE_MAX_BLOCKS 2 + +struct QwenConvNeXtBlock { + struct ggml_tensor * dwconv_w; // [K=7, 1, C] depthwise weight + struct ggml_tensor * dwconv_b; // [C] + struct ggml_tensor * norm_w; // LayerNorm gain [C] + struct ggml_tensor * norm_b; // LayerNorm bias [C] + struct ggml_tensor * pwconv1_w; // [C, 4*C] (in, out) + struct ggml_tensor * pwconv1_b; // [4*C] + struct ggml_tensor * pwconv2_w; // [4*C, C] + struct ggml_tensor * pwconv2_b; // [C] + struct ggml_tensor * gamma; // LayerScale [C] +}; + +struct QwenUpsampleStage { + int num_blocks; // 2 + int channels; // 1024 (= latent_dim) + int upsample_ratio; // 2 per block, 4x total + int dwconv_kernel; // 7 + + struct ggml_tensor * transconv_w[QWEN_UPSAMPLE_MAX_BLOCKS]; // pre-permuted [IC, K*OC] + struct ggml_tensor * transconv_b[QWEN_UPSAMPLE_MAX_BLOCKS]; // [OC] + QwenConvNeXtBlock convnext[QWEN_UPSAMPLE_MAX_BLOCKS]; + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +// Read upsample hyperparameters from GGUF metadata, allocate every weight +// tensor on the backend, and bind tensor pointers in the struct. +static bool qwen_upsample_stage_load(QwenUpsampleStage * stage, const GGUFModel & gf, ggml_backend_t backend) { + stage->channels = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.latent_dim"); + stage->dwconv_kernel = 7; + stage->upsample_ratio = 2; + stage->num_blocks = 2; + + if (stage->num_blocks > QWEN_UPSAMPLE_MAX_BLOCKS) { + fprintf(stderr, "[Upsample] FATAL: %d blocks exceeds compile-time max %d\n", stage->num_blocks, + QWEN_UPSAMPLE_MAX_BLOCKS); + return false; + } + + int n_tensors = stage->num_blocks * 11 + 4; // 2 transconv + 9 convnext per block, plus headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + for (int i = 0; i < stage->num_blocks; i++) { + char name[160]; + + snprintf(name, sizeof(name), "tok_dec.upsample.%d.conv.weight", i); + stage->transconv_w[i] = qwen_load_ctw_f32(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.conv.bias", i); + stage->transconv_b[i] = gf_load_tensor(&wctx, gf, name); + + QwenConvNeXtBlock & cn = stage->convnext[i]; + snprintf(name, sizeof(name), "tok_dec.upsample.%d.dwconv.weight", i); + cn.dwconv_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.dwconv.bias", i); + cn.dwconv_b = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.norm.weight", i); + cn.norm_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.norm.bias", i); + cn.norm_b = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.pwconv1.weight", i); + cn.pwconv1_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.pwconv1.bias", i); + cn.pwconv1_b = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.pwconv2.weight", i); + cn.pwconv2_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.pwconv2.bias", i); + cn.pwconv2_b = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.upsample.%d.gamma", i); + cn.gamma = gf_load_tensor(&wctx, gf, name); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[Upsample] FATAL: backend allocation failed\n"); + return false; + } + stage->weight_ctx = wctx.ctx; + stage->weight_buf = wctx.buffer; + + fprintf(stderr, + "[Upsample] Loaded: %d blocks (%dx ratio per block), channels %d, " + "dwconv kernel %d\n", + stage->num_blocks, stage->upsample_ratio, stage->channels, stage->dwconv_kernel); + return true; +} + +static void qwen_upsample_stage_free(QwenUpsampleStage * stage) { + if (stage->weight_buf) { + ggml_backend_buffer_free(stage->weight_buf); + stage->weight_buf = NULL; + } + if (stage->weight_ctx) { + ggml_free(stage->weight_ctx); + stage->weight_ctx = NULL; + } +} + +// One ConvNeXt block forward. +// x : [T, C] f32 T-first +// returns [T, C] f32 T-first +static struct ggml_tensor * qwen_convnext_block_forward(struct ggml_context * ctx, + const QwenConvNeXtBlock & block, + struct ggml_tensor * x, + int kernel) { + int T = (int) x->ne[0]; + int C = (int) x->ne[1]; + + struct ggml_tensor * residual = x; + + // dwconv : depthwise causal Conv1d. ggml_conv_1d_dw expects [T, C, B=1]. + // Pre-pad left by (kernel-1) zeros for causal behavior, no internal padding. + struct ggml_tensor * y = ggml_reshape_3d(ctx, x, T, C, 1); + y = ggml_pad_ext(ctx, y, kernel - 1, 0, 0, 0, 0, 0, 0, 0); + y = ggml_conv_1d_dw(ctx, block.dwconv_w, y, 1, 0, 1); // [T, C, 1] + y = ggml_reshape_2d(ctx, y, T, C); + if (block.dwconv_b) { + struct ggml_tensor * b2d = ggml_reshape_2d(ctx, block.dwconv_b, 1, C); // (1, C) broadcasts on T + y = ggml_add(ctx, y, b2d); + } + + // LayerNorm wants the channel dim on ne[0] : transpose to [C, T]. + y = ggml_cont(ctx, ggml_transpose(ctx, y)); + y = ggml_norm(ctx, y, 1e-6f); + y = ggml_mul(ctx, y, block.norm_w); + y = ggml_add(ctx, y, block.norm_b); + + // pwconv1 : Linear C -> 4*C. mul_mat contracts ne[0]=C of weight against + // ne[0]=C of input. + y = ggml_mul_mat(ctx, block.pwconv1_w, y); + y = ggml_add(ctx, y, block.pwconv1_b); + + y = ggml_gelu(ctx, y); + + // pwconv2 : Linear 4*C -> C + y = ggml_mul_mat(ctx, block.pwconv2_w, y); + y = ggml_add(ctx, y, block.pwconv2_b); + + // LayerScale gamma broadcast over T axis (ne[1]). + y = ggml_mul(ctx, y, block.gamma); + + // Back to T-first [T, C] to match the residual layout. + y = ggml_cont(ctx, ggml_transpose(ctx, y)); + + y = ggml_add(ctx, y, residual); + return y; +} + +// Full upsample stage forward : 2 (CausalTransConv + ConvNeXt) blocks. +// x : [T, C] f32 T-first +// returns [T * 4, C] f32 T-first +// +// The top-level upsample stage uses kernel == stride (no causal trim). +// The DAC decoder blocks (separate header) use kernel == 2 * stride +// with a stride-frame causal trim. +static struct ggml_tensor * qwen_upsample_stage_forward(struct ggml_context * ctx, + const QwenUpsampleStage * stage, + struct ggml_tensor * x) { + int kernel = stage->upsample_ratio; + for (int i = 0; i < stage->num_blocks; i++) { + x = qwen_causal_trans_conv1d(ctx, stage->transconv_w[i], stage->transconv_b[i], x, stage->upsample_ratio, + kernel, stage->channels); + x = qwen_convnext_block_forward(ctx, stage->convnext[i], x, stage->dwconv_kernel); + } + return x; +} diff --git a/src/dac-decoder-v2.h b/src/dac-decoder-v2.h new file mode 100644 index 0000000..0298f3a --- /dev/null +++ b/src/dac-decoder-v2.h @@ -0,0 +1,270 @@ +#pragma once +// dac-decoder-v2.h: DAC acoustic decoder for the Qwen3-TTS 12Hz tokenizer. +// +// Layout: T-first [T, C] from conv_pre to conv_post. The fused SNAKE op +// requires ne[0]=T and ne[1]=C, so the whole DAC pipeline matches that +// convention. ggml_conv_1d and ggml_conv_1d_dw are T-first natively, and +// the only mul_mat lives inside qwen_causal_trans_conv1d which transposes +// internally. +// +// Pipeline (input [T, 1024] -> audio [T*1920, 1] @ 24 kHz mono): +// conv_pre k=7 (1024 -> 1536, causal) +// 4 blocks: SnakeBeta -> CausalTransConv k=2*stride -> 3 ResUnits +// strides 8 / 5 / 4 / 3, channels 1536 -> 768 -> 384 -> 192 -> 96 +// ResUnit: SnakeBeta -> conv k=7 dilation -> SnakeBeta -> conv k=1 + skip +// dilations 1 / 3 / 9 +// snake_post (96) +// conv_post k=7 (96 -> 1, causal) +// +// SnakeBeta: the Qwen3-TTS reference applies exp() to alpha and beta on +// every forward. Both factors are precomputed CPU-side at load time and +// stored as a = exp(alpha) and inv_b = 1 / (exp(beta) + 1e-9), so the +// runtime kernel stays a single fused ggml_snake op with semantics +// y = x + sin^2(a * x) * inv_b. + +#include "causal-trans-conv.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include + +#define QWEN_DAC_NUM_BLOCKS 4 +#define QWEN_DAC_RES_UNITS 3 + +// SnakeBeta runtime parameters with exp() folded in: a holds exp(alpha) +// and inv_b holds 1 / (exp(beta) + 1e-9). Layout [1, C] f32 matches the +// broadcast convention of ggml_snake. +struct QwenDACSnake { + struct ggml_tensor * a; // [1, C] f32 + struct ggml_tensor * inv_b; // [1, C] f32 +}; + +// One residual unit: act1 -> conv1 (k=7, dilation d, causal) -> act2 -> +// conv2 (k=1) + skip. +struct QwenDACResUnit { + QwenDACSnake act1; + QwenDACSnake act2; + struct ggml_tensor * c1w; // [7, C, C] f32, stored (K, IC, OC) + struct ggml_tensor * c1b; // [C] f32 + struct ggml_tensor * c2w; // [1, C, C] f32 + struct ggml_tensor * c2b; // [C] f32 + int dilation; +}; + +// One DAC block: snake1 -> causal transposed conv (upsample) -> 3 res +// units. The transconv weight is pre-permuted to [IC, K*OC] at load time +// so the col2im_1d path stays a fused mul_mat. +struct QwenDACBlock { + QwenDACSnake snake1; + struct ggml_tensor * tcw; // [IC, K*OC] f32, pre-permuted from (IC, OC, K) + struct ggml_tensor * tcb; // [OC] f32 + QwenDACResUnit ru[QWEN_DAC_RES_UNITS]; + int in_ch; + int out_ch; + int stride; + int kernel; // 2 * stride +}; + +struct QwenDACDecoder { + // initial conv: 1024 -> 1536, k=7, causal + struct ggml_tensor * conv_pre_w; // [7, 1024, 1536] f32 + struct ggml_tensor * conv_pre_b; // [1536] f32 + + QwenDACBlock blk[QWEN_DAC_NUM_BLOCKS]; + + QwenDACSnake snake_post; // 96 channels + + // final conv: 96 -> 1, k=7, causal + struct ggml_tensor * conv_post_w; // [7, 96, 1] f32 + struct ggml_tensor * conv_post_b; // [1] f32 + + int channels[QWEN_DAC_NUM_BLOCKS + 1]; // 1536, 768, 384, 192, 96 + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +// Read alpha and beta from the GGUF, fold exp() and reciprocal CPU-side, +// and bind two [1, C] f32 tensors on the backend ctx as a and inv_b. +static void qwen_dac_load_snakebeta(WeightCtx * wctx, + const GGUFModel & gf, + QwenDACSnake * s, + const std::string & alpha_name, + const std::string & beta_name) { + struct ggml_tensor * alpha_meta = ggml_get_tensor(gf.meta, alpha_name.c_str()); + struct ggml_tensor * beta_meta = ggml_get_tensor(gf.meta, beta_name.c_str()); + if (!alpha_meta || !beta_meta) { + fprintf(stderr, "[DAC] FATAL: snake tensor '%s' or '%s' not found\n", alpha_name.c_str(), beta_name.c_str()); + exit(1); + } + if (alpha_meta->type != GGML_TYPE_F32 || beta_meta->type != GGML_TYPE_F32) { + fprintf(stderr, "[DAC] FATAL: snake '%s' expects F32 alpha/beta\n", alpha_name.c_str()); + exit(1); + } + int C = (int) alpha_meta->ne[0]; + if ((int) beta_meta->ne[0] != C) { + fprintf(stderr, "[DAC] FATAL: snake '%s' alpha/beta size mismatch (%d vs %d)\n", alpha_name.c_str(), C, + (int) beta_meta->ne[0]); + exit(1); + } + + s->a = ggml_new_tensor_2d(wctx->ctx, GGML_TYPE_F32, 1, C); + s->inv_b = ggml_new_tensor_2d(wctx->ctx, GGML_TYPE_F32, 1, C); + ggml_set_name(s->a, alpha_name.c_str()); + ggml_set_name(s->inv_b, beta_name.c_str()); + + const float * alpha_src = (const float *) gf_get_data(gf, alpha_name.c_str()); + const float * beta_src = (const float *) gf_get_data(gf, beta_name.c_str()); + + auto a_buf = std::make_unique((size_t) C); + auto inv_b_buf = std::make_unique((size_t) C); + for (int c = 0; c < C; c++) { + a_buf[c] = expf(alpha_src[c]); + inv_b_buf[c] = 1.0f / (expf(beta_src[c]) + 1e-9f); + } + + wctx->pending.push_back({ s->a, a_buf.get(), (size_t) C * sizeof(float), 0 }); + wctx->pending.push_back({ s->inv_b, inv_b_buf.get(), (size_t) C * sizeof(float), 0 }); + wctx->staging.push_back(std::move(a_buf)); + wctx->staging.push_back(std::move(inv_b_buf)); +} + +// Allocate every weight tensor on the backend, copy from the GGUF mapping +// with per-tensor transforms (snake exp/reciprocal, transconv permute). +static bool qwen_dac_decoder_load(QwenDACDecoder * d, const GGUFModel & gf, ggml_backend_t backend) { + static const int strides[QWEN_DAC_NUM_BLOCKS] = { 8, 5, 4, 3 }; + static const int chs[QWEN_DAC_NUM_BLOCKS + 1] = { 1536, 768, 384, 192, 96 }; + static const int dilations[QWEN_DAC_RES_UNITS] = { 1, 3, 9 }; + + for (int i = 0; i <= QWEN_DAC_NUM_BLOCKS; i++) { + d->channels[i] = chs[i]; + } + + int n_tensors = 132; // 118 actual + headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + d->conv_pre_w = gf_load_tensor(&wctx, gf, "tok_dec.dec.0.conv.weight"); + d->conv_pre_b = gf_load_tensor(&wctx, gf, "tok_dec.dec.0.conv.bias"); + + for (int i = 0; i < QWEN_DAC_NUM_BLOCKS; i++) { + QwenDACBlock & b = d->blk[i]; + b.in_ch = chs[i]; + b.out_ch = chs[i + 1]; + b.stride = strides[i]; + b.kernel = 2 * strides[i]; + + // Python ModuleList places blocks at indices 1, 2, 3, 4. Index 0 + // is the entry conv (loaded above), indices 5 and 6 are the post + // snake and final conv (loaded below). + int py_idx = i + 1; + char prefix[64]; + snprintf(prefix, sizeof(prefix), "tok_dec.dec.%d", py_idx); + + qwen_dac_load_snakebeta(&wctx, gf, &b.snake1, std::string(prefix) + ".snake.alpha", + std::string(prefix) + ".snake.beta"); + + b.tcw = qwen_load_ctw_f32(&wctx, gf, std::string(prefix) + ".conv_t.weight"); + b.tcb = gf_load_tensor(&wctx, gf, std::string(prefix) + ".conv_t.bias"); + + for (int r = 0; r < QWEN_DAC_RES_UNITS; r++) { + QwenDACResUnit & ru = b.ru[r]; + ru.dilation = dilations[r]; + + char rp[96]; + snprintf(rp, sizeof(rp), "%s.res.%d", prefix, r); + + qwen_dac_load_snakebeta(&wctx, gf, &ru.act1, std::string(rp) + ".act1.alpha", + std::string(rp) + ".act1.beta"); + ru.c1w = gf_load_tensor(&wctx, gf, std::string(rp) + ".conv1.weight"); + ru.c1b = gf_load_tensor(&wctx, gf, std::string(rp) + ".conv1.bias"); + qwen_dac_load_snakebeta(&wctx, gf, &ru.act2, std::string(rp) + ".act2.alpha", + std::string(rp) + ".act2.beta"); + ru.c2w = gf_load_tensor(&wctx, gf, std::string(rp) + ".conv2.weight"); + ru.c2b = gf_load_tensor(&wctx, gf, std::string(rp) + ".conv2.bias"); + } + } + + qwen_dac_load_snakebeta(&wctx, gf, &d->snake_post, "tok_dec.dec.5.snake.alpha", "tok_dec.dec.5.snake.beta"); + d->conv_post_w = gf_load_tensor(&wctx, gf, "tok_dec.dec.6.conv.weight"); + d->conv_post_b = gf_load_tensor(&wctx, gf, "tok_dec.dec.6.conv.bias"); + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[DAC] FATAL: backend allocation failed\n"); + return false; + } + d->weight_ctx = wctx.ctx; + d->weight_buf = wctx.buffer; + + fprintf(stderr, "[DAC] Loaded: %d blocks (strides 8/5/4/3), 24 kHz mono out, weights %.1f MB\n", + QWEN_DAC_NUM_BLOCKS, (float) ggml_backend_buffer_get_size(d->weight_buf) / (1024.0f * 1024.0f)); + return true; +} + +static void qwen_dac_decoder_free(QwenDACDecoder * d) { + if (d->weight_buf) { + ggml_backend_buffer_free(d->weight_buf); + d->weight_buf = NULL; + } + if (d->weight_ctx) { + ggml_free(d->weight_ctx); + d->weight_ctx = NULL; + } +} + +// SnakeBeta forward: y = x + sin^2(a * x) * inv_b. Written as primitive +// ops so the backend graph optimiser fuses them into a dedicated snake +// kernel where one is available, and falls back to a plain CPU/GPU op +// chain otherwise. x [T, C] T-first, a and inv_b broadcast on the C axis. +static struct ggml_tensor * qwen_dac_snake(struct ggml_context * ctx, struct ggml_tensor * x, const QwenDACSnake & s) { + struct ggml_tensor * t = ggml_mul(ctx, x, s.a); // a * x (broadcast over T) + t = ggml_sin(ctx, t); // sin(a * x) + t = ggml_sqr(ctx, t); // sin^2(a * x) + t = ggml_mul(ctx, t, s.inv_b); // sin^2(a * x) * inv_b + return ggml_add(ctx, x, t); // x + sin^2(a * x) * inv_b +} + +// Causal Conv1d helper qwen_causal_conv1d lives in causal-trans-conv.h +// alongside qwen_causal_trans_conv1d so that pre_conv (in pipeline-codec) +// and the DAC share the same primitive. + +// Residual unit forward: skip + conv2(snake(conv1(snake(x)))). +static struct ggml_tensor * qwen_dac_res_unit(struct ggml_context * ctx, + const QwenDACResUnit * ru, + struct ggml_tensor * x) { + struct ggml_tensor * skip = x; + x = qwen_dac_snake(ctx, x, ru->act1); + x = qwen_causal_conv1d(ctx, ru->c1w, ru->c1b, x, 7, ru->dilation); + x = qwen_dac_snake(ctx, x, ru->act2); + x = qwen_causal_conv1d(ctx, ru->c2w, ru->c2b, x, 1, 1); + return ggml_add(ctx, skip, x); +} + +// Full DAC forward graph. +// x: [T, 1024] f32 T-first +// returns [T * 1920, 1] f32 T-first (raw audio samples @ 24 kHz mono, +// without final clamp; the orchestration layer clips to [-1, 1]). +static struct ggml_tensor * qwen_dac_decoder_forward(struct ggml_context * ctx, + const QwenDACDecoder * d, + struct ggml_tensor * x) { + x = qwen_causal_conv1d(ctx, d->conv_pre_w, d->conv_pre_b, x, 7, 1); + + for (int i = 0; i < QWEN_DAC_NUM_BLOCKS; i++) { + const QwenDACBlock & b = d->blk[i]; + x = qwen_dac_snake(ctx, x, b.snake1); + x = qwen_causal_trans_conv1d(ctx, b.tcw, b.tcb, x, b.stride, b.kernel, b.out_ch); + for (int r = 0; r < QWEN_DAC_RES_UNITS; r++) { + x = qwen_dac_res_unit(ctx, &b.ru[r], x); + } + } + + x = qwen_dac_snake(ctx, x, d->snake_post); + x = qwen_causal_conv1d(ctx, d->conv_post_w, d->conv_post_b, x, 7, 1); + return x; +} diff --git a/src/debug.h b/src/debug.h new file mode 100644 index 0000000..996d417 --- /dev/null +++ b/src/debug.h @@ -0,0 +1,118 @@ +#pragma once +// debug.h : tensor dump and compare helpers for Python vs GGML validation. +// Dumps raw f32 arrays to binary files, both backends convert to f32 before +// dump. +// File format : [int32 ndims] [int32 dim0] [int32 dim1] ... [float data...] + +#include +#include +#include +#include + +struct DebugDumper { + char dir[512]; + bool enabled; +}; + +static void debug_init(DebugDumper * d, const char * dir) { + d->enabled = (dir != nullptr); + if (d->enabled) { + snprintf(d->dir, sizeof(d->dir), "%s", dir); + } +} + +// Dump f32 tensor to binary file. +// Format : [ndims:i32] [shape:i32 x ndims] [data:f32 x numel] +static void debug_dump(const DebugDumper * d, const char * name, const float * data, const int * shape, int ndims) { + if (!d->enabled) { + return; + } + char path[1024]; + snprintf(path, sizeof(path), "%s/%s.bin", d->dir, name); + + int numel = 1; + for (int i = 0; i < ndims; i++) { + numel *= shape[i]; + } + FILE * f = fopen(path, "wb"); + if (!f) { + fprintf(stderr, "[Debug] cannot write %s\n", path); + return; + } + fwrite(&ndims, sizeof(int32_t), 1, f); + fwrite(shape, sizeof(int32_t), ndims, f); + fwrite(data, sizeof(float), numel, f); + fclose(f); + + // First 4 values for quick sanity check on stderr. + fprintf(stderr, "[Debug] %s: [", name); + for (int i = 0; i < ndims; i++) { + fprintf(stderr, "%s%d", i ? ", " : "", shape[i]); + } + fprintf(stderr, "] first4:"); + for (int i = 0; i < 4 && i < numel; i++) { + fprintf(stderr, " %.6f", data[i]); + } + fprintf(stderr, "\n"); +} + +// Convenience : dump 1D tensor [n]. +static void debug_dump_1d(const DebugDumper * d, const char * name, const float * data, int n) { + debug_dump(d, name, data, &n, 1); +} + +// Convenience : dump 2D tensor [rows, cols]. +static void debug_dump_2d(const DebugDumper * d, const char * name, const float * data, int dim0, int dim1) { + int shape[2] = { dim0, dim1 }; + debug_dump(d, name, data, shape, 2); +} + +// Convenience : dump 3D tensor [d0, d1, d2]. +static void debug_dump_3d(const DebugDumper * d, const char * name, const float * data, int d0, int d1, int d2) { + int shape[3] = { d0, d1, d2 }; + debug_dump(d, name, data, shape, 3); +} + +// Convenience : dump 4D tensor [d0, d1, d2, d3]. +static void +debug_dump_4d(const DebugDumper * d, const char * name, const float * data, int d0, int d1, int d2, int d3) { + int shape[4] = { d0, d1, d2, d3 }; + debug_dump(d, name, data, shape, 4); +} + +// Cast a stream of int32 values to f32 in place into a temporary buffer and +// dump under the given name. Token comparisons are then expressed as cossim +// over float values, with exact match recoverable via integer compare on the +// loader side. +static void debug_dump_i32_as_f32(const DebugDumper * d, + const char * name, + const int32_t * data, + const int * shape, + int ndims) { + if (!d->enabled) { + return; + } + int numel = 1; + for (int i = 0; i < ndims; i++) { + numel *= shape[i]; + } + std::vector buf((size_t) numel); + for (int i = 0; i < numel; i++) { + buf[i] = (float) data[i]; + } + debug_dump(d, name, buf.data(), shape, ndims); +} + +// Cosine similarity between two f32 arrays. +static double debug_cosine_sim(const float * a, const float * b, int n) { + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < n; i++) { + dot += (double) a[i] * (double) b[i]; + na += (double) a[i] * (double) a[i]; + nb += (double) b[i] * (double) b[i]; + } + if (na < 1e-30 || nb < 1e-30) { + return 0.0; + } + return dot / (sqrt(na) * sqrt(nb)); +} diff --git a/src/encoder-downsample.h b/src/encoder-downsample.h new file mode 100644 index 0000000..b51a53d --- /dev/null +++ b/src/encoder-downsample.h @@ -0,0 +1,67 @@ +#pragma once +// encoder-downsample.h: final conv k=4 stride=2 between the encoder +// transformer output and the RVQ. Brings the rate from 25 Hz to 12.5 Hz +// while preserving the 512-channel hidden dimension. +// +// The Python upstream defines this as a MimiConv1d with use_causal_conv, +// no bias, kernel 4, stride 2. The shape in the GGUF is (4, 512, 512). + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include + +struct QwenEncoderDownsample { + struct ggml_tensor * weight; // [4, 512, 512] f32 + int in_ch; + int out_ch; + int kernel; + int stride; + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +static bool qwen_encoder_downsample_load(QwenEncoderDownsample * d, const GGUFModel & gf, ggml_backend_t backend) { + d->kernel = 4; + d->stride = 2; + + WeightCtx wctx; + wctx_init(&wctx, 4); + d->weight = gf_load_tensor(&wctx, gf, "tok_enc.downsample.weight"); + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[EncDownsample] FATAL: backend allocation failed\n"); + return false; + } + d->weight_ctx = wctx.ctx; + d->weight_buf = wctx.buffer; + + d->in_ch = (int) d->weight->ne[1]; + d->out_ch = (int) d->weight->ne[2]; + + fprintf(stderr, "[EncDownsample] Loaded: k=%d stride=%d, %d -> %d channels, weights %.1f MB\n", d->kernel, + d->stride, d->in_ch, d->out_ch, (float) ggml_backend_buffer_get_size(d->weight_buf) / (1024.0f * 1024.0f)); + return true; +} + +static void qwen_encoder_downsample_free(QwenEncoderDownsample * d) { + if (d->weight_buf) { + ggml_backend_buffer_free(d->weight_buf); + d->weight_buf = NULL; + } + if (d->weight_ctx) { + ggml_free(d->weight_ctx); + d->weight_ctx = NULL; + } +} + +// Forward: causal Conv1d k=4 stride=2, no bias. +// x: [T, 512] f32 T-first +// Returns [ceil(T/2), 512] f32 T-first. +static struct ggml_tensor * qwen_encoder_downsample_forward(struct ggml_context * ctx, + const QwenEncoderDownsample * d, + struct ggml_tensor * x) { + return qwen_causal_conv1d(ctx, d->weight, NULL, x, d->kernel, 1, d->stride); +} diff --git a/src/encoder-transformer.h b/src/encoder-transformer.h new file mode 100644 index 0000000..a25bf62 --- /dev/null +++ b/src/encoder-transformer.h @@ -0,0 +1,249 @@ +#pragma once +// encoder-transformer.h: 8-layer Mimi-style transformer for the Qwen3-TTS +// encoder. Operates on the SEANet output stream at 25 Hz, 512 channels. +// +// Differs from the Qwen3-style decoder transformer on several points: +// - LayerNorm with bias (not RMSNorm) +// - Plain MLP fc1 -> GELU -> fc2 (not SwiGLU) +// - No biases on q/k/v/o projections +// - 8 attention heads instead of 16 +// - intermediate_size 2048 instead of 1024 +// - Causal sliding window 250 frames instead of 72 +// - No top-level input_proj / output_proj brackets: the SEANet output +// already has hidden_size channels +// +// Common with the decoder side: +// - RoPE NEOX style with theta 10000 +// - LayerScale per channel post-attention and post-MLP +// - Pre-norm residual on both attention and MLP + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include +#include + +#define QWEN_ENCODER_TRANSFORMER_MAX_LAYERS 16 + +struct QwenEncoderTransformerLayer { + // Pre-attention LayerNorm + struct ggml_tensor * input_norm_w; + struct ggml_tensor * input_norm_b; + // Attention projections, no bias + struct ggml_tensor * q_proj_w; + struct ggml_tensor * k_proj_w; + struct ggml_tensor * v_proj_w; + struct ggml_tensor * o_proj_w; + // Post-attention LayerScale per channel + struct ggml_tensor * attn_scale; + // Post-attention pre-MLP LayerNorm + struct ggml_tensor * post_attn_norm_w; + struct ggml_tensor * post_attn_norm_b; + // MLP + struct ggml_tensor * fc1_w; + struct ggml_tensor * fc2_w; + // Post-MLP LayerScale per channel + struct ggml_tensor * mlp_scale; +}; + +struct QwenEncoderTransformer { + int hidden_size; + int num_layers; + int num_attention_heads; + int num_kv_heads; + int head_dim; + int intermediate_size; + int sliding_window; + float rope_theta; + float norm_eps; + + QwenEncoderTransformerLayer layers[QWEN_ENCODER_TRANSFORMER_MAX_LAYERS]; + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +static bool qwen_encoder_transformer_load(QwenEncoderTransformer * tr, const GGUFModel & gf, ggml_backend_t backend) { + tr->hidden_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.hidden_size"); + tr->num_layers = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_hidden_layers"); + tr->num_attention_heads = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_attention_heads"); + tr->num_kv_heads = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_key_value_heads"); + tr->head_dim = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.head_dim"); + tr->intermediate_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.intermediate_size"); + tr->sliding_window = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.sliding_window"); + tr->rope_theta = gf_get_f32(gf, "qwen3-tts-tokenizer.encoder.rope_theta"); + tr->norm_eps = gf_get_f32(gf, "qwen3-tts-tokenizer.encoder.norm_eps"); + + if (tr->num_layers > QWEN_ENCODER_TRANSFORMER_MAX_LAYERS) { + fprintf(stderr, "[EncTransformer] FATAL: %d layers exceeds compile-time max %d\n", tr->num_layers, + QWEN_ENCODER_TRANSFORMER_MAX_LAYERS); + return false; + } + + int n_tensors = tr->num_layers * 12 + 4; // 12 tensors per layer + headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + for (int l = 0; l < tr->num_layers; l++) { + QwenEncoderTransformerLayer & ly = tr->layers[l]; + char prefix[96]; + snprintf(prefix, sizeof(prefix), "tok_enc.blk.%d", l); + std::string p(prefix); + + ly.input_norm_w = gf_load_tensor(&wctx, gf, p + ".attn_norm.weight"); + ly.input_norm_b = gf_load_tensor(&wctx, gf, p + ".attn_norm.bias"); + ly.q_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_q.weight"); + ly.k_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_k.weight"); + ly.v_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_v.weight"); + ly.o_proj_w = gf_load_tensor(&wctx, gf, p + ".attn_output.weight"); + ly.attn_scale = gf_load_tensor(&wctx, gf, p + ".attn_scale"); + ly.post_attn_norm_w = gf_load_tensor(&wctx, gf, p + ".ffn_norm.weight"); + ly.post_attn_norm_b = gf_load_tensor(&wctx, gf, p + ".ffn_norm.bias"); + ly.fc1_w = gf_load_tensor(&wctx, gf, p + ".ffn_up.weight"); + ly.fc2_w = gf_load_tensor(&wctx, gf, p + ".ffn_down.weight"); + ly.mlp_scale = gf_load_tensor(&wctx, gf, p + ".ffn_scale"); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[EncTransformer] FATAL: backend allocation failed\n"); + return false; + } + tr->weight_ctx = wctx.ctx; + tr->weight_buf = wctx.buffer; + + fprintf(stderr, + "[EncTransformer] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, " + "FFN %d, RoPE theta %.0f, sliding window %d\n", + tr->num_layers, tr->hidden_size, tr->num_attention_heads, tr->num_kv_heads, tr->head_dim, + tr->intermediate_size, tr->rope_theta, tr->sliding_window); + return true; +} + +static void qwen_encoder_transformer_free(QwenEncoderTransformer * tr) { + if (tr->weight_buf) { + ggml_backend_buffer_free(tr->weight_buf); + tr->weight_buf = NULL; + } + if (tr->weight_ctx) { + ggml_free(tr->weight_ctx); + tr->weight_ctx = NULL; + } +} + +// Build a [T, T] additive mask (0 where allowed, -inf where masked) +// for a causal sliding window of width `window`. Allowed when k <= q +// and (q - k) < window. +static void qwen_encoder_build_causal_sliding_mask(int T, int window, std::vector & dst) { + dst.assign((size_t) T * (size_t) T, -INFINITY); + for (int q = 0; q < T; q++) { + int k_min = q - window + 1; + if (k_min < 0) { + k_min = 0; + } + for (int k = k_min; k <= q; k++) { + dst[(size_t) q * (size_t) T + (size_t) k] = 0.0f; + } + } +} + +static void qwen_encoder_build_positions(int T, std::vector & dst) { + dst.resize((size_t) T); + for (int i = 0; i < T; i++) { + dst[i] = i; + } +} + +// One Mimi transformer layer. Pre-LayerNorm with bias, attention without +// q/k/v/o biases, MLP with fc1 -> GELU -> fc2 (no SwiGLU), LayerScale on +// both residual paths. +// x : [hidden, T] f32 C-first +// positions : [T] i32 +// mask : [T, T] f32 additive +// Returns [hidden, T] f32 C-first. +static struct ggml_tensor * qwen_encoder_transformer_layer_forward(struct ggml_context * ctx, + const QwenEncoderTransformer * tr, + const QwenEncoderTransformerLayer & layer, + struct ggml_tensor * x, + struct ggml_tensor * positions, + struct ggml_tensor * mask, + int T) { + int hidden = tr->hidden_size; + int n_q_heads = tr->num_attention_heads; + int n_kv = tr->num_kv_heads; + int hd = tr->head_dim; + + // Pre-LayerNorm with affine (weight + bias). ggml_norm normalizes on ne[0] + // which is `hidden` here, matching PyTorch nn.LayerNorm(hidden). + struct ggml_tensor * ln1 = ggml_norm(ctx, x, tr->norm_eps); + ln1 = ggml_mul(ctx, ln1, layer.input_norm_w); + ln1 = ggml_add(ctx, ln1, layer.input_norm_b); + + struct ggml_tensor * q = ggml_mul_mat(ctx, layer.q_proj_w, ln1); + struct ggml_tensor * k = ggml_mul_mat(ctx, layer.k_proj_w, ln1); + struct ggml_tensor * v = ggml_mul_mat(ctx, layer.v_proj_w, ln1); + + q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T); + k = ggml_reshape_3d(ctx, k, hd, n_kv, T); + v = ggml_reshape_3d(ctx, v, hd, n_kv, T); + + q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + + struct ggml_tensor * q_p = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + struct ggml_tensor * k_p = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + struct ggml_tensor * v_p = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); + + struct ggml_tensor * scores = ggml_mul_mat(ctx, k_p, q_p); + float scale = 1.0f / sqrtf((float) hd); + scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); + + struct ggml_tensor * attn = ggml_mul_mat(ctx, v_p, scores); + attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3)); + attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T); + + struct ggml_tensor * o = ggml_mul_mat(ctx, layer.o_proj_w, attn); + + o = ggml_mul(ctx, o, layer.attn_scale); + x = ggml_add(ctx, x, o); + + // MLP block: pre-LayerNorm (with bias) + fc1 -> GELU(erf) -> fc2 + LayerScale + residual. + struct ggml_tensor * ln2 = ggml_norm(ctx, x, tr->norm_eps); + ln2 = ggml_mul(ctx, ln2, layer.post_attn_norm_w); + ln2 = ggml_add(ctx, ln2, layer.post_attn_norm_b); + + struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.fc1_w, ln2); + mlp = ggml_gelu_erf(ctx, mlp); + mlp = ggml_mul_mat(ctx, layer.fc2_w, mlp); + + mlp = ggml_mul(ctx, mlp, layer.mlp_scale); + x = ggml_add(ctx, x, mlp); + + (void) hidden; + return x; +} + +// Full encoder transformer forward. No top-level input_proj or output_proj +// brackets: the SEANet output already has hidden_size channels. +// x : [hidden, T] f32 C-first +// positions : [T] i32 +// mask : [T, T] f32 additive +// Returns [hidden, T] f32 C-first. +static struct ggml_tensor * qwen_encoder_transformer_forward(struct ggml_context * ctx, + const QwenEncoderTransformer * tr, + struct ggml_tensor * x, + struct ggml_tensor * positions, + struct ggml_tensor * mask) { + int T = (int) x->ne[1]; + for (int l = 0; l < tr->num_layers; l++) { + x = qwen_encoder_transformer_layer_forward(ctx, tr, tr->layers[l], x, positions, mask, T); + } + return x; +} diff --git a/src/gguf-weights.h b/src/gguf-weights.h new file mode 100644 index 0000000..2708099 --- /dev/null +++ b/src/gguf-weights.h @@ -0,0 +1,464 @@ +#pragma once +// gguf-weights.h: load model weights from GGUF files +// +// GGUF weight loader for all model components (LM, DiT, CondEncoder, TextEncoder, Detokenizer, VAE). +// All components use GGUF bf16 files generated by convert.py. +// +// Usage: +// GGUFModel gf; +// if (!gf_load(&gf, "model.gguf")) { error; } +// WeightCtx wctx; +// wctx_init(&wctx, n_tensors); +// ggml_tensor * w = gf_load_tensor(&wctx, gf, "layer.0.weight"); +// wctx_alloc(&wctx, backend); +// gf_close(&gf); // safe after wctx_alloc copied data to GPU + +#include "gguf.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include + +#ifdef _WIN32 +# define NOMINMAX +# include +#else +# include +# include +# include +# include +#endif + +struct GGUFModel { + struct gguf_context * gguf; // parsed header (KV + tensor metadata) + struct ggml_context * meta; // tensor descriptors (no data) + uint8_t * mapping; // mmapped file + size_t file_size; + size_t data_offset; // gguf_get_data_offset(gguf) +#ifdef _WIN32 + HANDLE fh; + HANDLE mh; +#else + int fd; +#endif +}; + +static void gf_close(GGUFModel * gf) { + if (gf->gguf) { + gguf_free(gf->gguf); + } + if (gf->meta) { + ggml_free(gf->meta); + } +#ifdef _WIN32 + if (gf->mapping) { + UnmapViewOfFile(gf->mapping); + } + if (gf->mh) { + CloseHandle(gf->mh); + } + if (gf->fh && gf->fh != INVALID_HANDLE_VALUE) { + CloseHandle(gf->fh); + } +#else + if (gf->mapping) { + munmap(gf->mapping, gf->file_size); + } + if (gf->fd >= 0) { + close(gf->fd); + } +#endif + *gf = {}; +} + +static bool gf_load(GGUFModel * gf, const char * path) { + *gf = {}; + + // mmap the file +#ifdef _WIN32 + gf->fh = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (gf->fh == INVALID_HANDLE_VALUE) { + fprintf(stderr, "[GGUF] Cannot open %s\n", path); + return false; + } + LARGE_INTEGER li; + GetFileSizeEx(gf->fh, &li); + gf->file_size = (size_t) li.QuadPart; + gf->mh = CreateFileMappingA(gf->fh, NULL, PAGE_READONLY, 0, 0, NULL); + if (!gf->mh) { + CloseHandle(gf->fh); + fprintf(stderr, "[GGUF] CreateFileMapping failed %s\n", path); + return false; + } + gf->mapping = (uint8_t *) MapViewOfFile(gf->mh, FILE_MAP_READ, 0, 0, 0); + if (!gf->mapping) { + CloseHandle(gf->mh); + CloseHandle(gf->fh); + fprintf(stderr, "[GGUF] MapViewOfFile failed %s\n", path); + return false; + } +#else + gf->fd = open(path, O_RDONLY); + if (gf->fd < 0) { + fprintf(stderr, "[GGUF] Cannot open %s\n", path); + return false; + } + struct stat sb; + fstat(gf->fd, &sb); + gf->file_size = (size_t) sb.st_size; + gf->mapping = (uint8_t *) mmap(NULL, gf->file_size, PROT_READ, MAP_PRIVATE, gf->fd, 0); + if (gf->mapping == MAP_FAILED) { + close(gf->fd); + gf->mapping = NULL; + fprintf(stderr, "[GGUF] Mmap failed %s\n", path); + return false; + } +#endif + + // Parse GGUF header, create tensor metadata context + struct ggml_context * meta = NULL; + struct gguf_init_params params = { /*no_alloc=*/true, /*ctx=*/&meta }; + gf->gguf = gguf_init_from_file(path, params); + if (!gf->gguf) { + fprintf(stderr, "[GGUF] Failed to parse %s\n", path); + gf_close(gf); + return false; + } + gf->meta = meta; + gf->data_offset = gguf_get_data_offset(gf->gguf); + + int64_t n = gguf_get_n_tensors(gf->gguf); + + // Verify every tensor fits inside the mapped file. Catches truncated + // downloads early with a clear message instead of a segfault deep in + // cuMemcpyHtoDAsync when the backend reads past the mmap. + for (int64_t i = 0; i < n; i++) { + const char * tname = gguf_get_tensor_name(gf->gguf, i); + struct ggml_tensor * t = ggml_get_tensor(gf->meta, tname); + size_t toff = gguf_get_tensor_offset(gf->gguf, i); + size_t tsize = ggml_nbytes(t); + size_t end = gf->data_offset + toff + tsize; + if (end > gf->file_size) { + fprintf(stderr, + "[GGUF] FATAL: '%s' is truncated or corrupt.\n" + " tensor '%s' needs bytes [%zu..%zu) but file is only %zu bytes.\n" + " Re-download the file and verify its size or checksum.\n", + path, tname, gf->data_offset + toff, end, gf->file_size); + gf_close(gf); + return false; + } + } + + fprintf(stderr, "[GGUF] %s: %lld tensors, data at offset %zu\n", path, (long long) n, gf->data_offset); + return true; +} + +// Load a tensor from GGUF into the weight context. +// Returns ggml_tensor (not yet backed by memory; call wctx_alloc after all loads). +// Tensor shapes are already in ggml order (ne[0]=innermost). +static struct ggml_tensor * gf_load_tensor(WeightCtx * wctx, + const GGUFModel & gf, + const std::string & name, + const int64_t * shape_override = nullptr, + int n_dims_override = 0) { + int64_t idx = gguf_find_tensor(gf.gguf, name.c_str()); + if (idx < 0) { + fprintf(stderr, "[GGUF] FATAL: tensor '%s' not found\n", name.c_str()); + exit(1); + } + + // Get metadata from the context populated by gguf_init_from_file + struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str()); + if (!src) { + fprintf(stderr, "[GGUF] FATAL: tensor '%s' not in meta context\n", name.c_str()); + exit(1); + } + + int n_dims; + int64_t ne[4] = { 1, 1, 1, 1 }; + + if (shape_override && n_dims_override > 0) { + n_dims = n_dims_override; + for (int i = 0; i < n_dims; i++) { + ne[i] = shape_override[i]; + } + } else { + n_dims = ggml_n_dims(src); + for (int i = 0; i < n_dims; i++) { + ne[i] = src->ne[i]; + } + } + + struct ggml_tensor * tensor = ggml_new_tensor(wctx->ctx, src->type, n_dims, ne); + ggml_set_name(tensor, name.c_str()); + + size_t offset = gguf_get_tensor_offset(gf.gguf, idx); + const void * data = gf.mapping + gf.data_offset + offset; + size_t nbytes = ggml_nbytes(src); + + wctx->pending.push_back({ tensor, data, nbytes, 0 }); + return tensor; +} + +// Try to load, returns nullptr if not found (no exit) +static struct ggml_tensor * gf_try_load_tensor(WeightCtx * wctx, const GGUFModel & gf, const std::string & name) { + int64_t idx = gguf_find_tensor(gf.gguf, name.c_str()); + if (idx < 0) { + return nullptr; + } + return gf_load_tensor(wctx, gf, name); +} + +// Load tensor, converting to F32 at load time (eliminates runtime cast nodes). +// Best for small tensors: norms [H], QK-norms [D], scale_shift_table [H,6], biases. +static struct ggml_tensor * gf_load_tensor_f32(WeightCtx * wctx, const GGUFModel & gf, const std::string & name) { + int64_t idx = gguf_find_tensor(gf.gguf, name.c_str()); + if (idx < 0) { + fprintf(stderr, "[GGUF] FATAL: tensor '%s' not found\n", name.c_str()); + exit(1); + } + struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str()); + int n_dims = ggml_n_dims(src); + int64_t ne[4] = { 1, 1, 1, 1 }; + for (int i = 0; i < n_dims; i++) { + ne[i] = src->ne[i]; + } + + // If already F32, just load normally + if (src->type == GGML_TYPE_F32) { + return gf_load_tensor(wctx, gf, name); + } + + // Bail early on unsupported types (before creating tensor in ctx) + if (src->type != GGML_TYPE_BF16 && src->type != GGML_TYPE_F16) { + fprintf(stderr, "[GGUF] WARNING: gf_load_tensor_f32 unsupported type %d for '%s', loading as-is\n", src->type, + name.c_str()); + return gf_load_tensor(wctx, gf, name); + } + + // Create F32 tensor + struct ggml_tensor * tensor = ggml_new_tensor(wctx->ctx, GGML_TYPE_F32, n_dims, ne); + ggml_set_name(tensor, name.c_str()); + + // Convert data into staging buffer. unique_ptr keeps .get() stable even + // when wctx->staging grows on subsequent calls. + size_t n = ggml_nelements(src); + auto buf = std::make_unique(n); + float * data = buf.get(); + + size_t offset = gguf_get_tensor_offset(gf.gguf, idx); + const void * raw = gf.mapping + gf.data_offset + offset; + + if (src->type == GGML_TYPE_BF16) { + const uint16_t * p = (const uint16_t *) raw; + for (size_t i = 0; i < n; i++) { + data[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *) &p[i]); + } + } else { + ggml_fp16_to_fp32_row((const ggml_fp16_t *) raw, data, (int) n); + } + + wctx->pending.push_back({ tensor, data, n * sizeof(float), 0 }); + wctx->staging.push_back(std::move(buf)); + return tensor; +} + +// Get raw pointer to tensor data in the mmapped file. +// Useful for CPU-side operations (e.g. bf16 embed lookup for lyrics). +// Returns NULL if not found. +static const void * gf_get_data(const GGUFModel & gf, const char * name) { + int64_t idx = gguf_find_tensor(gf.gguf, name); + if (idx < 0) { + return NULL; + } + size_t offset = gguf_get_tensor_offset(gf.gguf, idx); + return gf.mapping + gf.data_offset + offset; +} + +// Look up the native ggml type of a tensor stored in the GGUF, so callers +// can mirror it on the backend allocation and let dtype-agnostic memcpy +// handle the bytes. Aborts if the tensor is missing. +static enum ggml_type gf_get_type(const GGUFModel & gf, const std::string & name) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str()); + if (!src) { + fprintf(stderr, "[GGUF] FATAL: tensor '%s' not in meta context\n", name.c_str()); + exit(1); + } + return src->type; +} + +// Load a Conv1d weight onto an F16 backend tensor regardless of the source +// dtype. Mandatory on ARM aarch64 : the CPU im2col op asserts src0 is F16, +// while x86 silently accepts BF16 / F32. F16 source memcpy passes through ; +// F32 / BF16 widen ; Q8_0 / Q4_K / Q5_K / Q6_K dequantize via type traits. +// The destination tensor must be allocated as GGML_TYPE_F16. +static void gf_load_conv_f16(struct ggml_tensor * dst, const GGUFModel & gf, const std::string & name) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str()); + if (!src) { + fprintf(stderr, "[GGUF] FATAL: tensor '%s' not in meta context\n", name.c_str()); + exit(1); + } + GGML_ASSERT(dst->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_nelements(dst) == ggml_nelements(src)); + + const void * raw = gf_get_data(gf, name.c_str()); + size_t n = (size_t) ggml_nelements(src); + + // F16 source : direct memcpy, no conversion needed. + if (src->type == GGML_TYPE_F16) { + ggml_backend_tensor_set(dst, raw, 0, ggml_nbytes(dst)); + return; + } + + // All other types widen / dequantize to F32, then cast down to F16. + std::vector f32(n); + if (src->type == GGML_TYPE_F32) { + memcpy(f32.data(), raw, n * sizeof(float)); + } else if (src->type == GGML_TYPE_BF16) { + const uint16_t * p = (const uint16_t *) raw; + for (size_t i = 0; i < n; i++) { + f32[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *) &p[i]); + } + } else { + const struct ggml_type_traits * tr = ggml_get_type_traits(src->type); + if (!tr || !tr->to_float) { + fprintf(stderr, "[GGUF] FATAL: unsupported conv weight type %s for '%s'\n", ggml_type_name(src->type), + name.c_str()); + exit(1); + } + tr->to_float(raw, f32.data(), (int64_t) n); + } + + std::vector f16(n); + ggml_fp32_to_fp16_row(f32.data(), f16.data(), (int) n); + ggml_backend_tensor_set(dst, f16.data(), 0, n * sizeof(ggml_fp16_t)); +} + +// Fuse Q, K, V projection weights into a single tensor [ne0, q_ne1 + k_ne1 + v_ne1]. +// Works for any quantized type since quantization is per-row (along ne[0]). +// The fused tensor data is q rows || k rows || v rows (contiguous). +static struct ggml_tensor * gf_load_qkv_fused(WeightCtx * wctx, + const GGUFModel & gf, + const std::string & q_name, + const std::string & k_name, + const std::string & v_name) { + struct ggml_tensor * q_src = ggml_get_tensor(gf.meta, q_name.c_str()); + struct ggml_tensor * k_src = ggml_get_tensor(gf.meta, k_name.c_str()); + struct ggml_tensor * v_src = ggml_get_tensor(gf.meta, v_name.c_str()); + if (!q_src || !k_src || !v_src) { + fprintf(stderr, "[GGUF] FATAL: QKV tensor not found: %s / %s / %s\n", q_name.c_str(), k_name.c_str(), + v_name.c_str()); + exit(1); + } + // All must share ne[0] (input dim) and type - otherwise can't fuse + GGML_ASSERT(q_src->ne[0] == k_src->ne[0] && k_src->ne[0] == v_src->ne[0]); + if (q_src->type != k_src->type || k_src->type != v_src->type) { + return NULL; // caller should fall back to separate loads + } + + int64_t ne0 = q_src->ne[0]; + int64_t fused_ne1 = q_src->ne[1] + k_src->ne[1] + v_src->ne[1]; + int64_t ne[2] = { ne0, fused_ne1 }; + struct ggml_tensor * fused = ggml_new_tensor(wctx->ctx, q_src->type, 2, ne); + + size_t row_size = ggml_row_size(q_src->type, ne0); + size_t q_bytes = q_src->ne[1] * row_size; + size_t k_bytes = k_src->ne[1] * row_size; + size_t v_bytes = v_src->ne[1] * row_size; + + auto get_data = [&](const std::string & name) -> const void * { + int64_t idx = gguf_find_tensor(gf.gguf, name.c_str()); + size_t off = gguf_get_tensor_offset(gf.gguf, idx); + return gf.mapping + gf.data_offset + off; + }; + + wctx->pending.push_back({ fused, get_data(q_name), q_bytes, 0 }); + wctx->pending.push_back({ fused, get_data(k_name), k_bytes, q_bytes }); + wctx->pending.push_back({ fused, get_data(v_name), v_bytes, q_bytes + k_bytes }); + return fused; +} + +// Fuse two projection weights [ne0, a_ne1 + b_ne1] when types match. +// Returns NULL if types differ. +static struct ggml_tensor * gf_load_pair_fused(WeightCtx * wctx, + const GGUFModel & gf, + const std::string & a_name, + const std::string & b_name) { + struct ggml_tensor * a_src = ggml_get_tensor(gf.meta, a_name.c_str()); + struct ggml_tensor * b_src = ggml_get_tensor(gf.meta, b_name.c_str()); + if (!a_src || !b_src) { + return NULL; + } + if (a_src->ne[0] != b_src->ne[0] || a_src->type != b_src->type) { + return NULL; + } + + int64_t ne0 = a_src->ne[0]; + int64_t ne[2] = { ne0, a_src->ne[1] + b_src->ne[1] }; + struct ggml_tensor * fused = ggml_new_tensor(wctx->ctx, a_src->type, 2, ne); + + size_t row_size = ggml_row_size(a_src->type, ne0); + size_t a_bytes = a_src->ne[1] * row_size; + size_t b_bytes = b_src->ne[1] * row_size; + + auto get_data = [&](const std::string & name) -> const void * { + int64_t idx = gguf_find_tensor(gf.gguf, name.c_str()); + size_t off = gguf_get_tensor_offset(gf.gguf, idx); + return gf.mapping + gf.data_offset + off; + }; + + wctx->pending.push_back({ fused, get_data(a_name), a_bytes, 0 }); + wctx->pending.push_back({ fused, get_data(b_name), b_bytes, a_bytes }); + return fused; +} + +// Read a uint32 KV value (returns 0 if not found) +static uint32_t gf_get_u32(const GGUFModel & gf, const char * key) { + int64_t idx = gguf_find_key(gf.gguf, key); + if (idx < 0) { + return 0; + } + return gguf_get_val_u32(gf.gguf, idx); +} + +// Read a float32 KV value (returns 0 if not found) +static float gf_get_f32(const GGUFModel & gf, const char * key) { + int64_t idx = gguf_find_key(gf.gguf, key); + if (idx < 0) { + return 0.0f; + } + return gguf_get_val_f32(gf.gguf, idx); +} + +// Read a string KV value (returns "" if not found) +static const char * gf_get_str(const GGUFModel & gf, const char * key) { + int64_t idx = gguf_find_key(gf.gguf, key); + if (idx < 0) { + return ""; + } + return gguf_get_val_str(gf.gguf, idx); +} + +// Read a bool KV value (returns false if not found) +static bool gf_get_bool(const GGUFModel & gf, const char * key) { + int64_t idx = gguf_find_key(gf.gguf, key); + if (idx < 0) { + return false; + } + return gguf_get_val_bool(gf.gguf, idx); +} + +// Read an array of uint32 KV values (returns empty vector if not found) +static std::vector gf_get_array_u32(const GGUFModel & gf, const char * key) { + int64_t idx = gguf_find_key(gf.gguf, key); + if (idx < 0) { + return {}; + } + size_t n = gguf_get_arr_n(gf.gguf, idx); + const uint32_t * raw = (const uint32_t *) gguf_get_arr_data(gf.gguf, idx); + std::vector out(raw, raw + n); + return out; +} diff --git a/src/philox.h b/src/philox.h new file mode 100644 index 0000000..750dc67 --- /dev/null +++ b/src/philox.h @@ -0,0 +1,165 @@ +#pragma once +// philox.h Philox4x32-10 PRNG + Box-Muller normal distribution +// +// Matches PyTorch CUDA torch.randn() output (cuRAND Philox4_32_10). +// Zero dependencies beyond , , . +// +// CUDA kernel mapping (normal distribution): +// element[k] = philox_normal4(seed, subsequence=k, offset=0)[0] +// vals[1..3] discarded (one thread per element, one normal per thread). + +#include +#include +#include + +// Philox constants (same as cuRAND / Random123) +static constexpr uint32_t PHILOX_M0 = 0xD2511F53u; +static constexpr uint32_t PHILOX_M1 = 0xCD9E8D57u; +static constexpr uint32_t PHILOX_W0 = 0x9E3779B9u; +static constexpr uint32_t PHILOX_W1 = 0xBB67AE85u; + +// cuRAND uniform conversion +static constexpr float CURAND_2POW32_INV = 2.3283064365386963e-10f; // 1 / 2^32 +static constexpr float CURAND_2POW32_INV_2PI = 1.4629180792671596e-09f; // 2*PI / 2^32 + +struct Philox4 { + uint32_t x, y, z, w; +}; + +// 32x32 -> (hi32, lo32) +static inline void mulhilo32(uint32_t a, uint32_t b, uint32_t * hi, uint32_t * lo) { + uint64_t prod = (uint64_t) a * (uint64_t) b; + *lo = (uint32_t) prod; + *hi = (uint32_t) (prod >> 32); +} + +// Single Philox round +static inline Philox4 philox_round(Philox4 ctr, uint32_t k0, uint32_t k1) { + uint32_t hi0, lo0, hi1, lo1; + mulhilo32(PHILOX_M0, ctr.x, &hi0, &lo0); + mulhilo32(PHILOX_M1, ctr.z, &hi1, &lo1); + return { + hi1 ^ ctr.y ^ k0, + lo1, + hi0 ^ ctr.w ^ k1, + lo0, + }; +} + +// Philox4x32-10: 10 rounds +static inline Philox4 philox4x32_10(Philox4 ctr, uint32_t seed_lo, uint32_t seed_hi) { + uint32_t k0 = seed_lo; + uint32_t k1 = seed_hi; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + k0 += PHILOX_W0; + k1 += PHILOX_W1; + ctr = philox_round(ctr, k0, k1); + return ctr; +} + +// cuRAND Box-Muller: 2 uint32 -> 2 N(0,1) +static inline void box_muller(uint32_t u0, uint32_t u1, float * n0, float * n1) { + float u = (float) u0 * CURAND_2POW32_INV + (CURAND_2POW32_INV * 0.5f); + float v = (float) u1 * CURAND_2POW32_INV_2PI + (CURAND_2POW32_INV_2PI * 0.5f); + float s = sqrtf(-2.0f * logf(u)); + *n0 = s * sinf(v); + *n1 = s * cosf(v); +} + +// Generate 4 N(0,1) for (seed, subsequence, offset) +// counter = [offset_lo, offset_hi, subseq_lo, subseq_hi] +static inline void philox_normal4(int64_t seed, int64_t subsequence, int64_t offset, float out[4]) { + Philox4 ctr = { + (uint32_t) (offset), + (uint32_t) (offset >> 32), + (uint32_t) (subsequence), + (uint32_t) (subsequence >> 32), + }; + uint32_t slo = (uint32_t) (seed); + uint32_t shi = (uint32_t) ((uint64_t) seed >> 32); + Philox4 r = philox4x32_10(ctr, slo, shi); + box_muller(r.x, r.y, &out[0], &out[1]); + box_muller(r.z, r.w, &out[2], &out[3]); +} + +// bf16 round-trip (match torch.bfloat16 precision) +static inline float f32_to_bf16_to_f32(float x) { + uint32_t bits; + memcpy(&bits, &x, 4); + bits += 0x7FFF + ((bits >> 16) & 1); // round-to-nearest-even + bits &= 0xFFFF0000u; + float y; + memcpy(&y, &bits, 4); + return y; +} + +// Fill array with N(0,1) matching torch.randn() on CUDA with bf16. +// +// Reproduces: +// gen = torch.Generator(device="cuda").manual_seed(seed) +// torch.randn([...], generator=gen, device="cuda", dtype=torch.bfloat16) +// +// PyTorch CUDA normal distribution: each element k gets its own Philox +// subsequence and uses only the first Box-Muller output (val[0]). +// vals[1..3] are discarded. This matches the CUDA kernel behavior where +// grid = ceil(n / block_size), one element per thread. +static inline void philox_randn(int64_t seed, float * out, int n, bool bf16_round = true) { + for (int k = 0; k < n; k++) { + float vals[4]; + philox_normal4(seed, k, 0, vals); + out[k] = bf16_round ? f32_to_bf16_to_f32(vals[0]) : vals[0]; + } +} + +// Fill array with uniform [0, 1) drawn from Philox4x32-10. Matches +// PyTorch CUDA torch.rand kernels. +// +// Required by the multinomial sampler used during stochastic +// generation : torch.multinomial(probs, 1) decomposes mathematically +// as u ~ Uniform[0, 1) ; cdf = cumsum(probs) ; argmin{ i : cdf[i] >= u }. +// To stay byte-exact with the upstream Python pipeline, both sides +// must consume the same u from the same Philox state, hence we need a +// uniform draw, not a Box-Muller normal. +// +// Convention : +// key = seed +// subseq = subseq_start + k (one subsequence per element) +// ctr = (ctr_lo, 0, subseq_lo, subseq_hi) +// output = (r.x + 0.5) * 2^-32 in (0, 1) +// +// ctr_lo is the cumulative Philox block counter the caller advances +// across kernels. On the first call after manual_seed, ctr_lo = 0. +static inline void philox_uniform_fill(int64_t seed, int64_t subseq_start, uint32_t ctr_lo, float * out, int n) { + uint32_t slo = (uint32_t) seed; + uint32_t shi = (uint32_t) ((uint64_t) seed >> 32); + for (int k = 0; k < n; k++) { + uint64_t s = (uint64_t) (subseq_start + k); + Philox4 ctr = { ctr_lo, 0u, (uint32_t) s, (uint32_t) (s >> 32) }; + Philox4 r = philox4x32_10(ctr, slo, shi); + out[k] = ((float) r.x + 0.5f) * CURAND_2POW32_INV; + } +} diff --git a/src/pipeline-codec.cpp b/src/pipeline-codec.cpp new file mode 100644 index 0000000..6194c1a --- /dev/null +++ b/src/pipeline-codec.cpp @@ -0,0 +1,375 @@ +// pipeline-codec.cpp: load + decode for the Qwen3-TTS 12Hz codec. +// +// load chains the four module loaders (quantizer, transformer, upsample, +// DAC) and then loads the two pre_conv tensors into a dedicated wctx. +// decode builds the full forward graph in a per-call context, lets the +// scheduler allocate intermediates, uploads codes/positions/mask, runs +// graph_compute, and pulls the audio buffer back to host. + +#include "pipeline-codec.h" + +#include "causal-trans-conv.h" +#include "qt-error.h" + +#include +#include +#include +#include +#include + +bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair bp) { + pc->bp = bp; + pc->backend = bp.backend; + + if (!gf_load(&pc->gguf, gguf_path)) { + qt_log(QT_LOG_ERROR, "[Pipeline] failed to load %s", gguf_path); + return false; + } + + if (!qwen_quantizer_decoder_load(&pc->qdec, pc->gguf, pc->backend)) { + gf_close(&pc->gguf); + return false; + } + + if (!qwen_tokenizer_transformer_load(&pc->transformer, pc->gguf, pc->backend)) { + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + if (!qwen_upsample_stage_load(&pc->upsample, pc->gguf, pc->backend)) { + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + if (!qwen_dac_decoder_load(&pc->dac, pc->gguf, pc->backend)) { + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + // pre_conv: 2 tensors, dedicated wctx + { + WeightCtx wctx; + wctx_init(&wctx, 4); + pc->pre_conv_w = gf_load_tensor(&wctx, pc->gguf, "tok_dec.pre_conv.weight"); + pc->pre_conv_b = gf_load_tensor(&wctx, pc->gguf, "tok_dec.pre_conv.bias"); + if (!wctx_alloc(&wctx, pc->backend)) { + qt_log(QT_LOG_ERROR, "[Pipeline] pre_conv backend allocation failed"); + qwen_dac_decoder_free(&pc->dac); + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + pc->pre_conv_ctx = wctx.ctx; + pc->pre_conv_buf = wctx.buffer; + } + + if (!qwen_seanet_encoder_load(&pc->seanet, pc->gguf, pc->backend)) { + ggml_backend_buffer_free(pc->pre_conv_buf); + ggml_free(pc->pre_conv_ctx); + qwen_dac_decoder_free(&pc->dac); + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + if (!qwen_encoder_transformer_load(&pc->enc_transformer, pc->gguf, pc->backend)) { + qwen_seanet_encoder_free(&pc->seanet); + ggml_backend_buffer_free(pc->pre_conv_buf); + ggml_free(pc->pre_conv_ctx); + qwen_dac_decoder_free(&pc->dac); + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + if (!qwen_encoder_downsample_load(&pc->enc_downsample, pc->gguf, pc->backend)) { + qwen_encoder_transformer_free(&pc->enc_transformer); + qwen_seanet_encoder_free(&pc->seanet); + ggml_backend_buffer_free(pc->pre_conv_buf); + ggml_free(pc->pre_conv_ctx); + qwen_dac_decoder_free(&pc->dac); + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + if (!qwen_quantizer_encode_load(&pc->qenc, pc->gguf, pc->backend)) { + qwen_encoder_downsample_free(&pc->enc_downsample); + qwen_encoder_transformer_free(&pc->enc_transformer); + qwen_seanet_encoder_free(&pc->seanet); + ggml_backend_buffer_free(pc->pre_conv_buf); + ggml_free(pc->pre_conv_ctx); + qwen_dac_decoder_free(&pc->dac); + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + gf_close(&pc->gguf); + return false; + } + + pc->sched = backend_sched_new(bp, 4096); + + qt_log(QT_LOG_INFO, "[Pipeline] Ready: hop %d samples @ %d Hz mono, %d codebooks @ 12.5 Hz", + QWEN_TOKENIZER_HOP_LENGTH, QWEN_TOKENIZER_SAMPLE_RATE, QWEN_TOKENIZER_NUM_CODEBOOKS); + return true; +} + +std::vector pipeline_codec_decode(PipelineCodec * pc, const int32_t * codes, int K, int T) { + if (K != QWEN_TOKENIZER_NUM_CODEBOOKS) { + qt_log(QT_LOG_ERROR, "[Pipeline] codes have %d codebooks, expected %d", K, QWEN_TOKENIZER_NUM_CODEBOOKS); + return {}; + } + if (T <= 0) { + qt_log(QT_LOG_ERROR, "[Pipeline] T must be > 0 (got %d)", T); + return {}; + } + + // Per-call graph context: tensor descriptors only, allocation is + // delegated to the scheduler. + const int n_max_nodes = 4096; + const size_t graph_ctx_size = + ggml_tensor_overhead() * (size_t) n_max_nodes + ggml_graph_overhead_custom((size_t) n_max_nodes, false); + + struct ggml_init_params gp = { graph_ctx_size, NULL, /*no_alloc=*/true }; + struct ggml_context * gctx = ggml_init(gp); + if (!gctx) { + qt_log(QT_LOG_ERROR, "[Pipeline] ggml_init failed for graph ctx"); + return {}; + } + + // Inputs: codes [T, K] i32, positions [T] i32, mask [T, T] f32. + struct ggml_tensor * codes_in = ggml_new_tensor_2d(gctx, GGML_TYPE_I32, T, K); + ggml_set_name(codes_in, "codes_in"); + ggml_set_input(codes_in); + + struct ggml_tensor * positions = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T); + ggml_set_name(positions, "positions"); + ggml_set_input(positions); + + struct ggml_tensor * mask = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, T, T); + ggml_set_name(mask, "mask"); + ggml_set_input(mask); + + // Build forward graph. Layout transitions are explicit ggml_cont(ggml_transpose(...)) + // calls: 3 transposes total at the natural module boundaries. + struct ggml_tensor * h = qwen_quantizer_decode(gctx, &pc->qdec, codes_in); // [512, T] C-first + h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [T, 512] T-first + h = qwen_causal_conv1d(gctx, pc->pre_conv_w, pc->pre_conv_b, h, 3, 1); // [T, 1024] T-first + h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [1024, T] C-first + h = qwen_tokenizer_transformer_forward(gctx, &pc->transformer, h, positions, mask); // [1024, T] + h = ggml_cont(gctx, ggml_transpose(gctx, h)); // [T, 1024] T-first + h = qwen_upsample_stage_forward(gctx, &pc->upsample, h); // [T*4, 1024] + h = qwen_dac_decoder_forward(gctx, &pc->dac, h); // [T*1920, 1] + h = ggml_clamp(gctx, h, -1.0f, 1.0f); + + ggml_set_name(h, "audio_out"); + ggml_set_output(h); + + struct ggml_cgraph * graph = ggml_new_graph_custom(gctx, n_max_nodes, false); + ggml_build_forward_expand(graph, h); + + if (!ggml_backend_sched_alloc_graph(pc->sched, graph)) { + qt_log(QT_LOG_ERROR, "[Pipeline] sched_alloc_graph failed"); + ggml_backend_sched_reset(pc->sched); + ggml_free(gctx); + return {}; + } + + // Upload inputs + ggml_backend_tensor_set(codes_in, codes, 0, (size_t) T * (size_t) K * sizeof(int32_t)); + + std::vector pos_buf; + qwen_build_positions(T, pos_buf); + ggml_backend_tensor_set(positions, pos_buf.data(), 0, pos_buf.size() * sizeof(int32_t)); + + std::vector mask_buf; + qwen_build_causal_sliding_mask(T, pc->transformer.sliding_window, mask_buf); + ggml_backend_tensor_set(mask, mask_buf.data(), 0, mask_buf.size() * sizeof(float)); + + // Compute + enum ggml_status st = ggml_backend_sched_graph_compute(pc->sched, graph); + if (st != GGML_STATUS_SUCCESS) { + qt_log(QT_LOG_ERROR, "[Pipeline] graph_compute status=%d", (int) st); + ggml_backend_sched_reset(pc->sched); + ggml_free(gctx); + return {}; + } + + // Fetch audio output + const int n_samples = T * QWEN_TOKENIZER_HOP_LENGTH; + std::vector audio((size_t) n_samples); + ggml_backend_tensor_get(h, audio.data(), 0, (size_t) n_samples * sizeof(float)); + + ggml_backend_sched_reset(pc->sched); + ggml_free(gctx); + return audio; +} + +std::vector pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples) { + if (n_samples <= 0 || (n_samples % QWEN_TOKENIZER_HOP_LENGTH) != 0) { + qt_log(QT_LOG_ERROR, "[Pipeline] n_samples must be a positive multiple of %d (got %d)", + QWEN_TOKENIZER_HOP_LENGTH, n_samples); + return {}; + } + int T = n_samples / QWEN_TOKENIZER_HOP_LENGTH; + + // Lazy-load CPU mirror of the RVQ encode codebooks on first call. + if (!pc->qenc_host_ready) { + qwen_quantizer_encode_host_load(&pc->qenc_sem_host, pc->qenc.semantic, pc->qenc.codebook_size, + pc->qenc.codebook_dim, pc->qenc.hidden_size); + qwen_quantizer_encode_host_load(&pc->qenc_aco_host, pc->qenc.acoustic, pc->qenc.codebook_size, + pc->qenc.codebook_dim, pc->qenc.hidden_size); + pc->qenc_host_ready = true; + } + + const int n_max_nodes = 4096; + const size_t graph_ctx_size = + ggml_tensor_overhead() * (size_t) n_max_nodes + ggml_graph_overhead_custom((size_t) n_max_nodes, false); + + struct ggml_init_params gp = { graph_ctx_size, NULL, /*no_alloc=*/true }; + struct ggml_context * gctx = ggml_init(gp); + if (!gctx) { + qt_log(QT_LOG_ERROR, "[Pipeline] ggml_init failed for encode graph ctx"); + return {}; + } + + // SEANet input shape: [T_audio, 1] f32 T-first (mono waveform). + struct ggml_tensor * audio_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, n_samples, 1); + ggml_set_name(audio_in, "audio_in"); + ggml_set_input(audio_in); + + // Encoder transformer mask is built on the post-SEANet T = n_samples / 960. + int T_emb = n_samples / 960; + struct ggml_tensor * positions = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T_emb); + ggml_set_name(positions, "enc_positions"); + ggml_set_input(positions); + + struct ggml_tensor * mask = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, T_emb, T_emb); + ggml_set_name(mask, "enc_mask"); + ggml_set_input(mask); + + // Forward chain. + struct ggml_tensor * h_seanet = qwen_seanet_encoder_forward(gctx, &pc->seanet, audio_in); // [T_emb, 512] + struct ggml_tensor * h = ggml_cont(gctx, ggml_transpose(gctx, h_seanet)); // [512, T_emb] + struct ggml_tensor * h_et = + qwen_encoder_transformer_forward(gctx, &pc->enc_transformer, h, positions, mask); // [512, T_emb] + h = ggml_cont(gctx, ggml_transpose(gctx, h_et)); // [T_emb, 512] + h = qwen_encoder_downsample_forward(gctx, &pc->enc_downsample, h); // [T, 512] + + // The CPU RVQ encode loop expects the hidden buffer as [T, hidden] + // row-major (hidden fast in memory). The downsample output ne=(T, 512) + // walks T fast in ggml memory, which is [hidden, T] in numpy terms. + // Transpose to get the buffer layout we want once read back to host. + h = ggml_cont(gctx, ggml_transpose(gctx, h)); // ne=(512, T) + + const char * dump_dir = getenv("QWENTTS_DEBUG_DUMP"); + if (dump_dir) { + ggml_set_output(h_seanet); + ggml_set_name(h_seanet, "seanet_out"); + ggml_set_output(h_et); + ggml_set_name(h_et, "enc_transformer_out"); + } + + ggml_set_name(h, "enc_hidden_out"); + ggml_set_output(h); + + struct ggml_cgraph * graph = ggml_new_graph_custom(gctx, n_max_nodes, false); + ggml_build_forward_expand(graph, h); + + if (!ggml_backend_sched_alloc_graph(pc->sched, graph)) { + qt_log(QT_LOG_ERROR, "[Pipeline] encode sched_alloc_graph failed"); + ggml_backend_sched_reset(pc->sched); + ggml_free(gctx); + return {}; + } + + ggml_backend_tensor_set(audio_in, audio, 0, (size_t) n_samples * sizeof(float)); + + std::vector pos_buf; + qwen_encoder_build_positions(T_emb, pos_buf); + ggml_backend_tensor_set(positions, pos_buf.data(), 0, pos_buf.size() * sizeof(int32_t)); + + std::vector mask_buf; + qwen_encoder_build_causal_sliding_mask(T_emb, pc->enc_transformer.sliding_window, mask_buf); + ggml_backend_tensor_set(mask, mask_buf.data(), 0, mask_buf.size() * sizeof(float)); + + enum ggml_status st = ggml_backend_sched_graph_compute(pc->sched, graph); + if (st != GGML_STATUS_SUCCESS) { + qt_log(QT_LOG_ERROR, "[Pipeline] encode graph_compute status=%d", (int) st); + ggml_backend_sched_reset(pc->sched); + ggml_free(gctx); + return {}; + } + + if (dump_dir) { + auto dump = [&](const char * fname, struct ggml_tensor * t) { + size_t n = ggml_nelements(t); + std::vector buf(n); + ggml_backend_tensor_get(t, buf.data(), 0, n * sizeof(float)); + char path[512]; + snprintf(path, sizeof(path), "%s/%s.f32", dump_dir, fname); + FILE * f = fopen(path, "wb"); + if (f) { + fwrite(buf.data(), sizeof(float), n, f); + fclose(f); + qt_log(QT_LOG_INFO, "[Pipeline] Dumped %s: %zu floats, ne=(%lld, %lld, %lld, %lld)", path, n, + (long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], (long long) t->ne[3]); + } + }; + dump("seanet_out", h_seanet); + dump("enc_transformer_out", h_et); + dump("enc_downsample_out", h); + } + + // Read back the post-downsample hidden buffer for CPU-side RVQ encode. + // Layout in ggml is [T, hidden] with T on ne[0]. The contiguous memory + // walks T fast, hidden slow, which matches the `[T, hidden] row-major + // index = t*hidden + c` convention expected by qwen_quantizer_encode_cpu. + std::vector hidden_host((size_t) T * (size_t) pc->qenc.hidden_size); + ggml_backend_tensor_get(h, hidden_host.data(), 0, hidden_host.size() * sizeof(float)); + + ggml_backend_sched_reset(pc->sched); + ggml_free(gctx); + + return qwen_quantizer_encode_cpu(&pc->qenc_sem_host, &pc->qenc_aco_host, hidden_host.data(), T); +} + +void pipeline_codec_free(PipelineCodec * pc) { + if (pc->sched) { + ggml_backend_sched_free(pc->sched); + pc->sched = NULL; + } + qwen_quantizer_encode_free(&pc->qenc); + qwen_encoder_downsample_free(&pc->enc_downsample); + qwen_encoder_transformer_free(&pc->enc_transformer); + qwen_seanet_encoder_free(&pc->seanet); + if (pc->pre_conv_buf) { + ggml_backend_buffer_free(pc->pre_conv_buf); + pc->pre_conv_buf = NULL; + } + if (pc->pre_conv_ctx) { + ggml_free(pc->pre_conv_ctx); + pc->pre_conv_ctx = NULL; + } + qwen_dac_decoder_free(&pc->dac); + qwen_upsample_stage_free(&pc->upsample); + qwen_tokenizer_transformer_free(&pc->transformer); + qwen_quantizer_decoder_free(&pc->qdec); + if (pc->gguf.gguf) { + gf_close(&pc->gguf); + } +} diff --git a/src/pipeline-codec.h b/src/pipeline-codec.h new file mode 100644 index 0000000..eda13ae --- /dev/null +++ b/src/pipeline-codec.h @@ -0,0 +1,97 @@ +#pragma once +// pipeline-codec.h: codec decode pipeline for the Qwen3-TTS 12Hz tokenizer. +// Loads a codec GGUF (quantizer + pre_conv + pre_transformer + upsample + +// DAC), holds every weight on the backend, and exposes a one-shot decode: +// +// codes [num_codebooks, T] i32 -> audio [T * 1920] f32 mono 24 kHz +// +// Layout flow inside the graph: +// codes [T, K] i32 T-first +// v quantizer.decode +// hidden [512, T] f32 C-first +// v transpose to T-first +// v pre_conv (causal Conv1d k=3, 512 -> 1024) +// hidden [T, 1024] f32 T-first +// v transpose to C-first +// v pre_transformer (8 layers Qwen3, sliding window 72 causal) +// hidden [1024, T] f32 C-first +// v transpose to T-first +// v upsample stage (4x) +// hidden [T*4, 1024] f32 T-first +// v DAC decoder (480x) +// audio [T*1920, 1] f32 T-first +// v ggml_clamp(-1, 1) +// audio_out [T*1920, 1] f32 T-first + +#include "backend.h" +#include "convnext-block.h" +#include "dac-decoder-v2.h" +#include "encoder-downsample.h" +#include "encoder-transformer.h" +#include "ggml-backend.h" +#include "gguf-weights.h" +#include "quantizer-decode.h" +#include "quantizer-encode.h" +#include "seanet-encoder.h" +#include "tokenizer-transformer.h" +#include "weight-ctx.h" + +#include +#include + +#define QWEN_TOKENIZER_HOP_LENGTH 1920 +#define QWEN_TOKENIZER_SAMPLE_RATE 24000 +#define QWEN_TOKENIZER_NUM_CODEBOOKS 16 +#define QWEN_TOKENIZER_CODE_BITS 11 + +struct PipelineCodec { + GGUFModel gguf; + + // Decode side modules + QwenQuantizerDecoder qdec; + QwenTokenizerTransformer transformer; + QwenUpsampleStage upsample; + QwenDACDecoder dac; + + // pre_conv: causal Conv1d k=3, 512 -> 1024. Loaded into a dedicated + // weight ctx because it is the only module that does not own one. + struct ggml_tensor * pre_conv_w; // [3, 512, 1024] f32 + struct ggml_tensor * pre_conv_b; // [1024] f32 + struct ggml_context * pre_conv_ctx; + ggml_backend_buffer_t pre_conv_buf; + + // Encode side modules + QwenSEANetEncoder seanet; + QwenEncoderTransformer enc_transformer; + QwenEncoderDownsample enc_downsample; + QwenQuantizerEncode qenc; + + // CPU mirror of the RVQ encode side, lazy-loaded on first encode call. + QwenQuantizerEncodeHost qenc_sem_host; + QwenQuantizerEncodeHost qenc_aco_host; + bool qenc_host_ready; + + BackendPair bp; + ggml_backend_t backend; + ggml_backend_sched_t sched; +}; + +// Open the GGUF, load every module on the backend, build the scheduler. +// On failure leaves the struct in a clean state and returns false. +bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair bp); + +// Decode RVQ codes into a 24 kHz mono waveform. +// codes: flat int32 buffer, [K, T] row-major (T fastest). +// Returns audio of length T * QWEN_TOKENIZER_HOP_LENGTH, empty on failure. +std::vector pipeline_codec_decode(PipelineCodec * pc, const int32_t * codes, int K, int T); + +// Encode a 24 kHz mono waveform into RVQ codes. +// audio : [n_samples] f32 mono 24 kHz. Must be a multiple of +// QWEN_TOKENIZER_HOP_LENGTH (1920); the caller is expected +// to pad with zeros if needed. +// Returns codes flat as [K, T] row-major, K = QWEN_TOKENIZER_NUM_CODEBOOKS, +// T = n_samples / 1920. Empty on failure. +std::vector pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples); + +// Free every backend buffer and ggml context. Safe to call on a zeroed struct. +void pipeline_codec_free(PipelineCodec * pc); diff --git a/src/pipeline-tts.cpp b/src/pipeline-tts.cpp new file mode 100644 index 0000000..1b8a442 --- /dev/null +++ b/src/pipeline-tts.cpp @@ -0,0 +1,502 @@ +// pipeline-tts.cpp : load and verify both GGUF files (talker + codec) +// onto the same shared backend, parse all metadata into typed structs, +// and provide a structured load-time summary for --load-only mode. + +#include "pipeline-tts.h" + +#include "audio-io.h" +#include "bpe.h" +#include "code-predictor-forward.h" +#include "debug.h" +#include "ggml.h" +#include "pipeline-codec.h" +#include "prompt-builder.h" +#include "qt-error.h" +#include "sampling.h" +#include "speaker-encoder-extract.h" +#include "talker-forward.h" + +#include +#include +#include + +static void parse_codec_specials(const GGUFModel & gf, CodecSpecials & cs) { + cs.pad_id = (int) gf_get_u32(gf, "qwen3-tts.codec.pad_id"); + cs.bos_id = (int) gf_get_u32(gf, "qwen3-tts.codec.bos_id"); + cs.eos_id = (int) gf_get_u32(gf, "qwen3-tts.codec.eos_id"); + cs.think_id = (int) gf_get_u32(gf, "qwen3-tts.codec.think_id"); + cs.nothink_id = (int) gf_get_u32(gf, "qwen3-tts.codec.nothink_id"); + cs.think_bos_id = (int) gf_get_u32(gf, "qwen3-tts.codec.think_bos_id"); + cs.think_eos_id = (int) gf_get_u32(gf, "qwen3-tts.codec.think_eos_id"); +} + +static void parse_text_specials(const GGUFModel & gf, TextSpecials & ts) { + ts.im_start_id = (int) gf_get_u32(gf, "qwen3-tts.text.im_start_id"); + ts.im_end_id = (int) gf_get_u32(gf, "qwen3-tts.text.im_end_id"); + ts.tts_pad_id = (int) gf_get_u32(gf, "qwen3-tts.text.tts_pad_id"); + ts.tts_bos_id = (int) gf_get_u32(gf, "qwen3-tts.text.tts_bos_id"); + ts.tts_eos_id = (int) gf_get_u32(gf, "qwen3-tts.text.tts_eos_id"); +} + +static void parse_languages(const GGUFModel & gf, std::vector & out) { + int64_t name_idx = gguf_find_key(gf.gguf, "qwen3-tts.codec.language_names"); + int64_t id_idx = gguf_find_key(gf.gguf, "qwen3-tts.codec.language_ids"); + if (name_idx < 0 || id_idx < 0) { + return; + } + size_t n_names = gguf_get_arr_n(gf.gguf, name_idx); + size_t n_ids = gguf_get_arr_n(gf.gguf, id_idx); + if (n_names != n_ids) { + fprintf(stderr, "[Pipeline] WARNING: language arrays size mismatch (names=%zu, ids=%zu)\n", n_names, n_ids); + return; + } + const uint32_t * ids = (const uint32_t *) gguf_get_arr_data(gf.gguf, id_idx); + out.reserve(n_names); + for (size_t i = 0; i < n_names; i++) { + LanguageEntry e; + e.name = gguf_get_arr_str(gf.gguf, name_idx, i); + e.id = (int) ids[i]; + out.push_back(e); + } +} + +// Parse the speaker table for CustomVoice variants. Three parallel arrays +// produced by convert.py : speaker_names, speaker_ids, speaker_dialects. +// Empty dialect string means the speaker keeps the user supplied language. +// Skipped silently when the GGUF carries no speaker table (Base / VoiceDesign). +static void parse_speakers(const GGUFModel & gf, std::vector & out) { + int64_t name_idx = gguf_find_key(gf.gguf, "qwen3-tts.codec.speaker_names"); + int64_t id_idx = gguf_find_key(gf.gguf, "qwen3-tts.codec.speaker_ids"); + int64_t dialect_idx = gguf_find_key(gf.gguf, "qwen3-tts.codec.speaker_dialects"); + if (name_idx < 0 || id_idx < 0 || dialect_idx < 0) { + return; + } + size_t n_names = gguf_get_arr_n(gf.gguf, name_idx); + size_t n_ids = gguf_get_arr_n(gf.gguf, id_idx); + size_t n_dialects = gguf_get_arr_n(gf.gguf, dialect_idx); + if (n_names != n_ids || n_names != n_dialects) { + fprintf(stderr, "[Pipeline] WARNING: speaker arrays size mismatch (names=%zu, ids=%zu, dialects=%zu)\n", + n_names, n_ids, n_dialects); + return; + } + const uint32_t * ids = (const uint32_t *) gguf_get_arr_data(gf.gguf, id_idx); + out.reserve(n_names); + for (size_t i = 0; i < n_names; i++) { + SpeakerEntry e; + e.name = gguf_get_arr_str(gf.gguf, name_idx, i); + e.id = (int) ids[i]; + e.dialect = gguf_get_arr_str(gf.gguf, dialect_idx, i); + out.push_back(e); + } +} + +static void parse_generation_defaults(const GGUFModel & gf, GenerationDefaults & g) { + g.do_sample = gf_get_bool(gf, "generation.do_sample"); + g.top_k = (int) gf_get_u32(gf, "generation.top_k"); + g.top_p = gf_get_f32(gf, "generation.top_p"); + g.temperature = gf_get_f32(gf, "generation.temperature"); + g.repetition_penalty = gf_get_f32(gf, "generation.repetition_penalty"); + g.subtalker_do_sample = gf_get_bool(gf, "generation.subtalker_do_sample"); + g.subtalker_top_k = (int) gf_get_u32(gf, "generation.subtalker_top_k"); + g.subtalker_top_p = gf_get_f32(gf, "generation.subtalker_top_p"); + g.subtalker_temperature = gf_get_f32(gf, "generation.subtalker_temperature"); + g.max_new_tokens = (int) gf_get_u32(gf, "generation.max_new_tokens"); +} + +bool pipeline_tts_load(PipelineTTS * pt, const char * talker_gguf_path, const char * codec_gguf_path, BackendPair bp) { + pt->bp = bp; + pt->backend = bp.backend; + pt->sched = NULL; + pt->has_speaker_encoder = false; + + if (!gf_load(&pt->gguf_talker, talker_gguf_path)) { + qt_log(QT_LOG_ERROR, "[Pipeline] failed to load talker GGUF: %s", talker_gguf_path); + return false; + } + + const char * arch = gf_get_str(pt->gguf_talker, "general.architecture"); + if (!arch || std::strcmp(arch, "qwen3-tts") != 0) { + qt_log(QT_LOG_ERROR, "[Pipeline] talker GGUF has wrong architecture '%s', expected 'qwen3-tts'", + arch ? arch : ""); + gf_close(&pt->gguf_talker); + return false; + } + + pt->tokenizer_type = gf_get_str(pt->gguf_talker, "qwen3-tts.tokenizer_type"); + pt->model_size = gf_get_str(pt->gguf_talker, "qwen3-tts.model_size"); + pt->model_type = gf_get_str(pt->gguf_talker, "qwen3-tts.model_type"); + pt->num_code_groups = (int) gf_get_u32(pt->gguf_talker, "qwen3-tts.num_code_groups"); + + parse_codec_specials(pt->gguf_talker, pt->codec_specials); + parse_text_specials(pt->gguf_talker, pt->text_specials); + parse_languages(pt->gguf_talker, pt->languages); + parse_speakers(pt->gguf_talker, pt->speakers); + parse_generation_defaults(pt->gguf_talker, pt->gen_defaults); + + if (!talker_weights_load(&pt->talker, pt->gguf_talker, pt->backend)) { + gf_close(&pt->gguf_talker); + return false; + } + + if (!code_predictor_weights_load(&pt->code_predictor, pt->gguf_talker, pt->backend)) { + talker_weights_free(&pt->talker); + gf_close(&pt->gguf_talker); + return false; + } + + // Speaker encoder is only present in Base checkpoints. Treat absence + // as a soft condition : voice clone path stays disabled, base-direct + // synthesis still works. + if (pt->model_type == "base") { + if (!speaker_encoder_weights_load(&pt->speaker_encoder, pt->gguf_talker, pt->backend)) { + code_predictor_weights_free(&pt->code_predictor); + talker_weights_free(&pt->talker); + gf_close(&pt->gguf_talker); + return false; + } + pt->has_speaker_encoder = (pt->speaker_encoder.weight_buf != NULL); + } + + if (!pipeline_codec_load(&pt->codec, codec_gguf_path, bp)) { + if (pt->has_speaker_encoder) { + speaker_encoder_weights_free(&pt->speaker_encoder); + } + code_predictor_weights_free(&pt->code_predictor); + talker_weights_free(&pt->talker); + gf_close(&pt->gguf_talker); + return false; + } + + // Scheduler shared by talker_forward_prefill and code_predictor_step. + // Routes ops the GPU backend cannot run (typical case : K-quant + // get_rows on CUDA) to the CPU backend. 4096 nodes covers the 28L + // Qwen3 talker graph (~32 ops per layer + heads + dump taps) with + // headroom ; the 5L code predictor uses a fraction of that. + pt->sched = backend_sched_new(bp, 4096); + if (!pt->sched) { + pipeline_codec_free(&pt->codec); + if (pt->has_speaker_encoder) { + speaker_encoder_weights_free(&pt->speaker_encoder); + } + code_predictor_weights_free(&pt->code_predictor); + talker_weights_free(&pt->talker); + gf_close(&pt->gguf_talker); + return false; + } + + qt_log(QT_LOG_INFO, + "[Pipeline] Loaded: arch=%s variant=%s tokenizer=%s codebooks=%d speaker_encoder=%s speakers=%zu", + pt->model_size.c_str(), pt->model_type.c_str(), pt->tokenizer_type.c_str(), pt->num_code_groups, + pt->has_speaker_encoder ? "loaded" : "absent", pt->speakers.size()); + return true; +} + +void pipeline_tts_free(PipelineTTS * pt) { + if (pt->sched) { + ggml_backend_sched_free(pt->sched); + pt->sched = NULL; + } + pipeline_codec_free(&pt->codec); + if (pt->has_speaker_encoder) { + speaker_encoder_weights_free(&pt->speaker_encoder); + } + code_predictor_weights_free(&pt->code_predictor); + talker_weights_free(&pt->talker); + gf_close(&pt->gguf_talker); + pt->backend = NULL; + pt->bp = {}; + pt->has_speaker_encoder = false; +} + +// Pull one row of an embedding table directly from the GGUF mmap. Used +// in the generation loop to assemble the next-token embedding (sum of 16 +// codebook embeddings) without paying for a backend round-trip per row. +static void embed_row_from_gguf(const GGUFModel & gf, const char * tensor_name, int row_id, int hidden, float * dst) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name); + if (!src) { + qt_throw("[Pipeline] tensor not found in GGUF: %s", tensor_name); + } + const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name); + if (!base) { + qt_throw("[Pipeline] tensor data missing in GGUF: %s", tensor_name); + } + const size_t row_bytes = ggml_row_size(src->type, hidden); + const void * row = base + (size_t) row_id * row_bytes; + if (src->type == GGML_TYPE_F32) { + std::memcpy(dst, row, (size_t) hidden * sizeof(float)); + return; + } + const struct ggml_type_traits * tt = ggml_get_type_traits(src->type); + if (!tt || !tt->to_float) { + qt_throw("[Pipeline] unsupported codec_embedding dtype %d for %s", (int) src->type, tensor_name); + } + tt->to_float(row, dst, hidden); +} + +bool pipeline_tts_synthesize(PipelineTTS * pt, + BPETokenizer * tok, + const PipelineTTSSynthesizeParams & params, + PipelineTTSSynthesizeOutput * out) { + out->audio.clear(); + out->sample_rate = QWEN_TOKENIZER_SAMPLE_RATE; + + PromptBuilderOutput prompt; + const std::string instruct = params.instruct ? params.instruct : ""; + const std::string speaker = params.speaker ? params.speaker : ""; + const std::string ref_text = params.ref_text ? params.ref_text : ""; + + // Voice clone mode A : if ref_audio is given, run the speaker + // encoder on the WAV and feed the resulting embedding straight into + // the prompt builder. Mutually exclusive with --speaker. + std::vector ref_spk_emb; + const float * ref_spk_emb_ptr = NULL; + if (params.ref_audio && params.ref_audio[0]) { + if (!pt->has_speaker_encoder) { + fprintf(stderr, + "[Pipeline] FATAL: --ref-audio requires a model with a loaded speaker encoder (Base only)\n"); + return false; + } + if (!speaker_encoder_extract(&pt->speaker_encoder, pt->sched, params.ref_audio, ref_spk_emb)) { + return false; + } + if ((int) ref_spk_emb.size() != pt->talker.hidden_size) { + fprintf(stderr, "[Pipeline] FATAL: speaker embedding size %zu mismatches talker hidden %d\n", + ref_spk_emb.size(), pt->talker.hidden_size); + return false; + } + ref_spk_emb_ptr = ref_spk_emb.data(); + } + + // Voice clone mode B : if ref_text is also given, encode the + // reference audio into 16 codebook indices via the codec encoder. + // Layout returned by pipeline_codec_encode is [num_codebooks, T_codec] + // row major, matching what the prompt builder expects for the ICL + // sum loop. + std::vector ref_codes; + int ref_codes_T = 0; + if (!ref_text.empty()) { + if (!params.ref_audio || !params.ref_audio[0]) { + fprintf(stderr, "[Pipeline] FATAL: --ref-text requires --ref-audio\n"); + return false; + } + // audio_read_mono returns f32 mono at the codec sample rate. The + // codec hop is 1920 samples at 24 kHz so n_samples must be a + // multiple of 1920. Truncate to the nearest hop boundary. + int T_codec_audio = 0; + float * raw = audio_read_mono(params.ref_audio, QWEN_TOKENIZER_SAMPLE_RATE, &T_codec_audio); + if (!raw || T_codec_audio < QWEN_TOKENIZER_HOP_LENGTH) { + fprintf(stderr, "[Pipeline] FATAL: cannot read ref_audio for ICL '%s'\n", params.ref_audio); + if (raw) { + std::free(raw); + } + return false; + } + int aligned_T = (T_codec_audio / QWEN_TOKENIZER_HOP_LENGTH) * QWEN_TOKENIZER_HOP_LENGTH; + ref_codes = pipeline_codec_encode(&pt->codec, raw, aligned_T); + std::free(raw); + if (ref_codes.empty()) { + fprintf(stderr, "[Pipeline] FATAL: pipeline_codec_encode returned empty codes\n"); + return false; + } + ref_codes_T = (int) ref_codes.size() / pt->num_code_groups; + fprintf(stderr, "[Pipeline] ICL ref_codes: %d frames at 12.5 Hz (%d audio samples)\n", ref_codes_T, aligned_T); + } + + if (!prompt_builder_build(pt, tok, params.text, params.lang, instruct, speaker, ref_spk_emb_ptr, ref_text, + ref_codes_T > 0 ? ref_codes.data() : NULL, ref_codes_T, &prompt)) { + return false; + } + + if (params.dump_dir) { + DebugDumper d; + debug_init(&d, params.dump_dir); + std::vector ids32(prompt.prompt_ids.begin(), prompt.prompt_ids.end()); + int n_ids = (int) ids32.size(); + debug_dump_i32_as_f32(&d, "prompt-ids", ids32.data(), &n_ids, 1); + debug_dump_2d(&d, "talker-input-embed", prompt.input_embed.data(), prompt.T_ctx, prompt.hidden); + debug_dump_2d(&d, "trailing-text-hidden", prompt.trailing_text_hidden.data(), prompt.T_trailing, prompt.hidden); + debug_dump_1d(&d, "tts-pad-embed", prompt.tts_pad_embed.data(), prompt.hidden); + } + + // Generation loop : at each step we recompute the full Talker prefix + // (no KV cache yet) over the prompt prefix concatenated with all the + // next-token embeddings produced so far, sample c0, run the code + // predictor for the 15 acoustic codes, build the next-token + // embedding by summing the 16 codebook embeddings and the matching + // trailing-text overlay, and append it to the running context. We + // stop on codec_eos or when max_new_tokens is reached. + const int hidden = prompt.hidden; + const int codec_eos_id = pt->codec_specials.eos_id; + const int num_codebooks = pt->num_code_groups; + const int talker_vocab = pt->talker.vocab_size; + + // Greedy collapses to temperature <= 0 in sample_top_k_p. + float talker_T = params.do_sample ? params.temperature : 0.0f; + float subtk_T = params.subtalker_do_sample ? params.subtalker_temperature : 0.0f; + float talker_rp = params.repetition_penalty; + + std::vector ctx_embed; + ctx_embed.reserve((size_t) (prompt.T_ctx + params.max_new_tokens) * (size_t) hidden); + ctx_embed.insert(ctx_embed.end(), prompt.input_embed.begin(), prompt.input_embed.end()); + int T_ctx = prompt.T_ctx; + + std::vector> all_codes; + all_codes.reserve((size_t) params.max_new_tokens); + + // c0 codes already emitted, fed to repetition penalty. + std::vector talker_history; + talker_history.reserve((size_t) params.max_new_tokens); + + // Global Philox subsequence counter advances once per primitive + // sample (one for c0 of each step, then 15 for the predictor codes). + int64_t subseq_counter = 0; + + for (int step = 0; step < params.max_new_tokens; step++) { + TalkerForwardOutput fw; + const char * step_dump = (params.dump_dir && step == 0) ? params.dump_dir : NULL; + if (!talker_forward_prefill(&pt->talker, pt->sched, ctx_embed.data(), T_ctx, step_dump, &fw)) { + return false; + } + + // Bisection dump : the talker hidden_last at step 1 is the input + // the code predictor consumes after consuming the next-emb of + // step 0. Pairing it byte for byte with the Python hook tells us + // whether the next-emb composition + talker re-prefill round + // trip is bit exact end to end. + if (params.dump_dir && step == 1) { + DebugDumper d; + debug_init(&d, params.dump_dir); + debug_dump_1d(&d, "talker-hidden-step1", fw.hidden_last.data(), hidden); + } + + // Apply codec suppression : forbid [vocab - 1024, vocab) except + // codec_eos. Then run the upstream sampling chain. + apply_suppress(fw.logits_last.data(), talker_vocab, talker_vocab - 1024, talker_vocab, codec_eos_id); + float u_c0 = 0.0f; + int c0 = sample_top_k_p(fw.logits_last.data(), talker_vocab, talker_T, params.top_k, params.top_p, talker_rp, + talker_history.data(), (int) talker_history.size(), params.seed, subseq_counter, &u_c0); + subseq_counter++; + if (c0 < 0) { + qt_log(QT_LOG_ERROR, "[Pipeline] c0 sample returned no candidate"); + return false; + } + + // Trace the first 32 samples unconditionally so [Sample] lines + // up with [Sample-PY] / [Sample-CP] across the 16 codes of step + // 0 and step 1 the Python harness emits. + if ((subseq_counter - 1) < 32) { + fprintf(stderr, "[Sample] step=%d c0=%d u=%.10f subseq=%lld\n", step, c0, (double) u_c0, + (long long) (subseq_counter - 1)); + } + + if (c0 == codec_eos_id) { + qt_log(QT_LOG_INFO, "[Pipeline] EOS at step %d, stopping", step); + break; + } + + CodePredictorOutput cp; + const char * cp_dump = (params.dump_dir && step == 0) ? params.dump_dir : NULL; + if (!code_predictor_step(&pt->talker, &pt->code_predictor, pt->sched, fw.hidden_last.data(), c0, subtk_T, + params.subtalker_top_k, params.subtalker_top_p, params.seed, subseq_counter - 1, + cp_dump, &cp)) { + return false; + } + // Predictor consumed (num_codebooks - 1) subsequences after the + // c0 one (subseq_base + 1 .. subseq_base + 15). + subseq_counter += (num_codebooks - 1); + + all_codes.push_back(cp.codes); + talker_history.push_back(c0); + + // Build next-token embedding : sum of 16 codebook embeddings. + // codebook 0 uses talker.codec_embedding, the 15 acoustic + // codebooks use the predictor's private embedding tables. + std::vector next_emb((size_t) hidden, 0.0f); + std::vector tmp((size_t) hidden); + + embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", c0, hidden, tmp.data()); + for (int i = 0; i < hidden; i++) { + next_emb[(size_t) i] += tmp[(size_t) i]; + } + for (int g = 0; g < num_codebooks - 1; g++) { + int cg = cp.codes[(size_t) (g + 1)]; + char name[64]; + snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g); + embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data()); + for (int i = 0; i < hidden; i++) { + next_emb[(size_t) i] += tmp[(size_t) i]; + } + } + + // Trailing text overlay : while we still have utterance text + // hiddens to consume, add the next one ; otherwise add the + // tts_pad embedding. + const float * overlay = (step < prompt.T_trailing) ? + prompt.trailing_text_hidden.data() + (size_t) step * (size_t) hidden : + prompt.tts_pad_embed.data(); + for (int i = 0; i < hidden; i++) { + next_emb[(size_t) i] += overlay[(size_t) i]; + } + + ctx_embed.insert(ctx_embed.end(), next_emb.begin(), next_emb.end()); + T_ctx++; + + // Bisection dump : the next-token embedding produced at step 0 + // is the only thing controlling the talker forward at step 1, so + // matching it bit-exact against Python pinpoints any drift in + // the codebook embedding sums or the trailing text overlay. + if (params.dump_dir && step == 0) { + DebugDumper d; + debug_init(&d, params.dump_dir); + debug_dump_1d(&d, "next-emb-step0", next_emb.data(), hidden); + } + + if (((step + 1) % 8) == 0) { + qt_log(QT_LOG_INFO, "[Pipeline] Generated %d frames", step + 1); + } + } + + qt_log(QT_LOG_INFO, "[Pipeline] Generation done : %zu frames", all_codes.size()); + + if (params.dump_dir && !all_codes.empty()) { + DebugDumper d; + debug_init(&d, params.dump_dir); + int T_frames = (int) all_codes.size(); + std::vector flat((size_t) T_frames * (size_t) num_codebooks); + for (int t = 0; t < T_frames; t++) { + for (int k = 0; k < num_codebooks; k++) { + flat[(size_t) t * (size_t) num_codebooks + (size_t) k] = all_codes[(size_t) t][(size_t) k]; + } + } + int shape[2] = { T_frames, num_codebooks }; + debug_dump_i32_as_f32(&d, "codes-full", flat.data(), shape, 2); + } + + // Codec decode : transpose codes from [T_frames, K] to [K, T_frames] + // because pipeline_codec_decode expects K-major layout (codebooks + // first, frames second), then return the 24 kHz mono audio. + if (all_codes.empty()) { + return true; + } + + int T_frames = (int) all_codes.size(); + std::vector codes_kt((size_t) num_codebooks * (size_t) T_frames); + for (int t = 0; t < T_frames; t++) { + for (int k = 0; k < num_codebooks; k++) { + codes_kt[(size_t) k * (size_t) T_frames + (size_t) t] = all_codes[(size_t) t][(size_t) k]; + } + } + out->audio = pipeline_codec_decode(&pt->codec, codes_kt.data(), num_codebooks, T_frames); + if (out->audio.empty()) { + qt_log(QT_LOG_ERROR, "[Pipeline] codec decode returned no audio"); + return false; + } + + if (params.dump_dir) { + DebugDumper d; + debug_init(&d, params.dump_dir); + debug_dump_1d(&d, "output-audio", out->audio.data(), (int) out->audio.size()); + } + + return true; +} diff --git a/src/pipeline-tts.h b/src/pipeline-tts.h new file mode 100644 index 0000000..d22dbb6 --- /dev/null +++ b/src/pipeline-tts.h @@ -0,0 +1,140 @@ +#pragma once +// pipeline-tts.h : full TTS pipeline composition (Talker LM + code +// predictor MTP head + optional speaker encoder + 12Hz codec decoder). +// +// Phase 2.0 covers load-only: parse hyperparameters from both GGUF +// files, load every weight tensor on the configured backend, and +// expose the metadata needed to build forward graphs in later phases. +// No graph construction or sampling is wired here yet. + +#include "backend.h" +#include "code-predictor-weights.h" +#include "ggml-backend.h" +#include "gguf-weights.h" +#include "pipeline-codec.h" +#include "speaker-encoder-weights.h" +#include "talker-weights.h" + +#include +#include +#include + +struct CodecSpecials { + int pad_id; + int bos_id; + int eos_id; + int think_id; + int nothink_id; + int think_bos_id; + int think_eos_id; +}; + +struct TextSpecials { + int im_start_id; + int im_end_id; + int tts_pad_id; + int tts_bos_id; + int tts_eos_id; +}; + +struct LanguageEntry { + std::string name; + int id; +}; + +// Speaker entry for CustomVoice models. id is the codec embedding row id +// inserted in the talker prefix, dialect is empty unless the speaker +// overrides the user supplied language with a dialect lang_id (eric -> +// sichuan_dialect, dylan -> beijing_dialect on the upstream checkpoint). +struct SpeakerEntry { + std::string name; + int id; + std::string dialect; +}; + +struct GenerationDefaults { + bool do_sample; + int top_k; + float top_p; + float temperature; + float repetition_penalty; + bool subtalker_do_sample; + int subtalker_top_k; + float subtalker_top_p; + float subtalker_temperature; + int max_new_tokens; +}; + +struct PipelineTTS { + GGUFModel gguf_talker; + TalkerWeights talker; + CodePredictorWeights code_predictor; + SpeakerEncoderWeights speaker_encoder; + bool has_speaker_encoder; + + PipelineCodec codec; + + std::string tokenizer_type; + std::string model_size; + std::string model_type; + int num_code_groups; + + CodecSpecials codec_specials; + TextSpecials text_specials; + std::vector languages; + std::vector speakers; + GenerationDefaults gen_defaults; + + BackendPair bp; + ggml_backend_t backend; + ggml_backend_sched_t sched; +}; + +// Open the talker GGUF and the codec GGUF, load every module on the +// shared backend. Aborts with a logged error on any missing tensor or +// invalid metadata. Caller frees with pipeline_tts_free. +bool pipeline_tts_load(PipelineTTS * pt, const char * talker_gguf_path, const char * codec_gguf_path, BackendPair bp); + +void pipeline_tts_free(PipelineTTS * pt); + +struct BPETokenizer; + +// Parameters for one synthesis call. Lifetime constraint : text and lang +// are borrowed pointers, must outlive the call. dump_dir, when non-NULL, +// captures step 0 prefill activations plus the codes-full / output-audio +// dumps under the named directory ; debug only, slows the run. +struct PipelineTTSSynthesizeParams { + const char * text; + const char * lang; + const char * instruct; + const char * speaker; + const char * ref_audio; + const char * ref_text; + int64_t seed; + int max_new_tokens; + bool do_sample; + float temperature; + int top_k; + float top_p; + float repetition_penalty; + bool subtalker_do_sample; + float subtalker_temperature; + int subtalker_top_k; + float subtalker_top_p; + const char * dump_dir; +}; + +// Output of one synthesis call. audio is a 24 kHz mono F32 PCM buffer +// already decoded through the codec ; the caller writes it to disk. +struct PipelineTTSSynthesizeOutput { + std::vector audio; + int sample_rate; +}; + +// Run the full TTS pipeline : prompt assembly, prefill, frame loop with +// sampling, codec decode. Returns false on any failure with a diagnostic +// already routed through qt_log / qt_set_error. +bool pipeline_tts_synthesize(PipelineTTS * pt, + BPETokenizer * tok, + const PipelineTTSSynthesizeParams & params, + PipelineTTSSynthesizeOutput * out); diff --git a/src/prompt-builder.cpp b/src/prompt-builder.cpp new file mode 100644 index 0000000..93d6fb2 --- /dev/null +++ b/src/prompt-builder.cpp @@ -0,0 +1,628 @@ +// prompt-builder.cpp : CPU-side construction of the talker prefix +// input embedding. Mirrors generate() in qwen_tts/core/models/modeling_qwen3_tts.py +// for the strict subset {non-streaming, no voice clone}. +// +// Two streams are aligned then summed : +// text stream : text_projection(text_embedding(ids)) 151936 -> 2048 -> 1024 +// codec stream : codec_embedding(ids) 3072 -> 1024 +// +// Layout (lang_id != none, no speaker, no instruct) : +// +// role text(input_id[0:3]) 3 vecs +// prefill_lhs tts_pad x4 + tts_bos 5 vecs +// + codec_emb([think, think_bos, lang_id, think_eos, codec_pad]) +// trailing_lhs text(input_id[3:-5]) + tts_eos N_text + 1 vecs +// + codec_emb([codec_pad x (N_text + 1)]) +// trailing_rhs tts_pad + codec_emb([codec_bos]) 1 vec +// +// CustomVoice inserts the speaker codec embedding row between think_eos +// and codec_pad in the prefill, growing the prefill by one vector and +// substituting one tts_pad with another in the text stream alignment. +// +// VoiceDesign / CustomVoice may also prepend an instruct segment built +// from text_projection(text_embedding(<|im_start|>user\n{instruct}<|im_end|>\n)) +// laid out as N_instruct standalone vectors before the role. +// +// All math is f32. text_embedding and codec_embedding are read from +// the mmapped GGUF in their stored dtype (bf16 by default) and cast +// row by row. The 2-layer ResizeMLP runs as two GEMMs with a SiLU in +// between, with bias on both linear layers. + +#include "prompt-builder.h" + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +// Convert one row of an embedding matrix W [vocab, dim] to f32. Uses the +// ggml type traits to_float dispatch so every dtype shipped by the +// quantizer is supported (F32, BF16, F16, Q8_0, Q4_K_M, etc). The row +// stride is the type block size, computed via ggml_row_size. +static void embed_row_to_f32(const GGUFModel & gf, const char * tensor_name, int row_id, int dim, float * dst) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name); + if (!src) { + fprintf(stderr, "[Prompt] FATAL: tensor '%s' not in meta context\n", tensor_name); + std::exit(1); + } + if (src->ne[0] != dim) { + fprintf(stderr, "[Prompt] FATAL: tensor '%s' dim mismatch %lld vs %d\n", tensor_name, (long long) src->ne[0], + dim); + std::exit(1); + } + if (row_id < 0 || row_id >= (int) src->ne[1]) { + fprintf(stderr, "[Prompt] FATAL: row %d out of range for '%s' (vocab=%lld)\n", row_id, tensor_name, + (long long) src->ne[1]); + std::exit(1); + } + + const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name); + if (!base) { + fprintf(stderr, "[Prompt] FATAL: tensor '%s' has no data\n", tensor_name); + std::exit(1); + } + + const size_t row_bytes = ggml_row_size(src->type, dim); + const void * row = base + (size_t) row_id * row_bytes; + + if (src->type == GGML_TYPE_F32) { + std::memcpy(dst, row, (size_t) dim * sizeof(float)); + return; + } + + const struct ggml_type_traits * tt = ggml_get_type_traits(src->type); + if (!tt || !tt->to_float) { + fprintf(stderr, "[Prompt] FATAL: unsupported dtype %d for '%s'\n", (int) src->type, tensor_name); + std::exit(1); + } + tt->to_float(row, dst, dim); +} + +// Read a full small tensor (bias, projection weight) into an f32 buffer. +// Allocates dst.resize internally. Routed through ggml_get_type_traits so +// quants are accepted, same as embed_row_to_f32 above. +static void read_tensor_f32(const GGUFModel & gf, const char * tensor_name, std::vector & dst) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, tensor_name); + if (!src) { + fprintf(stderr, "[Prompt] FATAL: tensor '%s' not in meta context\n", tensor_name); + std::exit(1); + } + int64_t n = ggml_nelements(src); + const uint8_t * base = (const uint8_t *) gf_get_data(gf, tensor_name); + dst.resize((size_t) n); + + if (src->type == GGML_TYPE_F32) { + std::memcpy(dst.data(), base, (size_t) n * sizeof(float)); + return; + } + + const struct ggml_type_traits * tt = ggml_get_type_traits(src->type); + if (!tt || !tt->to_float) { + fprintf(stderr, "[Prompt] FATAL: unsupported dtype %d for '%s'\n", (int) src->type, tensor_name); + std::exit(1); + } + tt->to_float(base, dst.data(), (int64_t) n); +} + +// y = W @ x + b +// x [in_dim] f32, W [out_dim, in_dim] row-major f32, b [out_dim] f32 +// y [out_dim] f32 +// Naive dot-product GEMV, fine for small (≤2048) inputs at build time. +static void linear_f32(const float * x, const float * W, const float * b, int in_dim, int out_dim, float * y) { + for (int o = 0; o < out_dim; o++) { + const float * row = W + (size_t) o * (size_t) in_dim; + float acc = b ? b[o] : 0.0f; + for (int i = 0; i < in_dim; i++) { + acc += row[i] * x[i]; + } + y[o] = acc; + } +} + +static inline float silu(float v) { + return v / (1.0f + std::exp(-v)); +} + +// Apply text_projection : F1 (text_hidden -> text_hidden) -> SiLU -> F2 +// (text_hidden -> hidden), both with bias. +struct TextProjection { + int in_dim; // text_hidden_size + int hid_dim; // intermediate (= text_hidden_size in 0.6B) + int out_dim; // hidden_size + std::vector fc1_w; // [hid_dim, in_dim] + std::vector fc1_b; // [hid_dim] + std::vector fc2_w; // [out_dim, hid_dim] + std::vector fc2_b; // [out_dim] +}; + +static void text_projection_load(TextProjection * tp, const GGUFModel & gf, int text_hidden_size, int hidden_size) { + tp->in_dim = text_hidden_size; + tp->hid_dim = text_hidden_size; + tp->out_dim = hidden_size; + read_tensor_f32(gf, "talker.text_proj.fc1.weight", tp->fc1_w); + read_tensor_f32(gf, "talker.text_proj.fc1.bias", tp->fc1_b); + read_tensor_f32(gf, "talker.text_proj.fc2.weight", tp->fc2_w); + read_tensor_f32(gf, "talker.text_proj.fc2.bias", tp->fc2_b); +} + +static void text_projection_apply(const TextProjection * tp, const float * x, float * y) { + std::vector h((size_t) tp->hid_dim); + linear_f32(x, tp->fc1_w.data(), tp->fc1_b.data(), tp->in_dim, tp->hid_dim, h.data()); + for (int i = 0; i < tp->hid_dim; i++) { + h[(size_t) i] = silu(h[(size_t) i]); + } + linear_f32(h.data(), tp->fc2_w.data(), tp->fc2_b.data(), tp->hid_dim, tp->out_dim, y); +} + +// Compute text_proj(text_embedding(ids[start:end])) row by row, append +// to dst (which already holds previous rows). Each output row is one +// hidden-dim vector. +static void embed_text_range(const GGUFModel & gf, + const TextProjection * tp, + const int32_t * ids, + int start, + int end, + int text_hidden_size, + int hidden_size, + std::vector & dst) { + std::vector e((size_t) text_hidden_size); + std::vector y((size_t) hidden_size); + for (int i = start; i < end; i++) { + embed_row_to_f32(gf, "talker.text_embd.weight", ids[i], text_hidden_size, e.data()); + text_projection_apply(tp, e.data(), y.data()); + dst.insert(dst.end(), y.begin(), y.end()); + } +} + +// Append codec_embedding(id) to dst (one hidden-dim vector). +static void embed_codec(const GGUFModel & gf, int id, int hidden_size, std::vector & dst) { + size_t old = dst.size(); + dst.resize(old + (size_t) hidden_size); + embed_row_to_f32(gf, "talker.codec_embd.weight", id, hidden_size, dst.data() + old); +} + +// Vector add : a += b, length n. +static void vec_add(float * a, const float * b, int n) { + for (int i = 0; i < n; i++) { + a[i] += b[i]; + } +} + +bool prompt_builder_build(const PipelineTTS * pt, + const BPETokenizer * tok, + const std::string & utterance_text, + const std::string & language, + const std::string & instruct_text, + const std::string & speaker_name, + const float * ref_spk_emb, + const std::string & ref_text, + const int32_t * ref_codes, + int ref_codes_T, + PromptBuilderOutput * out) { + const int hidden = pt->talker.hidden_size; + const int text_hid = pt->talker.text_hidden_size; + + if (!speaker_name.empty() && ref_spk_emb != NULL) { + fprintf(stderr, "[Prompt] FATAL: speaker_name and ref_spk_emb are mutually exclusive\n"); + return false; + } + + // Voice clone mode B : ref_text and ref_codes drive an ICL prefix. + // Mode B requires ref_spk_emb so the speaker slot is also filled. + const bool icl = !ref_text.empty() && ref_codes != NULL && ref_codes_T > 0; + if (icl && ref_spk_emb == NULL) { + fprintf(stderr, "[Prompt] FATAL: ICL mode requires ref_spk_emb (no --ref-audio?)\n"); + return false; + } + + // Build the chat-templated prompt fed to the BPE tokenizer. + // Same wrap as the upstream demos : assistant role + utterance + + // im_end + newline + assistant role. + std::string full_text; + full_text.reserve(utterance_text.size() + 64); + full_text = "<|im_start|>assistant\n"; + full_text += utterance_text; + full_text += "<|im_end|>\n<|im_start|>assistant\n"; + + std::vector ids = bpe_encode(tok, full_text, /*add_eos=*/false); + if ((int) ids.size() < 8) { + fprintf(stderr, "[Prompt] FATAL: tokenized prompt too short (%d tokens)\n", (int) ids.size()); + return false; + } + + out->prompt_ids.assign(ids.begin(), ids.end()); + const int N = (int) ids.size(); + const int N_text = N - 3 - 5; + if (N_text <= 0) { + fprintf(stderr, "[Prompt] FATAL: no utterance text in prompt (N=%d)\n", N); + return false; + } + + // Resolve language : "auto" -> no language id, prefill is 3 codec + // tokens (nothink, think_bos, think_eos). Otherwise insert the + // configured language id between think_bos and think_eos. + int language_id = -1; + { + std::string lang_lc = language; + for (char & c : lang_lc) { + c = (char) std::tolower((unsigned char) c); + } + if (lang_lc != "auto") { + for (const LanguageEntry & e : pt->languages) { + if (e.name == lang_lc) { + language_id = e.id; + break; + } + } + if (language_id < 0) { + fprintf(stderr, "[Prompt] FATAL: unknown language '%s'\n", language.c_str()); + return false; + } + } + } + + // Resolve speaker : empty name -> no speaker. Otherwise lookup case + // insensitively in pt->speakers and override the language id with the + // dialect entry when the user supplied language is chinese or auto, + // mirroring modeling_qwen3_tts.py lines 2118 to 2122. + int speaker_id = -1; + if (!speaker_name.empty()) { + std::string spk_lc = speaker_name; + for (char & c : spk_lc) { + c = (char) std::tolower((unsigned char) c); + } + const SpeakerEntry * found = NULL; + for (const SpeakerEntry & e : pt->speakers) { + if (e.name == spk_lc) { + found = &e; + break; + } + } + if (!found) { + fprintf(stderr, "[Prompt] FATAL: unknown speaker '%s'\n", speaker_name.c_str()); + return false; + } + speaker_id = found->id; + + // Dialect override : applied only when the user supplied language + // is chinese or auto, the dialect string is non empty, and the + // dialect resolves to a known language id. + if (!found->dialect.empty()) { + std::string lang_lc = language; + for (char & c : lang_lc) { + c = (char) std::tolower((unsigned char) c); + } + if (lang_lc == "chinese" || lang_lc == "auto") { + int dialect_id = -1; + for (const LanguageEntry & e : pt->languages) { + if (e.name == found->dialect) { + dialect_id = e.id; + break; + } + } + if (dialect_id < 0) { + fprintf(stderr, "[Prompt] FATAL: dialect '%s' not in language table\n", found->dialect.c_str()); + return false; + } + language_id = dialect_id; + } + } + } + + // Load the small tensors needed for the builder onto the host side. + TextProjection tp; + text_projection_load(&tp, pt->gguf_talker, text_hid, hidden); + + // Special embeds (tts_bos, tts_eos, tts_pad, codec_pad, codec_bos) + // computed once. + std::vector tts_bos_emb((size_t) hidden); + std::vector tts_eos_emb((size_t) hidden); + std::vector tts_pad_emb((size_t) hidden); + { + std::vector e((size_t) text_hid); + embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_bos_id, text_hid, e.data()); + text_projection_apply(&tp, e.data(), tts_bos_emb.data()); + embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_eos_id, text_hid, e.data()); + text_projection_apply(&tp, e.data(), tts_eos_emb.data()); + embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", pt->text_specials.tts_pad_id, text_hid, e.data()); + text_projection_apply(&tp, e.data(), tts_pad_emb.data()); + } + + std::vector codec_pad_emb((size_t) hidden); + embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.pad_id, hidden, + codec_pad_emb.data()); + + // Codec prefill list : 3 ids if auto (no language), 4 otherwise. + // Speaker insertion : if a speaker id is set, the codec embedding row + // for that speaker slips between think_eos and codec_pad in the codec + // stream, mirroring modeling_qwen3_tts.py lines 2167 to 2172. + std::vector codec_prefill; + if (language_id < 0) { + codec_prefill = { pt->codec_specials.nothink_id, pt->codec_specials.think_bos_id, + pt->codec_specials.think_eos_id }; + } else { + codec_prefill = { pt->codec_specials.think_id, pt->codec_specials.think_bos_id, language_id, + pt->codec_specials.think_eos_id }; + } + if (speaker_id >= 0) { + codec_prefill.push_back(speaker_id); + } else if (ref_spk_emb != NULL) { + // Sentinel : the codec_left builder below copies ref_spk_emb in + // place of an embedding lookup whenever it sees -2. + codec_prefill.push_back(-2); + } + const int n_prefill = (int) codec_prefill.size(); + const int T_codec_prefix = n_prefill + 2; // + codec_pad + codec_bos + const int n_pad_pre = T_codec_prefix - 2; + + // Tokenize the instruct segment when non empty. The wrapper mirrors + // _build_instruct_text upstream : <|im_start|>user\n{instruct}<|im_end|>\n + // The result is a flat list of text token ids that will be projected + // and placed as standalone vectors at the head of the input embed, + // with no codec stream contribution. + std::vector instruct_ids; + if (!instruct_text.empty()) { + std::string wrapped; + wrapped.reserve(instruct_text.size() + 32); + wrapped = "<|im_start|>user\n"; + wrapped += instruct_text; + wrapped += "<|im_end|>\n"; + instruct_ids = bpe_encode(tok, wrapped, /*add_eos=*/false); + } + const int N_instruct = (int) instruct_ids.size(); + + // Tokenize the reference utterance when ICL is active. The wrap is + // identical to the main utterance : assistant role + ref_text + + // im_end + newline + assistant role. We slice [3:-5] later to keep + // only the inner text body, mirroring input_id[:, 3:-5] upstream. + std::vector ref_ids; + int N_ref_text = 0; + if (icl) { + std::string ref_full; + ref_full.reserve(ref_text.size() + 64); + ref_full = "<|im_start|>assistant\n"; + ref_full += ref_text; + ref_full += "<|im_end|>\n<|im_start|>assistant\n"; + ref_ids = bpe_encode(tok, ref_full, /*add_eos=*/false); + if ((int) ref_ids.size() < 8) { + fprintf(stderr, "[Prompt] FATAL: ref_text tokenized too short (%d tokens)\n", (int) ref_ids.size()); + return false; + } + // ref_ids[3 : -5] is the inner ref text body without role tokens + N_ref_text = (int) ref_ids.size() - 3 - 5; + if (N_ref_text <= 0) { + fprintf(stderr, "[Prompt] FATAL: empty ref_text body\n"); + return false; + } + } + + // ICL geometry. text_lens = N_ref_text + N_text + 1 (tts_eos). + // codec_lens = 1 (codec_bos) + ref_codes_T. The non_streaming_mode + // branch upstream pads the shorter stream so they end up the same + // length, except when text > codec where trailing_text_hidden carries + // the leftover text rows. + const int text_lens_icl = icl ? (N_ref_text + N_text + 1) : 0; + const int codec_lens_icl = icl ? (1 + ref_codes_T) : 0; + const int icl_T = icl ? (text_lens_icl > codec_lens_icl ? codec_lens_icl : codec_lens_icl) : 0; + + // Allocate the full output buffer. + // Standard layout : N_instruct + 3 (role) + (n_pad_pre + 1) + N_text + 1 (eos) + 1 (final) + // ICL layout : N_instruct + 3 (role) + (n_pad_pre + 1) + icl_T + const int T_ctx = + icl ? (N_instruct + 3 + (n_pad_pre + 1) + icl_T) : (N_instruct + 3 + (n_pad_pre + 1) + N_text + 1 + 1); + out->T_ctx = T_ctx; + out->hidden = hidden; + out->input_embed.assign((size_t) T_ctx * (size_t) hidden, 0.0f); + out->N_text = N_text; + + int row = 0; + auto row_ptr = [&](int r) { + return out->input_embed.data() + (size_t) r * (size_t) hidden; + }; + + // Instruct prefix : text_proj(text_embed(instruct_ids)). Standalone + // vectors with no codec stream (zero pad_id sum, ie nothing added). + if (N_instruct > 0) { + std::vector dst; + embed_text_range(pt->gguf_talker, &tp, instruct_ids.data(), 0, N_instruct, text_hid, hidden, dst); + std::memcpy(row_ptr(row), dst.data(), dst.size() * sizeof(float)); + row += N_instruct; + } + + // Role : text_proj(text_embed(ids[0:3])) + { + std::vector dst; + dst.reserve((size_t) 3 * (size_t) hidden); + embed_text_range(pt->gguf_talker, &tp, ids.data(), 0, 3, text_hid, hidden, dst); + std::memcpy(row_ptr(row), dst.data(), dst.size() * sizeof(float)); + row += 3; + } + + // Codec prefix : tts_pad x n_pad_pre + tts_bos, summed with + // codec_emb([codec_prefill_list[:-1]] + codec_pad). The Python code + // takes codec_input_embedding[:, :-1] which drops the codec_bos, + // leaving [codec_prefill_list..., codec_pad]. + { + std::vector codec_left = codec_prefill; + codec_left.push_back(pt->codec_specials.pad_id); + for (int i = 0; i < (int) codec_left.size(); i++) { + float * r = row_ptr(row + i); + // text stream : tts_pad * (n - 1) then tts_bos at the end + const float * text_vec = (i == (int) codec_left.size() - 1) ? tts_bos_emb.data() : tts_pad_emb.data(); + std::memcpy(r, text_vec, (size_t) hidden * sizeof(float)); + // codec stream : either an embedding lookup or, when the + // sentinel -2 marks the speaker slot, a direct copy of the + // user supplied ref_spk_emb (voice clone mode A). + std::vector ce((size_t) hidden); + if (codec_left[(size_t) i] == -2) { + std::memcpy(ce.data(), ref_spk_emb, (size_t) hidden * sizeof(float)); + } else { + embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", codec_left[(size_t) i], hidden, + ce.data()); + } + vec_add(r, ce.data(), hidden); + } + row += (int) codec_left.size(); + } + + // From here, two paths : standard (no ICL) builds the trailing + // utterance text + tts_eos + final_pad, ICL builds an aligned + // text/codec block that replaces those rows entirely. + if (!icl) { + // Standard layout : trailing utterance text + tts_eos rows summed + // with codec_pad, then a final tts_pad + codec_bos row. + for (int i = 0; i < N_text; i++) { + std::vector e((size_t) text_hid); + std::vector y((size_t) hidden); + embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ids[3 + i], text_hid, e.data()); + text_projection_apply(&tp, e.data(), y.data()); + float * r = row_ptr(row); + std::memcpy(r, y.data(), (size_t) hidden * sizeof(float)); + vec_add(r, codec_pad_emb.data(), hidden); + row++; + } + { + float * r = row_ptr(row); + std::memcpy(r, tts_eos_emb.data(), (size_t) hidden * sizeof(float)); + vec_add(r, codec_pad_emb.data(), hidden); + row++; + } + { + float * r = row_ptr(row); + std::memcpy(r, tts_pad_emb.data(), (size_t) hidden * sizeof(float)); + std::vector ce((size_t) hidden); + embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.bos_id, hidden, ce.data()); + vec_add(r, ce.data(), hidden); + row++; + } + } else { + // ICL layout : compute the text stream and the codec stream + // separately then add them. The text stream is text_proj of + // [ref_text_ids ; utterance_text_ids] followed by tts_eos. The + // codec stream is codec_emb(codec_bos) followed by sum over the + // 16 codebook embeddings of ref_codes[i, t] for each frame t. + // Both streams are aligned to length icl_T per the upstream + // non_streaming_mode=False branch of generate_icl_prompt. + const int T_icl = codec_lens_icl; // text_lens > codec : truncate to codec, else pad text up to codec + + // Build the codec stream [T_icl, hidden]. Row 0 : codec_emb(codec_bos). + // Row 1..ref_codes_T : sum over k=0..15 of codebook_k_emb(ref_codes[k, t]). + std::vector codec_stream((size_t) T_icl * (size_t) hidden, 0.0f); + { + // Row 0 : codec_bos lookup. + embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", pt->codec_specials.bos_id, hidden, + codec_stream.data()); + // Row 1..ref_codes_T : sum over codebooks. + std::vector tmp((size_t) hidden); + for (int t = 0; t < ref_codes_T; t++) { + float * dst = codec_stream.data() + (size_t) (1 + t) * (size_t) hidden; + // codebook 0 lives in talker.codec_embd + int code0 = ref_codes[(size_t) 0 * (size_t) ref_codes_T + (size_t) t]; + embed_row_to_f32(pt->gguf_talker, "talker.codec_embd.weight", code0, hidden, dst); + // codebooks 1..15 live in code_pred.codec_embd.{i-1} + for (int i = 1; i < pt->num_code_groups; i++) { + int code = ref_codes[(size_t) i * (size_t) ref_codes_T + (size_t) t]; + char tname[64]; + std::snprintf(tname, sizeof(tname), "code_pred.codec_embd.%d.weight", i - 1); + embed_row_to_f32(pt->gguf_talker, tname, code, hidden, tmp.data()); + vec_add(dst, tmp.data(), hidden); + } + } + } + + // Build the text stream [text_lens_icl, hidden] = text_proj of + // [ref_text ; utterance_text] then tts_eos. + std::vector text_stream((size_t) text_lens_icl * (size_t) hidden, 0.0f); + for (int i = 0; i < N_ref_text; i++) { + std::vector e((size_t) text_hid); + float * r = text_stream.data() + (size_t) i * (size_t) hidden; + embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ref_ids[3 + i], text_hid, e.data()); + text_projection_apply(&tp, e.data(), r); + } + for (int i = 0; i < N_text; i++) { + std::vector e((size_t) text_hid); + float * r = text_stream.data() + (size_t) (N_ref_text + i) * (size_t) hidden; + embed_row_to_f32(pt->gguf_talker, "talker.text_embd.weight", ids[3 + i], text_hid, e.data()); + text_projection_apply(&tp, e.data(), r); + } + // Append tts_eos at the end of the text stream. + std::memcpy(text_stream.data() + (size_t) (text_lens_icl - 1) * (size_t) hidden, tts_eos_emb.data(), + (size_t) hidden * sizeof(float)); + + // Align the two streams to T_icl. text_lens > codec : truncate + // text and stash the leftover into trailing_text_hidden. text_lens + // <= codec : pad text with tts_pad up to codec, trailing reduces + // to tts_pad. + std::vector aligned_text((size_t) T_icl * (size_t) hidden, 0.0f); + if (text_lens_icl >= T_icl) { + // truncate text to T_icl rows, leftover goes into trailing + std::memcpy(aligned_text.data(), text_stream.data(), (size_t) T_icl * (size_t) hidden * sizeof(float)); + const int trailing_n = text_lens_icl - T_icl; + out->T_trailing = trailing_n > 0 ? trailing_n : 1; + out->trailing_text_hidden.assign((size_t) out->T_trailing * (size_t) hidden, 0.0f); + if (trailing_n > 0) { + std::memcpy(out->trailing_text_hidden.data(), text_stream.data() + (size_t) T_icl * (size_t) hidden, + (size_t) trailing_n * (size_t) hidden * sizeof(float)); + } else { + std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float)); + } + } else { + // pad text with tts_pad up to T_icl, trailing = single tts_pad row + std::memcpy(aligned_text.data(), text_stream.data(), + (size_t) text_lens_icl * (size_t) hidden * sizeof(float)); + for (int i = text_lens_icl; i < T_icl; i++) { + std::memcpy(aligned_text.data() + (size_t) i * (size_t) hidden, tts_pad_emb.data(), + (size_t) hidden * sizeof(float)); + } + out->T_trailing = 1; + out->trailing_text_hidden.assign((size_t) hidden, 0.0f); + std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float)); + } + + // Sum aligned_text + codec_stream into the input embed at the + // current row offset. + for (int i = 0; i < T_icl; i++) { + float * r = row_ptr(row + i); + std::memcpy(r, aligned_text.data() + (size_t) i * (size_t) hidden, (size_t) hidden * sizeof(float)); + vec_add(r, codec_stream.data() + (size_t) i * (size_t) hidden, hidden); + } + row += T_icl; + } + + if (row != T_ctx) { + fprintf(stderr, "[Prompt] FATAL: layout error row=%d expected T_ctx=%d\n", row, T_ctx); + return false; + } + + // Trailing text hidden : non streaming mode (no ICL) collapses the + // overlay to a single row equal to tts_pad_embed + // (modeling_qwen3_tts.py line 2227). The full utterance text is + // already integrated into the prefill above as codec_pad summed text + // rows + tts_eos, so the overlay loop only ever needs tts_pad : step + // 0 reads trailing_text_hidden[0] which is tts_pad, every later step + // falls through to the else branch and reads tts_pad_embed. One row, + // bit exact with the Python hook dump. + // + // ICL mode populates out->trailing_text_hidden directly inside the + // ICL branch above, so we only set the default here for non ICL. + if (!icl) { + out->T_trailing = 1; + out->trailing_text_hidden.assign((size_t) hidden, 0.0f); + std::memcpy(out->trailing_text_hidden.data(), tts_pad_emb.data(), (size_t) hidden * sizeof(float)); + } + + out->tts_pad_embed = tts_pad_emb; + + fprintf(stderr, + "[Prompt] Built: %d ids, N_text=%d, N_instruct=%d, T_ctx=%d, hidden=%d, lang=%s (id=%d), speaker=%s " + "(id=%d) ref_spk_emb=%s icl=%s\n", + N, N_text, N_instruct, T_ctx, hidden, language.c_str(), language_id, + speaker_name.empty() ? "none" : speaker_name.c_str(), speaker_id, ref_spk_emb ? "yes" : "no", + icl ? "yes" : "no"); + + return true; +} diff --git a/src/prompt-builder.h b/src/prompt-builder.h new file mode 100644 index 0000000..4dbe3b5 --- /dev/null +++ b/src/prompt-builder.h @@ -0,0 +1,74 @@ +#pragma once +// prompt-builder.h : assemble the talker prefix input embedding from +// a tokenized text plus a language tag, mirroring the upstream +// generate() function of Qwen3-TTS. +// +// Output shape : [T_ctx, hidden_size] f32 row-major. Two pad-aligned +// streams (text and codec) are summed at the granularity of single +// vectors. The trailing text hidden buffer is also produced for the +// streaming-text overlay used during generation. +// +// Modes : +// base text only, no instruct, no speaker +// voice_design text + instruct (style description), no speaker +// custom_voice text + speaker, optional instruct +// +// Empty / NULL strings disable the corresponding stream. The builder runs +// CPU-side using the BF16 weight blocks mmapped from the talker GGUF, no +// backend allocation, no graph compute. + +#include "bpe.h" +#include "pipeline-tts.h" + +#include +#include +#include + +struct PromptBuilderOutput { + // Final talker input embedding [T_ctx, hidden] f32 row-major + std::vector input_embed; + int T_ctx; + int hidden; + + // Trailing text overlay : added on top of the next-token-input during + // the autoregressive loop, one vector per generated frame until + // exhausted, then tts_pad_embed for every following frame. + std::vector trailing_text_hidden; + int T_trailing; + + // tts_pad_embed [hidden] f32, kept around so the generation loop + // can fall back on it once the trailing text runs out. + std::vector tts_pad_embed; + + // Token ids fed through the tokenizer (kept for debug parity with + // the Python prompt-ids.bin dump). + std::vector prompt_ids; + + // Length of the text segment used as utterance text, ie input_id[3:-5]. + int N_text; +}; + +// Assemble the prefix. instruct_text is the raw user style instruction +// (empty for none). speaker_name is the lowercased speaker key looked up +// in pt->speakers (empty for none). ref_spk_emb is an optional pointer to +// an [hidden] f32 vector that takes the place of the speaker preset row +// for voice clone mode A : when non NULL it is inserted between think_eos +// and codec_pad in the codec stream, mutually exclusive with speaker_name. +// ref_text and ref_codes activate voice clone mode B (ICL) : the prompt +// becomes [icl_text + tts_eos] aligned with [codec_bos + ref_codes_summed], +// where ref_codes is a flat [num_code_groups, T_codec] int32 buffer +// produced by pipeline_codec_encode on the resampled reference audio. +// Returns false if BPE encoding produces fewer than the expected +// role/footer tokens, the language is unknown, speaker_name is set but +// not found, or speaker_name and ref_spk_emb are both set. +bool prompt_builder_build(const PipelineTTS * pt, + const BPETokenizer * tok, + const std::string & utterance_text, + const std::string & language, + const std::string & instruct_text, + const std::string & speaker_name, + const float * ref_spk_emb, + const std::string & ref_text, + const int32_t * ref_codes, + int ref_codes_T, + PromptBuilderOutput * out); diff --git a/src/qt-error.cpp b/src/qt-error.cpp new file mode 100644 index 0000000..181283e --- /dev/null +++ b/src/qt-error.cpp @@ -0,0 +1,116 @@ +// qt-error.cpp : implementation of the qt_log / qt_set_error / qt_throw +// helpers declared in qt-error.h. Storage is thread_local for the error +// slot, atomic for the log callback so qt_log_set is wait-free. + +#include "qt-error.h" + +#include +#include +#include +#include +#include + +// Thread-local backing store for qt_last_error(). std::string sized once +// per thread, grows on demand, never freed across calls : the std runtime +// reclaims it on thread exit. An empty string means "no error recorded +// on this thread yet", which qt_last_error() exposes as "". +static thread_local std::string g_last_error; + +void qt_set_error_v(const char * fmt, va_list ap) { + if (!fmt) { + g_last_error.clear(); + return; + } + // Two-pass vsnprintf : first call sizes the buffer, second writes the + // message. va_copy keeps the original ap valid for the second pass. + va_list ap2; + va_copy(ap2, ap); + int needed = std::vsnprintf(nullptr, 0, fmt, ap2); + va_end(ap2); + if (needed < 0) { + g_last_error = "qt_set_error : vsnprintf failed"; + return; + } + g_last_error.resize(static_cast(needed)); + std::vsnprintf(g_last_error.data(), static_cast(needed) + 1, fmt, ap); +} + +void qt_set_error(const char * fmt, ...) { + va_list ap; + va_start(ap, fmt); + qt_set_error_v(fmt, ap); + va_end(ap); +} + +const char * qt_last_error(void) { + return g_last_error.c_str(); +} + +// Formats a message with printf semantics and throws std::runtime_error. +// The catch site at the binary entry inspects the what() string and feeds +// it into qt_set_error so the user-visible diagnostic is identical +// whether the failure used the bool-return path or the throw path. +void qt_throw(const char * fmt, ...) { + char buf[1024]; + if (fmt) { + va_list ap; + va_start(ap, fmt); + std::vsnprintf(buf, sizeof(buf), fmt, ap); + va_end(ap); + } else { + buf[0] = '\0'; + } + throw std::runtime_error(buf); +} + +// Process-wide log callback. Atomic so qt_log_set can replace it without +// locking : write happens with memory_order_release, every reader sees a +// fully published callback pointer paired with its user_data slot. +// std::atomic on a function pointer is lock-free on every platform we +// target. user_data is a plain pointer because it is only ever published +// alongside cb under the same release ordering. +static std::atomic g_log_cb{ nullptr }; +static void * g_log_cb_user = nullptr; + +void qt_log_set(qt_log_cb cb, void * user_data) { + g_log_cb_user = user_data; + g_log_cb.store(cb, std::memory_order_release); +} + +// Routes one log line to the installed callback or to stderr. Two-pass +// vsnprintf sizes the heap buffer when the message exceeds the stack +// scratchpad, which keeps the common case allocation-free. +void qt_log(enum qt_log_level level, const char * fmt, ...) { + if (!fmt) { + return; + } + char stackbuf[512]; + char * buf = stackbuf; + int needed = 0; + va_list ap; + va_start(ap, fmt); + { + va_list ap2; + va_copy(ap2, ap); + needed = std::vsnprintf(stackbuf, sizeof(stackbuf), fmt, ap2); + va_end(ap2); + } + if (needed < 0) { + va_end(ap); + return; + } + std::string heapbuf; + if ((size_t) needed >= sizeof(stackbuf)) { + heapbuf.resize((size_t) needed); + std::vsnprintf(heapbuf.data(), (size_t) needed + 1, fmt, ap); + buf = heapbuf.data(); + } + va_end(ap); + + qt_log_cb cb = g_log_cb.load(std::memory_order_acquire); + if (cb) { + cb(level, buf, g_log_cb_user); + } else { + std::fprintf(stderr, "%s\n", buf); + } +} diff --git a/src/qt-error.h b/src/qt-error.h new file mode 100644 index 0000000..f631443 --- /dev/null +++ b/src/qt-error.h @@ -0,0 +1,68 @@ +#pragma once +// qt-error.h : internal helpers backing the (future) public qt_last_error +// entry and the qt_log callback routing. +// +// Storage is thread_local so concurrent synthesize calls on different +// threads never race on each other's messages. The setter is variadic +// with printf semantics ; messages longer than the internal buffer are +// truncated, never split. Passing NULL as fmt clears the slot. +// +// qt_throw is the load-path counterpart : functions deep inside the GGUF +// reader and the codec load chain cannot return false up dozens of call +// sites without a massive cascade. They throw a std::runtime_error +// instead, which the binary entry point (main, or a future ABI boundary) +// catches and converts to qt_set_error plus a non-zero exit. Exceptions +// never cross any future C ABI. +// +// qt_log routes a formatted message to the user-installed qt_log_cb, or +// to stderr when no callback is installed. Used by every translation +// unit in the lib that wants its diagnostics to be redirectable from a +// wrapper (Python logging, Rust tracing, ...). + +#include + +enum qt_log_level { + QT_LOG_DEBUG = 0, + QT_LOG_INFO = 1, + QT_LOG_WARN = 2, + QT_LOG_ERROR = 3, +}; + +typedef void (*qt_log_cb)(enum qt_log_level level, const char * msg, void * user_data); + +void qt_set_error(const char * fmt, ...) +#if defined(__GNUC__) || defined(__clang__) + __attribute__((format(printf, 1, 2))) +#endif + ; + +void qt_set_error_v(const char * fmt, va_list ap); + +// Throws std::runtime_error formatted with printf semantics. Tagged +// noreturn so the compiler can prune unreachable branches at the call +// site. Designed for the GGUF / codec load path where any failure means +// the model is unusable and unwinding to the boundary is the only sane +// recovery. +[[noreturn]] void qt_throw(const char * fmt, ...) +#if defined(__GNUC__) || defined(__clang__) + __attribute__((format(printf, 1, 2))) +#endif + ; + +// Routes a formatted message at the requested level to the installed +// callback, or to stderr when none is set. The message is the full line +// without trailing newline ; routing layers add their own framing. +void qt_log(enum qt_log_level level, const char * fmt, ...) +#if defined(__GNUC__) || defined(__clang__) + __attribute__((format(printf, 2, 3))) +#endif + ; + +// Install a process-wide log callback. Pass NULL to revert to stderr. +// user_data is opaque, forwarded as-is to every callback invocation. +void qt_log_set(qt_log_cb cb, void * user_data); + +// Returns the most recent error message recorded on the calling thread. +// Returns "" if no error has been set on this thread. The pointer stays +// valid until the next qt_set_error call on the same thread. +const char * qt_last_error(void); diff --git a/src/quantizer-decode.h b/src/quantizer-decode.h new file mode 100644 index 0000000..351e544 --- /dev/null +++ b/src/quantizer-decode.h @@ -0,0 +1,168 @@ +#pragma once +// quantizer-decode.h : split RVQ decode for the Qwen3-TTS 12Hz tokenizer +// (GGML). +// Reads 16 codebooks (1 semantic + 15 acoustic) of 2048 entries with +// internal dim 256, and produces a 512-channel hidden representation. +// +// Decode side : codes [T, 16] i32 -> hidden [T, 512] f32 by summing +// F.embedding(codes[:, k], codebook_k) within each split, then applying +// a per-split output_proj 1x1 conv (256 -> 512), then summing the two +// splits. + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include + +#define QWEN_RVQ_MAX_CODEBOOKS_PER_GROUP 15 + +struct QwenRVQGroup { + int num_codebooks; + struct ggml_tensor * embed[QWEN_RVQ_MAX_CODEBOOKS_PER_GROUP]; // each [256, 2048] f32 + struct ggml_tensor * out_proj_w; // [256, 512] f32 (Conv1d 1x1 reshaped) +}; + +struct QwenQuantizerDecoder { + int num_quantizers; // total RVQ stages, 16 + int num_semantic_quantizers; // 1 + int num_acoustic_quantizers; // 15 + int codebook_size; // 2048 + int codebook_dim_internal; // 256 + int hidden; // 512 + + QwenRVQGroup semantic; + QwenRVQGroup acoustic; + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +// Load a Conv1d 1x1 weight stored on disk as (1, in, out) and present it +// as a 2D [in, out] f32 tensor suitable for ggml_mul_mat. +static struct ggml_tensor * qwen_load_proj_1x1(WeightCtx * wctx, const GGUFModel & gf, const std::string & name) { + struct ggml_tensor * src = ggml_get_tensor(gf.meta, name.c_str()); + if (!src) { + fprintf(stderr, "[Quantizer] FATAL: tensor '%s' not found\n", name.c_str()); + exit(1); + } + if (src->ne[0] != 1) { + fprintf(stderr, "[Quantizer] FATAL: '%s' expected kernel=1 on ne[0], got %lld\n", name.c_str(), + (long long) src->ne[0]); + exit(1); + } + int64_t shape2d[2] = { src->ne[1], src->ne[2] }; // (in_dim, out_dim) in ggml row-major + return gf_load_tensor(wctx, gf, name, shape2d, 2); +} + +// Build the on-backend weights of the split RVQ decoder from a loaded GGUF. +// Mutates dec->weight_ctx and dec->weight_buf, and binds every group +// tensor pointer to a backend allocation. +static bool qwen_quantizer_decoder_load(QwenQuantizerDecoder * dec, const GGUFModel & gf, ggml_backend_t backend) { + dec->num_quantizers = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.num_quantizers"); + dec->num_semantic_quantizers = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.num_semantic_quantizers"); + dec->num_acoustic_quantizers = dec->num_quantizers - dec->num_semantic_quantizers; + dec->codebook_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.codebook_size"); + dec->codebook_dim_internal = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.codebook_dim_internal"); + dec->hidden = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.vector_quantization_hidden_dim"); + + if (dec->num_acoustic_quantizers > QWEN_RVQ_MAX_CODEBOOKS_PER_GROUP) { + fprintf(stderr, "[Quantizer] FATAL: %d acoustic codebooks exceeds compile-time max %d\n", + dec->num_acoustic_quantizers, QWEN_RVQ_MAX_CODEBOOKS_PER_GROUP); + return false; + } + + int n_tensors = (dec->num_semantic_quantizers + 1) // semantic codebooks + out_proj + + (dec->num_acoustic_quantizers + 1) // acoustic codebooks + out_proj + + 4; // headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + dec->semantic.num_codebooks = dec->num_semantic_quantizers; + dec->semantic.out_proj_w = qwen_load_proj_1x1(&wctx, gf, "tok_dec.vq_first.output_proj.weight"); + for (int k = 0; k < dec->semantic.num_codebooks; k++) { + char name[128]; + snprintf(name, sizeof(name), "tok_dec.vq_first.%d.codebook", k); + dec->semantic.embed[k] = gf_load_tensor(&wctx, gf, name); + } + + dec->acoustic.num_codebooks = dec->num_acoustic_quantizers; + dec->acoustic.out_proj_w = qwen_load_proj_1x1(&wctx, gf, "tok_dec.vq_rest.output_proj.weight"); + for (int k = 0; k < dec->acoustic.num_codebooks; k++) { + char name[128]; + snprintf(name, sizeof(name), "tok_dec.vq_rest.%d.codebook", k); + dec->acoustic.embed[k] = gf_load_tensor(&wctx, gf, name); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[Quantizer] FATAL: backend allocation failed\n"); + return false; + } + dec->weight_ctx = wctx.ctx; + dec->weight_buf = wctx.buffer; + + fprintf(stderr, + "[Quantizer] Loaded: %d codebooks (%d semantic + %d acoustic), " + "%d entries x %d dim, hidden %d\n", + dec->num_quantizers, dec->num_semantic_quantizers, dec->num_acoustic_quantizers, dec->codebook_size, + dec->codebook_dim_internal, dec->hidden); + return true; +} + +static void qwen_quantizer_decoder_free(QwenQuantizerDecoder * dec) { + if (dec->weight_buf) { + ggml_backend_buffer_free(dec->weight_buf); + dec->weight_buf = NULL; + } + if (dec->weight_ctx) { + ggml_free(dec->weight_ctx); + dec->weight_ctx = NULL; + } +} + +// Sum F.embedding(codes[:, k], embed_k) across the codebooks of one +// split, then project from internal_dim (256) to hidden (512) via a +// Conv1d 1x1 (mat_mul against out_proj_w). +// +// codes_split : [T, K] i32, K is the codebook count of this split +// returns : [hidden, T] f32 +static struct ggml_tensor * qwen_rvq_group_decode(struct ggml_context * ctx, + const QwenRVQGroup & g, + struct ggml_tensor * codes_split, + int T) { + struct ggml_tensor * sum = NULL; + for (int k = 0; k < g.num_codebooks; k++) { + struct ggml_tensor * idx = ggml_view_1d(ctx, codes_split, T, (size_t) k * codes_split->nb[1]); + struct ggml_tensor * emb = ggml_get_rows(ctx, g.embed[k], idx); + sum = (sum == NULL) ? emb : ggml_add(ctx, sum, emb); + } + // sum : [internal_dim=256, T] + // out_proj_w : [internal_dim=256, hidden=512] + // ggml_mul_mat returns [hidden=512, T] + return ggml_mul_mat(ctx, g.out_proj_w, sum); +} + +// codes : [T, num_quantizers=16] i32 +// returns : [hidden=512, T] f32 +static struct ggml_tensor * qwen_quantizer_decode(struct ggml_context * ctx, + const QwenQuantizerDecoder * dec, + struct ggml_tensor * codes) { + int T = (int) codes->ne[0]; + if ((int) codes->ne[1] != dec->num_quantizers) { + fprintf(stderr, "[Quantizer] FATAL: codes ne[1]=%lld != num_quantizers=%d\n", (long long) codes->ne[1], + dec->num_quantizers); + return NULL; + } + + struct ggml_tensor * codes_sem = ggml_view_2d(ctx, codes, T, dec->num_semantic_quantizers, codes->nb[1], 0); + size_t aco_off = (size_t) dec->num_semantic_quantizers * codes->nb[1]; + struct ggml_tensor * codes_aco = ggml_view_2d(ctx, codes, T, dec->num_acoustic_quantizers, codes->nb[1], aco_off); + + struct ggml_tensor * h_sem = qwen_rvq_group_decode(ctx, dec->semantic, codes_sem, T); + struct ggml_tensor * h_aco = qwen_rvq_group_decode(ctx, dec->acoustic, codes_aco, T); + + return ggml_add(ctx, h_sem, h_aco); +} diff --git a/src/quantizer-encode.h b/src/quantizer-encode.h new file mode 100644 index 0000000..cd811a5 --- /dev/null +++ b/src/quantizer-encode.h @@ -0,0 +1,258 @@ +#pragma once +// quantizer-encode.h: RVQ encode for the Qwen3-TTS 12Hz tokenizer. +// +// Inputs: hidden state [T, 512] f32 T-first, output of the encoder +// downsample. Outputs: codes [16, T] i32, with codebook 0 carrying the +// semantic stream and codebooks 1..15 carrying the acoustic residual +// stream. +// +// Each side has the same shape: +// input_proj : Conv1d k=1, 512 -> 256 (linear projection on channels) +// codebooks : list of [2048, 256] f32 entries used as kNN centroids +// output_proj : Conv1d k=1, 256 -> 512 (used only inside the residual loop) +// +// At encode time we run, for each side: +// y = input_proj(x) +// res = y +// codes = [] +// for layer in layers: +// idx = argmin_e ||res - codebook_e||^2 +// q = codebook[idx] +// res = res - q +// codes += [idx] +// +// The semantic side has 1 codebook, the acoustic side has 15 codebooks, +// concatenated to produce the final 16-codebook stream. + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include + +#define QWEN_ENC_QUANT_NUM_SEMANTIC 1 +#define QWEN_ENC_QUANT_NUM_ACOUSTIC 15 +#define QWEN_ENC_QUANT_TOTAL (QWEN_ENC_QUANT_NUM_SEMANTIC + QWEN_ENC_QUANT_NUM_ACOUSTIC) + +struct QwenQuantizerEncodeSide { + struct ggml_tensor * input_proj_w; // [1, 512, 256] f32, k=1 conv + struct ggml_tensor * output_proj_w; // [1, 256, 512] f32, k=1 conv + int num_layers; + struct ggml_tensor * codebooks[QWEN_ENC_QUANT_NUM_ACOUSTIC]; // [256, 2048] each +}; + +struct QwenQuantizerEncode { + QwenQuantizerEncodeSide semantic; // 1 codebook + QwenQuantizerEncodeSide acoustic; // 15 codebooks + + int codebook_size; // 2048 + int codebook_dim; // 256 + int hidden_size; // 512 + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +static bool qwen_quantizer_encode_load(QwenQuantizerEncode * q, const GGUFModel & gf, ggml_backend_t backend) { + q->codebook_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.codebook_size"); + q->codebook_dim = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.vector_quantization_hidden_dim"); + q->hidden_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.hidden_size"); + + int n_tensors = 4 + QWEN_ENC_QUANT_TOTAL + 4; // 4 proj + 16 codebooks + headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + q->semantic.num_layers = QWEN_ENC_QUANT_NUM_SEMANTIC; + q->semantic.input_proj_w = gf_load_tensor(&wctx, gf, "tok_enc.vq_semantic.input_proj.weight"); + q->semantic.output_proj_w = gf_load_tensor(&wctx, gf, "tok_enc.vq_semantic.output_proj.weight"); + for (int i = 0; i < QWEN_ENC_QUANT_NUM_SEMANTIC; i++) { + char name[96]; + snprintf(name, sizeof(name), "tok_enc.vq_semantic.%d.codebook", i); + q->semantic.codebooks[i] = gf_load_tensor(&wctx, gf, name); + } + + q->acoustic.num_layers = QWEN_ENC_QUANT_NUM_ACOUSTIC; + q->acoustic.input_proj_w = gf_load_tensor(&wctx, gf, "tok_enc.vq_acoustic.input_proj.weight"); + q->acoustic.output_proj_w = gf_load_tensor(&wctx, gf, "tok_enc.vq_acoustic.output_proj.weight"); + for (int i = 0; i < QWEN_ENC_QUANT_NUM_ACOUSTIC; i++) { + char name[96]; + snprintf(name, sizeof(name), "tok_enc.vq_acoustic.%d.codebook", i); + q->acoustic.codebooks[i] = gf_load_tensor(&wctx, gf, name); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[EncQuantizer] FATAL: backend allocation failed\n"); + return false; + } + q->weight_ctx = wctx.ctx; + q->weight_buf = wctx.buffer; + + fprintf(stderr, + "[EncQuantizer] Loaded: %d codebooks (%d semantic + %d acoustic), %d entries x %d dim, " + "hidden %d, weights %.1f MB\n", + QWEN_ENC_QUANT_TOTAL, QWEN_ENC_QUANT_NUM_SEMANTIC, QWEN_ENC_QUANT_NUM_ACOUSTIC, q->codebook_size, + q->codebook_dim, q->hidden_size, (float) ggml_backend_buffer_get_size(q->weight_buf) / (1024.0f * 1024.0f)); + return true; +} + +static void qwen_quantizer_encode_free(QwenQuantizerEncode * q) { + if (q->weight_buf) { + ggml_backend_buffer_free(q->weight_buf); + q->weight_buf = NULL; + } + if (q->weight_ctx) { + ggml_free(q->weight_ctx); + q->weight_ctx = NULL; + } +} + +// CPU-side host buffers for the RVQ encode loop. Read once from the +// backend at first use and reused across encode calls. The codebooks +// stay on the backend for any GPU-side use, this is just a CPU mirror +// for the per-frame argmin. +struct QwenQuantizerEncodeHost { + int num_layers; + int codebook_size; + int codebook_dim; + int hidden_size; + // input_proj as a row-major [in=hidden_size, out=codebook_dim] Linear. + std::vector input_proj; + // codebooks[l]: [codebook_size, codebook_dim] row-major + std::vector> codebooks; + // sqsum[l]: [codebook_size], precomputed ||c_e||^2 for the argmin trick + std::vector> sqsum; + // output_proj as a row-major [in=codebook_dim, out=hidden_size] Linear. + std::vector output_proj; +}; + +static void qwen_quantizer_encode_host_load(QwenQuantizerEncodeHost * h, + const QwenQuantizerEncodeSide & side, + int codebook_size, + int codebook_dim, + int hidden_size) { + h->num_layers = side.num_layers; + h->codebook_size = codebook_size; + h->codebook_dim = codebook_dim; + h->hidden_size = hidden_size; + + // input_proj weight has ggml shape [1, hidden_size, codebook_dim] (k=1 conv). + // The contiguous memory is [in=hidden_size, out=codebook_dim] row-major. + h->input_proj.resize((size_t) hidden_size * (size_t) codebook_dim); + ggml_backend_tensor_get(side.input_proj_w, h->input_proj.data(), 0, h->input_proj.size() * sizeof(float)); + + h->output_proj.resize((size_t) codebook_dim * (size_t) hidden_size); + ggml_backend_tensor_get(side.output_proj_w, h->output_proj.data(), 0, h->output_proj.size() * sizeof(float)); + + h->codebooks.resize((size_t) side.num_layers); + h->sqsum.resize((size_t) side.num_layers); + for (int l = 0; l < side.num_layers; l++) { + h->codebooks[l].resize((size_t) codebook_size * (size_t) codebook_dim); + ggml_backend_tensor_get(side.codebooks[l], h->codebooks[l].data(), 0, h->codebooks[l].size() * sizeof(float)); + + h->sqsum[l].resize((size_t) codebook_size); + for (int e = 0; e < codebook_size; e++) { + float s = 0.0f; + const float * c = h->codebooks[l].data() + (size_t) e * (size_t) codebook_dim; + for (int d = 0; d < codebook_dim; d++) { + s += c[d] * c[d]; + } + h->sqsum[l][e] = s; + } + } +} + +// Apply a Conv1d k=1 weight (PyTorch shape [out, in, 1], ggml ne=(1, in, out)) +// to a [N, in] row-major buffer, producing a [N, out] row-major buffer. +// The contiguous memory of the ggml weight walks `in` fast and `out` slow, +// matching the numpy view as [out, in] row-major. So `w[o*in + i]` selects +// row o, column i of the underlying [out, in] matrix. +static void qwen_quantizer_encode_linear(const float * w, int in_dim, int out_dim, const float * x, int N, float * y) { + for (int n = 0; n < N; n++) { + const float * xn = x + (size_t) n * (size_t) in_dim; + float * yn = y + (size_t) n * (size_t) out_dim; + for (int o = 0; o < out_dim; o++) { + float acc = 0.0f; + const float * w_row = w + (size_t) o * (size_t) in_dim; + for (int i = 0; i < in_dim; i++) { + acc += xn[i] * w_row[i]; + } + yn[o] = acc; + } + } +} + +// One RVQ side encode loop. Mutates `res` in-place as the residual stream. +// Appends T frames of codebook indices for each of side.num_layers, in +// the order: layer_0[0..T], layer_1[0..T], ..., layer_{L-1}[0..T]. +static void qwen_quantizer_encode_side_loop(const QwenQuantizerEncodeHost * h, + std::vector & res, + int T, + std::vector & codes_out) { + int D = h->codebook_dim; + int E = h->codebook_size; + + std::vector layer_codes((size_t) T); + for (int l = 0; l < h->num_layers; l++) { + const float * codebook = h->codebooks[l].data(); + const float * cb_sqsum = h->sqsum[l].data(); + + // For each frame t: idx = argmin_e (||c_e||^2 - 2 ) + for (int t = 0; t < T; t++) { + const float * r = res.data() + (size_t) t * (size_t) D; + int best_e = 0; + float best_sc = INFINITY; + for (int e = 0; e < E; e++) { + const float * c = codebook + (size_t) e * (size_t) D; + float dot = 0.0f; + for (int d = 0; d < D; d++) { + dot += r[d] * c[d]; + } + float sc = cb_sqsum[e] - 2.0f * dot; + if (sc < best_sc) { + best_sc = sc; + best_e = e; + } + } + layer_codes[t] = best_e; + // Subtract centroid from residual in place + const float * c = codebook + (size_t) best_e * (size_t) D; + for (int d = 0; d < D; d++) { + res[(size_t) t * (size_t) D + (size_t) d] -= c[d]; + } + } + codes_out.insert(codes_out.end(), layer_codes.begin(), layer_codes.end()); + } +} + +// Full RVQ encode. Takes the post-downsample hidden [T, hidden_size] f32 +// row-major buffer and returns flat codes [K, T] row-major, where K is +// QWEN_ENC_QUANT_TOTAL = 16. +// hidden : [T, hidden_size] f32 row-major (T fast in pseudo, but here +// row-major means index = t*hidden + c, t slow, c fast) +// +// Returns codes flat as [16, T] row-major: codes[k*T + t]. +static std::vector qwen_quantizer_encode_cpu(const QwenQuantizerEncodeHost * sem, + const QwenQuantizerEncodeHost * aco, + const float * hidden, + int T) { + std::vector codes; + codes.reserve((size_t) QWEN_ENC_QUANT_TOTAL * (size_t) T); + + // Project hidden to codebook_dim for each side independently. + int D = sem->codebook_dim; + + // Semantic side + std::vector proj_sem((size_t) T * (size_t) D); + qwen_quantizer_encode_linear(sem->input_proj.data(), sem->hidden_size, D, hidden, T, proj_sem.data()); + qwen_quantizer_encode_side_loop(sem, proj_sem, T, codes); + + // Acoustic side + std::vector proj_aco((size_t) T * (size_t) D); + qwen_quantizer_encode_linear(aco->input_proj.data(), aco->hidden_size, D, hidden, T, proj_aco.data()); + qwen_quantizer_encode_side_loop(aco, proj_aco, T, codes); + + return codes; +} diff --git a/src/sampling.h b/src/sampling.h new file mode 100644 index 0000000..fb53868 --- /dev/null +++ b/src/sampling.h @@ -0,0 +1,201 @@ +#pragma once +// sampling.h: token sampling for the Talker LM and the CodePredictor MTP +// head. Pipeline matches the HuggingFace generate() chain in F32 : +// repetition_penalty -> temperature -> top_k -> top_p -> softmax -> multinomial +// The multinomial uniform draw comes from philox_uniform_fill so the +// sequence stays byte for byte aligned with the patched torch.multinomial +// in tests/debug-tts-cossim.py. +// +// apply_suppress is exposed separately so callers can mask the codec +// reserved range before invoking the sampler. The Talker masks +// [vocab - 1024, vocab) except codec_eos before calling sample_top_k_p, +// the CodePredictor does not need any suppression. + +#include "philox.h" + +#include +#include +#include +#include +#include + +struct TokenProb { + int id; + float prob; +}; + +// Mask logits in [lo, hi) to -inf, except keep is left untouched. +static inline void apply_suppress(float * logits, int V, int lo, int hi, int keep) { + if (lo < 0) { + lo = 0; + } + if (hi > V) { + hi = V; + } + for (int i = lo; i < hi; i++) { + if (i == keep) { + continue; + } + logits[i] = -INFINITY; + } +} + +// Repetition penalty over unique tokens in history (HF rule): +// if score >= 0 -> score / penalty +// if score < 0 -> score * penalty +// Each token is touched at most once per call. +static inline void apply_repetition_penalty(float * logits, + int V, + const int32_t * history, + int n_history, + float penalty) { + if (penalty == 1.0f || n_history <= 0) { + return; + } + static thread_local std::vector seen_buf; + seen_buf.assign((size_t) V, 0); + for (int h = 0; h < n_history; h++) { + int32_t tok = history[h]; + if (tok < 0 || tok >= V) { + continue; + } + if (seen_buf[(size_t) tok]) { + continue; + } + seen_buf[(size_t) tok] = 1; + float s = logits[tok]; + logits[tok] = (s < 0.0f) ? s * penalty : s / penalty; + } +} + +// Stochastic sampler. Pipeline mirrors HF generate() in F32 : +// 1. repetition_penalty(history) +// 2. temperature divide +// 3. top_k mask (skipped when k <= 0 or k >= V) +// 4. top_p nucleus mask (skipped when p >= 1.0) +// 5. softmax +// 6. multinomial via philox_uniform_fill(seed, philox_subseq, 0) +// +// Greedy path : temperature <= 0 returns argmax over the suppressed +// logits, no rep_pen, no philox draw. +// +// Buffers are thread_local to avoid alloc per token. +static int sample_top_k_p(float * logits, + int V, + float temperature, + int top_k, + float top_p, + float rep_pen, + const int32_t * history, + int n_history, + int64_t seed, + int64_t philox_subseq, + float * dump_u_out) { + if (temperature <= 0.0f) { + if (dump_u_out) { + *dump_u_out = -1.0f; + } + return (int) (std::max_element(logits, logits + V) - logits); + } + + apply_repetition_penalty(logits, V, history, n_history, rep_pen); + + float inv_temp = 1.0f / temperature; + for (int i = 0; i < V; i++) { + logits[i] *= inv_temp; + } + + static thread_local std::vector tmp_buf; + static thread_local std::vector sorted_buf; + + if (top_k > 0 && top_k < V) { + tmp_buf.resize((size_t) V); + std::memcpy(tmp_buf.data(), logits, (size_t) V * sizeof(float)); + std::nth_element(tmp_buf.begin(), tmp_buf.begin() + (top_k - 1), tmp_buf.end(), std::greater()); + float threshold = tmp_buf[(size_t) (top_k - 1)]; + for (int i = 0; i < V; i++) { + if (logits[i] < threshold) { + logits[i] = -INFINITY; + } + } + } + + if (top_p > 0.0f && top_p < 1.0f) { + // Full vocab softmax for an exact cumsum boundary (matches HF + // TopPLogitsWarper which softmaxes the sorted tensor in place). + float max_logit = -INFINITY; + for (int i = 0; i < V; i++) { + if (logits[i] > max_logit) { + max_logit = logits[i]; + } + } + float sum_exp = 0.0f; + for (int i = 0; i < V; i++) { + sum_exp += expf(logits[i] - max_logit); + } + float inv_sum = 1.0f / sum_exp; + + // Compact only tokens above a relative cutoff. exp(-16) is ~1e-7 + // so the dropped mass stays well below any reachable nucleus + // boundary, and the sort runs on a small set (typically less + // than top_k entries). + float cutoff = max_logit - 16.0f; + sorted_buf.clear(); + for (int i = 0; i < V; i++) { + if (logits[i] >= cutoff) { + float prob = expf(logits[i] - max_logit) * inv_sum; + sorted_buf.push_back({ i, prob }); + } else { + logits[i] = -INFINITY; + } + } + + int K = (int) sorted_buf.size(); + if (K > 0) { + std::sort(sorted_buf.begin(), sorted_buf.end(), + [](const TokenProb & a, const TokenProb & b) { return a.prob > b.prob; }); + // HF convention : keep tokens until the cumulative probability + // crosses top_p, drop the rest. Test before accumulate so the + // first crossing entry is kept. + float cum = 0.0f; + for (int i = 0; i < K; i++) { + if (i > 0 && cum >= top_p) { + logits[sorted_buf[(size_t) i].id] = -INFINITY; + } + cum += sorted_buf[(size_t) i].prob; + } + } + } + + // Final softmax then philox driven multinomial. We keep the + // unnormalized exponentials and draw u in [0, sum) which is + // mathematically identical to softmax + u in [0, 1) but spares a + // pass. + float max_val = -INFINITY; + for (int i = 0; i < V; i++) { + if (logits[i] > max_val) { + max_val = logits[i]; + } + } + float sum = 0.0f; + for (int i = 0; i < V; i++) { + logits[i] = expf(logits[i] - max_val); + sum += logits[i]; + } + + float u = 0.0f; + philox_uniform_fill(seed, philox_subseq, 0u, &u, 1); + if (dump_u_out) { + *dump_u_out = u; + } + + float r = u * sum; + float acc = 0.0f; + for (int i = 0; i < V; i++) { + acc += logits[i]; + if (acc >= r) { + return i; + } + } + return V - 1; +} diff --git a/src/seanet-encoder.h b/src/seanet-encoder.h new file mode 100644 index 0000000..7a03fa7 --- /dev/null +++ b/src/seanet-encoder.h @@ -0,0 +1,206 @@ +#pragma once +// seanet-encoder.h: SEANet conv stack for the Qwen3-TTS encoder side. +// +// Mimi-style structure adapted to T-first ggml. The audio waveform enters +// at 24 kHz mono, gets downsampled by [4, 5, 6, 8] (cumulative 960x), and +// exits at 512 channels @ 25 Hz. A final downsample conv (k=4 stride=2) +// lives in encoder-downsample.h and brings the rate to 12.5 Hz. +// +// Structure: +// init : MimiConv1d k=7, 1 -> 64, causal stride=1 +// for ratio in [4, 5, 6, 8] (Python loop reversed: ratios 8,6,5,4 in +// downsampling order applied to the audio): +// resnet block: ELU -> Conv1d k=3 d=1 dim/2 -> ELU -> Conv1d k=1 dim +// ELU +// Conv1d k=2*ratio, stride=ratio, channels x2 +// last : MimiConv1d k=3, 1024 -> 512, causal stride=1 +// +// Apply order on a 24 kHz waveform: +// audio -> init (1->64) -> stage 0 (8x, 64->128) -> stage 1 (6x, 128->256) +// -> stage 2 (5x, 256->512) -> stage 3 (4x, 512->1024) -> last (1024->512) +// Total downsample = 8 * 6 * 5 * 4 = 960. The 12.5 Hz rate is reached after +// the final downsample conv (factor 2 more in encoder-downsample.h). + +#include "causal-trans-conv.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include + +#define QWEN_SEANET_NUM_STAGES 4 + +struct QwenSEANetResNet { + // First conv inside the residual: depthwise reduction by config.compress + // block_0: k=residual_kernel_size (3), stride=1, dilation=1, dim -> dim/2 + struct ggml_tensor * c0_w; + struct ggml_tensor * c0_b; + // Second conv: pointwise back to dim + // block_1: k=1, dim/2 -> dim + struct ggml_tensor * c1_w; + struct ggml_tensor * c1_b; +}; + +struct QwenSEANetStage { + QwenSEANetResNet resnet; + // Downsampling conv: k=2*ratio, stride=ratio, dim -> dim*2 + struct ggml_tensor * down_w; + struct ggml_tensor * down_b; + int ratio; + int in_ch; + int out_ch; +}; + +struct QwenSEANetEncoder { + // Initial conv: k=7, 1 -> num_filters (64), causal stride=1 + struct ggml_tensor * init_w; + struct ggml_tensor * init_b; + + QwenSEANetStage stages[QWEN_SEANET_NUM_STAGES]; + + // Last conv: k=last_kernel_size (3), final_dim -> hidden_size (512) + struct ggml_tensor * last_w; + struct ggml_tensor * last_b; + + int kernel_size; // 7 + int residual_kernel_size; // 3 + int last_kernel_size; // 3 + int num_filters; // 64 + int compress; // 2 + int hidden_size; // 512 + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +// Read encoder hyperparameters and bind every SEANet tensor on the backend. +static bool qwen_seanet_encoder_load(QwenSEANetEncoder * s, const GGUFModel & gf, ggml_backend_t backend) { + s->kernel_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.kernel_size"); + s->residual_kernel_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.residual_kernel_size"); + s->last_kernel_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.last_kernel_size"); + s->num_filters = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.num_filters"); + s->compress = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.compress"); + s->hidden_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.encoder.hidden_size"); + + // The config stores upsampling_ratios in upsample order [8, 6, 5, 4]. + // Python downsampling iterates `reversed(upsampling_ratios)` = [4, 5, 6, 8], + // so stage 0 applies ratio 4, stage 1 ratio 5, stage 2 ratio 6, stage 3 + // ratio 8. Cumulative downsample is 4*5*6*8 = 960. + int ratios[QWEN_SEANET_NUM_STAGES]; + { + const auto & arr = gf_get_array_u32(gf, "qwen3-tts-tokenizer.encoder.upsampling_ratios"); + if ((int) arr.size() != QWEN_SEANET_NUM_STAGES) { + fprintf(stderr, "[SEANet] FATAL: upsampling_ratios has %d entries, expected %d\n", (int) arr.size(), + QWEN_SEANET_NUM_STAGES); + return false; + } + for (int i = 0; i < QWEN_SEANET_NUM_STAGES; i++) { + ratios[i] = (int) arr[QWEN_SEANET_NUM_STAGES - 1 - i]; + } + } + + int n_tensors = 4 // init wb + last wb + + QWEN_SEANET_NUM_STAGES * 6 // 4 resnet wb + 2 down wb per stage + + 4; // headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + s->init_w = gf_load_tensor(&wctx, gf, "tok_enc.conv.0.weight"); + s->init_b = gf_load_tensor(&wctx, gf, "tok_enc.conv.0.bias"); + + // Stage indexing follows the Python ModuleList layout : + // res block at py idx {1, 4, 7, 10}, two convs at sub-idx 1 and 3 + // downsample conv at py idx {3, 6, 9, 12} + static const int RES_PY_IDX[] = { 1, 4, 7, 10 }; + static const int DOWN_PY_IDX[] = { 3, 6, 9, 12 }; + + int dim = s->num_filters; + for (int i = 0; i < QWEN_SEANET_NUM_STAGES; i++) { + QwenSEANetStage & stg = s->stages[i]; + stg.ratio = ratios[i]; + stg.in_ch = dim; + stg.out_ch = dim * 2; + + char name[80]; + snprintf(name, sizeof(name), "tok_enc.res.%d.blk.1.weight", RES_PY_IDX[i]); + stg.resnet.c0_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_enc.res.%d.blk.1.bias", RES_PY_IDX[i]); + stg.resnet.c0_b = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_enc.res.%d.blk.3.weight", RES_PY_IDX[i]); + stg.resnet.c1_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_enc.res.%d.blk.3.bias", RES_PY_IDX[i]); + stg.resnet.c1_b = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "tok_enc.conv.%d.weight", DOWN_PY_IDX[i]); + stg.down_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_enc.conv.%d.bias", DOWN_PY_IDX[i]); + stg.down_b = gf_load_tensor(&wctx, gf, name); + + dim = stg.out_ch; + } + + s->last_w = gf_load_tensor(&wctx, gf, "tok_enc.conv.14.weight"); + s->last_b = gf_load_tensor(&wctx, gf, "tok_enc.conv.14.bias"); + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[SEANet] FATAL: backend allocation failed\n"); + return false; + } + s->weight_ctx = wctx.ctx; + s->weight_buf = wctx.buffer; + + fprintf(stderr, + "[SEANet] Loaded: 4 stages (ratios %d/%d/%d/%d), %d -> %d channels, " + "weights %.1f MB\n", + ratios[0], ratios[1], ratios[2], ratios[3], s->num_filters, s->hidden_size, + (float) ggml_backend_buffer_get_size(s->weight_buf) / (1024.0f * 1024.0f)); + return true; +} + +static void qwen_seanet_encoder_free(QwenSEANetEncoder * s) { + if (s->weight_buf) { + ggml_backend_buffer_free(s->weight_buf); + s->weight_buf = NULL; + } + if (s->weight_ctx) { + ggml_free(s->weight_ctx); + s->weight_ctx = NULL; + } +} + +// SEANet ResNet forward: skip; ELU; conv k=3,s=1,d=1, dim->dim/2; ELU; +// conv k=1, dim/2->dim; add(skip). +static struct ggml_tensor * qwen_seanet_resnet_forward(struct ggml_context * ctx, + const QwenSEANetResNet * ru, + struct ggml_tensor * x, + int residual_kernel_size) { + struct ggml_tensor * skip = x; + x = ggml_elu(ctx, x); + x = qwen_causal_conv1d(ctx, ru->c0_w, ru->c0_b, x, residual_kernel_size, 1, 1); + x = ggml_elu(ctx, x); + x = qwen_causal_conv1d(ctx, ru->c1_w, ru->c1_b, x, 1, 1, 1); + return ggml_add(ctx, skip, x); +} + +// Full SEANet forward. +// x: [T_audio, 1] f32 T-first (mono waveform) +// returns [T_audio / 960, 512] f32 T-first. +static struct ggml_tensor * qwen_seanet_encoder_forward(struct ggml_context * ctx, + const QwenSEANetEncoder * s, + struct ggml_tensor * x) { + x = qwen_causal_conv1d(ctx, s->init_w, s->init_b, x, s->kernel_size, 1, 1); + + for (int i = 0; i < QWEN_SEANET_NUM_STAGES; i++) { + const QwenSEANetStage & stg = s->stages[i]; + x = qwen_seanet_resnet_forward(ctx, &stg.resnet, x, s->residual_kernel_size); + x = ggml_elu(ctx, x); + x = qwen_causal_conv1d(ctx, stg.down_w, stg.down_b, x, 2 * stg.ratio, 1, stg.ratio); + } + + x = ggml_elu(ctx, x); + x = qwen_causal_conv1d(ctx, s->last_w, s->last_b, x, s->last_kernel_size, 1, 1); + return x; +} diff --git a/src/speaker-encoder-extract.h b/src/speaker-encoder-extract.h new file mode 100644 index 0000000..0eeb0e9 --- /dev/null +++ b/src/speaker-encoder-extract.h @@ -0,0 +1,156 @@ +#pragma once +// speaker-encoder-extract.h : end to end speaker embedding extraction +// from a WAV path. Loads, mono-mixes and resamples to 24 kHz, reflect +// pads by (n_fft - hop) / 2 = 384 samples, builds the fused mel + ECAPA +// graph and returns the f32 [enc_dim] embedding. +// +// Mirrors qwen_tts.core.models.modeling_qwen3_tts.extract_speaker_embedding : +// +// audio = librosa.load(path, sr=None, mono=True)[0] +// audio_24k = librosa.resample(audio, orig_sr=sr, target_sr=24000) +// mels = mel_spectrogram(audio_24k, n_fft=1024, n_mels=128, sr=24000, +// hop=256, win=1024, fmin=0, fmax=12000, center=False) +// spk_emb = speaker_encoder(mels)[0] +// +// Memory layout : the audio waveform input is passed as a regular ggml +// input tensor [T_pad] f32 living on the talker backend. Caller owns the +// returned vector. The graph context is freed after each call. + +#include "audio-io.h" +#include "audio-mel.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "speaker-encoder-forward.h" +#include "speaker-encoder-weights.h" + +#include +#include +#include +#include + +// Public entry point. Returns true on success, fills emb_out with the +// 2048-dim f32 embedding. Returns false on any IO or graph failure. +static bool speaker_encoder_extract(const SpeakerEncoderWeights * sw, + ggml_backend_sched_t sched, + const char * wav_path, + std::vector & emb_out) { + if (sw->weight_buf == NULL) { + fprintf(stderr, "[SpkExtract] FATAL: speaker encoder weights not loaded\n"); + return false; + } + + AudioMelConfig mel_cfg; + mel_cfg.sample_rate = sw->sample_rate; + mel_cfg.n_fft = 1024; + mel_cfg.hop = 256; + mel_cfg.n_mels = sw->mel_dim; + mel_cfg.fmin = 0.0f; + mel_cfg.fmax = 12000.0f; + + // Load WAV, mono mix, resample to 24 kHz. audio_read_mono allocates + // with malloc, wrap in a unique_ptr for clean release. + int T_in = 0; + float * raw = audio_read_mono(wav_path, sw->sample_rate, &T_in); + if (!raw || T_in <= 0) { + fprintf(stderr, "[SpkExtract] FATAL: cannot read WAV '%s'\n", wav_path); + if (raw) { + std::free(raw); + } + return false; + } + std::unique_ptr raw_holder(raw, std::free); + + const int pad = (mel_cfg.n_fft - mel_cfg.hop) / 2; // 384 + const int T_pad = T_in + 2 * pad; + if (T_in < pad + 1) { + fprintf(stderr, "[SpkExtract] FATAL: audio too short (%d samples) for reflect pad %d\n", T_in, pad); + return false; + } + + // Reflect pad on the host so the graph just consumes a flat [T_pad] + // input. PyTorch reflect copies samples [1..pad] then [T-2..T-pad-1] + // into the padded edges (the boundary sample itself is not duplicated). + std::vector audio_padded((size_t) T_pad); + for (int i = 0; i < pad; i++) { + audio_padded[(size_t) i] = raw[pad - i]; + } + std::memcpy(audio_padded.data() + pad, raw, (size_t) T_in * sizeof(float)); + for (int i = 0; i < pad; i++) { + audio_padded[(size_t) (pad + T_in + i)] = raw[T_in - 2 - i]; + } + + // Bake CPU constants once per call : Hann, DFT, mel basis. The cost + // is dominated by the DFT precompute which is 524 KB of f32. + AudioMelConstants mel_c; + audio_mel_compute_constants(mel_cfg, mel_c); + + // Build the graph context. mel + ECAPA accounts for ~150 nodes per + // SE-Res2Net block + 30 for the mel front end + 60 for ASP and FC. + // 2048 nodes is a comfortable upper bound. + const size_t mem_size = ggml_tensor_overhead() * 4096 + ggml_graph_overhead_custom(2048, false); + ggml_init_params init = {}; + init.mem_size = mem_size; + init.mem_buffer = NULL; + init.no_alloc = true; + struct ggml_context * gctx = ggml_init(init); + + // Graph inputs : audio waveform and 4 mel constants. + struct ggml_tensor * audio_in = ggml_new_tensor_1d(gctx, GGML_TYPE_F32, T_pad); + struct ggml_tensor * hann_in = ggml_new_tensor_1d(gctx, GGML_TYPE_F32, mel_cfg.n_fft); + struct ggml_tensor * dft_re_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, mel_cfg.n_fft, mel_c.n_freq); + struct ggml_tensor * dft_im_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, mel_cfg.n_fft, mel_c.n_freq); + struct ggml_tensor * mel_b_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, mel_c.n_freq, mel_cfg.n_mels); + ggml_set_name(audio_in, "spk.audio_padded"); + ggml_set_name(hann_in, "spk.hann"); + ggml_set_name(dft_re_in, "spk.dft_real"); + ggml_set_name(dft_im_in, "spk.dft_imag"); + ggml_set_name(mel_b_in, "spk.mel_basis"); + // Mark as graph inputs so the scheduler routes them through the + // CPU backend then uploads on demand to the compute backend. + ggml_set_input(audio_in); + ggml_set_input(hann_in); + ggml_set_input(dft_re_in); + ggml_set_input(dft_im_in); + ggml_set_input(mel_b_in); + + struct ggml_tensor * emb = + speaker_encoder_forward(gctx, sw, audio_in, hann_in, dft_re_in, dft_im_in, mel_b_in, mel_cfg); + ggml_set_output(emb); + + struct ggml_cgraph * graph = ggml_new_graph_custom(gctx, 2048, false); + ggml_build_forward_expand(graph, emb); + + // Reset the shared sched before allocating : the talker may have left + // a residual graph state from a previous synthesis call. + ggml_backend_sched_reset(sched); + if (!ggml_backend_sched_alloc_graph(sched, graph)) { + fprintf(stderr, "[SpkExtract] FATAL: graph allocation failed\n"); + ggml_free(gctx); + return false; + } + + // Upload inputs to backend. + ggml_backend_tensor_set(audio_in, audio_padded.data(), 0, (size_t) T_pad * sizeof(float)); + ggml_backend_tensor_set(hann_in, mel_c.hann.data(), 0, mel_c.hann.size() * sizeof(float)); + ggml_backend_tensor_set(dft_re_in, mel_c.dft_real.data(), 0, mel_c.dft_real.size() * sizeof(float)); + ggml_backend_tensor_set(dft_im_in, mel_c.dft_imag.data(), 0, mel_c.dft_imag.size() * sizeof(float)); + ggml_backend_tensor_set(mel_b_in, mel_c.mel_basis.data(), 0, mel_c.mel_basis.size() * sizeof(float)); + + if (ggml_backend_sched_graph_compute(sched, graph) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "[SpkExtract] FATAL: graph compute failed\n"); + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return false; + } + + emb_out.assign((size_t) sw->enc_dim, 0.0f); + ggml_backend_tensor_get(emb, emb_out.data(), 0, (size_t) sw->enc_dim * sizeof(float)); + + ggml_backend_sched_reset(sched); + ggml_free(gctx); + + fprintf(stderr, "[SpkExtract] Extracted %d-dim embedding from %s (%d samples, padded %d)\n", sw->enc_dim, wav_path, + T_in, T_pad); + return true; +} diff --git a/src/speaker-encoder-forward.h b/src/speaker-encoder-forward.h new file mode 100644 index 0000000..add30a6 --- /dev/null +++ b/src/speaker-encoder-forward.h @@ -0,0 +1,312 @@ +#pragma once +// speaker-encoder-forward.h : ECAPA-TDNN forward graph in GGML. +// +// Mirrors qwen_tts.core.models.modeling_qwen3_tts.Qwen3TTSSpeakerEncoder +// for the single utterance unbatched path. The forward fuses the mel +// spectrogram extraction so the whole pipeline lives in one graph : +// +// audio [T_pad] f32 +// -> mel [128, T_frames] (audio-mel.h) +// -> conv0 TDNN k=5 + ReLU [512, T_frames] +// -> SE-Res2Net dil=2 [512, T_frames] +// -> SE-Res2Net dil=3 [512, T_frames] +// -> SE-Res2Net dil=4 [512, T_frames] +// -> cat blk[1..3] + MFA k=1 + ReLU [1536, T_frames] +// -> ASP attentive pooling [3072, 1] +// -> FC k=1 [2048, 1] +// -> squeeze [2048] +// +// Tensor convention : [C, T] inside the graph (ne[0]=C, ne[1]=T) so that +// ggml_im2col reads each Conv1d along the time axis and ggml_mul_mat +// contracts over the input channel axis. This matches the layout the +// upstream PyTorch code uses after its (1, 2) transpose. + +#include "audio-mel.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "speaker-encoder-weights.h" + +#include +#include +#include + +// Conv1d k=K with padding="same" mode="reflect" + bias add. The weight +// tensor lives in upstream layout [K, in_c, out_c]. We implement it with +// reflect pad + im2col + matmul. +// +// x [in_c, T] input +// w [K, in_c, out_c] weights +// b [out_c] bias (broadcast over T) +// Returns [out_c, T] +// +// Padding for "same" with kernel K and dilation d is (K - 1) * d / 2 on +// each side (PyTorch convention, kernel size always odd here so the +// division is exact). +static struct ggml_tensor * spk_conv1d_same(struct ggml_context * ctx, + struct ggml_tensor * x, + struct ggml_tensor * w, + struct ggml_tensor * b, + int dilation) { + const int K = (int) w->ne[0]; + const int OC = (int) w->ne[2]; + const int pad = ((K - 1) * dilation) / 2; + + // ggml_pad_reflect_1d pads ne[0]. Our temporal axis is ne[1], so + // transpose first, pad, transpose back. + struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, in_c] + if (pad > 0) { + x_t = ggml_pad_reflect_1d(ctx, x_t, pad, pad); // [T+2*pad, in_c] + } + x_t = ggml_cont(ctx, ggml_transpose(ctx, x_t)); // [in_c, T+2*pad] + + // Reshape as [W=T_pad, H=1, IC=in_c, N=1] for ggml_im2col 1D. + struct ggml_tensor * x4d = ggml_reshape_4d(ctx, x_t, x_t->ne[1], 1, x_t->ne[0], 1); + struct ggml_tensor * dummy = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, K, 1, x_t->ne[0], 1); + ggml_set_name(dummy, "spk.im2col_kernel"); + + // im2col output shape with is_2D=false : ne = (a.ne[1]*a.ne[0], OW, b.ne[2], 1) + // = (K, T_out, in_c, 1) here. To matmul against the [K*in_c, out_c] + // weight we need [K*in_c, T_out] which means permuting (K, IC, T) + // before flattening. ggml_permute(0, 2, 1, 3) swaps axes 1 and 2. + struct ggml_tensor * col = ggml_im2col(ctx, dummy, x4d, 1, 1, 0, 0, dilation, 1, false, GGML_TYPE_F32); + int T_out = (int) col->ne[1]; + int IC = (int) col->ne[2]; + col = ggml_cont(ctx, ggml_permute(ctx, col, 0, 2, 1, 3)); // [K, IC, T, 1] + col = ggml_reshape_2d(ctx, col, K * IC, T_out); + + // weight [K, in_c, out_c] reshape as [K * in_c, out_c]. mul_mat + // returns [out_c, T_out]. + struct ggml_tensor * w2d = ggml_reshape_2d(ctx, w, K * (int) w->ne[1], OC); + struct ggml_tensor * y = ggml_mul_mat(ctx, w2d, col); + ggml_mul_mat_set_prec(y, GGML_PREC_F32); + + // Add bias broadcast over T_out. b is [out_c], reshape [out_c, 1]. + struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, OC, 1); + y = ggml_add(ctx, y, b2d); + return y; +} + +// TDNN block : Conv1d(same, reflect) + ReLU. Used both as the conv0 +// frontend (k=5) and inside SE-Res2Net (k=1) and the MFA / ASP TDNNs. +static struct ggml_tensor * spk_tdnn(struct ggml_context * ctx, + const SpkEncTDNN & t, + struct ggml_tensor * x, + int dilation) { + struct ggml_tensor * y = spk_conv1d_same(ctx, x, t.weight, t.bias, dilation); + y = ggml_relu(ctx, y); + return y; +} + +// Res2Net block : split the channel axis in 8 chunks. chunk 0 passes +// through, chunk 1 goes through TDNN[0], chunks 2..7 mix with the +// previous chunk output before going through TDNN[i-1]. The 7 TDNN +// branches share dilation but operate on hidden / 8 channels each. +// +// x [C, T] +// Returns [C, T] +static struct ggml_tensor * spk_res2net(struct ggml_context * ctx, + const SpkEncRes2Net & rn, + struct ggml_tensor * x, + int dilation, + int scale) { + const int C = (int) x->ne[0]; + const int T = (int) x->ne[1]; + const int Cs = C / scale; + + std::vector outs; + outs.reserve(scale); + + // chunk i is the slice along ne[0] of width Cs starting at i * Cs. + auto chunk = [&](int i) -> struct ggml_tensor * { + return ggml_view_2d(ctx, x, Cs, T, x->nb[1], (size_t) (i * Cs) * x->nb[0]); + }; + + struct ggml_tensor * prev = NULL; + for (int i = 0; i < scale; i++) { + struct ggml_tensor * c = ggml_cont(ctx, chunk(i)); + if (i == 0) { + outs.push_back(c); + continue; + } + struct ggml_tensor * inp = c; + if (i >= 2) { + inp = ggml_add(ctx, c, prev); + } + struct ggml_tensor * y = spk_conv1d_same(ctx, inp, rn.weight[i - 1], rn.bias[i - 1], dilation); + y = ggml_relu(ctx, y); + outs.push_back(y); + prev = y; + } + + // Concat along ne[0]. ggml_concat with dim=0 stacks along the + // fastest axis. Build the concat tree iteratively. + struct ggml_tensor * acc = outs[0]; + for (int i = 1; i < scale; i++) { + acc = ggml_concat(ctx, acc, outs[i], 0); + } + return acc; +} + +// Squeeze and Excitation : compute the temporal mean per channel, +// project down to se_c with a 1x1 conv + ReLU, project back up to +// out_c with a 1x1 conv + sigmoid, then scale the input by the gate +// broadcast over T. +static struct ggml_tensor * spk_se(struct ggml_context * ctx, const SpkEncSE & se, struct ggml_tensor * x) { + const int T = (int) x->ne[1]; + + // Mean over T, keep dim. ggml_mean reduces along ne[0], so transpose + // to put T on ne[0], reduce, transpose back. + struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, C] + struct ggml_tensor * mean = ggml_mean(ctx, x_t); // [1, C] + mean = ggml_cont(ctx, ggml_transpose(ctx, mean)); // [C, 1] + + // conv1 1x1 reduces C -> se_c. dilation 1, padding "same" trivial + // since k=1. + struct ggml_tensor * h = spk_conv1d_same(ctx, mean, se.conv1_w, se.conv1_b, 1); + h = ggml_relu(ctx, h); + h = spk_conv1d_same(ctx, h, se.conv2_w, se.conv2_b, 1); + // Sigmoid over [out_c, 1]. + h = ggml_sigmoid(ctx, h); + + // Scale x by the gate. h is [C, 1], x is [C, T]. ggml_mul broadcasts + // ne[1]=1 to T. + struct ggml_tensor * y = ggml_mul(ctx, x, h); + (void) T; + return y; +} + +// SE-Res2Net block : tdnn1 (1x1) -> Res2Net -> tdnn2 (1x1) -> SE plus +// a residual add over the whole stack. +static struct ggml_tensor * spk_block(struct ggml_context * ctx, + const SpkEncBlock & blk, + struct ggml_tensor * x, + int res2net_scale) { + struct ggml_tensor * residual = x; + struct ggml_tensor * h = spk_tdnn(ctx, blk.tdnn1, x, 1); + h = spk_res2net(ctx, blk.res2net, h, blk.dilation, res2net_scale); + h = spk_tdnn(ctx, blk.tdnn2, h, 1); + h = spk_se(ctx, blk.se, h); + return ggml_add(ctx, h, residual); +} + +// Attentive Statistical Pooling : compute global mean and std along T, +// concat with x, run an attention TDNN + tanh + 1x1 conv, softmax along +// T, recompute weighted mean and std, return the [2C, 1] concat. +// +// x [C, T] +// Returns [2C, 1] +static struct ggml_tensor * spk_asp(struct ggml_context * ctx, const SpkEncASP & asp, struct ggml_tensor * x) { + const int C = (int) x->ne[0]; + const int T = (int) x->ne[1]; + + // Mean and std over T axis. The mask reduction is uniform 1/T. + // mean : [C, 1] + struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); + struct ggml_tensor * mean = ggml_mean(ctx, x_t); + mean = ggml_cont(ctx, ggml_transpose(ctx, mean)); + + // var = mean( (x - mean)^2 ) over T, then std = sqrt(clamp(var, eps)). + // Broadcasting (x - mean) requires mean repeated to T. ggml_repeat + // handles this when shapes are compatible. + struct ggml_tensor * mean_T = ggml_repeat(ctx, mean, x); + struct ggml_tensor * centered = ggml_sub(ctx, x, mean_T); + struct ggml_tensor * var_t = ggml_cont(ctx, ggml_transpose(ctx, ggml_sqr(ctx, centered))); + struct ggml_tensor * var = ggml_mean(ctx, var_t); + var = ggml_cont(ctx, ggml_transpose(ctx, var)); + var = ggml_scale_bias(ctx, var, 1.0f, 1e-12f); + struct ggml_tensor * std = ggml_sqrt(ctx, var); + + // Build [x, mean_repeat, std_repeat] concat along channel axis. + struct ggml_tensor * std_T = ggml_repeat(ctx, std, x); + struct ggml_tensor * cat = ggml_concat(ctx, x, mean_T, 0); + cat = ggml_concat(ctx, cat, std_T, 0); // [3C, T] + + // Attention TDNN : 3C -> attn_c, ReLU, then tanh, then 1x1 conv + // attn_c -> C. Upstream applies tanh on the TDNN output before the + // second conv ; the TDNN itself already runs ReLU so the order is + // ReLU then tanh which is unusual but mirrored faithfully. + struct ggml_tensor * a = spk_tdnn(ctx, asp.tdnn, cat, 1); + a = ggml_tanh(ctx, a); + a = spk_conv1d_same(ctx, a, asp.conv_w, asp.conv_b, 1); + + // Softmax along T axis. ggml_soft_max reduces ne[0], transpose first. + struct ggml_tensor * a_t = ggml_cont(ctx, ggml_transpose(ctx, a)); // [T, C] + struct ggml_tensor * w_t = ggml_soft_max(ctx, a_t); + struct ggml_tensor * w = ggml_cont(ctx, ggml_transpose(ctx, w_t)); // [C, T] + + // Weighted mean : sum(w * x) over T, w already sums to 1 over T. + struct ggml_tensor * wx = ggml_mul(ctx, w, x); + struct ggml_tensor * wx_t = ggml_cont(ctx, ggml_transpose(ctx, wx)); + // ggml_mean averages over ne[0]=T, giving 1/T scaling. We want the + // un-normalized sum since w already encodes the soft selection + // probability, so multiply back by T. + struct ggml_tensor * w_mean = ggml_mean(ctx, wx_t); + w_mean = ggml_scale(ctx, w_mean, (float) T); + w_mean = ggml_cont(ctx, ggml_transpose(ctx, w_mean)); // [C, 1] + + // Weighted std : sum(w * (x - w_mean)^2) over T. + struct ggml_tensor * w_mean_T = ggml_repeat(ctx, w_mean, x); + struct ggml_tensor * dev = ggml_sub(ctx, x, w_mean_T); + struct ggml_tensor * w_var_in = ggml_mul(ctx, w, ggml_sqr(ctx, dev)); + struct ggml_tensor * w_var_t = ggml_cont(ctx, ggml_transpose(ctx, w_var_in)); + struct ggml_tensor * w_var = ggml_mean(ctx, w_var_t); + w_var = ggml_scale(ctx, w_var, (float) T); + w_var = ggml_cont(ctx, ggml_transpose(ctx, w_var)); + w_var = ggml_scale_bias(ctx, w_var, 1.0f, 1e-12f); + struct ggml_tensor * w_std = ggml_sqrt(ctx, w_var); + + // Stack [w_mean, w_std] along channel -> [2C, 1]. Time axis already + // collapsed. + struct ggml_tensor * stats = ggml_concat(ctx, w_mean, w_std, 0); + (void) C; + return stats; +} + +// Full speaker encoder forward graph. Assumes the audio waveform has +// already been resampled to sr=24000 and reflect padded by +// (n_fft - hop) / 2 on each side. The padded buffer must outlive the +// graph compute call. +// +// Inputs : +// audio_padded [T_pad] f32, host or backend tensor +// mel constants hann/dft_real/dft_imag/mel_basis backend tensors +// Output : [enc_dim] f32, the speaker embedding (typically 2048 dims). +static struct ggml_tensor * speaker_encoder_forward(struct ggml_context * ctx, + const SpeakerEncoderWeights * sw, + struct ggml_tensor * audio_padded, + struct ggml_tensor * hann, + struct ggml_tensor * dft_real, + struct ggml_tensor * dft_imag, + struct ggml_tensor * mel_basis, + const AudioMelConfig & mel_cfg) { + // Mel : [n_mels=128, T_frames] + struct ggml_tensor * mel = audio_mel_build_graph(ctx, audio_padded, hann, dft_real, dft_imag, mel_basis, mel_cfg); + + // Frontend conv0 TDNN k=5 + ReLU : 128 -> 512, T preserved. + struct ggml_tensor * h = spk_tdnn(ctx, sw->conv0, mel, 1); + + // Three SE-Res2Net blocks at dilations 2, 3, 4. + struct ggml_tensor * b1 = spk_block(ctx, sw->blocks[0], h, sw->res2net_scale); + struct ggml_tensor * b2 = spk_block(ctx, sw->blocks[1], b1, sw->res2net_scale); + struct ggml_tensor * b3 = spk_block(ctx, sw->blocks[2], b2, sw->res2net_scale); + + // Multi-layer feature aggregation : cat blk1..3 then 1x1 TDNN + ReLU. + struct ggml_tensor * cat = ggml_concat(ctx, b1, b2, 0); + cat = ggml_concat(ctx, cat, b3, 0); // [1536, T] + struct ggml_tensor * mfa = spk_tdnn(ctx, sw->mfa, cat, 1); // [1536, T] + + // Attentive statistical pooling : [1536, T] -> [3072, 1]. + struct ggml_tensor * stats = spk_asp(ctx, sw->asp, mfa); + + // Final FC k=1 : [3072, 1] -> [enc_dim, 1]. + struct ggml_tensor * emb = spk_conv1d_same(ctx, stats, sw->fc_w, sw->fc_b, 1); + + // Squeeze T axis, return [enc_dim]. ggml_cont is required so the sched + // assigns a fresh backend buffer to the graph output rather than + // forwarding a view of the FC bias add. + emb = ggml_reshape_1d(ctx, emb, sw->enc_dim); + emb = ggml_cont(ctx, emb); + ggml_set_name(emb, "spk.embedding"); + return emb; +} diff --git a/src/speaker-encoder-weights.h b/src/speaker-encoder-weights.h new file mode 100644 index 0000000..c0a90bc --- /dev/null +++ b/src/speaker-encoder-weights.h @@ -0,0 +1,208 @@ +#pragma once +// speaker-encoder-weights.h : ECAPA-TDNN x-vector extractor used by the +// Base checkpoint to condition the Talker on a reference voice. +// +// Topology (from qwen_tts.core.models.modeling_qwen3_tts) : +// +// conv0 TDNN k=5, 128 -> 512 (initial frontend) +// blk[1..3] SE-Res2Net (TDNN1 + Res2Net 8-branch + TDNN2 + SE) +// mfa TDNN k=1, 1536 -> 1536 (cat of blk1..3) +// asp attentive statistical pooling, 1536 -> 3072 +// fc Conv1d k=1, 3072 -> 2048 +// +// Weights live on the talker backend buffer next to the talker LM. +// All tensors are stored F32 in the source GGUF and stay F32 when +// quantizing because should_quantize keeps spk_enc as is (small +// channel counts make quantization meaningless here). +// +// Constants : enc_dim 2048 (size of the speaker embedding fed into +// the codec_prefill slot), input mel_dim 128, ECAPA hidden 512, +// res2net scale 8 -> 7 dilated TDNN branches, se hidden 128, +// asp attention 128. + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include + +// Initial TDNN block : Conv1d(in=128, out=512, k=5, padding=same, reflect) +// followed by ReLU. Stored as 3D tensor [k, in_c, out_c] in the GGUF. +struct SpkEncTDNN { + struct ggml_tensor * weight; // [k, in_c, out_c] + struct ggml_tensor * bias; // [out_c] + int k; + int in_c; + int out_c; +}; + +// Squeeze-Excitation attention : conv1 (out -> se), conv2 (se -> out), +// k=1 padding=same. Operates on the temporal mean of the input then +// broadcasts a sigmoid gate back over the time axis. +struct SpkEncSE { + struct ggml_tensor * conv1_w; // [1, out_c, se_c] + struct ggml_tensor * conv1_b; // [se_c] + struct ggml_tensor * conv2_w; // [1, se_c, out_c] + struct ggml_tensor * conv2_b; // [out_c] +}; + +// Res2Net branch : 7 dilated TDNN k=3 conv1d, dilation comes from the +// parent SE-Res2Net block. We keep flat arrays since enc_res2net_scale +// is 8 (which yields scale - 1 = 7 branches). +struct SpkEncRes2Net { + struct ggml_tensor * weight[7]; // each [3, in_c/8, out_c/8] + struct ggml_tensor * bias[7]; // each [out_c/8] +}; + +// SE-Res2Net block : tdnn1 (k=1) -> Res2Net (k=3, dil=d) -> tdnn2 (k=1) +// -> SE attention, plus a residual add over the whole stack. +struct SpkEncBlock { + SpkEncTDNN tdnn1; + SpkEncRes2Net res2net; + SpkEncTDNN tdnn2; + SpkEncSE se; + int dilation; +}; + +// Attentive Statistical Pooling : tdnn maps from 3*1536 to 128 (channels +// concat of [x, mean, std]), conv maps 128 back to 1536. The mask +// branch reduces to a no-op for unbatched single-utterance inference, +// which is the only path the C++ side exposes. +struct SpkEncASP { + SpkEncTDNN tdnn; // [1, 3*hidden, attn] + struct ggml_tensor * conv_w; // [1, attn, hidden] + struct ggml_tensor * conv_b; // [hidden] +}; + +struct SpeakerEncoderWeights { + // Topology constants, sourced from upstream Qwen3TTSSpeakerEncoderConfig. + int enc_dim; // 2048 + int sample_rate; // 24000 + int mel_dim; // 128 + int hidden; // 512 + int mfa_hidden; // 1536 + int asp_attn; // 128 + int se_channels; // 128 + int res2net_scale; // 8 + + // Forward path tensors. blocks[0] is the conv0 TDNN frontend held in + // its own slot for clarity. The three SE-Res2Net stacks live in + // blocks[1..3]. + SpkEncTDNN conv0; // [5, 128, 512] + SpkEncBlock blocks[3]; // SE-Res2Net at dilations 2, 3, 4 + SpkEncTDNN mfa; // [1, 1536, 1536] + SpkEncASP asp; + struct ggml_tensor * fc_w; // [1, 3072, 2048] + struct ggml_tensor * fc_b; // [2048] + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +// Helpers to load each tensor by upstream name. The generic gf_load_tensor +// pulls the row major data, the conv weights are 3D so they keep their +// natural [k, in_c, out_c] layout which matches what ggml_im2col + matmul +// expects when we treat ne[0]=k as the spatial filter axis. +static struct ggml_tensor * spk_load(WeightCtx * wctx, const GGUFModel & gf, const std::string & name) { + return gf_load_tensor(wctx, gf, name); +} + +static void spk_load_tdnn(WeightCtx * wctx, const GGUFModel & gf, const std::string & prefix, SpkEncTDNN & t) { + t.weight = spk_load(wctx, gf, prefix + ".weight"); + t.bias = spk_load(wctx, gf, prefix + ".bias"); + t.k = (int) t.weight->ne[0]; + t.in_c = (int) t.weight->ne[1]; + t.out_c = (int) t.weight->ne[2]; +} + +static void spk_load_se(WeightCtx * wctx, const GGUFModel & gf, const std::string & prefix, SpkEncSE & se) { + se.conv1_w = spk_load(wctx, gf, prefix + ".conv1.weight"); + se.conv1_b = spk_load(wctx, gf, prefix + ".conv1.bias"); + se.conv2_w = spk_load(wctx, gf, prefix + ".conv2.weight"); + se.conv2_b = spk_load(wctx, gf, prefix + ".conv2.bias"); +} + +static void spk_load_res2net(WeightCtx * wctx, const GGUFModel & gf, const std::string & prefix, SpkEncRes2Net & rn) { + for (int i = 0; i < 7; i++) { + char p[64]; + std::snprintf(p, sizeof(p), "%s.%d", prefix.c_str(), i); + rn.weight[i] = spk_load(wctx, gf, std::string(p) + ".weight"); + rn.bias[i] = spk_load(wctx, gf, std::string(p) + ".bias"); + } +} + +static void spk_load_block(WeightCtx * wctx, const GGUFModel & gf, int idx, int dilation, SpkEncBlock & blk) { + char p[64]; + std::snprintf(p, sizeof(p), "spk_enc.blk.%d", idx); + spk_load_tdnn(wctx, gf, std::string(p) + ".tdnn1", blk.tdnn1); + spk_load_res2net(wctx, gf, std::string(p) + ".res2net", blk.res2net); + spk_load_tdnn(wctx, gf, std::string(p) + ".tdnn2", blk.tdnn2); + spk_load_se(wctx, gf, std::string(p) + ".se", blk.se); + blk.dilation = dilation; +} + +static bool speaker_encoder_weights_load(SpeakerEncoderWeights * sw, const GGUFModel & gf, ggml_backend_t backend) { + sw->enc_dim = (int) gf_get_u32(gf, "qwen3-tts.spk_enc.embedding_length"); + sw->sample_rate = (int) gf_get_u32(gf, "qwen3-tts.spk_enc.sample_rate"); + sw->mel_dim = 128; + sw->hidden = 512; + sw->mfa_hidden = 1536; + sw->asp_attn = 128; + sw->se_channels = 128; + sw->res2net_scale = 8; + + // Probe : Base GGUFs ship the speaker encoder, CustomVoice and + // VoiceDesign do not. A missing conv0.weight aborts cleanly. + if (gguf_find_tensor(gf.gguf, "spk_enc.conv0.weight") < 0) { + fprintf(stderr, "[SpeakerEncoder] No spk_enc.conv0.weight, base/clone mode unavailable\n"); + sw->weight_ctx = NULL; + sw->weight_buf = NULL; + return true; + } + + // Roughly 80 tensors total : 1 conv0 + 3 * (2 tdnn + 7 res2net + 4 se) + 1 mfa + // + asp.tdnn + asp.conv + fc, with weight + bias each. Allocate 100 slots + // for safety. + WeightCtx wctx; + wctx_init(&wctx, 100); + + spk_load_tdnn(&wctx, gf, "spk_enc.conv0", sw->conv0); + spk_load_block(&wctx, gf, 1, 2, sw->blocks[0]); + spk_load_block(&wctx, gf, 2, 3, sw->blocks[1]); + spk_load_block(&wctx, gf, 3, 4, sw->blocks[2]); + spk_load_tdnn(&wctx, gf, "spk_enc.mfa", sw->mfa); + spk_load_tdnn(&wctx, gf, "spk_enc.asp.tdnn", sw->asp.tdnn); + sw->asp.conv_w = spk_load(&wctx, gf, "spk_enc.asp.conv.weight"); + sw->asp.conv_b = spk_load(&wctx, gf, "spk_enc.asp.conv.bias"); + sw->fc_w = spk_load(&wctx, gf, "spk_enc.fc.weight"); + sw->fc_b = spk_load(&wctx, gf, "spk_enc.fc.bias"); + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[SpeakerEncoder] FATAL: backend allocation failed\n"); + return false; + } + sw->weight_ctx = wctx.ctx; + sw->weight_buf = wctx.buffer; + + fprintf(stderr, + "[SpeakerEncoder] Loaded: enc_dim=%d sr=%d mel_dim=%d hidden=%d mfa=%d asp_attn=%d se=%d scale=%d\n", + sw->enc_dim, sw->sample_rate, sw->mel_dim, sw->hidden, sw->mfa_hidden, sw->asp_attn, sw->se_channels, + sw->res2net_scale); + return true; +} + +static void speaker_encoder_weights_free(SpeakerEncoderWeights * sw) { + if (sw->weight_buf) { + ggml_backend_buffer_free(sw->weight_buf); + sw->weight_buf = NULL; + } + if (sw->weight_ctx) { + ggml_free(sw->weight_ctx); + sw->weight_ctx = NULL; + } +} diff --git a/src/talker-forward.cpp b/src/talker-forward.cpp new file mode 100644 index 0000000..5361c1c --- /dev/null +++ b/src/talker-forward.cpp @@ -0,0 +1,288 @@ +// talker-forward.cpp : eager prefill graph for the Talker LM. +// +// Mirrors Qwen3TTSTalkerDecoderLayer for TTS-only operation : +// pre-norm, GQA attention with per-head QK-norm, mrope collapsed to +// 1D NEOX (since the three multimodal axes share position ids in TTS +// mode), SwiGLU MLP, two residuals, repeated 28 times, then final +// RMSNorm and codec_head. Eager softmax in F32, no flash-attention, +// no KV cache. +// +// Tensor shapes follow ggml row-major convention : ne[0] is the fastest +// axis. Our input embedding lives as [hidden, T] inside the graph and +// the loader feeds it from a [T, hidden] f32 row-major host buffer +// (which becomes [hidden, T] in ggml after a 2d view because rows on +// the host are contiguous along the hidden axis). + +#include "talker-forward.h" + +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +// Bisect layers dumped when a dump_dir is set. Match the Python hook list +// in tests/debug-tts-cossim.py : 0, 7, 14, 21, 27. +static const int BISECT_LAYERS[] = { 0, 7, 14, 21, 27 }; +static const int N_BISECT_LAYERS = (int) (sizeof(BISECT_LAYERS) / sizeof(BISECT_LAYERS[0])); + +static bool is_bisect_layer(int l) { + for (int i = 0; i < N_BISECT_LAYERS; i++) { + if (BISECT_LAYERS[i] == l) { + return true; + } + } + return false; +} + +// Build the per-layer block. Returns the layer output [hidden, T]. +static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx, + const TalkerWeights * tw, + const TalkerLayer & layer, + struct ggml_tensor * x, + struct ggml_tensor * positions, + struct ggml_tensor * mask, + int T) { + const int n_q_heads = tw->num_attention_heads; + const int n_kv = tw->num_key_value_heads; + const int hd = tw->head_dim; + const float eps = tw->rms_norm_eps; + + // Pre-norm + struct ggml_tensor * h = ggml_rms_norm(ctx, x, eps); + h = ggml_mul(ctx, h, layer.input_norm_w); + + // Q/K/V projections + struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, h); // [n_q_heads*hd, T] + struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, h); // [n_kv*hd, T] + struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, h); // [n_kv*hd, T] + + q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T); // [hd, n_q_heads, T] + k = ggml_reshape_3d(ctx, k, hd, n_kv, T); + v = ggml_reshape_3d(ctx, v, hd, n_kv, T); + + // Per-head QK-norm : RMS over hd, then multiply by [hd] gain. The + // norm operates on ne[0] = hd, identical layout for q (16 heads) and + // k (8 heads), so the same code path covers both. + q = ggml_rms_norm(ctx, q, eps); + q = ggml_mul(ctx, q, layer.attn.q_norm_w); + k = ggml_rms_norm(ctx, k, eps); + k = ggml_mul(ctx, k, layer.attn.k_norm_w); + + // RoPE NEOX (half-split). In TTS-only mode the three mrope axes share + // position ids, so the multimodal interleaved cos/sin collapses to + // plain 1D rotate_half with the same freq base. + q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tw->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + + // Attention layout : [hd, T, n_heads]. ggml_mul_mat broadcasts on + // dims 2/3 when source has fewer heads than destination, which is + // exactly the GQA case with n_q_heads = n_kv * n_rep, no explicit + // repeat_kv needed. + struct ggml_tensor * q_p = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + struct ggml_tensor * k_p = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + struct ggml_tensor * v_p = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [T, n_kv, hd] + + struct ggml_tensor * scores = ggml_mul_mat(ctx, k_p, q_p); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + + float scale = 1.0f / sqrtf((float) hd); + scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); + + struct ggml_tensor * attn = ggml_mul_mat(ctx, v_p, scores); + attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3)); + attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T); + + struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn); + + x = ggml_add(ctx, x, o); + + // MLP block : pre-norm + SwiGLU + residual + struct ggml_tensor * h2 = ggml_rms_norm(ctx, x, eps); + h2 = ggml_mul(ctx, h2, layer.post_attn_norm_w); + + struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, h2); + struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, h2); + gate = ggml_silu(ctx, gate); + struct ggml_tensor * gu = ggml_mul(ctx, gate, up); + struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu); + + x = ggml_add(ctx, x, mlp); + return x; +} + +bool talker_forward_prefill(const TalkerWeights * tw, + ggml_backend_sched_t sched, + const float * input_embed, + int T, + const char * dump_dir, + TalkerForwardOutput * out) { + const int hidden = tw->hidden_size; + const int n_layers = tw->num_hidden_layers; + const int vocab = tw->vocab_size; + + // Dedicated context for graph + IO tensors. Counts approximate : + // per layer : ~30 ops -> ~30 graph nodes + // IO : 4 tensors (input embed, positions, mask, output norm) + // final : norm + codec_head + ~3 dump branches + const int max_nodes = 32 * n_layers + 256; + const size_t graph_arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false); + + struct ggml_init_params gparams = { + graph_arena_bytes, + NULL, + true, + }; + struct ggml_context * gctx = ggml_init(gparams); + if (!gctx) { + fprintf(stderr, "[TalkerForward] FATAL: ggml_init failed\n"); + return false; + } + + // IO tensors : input embedding, positions, causal mask + struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, T); + struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T); + struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, T, T); + ggml_set_name(x_in, "input_embed"); + ggml_set_name(pos_in, "positions"); + ggml_set_name(mask_in, "causal_mask"); + + // Build the layer stack. We collect bisect taps into a vector so the + // graph keeps them as roots and the scheduler does not optimize them + // away. + struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false); + + struct ggml_tensor * h = x_in; + std::vector taps(N_BISECT_LAYERS, NULL); + for (int l = 0; l < n_layers; l++) { + h = talker_layer_forward(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, T); + if (dump_dir && is_bisect_layer(l)) { + for (int i = 0; i < N_BISECT_LAYERS; i++) { + if (BISECT_LAYERS[i] == l) { + char tap_name[64]; + snprintf(tap_name, sizeof(tap_name), "tap_l%d", l); + ggml_set_name(h, tap_name); + ggml_set_output(h); + taps[(size_t) i] = h; + break; + } + } + } + } + + struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, tw->rms_norm_eps); + h_final = ggml_mul(gctx, h_final, tw->norm_w); + ggml_set_name(h_final, "hidden_final"); + ggml_set_output(h_final); + + // codec_head : [hidden, vocab]. ggml_mul_mat returns [vocab, T]. + struct ggml_tensor * logits = ggml_mul_mat(gctx, tw->codec_head_w, h_final); + ggml_set_name(logits, "logits"); + + // Wire dump roots and final outputs into the graph. + if (dump_dir) { + for (int i = 0; i < N_BISECT_LAYERS; i++) { + if (taps[(size_t) i]) { + ggml_build_forward_expand(gf, taps[(size_t) i]); + } + } + ggml_build_forward_expand(gf, h_final); + } + ggml_build_forward_expand(gf, logits); + + // Allocate intermediate buffers via the scheduler. The scheduler + // partitions ops between the GPU and the CPU fallback before + // alloc_graph wires up the buffers, so K-quant get_rows that the + // GPU backend cannot run land on the CPU instead of asserting. + if (!ggml_backend_sched_alloc_graph(sched, gf)) { + fprintf(stderr, "[TalkerForward] FATAL: graph allocation failed\n"); + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return false; + } + + // Upload input embedding (host [T, hidden] -> ggml [hidden, T]). + // The host buffer is row-major with hidden as the inner axis, which + // matches ggml's ne[0]-fastest layout exactly. + ggml_backend_tensor_set(x_in, input_embed, 0, (size_t) T * (size_t) hidden * sizeof(float)); + + // Positions 0..T-1 + { + std::vector pos((size_t) T); + for (int i = 0; i < T; i++) { + pos[(size_t) i] = i; + } + ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t)); + } + + // Causal mask : 0 on/below diagonal, -inf above. Stored row-major + // [T_q, T_k] with T_k as the fast axis (ne[0]). + { + std::vector mask((size_t) T * (size_t) T, -INFINITY); + for (int q = 0; q < T; q++) { + for (int k = 0; k <= q; k++) { + mask[(size_t) q * (size_t) T + (size_t) k] = 0.0f; + } + } + ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(float)); + } + + // Compute through the scheduler so cross-backend op placement stays + // honoured for every node in the graph. + if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "[TalkerForward] FATAL: graph compute failed\n"); + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return false; + } + + // Bisect dumps : pull each tap [hidden, T] back to host as [T, hidden] + // f32 row-major (which is exactly the ggml byte layout). + if (dump_dir) { + DebugDumper d; + debug_init(&d, dump_dir); + std::vector buf((size_t) T * (size_t) hidden); + for (int i = 0; i < N_BISECT_LAYERS; i++) { + if (!taps[(size_t) i]) { + continue; + } + ggml_backend_tensor_get(taps[(size_t) i], buf.data(), 0, buf.size() * sizeof(float)); + char name[64]; + snprintf(name, sizeof(name), "talker-hidden-prefill-l%d", BISECT_LAYERS[i]); + debug_dump_2d(&d, name, buf.data(), T, hidden); + } + ggml_backend_tensor_get(h_final, buf.data(), 0, buf.size() * sizeof(float)); + debug_dump_2d(&d, "talker-hidden-prefill-final", buf.data(), T, hidden); + } + + // Pull the last position : final hidden + logits + out->hidden = hidden; + out->vocab = vocab; + out->hidden_last.assign((size_t) hidden, 0.0f); + out->logits_last.assign((size_t) vocab, 0.0f); + { + // Logits ggml shape : [vocab, T]. We want the row at column T-1. + // Layout : data[t * vocab + v], so the offset for t = T-1 is + // (T-1) * vocab * 4 bytes. + size_t row_bytes = (size_t) vocab * sizeof(float); + ggml_backend_tensor_get(logits, out->logits_last.data(), (size_t) (T - 1) * row_bytes, row_bytes); + + size_t hrow_bytes = (size_t) hidden * sizeof(float); + ggml_backend_tensor_get(h_final, out->hidden_last.data(), (size_t) (T - 1) * hrow_bytes, hrow_bytes); + } + + if (dump_dir) { + DebugDumper d; + debug_init(&d, dump_dir); + debug_dump_1d(&d, "talker-logits-prefill", out->logits_last.data(), vocab); + } + + ggml_backend_sched_reset(sched); + ggml_free(gctx); + return true; +} diff --git a/src/talker-forward.h b/src/talker-forward.h new file mode 100644 index 0000000..81b7d53 --- /dev/null +++ b/src/talker-forward.h @@ -0,0 +1,46 @@ +#pragma once +// talker-forward.h : prefill forward of the Talker LM. +// +// Takes a precomputed input embedding [T, hidden] f32 row-major and runs +// the 28-layer Qwen3 decoder stack with multimodal RoPE collapsed to 1D +// NEOX, GQA attention with per-head QK-norm, and SwiGLU MLP. The final +// hidden state is RMS-normalised and projected through codec_head to +// produce codebook 0 logits over a 3072-entry vocab. +// +// Phase 4.1+ : eager attention, full F32 compute, no KV cache. The +// graph is built from scratch at every call ; a generation loop will +// later wrap this with a sliding KV cache. +// +// Optional dump_dir captures bisect-layer activations and the final +// logits in the same f32 binary format the Python reference produces, +// for stage-by-stage cossim validation. + +#include "backend.h" +#include "debug.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "talker-weights.h" + +#include +#include + +struct TalkerForwardOutput { + // Final hidden state for the last position [hidden] f32 (post final norm). + std::vector hidden_last; + + // Codec head logits for the last position [vocab] f32. + std::vector logits_last; + + int hidden; + int vocab; +}; + +// Build and run the prefill graph. The scheduler is responsible for op +// placement (GPU primary, CPU fallback) ; the caller owns its lifetime. +// dump_dir may be NULL. +bool talker_forward_prefill(const TalkerWeights * tw, + ggml_backend_sched_t sched, + const float * input_embed, // [T, hidden] row-major + int T, + const char * dump_dir, + TalkerForwardOutput * out); diff --git a/src/talker-weights.h b/src/talker-weights.h new file mode 100644 index 0000000..ec5a8bc --- /dev/null +++ b/src/talker-weights.h @@ -0,0 +1,207 @@ +#pragma once +// talker-weights.h : Qwen3-style autoregressive Talker LM weights. +// +// Carries 28 decoder layers in 0.6B (36 in 1.7B), each with pre-norm +// attention plus pre-norm SwiGLU MLP. Attention is multi-head with GQA +// (16 query heads, 8 kv heads, head_dim 128) and per-head QK-norm. RoPE +// is mrope-interleaved with sections [24, 20, 20] and freq base 1e6 ; +// in TTS-only mode this collapses to plain interleaved 1D RoPE since +// the three multimodal axes carry the same position index. +// +// Top-level the talker holds two embedding tables (codec vocab 3072, +// text vocab 151936 with hidden 2048), a 2-layer ResizeMLP that +// projects text embeddings down to hidden 1024, the final RMSNorm and +// a codec_head Linear 1024 -> 3072 that emits the codebook 0 logits. +// +// Tensor naming follows convert.py output (flat, talker.*) : +// talker.codec_embedding.weight [3072, 1024] +// talker.text_embedding.weight [151936, 2048] +// talker.text_projection.fc1.{weight,bias} [2048, 2048] / [2048] +// talker.text_projection.fc2.{weight,bias} [2048, 1024] / [1024] +// talker.codec_head.weight [3072, 1024] +// talker.norm.weight [1024] +// talker.layers.{0..N-1}.input_layernorm.weight [hidden] +// talker.layers.{0..N-1}.post_attention_layernorm.weight [hidden] +// talker.layers.{0..N-1}.attn.q_proj.weight [hidden, n_heads*head_dim] +// talker.layers.{0..N-1}.attn.k_proj.weight [hidden, n_kv_heads*head_dim] +// talker.layers.{0..N-1}.attn.v_proj.weight [hidden, n_kv_heads*head_dim] +// talker.layers.{0..N-1}.attn.o_proj.weight [n_heads*head_dim, hidden] +// talker.layers.{0..N-1}.attn.q_norm.weight [head_dim] +// talker.layers.{0..N-1}.attn.k_norm.weight [head_dim] +// talker.layers.{0..N-1}.mlp.gate_proj.weight [hidden, intermediate] +// talker.layers.{0..N-1}.mlp.up_proj.weight [hidden, intermediate] +// talker.layers.{0..N-1}.mlp.down_proj.weight [intermediate, hidden] + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include + +struct TalkerAttention { + struct ggml_tensor * q_proj_w; + struct ggml_tensor * k_proj_w; + struct ggml_tensor * v_proj_w; + struct ggml_tensor * o_proj_w; + struct ggml_tensor * q_norm_w; + struct ggml_tensor * k_norm_w; +}; + +struct TalkerMLP { + struct ggml_tensor * gate_proj_w; + struct ggml_tensor * up_proj_w; + struct ggml_tensor * down_proj_w; +}; + +struct TalkerLayer { + struct ggml_tensor * input_norm_w; + TalkerAttention attn; + struct ggml_tensor * post_attn_norm_w; + TalkerMLP mlp; +}; + +struct TalkerWeights { + int hidden_size; + int intermediate_size; + int num_hidden_layers; + int num_attention_heads; + int num_key_value_heads; + int head_dim; + int vocab_size; + int text_vocab_size; + int text_hidden_size; + int max_position_embeddings; + int position_id_per_seconds; + float rope_theta; + float rms_norm_eps; + int mrope_section_t; + int mrope_section_h; + int mrope_section_w; + bool mrope_interleaved; + + struct ggml_tensor * codec_embedding; + struct ggml_tensor * text_embedding; + struct ggml_tensor * text_proj_fc1_w; + struct ggml_tensor * text_proj_fc1_b; + struct ggml_tensor * text_proj_fc2_w; + struct ggml_tensor * text_proj_fc2_b; + struct ggml_tensor * codec_head_w; + struct ggml_tensor * norm_w; + + std::vector layers; + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +static bool talker_weights_load(TalkerWeights * tw, const GGUFModel & gf, ggml_backend_t backend) { + tw->hidden_size = (int) gf_get_u32(gf, "qwen3-tts.talker.embedding_length"); + tw->intermediate_size = (int) gf_get_u32(gf, "qwen3-tts.talker.feed_forward_length"); + tw->num_hidden_layers = (int) gf_get_u32(gf, "qwen3-tts.talker.block_count"); + tw->num_attention_heads = (int) gf_get_u32(gf, "qwen3-tts.talker.attention.head_count"); + tw->num_key_value_heads = (int) gf_get_u32(gf, "qwen3-tts.talker.attention.head_count_kv"); + tw->head_dim = (int) gf_get_u32(gf, "qwen3-tts.talker.attention.key_length"); + tw->vocab_size = (int) gf_get_u32(gf, "qwen3-tts.talker.vocab_size"); + tw->text_vocab_size = (int) gf_get_u32(gf, "qwen3-tts.talker.text_vocab_size"); + tw->text_hidden_size = (int) gf_get_u32(gf, "qwen3-tts.talker.text_hidden_size"); + tw->max_position_embeddings = (int) gf_get_u32(gf, "qwen3-tts.talker.context_length"); + tw->position_id_per_seconds = (int) gf_get_u32(gf, "qwen3-tts.talker.position_id_per_seconds"); + tw->rope_theta = gf_get_f32(gf, "qwen3-tts.talker.rope.freq_base"); + tw->rms_norm_eps = gf_get_f32(gf, "qwen3-tts.talker.attention.layer_norm_rms_epsilon"); + tw->mrope_interleaved = gf_get_bool(gf, "qwen3-tts.talker.rope.mrope_interleaved"); + + std::vector mrope = gf_get_array_u32(gf, "qwen3-tts.talker.rope.mrope_section"); + if (mrope.size() == 3) { + tw->mrope_section_t = (int) mrope[0]; + tw->mrope_section_h = (int) mrope[1]; + tw->mrope_section_w = (int) mrope[2]; + } else { + tw->mrope_section_t = tw->mrope_section_h = tw->mrope_section_w = 0; + } + + if (tw->num_hidden_layers <= 0 || tw->hidden_size <= 0) { + fprintf(stderr, "[Talker] FATAL: invalid hyperparameters in GGUF (layers=%d hidden=%d)\n", + tw->num_hidden_layers, tw->hidden_size); + return false; + } + + tw->layers.resize((size_t) tw->num_hidden_layers); + + // 8 top-level + per layer (2 norms + 4 attn + 2 qk norms + 3 mlp) = 11 + int n_tensors = 8 + tw->num_hidden_layers * 11 + 8; + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + tw->codec_embedding = gf_load_tensor(&wctx, gf, "talker.codec_embd.weight"); + tw->text_embedding = gf_load_tensor(&wctx, gf, "talker.text_embd.weight"); + tw->text_proj_fc1_w = gf_load_tensor(&wctx, gf, "talker.text_proj.fc1.weight"); + tw->text_proj_fc1_b = gf_load_tensor(&wctx, gf, "talker.text_proj.fc1.bias"); + tw->text_proj_fc2_w = gf_load_tensor(&wctx, gf, "talker.text_proj.fc2.weight"); + tw->text_proj_fc2_b = gf_load_tensor(&wctx, gf, "talker.text_proj.fc2.bias"); + tw->codec_head_w = gf_load_tensor(&wctx, gf, "talker.codec_head.weight"); + tw->norm_w = gf_load_tensor(&wctx, gf, "talker.output_norm.weight"); + + for (int l = 0; l < tw->num_hidden_layers; l++) { + TalkerLayer & layer = tw->layers[(size_t) l]; + char name[160]; + + snprintf(name, sizeof(name), "talker.blk.%d.attn_norm.weight", l); + layer.input_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "talker.blk.%d.ffn_norm.weight", l); + layer.post_attn_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "talker.blk.%d.attn_q.weight", l); + layer.attn.q_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "talker.blk.%d.attn_k.weight", l); + layer.attn.k_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "talker.blk.%d.attn_v.weight", l); + layer.attn.v_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "talker.blk.%d.attn_output.weight", l); + layer.attn.o_proj_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "talker.blk.%d.attn_q_norm.weight", l); + layer.attn.q_norm_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "talker.blk.%d.attn_k_norm.weight", l); + layer.attn.k_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "talker.blk.%d.ffn_gate.weight", l); + layer.mlp.gate_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "talker.blk.%d.ffn_up.weight", l); + layer.mlp.up_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "talker.blk.%d.ffn_down.weight", l); + layer.mlp.down_proj_w = gf_load_tensor(&wctx, gf, name); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[Talker] FATAL: backend allocation failed\n"); + return false; + } + tw->weight_ctx = wctx.ctx; + tw->weight_buf = wctx.buffer; + + fprintf(stderr, + "[Talker] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, " + "FFN %d, RoPE theta %.0f, mrope sections [%d,%d,%d] interleaved=%d\n", + tw->num_hidden_layers, tw->hidden_size, tw->num_attention_heads, tw->num_key_value_heads, tw->head_dim, + tw->intermediate_size, (double) tw->rope_theta, tw->mrope_section_t, tw->mrope_section_h, + tw->mrope_section_w, (int) tw->mrope_interleaved); + return true; +} + +static void talker_weights_free(TalkerWeights * tw) { + if (tw->weight_buf) { + ggml_backend_buffer_free(tw->weight_buf); + tw->weight_buf = NULL; + } + if (tw->weight_ctx) { + ggml_free(tw->weight_ctx); + tw->weight_ctx = NULL; + } + tw->layers.clear(); +} diff --git a/src/text-chunker.h b/src/text-chunker.h new file mode 100644 index 0000000..c3af0c4 --- /dev/null +++ b/src/text-chunker.h @@ -0,0 +1,382 @@ +#pragma once +// text-chunker.h: punctuation-aware long-form text splitter for TTS +// +// chunk_text_punctuation: splits text on sentence-ending punctuation +// (skipping abbreviation periods), then merges sentences into chunks of +// at most chunk_len UTF-8 codepoints. Optional min_chunk_len merges +// undersized chunks into a neighbour. Strings are UTF-8 in, UTF-8 out. +// Comparison and length are codepoint-based, matching Python str +// semantics. +// Math reference : omnivoice/utils/text.py (1:1 port). + +#include +#include +#include + +// Returns the byte length of the UTF-8 codepoint starting at b (1, 2, 3 or 4). +// Falls back to 1 on invalid first bytes so iteration always advances. +static inline int chunker_utf8_len(unsigned char b) { + if ((b & 0x80) == 0x00) { + return 1; + } + + if ((b & 0xE0) == 0xC0) { + return 2; + } + + if ((b & 0xF0) == 0xE0) { + return 3; + } + + if ((b & 0xF8) == 0xF0) { + return 4; + } + + return 1; +} + +// Sentence-ending punctuation. Mirrors SPLIT_PUNCTUATION in text.py. +static const std::set & chunker_split_punctuation() { + static const std::set s = { + ".", + ",", + ";", + ":", + "!", + "?", + "\xe3\x80\x82", // U+3002 ideographic full stop + "\xef\xbc\x8c", // U+FF0C fullwidth comma + "\xef\xbc\x9b", // U+FF1B fullwidth semicolon + "\xef\xbc\x9a", // U+FF1A fullwidth colon + "\xef\xbc\x81", // U+FF01 fullwidth exclamation mark + "\xef\xbc\x9f", // U+FF1F fullwidth question mark + }; + return s; +} + +// Closing marks attach to the preceding sentence. Mirrors CLOSING_MARKS. +static const std::set & chunker_closing_marks() { + static const std::set s = { + "\"", "'", "]", ">", + "\xe2\x80\x9c", // U+201C left double quotation mark + "\xe2\x80\x9d", // U+201D right double quotation mark + "\xe2\x80\x98", // U+2018 left single quotation mark + "\xe2\x80\x99", // U+2019 right single quotation mark + "\xef\xbc\x89", // U+FF09 fullwidth right parenthesis + "\xe3\x80\x8b", // U+300B right double angle bracket + "\xe3\x80\x8d", // U+300D right corner bracket + "\xe3\x80\x91", // U+3011 right black lenticular bracket + }; + return s; +} + +// Abbreviations that suppress the period as a sentence break. ASCII only, +// matched on the last whitespace-delimited word ending with the period. +// Mirrors ABBREVIATIONS in text.py. +static const std::set & chunker_abbreviations() { + static const std::set s = { + "Mr.", "Mrs.", "Ms.", "Dr.", "Prof.", "Sr.", "Jr.", "Rev.", "Fr.", "Hon.", "Pres.", + "Gov.", "Capt.", "Gen.", "Sen.", "Rep.", "Col.", "Maj.", "Lt.", "Cmdr.", "Sgt.", "Cpl.", + "Co.", "Corp.", "Inc.", "Ltd.", "Est.", "Dept.", "St.", "Ave.", "Blvd.", "Rd.", "Mt.", + "Ft.", "No.", "Jan.", "Feb.", "Mar.", "Apr.", "Aug.", "Sep.", "Sept.", "Oct.", "Nov.", + "Dec.", "i.e.", "e.g.", "vs.", "Vs.", "Etc.", "approx.", "fig.", "def.", + }; + return s; +} + +// Returns the last whitespace-delimited word of s, or s itself if no space. +// Used to detect abbreviation periods. +static std::string chunker_last_word(const std::string & s) { + size_t e = s.find_last_not_of(" \t\n\r"); + if (e == std::string::npos) { + return std::string(); + } + + std::string trimmed = s.substr(0, e + 1); + size_t sp = trimmed.find_last_of(" \t\n\r"); + if (sp == std::string::npos) { + return trimmed; + } + + return trimmed.substr(sp + 1); +} + +// Strips leading and trailing ASCII whitespace from s. +static std::string chunker_strip(const std::string & s) { + size_t a = s.find_first_not_of(" \t\n\r"); + if (a == std::string::npos) { + return std::string(); + } + + size_t b = s.find_last_not_of(" \t\n\r"); + return s.substr(a, b - a + 1); +} + +// Splits text on sentence-ending punctuation (skipping abbreviations) and +// merges sentences into chunks of at most chunk_len codepoints. If +// min_chunk_len > 0, undersized chunks are merged with a neighbour. +// Returns a list of stripped chunk strings (UTF-8). Empty chunks are dropped. +// +// Strict 1:1 port of chunk_text_punctuation in omnivoice/utils/text.py. +static std::vector chunk_text_punctuation(const std::string & text, int chunk_len, int min_chunk_len) { + // Step 1: tokenise into UTF-8 codepoints, then split on punctuation. + // sentences holds vectors of codepoints (each codepoint is a std::string). + std::vector> sentences; + std::vector current; + + const std::set & split_set = chunker_split_punctuation(); + const std::set & closing_set = chunker_closing_marks(); + const std::set & abbrev_set = chunker_abbreviations(); + + const unsigned char * p = (const unsigned char *) text.data(); + const unsigned char * end = p + text.size(); + + while (p < end) { + int n = chunker_utf8_len(*p); + if (p + n > end) { + n = (int) (end - p); + } + + std::string cp((const char *) p, (size_t) n); + p += n; + + bool is_split = split_set.count(cp) > 0; + bool is_closing = closing_set.count(cp) > 0; + + // Leading punctuation glues onto the previous sentence. + if (current.empty() && !sentences.empty() && (is_split || is_closing)) { + sentences.back().push_back(cp); + continue; + } + + current.push_back(cp); + + if (!is_split) { + continue; + } + + // Period after an abbreviation does not break the sentence. + bool is_abbreviation = false; + if (cp == ".") { + std::string joined; + for (const auto & c : current) { + joined += c; + } + + std::string last = chunker_last_word(joined); + if (!last.empty() && abbrev_set.count(last) > 0) { + is_abbreviation = true; + } + } + + if (!is_abbreviation) { + sentences.push_back(current); + current.clear(); + } + } + + if (!current.empty()) { + sentences.push_back(current); + } + + // Step 2: greedy merge of sentences into chunks of at most chunk_len + // codepoints. A sentence that does not fit starts a new chunk by itself, + // even if it is longer than chunk_len. + std::vector> merged; + std::vector cur_chunk; + + for (const auto & sent : sentences) { + if ((int) (cur_chunk.size() + sent.size()) <= chunk_len) { + for (const auto & c : sent) { + cur_chunk.push_back(c); + } + } else { + if (!cur_chunk.empty()) { + merged.push_back(cur_chunk); + } + + cur_chunk = sent; + } + } + + if (!cur_chunk.empty()) { + merged.push_back(cur_chunk); + } + + // Step 3: merge undersized chunks. The first chunk, if short, is folded + // into the second. Subsequent short chunks fold into the previous one. + std::vector> finals; + if (min_chunk_len > 0) { + bool first_short = !merged.empty() && (int) merged[0].size() < min_chunk_len; + + for (size_t i = 0; i < merged.size(); i++) { + const auto & chunk = merged[i]; + + if (i == 1 && first_short) { + for (const auto & c : chunk) { + finals.back().push_back(c); + } + continue; + } + + if ((int) chunk.size() >= min_chunk_len) { + finals.push_back(chunk); + continue; + } + + if (finals.empty()) { + finals.push_back(chunk); + } else { + for (const auto & c : chunk) { + finals.back().push_back(c); + } + } + } + } else { + finals = merged; + } + + // Step 4: join codepoints, strip whitespace, drop empty. + std::vector result; + result.reserve(finals.size()); + + for (const auto & chunk : finals) { + std::string joined; + for (const auto & c : chunk) { + joined += c; + } + + std::string stripped = chunker_strip(joined); + if (!stripped.empty()) { + result.push_back(stripped); + } + } + + return result; +} + +// Counts UTF-8 codepoints in s. Used to derive the per-chunk character budget +// from the average tokens-per-character of the full text, matching Python's +// len(text) which counts codepoints. +static int chunker_utf8_count(const std::string & text) { + const unsigned char * p = (const unsigned char *) text.data(); + const unsigned char * end = p + text.size(); + int n = 0; + + while (p < end) { + int s = chunker_utf8_len(*p); + if (p + s > end) { + s = (int) (end - p); + } + + p += s; + n += 1; + } + + return n; +} + +// Punctuation considered "terminal" by add_punctuation. Mirrors END_PUNCTUATION +// in text.py. ASCII first, then UTF-8 byte sequences for fancy quotes, ellipsis +// and Chinese variants. +static const std::set & chunker_end_punctuation() { + static const std::set s = { + ";", + ":", + ",", + ".", + "!", + "?", + ")", + "]", + "}", + "\"", + "'", + "\xe2\x80\xa6", // U+2026 horizontal ellipsis + "\xe2\x80\x9c", // U+201C left double quotation mark + "\xe2\x80\x9d", // U+201D right double quotation mark + "\xe2\x80\x98", // U+2018 left single quotation mark + "\xe2\x80\x99", // U+2019 right single quotation mark + "\xef\xbc\x9b", // U+FF1B fullwidth semicolon + "\xef\xbc\x9a", // U+FF1A fullwidth colon + "\xef\xbc\x8c", // U+FF0C fullwidth comma + "\xe3\x80\x82", // U+3002 ideographic full stop + "\xef\xbc\x81", // U+FF01 fullwidth exclamation mark + "\xef\xbc\x9f", // U+FF1F fullwidth question mark + "\xe3\x80\x81", // U+3001 ideographic comma + "\xef\xbc\x89", // U+FF09 fullwidth right parenthesis + "\xe3\x80\x91", // U+3011 right black lenticular bracket + }; + return s; +} + +// Returns the last UTF-8 codepoint of s as a std::string, or empty if s is +// empty. Walks the byte sequence to find the start of the last codepoint. +static std::string chunker_last_codepoint(const std::string & s) { + if (s.empty()) { + return std::string(); + } + + size_t i = s.size(); + while (i > 0) { + unsigned char b = (unsigned char) s[i - 1]; + if ((b & 0xC0) != 0x80) { + return s.substr(i - 1); + } + + i--; + } + + return s; +} + +// Returns true if any codepoint of s falls inside the CJK Unified Ideographs +// block (U+4E00..U+9FFF). Mirrors the Chinese-detection heuristic in +// add_punctuation upstream. +static bool chunker_contains_chinese(const std::string & s) { + const unsigned char * p = (const unsigned char *) s.data(); + const unsigned char * end = p + s.size(); + + while (p < end) { + int n = chunker_utf8_len(*p); + if (p + n > end) { + return false; + } + + if (n == 3) { + uint32_t cp = + ((uint32_t) (p[0] & 0x0F) << 12) | ((uint32_t) (p[1] & 0x3F) << 6) | ((uint32_t) (p[2] & 0x3F)); + if (cp >= 0x4E00 && cp <= 0x9FFF) { + return true; + } + } + + p += n; + } + + return false; +} + +// Strips text and appends a terminal punctuation if missing. Mirrors +// add_punctuation in omnivoice/utils/text.py: appends "." for non-Chinese +// text, and the ideographic full stop "。" for text containing CJK. +static std::string add_punctuation(const std::string & text) { + std::string s = chunker_strip(text); + if (s.empty()) { + return s; + } + + std::string last = chunker_last_codepoint(s); + const auto & end_set = chunker_end_punctuation(); + if (end_set.count(last) > 0) { + return s; + } + + if (chunker_contains_chinese(s)) { + s += "\xe3\x80\x82"; // U+3002 + } else { + s += "."; + } + + return s; +} diff --git a/src/tokenizer-transformer.h b/src/tokenizer-transformer.h new file mode 100644 index 0000000..be1ba5f --- /dev/null +++ b/src/tokenizer-transformer.h @@ -0,0 +1,295 @@ +#pragma once +// tokenizer-transformer.h : 8-layer Qwen3-style local-causal transformer +// for the Qwen3-TTS 12Hz tokenizer decoder. +// +// Hidden size 512, head_dim 64, 16 query and 16 KV heads (no GQA), FFN +// intermediate 1024, RoPE NEOX style with theta 10000, RMSNorm eps 1e-5, +// sliding window 72 frames causal attention, LayerScale post-attention +// and post-MLP. SwiGLU MLP. No biases on q/k/v/o or gate/up/down +// projections. Top-level input_proj 1024 -> 512 (with bias) and +// output_proj 512 -> 1024 (with bias) bracket the transformer stack and +// connect to the latent stream of the codec. + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf-weights.h" +#include "weight-ctx.h" + +#include +#include +#include +#include +#include +#include + +#define QWEN_TOKENIZER_TRANSFORMER_MAX_LAYERS 16 + +struct QwenTransformerAttention { + struct ggml_tensor * q_proj_w; // [hidden, num_q_heads * head_dim] + struct ggml_tensor * k_proj_w; // [hidden, num_kv_heads * head_dim] + struct ggml_tensor * v_proj_w; + struct ggml_tensor * o_proj_w; // [num_q_heads * head_dim, hidden] +}; + +struct QwenTransformerMLP { + struct ggml_tensor * gate_proj_w; // [hidden, intermediate] + struct ggml_tensor * up_proj_w; + struct ggml_tensor * down_proj_w; // [intermediate, hidden] +}; + +struct QwenTransformerLayer { + struct ggml_tensor * input_norm_w; // RMSNorm gain [hidden] + QwenTransformerAttention attn; + struct ggml_tensor * attn_scale; // LayerScale per-channel [hidden] + struct ggml_tensor * post_attn_norm_w; + QwenTransformerMLP mlp; + struct ggml_tensor * mlp_scale; +}; + +struct QwenTokenizerTransformer { + int hidden_size; + int latent_dim; + int num_layers; + int num_attention_heads; + int num_kv_heads; + int head_dim; + int intermediate_size; + int sliding_window; + float rope_theta; + float rms_norm_eps; + + struct ggml_tensor * input_proj_w; // [latent_dim, hidden] + struct ggml_tensor * input_proj_b; // [hidden] + QwenTransformerLayer layers[QWEN_TOKENIZER_TRANSFORMER_MAX_LAYERS]; + struct ggml_tensor * norm_w; // [hidden] + struct ggml_tensor * output_proj_w; // [hidden, latent_dim] + struct ggml_tensor * output_proj_b; // [latent_dim] + + struct ggml_context * weight_ctx; + ggml_backend_buffer_t weight_buf; +}; + +// Read decoder hyperparameters from GGUF metadata, allocate every weight +// tensor on the backend, and bind tensor pointers in the struct. Returns +// true on success. +static bool qwen_tokenizer_transformer_load(QwenTokenizerTransformer * tr, + const GGUFModel & gf, + ggml_backend_t backend) { + tr->hidden_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.hidden_size"); + tr->latent_dim = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.latent_dim"); + tr->num_layers = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.num_hidden_layers"); + tr->num_attention_heads = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.num_attention_heads"); + tr->num_kv_heads = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.num_key_value_heads"); + tr->head_dim = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.head_dim"); + tr->intermediate_size = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.intermediate_size"); + tr->sliding_window = (int) gf_get_u32(gf, "qwen3-tts-tokenizer.decoder.sliding_window"); + tr->rope_theta = gf_get_f32(gf, "qwen3-tts-tokenizer.decoder.rope_theta"); + tr->rms_norm_eps = gf_get_f32(gf, "qwen3-tts-tokenizer.decoder.rms_norm_eps"); + + if (tr->num_layers > QWEN_TOKENIZER_TRANSFORMER_MAX_LAYERS) { + fprintf(stderr, "[Transformer] FATAL: %d layers exceeds compile-time max %d\n", tr->num_layers, + QWEN_TOKENIZER_TRANSFORMER_MAX_LAYERS); + return false; + } + + int n_tensors = 6 // input_proj wb + norm + output_proj wb + + tr->num_layers * 11 // 2 norms + 4 attn + 3 mlp + 2 layer scales + + 4; // headroom + WeightCtx wctx; + wctx_init(&wctx, n_tensors); + + tr->input_proj_w = gf_load_tensor(&wctx, gf, "tok_dec.pre_tfm.input_proj.weight"); + tr->input_proj_b = gf_load_tensor(&wctx, gf, "tok_dec.pre_tfm.input_proj.bias"); + tr->norm_w = gf_load_tensor(&wctx, gf, "tok_dec.pre_tfm.norm.weight"); + tr->output_proj_w = gf_load_tensor(&wctx, gf, "tok_dec.pre_tfm.output_proj.weight"); + tr->output_proj_b = gf_load_tensor(&wctx, gf, "tok_dec.pre_tfm.output_proj.bias"); + + for (int l = 0; l < tr->num_layers; l++) { + QwenTransformerLayer & layer = tr->layers[l]; + char name[160]; + + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.attn_norm.weight", l); + layer.input_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.attn_q.weight", l); + layer.attn.q_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.attn_k.weight", l); + layer.attn.k_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.attn_v.weight", l); + layer.attn.v_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.attn_output.weight", l); + layer.attn.o_proj_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.attn_scale", l); + layer.attn_scale = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.ffn_norm.weight", l); + layer.post_attn_norm_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.ffn_gate.weight", l); + layer.mlp.gate_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.ffn_up.weight", l); + layer.mlp.up_proj_w = gf_load_tensor(&wctx, gf, name); + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.ffn_down.weight", l); + layer.mlp.down_proj_w = gf_load_tensor(&wctx, gf, name); + + snprintf(name, sizeof(name), "tok_dec.pre_tfm.blk.%d.ffn_scale", l); + layer.mlp_scale = gf_load_tensor(&wctx, gf, name); + } + + if (!wctx_alloc(&wctx, backend)) { + fprintf(stderr, "[Transformer] FATAL: backend allocation failed\n"); + return false; + } + tr->weight_ctx = wctx.ctx; + tr->weight_buf = wctx.buffer; + + fprintf(stderr, + "[Transformer] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, " + "FFN %d, RoPE theta %.0f, sliding window %d\n", + tr->num_layers, tr->hidden_size, tr->num_attention_heads, tr->num_kv_heads, tr->head_dim, + tr->intermediate_size, tr->rope_theta, tr->sliding_window); + return true; +} + +static void qwen_tokenizer_transformer_free(QwenTokenizerTransformer * tr) { + if (tr->weight_buf) { + ggml_backend_buffer_free(tr->weight_buf); + tr->weight_buf = NULL; + } + if (tr->weight_ctx) { + ggml_free(tr->weight_ctx); + tr->weight_ctx = NULL; + } +} + +// Fill a [T, T] f32 mask with 0 where attention is allowed and -inf +// elsewhere. Storage is row-major with k (key index) on the fast axis : +// dst[q * T + k] is the additive bias for query q attending to key k. +// Causal sliding window : mask[k, q] = 0 if (k <= q AND q - k < window), +// else -inf. +static void qwen_build_causal_sliding_mask(int T, int sliding_window, std::vector & dst) { + dst.assign((size_t) T * (size_t) T, -INFINITY); + for (int q = 0; q < T; q++) { + int k_min = q - sliding_window + 1; + if (k_min < 0) { + k_min = 0; + } + for (int k = k_min; k <= q; k++) { + dst[(size_t) q * (size_t) T + (size_t) k] = 0.0f; + } + } +} + +static void qwen_build_positions(int T, std::vector & dst) { + dst.resize((size_t) T); + for (int i = 0; i < T; i++) { + dst[i] = i; + } +} + +// One transformer layer : attention block then MLP block, both with +// pre-RMSNorm, post-LayerScale and residual connection. +static struct ggml_tensor * qwen_transformer_layer_forward(struct ggml_context * ctx, + const QwenTokenizerTransformer * tr, + const QwenTransformerLayer & layer, + struct ggml_tensor * x, + struct ggml_tensor * positions, + struct ggml_tensor * mask, + int T) { + int hidden = tr->hidden_size; + int n_q_heads = tr->num_attention_heads; + int n_kv = tr->num_kv_heads; + int hd = tr->head_dim; + + // Attention block : pre-RMSNorm + project Q/K/V + RoPE + scaled dot product + // + softmax with causal sliding mask + V combine + o_proj. + struct ggml_tensor * ln1 = ggml_rms_norm(ctx, x, tr->rms_norm_eps); + ln1 = ggml_mul(ctx, ln1, layer.input_norm_w); + + struct ggml_tensor * q = ggml_mul_mat(ctx, layer.attn.q_proj_w, ln1); // [n_q_heads*hd, T] + struct ggml_tensor * k = ggml_mul_mat(ctx, layer.attn.k_proj_w, ln1); // [n_kv*hd, T] + struct ggml_tensor * v = ggml_mul_mat(ctx, layer.attn.v_proj_w, ln1); // [n_kv*hd, T] + + q = ggml_reshape_3d(ctx, q, hd, n_q_heads, T); // [hd, n_q_heads, T] + k = ggml_reshape_3d(ctx, k, hd, n_kv, T); + v = ggml_reshape_3d(ctx, v, hd, n_kv, T); + + q = ggml_rope_ext(ctx, q, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + k = ggml_rope_ext(ctx, k, positions, NULL, hd, GGML_ROPE_TYPE_NEOX, 0, tr->rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, + 0.0f); + + // Permute to head-as-batch layout : [hd, T, n_heads] + struct ggml_tensor * q_p = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + struct ggml_tensor * k_p = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + struct ggml_tensor * v_p = ggml_permute(ctx, v, 1, 2, 0, 3); // [T, n_kv, hd] for V mul_mat + v_p = ggml_cont(ctx, v_p); + + // Scores : mul_mat(K, Q) -> [T_k, T_q, n_heads] + struct ggml_tensor * scores = ggml_mul_mat(ctx, k_p, q_p); + + float scale = 1.0f / sqrtf((float) hd); + scores = ggml_soft_max_ext(ctx, scores, mask, scale, 0.0f); + + // Attention output : mul_mat(V_T, scores) -> [hd, T_q, n_heads] + // V_T has T_k as ne[0], hd as ne[1], n_heads as ne[2]. + struct ggml_tensor * attn = ggml_mul_mat(ctx, v_p, scores); + + // Permute back to [hd, n_heads, T] then reshape to [n_heads*hd, T] + attn = ggml_cont(ctx, ggml_permute(ctx, attn, 0, 2, 1, 3)); + attn = ggml_reshape_2d(ctx, attn, n_q_heads * hd, T); + + struct ggml_tensor * o = ggml_mul_mat(ctx, layer.attn.o_proj_w, attn); // [hidden, T] + + // LayerScale + residual + o = ggml_mul(ctx, o, layer.attn_scale); + x = ggml_add(ctx, x, o); + + // MLP block : pre-RMSNorm + SwiGLU + LayerScale + residual. + struct ggml_tensor * ln2 = ggml_rms_norm(ctx, x, tr->rms_norm_eps); + ln2 = ggml_mul(ctx, ln2, layer.post_attn_norm_w); + + struct ggml_tensor * gate = ggml_mul_mat(ctx, layer.mlp.gate_proj_w, ln2); // [intermediate, T] + struct ggml_tensor * up = ggml_mul_mat(ctx, layer.mlp.up_proj_w, ln2); + gate = ggml_silu(ctx, gate); + struct ggml_tensor * gu = ggml_mul(ctx, gate, up); + struct ggml_tensor * mlp = ggml_mul_mat(ctx, layer.mlp.down_proj_w, gu); // [hidden, T] + + mlp = ggml_mul(ctx, mlp, layer.mlp_scale); + x = ggml_add(ctx, x, mlp); + + (void) hidden; + return x; +} + +// Full forward pass : input_proj, 8 layers, final norm, output_proj. +// +// x : [latent_dim, T] f32 +// positions : [T] i32 +// mask : [T, T] f32, additive (-inf where masked) +// returns : [latent_dim, T] f32 +static struct ggml_tensor * qwen_tokenizer_transformer_forward(struct ggml_context * ctx, + const QwenTokenizerTransformer * tr, + struct ggml_tensor * x, + struct ggml_tensor * positions, + struct ggml_tensor * mask) { + int T = (int) x->ne[1]; + + // input_proj : [latent_dim, T] -> [hidden, T] + struct ggml_tensor * h = ggml_mul_mat(ctx, tr->input_proj_w, x); + h = ggml_add(ctx, h, tr->input_proj_b); + + for (int l = 0; l < tr->num_layers; l++) { + h = qwen_transformer_layer_forward(ctx, tr, tr->layers[l], h, positions, mask, T); + } + + h = ggml_rms_norm(ctx, h, tr->rms_norm_eps); + h = ggml_mul(ctx, h, tr->norm_w); + + // output_proj : [hidden, T] -> [latent_dim, T] + h = ggml_mul_mat(ctx, tr->output_proj_w, h); + h = ggml_add(ctx, h, tr->output_proj_b); + + return h; +} diff --git a/src/wav.h b/src/wav.h new file mode 100644 index 0000000..56a3921 --- /dev/null +++ b/src/wav.h @@ -0,0 +1,183 @@ +#pragma once +// wav.h: minimal WAV reader +// +// read_wav_buf: PCM16 / PCM24 / float32, classic or WAVE_FORMAT_EXTENSIBLE, +// mono or stereo, any rate -> interleaved [T, 2] float + +#include +#include +#include +#include +#include + +static uint16_t wav_read_u16le(const uint8_t * p) { + return (uint16_t) (p[0] | (p[1] << 8)); +} + +static uint32_t wav_read_u32le(const uint8_t * p) { + return (uint32_t) p[0] | ((uint32_t) p[1] << 8) | ((uint32_t) p[2] << 16) | ((uint32_t) p[3] << 24); +} + +static int32_t wav_read_s24le(const uint8_t * p) { + uint32_t u = (uint32_t) p[0] | ((uint32_t) p[1] << 8) | ((uint32_t) p[2] << 16); + + if (u & 0x00800000u) { + u |= 0xff000000u; + } + + return (int32_t) u; +} + +static float wav_read_f32le(const uint8_t * p) { + uint32_t u = wav_read_u32le(p); + float f; + memcpy(&f, &u, 4); + return f; +} + +// Read WAV from memory buffer. +// Returns interleaved float [T, 2]. Sets *T_audio, *sr. Caller frees. +static float * read_wav_buf(const uint8_t * data, size_t size, int * T_audio, int * sr) { + *T_audio = 0; + *sr = 0; + + if (size < 12 || memcmp(data, "RIFF", 4) != 0 || memcmp(data + 8, "WAVE", 4) != 0) { + fprintf(stderr, "[WAV] Not a valid WAV buffer\n"); + return NULL; + } + + int n_channels = 0; + int sample_rate = 0; + int bits_per_sample = 0; + uint16_t audio_format = 0; + uint16_t extensible_subformat = 0; + float * audio = NULL; + int n_samples = 0; + size_t pos = 12; + + while (pos + 8 <= size) { + const uint8_t * chunk_id = data + pos; + uint32_t chunk_size = wav_read_u32le(data + pos + 4); + pos += 8; + + if (pos + (size_t) chunk_size > size) { + chunk_size = (uint32_t) (size - pos); + } + + if (memcmp(chunk_id, "fmt ", 4) == 0 && chunk_size >= 16) { + audio_format = wav_read_u16le(data + pos + 0); + n_channels = (int) wav_read_u16le(data + pos + 2); + sample_rate = (int) wav_read_u32le(data + pos + 4); + bits_per_sample = (int) wav_read_u16le(data + pos + 14); + + extensible_subformat = 0; + if (audio_format == 0xfffe && chunk_size >= 40) { + extensible_subformat = wav_read_u16le(data + pos + 24); + + // collapse extensible to its effective sample format + // 1 -> PCM int, 3 -> IEEE float + if (extensible_subformat == 1 || extensible_subformat == 3) { + audio_format = extensible_subformat; + } + } + + pos += (size_t) chunk_size; + + } else if (memcmp(chunk_id, "data", 4) == 0 && n_channels > 0) { + size_t data_bytes = (size_t) chunk_size; + + if (audio_format == 1 && bits_per_sample == 16) { + n_samples = (int) (data_bytes / ((size_t) n_channels * 2)); + audio = (float *) malloc((size_t) n_samples * 2 * sizeof(float)); + if (!audio) { + fprintf(stderr, "[WAV] OOM allocating PCM16 buffer for %d samples\n", n_samples); + return NULL; + } + const uint8_t * p = data + pos; + + for (int t = 0; t < n_samples; t++) { + if (n_channels == 1) { + int16_t s = (int16_t) wav_read_u16le(p + t * 2); + float f = (float) s / 32768.0f; + audio[t * 2 + 0] = f; + audio[t * 2 + 1] = f; + } else { + const uint8_t * frame = p + (size_t) t * n_channels * 2; + int16_t l = (int16_t) wav_read_u16le(frame + 0); + int16_t r = (int16_t) wav_read_u16le(frame + 2); + audio[t * 2 + 0] = (float) l / 32768.0f; + audio[t * 2 + 1] = (float) r / 32768.0f; + } + } + } else if (audio_format == 1 && bits_per_sample == 24) { + n_samples = (int) (data_bytes / ((size_t) n_channels * 3)); + audio = (float *) malloc((size_t) n_samples * 2 * sizeof(float)); + if (!audio) { + fprintf(stderr, "[WAV] OOM allocating PCM24 buffer for %d samples\n", n_samples); + return NULL; + } + const uint8_t * p = data + pos; + + for (int t = 0; t < n_samples; t++) { + if (n_channels == 1) { + int32_t s = wav_read_s24le(p + t * 3); + float f = (float) s / 8388608.0f; + audio[t * 2 + 0] = f; + audio[t * 2 + 1] = f; + } else { + const uint8_t * frame = p + (size_t) t * n_channels * 3; + int32_t l = wav_read_s24le(frame + 0); + int32_t r = wav_read_s24le(frame + 3); + audio[t * 2 + 0] = (float) l / 8388608.0f; + audio[t * 2 + 1] = (float) r / 8388608.0f; + } + } + } else if (audio_format == 3 && bits_per_sample == 32) { + n_samples = (int) (data_bytes / ((size_t) n_channels * 4)); + audio = (float *) malloc((size_t) n_samples * 2 * sizeof(float)); + if (!audio) { + fprintf(stderr, "[WAV] OOM allocating F32 buffer for %d samples\n", n_samples); + return NULL; + } + const uint8_t * p = data + pos; + + for (int t = 0; t < n_samples; t++) { + if (n_channels == 1) { + float s = wav_read_f32le(p + t * 4); + audio[t * 2 + 0] = s; + audio[t * 2 + 1] = s; + } else { + const uint8_t * frame = p + (size_t) t * n_channels * 4; + float l = wav_read_f32le(frame + 0); + float r = wav_read_f32le(frame + 4); + audio[t * 2 + 0] = l; + audio[t * 2 + 1] = r; + } + } + } else { + fprintf(stderr, "[WAV] Unsupported: format=%u bits=%d subformat=%u\n", (unsigned) audio_format, + bits_per_sample, (unsigned) extensible_subformat); + return NULL; + } + + break; + } else { + pos += (size_t) chunk_size; + } + + if (chunk_size & 1) { + pos += 1; + } + } + + if (!audio) { + fprintf(stderr, "[WAV] No audio data in buffer\n"); + return NULL; + } + + *T_audio = n_samples; + *sr = sample_rate; + fprintf(stderr, "[WAV] Read buffer: %d samples, %d Hz, %d ch, %d bit\n", n_samples, sample_rate, n_channels, + bits_per_sample); + return audio; +} diff --git a/src/weight-ctx.h b/src/weight-ctx.h new file mode 100644 index 0000000..e43f573 --- /dev/null +++ b/src/weight-ctx.h @@ -0,0 +1,83 @@ +#pragma once +// weight-ctx.h: format-independent weight loading context for ggml backends +// +// Manages a ggml_context for weight tensors + their backend buffer. +// Used by gguf-weights.h for all model loaders. +// +// Usage: +// WeightCtx wctx; +// wctx_init(&wctx, n_tensors); +// ggml_tensor * w = _load_tensor(&wctx, source, "name"); +// wctx_alloc(&wctx, backend); + +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include +#include + +struct WeightCtx { + struct ggml_context * ctx; + ggml_backend_buffer_t buffer; + + struct PendingCopy { + struct ggml_tensor * tensor; + const void * src; + size_t nbytes; + size_t offset; // byte offset into dst tensor (0 for regular loads) + }; + + std::vector pending; + + // Staging buffers for type-converted data, kept alive until wctx_alloc. + // unique_ptr keeps the data address stable even when the outer vector grows, + // so src pointers stored in pending stay valid across staging.push_back(). + std::vector> staging; +}; + +static void wctx_init(WeightCtx * wctx, int n_tensors) { + size_t ctx_size = (size_t) n_tensors * ggml_tensor_overhead() + 1024; + struct ggml_init_params params = { + /*.mem_size =*/ctx_size, + /*.mem_buffer =*/NULL, + /*.no_alloc =*/true, + }; + wctx->ctx = ggml_init(params); + wctx->buffer = NULL; + wctx->pending.clear(); + wctx->pending.reserve(n_tensors); +} + +static bool wctx_alloc(WeightCtx * wctx, ggml_backend_t backend) { + wctx->buffer = ggml_backend_alloc_ctx_tensors(wctx->ctx, backend); + if (!wctx->buffer) { + fprintf(stderr, "[WeightCtx] FATAL: failed to allocate backend buffer\n"); + return false; + } + // Mark as weight buffer so ggml_backend_sched assigns ops to the correct + // backend based on weight location (avoids fallback through expansion). + ggml_backend_buffer_set_usage(wctx->buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + size_t total = 0; + for (auto & pc : wctx->pending) { + ggml_backend_tensor_set(pc.tensor, pc.src, pc.offset, pc.nbytes); + total += pc.nbytes; + } + fprintf(stderr, "[WeightCtx] Loaded %zu tensors, %.1f MB into backend\n", wctx->pending.size(), + (float) total / (1024 * 1024)); + wctx->pending.clear(); + wctx->staging.clear(); + return true; +} + +static void wctx_free(WeightCtx * wctx) { + if (wctx->buffer) { + ggml_backend_buffer_free(wctx->buffer); + } + if (wctx->ctx) { + ggml_free(wctx->ctx); + } + wctx->buffer = NULL; + wctx->ctx = NULL; +} diff --git a/tests/cossim_common.py b/tests/cossim_common.py new file mode 100755 index 0000000..1584d1c --- /dev/null +++ b/tests/cossim_common.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +"""Shared helpers for the qwentts.cpp cossim debug scripts. + +Provides Philox uniform stream, dump load and save, install_hooks for the +talker submodel, the standard stage list and the metric helpers used by +debug-base / debug-tts / debug-customvoice / debug-clone cossim scripts. + +Importing this module patches sys.path so qwen_tts upstream loads without +the V1 25Hz tokenizer (sox dependency stubbed out), and forces TF32 off on +every torch CUDA matmul path so Python results stay bit comparable across +runs and across machines. +""" + +import os +import struct +import sys +import types + +os.environ["NVIDIA_TF32_OVERRIDE"] = "0" + +import numpy as np +import torch + +torch.backends.cuda.matmul.allow_tf32 = False +torch.backends.cudnn.allow_tf32 = False +torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False +torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False +torch.set_float32_matmul_precision("highest") + +# Stub the V1 25Hz tokenizer so qwen_tts.core imports without sox. +UPSTREAM_ROOT = "/mnt/workspace/Qwen3-TTS" +sys.path.insert(0, UPSTREAM_ROOT) + +class _StubV1Config: + pass + +class _StubV1Model: + pass + +_qwen_pkg = types.ModuleType("qwen_tts") +_qwen_pkg.__path__ = [os.path.join(UPSTREAM_ROOT, "qwen_tts")] +sys.modules["qwen_tts"] = _qwen_pkg + +_core_pkg = types.ModuleType("qwen_tts.core") +_core_pkg.__path__ = [os.path.join(UPSTREAM_ROOT, "qwen_tts", "core")] +from qwen_tts.core.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config +from qwen_tts.core.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Model +_core_pkg.Qwen3TTSTokenizerV1Config = _StubV1Config +_core_pkg.Qwen3TTSTokenizerV1Model = _StubV1Model +_core_pkg.Qwen3TTSTokenizerV2Config = Qwen3TTSTokenizerV2Config +_core_pkg.Qwen3TTSTokenizerV2Model = Qwen3TTSTokenizerV2Model +sys.modules["qwen_tts.core"] = _core_pkg + +from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration +from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig +from qwen_tts.core.models.processing_qwen3_tts import Qwen3TTSProcessor +from transformers import AutoConfig, AutoModel, AutoProcessor + +# Register the Qwen3-TTS classes once per process. Calling twice raises a +# ValueError inside transformers, hence the guard. +_REGISTERED = {"done": False} + +def register_qwen3_tts(): + if _REGISTERED["done"]: + return + AutoConfig.register("qwen3_tts", Qwen3TTSConfig) + AutoModel.register(Qwen3TTSConfig, Qwen3TTSForConditionalGeneration) + AutoProcessor.register(Qwen3TTSConfig, Qwen3TTSProcessor) + _REGISTERED["done"] = True + +# Path to the C++ qwen-tts binary, relative to tests/. +BIN = "../build/qwen-tts" + +# Standard stage list shared by every cossim script. Mode-specific scripts +# may extend this list before iterating (eg. clone adds SpeakerEmb / RefCodes). +STAGES_STANDARD = [ + ("Embed", "talker-input-embed.bin"), + ("TrailingText", "trailing-text-hidden.bin"), + ("TTSPadEmbed", "tts-pad-embed.bin"), + ("L0", "talker-hidden-prefill-l0.bin"), + ("L7", "talker-hidden-prefill-l7.bin"), + ("L14", "talker-hidden-prefill-l14.bin"), + ("L21", "talker-hidden-prefill-l21.bin"), + ("L27", "talker-hidden-prefill-l27.bin"), + ("Final", "talker-hidden-prefill-final.bin"), + ("Logits", "talker-logits-prefill.bin"), + ("NextEmbStep0", "next-emb-step0.bin"), + ("TalkerHiddenStep1", "talker-hidden-step1.bin"), +] + +# Philox4x32-10 mirror of src/philox.h. Returns the same float u that +# philox_uniform_fill(seed, subseq, ctr_lo=0) gives for n=1. +PHILOX_M0 = np.uint32(0xD2511F53) +PHILOX_M1 = np.uint32(0xCD9E8D57) +PHILOX_W0 = np.uint32(0x9E3779B9) +PHILOX_W1 = np.uint32(0xBB67AE85) +TWO_POW32_INV = np.float32(2.3283064365386963e-10) + +def _mulhilo32(a, b): + p = np.uint64(a) * np.uint64(b) + lo = np.uint32(p & np.uint64(0xFFFFFFFF)) + hi = np.uint32(p >> np.uint64(32)) + return hi, lo + +def _philox_round(ctr, k0, k1): + hi0, lo0 = _mulhilo32(PHILOX_M0, ctr[0]) + hi1, lo1 = _mulhilo32(PHILOX_M1, ctr[2]) + return (np.uint32(hi1 ^ ctr[1] ^ k0), + np.uint32(lo1), + np.uint32(hi0 ^ ctr[3] ^ k1), + np.uint32(lo0)) + +def _philox4x32_10(ctr, k0, k1): + mask = np.uint64(0xFFFFFFFF) + for _ in range(9): + ctr = _philox_round(ctr, k0, k1) + k0 = np.uint32((np.uint64(k0) + np.uint64(PHILOX_W0)) & mask) + k1 = np.uint32((np.uint64(k1) + np.uint64(PHILOX_W1)) & mask) + ctr = _philox_round(ctr, k0, k1) + return ctr + +def philox_uniform(seed, subseq, ctr_lo=0): + slo = np.uint32(np.uint64(seed) & np.uint64(0xFFFFFFFF)) + shi = np.uint32(np.uint64(seed) >> np.uint64(32)) + ctr = (np.uint32(ctr_lo), + np.uint32(0), + np.uint32(np.uint64(subseq) & np.uint64(0xFFFFFFFF)), + np.uint32(np.uint64(subseq) >> np.uint64(32))) + r = _philox4x32_10(ctr, slo, shi) + return (np.float32(r[0]) + np.float32(0.5)) * TWO_POW32_INV + +# Globals advanced exactly once per multinomial sample, mirroring the C++ +# side which advances subseq_counter at every sample_top_k_p call. +_subseq_counter = [0] +_seed = [42] +_trace_samples = [False] + +def reset_philox(seed): + _subseq_counter[0] = 0 + _seed[0] = int(seed) + +def set_trace(flag): + _trace_samples[0] = bool(flag) + +def patched_multinomial(input, num_samples, replacement=False, generator=None, out=None): + """Drop in replacement for torch.multinomial(num_samples=1) that pulls + the uniform draw from our Philox stream and walks the F32 cumulative + sum the same way src/sampling.h does.""" + assert num_samples == 1, "patched_multinomial only handles num_samples=1" + probs = input + if probs.dim() == 1: + probs = probs.unsqueeze(0) + bsz, vocab = probs.shape + out_ids = torch.zeros((bsz, 1), dtype=torch.long, device=probs.device) + for b in range(bsz): + u = philox_uniform(_seed[0], _subseq_counter[0], 0) + seq = _subseq_counter[0] + _subseq_counter[0] += 1 + row = probs[b].to(torch.float32).cpu().numpy() + s = float(row.sum()) + # The C++ sampler draws u in [0, 1) and compares against acc/sum + # implicitly via acc >= u*sum. We replicate that exact arithmetic. + target = float(u) * s + acc = 0.0 + idx = vocab - 1 + for i in range(vocab): + acc += float(row[i]) + if acc >= target: + idx = i + break + out_ids[b, 0] = idx + if _trace_samples[0] and seq < 32: + print(f"[Sample-PY] subseq={seq} u={float(u):.10f} idx={idx} top_prob={float(row.max()):.6f}") + if input.dim() == 1: + return out_ids.squeeze(0) + return out_ids + +def ensure_dir(path): + os.makedirs(path, exist_ok=True) + +def save_dump(path, data): + if isinstance(data, torch.Tensor): + data = data.detach().to(torch.float32).cpu().numpy() + data = np.ascontiguousarray(data.astype(np.float32)) + shape = data.shape + with open(path, "wb") as f: + f.write(struct.pack("i", len(shape))) + for s in shape: + f.write(struct.pack("i", s)) + f.write(data.tobytes()) + +def save_dump_i32(path, data): + if isinstance(data, torch.Tensor): + data = data.detach().to(torch.int64).cpu().numpy() + data = np.ascontiguousarray(data.astype(np.int64)) + shape = data.shape + fdata = data.astype(np.float32) + with open(path, "wb") as f: + f.write(struct.pack("i", len(shape))) + for s in shape: + f.write(struct.pack("i", s)) + f.write(fdata.tobytes()) + +def load_dump(path): + raw = np.fromfile(path, dtype=np.uint8) + ndim = int(np.frombuffer(raw[0:4], dtype=np.int32)[0]) + shape = tuple(int(x) for x in np.frombuffer(raw[4:4 + 4 * ndim], dtype=np.int32)) + body = np.frombuffer(raw[4 + 4 * ndim:], dtype=np.float32) + return body.reshape(shape), shape + +def cos(a, b): + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + n = min(len(a), len(b)) + a, b = a[:n], b[:n] + d = float(np.linalg.norm(a) * np.linalg.norm(b)) + return float(np.dot(a, b) / d) if d > 1e-10 else 0.0 + +def stft_cos(a, b, win=2048, hop=512): + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + n = min(len(a), len(b)) + a, b = a[:n], b[:n] + window = np.hanning(win) + frames = (n - win) // hop + 1 + if frames <= 0: + return 0.0 + sa = np.zeros((frames, win // 2 + 1)) + sb = np.zeros((frames, win // 2 + 1)) + for i in range(frames): + s = i * hop + sa[i] = np.abs(np.fft.rfft(a[s:s + win] * window)) + sb[i] = np.abs(np.fft.rfft(b[s:s + win] * window)) + return cos(sa.ravel(), sb.ravel()) + +def install_hooks(model, dump_dir, bisect_layers=(0, 7, 14, 21, 27)): + """Capture every intermediate tensor we can pin against the C++ side. + Layout : bisection layers, final norm, prefill logits, the input embed + fed to the talker and the trailing-text overlay tensors that drive the + next-token embedding sums during generation. Counters track how many + times the talker submodel has run so step 1 (first single-token forward + after prefill) gets its own dump.""" + seen_prefill = {"done": False} + seen_codes = {"done": False} + # talker_step counts how many times talker_model.forward has been called + # after the prefill. Prefill itself is recorded as 0, the first single + # token forward is 1, and so on. Single token forwards are detected by + # inputs_embeds.shape[1] == 1 in the pre hook. + talker_step = {"n": 0} + + talker_model = model.talker.model + talker_lm = model.talker + + seen_layers = {idx: False for idx in bisect_layers} + def make_layer_hook(layer_idx): + def hook(module, inputs, output): + if seen_layers[layer_idx]: + return + h = output[0] if isinstance(output, tuple) else output + if h.dim() == 3: + save_dump(os.path.join(dump_dir, f"talker-hidden-prefill-l{layer_idx}.bin"), h[0]) + seen_layers[layer_idx] = True + return hook + for layer_idx in bisect_layers: + talker_model.layers[layer_idx].register_forward_hook(make_layer_hook(layer_idx)) + + seen_norm = {"done": False} + def norm_hook(module, inputs, output): + if seen_norm["done"]: + return + if output.dim() == 3 and output.shape[1] > 1: + save_dump(os.path.join(dump_dir, "talker-hidden-prefill-final.bin"), output[0]) + seen_norm["done"] = True + talker_model.norm.register_forward_hook(norm_hook) + + # Pre-hook on the talker submodel : sees inputs_embeds whether the outer + # talker forward was invoked with input_ids (single token step) or + # inputs_embeds (prefill). The submodel always receives inputs_embeds + # because the wrapper rebuilds it before calling self.model. + def talker_model_pre_hook(module, args, kwargs): + ie = kwargs.get("inputs_embeds", None) + if ie is None: + return + if ie.dim() != 3: + return + if ie.shape[1] > 1: + return + if talker_step["n"] == 0: + save_dump(os.path.join(dump_dir, "next-emb-step0.bin"), ie[0, 0]) + talker_step["n"] += 1 + talker_model.register_forward_pre_hook(talker_model_pre_hook, with_kwargs=True) + + # Post-hook on the talker submodel : captures last_hidden_state at step + # 1 (first single token forward). That tensor is what feeds the code + # predictor at step 1, so any drift between Python and C++ tells us the + # next-emb-step0 changed the talker forward result. + talker_post_step = {"n": 0} + def talker_model_post_hook(module, inputs, output): + last = output.last_hidden_state if hasattr(output, "last_hidden_state") else output[0] + if last.dim() != 3 or last.shape[1] != 1: + return + if talker_post_step["n"] == 0: + save_dump(os.path.join(dump_dir, "talker-hidden-step1.bin"), last[0, -1]) + talker_post_step["n"] += 1 + talker_model.register_forward_hook(talker_model_post_hook) + + # Talker LM wrapper hook : captures the prefill input embed (the talker + # codec_embedding sum + text projection that mirrors what + # prompt_builder_build produces in C++), the prefill logits, and the + # trailing_text_hidden / tts_pad_embed overlay tensors carried by the + # output dataclass at every step (we only dump them once). + seen_overlay = {"done": False} + orig_talker_forward = talker_lm.forward + def hooked_talker_forward(*args, **kwargs): + inputs_embeds = kwargs.get("inputs_embeds", None) + if (inputs_embeds is not None and inputs_embeds.dim() == 3 + and inputs_embeds.shape[1] > 1 and not seen_prefill["done"]): + save_dump(os.path.join(dump_dir, "talker-input-embed.bin"), inputs_embeds[0]) + seen_prefill["done"] = True + out = orig_talker_forward(*args, **kwargs) + if (out is not None and getattr(out, "logits", None) is not None + and not seen_codes["done"]): + logits = out.logits + if logits.dim() == 3 and logits.shape[1] > 1: + save_dump(os.path.join(dump_dir, "talker-logits-prefill.bin"), logits[0, -1]) + seen_codes["done"] = True + if (out is not None and not seen_overlay["done"] + and getattr(out, "trailing_text_hidden", None) is not None + and getattr(out, "tts_pad_embed", None) is not None): + tth = out.trailing_text_hidden + tpe = out.tts_pad_embed + if tth.dim() == 3 and tpe.dim() == 3: + save_dump(os.path.join(dump_dir, "trailing-text-hidden.bin"), tth[0]) + save_dump(os.path.join(dump_dir, "tts-pad-embed.bin"), tpe[0, 0]) + seen_overlay["done"] = True + return out + talker_lm.forward = hooked_talker_forward + +def pair(name, dump_cpp, dump_pt): + a, _ = load_dump(os.path.join(dump_cpp, name)) + b, _ = load_dump(os.path.join(dump_pt, name)) + return a, b + +def metric(a, b): + n = min(a.size, b.size) + af = a.astype(np.float64).ravel()[:n] + bf = b.astype(np.float64).ravel()[:n] + d = np.abs(af - bf) + nrm_a = float(np.linalg.norm(af)) + nrm_b = float(np.linalg.norm(bf)) + c = float(np.dot(af, bf) / (nrm_a * nrm_b)) if nrm_a > 1e-10 and nrm_b > 1e-10 else 0.0 + return c, float(d.max()), float(d.mean()) + +def compare_stages(stages, dump_cpp, dump_pt): + """Iterate the stages list and print one line per pair. Skips silently + when a dump file is missing (eg. a mode that does not produce a given + intermediate).""" + for label, name in stages: + try: + a, b = pair(name, dump_cpp, dump_pt) + except FileNotFoundError: + print(f"[Cossim] {label} skipped (missing dump)") + continue + c, mx, mn = metric(a, b) + print(f"[Cossim] {label} cos: {c:.6f} max: {mx:.4e} mean: {mn:.4e}") + +def compare_exact_i32(name, dump_cpp, dump_pt, label): + """Compare two int dumps stored as f32 (the encoding path used by both + save_dump_i32 in Python and debug_dump_i32_as_f32 in C++). Prints an + exact match percentage. Returns the percentage as a float.""" + a, b = pair(name, dump_cpp, dump_pt) + ai = a.astype(np.int64).ravel() + bi = b.astype(np.int64).ravel() + n = min(ai.size, bi.size) + pct = 100.0 * float(np.mean(ai[:n] == bi[:n])) + print(f"[Cossim] {label} exact: {pct:.2f}% ({n} values)") + return pct diff --git a/tests/debug-base-cossim.py b/tests/debug-base-cossim.py new file mode 100755 index 0000000..eb9dcb4 --- /dev/null +++ b/tests/debug-base-cossim.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Cossim debug : C++ qwen-tts vs Python Qwen3-TTS on the Base 1.7B path. + +Inputs (relative to CWD = tests/) : + ../examples/prompt.txt target text fed to both pipelines + +Default mode is greedy (do_sample=False on both sides). The forward +chain is dumped layer by layer and compared paired with the Python +upstream hooks installed by cossim_common.install_hooks. Both pipelines +run on CUDA by default, the wrapper shell sweeps backends and quants. + +Dumps land in cpp/base/ (C++) and python/base/ (Python). The script +compares each matching .bin pair via cosine similarity over the f32 +payload, plus exact match rate for tensors that originated as int +(codec codes, prompt ids). +""" + +import argparse +import os +import subprocess +import sys + +import numpy as np +import soundfile as sf +import torch + +import cossim_common as cc + +MODEL_T = "../models/qwen-talker-1.7b-base-{q}.gguf" +MODEL_CDC_T = "../models/qwen-tokenizer-12hz-{q}.gguf" +CKPT = "../checkpoints/Qwen3-TTS-12Hz-1.7B-Base" +DUMP_CPP = "cpp/base" +DUMP_PT = "python/base" + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--prompt", default="../examples/prompt.txt") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--lang", default="english") + ap.add_argument("--quant", default="F32", + help="GGUF quantization suffix (F32, BF16, Q8_0, Q4_K_M)") + ap.add_argument("--out-pt", default=os.path.join(DUMP_PT, "base-python.wav")) + ap.add_argument("--out-cpp", default=os.path.join(DUMP_CPP, "base-cpp.wav")) + ap.add_argument("--max-new-tokens", type=int, default=64) + ap.add_argument("--trace", action="store_true", + help="print per sample u and idx for the first 32 samples") + args = ap.parse_args() + + cc.ensure_dir(DUMP_PT) + cc.ensure_dir(DUMP_CPP) + os.makedirs(os.path.dirname(args.out_pt) or ".", exist_ok=True) + + with open(args.prompt, "r", encoding="utf-8") as f: + text = f.read().strip() + print(f"[Input] Prompt: {len(text)} chars: {text[:60]}{'...' if len(text) > 60 else ''}") + print(f"[Input] Lang: {args.lang} Seed: {args.seed} MaxNewTokens: {args.max_new_tokens}") + print(f"[Input] Mode: greedy") + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + cc.set_trace(args.trace) + + cc.register_qwen3_tts() + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"[Python] Device: {device}") + model = cc.AutoModel.from_pretrained( + CKPT, + device_map=device, + dtype=torch.float32, + attn_implementation="eager", + ).eval() + processor = cc.AutoProcessor.from_pretrained(CKPT, fix_mistral_regex=True) + + assistant_text = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + inp = processor(text=assistant_text, return_tensors="pt", padding=True) + input_ids = inp["input_ids"].to(device) + if input_ids.dim() == 1: + input_ids = input_ids.unsqueeze(0) + print(f"[Python] InputIds shape: {tuple(input_ids.shape)}") + cc.save_dump_i32(os.path.join(DUMP_PT, "prompt-ids.bin"), input_ids[0]) + + cc.install_hooks(model, DUMP_PT) + + # Custom subtalker_* kwargs are forwarded to talker.forward but not + # declared on GenerationMixin, so transformers 4.57 rejects them under + # the strict validator. Disable it on the talker only. + model.talker._validate_model_kwargs = lambda *a, **k: None + + # Greedy hardcoded : argmax on both talker c0 and code predictor sub + # codes. Stochastic mode is not exercised here because the F32 drift + # between torch CUDA cuBLAS and ggml CUDA matmul on Qwen3 norm_w + # inflated activations propagates through the FFN and flips multinomial + # picks in flat distributions, breaking bit exactness. Argmax is robust + # to that drift, so greedy gives 100 percent CodesFull match and + # validates the full forward + sampling chain. + gen_kwargs = dict( + do_sample = False, + top_k = 1, + top_p = 1.0, + temperature = 1.0, + subtalker_dosample = False, + subtalker_top_k = 1, + subtalker_top_p = 1.0, + subtalker_temperature = 1.0, + repetition_penalty = 1.0, + ) + + talker_codes_list, _ = model.generate( + input_ids=[input_ids], + languages=[args.lang], + non_streaming_mode=True, + max_new_tokens=args.max_new_tokens, + **gen_kwargs, + ) + codes = talker_codes_list[0] + print(f"[Python] Codes shape: {tuple(codes.shape)} (T_frames, num_code_groups)") + cc.save_dump_i32(os.path.join(DUMP_PT, "codes-full.bin"), codes) + cc.save_dump_i32(os.path.join(DUMP_PT, "codes-step0.bin"), codes[0]) + + wavs, fs = model.speech_tokenizer.decode([{"audio_codes": codes}]) + audio_pt = np.asarray(wavs[0], dtype=np.float32) + sf.write(args.out_pt, audio_pt, fs, subtype="FLOAT") + cc.save_dump(os.path.join(DUMP_PT, "output-audio.bin"), audio_pt) + print(f"[Python] Audio: {audio_pt.shape[0]} samples {fs} Hz {audio_pt.shape[0]/fs:.2f}s -> {args.out_pt}") + + if not os.path.isfile(cc.BIN): + print(f"[Cossim] FATAL: {cc.BIN} not found, build qwen-tts first") + sys.exit(1) + model_lm = MODEL_T.format(q=args.quant) + model_cdc = MODEL_CDC_T.format(q=args.quant) + for p in (model_lm, model_cdc): + if not os.path.isfile(p): + print(f"[Cossim] FATAL: GGUF not found: {p}") + sys.exit(1) + print(f"[Quant] {args.quant} -> {model_lm} + {model_cdc}") + + del model + if device == "cuda": + torch.cuda.empty_cache() + + cmd = [ + cc.BIN, + "--model", model_lm, + "--codec", model_cdc, + "--seed", str(args.seed), + "--text", text, + "--lang", args.lang, + "--max-new", str(args.max_new_tokens), + "--dump", DUMP_CPP, + "-o", args.out_cpp, + "--greedy", + ] + print(f"[GGML] Cmd: {' '.join(cmd)}") + r = subprocess.run(cmd) + if r.returncode != 0: + sys.exit(r.returncode) + + audio_cpp, sr = sf.read(args.out_cpp) + if audio_cpp.ndim > 1: + audio_cpp = audio_cpp[:, 0] + audio_cpp = audio_cpp.astype(np.float32) + print(f"[GGML] Audio: {audio_cpp.shape[0]} samples {sr} Hz {audio_cpp.shape[0]/sr:.2f}s -> {args.out_cpp}") + + cc.compare_exact_i32("prompt-ids.bin", DUMP_CPP, DUMP_PT, "PromptIDs") + cc.compare_stages(cc.STAGES_STANDARD, DUMP_CPP, DUMP_PT) + cc.compare_exact_i32("codes-full.bin", DUMP_CPP, DUMP_PT, "CodesFull") + + aa, ab = cc.pair("output-audio.bin", DUMP_CPP, DUMP_PT) + print(f"[Cossim] Audio cos: {cc.cos(aa, ab):.6f}") + + n = min(audio_cpp.size, audio_pt.size) + print(f"[Cossim] WAV stft_cos: {cc.stft_cos(audio_cpp[:n], audio_pt[:n]):.6f} samples: {n}") + +if __name__ == "__main__": + main() diff --git a/tests/debug-base-cossim.sh b/tests/debug-base-cossim.sh new file mode 100755 index 0000000..71fa9be --- /dev/null +++ b/tests/debug-base-cossim.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +for backend in CUDA0 Vulkan0 CPU; do + for quant in F32 BF16 Q8_0 Q4_K_M; do + GGML_BACKEND=$backend ./debug-base-cossim.py --quant $quant \ + 2>&1 | tee base-${backend}-${quant}.log + done +done diff --git a/tools/quantize.cpp b/tools/quantize.cpp new file mode 100644 index 0000000..85d6692 --- /dev/null +++ b/tools/quantize.cpp @@ -0,0 +1,473 @@ +// quantize.cpp : GGUF requantizer for qwen +// Reads BF16 GGUF, writes quantized GGUF with mixed-precision K-quant policy. +// Policy mirrors llama-quantize: important tensors (v_proj, down_proj) get +// bumped in S/M variants, embed_tokens always Q6_K, norms promoted to F32. +// Streaming write: one tensor at a time, low memory footprint for small configs. +// +// Usage: quantize +// Types: Q2_K Q3_K_S Q3_K_M Q3_K_L Q4_K_S Q4_K_M Q5_K_S Q5_K_M Q6_K Q8_0 + +#include +#include +#include +#include + +#ifdef _WIN32 +# define NOMINMAX +# include +# define strcasecmp _stricmp +#else +# include +# include +# include +# include +#endif + +#include "ggml.h" +#include "gguf.h" +#include "version.h" + +// Quant variant: base type + optional bump rules for important tensors +struct QuantVariant { + const char * name; + enum ggml_type base; + enum ggml_type bump; // type for "important" tensors (or COUNT = no bump) + enum ggml_type embed; // type for embed_tokens (or COUNT = same as base) + // bump_mode: 0=none, 1=first N layers, 2=first+last+every 3rd, 3=all important + int bump_mode; + int bump_n; // for mode 1: number of layers to bump +}; + +static const QuantVariant VARIANTS[] = { + // name base bump embed mode n + { "BF16", GGML_TYPE_BF16, GGML_TYPE_COUNT, GGML_TYPE_BF16, 0, 0 }, + { "Q2_K", GGML_TYPE_Q2_K, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, 1, 4 }, + { "Q3_K_S", GGML_TYPE_Q3_K, GGML_TYPE_COUNT, GGML_TYPE_Q6_K, 0, 0 }, + { "Q3_K_M", GGML_TYPE_Q3_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, 2, 0 }, + { "Q3_K_L", GGML_TYPE_Q3_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, 3, 0 }, + { "Q4_K_S", GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, 1, 4 }, + { "Q4_K_M", GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, 2, 0 }, + { "Q5_K_S", GGML_TYPE_Q5_K, GGML_TYPE_COUNT, GGML_TYPE_Q6_K, 0, 0 }, + { "Q5_K_M", GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, 2, 0 }, + { "Q6_K", GGML_TYPE_Q6_K, GGML_TYPE_COUNT, GGML_TYPE_Q6_K, 0, 0 }, + { "Q8_0", GGML_TYPE_Q8_0, GGML_TYPE_COUNT, GGML_TYPE_Q8_0, 0, 0 }, +}; + +static const QuantVariant * find_variant(const char * s) { + for (const auto & v : VARIANTS) { + if (strcasecmp(s, v.name) == 0) { + return &v; + } + } + return nullptr; +} + +// Extract layer index from HF tensor name: model.layers.N.xxx -> N, else -1 +static int extract_layer(const char * name) { + const char * p = strstr(name, "layers."); + if (!p) { + return -1; + } + return atoi(p + 7); +} + +// Important tensors for S/M: v_proj + down_proj +static bool is_important_sm(const char * name) { + return (strstr(name, "v_proj.weight") != nullptr) || (strstr(name, "down_proj.weight") != nullptr); +} + +// Important tensors for L: v_proj + down_proj + o_proj +static bool is_important_l(const char * name) { + return is_important_sm(name) || (strstr(name, "o_proj.weight") != nullptr); +} + +// Tensors accessed via ggml_get_rows (text token embeddings, audio token +// embeddings, codebook lookups). These must use a type the CUDA get_rows +// kernel supports : F32, F16, BF16, Q4_0, Q4_1, Q5_0, Q5_1, Q8_0. K-quants +// are NOT supported. +// +// In the qwen3-tts naming convention : +// token_embd / output LM standard names +// talker.codec_embd / talker.text_embd TTS specific embeddings +// code_pred.codec_embd.{i} per codebook embeddings on the predictor +// code_pred.lm_head.{i} per codebook output heads +// tok_enc.vq_*.{i}.codebook RVQ codebook tables, nearest neighbor lookup +// tok_dec.vq_*.{i}.codebook RVQ codebook tables, F.embedding lookup +static bool is_embed(const char * name) { + return strstr(name, "token_embd") != nullptr || strstr(name, "embed_tokens.weight") != nullptr || + strstr(name, "audio_embeddings.weight") != nullptr || strstr(name, ".codec_embd") != nullptr || + strstr(name, ".text_embd") != nullptr; +} + +// Should this tensor be quantized at all? +// +// Single source of truth for the quantization policy. Applies to EVERY +// variant (BF16, Q8_0, Q6_K, Q5_K_M, Q4_K_M, ...) : tensors that return +// false here keep their source dtype (F32) regardless of the requested +// type. Conv weights pass through the main loop and fall back to F16 when +// the row width does not divide the variant block size (kernel K=7,3,1,...). +// gf_load_conv_f16 then memcpys F16 source straight to the F16 backend +// tensor (ARM im2col strict requirement, see src/gguf-weights.h). +// +// Sensitive tensors that MUST stay in full precision : +// tok_enc.vq_*.{i}.codebook RVQ codebook tables, encoder side +// tok_dec.vq_*.{i}.codebook RVQ codebook tables, decoder side +// tok_enc.vq_*.input_proj linear wrapping the RVQ encode loop +// tok_enc.vq_*.output_proj linear wrapping the RVQ encode loop +// tok_dec.vq_*.output_proj linear wrapping the RVQ decode loop +// Nearest-neighbor lookup is sensitive to per-row quantization noise ; +// even BF16 destroys the mantissa enough to mis-select codes and break +// voice cloning. Same philosophy as acestep.cpp keeping VAE-critical +// paths in full precision. +static bool should_quantize(const char * name, int n_dims, const char * arch) { + if (strstr(arch, "vae")) { + return false; + } + if (n_dims < 2) { + return false; + } + if (strstr(arch, "text-enc") && strstr(name, "embed_tokens")) { + return false; + } + if (strstr(name, "silence_latent")) { + return false; + } + if (strstr(name, "scale_shift_table")) { + return false; + } + if (strstr(name, "null_condition_emb")) { + return false; + } + // Snake activation parameters : stored as per-channel floats, are + // activation parameters, not weights. The DAC loaders widen them to + // F32 on the backend with a reciprocal transform, no other dtype + // path. Keep them source-dtype in every variant. Both the legacy + // snake1/snake2 names from the old layout and the standard tok_dec + // snake naming are matched. + if (strstr(name, ".snake.alpha") || strstr(name, ".snake.beta") || strstr(name, ".act1.alpha") || + strstr(name, ".act1.beta") || strstr(name, ".act2.alpha") || strstr(name, ".act2.beta") || + strstr(name, ".snake1.alpha") || strstr(name, ".snake2.alpha")) { + return false; + } + // RVQ codebooks and the linear projections wrapping them : nearest + // neighbor lookup is sensitive to per-row quantization noise. Q8_0 + // and K-quants break reference audio encoding and tank voice cloning ; + // BF16 already loses enough mantissa to drift codes. Keep at F32 in + // every variant. + if (strstr(name, "tok_enc.vq_") || strstr(name, "tok_dec.vq_")) { + return false; + } + // Speaker encoder final FC, semantically equivalent to the old + // top level fc.weight / fc2.weight in earlier layouts. + if (strstr(name, "spk_enc.fc.weight")) { + return false; + } + return true; +} + +// Decide target type for a single tensor given the variant + layer info +static enum ggml_type pick_type(const char * name, + int n_dims, + const char * arch, + const QuantVariant & v, + int n_layers) { + if (!should_quantize(name, n_dims, arch)) { + return GGML_TYPE_COUNT; + } + + // embed_tokens in LM: use embed type + if (is_embed(name) && !strstr(arch, "text-enc")) { + return (v.embed != GGML_TYPE_COUNT) ? v.embed : v.base; + } + + // Important tensor bump logic + bool important = (v.bump_mode == 3) ? is_important_l(name) : is_important_sm(name); + + if (important && v.bump != GGML_TYPE_COUNT) { + int layer = extract_layer(name); + bool bumped = false; + switch (v.bump_mode) { + case 1: // first N layers only + bumped = (layer >= 0 && layer < v.bump_n); + break; + case 2: + { // M variant: first few + last few + every 3rd + int ql = n_layers; + bumped = (layer >= 0) && (layer < ql / 9 || layer >= ql - ql / 7 || layer % 3 == 0); + break; + } + case 3: // L variant: all important tensors (v+down+o_proj) + bumped = true; + break; + } + if (bumped) { + return v.bump; + } + } + + return v.base; +} + +// Convert source data to F32 +static bool to_f32(const void * src, float * dst, int64_t n, enum ggml_type type) { + switch (type) { + case GGML_TYPE_BF16: + ggml_bf16_to_fp32_row((const ggml_bf16_t *) src, dst, n); + return true; + case GGML_TYPE_F16: + ggml_fp16_to_fp32_row((const ggml_fp16_t *) src, dst, n); + return true; + case GGML_TYPE_F32: + memcpy(dst, src, (size_t) n * sizeof(float)); + return true; + default: + return false; + } +} + +int main(int argc, char ** argv) { + if (argc != 4) { + fprintf(stderr, "qwentts.cpp %s\n\n", QWEN_VERSION); + fprintf(stderr, "Usage: %s \n", argv[0]); + fprintf(stderr, "Types:"); + for (const auto & v : VARIANTS) { + fprintf(stderr, " %s", v.name); + } + fprintf(stderr, "\n"); + return 1; + } + + const char * inp_path = argv[1]; + const char * out_path = argv[2]; + const QuantVariant * variant = find_variant(argv[3]); + + if (!variant) { + fprintf(stderr, "[Quantize] Unknown type: %s\n", argv[3]); + return 1; + } + + fprintf(stderr, "[Quantize] %s -> %s (%s)\n", inp_path, out_path, variant->name); + + // Mmap input file +#ifdef _WIN32 + HANDLE fh = CreateFileA(inp_path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); + if (fh == INVALID_HANDLE_VALUE) { + fprintf(stderr, "[Quantize] Failed to open %s\n", inp_path); + return 1; + } + HANDLE mh = CreateFileMappingA(fh, NULL, PAGE_READONLY, 0, 0, NULL); + if (!mh) { + fprintf(stderr, "[Quantize] CreateFileMapping failed %s\n", inp_path); + CloseHandle(fh); + return 1; + } + void * mapping = MapViewOfFile(mh, FILE_MAP_READ, 0, 0, 0); + if (!mapping) { + fprintf(stderr, "[Quantize] MapViewOfFile failed %s\n", inp_path); + CloseHandle(mh); + CloseHandle(fh); + return 1; + } +#else + int fd = open(inp_path, O_RDONLY); + if (fd < 0) { + perror("open"); + return 1; + } + struct stat st; + fstat(fd, &st); + size_t file_size = (size_t) st.st_size; + void * mapping = mmap(nullptr, file_size, PROT_READ, MAP_PRIVATE, fd, 0); + if (mapping == MAP_FAILED) { + perror("mmap"); + close(fd); + return 1; + } +#endif + + // Parse input GGUF + struct gguf_init_params params = { /*no_alloc=*/true, /*ctx=*/nullptr }; + struct ggml_context * meta = nullptr; + params.ctx = &meta; + + struct gguf_context * inp = gguf_init_from_file(inp_path, params); + if (!inp) { + fprintf(stderr, "[Quantize] Failed to read %s\n", inp_path); +#ifdef _WIN32 + UnmapViewOfFile(mapping); + CloseHandle(mh); + CloseHandle(fh); +#else + munmap(mapping, file_size); + close(fd); +#endif + return 1; + } + + const size_t data_off = gguf_get_data_offset(inp); + const int n_tensors = (int) gguf_get_n_tensors(inp); + + // Read architecture + char arch[64] = "unknown"; + { + int64_t idx = gguf_find_key(inp, "general.architecture"); + if (idx >= 0) { + const char * s = gguf_get_val_str(inp, (int) idx); + snprintf(arch, sizeof(arch), "%s", s); + } + } + + // Read block count for bump policy. Standard archs publish + // {arch}.block_count, but multi LM archs like qwen3-tts namespace + // it under a sub-component (talker) so we try a small list of + // fallbacks before giving up. + int n_layers = 0; + { + const char * candidates[] = { + ".block_count", + ".talker.block_count", + ".decoder.num_hidden_layers", + }; + for (const char * suffix : candidates) { + char key[160]; + snprintf(key, sizeof(key), "%s%s", arch, suffix); + int64_t idx = gguf_find_key(inp, key); + if (idx >= 0) { + n_layers = (int) gguf_get_val_u32(inp, (int) idx); + break; + } + } + } + + fprintf(stderr, "[Quantize] Arch=%s Layers=%d\n", arch, n_layers); + + // Create output GGUF: copy KV metadata + struct gguf_context * out = gguf_init_empty(); + gguf_set_kv(out, inp); + gguf_set_val_u32(out, "general.quantization_version", 2); + gguf_set_val_str(out, "general.file_type", variant->name); + + // Plan: for each tensor, decide target type + struct TensorPlan { + enum ggml_type target; + bool quantize; + }; + + std::vector plans((size_t) n_tensors); + + for (int i = 0; i < n_tensors; i++) { + const char * name = gguf_get_tensor_name(inp, i); + struct ggml_tensor * t = ggml_get_tensor(meta, name); + const int n_dims = ggml_n_dims(t); + + gguf_add_tensor(out, t); + plans[(size_t) i] = { GGML_TYPE_COUNT, false }; + + enum ggml_type target = pick_type(name, n_dims, arch, *variant, n_layers); + + if (target == GGML_TYPE_COUNT) { + continue; + } + + bool can_convert = (t->type == GGML_TYPE_BF16 || t->type == GGML_TYPE_F16 || t->type == GGML_TYPE_F32); + bool aligned = (t->ne[0] % ggml_blck_size(target) == 0); + + // Conv kernels (K=7,3,1,...) cannot fit a block-quant row : fall back + // to F16. F16 has no block size, 10-bit mantissa beats BF16 (7) and + // Q* effective on these weights, and gf_load_conv_f16 memcpys F16 + // source straight to the F16 backend tensor at load time. + if (can_convert && !aligned) { + target = GGML_TYPE_F16; + aligned = true; + } + + if (can_convert && aligned) { + gguf_set_tensor_type(out, name, target); + plans[(size_t) i] = { target, true }; + } + } + + // Write metadata only (header + tensor info, no data) + bool ok = gguf_write_to_file(out, out_path, true); + if (!ok) { + fprintf(stderr, "[Quantize] Failed to write metadata %s\n", out_path); + return 1; + } + + // Stream tensor data one at a time (low memory) + FILE * fout = fopen(out_path, "ab"); + if (!fout) { + fprintf(stderr, "[Quantize] Failed to open %s for append\n", out_path); + return 1; + } + + const size_t alignment = gguf_get_alignment(out); + int n_quantized = 0; + int64_t bytes_in = 0, bytes_out = 0; + size_t data_pos = 0; + + for (int i = 0; i < n_tensors; i++) { + const char * name = gguf_get_tensor_name(inp, i); + struct ggml_tensor * t = ggml_get_tensor(meta, name); + const int64_t nel = ggml_nelements(t); + const size_t src_size = ggml_nbytes(t); + const size_t t_off = gguf_get_tensor_offset(inp, i); + const void * src = (const uint8_t *) mapping + data_off + t_off; + + bytes_in += (int64_t) src_size; + + // Pad to alignment boundary + size_t pad = (alignment - (data_pos % alignment)) % alignment; + if (pad > 0) { + uint8_t zeros[64] = {}; + fwrite(zeros, 1, pad, fout); + data_pos += pad; + } + + const TensorPlan & plan = plans[(size_t) i]; + + if (plan.quantize) { + // Quantize: src -> f32 -> target + std::vector f32((size_t) nel); + to_f32(src, f32.data(), nel, t->type); + + const int64_t n_per_row = t->ne[0]; + const int64_t nrows = nel / n_per_row; + const size_t qsize = ggml_row_size(plan.target, n_per_row) * (size_t) nrows; + + std::vector qbuf(qsize); + ggml_quantize_chunk(plan.target, f32.data(), qbuf.data(), 0, nrows, n_per_row, nullptr); + + fwrite(qbuf.data(), 1, qsize, fout); + data_pos += qsize; + bytes_out += (int64_t) qsize; + n_quantized++; + } else { + // Keep as-is + fwrite(src, 1, src_size, fout); + data_pos += src_size; + bytes_out += (int64_t) src_size; + } + } + + fclose(fout); + + fprintf(stderr, "[Quantize] Quantized %d/%d tensors\n", n_quantized, n_tensors); + fprintf(stderr, "[Quantize] %.1f GB -> %.1f GB (%.1fx)\n", (double) bytes_in / 1e9, (double) bytes_out / 1e9, + bytes_out > 0 ? (double) bytes_in / (double) bytes_out : 0.0); + fprintf(stderr, "[Quantize] Wrote %s\n", out_path); + + gguf_free(out); + gguf_free(inp); + ggml_free(meta); +#ifdef _WIN32 + UnmapViewOfFile(mapping); + CloseHandle(mh); + CloseHandle(fh); +#else + munmap(mapping, file_size); + close(fd); +#endif + + return 0; +} diff --git a/tools/qwen-codec.cpp b/tools/qwen-codec.cpp new file mode 100644 index 0000000..e483542 --- /dev/null +++ b/tools/qwen-codec.cpp @@ -0,0 +1,283 @@ +// qwen-codec.cpp: codec CLI for Qwen3-TTS. +// +// Encode a 24 kHz mono WAV into RVQ codes (.rvq), or decode RVQ codes +// back into a 24 kHz mono float32 WAV. Mode is inferred from the input +// file extension: .wav in -> encode, .rvq in -> decode. Output is +// auto-named next to the input file by swapping the extension. +// +// File format (.rvq): flat code stream packed at 11 bits per code, +// LSB-first, no header. Layout is [K, T] row-major. K is fixed by the +// codec config in the GGUF (16 codebooks for the 12Hz tokenizer, +// codebook_size = 2048). T is the frame count derived from filesize. + +#include "audio-io.h" +#include "backend.h" +#include "pipeline-codec.h" +#include "version.h" + +#include +#include +#include +#include +#include +#include + +static const uint32_t QWEN_RVQ_CODE_MASK = (1u << QWEN_TOKENIZER_CODE_BITS) - 1u; + +static void print_usage(const char * prog) { + fprintf(stderr, "qwentts.cpp %s\n\n", QWEN_VERSION); + fprintf(stderr, + "Usage: %s --model [-i ] [--format ]\n\n" + "Required:\n" + " --model Codec GGUF (qwen-tokenizer-12hz-*.gguf)\n\n" + "Optional:\n" + " -i Input. WAV -> encode, .rvq -> decode\n" + " --format WAV output format: wav16, wav24, wav32 (default: wav16)\n\n" + "Output is auto-named next to input : clip.wav -> clip.rvq, clip.rvq -> clip.wav.\n" + "When -i is omitted, runs a load self-test of the codec GGUF.\n", + prog); +} + +// Symmetric unpack: reads N codes from packed bytes (11 bits LSB-first). +static std::vector unpack_codes(const std::vector & in, size_t n_codes) { + std::vector out(n_codes); + uint64_t acc = 0; + int bits_in_acc = 0; + size_t in_pos = 0; + for (size_t i = 0; i < n_codes; i++) { + while (bits_in_acc < QWEN_TOKENIZER_CODE_BITS && in_pos < in.size()) { + acc |= ((uint64_t) in[in_pos++]) << bits_in_acc; + bits_in_acc += 8; + } + out[i] = (int32_t) (acc & QWEN_RVQ_CODE_MASK); + acc >>= QWEN_TOKENIZER_CODE_BITS; + bits_in_acc -= QWEN_TOKENIZER_CODE_BITS; + } + return out; +} + +// Pack flat int32 codes into 11-bit LSB-first packed bytes. Output size is +// ceil(N * 11 / 8) bytes. +static std::vector pack_codes(const std::vector & codes) { + const size_t total_bits = codes.size() * (size_t) QWEN_TOKENIZER_CODE_BITS; + std::vector out((total_bits + 7) / 8, 0); + uint64_t acc = 0; + int bits_in_acc = 0; + size_t out_pos = 0; + for (size_t i = 0; i < codes.size(); i++) { + acc |= ((uint64_t) ((uint32_t) codes[i] & QWEN_RVQ_CODE_MASK)) << bits_in_acc; + bits_in_acc += QWEN_TOKENIZER_CODE_BITS; + while (bits_in_acc >= 8) { + out[out_pos++] = (uint8_t) (acc & 0xFF); + acc >>= 8; + bits_in_acc -= 8; + } + } + if (bits_in_acc > 0) { + out[out_pos++] = (uint8_t) (acc & 0xFF); + } + return out; +} + +// Read a .rvq file and unpack it into K*T codes. T is inferred from the +// file size: T = (filesize * 8) / (K * QWEN_TOKENIZER_CODE_BITS). +static bool read_rvq(const char * path, int K, std::vector & codes, int * n_frames) { + FILE * f = fopen(path, "rb"); + if (!f) { + fprintf(stderr, "[Codec] FATAL: cannot open %s\n", path); + return false; + } + fseek(f, 0, SEEK_END); + long sz = ftell(f); + fseek(f, 0, SEEK_SET); + if (sz <= 0) { + fprintf(stderr, "[Codec] FATAL: %s is empty\n", path); + fclose(f); + return false; + } + std::vector buf((size_t) sz); + if (fread(buf.data(), 1, buf.size(), f) != buf.size()) { + fprintf(stderr, "[Codec] FATAL: short read on %s\n", path); + fclose(f); + return false; + } + fclose(f); + + const size_t total_bits = (size_t) sz * 8; + const size_t n_codes = total_bits / (size_t) QWEN_TOKENIZER_CODE_BITS; + if (n_codes == 0 || (n_codes % (size_t) K) != 0) { + fprintf(stderr, "[Codec] FATAL: %s yields %zu codes, not a multiple of K=%d\n", path, n_codes, K); + return false; + } + codes = unpack_codes(buf, n_codes); + *n_frames = (int) (n_codes / (size_t) K); + return true; +} + +// Pack and write a .rvq file. +static bool write_rvq(const char * path, const std::vector & codes) { + std::vector packed = pack_codes(codes); + FILE * f = fopen(path, "wb"); + if (!f) { + fprintf(stderr, "[Codec] FATAL: cannot open %s for write\n", path); + return false; + } + if (fwrite(packed.data(), 1, packed.size(), f) != packed.size()) { + fprintf(stderr, "[Codec] FATAL: short write on %s\n", path); + fclose(f); + return false; + } + fclose(f); + return true; +} + +// Replace or append extension on a path string. +static std::string swap_ext(const std::string & path, const char * ext) { + size_t dot = path.find_last_of('.'); + size_t sep = path.find_last_of("/\\"); + if (dot != std::string::npos && (sep == std::string::npos || dot > sep)) { + return path.substr(0, dot) + ext; + } + return path + ext; +} + +// 0: unsupported, 1: encode (.wav in), 2: decode (.rvq in). +static int infer_mode(const char * path) { + size_t n = strlen(path); + if (n >= 4 && strcmp(path + n - 4, ".wav") == 0) { + return 1; + } + if (n >= 4 && strcmp(path + n - 4, ".rvq") == 0) { + return 2; + } + return 0; +} + +int main(int argc, char ** argv) { + if (argc <= 1) { + print_usage(argv[0]); + return 0; + } + + const char * model_path = NULL; + const char * input_path = NULL; + WavFormat wav_fmt = WAV_S16; + + for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "--model") == 0 && i + 1 < argc) { + model_path = argv[++i]; + } else if (strcmp(argv[i], "-i") == 0 && i + 1 < argc) { + input_path = argv[++i]; + } else if (strcmp(argv[i], "--format") == 0 && i + 1 < argc) { + if (!audio_parse_format(argv[++i], wav_fmt)) { + fprintf(stderr, "[CLI] ERROR: unknown format: %s\n", argv[i]); + print_usage(argv[0]); + return 1; + } + } else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { + print_usage(argv[0]); + return 0; + } else { + fprintf(stderr, "[CLI] ERROR: unknown arg: %s\n", argv[i]); + print_usage(argv[0]); + return 1; + } + } + + if (!model_path) { + print_usage(argv[0]); + return 1; + } + + int mode = 0; + if (input_path) { + mode = infer_mode(input_path); + if (mode == 0) { + fprintf(stderr, "[CLI] ERROR: %s: unsupported extension (expect .wav or .rvq)\n", input_path); + return 1; + } + } + + BackendPair bp = backend_init("Codec"); + if (!bp.backend) { + fprintf(stderr, "[Codec] FATAL: backend init failed\n"); + return 1; + } + + PipelineCodec pc = {}; + if (!pipeline_codec_load(&pc, model_path, bp)) { + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + + int rc = 0; + + if (!input_path) { + fprintf(stderr, "[Codec] Load self-test passed\n"); + } else if (mode == 1) { + // Encode .wav -> .rvq + const std::string out_str = swap_ext(input_path, ".rvq"); + + int T_in = 0; + float * audio_in = audio_read_mono(input_path, QWEN_TOKENIZER_SAMPLE_RATE, &T_in); + if (!audio_in || T_in <= 0) { + fprintf(stderr, "[Codec] FATAL: cannot read %s\n", input_path); + free(audio_in); + rc = 1; + } else { + // Pad to a multiple of HOP_LENGTH so the RVQ frame count is integral. + int hop = QWEN_TOKENIZER_HOP_LENGTH; + int T_padded = ((T_in + hop - 1) / hop) * hop; + int T_frames = T_padded / hop; + + std::vector audio_buf((size_t) T_padded, 0.0f); + memcpy(audio_buf.data(), audio_in, (size_t) T_in * sizeof(float)); + free(audio_in); + + fprintf(stderr, "[Codec] Encode: %s, %d samples @ %d Hz, padded to %d (%d frames @ 12.5 Hz, %.2f s)\n", + input_path, T_in, QWEN_TOKENIZER_SAMPLE_RATE, T_padded, T_frames, + (double) T_padded / (double) QWEN_TOKENIZER_SAMPLE_RATE); + + std::vector codes = pipeline_codec_encode(&pc, audio_buf.data(), T_padded); + if (codes.empty()) { + fprintf(stderr, "[Codec] FATAL: encode failed\n"); + rc = 1; + } else if (!write_rvq(out_str.c_str(), codes)) { + rc = 1; + } else { + fprintf(stderr, "[Codec] Wrote %s: K=%d T=%d, %zu codes -> %zu packed bytes\n", out_str.c_str(), + QWEN_TOKENIZER_NUM_CODEBOOKS, T_frames, codes.size(), + (codes.size() * (size_t) QWEN_TOKENIZER_CODE_BITS + 7) / 8); + } + } + } else { + // Decode .rvq -> .wav + const std::string out_str = swap_ext(input_path, ".wav"); + + std::vector codes; + int T = 0; + if (!read_rvq(input_path, QWEN_TOKENIZER_NUM_CODEBOOKS, codes, &T)) { + rc = 1; + } else { + fprintf(stderr, "[Codec] Decode: %s, K=%d T=%d (%.2f s)\n", input_path, QWEN_TOKENIZER_NUM_CODEBOOKS, T, + (double) (T * QWEN_TOKENIZER_HOP_LENGTH) / (double) QWEN_TOKENIZER_SAMPLE_RATE); + + std::vector audio = pipeline_codec_decode(&pc, codes.data(), QWEN_TOKENIZER_NUM_CODEBOOKS, T); + if (audio.empty()) { + fprintf(stderr, "[Codec] FATAL: decode failed\n"); + rc = 1; + } else if (!audio_write_wav(out_str.c_str(), audio.data(), (int) audio.size(), QWEN_TOKENIZER_SAMPLE_RATE, + wav_fmt)) { + fprintf(stderr, "[Codec] FATAL: cannot write %s\n", out_str.c_str()); + rc = 1; + } else { + fprintf(stderr, "[Codec] Wrote %s: %d samples @ %d Hz, %.2f s\n", out_str.c_str(), (int) audio.size(), + QWEN_TOKENIZER_SAMPLE_RATE, (double) audio.size() / (double) QWEN_TOKENIZER_SAMPLE_RATE); + } + } + } + + pipeline_codec_free(&pc); + backend_release(bp.backend, bp.cpu_backend); + return rc; +} diff --git a/tools/qwen-tts.cpp b/tools/qwen-tts.cpp new file mode 100644 index 0000000..a3396b5 --- /dev/null +++ b/tools/qwen-tts.cpp @@ -0,0 +1,340 @@ +// qwen-tts.cpp : thin CLI wrapper around the Qwen3-TTS synthesis +// pipeline. Parses arguments, loads the talker + codec GGUFs, hands +// off to pipeline_tts_synthesize and writes the resulting waveform as +// a WAV file. All heavy lifting lives in src/pipeline-tts.cpp. +// +// Talker variants : 0.6B-Base / 0.6B-CustomVoice / 1.7B-Base / +// 1.7B-CustomVoice / 1.7B-VoiceDesign. The decoder path is selected +// from GGUF metadata at load time. The CLI surface mirrors the +// omnivoice.cpp tooling : kebab-case flags, --format wav16/wav24/wav32, +// -o '-' streams to stdout, --seed -1 means non deterministic, the +// utterance text comes from --text or stdin if --text is absent. + +#include "audio-io.h" +#include "backend.h" +#include "bpe.h" +#include "pipeline-tts.h" +#include "version.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +static void print_usage(const char * prog) { + fprintf(stderr, "qwentts.cpp %s\n\n", QWEN_VERSION); + fprintf(stderr, + "Usage: %s --model --codec [options] -o \n\n" + "Required:\n" + " --model Talker LM GGUF (qwen-talker-*.gguf)\n" + " --codec Tokenizer GGUF (qwen-tokenizer-*.gguf)\n" + " -o Output WAV. '-' streams to stdout (pipe friendly).\n\n" + "Input:\n" + " --text Utterance text. If absent, stdin is read fully.\n\n" + "Synthesis options:\n" + " --lang Language (auto, english, chinese, ...) (default: english)\n" + " --instruct Style instruction. Required for VoiceDesign, optional\n" + " for CustomVoice. Rejected for Base.\n" + " --speaker Speaker name. Only valid for CustomVoice.\n" + " --ref-audio Reference WAV path for voice clone (Base only). Mutually\n" + " exclusive with --speaker. Mode A (x_vector_only) extracts\n" + " a speaker embedding via the ECAPA-TDNN encoder.\n" + " --ref-text Reference transcript for voice clone ICL mode (Base only,\n" + " requires --ref-audio). Switches the prompt to ICL mode B\n" + " where the talker conditions on the reference codec codes.\n" + " --max-new Max new audio frames (default: 2048)\n" + " --format WAV output format: wav16, wav24, wav32 (default: wav16)\n\n" + "Sampling options:\n" + " --seed Sampling seed, -1 for random (default: -1)\n" + " --greedy Disable stochastic sampling on both stacks\n" + " --temp Talker temperature (default: 0.9)\n" + " --top-k Talker top-k (default: 50, 0 = disabled)\n" + " --top-p Talker top-p (default: 1.0)\n" + " --rep-pen Talker repetition penalty (default: 1.05)\n" + " --sub-temp Sub-talker temperature (default: 0.9)\n" + " --sub-top-k Sub-talker top-k (default: 50)\n" + " --sub-top-p Sub-talker top-p (default: 1.0)\n\n" + "Debug:\n" + " --dump Dump intermediate tensors for cossim debug\n", + prog); +} + +struct Args { + const char * model; + const char * codec; + const char * text; + const char * lang; + const char * instruct; + const char * speaker; + const char * ref_audio; + const char * ref_text; + const char * dump_dir; + const char * out_wav; + const char * format; + int max_new_tokens; + int64_t seed; + bool do_sample; + float temperature; + int top_k; + float top_p; + float repetition_penalty; + int subtalker_top_k; + float subtalker_top_p; + float subtalker_temperature; + bool subtalker_do_sample; +}; + +// Read all of stdin into a string. Trims trailing newlines so a piped +// text file behaves like a clean --text argument. +static std::string read_stdin_text() { + std::ostringstream ss; + ss << std::cin.rdbuf(); + std::string s = ss.str(); + while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) { + s.pop_back(); + } + return s; +} + +static bool parse_args(int argc, char ** argv, Args & a) { + a = {}; + a.lang = "english"; + a.format = "wav16"; + a.max_new_tokens = 2048; + a.seed = -1; + a.do_sample = true; + a.temperature = 0.9f; + a.top_k = 50; + a.top_p = 1.0f; + a.repetition_penalty = 1.05f; + a.subtalker_do_sample = true; + a.subtalker_top_k = 50; + a.subtalker_top_p = 1.0f; + a.subtalker_temperature = 0.9f; + for (int i = 1; i < argc; i++) { + const char * arg = argv[i]; + if (std::strcmp(arg, "-h") == 0 || std::strcmp(arg, "--help") == 0) { + return false; + } + if (std::strcmp(arg, "--model") == 0 && i + 1 < argc) { + a.model = argv[++i]; + } else if (std::strcmp(arg, "--codec") == 0 && i + 1 < argc) { + a.codec = argv[++i]; + } else if (std::strcmp(arg, "--text") == 0 && i + 1 < argc) { + a.text = argv[++i]; + } else if (std::strcmp(arg, "--lang") == 0 && i + 1 < argc) { + a.lang = argv[++i]; + } else if (std::strcmp(arg, "--instruct") == 0 && i + 1 < argc) { + a.instruct = argv[++i]; + } else if (std::strcmp(arg, "--speaker") == 0 && i + 1 < argc) { + a.speaker = argv[++i]; + } else if (std::strcmp(arg, "--ref-audio") == 0 && i + 1 < argc) { + a.ref_audio = argv[++i]; + } else if (std::strcmp(arg, "--ref-text") == 0 && i + 1 < argc) { + a.ref_text = argv[++i]; + } else if (std::strcmp(arg, "--format") == 0 && i + 1 < argc) { + a.format = argv[++i]; + } else if (std::strcmp(arg, "--dump") == 0 && i + 1 < argc) { + a.dump_dir = argv[++i]; + } else if (std::strcmp(arg, "--max-new") == 0 && i + 1 < argc) { + a.max_new_tokens = std::atoi(argv[++i]); + } else if (std::strcmp(arg, "--seed") == 0 && i + 1 < argc) { + a.seed = (int64_t) std::atoll(argv[++i]); + } else if (std::strcmp(arg, "--greedy") == 0) { + // Greedy mode : argmax sampling on both stacks. The sampling + // fast path in sampling.h uses temperature <= 0 to short + // circuit to argmax, bypassing rep penalty and top-k/p + // truncation, which exactly mirrors the Python reference + // greedy behaviour used by tests/debug-tts-cossim.py. + a.do_sample = false; + a.subtalker_do_sample = false; + } else if (std::strcmp(arg, "--temp") == 0 && i + 1 < argc) { + a.temperature = (float) std::atof(argv[++i]); + } else if (std::strcmp(arg, "--top-k") == 0 && i + 1 < argc) { + a.top_k = std::atoi(argv[++i]); + } else if (std::strcmp(arg, "--top-p") == 0 && i + 1 < argc) { + a.top_p = (float) std::atof(argv[++i]); + } else if (std::strcmp(arg, "--rep-pen") == 0 && i + 1 < argc) { + a.repetition_penalty = (float) std::atof(argv[++i]); + } else if (std::strcmp(arg, "--sub-temp") == 0 && i + 1 < argc) { + a.subtalker_temperature = (float) std::atof(argv[++i]); + } else if (std::strcmp(arg, "--sub-top-k") == 0 && i + 1 < argc) { + a.subtalker_top_k = std::atoi(argv[++i]); + } else if (std::strcmp(arg, "--sub-top-p") == 0 && i + 1 < argc) { + a.subtalker_top_p = (float) std::atof(argv[++i]); + } else if (std::strcmp(arg, "-o") == 0 && i + 1 < argc) { + a.out_wav = argv[++i]; + } else { + fprintf(stderr, "[CLI] ERROR: unknown or incomplete argument: %s\n", arg); + return false; + } + } + return a.model && a.codec; +} + +static int run(const Args & a) { + BackendPair bp = backend_init("Talker"); + + PipelineTTS pt; + if (!pipeline_tts_load(&pt, a.model, a.codec, bp)) { + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + + // Validate mode flag combination against the loaded model_type. The + // upstream Python raises ValueError when generate_voice_design is + // called on a non voice_design model and the same shape applies to + // generate_custom_voice. We mirror that here, explicit and KISS, so + // the user never gets a silently wrong synthesis. + const std::string mt = pt.model_type; + if (a.speaker && mt != "custom_voice") { + fprintf(stderr, "[CLI] ERROR: --speaker is only valid for custom_voice models (loaded: %s)\n", mt.c_str()); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + if (a.instruct && mt == "base") { + fprintf(stderr, "[CLI] ERROR: --instruct is not supported for base models\n"); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + if (mt == "custom_voice" && !a.speaker) { + fprintf(stderr, "[CLI] ERROR: custom_voice models require --speaker\n"); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + if (mt == "voice_design" && (!a.instruct || a.instruct[0] == '\0')) { + fprintf(stderr, "[CLI] ERROR: voice_design models require --instruct\n"); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + if (a.ref_audio && mt != "base") { + fprintf(stderr, "[CLI] ERROR: --ref-audio is only valid for base models (loaded: %s)\n", mt.c_str()); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + if (a.speaker && a.ref_audio) { + fprintf(stderr, "[CLI] ERROR: --speaker and --ref-audio are mutually exclusive\n"); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + if (a.ref_text && !a.ref_audio) { + fprintf(stderr, "[CLI] ERROR: --ref-text requires --ref-audio\n"); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + + // Resolve output WAV format string : wav16 / wav24 / wav32. Default + // wav16 mirrors the omnivoice.cpp default. + WavFormat wav_fmt; + if (!audio_parse_format(a.format, wav_fmt)) { + fprintf(stderr, "[CLI] ERROR: invalid --format '%s' (expected wav16, wav24, wav32)\n", a.format); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + + // Resolve utterance text : explicit --text wins, otherwise read stdin + // fully. Empty stdin combined with no --text triggers a clean error. + std::string text_buf; + const char * text = a.text; + if (!text) { + text_buf = read_stdin_text(); + if (text_buf.empty()) { + fprintf(stderr, "[CLI] ERROR: no --text and stdin is empty\n"); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + text = text_buf.c_str(); + } + + // Resolve seed : -1 means non deterministic, sample from a hardware + // random_device. Anything else is taken verbatim, including negative + // values reaching int64 range, so reproducibility is one --seed away. + int64_t seed = a.seed; + if (seed < 0) { + std::random_device rd; + seed = (int64_t) (((uint64_t) rd() << 32) ^ (uint64_t) rd()); + } + + BPETokenizer tok = {}; + if (!load_bpe_from_gguf(&tok, a.model)) { + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + const char * specials_keys[] = { + "qwen3-tts.text.im_start_id", "qwen3-tts.text.im_end_id", "qwen3-tts.text.tts_pad_id", + "qwen3-tts.text.tts_bos_id", "qwen3-tts.text.tts_eos_id", + }; + bpe_load_specials_from_keys(&tok, a.model, specials_keys, 5); + + PipelineTTSSynthesizeParams p = {}; + p.text = text; + p.lang = a.lang; + p.instruct = a.instruct; + p.speaker = a.speaker; + p.ref_audio = a.ref_audio; + p.ref_text = a.ref_text; + p.seed = seed; + p.max_new_tokens = a.max_new_tokens; + p.do_sample = a.do_sample; + p.temperature = a.temperature; + p.top_k = a.top_k; + p.top_p = a.top_p; + p.repetition_penalty = a.repetition_penalty; + p.subtalker_do_sample = a.subtalker_do_sample; + p.subtalker_temperature = a.subtalker_temperature; + p.subtalker_top_k = a.subtalker_top_k; + p.subtalker_top_p = a.subtalker_top_p; + p.dump_dir = a.dump_dir; + + PipelineTTSSynthesizeOutput out; + if (!pipeline_tts_synthesize(&pt, &tok, p, &out)) { + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + + if (!out.audio.empty()) { + const char * out_path = a.out_wav ? a.out_wav : "out.wav"; + if (!audio_write_wav(out_path, out.audio.data(), (int) out.audio.size(), out.sample_rate, wav_fmt)) { + fprintf(stderr, "[Pipeline] FATAL: WAV write failed for %s\n", out_path); + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 1; + } + qt_log(QT_LOG_INFO, "[Pipeline] Wrote %zu samples (%.2f s) -> %s", out.audio.size(), + (double) out.audio.size() / (double) out.sample_rate, out_path); + } + + pipeline_tts_free(&pt); + backend_release(bp.backend, bp.cpu_backend); + return 0; +} + +int main(int argc, char ** argv) { + Args a; + if (!parse_args(argc, argv, a)) { + print_usage(argv[0]); + return 1; + } + try { + return run(a); + } catch (const std::runtime_error & e) { + qt_set_error("%s", e.what()); + qt_log(QT_LOG_ERROR, "%s", e.what()); + return 1; + } +} diff --git a/tools/version.cmake b/tools/version.cmake new file mode 100644 index 0000000..7ec6b5c --- /dev/null +++ b/tools/version.cmake @@ -0,0 +1,38 @@ +# Generate version.h with the current git commit hash and date. +# Only rewrites the file if the content changed (avoids rebuild cascade). +# Usage: cmake -DSRC_DIR=... -DOUTPUT=... -P version.cmake + +execute_process( + COMMAND git rev-parse --short HEAD + WORKING_DIRECTORY "${SRC_DIR}" + OUTPUT_VARIABLE GIT_HASH + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + RESULT_VARIABLE GIT_RESULT +) +if(NOT GIT_RESULT EQUAL 0) + set(GIT_HASH "unknown") +endif() + +execute_process( + COMMAND git show -s --format=%cs HEAD + WORKING_DIRECTORY "${SRC_DIR}" + OUTPUT_VARIABLE GIT_DATE + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + RESULT_VARIABLE DATE_RESULT +) +if(NOT DATE_RESULT EQUAL 0) + set(GIT_DATE "unknown") +endif() + +set(CONTENT "#pragma once\n#define QWEN_VERSION \"${GIT_HASH} (${GIT_DATE})\"\n") + +if(EXISTS "${OUTPUT}") + file(READ "${OUTPUT}" EXISTING) + if("${EXISTING}" STREQUAL "${CONTENT}") + return() + endif() +endif() + +file(WRITE "${OUTPUT}" "${CONTENT}") diff --git a/update.cmd b/update.cmd new file mode 100644 index 0000000..9941135 --- /dev/null +++ b/update.cmd @@ -0,0 +1,6 @@ +@echo off + +cd ggml +git pull --rebase +cd .. +git pull --rebase diff --git a/update.sh b/update.sh new file mode 100755 index 0000000..ae29614 --- /dev/null +++ b/update.sh @@ -0,0 +1,6 @@ +#!/bin/bash + +cd ggml +git pull --rebase +cd .. +git pull --rebase