#pragma once
// adapter-merge.h: runtime adapter merge into GGUF weights before QKV fusion.
//
// Supported algorithms:
//   LoRA (low rank adaptation):  delta = (alpha / rank) * B @ A
//   LoKr (low rank Kronecker):   delta = (alpha / rank) * kron(w1, w2_a @ w2_b)
//                                optional DoRA magnitude rescale per output row
//                                when a dora_scale tensor is present.
//
// Called after individual GGUF projection tensors are loaded into WeightCtx
// but BEFORE wctx_alloc uploads to GPU and BEFORE QKV fusion concatenation.
//
// Each projection (q_proj, k_proj, v_proj, o_proj) has its own PendingCopy
// even when destined for a fused QKV tensor. We patch each one separately,
// so fusion proceeds normally on already merged data.
//
// Performance: one backend graph per tensor runs the full pipeline. The base
// weight is uploaded in its native GGUF type directly from the mmap, the
// backend dequantizes it to F32 via ggml_cast, the adapter delta is built
// from adapter factors, BF16 rounded in a cast chain, added to base, DoRA
// rescaled when requested, then cast back to the native GGUF type when the
// backend supports that encode direction. The result is downloaded straight
// into the PendingCopy staging buffer.
//
// ACE-Step GGUFs ship in BF16, Q8_0, Q4_K_M, Q5_K_M, and Q6_K. On CUDA the
// backend encode cast handles F32 -> BF16 and F32 -> Q8_0 directly; the
// K-quants have no F32 -> native kernel so the graph terminates at F32 and
// ggml_quantize_chunk completes the job on host. Either way the base upload
// PCIe is cut vs a prior F32 upload, and the dequant, add, BF16 round and
// DoRA rescale all run on the backend.
// PendingCopy lookup is O(1) via hashmap.

#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "gguf-weights.h"
#include "safetensors.h"
#include "weight-ctx.h"
#include "weight-source.h"
#include "yyjson.h"
#include "hot-step-params.h"  // HOT-Step: per-group adapter scales via sideband

#include <sys/stat.h>
#ifdef _WIN32
#    ifndef S_ISDIR
#        define S_ISDIR(m) (((m) & _S_IFMT) == _S_IFDIR)
#    endif
#endif

#include <cmath>
#include <cstdio>
#include <cstring>
#include <functional>
#include <map>
#include <string>
#include <unordered_map>
#include <vector>

// Convert safetensors tensor data to F32 based on dtype string.
// Handles "F32", "BF16", "F16". Returns false for unknown dtypes.
static bool adapter_to_f32(const void * src, float * dst, int64_t n, const std::string & dtype) {
    if (dtype == "F32") {
        memcpy(dst, src, (size_t) n * sizeof(float));
    } else if (dtype == "BF16") {
        ggml_bf16_to_fp32_row((const ggml_bf16_t *) src, dst, n);
    } else if (dtype == "F16") {
        ggml_fp16_to_fp32_row((const ggml_fp16_t *) src, dst, n);
    } else {
        return false;
    }
    return true;
}

// Map a LoRA safetensors key to the GGUF base tensor name.
//
// Supported key formats (all map to GGUF "decoder.layers.0.self_attn.q_proj.weight"):
//
//   PEFT adapter_model.safetensors:
//     base_model.model.layers.0.self_attn.q_proj.lora_A.default.weight
//     base_model.model.layers.0.self_attn.q_proj.lora_A.weight
//
//   ComfyUI single-file (no prefix):
//     layers.0.self_attn.q_proj.lora_A.weight
//
//   ComfyUI single-file (diffusion_model prefix):
//     diffusion_model.layers.0.self_attn.q_proj.lora_A.weight
//
//   PEFT raw export (wrapped modules):
//     base_model.model.layers.0.mlp.original_mlp.down_proj.lora_A.default.weight
//     (the ".original_mlp." infix is stripped before matching)
//
// Steps: strip known prefix, extract module path before ".lora_",
// prepend "decoder." if needed, append ".weight".
static std::string lora_base_name(const std::string & key) {
    std::string s = key;

    // strip known prefixes (PEFT, ComfyUI)
    static const char * prefixes[] = {
        "base_model.model.",  // PEFT
        "diffusion_model.",   // ComfyUI official ACE-Step format
    };
    for (const char * pfx : prefixes) {
        size_t pfx_len = strlen(pfx);
        if (s.compare(0, pfx_len, pfx) == 0) {
            s = s.substr(pfx_len);
            break;
        }
    }

    // strip known PEFT module wrapper infixes.
    // HuggingFace PEFT wraps modules as e.g. "mlp.original_mlp.proj" instead
    // of "mlp.proj".  The ".original_<module>." segment is redundant.
    static const char * peft_infixes[] = {
        ".original_mlp.",
        ".original_self_attn.",
        ".original_cross_attn.",
    };
    for (const char * infix : peft_infixes) {
        size_t ipos = s.find(infix);
        if (ipos != std::string::npos) {
            s.erase(ipos, strlen(infix) - 1);  // erase infix, keep trailing dot
            break;
        }
    }

    // everything before ".lora_" is the module path
    size_t pos = s.find(".lora_");
    if (pos == std::string::npos) {
        return "";
    }
    s = s.substr(0, pos);

    // ensure decoder prefix (PEFT wraps the decoder directly,
    // so the internal path starts at "layers." not "decoder.layers.")
    if (s.compare(0, 8, "decoder.") != 0) {
        s = "decoder." + s;
    }

    return s + ".weight";
}

// Check whether a safetensors key is a lora_A/down or lora_B/up weight.
// PEFT uses .lora_A. / .lora_B., ComfyUI single-file uses .lora_down. / .lora_up.
static bool lora_is_a(const std::string & key) {
    return key.find(".lora_A.") != std::string::npos || key.find(".lora_down.") != std::string::npos;
}

static bool lora_is_b(const std::string & key) {
    return key.find(".lora_B.") != std::string::npos || key.find(".lora_up.") != std::string::npos;
}

// Read adapter_config.json for alpha. Returns alpha or 0 if not found.
// Rank is always read from the actual tensor shapes (more reliable).
static int adapter_read_alpha(const char * dir) {
    std::string path = std::string(dir) + "/adapter_config.json";

    FILE * f = fopen(path.c_str(), "rb");
    if (!f) {
        return 0;
    }

    fseek(f, 0, SEEK_END);
    long len = ftell(f);
    fseek(f, 0, SEEK_SET);

    std::vector<char> buf((size_t) len + 1);
    size_t            nr = fread(buf.data(), 1, (size_t) len, f);
    fclose(f);
    if (nr != (size_t) len) {
        return 0;
    }
    buf[(size_t) len] = '\0';

    const char * json  = buf.data();
    int          alpha = 0;

    // look for "lora_alpha": <int>
    const char * p = strstr(json, "\"lora_alpha\"");
    if (p) {
        p = strchr(p + 12, ':');
        if (p) {
            alpha = atoi(p + 1);
        }
    }

    // fallback: try "alpha": <int> (some configs use this)
    if (alpha == 0) {
        p = strstr(json, "\"alpha\"");
        if (p) {
            p = strchr(p + 7, ':');
            if (p) {
                alpha = atoi(p + 1);
            }
        }
    }

    if (alpha > 0) {
        fprintf(stderr, "[Adapter] adapter_config.json: alpha=%d\n", alpha);
    }
    return alpha;
}

