// pipeline-synth-ops.cpp: primitive operations of the synthesis pipeline
//
// Each op takes AceSynth (the pipeline context) and SynthState (the transient
// job state). See pipeline-synth-ops.h for the per-op contract and
// pipeline-synth-impl.h for the struct layouts.

#include "pipeline-synth-ops.h"

#include "hot-step-sampler.h"
#include "hot-step-sampler-trt.h"
#include "adapter-trt.h"
#include "philox.h"
#include "pipeline-synth-impl.h"
#include "task-types.h"
#include "vae-enc.h"

#include <charconv>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <system_error>
#include <vector>

static const int FRAMES_PER_SECOND = 25;

// ─── Determinism diagnostic ─────────────────────────────────────────────────
static void diag_stats_f32(const char * label, const float * data, size_t n) {
    if (!data || n == 0) { return; }
    double sum = 0.0, sum_sq = 0.0;
    float  mn = data[0], mx = data[0];
    for (size_t i = 0; i < n; i++) {
        double v = (double) data[i];
        sum += v;
        sum_sq += v * v;
        if (data[i] < mn) { mn = data[i]; }
        if (data[i] > mx) { mx = data[i]; }
    }
    double mean = sum / (double) n;
    double rms  = sqrt(sum_sq / (double) n);
    fprintf(stderr, "[DIAG] %s: n=%zu mean=%.8f rms=%.8f min=%.6f max=%.6f sum=%.10f\n",
            label, n, mean, rms, mn, mx, sum);
}

// CSV list parser tolerant to any whitespace around commas. Locale-immune via
// std::from_chars (C++17 charconv) for integers. Float parsing uses strtof
// for portability (Apple libc++ does not implement from_chars for floats).
// Used for audio_codes (int) and custom_timesteps (float). Bails on first
// parse error or overflow, returning the values consumed so far.
template <typename T> static std::vector<T> parse_csv(const std::string & s) {
    std::vector<T> out;
    const char *   first = s.data();
    const char *   last  = first + s.size();
    while (first < last) {
        while (first < last && (*first == ',' || *first == ' ')) {
            first++;
        }
        if (first == last) {
            break;
        }
        T    v{};
        auto r = std::from_chars(first, last, v);
        if (r.ec != std::errc{}) {
            break;
        }
        out.push_back(v);
        first = r.ptr;
    }
    return out;
}

// Float specialization: strtof-based (Apple libc++ lacks from_chars<float>)
template <> std::vector<float> parse_csv<float>(const std::string & s) {
    std::vector<float> out;
    const char *       first = s.data();
    const char *       last  = first + s.size();
    while (first < last) {
        while (first < last && (*first == ',' || *first == ' ')) {
            first++;
        }
        if (first == last) {
            break;
        }
        char * end = nullptr;
        float  v   = std::strtof(first, &end);
        if (end == first) {
            break;  // no progress — parse error
        }
        out.push_back(v);
        first = end;
    }
    return out;
}

int ops_encode_src(const AceSynth * ctx,
                   const float *    src_audio,
                   int              src_len,
                   const float *    src_latents,
                   int              src_T_latent,
                   SynthState &     s) {
    // Cover mode: ingest source either as pre-encoded latents (zero VAE work)
    // or by acquiring the VAE encoder and running it on src_audio. When both
    // are provided latents win: they were either produced by a previous run
    // or supplied verbatim by the client and need no further processing.
    s.have_cover = false;
    s.T_cover    = 0;
    if (src_latents && src_T_latent > 0) {
        s.cover_latents.assign(src_latents, src_latents + (size_t) src_T_latent * 64);
        s.T_cover    = src_T_latent;
        s.have_cover = true;
        fprintf(stderr, "[Encode-Src] Latents in: T_cover=%d (%.2fs), VAE encode skipped\n", s.T_cover,
                (float) s.T_cover * 1920.0f / 48000.0f);
        return 0;
    }
    if (src_audio && src_len > 0) {
        s.timer.reset();
        int T_audio = src_len;

        VAEEncoder * vae_enc = store_require_vae_enc(ctx->store, ctx->vae_enc_key);
        if (!vae_enc) {
            fprintf(stderr, "[Encode-Src] FATAL: store_require_vae_enc failed\n");
            return -1;
        }
        ModelHandle vae_enc_guard(ctx->store, vae_enc);

        int max_T_lat = (T_audio / 1920) + 64;
        s.cover_latents.resize(max_T_lat * 64);

        s.T_cover = vae_enc_encode_tiled(vae_enc, src_audio, T_audio, s.cover_latents.data(), max_T_lat,
                                         ctx->params.vae_chunk, ctx->params.vae_overlap);
        if (s.T_cover < 0) {
            fprintf(stderr, "[Encode-Src] FATAL: encode failed\n");
            return -1;
        }
        s.cover_latents.resize(s.T_cover * 64);
        fprintf(stderr, "[Encode-Src] Encoded: T_cover=%d (%.2fs), %.1f ms\n", s.T_cover,
                (float) s.T_cover * 1920.0f / 48000.0f, s.timer.ms());

        s.have_cover = true;
    }

    return 0;
}

void ops_fsq_roundtrip(const AceSynth * ctx, SynthState & s) {
    // FSQ roundtrip for cover: tokenize (25Hz->5Hz) + detokenize (5Hz->25Hz).
    // The lossy 5:1 temporal compression destroys micro-timings, ornaments and
    // transients. The DiT receives degraded latents and diverges from the source,
    // producing a free reinterpretation rather than a close remix.
    // cover-nofsq skips this call and feeds clean 25Hz VAE latents directly,
    // producing remixes that stay close to the source.
    // Other tasks (lego, extract, repaint, complete) also use clean latents.
    if (!s.have_cover) {
        return;
    }
    s.timer.reset();
    int              T_5Hz = (s.T_cover + 4) / 5;
    std::vector<int> codes(T_5Hz);

    // Tokenizer scope: acquire, encode 25Hz latents to 5Hz codes, release.
    int T_5Hz_actual;
    {
        TokGGML * tok = store_require_fsq_tok(ctx->store, ctx->fsq_tok_key);
        if (!tok) {
            fprintf(stderr, "[FSQ-Roundtrip] FATAL: store_require_fsq_tok failed\n");
            return;
        }
        ModelHandle tok_guard(ctx->store, tok);
        if (!ctx->params.use_fa) {
            tok->use_flash_attn = false;
        }

        T_5Hz_actual =
            tok_ggml_encode(tok, s.cover_latents.data(), s.T_cover, codes.data(), ctx->meta->silence_full.data());
    }
    if (T_5Hz_actual <= 0) {
        return;
    }

    // Detokenizer scope: acquire, decode 5Hz codes back to 25Hz latents, release.
    int                T_25Hz_rt = T_5Hz_actual * 5;
    std::vector<float> rt_latents(T_25Hz_rt * 64);
    int                ret;
    {
        DetokGGML * detok = store_require_fsq_detok(ctx->store, ctx->fsq_detok_key);
        if (!detok) {
            fprintf(stderr, "[FSQ-Roundtrip] FATAL: store_require_fsq_detok failed\n");
            return;
        }
        ModelHandle detok_guard(ctx->store, detok);
        if (!ctx->params.use_fa) {
            detok->use_flash_attn = false;
        }

        ret = detok_ggml_decode(detok, codes.data(), T_5Hz_actual, rt_latents.data());
    }
    if (ret < 0) {
        return;
    }
    int copy_T = T_25Hz_rt < s.T_cover ? T_25Hz_rt : s.T_cover;
    memcpy(s.cover_latents.data(), rt_latents.data(), (size_t) copy_T * 64 * sizeof(float));
    fprintf(stderr, "[FSQ-Roundtrip] %d->%d->%d frames, %.1f ms\n", s.T_cover, T_5Hz_actual, copy_T, s.timer.ms());
}

int ops_resolve_params(const AceSynth * ctx, const AceRequest * reqs, int batch_n, SynthState & s) {
    // Extract shared params from first request
    s.duration = s.rr.duration > 0 ? s.rr.duration : 30.0f;

    // Resolve DiT sampling params: 0 = auto-detect from model type.
    // Turbo: 8 steps, guidance=1.0, s.shift=3.0
    // Base/SFT: 50 steps, guidance=1.0, s.shift=1.0
    s.num_steps      = s.rr.inference_steps;
    s.guidance_scale = s.rr.guidance_scale;
    s.shift          = s.rr.shift;

    if (s.num_steps <= 0) {
        s.num_steps = ctx->meta->is_turbo ? 8 : 50;
    }
    if (s.num_steps > 200) {
        fprintf(stderr, "[Resolve-Params] WARNING: inference_steps %d clamped to 200\n", s.num_steps);
        s.num_steps = 200;
    }

    if (s.guidance_scale <= 0.0f) {
        s.guidance_scale = 1.0f;
    } else if (ctx->meta->is_turbo && s.guidance_scale > 1.0f) {
        fprintf(stderr,
                "[Resolve-Params] WARNING: guidance_scale=%.1f on turbo model. "
                "Distilled turbo internalizes guidance; stacking CFG on top usually "
                "oversaturates. Use it only if you know what you do (typically a merge).\n",
                s.guidance_scale);
    }

    if (s.shift <= 0.0f) {
        s.shift = ctx->meta->is_turbo ? 3.0f : 1.0f;
    }

    // Audio codes: parse once per request and stash in s.per_codes so
    // ops_build_context does not have to re-parse. Also records the longest
    // set (drives s.T) and whether any batch item carries codes at all.
    // Shorter code sets are padded with silence, longer ones are never truncated.
    s.per_codes.assign(batch_n, {});
    s.max_codes_len = 0;
    s.have_codes    = false;
    for (int b = 0; b < batch_n; b++) {
        s.per_codes[b] = parse_csv<int>(reqs[b].audio_codes);
        int sz         = (int) s.per_codes[b].size();
        if (sz > s.max_codes_len) {
            s.max_codes_len = sz;
        }
        if (sz > 0) {
            s.have_codes = true;
        }
    }
    if (s.have_codes) {
        fprintf(stderr, "[Resolve-Params] max audio codes across batch: %d (%.1fs @ 5Hz)\n", s.max_codes_len,
                (float) s.max_codes_len / 5.0f);
    }

    return 0;
}

void ops_build_schedule(SynthState & s) {
    // Custom timesteps override: CSV floats like
    // "0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0". Last value is the x0
    // endpoint handled implicitly by the sampler, so we drop it and take
    // schedule = first N-1 entries, num_steps = N-1.
    if (!s.rr.custom_timesteps.empty()) {
        std::vector<float> ts = parse_csv<float>(s.rr.custom_timesteps);
        if (ts.size() >= 2) {
            s.num_steps = (int) ts.size() - 1;
            s.schedule.assign(ts.begin(), ts.end() - 1);
            fprintf(stderr, "[Build-Schedule] Custom timesteps: %d steps\n", s.num_steps);
            return;
        }
        fprintf(stderr, "[Build-Schedule] WARN: custom_timesteps needs >= 2 values, falling back to shift\n");
    }
    // Default: t_i = shift * t / (1 + (shift-1)*t) with t = 1 - i/steps
    s.schedule.resize(s.num_steps);
    for (int i = 0; i < s.num_steps; i++) {
        float t       = 1.0f - (float) i / (float) s.num_steps;
        s.schedule[i] = s.shift * t / (1.0f + (s.shift - 1.0f) * t);
    }
}

