o }oi< @sLddlZddlmZdejdedejdejfddZed d d d d ZdS) N)jitDrate phi_advance scale_buffercCsztjd|jd|tjd}tj|jdt|f|jdd}t|dddf}tj|ddgd d }t ||||||}|S) a Optimized implementation of phase vocoder from Librosa. Reference implementation: - https://librosa.github.io/librosa/generated/librosa.core.phase_vocoder.html Args: D: Complex spectograms of shape [d, t, complex=2]. rate: Speed rate, must be float greater than 0. phi_advance: Precomputed phase advance buffer array of length [n_fft + 1] scale_buffer: Precomputed numpy buffer array of length [n_fft + 1] Returns: Complex64 ndarray of shape [d, t / rate, complex=2] r)dtypeF)orderN)rr)rconstant)mode) nparangeshapefloat64zeroslenranglepad_phase_vocoder_kernel)rrrr time_steps d_stretch phase_accr`/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/utils/numba_utils.py phase_vocoders  rT)nopythonnogilcCsdtj}t|D]l\}}|ddt|t|df} | dddf} | dddf} t|d} d| t| | t| } | td||dd|f<t| t| |}||}t|d||||}|||7}q |S)a Numba optimized kernel to compute the phase vocoder step. Args: D: Complex spectograms of shape [d, t, complex=2]. rate: Speed rate, must be float greater than 0. time_steps: Numpy ndarray of linearly spaced time steps, shape = [t] phi_advance: Precomputed phase advance buffer array of length [n_fft + 1] d_stretch: Output complex matrix of shape [d, t / rate, complex=2] phase_acc: Phase accumulator initialized to first sample of shape [d, complex=2] scale_buffer: Precomputed numpy buffer array of length [n_fft + 1] Returns: Complex64 ndarray of shape [d, t / rate, complex=2] g@Nr rrg?y?) rpi enumerateintmodabsexprround)rrrrrrtwo_pitstepcolumns columns_0 columns_1alphamagdphasescalerrrr/s     r)numpyrnumbarndarrayfloatrrrrrrs