
    hO                     B    d dl Z d dlmZ d dlmZ d dlZ G d de      Zy)    N)SpectralGate)	TorchGatec                   P     e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d fd	Zd Z xZS )StreamedTorchGatez)
    Run interface with noisereduce.
    c                 ,   t         |   ||||||||||||||       || _        ||j                  d   |j                  d   kD  r|r|d |j                  d    }t	        j
                  |      j                  |      }t        |j                        dk(  r|j                  d      }|| _	        t        || ||	d|
z  t        || j                  z  |z        || j                  | j                  | j                  ||      j                  |      | _        y )N)ysr
chunk_sizepaddingn_fft
win_length
hop_lengthtime_constant_sfreq_mask_smooth_hztime_mask_smooth_ms
tmp_folderprop_decreaseuse_tqdmn_jobs   r   )r	   nonstationaryn_std_thresh_stationaryn_thresh_nonstationarytemp_coeff_nonstationaryn_movemean_nonstationaryr   r   r   r   r   r   )super__init__deviceshapetorch
from_numpytolen	unsqueezey_noiseTGint_hop_length_n_fft_win_lengthtg)selfr   r	   
stationaryr&   r   r   r   r   thresh_n_mult_nonstationarysigmoid_slope_nonstationaryr   r   r
   r   r   r   r   clip_noise_stationaryr   r   r   	__class__s                         f/var/www/html/engine/venv/lib/python3.12/site-packages/noisereduce/spectralgate/streamed_torch_gate.pyr   zStreamedTorchGate.__init__   s2   0 	!!!+ 3 3!' 	 	
"  }}R 1772;.3H!-AGGBK0&&w/226:G7==!Q&!++A. (.$;#>%&)D%D%(4;K;K)Kb)P%Q'++'''' 3 3
 "V* 	    c                 &   t        |      t        j                  u r.t        j                  |      j                  | j                        }| j                  || j                        }|j                         j                         j                         S )zDo the actual filtering)xxn)typenpndarrayr!   r"   r#   r   r,   r&   cpudetachnumpy)r-   chunkchunk_filtereds      r3   
_do_filterzStreamedTorchGate._do_filterQ   sj     ;"**$$$U+..t{{;E5T\\:!!#**,2244r4   )FNg      ?g       @i  2      
   g      ?Ni'	 i0u  i   NNTFr   cuda)__name__
__module____qualname____doc__r   r@   __classcell__)r2   s   @r3   r   r      sO      # "()(*$'"&-CJ5r4   r   )	r!   noisereduce.spectralgate.baser   noisereduce.torchgater   r'   r=   r9   r    r4   r3   <module>rM      s     6 1 P5 P5r4   