o rq i@s`ddlZddlmZddlmZmZddlmZeZddZ eddd Z d d Z d d Z dS)N)contextmanager)Devices Frameworks) get_loggercCsd}|dur |dksJ||}|d}t|dks!J||dus'J|ddvs1J||d}d}t|dkrCt|d}|d krJtj}|tjkrU|durUd}||fS) a^ Verify device is valid, device should be either cpu, cuda, gpu, cuda:X or gpu:X. Args: device (str): device str, should be either cpu, cuda, gpu, gpu:X or cuda:X where X is the ordinal for gpu device. Return: device info (tuple): device_type and device_id, if device_id is not set, will use 0 as default. z^device should be either cpu, cuda, gpu, gpu:X or cuda:X where X is the ordinal for gpu device.N:r)cpucudagpur )lowersplitlenintrr ) device_nameerr_msgZeles device_type device_idr^/home/app/PaddleOCR-VL/.venv_paddleocr/lib/python3.10/site-packages/modelscope/utils/device.py verify_device s      rgpu:0ccs&t|\}}|tjkrgddl}|tjkr|jst dtj }|tj krB| d dVWddS1s;wYdS|tjkre| d| dVWddS1s^wYdSdS|tj krddl }|tjkr|j r|j d|nt ddVdSdVdS)a Device placement function, allow user to specify which device to place model or tensor Args: framework (str): tensorflow or pytorch. device (str): gpu or cpu to use, if you want to specify certain gpu, use gpu:$gpu_id or cuda:$gpu_id. Returns: Context manager Examples: >>> # Requests for using model on cuda:0 for gpu >>> with device_placement('pytorch', device='gpu:0'): >>> model = Model.from_pretrained(...) rNz5tensorflow: cuda is not available, using cpu instead.z/CPU:0z /device:gpu:cuda:z2pytorch: cuda is not available, using cpu instead.)rrtfZ tensorflowrr testZis_gpu_availableloggerdebugr devicetorchr is_availableZ set_device)Z frameworkrrrrrrrrdevice_placement's8    " "     r!cCsdddl}t|\}}d}|tjkrd}|jstdd}|r+|d|}|S|d}|S)z] create torch device Args: device_name (str): cpu, gpu, gpu:0, cuda:0 etc. rNFTz)cuda is not available, using cpu instead.rr ) rrrr r r rinfor)rrrrZuse_cudarrrr create_deviceUs     r#cCs^ddl}ddlm}|jr(|r%|r%dtjvr%dtjd}nd}nd}||S)Nr) distributedZ LOCAL_RANKrzcuda:0r )rr$r r Zis_initializedosenvironr)rdistrrrr get_devicels    r()r) r% contextlibrZmodelscope.utils.constantrrZmodelscope.utils.loggerrrrr!r#r(rrrrs   -