// ace-server.cpp: HTTP server for ACE-Step music generation
//
// Single binary, one port. All compute endpoints (POST /lm, POST /synth,
// POST /understand) are asynchronous: they validate the request, create a
// job, push it to a FIFO queue, and return the job ID immediately.
// A single worker thread processes jobs in order.
// Clients poll GET /job?id=N for status and fetch results with
// GET /job?id=N&result=1. POST /job?id=N&cancel=1 cancels a job.
//
// Job IDs are random 64-bit hex strings (non-predictable).
// Completed jobs are evicted FIFO when the pool exceeds MAX_JOBS.
// Running jobs are never evicted.
//
// Models are discovered by scanning --models directory at startup
// (reads GGUF metadata only, no weights loaded).
// Each request loads the model, executes, and frees it. No model persists
// in VRAM between requests unless --keep-loaded is set. GPU access is
// serialized by the single worker thread (no mutex needed).
//
// Available models are classified by their GGUF general.architecture:
//   acestep-lm       -> lm bucket
//   acestep-dit      -> dit bucket
//   acestep-text-enc -> text-enc bucket (singleton, first entry used)
//   acestep-vae      -> vae bucket      (singleton, first entry used)
//
// Endpoint requirements:
//   /lm         LM
//   /synth      DiT + Text-Enc + VAE
//   /understand LM + DiT + VAE

#include "audio-io.h"
#include "model-registry.h"
#include "model-store.h"
#include "pipeline-lm.h"
#include "pipeline-synth.h"
#include "pipeline-understand.h"
#include "request.h"
#include "synth-batch-runner.h"
#include "task-types.h"
#include "vae.h"
#include "version.h"
#include "yyjson.h"

// embedded webui (generated by xxd.cmake from tools/webui/public/index.html.gz)
#include "index.html.gz.hpp"

// suppress warnings in third-party headers
#ifdef __GNUC__
#    pragma GCC diagnostic push
#    pragma GCC diagnostic ignored "-Wshadow"
#endif
#include "httplib.h"
#ifdef __GNUC__
#    pragma GCC diagnostic pop
#endif

#include <atomic>
#include <condition_variable>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <functional>
#include <memory>
#include <mutex>
#include <random>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#ifdef _WIN32
#    include <fcntl.h>
#    include <io.h>
#    ifndef STDERR_FILENO
#        define STDERR_FILENO 2
#    endif
#else
#    include <unistd.h>
#endif

// portable fd wrappers. avoids macros that collide with C++ method names
// (e.g. sink.write() in httplib would be eaten by a write() macro).
#ifdef _WIN32
static int fd_pipe(int fd[2]) {
    return _pipe(fd, 4096, _O_BINARY);
}

static int fd_dup(int fd) {
    return _dup(fd);
}

static int fd_dup2(int src, int dst) {
    return _dup2(src, dst);
}

static int fd_read(int fd, void * buf, size_t n) {
    return _read(fd, buf, (unsigned) n);
}

static int fd_write(int fd, const void * buf, size_t n) {
    return _write(fd, buf, (unsigned) n);
}

static void fd_close(int fd) {
    _close(fd);
}
#else
static int fd_pipe(int fd[2]) {
    return pipe(fd);
}

static int fd_dup(int fd) {
    return dup(fd);
}

static int fd_dup2(int src, int dst) {
    return dup2(src, dst);
}

static int fd_read(int fd, void * buf, size_t n) {
    return (int) read(fd, buf, n);
}

static int fd_write(int fd, const void * buf, size_t n) {
    return (int) write(fd, buf, n);
}

static void fd_close(int fd) {
    close(fd);
}
#endif

// server instance pointer for the signal handler
static httplib::Server * g_svr = nullptr;

static void on_signal(int) {
    if (g_svr) {
        g_svr->stop();
    }
}

// work queue: all GPU jobs go through a single FIFO queue processed
// by one worker thread. GPU access is serialized by construction.
static std::deque<std::function<void()>> g_work_queue;
static std::mutex                        mtx_work;
static std::condition_variable           cv_work;
static bool                              g_work_stop = false;

static void work_push(std::function<void()> fn) {
    std::lock_guard<std::mutex> lock(mtx_work);
    g_work_queue.push_back(std::move(fn));
    cv_work.notify_one();
}

// worker thread: consume jobs in FIFO order until shutdown.
// on stop: finishes the current job, discards pending ones.
static void worker_main() {
    for (;;) {
        std::function<void()> fn;
        {
            std::unique_lock<std::mutex> lock(mtx_work);
            cv_work.wait(lock, [] { return g_work_stop || !g_work_queue.empty(); });
            if (g_work_stop) {
                break;
            }
            fn = std::move(g_work_queue.front());
            g_work_queue.pop_front();
        }
        fn();
    }
}

// central GGML module store shared across pipelines. Policy picked at startup
// from --keep-loaded: STRICT by default (one GPU module resident at a time),
// NEVER when the flag is set (accumulate across requests).
static ModelStore * g_store = nullptr;

// model registry (populated at startup from GGUF metadata)
static ModelRegistry g_registry;

// loaded model names (empty = nothing loaded)
static std::string g_loaded_lm;
static std::string g_loaded_dit;
static std::string g_loaded_adapter;
static float       g_loaded_adapter_scale = 1.0f;
static std::string g_loaded_und_dit;
static std::string g_loaded_vae;

// pipeline params (rebuilt from registry paths on each load)
static AceLmParams         g_lm_params;
static AceSynthParams      g_synth_params;
static AceUnderstandParams g_und_params;

// limits
static int  g_max_batch   = 1;
static bool g_keep_loaded = false;

// job system: all compute endpoints create a job and return its ID
// immediately. the worker thread processes jobs in FIFO order, stores
// the result. the client polls GET /job?id=N until done, then fetches
// the result with GET /job?id=N&result=1.
// cancel: POST /job?id=N&cancel=1 sets the per-job flag.
enum class JobStatus : int {
    RUNNING   = 0,
    DONE      = 1,
    FAILED    = 2,
    CANCELLED = 3,
};

struct Job {
    std::string            id;
    std::atomic<JobStatus> status{ JobStatus::RUNNING };
    std::string            result_body;
    std::string            result_mime;
    std::atomic<bool>      cancel{ false };

    // memory ordering contract: result_body and result_mime are written
    // before status is stored (seq_cst). the client loads status (seq_cst)
    // and only reads result fields after seeing done/failed. this guarantees
    // visibility without an explicit mutex on the result fields.
};

static std::mutex                                            mtx_jobs;
static std::unordered_map<std::string, std::shared_ptr<Job>> g_jobs;
static std::deque<std::string>                               g_job_order;
static const int                                             MAX_JOBS = 32;

// Source latent cap: matches the silence_latent tensor baked into the DiT
// GGUF, which is fixed at [15000, 64] f32. The pipeline indexes into it
// directly when padding context, so any T_latent above 15000 would walk
// past the buffer. Same hard limit ops_resolve_T enforces on s.T.
static const int MAX_T_LATENT = 15000;

// Latent payload format: raw f32 [T * 64] little-endian, no header. Same
// layout neural-codec emits in --encode -f f32 mode and what /synth,
// /understand, /vae return as the latent multipart part or raw body.
static const int LATENT_CHANNELS    = 64;
static const int LATENT_FRAME_BYTES = LATENT_CHANNELS * (int) sizeof(float);

// Validate a latent payload coming from the wire: size must be a strict
// multiple of one frame, T must be in (0, MAX_T_LATENT]. Returns the frame
// count or -1 on failure (with the HTTP code the caller should reply).
static int latent_payload_validate(size_t size, int * http_code_out) {
    if (size == 0 || (size % LATENT_FRAME_BYTES) != 0) {
        if (http_code_out) {
            *http_code_out = 400;
        }
        return -1;
    }
    int T = (int) (size / (size_t) LATENT_FRAME_BYTES);
    if (T <= 0) {
        if (http_code_out) {
            *http_code_out = 400;
        }
        return -1;
    }
    if (T > MAX_T_LATENT) {
        if (http_code_out) {
            *http_code_out = 413;
        }
        return -1;
    }
    return T;
}

// Build a multipart/mixed body that bundles the primary payload with its
// latents. The audio variant pairs one audio part with one latent part per
// track, the JSON variant carries a single payload and one optional latent.
// The boundary is fixed and matches the existing batch format; the client
// splits on it the same way for every endpoint.
static const char * MULTIPART_BOUNDARY = "ace-batch-boundary";

static std::string multipart_build_audio_latent(const std::vector<std::string> &        audio_parts,
                                                const char *                            audio_mime,
                                                const std::vector<std::vector<float>> & latents) {
    std::string body;
    for (size_t i = 0; i < audio_parts.size(); i++) {
        if (audio_parts[i].empty()) {
            continue;
        }
        body += "--";
        body += MULTIPART_BOUNDARY;
        body += "\r\nContent-Type: ";
        body += audio_mime;
        body += "\r\n\r\n";
        body += audio_parts[i];
        body += "\r\n";
        body += "--";
        body += MULTIPART_BOUNDARY;
        body += "\r\nContent-Type: application/octet-stream\r\n";
        body += "Content-Disposition: form-data; name=\"latent\"\r\n\r\n";
        body.append(reinterpret_cast<const char *>(latents[i].data()), latents[i].size() * sizeof(float));
        body += "\r\n";
    }
    body += "--";
    body += MULTIPART_BOUNDARY;
    body += "--\r\n";
    return body;
}

