87 lines
3.9 KiB
Plaintext
87 lines
3.9 KiB
Plaintext
#pragma once
|
||
#include <ggml.h>
|
||
#include <gguf.h>
|
||
#include <ggml-backend.h>
|
||
#include <string>
|
||
#include <vector>
|
||
#include <unordered_map>
|
||
|
||
namespace kazeia::tts {
|
||
|
||
struct DecoderConfig {
|
||
int latent_dim; // 1024
|
||
int codebook_dim; // 512
|
||
int codebook_size; // 2048
|
||
int decoder_dim; // 1536
|
||
int hidden_size; // 512
|
||
int intermediate_size; // 1024
|
||
int num_attention_heads; // 16
|
||
int head_dim; // 64
|
||
int num_hidden_layers; // 8
|
||
int num_quantizers; // 16
|
||
std::vector<int> upsample_rates; // {8, 5, 4, 3}
|
||
std::vector<int> upsampling_ratios; // {2, 2}
|
||
float layer_scale_initial; // 0.01
|
||
float rope_theta; // 10000.0
|
||
float rms_norm_eps; // 1e-5
|
||
int total_upsample; // = prod(upsample_rates) * prod(upsampling_ratios) = 1920
|
||
};
|
||
|
||
struct Decoder {
|
||
DecoderConfig cfg;
|
||
|
||
// Tensor lookup table — populated from GGUF on load.
|
||
std::unordered_map<std::string, struct ggml_tensor*> tensors;
|
||
|
||
// ggml backend + context (CPU baseline ; Vulkan/Hexagon en Phase 3-4).
|
||
// Loader créé un ctx_w avec no_alloc=true puis place les data sur le
|
||
// backend buffer `buf_w` (CPU par défaut). Cette indirection est
|
||
// requise pour que le scheduler ggml_backend_sched puisse transférer
|
||
// les poids vers Hexagon/Vulkan à la demande.
|
||
struct ggml_context * ctx_w = nullptr;
|
||
struct ggml_context * ctx_g = nullptr;
|
||
struct gguf_context * gguf = nullptr;
|
||
struct ggml_backend_buffer * buf_w = nullptr;
|
||
|
||
// Backend list + sched. backends[0] est prioritaire (HTP si dispo, sinon CPU).
|
||
// backends[end] doit TOUJOURS être CPU (fallback pour ops non-supportées).
|
||
// load_with_backends() peuple ça, et compute_graph() exécute via le sched.
|
||
std::vector<ggml_backend_t> backends;
|
||
struct ggml_backend_sched * sched = nullptr;
|
||
bool owns_backends = false; // si true, unload() libère aussi backends[]
|
||
|
||
// Snake constants pré-calculées au load_with_backends pour rendre
|
||
// stage_bigvgan sched-friendly. Sans ça, precompute_snake host-side ne
|
||
// marche pas avec ctx no_alloc=true. Map clé "{prefix}.alpha" -> (a_eff, inv_b)
|
||
// avec a_eff = exp(α), inv_b = 1 / (exp(β) + eps).
|
||
// Alloués dans ctx_const + buf_const sur le même backend que buf_w.
|
||
struct ggml_context * ctx_const = nullptr;
|
||
struct ggml_backend_buffer * buf_const = nullptr;
|
||
std::unordered_map<std::string, std::pair<struct ggml_tensor*, struct ggml_tensor*>> snake_consts;
|
||
|
||
// load() utilise le backend CPU par défaut. Pour Hexagon/Vulkan, passer
|
||
// un buft explicite via load_with_buft(), OU une liste de backends via
|
||
// load_with_backends() (recommandé pour HTP : crée le sched + buft auto).
|
||
bool load(const std::string & gguf_path);
|
||
bool load_with_buft(const std::string & gguf_path, ggml_backend_buffer_type_t buft);
|
||
bool load_with_backends(const std::string & gguf_path,
|
||
const std::vector<ggml_backend_t> & devs);
|
||
void unload();
|
||
|
||
// Helper : compute un graphe via le sched (split CPU/HTP auto).
|
||
// Si sched=NULL (chargement CPU pur via load_with_buft), tombe sur
|
||
// ggml_graph_compute_with_ctx pour rester compatible.
|
||
void compute_graph(struct ggml_context * ctx, struct ggml_cgraph * gf, int n_threads);
|
||
|
||
// Forward : codes [B, num_quantizers=16, T] long → wav [B, 1, T*total_upsample] f32
|
||
// Returns owned buffer, length = T * total_upsample.
|
||
std::vector<float> forward(const std::vector<int32_t>& codes, int T);
|
||
|
||
// ---- Tests par étape ----
|
||
// Quantizer seul : codes [16, T] → hidden [512, T]. Permet validation
|
||
// tensor-par-tensor vs scripts/dump_reference.py.
|
||
std::vector<float> test_quantizer(const std::vector<int32_t>& codes_flat, int T);
|
||
};
|
||
|
||
} // namespace kazeia::tts
|