| 697 | |
| 698 | std::vector<float> build_band_features( |
| 699 | const engine::audio::AudioTensor & stft, |
| 700 | const RoformerArchitectureConfig & config) { |
| 701 | if (stft.shape.size() != 4) { |
| 702 | throw std::runtime_error("RoFormer STFT tensor must be rank-4"); |
| 703 | } |
| 704 | const int64_t channels = stft.shape[0]; |
| 705 | const int64_t freq_bins = stft.shape[1]; |
| 706 | const int64_t frames = stft.shape[2]; |
| 707 | if (channels != config.channels || freq_bins != config.stft_freq_bins || frames != config.chunk_frames || stft.shape[3] != 2) { |
| 708 | throw std::runtime_error("RoFormer STFT tensor shape mismatch"); |
| 709 | } |
| 710 | std::vector<float> features(static_cast<size_t>(frames * config.total_band_input_dim), 0.0f); |
| 711 | const int64_t merged = static_cast<int64_t>(config.merged_freq_indices.size()); |
| 712 | #ifdef _OPENMP |
| 713 | #pragma omp parallel for collapse(2) if(frames * merged >= 4096) |
| 714 | #endif |
| 715 | for (int64_t t = 0; t < frames; ++t) { |
| 716 | for (int64_t m = 0; m < merged; ++m) { |
| 717 | const int64_t merged_index = config.merged_freq_indices[static_cast<size_t>(m)]; |
| 718 | const int64_t channel = merged_index % config.channels; |
| 719 | const int64_t freq = merged_index / config.channels; |
| 720 | const size_t stft_base = static_cast<size_t>((((channel * freq_bins) + freq) * frames + t) * 2); |
| 721 | const size_t dst = static_cast<size_t>(t * config.total_band_input_dim + m * 2); |
| 722 | features[dst] = stft.values[stft_base]; |
| 723 | features[dst + 1] = stft.values[stft_base + 1]; |
| 724 | } |
| 725 | } |
| 726 | return features; |
| 727 | } |
| 728 | |
| 729 | std::vector<float> apply_masks_to_stft( |
| 730 | const std::vector<float> & raw_masks, |
| 731 | const engine::audio::AudioTensor & stft, |
| 732 | const RoformerArchitectureConfig & config) { |
| 733 | const int64_t channels = config.channels; |
| 734 | const int64_t freq_bins = config.stft_freq_bins; |
| 735 | const int64_t frames = config.chunk_frames; |
| 736 | std::vector<float> averaged_masks(static_cast<size_t>(channels * freq_bins * frames * 2), 0.0f); |
| 737 | const int64_t merged = static_cast<int64_t>(config.merged_freq_indices.size()); |
| 738 | |
| 739 | #ifdef _OPENMP |
| 740 | #pragma omp parallel for if(merged * frames >= 4096) |
| 741 | #endif |
| 742 | for (int64_t m = 0; m < merged; ++m) { |
| 743 | const int64_t merged_index = config.merged_freq_indices[static_cast<size_t>(m)]; |
| 744 | for (int64_t t = 0; t < frames; ++t) { |
| 745 | const size_t src = static_cast<size_t>(t * config.total_band_input_dim + m * 2); |
| 746 | const size_t dst = static_cast<size_t>(((merged_index * frames) + t) * 2); |
| 747 | averaged_masks[dst] += raw_masks[src]; |
| 748 | averaged_masks[dst + 1] += raw_masks[src + 1]; |
| 749 | } |
| 750 | } |
| 751 | |
| 752 | #ifdef _OPENMP |
| 753 | #pragma omp parallel for if(channels * freq_bins >= 512) |
| 754 | #endif |
| 755 | for (int64_t merged_index = 0; merged_index < channels * freq_bins; ++merged_index) { |
| 756 | const float denom = static_cast<float>(std::max<int64_t>(1, config.merged_band_counts[static_cast<size_t>(merged_index)])); |
no test coverage detected