// Same shape, JSON primary instead of audio. Used by /understand.
static std::string multipart_build_json_latent(const std::string &        json_part,
                                               const std::vector<float> & latent,
                                               int                        T_latent) {
    std::string body;
    body += "--";
    body += MULTIPART_BOUNDARY;
    body += "\r\nContent-Type: application/json\r\n\r\n";
    body += json_part;
    body += "\r\n";
    if (T_latent > 0 && !latent.empty()) {
        body += "--";
        body += MULTIPART_BOUNDARY;
        body += "\r\nContent-Type: application/octet-stream\r\n";
        body += "Content-Disposition: form-data; name=\"latent\"\r\n\r\n";
        body.append(reinterpret_cast<const char *>(latent.data()), (size_t) T_latent * LATENT_FRAME_BYTES);
        body += "\r\n";
    }
    body += "--";
    body += MULTIPART_BOUNDARY;
    body += "--\r\n";
    return body;
}

static const std::string MULTIPART_MIME = std::string("multipart/mixed; boundary=") + MULTIPART_BOUNDARY;

// generate a random hex ID (64 bits of entropy, non-predictable)
static std::string job_make_id() {
    static std::mt19937_64      rng(std::random_device{}());
    static std::mutex           mtx_rng;
    std::lock_guard<std::mutex> lock(mtx_rng);
    char                        buf[17];
    snprintf(buf, sizeof(buf), "%016llx", (unsigned long long) rng());
    return buf;
}

static std::shared_ptr<Job> job_create() {
    std::lock_guard<std::mutex> lock(mtx_jobs);
    auto                        job = std::make_shared<Job>();
    job->id                         = job_make_id();
    g_jobs[job->id]                 = job;
    g_job_order.push_back(job->id);

    // evict oldest completed jobs to stay under MAX_JOBS.
    // running jobs (status 0) are never evicted.
    while ((int) g_job_order.size() > MAX_JOBS) {
        bool evicted = false;
        for (auto it = g_job_order.begin(); it != g_job_order.end(); ++it) {
            auto jit = g_jobs.find(*it);
            if (jit == g_jobs.end() || jit->second->status.load() != JobStatus::RUNNING) {
                if (jit != g_jobs.end()) {
                    g_jobs.erase(jit);
                }
                g_job_order.erase(it);
                evicted = true;
                break;
            }
        }
        if (!evicted) {
            break;
        }
    }
    return job;
}

static std::shared_ptr<Job> job_find(const std::string & id) {
    std::lock_guard<std::mutex> lock(mtx_jobs);
    auto                        it = g_jobs.find(id);
    return it != g_jobs.end() ? it->second : nullptr;
}

static const char * job_status_str(JobStatus s) {
    switch (s) {
        case JobStatus::RUNNING:
            return "running";
        case JobStatus::DONE:
            return "done";
        case JobStatus::FAILED:
            return "failed";
        case JobStatus::CANCELLED:
            return "cancelled";
    }
    return "unknown";
}

// log capture: intercept stderr via pipe, forward to terminal + ring buffer.
// SSE clients connect to /logs and receive lines in real time.
#define LOG_RING_BITS 9
#define LOG_RING_SIZE (1 << LOG_RING_BITS)
#define LOG_RING_MASK (LOG_RING_SIZE - 1)

static std::mutex              mtx_log;
static std::condition_variable cv_log;
static std::string             log_ring[LOG_RING_SIZE];
static uint64_t                log_seq = 0;

static int g_real_stderr_fd = -1;
static int g_pipe_read_fd   = -1;

static void setup_log_capture() {
    g_real_stderr_fd = fd_dup(STDERR_FILENO);
    int pipefd[2];
    if (fd_pipe(pipefd) != 0) {
        g_real_stderr_fd = -1;
        return;
    }
    g_pipe_read_fd = pipefd[0];
    fd_dup2(pipefd[1], STDERR_FILENO);
    fd_close(pipefd[1]);
}

// reader thread: drain pipe, forward to real stderr, push lines to ring.
// exits when the write end of the pipe is closed (fd_dup2 restores real stderr).
static void log_reader_main() {
    char        buf[4096];
    std::string partial;
    for (;;) {
        int n = fd_read(g_pipe_read_fd, buf, sizeof(buf));
        if (n <= 0) {
            break;
        }
        fd_write(g_real_stderr_fd, buf, (size_t) n);
        partial.append(buf, (size_t) n);
        size_t pos;
        while ((pos = partial.find('\n')) != std::string::npos) {
            std::lock_guard<std::mutex> lock(mtx_log);
            log_ring[log_seq & LOG_RING_MASK] = partial.substr(0, pos);
            log_seq++;
            cv_log.notify_all();
            partial.erase(0, pos + 1);
        }
    }
    if (!partial.empty()) {
        std::lock_guard<std::mutex> lock(mtx_log);
        log_ring[log_seq & LOG_RING_MASK] = std::move(partial);
        log_seq++;
        cv_log.notify_all();
    }
    fd_close(g_pipe_read_fd);
}

static void teardown_log_capture() {
    if (g_real_stderr_fd < 0) {
        return;
    }
    fflush(stderr);
    fd_dup2(g_real_stderr_fd, STDERR_FILENO);
    // g_real_stderr_fd stays open: the reader thread writes to it
}

// RAII: captures stderr on construction, restores + joins reader on destruction.
// safe on any exit path (early arg errors, model load failures, normal shutdown).
struct LogCapture {
    std::thread reader;

    LogCapture() {
        setup_log_capture();
        reader = std::thread(log_reader_main);
    }

    ~LogCapture() {
        teardown_log_capture();
        cv_log.notify_all();
        if (reader.joinable()) {
            reader.join();
        }

        // reader is done draining the pipe, safe to close
        if (g_real_stderr_fd >= 0) {
            fd_close(g_real_stderr_fd);
            g_real_stderr_fd = -1;
        }
    }
};

// GET /logs: SSE stream of stderr lines.
// sends backlog (up to LOG_RING_SIZE) then streams new lines in real time.
static void handle_logs(const httplib::Request &, httplib::Response & res) {
    res.set_header("Cache-Control", "no-cache");
    res.set_header("X-Accel-Buffering", "no");
    res.set_chunked_content_provider(
        "text/event-stream", [cursor = uint64_t(0), init = false](size_t, httplib::DataSink & sink) mutable -> bool {
            std::unique_lock<std::mutex> lock(mtx_log);
            if (!init) {
                uint64_t avail = log_seq < LOG_RING_SIZE ? log_seq : (uint64_t) LOG_RING_SIZE;
                cursor         = log_seq - avail;
                while (cursor < log_seq) {
                    std::string ev = "data: " + log_ring[cursor & LOG_RING_MASK] + "\n\n";
                    cursor++;
                    lock.unlock();
                    if (!sink.write(ev.c_str(), ev.size())) {
                        return false;
                    }
                    lock.lock();
                }
                init = true;
            }
            cv_log.wait_for(lock, std::chrono::seconds(2));
            while (cursor < log_seq) {
                std::string ev = "data: " + log_ring[cursor & LOG_RING_MASK] + "\n\n";
                cursor++;
                lock.unlock();
                if (!sink.write(ev.c_str(), ev.size())) {
                    return false;
                }
                lock.lock();
            }
            return true;
        });
}

// cancel callback: checks the per-job cancel flag.
static bool server_cancel_job(void * data) {
    auto * flag = (const std::atomic<bool> *) data;
    return flag && flag->load(std::memory_order_relaxed);
}

// helper: set a JSON error response
static void json_error(httplib::Response & res, int status, const char * msg) {
    yyjson_mut_doc * doc  = yyjson_mut_doc_new(NULL);
    yyjson_mut_val * root = yyjson_mut_obj(doc);
    yyjson_mut_doc_set_root(doc, root);
    yyjson_mut_obj_add_str(doc, root, "error", msg);
    char * json = yyjson_mut_write(doc, 0, NULL);
    yyjson_mut_doc_free(doc);
    res.status = status;
    res.set_content(json, "application/json");
    free(json);
}

// resolve model name: explicit request > already loaded > first in bucket
static std::string resolve_name(const std::vector<ModelEntry> & bucket,
                                const std::string &             requested,
                                const std::string &             loaded) {
    if (!requested.empty()) {
        return requested;
    }
    if (!loaded.empty()) {
        return loaded;
    }
    if (!bucket.empty()) {
        return bucket[0].name;
    }
    return "";
}

