Function generate_heartmula_frames

src/models/heartmula/generator.cpp:362–430

}  // namespace

HeartMuLaGeneratedFrames generate_heartmula_frames(
    const HeartMuLaPromptRequest & request,
    const HeartMuLaTextTokenizer & tokenizer,
    const HeartMuLaWeightsRuntime & mula,
    uint64_t seed) {
    const auto & assets = mula.assets();
    const auto & config = assets.mula_config;
    const auto encoding = tokenizer.encode_prompt(request);
    const bool use_cfg = request.options.guidance_scale > 1.0F && encoding.batch_size > 1 && encoding.batch_size % 2 == 0;
    auto prompt_embeddings = mula.merge_frame_embeddings(prompt_embedding_inputs(encoding, config, use_cfg));
    auto backbone = mula.backbone_prefill_embeddings(
        prompt_embeddings.values,
        prompt_embeddings.batch_size,
        prompt_embeddings.steps);
    HeartMuLaBackboneCachedState backbone_state;
    mula.reset_backbone_cached_state(backbone_state, std::move(backbone.state));

    const TorchCudaSamplingPolicy policy = engine::sampling::resolve_torch_cuda_sampling_policy(
        mula.backend_type(),
        mula.device(),
        "heartmula.cuda_sampling_policy",
        "HeartMuLa",
        engine::sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda);
    TopKSamplerScratch scratch;
    uint64_t sample_call_index = 0;
    HeartMuLaGeneratedFrames out;
    out.codebooks = config.audio_num_codebooks;
    auto current_frame = generate_frame_from_backbone(
        backbone.result,
        mula,
        encoding.batch_size,
        request.options,
        seed,
        sample_call_index,
        policy,
        scratch);
    append_first_branch_frame(out, current_frame, config);

    const int64_t max_audio_frames = static_cast<int64_t>(request.options.duration_seconds * 1000.0F) / 80;
    for (int64_t frame_index = 0; frame_index < max_audio_frames; ++frame_index) {
        auto next_embeddings = mula.merge_frame_embeddings(next_frame_embedding_inputs(current_frame, encoding.batch_size, config));
        auto next_backbone = mula.backbone_cached_step(
            next_embeddings.values,
            encoding.batch_size,
            backbone_state,
            encoding.prompt_len + max_audio_frames + 1);
        current_frame = generate_frame_from_backbone(
            next_backbone,
            mula,
            encoding.batch_size,
            request.options,
            seed,
            sample_call_index,
            policy,
            scratch);
        if (first_branch_has_eos(current_frame, assets)) {
            break;
        }