Function build_rf_layer

src/models/stable_audio/rf_dit.cpp:343–387

}

core::TensorValue build_rf_layer(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & cross_context,
    const core::TensorValue & global_six,
    const core::TensorValue & local_add,
    const core::TensorValue & positions,
    const core::TensorValue & padding_mask,
    const StableAudioRfLayerWeights & weights,
    const StableAudioConfig & config) {
    auto gate_bias = core::reshape_tensor(ctx, weights.scale_shift_gate, core::TensorShape::from_dims({1, 6 * config.embed_dim}));
    gate_bias = modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], 6 * config.embed_dim})}).build(ctx, gate_bias);
    auto adaln = modules::AddModule{}.build(ctx, global_six, gate_bias);
    adaln = core::reshape_tensor(ctx, adaln, core::TensorShape::from_dims({input.shape.dims[0], 1, 6 * config.embed_dim}));
    auto scale_self = modules::SliceModule({2, 0, config.embed_dim}).build(ctx, adaln);
    auto shift_self = modules::SliceModule({2, config.embed_dim, config.embed_dim}).build(ctx, adaln);
    auto gate_self = modules::SliceModule({2, 2 * config.embed_dim, config.embed_dim}).build(ctx, adaln);
    auto scale_ff = modules::SliceModule({2, 3 * config.embed_dim, config.embed_dim}).build(ctx, adaln);
    auto shift_ff = modules::SliceModule({2, 4 * config.embed_dim, config.embed_dim}).build(ctx, adaln);
    auto gate_ff = modules::SliceModule({2, 5 * config.embed_dim, config.embed_dim}).build(ctx, adaln);
    scale_self = modules::RepeatModule({input.shape}).build(ctx, scale_self);
    shift_self = modules::RepeatModule({input.shape}).build(ctx, shift_self);
    gate_self = modules::RepeatModule({input.shape}).build(ctx, gate_self);
    scale_ff = modules::RepeatModule({input.shape}).build(ctx, scale_ff);
    shift_ff = modules::RepeatModule({input.shape}).build(ctx, shift_ff);
    gate_ff = modules::RepeatModule({input.shape}).build(ctx, gate_ff);

    auto hidden = rms_norm(ctx, input, weights.pre_norm_gamma, kTransformerNormEps);
    hidden = modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, hidden, one_plus(ctx, scale_self)), shift_self);
    auto attn = self_attention(ctx, hidden, positions, padding_mask, weights.self_attn, config);
    attn = modules::MulModule{}.build(ctx, attn, sigmoid_one_minus(ctx, gate_self));
    hidden = modules::AddModule{}.build(ctx, input, attn);

    auto cross_norm = rms_norm(ctx, hidden, weights.cross_attend_norm_gamma, kTransformerNormEps);
    auto cross = cross_attention(ctx, cross_norm, cross_context, weights.cross_attn, config);
    hidden = modules::AddModule{}.build(ctx, hidden, cross);
    auto local = build_local_embedding(ctx, hidden, local_add, weights, config);
    hidden = modules::AddModule{}.build(ctx, hidden, local);

    auto ff_in = rms_norm(ctx, hidden, weights.ff_norm_gamma, kTransformerNormEps);
    ff_in = modules::AddModule{}.build(ctx, modules::MulModule{}.build(ctx, ff_in, one_plus(ctx, scale_ff)), shift_ff);
    auto ff = swiglu_ff(ctx, ff_in, weights.ff, config);
    ff = modules::MulModule{}.build(ctx, ff, sigmoid_one_minus(ctx, gate_ff));
    return modules::AddModule{}.build(ctx, hidden, ff);
}

StableAudioRfAttentionWeights load_self_attention(
    core::BackendWeightStore & store,