Function flow_transformer_block

src/models/heartmula/codec.cpp:847–874

    int64_t hidden_size) {
    const int64_t batch = input.shape.dims[0];
    const int64_t frames = input.shape.dims[1];
    auto table = core::reshape_tensor(ctx, scale_shift_table, core::TensorShape::from_dims({1, 2, hidden_size}));
    table = modules::RepeatModule({core::TensorShape::from_dims({batch, 2, hidden_size})}).build(ctx, table);
    auto embedded = core::reshape_tensor(ctx, embedded_timestep, core::TensorShape::from_dims({batch, 1, hidden_size}));
    embedded = modules::RepeatModule({core::TensorShape::from_dims({batch, 2, hidden_size})}).build(ctx, embedded);
    auto parts = modules::AddModule{}.build(ctx, embedded, table);
    auto shift = expand_batch_token(ctx, modules::SliceModule({1, 0, 1}).build(ctx, parts), frames, hidden_size);
    auto scale_part = expand_batch_token(ctx, modules::SliceModule({1, 1, 1}).build(ctx, parts), frames, hidden_size);
    auto normalized = modules::LayerNormModule({hidden_size, 1.0e-6F, false, false}).build(ctx, input, {});
    return modules::AddModule{}.build(
        ctx,
        modules::MulModule{}.build(ctx, normalized, add_one(ctx, scale_part)),
        shift);
}

FlowEstimatorBuildResult flow_estimator(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & hidden_states,
    const core::TensorValue & timesteps,
    const core::TensorValue & freqs,
    const core::TensorValue & positions,
    const HeartCodecFlowEstimatorWeights & weights,
    const HeartCodecConfig & config) {
    const int64_t estimator_dim = config.num_attention_heads * config.attention_head_dim;
    const int64_t estimator_dim_2 = 2 * estimator_dim;
    auto s = project_layer(ctx, hidden_states, weights.proj_in, config.in_channels, estimator_dim);
    auto embedded_timestep = adaln_embedding(ctx, timesteps, freqs, weights.adaln_single, estimator_dim);
    auto timestep_mod = modules::LinearModule({estimator_dim, 6 * estimator_dim, true, GGML_PREC_F32})
                            .build(ctx, modules::SiluModule{}.build(ctx, embedded_timestep), weights.adaln_single.linear);
    for (size_t index = 0; index < weights.transformer_blocks.size(); ++index) {
        const auto & block = weights.transformer_blocks[index];