// LM worker: generates metadata + lyrics + codes, stores JSON result in job.
static void lm_worker(std::shared_ptr<Job> job, std::vector<AceRequest> ace_reqs, int mode) {
    if (job->cancel.load()) {
        job->status.store(JobStatus::CANCELLED);
        return;
    }

    const int N = (int) ace_reqs.size();

    // Resolve model name and build per-request params from the template.
    std::string        lm_name = resolve_name(g_registry.lm, ace_reqs[0].lm_model, g_loaded_lm);
    const ModelEntry * entry   = registry_find(g_registry.lm, lm_name.c_str());
    if (!entry) {
        fprintf(stderr, "[Server] LM not found: %s\n", lm_name.c_str());
        job->status.store(JobStatus::FAILED);
        return;
    }
    AceLmParams p = g_lm_params;
    p.model_path  = entry->path.c_str();

    // Acquire a fresh LM ctx from the shared store. Under EVICT_STRICT the
    // module is reloaded if another pipeline evicted it; under EVICT_NEVER
    // the store returns the cached instance.
    AceLm * ctx = ace_lm_load(g_store, &p);
    if (!ctx) {
        fprintf(stderr, "[Server] FATAL: LM load failed\n");
        job->status.store(JobStatus::FAILED);
        return;
    }

    // Execute and always free the ctx, success or failure: the store decides
    // whether the underlying GPU module stays resident.
    std::vector<AceRequest> out(N);
    int rc = ace_lm_generate(ctx, ace_reqs.data(), N, out.data(), NULL, NULL, server_cancel_job, (void *) &job->cancel,
                             mode);
    ace_lm_free(ctx);

    if (rc != 0) {
        job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::FAILED);
        return;
    }

    // Sticky name hint for resolve_name under --keep-loaded. Master clears it
    // in the default mode since the ctx is gone; we match that behavior.
    if (g_keep_loaded) {
        g_loaded_lm = lm_name;
    } else {
        g_loaded_lm.clear();
    }

    // serialize output as a JSON array
    std::string body = "[";
    for (int i = 0; i < N; i++) {
        if (i > 0) {
            body += ",";
        }
        body += request_to_json(&out[i]);
    }
    body += "]";

    job->result_body = std::move(body);
    job->result_mime = "application/json";
    job->status.store(JobStatus::DONE);
    fprintf(stderr, "[Server] Job %s done (LM, %d results)\n", job->id.c_str(), N);
}

// POST /lm
// accepts: AceRequest JSON or [AceRequest, ...] (lm_mode selects the mode).
// A single request expands into lm_batch_size seed variants of one prompt.
// An array submits independent requests generated once each in one GPU
// batch; items must share lm_mode and lm_model, lm_batch_size is ignored.
// returns: JSON {"id":"N"} immediately. result is a JSON array of enriched
// AceRequests, one per generated variant or array item.
// modes (AceRequest.lm_mode):
//   generate  metadata + lyrics + audio_codes  (full composer pass)
//   inspire   metadata + lyrics                (audio_codes stays empty)
//   format    metadata + lyrics                (audio_codes stays empty)
static void handle_lm(const httplib::Request & req, httplib::Response & res) {
    if (g_registry.lm.empty()) {
        json_error(res, 501, "No LM models in registry");
        return;
    }

    // parse request: single object {} or array [{}, ...]
    std::vector<AceRequest> ace_reqs;
    if (!request_parse_json_array(req.body.c_str(), &ace_reqs)) {
        json_error(res, 400, "Invalid JSON");
        return;
    }
    if ((int) ace_reqs.size() > g_max_batch) {
        json_error(res, 400, "Request array exceeds max_batch");
        return;
    }
    for (size_t i = 0; i < ace_reqs.size(); i++) {
        if (ace_reqs[i].caption.empty()) {
            json_error(res, 400, "Caption is required");
            return;
        }
        if (i > 0 && (ace_reqs[i].lm_mode != ace_reqs[0].lm_mode || ace_reqs[i].lm_model != ace_reqs[0].lm_model)) {
            json_error(res, 400, "Array items must share lm_mode and lm_model");
            return;
        }
    }

    // Resolve lm_mode string to integer mode used by ace_lm_generate.
    int mode;
    if (ace_reqs[0].lm_mode == LM_MODE_NAME_GENERATE) {
        mode = LM_MODE_GENERATE;
    } else if (ace_reqs[0].lm_mode == LM_MODE_NAME_INSPIRE) {
        mode = LM_MODE_INSPIRE;
    } else if (ace_reqs[0].lm_mode == LM_MODE_NAME_FORMAT) {
        mode = LM_MODE_FORMAT;
    } else {
        json_error(res, 400, "Invalid lm_mode (use: generate, inspire, format)");
        return;
    }

    if (ace_reqs.size() == 1) {
        // Single request: expand into lm_batch_size seed variants of one
        // prompt, clamped to [1, max_batch]. The pipeline detects the shared
        // prompt and prefills once.
        int lm_batch_size = ace_reqs[0].lm_batch_size;
        if (lm_batch_size < 1) {
            lm_batch_size = 1;
        }
        if (lm_batch_size > g_max_batch) {
            lm_batch_size = g_max_batch;
        }
        AceRequest base = ace_reqs[0];
        request_resolve_lm_seed(&base);
        request_resolve_seed(&base);
        ace_reqs.assign(lm_batch_size, base);
        for (int b = 0; b < lm_batch_size; b++) {
            ace_reqs[b].lm_seed = base.lm_seed + b;
            ace_reqs[b].seed    = base.seed + b;
        }
    } else {
        // Array request: each item is generated once with its own seeds.
        for (auto & r : ace_reqs) {
            request_resolve_lm_seed(&r);
            request_resolve_seed(&r);
        }
    }

    auto job = job_create();
    fprintf(stderr, "[Server] Job %s created (LM, mode=%d, N=%zu)\n", job->id.c_str(), mode, ace_reqs.size());

    work_push([job, reqs = std::move(ace_reqs), mode]() mutable { lm_worker(job, std::move(reqs), mode); });

    std::string body = "{\"id\":\"" + job->id + "\"}";
    res.set_content(body, "application/json");
}

// synth worker: processes synth request, stores audio result in job.
static void synth_worker(std::shared_ptr<Job>    job,
                         std::vector<AceRequest> ace_reqs,
                         float *                 src_interleaved,
                         int                     src_len,
                         std::vector<float>      src_latents,
                         int                     src_T_latent,
                         float *                 ref_interleaved,
                         int                     ref_len,
                         std::vector<float>      ref_latents,
                         int                     ref_T_latent,
                         bool                    output_wav,
                         WavFormat               wav_fmt,
                         int                     peak_clip) {
    // Generate every request in one DiT batch. synth_batch_size expands each
    // request into per-seed variants. Total clamped to DiT max 9.
    const int batch_n     = (int) ace_reqs.size();
    int       total_alloc = 0;
    for (int ri = 0; ri < batch_n; ri++) {
        int sbs = ace_reqs[ri].synth_batch_size;
        total_alloc += sbs < 1 ? 1 : (sbs > 9 ? 9 : sbs);
    }
    if (total_alloc > 9) {
        fprintf(stderr, "[Server] Batch %d exceeds DiT max 9, clamping\n", total_alloc);
        total_alloc = 9;
    }
    std::vector<AceAudio> audio(total_alloc);

    if (job->cancel.load()) {
        free(src_interleaved);
        free(ref_interleaved);
        job->status.store(JobStatus::CANCELLED);
        return;
    }

    // Resolve DiT, adapter, VAE and the text-encoder singleton.
    std::string        dit_name = resolve_name(g_registry.dit, ace_reqs[0].synth_model, g_loaded_dit);
    const ModelEntry * dit      = registry_find(g_registry.dit, dit_name.c_str());
    if (!dit) {
        fprintf(stderr, "[Server] DiT not found: %s\n", dit_name.c_str());
        free(src_interleaved);
        free(ref_interleaved);
        job->status.store(JobStatus::FAILED);
        return;
    }
    if (g_registry.text_enc.empty() || g_registry.vae.empty()) {
        fprintf(stderr, "[Server] Missing Text-Enc or VAE in registry\n");
        free(src_interleaved);
        free(ref_interleaved);
        job->status.store(JobStatus::FAILED);
        return;
    }
    std::string        vae_name = resolve_name(g_registry.vae, ace_reqs[0].vae, g_loaded_vae);
    const ModelEntry * vae      = registry_find(g_registry.vae, vae_name.c_str());
    if (!vae) {
        fprintf(stderr, "[Server] VAE not found: %s\n", vae_name.c_str());
        free(src_interleaved);
        free(ref_interleaved);
        job->status.store(JobStatus::FAILED);
        return;
    }

    AceSynthParams p    = g_synth_params;
    p.text_encoder_path = g_registry.text_enc[0].path.c_str();
    p.dit_path          = dit->path.c_str();
    p.vae_path          = vae->path.c_str();
    p.adapter_path      = nullptr;
    p.adapter_scale     = 1.0f;
    if (!ace_reqs[0].adapter.empty()) {
        const AdapterEntry * adapter = registry_find_adapter(g_registry, ace_reqs[0].adapter.c_str());
        if (!adapter) {
            fprintf(stderr, "[Server] Adapter not found: %s\n", ace_reqs[0].adapter.c_str());
            free(src_interleaved);
            free(ref_interleaved);
            job->status.store(JobStatus::FAILED);
            return;
        }
        p.adapter_path  = adapter->path.c_str();
        p.adapter_scale = ace_reqs[0].adapter_scale;
    }
    fprintf(stderr, "[Server] Loading synth: DiT=%s VAE=%s%s%s\n", dit_name.c_str(), vae_name.c_str(),
            ace_reqs[0].adapter.empty() ? "" : " Adapter=", ace_reqs[0].adapter.c_str());

    AceSynth * ctx = ace_synth_load(g_store, &p);
    if (!ctx) {
        fprintf(stderr, "[Server] FATAL: synth load failed\n");
        free(src_interleaved);
        free(ref_interleaved);
        job->status.store(JobStatus::FAILED);
        return;
    }

    // Build the flat batch. Seeds are resolved per original request, then
    // synth_batch_size is expanded into per-seed variants in groups[0].
    std::vector<std::vector<AceRequest>> groups(1);
    groups[0].reserve(total_alloc);
    int off = 0;
    for (int ri = 0; ri < batch_n && off < total_alloc; ri++) {
        auto & r   = ace_reqs[ri];
        int    sbs = r.synth_batch_size;
        if (sbs < 1) {
            sbs = 1;
        }
        if (sbs > 9) {
            sbs = 9;
        }
        if (off + sbs > total_alloc) {
            sbs = total_alloc - off;
        }
        request_resolve_seed(&r);
        const long long base_seed = r.seed;

        for (int i = 0; i < sbs; i++) {
            AceRequest v = r;
            v.seed       = base_seed + i;
            groups[0].push_back(v);
        }
        off += sbs;
    }

    if (total_alloc > 1) {
        fprintf(stderr, "[Server] Batch: %d track(s) from %d request(s)\n", total_alloc, batch_n);
    }

    // Two-phase run. The store acquires and releases GPU modules around each
    // op (STRICT) or keeps them across ops (NEVER). The synth ctx is always
    // freed at the end of this handler. The runner ingests src and ref as
    // either audio or latents (latents win when both are set per side) and
    // captures one post-DiT latent per generated track for the multipart
    // response.
    const float *                   src_lat_ptr = src_latents.empty() ? nullptr : src_latents.data();
    const float *                   ref_lat_ptr = ref_latents.empty() ? nullptr : ref_latents.data();
    std::vector<std::vector<float>> captured_latents;
    const int rc = synth_batch_run(ctx, groups, src_interleaved, src_len, src_lat_ptr, src_T_latent, ref_interleaved,
                                   ref_len, ref_lat_ptr, ref_T_latent, audio.data(), &captured_latents,
                                   server_cancel_job, (void *) &job->cancel);
    ace_synth_free(ctx);
    free(src_interleaved);
    free(ref_interleaved);

    if (rc != 0) {
        for (auto & a : audio) {
            ace_audio_free(&a);
        }
        job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::FAILED);
        return;
    }

    // Sticky name hints for resolve_name under --keep-loaded. Master clears
    // them in the default mode since the ctx is gone; we match that behavior.
    if (g_keep_loaded) {
        g_loaded_dit           = dit_name;
        g_loaded_adapter       = ace_reqs[0].adapter;
        g_loaded_adapter_scale = ace_reqs[0].adapter_scale;
        g_loaded_vae           = vae_name;
    } else {
        g_loaded_dit.clear();
        g_loaded_adapter.clear();
        g_loaded_adapter_scale = 1.0f;
        g_loaded_vae.clear();
    }

    const int total_tracks = total_alloc;

    // encode each track (peak normalize + encode)
    const char * mime = output_wav ? "audio/wav" : "audio/mpeg";

    std::vector<std::string> encoded(total_tracks);
    for (int b = 0; b < total_tracks; b++) {
        if (!audio[b].samples) {
            continue;
        }
        if (!output_wav || wav_fmt != WAV_F32) {
            audio_normalize(audio[b].samples, audio[b].n_samples * 2, peak_clip);
        }
        if (output_wav) {
            encoded[b] = audio_encode_wav(audio[b].samples, audio[b].n_samples, 48000, wav_fmt);
        } else {
            encoded[b] = audio_encode_mp3(audio[b].samples, audio[b].n_samples, 48000, groups[0][b].mp3_bitrate,
                                          server_cancel_job, (void *) &job->cancel);
        }
        ace_audio_free(&audio[b]);
    }

    // store result in job: every synth response is multipart, with one audio
    // part and one latent part per generated track, paired in wire order.
    // The audio mime is per-part so the client knows wav vs mp3 without a
    // query. The latent reproduces its track's audio when fed back to /vae.
    job->result_body = multipart_build_audio_latent(encoded, mime, captured_latents);
    job->result_mime = MULTIPART_MIME;

    job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::DONE);
    fprintf(stderr, "[Server] Job %s done (%d tracks)\n", job->id.c_str(), total_tracks);
}