int ops_resolve_T(const AceSynth * ctx, SynthState & s) {
    // s.T = number of 25Hz latent frames for DiT
    // Source tasks: from source audio. Codes: from code count. Else: from s.duration.
    if (s.use_source_context && s.have_cover) {
        s.T        = s.T_cover;
        // s.duration in metas must match actual source length, not JSON default
        s.duration = (float) s.T_cover / (float) FRAMES_PER_SECOND;
    } else if (s.have_codes) {
        s.T = s.max_codes_len * 5;
    } else if (s.use_source_context) {
        // source context requested but neither cover_latents nor codes available.
        // duration fallthrough would produce a meaningless T for source tasks.
        fprintf(stderr, "[Resolve-T] FATAL: use_source_context but no cover_latents and no audio_codes\n");
        return -1;
    } else {
        s.T = (int) (s.duration * FRAMES_PER_SECOND);
    }
    s.T     = ((s.T + ctx->meta->cfg.patch_size - 1) / ctx->meta->cfg.patch_size) * ctx->meta->cfg.patch_size;
    s.S     = s.T / ctx->meta->cfg.patch_size;
    s.enc_S = 0;

    fprintf(stderr, "[Resolve-T] T=%d, S=%d\n", s.T, s.S);
    fprintf(stderr, "[Resolve-T] seed=%lld, steps=%d, guidance=%.1f, shift=%.1f, duration=%.1fs\n",
            (long long) s.rr.seed, s.num_steps, s.guidance_scale, s.shift, s.duration);

    if (s.T > 15000) {
        fprintf(stderr, "[Resolve-T] ERROR: T=%d exceeds silence_latent max 15000, skipping\n", s.T);
        return -1;
    }

    return 0;
}

void ops_encode_timbre(const AceSynth * ctx,
                       const float *    ref_audio,
                       int              ref_len,
                       const float *    ref_latents,
                       int              ref_T_latent,
                       SynthState &     s) {
    // Timbre features from ref_audio or ref_latents (independent of src).
    // Two paths converge into s.timbre_feats: pre-encoded latents skip the
    // VAE encoder entirely, raw audio takes the encoder path. Latents win
    // when both are set. Without either input the timbre falls back to a
    // single silence frame, disabling timbre conditioning.
    if (ref_latents && ref_T_latent > 0) {
        s.S_ref_timbre = ref_T_latent;
        s.timbre_feats.assign(ref_latents, ref_latents + (size_t) ref_T_latent * 64);
        fprintf(stderr, "[Encode-Timbre] Latents in: %d frames (%.1fs), VAE encode skipped\n", ref_T_latent,
                (float) ref_T_latent / 25.0f);
        return;
    }
    if (ref_audio && ref_len > 0) {
        s.timer.reset();
        VAEEncoder * ref_vae = store_require_vae_enc(ctx->store, ctx->vae_enc_key);
        if (!ref_vae) {
            fprintf(stderr, "[Encode-Timbre] WARNING: store_require_vae_enc failed, using silence\n");
            s.S_ref_timbre = 1;
            s.timbre_feats.assign(ctx->meta->silence_full.data(), ctx->meta->silence_full.data() + 64);
            return;
        }
        ModelHandle ref_vae_guard(ctx->store, ref_vae);

        int                max_T_ref = (ref_len / 1920) + 64;
        std::vector<float> ref_lat_buf(max_T_ref * 64);
        int                T_ref = vae_enc_encode_tiled(ref_vae, ref_audio, ref_len, ref_lat_buf.data(), max_T_ref,
                                                        ctx->params.vae_chunk, ctx->params.vae_overlap);
        if (T_ref < 0) {
            fprintf(stderr, "[Encode-Timbre] WARNING: ref_audio encode failed, using silence\n");
            s.S_ref_timbre = 1;
            s.timbre_feats.assign(ctx->meta->silence_full.data(), ctx->meta->silence_full.data() + 64);
        } else {
            s.S_ref_timbre = T_ref;
            s.timbre_feats.assign(ref_lat_buf.data(), ref_lat_buf.data() + (size_t) T_ref * 64);
            fprintf(stderr, "[Encode-Timbre] ref_audio: %d frames (%.1fs), %.1f ms\n", T_ref, (float) T_ref / 25.0f,
                    s.timer.ms());
        }
    } else {
        s.S_ref_timbre = 1;
        s.timbre_feats.assign(ctx->meta->silence_full.data(), ctx->meta->silence_full.data() + 64);
    }
}

// Per-batch CPU-resident forward from the text encoder.
// Lives in a local array between the text_enc and cond_enc phases so the
// two GPU modules never need to coexist under EVICT_STRICT.
struct TextEncForward {
    std::vector<float> text_hidden;  // [S_text * H_text] f32
    std::vector<float> lyric_embed;  // [S_lyric * H_text] f32
    int                S_text;
    int                S_lyric;
};

// Build the text/lyric prompt pair that feeds the text encoder for one batch
// element. instruction is the DiT instruction header (main or non-cover).
static void build_prompt_strings(const AceRequest &  rb,
                                 const std::string & instruction,
                                 float               duration,
                                 std::string &       text_out,
                                 std::string &       lyric_out) {
    char bpm_b[16] = "N/A";
    if (rb.bpm > 0) {
        snprintf(bpm_b, sizeof(bpm_b), "%d", rb.bpm);
    }
    const char * keyscale_b = rb.keyscale.empty() ? "N/A" : rb.keyscale.c_str();
    const char * timesig_b  = rb.timesignature.empty() ? "N/A" : rb.timesignature.c_str();
    const char * language_b = rb.vocal_language.empty() ? "unknown" : rb.vocal_language.c_str();

    char metas_b[512];
    snprintf(metas_b, sizeof(metas_b), "- bpm: %s\n- timesignature: %s\n- keyscale: %s\n- duration: %d seconds\n",
             bpm_b, timesig_b, keyscale_b, (int) duration);
    text_out = std::string("# Instruction\n") + instruction + "\n\n" + "# Caption\n" + rb.caption + "\n\n" +
               "# Metas\n" + metas_b + "<|endoftext|>\n";
    lyric_out = std::string("# Languages\n") + language_b + "\n\n# Lyric\n" + rb.lyrics + "<|endoftext|>";
}

