o
    i                     @   sp   d dl mZmZmZmZmZ d dlZd dlZd dlm	Z	 d dl
mZ d dlmZ d dlmZ G dd deZdS )	    )AnyDictOptionalTupleUnionN)check_argument_types)LogMel)Stft)AbsFeatsExtractc                       s   e Zd ZdZ													
		d$deeef dedededee dededededee dee dedee	 f fddZ
defddZdeeef fddZ	d%d ejd!ejdeejejf fd"d#Z  ZS )&LogMelFbankzZConventional frontend structure for TTS.

    Stft -> amplitude-spec -> Log-Mel-Fbank
    >     N   hannTFP           $@fsn_fft
win_length
hop_lengthwindowcenter
normalizedonesidedn_melsfminfmaxhtklog_basec              	      s   t  sJ t   t|trt|}|| _|	| _|| _	|| _
|| _|| _|
| _|| _t|||||||d| _t|||	|
|||d| _d S )N)r   r   r   r   r   r   r   )r   r   r   r   r   r   r   )r   super__init__
isinstancestrhumanfriendly
parse_sizer   r   r   r   r   r   r   r   r	   stftr   logmel)selfr   r   r   r   r   r   r   r   r   r   r   r   r   	__class__ [/home/ubuntu/.local/lib/python3.10/site-packages/espnet2/tts/feats_extract/log_mel_fbank.pyr!      s<   




zLogMelFbank.__init__returnc                 C   s   | j S N)r   r(   r+   r+   r,   output_sizeD   s   zLogMelFbank.output_sizec              
   C   s(   t | j| j| j| j| j| j| j| jdS )z)Return the parameters required by Vocoder)r   r   n_shiftr   r   r   r   r   )	dictr   r   r   r   r   r   r   r   r/   r+   r+   r,   get_parametersG   s   zLogMelFbank.get_parametersinputinput_lengthsc           	      C   s   |  ||\}}| dksJ |j|jd dksJ |j|d d |d d  }ttj|dd}| ||\}}||fS )N      ).r   ).   g|=)min)r&   dimshapetorchsqrtclampr'   )	r(   r4   r5   
input_stft
feats_lensinput_power	input_ampinput_feats_r+   r+   r,   forwardT   s   zLogMelFbank.forward)r   r   Nr   r   TFTr   r   r   Fr   r.   )__name__
__module____qualname____doc__r   intr#   r   boolfloatr!   r0   r   r   r3   r=   Tensorr   rF   __classcell__r+   r+   r)   r,   r      sh    
	
2r   )typingr   r   r   r   r   r$   r=   	typeguardr   espnet2.layers.log_melr   espnet2.layers.stftr	   +espnet2.tts.feats_extract.abs_feats_extractr
   r   r+   r+   r+   r,   <module>   s    