
    `i                     P    S SK Jr  S SKJr  S SKJrJr  \" \5      r " S S5      r	g)    )	getLogger)	OnnxModel)UnionListc            
       J    \ rS rSr S
S\S\S\\\\   4   S\4S jjrS r	S	r
g)Fusion   Nmodelfused_op_typesearch_op_typesdescriptionc                     [        U[        5      (       a  U/OUU l        X l        U(       a  U SU S3OUU l        Xl        / U l        / U l        SU l        g )N()F)	
isinstancestrr   r   r   r
   nodes_to_removenodes_to_addprune_graph)selfr
   r   r   r   s        i/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/onnxruntime_tools/transformers/fusion_base.py__init__Fusion.__init__   s^    
 @J/[^?_?_?*;et"/EP=/;-q AVc %
%'"$!&    c                    [         R                  SU R                   S35        U R                  R	                  5       nU R                  R                  5       nU R                   H7  nU R                  R                  U5       H  nU R                  XAU5        M     M9     U R                   Vs/ s H  oDR                  PM     nnUR                  U R                  5      nUS:  a%  [         R                  SU R                   SU 35        U R                  R                  U R                  5        U R                  R!                  U R                  5        U R"                  (       a  U R                  R#                  5         g U R                  (       d  U R                  (       a  U R                  R%                  5         g g s  snf )Nzstart z
 fusion...r   zFused z count: )loggerdebugr   r
   input_name_to_nodesoutput_name_to_noder   get_nodes_by_op_typefuser   op_typecountr   inforemove_nodesr   	add_nodesr   update_graph)r   r   r   search_op_typenodeop_listr#   s          r   applyFusion.apply   sN   vd../z:;"jj<<>"jj<<> #22N

77G		$5HI H 3 -1,=,=>,=D<<,=>d00119KK&!1!1 2(5'BC

 4 45

T../JJ""$!!T%6%6JJ##% &7 ?s   -F>)r   r   r
   r   r   r   r   )N)__name__
__module____qualname____firstlineno__r   r   r   r   r   r+   __static_attributes__ r   r   r   r      sB    
 %)	'!' #' #(T#Y"7' "	'&r   r   N)
loggingr   
onnx_modelr   typingr   r   r-   r   r   r2   r   r   <module>r6      s$   
    	8	#& #&r   