From 8f23b8d81ee5f16ba5e213f7b8c0340c1b771636 Mon Sep 17 00:00:00 2001 From: Concedo <39025047+LostRuins@users.noreply.github.com> Date: Thu, 12 Mar 2026 23:46:10 +0800 Subject: [PATCH] wip on ref audio, but it compiles --- expose.h | 1 + koboldcpp.py | 6 +- otherarch/acestep/ace-qwen3.cpp | 4 +- otherarch/acestep/dit-vae.cpp | 154 ++++++++++--- otherarch/acestep/dit.h | 17 +- otherarch/acestep/request.cpp | 5 + otherarch/acestep/request.h | 3 + otherarch/acestep/vae-enc.h | 392 ++++++++++++++++++++++++++++++++ 8 files changed, 549 insertions(+), 33 deletions(-) create mode 100644 otherarch/acestep/vae-enc.h diff --git a/expose.h b/expose.h index 9d36fcd77..9b7347219 100644 --- a/expose.h +++ b/expose.h @@ -351,6 +351,7 @@ struct music_generation_inputs const bool gen_codes = false; const bool rewrite_caption = true; const char * input_json = nullptr; + const char * music_reference_audio_data = nullptr; }; struct music_generation_outputs { diff --git a/koboldcpp.py b/koboldcpp.py index 4ba4c61da..c5508d26e 100755 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -464,7 +464,8 @@ class music_generation_inputs(ctypes.Structure): ("use_mp3", ctypes.c_bool), ("gen_codes", ctypes.c_bool), ("rewrite_caption", ctypes.c_bool), - ("input_json", ctypes.c_char_p)] + ("input_json", ctypes.c_char_p), + ("music_reference_audio_data", ctypes.c_char_p)] class music_generation_outputs(ctypes.Structure): _fields_ = [("status", ctypes.c_int), @@ -2490,6 +2491,7 @@ def music_generate_codes(genparams): inputs.gen_codes = genparams.get('gen_codes', False) inputs.rewrite_caption = genparams.get('rewrite_caption', True) inputs.input_json = input_json.encode("UTF-8") + inputs.music_reference_audio_data = "".encode("UTF-8") ret = handle.music_generate(inputs) outstr = "" if ret.status==1: @@ -2507,6 +2509,8 @@ def music_generate_audio(genparams): inputs.gen_codes = genparams.get('gen_codes', False) inputs.rewrite_caption = genparams.get('rewrite_caption', True) inputs.input_json = input_json.encode("UTF-8") + refaudio = genparams.get('music_reference_audio_data', None) + inputs.music_reference_audio_data = (refaudio.encode("UTF-8") if (refaudio and refaudio!="") else "".encode("UTF-8")) ret = handle.music_generate(inputs) outstr = "" if ret.status==1: diff --git a/otherarch/acestep/ace-qwen3.cpp b/otherarch/acestep/ace-qwen3.cpp index d3d34aff2..4100b3499 100644 --- a/otherarch/acestep/ace-qwen3.cpp +++ b/otherarch/acestep/ace-qwen3.cpp @@ -528,10 +528,10 @@ struct MetadataFSM { language_name = bpe_encode(&bpe, "language:", false); timesig_name = bpe_encode(&bpe, "timesignature:", false); - // BPM 30-300 + // BPM 30-195 { std::vector vals; - for (int v = 30; v <= 300; v++) vals.push_back(std::to_string(v)); + for (int v = 30; v <= 195; v++) vals.push_back(std::to_string(v)); build_value_tree(bpe, bpm_tree, "bpm:", vals); } // Duration 40-450 diff --git a/otherarch/acestep/dit-vae.cpp b/otherarch/acestep/dit-vae.cpp index 0c9e6f98f..8ece153a2 100644 --- a/otherarch/acestep/dit-vae.cpp +++ b/otherarch/acestep/dit-vae.cpp @@ -21,6 +21,8 @@ #include "./bpe.h" #include "./debug.h" #include "./request.h" +#include "./vae-enc.h" +#include "otherarch/utils.h" // Minimal WAV writer (16-bit PCM stereo) static bool write_wav(const char * path, const float * audio, int T_audio, int sr) { @@ -580,6 +582,7 @@ static DiTGGMLConfig music_dit_cfg; static Timer music_dit_timer; static bool is_turbo = false; static VAEGGML vae = {}; +static VAEEncoder vae_enc = {}; static BPETokenizer music_tok; static Qwen3GGML music_text_enc = {}; static GGUFModel gf_te = {}; @@ -606,6 +609,7 @@ void unload_acestep_dit_others() if(acestep_dit_others_loaded) { acestep_dit_others_loaded = false; + vae_enc_free(&vae_enc); vae_ggml_free(&vae); gf_close(&gf_te); cond_ggml_free(&music_cond); @@ -672,6 +676,10 @@ bool load_acestep_dit(std::string music_embd_path, std::string music_dit_path, s vae_ggml_load(&vae, vae_gguf); fprintf(stderr, "[Load] VAE weights: %.1f ms\n", music_dit_timer.ms()); + music_dit_timer.reset(); + vae_enc_load(&vae_enc, vae_gguf); + fprintf(stderr, "[Load] VAE Enc weights: %.1f ms\n", music_dit_timer.ms()); + music_dit_timer.reset(); if (!load_bpe_from_gguf(&music_tok, text_enc_gguf)) { fprintf(stderr, "FATAL: failed to load music tokenizer from %s\n", text_enc_gguf); @@ -733,6 +741,50 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) } } + const int FRAMES_PER_SECOND = 25; + int Oc = music_dit_cfg.out_channels; // 64 + int ctx_ch = music_dit_cfg.in_channels - Oc; // 128 + int batch_n = 1; + int vae_chunk = 256; + int vae_overlap = 64; + + // Cover mode: load VAE encoder and encode source audio + bool have_cover = false; + std::vector cover_latents; // [T_cover, 64] time-major + int T_cover = 0; + std::string custom_reference_audio_str = inputs.music_reference_audio_data; + if (custom_reference_audio_str!="") { + music_dit_timer.reset(); + int T_audio = 0, wav_sr = 0; + + std::vector media_data_buffer = kcpp_base64_decode(custom_reference_audio_str); + std::vector custom_reference_audio_pcmf32; + bool ok = kcpp_decode_audio_from_buf(media_data_buffer.data(), media_data_buffer.size(), 48000, custom_reference_audio_pcmf32); + if (!ok) { + printf("\nError: Cannot read input audio file.\n"); + return ""; + } + + wav_sr = 48000; + T_audio = custom_reference_audio_pcmf32.size(); + float * wav_data = custom_reference_audio_pcmf32.data(); + + fprintf(stderr, "[Cover] Source audio: %.2fs\n", (float)T_audio / (float)(wav_sr > 0 ? wav_sr : 48000)); + int max_T_lat = (T_audio / 1920) + 64; + cover_latents.resize(max_T_lat * 64); + T_cover = vae_enc_encode_tiled(&vae_enc, wav_data, T_audio, + cover_latents.data(), max_T_lat, + vae_chunk, vae_overlap); + if (T_cover < 0) { + fprintf(stderr, "FATAL: VAE encode of src_audio failed\n"); + return ""; + } + cover_latents.resize(T_cover * 64); + fprintf(stderr, "[Cover] Encoded: T_cover=%d (%.2fs), %.1f ms\n", + T_cover, (float)T_cover * 1920.0f / 48000.0f, music_dit_timer.ms()); + have_cover = true; + } + // Parse request JSON AceRequest req; std::string injson = inputs.input_json; @@ -748,13 +800,6 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) req.inference_steps = (req.inference_steps>100?100:req.inference_steps); //clamp to 100 req.duration = (req.duration>420?420:req.duration); //clamp to 7 min - const int FRAMES_PER_SECOND = 25; - int Oc = music_dit_cfg.out_channels; // 64 - int ctx_ch = music_dit_cfg.in_channels - Oc; // 128 - int batch_n = 1; - int vae_chunk = 256; - int vae_overlap = 64; - // Extract params const char * caption = req.caption.c_str(); const char * lyrics = req.lyrics.empty() ? "[Instrumental]" : req.lyrics.c_str(); @@ -780,8 +825,6 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) { seed = (((uint32_t)time(NULL)) % 1000000u); } - fprintf(stderr, "[Pipeline] seed=%lld, steps=%d, guidance=%.1f, shift=%.1f, duration=%.1fs\n", - seed, num_steps, guidance_scale, shift, duration); // Parse audio codes from request std::vector codes_vec = parse_codes_string(req.audio_codes); @@ -797,16 +840,24 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) } // 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; + // Cover: from source audio. Codes: from code count. Else: from duration. + int T; + if (have_cover) { + T = T_cover; + // duration in metas must match actual source length, not JSON default + duration = (float)T_cover / (float)FRAMES_PER_SECOND; + } else if (!codes_vec.empty()) { + T = (int)codes_vec.size() * 5; + } else { + T = (int)(duration * FRAMES_PER_SECOND); + } T = ((T + music_dit_cfg.patch_size - 1) / music_dit_cfg.patch_size) * music_dit_cfg.patch_size; int S = T / music_dit_cfg.patch_size; int enc_S = 0; fprintf(stderr, "[Pipeline] T=%d, S=%d\n", T, S); + fprintf(stderr, "[Pipeline] seed=%lld, steps=%d, guidance=%.1f, shift=%.1f, duration=%.1fs\n", + seed, num_steps, guidance_scale, shift, duration); if (T > 15000) { fprintf(stderr, "ERROR: T=%d exceeds silence_latent max 15000, skipping\n", T); @@ -818,7 +869,8 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) music_dit_timer.reset(); // 2. Build formatted prompts - const char * instruction = "Generate audio semantic tokens based on the given conditions:"; + // Same instruction for all modes. Cover differs only by context content (audio vs silence). + const char * instruction = "Fill the audio semantic mask 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", @@ -873,10 +925,10 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) // std::vector silence(Oc * T); // memcpy(silence.data(), silence_full.data(), (size_t)(Oc * T) * sizeof(float)); - // Decode audio codes if provided + // Decode audio codes if provided (passthrough mode only, NOT cover) int decoded_T = 0; std::vector decoded_latents; - if (!codes_vec.empty()) { + if (!have_cover && !codes_vec.empty()) { int T_5Hz = (int)codes_vec.size(); int T_25Hz_codes = T_5Hz * 5; decoded_latents.resize(T_25Hz_codes * Oc); @@ -892,16 +944,31 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) decoded_T = T_25Hz_codes < T ? T_25Hz_codes : T; } - // Build single context: [T, ctx_ch] = src_latents[64] + mask_ones[64] + // Build context: [T, ctx_ch] = src_latents[64] + mask_ones[64] + // Cover: VAE latents directly (matching Python: is_covers=False, raw latents as context) + // Passthrough: detokenized FSQ codes + silence padding + // Text2music: silence only std::vector context_single(T * ctx_ch); - for (int t = 0; t < T; t++) { - const float * src = (t < decoded_T) - ? decoded_latents.data() + t * Oc - : silence_full.data() + (t - decoded_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; + if (have_cover) { + for (int t = 0; t < T; t++) { + const float * src = (t < T_cover) + ? cover_latents.data() + t * Oc + : silence_full.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; + } + } else { + for (int t = 0; t < T; t++) { + const float * src = (t < decoded_T) + ? decoded_latents.data() + t * Oc + : silence_full.data() + (t - decoded_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) @@ -911,6 +978,32 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) memcpy(context.data() + b * T * ctx_ch, context_single.data(), T * ctx_ch * sizeof(float)); } + // Cover mode: build silence context for audio_cover_strength switching + // When step >= cover_steps, DiT switches from cover context to silence context + std::vector context_silence; + int cover_steps = -1; + if (have_cover) { + float cover_strength = req.audio_cover_strength; + if (cover_strength < 1.0f) { + // Build silence context: all frames use silence_latent + std::vector silence_single(T * ctx_ch); + for (int t = 0; t < T; t++) { + const float * src = silence_full.data() + t * Oc; + for (int c = 0; c < Oc; c++) + silence_single[t * ctx_ch + c] = src[c]; + for (int c = 0; c < Oc; c++) + silence_single[t * ctx_ch + Oc + c] = 1.0f; + } + context_silence.resize(batch_n * T * ctx_ch); + for (int b = 0; b < batch_n; b++) + memcpy(context_silence.data() + b * T * ctx_ch, + silence_single.data(), T * ctx_ch * sizeof(float)); + cover_steps = (int)((float)num_steps * cover_strength); + fprintf(stderr, "[Cover] audio_cover_strength=%.2f -> switch at step %d/%d\n", + cover_strength, cover_steps, num_steps); + } + } + // Generate N noise samples std::vector noise(batch_n * Oc * T); @@ -929,13 +1022,16 @@ std::string acestep_generate_audio(const music_generation_inputs inputs) // DiT Generate std::vector output(batch_n * Oc * T); - fprintf(stderr, "[DiT] Starting: T=%d, S=%d, enc_S=%d, steps=%d, batch=%d\n", - T, S, enc_S, num_steps, batch_n); + fprintf(stderr, "[DiT] Starting: T=%d, S=%d, enc_S=%d, steps=%d, batch=%d%s\n", + T, S, enc_S, num_steps, batch_n, + have_cover ? " (cover)" : ""); music_dit_timer.reset(); dit_ggml_generate(&acestep_dit, noise.data(), context.data(), enc_hidden.data(), enc_S, T, batch_n, num_steps, schedule.data(), output.data(), - guidance_scale); + guidance_scale, nullptr, + context_silence.empty() ? nullptr : context_silence.data(), + cover_steps); fprintf(stderr, "[DiT] Total generation: %.1f ms (%.1f ms/sample)\n", music_dit_timer.ms(), music_dit_timer.ms() / batch_n); diff --git a/otherarch/acestep/dit.h b/otherarch/acestep/dit.h index 4210aebb6..8fd57e775 100644 --- a/otherarch/acestep/dit.h +++ b/otherarch/acestep/dit.h @@ -1084,7 +1084,9 @@ static void dit_ggml_generate( const float * schedule, float * output, float guidance_scale = 1.0f, - const DebugDumper * dbg = nullptr) { + const DebugDumper * dbg = nullptr, + const float * context_switch = nullptr, + int cover_steps = -1) { DiTGGMLConfig & c = model->cfg; int Oc = c.out_channels; // 64 @@ -1241,9 +1243,22 @@ static void dit_ggml_generate( struct ggml_tensor * t_t = ggml_graph_get_tensor(gf, "t"); // Flow matching loop + bool switched_cover = false; for (int step = 0; step < num_steps; step++) { float t_curr = schedule[step]; + // Cover mode: switch context from cover to non-cover at cover_steps + if (context_switch && cover_steps >= 0 && step >= cover_steps && !switched_cover) { + switched_cover = true; + for (int b = 0; b < N; b++) + for (int t = 0; t < T; t++) + memcpy(&input_buf[b * T * in_ch + t * in_ch], + &context_switch[b * T * ctx_ch + t * ctx_ch], + ctx_ch * sizeof(float)); + fprintf(stderr, "[DiT] Cover: switched to non-cover context at step %d/%d\n", + step, num_steps); + } + // Set timestep (changes each step) if (t_t) { ggml_backend_tensor_set(t_t, &t_curr, 0, sizeof(float)); diff --git a/otherarch/acestep/request.cpp b/otherarch/acestep/request.cpp index a3d0b3696..8990fd0e1 100644 --- a/otherarch/acestep/request.cpp +++ b/otherarch/acestep/request.cpp @@ -36,6 +36,7 @@ void request_init(AceRequest * r) { r->inference_steps = 8; r->guidance_scale = 1.0f; r->shift = 3.0f; + r->audio_cover_strength = 0.5f; } // JSON string escape / unescape @@ -250,6 +251,7 @@ bool request_parse_from_str(AceRequest * r, std::string json) { 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()); + else if (k == "audio_cover_strength") r->audio_cover_strength = (float)atof(v.c_str()); else if (k == "codes_temperature") r->codes_temperature = (float)atof(v.c_str()); else if (k == "codes_top_p") r->codes_top_p = (float)atof(v.c_str()); @@ -297,6 +299,7 @@ bool request_write(const AceRequest * r, const char * path) { 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); + fprintf(f, " \"audio_cover_strength\": %.2f,\n", r->audio_cover_strength); // audio_codes last (no trailing comma) fprintf(f, " \"audio_codes\": \"%s\"\n", json_escape(r->audio_codes).c_str()); fprintf(f, "}\n"); @@ -321,6 +324,8 @@ void request_dump(const AceRequest * r, FILE * f) { r->lm_temperature, r->lm_top_p, r->lm_top_k); fprintf(f, " dit: steps=%d guidance=%.1f shift=%.1f\n", r->inference_steps, r->guidance_scale, r->shift); + if (r->audio_cover_strength != 0.5f) + fprintf(f, " cover: strength=%.2f\n", r->audio_cover_strength); fprintf(f, " audio_codes: %s\n", r->audio_codes.empty() ? "(none)" : "(present)"); } diff --git a/otherarch/acestep/request.h b/otherarch/acestep/request.h index a71ab62db..014d62ec2 100644 --- a/otherarch/acestep/request.h +++ b/otherarch/acestep/request.h @@ -47,6 +47,9 @@ struct AceRequest { int inference_steps; // 8 float guidance_scale; // 7.0 float shift; // 1.0 + + // cover mode (active when --src-audio is provided on CLI) + float audio_cover_strength; // 0.5 (0-1, fraction of DiT steps using source context) }; // Initialize all fields to defaults (matches Python GenerationParams defaults) diff --git a/otherarch/acestep/vae-enc.h b/otherarch/acestep/vae-enc.h new file mode 100644 index 000000000..0e4ca7320 --- /dev/null +++ b/otherarch/acestep/vae-enc.h @@ -0,0 +1,392 @@ +// vae-enc.h: AutoencoderOobleck encoder (audio -> latent) via ggml +// +// Mirror of vae.h decoder. Reuses VAEResUnit, load helpers, graph ops. +// Architecture: conv1(2->128,k=7) -> 5x(3xresunit+snake+strided_conv) -> snake+conv2(2048->128,k=3) +// Output 128ch = mean[64] + scale[64]. Deterministic encode returns mean. +// Downsample: 2x4x4x6x10 = 1920x (matches decoder upsample) + +#pragma once +#include "vae.h" + +// Encoder block: 3xResUnit(in_ch) -> snake(in_ch) -> strided Conv1d(in_ch -> out_ch) +// Decoder block is the mirror: snake(in_ch) -> ConvT(in_ch -> out_ch) -> 3xResUnit(out_ch) +struct VAEEncBlock { + VAEResUnit ru[3]; + struct ggml_tensor * sa, * sb; // snake [1, in_ch] + struct ggml_tensor * dw, * db; // strided conv [K, in_ch, out_ch], bias [out_ch] + int in_ch, out_ch, stride, kernel, padding; +}; + +struct VAEEncoder { + struct ggml_tensor * c1w, * c1b; // conv1 [7, 2, 128], bias [128] + VAEEncBlock blk[5]; + struct ggml_tensor * sa, * sb; // final snake [1, 2048] + struct ggml_tensor * c2w, * c2b; // conv2 [3, 2048, 128], bias [128] + + ggml_backend_t backend; + ggml_backend_t cpu_backend; + ggml_backend_sched_t sched; + ggml_backend_buffer_t buf; + struct ggml_context * weight_ctx; + + // graph cache (rebuilt when T_audio changes) + struct ggml_context * graph_ctx; + uint8_t * graph_buf; + struct ggml_cgraph * graph; + struct ggml_tensor * graph_input; // [T_audio, 2] + struct ggml_tensor * graph_output; // [T_latent, 128] + int graph_T; // cached T_audio (0 = no cache) + + std::vector scratch_in; // transposed input [2 * T_audio] +}; + +// Load encoder weights from the same VAE GGUF (encoder.* tensors) +static void vae_enc_load(VAEEncoder * m, const char * path) { + GGUFModel gf = {}; + if (!gf_load(&gf, path)) { + fprintf(stderr, "[VAE-Enc] FATAL: cannot load %s\n", path); + exit(1); + } + + // Encoder channel layout (mirror of decoder, bottom-up): + // conv1: 2 -> 128 + // block: [128->128, 128->256, 256->512, 512->1024, 1024->2048] + // conv2: 2048 -> 128 (split: mean[64] + scale[64]) + // ResUnits run at in_ch (before downsample), unlike decoder (at out_ch, after upsample). + static const int in_ch[] = {128, 128, 256, 512, 1024}; + static const int out_ch[] = {128, 256, 512, 1024, 2048}; + static const int strides[] = { 2, 4, 4, 6, 10}; + static const int dilations[] = {1, 3, 9}; + + // Phase 1: create weight tensors + size_t ctx_size = ggml_tensor_overhead() * 256; + 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, 2, 128); + m->c1b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + + for (int i = 0; i < 5; i++) { + VAEEncBlock & 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; + b.padding = (strides[i] + 1) / 2; // ceil(stride / 2) + int C = in_ch[i]; // res_units + snake at in_ch + + // 3 res units at in_ch + 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); + } + + // snake at in_ch (before downsample conv) + b.sa = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C); + b.sb = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C); + + // strided conv1d: [K, in_ch, out_ch] + b.dw = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, b.kernel, in_ch[i], out_ch[i]); + b.db = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_ch[i]); + } + + m->sa = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 2048); + m->sb = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 2048); + m->c2w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 3, 2048, 128); + m->c2b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + + // Phase 2: allocate backend buffer + BackendPair bp = backend_init("VAE-Enc"); + 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); + if (!m->buf) { + fprintf(stderr, "[VAE-Enc] FATAL: failed to allocate weight buffer\n"); + exit(1); + } + fprintf(stderr, "[VAE-Enc] 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 and fuse weights + vae_fuse_wn(m->c1w, gf, "encoder.conv1"); + vae_load_bias(m->c1b, gf, "encoder.conv1.bias"); + + for (int i = 0; i < 5; i++) { + VAEEncBlock & b = m->blk[i]; + std::string blk_pfx = "encoder.block." + std::to_string(i); + + // res_units first (same load pattern as decoder) + 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"); + } + + // snake + strided downsample conv (regular conv1d, NOT transposed) + 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.dw, gf, blk_pfx + ".conv1"); + vae_load_bias(b.db, gf, blk_pfx + ".conv1.bias"); + } + + vae_load_snake(m->sa, gf, "encoder.snake1.alpha"); + vae_load_snake_inv(m->sb, gf, "encoder.snake1.beta"); + vae_fuse_wn(m->c2w, gf, "encoder.conv2"); + vae_load_bias(m->c2b, gf, "encoder.conv2.bias"); + + fprintf(stderr, "[VAE-Enc] Loaded: 5 blocks, downsample=1920x, F32 activations\n"); + gf_close(&gf); +} + +// Build encoder graph: audio [T_audio, 2] -> [T_latent, 128] +static struct ggml_tensor * vae_enc_build_graph( + struct ggml_context * ctx, + VAEEncoder * m, + struct ggml_tensor * audio) { // [T, 2] + + // conv1: [T, 2] -> [T, 128] + struct ggml_tensor * x = vae_conv1d(ctx, m->c1w, m->c1b, audio, 1, 3, 1); + + // 5 encoder blocks: resunits(in_ch) -> snake(in_ch) -> strided conv(in_ch -> out_ch) + for (int i = 0; i < 5; i++) { + VAEEncBlock & b = m->blk[i]; + for (int r = 0; r < 3; r++) + x = vae_res_unit(ctx, &b.ru[r], x); + x = vae_snake(ctx, x, b.sa, b.sb); + x = vae_conv1d(ctx, b.dw, b.db, x, b.stride, b.padding, 1); + } + + // Final: snake(2048) -> conv2(2048 -> 128, k=3, pad=1) + x = vae_snake(ctx, x, m->sa, m->sb); + x = vae_conv1d(ctx, m->c2w, m->c2b, x, 1, 1, 1); + + return x; // [T_latent, 128] +} + +// Core compute: build/cache graph, set input, run. Returns T_latent or -1. +// Output stays in m->graph_output for caller to read. +static int vae_enc_compute( + VAEEncoder * m, + const float * audio, // [T_audio, 2] time-major interleaved stereo + int T_audio) { + + // Rebuild graph when T_audio changes + if (m->graph_T != T_audio) { + 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_audio, 2); + ggml_set_name(m->graph_input, "enc_input"); + ggml_set_input(m->graph_input); + + m->graph_output = vae_enc_build_graph(ctx, m, m->graph_input); + ggml_set_name(m->graph_output, "enc_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-Enc] FATAL: graph alloc failed for T=%d\n", T_audio); + 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_audio; + fprintf(stderr, "[VAE-Enc] Graph: %d nodes, T_audio=%d\n", + ggml_graph_n_nodes(m->graph), T_audio); + } + + // Transpose: [T, 2] time-major -> ggml [T, 2] channel-contiguous + // ggml ne[0]=T is the contiguous dim, so we write all T samples per channel + size_t in_size = (size_t)2 * T_audio; + if (m->scratch_in.size() < in_size) + m->scratch_in.resize(in_size); + for (int c = 0; c < 2; c++) + for (int t = 0; t < T_audio; t++) + m->scratch_in[c * T_audio + t] = audio[t * 2 + 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]; // T_latent +} + +// Encode API: audio [T_audio, 2] -> latent_out [T_latent, 64] (mean only, deterministic) +// Returns T_latent (or -1 on error). +// latent_out must hold at least (T_audio / 1920) * 64 floats. +static int vae_enc_encode( + VAEEncoder * m, + const float * audio, // [T_audio, 2] interleaved stereo + int T_audio, + float * latent_out, // [T_latent, 64] output, time-major + int max_T_latent) { + + int T_latent = vae_enc_compute(m, audio, T_audio); + if (T_latent < 0) return -1; + + if (T_latent > max_T_latent) { + fprintf(stderr, "[VAE-Enc] T_latent %d exceeds max %d\n", T_latent, max_T_latent); + return -1; + } + + // Graph output is [ne0=T_latent, ne1=128] in ggml, channel-contiguous. + // Channels 0..63 = mean, 64..127 = scale. We only read mean. + // ggml layout: data[c * T_latent + t] for channel c, time t. + // We write time-major: latent_out[t * 64 + c] = data[c * T_latent + t] + // + // Read the full 128ch output once, extract mean channels 0..63 + size_t out_bytes = (size_t)128 * T_latent * sizeof(float); + std::vector raw(128 * T_latent); + ggml_backend_tensor_get(m->graph_output, raw.data(), 0, out_bytes); + + for (int t = 0; t < T_latent; t++) + for (int c = 0; c < 64; c++) + latent_out[t * 64 + c] = raw[c * T_latent + t]; + + fprintf(stderr, "[VAE-Enc] Encode: T_audio=%d -> T_latent=%d (%.2fs @ 48kHz)\n", + T_audio, T_latent, (float)T_audio / 48000.0f); + + return T_latent; +} + +// Tiled encode for long audio (same chunking strategy as decoder) +// chunk_size: latent frames per tile, overlap: context frames on each side +static int vae_enc_encode_tiled( + VAEEncoder * m, + const float * audio, // [T_audio, 2] interleaved stereo + int T_audio, + float * latent_out, // [T_latent, 64] output, time-major + int max_T_latent, + int chunk_size = 256, + int overlap = 64) { + + // Work in audio-sample space. Each latent frame = 1920 audio samples. + int audio_chunk = chunk_size * 1920; + int audio_overlap = overlap * 1920; + + // Shrink overlap until stride is positive + while (audio_chunk - 2 * audio_overlap <= 0 && audio_overlap > 0) + audio_overlap /= 2; + + // Short audio: encode directly + if (T_audio <= audio_chunk) + return vae_enc_encode(m, audio, T_audio, latent_out, max_T_latent); + + int audio_stride = audio_chunk - 2 * audio_overlap; + int num_steps = (T_audio + audio_stride - 1) / audio_stride; + + fprintf(stderr, "[VAE-Enc] Tiled encode: %d tiles (chunk=%d, overlap=%d, stride=%d audio samples)\n", + num_steps, audio_chunk, audio_overlap, audio_stride); + + float downsample_factor = 0.0f; + int latent_write_pos = 0; + + for (int i = 0; i < num_steps; i++) { + // Core range in audio samples (the part we keep) + int core_start = i * audio_stride; + int core_end = core_start + audio_stride; + if (core_end > T_audio) core_end = T_audio; + + // Window with overlap context + int win_start = core_start - audio_overlap; + if (win_start < 0) win_start = 0; + int win_end = core_end + audio_overlap; + if (win_end > T_audio) win_end = T_audio; + int win_len = win_end - win_start; + + // Encode this window + int tile_T = vae_enc_compute(m, audio + win_start * 2, win_len); + if (tile_T < 0) { + fprintf(stderr, "[VAE-Enc] FATAL: tile %d encode failed\n", i); + return -1; + } + + // Determine downsample factor from first tile + if (i == 0) { + downsample_factor = (float)tile_T / (float)win_len; + fprintf(stderr, "[VAE-Enc] Downsample factor: %.6f (expected ~1/1920)\n", + downsample_factor); + } + + // Trim in latent frames (mirror of decoder trim logic) + int added_start = core_start - win_start; + int trim_start = (int)roundf((float)added_start * downsample_factor); + int added_end = win_end - core_end; + int trim_end = (int)roundf((float)added_end * downsample_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; + + if (latent_write_pos + core_len > max_T_latent) { + fprintf(stderr, "[VAE-Enc] FATAL: tiled output exceeds max_T_latent\n"); + return -1; + } + + // Read tile output [ne0=tile_T, ne1=128], extract mean (ch 0..63), transpose + // Only read the first 64 channels (mean), skip scale channels 64..127 + size_t out_bytes = (size_t)128 * tile_T * sizeof(float); + std::vector raw(128 * tile_T); + ggml_backend_tensor_get(m->graph_output, raw.data(), 0, out_bytes); + + for (int t = 0; t < core_len; t++) + for (int c = 0; c < 64; c++) + latent_out[(latent_write_pos + t) * 64 + c] = + raw[c * tile_T + (trim_start + t)]; + + latent_write_pos += core_len; + } + + fprintf(stderr, "[VAE-Enc] Tiled encode done: %d tiles -> T_latent=%d (%.2fs @ 48kHz)\n", + num_steps, latent_write_pos, (float)T_audio / 48000.0f); + + return latent_write_pos; +} + +// Free all resources +static void vae_enc_free(VAEEncoder * 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 = {}; +}