int ops_encode_text(const AceSynth * ctx, const AceRequest * reqs, int batch_n, SynthState & s) {
    // Per-batch text encoding in two GPU phases to keep EVICT_STRICT at one
    // module resident at a time:
    //   Phase A: acquire text_enc, run all qwen3 forwards (main and optional
    //            non-cover) into CPU-resident TextEncForward, release.
    //   Phase B: acquire cond_enc, run all cond_ggml forwards consuming those
    //            cached hidden states into s.per_enc / s.per_enc_nc, release.
    // Intermediate CPU footprint peaks at roughly batch_n * 2 * S * H_text * 4
    // bytes (a few MB), negligible next to the GPU modules.
    if (s.instruction_str.empty()) {
        fprintf(stderr, "[Encode-Text] FATAL: instruction_str is empty (unknown task or orchestrator bug)\n");
        return -1;
    }

    s.need_enc_switch = s.use_source_context && !s.is_repaint && !s.is_lego_region && s.rr.audio_cover_strength < 1.0f;

    // BPE tokenizer: when text_encoder_path points at an .onnx file,
    // vocab.json + merges.txt live in the same directory, not inside the file.
    std::string bpe_dir;
    const char * bpe_path = ctx->params.text_encoder_path;
    if (ctx->is_onnx_pipeline) {
        bpe_dir = ctx->text_enc_ort_key.path;
        auto slash = bpe_dir.find_last_of("/\\");
        if (slash != std::string::npos) bpe_dir = bpe_dir.substr(0, slash);
        bpe_path = bpe_dir.c_str();
    }
    BPETokenizer * bpe = store_bpe(ctx->store, bpe_path);
    if (!bpe) {
        fprintf(stderr, "[Encode-Text] FATAL: store_bpe failed (path=%s)\n", bpe_path);
        return -1;
    }

    std::vector<TextEncForward> main_fwd(batch_n);
    std::vector<TextEncForward> nc_fwd(s.need_enc_switch ? batch_n : 0);
    int                         H_text = 0;
    int                         H_cond = 0;

    // ═══════════════════════════════════════════════════════════════════
    // ORT PATH: text encoder + condition encoder via ONNX Runtime
    // ═══════════════════════════════════════════════════════════════════
    if (ctx->is_onnx_pipeline) {
        fprintf(stderr, "[Encode-Text] Using ONNX pipeline (TextEnc-ORT + CondEnc-ORT)\n");

        // Phase A(ORT): text encoder
        {
            TextEncOrt * te = store_require_text_enc_ort(ctx->store, ctx->text_enc_ort_key);
            if (!te) {
                fprintf(stderr, "[Encode-Text] FATAL: store_require_text_enc_ort failed\n");
                return -1;
            }
            ModelHandle te_guard(ctx->store, te);
            H_text = te->hidden_size;

            for (int b = 0; b < batch_n; b++) {
                std::string text_str;
                std::string lyric_str;
                build_prompt_strings(reqs[b], s.instruction_str, s.duration, text_str, lyric_str);

                // Determinism diagnostic
                {
                    auto str_hash = [](const std::string & s) -> uint64_t {
                        uint64_t h = 14695981039346656037ULL;
                        for (char c : s) { h ^= (uint8_t)c; h *= 1099511628211ULL; }
                        return h;
                    };
                    fprintf(stderr, "[DIAG] enc_input_b%d: instr=\"%s\" caption_hash=%016llx lyrics_hash=%016llx text_len=%zu lyric_len=%zu\n",
                            b, s.instruction_str.c_str(),
                            (unsigned long long)str_hash(reqs[b].caption),
                            (unsigned long long)str_hash(reqs[b].lyrics),
                            text_str.size(), lyric_str.size());
                }

                auto text_ids  = bpe_encode(bpe, text_str.c_str(), true);
                auto lyric_ids = bpe_encode(bpe, lyric_str.c_str(), true);
                int  S_text    = (int) text_ids.size();
                int  S_lyric   = (int) lyric_ids.size();

                // LRC capture (batch 0 only)
                if (b == 0 && reqs[0].get_lrc) {
                    s.get_lrc         = true;
                    s.lyric_token_ids = lyric_ids;
                    s.vocal_language  = reqs[0].vocal_language.empty() ? "en" : reqs[0].vocal_language;

                    const char * lang_b = s.vocal_language.c_str();
                    std::string  hdr    = std::string("# Languages\n") + lang_b + "\n\n# Lyric\n";
                    auto         hdr_ids = bpe_encode(bpe, hdr.c_str(), false);
                    s.lyric_start_idx = (int) hdr_ids.size();

                    s.lyric_end_idx = (int) lyric_ids.size();
                    for (int ti = 0; ti < (int) lyric_ids.size(); ti++) {
                        if (lyric_ids[ti] == 151643) { s.lyric_end_idx = ti; break; }
                    }

                    int pure_n = s.lyric_end_idx - s.lyric_start_idx;
                    s.lyric_token_texts.resize(pure_n);
                    if (pure_n > 0) {
                        std::string prev_full;
                        for (int ti = s.lyric_start_idx; ti < s.lyric_end_idx; ti++) {
                            std::vector<int> prefix(lyric_ids.begin() + s.lyric_start_idx,
                                                    lyric_ids.begin() + ti + 1);
                            std::string full;
                            for (int pid : prefix) {
                                if (pid >= 0 && pid < bpe->n_vocab) {
                                    const std::string & bpe_str = bpe->id_to_str[pid];
                                    for (size_t ci = 0; ci < bpe_str.size(); ) {
                                        int adv;
                                        int cp = utf8_codepoint(bpe_str.c_str() + ci, &adv);
                                        bool found = false;
                                        for (int by = 0; by < 256; by++) {
                                            int a2;
                                            int cp2 = utf8_codepoint(bpe->byte2str[by].c_str(), &a2);
                                            if (cp2 == cp && (int) bpe->byte2str[by].size() == adv) {
                                                full += (char)(unsigned char) by;
                                                found = true;
                                                break;
                                            }
                                        }
                                        if (!found) full += '?';
                                        ci += adv;
                                    }
                                }
                            }
                            int idx = ti - s.lyric_start_idx;
                            if (full.size() > prev_full.size()) {
                                s.lyric_token_texts[idx] = full.substr(prev_full.size());
                            } else {
                                s.lyric_token_texts[idx] = "";
                            }
                            prev_full = full;
                        }
                    }

                    fprintf(stderr, "[Encode-Text] LRC: captured %d lyric tokens [%d..%d) of %d total\n",
                            pure_n, s.lyric_start_idx, s.lyric_end_idx, (int) lyric_ids.size());
                }

                main_fwd[b].S_text  = S_text;
                main_fwd[b].S_lyric = S_lyric;

                // ORT text encoder forward
                std::vector<float> text_hidden_ort;
                if (text_enc_ort_forward(te, text_ids.data(), S_text, text_hidden_ort) != 0) {
                    fprintf(stderr, "[Encode-Text] FATAL: text_enc_ort_forward failed\n");
                    return -1;
                }
                main_fwd[b].text_hidden = std::move(text_hidden_ort);

                diag_stats_f32("text_hidden_b0", main_fwd[b].text_hidden.data(),
                               main_fwd[b].text_hidden.size());

                // ORT embed lookup for lyrics
                std::vector<float> lyric_embed_ort;
                if (text_enc_ort_embed_lookup(te, lyric_ids.data(), S_lyric, lyric_embed_ort) != 0) {
                    fprintf(stderr, "[Encode-Text] FATAL: text_enc_ort_embed_lookup failed\n");
                    return -1;
                }
                main_fwd[b].lyric_embed = std::move(lyric_embed_ort);
            }

            if (s.need_enc_switch) {
                for (int b = 0; b < batch_n; b++) {
                    std::string text_str;
                    std::string lyric_str;
                    build_prompt_strings(reqs[b], DIT_INSTR_TEXT2MUSIC, s.duration, text_str, lyric_str);

                    auto text_ids  = bpe_encode(bpe, text_str.c_str(), true);
                    auto lyric_ids = bpe_encode(bpe, lyric_str.c_str(), true);
                    int  S_text    = (int) text_ids.size();
                    int  S_lyric   = (int) lyric_ids.size();

                    nc_fwd[b].S_text  = S_text;
                    nc_fwd[b].S_lyric = S_lyric;
                    std::vector<float> tmp_text;
                    text_enc_ort_forward(te, text_ids.data(), S_text, tmp_text);
                    nc_fwd[b].text_hidden = std::move(tmp_text);
                    std::vector<float> tmp_lyric;
                    text_enc_ort_embed_lookup(te, lyric_ids.data(), S_lyric, tmp_lyric);
                    nc_fwd[b].lyric_embed = std::move(tmp_lyric);
                }
            }

            // Negative prompt encoding
            if (!reqs[0].negative_prompt.empty()) {
                std::string neg_text_str, neg_lyric_str;
                {
                    AceRequest neg_req = reqs[0];
                    neg_req.caption    = reqs[0].negative_prompt;
                    neg_req.lyrics     = "";
                    build_prompt_strings(neg_req, s.instruction_str, s.duration, neg_text_str, neg_lyric_str);
                }
                auto neg_text_ids = bpe_encode(bpe, neg_text_str.c_str(), true);
                s.neg_S_text = (int) neg_text_ids.size();
                text_enc_ort_forward(te, neg_text_ids.data(), s.neg_S_text, s.neg_text_hidden);
                fprintf(stderr, "[Encode-Text] negative_prompt text encoded (ORT): %d tokens\n", s.neg_S_text);
            }

            debug_dump_2d(&s.dbg, "text_hidden", main_fwd[0].text_hidden.data(), main_fwd[0].S_text, H_text);
            debug_dump_2d(&s.dbg, "lyric_embed", main_fwd[0].lyric_embed.data(), main_fwd[0].S_lyric, H_text);
        }

        // Phase B(ORT): condition encoder
        s.per_enc.resize(batch_n);
        s.per_enc_S.resize(batch_n);
        s.per_enc_nc.resize(batch_n);
        s.per_enc_S_nc.assign(batch_n, 0);
        {
            CondEncOrt * ce = store_require_cond_enc_ort(ctx->store, ctx->cond_enc_ort_key);
            if (!ce) {
                fprintf(stderr, "[Encode-Text] FATAL: store_require_cond_enc_ort failed\n");
                return -1;
            }
            ModelHandle ce_guard(ctx->store, ce);
            H_cond = ce->hidden_size;

            // null_condition_emb from the ORT model
            s.null_cond_vec.resize(H_cond);
            if (!ce->null_cond_emb.empty()) {
                memcpy(s.null_cond_vec.data(), ce->null_cond_emb.data(), H_cond * sizeof(float));
            }

            // Negative prompt encoding through cond encoder
            if (!s.neg_text_hidden.empty() && s.neg_S_text > 0) {
                std::vector<float> neg_enc;
                int                neg_enc_S = 0;
                cond_enc_ort_forward(ce, s.neg_text_hidden.data(), s.neg_S_text,
                                     nullptr, 0,
                                     s.timbre_feats.data(), s.S_ref_timbre, neg_enc, &neg_enc_S);
                if (neg_enc_S > 0 && !neg_enc.empty()) {
                    s.null_cond_vec.assign(H_cond, 0.0f);
                    for (int si = 0; si < neg_enc_S; si++)
                        for (int h = 0; h < H_cond; h++)
                            s.null_cond_vec[h] += neg_enc[(size_t)si * H_cond + h];
                    float inv = 1.0f / (float)neg_enc_S;
                    for (int h = 0; h < H_cond; h++) s.null_cond_vec[h] *= inv;
                    fprintf(stderr, "[Encode-Text] negative_prompt encoded (ORT): enc_S=%d\n", neg_enc_S);
                }
            }

            for (int b = 0; b < batch_n; b++) {
                s.timer.reset();
                cond_enc_ort_forward(ce, main_fwd[b].text_hidden.data(), main_fwd[b].S_text,
                                     main_fwd[b].lyric_embed.data(), main_fwd[b].S_lyric,
                                     s.timbre_feats.data(), s.S_ref_timbre,
                                     s.per_enc[b], &s.per_enc_S[b]);
                fprintf(stderr, "[Encode-Text(ORT) Batch%d] %d+%d tokens -> enc_S=%d, %.1f ms\n",
                        b, main_fwd[b].S_text, main_fwd[b].S_lyric, s.per_enc_S[b], s.timer.ms());
            }
            debug_dump_2d(&s.dbg, "enc_hidden", s.per_enc[0].data(), s.per_enc_S[0], H_cond);

            if (s.need_enc_switch) {
                for (int b = 0; b < batch_n; b++) {
                    cond_enc_ort_forward(ce, nc_fwd[b].text_hidden.data(), nc_fwd[b].S_text,
                                         nc_fwd[b].lyric_embed.data(), nc_fwd[b].S_lyric,
                                         s.timbre_feats.data(), s.S_ref_timbre,
                                         s.per_enc_nc[b], &s.per_enc_S_nc[b]);
                    fprintf(stderr, "[Encode-Text(ORT) Batch%d] non-cover: %d+%d tokens -> enc_S=%d\n",
                            b, nc_fwd[b].S_text, nc_fwd[b].S_lyric, s.per_enc_S_nc[b]);
                }
            }
        }

    } else {
    // ═══════════════════════════════════════════════════════════════════
    // GGML PATH: existing text encoder + condition encoder via GGML
    // ═══════════════════════════════════════════════════════════════════

    // Phase A: text encoder.
    {
        Qwen3GGML * te = store_require_text_enc(ctx->store, ctx->text_enc_key);
        if (!te) {
            fprintf(stderr, "[Encode-Text] FATAL: store_require_text_enc failed\n");
            return -1;
        }
        ModelHandle te_guard(ctx->store, te);
        if (!ctx->params.use_fa) {
            te->use_flash_attn = false;
        }

        H_text = te->cfg.hidden_size;

        for (int b = 0; b < batch_n; b++) {
            std::string text_str;
            std::string lyric_str;
            build_prompt_strings(reqs[b], s.instruction_str, s.duration, text_str, lyric_str);

            // Determinism diagnostic: hash the text inputs to detect caption/lyrics/instruction changes
            {
                auto str_hash = [](const std::string & s) -> uint64_t {
                    uint64_t h = 14695981039346656037ULL;
                    for (char c : s) { h ^= (uint8_t)c; h *= 1099511628211ULL; }
                    return h;
                };
                fprintf(stderr, "[DIAG] enc_input_b%d: instr=\"%s\" caption_hash=%016llx lyrics_hash=%016llx text_len=%zu lyric_len=%zu\n",
                        b, s.instruction_str.c_str(),
                        (unsigned long long)str_hash(reqs[b].caption),
                        (unsigned long long)str_hash(reqs[b].lyrics),
                        text_str.size(), lyric_str.size());
            }

            auto text_ids  = bpe_encode(bpe, text_str.c_str(), true);
            auto lyric_ids = bpe_encode(bpe, lyric_str.c_str(), true);
            int  S_text    = (int) text_ids.size();
            int  S_lyric   = (int) lyric_ids.size();

            // Capture lyric token IDs for LRC alignment (batch 0 only)
            if (b == 0 && reqs[0].get_lrc) {
                s.get_lrc         = true;
                s.lyric_token_ids = lyric_ids;
                s.vocal_language  = reqs[0].vocal_language.empty() ? "en" : reqs[0].vocal_language;

                // Build header to find pure lyric boundary
                const char * lang_b = s.vocal_language.c_str();
                std::string  hdr    = std::string("# Languages\n") + lang_b + "\n\n# Lyric\n";
                auto         hdr_ids = bpe_encode(bpe, hdr.c_str(), false);
                s.lyric_start_idx = (int) hdr_ids.size();

                // Find <|endoftext|> (151643) for end boundary
                s.lyric_end_idx = (int) lyric_ids.size();
                for (int ti = 0; ti < (int) lyric_ids.size(); ti++) {
                    if (lyric_ids[ti] == 151643) { s.lyric_end_idx = ti; break; }
                }

                // Incremental per-token text decode (matches Python _decode_tokens_incrementally)
                int pure_n = s.lyric_end_idx - s.lyric_start_idx;
                s.lyric_token_texts.resize(pure_n);
                if (pure_n > 0) {
                    // Use byte-level BPE reverse: each token's id_to_str gives its BPE-encoded
                    // string. The GPT-2 byte encoder maps bytes→unicode codepoints; we must reverse
                    // that to get raw UTF-8 bytes. For the alignment, the key is just whether a
                    // token contains a newline — so we do a simpler decode via id_to_str.
                    std::string prev_full;
                    for (int ti = s.lyric_start_idx; ti < s.lyric_end_idx; ti++) {
                        // Decode all tokens from start to current position
                        std::vector<int> prefix(lyric_ids.begin() + s.lyric_start_idx,
                                                lyric_ids.begin() + ti + 1);
                        // Use byte-level reverse to get text
                        std::string full;
                        for (int pid : prefix) {
                            if (pid >= 0 && pid < bpe->n_vocab) {
                                const std::string & bpe_str = bpe->id_to_str[pid];
                                // Reverse GPT-2 byte encoding: each BPE char → original byte
                                for (size_t ci = 0; ci < bpe_str.size(); ) {
                                    int adv;
                                    int cp = utf8_codepoint(bpe_str.c_str() + ci, &adv);
                                    // Find which byte maps to this codepoint
                                    bool found = false;
                                    for (int by = 0; by < 256; by++) {
                                        int a2;
                                        int cp2 = utf8_codepoint(bpe->byte2str[by].c_str(), &a2);
                                        if (cp2 == cp && (int) bpe->byte2str[by].size() == adv) {
                                            full += (char)(unsigned char) by;
                                            found = true;
                                            break;
                                        }
                                    }
                                    if (!found) full += '?';
                                    ci += adv;
                                }
                            }
                        }
                        // Token contribution = new bytes since last prefix
                        int idx = ti - s.lyric_start_idx;
                        if (full.size() > prev_full.size()) {
                            s.lyric_token_texts[idx] = full.substr(prev_full.size());
                        } else {
                            s.lyric_token_texts[idx] = "";
                        }
                        prev_full = full;
                    }
                }

                fprintf(stderr, "[Encode-Text] LRC: captured %d lyric tokens [%d..%d) of %d total\n",
                        pure_n, s.lyric_start_idx, s.lyric_end_idx, (int) lyric_ids.size());
            }

            main_fwd[b].S_text  = S_text;
            main_fwd[b].S_lyric = S_lyric;
            main_fwd[b].text_hidden.resize((size_t) H_text * S_text);
            qwen3_forward(te, text_ids.data(), S_text, main_fwd[b].text_hidden.data());

            // Determinism diagnostic: text_hidden stats after qwen3_forward
            diag_stats_f32("text_hidden_b0", main_fwd[b].text_hidden.data(),
                           main_fwd[b].text_hidden.size());

            main_fwd[b].lyric_embed.resize((size_t) H_text * S_lyric);
            qwen3_embed_lookup(te, lyric_ids.data(), S_lyric, main_fwd[b].lyric_embed.data());
        }

        if (s.need_enc_switch) {
            for (int b = 0; b < batch_n; b++) {
                std::string text_str;
                std::string lyric_str;
                build_prompt_strings(reqs[b], DIT_INSTR_TEXT2MUSIC, s.duration, text_str, lyric_str);

                auto text_ids  = bpe_encode(bpe, text_str.c_str(), true);
                auto lyric_ids = bpe_encode(bpe, lyric_str.c_str(), true);
                int  S_text    = (int) text_ids.size();
                int  S_lyric   = (int) lyric_ids.size();

                nc_fwd[b].S_text  = S_text;
                nc_fwd[b].S_lyric = S_lyric;
                nc_fwd[b].text_hidden.resize((size_t) H_text * S_text);
                qwen3_forward(te, text_ids.data(), S_text, nc_fwd[b].text_hidden.data());
                nc_fwd[b].lyric_embed.resize((size_t) H_text * S_lyric);
                qwen3_embed_lookup(te, lyric_ids.data(), S_lyric, nc_fwd[b].lyric_embed.data());
            }
        }

        // Debug dump of sample 0 while text_hidden and lyric_embed are live.
        // NEGATIVE PROMPT Phase A: encode negative text while text encoder is loaded
        if (!reqs[0].negative_prompt.empty()) {
            std::string neg_text_str, neg_lyric_str;
            {
                AceRequest neg_req = reqs[0];
                neg_req.caption    = reqs[0].negative_prompt;
                neg_req.lyrics     = "";
                build_prompt_strings(neg_req, s.instruction_str, s.duration, neg_text_str, neg_lyric_str);
            }
            auto neg_text_ids = bpe_encode(bpe, neg_text_str.c_str(), true);
            s.neg_S_text = (int) neg_text_ids.size();
            s.neg_text_hidden.resize((size_t) H_text * s.neg_S_text);
            qwen3_forward(te, neg_text_ids.data(), s.neg_S_text, s.neg_text_hidden.data());
            fprintf(stderr, "[Encode-Text] negative_prompt text encoded: %d tokens\n", s.neg_S_text);
        }
        debug_dump_2d(&s.dbg, "text_hidden", main_fwd[0].text_hidden.data(), main_fwd[0].S_text, H_text);
        debug_dump_2d(&s.dbg, "lyric_embed", main_fwd[0].lyric_embed.data(), main_fwd[0].S_lyric, H_text);
    }

    // Phase B: condition encoder.
    s.per_enc.resize(batch_n);
    s.per_enc_S.resize(batch_n);
    s.per_enc_nc.resize(batch_n);
    s.per_enc_S_nc.assign(batch_n, 0);
    {
        CondGGML * ce = store_require_cond_enc(ctx->store, ctx->cond_enc_key);
        if (!ce) {
            fprintf(stderr, "[Encode-Text] FATAL: store_require_cond_enc failed\n");
            return -1;
        }
        ModelHandle ce_guard(ctx->store, ce);
        if (!ctx->params.use_fa) {
            ce->use_flash_attn = false;
        }
        ce->clamp_fp16 = ctx->params.clamp_fp16;

        H_cond = ce->lyric_cfg.hidden_size;

        // null_condition_emb lives on the DiTMeta. Empty when the model has none.
        s.null_cond_vec.resize(H_cond);
        if (!ctx->meta->null_cond_cpu.empty()) {
            memcpy(s.null_cond_vec.data(), ctx->meta->null_cond_cpu.data(), H_cond * sizeof(float));
        }

        // NEGATIVE PROMPT Phase B: cond-encode neg text, mean-pool into null_cond_vec
        if (!s.neg_text_hidden.empty() && s.neg_S_text > 0) {
            std::vector<float> neg_enc;
            int                neg_enc_S = 0;
            cond_ggml_forward(ce, s.neg_text_hidden.data(), s.neg_S_text,
                              nullptr, 0,
                              s.timbre_feats.data(), s.S_ref_timbre, neg_enc, &neg_enc_S);
            if (neg_enc_S > 0 && !neg_enc.empty()) {
                s.null_cond_vec.assign(H_cond, 0.0f);
                for (int si = 0; si < neg_enc_S; si++)
                    for (int h = 0; h < H_cond; h++)
                        s.null_cond_vec[h] += neg_enc[(size_t)si * H_cond + h];
                float inv = 1.0f / (float)neg_enc_S;
                for (int h = 0; h < H_cond; h++) s.null_cond_vec[h] *= inv;
                fprintf(stderr, "[Encode-Text] negative_prompt encoded: enc_S=%d\n", neg_enc_S);
            }
        }
        for (int b = 0; b < batch_n; b++) {
            s.timer.reset();
            cond_ggml_forward(ce, main_fwd[b].text_hidden.data(), main_fwd[b].S_text, main_fwd[b].lyric_embed.data(),
                              main_fwd[b].S_lyric, s.timbre_feats.data(), s.S_ref_timbre, s.per_enc[b],
                              &s.per_enc_S[b]);
            fprintf(stderr, "[Encode-Text Batch%d] %d+%d tokens -> enc_S=%d, %.1f ms\n", b, main_fwd[b].S_text,
                    main_fwd[b].S_lyric, s.per_enc_S[b], s.timer.ms());
        }
        debug_dump_2d(&s.dbg, "enc_hidden", s.per_enc[0].data(), s.per_enc_S[0], H_cond);

        if (s.need_enc_switch) {
            for (int b = 0; b < batch_n; b++) {
                cond_ggml_forward(ce, nc_fwd[b].text_hidden.data(), nc_fwd[b].S_text, nc_fwd[b].lyric_embed.data(),
                                  nc_fwd[b].S_lyric, s.timbre_feats.data(), s.S_ref_timbre, s.per_enc_nc[b],
                                  &s.per_enc_S_nc[b]);
                fprintf(stderr, "[Encode-Text Batch%d] non-cover: %d+%d tokens -> enc_S=%d\n", b, nc_fwd[b].S_text,
                        nc_fwd[b].S_lyric, s.per_enc_S_nc[b]);
            }
        }
    }
    } // end GGML else-branch

    // find max s.enc_S across both encodings (cover + text2music),
    // pad shorter encodings with null_cond, stack into [H, s.max_enc_S, N]
    s.max_enc_S = 0;
    for (int b = 0; b < batch_n; b++) {
        if (s.per_enc_S[b] > s.max_enc_S) {
            s.max_enc_S = s.per_enc_S[b];
        }
        if (s.need_enc_switch && s.per_enc_S_nc[b] > s.max_enc_S) {
            s.max_enc_S = s.per_enc_S_nc[b];
        }
    }
    s.enc_S = s.max_enc_S;

    s.enc_hidden.resize((size_t) H_cond * s.max_enc_S * batch_n);
    for (int b = 0; b < batch_n; b++) {
        float * dst = s.enc_hidden.data() + (size_t) b * s.max_enc_S * H_cond;
        memcpy(dst, s.per_enc[b].data(), (size_t) s.per_enc_S[b] * H_cond * sizeof(float));
        for (int si = s.per_enc_S[b]; si < s.max_enc_S; si++) {
            memcpy(dst + si * H_cond, s.null_cond_vec.data(), H_cond * sizeof(float));
        }
    }

    // pad and stack text2music encoding (same s.max_enc_S for graph compatibility)
    if (s.need_enc_switch) {
        s.enc_hidden_nc.resize((size_t) H_cond * s.max_enc_S * batch_n);
        s.per_enc_S_nc_final.resize(batch_n);
        for (int b = 0; b < batch_n; b++) {
            float * dst = s.enc_hidden_nc.data() + (size_t) b * s.max_enc_S * H_cond;
            memcpy(dst, s.per_enc_nc[b].data(), (size_t) s.per_enc_S_nc[b] * H_cond * sizeof(float));
            for (int si = s.per_enc_S_nc[b]; si < s.max_enc_S; si++) {
                memcpy(dst + si * H_cond, s.null_cond_vec.data(), H_cond * sizeof(float));
            }
            s.per_enc_S_nc_final[b] = s.per_enc_S_nc[b];
        }
    }

    if (batch_n > 1) {
        fprintf(stderr, "[Encode-Text] Per-batch encoding done: max_enc_S=%d\n", s.max_enc_S);
    }

    return 0;
}

