Method FlowLMStepRuntime

src/models/pocket_tts/flow_lm.cpp:190–417

    core::BackendType backend_type_ = core::BackendType::Cpu;
};

class FlowLMStepRuntime {
public:
    FlowLMStepRuntime(
        const FlowLMConfig & config,
        ggml_backend_t backend,
        int threads,
        std::shared_ptr<const FlowLMWeightsRuntime> weights_runtime,
        int64_t cache_steps,
        int64_t prompt_steps,
        int64_t prompt_prefix_steps,
        size_t graph_context_bytes)
        : config_(config),
          backend_(backend),
          threads_(threads),
          weights_runtime_(std::move(weights_runtime)),
          cache_steps_(std::max<int64_t>(1, cache_steps)),
          prompt_steps_(std::max<int64_t>(0, prompt_steps)),
          prompt_prefix_steps_(std::max<int64_t>(0, prompt_prefix_steps)),
          head_dim_(config.hidden_size / config.num_heads) {
        if (weights_runtime_ == nullptr) {
            throw std::runtime_error("FlowLM step runtime requires weights runtime");
        }
        const auto & weights = weights_runtime_->weights();
        ggml_ctx_ = ggml_init({graph_context_bytes, nullptr, true});
        if (ggml_ctx_ == nullptr) {
            throw std::runtime_error("Failed to initialize FlowLM step ggml context");
        }

        core::ModuleBuildContext ctx{ggml_ctx_, "flow_lm_step_runtime", weights_runtime_->backend_type()};

        if (prompt_steps_ > 0) {
            prompt_positions_buffer_.resize(static_cast<size_t>(prompt_steps_));
            prompt_embeddings_buffer_.assign(
                static_cast<size_t>(prompt_steps_ * config_.hidden_size),
                0.0F);
            prompt_embeddings_ = core::make_tensor(
                ctx,
                GGML_TYPE_F32,
                core::TensorShape::from_dims({1, prompt_steps_, config_.hidden_size}));
            prompt_positions_ = core::make_tensor(
                ctx,
                GGML_TYPE_I32,
                core::TensorShape::from_dims({prompt_steps_}));
            write_prompt_attention_mask(
                prompt_attention_mask_buffer_,
                0,
                prompt_prefix_steps_,
                0,
                prompt_steps_);
            prompt_attention_mask_ = core::make_tensor(
                ctx,
                GGML_TYPE_F32,
                core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_prefix_steps_ + prompt_steps_}));
        }

        latent_ = core::make_tensor(
            ctx,