
"""simple docstring"""


import unittest

from transformers import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING, is_vision_available, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
		from PIL import Image
else:







		class 					lowerCAmelCase_					:



								'''simple docstring'''


								@staticmethod
								def     _SCREAMING_SNAKE_CASE   (			*A_  :   Any	,**A_  :   int   )  -> List[Any]:
									pass







@is_pipeline_test
@require_vision
@require_torch
class 					lowerCAmelCase_					(   unittest.TestCase     ):



						'''simple docstring'''


						_lowerCamelCase:     Optional[Any]   =     MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING


						def     _SCREAMING_SNAKE_CASE   (			self  :   Dict	,A_  :   Tuple	,A_  :   Union[str, Any]	,A_  :   str   )  -> Dict:
							A     = pipeline(
							    'zero-shot-object-detection'	,model='hf-internal-testing/tiny-random-owlvit-object-detection'   )

							A     = [
							    {
							        'image': './tests/fixtures/tests_samples/COCO/000000039769.png',
							        'candidate_labels': ['cat', 'remote', 'couch'],
							    }
							]
							return object_detector, examples


						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]	,A_  :   Union[str, Any]	,A_  :   List[str]   )  -> List[Any]:
							A     = object_detector(examples[0]	,threshold=0.0   )

							A     = len(A_   )
							self.assertGreater(A_	,0   )
							self.assertEqual(
							    A_	,[
							        {
							            'score': ANY(A_   ),
							            'label': ANY(A_   ),
							            'box': {'xmin': ANY(A_   ), 'ymin': ANY(A_   ), 'xmax': ANY(A_   ), 'ymax': ANY(A_   )},
							        }
							        for i in range(A_   )
							    ]	,)


						@require_tf
						@unittest.skip('Zero Shot Object Detection not implemented in TF'   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> Dict:
							pass


						@require_torch
						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> Union[str, Any]:
							A     = pipeline(
							    'zero-shot-object-detection'	,model='hf-internal-testing/tiny-random-owlvit-object-detection'   )

							A     = object_detector(
							    './tests/fixtures/tests_samples/COCO/000000039769.png'	,candidate_labels=['cat', 'remote', 'couch']	,threshold=0.64	,)

							self.assertEqual(
							    nested_simplify(A_	,decimals=4   )	,[
							        {'score': 0.72_35, 'label': 'cat', 'box': {'xmin': 204, 'ymin': 167, 'xmax': 232, 'ymax': 190}},
							        {'score': 0.72_18, 'label': 'remote', 'box': {'xmin': 204, 'ymin': 167, 'xmax': 232, 'ymax': 190}},
							        {'score': 0.71_84, 'label': 'couch', 'box': {'xmin': 204, 'ymin': 167, 'xmax': 232, 'ymax': 190}},
							        {'score': 0.67_48, 'label': 'remote', 'box': {'xmin': 571, 'ymin': 83, 'xmax': 598, 'ymax': 103}},
							        {'score': 0.66_56, 'label': 'cat', 'box': {'xmin': 571, 'ymin': 83, 'xmax': 598, 'ymax': 103}},
							        {'score': 0.66_14, 'label': 'couch', 'box': {'xmin': 571, 'ymin': 83, 'xmax': 598, 'ymax': 103}},
							        {'score': 0.64_56, 'label': 'remote', 'box': {'xmin': 494, 'ymin': 105, 'xmax': 521, 'ymax': 127}},
							        {'score': 0.6_42, 'label': 'remote', 'box': {'xmin': 67, 'ymin': 274, 'xmax': 93, 'ymax': 297}},
							        {'score': 0.64_19, 'label': 'cat', 'box': {'xmin': 494, 'ymin': 105, 'xmax': 521, 'ymax': 127}},
							    ]	,)

							A     = object_detector(
							    [
							        {
							            'image': './tests/fixtures/tests_samples/COCO/000000039769.png',
							            'candidate_labels': ['cat', 'remote', 'couch'],
							        }
							    ]	,threshold=0.64	,)

							self.assertEqual(
							    nested_simplify(A_	,decimals=4   )	,[
							        [
							            {'score': 0.72_35, 'label': 'cat', 'box': {'xmin': 204, 'ymin': 167, 'xmax': 232, 'ymax': 190}},
							            {'score': 0.72_18, 'label': 'remote', 'box': {'xmin': 204, 'ymin': 167, 'xmax': 232, 'ymax': 190}},
							            {'score': 0.71_84, 'label': 'couch', 'box': {'xmin': 204, 'ymin': 167, 'xmax': 232, 'ymax': 190}},
							            {'score': 0.67_48, 'label': 'remote', 'box': {'xmin': 571, 'ymin': 83, 'xmax': 598, 'ymax': 103}},
							            {'score': 0.66_56, 'label': 'cat', 'box': {'xmin': 571, 'ymin': 83, 'xmax': 598, 'ymax': 103}},
							            {'score': 0.66_14, 'label': 'couch', 'box': {'xmin': 571, 'ymin': 83, 'xmax': 598, 'ymax': 103}},
							            {'score': 0.64_56, 'label': 'remote', 'box': {'xmin': 494, 'ymin': 105, 'xmax': 521, 'ymax': 127}},
							            {'score': 0.6_42, 'label': 'remote', 'box': {'xmin': 67, 'ymin': 274, 'xmax': 93, 'ymax': 297}},
							            {'score': 0.64_19, 'label': 'cat', 'box': {'xmin': 494, 'ymin': 105, 'xmax': 521, 'ymax': 127}},
							        ]
							    ]	,)


						@require_torch
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> Optional[int]:
							A     = pipeline('zero-shot-object-detection'   )

							A     = object_detector(
							    'http://images.cocodataset.org/val2017/000000039769.jpg'	,candidate_labels=['cat', 'remote', 'couch']	,)
							self.assertEqual(
							    nested_simplify(A_	,decimals=4   )	,[
							        {'score': 0.28_68, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}},
							        {'score': 0.2_77, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 177, 'ymax': 115}},
							        {'score': 0.25_37, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 315, 'ymax': 472}},
							        {'score': 0.14_74, 'label': 'remote', 'box': {'xmin': 335, 'ymin': 74, 'xmax': 371, 'ymax': 187}},
							        {'score': 0.12_08, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 642, 'ymax': 476}},
							    ]	,)

							A     = object_detector(
							    [
							        {
							            'image': 'http://images.cocodataset.org/val2017/000000039769.jpg',
							            'candidate_labels': ['cat', 'remote', 'couch'],
							        },
							        {
							            'image': 'http://images.cocodataset.org/val2017/000000039769.jpg',
							            'candidate_labels': ['cat', 'remote', 'couch'],
							        },
							    ]	,)
							self.assertEqual(
							    nested_simplify(A_	,decimals=4   )	,[
							        [
							            {'score': 0.28_68, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}},
							            {'score': 0.2_77, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 177, 'ymax': 115}},
							            {'score': 0.25_37, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 315, 'ymax': 472}},
							            {'score': 0.14_74, 'label': 'remote', 'box': {'xmin': 335, 'ymin': 74, 'xmax': 371, 'ymax': 187}},
							            {'score': 0.12_08, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 642, 'ymax': 476}},
							        ],
							        [
							            {'score': 0.28_68, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}},
							            {'score': 0.2_77, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 177, 'ymax': 115}},
							            {'score': 0.25_37, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 315, 'ymax': 472}},
							            {'score': 0.14_74, 'label': 'remote', 'box': {'xmin': 335, 'ymin': 74, 'xmax': 371, 'ymax': 187}},
							            {'score': 0.12_08, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 642, 'ymax': 476}},
							        ],
							    ]	,)


						@require_tf
						@unittest.skip('Zero Shot Object Detection not implemented in TF'   )
						def     _SCREAMING_SNAKE_CASE   (			self  :   int   )  -> List[str]:
							pass


						@require_torch
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> List[Any]:
							A     = 0.2
							A     = pipeline('zero-shot-object-detection'   )

							A     = object_detector(
							    'http://images.cocodataset.org/val2017/000000039769.jpg'	,candidate_labels=['cat', 'remote', 'couch']	,threshold=A_	,)
							self.assertEqual(
							    nested_simplify(A_	,decimals=4   )	,[
							        {'score': 0.28_68, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}},
							        {'score': 0.2_77, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 177, 'ymax': 115}},
							        {'score': 0.25_37, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 315, 'ymax': 472}},
							    ]	,)




						@require_torch
						@slow
						def     _SCREAMING_SNAKE_CASE   (			self  :   Tuple   )  -> Dict:
							A     = 2
							A     = pipeline('zero-shot-object-detection'   )

							A     = object_detector(
							    'http://images.cocodataset.org/val2017/000000039769.jpg'	,candidate_labels=['cat', 'remote', 'couch']	,top_k=A_	,)
							self.assertEqual(
							    nested_simplify(A_	,decimals=4   )	,[
							        {'score': 0.28_68, 'label': 'cat', 'box': {'xmin': 324, 'ymin': 20, 'xmax': 640, 'ymax': 373}},
							        {'score': 0.2_77, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 177, 'ymax': 115}},
							    ]	,)