int ops_build_context(const AceSynth * ctx, const AceRequest * reqs, int batch_n, SynthState & s) {
    // Build s.context: [batch_n, s.T, s.ctx_ch] = src_latents[64] + chunk_mask[64]
    // Cover/Lego/Repaint: shared s.context replicated (s.cover_latents from src_audio).
    // Passthrough: per-batch detokenized FSQ codes + silence padding, mask = 1.0.
    // Text2music: silence only, mask = 1.0.
    s.repaint_t0 = 0, s.repaint_t1 = 0;
    if (s.is_repaint) {
        s.repaint_t0 = (int) (s.rs * 48000.0f / 1920.0f);
        s.repaint_t1 = (int) (s.re * 48000.0f / 1920.0f);
        if (s.repaint_t0 < 0) {
            s.repaint_t0 = 0;
        }
        if (s.repaint_t1 > s.T) {
            s.repaint_t1 = s.T;
        }
        if (s.repaint_t0 > s.T) {
            s.repaint_t0 = s.T;
        }
        fprintf(stderr, "[Build-Context] Latent frames: [%d, %d) / %d\n", s.repaint_t0, s.repaint_t1, s.T);
    }

    s.context.resize(batch_n * s.T * s.ctx_ch);

    if (s.use_source_context && s.have_cover) {
        // Cover/Lego/Repaint: build once, replicate (s.cover_latents are shared)
        std::vector<float> context_single(s.T * s.ctx_ch);
        for (int t = 0; t < s.T; t++) {
            bool          in_region = (s.is_repaint || s.is_lego_region) && t >= s.repaint_t0 && t < s.repaint_t1;
            // repaint silences the zone (DiT generates fresh there).
            // lego keeps full cover everywhere (DiT hears the whole backing track).
            const float * src;
            if (s.is_repaint && in_region) {
                src = ctx->meta->silence_full.data() + t * s.Oc;
            } else {
                src = (t < s.T_cover) ? s.cover_latents.data() + t * s.Oc : ctx->meta->silence_full.data() + t * s.Oc;
            }
            // region tasks: explicit 0/1 mask. all others: 1.0 (training distribution).
            float mask_val;
            if (s.is_repaint || s.is_lego_region) {
                mask_val = in_region ? 1.0f : 0.0f;
            } else {
                mask_val = 1.0f;  // training distribution: only 0/1 seen during training
            }
            for (int c = 0; c < s.Oc; c++) {
                context_single[t * s.ctx_ch + c] = src[c];
            }
            for (int c = 0; c < s.Oc; c++) {
                context_single[t * s.ctx_ch + s.Oc + c] = mask_val;
            }
        }
        for (int b = 0; b < batch_n; b++) {
            memcpy(s.context.data() + b * s.T * s.ctx_ch, context_single.data(), s.T * s.ctx_ch * sizeof(float));
        }
    } else {
        // Per-batch context from audio_codes or silence (text2music).
        // use_source_context with neither cover nor codes is an invalid state:
        // the orchestrator promised source context but provided nothing to condition on.
        if (s.use_source_context && !s.have_codes) {
            fprintf(stderr, "[Build-Context] FATAL: use_source_context but no cover_latents and no audio_codes\n");
            return -1;
        }

        // Decode batch items with audio codes through the FSQ detokenizer, cached
        // in CPU buffers before the fill loop so the detokenizer is held for the
        // shortest possible window under STRICT.
        std::vector<std::vector<float>> decoded_per_b(batch_n);
        std::vector<int>                decoded_T_per_b(batch_n, 0);

        // s.have_codes is already posed by ops_resolve_params over the same batch.
        bool any_codes = s.have_codes;

        if (any_codes) {
            DetokGGML * detok = store_require_fsq_detok(ctx->store, ctx->fsq_detok_key);
            if (!detok) {
                fprintf(stderr, "[Build-Context] FATAL: store_require_fsq_detok failed\n");
                return -1;
            }
            ModelHandle detok_guard(ctx->store, detok);
            if (!ctx->params.use_fa) {
                detok->use_flash_attn = false;
            }

            for (int b = 0; b < batch_n; b++) {
                const std::vector<int> & codes_b = s.per_codes[b];
                if (codes_b.empty()) {
                    continue;
                }
                s.timer.reset();
                int T_5Hz        = (int) codes_b.size();
                int T_25Hz_codes = T_5Hz * 5;
                decoded_per_b[b].resize((size_t) T_25Hz_codes * s.Oc);

                int ret = detok_ggml_decode(detok, codes_b.data(), T_5Hz, decoded_per_b[b].data());
                if (ret < 0) {
                    fprintf(stderr, "[Build-Context Batch%d] FATAL: detokenizer decode failed\n", b);
                    return -1;
                }
                fprintf(stderr, "[Build-Context Batch%d] Detokenizer: %.1f ms, %d codes\n", b, s.timer.ms(), T_5Hz);

                decoded_T_per_b[b] = T_25Hz_codes < s.T ? T_25Hz_codes : s.T;
                if (b == 0) {
                    debug_dump_2d(&s.dbg, "detok_output", decoded_per_b[b].data(), T_25Hz_codes, s.Oc);
                }
            }
        }

        // Fill s.context: decoded latents then silence, mask = 1.0 (training distribution).
        // The detokenizer is already released at this point; the CPU buffers in
        // decoded_per_b carry everything we need.
        for (int b = 0; b < batch_n; b++) {
            float *       ctx_dst   = s.context.data() + b * s.T * s.ctx_ch;
            const float * decoded   = decoded_per_b[b].data();
            int           decoded_T = decoded_T_per_b[b];

            for (int t = 0; t < s.T; t++) {
                const float * src =
                    (t < decoded_T) ? decoded + t * s.Oc : ctx->meta->silence_full.data() + (t - decoded_T) * s.Oc;
                for (int c = 0; c < s.Oc; c++) {
                    ctx_dst[t * s.ctx_ch + c] = src[c];
                }
                for (int c = 0; c < s.Oc; c++) {
                    ctx_dst[t * s.ctx_ch + s.Oc + c] = 1.0f;
                }
            }
        }
    }

    return 0;
}

