| 275 | } |
| 276 | |
| 277 | std::vector<int32_t> generate_frame_from_backbone( |
| 278 | const HeartMuLaBackboneResult & backbone_result, |
| 279 | const HeartMuLaWeightsRuntime & mula, |
| 280 | int64_t batch, |
| 281 | const HeartMuLaGenerationOptions & options, |
| 282 | uint64_t seed, |
| 283 | uint64_t & sample_call_index, |
| 284 | const TorchCudaSamplingPolicy & policy, |
| 285 | TopKSamplerScratch & scratch) { |
| 286 | const auto & config = mula.assets().mula_config; |
| 287 | auto c0_sample = sample_with_cfg( |
| 288 | backbone_result.logits.values, |
| 289 | batch, |
| 290 | backbone_result.logits.vocab_size, |
| 291 | options.guidance_scale, |
| 292 | options.top_k, |
| 293 | options.temperature, |
| 294 | seed, |
| 295 | sample_call_index, |
| 296 | policy, |
| 297 | scratch); |
| 298 | std::vector<int32_t> frame(static_cast<size_t>(batch * config.audio_num_codebooks), 0); |
| 299 | for (int64_t b = 0; b < batch; ++b) { |
| 300 | frame[static_cast<size_t>(b * config.audio_num_codebooks)] = c0_sample[static_cast<size_t>(b)]; |
| 301 | } |
| 302 | auto c0_embedding = mula.merge_frame_embeddings(audio_frame_embedding_inputs(c0_sample, batch, 0, config)); |
| 303 | const auto decoder_input = decoder_prefill_input(backbone_result.last_hidden, c0_embedding); |
| 304 | HeartMuLaDecoderCachedState decoder_state; |
| 305 | auto decoder = mula.decoder_prefill_embeddings( |
| 306 | decoder_input, |
| 307 | batch, |
| 308 | 2, |
| 309 | 0); |
| 310 | mula.reset_decoder_cached_state(decoder_state, std::move(decoder.state)); |
| 311 | for (int64_t codebook = 1; codebook < config.audio_num_codebooks; ++codebook) { |
| 312 | const auto sample = sample_with_cfg( |
| 313 | decoder.result.logits.values, |
| 314 | batch, |
| 315 | decoder.result.logits.vocab_size, |
| 316 | options.guidance_scale, |
| 317 | options.top_k, |
| 318 | options.temperature, |
| 319 | seed, |
| 320 | sample_call_index, |
| 321 | policy, |
| 322 | scratch); |
| 323 | for (int64_t b = 0; b < batch; ++b) { |
| 324 | frame[static_cast<size_t>(b * config.audio_num_codebooks + codebook)] = sample[static_cast<size_t>(b)]; |
| 325 | } |
| 326 | if (codebook + 1 >= config.audio_num_codebooks) { |
| 327 | break; |
| 328 | } |
| 329 | auto ci_embedding = mula.merge_frame_embeddings(audio_frame_embedding_inputs(sample, batch, codebook, config)); |
| 330 | decoder.result = mula.decoder_cached_step( |
| 331 | ci_embedding.values, |
| 332 | batch, |
| 333 | codebook, |
| 334 | decoder_state, |
no test coverage detected