Function build_snake1d_exact_bct

src/models/ace_step/vae_decoder.cpp:196–221

}

core::TensorValue build_snake1d_exact_bct(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const SnakeExactWeights & weights,
    int64_t channels) {
    const auto input_f32 = ensure_f32(ctx, ensure_contiguous_nontransposed(ctx, input));
    auto alpha = ensure_f32(ctx, weights.alpha);
    auto beta_inv = ensure_f32(ctx, weights.beta_inv);
    if (alpha.shape.rank != 3 || alpha.shape.dims[0] != 1 || alpha.shape.dims[1] != channels || alpha.shape.dims[2] != 1 ||
        beta_inv.shape.rank != 3 || beta_inv.shape.dims[0] != 1 || beta_inv.shape.dims[1] != channels || beta_inv.shape.dims[2] != 1) {
        throw std::runtime_error("ACE-Step VAE Snake1d weight shape mismatch");
    }
    if (input_f32.shape.rank != 3) {
        throw std::runtime_error("ACE-Step VAE Snake1d expects rank-3 BCT input");
    }
    auto periodic = core::wrap_tensor(
        ggml_sqr(ctx.ggml, ggml_sin(ctx.ggml, ggml_mul(ctx.ggml, input_f32.tensor, alpha.tensor))),
        input_f32.shape,
        GGML_TYPE_F32);
    const auto beta_inv_broadcast = core::wrap_tensor(
        ggml_repeat(ctx.ggml, beta_inv.tensor, input_f32.tensor),
        input_f32.shape,
        GGML_TYPE_F32);
    auto frac = modules::MulModule{}.build(ctx, periodic, beta_inv_broadcast);
    return modules::AddModule{}.build(ctx, input_f32, frac);
}

SnakeExactWeights load_snake_exact(
    core::BackendWeightStore & store,