void ops_build_context_silence(const AceSynth * ctx, int batch_n, SynthState & s) {
    // Cover mode: build silence s.context for audio_cover_strength switching
    // When step >= s.cover_steps, DiT switches from cover s.context to silence s.context
    // Repaint/lego_region: mask handles region; s.context switch never applies
    s.cover_steps = -1;
    if (s.use_source_context && !s.is_repaint && !s.is_lego_region) {
        float cover_strength = s.rr.audio_cover_strength;
        if (cover_strength < 1.0f) {
            // Build silence s.context: all frames use silence_latent
            std::vector<float> silence_single(s.T * s.ctx_ch);
            for (int t = 0; t < s.T; t++) {
                const float * src = ctx->meta->silence_full.data() + t * s.Oc;
                for (int c = 0; c < s.Oc; c++) {
                    silence_single[t * s.ctx_ch + c] = src[c];
                }
                for (int c = 0; c < s.Oc; c++) {
                    silence_single[t * s.ctx_ch + s.Oc + c] = 1.0f;
                }
            }
            s.context_silence.resize(batch_n * s.T * s.ctx_ch);
            for (int b = 0; b < batch_n; b++) {
                memcpy(s.context_silence.data() + b * s.T * s.ctx_ch, silence_single.data(),
                       s.T * s.ctx_ch * sizeof(float));
            }
            s.cover_steps = (int) ((float) s.num_steps * cover_strength);
            fprintf(stderr, "[Context-Silence] audio_cover_strength=%.2f -> switch at step %d/%d\n", cover_strength,
                    s.cover_steps, s.num_steps);
        }
    }
}

void ops_init_noise(const AceSynth * ctx, const AceRequest * reqs, int batch_n, SynthState & s) {
    // Generate N s.noise samples (Philox4x32-10, matches torch.randn on CUDA with bf16).
    // Each batch item uses its own seed from the request.
    s.noise.resize(batch_n * s.Oc * s.T);
    s.seeds.resize(batch_n);
    for (int b = 0; b < batch_n; b++) {
        float * dst = s.noise.data() + b * s.Oc * s.T;
        s.seeds[b]  = reqs[b].seed;
        philox_randn(reqs[b].seed, dst, s.Oc * s.T, /*bf16_round=*/true);
        fprintf(stderr, "[Init-Noise Batch%d] Philox noise seed=%lld, [%d, %d]%s\n", b, (long long) reqs[b].seed, s.T,
                s.Oc, s.use_sde ? " (SDE)" : "");
    }

    // cover_noise_strength: blend initial noise with clean source latents.
    // xt = nearest_t * noise + (1 - nearest_t) * clean_latents, then truncate schedule.
    // the FSQ roundtrip degrades cover_latents for context conditioning, but noise
    // blending needs the original clean VAE latents. noise_blend_latents holds the
    // clean copy when FSQ was applied; otherwise fall back to cover_latents (already clean).
    if (s.use_source_context && s.have_cover && s.rr.cover_noise_strength > 0.0f) {
        const std::vector<float> & blend_src = s.noise_blend_latents.empty() ? s.cover_latents : s.noise_blend_latents;
        float                      effective_noise_level = 1.0f - s.rr.cover_noise_strength;
        // find nearest timestep in s.schedule
        int                        start_idx             = 0;
        float                      best_dist             = fabsf(s.schedule[0] - effective_noise_level);
        for (int i = 1; i < s.num_steps; i++) {
            float dist = fabsf(s.schedule[i] - effective_noise_level);
            if (dist < best_dist) {
                best_dist = dist;
                start_idx = i;
            }
        }
        float nearest_t = s.schedule[start_idx];
        // blend: xt = nearest_t * s.noise + (1 - nearest_t) * clean_latents
        for (int b = 0; b < batch_n; b++) {
            float * n = s.noise.data() + b * s.Oc * s.T;
            for (int t = 0; t < s.T; t++) {
                int           t_src = t < s.T_cover ? t : s.T_cover - 1;
                const float * src   = blend_src.data() + t_src * s.Oc;
                for (int c = 0; c < s.Oc; c++) {
                    int idx = t * s.Oc + c;
                    n[idx]  = nearest_t * n[idx] + (1.0f - nearest_t) * src[c];
                }
            }
        }
        // truncate s.schedule
        bool use_rescale = (s.rr.cover_noise_method == "rescale");

        if (use_rescale) {
            // RESCALE: rebuild schedule with full step count in [start_sigma, 0] range.
            // Preserves the shift distribution within the reduced range.
            float start_sigma = nearest_t;
            float sh = s.shift;
            for (int i = 0; i < s.num_steps; i++) {
                float t = 1.0f - (float) i / (float) s.num_steps;
                float sigma = sh * t / (1.0f + (sh - 1.0f) * t);
                s.schedule[i] = sigma * start_sigma;
            }
            // num_steps unchanged — full step budget
            if (s.cover_steps >= 0) {
                s.cover_steps = (int) ((float) s.num_steps * s.rr.audio_cover_strength);
            }
            fprintf(stderr, "[Noise] Rescale: %d steps in [%.4f -> 0], shift=%.1f\n",
                    s.num_steps, start_sigma, sh);
        } else {
            // TRUNCATE (original behavior): remove early steps
            s.schedule.erase(s.schedule.begin(), s.schedule.begin() + start_idx);
            s.num_steps = (int) s.schedule.size();
            if (s.cover_steps >= 0) {
                s.cover_steps = (int) ((float) s.num_steps * s.rr.audio_cover_strength);
            }
            fprintf(stderr, "[Noise] Truncate: %d steps remaining (removed %d)\n",
                    s.num_steps, start_idx);
        }
        fprintf(stderr,
                "[Init-Noise] cover_noise_strength=%.2f -> noise_level=%.4f, nearest_t=%.4f, remaining_steps=%d\n",
                s.rr.cover_noise_strength, effective_noise_level, nearest_t, s.num_steps);
    }

    // DiT Generate
    s.output.resize(batch_n * s.Oc * s.T);

    // Per-batch sequence lengths for attention padding masks.
    // Within a synth_batch_size group, all elements share the same s.T (same codes),
    // so s.per_S[b] = s.S for all b. The s.per_enc_S[] array has real encoder lengths
    // from per-batch text encoding above.
    // These become meaningful when the server/CLI batches requests with different s.T.
    s.per_S.assign(batch_n, s.S);

    // Debug dumps (sample 0)
    debug_dump_2d(&s.dbg, "noise", s.noise.data(), s.T, s.Oc);
    debug_dump_2d(&s.dbg, "context", s.context.data(), s.T, s.ctx_ch);

    fprintf(stderr, "[Init-Noise] Starting: T=%d, S=%d, enc_S=%d, steps=%d, batch=%d%s\n", s.T, s.S, s.enc_S,
            s.num_steps, batch_n, s.use_source_context ? " (cover)" : "");
}

