Function flow_estimator

src/models/heartmula/codec.cpp:898–930

}  // namespace

class HeartCodecFlowEstimatorGraph {
public:
    HeartCodecFlowEstimatorGraph(
        const HeartCodecWeightsRuntime & runtime,
        int64_t batch_size,
        int64_t frames)
        : runtime_(&runtime),
          batch_size_(batch_size),
          frames_(frames) {
        if (batch_size_ <= 0 || frames_ <= 0) {
            throw std::runtime_error("HeartCodec flow estimator graph shape is invalid");
        }
        const auto & config = runtime_->assets().codec_config;
        ggml_init_params params{runtime_->flow_estimator_graph_arena_bytes(), nullptr, true};
        ctx_.reset(ggml_init(params));
        if (ctx_ == nullptr) {
            throw std::runtime_error("failed to initialize HeartCodec flow estimator graph context");
        }
        core::ModuleBuildContext ctx{ctx_.get(), "heartcodec.flow_estimator", runtime_->backend_type()};
        auto input = core::make_tensor(
            ctx,
            GGML_TYPE_F32,
            core::TensorShape::from_dims({batch_size_, frames_, config.in_channels}));
        input_ = input.tensor;
        auto timesteps = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({batch_size_}));
        timesteps_ = timesteps.tensor;
        freqs_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, 256, 1);
        auto freqs = core::wrap_tensor(freqs_, core::TensorShape::from_dims({1, 256}), GGML_TYPE_F32);
        positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, frames_);
        auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({frames_}), GGML_TYPE_I32);
        auto out = flow_estimator(
            ctx,
            input,
            timesteps,
            freqs,
            positions,