// Read linear_dim from LyCORIS __metadata__.lokr_config for LoKr payloads.
// LyCORIS stores lokr_config as a serialized JSON string inside the safetensors
// header __metadata__ object. This dim is the LoRA rank used for the w2
// factorization, and is the only place it lives when w2 is stored as a single
// monolithic tensor (LyCORIS use_w2 path, factor=-1 with linear_dim bigger
// than the factorized budget). Returns 0 when absent or unparseable.
static int adapter_read_lokr_dim(const STFile & st) {
    if (st.data_offset <= 8) {
        return 0;
    }
    size_t       hdr_len = st.data_offset - 8;
    const char * hdr     = (const char *) st.mapping + 8;

    yyjson_doc * doc = yyjson_read(hdr, hdr_len, 0);
    if (!doc) {
        return 0;
    }
    int          dim  = 0;
    yyjson_val * root = yyjson_doc_get_root(doc);
    yyjson_val * meta = yyjson_obj_get(root, "__metadata__");
    if (meta && yyjson_is_obj(meta)) {
        yyjson_val * cfg = yyjson_obj_get(meta, "lokr_config");
        if (cfg && yyjson_is_str(cfg)) {
            yyjson_doc * sub = yyjson_read(yyjson_get_str(cfg), yyjson_get_len(cfg), 0);
            if (sub) {
                yyjson_val * ld = yyjson_obj_get(yyjson_doc_get_root(sub), "linear_dim");
                if (ld && yyjson_is_int(ld)) {
                    dim = (int) yyjson_get_int(ld);
                }
                yyjson_doc_free(sub);
            }
        }
    }
    yyjson_doc_free(doc);
    return dim;
}

// Requant F32 data back to original type. Writes into dst buffer.
// Returns the number of bytes written. Used only as host fallback when the
// backend lacks an F32 -> native cast kernel (K-quants on CUDA today).
static size_t adapter_requant(const float * src, void * dst, int64_t nel, int64_t n_per_row, enum ggml_type type) {
    if (type == GGML_TYPE_F32) {
        size_t nb = (size_t) nel * sizeof(float);
        memcpy(dst, src, nb);
        return nb;
    }

    const struct ggml_type_traits * traits = ggml_get_type_traits(type);

    if (traits->is_quantized) {
        // quantized types: use ggml_quantize_chunk (handles block alignment)
        int64_t nrows = nel / n_per_row;
        size_t  qsize = ggml_row_size(type, n_per_row) * (size_t) nrows;
        ggml_quantize_chunk(type, src, dst, 0, nrows, n_per_row, NULL);
        return qsize;
    }

    // non quantized (BF16, F16): use from_float_ref
    if (traits->from_float_ref) {
        size_t nb = (size_t) nel * traits->type_size;
        traits->from_float_ref(src, dst, nel);
        return nb;
    }

    fprintf(stderr, "[Adapter] WARNING: no requant for type %d\n", type);
    return 0;
}

// True when the backend has an F32 -> native encode cast kernel for this type.
// Probed via ggml_backend_supports_op on a throwaway GGML_OP_CPY tensor: the
// backend only inspects src[0]->type and src[1]->type so no buffer allocation
// or data upload is needed. When false, the caller downloads F32 and calls
// adapter_requant on host.
static bool adapter_backend_can_encode(ggml_backend_t backend, enum ggml_type type) {
    if (type == GGML_TYPE_F32) {
        return true;
    }
    size_t                  meta   = ggml_tensor_overhead() * 4 + 1024;
    struct ggml_init_params params = { meta, NULL, true };
    struct ggml_context *   ctx    = ggml_init(params);
    struct ggml_tensor *    src    = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 32);
    struct ggml_tensor *    dst    = ggml_cast(ctx, src, type);
    bool                    ok     = ggml_backend_supports_op(backend, dst);
    ggml_free(ctx);
    return ok;
}

// Build the reverse map from LyCORIS key prefix to GGUF tensor name.
// LyCORIS stores adapter tensors as "lycoris_<path_with_underscores>.<suffix>",
// where the torch module path has all dots flattened to underscores. We cannot
// safely reverse that transform blindly (names like "cross_attn" contain real
// underscores), so we enumerate the GGUF decoder .weight tensors we already
// have and build the mapping from them.
//
// Example:
//   GGUF tensor "decoder.layers.0.self_attn.q_proj.weight"
//   -> "lycoris_layers_0_self_attn_q_proj" -> "decoder.layers.0.self_attn.q_proj.weight"
static std::unordered_map<std::string, std::string> lokr_build_reverse_map(const WeightSource & ws) {
    std::unordered_map<std::string, std::string> out;
    auto names = ws.tensor_names();
    for (const auto & s : names) {
        // only decoder.*.weight tensors qualify as LoKr targets
        static const char * suffix = ".weight";
        size_t slen = strlen(suffix);
        if (s.size() <= slen || s.compare(s.size() - slen, slen, suffix) != 0) continue;
        if (s.compare(0, 8, "decoder.") != 0) continue;
        std::string path = s.substr(8, s.size() - 8 - slen);
        for (char & c : path) { if (c == '.') c = '_'; }
        out["lycoris_" + path] = s;
    }
    return out;
}

// Split a safetensors key on its last dot. Returns false when no dot exists.
// Example: "lycoris_layers_0_self_attn_q_proj.lokr_w1"
//   -> prefix "lycoris_layers_0_self_attn_q_proj", suffix "lokr_w1"
static bool adapter_split_suffix(const std::string & key, std::string * prefix, std::string * suffix) {
    size_t dot = key.rfind('.');
    if (dot == std::string::npos) {
        return false;
    }
    *prefix = key.substr(0, dot);
    *suffix = key.substr(dot + 1);
    return true;
}

// Detect whether the safetensors payload is a LyCORIS LoKr pack.
// Any tensor named "*.lokr_w1" or "*.lokr_w2*" is a LoKr marker that LoRA
// payloads never produce, so a single match is sufficient.
static bool adapter_detect_lokr(const STFile & st) {
    for (const auto & e : st.entries) {
        if (e.name.find(".lokr_w1") != std::string::npos || e.name.find(".lokr_w2") != std::string::npos) {
            return true;
        }
    }
    return false;
}

// Merge one base tensor with an adapter delta entirely inside a backend graph.
//
// Responsibilities split between caller and helper:
//   caller: provide the delta builder lambda that, given the graph ctx,
//           returns a tdelta tensor of shape ggml ne=(in, out) in F32 row
//           major (out_feat, in_feat). The lambda is also responsible for
//           uploading its own adapter factors to the backend inside the
//           helper's alloc_ctx_tensors sweep (see pattern below).
//   helper: upload base in native type from mmap, dequant to F32 on backend,
//           BF16 round delta, add base + delta (plain or DoRA), encode back
//           to native on backend when supported, download into staging.
//
// Upload pattern for caller-owned tensors: the lambda allocates tensors in
// the shared ctx. The helper calls alloc_ctx_tensors once after the lambda
// runs, then the lambda's returned populator function is invoked to do the
// actual ggml_backend_tensor_set calls.
//
// build_delta signature expanded: pair<tdelta, populator> so the caller can
// defer its uploads until after the single alloc sweep. This keeps one graph,
// one buffer alloc, one compute, one download per merge.
struct adapter_delta_build {
    struct ggml_tensor *  tdelta;
    std::function<void()> upload;
};

// Execute the unified merge graph for one tensor. Returns true on success.
// See adapter_delta_build for the caller contract.

