Function load_layer

src/models/stable_audio/rf_dit.cpp:467–498

}

StableAudioRfLayerWeights load_layer(
    core::BackendWeightStore & store,
    const assets::TensorSource & source,
    const std::string & prefix,
    const StableAudioConfig & config,
    assets::TensorStorageType storage_type) {
    StableAudioRfLayerWeights weights;
    weights.pre_norm_gamma = store.load_f32_tensor(source, prefix + ".pre_norm.gamma", {config.embed_dim});
    weights.cross_attend_norm_gamma = store.load_f32_tensor(source, prefix + ".cross_attend_norm.gamma", {config.embed_dim});
    weights.ff_norm_gamma = store.load_f32_tensor(source, prefix + ".ff_norm.gamma", {config.embed_dim});
    weights.self_attn = load_self_attention(store, source, prefix + ".self_attn", config, storage_type);
    weights.cross_attn = load_cross_attention(store, source, prefix + ".cross_attn", config, storage_type);
    weights.ff = load_feed_forward(store, source, prefix + ".ff", config, storage_type);
    weights.local_embed_in = load_linear(
        store,
        source,
        prefix + ".to_local_embed.0",
        storage_type,
        config.embed_dim,
        config.local_add_cond_dim,
        true);
    weights.local_embed_out = load_linear(
        store,
        source,
        prefix + ".to_local_embed.2",
        storage_type,
        config.embed_dim,
        config.embed_dim,
        true);
    weights.scale_shift_gate = store.load_f32_tensor(source, prefix + ".to_scale_shift_gate", {6 * config.embed_dim});
    return weights;
}

}  // namespace