int ops_dit_generate(const AceSynth * ctx, int batch_n, SynthState & s, bool (*cancel)(void *), void * cancel_data) {
#ifdef HOT_STEP_TRT
    // TRT path: if DiT model is ONNX (file or directory), use TensorRT acceleration
    if (dit_ends_with_onnx(ctx->dit_key.path.c_str())) {

        fprintf(stderr, "[DiT-Generate] Using TRT path: %s\n", ctx->dit_key.path.c_str());

        // Static TRT context — built once, reused across requests
        static DitTrt s_trt;
        static bool   s_trt_ready = false;
        static std::string s_trt_onnx_path;

        if (!s_trt_ready || s_trt_onnx_path != ctx->dit_key.path) {
            // Need to build or load the TRT engine
            if (s_trt_ready) {
                dit_trt_free(&s_trt);
                // dit_trt_free already nulls pointers and clears base_weights
                s_trt.current_adapter.clear();
                s_trt_ready = false;
            }

            // Resolve actual ONNX file: dit_key.path may be a directory or .onnx file
            std::string onnx_path;
            {
                const std::string & p = ctx->dit_key.path;
                size_t plen = p.size();
                if (plen >= 5 && p.compare(plen - 5, 5, ".onnx") == 0) {
                    onnx_path = p;  // already a .onnx file path
                } else {
                    // Directory: find the dit*.onnx file inside
#ifdef _WIN32
                    std::string pattern = p + "\\*.onnx";
                    WIN32_FIND_DATAA fd;
                    HANDLE h = FindFirstFileA(pattern.c_str(), &fd);
                    if (h != INVALID_HANDLE_VALUE) {
                        do {
                            std::string fname(fd.cFileName);
                            std::string lower = fname;
                            for (auto & c : lower) c = (char)tolower((unsigned char)c);
                            if (lower.find("dit") != std::string::npos) {
                                onnx_path = p + "\\" + fname;
                                break;
                            }
                        } while (FindNextFileA(h, &fd));
                        FindClose(h);
                    }
#else
                    DIR * d = opendir(p.c_str());
                    if (d) {
                        struct dirent * ent;
                        while ((ent = readdir(d)) != nullptr) {
                            std::string fname(ent->d_name);
                            std::string lower = fname;
                            for (auto & c : lower) c = (char)tolower((unsigned char)c);
                            if (lower.find("dit") != std::string::npos &&
                                lower.size() >= 5 && lower.compare(lower.size() - 5, 5, ".onnx") == 0) {
                                onnx_path = p + "/" + fname;
                                break;
                            }
                        }
                        closedir(d);
                    }
#endif
                    if (onnx_path.empty()) {
                        fprintf(stderr, "[DiT-Generate] FATAL: no dit*.onnx found in %s\n", p.c_str());
                        return -1;
                    }
                }
            }

            // Engine path: same directory, same name but .engine extension
            std::string engine_path = onnx_path.substr(0, onnx_path.size() - 5) + ".engine";

            // Check if engine exists
            FILE * ef = fopen(engine_path.c_str(), "rb");
            if (ef) {
                fclose(ef);
                fprintf(stderr, "[DiT-Generate] Loading cached TRT engine: %s\n", engine_path.c_str());
            } else {
                // Build engine from ONNX (slow, first run only)
                fprintf(stderr, "[DiT-Generate] Building TRT engine from ONNX...\n");
                if (!dit_trt_build(onnx_path.c_str(), engine_path.c_str())) {
                    fprintf(stderr, "[DiT-Generate] FATAL: TRT engine build failed\n");
                    return -1;
                }
            }

            // Load engine + refit with base weights
            if (!dit_trt_load(&s_trt, engine_path.c_str(), onnx_path.c_str())) {
                fprintf(stderr, "[DiT-Generate] FATAL: TRT engine load failed\n");
                return -1;
            }
            s_trt_onnx_path = onnx_path;
            s_trt_ready = true;
        }

        // LoRA adapter refitting: apply/revert/switch as needed
        // Non-bf16 engines (fp32 from FP8, etc.) need different merge math.
        if (s_trt.io_dtype != DitTrt::IO_BF16) {
            if (!ctx->dit_key.adapter_path.empty()) {
                fprintf(stderr, "[Adapter-TRT] WARNING: LoRA adapters not yet supported with this engine I/O dtype, skipping\n");
                fflush(stderr);
            }
        } else {
            auto t_adapter_wall = std::chrono::steady_clock::now();
            const std::string& want_adapter = ctx->dit_key.adapter_path;
            const std::string& have_adapter = s_trt.current_adapter;

            if (want_adapter.empty() && !have_adapter.empty()) {
                // Revert to base model
                fprintf(stderr, "[Adapter-TRT] Reverting adapter '%s' → base\n",
                        have_adapter.c_str());
                fflush(stderr);
                int64_t ms = adapter_trt_revert(&s_trt);
                if (ms >= 0) {
                    fprintf(stderr, "[Adapter-TRT] Reverted in %lld ms\n", (long long)ms);
                    fflush(stderr);
                }
            } else if (!want_adapter.empty() && want_adapter != have_adapter) {
                // Apply new adapter (or switch from one to another)
                if (!have_adapter.empty()) {
                    fprintf(stderr, "[Adapter-TRT] Switching adapter '%s' → '%s'\n",
                            have_adapter.c_str(), want_adapter.c_str());
                    fflush(stderr);
                    // Revert to base first, then apply new
                    adapter_trt_revert(&s_trt);
                }
                fprintf(stderr, "[Adapter-TRT] Applying adapter: %s (scale=%.2f)\n",
                        want_adapter.c_str(), ctx->dit_key.adapter_scale);
                fflush(stderr);
                int64_t ms = adapter_trt_apply(&s_trt, want_adapter.c_str(),
                                               ctx->dit_key.adapter_scale);
                if (ms < 0) {
                    fprintf(stderr, "[Adapter-TRT] WARNING: adapter apply failed, continuing with base\n");
                } else {
                    fprintf(stderr, "[Adapter-TRT] Applied in %lld ms\n", (long long)ms);
                }
                fflush(stderr);
            }
            // else: want == have, no change needed

            auto t_adapter_end = std::chrono::steady_clock::now();
            auto adapter_wall_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
                t_adapter_end - t_adapter_wall).count();
            if (adapter_wall_ms > 50) {
                fprintf(stderr, "[Adapter-TRT] Wall clock total: %lld ms\n", (long long)adapter_wall_ms);
                fflush(stderr);
            }
        }

        s.timer.reset();
        int dit_rc = dit_trt_generate(
            &s_trt, s.noise.data(), s.context.data(), s.enc_hidden.data(), s.enc_S, s.T, batch_n, s.num_steps,
            s.schedule.data(), s.output.data(), s.guidance_scale, &s.dbg,
            s.context_silence.empty() ? nullptr : s.context_silence.data(), s.cover_steps, cancel, cancel_data,
            s.per_S.data(), s.per_enc_S.data(), s.enc_hidden_nc.empty() ? nullptr : s.enc_hidden_nc.data(),
            s.per_enc_S_nc_final.empty() ? nullptr : s.per_enc_S_nc_final.data(), s.use_sde, s.seeds.data(),
            ctx->params.use_batch_cfg,
            s.null_cond_vec.empty() ? nullptr : s.null_cond_vec.data());
        if (dit_rc != 0) {
            return -1;
        }
        fprintf(stderr, "[DiT-Generate] TRT Total: %.1f ms (%.1f ms/sample)\n", s.timer.ms(), s.timer.ms() / batch_n);
        fflush(stderr);

        // Release TRT engine from VRAM if eviction policy says so.
        // EVICT_STRICT (default) = unload after use, like GGML's store_release.
        // EVICT_NEVER ("Keep DiT & VAE loaded") = keep engine in VRAM.
        if (store_get_policy(ctx->store) == EVICT_STRICT) {
            dit_trt_free(&s_trt);
            s_trt_ready = false;
            s_trt_onnx_path.clear();
            fprintf(stderr, "[DiT-TRT] Engine unloaded from VRAM\n");
        }

        // (fall through to latent post-processing below)
        goto post_dit;
    }
#endif // HOT_STEP_TRT

    // GGML path (default)
    {
        DiTGGML * dit = store_require_dit(ctx->store, ctx->dit_key);
        if (!dit) {
            fprintf(stderr, "[DiT-Generate] FATAL: store_require_dit failed\n");
            return -1;
        }
        ModelHandle dit_guard(ctx->store, dit);
        if (!ctx->params.use_fa) {
            dit->use_flash_attn = false;
        }

        s.timer.reset();
        int dit_rc = dit_ggml_generate(
            dit, s.noise.data(), s.context.data(), s.enc_hidden.data(), s.enc_S, s.T, batch_n, s.num_steps,
            s.schedule.data(), s.output.data(), s.guidance_scale, &s.dbg,
            s.context_silence.empty() ? nullptr : s.context_silence.data(), s.cover_steps, cancel, cancel_data,
            s.per_S.data(), s.per_enc_S.data(), s.enc_hidden_nc.empty() ? nullptr : s.enc_hidden_nc.data(),
            s.per_enc_S_nc_final.empty() ? nullptr : s.per_enc_S_nc_final.data(), s.use_sde, s.seeds.data(),
            ctx->params.use_batch_cfg,
            s.null_cond_vec.empty() ? nullptr : s.null_cond_vec.data());
        if (dit_rc != 0) {
            return -1;
        }
        fprintf(stderr, "[DiT-Generate] Total: %.1f ms (%.1f ms/sample)\n", s.timer.ms(), s.timer.ms() / batch_n);

        // LRC alignment: run while the DiT is still held (dit_guard in scope).
        if (s.get_lrc) {
            ops_lrc_extract(ctx, dit, batch_n, s);
        }
    }

#ifdef HOT_STEP_TRT
post_dit:
#endif

    // Latent post-processing before VAE decode: pred = pred * rescale + shift.
    // Skipped at defaults (1.0 / 0.0).
    if (s.rr.latent_rescale != 1.0f || s.rr.latent_shift != 0.0f) {
        fprintf(stderr, "[DiT-Generate] Latent post: shift=%.3f rescale=%.3f\n", s.rr.latent_shift,
                s.rr.latent_rescale);
        const int n = (int) s.output.size();
        for (int i = 0; i < n; i++) {
            s.output[i] = s.output[i] * s.rr.latent_rescale + s.rr.latent_shift;
        }
    }

    // ── Latent output RMS guard ─────────────────────────────────────────
    // Detect latent blowup (common with XL models at high noise + high CFG)
    // and auto-rescale to prevent pure-noise/static VAE output.
    // NOTE: Turbo models naturally output RMS ~1.0-1.1 (guidance=1.0 distillation).
    // Threshold must be high enough to avoid false positives on turbo output.
    {
        const int  n       = (int) s.output.size();
        double     sum_sq  = 0.0;
        for (int i = 0; i < n; i++) {
            sum_sq += (double) s.output[i] * s.output[i];
        }
        float out_rms = (float) sqrt(sum_sq / (double) n);
        const float RMS_THRESHOLD = 3.0f;
        const float TARGET_RMS    = 0.3f;
        if (out_rms > RMS_THRESHOLD) {
            float gain = TARGET_RMS / out_rms;
            fprintf(stderr,
                    "[DiT-Generate] WARNING: output RMS=%.3f exceeds safe range (>%.1f), "
                    "auto-rescaling by %.4f to target RMS=%.1f. "
                    "This may indicate XL model divergence at high noise + high CFG.\n",
                    out_rms, RMS_THRESHOLD, gain, TARGET_RMS);
            for (int i = 0; i < n; i++) {
                s.output[i] *= gain;
            }
        }
    }

    debug_dump_2d(&s.dbg, "dit_output", s.output.data(), s.T, s.Oc);

    // LRC alignment: run while the DiT is still held (dit_guard in scope).
    // This avoids a full DiT eviction + reload + adapter merge that would
    // happen if we ran LRC as a separate phase under EVICT_STRICT.
    // NOTE: LRC requires the GGML DiT model for cross-attention maps.
    //       Not supported in TRT path yet.