// POST /synth
// returns JSON {"id":"N"} immediately.
// input:
//   application/json body        -> single request {} or batch [{req0}, {req1}, ...]
//   multipart/form-data          -> single request + optional audio or latents
//     part "request":     JSON text (model selection, output_format, etc.)
//     part "audio":       source audio (WAV or MP3)
//     part "src_latents": pre-encoded source latents (raw f32, [T*64]), wins over "audio"
//     part "ref_audio":   timbre reference audio (WAV or MP3), optional
//     part "ref_latents": pre-encoded timbre latents (raw f32, [T*64]), wins over "ref_audio"
// output: multipart/mixed
//   one audio part (audio/mpeg or audio/wav per request output_format) and
//   one latent part (application/octet-stream, raw f32 [T*64]) per generated
//   track, paired in wire order. The latent reproduces its track's audio when
//   fed back to /vae decode.
// Batch size = number of JSON objects (after synth_batch_size expansion, clamped to 9).
// Metadata (seed, duration, etc) is already in the request JSON from /lm.
static void handle_synth(const httplib::Request & req, httplib::Response & res) {
    if (g_registry.dit.empty() || g_registry.text_enc.empty() || g_registry.vae.empty()) {
        json_error(res, 501, "No synth models in registry (need dit + text-encoder + vae)");
        return;
    }

    // parse request: plain JSON (single or array) or multipart (JSON + audio file or src_latents).
    // synth_model, lm_model, adapter, adapter_scale travel inside AceRequest now.
    std::vector<AceRequest> ace_reqs;
    float *                 src_interleaved = nullptr;
    int                     src_len         = 0;
    std::vector<float>      src_latents;
    int                     src_T_latent    = 0;
    float *                 ref_interleaved = nullptr;
    int                     ref_len         = 0;
    std::vector<float>      ref_latents;
    int                     ref_T_latent = 0;

    if (req.is_multipart_form_data()) {
        // multipart mode: single request + optional audio files or src_latents
        AceRequest ace_req;

        std::string json_body;
        if (req.form.has_file("request")) {
            json_body = req.form.get_file("request").content;
        } else if (req.form.has_field("request")) {
            json_body = req.form.get_field("request");
        } else {
            json_error(res, 400, "Multipart: missing 'request' part");
            return;
        }
        if (!request_parse_json(&ace_req, json_body.c_str())) {
            json_error(res, 400, "Multipart: invalid JSON in 'request' part");
            return;
        }

        // src_latents wins over audio when both are sent: a client that
        // already cached the latent skips the VAE encode regardless of
        // whether it also attached the original audio for fallback.
        if (req.form.has_file("src_latents")) {
            const auto & file      = req.form.get_file("src_latents");
            int          http_code = 0;
            int          T         = latent_payload_validate(file.content.size(), &http_code);
            if (T < 0) {
                json_error(res, http_code,
                           http_code == 413 ? "src_latents exceeds max frames" :
                                              "src_latents size not a multiple of 64*4 bytes");
                return;
            }
            src_latents.assign(reinterpret_cast<const float *>(file.content.data()),
                               reinterpret_cast<const float *>(file.content.data()) + (size_t) T * LATENT_CHANNELS);
            src_T_latent = T;
            fprintf(stderr, "[Server] Source latents: %d frames (%.2fs), VAE encode skipped\n", T,
                    (float) T * 1920.0f / 48000.0f);
        }

        if (req.form.has_file("audio")) {
            auto file = req.form.get_file("audio");
            if (file.content.empty()) {
                json_error(res, 400, "Multipart: empty 'audio' part");
                return;
            }
            int     T_audio = 0;
            float * planar  = audio_read_48k_buf((const uint8_t *) file.content.data(), file.content.size(), &T_audio);
            if (!planar || T_audio <= 0) {
                json_error(res, 400, "Failed to decode audio");
                return;
            }
            fprintf(stderr, "[Server] Source audio: %.2fs @ 48kHz\n", (float) T_audio / 48000.0f);
            src_interleaved = audio_planar_to_interleaved(planar, T_audio);
            free(planar);
            src_len = T_audio;
        }

        if (req.form.has_file("ref_latents")) {
            const auto & file      = req.form.get_file("ref_latents");
            int          http_code = 0;
            int          T         = latent_payload_validate(file.content.size(), &http_code);
            if (T < 0) {
                json_error(res, http_code,
                           http_code == 413 ? "ref_latents exceeds max frames" :
                                              "ref_latents size not a multiple of 64*4 bytes");
                return;
            }
            ref_latents.assign(reinterpret_cast<const float *>(file.content.data()),
                               reinterpret_cast<const float *>(file.content.data()) + (size_t) T * LATENT_CHANNELS);
            ref_T_latent = T;
            fprintf(stderr, "[Server] Reference latents: %d frames (%.2fs), VAE encode skipped\n", T,
                    (float) T * 1920.0f / 48000.0f);
        }

        if (req.form.has_file("ref_audio")) {
            auto file = req.form.get_file("ref_audio");
            if (!file.content.empty()) {
                int     T_audio = 0;
                float * planar =
                    audio_read_48k_buf((const uint8_t *) file.content.data(), file.content.size(), &T_audio);
                if (planar && T_audio > 0) {
                    fprintf(stderr, "[Server] Reference audio: %.2fs @ 48kHz\n", (float) T_audio / 48000.0f);
                    ref_interleaved = audio_planar_to_interleaved(planar, T_audio);
                    free(planar);
                    ref_len = T_audio;
                } else {
                    fprintf(stderr, "[Server] WARNING: failed to decode ref_audio, ignoring\n");
                }
            }
        }
        ace_reqs.push_back(ace_req);
    } else {
        // plain JSON body: single object {} or array [{}, ...]
        if (!request_parse_json_array(req.body.c_str(), &ace_reqs)) {
            json_error(res, 400, "Invalid JSON");
            return;
        }
    }

    if (ace_reqs.empty()) {
        json_error(res, 400, "Empty request");
        return;
    }
    if (ace_reqs[0].caption.empty() && ace_reqs[0].task_type != TASK_LEGO && ace_reqs[0].task_type != TASK_EXTRACT &&
        ace_reqs[0].task_type != TASK_COMPLETE) {
        json_error(res, 400, "Caption is required");
        return;
    }

    // Output format from AceRequest.output_format. Converts the string to
    // (output_wav, wav_fmt) using the same parser the CLI uses.
    bool      output_wav = false;
    WavFormat wav_fmt    = WAV_S16;
    {
        bool is_mp3 = true;
        if (!audio_parse_format(ace_reqs[0].output_format.c_str(), is_mp3, wav_fmt)) {
            json_error(res, 400, "Invalid output_format (use: mp3, wav16, wav24, wav32)");
            return;
        }
        output_wav = !is_mp3;
    }
    int peak_clip = ace_reqs[0].peak_clip;

    // create job, spawn worker, return ID
    auto job = job_create();
    fprintf(stderr, "[Server] Job %s created (%d requests)\n", job->id.c_str(), (int) ace_reqs.size());

    work_push([job, reqs = std::move(ace_reqs), src_interleaved, src_len, src_lat = std::move(src_latents),
               src_T_latent, ref_interleaved, ref_len, ref_lat = std::move(ref_latents), ref_T_latent, output_wav,
               wav_fmt, peak_clip]() mutable {
        synth_worker(job, std::move(reqs), src_interleaved, src_len, std::move(src_lat), src_T_latent, ref_interleaved,
                     ref_len, std::move(ref_lat), ref_T_latent, output_wav, wav_fmt, peak_clip);
    });

    // return job ID immediately
    std::string body = "{\"id\":\"" + job->id + "\"}";
    res.set_content(body, "application/json");
}

