Function dit_layer

src/models/ace_step/diffusion.cpp:328–396

        encoder_hidden_states);
    x = modules::AddModule{}.build(ctx, x, cross_out);

    auto mlp_norm = apply_modulated_rms_norm(
        ctx,
        x,
        weights.mlp_norm,
        one,
        c_shift_msa,
        c_scale_msa,
        config.rms_norm_eps);
    auto ff = build_mlp(ctx, mlp_norm, weights.mlp_gate, weights.mlp_up, weights.mlp_down, config);
    ff = modules::MulModule{}.build(ctx, ff, c_gate_msa);
    return modules::AddModule{}.build(ctx, x, ff);
}

std::vector<float> build_cross_attention_mask_values(
    int64_t seq_len,
    const std::vector<int32_t> & encoder_attention_mask,
    int64_t valid_seq_len) {
    const int64_t encoder_tokens = static_cast<int64_t>(encoder_attention_mask.size());
    if (seq_len <= 0 || encoder_tokens <= 0 || valid_seq_len <= 0 || valid_seq_len > seq_len) {
        throw std::runtime_error("ACE-Step diffusion cross attention mask shape is invalid");
    }
    std::vector<float> values(static_cast<size_t>(seq_len * encoder_tokens), 0.0F);
    const float masked = std::numeric_limits<float>::lowest();
    for (int64_t q = 0; q < seq_len; ++q) {
        for (int64_t k = 0; k < encoder_tokens; ++k) {
            if (q >= valid_seq_len) {
                values[static_cast<size_t>(q * encoder_tokens + k)] = masked;
            }
        }
    }
    return values;
}

std::vector<float> build_sliding_mask_values(int64_t tokens, int64_t sliding_window) {
    std::vector<float> values(static_cast<size_t>(tokens * tokens), 0.0F);
    const float masked = std::numeric_limits<float>::lowest();
    for (int64_t q = 0; q < tokens; ++q) {
        for (int64_t k = 0; k < tokens; ++k) {
            if (std::llabs(q - k) > sliding_window) {
                values[static_cast<size_t>(q * tokens + k)] = masked;
            }
        }
    }
    return values;
}

std::vector<float> build_self_attention_padding_mask_values(int64_t seq_len, int64_t valid_seq_len) {
    if (seq_len <= 0 || valid_seq_len <= 0 || valid_seq_len > seq_len) {
        throw std::runtime_error("ACE-Step diffusion self attention padding mask is invalid");
    }
    std::vector<float> values(static_cast<size_t>(seq_len * seq_len), 0.0F);
    const float masked = std::numeric_limits<float>::lowest();
    for (int64_t q = 0; q < valid_seq_len; ++q) {
        for (int64_t k = valid_seq_len; k < seq_len; ++k) {
            values[static_cast<size_t>(q * seq_len + k)] = masked;
        }
    }