104 lines
3.9 KiB
C++
104 lines
3.9 KiB
C++
#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 <cstdint>
|
|
#include <vector>
|
|
|
|
#define TOKENIZER_HOP_LENGTH 1920
|
|
#define TOKENIZER_SAMPLE_RATE 24000
|
|
#define TOKENIZER_NUM_CODEBOOKS 16
|
|
#define 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. Owns a dedicated
|
|
// WeightCtx because it is the only module without one.
|
|
struct ggml_tensor * pre_conv_w; // [3, 512, 1024] f32
|
|
struct ggml_tensor * pre_conv_b; // [1024] f32
|
|
WeightCtx pre_conv_wctx;
|
|
|
|
// 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 * TOKENIZER_HOP_LENGTH, empty on failure.
|
|
std::vector<float> 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
|
|
// TOKENIZER_HOP_LENGTH (1920); the caller is expected
|
|
// to pad with zeros if needed.
|
|
// dump_dir: optional path. When non NULL, dumps the SEANet, encoder
|
|
// transformer and post-downsample (pre-FSQ latents) buffers
|
|
// into seanet-out.bin, enc-transformer-out.bin and
|
|
// codec-pre-fsq.bin under that directory. Quiet otherwise.
|
|
// Returns codes flat as [K, T] row-major, K = TOKENIZER_NUM_CODEBOOKS,
|
|
// T = n_samples / 1920. Empty on failure.
|
|
std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc,
|
|
const float * audio,
|
|
int n_samples,
|
|
const char * dump_dir = NULL);
|
|
|
|
// Free every backend buffer and ggml context. Safe to call on a zeroed struct.
|
|
void pipeline_codec_free(PipelineCodec * pc);
|