| 50 | } |
| 51 | |
| 52 | void HTDemucsPostprocessor::combine_chunk_into( |
| 53 | const float * freq_output, |
| 54 | const float * time_output, |
| 55 | int64_t input_samples, |
| 56 | float freq_mean, |
| 57 | float freq_std, |
| 58 | float time_mean, |
| 59 | float time_std, |
| 60 | float * output) { |
| 61 | if (input_samples <= 0 || input_samples > config_.segment_samples) { |
| 62 | throw std::runtime_error("HTDemucs postprocessor input sample count is invalid"); |
| 63 | } |
| 64 | if (freq_output == nullptr || time_output == nullptr || output == nullptr) { |
| 65 | throw std::runtime_error("HTDemucs postprocessor received null buffers"); |
| 66 | } |
| 67 | |
| 68 | const int64_t model_samples = config_.segment_samples; |
| 69 | const int64_t le = static_cast<int64_t>(std::ceil(static_cast<double>(model_samples) / static_cast<double>(config_.hop_length))); |
| 70 | const int64_t stft_frames = config_.stft_frames; |
| 71 | const int64_t pad = (config_.hop_length / 2) * 3; |
| 72 | const int64_t istft_length = config_.hop_length * le + 2 * pad; |
| 73 | const int64_t full_freq_bins = config_.stft_freq_bins + 1; |
| 74 | const int64_t full_frames = stft_frames + 4; |
| 75 | const int64_t batch = static_cast<int64_t>(config_.sources.size() * config_.audio_channels); |
| 76 | float * spectrum_data = reinterpret_cast<float *>(spectrum_.data()); |
| 77 | const float freq_scale = freq_std * fft_scale_; |
| 78 | const float freq_bias = freq_mean * fft_scale_; |
| 79 | const int64_t source_count = static_cast<int64_t>(config_.sources.size()); |
| 80 | |
| 81 | #ifdef _OPENMP |
| 82 | #pragma omp parallel for collapse(2) if(static_cast<int64_t>(config_.sources.size()) * config_.audio_channels >= 4) |
| 83 | #endif |
| 84 | for (int64_t source = 0; source < source_count; ++source) { |
| 85 | for (int ch = 0; ch < config_.audio_channels; ++ch) { |
| 86 | const int64_t batch_index = static_cast<int64_t>(source * config_.audio_channels + ch); |
| 87 | float * batch_dst = spectrum_data + |
| 88 | static_cast<size_t>(batch_index * full_frames * full_freq_bins * 2); |
| 89 | const int64_t real_channel = static_cast<int64_t>(source * config_.audio_channels * 2 + ch * 2); |
| 90 | const int64_t imag_channel = real_channel + 1; |
| 91 | for (int64_t f = 0; f < config_.stft_freq_bins; ++f) { |
| 92 | for (int64_t t = 0; t < stft_frames; ++t) { |
| 93 | const size_t src_real = static_cast<size_t>(((real_channel * config_.stft_freq_bins + f) * stft_frames + t)); |
| 94 | const size_t src_imag = static_cast<size_t>(((imag_channel * config_.stft_freq_bins + f) * stft_frames + t)); |
| 95 | float * dst = batch_dst + static_cast<size_t>((((t + 2) * full_freq_bins) + f) * 2); |
| 96 | dst[0] = freq_output[src_real] * freq_scale + freq_bias; |
| 97 | dst[1] = freq_output[src_imag] * freq_scale + freq_bias; |
| 98 | } |
| 99 | } |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | const int64_t padded_samples = istft_length + config_.n_fft; |
| 104 | std::fill(accum_.begin(), accum_.end(), 0.0f); |
| 105 | |
| 106 | const engine::audio::TensorShape output_shape{ |
| 107 | static_cast<size_t>(batch), |
| 108 | static_cast<size_t>(full_frames), |
| 109 | static_cast<size_t>(config_.n_fft), |