// understand worker: load LM + tokenizer, run understand, store JSON result + latent in job.
static void understand_worker(std::shared_ptr<Job> job,
                              AceRequest           ace_req,
                              float *              src_interleaved,
                              int                  src_len,
                              std::vector<float>   src_latents,
                              int                  src_T_latent) {
    if (job->cancel.load()) {
        free(src_interleaved);
        job->status.store(JobStatus::CANCELLED);
        return;
    }

    // Resolve LM + DiT (the DiT path carries the tokenizer weights) + VAE.
    std::string        lm_name   = resolve_name(g_registry.lm, ace_req.lm_model, g_loaded_lm);
    std::string        dit_name  = resolve_name(g_registry.dit, ace_req.synth_model, g_loaded_dit);
    std::string        vae_name  = resolve_name(g_registry.vae, ace_req.vae, g_loaded_vae);
    const ModelEntry * lm_entry  = registry_find(g_registry.lm, lm_name.c_str());
    const ModelEntry * dit       = registry_find(g_registry.dit, dit_name.c_str());
    const ModelEntry * vae_entry = registry_find(g_registry.vae, vae_name.c_str());
    if (!lm_entry || !dit || !vae_entry) {
        fprintf(stderr, "[Server] LM, DiT or VAE not found: lm=%s dit=%s vae=%s\n", lm_name.c_str(), dit_name.c_str(),
                vae_name.c_str());
        free(src_interleaved);
        job->status.store(JobStatus::FAILED);
        return;
    }

    AceUnderstandParams p = g_und_params;
    p.model_path          = lm_entry->path.c_str();
    p.dit_path            = dit->path.c_str();
    p.vae_path            = vae_entry->path.c_str();

    AceUnderstand * ctx = ace_understand_load(g_store, &p);
    if (!ctx) {
        fprintf(stderr, "[Server] FATAL: understand load failed\n");
        free(src_interleaved);
        job->status.store(JobStatus::FAILED);
        return;
    }

    AceRequest         out;
    std::vector<float> captured_latent;
    int                captured_T_latent = 0;
    const float *      src_lat_ptr       = src_latents.empty() ? nullptr : src_latents.data();
    int rc = ace_understand_generate(ctx, src_interleaved, src_len, src_lat_ptr, src_T_latent, &ace_req, &out,
                                     &captured_latent, &captured_T_latent, server_cancel_job, (void *) &job->cancel);
    ace_understand_free(ctx);
    free(src_interleaved);

    if (rc != 0) {
        job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::FAILED);
        return;
    }

    // Sticky name hints for resolve_name under --keep-loaded. Master clears
    // them in the default mode since the ctx is gone; we match that behavior.
    if (g_keep_loaded) {
        g_loaded_lm      = lm_name;
        g_loaded_und_dit = dit_name;
        g_loaded_vae     = vae_name;
    } else {
        g_loaded_lm.clear();
        g_loaded_und_dit.clear();
        g_loaded_vae.clear();
    }

    std::string json_part = "[" + request_to_json(&out) + "]";
    job->result_body      = multipart_build_json_latent(json_part, captured_latent, captured_T_latent);
    job->result_mime      = MULTIPART_MIME;
    job->status.store(JobStatus::DONE);
    fprintf(stderr, "[Server] Job %s done (understand)\n", job->id.c_str());
}

// POST /understand
// multipart/form-data: full pipeline (audio + optional JSON params)
//   part "request":     JSON text (optional, for model selection and sampling params)
//   part "audio":       WAV or MP3 file (required unless src_latents provided)
//   part "src_latents": pre-encoded latent bytes (raw f32, [T*64])
// returns: JSON {"id":"N"} immediately. Result is multipart/mixed of one
// JSON part and one optional latent part.
static void handle_understand(const httplib::Request & req, httplib::Response & res) {
    if (g_registry.lm.empty() || g_registry.dit.empty() || g_registry.vae.empty()) {
        json_error(res, 501, "Understand requires LM, DiT and VAE models");
        return;
    }

    if (!req.is_multipart_form_data()) {
        json_error(res, 400, "Understand requires multipart/form-data");
        return;
    }

    // parse multipart: required "audio" or "src_latents" part, optional "request" part for sampling params.
    // synth_model, lm_model, adapter, adapter_scale travel inside AceRequest.
    AceRequest ace_req;
    request_init(&ace_req);
    ace_req.lm_temperature = 0.3f;  // understand default: lower than generation
    ace_req.lm_top_p       = 1.0f;  // understand default: no nucleus sampling

    if (req.form.has_file("request")) {
        const std::string & json = req.form.get_file("request").content;
        if (!request_parse_json(&ace_req, json.c_str())) {
            json_error(res, 400, "Multipart: invalid JSON in 'request' part");
            return;
        }
    } else if (req.form.has_field("request")) {
        const std::string & json = req.form.get_field("request");
        if (!request_parse_json(&ace_req, json.c_str())) {
            json_error(res, 400, "Multipart: invalid JSON in 'request' part");
            return;
        }
    }

    std::vector<float> src_latents;
    int                src_T_latent = 0;
    if (req.form.has_file("src_latents")) {
        const auto & file      = req.form.get_file("src_latents");
        int          http_code = 0;
        int          T         = latent_payload_validate(file.content.size(), &http_code);
        if (T < 0) {
            json_error(
                res, http_code,
                http_code == 413 ? "src_latents exceeds max frames" : "src_latents size not a multiple of 64*4 bytes");
            return;
        }
        src_latents.assign(reinterpret_cast<const float *>(file.content.data()),
                           reinterpret_cast<const float *>(file.content.data()) + (size_t) T * LATENT_CHANNELS);
        src_T_latent = T;
        fprintf(stderr, "[Server] Understand source: %d latent frames (%.2fs)\n", T, (float) T * 1920.0f / 48000.0f);
    }

    float * src_interleaved = nullptr;
    int     src_len         = 0;
    if (req.form.has_file("audio")) {
        auto file = req.form.get_file("audio");
        if (file.content.empty()) {
            json_error(res, 400, "Multipart: empty 'audio' part");
            return;
        }
        // decode directly from multipart buffer (WAV/MP3 auto-detected)
        int     T_audio = 0;
        float * planar  = audio_read_48k_buf((const uint8_t *) file.content.data(), file.content.size(), &T_audio);
        if (!planar || T_audio <= 0) {
            json_error(res, 400, "Failed to decode audio");
            return;
        }
        fprintf(stderr, "[Server] Understand source: %.2fs @ 48kHz\n", (float) T_audio / 48000.0f);
        // convert planar [L:T][R:T] to interleaved [L0,R0,L1,R1,...] for pipeline
        src_interleaved = audio_planar_to_interleaved(planar, T_audio);
        free(planar);
        src_len = T_audio;
    }

    if (!src_interleaved && src_T_latent == 0) {
        json_error(res, 400, "Multipart: missing 'audio' or 'src_latents' part");
        return;
    }

    auto job = job_create();
    fprintf(stderr, "[Server] Job %s created (understand)\n", job->id.c_str());

    request_resolve_lm_seed(&ace_req);

    work_push([job, ace_req, src_interleaved, src_len, latents = std::move(src_latents), src_T_latent]() mutable {
        understand_worker(job, ace_req, src_interleaved, src_len, std::move(latents), src_T_latent);
    });

    std::string body = "{\"id\":\"" + job->id + "\"}";
    res.set_content(body, "application/json");
}