#ifdef HOT_STEP_TRT
    if (s.get_lrc && ctx->dit_key.path.size() > 5 &&
        ctx->dit_key.path.compare(ctx->dit_key.path.size() - 5, 5, ".onnx") == 0) {
        fprintf(stderr, "[DiT-Generate] LRC not yet supported with TRT DiT, skipping\n");
        s.get_lrc = false;
    }
#endif

    return 0;
}

int ops_vae_decode(const AceSynth * ctx,
                   int              batch_n,
                   AceAudio *       out,
                   SynthState &     s,
                   bool (*cancel)(void *),
                   void * cancel_data) {
    // ── Decide: ORT path or GGML path ──────────────────────────────
    bool use_ort = s.rr.use_ort_vae && !ctx->onnx_vae_path.empty();

    // Acquire the appropriate decoder module.
    // We try ORT first; on failure, fall back to GGML.
    VAEGGML * vae_ggml = nullptr;
    VaeOrt  * vae_ort  = nullptr;

    if (use_ort) {
        vae_ort = store_require_vae_dec_ort(ctx->store, ctx->vae_dec_ort_key);
        if (!vae_ort) {
            fprintf(stderr, "[VAE-Decode] WARNING: ORT VAE load failed, falling back to GGML\n");
            use_ort = false;
        } else {
            fprintf(stderr, "[VAE-Decode] Using ORT VAE decoder\n");
        }
    }
    if (!use_ort) {
        vae_ggml = store_require_vae_dec(ctx->store, ctx->vae_dec_key);
        if (!vae_ggml) {
            fprintf(stderr, "[VAE-Decode] FATAL: store_require_vae_dec failed\n");
            return -1;
        }
    }

    // RAII guard: whichever module we loaded, release it on scope exit
    ModelHandle vae_guard(ctx->store, use_ort ? (void *)vae_ort : (void *)vae_ggml);
    // Latent splice for repaint/lego: keep s.output inside [t0, t1), copy
    // s.cover_latents outside. Hard cut at frame boundary, the VAE tiled
    // decoder smooths the seam in the waveform.
    bool have_region = s.is_repaint || s.is_lego_region;
    if (have_region && s.have_cover && s.T_cover > 0) {
        int copy_T = s.T_cover < s.T ? s.T_cover : s.T;
        for (int b = 0; b < batch_n; b++) {
            float * dst = s.output.data() + (size_t) b * s.Oc * s.T;
            for (int t = 0; t < copy_T; t++) {
                if (t < s.repaint_t0 || t >= s.repaint_t1) {
                    memcpy(dst + (size_t) t * 64, s.cover_latents.data() + (size_t) t * 64, 64 * sizeof(float));
                }
            }
        }
        fprintf(stderr, "[Latent-Splice] kept generated frames [%d, %d) / %d, source elsewhere\n", s.repaint_t0,
                s.repaint_t1, s.T);
    }

    int                T_latent    = s.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 = s.output.data() + b * s.Oc * s.T;

        s.timer.reset();
        int T_audio;
        if (use_ort) {
            T_audio = vae_ort_decode_tiled(vae_ort, dit_out, T_latent, audio.data(), T_audio_max,
                                            ctx->params.vae_chunk, ctx->params.vae_overlap);
        } else {
            T_audio = vae_ggml_decode_tiled(vae_ggml, dit_out, T_latent, audio.data(), T_audio_max,
                                            ctx->params.vae_chunk, ctx->params.vae_overlap, cancel, cancel_data);
        }
        if (T_audio < 0) {
            if (cancel && cancel(cancel_data)) {
                fprintf(stderr, "[VAE-Decode Batch%d] Cancelled\n", b);
                return -1;
            }
            fprintf(stderr, "[VAE-Decode Batch%d] ERROR: decode failed\n", b);
            out[b].samples     = NULL;
            out[b].n_samples   = 0;
            out[b].sample_rate = 48000;
            continue;
        }
        fprintf(stderr, "[VAE-Decode Batch%d] Decode: %.1f ms (%s)\n", b, s.timer.ms(),
                use_ort ? "ORT" : "GGML");

        if (b == 0) {
            debug_dump_2d(&s.dbg, "vae_audio", audio.data(), 2, T_audio);
        }

        int n_total    = 2 * T_audio;
        out[b].samples = (float *) malloc((size_t) n_total * sizeof(float));
        if (!out[b].samples) {
            fprintf(stderr, "[VAE-Decode Batch%d] ERROR: OOM allocating output (%d samples)\n", b, n_total);
            out[b].n_samples   = 0;
            out[b].sample_rate = 48000;
            continue;
        }
        memcpy(out[b].samples, audio.data(), (size_t) n_total * sizeof(float));
        out[b].n_samples   = T_audio;
        out[b].sample_rate = 48000;

        {
            char dlabel[64];
            snprintf(dlabel, sizeof(dlabel), "vae_audio_b%d", b);
            diag_stats_f32(dlabel, out[b].samples, (size_t) n_total);
        }

        // Waveform hard splice (audio path only): replace out-of-zone samples
        // with the original source PCM. The VAE roundtrip is lossy, so this
        // restores bit-exact fidelity outside the region. Skipped on the
        // latent path (no original waveform available, single VAE decode is
        // already the cleanest output we can produce).
        if (have_region && !s.padded_src.empty()) {
            int       T_src   = (int) (s.padded_src.size() / 2);
            const int start_s = (int) ((size_t) s.repaint_t0 * 1920);
            const int end_s   = (int) ((size_t) s.repaint_t1 * 1920);
            const int T_clip  = T_audio < T_src ? T_audio : T_src;
            for (int ch = 0; ch < 2; ch++) {
                float * pred = out[b].samples + (size_t) ch * T_audio;
                for (int si = 0; si < start_s && si < T_clip; si++) {
                    pred[si] = s.padded_src[(size_t) si * 2 + ch];
                }
                for (int si = end_s; si < T_clip; si++) {
                    pred[si] = s.padded_src[(size_t) si * 2 + ch];
                }
            }
            fprintf(stderr, "[WAV-Splice Batch%d] hard splice samples [%d, %d) / %d\n", b, start_s, end_s, T_clip);
        }
    }
    return 0;
}

// ─── Postprocess plugin VAE decode ─────────────────────────────────────────
//
// Routes latent→audio through a Lua postprocess plugin. The plugin controls
// tiling, overlap, crossfading, and any DSP post-processing. The engine
// provides a vae_decode callback that the plugin calls for each tile.
//
// Falls back to the built-in ops_vae_decode if the plugin is not found.

int ops_vae_decode_postprocess(const AceSynth * ctx,
                               int              batch_n,
                               AceAudio *       out,
                               SynthState &     s,
                               const char *     plugin_name,
                               bool (*cancel)(void *),
                               void * cancel_data) {
    // Look up the postprocess plugin
    auto & reg = PluginRegistry::instance();
    LuaPlugin * plugin = reg.postprocess_lookup(plugin_name);
    if (!plugin) {
        fprintf(stderr, "[Postprocess] WARNING: plugin '%s' not found, falling back to built-in decoder\n",
                plugin_name ? plugin_name : "(null)");
        return ops_vae_decode(ctx, batch_n, out, s, cancel, cancel_data);
    }

    // Acquire VAE decoder (same as ops_vae_decode)
    VAEGGML * vae = store_require_vae_dec(ctx->store, ctx->vae_dec_key);
    if (!vae) {
        fprintf(stderr, "[Postprocess] FATAL: store_require_vae_dec failed\n");
        return -1;
    }
    ModelHandle vae_guard(ctx->store, vae);

    // Latent splice for repaint/lego (identical to ops_vae_decode)
    bool have_region = s.is_repaint || s.is_lego_region;
    if (have_region && s.have_cover && s.T_cover > 0) {
        int copy_T = s.T_cover < s.T ? s.T_cover : s.T;
        for (int b = 0; b < batch_n; b++) {
            float * dst = s.output.data() + (size_t) b * s.Oc * s.T;
            for (int t = 0; t < copy_T; t++) {
                if (t < s.repaint_t0 || t >= s.repaint_t1) {
                    memcpy(dst + (size_t) t * 64, s.cover_latents.data() + (size_t) t * 64, 64 * sizeof(float));
                }
            }
        }
        fprintf(stderr, "[Latent-Splice] kept generated frames [%d, %d) / %d, source elsewhere\n", s.repaint_t0,
                s.repaint_t1, s.T);
    }

    int T_latent    = s.T;
    int T_audio_max = T_latent * 1920;

    fprintf(stderr, "[Postprocess] Using plugin '%s' for VAE decode (T_latent=%d)\n",
            plugin->name.c_str(), T_latent);

    for (int b = 0; b < batch_n; b++) {
        if (cancel && cancel(cancel_data)) {
            fprintf(stderr, "[Postprocess Batch%d] Cancelled\n", b);
            return -1;
        }

        float * dit_out = s.output.data() + b * s.Oc * s.T;

        // Build the VAE decode callback for the Lua plugin
        // This wraps vae_ggml_decode_tiled with the engine's VAE and chunk params
        PostprocessVaeDecodeFn decode_fn = [&](const float * latent, int T_lat, float * aud_out, int max_T) -> int {
            return vae_ggml_decode_tiled(vae, latent, T_lat, aud_out, max_T,
                                         ctx->params.vae_chunk, ctx->params.vae_overlap,
                                         cancel, cancel_data);
        };

        // Allocate output buffer
        std::vector<float> audio(2 * T_audio_max);

        s.timer.reset();
        int T_audio = lua_call_postprocess(
            *plugin, dit_out, T_latent, s.Oc, 2,
            audio.data(), T_audio_max, decode_fn,
            g_hotstep_params.plugin_params);

        if (T_audio < 0) {
            fprintf(stderr, "[Postprocess Batch%d] ERROR: plugin decode failed, falling back to built-in\n", b);
            // Fallback for this batch item
            T_audio = vae_ggml_decode_tiled(vae, dit_out, T_latent, audio.data(), T_audio_max,
                                            ctx->params.vae_chunk, ctx->params.vae_overlap,
                                            cancel, cancel_data);
            if (T_audio < 0) {
                out[b].samples     = NULL;
                out[b].n_samples   = 0;
                out[b].sample_rate = 48000;
                continue;
            }
        }
        fprintf(stderr, "[Postprocess Batch%d] Decode: %.1f ms (T_audio=%d)\n", b, s.timer.ms(), T_audio);

        int n_total    = 2 * T_audio;
        out[b].samples = (float *) malloc((size_t) n_total * sizeof(float));
        if (!out[b].samples) {
            fprintf(stderr, "[Postprocess Batch%d] ERROR: OOM allocating output (%d samples)\n", b, n_total);
            out[b].n_samples   = 0;
            out[b].sample_rate = 48000;
            continue;
        }
        memcpy(out[b].samples, audio.data(), (size_t) n_total * sizeof(float));
        out[b].n_samples   = T_audio;
        out[b].sample_rate = 48000;

        // Waveform hard splice for repaint/lego (audio path only)
        if (have_region && !s.padded_src.empty()) {
            int       T_src   = (int) (s.padded_src.size() / 2);
            const int start_s = (int) ((size_t) s.repaint_t0 * 1920);
            const int end_s   = (int) ((size_t) s.repaint_t1 * 1920);
            const int T_clip  = T_audio < T_src ? T_audio : T_src;
            for (int ch = 0; ch < 2; ch++) {
                float * pred = out[b].samples + (size_t) ch * T_audio;
                for (int si = 0; si < start_s && si < T_clip; si++) {
                    pred[si] = s.padded_src[(size_t) si * 2 + ch];
                }
                for (int si = end_s; si < T_clip; si++) {
                    pred[si] = s.padded_src[(size_t) si * 2 + ch];
                }
            }
            fprintf(stderr, "[WAV-Splice Batch%d] hard splice samples [%d, %d) / %d\n", b, start_s, end_s, T_clip);
        }
    }
    return 0;
}

