o
    H!i1                     @   s   d dl Z d dlmZ d dlmZmZ ddlmZ ddlm	Z	m
Z
mZ ddlmZmZmZ dd	lmZ dd
lmZ ddlmZ ddlmZmZ dd ZG dd deZdS )    N)deepcopy)DictList   )logging)get_default_deviceparse_deviceset_env_for_device_type)ENABLE_MKLDNN_BYDEFAULTUSE_PIR_TRTDISABLE_DEVICE_FALLBACK   )is_mkldnn_available)MKLDNN_BLOCKLIST)NEWIR_BLOCKLIST)TRT_CFG_SETTINGTRT_PRECISION_MAPc                 C   s.   | sdS |dkr
dS t rt r| tvrdS dS )Npaddlecpumkldnn)r
   r   r   )
model_namedevice_type r   m/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/paddlex/inference/utils/pp_option.pyget_default_run_mode   s   r   c                       s  e Zd ZdZdZdZ fddZdd Zdd	 Zd
d Z	dd Z
dd Zedd ZejdefddZedd Zejdd Zedd Zejdd Zedd Zejdd Zedd Zejdd Zed d! Zejd"efd#d!Zed$d% Zejd&efd'd%Zed(d) Zejd*efd+d)Zed,d- Zejd.d- Zed/d0 Zejd1d0 Zed2d3 Zejd4d3 Zed5d6 Zejd7eeeee  f fd8d6Zed9d: Zejd;eeee  f fd<d:Zed=d> Z!e!jd?efd@d>Z!edAdB Z"e"jdCdB Z"edDdE Z#e#jdFefdGdEZ#edHdI Z$e$jdJdI Z$dKefdLdMZ%dNdO Z&dPdQ Z'dRdS Z(dTdU Z)dVdW Z*dXdY Z+dZd[ Z,  Z-S )\PaddlePredictorOptionzPaddle Inference Engine Option)r   Zpaddle_fp32Zpaddle_fp16trt_fp32trt_fp16trt_int8r   Zmkldnn_bf16)gpur   ZnpuZxpuZmluZdcuZgcuZiluvatar_gpuc                    s$   t    i | _| jdi | d S )Nr   )super__init___cfg_init_option)selfkwargs	__class__r   r   r!   ;   s   
zPaddlePredictorOption.__init__c                 C   s,   t |  }t| j|_t| dr| j|_|S Ntrt_cfg_setting)typer   r"   hasattrr)   r$   objr   r   r   copy@   s
   

zPaddlePredictorOption.copyc                 K   sB   |  D ]\}}| |rt| || qt| d|   d S )Nz3 is not supported to set! The supported option is: )items_has_settersetattr	Exception_get_settable_attributes)r$   r%   kvr   r   r   r#   G   s   
z"PaddlePredictorOption._init_optionc                 C   s   |  | D ]\}}| j|| q| jdkr@dd l}|j r)|jj	 dks@t
r/tdtd d| _t|d| _d | _| jdv rat| }trSt| j |d< n	t| j |d d< || _d S d S )	Nr   r   zDevice fallback is disabled and the specified device (GPU) is not available. To fall back to CPU instead, unset the PADDLE_PDX_DISABLE_DEVICE_FALLBACK environment variable.zFThe specified device (GPU) is not available! Switching to CPU instead.r   )r   r   r   Zprecision_modeZenable_tensorrt_engine)_get_default_configr/   r"   
setdefaultr   r   deviceZis_compiled_with_cudacudaZdevice_countr   RuntimeErrorr   warningr   run_mode	device_idr   r   r   r)   )r$   r   r4   r5   r   r)   r   r   r   setdefault_by_model_nameP   s.   