// Compute one delta on GPU and store the result as F32.
// Uses a temporary graph: builds the delta subgraph (LoRA B@A or LoKr kron),
// computes on backend, downloads the result.
// Returns true on success. Used by the split merge path (NVFP4/MXFP4) and
// by adapter-runtime.h for precomputing runtime deltas.
static bool adapter_compute_delta(
    const std::function<adapter_delta_build(struct ggml_context *)> & build_delta,
    int64_t          ne0,
    int64_t          ne1,
    ggml_backend_t   backend,
    std::vector<float> & out_f32) {

    int64_t nel = ne0 * ne1;

    // Build a minimal graph: just the delta subgraph
    size_t                  meta   = ggml_tensor_overhead() * 64 + ggml_graph_overhead() + 32 * 1024;
    struct ggml_init_params params = { meta, NULL, true };
    struct ggml_context *   ctx    = ggml_init(params);
    if (!ctx) return false;

    adapter_delta_build db = build_delta(ctx);

    struct ggml_cgraph * graph = ggml_new_graph(ctx);
    ggml_build_forward_expand(graph, db.tdelta);

    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
    if (!buf) {
        ggml_free(ctx);
        return false;
    }

    // Upload adapter factor tensors
    db.upload();

    // Compute delta on GPU
    ggml_backend_graph_compute(backend, graph);

    // Download result
    out_f32.resize((size_t) nel);
    ggml_backend_tensor_get(db.tdelta, out_f32.data(), 0, (size_t) nel * sizeof(float));

    ggml_backend_buffer_free(buf);
    ggml_free(ctx);
    return true;
}

// Execute the unified merge graph for one tensor. Returns true on success.
// See adapter_delta_build for the caller contract.
static bool adapter_merge_on_backend(WeightCtx *                                                       wctx,
                                     std::unordered_map<const void *, size_t> &                        pending_idx,
                                     const void *                                                      base_ptr,
                                     enum ggml_type                                                    ttype,
                                     int64_t                                                           ne0,
                                     int64_t                                                           ne1,
                                     const float *                                                     ds,
                                     float                                                             user_scale,
                                     ggml_backend_t                                                    backend,
                                     const char *                                                      gguf_name,
                                     const std::function<adapter_delta_build(struct ggml_context *)> & build_delta,
                                     bool                                                              promote_f32 = false) {
    // torch.finfo(torch.bfloat16).eps, used verbatim in LyCORIS apply_weight_decompose
    const float eps = 7.8125e-3f;

    // locate the pending copy upfront: no point computing if we can't apply
    auto pc_it = pending_idx.find(base_ptr);
    if (pc_it == pending_idx.end()) {
        fprintf(stderr, "[Adapter] WARNING: no PendingCopy for %s, skipping\n", gguf_name);
        return false;
    }
    WeightCtx::PendingCopy * pc = &wctx->pending[pc_it->second];

    int64_t nel       = ne0 * ne1;
    size_t  base_nb   = ggml_row_size(ttype, ne0) * (size_t) ne1;

    // slack for the largest graph (DoRA + BF16 round + cast in/out + caller subgraph)
    size_t                  meta   = ggml_tensor_overhead() * 64 + ggml_graph_overhead() + 32 * 1024;
    struct ggml_init_params params = { meta, NULL, true };
    struct ggml_context *   ctx    = ggml_init(params);
    if (!ctx) {
        return false;
    }

    // Check if the backend can decode (dequant) from native to F32 via ggml_cast.
    // NVFP4/MXFP4 have dequant in mul_mat kernels but NOT in ggml_cuda_cpy,
    // so ggml_cast(NVFP4 → F32) fails on CUDA. For these types, use a split
    // path: compute delta on GPU, then merge + requant on host.
    bool decode_on_host = false;
    if (ttype != GGML_TYPE_F32 && ttype != GGML_TYPE_BF16 && ttype != GGML_TYPE_F16) {
        size_t                  pmeta   = ggml_tensor_overhead() * 4 + 1024;
        struct ggml_init_params pparams = { pmeta, NULL, true };
        struct ggml_context *   pctx    = ggml_init(pparams);
        struct ggml_tensor *    psrc    = ggml_new_tensor_1d(pctx, ttype, 64);
        struct ggml_tensor *    pdst    = ggml_cast(pctx, psrc, GGML_TYPE_F32);
        decode_on_host = !ggml_backend_supports_op(backend, pdst);
        ggml_free(pctx);
    }

    if (decode_on_host) {
        // ─── Split path for types without GPU cast (NVFP4, MXFP4) ───
        // Step 1: compute delta on GPU only (adapter factors are small, fast upload)
        std::vector<float> delta_f32;
        if (!adapter_compute_delta(build_delta, ne0, ne1, backend, delta_f32)) {
            ggml_free(ctx);
            return false;
        }

        // Step 2: dequant base weight on host
        std::vector<float> base_f32((size_t) nel);
        const struct ggml_type_traits * traits = ggml_get_type_traits(ttype);
        if (!traits->to_float) {
            fprintf(stderr, "[Adapter] WARNING: no host dequant for type %d, skipping %s\n", ttype, gguf_name);
            ggml_free(ctx);
            return false;
        }
        traits->to_float(base_ptr, base_f32.data(), nel);

        // Step 3: merge on host — base + delta (with user_scale and DoRA if present)
        if (!ds) {
            // plain: merged = base + user_scale * delta
            for (size_t i = 0; i < (size_t) nel; i++) {
                base_f32[i] += user_scale * delta_f32[i];
            }
        } else {
            // DoRA: merged = (base + delta) * dora_scale_per_row
            for (size_t i = 0; i < (size_t) nel; i++) {
                base_f32[i] += delta_f32[i];
            }
            // per-row DoRA rescale
            for (int64_t row = 0; row < ne1; row++) {
                float * rp = base_f32.data() + row * ne0;
                float sum_sq = 0.0f;
                for (int64_t j = 0; j < ne0; j++) { sum_sq += rp[j] * rp[j]; }
                float rn = sqrtf(sum_sq) + eps;
                float s = ds[row] / rn;
                if (user_scale != 1.0f) { s = s * user_scale + (1.0f - user_scale); }
                for (int64_t j = 0; j < ne0; j++) { rp[j] *= s; }
            }
        }
        delta_f32.clear();  // free delta memory

        // Step 4: promote to BF16 instead of requanting to native type.
        // NVFP4/MXFP4 quantization on host is extremely slow (~60s for 359 tensors).
        // BF16 conversion is a trivial bitshift and completes instantly.
        // VRAM cost: 2x native (e.g. 2.8 GB NVFP4 → 5.6 GB BF16), still reasonable.
        size_t bf16_nb = (size_t) nel * sizeof(ggml_bf16_t);
        size_t n_floats = (bf16_nb + sizeof(float) - 1) / sizeof(float);
        auto staging_buf = std::make_unique<float[]>(n_floats);
        ggml_fp32_to_bf16_row(base_f32.data(), (ggml_bf16_t *) staging_buf.get(), nel);
        base_f32.clear();  // free F32 buffer

        // Promote tensor type to BF16 and recalculate strides
        pc->tensor->type  = GGML_TYPE_BF16;
        pc->tensor->nb[0] = sizeof(ggml_bf16_t);
        for (int d = 1; d < GGML_MAX_DIMS; d++) {
            pc->tensor->nb[d] = pc->tensor->nb[d - 1] * pc->tensor->ne[d - 1];
        }
        pc->src    = staging_buf.get();
        pc->nbytes = bf16_nb;
        wctx->staging.push_back(std::move(staging_buf));
        ggml_free(ctx);
        return true;
    }

    // ─── Unified GPU graph path (BF16, F16, K-quants — types with GPU cast) ───

    struct ggml_tensor * tbase_native = ggml_new_tensor_2d(ctx, ttype, ne0, ne1);
    struct ggml_tensor * tbase_f32    = ggml_cast(ctx, tbase_native, GGML_TYPE_F32);

    // DoRA scale vector, one F32 per output row when dora_scale is set
    struct ggml_tensor * tds = ds ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, ne1) : NULL;

    // caller builds the delta subgraph and returns its upload closure
    adapter_delta_build db = build_delta(ctx);

    // BF16 round mirrors LyCORIS diff.to(base.dtype) and PEFT merge_and_unload
    // intermediate cast. Cast chain runs entirely on backend.
    // When promote_f32: skip the round-trip — the delta stays in full F32 precision,
    // and the merged result will be stored as F32 to avoid catastrophic cancellation
    // when a small adapter delta is added to a large base weight in limited precision.
    struct ggml_tensor * tdelta_use;
    if (promote_f32) {
        tdelta_use = db.tdelta;  // skip BF16 round-trip, preserve full precision
    } else {
        struct ggml_tensor * tdelta_bf = ggml_cast(ctx, db.tdelta, GGML_TYPE_BF16);
        tdelta_use = ggml_cast(ctx, tdelta_bf, GGML_TYPE_F32);
    }

    // merge: plain scaled add for no DoRA, weight decompose otherwise
    struct ggml_tensor * tmerged;
    if (!tds) {
        struct ggml_tensor * td_u = (user_scale != 1.0f) ? ggml_scale(ctx, tdelta_use, user_scale) : tdelta_use;
        tmerged                   = ggml_add(ctx, tbase_f32, td_u);
    } else {
        // DoRA per output row: scale = user * (ds / sqrt(sum_sq(m_pre)) + eps) + (1 - user)
        struct ggml_tensor * tm_pre  = ggml_add(ctx, tbase_f32, tdelta_use);
        struct ggml_tensor * tsq     = ggml_sqr(ctx, tm_pre);
        struct ggml_tensor * tss     = ggml_sum_rows(ctx, tsq);  // ne=(1, out)
        struct ggml_tensor * trn     = ggml_sqrt(ctx, tss);
        struct ggml_tensor * trn_eps = ggml_scale_bias(ctx, trn, 1.0f, eps);
        struct ggml_tensor * tscale  = ggml_div(ctx, tds, trn_eps);
        struct ggml_tensor * tscale_m =
            (user_scale != 1.0f) ? ggml_scale_bias(ctx, tscale, user_scale, 1.0f - user_scale) : tscale;
        // broadcast (in, out) * (1, out)
        tmerged = ggml_mul(ctx, tm_pre, tscale_m);
    }

    // output: promote_f32 keeps F32 (no encode back to native), otherwise native
    struct ggml_tensor * tout;
    if (promote_f32) {
        tout = tmerged;  // stay F32 — tensor type will be promoted below
    } else {
        bool encode_ok = adapter_backend_can_encode(backend, ttype);
        tout = encode_ok ? ggml_cast(ctx, tmerged, ttype) : tmerged;
    }

    struct ggml_cgraph * graph = ggml_new_graph(ctx);
    ggml_build_forward_expand(graph, tout);

    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
    if (!buf) {
        ggml_free(ctx);
        return false;
    }

    // helper-owned uploads: base in native type from mmap, ds if DoRA
    ggml_backend_tensor_set(tbase_native, base_ptr, 0, base_nb);
    if (tds) {
        ggml_backend_tensor_set(tds, ds, 0, (size_t) ne1 * sizeof(float));
    }
    // caller-owned uploads for adapter factors
    db.upload();

    ggml_backend_graph_compute(backend, graph);

    // allocate a staging slot and download the merged result
    if (promote_f32) {
        // F32 promotion: download merged F32, promote tensor type so wctx_alloc
        // allocates the right amount of backend memory for this tensor.
        size_t f32_nb = (size_t) nel * sizeof(float);
        auto staging_buf = std::make_unique<float[]>((size_t) nel);
        ggml_backend_tensor_get(tout, staging_buf.get(), 0, f32_nb);
        // Promote tensor type AND recalculate strides. ggml_nbytes uses nb[]
        // to compute buffer size, so leaving BF16 strides after type change
        // causes the backend to allocate a BF16-sized buffer for F32 data.
        pc->tensor->type  = GGML_TYPE_F32;
        pc->tensor->nb[0] = sizeof(float);
        for (int d = 1; d < GGML_MAX_DIMS; d++) {
            pc->tensor->nb[d] = pc->tensor->nb[d - 1] * pc->tensor->ne[d - 1];
        }
        pc->src    = staging_buf.get();
        pc->nbytes = f32_nb;
        wctx->staging.push_back(std::move(staging_buf));
    } else {
        bool encode_ok = adapter_backend_can_encode(backend, ttype);
        size_t n_floats    = (base_nb + sizeof(float) - 1) / sizeof(float);
        auto   staging_buf = std::make_unique<float[]>(n_floats);
        void * merged_buf  = staging_buf.get();

        if (encode_ok) {
            // download straight into staging in native type, zero host postprocess
            ggml_backend_tensor_get(tout, merged_buf, 0, base_nb);
            pc->src    = merged_buf;
            pc->nbytes = base_nb;
        } else {
            // download F32 then requant on host (K-quants on CUDA: Q4_K_M, Q5_K_M, Q6_K)
            std::vector<float> merged_f32((size_t) nel);
            ggml_backend_tensor_get(tout, merged_f32.data(), 0, (size_t) nel * sizeof(float));
            size_t merged_bytes = adapter_requant(merged_f32.data(), merged_buf, nel, ne0, ttype);
            if (merged_bytes == 0) {
                ggml_backend_buffer_free(buf);
                ggml_free(ctx);
                return false;
            }
            pc->src    = merged_buf;
            pc->nbytes = merged_bytes;
        }
        wctx->staging.push_back(std::move(staging_buf));
    }

    ggml_backend_buffer_free(buf);
    ggml_free(ctx);
    return true;
}