// decode worker: VAE decode only. Loads the requested VAE from the store,
// decodes the supplied latents to 48kHz stereo audio, encodes to the
// requested output format and stores the audio body in the job. The client
// already holds the latent it just uploaded, so the response carries audio
// only; no need to echo the input back.
static void decode_worker(std::shared_ptr<Job> job,
                          AceRequest           ace_req,
                          std::vector<float>   src_latents,
                          int                  src_T_latent,
                          bool                 output_wav,
                          WavFormat            wav_fmt,
                          int                  peak_clip) {
    if (job->cancel.load()) {
        job->status.store(JobStatus::CANCELLED);
        return;
    }

    std::string        vae_name  = resolve_name(g_registry.vae, ace_req.vae, g_loaded_vae);
    const ModelEntry * vae_entry = registry_find(g_registry.vae, vae_name.c_str());
    if (!vae_entry) {
        fprintf(stderr, "[Server] decode: VAE not found: %s\n", vae_name.c_str());
        job->status.store(JobStatus::FAILED);
        return;
    }

    ModelKey vae_key;
    vae_key.kind          = MODEL_VAE_DEC;
    vae_key.path          = vae_entry->path;
    vae_key.adapter_scale = 1.0f;

    Timer     t_dec;
    VAEGGML * vae = store_require_vae_dec(g_store, vae_key);
    if (!vae) {
        fprintf(stderr, "[Server] decode: store_require_vae_dec failed\n");
        job->status.store(JobStatus::FAILED);
        return;
    }
    ModelHandle vae_guard(g_store, vae);

    int                T_audio_max = (src_T_latent + 64) * 1920;
    std::vector<float> audio_buf((size_t) T_audio_max * 2);
    int T_audio = vae_ggml_decode_tiled(vae, src_latents.data(), src_T_latent, audio_buf.data(), T_audio_max,
                                        g_synth_params.vae_chunk, g_synth_params.vae_overlap);
    if (T_audio < 0) {
        fprintf(stderr, "[Server] decode: vae_ggml_decode_tiled failed\n");
        job->status.store(JobStatus::FAILED);
        return;
    }
    fprintf(stderr, "[Server] decode: %d latent frames -> %d audio samples (%.2fs), %.0fms\n", src_T_latent, T_audio,
            (float) T_audio / 48000.0f, t_dec.ms());

    if (g_keep_loaded) {
        g_loaded_vae = vae_name;
    } else {
        g_loaded_vae.clear();
    }

    // Encode the audio (peak normalize then mp3 or wav). vae_ggml_decode_tiled
    // writes interleaved stereo, audio_normalize and the encoders consume
    // the same layout the synth path uses.
    if (!output_wav || wav_fmt != WAV_F32) {
        audio_normalize(audio_buf.data(), T_audio * 2, peak_clip);
    }
    std::string  encoded;
    const char * mime = output_wav ? "audio/wav" : "audio/mpeg";
    if (output_wav) {
        encoded = audio_encode_wav(audio_buf.data(), T_audio, 48000, wav_fmt);
    } else {
        encoded = audio_encode_mp3(audio_buf.data(), T_audio, 48000, ace_req.mp3_bitrate, server_cancel_job,
                                   (void *) &job->cancel);
    }

    // Response: raw audio, single Content-Type. No latent in the body: the
    // client just uploaded it, echoing it back would only burn bandwidth.
    job->result_body = std::move(encoded);
    job->result_mime = mime;
    job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::DONE);
    fprintf(stderr, "[Server] Job %s done (decode)\n", job->id.c_str());
}

// encode worker: VAE encode only. Loads the requested VAE encoder, encodes
// the 48kHz interleaved stereo audio into latents [T_25Hz, 64] time-major
// and stores the raw f32 buffer in the job. No LM, no FSQ. Mirrors
// decode_worker, minus the audio codec at the output: the body is raw
// latent bytes the client caches alongside its source card.
static void encode_worker(std::shared_ptr<Job> job, AceRequest ace_req, float * src_interleaved, int src_len) {
    struct buf_guard {
        float * p;

        ~buf_guard() {
            if (p) {
                free(p);
            }
        }
    } buf{ src_interleaved };

    if (job->cancel.load()) {
        job->status.store(JobStatus::CANCELLED);
        return;
    }

    std::string        vae_name  = resolve_name(g_registry.vae, ace_req.vae, g_loaded_vae);
    const ModelEntry * vae_entry = registry_find(g_registry.vae, vae_name.c_str());
    if (!vae_entry) {
        fprintf(stderr, "[Server] encode: VAE not found: %s\n", vae_name.c_str());
        job->status.store(JobStatus::FAILED);
        return;
    }

    ModelKey vae_key;
    vae_key.kind          = MODEL_VAE_ENC;
    vae_key.path          = vae_entry->path;
    vae_key.adapter_scale = 1.0f;

    Timer        t_enc;
    VAEEncoder * vae = store_require_vae_enc(g_store, vae_key);
    if (!vae) {
        fprintf(stderr, "[Server] encode: store_require_vae_enc failed\n");
        job->status.store(JobStatus::FAILED);
        return;
    }
    ModelHandle vae_guard(g_store, vae);

    // 1 latent frame covers 1920 audio samples, plus a safety tile for the
    // tiled encoder boundary rounding. Capped at MAX_T_LATENT so oversized
    // audio never sneaks past the handler-side pre-check.
    int T_latent_max = src_len / 1920 + 64;
    if (T_latent_max > MAX_T_LATENT) {
        T_latent_max = MAX_T_LATENT;
    }
    std::vector<float> latent((size_t) T_latent_max * LATENT_CHANNELS);
    int                T_latent = vae_enc_encode_tiled(vae, src_interleaved, src_len, latent.data(), T_latent_max,
                                                       g_synth_params.vae_chunk, g_synth_params.vae_overlap);
    if (T_latent < 0) {
        fprintf(stderr, "[Server] encode: vae_enc_encode_tiled failed\n");
        job->status.store(JobStatus::FAILED);
        return;
    }
    fprintf(stderr, "[Server] encode: %d audio samples (%.2fs) -> %d latent frames, %.0fms\n", src_len,
            (float) src_len / 48000.0f, T_latent, t_enc.ms());

    if (g_keep_loaded) {
        g_loaded_vae = vae_name;
    } else {
        g_loaded_vae.clear();
    }

    // Response: raw f32 [T_latent, 64] time-major, no header. Single
    // Content-Type, no multipart: the client still holds the source audio
    // it just uploaded, only the fresh latents need to travel back.
    std::string body;
    body.resize((size_t) T_latent * LATENT_FRAME_BYTES);
    std::memcpy(body.data(), latent.data(), body.size());
    job->result_body = std::move(body);
    job->result_mime = "application/octet-stream";
    job->status.store(job->cancel.load() ? JobStatus::CANCELLED : JobStatus::DONE);
    fprintf(stderr, "[Server] Job %s done (encode)\n", job->id.c_str());
}