z.PaddlePredictorOption.setdefault_by_model_namec                 C   s   | j du rtt \}}|du rdn|d }n| j | j}}i dt||d|d|dddg d	|tvr8d
nddddi dd
dd
dddddddddd
dd}|S )zget default configNr   r<   r   r=   cpu_threads
   delete_passenable_new_irTFenable_cinnr)   trt_use_dynamic_shapestrt_collect_shape_range_info#trt_discard_cached_shape_range_infotrt_dynamic_shapestrt_dynamic_shape_input_datatrt_shape_range_info_pathtrt_allow_rebuild_at_runtimemkldnn_cache_capacity)r   r   r   r=   r   r   )r$   r   r   
device_idsr=   Zdefault_configr   r   r   r6   p   sL   

	
z)PaddlePredictorOption._get_default_configc                 C   s   || j |< d S N)r"   )r$   r4   r5   r   r   r   _update   s   zPaddlePredictorOption._updatec                 C      | j dS )Nr<   r"   getr$   r   r   r   r<         zPaddlePredictorOption.run_moder<   c                 C   sX   || j vrd| j }td| dt| d|dr$t s$td| d| dS )	set run mode, z`run_mode` must be , but received .r   zMKL-DNN is not availabler<   N)SUPPORT_RUN_MODEjoin
ValueErrorrepr
startswithr   rN   )r$   r<   support_run_mode_strr   r   r   r<      s   
c                 C   rO   )Nr   rP   rR   r   r   r   r      rS   z!PaddlePredictorOption.device_typec                 C   s^   || j vrd| j }td| dt| d| d| t| |dv r-dtjd< d S d S )	NrU   zThe device type must be one of rV   rW   r   )r   r   1ZFLAGS_enable_pir_api)SUPPORT_DEVICErY   rZ   r[   rN   r	   osenviron)r$   r   r]   r   r   r   r      s   
c                 C   rO   Nr=   rP   rR   r   r   r   r=      rS   zPaddlePredictorOption.device_idc                 C      |  d| d S rb   rN   )r$   r=   r   r   r   r=         c                 C   rO   )Nr?   rP   rR   r   r   r   r?      rS   z!PaddlePredictorOption.cpu_threadsc                 C   s(   t |tr	|dk rt | d| dS )zset cpu threadsr   r?   N)
isinstanceintr2   rN   )r$   r?   r   r   r   r?      s   c                 C   rO   NrA   rP   rR   r   r   r   rA      rS   z!PaddlePredictorOption.delete_passc                 C   rc   rh   rd   )r$   rA   r   r   r   rA      re   c                 C   rO   )NrB   rP   rR   r   r   r   rB      rS   z#PaddlePredictorOption.enable_new_irrB   c                 C      |  d| dS )rT   rB   Nrd   )r$   rB   r   r   r   rB         c                 C   rO   )NrC   rP   rR   r   r   r   rC      rS   z!PaddlePredictorOption.enable_cinnrC   c                 C   ri   )rT   rC   Nrd   )r$   rC   r   r   r   rC      rj   c                 C   rO   r(   rP   rR   r   r   r   r)      rS   z%PaddlePredictorOption.trt_cfg_settingconfigc                 C   s6   t |ttdfsJ dt| d| d| dS )zset trt configNz7The trt_cfg_setting must be `dict` type, but received `z` type!r)   )rf   dictr*   rN   )r$   rk   r   r   r   r)      s   c                 C   rO   NrD   rP   rR   r   r   r   rD      rS   z,PaddlePredictorOption.trt_use_dynamic_shapesc                 C   rc   rm   rd   )r$   rD   r   r   r   rD      re   c                 C   rO   NrE   rP   rR   r   r   r   rE      rS   z2PaddlePredictorOption.trt_collect_shape_range_infoc                 C   rc   rn   rd   )r$   rE   r   r   r   rE      re   c                 C   rO   NrF   rP   rR   r   r   r   rF      rS   z9PaddlePredictorOption.trt_discard_cached_shape_range_infoc                 C   rc   ro   rd   )r$   rF   r   r   r   rF      s   c                 C   rO   NrG   rP   rR   r   r   r   rG     rS   z(PaddlePredictorOption.trt_dynamic_shapesrG   c                 C   s:   t |tsJ |D ]}t || tsJ q	| d| d S rp   )rf   rl   listrN   )r$   rG   Zinput_kr   r   r   rG   	  s   c                 C   rO   NrH   rP   rR   r   r   r   rH     rS   z2PaddlePredictorOption.trt_dynamic_shape_input_datarH   c                 C   rc   rr   rd   )r$   rH   r   r   r   rH     s   c                 C   rO   )NrI   rP   rR   r   r   r   rI     rS   z/PaddlePredictorOption.trt_shape_range_info_pathrI   c                 C   ri   )zset shape info filenamerI   Nrd   )r$   rI   r   r   r   rI     rj   c                 C   rO   NrJ   rP   rR   r   r   r   rJ   #  rS   z2PaddlePredictorOption.trt_allow_rebuild_at_runtimec                 C   rc   rs   rd   )r$   rJ   r   r   r   rJ   '  re   c                 C   rO   NrK   rP   rR   r   r   r   rK   +  rS   z+PaddlePredictorOption.mkldnn_cache_capacitycapacityc                 C   rc   rt   rd   )r$   ru   r   r   r   rK   /  re   c                 C      | j S rM   rI   rR   r   r   r   shape_info_filename5     z)PaddlePredictorOption.shape_info_filenamec                 C   s
   || _ d S rM   rw   )r$   rx   r   r   r   rx   9  s   
