Function load_diffusion_weights

src/models/ace_step/dit_weights_runtime.cpp:228–289

}

std::shared_ptr<const AceStepDiffusionWeights> load_diffusion_weights(
    const std::shared_ptr<core::BackendWeightStore> & store,
    const AceStepAssets & assets,
    assets::TensorStorageType storage_type) {
    const auto & config = assets.config.diffusion;
    const auto & source = *assets.dit_weights;
    auto weights = std::make_shared<AceStepDiffusionWeights>();
    weights->store = store;
    weights->one = store->make_f32(core::TensorShape::from_dims({1}), std::vector<float>{1.0F});
    weights->proj_in = {
        store->load_tensor(source, "decoder.proj_in.1.weight", storage_type, {config.hidden_size, config.in_channels, config.patch_size}),
        store->load_tensor(source, "decoder.proj_in.1.bias", assets::TensorStorageType::F32, {config.hidden_size}),
    };
    weights->time_embed = load_time_embedding_weights(*store, source, "decoder.time_embed", storage_type, config.hidden_size);
    weights->time_embed_r = load_time_embedding_weights(*store, source, "decoder.time_embed_r", storage_type, config.hidden_size);
    weights->condition_embedder = {
        store->load_tensor(source, "decoder.condition_embedder.weight", storage_type, {config.hidden_size, config.hidden_size}),
        store->load_tensor(source, "decoder.condition_embedder.bias", assets::TensorStorageType::F32, {config.hidden_size}),
    };
    if (!config.is_turbo) {
        weights->null_condition_emb_host = source.require_f32("null_condition_emb", {1, 1, config.hidden_size});
    }
    weights->layers.reserve(static_cast<size_t>(config.num_hidden_layers));
    for (int64_t i = 0; i < config.num_hidden_layers; ++i) {
        const std::string prefix = "decoder.layers." + std::to_string(i);
        AceStepDiTLayerWeights layer;
        layer.self_attn_norm = store->load_f32_tensor(source, prefix + ".self_attn_norm.weight", {config.hidden_size});
        layer.self_attn = load_dit_attention_weights(*store, source, prefix + ".self_attn", storage_type, config);
        layer.cross_attn_norm = store->load_f32_tensor(source, prefix + ".cross_attn_norm.weight", {config.hidden_size});
        layer.cross_attn = load_dit_attention_weights(*store, source, prefix + ".cross_attn", storage_type, config);
        layer.mlp_norm = store->load_f32_tensor(source, prefix + ".mlp_norm.weight", {config.hidden_size});
        layer.mlp_gate = {
            store->load_tensor(source, prefix + ".mlp.gate_proj.weight", storage_type, {config.intermediate_size, config.hidden_size}),
            std::nullopt,
        };
        layer.mlp_up = {
            store->load_tensor(source, prefix + ".mlp.up_proj.weight", storage_type, {config.intermediate_size, config.hidden_size}),
            std::nullopt,
        };
        layer.mlp_down = {
            store->load_tensor(source, prefix + ".mlp.down_proj.weight", storage_type, {config.hidden_size, config.intermediate_size}),
            std::nullopt,
        };
        layer.scale_shift_table = store->load_tensor(
            source,
            prefix + ".scale_shift_table",
            storage_type,
            {1, 6, config.hidden_size});
        weights->layers.push_back(std::move(layer));
    }
    weights->norm_out = store->load_f32_tensor(source, "decoder.norm_out.weight", {config.hidden_size});
    weights->proj_out = {
        store->load_tensor(source, "decoder.proj_out.1.weight", storage_type, {config.hidden_size, config.latent_channels, config.patch_size}),
        store->load_tensor(source, "decoder.proj_out.1.bias", assets::TensorStorageType::F32, {config.latent_channels}),
    };
    weights->final_scale_shift_table = store->load_tensor(
        source,
        "decoder.scale_shift_table",