// ─── PP-VAE Re-encode ──────────────────────────────────────────────────────
//
// Round-trip audio through a separate post-processing VAE to clean up spectral
// artifacts from the ACE-Step VAE. The PP-VAE is a different autoencoder with
// superior reconstruction fidelity.
//
// Pipeline:
//   1. Measure input RMS + peak per batch item
//   2. Acquire PP-VAE encoder → tiled encode to latent [T_latent, 64] → release
//   3. Acquire PP-VAE decoder → tiled decode to PCM → release
//   4. RMS gain match: scale output to match input RMS, cap at input peak
//   5. Replace audio[b].samples with re-encoded PCM

static void pp_vae_rms(const float * samples, int n_samples, float * rms_out, float * peak_out) {
    // Planar stereo: [L0..LN, R0..RN], n_samples per channel
    double sum_sq = 0.0;
    float  peak   = 0.0f;
    int    total  = n_samples * 2;  // both channels
    for (int i = 0; i < total; i++) {
        float v = samples[i];
        sum_sq += (double) v * v;
        float av = fabsf(v);
        if (av > peak) {
            peak = av;
        }
    }
    *rms_out  = (float) sqrt(sum_sq / (double) total);
    *peak_out = peak;
}

int ops_pp_vae_reencode(const AceSynth * ctx, int batch_n, AceAudio * out, SynthState & s) {
    s.timer.reset();
    fprintf(stderr, "[PP-VAE] Re-encoding %d track(s)...\n", batch_n);

    // Pre-measure input RMS + peak for each batch item
    std::vector<float> in_rms(batch_n), in_peak(batch_n);
    for (int b = 0; b < batch_n; b++) {
        if (!out[b].samples || out[b].n_samples <= 0) {
            continue;
        }
        pp_vae_rms(out[b].samples, out[b].n_samples, &in_rms[b], &in_peak[b]);
        // DEBUG: input audio stats
        {
            float mn = out[b].samples[0], mx = out[b].samples[0];
            double sum = 0;
            int n_total = out[b].n_samples * 2;
            for (int i = 0; i < n_total; i++) {
                float v = out[b].samples[i];
                sum += v;
                if (v < mn) mn = v;
                if (v > mx) mx = v;
            }
            fprintf(stderr, "[PP-VAE Batch%d] Input: n=%d, rms=%.5f, peak=%.5f, min=%.5f, max=%.5f, mean=%.6f\n",
                    b, out[b].n_samples, in_rms[b], in_peak[b], mn, mx, (float)(sum / n_total));
        }
    }

    // Phase 1: Encode all batch items through PP-VAE encoder → latents
    // Prefers ORT/TRT encoder when available, falls back to GGML.
    // Encoder converts planar stereo PCM → interleaved → VAE latents [T_latent, 64]
    std::vector<std::vector<float>> latents(batch_n);
    std::vector<int>                T_latent(batch_n, 0);

    {
        bool use_ort_enc = !ctx->pp_vae_onnx_enc_path.empty();
        VaeEncOrt  * enc_ort  = nullptr;
        VAEEncoder * enc_ggml = nullptr;

        if (use_ort_enc) {
            enc_ort = store_require_vae_enc_ort(ctx->store, ctx->pp_vae_enc_ort_key);
            if (!enc_ort) {
                fprintf(stderr, "[PP-VAE] ORT encoder unavailable, falling back to GGML\n");
                use_ort_enc = false;
            }
        }
        if (!use_ort_enc) {
            enc_ggml = store_require_vae_enc(ctx->store, ctx->pp_vae_enc_key);
            if (!enc_ggml) {
                fprintf(stderr, "[PP-VAE] WARNING: encoder unavailable, skipping\n");
                return 0;  // non-fatal: just skip the re-encode
            }
        }
        ModelHandle enc_guard(ctx->store, use_ort_enc ? (void *)enc_ort : (void *)enc_ggml);

        fprintf(stderr, "[PP-VAE] Encoding via %s\n", use_ort_enc ? "ORT/TRT" : "GGML");

        for (int b = 0; b < batch_n; b++) {
            if (!out[b].samples || out[b].n_samples <= 0) {
                continue;
            }

            int   T_audio = out[b].n_samples;
            int   max_T   = (T_audio / 1920) + 64;
            latents[b].resize((size_t) max_T * 64);

            // vae_enc expects interleaved stereo [T*2]
            // out[b].samples is planar [L0..LN, R0..RN] → need to interleave
            std::vector<float> interleaved(T_audio * 2);
            const float * L = out[b].samples;
            const float * R = out[b].samples + T_audio;
            for (int i = 0; i < T_audio; i++) {
                interleaved[i * 2 + 0] = L[i];
                interleaved[i * 2 + 1] = R[i];
            }

            if (use_ort_enc) {
                T_latent[b] = vae_enc_ort_encode_tiled(enc_ort, interleaved.data(), T_audio,
                                                        latents[b].data(), max_T,
                                                        ctx->params.vae_chunk, ctx->params.vae_overlap);
            } else {
                T_latent[b] = vae_enc_encode_tiled(enc_ggml, interleaved.data(), T_audio,
                                                    latents[b].data(), max_T,
                                                    ctx->params.vae_chunk, ctx->params.vae_overlap);
            }

            if (T_latent[b] <= 0) {
                fprintf(stderr, "[PP-VAE Batch%d] WARNING: encode failed\n", b);
                T_latent[b] = 0;
            } else {
                // DEBUG: latent stats after encode
                float mn = latents[b][0], mx = latents[b][0];
                double sum = 0, sum_sq = 0;
                int n_el = T_latent[b] * 64;
                for (int i = 0; i < n_el; i++) {
                    float v = latents[b][i];
                    sum += v;
                    sum_sq += (double)v * v;
                    if (v < mn) mn = v;
                    if (v > mx) mx = v;
                }
                float lat_mean = (float)(sum / n_el);
                float lat_rms  = (float)sqrt(sum_sq / n_el);
                fprintf(stderr, "[PP-VAE Batch%d] Latent: T=%d, dim=64, n_el=%d, mean=%.5f, rms=%.5f, min=%.5f, max=%.5f\n",
                        b, T_latent[b], n_el, lat_mean, lat_rms, mn, mx);
            }
        }
    }
    fprintf(stderr, "[PP-VAE] Encode done: %.1f ms\n", s.timer.ms());

    // Phase 2: Decode all latents through PP-VAE decoder → PCM
    // Prefers ORT/TRT decoder when available, falls back to GGML.
    {
        s.timer.reset();
        bool use_ort_dec = !ctx->pp_vae_onnx_dec_path.empty();
        VaeOrt  * dec_ort  = nullptr;
        VAEGGML * dec_ggml = nullptr;

        if (use_ort_dec) {
            dec_ort = store_require_vae_dec_ort(ctx->store, ctx->pp_vae_dec_ort_key);
            if (!dec_ort) {
                fprintf(stderr, "[PP-VAE] ORT decoder unavailable, falling back to GGML\n");
                use_ort_dec = false;
            }
        }
        if (!use_ort_dec) {
            dec_ggml = store_require_vae_dec(ctx->store, ctx->pp_vae_dec_key);
            if (!dec_ggml) {
                fprintf(stderr, "[PP-VAE] WARNING: decoder unavailable, skipping\n");
                return 0;
            }
        }
        ModelHandle dec_guard(ctx->store, use_ort_dec ? (void *)dec_ort : (void *)dec_ggml);

        fprintf(stderr, "[PP-VAE] Decoding via %s\n", use_ort_dec ? "ORT/TRT" : "GGML");

        for (int b = 0; b < batch_n; b++) {
            if (T_latent[b] <= 0) {
                continue;
            }

            int                T_audio_max = T_latent[b] * 1920;
            std::vector<float> audio(2 * T_audio_max);

            int T_audio;
            if (use_ort_dec) {
                T_audio = vae_ort_decode_tiled(dec_ort, latents[b].data(), T_latent[b],
                                                audio.data(), T_audio_max,
                                                ctx->params.vae_chunk, ctx->params.vae_overlap);
            } else {
                T_audio = vae_ggml_decode_tiled(dec_ggml, latents[b].data(), T_latent[b],
                                                 audio.data(), T_audio_max,
                                                 ctx->params.vae_chunk, ctx->params.vae_overlap, NULL, NULL);
            }

            if (T_audio <= 0) {
                fprintf(stderr, "[PP-VAE Batch%d] WARNING: decode failed\n", b);
                continue;
            }

            // DEBUG: raw decode output stats (before gain)
            {
                float mn = audio[0], mx = audio[0];
                double sum = 0;
                int n_total = 2 * T_audio;
                for (int i = 0; i < n_total; i++) {
                    float v = audio[i];
                    sum += v;
                    if (v < mn) mn = v;
                    if (v > mx) mx = v;
                }
                float raw_rms, raw_peak;
                pp_vae_rms(audio.data(), T_audio, &raw_rms, &raw_peak);
                fprintf(stderr, "[PP-VAE Batch%d] Decode: T=%d, rms=%.5f, peak=%.5f, min=%.5f, max=%.5f, mean=%.6f\n",
                        b, T_audio, raw_rms, raw_peak, mn, mx, (float)(sum / n_total));
            }

            // RMS gain match: scale output to match input RMS, cap at input peak
            float out_rms, out_peak;
            pp_vae_rms(audio.data(), T_audio, &out_rms, &out_peak);

            float gain = 1.0f;
            if (out_rms > 1e-8f) {
                gain = in_rms[b] / out_rms;
                // Cap so we never exceed input peak
                if (out_peak * gain > in_peak[b] + 0.01f) {
                    gain = in_peak[b] / (out_peak + 1e-8f);
                }
            }

            // Apply gain
            int n_total = 2 * T_audio;
            for (int i = 0; i < n_total; i++) {
                audio[i] *= gain;
            }

            // Replace output samples
            free(out[b].samples);
            out[b].samples = (float *) malloc((size_t) n_total * sizeof(float));
            if (!out[b].samples) {
                fprintf(stderr, "[PP-VAE Batch%d] ERROR: OOM\n", b);
                out[b].n_samples = 0;
                continue;
            }
            memcpy(out[b].samples, audio.data(), (size_t) n_total * sizeof(float));
            out[b].n_samples = T_audio;

            fprintf(stderr, "[PP-VAE Batch%d] OK: gain=%.3f (in_rms=%.4f, out_rms=%.4f, in_peak=%.4f, out_peak=%.4f)\n",
                    b, gain, in_rms[b], out_rms, in_peak[b], out_peak);
        }
    }
    fprintf(stderr, "[PP-VAE] Decode done: %.1f ms total\n", s.timer.ms());

    return 0;
}