Function load_dit_attention_weights

src/models/ace_step/dit_weights_runtime.cpp:211–226

}

AceStepDiTAttentionWeights load_dit_attention_weights(
    core::BackendWeightStore & store,
    const assets::TensorSource & source,
    const std::string & prefix,
    assets::TensorStorageType storage_type,
    const AceStepDiffusionConfig & config) {
    const int64_t dim = ace_step_diffusion_attention_head_dim(config, "ACE-Step DiT");
    return {
        store.load_tensor(source, prefix + ".q_proj.weight", storage_type, {config.num_attention_heads * dim, config.hidden_size}),
        store.load_tensor(source, prefix + ".k_proj.weight", storage_type, {config.num_key_value_heads * dim, config.hidden_size}),
        store.load_tensor(source, prefix + ".v_proj.weight", storage_type, {config.num_key_value_heads * dim, config.hidden_size}),
        store.load_tensor(source, prefix + ".o_proj.weight", storage_type, {config.hidden_size, config.num_attention_heads * dim}),
        store.load_f32_tensor(source, prefix + ".q_norm.weight", {dim}),
        store.load_f32_tensor(source, prefix + ".k_norm.weight", {dim}),
    };
}

std::shared_ptr<const AceStepDiffusionWeights> load_diffusion_weights(
    const std::shared_ptr<core::BackendWeightStore> & store,