"""simple docstring"""


import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version('''>=''', FSDP_PYTORCH_VERSION):
		import torch.distributed.checkpoint as dist_cp
		from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
		from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
		from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
		from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_lowercase				 =					get_logger(__name__)






def    _snake_case					(		snake_case__ :  Optional[Any]						,   snake_case__ :  List[Any]						,   snake_case__ :  int						,   snake_case__ :  int						,   snake_case__ :  str=0  ):
	os.makedirs(snake_case__						,   exist_ok=snake_case__  )
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		A     = model.state_dict()
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A     = F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
			A     = os.path.join(snake_case__						,   snake_case__  )
			if accelerator.process_index == 0:
				logger.info(F'Saving model to {output_model_file}'  )
				torch.save(snake_case__						,   snake_case__  )
				logger.info(F'Model saved to {output_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A     = (
			    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
			    if model_index == 0
			    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
			)
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Saving model to {output_model_file}'  )
			torch.save(snake_case__						,   snake_case__  )
			logger.info(F'Model saved to {output_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A     = os.path.join(snake_case__						,   F'{MODEL_NAME}_{model_index}'  )
			os.makedirs(snake_case__						,   exist_ok=snake_case__  )
			logger.info(F'Saving model to {ckpt_dir}'  )
			A     = {'model': state_dict}

			dist_cp.save_state_dict(
			    state_dict=snake_case__						,   storage_writer=dist_cp.FileSystemWriter(snake_case__  )						,   planner=DefaultSavePlanner()						,   )
			logger.info(F'Model saved to {ckpt_dir}'  )






def    _snake_case					(		snake_case__ :  int						,   snake_case__ :  List[str]						,   snake_case__ :  str						,   snake_case__ :  str						,   snake_case__ :  Any=0  ):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if type(snake_case__  ) != FSDP and accelerator.process_index != 0:
				if not fsdp_plugin.sync_module_states:
					raise ValueError(
					    'Set the `sync_module_states` flag to `True` so that model states are synced across processes when '
					    'initializing FSDP object'  )
				return
			A     = F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Loading model from {input_model_file}'  )
			A     = torch.load(snake_case__  )
			logger.info(F'Model loaded from {input_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A     = (
			    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
			    if model_index == 0
			    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
			)
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Loading model from {input_model_file}'  )
			A     = torch.load(snake_case__  )
			logger.info(F'Model loaded from {input_model_file}'  )
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A     = (
			    os.path.join(snake_case__						,   F'{MODEL_NAME}_{model_index}'  )
			    if F'{MODEL_NAME}' not in input_dir
			    else input_dir
			)
			logger.info(F'Loading model from {ckpt_dir}'  )
			A     = {'model': model.state_dict()}
			dist_cp.load_state_dict(
			    state_dict=snake_case__						,   storage_reader=dist_cp.FileSystemReader(snake_case__  )						,   planner=DefaultLoadPlanner()						,   )
			A     = state_dict['model']
			logger.info(F'Model loaded from {ckpt_dir}'  )
		model.load_state_dict(snake_case__  )






def    _snake_case					(		snake_case__ :  Tuple						,   snake_case__ :  Optional[int]						,   snake_case__ :  Optional[int]						,   snake_case__ :  Dict						,   snake_case__ :  List[str]						,   snake_case__ :  Any=0  ):
	os.makedirs(snake_case__						,   exist_ok=snake_case__  )
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		A     = FSDP.optim_state_dict(snake_case__						,   snake_case__  )
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if accelerator.process_index == 0:
				A     = (
				    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
				)
				A     = os.path.join(snake_case__						,   snake_case__  )
				logger.info(F'Saving Optimizer state to {output_optimizer_file}'  )
				torch.save(snake_case__						,   snake_case__  )
				logger.info(F'Optimizer state saved in {output_optimizer_file}'  )
		else:
			A     = os.path.join(snake_case__						,   F'{OPTIMIZER_NAME}_{optimizer_index}'  )
			os.makedirs(snake_case__						,   exist_ok=snake_case__  )
			logger.info(F'Saving Optimizer state to {ckpt_dir}'  )
			dist_cp.save_state_dict(
			    state_dict={'optimizer': optim_state}						,   storage_writer=dist_cp.FileSystemWriter(snake_case__  )						,   planner=DefaultSavePlanner()						,   )
			logger.info(F'Optimizer state saved in {ckpt_dir}'  )






def    _snake_case					(		snake_case__ :  str						,   snake_case__ :  Optional[int]						,   snake_case__ :  List[str]						,   snake_case__ :  List[str]						,   snake_case__ :  int						,   snake_case__ :  Optional[int]=0  ):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__						,   fsdp_plugin.state_dict_type						,   fsdp_plugin.state_dict_config						,   fsdp_plugin.optim_state_dict_config  ):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A     = None
			# below check should work but currently it isn't working (mostly opytorch issue),
			# in the meantime disabling it at the cost of excess memory usage
			# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
			A     = (
			    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
			)
			A     = os.path.join(snake_case__						,   snake_case__  )
			logger.info(F'Loading Optimizer state from {input_optimizer_file}'  )
			A     = torch.load(snake_case__  )
			logger.info(F'Optimizer state loaded from {input_optimizer_file}'  )
		else:
			A     = (
			    os.path.join(snake_case__						,   F'{OPTIMIZER_NAME}_{optimizer_index}'  )
			    if F'{OPTIMIZER_NAME}' not in input_dir
			    else input_dir
			)
			logger.info(F'Loading Optimizer from {ckpt_dir}'  )
			A     = load_sharded_optimizer_state_dict(
			    model_state_dict=model.state_dict()						,   optimizer_key='optimizer'						,   storage_reader=dist_cp.FileSystemReader(snake_case__  )						,   )
			A     = optim_state['optimizer']
			logger.info(F'Optimizer loaded from {ckpt_dir}'  )
		A     = FSDP.optim_state_dict_to_load(snake_case__						,   snake_case__						,   snake_case__  )
		optimizer.load_state_dict(snake_case__  )