// In merge_hq mode, decide whether to skip a non-layer tensor.
// Layer weights (decoder.layers.N.*) are always merged.
// Non-layer weights are skipped when their group scale is 0.0 (the default for
// time_embed and cond_embed in merge_hq). Setting a non-zero scale re-enables them.
static bool adapter_hq_should_skip(const std::string & name) {
    // Layer projections — never skip
    if (name.find("decoder.layers.") == 0) return false;

    // Non-layer tensor: skip if its group scale is 0.0
    const std::string group = adapter_determine_group(name);
    float g_scale = adapter_group_scale_for(g_hotstep_params.adapter_group_scales, group);
    return g_scale == 0.0f;
}

// LoRA merge path. Matches PEFT merge_and_unload for PEFT payloads and ComfyUI
// merge semantics for single file LoRA:
//   delta = (alpha / rank) * scale * B @ A
// Applied to base weights in place. Alpha is read per tensor if present
// (ComfyUI baked), else from adapter_config.json, else defaults to rank.
static bool adapter_merge_lora(WeightCtx *            wctx,
                               const WeightSource &   ws,
                               const STFile &      st,
                               const std::string & cfg_dir,
                               float               scale,
                               ggml_backend_t      backend,
                               bool                promote_f32 = false) {
    int alpha_cfg = adapter_read_alpha(cfg_dir.c_str());

    // group lora_A and lora_B entries by their GGUF base tensor name.
    // also collect per tensor alpha scalars (ComfyUI baked format).
    std::map<std::string, const STEntry *> a_map, b_map;
    std::map<std::string, float>           alpha_map;
    for (const auto & e : st.entries) {
        // per tensor alpha: "base_model.model.layers.0.self_attn.q_proj.alpha"
        // scalar F32 with shape [] containing the baked alpha value
        const char * alpha_suffix = ".alpha";
        size_t       slen         = strlen(alpha_suffix);
        if (e.name.size() > slen && e.name.compare(e.name.size() - slen, slen, alpha_suffix) == 0 && e.dtype == "F32" &&
            e.n_dims == 0) {
            std::string fake_key = e.name.substr(0, e.name.size() - slen) + ".lora_.x";
            std::string base     = lora_base_name(fake_key);
            if (!base.empty()) {
                float val = 0.0f;
                memcpy(&val, st_data(st, e), sizeof(float));
                alpha_map[base] = val;
            }
            continue;
        }

        std::string base = lora_base_name(e.name);
        if (base.empty()) {
            continue;
        }
        if (lora_is_a(e.name)) {
            a_map[base] = &e;
        } else if (lora_is_b(e.name)) {
            b_map[base] = &e;
        }
    }

    std::unordered_map<const void *, size_t> pending_idx;
    pending_idx.reserve(wctx->pending.size());
    for (size_t i = 0; i < wctx->pending.size(); i++) {
        pending_idx[wctx->pending[i].src] = i;
    }

    int merged  = 0;
    int skipped = 0;

    for (const auto & kv : a_map) {
        const std::string & gguf_name = kv.first;
        const STEntry *     ea        = kv.second;

        auto it = b_map.find(gguf_name);
        if (it == b_map.end()) {
            fprintf(stderr, "[Adapter] WARNING: no lora_B for %s, skipping\n", gguf_name.c_str());
            skipped++;
            continue;
        }
        const STEntry * eb = it->second;

        if (!ws.exists(gguf_name.c_str())) {
            fprintf(stderr, "[Adapter] WARNING: tensor %s not in base model, skipping\n", gguf_name.c_str());
            skipped++;
            continue;
        }

        // merge_hq: skip non-layer weights to match runtime behavior
        if (promote_f32 && adapter_hq_should_skip(gguf_name)) {
            fprintf(stderr, "[Adapter-HQ] Skipping non-layer weight: %s\n", gguf_name.c_str());
            skipped++;
            continue;
        }
        enum ggml_type ttype = ws.type(gguf_name.c_str());
        int n_dims; int64_t ne_arr[4];
        ws.shape(gguf_name.c_str(), n_dims, ne_arr);
        int64_t ne0 = ne_arr[0];
        int64_t ne1 = ne_arr[1];
        ggml_type data_type;
        const void * base_ptr = ws.data(gguf_name.c_str(), data_type);

        // Conv1d tensors (e.g. proj_in): see comment in LoKr path
        if (n_dims >= 3 && pending_idx.find(base_ptr) == pending_idx.end()) {
            for (size_t pi = 0; pi < wctx->pending.size(); pi++) {
                auto & pc = wctx->pending[pi];
                if (pc.tensor && pc.tensor->name && gguf_name == pc.tensor->name) {
                    ne0      = pc.tensor->ne[0];
                    ne1      = pc.tensor->ne[1];
                    ttype    = pc.tensor->type;
                    base_ptr = pc.src;
                    fprintf(stderr, "[Adapter] Conv1d %s: using pre-permuted shape [%lld, %lld] (%s)\n",
                            gguf_name.c_str(), (long long) ne0, (long long) ne1,
                            ggml_type_name(ttype));
                    break;
                }
            }
        }

        // LoRA shapes (safetensors PyTorch convention, row major):
        //   A: [rank, in_features]
        //   B: [out_features, rank]
        int64_t rank     = ea->shape[0];
        int64_t in_feat  = ea->shape[1];
        int64_t out_feat = eb->shape[0];

        if (eb->shape[1] != rank) {
            fprintf(stderr, "[Adapter] WARNING: rank mismatch A=%lld vs B=%lld for %s\n", (long long) rank,
                    (long long) eb->shape[1], gguf_name.c_str());
            skipped++;
            continue;
        }
        // Conv1d expansion: adapter trained on [out, in_ch] but base is [in_ch*P, out]
        int64_t conv_expand = 1;
        if (in_feat == ne0 && out_feat == ne1) {
            // exact match
        } else if (out_feat == ne1 && ne0 > 0 && (ne0 % in_feat) == 0) {
            conv_expand = ne0 / in_feat;
            fprintf(stderr, "[Adapter] Conv1d %s: tiling LoRA delta [%lld, %lld] x%lld -> [%lld, %lld]\n",
                    gguf_name.c_str(), (long long) in_feat, (long long) out_feat,
                    (long long) conv_expand, (long long) ne0, (long long) ne1);
        } else {
            fprintf(stderr, "[Adapter] WARNING: shape mismatch for %s: LoRA [%lld,%lld] vs GGUF [%lld,%lld]\n",
                    gguf_name.c_str(), (long long) out_feat, (long long) in_feat, (long long) ne1, (long long) ne0);
            skipped++;
            continue;
        }

        // alpha: prefer per tensor (ComfyUI baked), then config, fallback to rank
        float alpha;
        auto  alpha_it = alpha_map.find(gguf_name);
        if (alpha_it != alpha_map.end()) {
            alpha = alpha_it->second;
        } else if (alpha_cfg > 0) {
            alpha = (float) alpha_cfg;
        } else {
            alpha = (float) rank;
        }
        float scaling = (alpha / (float) rank) * scale;

        // HOT-Step: apply per-group adapter scale from sideband
        float g_scale = adapter_group_scale_for(g_hotstep_params.adapter_group_scales,
                                                 adapter_determine_group(gguf_name));
        scaling *= g_scale;

        // load A and B to F32, PEFT rounds them through BF16 before the GEMM
        int64_t            a_nel = rank * in_feat;
        int64_t            b_nel = out_feat * rank;
        std::vector<float> a_f32((size_t) a_nel);
        std::vector<float> b_f32((size_t) b_nel);
        if (!adapter_to_f32(st_data(st, *ea), a_f32.data(), a_nel, ea->dtype)) {
            fprintf(stderr, "[Adapter] WARNING: unsupported dtype %s for lora_A\n", ea->dtype.c_str());
            skipped++;
            continue;
        }
        if (!adapter_to_f32(st_data(st, *eb), b_f32.data(), b_nel, eb->dtype)) {
            fprintf(stderr, "[Adapter] WARNING: unsupported dtype %s for lora_B\n", eb->dtype.c_str());
            skipped++;
            continue;
        }

        // delta = scaling * B @ A, built inside the unified merge graph.
        // A row major (rank, in_feat) stored as ggml ne=(in_feat, rank), transposed
        // so rank sits on ne[0] for mul_mat contraction. Result ne=(in_feat, out_feat).
        auto build = [&](struct ggml_context * ctx) {
            struct ggml_tensor * ta     = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, in_feat, rank);
            struct ggml_tensor * tb     = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rank, out_feat);
            // PEFT BF16 round on A and B before the GEMM, done on backend.
            struct ggml_tensor * ta_br  = ggml_cast(ctx, ggml_cast(ctx, ta, GGML_TYPE_BF16), GGML_TYPE_F32);
            struct ggml_tensor * tb_br  = ggml_cast(ctx, ggml_cast(ctx, tb, GGML_TYPE_BF16), GGML_TYPE_F32);
            struct ggml_tensor * ta_t   = ggml_cont(ctx, ggml_transpose(ctx, ta_br));
            struct ggml_tensor * tdelta = ggml_scale(ctx, ggml_mul_mat(ctx, ta_t, tb_br), scaling);

            // Conv1d expansion: tile delta [in_ch, out] -> [in_ch*P, out]
            if (conv_expand > 1) {
                struct ggml_tensor * ttile = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, in_feat * conv_expand, out_feat);
                tdelta = ggml_repeat(ctx, tdelta, ttile);
            }

            adapter_delta_build db;
            db.tdelta = tdelta;
            db.upload = [=]() {
                ggml_backend_tensor_set(ta, a_f32.data(), 0, (size_t) a_nel * sizeof(float));
                ggml_backend_tensor_set(tb, b_f32.data(), 0, (size_t) b_nel * sizeof(float));
            };
            return db;
        };

        if (!adapter_merge_on_backend(wctx, pending_idx, base_ptr, ttype, ne0, ne1, nullptr, 1.0f, backend,
                                      gguf_name.c_str(), build, promote_f32)) {
            skipped++;
            continue;
        }
        merged++;
    }

    const auto & gs = g_hotstep_params.adapter_group_scales;
    fprintf(stderr, "[Adapter] LoRA merged %d pairs (skipped %d), scale=%.2f\n", merged, skipped, scale);
    fprintf(stderr, "[Adapter] Group scales: self_attn=%.2f, cross_attn=%.2f, mlp=%.2f, cond_embed=%.2f, time_embed=%.2f, proj_in=%.2f\n",
            gs.self_attn, gs.cross_attn, gs.mlp, gs.cond_embed, gs.time_embed, gs.proj_in);
    return merged > 0;
}