r8   c                 C   sb   |sdS t |\}}|| _|dur|d nd}|| _|du s$t|dkr/td| d dS dS )z
set deviceNr   r   zThe device ID has been set to rW   )r   r   r=   lenr   debug)r$   r8   r   rL   r=   r   r   r   
set_device=  s   z PaddlePredictorOption.set_devicec                 C   rv   )zget supported run mode)rX   rR   r   r   r   get_support_run_modeH  ry   z*PaddlePredictorOption.get_support_run_modec                 C   rv   )zget supported device)r_   rR   r   r   r   get_support_deviceL  ry   z(PaddlePredictorOption.get_support_devicec                 C   s   d dd | j D S )Nz,  c                 S   s   g | ]\}}| d | qS )z: r   ).0r4   r5   r   r   r   
<listcomp>Q  s    z1PaddlePredictorOption.__str__.<locals>.<listcomp>)rY   r"   r/   rR   r   r   r   __str__P  s   zPaddlePredictorOption.__str__c                 C   s,   || j vrtd| d| j  | j |S )Nz	The key (z) is not found in cfg: 
 )r"   r2   rQ   )r$   keyr   r   r   __getattr__S  s   
z!PaddlePredictorOption.__getattr__c                 C   s   t |tr|j| jkS dS )NF)rf   r   r"   r,   r   r   r   __eq__X  s   
zPaddlePredictorOption.__eq__c                 C   s"   t | j|d }t|to|jd uS rM   )getattrr'   rf   propertyfset)r$   attrpropr   r   r   r0   ]  s   z!PaddlePredictorOption._has_setterc                 C   s   dd t | j D S )Nc                 S   s(   g | ]\}}t |tr|jd ur|qS rM   )rf   r   r   )r   namer   r   r   r   r   b  s    
zBPaddlePredictorOption._get_settable_attributes.<locals>.<listcomp>)varsr'   r/   rR   r   r   r   r3   a  s   z.PaddlePredictorOption._get_settable_attributes).__name__
__module____qualname____doc__rX   r_   r!   r.   r#   r>   r6   rN   r   r<   setterstrr   r=   r?   rA   rB   boolrC   r)   r   rD   rE   rF   rG   r   rg   rH   floatrI   rJ   rK   rx   r|   r}   r~   r   r   r   r0   r3   __classcell__r   r   r&   r   r   +   s    
	 


















 






r   )r`   r.   r   typingr   r   utilsr   Zutils.devicer   r   r	   Zutils.flagsr
   r   r   miscr   Zmkldnn_blocklistr   Znew_ir_blocklistr   Z
trt_configr   r   r   objectr   r   r   r   r   <module>   s   