// POST /vae
// multipart/form-data: single VAE entrypoint, dispatches on which side is
// supplied in the request body. Symmetric with /synth and /understand on
// the 'audio or src_latents' input contract, except here they are mutually
// exclusive (the direction of travel depends on which one you send).
//   part "audio":       WAV or MP3 source audio -> encode path, latents out.
//   part "src_latents": pre-encoded latent bytes -> decode path, audio out.
//   part "request":     JSON text (optional, for VAE selection, output
//                       format and peak_clip on the decode path).
// returns: JSON {"id":"N"} immediately. Result body is either raw .vae
// bytes (application/octet-stream, f32 [T*64] time-major, no header) when
// audio was sent, or audio/mpeg|audio/wav when src_latents was sent. No
// echo of the uploaded side: the client already holds it.
static void handle_vae(const httplib::Request & req, httplib::Response & res) {
    if (g_registry.vae.empty()) {
        json_error(res, 501, "VAE endpoint requires a VAE in the registry");
        return;
    }
    if (!req.is_multipart_form_data()) {
        json_error(res, 400, "VAE endpoint requires multipart/form-data");
        return;
    }

    AceRequest ace_req;
    request_init(&ace_req);

    if (req.form.has_file("request")) {
        const std::string & json = req.form.get_file("request").content;
        if (!request_parse_json(&ace_req, json.c_str())) {
            json_error(res, 400, "Multipart: invalid JSON in 'request' part");
            return;
        }
    } else if (req.form.has_field("request")) {
        const std::string & json = req.form.get_field("request");
        if (!request_parse_json(&ace_req, json.c_str())) {
            json_error(res, 400, "Multipart: invalid JSON in 'request' part");
            return;
        }
    }

    bool has_audio   = req.form.has_file("audio");
    bool has_latents = req.form.has_file("src_latents");
    if (has_audio == has_latents) {
        json_error(res, 400, "Multipart: provide exactly one of 'audio' (encode) or 'src_latents' (decode)");
        return;
    }

    if (has_audio) {
        // encode path: audio in -> raw latents out
        const auto & file = req.form.get_file("audio");
        if (file.content.empty()) {
            json_error(res, 400, "Multipart: empty 'audio' part");
            return;
        }
        int     T_audio = 0;
        float * planar  = audio_read_48k_buf((const uint8_t *) file.content.data(), file.content.size(), &T_audio);
        if (!planar || T_audio <= 0) {
            if (planar) {
                free(planar);
            }
            json_error(res, 400, "Failed to decode audio");
            return;
        }
        if ((int64_t) T_audio / 1920 >= (int64_t) MAX_T_LATENT) {
            free(planar);
            json_error(res, 413, "audio exceeds max duration (10 min)");
            return;
        }
        float * src_interleaved = audio_planar_to_interleaved(planar, T_audio);
        free(planar);
        int src_len = T_audio;

        auto job = job_create();
        fprintf(stderr, "[Server] Job %s created (vae encode, %.2fs audio)\n", job->id.c_str(),
                (float) src_len / 48000.0f);

        work_push([job, ace_req, src_interleaved, src_len]() mutable {
            encode_worker(job, ace_req, src_interleaved, src_len);
        });

        std::string body = "{\"id\":\"" + job->id + "\"}";
        res.set_content(body, "application/json");
        return;
    }

    // decode path: raw latents in -> audio out
    const auto & file      = req.form.get_file("src_latents");
    int          http_code = 0;
    int          T         = latent_payload_validate(file.content.size(), &http_code);
    if (T < 0) {
        json_error(
            res, http_code,
            http_code == 413 ? "src_latents exceeds max frames" : "src_latents size not a multiple of 64*4 bytes");
        return;
    }
    std::vector<float> src_latents(reinterpret_cast<const float *>(file.content.data()),
                                   reinterpret_cast<const float *>(file.content.data()) + (size_t) T * LATENT_CHANNELS);

    bool      output_wav = false;
    WavFormat wav_fmt    = WAV_S16;
    {
        bool is_mp3 = true;
        if (!audio_parse_format(ace_req.output_format.c_str(), is_mp3, wav_fmt)) {
            json_error(res, 400, "Invalid output_format (use: mp3, wav16, wav24, wav32)");
            return;
        }
        output_wav = !is_mp3;
    }
    int peak_clip = ace_req.peak_clip;

    auto job = job_create();
    fprintf(stderr, "[Server] Job %s created (vae decode, %d latent frames)\n", job->id.c_str(), T);

    work_push([job, ace_req, latents = std::move(src_latents), T, output_wav, wav_fmt, peak_clip]() mutable {
        decode_worker(job, ace_req, std::move(latents), T, output_wav, wav_fmt, peak_clip);
    });

    std::string body = "{\"id\":\"" + job->id + "\"}";
    res.set_content(body, "application/json");
}

// GET /props
// server configuration, available models, and default request.
// the webui reads this at boot to populate dropdowns and status indicators.
static void handle_props(const httplib::Request &, httplib::Response & res) {
    yyjson_mut_doc * doc  = yyjson_mut_doc_new(NULL);
    yyjson_mut_val * root = yyjson_mut_obj(doc);
    yyjson_mut_doc_set_root(doc, root);

    yyjson_mut_obj_add_str(doc, root, "version", ACE_VERSION);

    // helper: build a JSON array of model entry names
    auto add_names = [&](yyjson_mut_val * parent, const char * key, const std::vector<ModelEntry> & bucket) {
        yyjson_mut_val * arr = yyjson_mut_arr(doc);
        for (const auto & e : bucket) {
            yyjson_mut_arr_add_str(doc, arr, e.name.c_str());
        }
        yyjson_mut_obj_add_val(doc, parent, key, arr);
    };

    // models: available model names per bucket
    yyjson_mut_val * models = yyjson_mut_obj(doc);
    yyjson_mut_obj_add_val(doc, root, "models", models);
    add_names(models, "lm", g_registry.lm);
    add_names(models, "embedding", g_registry.text_enc);
    add_names(models, "dit", g_registry.dit);
    add_names(models, "vae", g_registry.vae);

    // adapters: available adapter names
    yyjson_mut_val * adapters_arr = yyjson_mut_arr(doc);
    for (const auto & e : g_registry.adapters) {
        yyjson_mut_arr_add_str(doc, adapters_arr, e.name.c_str());
    }
    yyjson_mut_obj_add_val(doc, root, "adapters", adapters_arr);

    // cli: server settings
    yyjson_mut_val * cli = yyjson_mut_obj(doc);
    yyjson_mut_obj_add_val(doc, root, "cli", cli);
    yyjson_mut_obj_add_int(doc, cli, "max_batch", g_max_batch);

    // default: full AceRequest with all defaults from request_init().
    // the webui reads this to populate LM placeholders.
    // DiT fields (inference_steps, guidance_scale, shift) are 0 = auto-detect;
    // their resolved placeholders come from presets below.
    AceRequest defaults;
    request_init(&defaults);
    std::string      defaults_str  = request_to_json(&defaults, false);
    yyjson_doc *     defaults_doc  = yyjson_read(defaults_str.c_str(), defaults_str.size(), 0);
    yyjson_mut_val * defaults_copy = yyjson_val_mut_copy(doc, yyjson_doc_get_root(defaults_doc));
    yyjson_mut_obj_add_val(doc, root, "default", defaults_copy);
    yyjson_doc_free(defaults_doc);

    // presets: auto-detect values for DiT sampling params.
    // the webui switches placeholders based on the selected DiT model.
    yyjson_mut_val * presets = yyjson_mut_obj(doc);
    yyjson_mut_obj_add_val(doc, root, "presets", presets);

    yyjson_mut_val * turbo = yyjson_mut_obj(doc);
    yyjson_mut_obj_add_int(doc, turbo, "inference_steps", 8);
    yyjson_mut_obj_add_real(doc, turbo, "guidance_scale", 1.0);
    yyjson_mut_obj_add_real(doc, turbo, "shift", 3.0);
    yyjson_mut_obj_add_val(doc, presets, "turbo", turbo);

    yyjson_mut_val * sft = yyjson_mut_obj(doc);
    yyjson_mut_obj_add_int(doc, sft, "inference_steps", 50);
    yyjson_mut_obj_add_real(doc, sft, "guidance_scale", 1.0);
    yyjson_mut_obj_add_real(doc, sft, "shift", 1.0);
    yyjson_mut_obj_add_val(doc, presets, "sft", sft);

    // serialize
    yyjson_write_flag flags = YYJSON_WRITE_PRETTY | YYJSON_WRITE_PRETTY_TWO_SPACES | YYJSON_WRITE_FP_TO_FIXED(2);
    char *            json  = yyjson_mut_write(doc, flags, NULL);
    yyjson_mut_doc_free(doc);
    res.set_content(json, "application/json");
    free(json);
}

static void usage(const char * prog) {
    AceLmParams    lm_d;
    AceSynthParams synth_d;
    ace_lm_default_params(&lm_d);
    ace_synth_default_params(&synth_d);

    fprintf(stderr, "acestep.cpp %s\n\n", ACE_VERSION);
    fprintf(stderr,
            "Usage: %s --models <dir> [options]\n"
            "\n"
            "Required:\n"
            "  --models <dir>          Directory of GGUF model files\n"
            "\n"
            "Adapter:\n"
            "  --adapters <dir>        Directory of adapters\n"
            "\n"
            "Memory control:\n"
            "  --keep-loaded           Keep models in VRAM between requests\n"
            "  --vae-chunk <N>         Latent frames per tile (default: %d)\n"
            "  --vae-overlap <N>       Overlap frames per side (default: %d)\n"
            "\n"
            "Server:\n"
            "  --host <addr>           Listen address (default: 127.0.0.1)\n"
            "  --port <N>              Listen port (default: 8080)\n"
            "  --max-batch <N>         LM batch limit (default: %d)\n"
            "  --max-seq <N>           KV cache size (default: %d)\n"
            "\n"
            "Debug:\n"
            "  --no-fsm                Disable FSM constrained decoding\n"
            "  --no-fa                 Disable flash attention\n"
            "  --no-batch-cfg          Split CFG into two separate forwards (LM + DiT)\n"
            "  --clamp-fp16            Clamp hidden states to FP16 range\n",
            prog, synth_d.vae_chunk, synth_d.vae_overlap, g_max_batch, lm_d.max_seq);
}

