
    
Ng                         d dl Zd dlmZ d dlmZ d dlmZ d dl	m
Z
 d ddZd Zdd
efdZdd
efdZ G d de          ZdS )    N)OpRun)_concat_from_sequence)_cfft)_sliceaxisc                 .    t          j        ||           S )Nr   )npconcatenate)r   argss     V/var/www/html/ai-engine/env/lib/python3.11/site-packages/onnx/reference/ops/op_stft.py_concatr      s    >$T****    c                     	 t          j        | |          S # t          $ rk t          |          dk    r+t          j        | t	          |          d                   cY S t          |          D ]}t          j        | |          } | cY S w xY w)Nr      r   )r
   expand_dims	TypeErrorlentuplereversed)ar   xs      r   
_unsqueezer      s    ~ad++++   t99>>>!%++a.999999$ 	* 	*Aqq)))AAs    AB *BBF
fft_lengthc                 
   t          | j                  dz
  }dg}dg}|j        d         }	g }
t          |          D ]}||z  }||	z   }t          | t	          j        |g          t	          j        |g          |          }|j        dd         }|	|d         z
  f}|j        dd         |z   |j        dd         z   }t	          j        || j                  }t          ||d          }t          ||          }|

                    |           t          |
dd	          }|j        }|dd         }t          d
 |D                       }g ||	dR }t	          j        ||          }||z  }t          ||||d          }|S )zApplies one dimensional FFT with window weights.

    torch defines the number of frames as:
    `n_frames = 1 + (len - n_fft) // hop_length`.
    r   r   Ndtyper   r   new_axisc              3      K   | ]}d V  dS )r   N ).0_s     r   	<genexpr>z_stft.<locals>.<genexpr>A   s"      77Aa777777r   F)onesided	normalize)r   shaperanger   r
   arrayzerosr    r   r   appendr   r   reshape_dft)r   r   
hop_lengthn_frameswindowr(   	last_axisr   axis2window_sizeseqfsbeginendsliced_xnew_dimmissing	new_shapecstpad_sliced_xun_sliced_xnew_xshape_xshape_x_shortshape_x_short_onewindow_shapeweightsweighted_new_xresults                                r   _stftrJ      s    AGq I4DDE,q/K CHoo    Zk!!RXug..#FF .B'+-N3B3''1HN2334GG	hy000x2666 !u55

; "#B;;;E kGCRCLM77777777&77Q77Lj..GW_N
IE  F Mr   c           
      x   dg}dg}dg}dg}| j         d         }	|||	dz
  z  z   }
g }g }g }t          |	          D ]Y}|}|dz   }t          j        t	          | t          j        |g          t          j        |g          |          |d                   }t          ||d|d          }t          |j                   }|dz
  }t	          |t          j        |          t          j        |          |g          }t          j        ||          }t          j        |j         d| j	        	          |z  }|j         d
d         }||z  }|j         d         }|
||z   z
  }g ||R }g ||R }t          j
        || j	                  }t          j
        || j	                  }t          |||d          } t          |||d          }!t	          |t          j        |          t          j        |          |g          }t          j        ||          }"t          ||"|d          }#|                    t          | d                     |                    t          |#d                     |                    t          |!d                     [t          |dd          }$t          |dd          }%t          |dd          }&|$                    dd          }'|%                    dd          }(|&                    dd          })|'|)z  }*|(|)z  }+t          j        |*d          },t          j        |+d          }-t          |,|-d          }.|.j         }/|.                    d          }0t          j        |0d          }1|/dd
         }2t          |2|d          }3|1                    |3          }4|4S )zReverses of `stft`.r   r      r   r   r   T)r   r(   r)   )
fill_valuer    Nr   r!   )r   keepdims)rL   r   )r   r   )r*   r+   r
   squeezer   r,   r0   r   fullr    r-   r   r.   r   r   sumr   r/   	transpose)5r   r   r1   r3   r(   zeroonetwoaxisfr2   expected_signal_lenseqrseqiseqcr8   r9   r:   frame_xiftn_dimsn_dims_1slicedytmpctmpshape_beginn_leftsizen_right
left_shaperight_shaperightleftyycitmpyiredrrediredcresrresirescrrrirr0ri0concresult_shapereshaped_result
transposedother_dimensionsfinal_shapefinals5                                                        r   _istftr   M   s   3D#C#CDEwr{H$zX\'BB DDDHoo &- &-1f*1bhw''3%%@@uQx
 
 

 7JR(dSSSSY A:RXd^^RXc]]XJGGz&x000wtzaqw???&Hj"ojz"~%$7+{+F++
--W--AG444x
!'222D$B///T4R000 RXc]]BHSMMH:FFz&x000T4R000 	Jqr***+++Jr+++,,,Jr+++,,,, !B;;;D B;;;D B;;;D 88a8((D88a8((D88a8((D	B	B .!
$
$
$C
.!
$
$
$C3!$$$D :Lll7++Oov66J#ABB'*Ca888K{++ELr   c                       e Zd ZddZdS )STFTNc                    |||j         d         }n|j         d         }|}|t          j        |f|j                  }d|j         d         |z
  |z  z   }t	          ||||||          }|                    |j                  fS )Nr   r   r   r   )r(   )r*   r
   onesr    rJ   astype)	selfr   
frame_stepr3   frame_lengthr(   r1   r2   ress	            r   _runz	STFT._run   s    ~ wr{%|A
>Wl_AG<<<Fl2zAAA|Z6HUUU

17##%%r   )NNN)__name__
__module____qualname__r   r$   r   r   r   r      s(        & & & & & &r   r   )F)numpyr
   onnx.reference.op_runr   *onnx.reference.ops.op_concat_from_sequencer   onnx.reference.ops.op_dftr   r0   onnx.reference.ops.op_slicer   r   r   intrJ   r   r   r$   r   r   <module>r      s       ' ' ' ' ' ' L L L L L L 3 3 3 3 3 3 . . . . . .  + + + + +	 	 	, , , , , ,^M M# M M M M`& & & & &5 & & & & &r   