




'''simple docstring'''





from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_torch_available():
		import torch

		from transformers.modeling_outputs import BaseModelOutput

		from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

_SCREAMING_SNAKE_CASE							:       List[Any]						=			logging.get_logger(__name__)




@add_end_docstrings(lowercase_							)
class 			_snake_case		(     lowercase_							):




			def __init__( self    ,							**a__							)    ->    Any:

										'''simple docstring'''




										super().__init__(**a__							)

										if self.framework == "tf":
																	raise ValueError(F'The {self.__class__} is only available in PyTorch.'							)

										requires_backends(self    ,							"vision"							)
										self.check_model_type(a__							)




			def __call__( self    ,							a__    ,							a__ = None    ,							**a__    ,							)    ->    Optional[int]:

										'''simple docstring'''




										if "text_queries" in kwargs:
																	snake_case_      =				kwargs.pop("text_queries"							)

										if isinstance(a__    ,							(str, Image.Image)							):
																	snake_case_      =				{"image": image, "candidate_labels": candidate_labels}
										else:
																	snake_case_      =				image
										snake_case_      =				super().__call__(a__    ,							**a__							)
										return results




			def        lowerCAmelCase__    ( self    ,							**a__							)    ->    int:

										'''simple docstring'''




										snake_case_      =				{}
										if "threshold" in kwargs:
																	snake_case_      =				kwargs["threshold"]
										if "top_k" in kwargs:
																	snake_case_      =				kwargs["top_k"]
										return {}, {}, postprocess_params




			def        lowerCAmelCase__    ( self    ,							a__							)    ->    int:

										'''simple docstring'''




										snake_case_      =				load_image(inputs["image"]							)
										snake_case_      =				inputs["candidate_labels"]
										if isinstance(a__    ,							a__							):
																	snake_case_      =				candidate_labels.split(","							)

										snake_case_      =				torch.tensor([[image.height, image.width]]    ,							dtype=torch.intaa							)
										for i, candidate_label in enumerate(a__							):
																	snake_case_      =				self.tokenizer(a__    ,							return_tensors=self.framework							)
																	snake_case_      =				self.image_processor(a__    ,							return_tensors=self.framework							)
																	yield {
																	    "is_last": i == len(a__							) - 1,
																	    "target_size": target_size,
																	    "candidate_label": candidate_label,
																	    **text_inputs,
																	    **image_features,
																	}




			def        lowerCAmelCase__    ( self    ,							a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										snake_case_      =				model_inputs.pop("target_size"							)
										snake_case_      =				model_inputs.pop("candidate_label"							)
										snake_case_      =				model_inputs.pop("is_last"							)

										snake_case_      =				self.model(**a__							)

										snake_case_      =				{"target_size": target_size, "candidate_label": candidate_label, "is_last": is_last, **outputs}
										return model_outputs




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__=0.1    ,							a__=None							)    ->    Tuple:

										'''simple docstring'''




										snake_case_      =				[]
										for model_output in model_outputs:
																	snake_case_      =				model_output["candidate_label"]
																	snake_case_      =				BaseModelOutput(a__							)
																	snake_case_      =				self.image_processor.post_process_object_detection(
																	    outputs=a__    ,							threshold=a__    ,							target_sizes=model_output["target_size"]							)[0]

																	for index in outputs["scores"].nonzero():
																								snake_case_      =				outputs["scores"][index].item()
																								snake_case_      =				self._get_bounding_box(outputs["boxes"][index][0]							)

																								snake_case_      =				{"score": score, "label": label, "box": box}
																								results.append(a__							)

										snake_case_      =				sorted(a__    ,							key=lambda a__							: x["score"]    ,							reverse=a__							)
										if top_k:
																	snake_case_      =				results[:top_k]

										return results




			def        lowerCAmelCase__    ( self    ,							a__							)    ->    Dict[str, int]:

										'''simple docstring'''




										if self.framework != "pt":
																	raise ValueError("The ZeroShotObjectDetectionPipeline is only available in PyTorch."							)
										snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      =				box.int().tolist()
										snake_case_      =				{
										    "xmin": xmin,
										    "ymin": ymin,
										    "xmax": xmax,
										    "ymax": ymax,
										}
										return bbox