// LoKr merge path. Matches the LyCORIS runtime forward for LoKr at
// lycoris/modules/lokr.py:551..566 (non bypass mode, scalar=1):
//   delta       = (alpha / rank) * kron(w1, W2)
//   no DoRA     : merged = base + delta * multiplier
//   DoRA present: merged = apply_weight_decompose(base + delta, multiplier)
//
// Two W2 variants are handled, selected per module by the suffixes present:
//   factorized  : W2 = w2_a @ w2_b, rank = w2_a.shape[1] = w2_b.shape[0]
//   monolithic  : W2 = w2 directly, rank = __metadata__.lokr_config.linear_dim
//
// LyCORIS picks monolithic when the factorized rank would not shrink the
// tensor (use_w2 path). qinglong SFT is fully factorized, garage-band is
// fully monolithic. w1 is always whole, not factorized (no w1_a). Tucker
// decomposition and convolutional LoKr are not implemented.
//
// Kron layout on GGML (same for both variants):
//   w1     row major (a, b)     -> ggml ne=(b, a)
//   W2     row major (c, d)     -> ggml ne=(d, c)
//   delta  row major (a*c, b*d) -> ggml ne=(b*d, a*c)
//
// Graph, factorized:
//   tw2 = mul_mat(cont(transpose(tw2b)), tw2a) = ne=(d, c)
// Graph, monolithic:
//   tw2 uploaded directly, ne=(d, c)
// Shared downstream:
//   tw1_s = scale(tw1, alpha / rank), scaling on the tiny side.
//   touter = mul_mat(reshape(tw1_s, 1, a*b), reshape(tw2, 1, c*d)) = ne=(a*b, c*d)
//   kron_p = permute(reshape_4d(touter, b, a, d, c), 1, 3, 0, 2) = ne=(d, b, c, a)
//   delta  = reshape_2d(cont(kron_p), b*d, a*c)
//
// ggml_permute axis_i positions src axis i AT new ne index axis_i
// (ggml.c:3781 sets ne[axis_i] = src.ne[i]), so mapping src (b, a, d, c)
// -> new (d, b, c, a) needs src axes (0, 1, 2, 3) to land at new positions
// (1, 3, 0, 2). The fast pair (d, b) then collapses into in_feat and the
// slow pair (c, a) into out_feat under reshape_2d. Net effect:
//   delta_rm[aa*c + cc, bb*d + dd] = W1[aa, bb] * W2[cc, dd]
static bool adapter_merge_lokr(WeightCtx *          wctx,
                               const WeightSource & ws,
                               const STFile &    st,
                               float             user_scale,
                               ggml_backend_t    backend,
                               bool              promote_f32 = false) {
    // group the per module tensors by LyCORIS prefix. Each module has either
    // w2 alone (monolithic) or w2_a + w2_b (factorized), never both.
    struct LoKrEntry {
        const STEntry * w1         = nullptr;
        const STEntry * w2         = nullptr;
        const STEntry * w2_a       = nullptr;
        const STEntry * w2_b       = nullptr;
        const STEntry * alpha      = nullptr;
        const STEntry * dora_scale = nullptr;
    };

    std::map<std::string, LoKrEntry> modules;

    for (const auto & e : st.entries) {
        std::string prefix, suffix;
        if (!adapter_split_suffix(e.name, &prefix, &suffix)) {
            continue;
        }
        if (prefix.compare(0, 8, "lycoris_") != 0) {
            continue;
        }
        LoKrEntry & m = modules[prefix];
        if (suffix == "lokr_w1") {
            m.w1 = &e;
        } else if (suffix == "lokr_w2") {
            m.w2 = &e;
        } else if (suffix == "lokr_w2_a") {
            m.w2_a = &e;
        } else if (suffix == "lokr_w2_b") {
            m.w2_b = &e;
        } else if (suffix == "alpha") {
            m.alpha = &e;
        } else if (suffix == "dora_scale") {
            m.dora_scale = &e;
        }
    }

    std::unordered_map<std::string, std::string> name_map = lokr_build_reverse_map(ws);

    std::unordered_map<const void *, size_t> pending_idx;
    pending_idx.reserve(wctx->pending.size());
    for (size_t i = 0; i < wctx->pending.size(); i++) {
        pending_idx[wctx->pending[i].src] = i;
    }

    // linear_dim from __metadata__.lokr_config, only needed by monolithic modules.
    // Factorized modules derive rank from w2_a / w2_b shapes and ignore this value.
    int lokr_dim = adapter_read_lokr_dim(st);

    int merged     = 0;
    int skipped    = 0;
    int dora_count = 0;
    int mono_count = 0;

    for (const auto & kv : modules) {
        const std::string & lyc_prefix = kv.first;
        const LoKrEntry &   m          = kv.second;

        // per module variant: factorized (w2_a + w2_b) XOR monolithic (w2)
        bool has_factor = (m.w2_a && m.w2_b);
        bool has_mono   = (m.w2 != nullptr);
        if (!m.w1 || !m.alpha || has_factor == has_mono) {
            fprintf(stderr, "[Adapter] WARNING: incomplete or ambiguous LoKr module %s, skipping\n",
                    lyc_prefix.c_str());
            skipped++;
            continue;
        }

        auto nm_it = name_map.find(lyc_prefix);
        if (nm_it == name_map.end()) {
            fprintf(stderr, "[Adapter] WARNING: no GGUF tensor for %s, skipping\n", lyc_prefix.c_str());
            skipped++;
            continue;
        }
        const std::string & gguf_name = nm_it->second;

        if (!ws.exists(gguf_name.c_str())) {
            fprintf(stderr, "[Adapter] WARNING: tensor %s not in base model, skipping\n", gguf_name.c_str());
            skipped++;
            continue;
        }

        // merge_hq: skip non-layer weights to match runtime behavior
        if (promote_f32 && adapter_hq_should_skip(gguf_name)) {
            fprintf(stderr, "[Adapter-HQ] Skipping non-layer weight: %s\n", gguf_name.c_str());
            skipped++;
            continue;
        }
        enum ggml_type ttype = ws.type(gguf_name.c_str());
        int n_dims; int64_t ne_arr[4];
        ws.shape(gguf_name.c_str(), n_dims, ne_arr);
        int64_t ne0 = ne_arr[0];
        int64_t ne1 = ne_arr[1];
        ggml_type data_type;
        const void * base_ptr = ws.data(gguf_name.c_str(), data_type);

        // Conv1d tensors (e.g. proj_in) are stored as 3D in GGUF [P, in_ch, H]
        // but dit_load_proj_in_w() pre-permutes them to 2D F32 [in_ch*P, H].
        // The raw GGUF pointer won't match pending_idx (which indexes the permuted
        // staging buffer), so we fall back to the PendingCopy's tensor shape.
        bool is_prepermuted = false;
        if (n_dims >= 3 && pending_idx.find(base_ptr) == pending_idx.end()) {
            // Search for the matching PendingCopy by tensor name
            for (size_t pi = 0; pi < wctx->pending.size(); pi++) {
                auto & pc = wctx->pending[pi];
                if (pc.tensor && pc.tensor->name && gguf_name == pc.tensor->name) {
                    // Use the permuted tensor's 2D shape and staging data pointer
                    ne0      = pc.tensor->ne[0];
                    ne1      = pc.tensor->ne[1];
                    ttype    = pc.tensor->type;
                    base_ptr = pc.src;
                    is_prepermuted = true;
                    fprintf(stderr, "[Adapter] Conv1d %s: using pre-permuted shape [%lld, %lld] (%s)\n",
                            gguf_name.c_str(), (long long) ne0, (long long) ne1,
                            ggml_type_name(ttype));
                    break;
                }
            }
        }

        // LoKr shapes (safetensors row major):
        //   w1 : (a, b)
        //   factorized : w2_a (c, rank), w2_b (rank, d)
        //   monolithic : w2 (c, d), rank read from lokr_config.linear_dim
        // Kronecker product yields (a*c, b*d) = (out_feat, in_feat) = (ne1, ne0).
        int64_t a = m.w1->shape[0];
        int64_t b = m.w1->shape[1];
        int64_t c;
        int64_t d;
        int64_t r;

        if (has_factor) {
            c          = m.w2_a->shape[0];
            r          = m.w2_a->shape[1];
            d          = m.w2_b->shape[1];
            int64_t r2 = m.w2_b->shape[0];
            if (r != r2) {
                fprintf(stderr, "[Adapter] WARNING: LoKr rank mismatch w2_a=%lld vs w2_b=%lld for %s\n", (long long) r,
                        (long long) r2, lyc_prefix.c_str());
                skipped++;
                continue;
            }
        } else {
            c = m.w2->shape[0];
            d = m.w2->shape[1];
            if (lokr_dim <= 0) {
                fprintf(stderr,
                        "[Adapter] WARNING: monolithic LoKr %s needs __metadata__.lokr_config.linear_dim, skipping\n",
                        lyc_prefix.c_str());
                skipped++;
                continue;
            }
            r = lokr_dim;
        }

        // Conv1d expansion: adapter trained on [out, in_ch] but base tensor is
        // [in_ch*P, out] (pre-permuted). Detect and tile delta across patch dim.
        int64_t conv_expand = 1;
        if (a * c == ne1 && b * d == ne0) {
            // exact match — normal 2D weight
        } else if (a * c == ne1 && ne0 > 0 && (ne0 % (b * d)) == 0) {
            conv_expand = ne0 / (b * d);
            fprintf(stderr, "[Adapter] Conv1d %s: tiling delta [%lld, %lld] x%lld -> [%lld, %lld]\n",
                    gguf_name.c_str(), (long long)(b * d), (long long)(a * c),
                    (long long) conv_expand, (long long) ne0, (long long) ne1);
        } else {
            fprintf(stderr,
                    "[Adapter] WARNING: LoKr shape mismatch for %s: kron(%lldx%lld, %lldx%lld) = %lldx%lld vs GGUF "
                    "out=%lld in=%lld\n",
                    gguf_name.c_str(), (long long) a, (long long) b, (long long) c, (long long) d, (long long) (a * c),
                    (long long) (b * d), (long long) ne1, (long long) ne0);
            skipped++;
            continue;
        }

        // alpha scalar, shape [] varies across trainers in dtype: F32, BF16, F16
        float alpha = 0.0f;
        if (!adapter_to_f32(st_data(st, *m.alpha), &alpha, 1, m.alpha->dtype)) {
            fprintf(stderr, "[Adapter] WARNING: unsupported alpha dtype %s for %s, skipping\n", m.alpha->dtype.c_str(),
                    lyc_prefix.c_str());
            skipped++;
            continue;
        }

        // load w1 always to F32
        int64_t            w1_nel = a * b;
        std::vector<float> w1_f32((size_t) w1_nel);
        if (!adapter_to_f32(st_data(st, *m.w1), w1_f32.data(), w1_nel, m.w1->dtype)) {
            fprintf(stderr, "[Adapter] WARNING: unsupported dtype %s for lokr_w1 %s\n", m.w1->dtype.c_str(),
                    lyc_prefix.c_str());
            skipped++;
            continue;
        }

        // load w2 payload: two buffers for factorized, one for monolithic
        int64_t            w2_nel  = 0;
        int64_t            w2a_nel = 0;
        int64_t            w2b_nel = 0;
        std::vector<float> w2_f32;
        std::vector<float> w2a_f32;
        std::vector<float> w2b_f32;
        if (has_factor) {
            w2a_nel = c * r;
            w2b_nel = r * d;
            w2a_f32.resize((size_t) w2a_nel);
            w2b_f32.resize((size_t) w2b_nel);
            if (!adapter_to_f32(st_data(st, *m.w2_a), w2a_f32.data(), w2a_nel, m.w2_a->dtype) ||
                !adapter_to_f32(st_data(st, *m.w2_b), w2b_f32.data(), w2b_nel, m.w2_b->dtype)) {
                fprintf(stderr, "[Adapter] WARNING: unsupported dtype in LoKr module %s, skipping\n",
                        lyc_prefix.c_str());
                skipped++;
                continue;
            }
        } else {
            w2_nel = c * d;
            w2_f32.resize((size_t) w2_nel);
            if (!adapter_to_f32(st_data(st, *m.w2), w2_f32.data(), w2_nel, m.w2->dtype)) {
                fprintf(stderr, "[Adapter] WARNING: unsupported dtype %s for lokr_w2 %s\n", m.w2->dtype.c_str(),
                        lyc_prefix.c_str());
                skipped++;
                continue;
            }
        }

        // optional DoRA scale vector, one F32 per output row
        const float *      ds_ptr = nullptr;
        std::vector<float> ds_f32;
        if (m.dora_scale) {
            int64_t ds_out = m.dora_scale->shape[0];
            if (ds_out != ne1) {
                fprintf(stderr, "[Adapter] WARNING: dora_scale dim0 %lld != out_feat %lld for %s\n", (long long) ds_out,
                        (long long) ne1, gguf_name.c_str());
                skipped++;
                continue;
            }
            ds_f32.resize((size_t) ds_out);
            if (!adapter_to_f32(st_data(st, *m.dora_scale), ds_f32.data(), ds_out, m.dora_scale->dtype)) {
                fprintf(stderr, "[Adapter] WARNING: unsupported dora_scale dtype %s for %s\n",
                        m.dora_scale->dtype.c_str(), gguf_name.c_str());
                skipped++;
                continue;
            }
            ds_ptr = ds_f32.data();
        }

        float scaling = alpha / (float) r;

        // HOT-Step: apply per-group adapter scale from sideband
        float g_scale = adapter_group_scale_for(g_hotstep_params.adapter_group_scales,
                                                 adapter_determine_group(gguf_name));
        float effective_user_scale = user_scale * g_scale;

        auto build = [&](struct ggml_context * ctx) {
            struct ggml_tensor * tw1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, b, a);

            // tw2 ne=(d, c): w2_a @ w2_b for factorized, direct upload for monolithic
            struct ggml_tensor * tw2;
            struct ggml_tensor * tw2_src = nullptr;
            struct ggml_tensor * tw2a    = nullptr;
            struct ggml_tensor * tw2b    = nullptr;
            if (has_factor) {
                tw2a                        = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, r, c);
                tw2b                        = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d, r);
                // transpose w2_b so rank sits on ne[0] contraction dim
                struct ggml_tensor * tw2b_T = ggml_cont(ctx, ggml_transpose(ctx, tw2b));
                tw2                         = ggml_mul_mat(ctx, tw2b_T, tw2a);
            } else {
                tw2_src = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d, c);
                tw2     = tw2_src;
            }

            // scale alpha / rank on the tiny (typically 4x4) side, cheapest placement
            struct ggml_tensor * tw1_s = ggml_scale(ctx, tw1, scaling);

            // kron via outer product + axis swap
            struct ggml_tensor * tw1_flat  = ggml_reshape_2d(ctx, tw1_s, 1, a * b);
            struct ggml_tensor * tw2_flat  = ggml_reshape_2d(ctx, tw2, 1, c * d);
            struct ggml_tensor * touter    = ggml_mul_mat(ctx, tw1_flat, tw2_flat);
            struct ggml_tensor * touter_4d = ggml_reshape_4d(ctx, touter, b, a, d, c);
            struct ggml_tensor * tkron_p   = ggml_permute(ctx, touter_4d, 1, 3, 0, 2);
            struct ggml_tensor * tkron_c   = ggml_cont(ctx, tkron_p);
            struct ggml_tensor * tdelta    = ggml_reshape_2d(ctx, tkron_c, b * d, a * c);

            // Conv1d expansion: tile delta [in_ch, out] -> [in_ch*P, out]
            // Each patch position receives the same adapter delta.
            if (conv_expand > 1) {
                struct ggml_tensor * ttile = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, b * d * conv_expand, a * c);
                tdelta = ggml_repeat(ctx, tdelta, ttile);
            }

            adapter_delta_build db;
            db.tdelta = tdelta;
            db.upload = [=]() {
                ggml_backend_tensor_set(tw1, w1_f32.data(), 0, (size_t) w1_nel * sizeof(float));
                if (has_factor) {
                    ggml_backend_tensor_set(tw2a, w2a_f32.data(), 0, (size_t) w2a_nel * sizeof(float));
                    ggml_backend_tensor_set(tw2b, w2b_f32.data(), 0, (size_t) w2b_nel * sizeof(float));
                } else {
                    ggml_backend_tensor_set(tw2_src, w2_f32.data(), 0, (size_t) w2_nel * sizeof(float));
                }
            };
            return db;
        };

        if (!adapter_merge_on_backend(wctx, pending_idx, base_ptr, ttype, ne0, ne1, ds_ptr, effective_user_scale, backend,
                                      gguf_name.c_str(), build, promote_f32)) {
            skipped++;
            continue;
        }
        if (ds_ptr) {
            dora_count++;
        }
        if (!has_factor) {
            mono_count++;
        }
        merged++;
    }

    const auto & gs2 = g_hotstep_params.adapter_group_scales;
    fprintf(stderr,
            "[Adapter] LoKr merged %d modules (%d factorized, %d monolithic, %d with DoRA, skipped %d), scale=%.2f\n",
            merged, merged - mono_count, mono_count, dora_count, skipped, user_scale);
    fprintf(stderr, "[Adapter] Group scales: self_attn=%.2f, cross_attn=%.2f, mlp=%.2f, cond_embed=%.2f, time_embed=%.2f, proj_in=%.2f\n",
            gs2.self_attn, gs2.cross_attn, gs2.mlp, gs2.cond_embed, gs2.time_embed, gs2.proj_in);
    return merged > 0;
}