int main(int argc, char ** argv) {
    ace_lm_default_params(&g_lm_params);
    ace_synth_default_params(&g_synth_params);

    const char * host         = "127.0.0.1";
    int          port         = 8080;
    const char * models_dir   = nullptr;
    const char * adapters_dir = nullptr;

    if (argc < 2) {
        usage(argv[0]);
        return 1;
    }

    for (int i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "--models") && i + 1 < argc) {
            models_dir = argv[++i];
        } else if (!strcmp(argv[i], "--adapters") && i + 1 < argc) {
            adapters_dir = argv[++i];
        } else if (!strcmp(argv[i], "--max-seq") && i + 1 < argc) {
            g_lm_params.max_seq = atoi(argv[++i]);

            // vae tiling
        } else if (!strcmp(argv[i], "--vae-chunk") && i + 1 < argc) {
            g_synth_params.vae_chunk = atoi(argv[++i]);
        } else if (!strcmp(argv[i], "--vae-overlap") && i + 1 < argc) {
            g_synth_params.vae_overlap = atoi(argv[++i]);
        } else if (!strcmp(argv[i], "--keep-loaded")) {
            g_keep_loaded = true;

            // server
        } else if (!strcmp(argv[i], "--host") && i + 1 < argc) {
            host = argv[++i];
        } else if (!strcmp(argv[i], "--port") && i + 1 < argc) {
            port = atoi(argv[++i]);
        } else if (!strcmp(argv[i], "--max-batch") && i + 1 < argc) {
            g_max_batch = atoi(argv[++i]);

            // debug
        } else if (!strcmp(argv[i], "--no-fsm")) {
            g_lm_params.use_fsm = false;
        } else if (!strcmp(argv[i], "--no-fa")) {
            g_lm_params.use_fa    = false;
            g_synth_params.use_fa = false;
        } else if (!strcmp(argv[i], "--no-batch-cfg")) {
            g_lm_params.use_batch_cfg    = false;
            g_synth_params.use_batch_cfg = false;
        } else if (!strcmp(argv[i], "--clamp-fp16")) {
            g_lm_params.clamp_fp16    = true;
            g_synth_params.clamp_fp16 = true;

        } else if (!strcmp(argv[i], "--help") || !strcmp(argv[i], "-h")) {
            usage(argv[0]);
            return 0;
        } else {
            fprintf(stderr, "Unknown option: %s\n", argv[i]);
            usage(argv[0]);
            return 1;
        }
    }

    // --models is required
    if (!models_dir) {
        fprintf(stderr, "[Server] ERROR: --models is required\n");
        usage(argv[0]);
        return 1;
    }

    // stderr capture for SSE /logs (must be after arg parsing so --help prints directly)
    LogCapture log_capture;

    // scan models directory (reads GGUF metadata only)
    fprintf(stderr, "[Server] Scanning models in %s\n", models_dir);
    if (!registry_scan(&g_registry, models_dir)) {
        fprintf(stderr, "[Server] ERROR: no GGUF models found in %s\n", models_dir);
        return 1;
    }

    // scan adapters directory (optional)
    if (adapters_dir) {
        fprintf(stderr, "[Server] Scanning adapters in %s\n", adapters_dir);
        registry_scan_adapters(&g_registry, adapters_dir);
    }

    // validate pipeline
    bool have_lm    = !g_registry.lm.empty();
    bool have_dit   = !g_registry.dit.empty();
    bool have_enc   = !g_registry.text_enc.empty();
    bool have_vae   = !g_registry.vae.empty();
    bool have_synth = have_dit && have_enc && have_vae;

    // partial synth: some components found but pipeline incomplete
    if (!have_synth && (have_dit || have_enc || have_vae)) {
        char missing[64];
        int  n = 0;
        if (!have_dit) {
            n += snprintf(missing + n, sizeof(missing) - n, "%sDiT", n ? ", " : "");
        }
        if (!have_enc) {
            n += snprintf(missing + n, sizeof(missing) - n, "%sText-Enc", n ? ", " : "");
        }
        if (!have_vae) {
            n += snprintf(missing + n, sizeof(missing) - n, "%sVAE", n ? ", " : "");
        }
        if (have_lm) {
            fprintf(stderr, "[Server] WARNING: /synth unavailable, missing: %s\n", missing);
        } else {
            fprintf(stderr, "[Server] ERROR: no usable pipeline, synth missing: %s\n", missing);
            return 1;
        }
    }

    // clamp max_batch
    if (g_max_batch < 1) {
        g_max_batch = 1;
    }
    if (g_max_batch > 9) {
        g_max_batch = 9;
    }
    g_lm_params.max_batch = g_max_batch;

    // init understand params (dit + vae resolved per-request)
    ace_understand_default_params(&g_und_params);
    g_und_params.use_fa      = g_lm_params.use_fa;
    g_und_params.use_fsm     = g_lm_params.use_fsm;
    g_und_params.max_seq     = g_lm_params.max_seq;         // must match ace_lm: part of the LM ModelKey
    g_und_params.max_batch   = g_lm_params.max_batch;       // must match ace_lm: part of the LM ModelKey
    g_und_params.vae_chunk   = g_synth_params.vae_chunk;    // share --vae-chunk with /synth
    g_und_params.vae_overlap = g_synth_params.vae_overlap;  // share --vae-overlap with /synth

    bool have_understand = have_lm && have_dit && have_vae;

    // central store: one policy for the whole server lifetime. STRICT keeps
    // at most one GPU module resident at a time; --keep-loaded flips it to
    // NEVER and lets the working set accumulate across requests.
    g_store = store_create(g_keep_loaded ? EVICT_NEVER : EVICT_STRICT);

    // setup HTTP server
    httplib::Server svr;
    g_svr = &svr;

    // per-operation socket idle timeout (httplib default is 5s).
    // generous margin for slow networks and large audio transfers.
    svr.set_read_timeout(600);
    svr.set_write_timeout(600);

    // SO_REUSEADDR: allow rebind after TIME_WAIT (normal restart).
    // no SO_REUSEPORT: fail if another process is actively listening.
    svr.set_socket_options([](socket_t sock) {
        int one = 1;
#ifdef _WIN32
        setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (const char *) &one, sizeof(one));
#else
        setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
#endif
    });

    // reject oversized bodies (256 MB: src + ref audio, up to 10min WAV each)
    svr.set_payload_max_length(256 * 1024 * 1024);

    // all endpoints are always registered. handlers return 501 when the
    // backing pipeline has no models in the registry.
    svr.Post("/lm", handle_lm);
    svr.Post("/synth", handle_synth);
    svr.Post("/understand", handle_understand);
    svr.Post("/vae", handle_vae);
    svr.Get("/health", [](const httplib::Request &, httplib::Response & res) {
        res.set_content("{\"status\":\"ok\"}", "application/json");
    });
    svr.Get("/props", handle_props);
    svr.Get("/logs", handle_logs);

    // job system endpoints
    svr.Get("/job", [](const httplib::Request & req, httplib::Response & res) {
        if (!req.has_param("id")) {
            json_error(res, 400, "Missing id parameter");
            return;
        }
        auto job = job_find(req.get_param_value("id"));
        if (!job) {
            json_error(res, 404, "Job not found");
            return;
        }
        // ?result=1: return result body
        if (req.has_param("result") && req.get_param_value("result") == "1") {
            if (job->status.load() != JobStatus::DONE) {
                json_error(res, 404, "Result not ready");
                return;
            }
            res.set_content(job->result_body, job->result_mime);
            return;
        }
        // default: return status JSON
        std::string body = "{\"status\":\"";
        body += job_status_str(job->status.load());
        body += "\"}";
        res.set_content(body, "application/json");
    });
    svr.Post("/job", [](const httplib::Request & req, httplib::Response & res) {
        if (!req.has_param("id")) {
            json_error(res, 400, "Missing id parameter");
            return;
        }
        auto job = job_find(req.get_param_value("id"));
        if (!job) {
            json_error(res, 404, "Job not found");
            return;
        }
        // ?cancel=1: cancel the job
        if (req.has_param("cancel") && req.get_param_value("cancel") == "1") {
            JobStatus status = job->status.load();
            if (status == JobStatus::RUNNING) {
                job->cancel.store(true);
                fprintf(stderr, "[Server] Cancel requested for job %s\n", job->id.c_str());
                status = JobStatus::CANCELLED;
            }
            std::string body = "{\"status\":\"";
            body += job_status_str(status);
            body += "\"}";
            res.set_content(body, "application/json");
            return;
        }
        json_error(res, 400, "Unknown action");
    });

    // embedded webui: gzipped single-page app (built by tools/webui/).
    // the browser decompresses transparently via Content-Encoding: gzip.
    // the .gz is committed to git so cloning + cmake + make gives a working UI.
    if (index_html_gz_len > 0) {
        svr.Get("/", [](const httplib::Request & req, httplib::Response & res) {
            if (req.get_header_value("Accept-Encoding").find("gzip") == std::string::npos) {
                res.set_content("Error: gzip is not supported by this browser", "text/plain");
            } else {
                res.set_header("Content-Encoding", "gzip");
                res.set_content(reinterpret_cast<const char *>(index_html_gz), index_html_gz_len,
                                "text/html; charset=utf-8");
            }
        });
    }

    // graceful shutdown on SIGINT/SIGTERM
    signal(SIGINT, on_signal);
    signal(SIGTERM, on_signal);

    // start FIFO worker thread (processes all GPU jobs in order)
    std::thread worker(worker_main);

    fprintf(stderr, "[Server] acestep.cpp %s\n", ACE_VERSION);
    fprintf(stderr, "[Server] Listening on %s:%d\n", host, port);
    fprintf(stderr, "[Server] Pipelines:%s%s%s\n", have_lm ? " /lm" : "", have_synth ? " /synth" : "",
            have_understand ? " /understand" : "");
    fprintf(stderr, "[Server] Models: %zu LM, %zu Text-Enc, %zu DiT, %zu VAE, %zu Adapter\n", g_registry.lm.size(),
            g_registry.text_enc.size(), g_registry.dit.size(), g_registry.vae.size(), g_registry.adapters.size());
    if (!svr.listen(host, port)) {
        fprintf(stderr, "[Server] FATAL: cannot bind %s:%d\n", host, port);
    }

    // stop worker thread (finishes current job, discards pending)
    {
        std::lock_guard<std::mutex> lock(mtx_work);
        g_work_stop = true;
    }
    cv_work.notify_one();
    worker.join();

    // cleanup
    fprintf(stderr, "[Server] Shutting down...\n");
    store_free(g_store);
    fprintf(stderr, "[Server] Done\n");

    return 0;
}