'''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


_SCREAMING_SNAKE_CASE							:       Tuple						=			get_logger(__name__)


def 					UpperCamelCase_(     snake_case							:	Optional[Any]							,    snake_case							:	Optional[Any]							,    snake_case							:	Optional[Any]							,    snake_case							:	int							,    snake_case							:	List[Any]=0					):

							'''simple docstring'''






							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					):
														snake_case_      =				model.state_dict()
														if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
																					snake_case_      =				f'{MODEL_NAME}.bin' if model_index == 0 else f'{MODEL_NAME}_{model_index}.bin'
																					snake_case_      =				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:
																					snake_case_      =				(
																					    f'{MODEL_NAME}_rank{accelerator.process_index}.bin'
																					    if model_index == 0
																					    else f'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
																					)
																					snake_case_      =				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:
																					snake_case_      =				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}'					)
																					snake_case_      =				{"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 					UpperCamelCase_(     snake_case							:	Optional[int]							,    snake_case							:	Optional[Any]							,    snake_case							:	int							,    snake_case							:	Union[str, Any]							,    snake_case							:	Any=0					):

							'''simple docstring'''






							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
																					snake_case_      =				f'{MODEL_NAME}.bin' if model_index == 0 else f'{MODEL_NAME}_{model_index}.bin'
																					snake_case_      =				os.path.join(snake_case							,    snake_case					)
																					logger.info(f'Loading model from {input_model_file}'					)
																					snake_case_      =				torch.load(snake_case					)
																					logger.info(f'Model loaded from {input_model_file}'					)
														elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
																					snake_case_      =				(
																					    f'{MODEL_NAME}_rank{accelerator.process_index}.bin'
																					    if model_index == 0
																					    else f'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
																					)
																					snake_case_      =				os.path.join(snake_case							,    snake_case					)
																					logger.info(f'Loading model from {input_model_file}'					)
																					snake_case_      =				torch.load(snake_case					)
																					logger.info(f'Model loaded from {input_model_file}'					)
														elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
																					snake_case_      =				(
																					    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}'					)
																					snake_case_      =				{"model": model.state_dict()}
																					dist_cp.load_state_dict(
																					    state_dict=snake_case							,    storage_reader=dist_cp.FileSystemReader(snake_case					)							,    planner=DefaultLoadPlanner()							,    )
																					snake_case_      =				state_dict["model"]
																					logger.info(f'Model loaded from {ckpt_dir}'					)
														model.load_state_dict(snake_case					)


def 					UpperCamelCase_(     snake_case							:	str							,    snake_case							:	List[str]							,    snake_case							:	Any							,    snake_case							:	Tuple							,    snake_case							:	Optional[Any]							,    snake_case							:	Tuple=0					):

							'''simple docstring'''






							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					):
														snake_case_      =				FSDP.optim_state_dict(snake_case							,    snake_case					)
														if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
																					if accelerator.process_index == 0:
																												snake_case_      =				(
																												    f'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else f'{OPTIMIZER_NAME}_{optimizer_index}.bin'
																												)
																												snake_case_      =				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:
																					snake_case_      =				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 					UpperCamelCase_(     snake_case							:	Optional[Any]							,    snake_case							:	List[str]							,    snake_case							:	Union[str, Any]							,    snake_case							:	int							,    snake_case							:	Optional[int]							,    snake_case							:	Union[str, Any]=0					):

							'''simple docstring'''






							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:
																					snake_case_      =				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:
																					snake_case_      =				(
																					    f'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else f'{OPTIMIZER_NAME}_{optimizer_index}.bin'
																					)
																					snake_case_      =				os.path.join(snake_case							,    snake_case					)
																					logger.info(f'Loading Optimizer state from {input_optimizer_file}'					)
																					snake_case_      =				torch.load(snake_case					)
																					logger.info(f'Optimizer state loaded from {input_optimizer_file}'					)
														else:
																					snake_case_      =				(
																					    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}'					)
																					snake_case_      =				load_sharded_optimizer_state_dict(
																					    model_state_dict=model.state_dict()							,    optimizer_key="optimizer"							,    storage_reader=dist_cp.FileSystemReader(snake_case					)							,    )
																					snake_case_      =				optim_state["optimizer"]
																					logger.info(f'Optimizer loaded from {ckpt_dir}'					)
														snake_case_      =				FSDP.optim_state_dict_to_load(snake_case							,    snake_case							,    snake_case					)
														optimizer.load_state_dict(snake_case					)