// Main adapter merge entry point.
//
// Call after all tensors are loaded into wctx->pending but before wctx_alloc.
// Detects the adapter algorithm from the safetensors payload and dispatches to
// the matching merge path. The adapter_path points to either:
//   PEFT directory  : a folder with adapter_model.safetensors + adapter_config.json
//   LyCORIS file    : a flat .safetensors file (LoRA ComfyUI or LoKr)
// Directories exist only for PEFT. LyCORIS ships as a single file for both LoRA
// and LoKr payloads.
//
// promote_f32: when true, merged weights stay F32 on the backend (avoids
// host requantization precision loss for K-quants). When false, the merge
// graph encodes back to the native type on GPU when possible, falling back
// to host requantization via ggml_quantize_chunk otherwise.
// Default AUTO: NVFP4/MXFP4 → false (saves ~13 GB VRAM, host requant is
// well-supported), everything else → true.
static bool adapter_merge(WeightCtx *          wctx,
                          const WeightSource & ws,
                          const char *         adapter_path,
                          float                scale,
                          ggml_backend_t       backend,
                          bool                 promote_f32 = true) {
    std::string sf_path;
    std::string cfg_dir;

    struct stat sb;
    if (stat(adapter_path, &sb) != 0) {
        fprintf(stderr, "[Adapter] path does not exist: %s\n", adapter_path);
        return false;
    }

    if (S_ISDIR(sb.st_mode)) {
        // PEFT directory: adapter_model.safetensors is mandatory
        sf_path = std::string(adapter_path) + "/adapter_model.safetensors";
        cfg_dir = adapter_path;
        if (stat(sf_path.c_str(), &sb) != 0) {
            fprintf(stderr, "[Adapter] directory %s is not a PEFT layout, missing adapter_model.safetensors\n",
                    adapter_path);
            return false;
        }
        // warn if adapter_config.json is missing, alpha lives there for PEFT so
        // the merge silently falls back to alpha=rank (scaling=1) otherwise
        std::string cfg_path = cfg_dir + "/adapter_config.json";
        if (stat(cfg_path.c_str(), &sb) != 0) {
            fprintf(stderr,
                    "[Adapter] WARNING: PEFT directory %s missing adapter_config.json, alpha falls back to rank "
                    "(scaling=1.0). If training used lora_alpha != rank, the merge will be under or over scaled.\n",
                    adapter_path);
        }
    } else {
        // LyCORIS flat file, LoRA or LoKr
        sf_path    = adapter_path;
        size_t sep = sf_path.find_last_of("/\\");
        cfg_dir    = (sep != std::string::npos) ? sf_path.substr(0, sep) : ".";
    }

    STFile st = {};
    if (!st_open(&st, sf_path.c_str())) {
        return false;
    }

    bool ok;
    if (adapter_detect_lokr(st)) {
        ok = adapter_merge_lokr(wctx, ws, st, scale, backend, promote_f32);
    } else {
        ok = adapter_merge_lora(wctx, ws, st, cfg_dir, scale, backend, promote_f32);
    }

    st_close(&st);
    return ok;
}