
    h&                     8   S SK r S SKrS SKrS SKJr  S SKrS SKrS SKJr  S SK	J
r
Jr  SSKJr  SSKJr  SSKJrJr  \ R&                  " \5      r             SS	\\-  \R.                  -  S-  S
\\-  S-  S\S\S\S\S\S\S\S\S\S\S-  S\SS4S jjrg)    N)Path)SymbolicShapeInference)extract_raw_data_from_modelhas_external_data   )ReplaceUpsampleWithResize)	ONNXModel)add_pre_process_metadata&save_and_reload_model_with_shape_inferinput_modeloutput_model_pathskip_optimizationskip_onnx_shapeskip_symbolic_shape
auto_mergeint_maxguess_output_rankverbosesave_as_external_dataall_tensors_to_one_fileexternal_data_locationexternal_data_size_thresholdreturnc           
      	   U c  UR                  SS5      n U c   eUc   S5       e[        R                  " SS9 n[        U5      nSnU(       df  [        R                  S5        [        U [        R                  5      (       a  U O[        R                  " U 5      n[        R                  " UUUUU5      nU(       Ga  UR                   Vs/ s H(  nUR                  (       a  UR                  S:X  d  M&  UPM*     nn[        U5      S:X  a  US	   R                  nUS
:  a  [!        [#        U5      U5      R%                  5         UR                  R'                  US	   5        SnUR                  R)                  [        R*                  R-                  SU5      /5        [        R.                  R1                  UU5      n[3        U5      nU(       GdP  U(       dH  [5        US-  5      n U	(       a  [        R6                  " UU SU
USS9  O[        R8                  " UU 5        Sn[5        US-  5      n [:        R<                  " 5       nUUl        [:        R@                  RB                  Ul"        [        U [        R                  5      (       a^  [G        U 5      (       a  [I        S5      e[K        U 5      u  nnURM                  [O        U5      [O        U5      5        U RQ                  5       n O U(       a  U	(       a  URS                  SS5        [:        RT                  " U US/S9nAUn U(       d  UbH  [5        US-  5      n U	(       a  [        R6                  " UU SU
USS9  O[        R8                  " UU 5        Sn[        U [        R                  5      (       a0  [5        [        U5      S-  5      n [        R6                  " UU SU
USS9  [5        US-  5      n[        R^                  Ra                  U U5        [        R                  " U5      nSSS5        Wc7  [        U [        R                  5      (       a  U O[        R                  " U 5      n[c        U5        U	(       a  [        R6                  " UUSU
UUSS9  g[        R8                  " UU5        gs  snf ! [V         aA    [        RY                  S5        [        RY                  [Z        R\                  " 5       5         GNf = f! , (       d  f       N= f)a(  Shape inference and model optimization, in preparation for quantization.

Args:
    input_model: Path to the input model file or ModelProto
    output_model_path: Path to the output model file
    skip_optimization: Skip model optimization step if true. This may result in ONNX shape
        inference failure for some models.
    skip_onnx_shape: Skip ONNX shape inference. Symbolic shape inference is most effective
        with transformer based models. Skipping all shape inferences may
        reduce the effectiveness of quantization, as a tensor with unknown
        shape can not be quantized.
    skip_symbolic_shape: Skip symbolic shape inference. Symbolic shape inference is most
        effective with transformer based models. Skipping all shape
        inferences may reduce the effectiveness of quantization, as a tensor
        with unknown shape can not be quantized.
    auto_merge: For symbolic shape inference, automatically merge symbolic dims when
        conflict happens.
    int_max: For symbolic shape inference, specify the maximum value for integer to be
        treated as boundless for ops like slice
    guess_output_rank: Guess output rank to be the same as input 0 for unknown ops
    verbose: Logs detailed info of inference, 0: turn off, 1: warnings, 3: detailed
    save_as_external_data: Saving an ONNX model to external data
    all_tensors_to_one_file: Saving all the external data to one file
    external_data_location: The file location to save the external file
    external_data_size_threshold: The size threshold for external data
Ninput_model_pathzoutput_model_path is required.z
pre.quant.)prefixz&Performing symbolic shape inference...zai.onnxr   r   
       zsymbolic_shape_inferred.onnxTF)r   r   size_thresholdconvert_attributezoptimized.onnxzModelProto has external data not loaded into memory, ORT cannot create session. Please load external data before calling this function. See https://onnx.ai/onnx/repo-docs/ExternalData.html for more information.z7session.optimized_model_external_initializers_file_namezoptimized.onnx.dataCPUExecutionProvider)	providerszYONNX Runtime Model Optimization Failed! Consider rerun with option `--skip_optimization'.zmodel_input.onnxzonnx_shape_inferred.onnx)r   r   locationr    r!   )2poptempfileTemporaryDirectoryr   loggerinfo
isinstanceonnx
ModelProtoloadr   infer_shapesopset_importdomainlenversionr   r	   applyremoveextendhelpermake_opsetidversion_converterconvert_versionr   str
save_modelsaveonnxruntimeSessionOptionsoptimized_model_filepathGraphOptimizationLevelORT_ENABLE_BASICgraph_optimization_levelr   
ValueErrorr   add_external_initializerslistSerializeToStringadd_session_config_entryInferenceSession	Exceptionerror	traceback
format_excshape_inferenceinfer_shapes_pathr
   )r   r   r   r   r   r   r   r   r   r   r   r   r   deprecated_kwargsquant_tmp_dir	temp_pathmodelloaded_modelopsetai_onnx_domainopset_versionopt_model_pathsess_optionexternal_namesexternal_valuessessinferred_model_paths                              b/var/www/fran/franai/venv/lib/python3.13/site-packages/onnxruntime/quantization/shape_inference.pyquant_pre_processr^      sE   V '++,>E"""(J*JJ(		$	$L	9]'	"KK@A*4[$//*R*R;X\XaXabmXnL*77!E 161C1Cu1C5<<[`[g[gkt[te1CNu>"a' .q 1 9 9 2%-i.>NTTV&&--nQ.?@$&M&&--t{{/G/GM/Z.[\ 22BB5-XEB5IE &!).L"LM(OO#.20G'C*/ IIe[1 -=!=>N5)88:7E47B7Y7Y7j7j4k4??;;(55(i 
 7RR]6^3NO99$~:NPTUdPef"-"?"?"AK )-B88QSh #33KYoXpq  )K
  !).L"LM(OO#.20G'C*/ IIe[1+t77!$}"58J"JK*.,C#?&+ #&i2L&L"M  22;@STII12Ei 
:l })+tGGTYYWbMcU#"&$;+7#	
 			%*+g vp  5o Y1134	5W 
:	9sL   BS
%R3R9D3S-C0RC'SSASSSS
S,)NNFFFFiFr   FFNi   )loggingr&   rK   pathlibr   r+   r=   &onnxruntime.tools.symbolic_shape_inferr   #onnxruntime.transformers.onnx_utilsr   r   fusionsr   
onnx_modelr	   quant_utilsr
   r   	getLogger__name__r(   r:   r,   boolintr^        r]   <module>rl      s         I ^ . ! Y			8	$ 8<+/#! %#"'$))-(,w,tdoo-4w,TzD(w, w, 	w,
 w, w, w, w, w,  w, "w,  $Jw, #&w, 
w,rk   