Function build_attention

src/models/ace_step/diffusion.cpp:166–226

    context = core::reshape_tensor(
        ctx,
        context,
        core::TensorShape::from_dims({hidden_states.shape.dims[0], hidden_states.shape.dims[1], config.hidden_size}));
    return modules::LinearModule({config.hidden_size, config.hidden_size, false, GGML_PREC_F32})
        .build(ctx, context, {weights.out_weight, std::nullopt});
}

core::TensorValue apply_modulated_rms_norm(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & norm_weight,
    const core::TensorValue & one,
    const core::TensorValue & shift,
    const core::TensorValue & scale,
    float rms_eps) {
    auto hidden = modules::RMSNormModule({input.shape.dims[input.shape.rank - 1], rms_eps, true, false})
                      .build(ctx, input, {norm_weight, std::nullopt});
    auto scaled = core::wrap_tensor(ggml_add(ctx.ggml, scale.tensor, one.tensor), scale.shape, GGML_TYPE_F32);
    hidden = modules::MulModule{}.build(ctx, hidden, scaled);
    return modules::AddModule{}.build(ctx, hidden, shift);
}

core::TensorValue expand_conditioning(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & conditioning,
    int64_t seq_len) {
    auto reshaped = core::reshape_tensor(
        ctx,
        conditioning,
        core::TensorShape::from_dims({conditioning.shape.dims[0], 1, conditioning.shape.dims[1]}));
    return modules::RepeatModule({core::TensorShape::from_dims({conditioning.shape.dims[0], seq_len, conditioning.shape.dims[1]})})
        .build(ctx, reshaped);
}

struct TimeEmbeddingOutputs {
    core::TensorValue temb;
    core::TensorValue timestep_proj;
};

core::TensorValue build_time_frequency_tensor(
    core::ModuleBuildContext & ctx,
    ggml_tensor * freqs_tensor,
    const core::TensorValue & timestep) {
    const int64_t batch_size = timestep.shape.dims[0];
    auto freqs = core::wrap_tensor(freqs_tensor, core::TensorShape::from_dims({128}), GGML_TYPE_F32);
    auto expanded_freqs = core::reshape_tensor(ctx, freqs, core::TensorShape::from_dims({1, 128}));
    expanded_freqs = modules::RepeatModule({core::TensorShape::from_dims({batch_size, 128})}).build(ctx, expanded_freqs);
    auto expanded_time = modules::RepeatModule({core::TensorShape::from_dims({batch_size, 128})}).build(ctx, timestep);
    return modules::MulModule{}.build(ctx, expanded_time, expanded_freqs);
}

TimeEmbeddingOutputs build_time_embedding(
    core::ModuleBuildContext & ctx,
    ggml_tensor * freqs_tensor,
    const core::TensorValue & timestep,
    const AceStepTimeEmbeddingWeights & weights,
    int64_t hidden_size) {
    auto scaled_t = core::wrap_tensor(ggml_scale(ctx.ggml, timestep.tensor, 1000.0F), timestep.shape, GGML_TYPE_F32);
    auto args = build_time_frequency_tensor(ctx, freqs_tensor, scaled_t);