mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-07-10 01:18:32 +00:00
initial files for ace step
This commit is contained in:
Concedo
parent
ac70ca35dd
commit
5cd6e50eab
17 changed files with 7346 additions and 1 deletions
6
Makefile
6
Makefile
|
|
@ -712,7 +712,7 @@ gpttype_adapter_vulkan_noavx2.o: $(GPTTYPE_ADAPTER)
|
|||
$(CXX) $(CXXFLAGS) $(FAILSAFE_FLAGS) $(VULKAN_FLAGS) -c $< -o $@
|
||||
|
||||
clean:
|
||||
rm -vf *.o main ttsmain sdmain whispermain quantize_gguf quantize_clip quantize_gpt2 quantize_gptj quantize_neox quantize_mpt vulkan-shaders-gen vulkan-shaders-gen-noext gguf-split mtmd-cli mainvk fitparams embedding embeddingvk embeddingvk.exe embedding.exe fitparams.exe mainvk.exe mtmd-cli.exe gguf-split.exe vulkan-shaders-gen.exe vulkan-shaders-gen-noext.exe main.exe ttsmain.exe sdmain.exe whispermain.exe quantize_clip.exe quantize_gguf.exe quantize_gptj.exe quantize_gpt2.exe quantize_neox.exe quantize_mpt.exe koboldcpp_default.dll koboldcpp_failsafe.dll koboldcpp_noavx2.dll koboldcpp_vulkan_failsafe.dll koboldcpp_cublas.dll koboldcpp_hipblas.dll koboldcpp_vulkan.dll koboldcpp_vulkan_noavx2.dll koboldcpp_default.so koboldcpp_failsafe.so koboldcpp_noavx2.so koboldcpp_vulkan_failsafe.so koboldcpp_cublas.so koboldcpp_hipblas.so koboldcpp_vulkan.so koboldcpp_vulkan_noavx2.so ggml/src/ggml-vulkan-shaders.cpp ggml/src/ggml-vulkan-shaders.hpp ggml/src/ggml-vulkan-shaders-noext.cpp ggml/src/ggml-vulkan-shaders-noext.hpp
|
||||
rm -vf *.o main ttsmain sdmain whispermain quantize_gguf quantize_clip quantize_gpt2 quantize_gptj quantize_neox quantize_mpt vulkan-shaders-gen vulkan-shaders-gen-noext gguf-split mtmd-cli mainvk fitparams embedding embeddingvk qwen3tts acestep-a acestep-b acestep-b.exe acestep-a.exe qwen3tts.exe embeddingvk.exe embedding.exe fitparams.exe mainvk.exe mtmd-cli.exe gguf-split.exe vulkan-shaders-gen.exe vulkan-shaders-gen-noext.exe main.exe ttsmain.exe sdmain.exe whispermain.exe quantize_clip.exe quantize_gguf.exe quantize_gptj.exe quantize_gpt2.exe quantize_neox.exe quantize_mpt.exe koboldcpp_default.dll koboldcpp_failsafe.dll koboldcpp_noavx2.dll koboldcpp_vulkan_failsafe.dll koboldcpp_cublas.dll koboldcpp_hipblas.dll koboldcpp_vulkan.dll koboldcpp_vulkan_noavx2.dll koboldcpp_default.so koboldcpp_failsafe.so koboldcpp_noavx2.so koboldcpp_vulkan_failsafe.so koboldcpp_cublas.so koboldcpp_hipblas.so koboldcpp_vulkan.so koboldcpp_vulkan_noavx2.so ggml/src/ggml-vulkan-shaders.cpp ggml/src/ggml-vulkan-shaders.hpp ggml/src/ggml-vulkan-shaders-noext.cpp ggml/src/ggml-vulkan-shaders-noext.hpp
|
||||
rm -vrf ggml/src/ggml-cuda/*.o
|
||||
rm -vrf ggml/src/ggml-cuda/template-instances/*.o
|
||||
rm -vrf llguidance
|
||||
|
|
@ -742,6 +742,10 @@ ttscppmain: otherarch/ttscpp/cli/cli.cpp otherarch/ttscpp/cli/playback.cpp other
|
|||
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
qwen3tts: otherarch/qwen3tts/q3ttsmain.cpp otherarch/qwen3tts/qwen3_tts.cpp otherarch/qwen3tts/text_tokenizer.cpp otherarch/qwen3tts/gguf_loader.cpp otherarch/qwen3tts/tts_transformer.cpp otherarch/qwen3tts/audio_tokenizer_decoder.cpp otherarch/qwen3tts/audio_tokenizer_encoder.cpp otherarch/qwen3tts/coreml_code_predictor_stub.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend_default.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
acestep-a: otherarch/acestep/ace-qwen3.cpp otherarch/acestep/request.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend_default.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
acestep-b: otherarch/acestep/dit-vae.cpp otherarch/acestep/request.cpp ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend_default.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
|
||||
ggml/src/ggml-vulkan-shaders.cpp:
|
||||
ifdef VULKAN_BUILD
|
||||
|
|
|
|||
1339
otherarch/acestep/ace-qwen3.cpp
Normal file
1339
otherarch/acestep/ace-qwen3.cpp
Normal file
File diff suppressed because it is too large
Load diff
49
otherarch/acestep/backend.h
Normal file
49
otherarch/acestep/backend.h
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
#pragma once
|
||||
// backend.h: shared GGML backend initialization
|
||||
//
|
||||
// All modules use the same pattern: load all backends, pick best GPU,
|
||||
// keep CPU as fallback. This avoids duplicating init logic across
|
||||
// qwen3.h, qwen3-lm.h, cond.h, dit.h, vae.h.
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml-cpu.h"
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <thread>
|
||||
|
||||
struct BackendPair {
|
||||
ggml_backend_t backend;
|
||||
ggml_backend_t cpu_backend;
|
||||
};
|
||||
|
||||
// Initialize backends: load all available (CUDA, Metal, Vulkan...),
|
||||
// pick the best one, keep CPU as fallback.
|
||||
// label: log prefix, e.g. "DiT", "VAE", "LM"
|
||||
static BackendPair backend_init(const char * label) {
|
||||
ggml_backend_load_all();
|
||||
BackendPair bp = {};
|
||||
bp.backend = ggml_backend_init_best();
|
||||
int n_threads = (int)std::thread::hardware_concurrency() / 2;
|
||||
if (n_threads < 1) n_threads = 1;
|
||||
// [GGML] If best backend is already CPU, reuse it (avoid 2 CPU instances
|
||||
// where only one gets the thread count)
|
||||
bool best_is_cpu = (strcmp(ggml_backend_name(bp.backend), "CPU") == 0);
|
||||
if (best_is_cpu) {
|
||||
bp.cpu_backend = bp.backend;
|
||||
ggml_backend_cpu_set_n_threads(bp.backend, n_threads);
|
||||
} else {
|
||||
bp.cpu_backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, NULL);
|
||||
ggml_backend_cpu_set_n_threads(bp.cpu_backend, n_threads);
|
||||
}
|
||||
fprintf(stderr, "[Load] %s backend: %s (CPU threads: %d)\n",
|
||||
label, ggml_backend_name(bp.backend), n_threads);
|
||||
return bp;
|
||||
}
|
||||
|
||||
// Create a scheduler from a backend pair.
|
||||
// max_nodes: graph size hint (4096 for small models, 8192 for large)
|
||||
static ggml_backend_sched_t backend_sched_new(BackendPair bp, int max_nodes) {
|
||||
ggml_backend_t backends[2] = { bp.backend, bp.cpu_backend };
|
||||
int n = (bp.backend == bp.cpu_backend) ? 1 : 2;
|
||||
return ggml_backend_sched_new(backends, NULL, n, max_nodes, false, true);
|
||||
}
|
||||
553
otherarch/acestep/bpe.h
Normal file
553
otherarch/acestep/bpe.h
Normal file
|
|
@ -0,0 +1,553 @@
|
|||
// bpe.h, Qwen3/GPT-2 byte-level BPE tokenizer (CPU-only, no dependencies)
|
||||
// Loads vocab.json + merges.txt from HuggingFace tokenizer directory.
|
||||
// Handles: byte-level encoding, GPT-2 regex pre-tokenizer, BPE merges,
|
||||
// special token recognition (<|endoftext|>).
|
||||
|
||||
#pragma once
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <unordered_map>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <cassert>
|
||||
#include <climits>
|
||||
|
||||
#include "gguf.h"
|
||||
|
||||
// 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); }
|
||||
*advance = 1; return c; // fallback
|
||||
}
|
||||
|
||||
// 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<std::string> gpt2_pre_tokenize(const std::string &text) {
|
||||
std::vector<std::string> 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 += 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<std::string, int> vocab; // token_str -> id
|
||||
std::unordered_map<std::string, int> merges; // "a b" -> rank
|
||||
std::string byte2str[256]; // byte -> GPT-2 UTF-8 string
|
||||
int eos_id; // <|endoftext|> = 151643
|
||||
int n_vocab;
|
||||
std::vector<std::string> id_to_str; // id -> token_str (reverse vocab)
|
||||
};
|
||||
|
||||
// Minimal JSON parser for vocab.json ({"str": int, ...})
|
||||
static bool load_vocab_json(const std::string &path, std::unordered_map<std::string, int> &vocab) {
|
||||
FILE *f = fopen(path.c_str(), "rb");
|
||||
if (!f) return false;
|
||||
fseek(f, 0, SEEK_END);
|
||||
size_t sz = ftell(f);
|
||||
fseek(f, 0, SEEK_SET);
|
||||
std::string data(sz, '\0');
|
||||
fread(&data[0], 1, sz, f);
|
||||
fclose(f);
|
||||
|
||||
size_t i = 0;
|
||||
auto skip_ws = [&]() { while (i < sz && (data[i] == ' ' || data[i] == '\n' || data[i] == '\r' || data[i] == '\t')) i++; };
|
||||
|
||||
skip_ws();
|
||||
if (data[i] != '{') return false;
|
||||
i++;
|
||||
|
||||
while (i < sz) {
|
||||
skip_ws();
|
||||
if (data[i] == '}') break;
|
||||
if (data[i] == ',') { i++; skip_ws(); }
|
||||
if (data[i] != '"') return false;
|
||||
i++;
|
||||
|
||||
// Parse key (JSON string with escape handling)
|
||||
std::string key;
|
||||
while (i < sz && data[i] != '"') {
|
||||
if (data[i] == '\\') {
|
||||
i++;
|
||||
if (i >= sz) return false;
|
||||
switch (data[i]) {
|
||||
case '"': key += '"'; break;
|
||||
case '\\': key += '\\'; break;
|
||||
case '/': key += '/'; break;
|
||||
case 'n': key += '\n'; break;
|
||||
case 'r': key += '\r'; break;
|
||||
case 't': key += '\t'; break;
|
||||
case 'u': {
|
||||
// \uXXXX -> UTF-8
|
||||
if (i + 4 >= sz) return false;
|
||||
char hex[5] = {data[i+1], data[i+2], data[i+3], data[i+4], 0};
|
||||
int cp = (int)strtol(hex, nullptr, 16);
|
||||
i += 4;
|
||||
if (cp < 0x80) { key += (char)cp; }
|
||||
else if (cp < 0x800) { key += (char)(0xC0 | (cp >> 6)); key += (char)(0x80 | (cp & 0x3F)); }
|
||||
else { key += (char)(0xE0 | (cp >> 12)); key += (char)(0x80 | ((cp >> 6) & 0x3F)); key += (char)(0x80 | (cp & 0x3F)); }
|
||||
break;
|
||||
}
|
||||
default: key += data[i]; break;
|
||||
}
|
||||
} else {
|
||||
key += data[i];
|
||||
}
|
||||
i++;
|
||||
}
|
||||
i++; // skip closing "
|
||||
|
||||
skip_ws();
|
||||
if (data[i] != ':') return false;
|
||||
i++;
|
||||
skip_ws();
|
||||
|
||||
// Parse integer value
|
||||
bool neg = false;
|
||||
if (data[i] == '-') { neg = true; i++; }
|
||||
int val = 0;
|
||||
while (i < sz && data[i] >= '0' && data[i] <= '9') {
|
||||
val = val * 10 + (data[i] - '0');
|
||||
i++;
|
||||
}
|
||||
if (neg) val = -val;
|
||||
|
||||
vocab[key] = val;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Load merges from merges.txt
|
||||
static bool load_merges(const std::string &path, std::unordered_map<std::string, int> &merges) {
|
||||
FILE *f = fopen(path.c_str(), "r");
|
||||
if (!f) return false;
|
||||
char line[512];
|
||||
int rank = 0;
|
||||
// Skip header line (#version: 0.2)
|
||||
if (fgets(line, sizeof(line), f) == nullptr) { fclose(f); return false; }
|
||||
while (fgets(line, sizeof(line), f)) {
|
||||
int len = (int)strlen(line);
|
||||
while (len > 0 && (line[len-1] == '\n' || line[len-1] == '\r')) line[--len] = '\0';
|
||||
if (len == 0) continue;
|
||||
merges[std::string(line)] = rank++;
|
||||
}
|
||||
fclose(f);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Load tokenizer from directory (vocab.json + merges.txt)
|
||||
static bool load_bpe_tokenizer(BPETokenizer *tok, const char *dir) {
|
||||
build_byte_encoder(tok->byte2str);
|
||||
|
||||
std::string vocab_path = std::string(dir) + "/vocab.json";
|
||||
std::string merges_path = std::string(dir) + "/merges.txt";
|
||||
|
||||
if (!load_vocab_json(vocab_path, tok->vocab)) {
|
||||
fprintf(stderr, "ERROR: failed to load %s\n", vocab_path.c_str());
|
||||
return false;
|
||||
}
|
||||
if (!load_merges(merges_path, tok->merges)) {
|
||||
fprintf(stderr, "ERROR: failed to load %s\n", merges_path.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
tok->n_vocab = (int)tok->vocab.size();
|
||||
tok->eos_id = 151643; // <|endoftext|>
|
||||
|
||||
// Build reverse vocab (id -> token string)
|
||||
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;
|
||||
}
|
||||
|
||||
fprintf(stderr, "[BPE] Loaded: %d vocab, %d merges\n", tok->n_vocab, (int)tok->merges.size());
|
||||
return true;
|
||||
}
|
||||
|
||||
// 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->eos_id = 151643;
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
fprintf(stderr, "[BPE] Loaded from GGUF: %d vocab, %d merges\n", tok->n_vocab, n_merges);
|
||||
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<std::string> bpe_merge(const std::unordered_map<std::string, int> &merge_rank,
|
||||
const std::vector<std::string> &symbols) {
|
||||
if (symbols.size() <= 1) return symbols;
|
||||
|
||||
std::vector<std::string> 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<int> &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<std::string> 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<std::string> 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
|
||||
// Handles special tokens (<|endoftext|>) embedded in text.
|
||||
// Appends EOS at end (post-processor behavior).
|
||||
static std::vector<int> bpe_encode(const BPETokenizer *tok, const std::string &text,
|
||||
bool add_eos = true) {
|
||||
std::vector<int> ids;
|
||||
const std::string special = "<|endoftext|>";
|
||||
|
||||
// Split on special tokens
|
||||
size_t pos = 0;
|
||||
while (pos < text.size()) {
|
||||
size_t found = text.find(special, pos);
|
||||
std::string segment = (found == std::string::npos)
|
||||
? text.substr(pos) : text.substr(pos, found - pos);
|
||||
|
||||
// Pre-tokenize and encode the segment
|
||||
if (!segment.empty()) {
|
||||
auto chunks = gpt2_pre_tokenize(segment);
|
||||
for (const auto &chunk : chunks) {
|
||||
encode_chunk(tok, chunk, ids);
|
||||
}
|
||||
}
|
||||
|
||||
if (found == std::string::npos) break;
|
||||
ids.push_back(tok->eos_id); // <|endoftext|>
|
||||
pos = found + special.size();
|
||||
}
|
||||
|
||||
if (add_eos) ids.push_back(tok->eos_id);
|
||||
return ids;
|
||||
}
|
||||
293
otherarch/acestep/cond.h
Normal file
293
otherarch/acestep/cond.h
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
// cond.h: ACEStep Condition Encoder via ggml
|
||||
//
|
||||
// Produces encoder_hidden_states [S_total, 2048] from (all arrays H-contiguous per token):
|
||||
// - text_hidden [S_text, 1024] : from Qwen3-Embedding text encoder
|
||||
// - lyric_embed [S_lyric, 1024] : from CPU vocab lookup of lyric tokens
|
||||
// - timbre_feats [S_ref, 64] : reference audio features (optional)
|
||||
//
|
||||
// Internal pipeline (ggml notation [ne0, ne1]):
|
||||
// text_hidden -> Linear(1024->2048) -> [2048, S_text]
|
||||
// lyric_embed -> Linear(1024->2048)+bias -> 8L -> norm -> [2048, S_lyric]
|
||||
// timbre_feats -> Linear(64->2048)+bias -> 4L -> norm -> take frame[0] -> [2048, 1]
|
||||
// Pack: cat(lyric, timbre[0:1], text_proj) -> [2048, S_total]
|
||||
|
||||
#pragma once
|
||||
#include "qwen3.h"
|
||||
|
||||
// Lyric/Timbre encoder configs
|
||||
static Qwen3Config qwen3_lyric_config() {
|
||||
return {
|
||||
/*hidden_size*/ 2048,
|
||||
/*intermediate_size*/ 6144,
|
||||
/*n_heads*/ 16,
|
||||
/*n_kv_heads*/ 8,
|
||||
/*head_dim*/ 128,
|
||||
/*n_layers*/ 8,
|
||||
/*rope_theta*/ 1000000.0f,
|
||||
/*rms_norm_eps*/ 1e-6f,
|
||||
/*is_causal*/ false,
|
||||
};
|
||||
}
|
||||
|
||||
static Qwen3Config qwen3_timbre_config() {
|
||||
return {
|
||||
/*hidden_size*/ 2048,
|
||||
/*intermediate_size*/ 6144,
|
||||
/*n_heads*/ 16,
|
||||
/*n_kv_heads*/ 8,
|
||||
/*head_dim*/ 128,
|
||||
/*n_layers*/ 4,
|
||||
/*rope_theta*/ 1000000.0f,
|
||||
/*rms_norm_eps*/ 1e-6f,
|
||||
/*is_causal*/ false,
|
||||
};
|
||||
}
|
||||
|
||||
// Struct
|
||||
struct CondGGML {
|
||||
// Lyric encoder (8L, H=2048)
|
||||
Qwen3Config lyric_cfg;
|
||||
Qwen3Layer lyric_layers[8];
|
||||
struct ggml_tensor * lyric_embed_w; // [2048, 1024] ggml = Linear(1024->2048)
|
||||
struct ggml_tensor * lyric_embed_b; // [2048]
|
||||
struct ggml_tensor * lyric_norm; // [2048]
|
||||
|
||||
// Timbre encoder (4L, H=2048)
|
||||
Qwen3Config timbre_cfg;
|
||||
Qwen3Layer timbre_layers[4];
|
||||
struct ggml_tensor * timbre_embed_w; // [2048, 64] ggml = Linear(64->2048)
|
||||
struct ggml_tensor * timbre_embed_b; // [2048]
|
||||
struct ggml_tensor * timbre_norm; // [2048]
|
||||
|
||||
// Text projector: Linear(1024->2048) no bias
|
||||
struct ggml_tensor * text_proj_w; // [2048, 1024] ggml
|
||||
|
||||
// Null condition embedding (for classifier-free guidance)
|
||||
struct ggml_tensor * null_cond_emb; // [2048, 1, 1]
|
||||
|
||||
// Backend
|
||||
ggml_backend_t backend;
|
||||
ggml_backend_t cpu_backend;
|
||||
ggml_backend_sched_t sched;
|
||||
WeightCtx wctx;
|
||||
};
|
||||
|
||||
// Init
|
||||
static void cond_ggml_init_backend(CondGGML * m) {
|
||||
BackendPair bp = backend_init("CondEncoder");
|
||||
m->backend = bp.backend;
|
||||
m->cpu_backend = bp.cpu_backend;
|
||||
m->sched = backend_sched_new(bp, 8192);
|
||||
}
|
||||
|
||||
// Load from ACEStep DiT GGUF
|
||||
// gguf_path: path to the DiT .gguf file
|
||||
// Tensors have prefix "encoder." for lyric/timbre, and "null_condition_emb"
|
||||
static bool cond_ggml_load(CondGGML * m, const char * gguf_path) {
|
||||
m->lyric_cfg = qwen3_lyric_config();
|
||||
m->timbre_cfg = qwen3_timbre_config();
|
||||
|
||||
GGUFModel gf;
|
||||
if (!gf_load(&gf, gguf_path)) {
|
||||
fprintf(stderr, "[Load] FATAL: cannot load %s\n", gguf_path);
|
||||
return false;
|
||||
}
|
||||
|
||||
// Count tensors:
|
||||
// lyric: embed_w(1) + embed_b(1) + 8 layers x 11(88) + norm(1) = 91
|
||||
// timbre: embed_w(1) + embed_b(1) + 4 layers x 11(44) + norm(1) = 47
|
||||
// text_proj(1) + null_cond(1) = 2
|
||||
// Total: 140
|
||||
int n_tensors = 91 + 47 + 2;
|
||||
wctx_init(&m->wctx, n_tensors);
|
||||
|
||||
// Lyric encoder
|
||||
m->lyric_embed_w = gf_load_tensor(&m->wctx, gf, "encoder.lyric_encoder.embed_tokens.weight");
|
||||
m->lyric_embed_b = gf_load_tensor_f32(&m->wctx, gf, "encoder.lyric_encoder.embed_tokens.bias");
|
||||
m->lyric_norm = gf_load_tensor_f32(&m->wctx, gf, "encoder.lyric_encoder.norm.weight");
|
||||
for (int i = 0; i < m->lyric_cfg.n_layers; i++) {
|
||||
char prefix[128];
|
||||
snprintf(prefix, sizeof(prefix), "encoder.lyric_encoder.layers.%d", i);
|
||||
qwen3_load_layer(&m->wctx, gf, &m->lyric_layers[i], prefix);
|
||||
}
|
||||
|
||||
// Timbre encoder
|
||||
m->timbre_embed_w = gf_load_tensor(&m->wctx, gf, "encoder.timbre_encoder.embed_tokens.weight");
|
||||
m->timbre_embed_b = gf_load_tensor_f32(&m->wctx, gf, "encoder.timbre_encoder.embed_tokens.bias");
|
||||
m->timbre_norm = gf_load_tensor_f32(&m->wctx, gf, "encoder.timbre_encoder.norm.weight");
|
||||
for (int i = 0; i < m->timbre_cfg.n_layers; i++) {
|
||||
char prefix[128];
|
||||
snprintf(prefix, sizeof(prefix), "encoder.timbre_encoder.layers.%d", i);
|
||||
qwen3_load_layer(&m->wctx, gf, &m->timbre_layers[i], prefix);
|
||||
}
|
||||
|
||||
// Text projector + null condition
|
||||
m->text_proj_w = gf_load_tensor(&m->wctx, gf, "encoder.text_projector.weight");
|
||||
m->null_cond_emb = gf_load_tensor(&m->wctx, gf, "null_condition_emb");
|
||||
|
||||
if (!wctx_alloc(&m->wctx, m->backend)) {
|
||||
gf_close(&gf);
|
||||
return false;
|
||||
}
|
||||
gf_close(&gf);
|
||||
|
||||
fprintf(stderr, "[Load] CondEncoder: lyric(%dL), timbre(%dL), text_proj, null_cond\n",
|
||||
m->lyric_cfg.n_layers, m->timbre_cfg.n_layers);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Forward
|
||||
//
|
||||
// Inputs (CPU float arrays):
|
||||
// text_hidden: [1024 * S_text] from text encoder (Qwen3-Embedding)
|
||||
// lyric_embed: [1024 * S_lyric] from CPU vocab lookup of lyric tokens
|
||||
// timbre_feats: [64 * S_ref] reference audio features (NULL if none)
|
||||
// S_text, S_lyric, S_ref sequence lengths
|
||||
//
|
||||
// Output:
|
||||
// enc_hidden: resized to [2048 * S_total] float
|
||||
// *out_enc_S: total sequence length
|
||||
//
|
||||
// Layout: all arrays in ggml order (ne[0]=dim contiguous, then sequence)
|
||||
static void cond_ggml_forward(CondGGML * m,
|
||||
const float * text_hidden, int S_text,
|
||||
const float * lyric_embed, int S_lyric,
|
||||
const float * timbre_feats, int S_ref,
|
||||
std::vector<float> & enc_hidden,
|
||||
int * out_enc_S) {
|
||||
int H = 2048;
|
||||
bool has_timbre = (timbre_feats != NULL && S_ref > 0);
|
||||
|
||||
// Build graph
|
||||
size_t ctx_size = 4096 * ggml_tensor_overhead() + ggml_graph_overhead();
|
||||
struct ggml_init_params gp = { ctx_size, NULL, true };
|
||||
struct ggml_context * ctx = ggml_init(gp);
|
||||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
|
||||
|
||||
// Positions for lyric (bidirectional, 0..S_lyric-1)
|
||||
struct ggml_tensor * lyric_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, S_lyric);
|
||||
ggml_set_name(lyric_pos, "lyric_pos");
|
||||
ggml_set_input(lyric_pos);
|
||||
|
||||
// Path 1: Lyric encoder
|
||||
struct ggml_tensor * t_lyric_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1024, S_lyric);
|
||||
ggml_set_name(t_lyric_in, "lyric_in");
|
||||
ggml_set_input(t_lyric_in);
|
||||
|
||||
// Linear embed: [1024, S_lyric] -> [2048, S_lyric]
|
||||
struct ggml_tensor * lyric_h = qwen3_linear_bias(ctx, m->lyric_embed_w,
|
||||
m->lyric_embed_b, t_lyric_in);
|
||||
// 8 layers + final norm (bidirectional: mask=NULL)
|
||||
lyric_h = qwen3_build_layers(ctx, m->lyric_cfg, m->lyric_layers, m->lyric_norm,
|
||||
lyric_h, lyric_pos, NULL, S_lyric);
|
||||
|
||||
ggml_set_name(lyric_h, "lyric_out");
|
||||
ggml_set_output(lyric_h);
|
||||
|
||||
// Path 2: Text projection
|
||||
struct ggml_tensor * t_text_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1024, S_text);
|
||||
ggml_set_name(t_text_in, "text_in");
|
||||
ggml_set_input(t_text_in);
|
||||
|
||||
// Linear: [1024, S_text] -> [2048, S_text]
|
||||
struct ggml_tensor * text_proj = qwen3_linear(ctx, m->text_proj_w, t_text_in);
|
||||
ggml_set_name(text_proj, "text_proj_out");
|
||||
ggml_set_output(text_proj);
|
||||
|
||||
// Path 3: Timbre encoder (optional)
|
||||
struct ggml_tensor * timbre_out = NULL;
|
||||
struct ggml_tensor * t_timbre_in = NULL;
|
||||
struct ggml_tensor * timbre_pos = NULL;
|
||||
|
||||
if (has_timbre) {
|
||||
timbre_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, S_ref);
|
||||
ggml_set_name(timbre_pos, "timbre_pos");
|
||||
ggml_set_input(timbre_pos);
|
||||
|
||||
t_timbre_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 64, S_ref);
|
||||
ggml_set_name(t_timbre_in, "timbre_in");
|
||||
ggml_set_input(t_timbre_in);
|
||||
|
||||
// Linear embed: [64, S_ref] -> [2048, S_ref]
|
||||
struct ggml_tensor * timbre_h = qwen3_linear_bias(ctx, m->timbre_embed_w,
|
||||
m->timbre_embed_b, t_timbre_in);
|
||||
// 4 layers + final norm
|
||||
timbre_h = qwen3_build_layers(ctx, m->timbre_cfg, m->timbre_layers, m->timbre_norm,
|
||||
timbre_h, timbre_pos, NULL, S_ref);
|
||||
|
||||
// Take first frame: [2048, S_ref] -> view [2048, 1]
|
||||
timbre_out = ggml_view_2d(ctx, timbre_h, H, 1,
|
||||
timbre_h->nb[1], 0);
|
||||
ggml_set_name(timbre_out, "timbre_out");
|
||||
ggml_set_output(timbre_out);
|
||||
}
|
||||
|
||||
// Build forward
|
||||
ggml_build_forward_expand(gf, lyric_h);
|
||||
ggml_build_forward_expand(gf, text_proj);
|
||||
if (timbre_out) ggml_build_forward_expand(gf, timbre_out);
|
||||
|
||||
// Allocate and set inputs
|
||||
ggml_backend_sched_alloc_graph(m->sched, gf);
|
||||
|
||||
ggml_backend_tensor_set(t_lyric_in, lyric_embed, 0, 1024 * S_lyric * sizeof(float));
|
||||
ggml_backend_tensor_set(t_text_in, text_hidden, 0, 1024 * S_text * sizeof(float));
|
||||
|
||||
// Lyric positions
|
||||
{
|
||||
std::vector<int> pos(S_lyric);
|
||||
for (int i = 0; i < S_lyric; i++) pos[i] = i;
|
||||
ggml_backend_tensor_set(lyric_pos, pos.data(), 0, S_lyric * sizeof(int));
|
||||
}
|
||||
|
||||
if (has_timbre) {
|
||||
ggml_backend_tensor_set(t_timbre_in, timbre_feats, 0, 64 * S_ref * sizeof(float));
|
||||
std::vector<int> pos(S_ref);
|
||||
for (int i = 0; i < S_ref; i++) pos[i] = i;
|
||||
ggml_backend_tensor_set(timbre_pos, pos.data(), 0, S_ref * sizeof(int));
|
||||
}
|
||||
|
||||
// Compute
|
||||
ggml_backend_sched_graph_compute(m->sched, gf);
|
||||
|
||||
// Read outputs and pack on CPU
|
||||
// Pack order: lyric, timbre[0:1], text_proj
|
||||
int S_timbre_out = has_timbre ? 1 : 0;
|
||||
int S_total = S_lyric + S_timbre_out + S_text;
|
||||
enc_hidden.resize(H * S_total);
|
||||
*out_enc_S = S_total;
|
||||
|
||||
int offset = 0;
|
||||
|
||||
// Lyric: [2048, S_lyric]
|
||||
ggml_backend_tensor_get(lyric_h, enc_hidden.data() + offset * H,
|
||||
0, H * S_lyric * sizeof(float));
|
||||
offset += S_lyric;
|
||||
|
||||
// Timbre: [2048, 1]
|
||||
if (timbre_out) {
|
||||
ggml_backend_tensor_get(timbre_out, enc_hidden.data() + offset * H,
|
||||
0, H * 1 * sizeof(float));
|
||||
offset += 1;
|
||||
}
|
||||
|
||||
// Text projection: [2048, S_text]
|
||||
ggml_backend_tensor_get(text_proj, enc_hidden.data() + offset * H,
|
||||
0, H * S_text * sizeof(float));
|
||||
offset += S_text;
|
||||
|
||||
fprintf(stderr, "[Encode] Packed: lyric=%d + timbre=%d + text=%d = %d tokens\n",
|
||||
S_lyric, S_timbre_out, S_text, S_total);
|
||||
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(ctx);
|
||||
}
|
||||
|
||||
// Free
|
||||
static void cond_ggml_free(CondGGML * m) {
|
||||
if (m->sched) ggml_backend_sched_free(m->sched);
|
||||
if (m->backend && m->backend != m->cpu_backend) ggml_backend_free(m->backend);
|
||||
if (m->cpu_backend) ggml_backend_free(m->cpu_backend);
|
||||
wctx_free(&m->wctx);
|
||||
*m = {};
|
||||
}
|
||||
115
otherarch/acestep/debug.h
Normal file
115
otherarch/acestep/debug.h
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
#pragma once
|
||||
// debug.h: tensor dump/compare utilities 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 <cstdio>
|
||||
#include <cstdint>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
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);
|
||||
|
||||
// Print first 4 values for quick sanity check
|
||||
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 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 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);
|
||||
}
|
||||
|
||||
// Load a debug dump, returns data and fills shape/ndims
|
||||
static std::vector<float> debug_load(const char *path, std::vector<int> &shape) {
|
||||
FILE *f = fopen(path, "rb");
|
||||
if (!f) { fprintf(stderr, "[Debug] cannot read %s\n", path); return {}; }
|
||||
|
||||
int32_t ndims;
|
||||
fread(&ndims, sizeof(int32_t), 1, f);
|
||||
|
||||
shape.resize(ndims);
|
||||
fread(shape.data(), sizeof(int32_t), ndims, f);
|
||||
|
||||
int numel = 1;
|
||||
for (int i = 0; i < ndims; i++) numel *= shape[i];
|
||||
|
||||
std::vector<float> data(numel);
|
||||
fread(data.data(), sizeof(float), numel, f);
|
||||
fclose(f);
|
||||
return data;
|
||||
}
|
||||
|
||||
// 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));
|
||||
}
|
||||
|
||||
// Max absolute error
|
||||
static double debug_max_abs_err(const float *a, const float *b, int n) {
|
||||
double mx = 0;
|
||||
for (int i = 0; i < n; i++) {
|
||||
double d = fabs((double)a[i] - (double)b[i]);
|
||||
if (d > mx) mx = d;
|
||||
}
|
||||
return mx;
|
||||
}
|
||||
|
||||
// Mean absolute error
|
||||
static double debug_mean_abs_err(const float *a, const float *b, int n) {
|
||||
double sum = 0;
|
||||
for (int i = 0; i < n; i++) {
|
||||
sum += fabs((double)a[i] - (double)b[i]);
|
||||
}
|
||||
return sum / n;
|
||||
}
|
||||
|
||||
581
otherarch/acestep/dit-vae.cpp
Normal file
581
otherarch/acestep/dit-vae.cpp
Normal file
|
|
@ -0,0 +1,581 @@
|
|||
// dit.cpp: ACEStep music generation via ggml (dit-vae binary)
|
||||
//
|
||||
// Usage: ./dit-vae [options]
|
||||
// See --help for full option list.
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
#include <chrono>
|
||||
#include <random>
|
||||
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
#include "./dit.h"
|
||||
#include "./vae.h"
|
||||
#include "./qwen3.h"
|
||||
#include "./tokenizer.h"
|
||||
#include "./cond.h"
|
||||
#include "./bpe.h"
|
||||
#include "./debug.h"
|
||||
#include "./request.h"
|
||||
|
||||
struct Timer {
|
||||
std::chrono::steady_clock::time_point t;
|
||||
Timer() : t(std::chrono::steady_clock::now()) {}
|
||||
double ms() const {
|
||||
return std::chrono::duration<double, std::milli>(
|
||||
std::chrono::steady_clock::now() - t).count();
|
||||
}
|
||||
void reset() { t = std::chrono::steady_clock::now(); }
|
||||
};
|
||||
|
||||
// Minimal WAV writer (16-bit PCM stereo)
|
||||
static bool write_wav(const char * path, const float * audio, int T_audio, int sr) {
|
||||
FILE * f = fopen(path, "wb");
|
||||
if (!f) return false;
|
||||
int n_channels = 2;
|
||||
int bits = 16;
|
||||
int byte_rate = sr * n_channels * (bits / 8);
|
||||
int block_align = n_channels * (bits / 8);
|
||||
int data_size = T_audio * n_channels * (bits / 8);
|
||||
int file_size = 36 + data_size;
|
||||
fwrite("RIFF", 1, 4, f);
|
||||
fwrite(&file_size, 4, 1, f);
|
||||
fwrite("WAVE", 1, 4, f);
|
||||
fwrite("fmt ", 1, 4, f);
|
||||
int fmt_size = 16; fwrite(&fmt_size, 4, 1, f);
|
||||
short audio_fmt = 1; fwrite(&audio_fmt, 2, 1, f);
|
||||
short nc = (short)n_channels; fwrite(&nc, 2, 1, f);
|
||||
fwrite(&sr, 4, 1, f);
|
||||
fwrite(&byte_rate, 4, 1, f);
|
||||
short ba = (short)block_align; fwrite(&ba, 2, 1, f);
|
||||
short bp = (short)bits; fwrite(&bp, 2, 1, f);
|
||||
fwrite("data", 1, 4, f);
|
||||
fwrite(&data_size, 4, 1, f);
|
||||
for (int t = 0; t < T_audio; t++) {
|
||||
for (int c = 0; c < 2; c++) {
|
||||
float s = audio[c * T_audio + t];
|
||||
s = s < -1.0f ? -1.0f : (s > 1.0f ? 1.0f : s);
|
||||
short v = (short)(s * 32767.0f);
|
||||
fwrite(&v, 2, 1, f);
|
||||
}
|
||||
}
|
||||
fclose(f);
|
||||
return true;
|
||||
}
|
||||
|
||||
static void print_usage(const char * prog) {
|
||||
fprintf(stderr,
|
||||
"Usage: %s --request <json...> --text-encoder <gguf> --dit <gguf> --vae <gguf> [options]\n\n"
|
||||
"Required:\n"
|
||||
" --request <json...> One or more request JSONs (from ace-qwen3 --request)\n"
|
||||
" --text-encoder <gguf> Text encoder GGUF file\n"
|
||||
" --dit <gguf> DiT GGUF file (from convert.py)\n"
|
||||
" --vae <gguf> VAE GGUF file\n\n"
|
||||
"Batch:\n"
|
||||
" --batch <n> DiT variations per request (default: 1, max 9)\n\n"
|
||||
"Audio:\n"
|
||||
" --noise-file <path> Load noise from bf16 file (Philox RNG dump, batch=1 only)\n\n"
|
||||
"VAE tiling (memory control):\n"
|
||||
" --vae-chunk <n> Latent frames per tile (default: 256)\n"
|
||||
" --vae-overlap <n> Overlap frames per side (default: 64)\n\n"
|
||||
"Output naming: input.json -> input0.wav, input1.wav, ... (last digit = batch index)\n\n"
|
||||
"Debug:\n"
|
||||
" --dump <dir> Dump intermediate tensors\n", prog);
|
||||
}
|
||||
|
||||
// Parse comma-separated codes string into vector
|
||||
static std::vector<int> parse_codes_string(const std::string & s) {
|
||||
std::vector<int> codes;
|
||||
if (s.empty()) return codes;
|
||||
const char * p = s.c_str();
|
||||
while (*p) {
|
||||
while (*p == ',' || *p == ' ') p++;
|
||||
if (!*p) break;
|
||||
codes.push_back(atoi(p));
|
||||
while (*p && *p != ',') p++;
|
||||
}
|
||||
return codes;
|
||||
}
|
||||
|
||||
int main(int argc, char ** argv) {
|
||||
if (argc < 2) { print_usage(argv[0]); return 1; }
|
||||
|
||||
std::vector<const char *> request_paths;
|
||||
const char * text_enc_gguf = NULL;
|
||||
const char * dit_gguf = NULL;
|
||||
const char * vae_gguf = NULL;
|
||||
const char * dump_dir = NULL;
|
||||
const char * noise_file = NULL;
|
||||
int batch_n = 1;
|
||||
int vae_chunk = 256;
|
||||
int vae_overlap = 64;
|
||||
|
||||
for (int i = 1; i < argc; i++) {
|
||||
if (strcmp(argv[i], "--request") == 0) {
|
||||
// Collect all following non-option args
|
||||
while (i+1 < argc && argv[i+1][0] != '-')
|
||||
request_paths.push_back(argv[++i]);
|
||||
}
|
||||
else if (strcmp(argv[i], "--text-encoder") == 0 && i+1 < argc) text_enc_gguf = argv[++i];
|
||||
else if (strcmp(argv[i], "--dit") == 0 && i+1 < argc) dit_gguf = argv[++i];
|
||||
else if (strcmp(argv[i], "--vae") == 0 && i+1 < argc) vae_gguf = argv[++i];
|
||||
else if (strcmp(argv[i], "--noise-file") == 0 && i+1 < argc) noise_file = argv[++i];
|
||||
else if (strcmp(argv[i], "--dump") == 0 && i+1 < argc) dump_dir = argv[++i];
|
||||
else if (strcmp(argv[i], "--batch") == 0 && i+1 < argc) batch_n = atoi(argv[++i]);
|
||||
else if (strcmp(argv[i], "--vae-chunk") == 0 && i+1 < argc) vae_chunk = atoi(argv[++i]);
|
||||
else if (strcmp(argv[i], "--vae-overlap") == 0 && i+1 < argc) vae_overlap = atoi(argv[++i]);
|
||||
else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) {
|
||||
print_usage(argv[0]); return 0;
|
||||
} else {
|
||||
fprintf(stderr, "Unknown option: %s\n", argv[i]);
|
||||
print_usage(argv[0]); return 1;
|
||||
}
|
||||
}
|
||||
|
||||
if (request_paths.empty()) {
|
||||
fprintf(stderr, "ERROR: --request required\n");
|
||||
print_usage(argv[0]); return 1;
|
||||
}
|
||||
if (batch_n < 1 || batch_n > 9) {
|
||||
fprintf(stderr, "ERROR: --batch must be 1..9\n"); return 1;
|
||||
}
|
||||
if (!dit_gguf) {
|
||||
fprintf(stderr, "ERROR: --dit required\n");
|
||||
print_usage(argv[0]); return 1;
|
||||
}
|
||||
if (!text_enc_gguf) {
|
||||
fprintf(stderr, "ERROR: --text-encoder required\n");
|
||||
print_usage(argv[0]); return 1;
|
||||
}
|
||||
|
||||
|
||||
const int FRAMES_PER_SECOND = 25;
|
||||
|
||||
DebugDumper dbg;
|
||||
debug_init(&dbg, dump_dir);
|
||||
|
||||
Timer timer;
|
||||
DiTGGMLConfig cfg;
|
||||
DiTGGML model = {};
|
||||
|
||||
// Load DiT model (once for all requests)
|
||||
dit_ggml_init_backend(&model);
|
||||
fprintf(stderr, "[Load] Backend init: %.1f ms\n", timer.ms());
|
||||
|
||||
timer.reset();
|
||||
if (!dit_ggml_load(&model, dit_gguf, cfg)) {
|
||||
fprintf(stderr, "FATAL: failed to load DiT model\n");
|
||||
return 1;
|
||||
}
|
||||
fprintf(stderr, "[Load] DiT weight load: %.1f ms\n", timer.ms());
|
||||
|
||||
// Read DiT GGUF metadata + silence_latent tensor (once)
|
||||
bool is_turbo = false;
|
||||
std::vector<float> silence_full; // [15000, 64] f32
|
||||
{
|
||||
GGUFModel gf = {};
|
||||
if (gf_load(&gf, dit_gguf)) {
|
||||
is_turbo = gf_get_bool(gf, "acestep.is_turbo");
|
||||
const void * sl_data = gf_get_data(gf, "silence_latent");
|
||||
if (sl_data) {
|
||||
silence_full.resize(15000 * 64);
|
||||
memcpy(silence_full.data(), sl_data, 15000 * 64 * sizeof(float));
|
||||
fprintf(stderr, "[Load] silence_latent: [15000, 64] from GGUF\n");
|
||||
} else {
|
||||
fprintf(stderr, "FATAL: silence_latent tensor not found in %s\n", dit_gguf);
|
||||
gf_close(&gf);
|
||||
dit_ggml_free(&model);
|
||||
return 1;
|
||||
}
|
||||
gf_close(&gf);
|
||||
} else {
|
||||
fprintf(stderr, "FATAL: cannot reopen %s for metadata\n", dit_gguf);
|
||||
dit_ggml_free(&model);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
int Oc = cfg.out_channels; // 64
|
||||
int ctx_ch = cfg.in_channels - Oc; // 128
|
||||
|
||||
// Load VAE model (once for all requests)
|
||||
VAEGGML vae = {};
|
||||
bool have_vae = false;
|
||||
if (vae_gguf) {
|
||||
timer.reset();
|
||||
vae_ggml_load(&vae, vae_gguf);
|
||||
fprintf(stderr, "[Load] VAE weights: %.1f ms\n", timer.ms());
|
||||
have_vae = true;
|
||||
}
|
||||
|
||||
// Process each request
|
||||
for (int ri = 0; ri < (int)request_paths.size(); ri++) {
|
||||
const char * rpath = request_paths[ri];
|
||||
fprintf(stderr, "[Request %d/%d] %s (batch=%d)\n",
|
||||
ri + 1, (int)request_paths.size(), rpath, batch_n);
|
||||
|
||||
// Compute output basename: strip .json suffix
|
||||
std::string basename(rpath);
|
||||
{
|
||||
size_t dot = basename.rfind(".json");
|
||||
if (dot != std::string::npos)
|
||||
basename = basename.substr(0, dot);
|
||||
}
|
||||
|
||||
// Parse request JSON
|
||||
AceRequest req;
|
||||
request_init(&req);
|
||||
if (!request_parse(&req, rpath)) {
|
||||
fprintf(stderr, "ERROR: failed to parse %s, skipping\n", rpath);
|
||||
continue;
|
||||
}
|
||||
if (req.caption.empty()) {
|
||||
fprintf(stderr, "ERROR: caption is empty in %s, skipping\n", rpath);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Extract params
|
||||
const char * caption = req.caption.c_str();
|
||||
const char * lyrics = req.lyrics.empty() ? "[Instrumental]" : req.lyrics.c_str();
|
||||
char bpm_str[16] = "N/A";
|
||||
if (req.bpm > 0) snprintf(bpm_str, sizeof(bpm_str), "%d", req.bpm);
|
||||
const char * bpm = bpm_str;
|
||||
const char * keyscale = req.keyscale.empty() ? "N/A" : req.keyscale.c_str();
|
||||
const char * timesig = req.timesignature.empty() ? "N/A" : req.timesignature.c_str();
|
||||
const char * language = req.vocal_language.empty() ? "en" : req.vocal_language.c_str();
|
||||
float duration = req.duration > 0 ? req.duration : 120.0f;
|
||||
int seed = req.seed;
|
||||
int num_steps = req.inference_steps > 0 ? req.inference_steps : 8;
|
||||
float guidance_scale = req.guidance_scale > 0 ? req.guidance_scale : 7.0f;
|
||||
float shift = req.shift > 0 ? req.shift : 1.0f;
|
||||
|
||||
if (is_turbo && guidance_scale > 1.0f) {
|
||||
fprintf(stderr, "[Pipeline] WARNING: turbo model, forcing guidance_scale=1.0 (was %.1f)\n",
|
||||
guidance_scale);
|
||||
guidance_scale = 1.0f;
|
||||
}
|
||||
|
||||
if (seed < 0) {
|
||||
std::random_device rd;
|
||||
seed = (int)(rd() & 0x7FFFFFFF);
|
||||
}
|
||||
fprintf(stderr, "[Pipeline] seed=%d, steps=%d, guidance=%.1f, shift=%.1f, duration=%.1fs\n",
|
||||
seed, num_steps, guidance_scale, shift, duration);
|
||||
|
||||
// Parse audio codes from request
|
||||
std::vector<int> codes_vec = parse_codes_string(req.audio_codes);
|
||||
if (!codes_vec.empty())
|
||||
fprintf(stderr, "[Pipeline] %zu audio codes (%.1fs @ 5Hz)\n",
|
||||
codes_vec.size(), (float)codes_vec.size() / 5.0f);
|
||||
|
||||
// Build schedule: t_i = shift * t / (1 + (shift-1)*t) where t = 1 - i/steps
|
||||
std::vector<float> schedule(num_steps);
|
||||
for (int i = 0; i < num_steps; i++) {
|
||||
float t = 1.0f - (float)i / (float)num_steps;
|
||||
schedule[i] = shift * t / (1.0f + (shift - 1.0f) * t);
|
||||
}
|
||||
|
||||
// T = number of 25Hz latent frames for DiT
|
||||
// When audio codes are present, T is determined by the codes.
|
||||
// Otherwise, T is derived from the requested duration.
|
||||
int T = codes_vec.empty()
|
||||
? (int)(duration * FRAMES_PER_SECOND)
|
||||
: (int)codes_vec.size() * 5;
|
||||
T = ((T + cfg.patch_size - 1) / cfg.patch_size) * cfg.patch_size;
|
||||
int S = T / cfg.patch_size;
|
||||
int enc_S = 0;
|
||||
|
||||
fprintf(stderr, "[Pipeline] T=%d, S=%d\n", T, S);
|
||||
|
||||
if (T > 15000) {
|
||||
fprintf(stderr, "ERROR: T=%d exceeds silence_latent max 15000, skipping\n", T);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Text encoding
|
||||
// 1. Load BPE tokenizer
|
||||
timer.reset();
|
||||
BPETokenizer tok;
|
||||
if (!load_bpe_from_gguf(&tok, text_enc_gguf)) {
|
||||
fprintf(stderr, "FATAL: failed to load tokenizer from %s\n", text_enc_gguf);
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
fprintf(stderr, "[Load] BPE tokenizer: %.1f ms\n", timer.ms());
|
||||
|
||||
// 2. Build formatted prompts
|
||||
const char * instruction = "Generate audio semantic tokens based on the given conditions:";
|
||||
char metas[512];
|
||||
snprintf(metas, sizeof(metas),
|
||||
"- bpm: %s\n- timesignature: %s\n- keyscale: %s\n- duration: %d seconds\n",
|
||||
bpm, timesig, keyscale, (int)duration);
|
||||
std::string text_str = std::string("# Instruction\n")
|
||||
+ instruction + "\n\n"
|
||||
+ "# Caption\n" + caption + "\n\n"
|
||||
+ "# Metas\n" + metas + "<|endoftext|>\n";
|
||||
|
||||
bool instrumental = (strcmp(lyrics, "[Instrumental]") == 0 || strcmp(lyrics, "[instrumental]") == 0);
|
||||
std::string lyric_str = std::string("# Languages\n") + language + "\n\n# Lyric\n"
|
||||
+ (instrumental ? "[Instrumental]" : lyrics) + "<|endoftext|>";
|
||||
|
||||
// 3. Tokenize
|
||||
auto text_ids = bpe_encode(&tok, text_str.c_str(), true);
|
||||
auto lyric_ids = bpe_encode(&tok, lyric_str.c_str(), true);
|
||||
int S_text = (int)text_ids.size();
|
||||
int S_lyric = (int)lyric_ids.size();
|
||||
fprintf(stderr, "[Pipeline] caption: %d tokens, lyrics: %d tokens\n", S_text, S_lyric);
|
||||
|
||||
// 4. Text encoder forward (caption only)
|
||||
timer.reset();
|
||||
Qwen3GGML text_enc = {};
|
||||
qwen3_init_backend(&text_enc);
|
||||
if (!qwen3_load_text_encoder(&text_enc, text_enc_gguf)) {
|
||||
fprintf(stderr, "FATAL: failed to load text encoder\n");
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
fprintf(stderr, "[Load] TextEncoder: %.1f ms\n", timer.ms());
|
||||
|
||||
int H_text = text_enc.cfg.hidden_size; // 1024
|
||||
std::vector<float> text_hidden(H_text * S_text);
|
||||
|
||||
timer.reset();
|
||||
qwen3_forward(&text_enc, text_ids.data(), S_text, text_hidden.data());
|
||||
fprintf(stderr, "[Encode] TextEncoder (%d tokens): %.1f ms\n", S_text, timer.ms());
|
||||
debug_dump_2d(&dbg, "text_hidden", text_hidden.data(), S_text, H_text);
|
||||
|
||||
// 5. Lyric embedding (CPU vocab lookup from text encoder embed table)
|
||||
timer.reset();
|
||||
std::vector<float> lyric_embed(H_text * S_lyric);
|
||||
{
|
||||
GGUFModel gf_te = {};
|
||||
if (!gf_load(&gf_te, text_enc_gguf)) {
|
||||
fprintf(stderr, "FATAL: cannot reopen text encoder GGUF for lyric embed\n");
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
const void * embed_data = gf_get_data(gf_te, "embed_tokens.weight");
|
||||
if (!embed_data) {
|
||||
fprintf(stderr, "FATAL: embed_tokens.weight not found\n");
|
||||
gf_close(&gf_te);
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
qwen3_cpu_embed_lookup(embed_data, H_text,
|
||||
lyric_ids.data(), S_lyric,
|
||||
lyric_embed.data());
|
||||
gf_close(&gf_te);
|
||||
}
|
||||
fprintf(stderr, "[Encode] Lyric vocab lookup (%d tokens): %.1f ms\n", S_lyric, timer.ms());
|
||||
debug_dump_2d(&dbg, "lyric_embed", lyric_embed.data(), S_lyric, H_text);
|
||||
|
||||
// 6. Condition encoder forward
|
||||
timer.reset();
|
||||
CondGGML cond = {};
|
||||
cond_ggml_init_backend(&cond);
|
||||
if (!cond_ggml_load(&cond, dit_gguf)) {
|
||||
fprintf(stderr, "FATAL: failed to load condition encoder\n");
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
fprintf(stderr, "[Load] ConditionEncoder: %.1f ms\n", timer.ms());
|
||||
|
||||
// Silence feats for timbre input: first 750 frames (30s @ 25Hz)
|
||||
const int S_ref = 750;
|
||||
std::vector<float> silence_feats(S_ref * 64);
|
||||
memcpy(silence_feats.data(), silence_full.data(), S_ref * 64 * sizeof(float));
|
||||
|
||||
timer.reset();
|
||||
std::vector<float> enc_hidden;
|
||||
cond_ggml_forward(&cond, text_hidden.data(), S_text,
|
||||
lyric_embed.data(), S_lyric,
|
||||
silence_feats.data(), S_ref,
|
||||
enc_hidden, &enc_S);
|
||||
fprintf(stderr, "[Encode] ConditionEncoder: %.1f ms, enc_S=%d\n", timer.ms(), enc_S);
|
||||
|
||||
qwen3_free(&text_enc);
|
||||
cond_ggml_free(&cond);
|
||||
|
||||
debug_dump_2d(&dbg, "enc_hidden", enc_hidden.data(), enc_S, 2048);
|
||||
|
||||
// Context building
|
||||
// Silence latent for this T
|
||||
std::vector<float> silence(Oc * T);
|
||||
memcpy(silence.data(), silence_full.data(), (size_t)(Oc * T) * sizeof(float));
|
||||
|
||||
// Decode audio codes if provided
|
||||
int decoded_T = 0;
|
||||
std::vector<float> decoded_latents;
|
||||
if (!codes_vec.empty()) {
|
||||
timer.reset();
|
||||
DetokGGML detok = {};
|
||||
if (!detok_ggml_load(&detok, dit_gguf, model.backend, model.cpu_backend)) {
|
||||
fprintf(stderr, "FATAL: failed to load detokenizer\n");
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
fprintf(stderr, "[Load] Detokenizer: %.1f ms\n", timer.ms());
|
||||
|
||||
int T_5Hz = (int)codes_vec.size();
|
||||
int T_25Hz_codes = T_5Hz * 5;
|
||||
decoded_latents.resize(T_25Hz_codes * Oc);
|
||||
|
||||
timer.reset();
|
||||
int ret = detok_ggml_decode(&detok, codes_vec.data(), T_5Hz, decoded_latents.data());
|
||||
if (ret < 0) {
|
||||
fprintf(stderr, "FATAL: detokenizer decode failed\n");
|
||||
dit_ggml_free(&model);
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
return 1;
|
||||
}
|
||||
fprintf(stderr, "[Context] Detokenizer: %.1f ms\n", timer.ms());
|
||||
|
||||
decoded_T = T_25Hz_codes < T ? T_25Hz_codes : T;
|
||||
debug_dump_2d(&dbg, "detok_output", decoded_latents.data(), T_25Hz_codes, Oc);
|
||||
detok_ggml_free(&detok);
|
||||
}
|
||||
|
||||
// Build single context: [T, ctx_ch] = src_latents[64] + mask_ones[64]
|
||||
std::vector<float> context_single(T * ctx_ch);
|
||||
for (int t = 0; t < T; t++) {
|
||||
const float * src = (t < decoded_T)
|
||||
? decoded_latents.data() + t * Oc
|
||||
: silence.data() + t * Oc;
|
||||
for (int c = 0; c < Oc; c++)
|
||||
context_single[t * ctx_ch + c] = src[c];
|
||||
for (int c = 0; c < Oc; c++)
|
||||
context_single[t * ctx_ch + Oc + c] = 1.0f;
|
||||
}
|
||||
|
||||
// Replicate context for N batch samples (all identical)
|
||||
std::vector<float> context(batch_n * T * ctx_ch);
|
||||
for (int b = 0; b < batch_n; b++)
|
||||
memcpy(context.data() + b * T * ctx_ch, context_single.data(),
|
||||
T * ctx_ch * sizeof(float));
|
||||
|
||||
// Generate N noise samples
|
||||
std::vector<float> noise(batch_n * Oc * T);
|
||||
|
||||
if (noise_file && batch_n == 1) {
|
||||
// Load pre-generated Philox noise from bf16 file (batch=1 only)
|
||||
FILE * nf = fopen(noise_file, "rb");
|
||||
if (!nf) { fprintf(stderr, "FATAL: cannot open noise file %s\n", noise_file); return 1; }
|
||||
fseek(nf, 0, SEEK_END);
|
||||
int total_bf16 = (int)(ftell(nf) / sizeof(uint16_t));
|
||||
fseek(nf, 0, SEEK_SET);
|
||||
int T_file = total_bf16 / Oc;
|
||||
if (T_file < T) {
|
||||
fprintf(stderr, "FATAL: noise file too short: T_file=%d < T=%d\n", T_file, T);
|
||||
fclose(nf); return 1;
|
||||
}
|
||||
std::vector<uint16_t> bf16_all(Oc * T);
|
||||
if (fread(bf16_all.data(), sizeof(uint16_t), Oc * T, nf) != (size_t)(Oc * T)) {
|
||||
fprintf(stderr, "FATAL: noise file read error\n");
|
||||
fclose(nf); return 1;
|
||||
}
|
||||
fclose(nf);
|
||||
for (int i = 0; i < Oc * T; i++) {
|
||||
uint32_t w = (uint32_t)bf16_all[i] << 16;
|
||||
float v; memcpy(&v, &w, 4);
|
||||
noise[i] = v;
|
||||
}
|
||||
fprintf(stderr, "[Context] loaded noise from %s: [%d, %d] bf16\n",
|
||||
noise_file, T, Oc);
|
||||
} else {
|
||||
// Generate N noise samples with seeds: seed, seed+1, ..., seed+N-1
|
||||
for (int b = 0; b < batch_n; b++) {
|
||||
std::mt19937 rng(seed + b);
|
||||
std::normal_distribution<float> normal(0.0f, 1.0f);
|
||||
float * dst = noise.data() + b * Oc * T;
|
||||
for (int i = 0; i < Oc * T; i++)
|
||||
dst[i] = normal(rng);
|
||||
fprintf(stderr, "[Context Batch%d] noise seed=%d\n", b, seed + b);
|
||||
}
|
||||
}
|
||||
|
||||
// DiT Generate
|
||||
std::vector<float> output(batch_n * Oc * T);
|
||||
|
||||
// Debug dumps (sample 0)
|
||||
debug_dump_2d(&dbg, "noise", noise.data(), T, Oc);
|
||||
debug_dump_2d(&dbg, "context", context.data(), T, ctx_ch);
|
||||
|
||||
fprintf(stderr, "[DiT] Starting: T=%d, S=%d, enc_S=%d, steps=%d, batch=%d\n",
|
||||
T, S, enc_S, num_steps, batch_n);
|
||||
|
||||
timer.reset();
|
||||
dit_ggml_generate(&model, noise.data(), context.data(), enc_hidden.data(),
|
||||
enc_S, T, batch_n, num_steps, schedule.data(), output.data(),
|
||||
guidance_scale, &dbg);
|
||||
fprintf(stderr, "[DiT] Total generation: %.1f ms (%.1f ms/sample)\n",
|
||||
timer.ms(), timer.ms() / batch_n);
|
||||
|
||||
debug_dump_2d(&dbg, "dit_output", output.data(), T, Oc);
|
||||
|
||||
// VAE Decode + Write WAVs
|
||||
if (have_vae) {
|
||||
int T_latent = T;
|
||||
int T_audio_max = T_latent * 1920;
|
||||
std::vector<float> audio(2 * T_audio_max);
|
||||
|
||||
for (int b = 0; b < batch_n; b++) {
|
||||
float * dit_out = output.data() + b * Oc * T;
|
||||
|
||||
timer.reset();
|
||||
int T_audio = vae_ggml_decode_tiled(&vae, dit_out, T_latent, audio.data(),
|
||||
T_audio_max, vae_chunk, vae_overlap);
|
||||
if (T_audio < 0) {
|
||||
fprintf(stderr, "[VAE Batch%d] ERROR: decode failed\n", b);
|
||||
continue;
|
||||
}
|
||||
fprintf(stderr, "[VAE Batch%d] Decode: %.1f ms\n", b, timer.ms());
|
||||
|
||||
// Peak normalization to -1.0 dB
|
||||
{
|
||||
float peak = 0.0f;
|
||||
int n_samples = 2 * T_audio;
|
||||
for (int i = 0; i < n_samples; i++) {
|
||||
float a = audio[i] < 0 ? -audio[i] : audio[i];
|
||||
if (a > peak) peak = a;
|
||||
}
|
||||
if (peak > 1e-6f) {
|
||||
const float target_amp = powf(10.0f, -1.0f / 20.0f);
|
||||
float gain = target_amp / peak;
|
||||
for (int i = 0; i < n_samples; i++)
|
||||
audio[i] *= gain;
|
||||
}
|
||||
}
|
||||
|
||||
// Write WAV: basename + batch_index + .wav
|
||||
char wav_path[1024];
|
||||
snprintf(wav_path, sizeof(wav_path), "%s%d.wav", basename.c_str(), b);
|
||||
|
||||
if (b == 0) debug_dump_2d(&dbg, "vae_audio", audio.data(), 2, T_audio);
|
||||
|
||||
if (write_wav(wav_path, audio.data(), T_audio, 48000)) {
|
||||
fprintf(stderr, "[VAE Batch%d] Wrote %s: %d samples (%.2fs @ 48kHz stereo)\n",
|
||||
b, wav_path, T_audio, (float)T_audio / 48000.0f);
|
||||
} else {
|
||||
fprintf(stderr, "[VAE Batch%d] FATAL: failed to write %s\n", b, wav_path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fprintf(stderr, "[Request %d/%d] Done\n", ri + 1, (int)request_paths.size());
|
||||
}
|
||||
|
||||
if (have_vae) vae_ggml_free(&vae);
|
||||
dit_ggml_free(&model);
|
||||
fprintf(stderr, "[Pipeline] All done\n");
|
||||
return 0;
|
||||
}
|
||||
1400
otherarch/acestep/dit.h
Normal file
1400
otherarch/acestep/dit.h
Normal file
File diff suppressed because it is too large
Load diff
343
otherarch/acestep/gguf_weights.h
Normal file
343
otherarch/acestep/gguf_weights.h
Normal file
|
|
@ -0,0 +1,343 @@
|
|||
#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 "weight_ctx.h"
|
||||
#include "gguf.h"
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <windows.h>
|
||||
#else
|
||||
#include <fcntl.h>
|
||||
#include <sys/mman.h>
|
||||
#include <sys/stat.h>
|
||||
#include <unistd.h>
|
||||
#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);
|
||||
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
|
||||
size_t n = ggml_nelements(src);
|
||||
wctx->staging.emplace_back(n);
|
||||
std::vector<float> & buf = wctx->staging.back();
|
||||
|
||||
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++)
|
||||
buf[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *)&p[i]);
|
||||
} else {
|
||||
ggml_fp16_to_fp32_row((const ggml_fp16_t *)raw, buf.data(), (int)n);
|
||||
}
|
||||
|
||||
wctx->pending.push_back({tensor, buf.data(), n * sizeof(float), 0});
|
||||
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;
|
||||
}
|
||||
|
||||
// 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);
|
||||
}
|
||||
379
otherarch/acestep/quantize-acestep.cpp
Normal file
379
otherarch/acestep/quantize-acestep.cpp
Normal file
|
|
@ -0,0 +1,379 @@
|
|||
// quantize.cpp : GGUF requantizer for ACE-Step
|
||||
// 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 <input.gguf> <output.gguf> <type>
|
||||
// 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 <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#ifdef _WIN32
|
||||
# include <windows.h>
|
||||
# define strcasecmp _stricmp
|
||||
#else
|
||||
# include <sys/mman.h>
|
||||
# include <sys/stat.h>
|
||||
# include <fcntl.h>
|
||||
# include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include "ggml.h"
|
||||
#include "gguf.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
|
||||
{"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);
|
||||
}
|
||||
|
||||
static bool is_embed(const char * name) {
|
||||
return strstr(name, "embed_tokens.weight") != nullptr;
|
||||
}
|
||||
|
||||
// Should this tensor be quantized at all?
|
||||
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;
|
||||
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;
|
||||
}
|
||||
|
||||
// Promote 1D tensors (norms/biases) to F32 for precision
|
||||
static bool should_promote_f32(int n_dims) {
|
||||
return n_dims < 2;
|
||||
}
|
||||
|
||||
// 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, "Usage: %s <input.gguf> <output.gguf> <type>\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
|
||||
int n_layers = 0;
|
||||
{
|
||||
char key[128];
|
||||
snprintf(key, sizeof(key), "%s.block_count", arch);
|
||||
int64_t idx = gguf_find_key(inp, key);
|
||||
if (idx >= 0) n_layers = (int)gguf_get_val_u32(inp, (int)idx);
|
||||
}
|
||||
|
||||
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; bool promote; };
|
||||
std::vector<TensorPlan> 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, false};
|
||||
|
||||
enum ggml_type target = pick_type(name, n_dims, arch, *variant, n_layers);
|
||||
|
||||
// Promote 1D norms/biases BF16/F16 -> F32
|
||||
if (target == GGML_TYPE_COUNT && should_promote_f32(n_dims) &&
|
||||
(t->type == GGML_TYPE_BF16 || t->type == GGML_TYPE_F16)) {
|
||||
gguf_set_tensor_type(out, name, GGML_TYPE_F32);
|
||||
plans[(size_t)i] = {GGML_TYPE_F32, false, true};
|
||||
continue;
|
||||
}
|
||||
|
||||
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);
|
||||
|
||||
if (can_convert && aligned) {
|
||||
gguf_set_tensor_type(out, name, target);
|
||||
plans[(size_t)i] = {target, true, false};
|
||||
}
|
||||
}
|
||||
|
||||
// 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, n_promoted = 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.promote) {
|
||||
// BF16/F16 -> F32
|
||||
std::vector<float> f32((size_t)nel);
|
||||
to_f32(src, f32.data(), nel, t->type);
|
||||
size_t out_size = (size_t)nel * sizeof(float);
|
||||
fwrite(f32.data(), 1, out_size, fout);
|
||||
data_pos += out_size;
|
||||
bytes_out += (int64_t)out_size;
|
||||
n_promoted++;
|
||||
} else if (plan.quantize) {
|
||||
// Quantize: src -> f32 -> target
|
||||
std::vector<float> 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<uint8_t> 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, promoted %d to F32\n",
|
||||
n_quantized, n_tensors, n_promoted);
|
||||
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;
|
||||
}
|
||||
728
otherarch/acestep/qwen3-lm.h
Normal file
728
otherarch/acestep/qwen3-lm.h
Normal file
|
|
@ -0,0 +1,728 @@
|
|||
// qwen3-lm.h : Qwen3 causal LM with KV cache (GGML)
|
||||
// Autoregressive text + audio code generation for ACE-Step
|
||||
// Loads from GGUF, supports prefill + decode, tied lm_head
|
||||
#pragma once
|
||||
|
||||
#include "qwen3.h" // Qwen3Layer, Qwen3Config, layer build helpers
|
||||
#include "bpe.h"
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cmath>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// LM config (superset of encoder config)
|
||||
struct Qwen3LMConfig {
|
||||
int vocab_size;
|
||||
int hidden_size;
|
||||
int intermediate_size;
|
||||
int n_heads;
|
||||
int n_kv_heads;
|
||||
int head_dim;
|
||||
int n_layers;
|
||||
float rope_theta;
|
||||
float rms_norm_eps;
|
||||
bool tie_embeddings;
|
||||
int max_seq_len; // KV cache capacity
|
||||
};
|
||||
|
||||
// KV cache set (one per CFG path: conditional + unconditional)
|
||||
#define QW3LM_MAX_KV_SETS 32 // batch N * 2 (cond + uncond CFG)
|
||||
#define QW3LM_MAX_LAYERS 64
|
||||
|
||||
struct Qwen3LM {
|
||||
Qwen3LMConfig cfg;
|
||||
|
||||
// Weights (on backend)
|
||||
struct ggml_tensor * embed_tokens; // [H, V] on GPU (used by mul_mat lm_head)
|
||||
struct ggml_tensor * final_norm; // [H]
|
||||
// lm_head = embed_tokens when tie_embeddings
|
||||
Qwen3Layer layers[QW3LM_MAX_LAYERS];
|
||||
|
||||
WeightCtx wctx;
|
||||
ggml_backend_t backend;
|
||||
ggml_backend_t cpu_backend;
|
||||
ggml_backend_sched_t sched;
|
||||
|
||||
// CPU-side embed lookup via mmap (avoids ggml_get_rows which lacks
|
||||
// CUDA K-quant support, preventing costly cross-backend tensor copies)
|
||||
GGUFModel gf_mmap;
|
||||
const void * embed_mmap_data;
|
||||
enum ggml_type embed_type;
|
||||
|
||||
// KV cache: per-set, per-layer [D, max_seq, Nkv] f16
|
||||
struct ggml_context * kv_ctx;
|
||||
ggml_backend_buffer_t kv_buf;
|
||||
struct ggml_tensor * kv_k[QW3LM_MAX_KV_SETS][QW3LM_MAX_LAYERS];
|
||||
struct ggml_tensor * kv_v[QW3LM_MAX_KV_SETS][QW3LM_MAX_LAYERS];
|
||||
int kv_pos[QW3LM_MAX_KV_SETS];
|
||||
int n_kv_sets;
|
||||
};
|
||||
|
||||
// Parse config.json integers, floats, bools
|
||||
static int qw3lm_json_int(const char * json, const char * key, int fb) {
|
||||
char needle[128];
|
||||
snprintf(needle, sizeof(needle), "\"%s\"", key);
|
||||
const char * p = strstr(json, needle);
|
||||
if (!p) return fb;
|
||||
p = strchr(p + strlen(needle), ':');
|
||||
if (!p) return fb;
|
||||
p++;
|
||||
while (*p == ' ' || *p == '\t') p++;
|
||||
return atoi(p);
|
||||
}
|
||||
|
||||
static float qw3lm_json_float(const char * json, const char * key, float fb) {
|
||||
char needle[128];
|
||||
snprintf(needle, sizeof(needle), "\"%s\"", key);
|
||||
const char * p = strstr(json, needle);
|
||||
if (!p) return fb;
|
||||
p = strchr(p + strlen(needle), ':');
|
||||
if (!p) return fb;
|
||||
p++;
|
||||
while (*p == ' ' || *p == '\t') p++;
|
||||
return (float)atof(p);
|
||||
}
|
||||
|
||||
static bool qw3lm_json_bool(const char * json, const char * key, bool fb) {
|
||||
char needle[128];
|
||||
snprintf(needle, sizeof(needle), "\"%s\"", key);
|
||||
const char * p = strstr(json, needle);
|
||||
if (!p) return fb;
|
||||
p = strchr(p + strlen(needle), ':');
|
||||
if (!p) return fb;
|
||||
p++;
|
||||
while (*p == ' ' || *p == '\t') p++;
|
||||
return (strncmp(p, "true", 4) == 0);
|
||||
}
|
||||
|
||||
// Load config from GGUF KV metadata (acestep.config_json)
|
||||
static Qwen3LMConfig qw3lm_load_config(const GGUFModel & gf) {
|
||||
// 0.6B defaults
|
||||
Qwen3LMConfig c = {
|
||||
/*vocab_size*/ 217204,
|
||||
/*hidden_size*/ 1024,
|
||||
/*intermediate_size*/ 3072,
|
||||
/*n_heads*/ 16,
|
||||
/*n_kv_heads*/ 8,
|
||||
/*head_dim*/ 128,
|
||||
/*n_layers*/ 28,
|
||||
/*rope_theta*/ 1000000.0f,
|
||||
/*rms_norm_eps*/ 1e-6f,
|
||||
/*tie_embeddings*/ true,
|
||||
/*max_seq_len*/ 8192,
|
||||
};
|
||||
|
||||
const char * j = gf_get_str(gf, "acestep.config_json");
|
||||
if (!j || !j[0]) {
|
||||
fprintf(stderr, "[LM-Config] No acestep.config_json, using 0.6B defaults\n");
|
||||
return c;
|
||||
}
|
||||
|
||||
c.vocab_size = qw3lm_json_int(j, "vocab_size", c.vocab_size);
|
||||
c.hidden_size = qw3lm_json_int(j, "hidden_size", c.hidden_size);
|
||||
c.intermediate_size = qw3lm_json_int(j, "intermediate_size", c.intermediate_size);
|
||||
c.n_heads = qw3lm_json_int(j, "num_attention_heads", c.n_heads);
|
||||
c.n_kv_heads = qw3lm_json_int(j, "num_key_value_heads", c.n_kv_heads);
|
||||
c.head_dim = qw3lm_json_int(j, "head_dim", c.head_dim);
|
||||
c.n_layers = qw3lm_json_int(j, "num_hidden_layers", c.n_layers);
|
||||
c.rope_theta = qw3lm_json_float(j, "rope_theta", c.rope_theta);
|
||||
c.rms_norm_eps = qw3lm_json_float(j, "rms_norm_eps", c.rms_norm_eps);
|
||||
c.tie_embeddings = qw3lm_json_bool(j, "tie_word_embeddings", c.tie_embeddings);
|
||||
|
||||
fprintf(stderr, "[LM-Config] %dL, H=%d, V=%d, Nh=%d, Nkv=%d, D=%d, tied=%d\n",
|
||||
c.n_layers, c.hidden_size, c.vocab_size, c.n_heads, c.n_kv_heads,
|
||||
c.head_dim, c.tie_embeddings);
|
||||
return c;
|
||||
}
|
||||
|
||||
// Init backend (same pattern as qwen3.h)
|
||||
static void qw3lm_init_backend(Qwen3LM * m) {
|
||||
BackendPair bp = backend_init("LM");
|
||||
m->backend = bp.backend;
|
||||
m->cpu_backend = bp.cpu_backend;
|
||||
m->sched = backend_sched_new(bp, 8192);
|
||||
}
|
||||
|
||||
// Allocate KV cache
|
||||
static void qw3lm_alloc_kv_cache(Qwen3LM * m, int n_sets) {
|
||||
const Qwen3LMConfig & c = m->cfg;
|
||||
int D = c.head_dim;
|
||||
int Nkv = c.n_kv_heads;
|
||||
int L = c.n_layers;
|
||||
int S = c.max_seq_len;
|
||||
|
||||
m->n_kv_sets = n_sets;
|
||||
|
||||
// Each KV tensor: [D, max_seq, Nkv] f16
|
||||
int n_tensors = n_sets * L * 2;
|
||||
size_t ctx_size = (size_t)n_tensors * ggml_tensor_overhead() + 1024;
|
||||
struct ggml_init_params gp = { ctx_size, NULL, true };
|
||||
m->kv_ctx = ggml_init(gp);
|
||||
|
||||
for (int s = 0; s < n_sets; s++) {
|
||||
for (int l = 0; l < L; l++) {
|
||||
m->kv_k[s][l] = ggml_new_tensor_3d(m->kv_ctx, GGML_TYPE_F16, D, S, Nkv);
|
||||
m->kv_v[s][l] = ggml_new_tensor_3d(m->kv_ctx, GGML_TYPE_F16, D, S, Nkv);
|
||||
char name[64];
|
||||
snprintf(name, sizeof(name), "kv_k_%d_%d", s, l);
|
||||
ggml_set_name(m->kv_k[s][l], name);
|
||||
snprintf(name, sizeof(name), "kv_v_%d_%d", s, l);
|
||||
ggml_set_name(m->kv_v[s][l], name);
|
||||
}
|
||||
m->kv_pos[s] = 0;
|
||||
}
|
||||
|
||||
m->kv_buf = ggml_backend_alloc_ctx_tensors(m->kv_ctx, m->backend);
|
||||
if (!m->kv_buf) {
|
||||
fprintf(stderr, "[LM-KV] FATAL: failed to allocate KV cache\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
size_t kv_bytes = (size_t)n_sets * L * 2 * D * S * Nkv * ggml_type_size(GGML_TYPE_F16);
|
||||
fprintf(stderr, "[LM-KV] Allocated %d sets x %d layers, %.1f MB\n",
|
||||
n_sets, L, (float)kv_bytes / (1024 * 1024));
|
||||
}
|
||||
|
||||
// Clear KV cache for a given set
|
||||
static void qw3lm_reset_kv(Qwen3LM * m, int kv_set) {
|
||||
m->kv_pos[kv_set] = 0;
|
||||
// No need to zero memory: kv_pos tracks valid range
|
||||
}
|
||||
|
||||
// Copy KV cache from one set to another (for batched prefill sharing)
|
||||
static void qw3lm_copy_kv(Qwen3LM * m, int src, int dst) {
|
||||
for (int l = 0; l < m->cfg.n_layers; l++) {
|
||||
ggml_backend_tensor_copy(m->kv_k[src][l], m->kv_k[dst][l]);
|
||||
ggml_backend_tensor_copy(m->kv_v[src][l], m->kv_v[dst][l]);
|
||||
}
|
||||
m->kv_pos[dst] = m->kv_pos[src];
|
||||
}
|
||||
|
||||
// Load model weights from GGUF
|
||||
static bool qw3lm_load(Qwen3LM * m, const char * gguf_path, int max_seq_len, int n_kv_sets) {
|
||||
*m = {};
|
||||
|
||||
qw3lm_init_backend(m);
|
||||
|
||||
GGUFModel gf;
|
||||
if (!gf_load(&gf, gguf_path)) {
|
||||
fprintf(stderr, "[LM-Load] FATAL: cannot load %s\n", gguf_path);
|
||||
return false;
|
||||
}
|
||||
|
||||
m->cfg = qw3lm_load_config(gf);
|
||||
if (max_seq_len > 0) m->cfg.max_seq_len = max_seq_len;
|
||||
const Qwen3LMConfig & c = m->cfg;
|
||||
|
||||
if (c.n_layers > QW3LM_MAX_LAYERS) {
|
||||
fprintf(stderr, "[LM-Load] FATAL: %d layers > max %d\n", c.n_layers, QW3LM_MAX_LAYERS);
|
||||
gf_close(&gf);
|
||||
return false;
|
||||
}
|
||||
|
||||
// embed(1) + layers * 11 + final_norm(1) = 2 + n_layers * 11
|
||||
int n_tensors = 2 + c.n_layers * 11;
|
||||
wctx_init(&m->wctx, n_tensors);
|
||||
|
||||
m->embed_tokens = gf_load_tensor(&m->wctx, gf, "model.embed_tokens.weight");
|
||||
m->final_norm = gf_load_tensor_f32(&m->wctx, gf, "model.norm.weight");
|
||||
|
||||
for (int i = 0; i < c.n_layers; i++) {
|
||||
char prefix[64];
|
||||
snprintf(prefix, sizeof(prefix), "model.layers.%d", i);
|
||||
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix);
|
||||
}
|
||||
|
||||
wctx_alloc(&m->wctx, m->backend);
|
||||
|
||||
// Keep mmap alive for CPU embed dequant lookup
|
||||
m->embed_mmap_data = gf_get_data(gf, "model.embed_tokens.weight");
|
||||
m->embed_type = m->embed_tokens->type;
|
||||
if (!m->embed_mmap_data) {
|
||||
fprintf(stderr, "[LM-Load] FATAL: embed_tokens not found in mmap\n");
|
||||
gf_close(&gf);
|
||||
return false;
|
||||
}
|
||||
m->gf_mmap = gf; // transfer ownership (no gf_close here)
|
||||
fprintf(stderr, "[LM-Load] CPU embed lookup: type=%s, row=%zu bytes\n",
|
||||
ggml_type_name(m->embed_type),
|
||||
ggml_row_size(m->embed_type, c.hidden_size));
|
||||
|
||||
// KV cache
|
||||
qw3lm_alloc_kv_cache(m, n_kv_sets > 0 ? n_kv_sets : 1);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
// Build self-attention with KV cache write + read
|
||||
// x: [H, n_tokens], positions: [n_tokens], mask: [kv_len, n_tokens] or NULL
|
||||
static struct ggml_tensor * qw3lm_build_attn(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_cgraph * gf,
|
||||
const Qwen3LMConfig & c,
|
||||
Qwen3Layer * ly,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * positions,
|
||||
struct ggml_tensor * mask,
|
||||
struct ggml_tensor * cache_k, // [D, max_seq, Nkv] f16
|
||||
struct ggml_tensor * cache_v, // [D, max_seq, Nkv] f16
|
||||
int kv_pos,
|
||||
int kv_len,
|
||||
int n_tokens) {
|
||||
|
||||
int D = c.head_dim;
|
||||
int Nh = c.n_heads;
|
||||
int Nkv = c.n_kv_heads;
|
||||
int S = n_tokens;
|
||||
|
||||
// QKV projections
|
||||
struct ggml_tensor * q = qwen3_linear(ctx, ly->q_proj, x); // [Nh*D, S]
|
||||
struct ggml_tensor * k = qwen3_linear(ctx, ly->k_proj, x); // [Nkv*D, S]
|
||||
struct ggml_tensor * v = qwen3_linear(ctx, ly->v_proj, x); // [Nkv*D, S]
|
||||
|
||||
// Reshape to heads: [X*D, S] -> [D, X, S]
|
||||
q = ggml_reshape_3d(ctx, q, D, Nh, S);
|
||||
k = ggml_reshape_3d(ctx, k, D, Nkv, S);
|
||||
v = ggml_reshape_3d(ctx, v, D, Nkv, S);
|
||||
|
||||
// QK-Norm
|
||||
q = ggml_rms_norm(ctx, q, c.rms_norm_eps);
|
||||
q = ggml_mul(ctx, q, qwen3_f32(ctx, ly->q_norm));
|
||||
k = ggml_rms_norm(ctx, k, c.rms_norm_eps);
|
||||
k = ggml_mul(ctx, k, qwen3_f32(ctx, ly->k_norm));
|
||||
|
||||
// RoPE (NEOX mode=2)
|
||||
q = ggml_rope_ext(ctx, q, positions, NULL,
|
||||
D, 2, 0, c.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx, k, positions, NULL,
|
||||
D, 2, 0, c.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// Permute for flash_attn: [D, X, S] -> [D, S, X]
|
||||
q = ggml_permute(ctx, q, 0, 2, 1, 3); // [D, S, Nh]
|
||||
k = ggml_permute(ctx, k, 0, 2, 1, 3); // [D, S, Nkv]
|
||||
v = ggml_permute(ctx, v, 0, 2, 1, 3); // [D, S, Nkv]
|
||||
|
||||
// Make contiguous for cpy to f16 cache
|
||||
k = ggml_cont(ctx, k);
|
||||
v = ggml_cont(ctx, v);
|
||||
|
||||
// Write K,V to cache at kv_pos
|
||||
// Cache layout: [D, max_seq, Nkv] f16
|
||||
size_t nb1 = (size_t)D * ggml_type_size(GGML_TYPE_F16);
|
||||
size_t nb2 = (size_t)D * c.max_seq_len * ggml_type_size(GGML_TYPE_F16);
|
||||
size_t off = (size_t)kv_pos * nb1;
|
||||
|
||||
struct ggml_tensor * k_dst = ggml_view_3d(ctx, cache_k, D, S, Nkv, nb1, nb2, off);
|
||||
struct ggml_tensor * v_dst = ggml_view_3d(ctx, cache_v, D, S, Nkv, nb1, nb2, off);
|
||||
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx, k, k_dst));
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx, v, v_dst));
|
||||
|
||||
// Read full KV from cache [0..kv_len]
|
||||
struct ggml_tensor * k_full = ggml_view_3d(ctx, cache_k, D, kv_len, Nkv, nb1, nb2, 0);
|
||||
struct ggml_tensor * v_full = ggml_view_3d(ctx, cache_v, D, kv_len, Nkv, nb1, nb2, 0);
|
||||
|
||||
// Flash attention
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q, k_full, v_full, mask, scale, 0.0f, 0.0f);
|
||||
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32); // F32 accumulation
|
||||
|
||||
// Reshape: [D, Nh, S] -> [Nh*D, S]
|
||||
attn = ggml_reshape_2d(ctx, attn, Nh * D, S);
|
||||
|
||||
// O projection
|
||||
return qwen3_linear(ctx, ly->o_proj, attn);
|
||||
}
|
||||
|
||||
// Forward pass: token_ids[n_tokens] -> logits[vocab_size] (last token only)
|
||||
// kv_set: which KV cache set to use (0=conditional, 1=unconditional for CFG)
|
||||
static void qw3lm_forward(Qwen3LM * m, const int * token_ids, int n_tokens,
|
||||
int kv_set, float * logits) {
|
||||
const Qwen3LMConfig & c = m->cfg;
|
||||
int H = c.hidden_size;
|
||||
int kv_pos = m->kv_pos[kv_set];
|
||||
int kv_len = kv_pos + n_tokens;
|
||||
|
||||
if (kv_len > c.max_seq_len) {
|
||||
fprintf(stderr, "[LM-Forward] FATAL: kv_len %d > max_seq %d\n", kv_len, c.max_seq_len);
|
||||
return;
|
||||
}
|
||||
|
||||
// Graph context
|
||||
size_t ctx_size = (size_t)16384 * ggml_tensor_overhead() + ggml_graph_overhead();
|
||||
struct ggml_init_params gp = { ctx_size, NULL, true };
|
||||
struct ggml_context * ctx = ggml_init(gp);
|
||||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
|
||||
|
||||
// Inputs
|
||||
struct ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
|
||||
ggml_set_name(positions, "positions");
|
||||
ggml_set_input(positions);
|
||||
|
||||
// Causal mask: needed for prefill (n_tokens > 1), not for decode (n_tokens == 1)
|
||||
struct ggml_tensor * mask = NULL;
|
||||
if (n_tokens > 1) {
|
||||
mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv_len, n_tokens);
|
||||
ggml_set_name(mask, "causal_mask");
|
||||
ggml_set_input(mask);
|
||||
}
|
||||
|
||||
// Embedding: CPU dequant from mmap, fed as F32 input.
|
||||
// This keeps embed_tokens out of get_rows (no CUDA K-quant support)
|
||||
// and only in mul_mat (lm_head) which has full K-quant CUDA support.
|
||||
struct ggml_tensor * embed_out = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, H, n_tokens);
|
||||
ggml_set_name(embed_out, "embed_out");
|
||||
ggml_set_input(embed_out);
|
||||
|
||||
struct ggml_tensor * hidden = embed_out;
|
||||
|
||||
// Transformer layers
|
||||
for (int l = 0; l < c.n_layers; l++) {
|
||||
Qwen3Layer * ly = &m->layers[l];
|
||||
|
||||
// Pre-attention norm
|
||||
struct ggml_tensor * norm = qwen3_rms_norm(ctx, hidden, ly->input_layernorm, c.rms_norm_eps);
|
||||
|
||||
// Self-attention with KV cache
|
||||
struct ggml_tensor * attn = qw3lm_build_attn(
|
||||
ctx, gf, c, ly, norm, positions, mask,
|
||||
m->kv_k[kv_set][l], m->kv_v[kv_set][l],
|
||||
kv_pos, kv_len, n_tokens);
|
||||
|
||||
// Residual
|
||||
hidden = ggml_add(ctx, hidden, attn);
|
||||
|
||||
// Post-attention norm + MLP
|
||||
norm = qwen3_rms_norm(ctx, hidden, ly->post_attn_layernorm, c.rms_norm_eps);
|
||||
struct ggml_tensor * mlp = qwen3_build_mlp(ctx, ly, norm, n_tokens);
|
||||
hidden = ggml_add(ctx, hidden, mlp);
|
||||
}
|
||||
|
||||
// Final norm
|
||||
hidden = qwen3_rms_norm(ctx, hidden, m->final_norm, c.rms_norm_eps);
|
||||
|
||||
// Extract last token hidden state: [H, n_tokens] -> [H, 1]
|
||||
if (n_tokens > 1) {
|
||||
hidden = ggml_view_1d(ctx, hidden, H,
|
||||
(int64_t)(n_tokens - 1) * H * sizeof(float));
|
||||
}
|
||||
|
||||
// LM head: logits = embed_tokens^T @ hidden -> [V, 1]
|
||||
struct ggml_tensor * lgt = ggml_mul_mat(ctx, m->embed_tokens, hidden);
|
||||
ggml_set_name(lgt, "logits");
|
||||
ggml_set_output(lgt);
|
||||
ggml_build_forward_expand(gf, lgt);
|
||||
|
||||
// Schedule + allocate
|
||||
ggml_backend_sched_alloc_graph(m->sched, gf);
|
||||
|
||||
// CPU-side embedding dequantization from mmap
|
||||
{
|
||||
const int64_t row_size = (int64_t)ggml_row_size(m->embed_type, H);
|
||||
const ggml_to_float_t to_float = ggml_get_type_traits(m->embed_type)->to_float;
|
||||
std::vector<float> embed_buf((size_t)H * n_tokens);
|
||||
for (int i = 0; i < n_tokens; i++) {
|
||||
const void * row = (const char *)m->embed_mmap_data + (int64_t)token_ids[i] * row_size;
|
||||
to_float(row, embed_buf.data() + (int64_t)i * H, H);
|
||||
}
|
||||
ggml_backend_tensor_set(embed_out, embed_buf.data(), 0,
|
||||
(size_t)H * n_tokens * sizeof(float));
|
||||
}
|
||||
|
||||
{
|
||||
std::vector<int> pos_data(n_tokens);
|
||||
for (int i = 0; i < n_tokens; i++) pos_data[i] = kv_pos + i;
|
||||
ggml_backend_tensor_set(positions, pos_data.data(), 0, n_tokens * sizeof(int));
|
||||
}
|
||||
|
||||
if (mask) {
|
||||
// Causal mask: [kv_len, n_tokens]
|
||||
// Row i (query at position kv_pos+i) can attend to columns [0..kv_pos+i]
|
||||
std::vector<uint16_t> mask_data((size_t)kv_len * n_tokens);
|
||||
for (int i = 0; i < n_tokens; i++) {
|
||||
int query_abs_pos = kv_pos + i;
|
||||
for (int j = 0; j < kv_len; j++) {
|
||||
float v = (j <= query_abs_pos) ? 0.0f : -INFINITY;
|
||||
mask_data[(size_t)i * kv_len + j] = ggml_fp32_to_fp16(v);
|
||||
}
|
||||
}
|
||||
ggml_backend_tensor_set(mask, mask_data.data(), 0,
|
||||
(size_t)kv_len * n_tokens * sizeof(uint16_t));
|
||||
}
|
||||
|
||||
// Compute
|
||||
ggml_backend_sched_graph_compute(m->sched, gf);
|
||||
|
||||
// Read logits [V]
|
||||
ggml_backend_tensor_get(lgt, logits, 0, c.vocab_size * sizeof(float));
|
||||
|
||||
// Advance KV position
|
||||
m->kv_pos[kv_set] += n_tokens;
|
||||
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(ctx);
|
||||
}
|
||||
|
||||
// Batched decode forward: N tokens (1 per sequence), batched weight matmuls.
|
||||
// kv_pos per element from m->kv_pos[kv_sets[i]], supports different prompt lengths.
|
||||
// kv_sets[N]: which KV set each token uses.
|
||||
// logits: [N * V] output, N logit vectors concatenated.
|
||||
static void qw3lm_forward_batch(Qwen3LM * m, const int * token_ids,
|
||||
const int * kv_sets, int N, float * logits) {
|
||||
const Qwen3LMConfig & c = m->cfg;
|
||||
int H = c.hidden_size;
|
||||
int D = c.head_dim;
|
||||
int Nh = c.n_heads;
|
||||
int Nkv = c.n_kv_heads;
|
||||
|
||||
// Per-element kv_pos (supports different prompt lengths)
|
||||
int max_kv_len = 0;
|
||||
for (int i = 0; i < N; i++) {
|
||||
int kl = m->kv_pos[kv_sets[i]] + 1;
|
||||
if (kl > max_kv_len) max_kv_len = kl;
|
||||
if (kl > c.max_seq_len) {
|
||||
fprintf(stderr, "[LM-Batch] FATAL: kv_len %d > max_seq %d (set %d)\n",
|
||||
kl, c.max_seq_len, kv_sets[i]);
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
// Exact tensor count for context allocation
|
||||
// qwen3_f32(t) creates 0 tensors if t is F32, 1 (ggml_cast) if not.
|
||||
// Count conditionally based on actual weight types.
|
||||
//
|
||||
// GLOBAL (2):
|
||||
// embed_out (new_tensor_2d) 1
|
||||
// positions (new_tensor_1d) 1
|
||||
//
|
||||
// PER LAYER fixed (17):
|
||||
// qwen3_rms_norm(input_layernorm):
|
||||
// ggml_rms_norm + ggml_mul 2
|
||||
// q_proj, k_proj, v_proj (qwen3_linear = mul_mat) 3
|
||||
// reshape_3d (q, k, v) 3
|
||||
// ggml_rms_norm(q) + ggml_mul(q, q_norm) 2
|
||||
// ggml_rms_norm(k) + ggml_mul(k, k_norm) 2
|
||||
// ggml_rope_ext (q, k) 2
|
||||
// ggml_cont (q, k, v) 3
|
||||
//
|
||||
// PER LAYER * N (16 each):
|
||||
// view_3d (qi, ki, vi) 3
|
||||
// permute (qi, ki, vi) 3
|
||||
// cont (ki, vi) 2
|
||||
// view_3d (k_dst, v_dst) 2
|
||||
// cpy (ki, vi) 2
|
||||
// view_3d (k_full, v_full) 2
|
||||
// flash_attn_ext 1
|
||||
// reshape_2d 1
|
||||
//
|
||||
// CONCATS: N-1 (first element reuses reshape_2d)
|
||||
//
|
||||
// PER LAYER post (9):
|
||||
// qwen3_linear(o_proj) 1
|
||||
// ggml_add (residual) 1
|
||||
// qwen3_rms_norm(post_attn_layernorm):
|
||||
// ggml_rms_norm + ggml_mul 2
|
||||
// qwen3_build_mlp:
|
||||
// gate_proj, up_proj (mul_mat) 2
|
||||
// ggml_swiglu_split 1
|
||||
// down_proj (mul_mat) 1
|
||||
// ggml_add (residual) 1
|
||||
//
|
||||
// PER LAYER conditional casts (0 to 4):
|
||||
// qwen3_f32(input_layernorm) 0 or 1
|
||||
// qwen3_f32(q_norm) 0 or 1
|
||||
// qwen3_f32(k_norm) 0 or 1
|
||||
// qwen3_f32(post_attn_norm) 0 or 1
|
||||
//
|
||||
// POST LAYERS (3):
|
||||
// qwen3_rms_norm(final_norm): ggml_rms_norm + ggml_mul 2
|
||||
// ggml_mul_mat (lm_head) 1
|
||||
//
|
||||
// POST conditional cast (0 or 1):
|
||||
// qwen3_f32(final_norm) 0 or 1
|
||||
//
|
||||
// TOTAL = 5 + n_layers * (25 + 17*N + casts_per_layer) + 3 + cast_final
|
||||
// = 8 + n_layers * (25 + 17*N + casts_per_layer) + cast_final
|
||||
|
||||
int casts_per_layer = 0;
|
||||
if (m->layers[0].input_layernorm->type != GGML_TYPE_F32) casts_per_layer++;
|
||||
if (m->layers[0].q_norm->type != GGML_TYPE_F32) casts_per_layer++;
|
||||
if (m->layers[0].k_norm->type != GGML_TYPE_F32) casts_per_layer++;
|
||||
if (m->layers[0].post_attn_layernorm->type != GGML_TYPE_F32) casts_per_layer++;
|
||||
int cast_final = (m->final_norm->type != GGML_TYPE_F32) ? 1 : 0;
|
||||
|
||||
size_t n_tensors = 8
|
||||
+ (size_t)c.n_layers * (25 + 17 * N + casts_per_layer)
|
||||
+ cast_final;
|
||||
size_t est = n_tensors * ggml_tensor_overhead()
|
||||
+ ggml_graph_overhead_custom(16384, false);
|
||||
struct ggml_init_params gp = { est, NULL, true };
|
||||
struct ggml_context * ctx = ggml_init(gp);
|
||||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
|
||||
|
||||
// Embedding: [H, N]
|
||||
struct ggml_tensor * embed_out = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, H, N);
|
||||
ggml_set_name(embed_out, "embed_out");
|
||||
ggml_set_input(embed_out);
|
||||
|
||||
// Positions: [N], per-element kv_pos
|
||||
struct ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N);
|
||||
ggml_set_name(positions, "positions");
|
||||
ggml_set_input(positions);
|
||||
|
||||
struct ggml_tensor * hidden = embed_out;
|
||||
|
||||
for (int l = 0; l < c.n_layers; l++) {
|
||||
Qwen3Layer * ly = &m->layers[l];
|
||||
|
||||
// Pre-attention norm [H, N]
|
||||
struct ggml_tensor * norm = qwen3_rms_norm(ctx, hidden, ly->input_layernorm, c.rms_norm_eps);
|
||||
|
||||
// Batched QKV projections (weights read once for N tokens)
|
||||
struct ggml_tensor * q = qwen3_linear(ctx, ly->q_proj, norm); // [Nh*D, N]
|
||||
struct ggml_tensor * k = qwen3_linear(ctx, ly->k_proj, norm); // [Nkv*D, N]
|
||||
struct ggml_tensor * v = qwen3_linear(ctx, ly->v_proj, norm); // [Nkv*D, N]
|
||||
|
||||
// Reshape to heads: [D, Heads, N]
|
||||
q = ggml_reshape_3d(ctx, q, D, Nh, N);
|
||||
k = ggml_reshape_3d(ctx, k, D, Nkv, N);
|
||||
v = ggml_reshape_3d(ctx, v, D, Nkv, N);
|
||||
|
||||
// QK-Norm (rms_norm on dim0=D, per head per seq)
|
||||
q = ggml_rms_norm(ctx, q, c.rms_norm_eps);
|
||||
q = ggml_mul(ctx, q, qwen3_f32(ctx, ly->q_norm));
|
||||
k = ggml_rms_norm(ctx, k, c.rms_norm_eps);
|
||||
k = ggml_mul(ctx, k, qwen3_f32(ctx, ly->k_norm));
|
||||
|
||||
// RoPE: positions [N] maps to dim 2 of [D, Heads, N]
|
||||
q = ggml_rope_ext(ctx, q, positions, NULL, D, 2, 0,
|
||||
c.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx, k, positions, NULL, D, 2, 0,
|
||||
c.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// Contiguous for clean slicing
|
||||
q = ggml_cont(ctx, q);
|
||||
k = ggml_cont(ctx, k);
|
||||
v = ggml_cont(ctx, v);
|
||||
|
||||
// Per-sequence attention with individual KV caches
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
size_t nb1_c = (size_t)D * ggml_type_size(GGML_TYPE_F16);
|
||||
size_t nb2_c = (size_t)D * c.max_seq_len * ggml_type_size(GGML_TYPE_F16);
|
||||
|
||||
struct ggml_tensor * attn_cat = NULL;
|
||||
|
||||
for (int i = 0; i < N; i++) {
|
||||
int set = kv_sets[i];
|
||||
int elem_kv_pos = m->kv_pos[set];
|
||||
int elem_kv_len = elem_kv_pos + 1;
|
||||
size_t off_c = (size_t)elem_kv_pos * nb1_c;
|
||||
|
||||
// Slice [D, Heads, 1] from [D, Heads, N]
|
||||
struct ggml_tensor * qi = ggml_view_3d(ctx, q, D, Nh, 1,
|
||||
q->nb[1], q->nb[2], (size_t)i * q->nb[2]);
|
||||
struct ggml_tensor * ki = ggml_view_3d(ctx, k, D, Nkv, 1,
|
||||
k->nb[1], k->nb[2], (size_t)i * k->nb[2]);
|
||||
struct ggml_tensor * vi = ggml_view_3d(ctx, v, D, Nkv, 1,
|
||||
v->nb[1], v->nb[2], (size_t)i * v->nb[2]);
|
||||
|
||||
// Permute [D, Heads, 1] to [D, 1, Heads]
|
||||
qi = ggml_permute(ctx, qi, 0, 2, 1, 3);
|
||||
ki = ggml_permute(ctx, ki, 0, 2, 1, 3);
|
||||
vi = ggml_permute(ctx, vi, 0, 2, 1, 3);
|
||||
ki = ggml_cont(ctx, ki);
|
||||
vi = ggml_cont(ctx, vi);
|
||||
|
||||
// Write K,V to cache at kv_pos
|
||||
struct ggml_tensor * k_dst = ggml_view_3d(ctx, m->kv_k[set][l],
|
||||
D, 1, Nkv, nb1_c, nb2_c, off_c);
|
||||
struct ggml_tensor * v_dst = ggml_view_3d(ctx, m->kv_v[set][l],
|
||||
D, 1, Nkv, nb1_c, nb2_c, off_c);
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx, ki, k_dst));
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx, vi, v_dst));
|
||||
|
||||
// Read full KV [0..elem_kv_len]
|
||||
struct ggml_tensor * k_full = ggml_view_3d(ctx, m->kv_k[set][l],
|
||||
D, elem_kv_len, Nkv, nb1_c, nb2_c, 0);
|
||||
struct ggml_tensor * v_full = ggml_view_3d(ctx, m->kv_v[set][l],
|
||||
D, elem_kv_len, Nkv, nb1_c, nb2_c, 0);
|
||||
|
||||
// Flash attention
|
||||
struct ggml_tensor * attn_i = ggml_flash_attn_ext(ctx, qi, k_full, v_full,
|
||||
NULL, scale, 0.0f, 0.0f);
|
||||
ggml_flash_attn_ext_set_prec(attn_i, GGML_PREC_F32); // F32 accumulation
|
||||
attn_i = ggml_reshape_2d(ctx, attn_i, Nh * D, 1);
|
||||
|
||||
if (i == 0) attn_cat = attn_i;
|
||||
else attn_cat = ggml_concat(ctx, attn_cat, attn_i, 1);
|
||||
}
|
||||
// attn_cat: [Nh*D, N]
|
||||
|
||||
// Batched O proj
|
||||
struct ggml_tensor * attn_out = qwen3_linear(ctx, ly->o_proj, attn_cat);
|
||||
hidden = ggml_add(ctx, hidden, attn_out);
|
||||
|
||||
// Batched FFN
|
||||
norm = qwen3_rms_norm(ctx, hidden, ly->post_attn_layernorm, c.rms_norm_eps);
|
||||
struct ggml_tensor * mlp = qwen3_build_mlp(ctx, ly, norm, N);
|
||||
hidden = ggml_add(ctx, hidden, mlp);
|
||||
}
|
||||
|
||||
// Final norm + LM head: [V, N]
|
||||
hidden = qwen3_rms_norm(ctx, hidden, m->final_norm, c.rms_norm_eps);
|
||||
struct ggml_tensor * lgt = ggml_mul_mat(ctx, m->embed_tokens, hidden);
|
||||
ggml_set_name(lgt, "logits");
|
||||
ggml_set_output(lgt);
|
||||
ggml_build_forward_expand(gf, lgt);
|
||||
|
||||
// Schedule + allocate
|
||||
ggml_backend_sched_alloc_graph(m->sched, gf);
|
||||
|
||||
// CPU-side embedding dequant
|
||||
{
|
||||
const int64_t row_size = (int64_t)ggml_row_size(m->embed_type, H);
|
||||
const ggml_to_float_t to_float = ggml_get_type_traits(m->embed_type)->to_float;
|
||||
std::vector<float> embed_buf((size_t)H * N);
|
||||
for (int i = 0; i < N; i++) {
|
||||
const void * row = (const char *)m->embed_mmap_data + (int64_t)token_ids[i] * row_size;
|
||||
to_float(row, embed_buf.data() + (int64_t)i * H, H);
|
||||
}
|
||||
ggml_backend_tensor_set(embed_out, embed_buf.data(), 0, (size_t)H * N * sizeof(float));
|
||||
}
|
||||
|
||||
// Positions: per-element kv_pos
|
||||
{
|
||||
std::vector<int> pos_data(N);
|
||||
for (int i = 0; i < N; i++)
|
||||
pos_data[i] = m->kv_pos[kv_sets[i]];
|
||||
ggml_backend_tensor_set(positions, pos_data.data(), 0, N * sizeof(int));
|
||||
}
|
||||
|
||||
// Compute
|
||||
ggml_backend_sched_graph_compute(m->sched, gf);
|
||||
|
||||
// Read logits [V, N]
|
||||
ggml_backend_tensor_get(lgt, logits, 0, (size_t)c.vocab_size * N * sizeof(float));
|
||||
|
||||
// Advance all KV positions
|
||||
for (int i = 0; i < N; i++)
|
||||
m->kv_pos[kv_sets[i]]++;
|
||||
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(ctx);
|
||||
}
|
||||
|
||||
// Free all resources
|
||||
static void qw3lm_free(Qwen3LM * m) {
|
||||
if (m->sched) ggml_backend_sched_free(m->sched);
|
||||
if (m->kv_buf) ggml_backend_buffer_free(m->kv_buf);
|
||||
if (m->kv_ctx) ggml_free(m->kv_ctx);
|
||||
if (m->backend && m->backend != m->cpu_backend) ggml_backend_free(m->backend);
|
||||
if (m->cpu_backend) ggml_backend_free(m->cpu_backend);
|
||||
wctx_free(&m->wctx);
|
||||
gf_close(&m->gf_mmap);
|
||||
*m = {};
|
||||
}
|
||||
384
otherarch/acestep/qwen3.h
Normal file
384
otherarch/acestep/qwen3.h
Normal file
|
|
@ -0,0 +1,384 @@
|
|||
// qwen3.h: Qwen3 transformer encoder via ggml
|
||||
//
|
||||
// Generic Qwen3 backbone used by:
|
||||
// Text encoder (Qwen3-Embedding-0.6B): 28L, H=1024, causal, vocab lookup
|
||||
// Lyric encoder (ACEStep cond): 8L, H=2048, bidir, linear(1024->2048)
|
||||
// Timbre encoder (ACEStep cond): 4L, H=2048, bidir, linear(64->2048)
|
||||
//
|
||||
// Architecture per layer:
|
||||
// RMSNorm -> Q/K/V proj -> QK-Norm -> RoPE -> GQA -> O proj -> +residual
|
||||
// RMSNorm -> gate/up proj -> SwiGLU -> down proj -> +residual
|
||||
// Final: RMSNorm
|
||||
|
||||
#pragma once
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
#include "backend.h"
|
||||
#include "gguf_weights.h"
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#define QWEN3_MAX_LAYERS 32
|
||||
|
||||
// Config
|
||||
struct Qwen3Config {
|
||||
int hidden_size; // H
|
||||
int intermediate_size; // FFN inner dim
|
||||
int n_heads; // Nh (query heads)
|
||||
int n_kv_heads; // Nkv (key/value heads, for GQA)
|
||||
int head_dim; // D = H / Nh
|
||||
int n_layers;
|
||||
float rope_theta;
|
||||
float rms_norm_eps;
|
||||
bool is_causal; // true for text encoder, false for lyric/timbre
|
||||
};
|
||||
|
||||
// Per-layer weights
|
||||
struct Qwen3Layer {
|
||||
struct ggml_tensor * input_layernorm; // [H]
|
||||
struct ggml_tensor * post_attn_layernorm; // [H]
|
||||
struct ggml_tensor * q_proj; // [H, Nh*D] ggml = [Nh*D, H] PyTorch
|
||||
struct ggml_tensor * k_proj; // [H, Nkv*D]
|
||||
struct ggml_tensor * v_proj; // [H, Nkv*D]
|
||||
struct ggml_tensor * o_proj; // [Nh*D, H]
|
||||
struct ggml_tensor * q_norm; // [D]
|
||||
struct ggml_tensor * k_norm; // [D]
|
||||
struct ggml_tensor * gate_proj; // [H, FFN]
|
||||
struct ggml_tensor * up_proj; // [H, FFN]
|
||||
struct ggml_tensor * down_proj; // [FFN, H]
|
||||
};
|
||||
|
||||
// Standalone model (text encoder)
|
||||
struct Qwen3GGML {
|
||||
Qwen3Config cfg;
|
||||
Qwen3Layer layers[QWEN3_MAX_LAYERS];
|
||||
|
||||
// Embedding: vocab lookup table [H, vocab_size] in ggml
|
||||
struct ggml_tensor * embed_tokens; // [H, V]
|
||||
struct ggml_tensor * final_norm; // [H]
|
||||
|
||||
// Backend
|
||||
ggml_backend_t backend;
|
||||
ggml_backend_t cpu_backend;
|
||||
ggml_backend_sched_t sched;
|
||||
WeightCtx wctx;
|
||||
};
|
||||
|
||||
// Helpers (pure graph ops, no side effects)
|
||||
static struct ggml_tensor * qwen3_f32(struct ggml_context * ctx, struct ggml_tensor * t) {
|
||||
if (t->type == GGML_TYPE_F32) return t;
|
||||
return ggml_cast(ctx, t, GGML_TYPE_F32);
|
||||
}
|
||||
|
||||
static struct ggml_tensor * qwen3_linear(struct ggml_context * ctx,
|
||||
struct ggml_tensor * w,
|
||||
struct ggml_tensor * x) {
|
||||
return ggml_mul_mat(ctx, w, x);
|
||||
}
|
||||
|
||||
static struct ggml_tensor * qwen3_linear_bias(struct ggml_context * ctx,
|
||||
struct ggml_tensor * w,
|
||||
struct ggml_tensor * b,
|
||||
struct ggml_tensor * x) {
|
||||
struct ggml_tensor * out = ggml_mul_mat(ctx, w, x);
|
||||
return ggml_add(ctx, out, qwen3_f32(ctx, b));
|
||||
}
|
||||
|
||||
static struct ggml_tensor * qwen3_rms_norm(struct ggml_context * ctx,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * w,
|
||||
float eps) {
|
||||
struct ggml_tensor * n = ggml_rms_norm(ctx, x, eps);
|
||||
return ggml_mul(ctx, n, qwen3_f32(ctx, w));
|
||||
}
|
||||
|
||||
// Graph builders
|
||||
// These build sub-graphs and return output tensors.
|
||||
// They operate on ggml layout: [H, S] for hidden states.
|
||||
|
||||
// Self-attention: norm_in [H, S] -> attn_out [H, S]
|
||||
static struct ggml_tensor * qwen3_build_self_attn(
|
||||
struct ggml_context * ctx,
|
||||
const Qwen3Config & c,
|
||||
Qwen3Layer * ly,
|
||||
struct ggml_tensor * x, // [H, S]
|
||||
struct ggml_tensor * positions, // [S] int32
|
||||
struct ggml_tensor * mask, // [S, S] or NULL
|
||||
int S) {
|
||||
|
||||
int D = c.head_dim;
|
||||
int Nh = c.n_heads;
|
||||
int Nkv = c.n_kv_heads;
|
||||
|
||||
// 1) Q/K/V projections
|
||||
struct ggml_tensor * q = qwen3_linear(ctx, ly->q_proj, x); // [Nh*D, S]
|
||||
struct ggml_tensor * k = qwen3_linear(ctx, ly->k_proj, x); // [Nkv*D, S]
|
||||
struct ggml_tensor * v = qwen3_linear(ctx, ly->v_proj, x); // [Nkv*D, S]
|
||||
|
||||
// 2) Reshape to heads: [X*D, S] -> [D, X, S]
|
||||
q = ggml_reshape_3d(ctx, q, D, Nh, S);
|
||||
k = ggml_reshape_3d(ctx, k, D, Nkv, S);
|
||||
v = ggml_reshape_3d(ctx, v, D, Nkv, S);
|
||||
|
||||
// 3) QK-Norm: per-head RMSNorm on D
|
||||
q = ggml_rms_norm(ctx, q, c.rms_norm_eps);
|
||||
q = ggml_mul(ctx, q, qwen3_f32(ctx, ly->q_norm));
|
||||
k = ggml_rms_norm(ctx, k, c.rms_norm_eps);
|
||||
k = ggml_mul(ctx, k, qwen3_f32(ctx, ly->k_norm));
|
||||
|
||||
// 4) RoPE
|
||||
// ggml pitfall: mode=2 (NEOX half-split [i, i+D/2]), NOT mode=0 (consecutive [2i, 2i+1])
|
||||
// Python ref: rope_batch_kernel pairs ptr[d] with ptr[d+half] = NEOX
|
||||
q = ggml_rope_ext(ctx, q, positions, NULL,
|
||||
D, 2, 0, c.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx, k, positions, NULL,
|
||||
D, 2, 0, c.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// 5) Permute for flash_attn_ext: [D, X, S] -> [D, S, X]
|
||||
q = ggml_permute(ctx, q, 0, 2, 1, 3);
|
||||
k = ggml_permute(ctx, k, 0, 2, 1, 3);
|
||||
v = ggml_permute(ctx, v, 0, 2, 1, 3);
|
||||
|
||||
// 6) Flash attention (handles GQA)
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
struct ggml_tensor * attn = ggml_flash_attn_ext(ctx, q, k, v, mask, scale, 0.0f, 0.0f);
|
||||
ggml_flash_attn_ext_set_prec(attn, GGML_PREC_F32); // F32 accumulation
|
||||
|
||||
// 7) Reshape back: [D, Nh, S] -> [Nh*D, S]
|
||||
attn = ggml_reshape_2d(ctx, attn, Nh * D, S);
|
||||
|
||||
// 8) O projection
|
||||
return qwen3_linear(ctx, ly->o_proj, attn);
|
||||
}
|
||||
|
||||
// MLP: SwiGLU
|
||||
static struct ggml_tensor * qwen3_build_mlp(
|
||||
struct ggml_context * ctx,
|
||||
Qwen3Layer * ly,
|
||||
struct ggml_tensor * x, // [H, S]
|
||||
int S) {
|
||||
(void)S;
|
||||
struct ggml_tensor * gate = qwen3_linear(ctx, ly->gate_proj, x);
|
||||
struct ggml_tensor * up = qwen3_linear(ctx, ly->up_proj, x);
|
||||
struct ggml_tensor * ff = ggml_swiglu_split(ctx, gate, up);
|
||||
return qwen3_linear(ctx, ly->down_proj, ff);
|
||||
}
|
||||
|
||||
// Single layer: input [H, S] -> output [H, S]
|
||||
static struct ggml_tensor * qwen3_build_layer(
|
||||
struct ggml_context * ctx,
|
||||
const Qwen3Config & c,
|
||||
Qwen3Layer * ly,
|
||||
struct ggml_tensor * hidden,
|
||||
struct ggml_tensor * positions,
|
||||
struct ggml_tensor * mask,
|
||||
int S) {
|
||||
|
||||
// Self-attention block
|
||||
struct ggml_tensor * norm = qwen3_rms_norm(ctx, hidden, ly->input_layernorm, c.rms_norm_eps);
|
||||
struct ggml_tensor * attn = qwen3_build_self_attn(ctx, c, ly, norm, positions, mask, S);
|
||||
hidden = ggml_add(ctx, hidden, attn);
|
||||
|
||||
// MLP block
|
||||
norm = qwen3_rms_norm(ctx, hidden, ly->post_attn_layernorm, c.rms_norm_eps);
|
||||
struct ggml_tensor * mlp = qwen3_build_mlp(ctx, ly, norm, S);
|
||||
hidden = ggml_add(ctx, hidden, mlp);
|
||||
|
||||
return hidden;
|
||||
}
|
||||
|
||||
// Full N-layer stack: input [H, S] -> output [H, S] (post final-norm)
|
||||
static struct ggml_tensor * qwen3_build_layers(
|
||||
struct ggml_context * ctx,
|
||||
const Qwen3Config & c,
|
||||
Qwen3Layer * layers,
|
||||
struct ggml_tensor * final_norm_w,
|
||||
struct ggml_tensor * hidden,
|
||||
struct ggml_tensor * positions,
|
||||
struct ggml_tensor * mask,
|
||||
int S) {
|
||||
|
||||
for (int i = 0; i < c.n_layers; i++) {
|
||||
hidden = qwen3_build_layer(ctx, c, &layers[i], hidden, positions, mask, S);
|
||||
}
|
||||
return qwen3_rms_norm(ctx, hidden, final_norm_w, c.rms_norm_eps);
|
||||
}
|
||||
|
||||
// Loading
|
||||
static void qwen3_load_layer(WeightCtx * wctx, const GGUFModel & gf,
|
||||
Qwen3Layer * ly, const std::string & prefix) {
|
||||
ly->input_layernorm = gf_load_tensor_f32(wctx, gf, prefix + ".input_layernorm.weight");
|
||||
ly->post_attn_layernorm = gf_load_tensor_f32(wctx, gf, prefix + ".post_attention_layernorm.weight");
|
||||
ly->q_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.q_proj.weight");
|
||||
ly->k_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.k_proj.weight");
|
||||
ly->v_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.v_proj.weight");
|
||||
ly->o_proj = gf_load_tensor(wctx, gf, prefix + ".self_attn.o_proj.weight");
|
||||
ly->q_norm = gf_load_tensor_f32(wctx, gf, prefix + ".self_attn.q_norm.weight");
|
||||
ly->k_norm = gf_load_tensor_f32(wctx, gf, prefix + ".self_attn.k_norm.weight");
|
||||
ly->gate_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.gate_proj.weight");
|
||||
ly->up_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.up_proj.weight");
|
||||
ly->down_proj = gf_load_tensor(wctx, gf, prefix + ".mlp.down_proj.weight");
|
||||
}
|
||||
|
||||
// Backend init
|
||||
static void qwen3_init_backend(Qwen3GGML * m) {
|
||||
BackendPair bp = backend_init("TextEncoder");
|
||||
m->backend = bp.backend;
|
||||
m->cpu_backend = bp.cpu_backend;
|
||||
m->sched = backend_sched_new(bp, 4096);
|
||||
}
|
||||
|
||||
// Load standalone text encoder (Qwen3-Embedding) from GGUF
|
||||
// gguf_path: path to the .gguf file
|
||||
static bool qwen3_load_text_encoder(Qwen3GGML * m, const char * gguf_path) {
|
||||
m->cfg = {
|
||||
/*hidden_size*/ 1024,
|
||||
/*intermediate_size*/ 3072,
|
||||
/*n_heads*/ 16,
|
||||
/*n_kv_heads*/ 8,
|
||||
/*head_dim*/ 128,
|
||||
/*n_layers*/ 28,
|
||||
/*rope_theta*/ 1000000.0f,
|
||||
/*rms_norm_eps*/ 1e-6f,
|
||||
/*is_causal*/ true,
|
||||
};
|
||||
|
||||
GGUFModel gf;
|
||||
if (!gf_load(&gf, gguf_path)) {
|
||||
fprintf(stderr, "[Load] FATAL: cannot load %s\n", gguf_path);
|
||||
return false;
|
||||
}
|
||||
|
||||
// embed(1) + 28 layers * 11 weights + final_norm(1) = 310
|
||||
int n_tensors = 1 + m->cfg.n_layers * 11 + 1;
|
||||
wctx_init(&m->wctx, n_tensors);
|
||||
|
||||
m->embed_tokens = gf_load_tensor(&m->wctx, gf, "embed_tokens.weight");
|
||||
m->final_norm = gf_load_tensor_f32(&m->wctx, gf, "norm.weight");
|
||||
|
||||
for (int i = 0; i < m->cfg.n_layers; i++) {
|
||||
char prefix[64];
|
||||
snprintf(prefix, sizeof(prefix), "layers.%d", i);
|
||||
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix);
|
||||
}
|
||||
|
||||
if (!wctx_alloc(&m->wctx, m->backend)) {
|
||||
gf_close(&gf);
|
||||
return false;
|
||||
}
|
||||
gf_close(&gf);
|
||||
|
||||
fprintf(stderr, "[Load] TextEncoder: %dL, H=%d, Nh=%d/%d\n",
|
||||
m->cfg.n_layers, m->cfg.hidden_size, m->cfg.n_heads, m->cfg.n_kv_heads);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Forward: token IDs -> hidden states
|
||||
// token_ids: [S] int32 (CPU)
|
||||
// output: [H * S] float (CPU, caller-allocated)
|
||||
// Returns hidden states in ggml layout: ne[0]=H contiguous, S rows.
|
||||
static void qwen3_forward(Qwen3GGML * m, const int * token_ids, int S, float * output) {
|
||||
const Qwen3Config & c = m->cfg;
|
||||
int H = c.hidden_size;
|
||||
|
||||
// Graph context
|
||||
size_t ctx_size = 2048 * ggml_tensor_overhead() + ggml_graph_overhead();
|
||||
struct ggml_init_params gp = { ctx_size, NULL, true };
|
||||
struct ggml_context * ctx = ggml_init(gp);
|
||||
|
||||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 4096, false);
|
||||
|
||||
// Input: token IDs [S]
|
||||
struct ggml_tensor * t_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, S);
|
||||
ggml_set_name(t_ids, "token_ids");
|
||||
ggml_set_input(t_ids);
|
||||
|
||||
// Embedding lookup: [H, V] x [S] -> [H, S]
|
||||
struct ggml_tensor * hidden = ggml_get_rows(ctx, m->embed_tokens, t_ids);
|
||||
|
||||
// Positions: [S] int32 (0, 1, 2, ...)
|
||||
struct ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, S);
|
||||
ggml_set_name(positions, "positions");
|
||||
ggml_set_input(positions);
|
||||
|
||||
// Attention mask: causal uses fp16 mask, bidirectional uses NULL (no mask)
|
||||
struct ggml_tensor * mask = NULL;
|
||||
if (c.is_causal) {
|
||||
mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, S, S);
|
||||
ggml_set_name(mask, "attn_mask");
|
||||
ggml_set_input(mask);
|
||||
}
|
||||
|
||||
// N layers + final norm
|
||||
struct ggml_tensor * out = qwen3_build_layers(ctx, c, m->layers, m->final_norm,
|
||||
hidden, positions, mask, S);
|
||||
ggml_set_name(out, "output");
|
||||
ggml_set_output(out);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
|
||||
// Allocate
|
||||
ggml_backend_sched_alloc_graph(m->sched, gf);
|
||||
|
||||
// Set inputs
|
||||
ggml_backend_tensor_set(t_ids, token_ids, 0, S * sizeof(int));
|
||||
|
||||
{
|
||||
std::vector<int> pos_data(S);
|
||||
for (int i = 0; i < S; i++) pos_data[i] = i;
|
||||
ggml_backend_tensor_set(positions, pos_data.data(), 0, S * sizeof(int));
|
||||
}
|
||||
|
||||
if (c.is_causal) {
|
||||
std::vector<uint16_t> mask_data(S * S);
|
||||
for (int i = 0; i < S; i++)
|
||||
for (int j = 0; j < S; j++) {
|
||||
float v = (j <= i) ? 0.0f : -INFINITY;
|
||||
mask_data[i * S + j] = ggml_fp32_to_fp16(v);
|
||||
}
|
||||
ggml_backend_tensor_set(mask, mask_data.data(), 0, S * S * sizeof(uint16_t));
|
||||
}
|
||||
|
||||
// Compute
|
||||
ggml_backend_sched_graph_compute(m->sched, gf);
|
||||
|
||||
// Read output [H, S]
|
||||
ggml_backend_tensor_get(out, output, 0, H * S * sizeof(float));
|
||||
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(ctx);
|
||||
}
|
||||
|
||||
// CPU vocab lookup utility
|
||||
// For lyric embedding: look up token IDs in text encoder's embed table (bf16 -> f32)
|
||||
// GGUF keeps mmapped data alive. Output: [H, S] float (H contiguous per token).
|
||||
//
|
||||
// embed_data: pointer to bf16 weight data [vocab, H] in PyTorch layout (H contiguous per row)
|
||||
// token_ids: [S] int32
|
||||
// output: [H * S] float (ggml layout: H contiguous, S tokens)
|
||||
static void qwen3_cpu_embed_lookup(const void * embed_data, int H,
|
||||
const int * token_ids, int S,
|
||||
float * output) {
|
||||
const uint16_t * bf16 = (const uint16_t *)embed_data;
|
||||
for (int s = 0; s < S; s++) {
|
||||
int tok = token_ids[s];
|
||||
const uint16_t * row = bf16 + (int64_t)tok * H;
|
||||
float * dst = output + (int64_t)s * H;
|
||||
for (int h = 0; h < H; h++) {
|
||||
// bf16 to f32: shift left 16 bits
|
||||
uint32_t bits = (uint32_t)row[h] << 16;
|
||||
memcpy(&dst[h], &bits, 4);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Free
|
||||
static void qwen3_free(Qwen3GGML * m) {
|
||||
if (m->sched) ggml_backend_sched_free(m->sched);
|
||||
if (m->backend && m->backend != m->cpu_backend) ggml_backend_free(m->backend);
|
||||
if (m->cpu_backend) ggml_backend_free(m->cpu_backend);
|
||||
wctx_free(&m->wctx);
|
||||
*m = {};
|
||||
}
|
||||
305
otherarch/acestep/request.cpp
Normal file
305
otherarch/acestep/request.cpp
Normal file
|
|
@ -0,0 +1,305 @@
|
|||
//
|
||||
// request.cpp - AceStep request JSON read/write
|
||||
//
|
||||
// Flat JSON only (no nested objects). Handles string escaping for lyrics etc.
|
||||
//
|
||||
|
||||
#include "./request.h"
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// Defaults (aligned with Python GenerationParams)
|
||||
void request_init(AceRequest * r) {
|
||||
r->caption = "";
|
||||
r->lyrics = "";
|
||||
r->instrumental = false;
|
||||
r->bpm = 0;
|
||||
r->duration = -1.0f;
|
||||
r->keyscale = "";
|
||||
r->timesignature = "";
|
||||
r->vocal_language = "unknown";
|
||||
r->task_type = "text2music";
|
||||
r->seed = -1;
|
||||
r->thinking = true;
|
||||
r->lm_temperature = 0.85f;
|
||||
r->lm_cfg_scale = 2.0f;
|
||||
r->lm_top_p = 0.9f;
|
||||
r->lm_negative_prompt = "NO USER INPUT";
|
||||
r->audio_codes = "";
|
||||
r->inference_steps = 8;
|
||||
r->guidance_scale = 1.0f;
|
||||
r->shift = 3.0f;
|
||||
}
|
||||
|
||||
// JSON string escape / unescape
|
||||
static std::string json_escape(const std::string & s) {
|
||||
std::string out;
|
||||
out.reserve(s.size() + 16);
|
||||
for (char c : s) {
|
||||
switch (c) {
|
||||
case '"': out += "\\\""; break;
|
||||
case '\\': out += "\\\\"; break;
|
||||
case '\n': out += "\\n"; break;
|
||||
case '\r': out += "\\r"; break;
|
||||
case '\t': out += "\\t"; break;
|
||||
default:
|
||||
if ((unsigned char)c < 0x20) {
|
||||
char buf[8];
|
||||
snprintf(buf, sizeof(buf), "\\u%04x", (unsigned char)c);
|
||||
out += buf;
|
||||
} else {
|
||||
out += c;
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
static std::string json_unescape(const char * s, size_t len) {
|
||||
std::string out;
|
||||
out.reserve(len);
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
if (s[i] == '\\' && i + 1 < len) {
|
||||
switch (s[++i]) {
|
||||
case '"': out += '"'; break;
|
||||
case '\\': out += '\\'; break;
|
||||
case '/': out += '/'; break;
|
||||
case 'n': out += '\n'; break;
|
||||
case 'r': out += '\r'; break;
|
||||
case 't': out += '\t'; break;
|
||||
case 'u':
|
||||
// \\uXXXX: parse 4 hex digits, emit as UTF-8 (ASCII subset only)
|
||||
if (i + 4 < len) {
|
||||
char hex[5] = {s[i+1], s[i+2], s[i+3], s[i+4], 0};
|
||||
unsigned cp = (unsigned)strtoul(hex, nullptr, 16);
|
||||
i += 4;
|
||||
if (cp < 0x80) {
|
||||
out += (char)cp;
|
||||
} else if (cp < 0x800) {
|
||||
out += (char)(0xC0 | (cp >> 6));
|
||||
out += (char)(0x80 | (cp & 0x3F));
|
||||
} else {
|
||||
out += (char)(0xE0 | (cp >> 12));
|
||||
out += (char)(0x80 | ((cp >> 6) & 0x3F));
|
||||
out += (char)(0x80 | (cp & 0x3F));
|
||||
}
|
||||
}
|
||||
break;
|
||||
default: out += s[i]; break;
|
||||
}
|
||||
} else {
|
||||
out += s[i];
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// Minimal flat JSON parser
|
||||
struct JsonPair {
|
||||
std::string key;
|
||||
std::string value; // raw value (unquoted strings are unescaped, numbers/bools as-is)
|
||||
bool is_string;
|
||||
};
|
||||
|
||||
static const char * skip_ws(const char * p) {
|
||||
while (*p == ' ' || *p == '\t' || *p == '\n' || *p == '\r') p++;
|
||||
return p;
|
||||
}
|
||||
|
||||
// Parse a JSON string starting at p (must point to opening '"').
|
||||
// Returns pointer past closing '"', or nullptr on error.
|
||||
static const char * parse_string(const char * p, std::string * out) {
|
||||
if (*p != '"') return nullptr;
|
||||
p++;
|
||||
const char * start = p;
|
||||
while (*p && *p != '"') {
|
||||
if (*p == '\\') {
|
||||
p++;
|
||||
if (!*p) return nullptr;
|
||||
}
|
||||
p++;
|
||||
}
|
||||
if (*p != '"') return nullptr;
|
||||
*out = json_unescape(start, (size_t)(p - start));
|
||||
return p + 1;
|
||||
}
|
||||
|
||||
// Parse a JSON value (string, number, bool, null).
|
||||
// Skips arrays/objects by bracket matching (for forward compat).
|
||||
static const char * parse_value(const char * p, std::string * out, bool * is_str) {
|
||||
*is_str = false;
|
||||
if (*p == '"') {
|
||||
*is_str = true;
|
||||
return parse_string(p, out);
|
||||
}
|
||||
if (*p == '[' || *p == '{') {
|
||||
// skip nested structure (not used, but don't choke)
|
||||
char open = *p, close = (*p == '[') ? ']' : '}';
|
||||
int depth = 1;
|
||||
const char * start = p;
|
||||
p++;
|
||||
while (*p && depth > 0) {
|
||||
if (*p == open) depth++;
|
||||
else if (*p == close) depth--;
|
||||
else if (*p == '"') {
|
||||
// skip strings inside nested structure
|
||||
std::string dummy;
|
||||
p = parse_string(p, &dummy);
|
||||
if (!p) return nullptr;
|
||||
continue;
|
||||
}
|
||||
p++;
|
||||
}
|
||||
*out = std::string(start, (size_t)(p - start));
|
||||
return p;
|
||||
}
|
||||
// number, bool, null
|
||||
const char * start = p;
|
||||
while (*p && *p != ',' && *p != '}' && *p != ']' &&
|
||||
*p != ' ' && *p != '\t' && *p != '\n' && *p != '\r') p++;
|
||||
*out = std::string(start, (size_t)(p - start));
|
||||
return p;
|
||||
}
|
||||
|
||||
static bool parse_json_flat(const char * json, std::vector<JsonPair> * pairs) {
|
||||
const char * p = skip_ws(json);
|
||||
if (*p != '{') return false;
|
||||
p = skip_ws(p + 1);
|
||||
|
||||
while (*p && *p != '}') {
|
||||
JsonPair kv;
|
||||
p = parse_string(p, &kv.key);
|
||||
if (!p) return false;
|
||||
p = skip_ws(p);
|
||||
if (*p != ':') return false;
|
||||
p = skip_ws(p + 1);
|
||||
p = parse_value(p, &kv.value, &kv.is_string);
|
||||
if (!p) return false;
|
||||
pairs->push_back(kv);
|
||||
p = skip_ws(p);
|
||||
if (*p == ',') p = skip_ws(p + 1);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// File I/O helpers
|
||||
static std::string read_file(const char * path) {
|
||||
FILE * f = fopen(path, "rb");
|
||||
if (!f) return "";
|
||||
fseek(f, 0, SEEK_END);
|
||||
long sz = ftell(f);
|
||||
fseek(f, 0, SEEK_SET);
|
||||
std::string buf((size_t)sz, '\0');
|
||||
size_t nr = fread(&buf[0], 1, (size_t)sz, f);
|
||||
fclose(f);
|
||||
if ((long)nr != sz) buf.resize(nr);
|
||||
return buf;
|
||||
}
|
||||
|
||||
// Public API
|
||||
bool request_parse(AceRequest * r, const char * path) {
|
||||
request_init(r);
|
||||
|
||||
std::string json = read_file(path);
|
||||
if (json.empty()) {
|
||||
fprintf(stderr, "[Request] ERROR: cannot read %s\n", path);
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<JsonPair> pairs;
|
||||
if (!parse_json_flat(json.c_str(), &pairs)) {
|
||||
fprintf(stderr, "[Request] ERROR: malformed JSON in %s\n", path);
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const auto & kv : pairs) {
|
||||
const std::string & k = kv.key;
|
||||
const std::string & v = kv.value;
|
||||
|
||||
// strings
|
||||
if (k == "caption") r->caption = v;
|
||||
else if (k == "lyrics") r->lyrics = v;
|
||||
else if (k == "keyscale") r->keyscale = v;
|
||||
else if (k == "timesignature") r->timesignature = v;
|
||||
else if (k == "vocal_language") r->vocal_language = v;
|
||||
else if (k == "task_type") r->task_type = v;
|
||||
else if (k == "audio_codes") r->audio_codes = v;
|
||||
else if (k == "lm_negative_prompt") r->lm_negative_prompt = v;
|
||||
|
||||
// ints
|
||||
else if (k == "bpm") r->bpm = atoi(v.c_str());
|
||||
else if (k == "seed") r->seed = atoi(v.c_str());
|
||||
|
||||
// floats
|
||||
else if (k == "duration") r->duration = (float)atof(v.c_str());
|
||||
else if (k == "lm_temperature") r->lm_temperature = (float)atof(v.c_str());
|
||||
else if (k == "lm_cfg_scale") r->lm_cfg_scale = (float)atof(v.c_str());
|
||||
else if (k == "lm_top_p") r->lm_top_p = (float)atof(v.c_str());
|
||||
else if (k == "inference_steps") r->inference_steps = atoi(v.c_str());
|
||||
else if (k == "guidance_scale") r->guidance_scale = (float)atof(v.c_str());
|
||||
else if (k == "shift") r->shift = (float)atof(v.c_str());
|
||||
|
||||
// bools
|
||||
else if (k == "thinking") r->thinking = (v == "true");
|
||||
else if (k == "instrumental") r->instrumental = (v == "true");
|
||||
// unknown keys: silently ignored (forward compat)
|
||||
}
|
||||
|
||||
fprintf(stderr, "[Request] parsed %s (%zu fields)\n", path, pairs.size());
|
||||
return true;
|
||||
}
|
||||
|
||||
bool request_write(const AceRequest * r, const char * path) {
|
||||
FILE * f = fopen(path, "w");
|
||||
if (!f) {
|
||||
fprintf(stderr, "[Request] ERROR: cannot write %s\n", path);
|
||||
return false;
|
||||
}
|
||||
|
||||
fprintf(f, "{\n");
|
||||
fprintf(f, " \"caption\": \"%s\",\n", json_escape(r->caption).c_str());
|
||||
fprintf(f, " \"lyrics\": \"%s\",\n", json_escape(r->lyrics).c_str());
|
||||
if (r->instrumental)
|
||||
fprintf(f, " \"instrumental\": true,\n");
|
||||
fprintf(f, " \"bpm\": %d,\n", r->bpm);
|
||||
fprintf(f, " \"duration\": %.1f,\n", r->duration);
|
||||
fprintf(f, " \"keyscale\": \"%s\",\n", json_escape(r->keyscale).c_str());
|
||||
fprintf(f, " \"timesignature\": \"%s\",\n", json_escape(r->timesignature).c_str());
|
||||
fprintf(f, " \"vocal_language\": \"%s\",\n", json_escape(r->vocal_language).c_str());
|
||||
fprintf(f, " \"task_type\": \"%s\",\n", json_escape(r->task_type).c_str());
|
||||
fprintf(f, " \"seed\": %d,\n", r->seed);
|
||||
fprintf(f, " \"thinking\": %s,\n", r->thinking ? "true" : "false");
|
||||
fprintf(f, " \"lm_temperature\": %.2f,\n", r->lm_temperature);
|
||||
fprintf(f, " \"lm_cfg_scale\": %.1f,\n", r->lm_cfg_scale);
|
||||
fprintf(f, " \"lm_top_p\": %.2f,\n", r->lm_top_p);
|
||||
fprintf(f, " \"lm_negative_prompt\": \"%s\",\n", json_escape(r->lm_negative_prompt).c_str());
|
||||
fprintf(f, " \"inference_steps\": %d,\n", r->inference_steps);
|
||||
fprintf(f, " \"guidance_scale\": %.1f,\n", r->guidance_scale);
|
||||
fprintf(f, " \"shift\": %.1f,\n", r->shift);
|
||||
// audio_codes last (no trailing comma)
|
||||
fprintf(f, " \"audio_codes\": \"%s\"\n", json_escape(r->audio_codes).c_str());
|
||||
fprintf(f, "}\n");
|
||||
|
||||
fclose(f);
|
||||
fprintf(stderr, "[Request] wrote %s\n", path);
|
||||
return true;
|
||||
}
|
||||
|
||||
void request_dump(const AceRequest * r, FILE * f) {
|
||||
fprintf(f, "[Request] task=%s thinking=%s seed=%d\n",
|
||||
r->task_type.c_str(), r->thinking ? "true" : "false", r->seed);
|
||||
fprintf(f, " caption: %.60s%s\n",
|
||||
r->caption.c_str(), r->caption.size() > 60 ? "..." : "");
|
||||
fprintf(f, " lyrics: %zu bytes\n", r->lyrics.size());
|
||||
fprintf(f, " bpm=%d dur=%.0f key=%s ts=%s lang=%s\n",
|
||||
r->bpm, r->duration, r->keyscale.c_str(),
|
||||
r->timesignature.c_str(), r->vocal_language.c_str());
|
||||
fprintf(f, " lm: temp=%.2f cfg=%.1f top_p=%.2f\n",
|
||||
r->lm_temperature, r->lm_cfg_scale, r->lm_top_p);
|
||||
fprintf(f, " dit: steps=%d guidance=%.1f shift=%.1f\n",
|
||||
r->inference_steps, r->guidance_scale, r->shift);
|
||||
fprintf(f, " audio_codes: %s\n",
|
||||
r->audio_codes.empty() ? "(none)" : "(present)");
|
||||
}
|
||||
57
otherarch/acestep/request.h
Normal file
57
otherarch/acestep/request.h
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
#pragma once
|
||||
//
|
||||
// request.h - AceStep generation request (JSON serialization)
|
||||
//
|
||||
// Pure data container + JSON read/write. Zero business logic.
|
||||
// Aligned with Python GenerationParams (inference.py:39) and API /release_task.
|
||||
//
|
||||
|
||||
#include <string>
|
||||
#include <cstdio>
|
||||
|
||||
struct AceRequest {
|
||||
// text content
|
||||
std::string caption; // ""
|
||||
std::string lyrics; // ""
|
||||
bool instrumental; // false
|
||||
|
||||
// metadata (user-provided or LLM-enriched)
|
||||
int bpm; // 0 = unset
|
||||
float duration; // -1 = unset
|
||||
std::string keyscale; // "" = unset
|
||||
std::string timesignature; // "" = unset
|
||||
std::string vocal_language; // "unknown"
|
||||
|
||||
// generation
|
||||
std::string task_type; // "text2music"
|
||||
int seed; // -1 = random
|
||||
|
||||
// LM control
|
||||
bool thinking; // true
|
||||
float lm_temperature; // 0.85
|
||||
float lm_cfg_scale; // 2.0
|
||||
float lm_top_p; // 0.9
|
||||
std::string lm_negative_prompt; // "NO USER INPUT"
|
||||
|
||||
// codes (Python-compatible string: "3101,11837,27514,...")
|
||||
// empty = text2music (silence context), non-empty = cover mode
|
||||
std::string audio_codes; // ""
|
||||
|
||||
// DiT control (Python: inference_steps, guidance_scale, shift)
|
||||
int inference_steps; // 8
|
||||
float guidance_scale; // 7.0
|
||||
float shift; // 1.0
|
||||
};
|
||||
|
||||
// Initialize all fields to defaults (matches Python GenerationParams defaults)
|
||||
void request_init(AceRequest * r);
|
||||
|
||||
// Parse JSON file into struct. Missing fields keep their defaults.
|
||||
// Returns false on file error or malformed JSON.
|
||||
bool request_parse(AceRequest * r, const char * path);
|
||||
|
||||
// Write struct to JSON file (overwrites). Returns false on file error.
|
||||
bool request_write(const AceRequest * r, const char * path);
|
||||
|
||||
// Dump human-readable summary to stream (debug)
|
||||
void request_dump(const AceRequest * r, FILE * f);
|
||||
221
otherarch/acestep/tokenizer.h
Normal file
221
otherarch/acestep/tokenizer.h
Normal file
|
|
@ -0,0 +1,221 @@
|
|||
// tokenizer.h: FSQ + Detokenizer via ggml
|
||||
//
|
||||
// Decodes LM audio codes into context_latents for DiT:
|
||||
// codes [T_5Hz] -> FSQ decode -> [T_5Hz, 6] -> project_out -> [T_5Hz, 2048]
|
||||
// -> detokenizer (per token): embed + special_tokens broadcast + 2L encoder + proj_out
|
||||
// -> [T_25Hz, 64] context_latents (T_25Hz = T_5Hz * 5)
|
||||
//
|
||||
// Weights live in the DiT GGUF (prefix "tokenizer." and "detokenizer.")
|
||||
// Detokenizer reuses Qwen3 encoder infrastructure from qwen3.h
|
||||
|
||||
#pragma once
|
||||
#include "qwen3.h"
|
||||
|
||||
// FSQ constants
|
||||
static const int FSQ_NDIMS = 6;
|
||||
static const int FSQ_LEVELS[6] = {8, 8, 8, 5, 5, 5};
|
||||
|
||||
// FSQ decode: integer index -> 6 normalized float values
|
||||
// Each dimension: level_idx / ((L-1)/2) - 1.0 (maps to [-1, 1])
|
||||
static void fsq_decode_index(int index, float * out) {
|
||||
int stride = 1;
|
||||
for (int d = 0; d < FSQ_NDIMS; d++) {
|
||||
int L = FSQ_LEVELS[d];
|
||||
int level_idx = (index / stride) % L;
|
||||
float half_L = (float)(L - 1) / 2.0f;
|
||||
out[d] = (float)level_idx / half_L - 1.0f;
|
||||
stride *= L;
|
||||
}
|
||||
}
|
||||
|
||||
// Detokenizer config: same Qwen3 arch as lyric/timbre encoders, 2 layers
|
||||
static Qwen3Config detok_config() {
|
||||
return {
|
||||
/*hidden_size*/ 2048,
|
||||
/*intermediate_size*/ 6144,
|
||||
/*n_heads*/ 16,
|
||||
/*n_kv_heads*/ 8,
|
||||
/*head_dim*/ 128,
|
||||
/*n_layers*/ 2,
|
||||
/*rope_theta*/ 1000000.0f,
|
||||
/*rms_norm_eps*/ 1e-6f,
|
||||
/*is_causal*/ false,
|
||||
};
|
||||
}
|
||||
|
||||
struct DetokGGML {
|
||||
// FSQ project_out: Linear(6, 2048) + bias
|
||||
struct ggml_tensor * fsq_proj_w; // [2048, 6]
|
||||
struct ggml_tensor * fsq_proj_b; // [2048]
|
||||
|
||||
// Detokenizer
|
||||
struct ggml_tensor * embed_w; // [2048, 2048]
|
||||
struct ggml_tensor * embed_b; // [2048]
|
||||
struct ggml_tensor * special_tok; // [2048, 5] (broadcast positional)
|
||||
Qwen3Config cfg;
|
||||
Qwen3Layer layers[2];
|
||||
struct ggml_tensor * norm; // [2048]
|
||||
struct ggml_tensor * proj_out_w; // [64, 2048]
|
||||
struct ggml_tensor * proj_out_b; // [64]
|
||||
|
||||
ggml_backend_t backend;
|
||||
ggml_backend_t cpu_backend;
|
||||
ggml_backend_sched_t sched;
|
||||
WeightCtx wctx;
|
||||
};
|
||||
|
||||
// Load from DiT GGUF
|
||||
static bool detok_ggml_load(DetokGGML * m, const char * gguf_path,
|
||||
ggml_backend_t backend, ggml_backend_t cpu_backend) {
|
||||
m->cfg = detok_config();
|
||||
m->backend = backend;
|
||||
m->cpu_backend = cpu_backend;
|
||||
|
||||
GGUFModel gf;
|
||||
if (!gf_load(&gf, gguf_path)) {
|
||||
fprintf(stderr, "[Load] FATAL: cannot load %s\n", gguf_path);
|
||||
return false;
|
||||
}
|
||||
|
||||
// FSQ(2) + embed(2) + special(1) + 2 layers x 11(22) + norm(1) + proj(2) = 30
|
||||
wctx_init(&m->wctx, 30);
|
||||
|
||||
m->fsq_proj_w = gf_load_tensor(&m->wctx, gf, "tokenizer.quantizer.project_out.weight");
|
||||
m->fsq_proj_b = gf_load_tensor(&m->wctx, gf, "tokenizer.quantizer.project_out.bias");
|
||||
|
||||
m->embed_w = gf_load_tensor(&m->wctx, gf, "detokenizer.embed_tokens.weight");
|
||||
m->embed_b = gf_load_tensor(&m->wctx, gf, "detokenizer.embed_tokens.bias");
|
||||
m->norm = gf_load_tensor(&m->wctx, gf, "detokenizer.norm.weight");
|
||||
m->proj_out_w = gf_load_tensor(&m->wctx, gf, "detokenizer.proj_out.weight");
|
||||
m->proj_out_b = gf_load_tensor(&m->wctx, gf, "detokenizer.proj_out.bias");
|
||||
|
||||
// special_tokens: GGUF [2048, 5, 1] (ggml order), reshape to [2048, 5]
|
||||
m->special_tok = gf_load_tensor(&m->wctx, gf, "detokenizer.special_tokens");
|
||||
|
||||
for (int i = 0; i < m->cfg.n_layers; i++) {
|
||||
char prefix[128];
|
||||
snprintf(prefix, sizeof(prefix), "detokenizer.layers.%d", i);
|
||||
qwen3_load_layer(&m->wctx, gf, &m->layers[i], prefix);
|
||||
}
|
||||
|
||||
if (!wctx_alloc(&m->wctx, backend)) {
|
||||
gf_close(&gf);
|
||||
return false;
|
||||
}
|
||||
gf_close(&gf);
|
||||
|
||||
// Scheduler
|
||||
ggml_backend_t backends[2] = { backend, cpu_backend };
|
||||
int n = (backend == cpu_backend) ? 1 : 2;
|
||||
m->sched = ggml_backend_sched_new(backends, NULL, n, 4096, false, true);
|
||||
|
||||
fprintf(stderr, "[Load] Detokenizer: FSQ(6->2048) + 2L encoder(S=5, 2048->64)\n");
|
||||
return true;
|
||||
}
|
||||
|
||||
// Decode LM audio codes -> context_latents
|
||||
// codes: [T_5Hz] integer array
|
||||
// context_out: [64 * T_25Hz] flat, caller allocates (T_25Hz = T_5Hz * 5)
|
||||
// ggml layout [64, T_25Hz]: element (c, t) = data[t * 64 + c]
|
||||
static int detok_ggml_decode(DetokGGML * m, const int * codes, int T_5Hz,
|
||||
float * context_out) {
|
||||
int T_25Hz = T_5Hz * 5;
|
||||
int H = 2048;
|
||||
int P = 5; // pool window: each 5Hz token -> 5 frames at 25Hz
|
||||
|
||||
// Step 1: FSQ decode all indices on CPU -> [T_5Hz, 6]
|
||||
std::vector<float> fsq_decoded(T_5Hz * FSQ_NDIMS);
|
||||
for (int g = 0; g < T_5Hz; g++)
|
||||
fsq_decode_index(codes[g], fsq_decoded.data() + g * FSQ_NDIMS);
|
||||
|
||||
// Step 2: build ggml graph for one token
|
||||
// input [6] -> project_out [2048] -> embed_tokens [2048]
|
||||
// -> broadcast + special_tokens [2048, 5] -> 2L encoder -> norm -> proj_out [64, 5]
|
||||
size_t ctx_size = ggml_tensor_overhead() * 512 + ggml_graph_overhead_custom(4096, false);
|
||||
std::vector<uint8_t> ctx_buf(ctx_size);
|
||||
struct ggml_init_params p = { ctx_size, ctx_buf.data(), true };
|
||||
struct ggml_context * ctx = ggml_init(p);
|
||||
|
||||
// Input: one FSQ-decoded vector [6]
|
||||
// ggml pitfall: [6] is ne[0]=6, matches project_out weight [2048, 6]
|
||||
struct ggml_tensor * fsq_in = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, FSQ_NDIMS);
|
||||
ggml_set_name(fsq_in, "fsq_in");
|
||||
ggml_set_input(fsq_in);
|
||||
|
||||
// project_out: [6] -> [2048]
|
||||
struct ggml_tensor * quantized = ggml_mul_mat(ctx, m->fsq_proj_w, fsq_in);
|
||||
quantized = ggml_add(ctx, quantized, qwen3_f32(ctx, m->fsq_proj_b));
|
||||
|
||||
// embed_tokens: [2048] -> [2048]
|
||||
struct ggml_tensor * embedded = ggml_mul_mat(ctx, m->embed_w,
|
||||
ggml_reshape_2d(ctx, quantized, H, 1));
|
||||
embedded = ggml_add(ctx, embedded, qwen3_f32(ctx, ggml_reshape_2d(ctx, m->embed_b, H, 1)));
|
||||
|
||||
// Broadcast [2048, 1] -> [2048, 5] + special_tokens [2048, 5]
|
||||
// ggml pitfall: special_tokens loaded as BF16, cast to F32 for add
|
||||
struct ggml_tensor * special_2d = ggml_reshape_2d(ctx, m->special_tok, H, P);
|
||||
struct ggml_tensor * special_f32 = qwen3_f32(ctx, special_2d);
|
||||
struct ggml_tensor * hidden = ggml_add(ctx, ggml_repeat(ctx, embedded, special_f32),
|
||||
special_f32);
|
||||
|
||||
// Position indices for RoPE: [0, 1, 2, 3, 4]
|
||||
struct ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, P);
|
||||
ggml_set_name(positions, "detok_pos");
|
||||
ggml_set_input(positions);
|
||||
|
||||
// 2L encoder + norm (non-causal, no mask needed at S=5)
|
||||
hidden = qwen3_build_layers(ctx, m->cfg, m->layers, m->norm,
|
||||
hidden, positions, NULL, P);
|
||||
|
||||
// proj_out: [2048, 5] -> [64, 5]
|
||||
struct ggml_tensor * output = ggml_mul_mat(ctx, m->proj_out_w, hidden);
|
||||
output = ggml_add(ctx, output, qwen3_f32(ctx, ggml_reshape_2d(ctx, m->proj_out_b, 64, 1)));
|
||||
ggml_set_name(output, "detok_out");
|
||||
ggml_set_output(output);
|
||||
|
||||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 4096, false);
|
||||
ggml_build_forward_expand(gf, output);
|
||||
|
||||
if (!ggml_backend_sched_alloc_graph(m->sched, gf)) {
|
||||
fprintf(stderr, "[Context] FATAL: graph alloc failed\n");
|
||||
ggml_free(ctx);
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Set positions once (constant: 0,1,2,3,4)
|
||||
int pos_data[5] = {0, 1, 2, 3, 4};
|
||||
ggml_backend_tensor_set(ggml_graph_get_tensor(gf, "detok_pos"),
|
||||
pos_data, 0, P * sizeof(int));
|
||||
|
||||
struct ggml_tensor * t_in = ggml_graph_get_tensor(gf, "fsq_in");
|
||||
struct ggml_tensor * t_out = ggml_graph_get_tensor(gf, "detok_out");
|
||||
|
||||
// Step 3: loop over T_5Hz tokens
|
||||
struct ggml_tensor * t_pos = ggml_graph_get_tensor(gf, "detok_pos");
|
||||
for (int g = 0; g < T_5Hz; g++) {
|
||||
// Re-set positions every iteration (allocator may share buffer with intermediates)
|
||||
ggml_backend_tensor_set(t_pos, pos_data, 0, P * sizeof(int));
|
||||
ggml_backend_tensor_set(t_in, fsq_decoded.data() + g * FSQ_NDIMS,
|
||||
0, FSQ_NDIMS * sizeof(float));
|
||||
ggml_backend_sched_graph_compute(m->sched, gf);
|
||||
|
||||
// output [64, 5]: 5 frames of 64 channels
|
||||
// context_out layout: [64, T_25Hz], frame t at offset t*64
|
||||
ggml_backend_tensor_get(t_out, context_out + g * P * 64,
|
||||
0, P * 64 * sizeof(float));
|
||||
}
|
||||
|
||||
fprintf(stderr, "[Context] Decoded: %d codes -> %d frames (%.1fs @ 25Hz)\n",
|
||||
T_5Hz, T_25Hz, (float)T_25Hz / 25.0f);
|
||||
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(ctx);
|
||||
return T_25Hz;
|
||||
}
|
||||
|
||||
// Free
|
||||
static void detok_ggml_free(DetokGGML * m) {
|
||||
if (m->sched) ggml_backend_sched_free(m->sched);
|
||||
wctx_free(&m->wctx);
|
||||
*m = {};
|
||||
}
|
||||
520
otherarch/acestep/vae.h
Normal file
520
otherarch/acestep/vae.h
Normal file
|
|
@ -0,0 +1,520 @@
|
|||
// vae.h: AutoencoderOobleck decoder (audio VAE) via ggml
|
||||
//
|
||||
// Architecture: conv1(64->2048,k=7) -> 5xblock(snake+convT+3xresunit) -> snake+conv2(128->2,k=7)
|
||||
// ResUnit(ch, dil): skip=x -> snake->conv(k=7,dil)->snake->conv(k=1)->+skip
|
||||
// Snake: x + sin^2(e^a * x) / e^b
|
||||
// Weight norm fused at load: w = g*v/||v||
|
||||
// Upsample: 10x6x4x4x2 = 1920x
|
||||
|
||||
#pragma once
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
#include "gguf_weights.h"
|
||||
#include "backend.h"
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// Structs
|
||||
struct VAEResUnit {
|
||||
struct ggml_tensor * s1a, * s1b; // snake1 exp(alpha), exp(beta) [1, C]
|
||||
struct ggml_tensor * c1w, * c1b; // conv1 fused [7, C, C], bias [C]
|
||||
struct ggml_tensor * s2a, * s2b; // snake2
|
||||
struct ggml_tensor * c2w, * c2b; // conv2 fused [1, C, C], bias [C]
|
||||
int dilation;
|
||||
};
|
||||
|
||||
struct VAEBlock {
|
||||
struct ggml_tensor * sa, * sb; // snake exp(a/b) [1, in_ch]
|
||||
struct ggml_tensor * ctw, * ctb; // conv_transpose fused [K, out_ch, in_ch], bias [out_ch]
|
||||
int in_ch, out_ch, stride, kernel;
|
||||
VAEResUnit ru[3];
|
||||
};
|
||||
|
||||
struct VAEGGML {
|
||||
struct ggml_tensor * c1w, * c1b; // conv1 [7, 64, 2048], bias [2048]
|
||||
VAEBlock blk[5];
|
||||
struct ggml_tensor * sa, * sb; // final snake [1, 128]
|
||||
struct ggml_tensor * c2w; // conv2 [7, 128, 2] (no bias)
|
||||
|
||||
ggml_backend_t backend;
|
||||
ggml_backend_t cpu_backend;
|
||||
ggml_backend_sched_t sched;
|
||||
ggml_backend_buffer_t buf;
|
||||
struct ggml_context * weight_ctx; // holds weight tensor metadata
|
||||
|
||||
// Graph cache for tiled decode (avoids rebuild per tile)
|
||||
struct ggml_context * graph_ctx;
|
||||
uint8_t * graph_buf; // heap-allocated backing for graph_ctx
|
||||
struct ggml_cgraph * graph;
|
||||
struct ggml_tensor * graph_input;
|
||||
struct ggml_tensor * graph_output;
|
||||
int graph_T; // cached T_latent (0 = no cache)
|
||||
|
||||
// Scratch buffer (reused across tiles, grown as needed)
|
||||
std::vector<float> scratch_in; // transposed input [64 * T]
|
||||
};
|
||||
|
||||
// Load helpers
|
||||
// Fuse weight_norm: w = g*v/||v||, write f32 into pre-allocated ggml_tensor
|
||||
// Works for both Conv1d [OC,IC,K] and ConvTranspose1d [IC,OC,K]:
|
||||
// weight_norm normalizes over dim=0 (shape[0]), regardless of semantics.
|
||||
static void vae_fuse_wn(struct ggml_tensor * dst, const GGUFModel & gf, const std::string & pfx) {
|
||||
struct ggml_tensor * mv = ggml_get_tensor(gf.meta, (pfx + ".weight_v").c_str());
|
||||
const uint16_t * g = (const uint16_t *)gf_get_data(gf, (pfx + ".weight_g").c_str());
|
||||
const uint16_t * v = (const uint16_t *)gf_get_data(gf, (pfx + ".weight_v").c_str());
|
||||
// PyTorch dim0 is ggml ne[n_dims-1]
|
||||
int n_dims = ggml_n_dims(mv);
|
||||
int dim0 = (int)mv->ne[n_dims - 1];
|
||||
int fan = (int)(ggml_nelements(mv) / dim0);
|
||||
std::vector<float> w(dim0 * fan);
|
||||
for (int d = 0; d < dim0; d++) {
|
||||
float gv = ggml_bf16_to_fp32(*(const ggml_bf16_t *)&g[d]);
|
||||
float nsq = 0;
|
||||
for (int i = 0; i < fan; i++) {
|
||||
float vv = ggml_bf16_to_fp32(*(const ggml_bf16_t *)&v[d * fan + i]);
|
||||
nsq += vv * vv;
|
||||
}
|
||||
float s = gv / (sqrtf(nsq) + 1e-12f);
|
||||
for (int i = 0; i < fan; i++) {
|
||||
float vv = ggml_bf16_to_fp32(*(const ggml_bf16_t *)&v[d * fan + i]);
|
||||
w[d * fan + i] = vv * s;
|
||||
}
|
||||
}
|
||||
if (dst->type == GGML_TYPE_F16) {
|
||||
std::vector<ggml_fp16_t> w16(w.size());
|
||||
ggml_fp32_to_fp16_row(w.data(), w16.data(), (int)w.size());
|
||||
ggml_backend_tensor_set(dst, w16.data(), 0, w16.size() * sizeof(ggml_fp16_t));
|
||||
} else {
|
||||
ggml_backend_tensor_set(dst, w.data(), 0, w.size() * sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
// Load bf16 snake param [1,C,1] -> exp -> f32 [1, C]
|
||||
static void vae_load_snake(struct ggml_tensor * dst, const GGUFModel & gf, const std::string & name) {
|
||||
struct ggml_tensor * mt = ggml_get_tensor(gf.meta, name.c_str());
|
||||
int C = (int)mt->ne[1]; // PyTorch [1,C,1] -> ggml ne=[1,C,1], middle dim
|
||||
const uint16_t * raw = (const uint16_t *)gf_get_data(gf, name.c_str());
|
||||
std::vector<float> d(C);
|
||||
for (int i = 0; i < C; i++)
|
||||
d[i] = expf(ggml_bf16_to_fp32(*(const ggml_bf16_t *)&raw[i]));
|
||||
ggml_backend_tensor_set(dst, d.data(), 0, C * sizeof(float));
|
||||
}
|
||||
|
||||
// Load bf16 snake param [1,C,1] -> 1/exp -> f32 [1, C] (reciprocal for mul fusion)
|
||||
static void vae_load_snake_inv(struct ggml_tensor * dst, const GGUFModel & gf, const std::string & name) {
|
||||
struct ggml_tensor * mt = ggml_get_tensor(gf.meta, name.c_str());
|
||||
int C = (int)mt->ne[1];
|
||||
const uint16_t * raw = (const uint16_t *)gf_get_data(gf, name.c_str());
|
||||
std::vector<float> d(C);
|
||||
for (int i = 0; i < C; i++)
|
||||
d[i] = 1.0f / expf(ggml_bf16_to_fp32(*(const ggml_bf16_t *)&raw[i]));
|
||||
ggml_backend_tensor_set(dst, d.data(), 0, C * sizeof(float));
|
||||
}
|
||||
|
||||
// Load bf16 bias [C] -> f32
|
||||
static void vae_load_bias(struct ggml_tensor * dst, const GGUFModel & gf, const std::string & name) {
|
||||
struct ggml_tensor * mt = ggml_get_tensor(gf.meta, name.c_str());
|
||||
int C = (int)mt->ne[0]; // 1D: ne[0] = C
|
||||
const uint16_t * raw = (const uint16_t *)gf_get_data(gf, name.c_str());
|
||||
std::vector<float> d(C);
|
||||
for (int i = 0; i < C; i++)
|
||||
d[i] = ggml_bf16_to_fp32(*(const ggml_bf16_t *)&raw[i]);
|
||||
ggml_backend_tensor_set(dst, d.data(), 0, C * sizeof(float));
|
||||
}
|
||||
|
||||
// Load model
|
||||
static void vae_ggml_load(VAEGGML * m, const char * path) {
|
||||
GGUFModel gf = {};
|
||||
if (!gf_load(&gf, path)) {
|
||||
fprintf(stderr, "[VAE] FATAL: cannot load %s\n", path);
|
||||
exit(1);
|
||||
}
|
||||
|
||||
static const int strides[] = {10, 6, 4, 4, 2};
|
||||
static const int in_ch[] = {2048, 1024, 512, 256, 128};
|
||||
static const int out_ch[] = {1024, 512, 256, 128, 128};
|
||||
static const int dilations[] = {1, 3, 9};
|
||||
|
||||
// Phase 1: create tensor metadata (no_alloc context)
|
||||
size_t ctx_size = ggml_tensor_overhead() * 200;
|
||||
struct ggml_init_params p = { ctx_size, NULL, true };
|
||||
m->weight_ctx = ggml_init(p);
|
||||
struct ggml_context * ctx = m->weight_ctx;
|
||||
|
||||
m->c1w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 7, 64, 2048);
|
||||
m->c1b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048);
|
||||
|
||||
for (int i = 0; i < 5; i++) {
|
||||
VAEBlock & b = m->blk[i];
|
||||
b.in_ch = in_ch[i];
|
||||
b.out_ch = out_ch[i];
|
||||
b.stride = strides[i];
|
||||
b.kernel = strides[i] * 2;
|
||||
int C = out_ch[i];
|
||||
b.sa = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, in_ch[i]);
|
||||
b.sb = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, in_ch[i]);
|
||||
b.ctw = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, b.kernel, out_ch[i], in_ch[i]);
|
||||
b.ctb = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_ch[i]);
|
||||
for (int r = 0; r < 3; r++) {
|
||||
VAEResUnit & ru = b.ru[r];
|
||||
ru.dilation = dilations[r];
|
||||
ru.s1a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
|
||||
ru.s1b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
|
||||
ru.c1w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 7, C, C);
|
||||
ru.c1b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
|
||||
ru.s2a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
|
||||
ru.s2b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
|
||||
ru.c2w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 1, C, C);
|
||||
ru.c2b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
|
||||
}
|
||||
}
|
||||
m->sa = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 128);
|
||||
m->sb = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 128);
|
||||
m->c2w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 7, 128, 2);
|
||||
|
||||
// Phase 2: allocate backend buffer (im2col grid Y fix enables long-sequence conv1d)
|
||||
BackendPair bp = backend_init("VAE");
|
||||
m->backend = bp.backend;
|
||||
m->cpu_backend = bp.cpu_backend;
|
||||
m->sched = backend_sched_new(bp, 8192);
|
||||
m->buf = ggml_backend_alloc_ctx_tensors(ctx, m->backend);
|
||||
fprintf(stderr, "[VAE] Backend: %s, Weight buffer: %.1f MB\n",
|
||||
ggml_backend_name(m->backend),
|
||||
(float)ggml_backend_buffer_get_size(m->buf) / (1024 * 1024));
|
||||
|
||||
// Phase 3: load & fuse weights
|
||||
vae_fuse_wn(m->c1w, gf, "decoder.conv1");
|
||||
vae_load_bias(m->c1b, gf, "decoder.conv1.bias");
|
||||
|
||||
for (int i = 0; i < 5; i++) {
|
||||
VAEBlock & b = m->blk[i];
|
||||
std::string blk_pfx = "decoder.block." + std::to_string(i);
|
||||
vae_load_snake(b.sa, gf, blk_pfx + ".snake1.alpha");
|
||||
vae_load_snake_inv(b.sb, gf, blk_pfx + ".snake1.beta");
|
||||
vae_fuse_wn(b.ctw, gf, blk_pfx + ".conv_t1");
|
||||
vae_load_bias(b.ctb, gf, blk_pfx + ".conv_t1.bias");
|
||||
for (int r = 0; r < 3; r++) {
|
||||
VAEResUnit & ru = b.ru[r];
|
||||
std::string rp = blk_pfx + ".res_unit" + std::to_string(r + 1);
|
||||
vae_load_snake(ru.s1a, gf, rp + ".snake1.alpha");
|
||||
vae_load_snake_inv(ru.s1b, gf, rp + ".snake1.beta");
|
||||
vae_fuse_wn(ru.c1w, gf, rp + ".conv1");
|
||||
vae_load_bias(ru.c1b, gf, rp + ".conv1.bias");
|
||||
vae_load_snake(ru.s2a, gf, rp + ".snake2.alpha");
|
||||
vae_load_snake_inv(ru.s2b, gf, rp + ".snake2.beta");
|
||||
vae_fuse_wn(ru.c2w, gf, rp + ".conv2");
|
||||
vae_load_bias(ru.c2b, gf, rp + ".conv2.bias");
|
||||
}
|
||||
}
|
||||
vae_load_snake(m->sa, gf, "decoder.snake1.alpha");
|
||||
vae_load_snake_inv(m->sb, gf, "decoder.snake1.beta");
|
||||
vae_fuse_wn(m->c2w, gf, "decoder.conv2");
|
||||
|
||||
fprintf(stderr, "[VAE] Loaded: 5 blocks, upsample=1920x\n");
|
||||
gf_close(&gf);
|
||||
}
|
||||
|
||||
// Graph building
|
||||
// Snake activation: x + sin^2(exp_a * x) * inv_b
|
||||
// x: [T, C], exp_a: [1, C], inv_b: [1, C] (pre-computed 1/exp(b))
|
||||
static struct ggml_tensor * vae_snake(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * exp_a,
|
||||
struct ggml_tensor * inv_b) {
|
||||
struct ggml_tensor * ax = ggml_mul(ctx, x, exp_a); // [T, C] (broadcast 1->T)
|
||||
struct ggml_tensor * s = ggml_sin(ctx, ax); // sin(e^a * x)
|
||||
struct ggml_tensor * s2 = ggml_sqr(ctx, s); // sin^2
|
||||
struct ggml_tensor * d = ggml_mul(ctx, s2, inv_b); // * 1/e^b
|
||||
return ggml_add(ctx, x, d); // x + ...
|
||||
}
|
||||
|
||||
// Conv1d + bias: data [T, IC] -> [T_out, OC]
|
||||
static struct ggml_tensor * vae_conv1d(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * w, // [K, IC, OC] (F16, pre-cast at load)
|
||||
struct ggml_tensor * b, // [OC] or NULL
|
||||
struct ggml_tensor * x, // [T, IC]
|
||||
int stride, int padding, int dilation) {
|
||||
struct ggml_tensor * y = ggml_conv_1d(ctx, w, x, stride, padding, dilation);
|
||||
// ggml_conv_1d returns [OL, OC, N=1], squeeze to 2d
|
||||
y = ggml_reshape_2d(ctx, y, y->ne[0], y->ne[1]);
|
||||
if (b) {
|
||||
// bias [OC] -> [1, OC] for broadcast over OL dimension
|
||||
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
|
||||
y = ggml_add(ctx, y, b2d);
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// ConvTranspose1d + bias
|
||||
static struct ggml_tensor * vae_conv_t1d(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * w, // [K, out_ch, in_ch]
|
||||
struct ggml_tensor * b, // [out_ch] or NULL
|
||||
struct ggml_tensor * x, // [T, in_ch]
|
||||
int stride, int padding) {
|
||||
// ggml_conv_transpose_1d asserts p0==0, so we crop manually
|
||||
struct ggml_tensor * y = ggml_conv_transpose_1d(ctx, w, x, stride, 0, 1);
|
||||
// y is 4d: [OL, OC, 1, 1], squeeze to 2d
|
||||
y = ggml_reshape_2d(ctx, y, y->ne[0], y->ne[1]);
|
||||
// Crop padding from both sides: [OL, OC] -> [OL - 2*pad, OC]
|
||||
if (padding > 0) {
|
||||
y = ggml_view_2d(ctx, y, y->ne[0] - 2 * padding, y->ne[1],
|
||||
y->nb[1], padding * sizeof(float));
|
||||
}
|
||||
if (b) {
|
||||
struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
|
||||
y = ggml_add(ctx, y, b2d);
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// ResUnit forward
|
||||
static struct ggml_tensor * vae_res_unit(
|
||||
struct ggml_context * ctx,
|
||||
VAEResUnit * ru,
|
||||
struct ggml_tensor * x) { // [T, C]
|
||||
struct ggml_tensor * skip = x;
|
||||
|
||||
// snake1 -> dilated conv(k=7) -> snake2 -> conv(k=1)
|
||||
int pad = 3 * ru->dilation; // (k-1)*dil/2 = 3*dil
|
||||
x = vae_snake(ctx, x, ru->s1a, ru->s1b);
|
||||
x = vae_conv1d(ctx, ru->c1w, ru->c1b, x, 1, pad, ru->dilation);
|
||||
x = vae_snake(ctx, x, ru->s2a, ru->s2b);
|
||||
x = vae_conv1d(ctx, ru->c2w, ru->c2b, x, 1, 0, 1);
|
||||
|
||||
return ggml_add(ctx, skip, x);
|
||||
}
|
||||
|
||||
// Build full VAE decode graph
|
||||
// latent: [T_latent, 64] -> audio: [T_audio, 2]
|
||||
static struct ggml_tensor * vae_ggml_build_graph(
|
||||
struct ggml_context * ctx,
|
||||
VAEGGML * m,
|
||||
struct ggml_tensor * latent) { // [T, 64] input
|
||||
|
||||
// conv1: [T, 64] -> [T, 2048]
|
||||
struct ggml_tensor * x = vae_conv1d(ctx, m->c1w, m->c1b, latent, 1, 3, 1);
|
||||
|
||||
// 5 decoder blocks
|
||||
for (int i = 0; i < 5; i++) {
|
||||
VAEBlock & b = m->blk[i];
|
||||
// snake -> conv_transpose (upsample)
|
||||
x = vae_snake(ctx, x, b.sa, b.sb);
|
||||
int pad = (b.kernel - b.stride) / 2;
|
||||
x = vae_conv_t1d(ctx, b.ctw, b.ctb, x, b.stride, pad);
|
||||
// 3 res units
|
||||
for (int r = 0; r < 3; r++)
|
||||
x = vae_res_unit(ctx, &b.ru[r], x);
|
||||
}
|
||||
|
||||
// Final: snake -> conv2(128->2, k=7, pad=3)
|
||||
x = vae_snake(ctx, x, m->sa, m->sb);
|
||||
x = vae_conv1d(ctx, m->c2w, NULL, x, 1, 3, 1);
|
||||
|
||||
return x; // [T_audio, 2]
|
||||
}
|
||||
|
||||
// Core compute: ensure graph cached, set input, run. Returns T_audio or -1.
|
||||
// Output remains in m->graph_output for caller to read as needed.
|
||||
static int vae_ggml_compute(
|
||||
VAEGGML * m,
|
||||
const float * latent, // [T_full, 64] time-major
|
||||
int T_latent, // window length to decode
|
||||
int win_start = 0) { // offset into latent
|
||||
|
||||
// Build graph only when T_latent changes
|
||||
if (m->graph_T != T_latent) {
|
||||
if (m->graph_ctx) {
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(m->graph_ctx);
|
||||
free(m->graph_buf);
|
||||
}
|
||||
size_t ctx_size = ggml_tensor_overhead() * 1024 + ggml_graph_overhead_custom(8192, false);
|
||||
m->graph_buf = (uint8_t *)malloc(ctx_size);
|
||||
struct ggml_init_params p = { ctx_size, m->graph_buf, true };
|
||||
struct ggml_context * ctx = ggml_init(p);
|
||||
|
||||
m->graph_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_latent, 64);
|
||||
ggml_set_name(m->graph_input, "vae_input");
|
||||
ggml_set_input(m->graph_input);
|
||||
|
||||
m->graph_output = vae_ggml_build_graph(ctx, m, m->graph_input);
|
||||
ggml_set_name(m->graph_output, "vae_output");
|
||||
ggml_set_output(m->graph_output);
|
||||
|
||||
m->graph = ggml_new_graph_custom(ctx, 8192, false);
|
||||
ggml_build_forward_expand(m->graph, m->graph_output);
|
||||
|
||||
if (!ggml_backend_sched_alloc_graph(m->sched, m->graph)) {
|
||||
fprintf(stderr, "[VAE] FATAL: graph alloc failed\n");
|
||||
ggml_free(ctx);
|
||||
free(m->graph_buf);
|
||||
m->graph_ctx = NULL;
|
||||
m->graph_buf = NULL;
|
||||
m->graph_T = 0;
|
||||
return -1;
|
||||
}
|
||||
m->graph_ctx = ctx;
|
||||
m->graph_T = T_latent;
|
||||
fprintf(stderr, "[VAE] Graph: %d nodes, T_latent=%d\n",
|
||||
ggml_graph_n_nodes(m->graph), T_latent);
|
||||
}
|
||||
|
||||
// Extract window + transpose: [T, 64] time-major -> ggml [T, 64] channel-major
|
||||
size_t in_size = 64 * T_latent;
|
||||
if (m->scratch_in.size() < in_size)
|
||||
m->scratch_in.resize(in_size);
|
||||
for (int c = 0; c < 64; c++)
|
||||
for (int t = 0; t < T_latent; t++)
|
||||
m->scratch_in[c * T_latent + t] = latent[(win_start + t) * 64 + c];
|
||||
ggml_backend_tensor_set(m->graph_input,
|
||||
m->scratch_in.data(), 0, in_size * sizeof(float));
|
||||
|
||||
ggml_backend_sched_graph_compute(m->sched, m->graph);
|
||||
|
||||
return (int)m->graph_output->ne[0];
|
||||
}
|
||||
|
||||
// Decode API: latent [T_latent, 64] -> audio [2, T_audio] flat.
|
||||
// Returns T_audio (or -1 on error).
|
||||
static int vae_ggml_decode(
|
||||
VAEGGML * m,
|
||||
const float * latent,
|
||||
int T_latent,
|
||||
float * audio_out,
|
||||
int max_T_audio) {
|
||||
|
||||
int T_audio = T_latent * 1920;
|
||||
if (T_audio > max_T_audio) {
|
||||
fprintf(stderr, "[VAE] T_audio %d exceeds max %d\n", T_audio, max_T_audio);
|
||||
return -1;
|
||||
}
|
||||
|
||||
int T_out = vae_ggml_compute(m, latent, T_latent, 0);
|
||||
if (T_out < 0) return -1;
|
||||
|
||||
ggml_backend_tensor_get(m->graph_output, audio_out, 0, T_out * 2 * sizeof(float));
|
||||
|
||||
fprintf(stderr, "[VAE] Decoded: T_latent=%d -> T_audio=%d (%.2fs @ 48kHz)\n",
|
||||
T_latent, T_out, (float)T_out / 48000.0f);
|
||||
return T_out;
|
||||
}
|
||||
|
||||
// Tiled decode: overlap-discard chunking for bounded VRAM usage.
|
||||
// Matches Python handler.py tiled_decode / _tiled_decode_gpu:
|
||||
// stride = chunk_size - 2*overlap
|
||||
// For each tile: decode latent window with overlap context, trim to core, concatenate.
|
||||
// Default chunk=256, overlap=64 matches Python handler.py fallback defaults.
|
||||
// Python auto-tunes chunk by VRAM: >=24GB->512, >=16GB->384, >=12GB->256, <12GB->128.
|
||||
// Returns T_audio (total samples per channel) or -1 on error.
|
||||
static int vae_ggml_decode_tiled(
|
||||
VAEGGML * m,
|
||||
const float * latent, // [T_latent, 64] flat time-major (DiT output layout)
|
||||
int T_latent,
|
||||
float * audio_out, // [2, T_audio] flat (caller allocs)
|
||||
int max_T_audio,
|
||||
int chunk_size = 256,
|
||||
int overlap = 64) {
|
||||
|
||||
// Ensure positive stride (matches Python effective_overlap reduction)
|
||||
while (chunk_size - 2 * overlap <= 0 && overlap > 0)
|
||||
overlap /= 2;
|
||||
|
||||
// Short sequence: decode directly
|
||||
if (T_latent <= chunk_size)
|
||||
return vae_ggml_decode(m, latent, T_latent, audio_out, max_T_audio);
|
||||
|
||||
int stride = chunk_size - 2 * overlap;
|
||||
int num_steps = (T_latent + stride - 1) / stride;
|
||||
|
||||
fprintf(stderr, "[VAE] Tiled decode: %d tiles (chunk=%d, overlap=%d, stride=%d)\n",
|
||||
num_steps, chunk_size, overlap, stride);
|
||||
|
||||
float upsample_factor = 0.0f;
|
||||
int audio_write_pos = 0;
|
||||
|
||||
for (int i = 0; i < num_steps; i++) {
|
||||
// Core range in latent frames (the part we keep)
|
||||
int core_start = i * stride;
|
||||
int core_end = core_start + stride;
|
||||
if (core_end > T_latent) core_end = T_latent;
|
||||
|
||||
// Window range with overlap context
|
||||
int win_start = core_start - overlap;
|
||||
if (win_start < 0) win_start = 0;
|
||||
int win_end = core_end + overlap;
|
||||
if (win_end > T_latent) win_end = T_latent;
|
||||
int win_len = win_end - win_start;
|
||||
|
||||
// Compute tile (graph cached, extract+transpose fused)
|
||||
int tile_T = vae_ggml_compute(m, latent, win_len, win_start);
|
||||
if (tile_T < 0) {
|
||||
fprintf(stderr, "[VAE] FATAL: tile %d decode failed\n", i);
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Determine upsample factor from first tile
|
||||
if (i == 0) {
|
||||
upsample_factor = (float)tile_T / (float)win_len;
|
||||
fprintf(stderr, "[VAE] Upsample factor: %.2f (expected ~1920)\n", upsample_factor);
|
||||
}
|
||||
|
||||
// Compute trim in audio samples (matches Python int(round(...)))
|
||||
int added_start = core_start - win_start;
|
||||
int trim_start = (int)roundf((float)added_start * upsample_factor);
|
||||
int added_end = win_end - core_end;
|
||||
int trim_end = (int)roundf((float)added_end * upsample_factor);
|
||||
|
||||
int end_idx = (trim_end > 0) ? (tile_T - trim_end) : tile_T;
|
||||
int core_len = end_idx - trim_start;
|
||||
if (core_len <= 0) continue;
|
||||
|
||||
// Check output bounds
|
||||
if (audio_write_pos + core_len > max_T_audio) {
|
||||
fprintf(stderr, "[VAE] FATAL: tiled output exceeds max_T_audio\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Read trimmed ch0 and ch1 directly from backend tensor into final audio_out
|
||||
// Layout: [ch0: tile_T floats, ch1: tile_T floats]
|
||||
ggml_backend_tensor_get(m->graph_output,
|
||||
audio_out + audio_write_pos,
|
||||
trim_start * sizeof(float),
|
||||
core_len * sizeof(float));
|
||||
ggml_backend_tensor_get(m->graph_output,
|
||||
audio_out + max_T_audio + audio_write_pos,
|
||||
(tile_T + trim_start) * sizeof(float),
|
||||
core_len * sizeof(float));
|
||||
audio_write_pos += core_len;
|
||||
}
|
||||
|
||||
// Compact ch1 from offset max_T_audio to offset audio_write_pos
|
||||
memmove(audio_out + audio_write_pos,
|
||||
audio_out + max_T_audio,
|
||||
audio_write_pos * sizeof(float));
|
||||
|
||||
fprintf(stderr, "[VAE] Tiled decode done: %d tiles -> T_audio=%d (%.2fs @ 48kHz)\n",
|
||||
num_steps, audio_write_pos, (float)audio_write_pos / 48000.0f);
|
||||
|
||||
return audio_write_pos;
|
||||
}
|
||||
|
||||
// Free
|
||||
static void vae_ggml_free(VAEGGML * m) {
|
||||
if (m->graph_ctx) {
|
||||
ggml_backend_sched_reset(m->sched);
|
||||
ggml_free(m->graph_ctx);
|
||||
free(m->graph_buf);
|
||||
}
|
||||
if (m->sched) ggml_backend_sched_free(m->sched);
|
||||
if (m->buf) ggml_backend_buffer_free(m->buf);
|
||||
if (m->weight_ctx) ggml_free(m->weight_ctx);
|
||||
if (m->backend && m->backend != m->cpu_backend) ggml_backend_free(m->backend);
|
||||
if (m->cpu_backend) ggml_backend_free(m->cpu_backend);
|
||||
*m = {};
|
||||
}
|
||||
74
otherarch/acestep/weight_ctx.h
Normal file
74
otherarch/acestep/weight_ctx.h
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
#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 = <loader>_load_tensor(&wctx, source, "name");
|
||||
// wctx_alloc(&wctx, backend);
|
||||
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
#include <cstdio>
|
||||
#include <cstddef>
|
||||
#include <vector>
|
||||
|
||||
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<PendingCopy> pending;
|
||||
|
||||
// Staging buffers for type-converted data (kept alive until wctx_alloc)
|
||||
std::vector<std::vector<float>> 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;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue