

"""simple docstring"""




from typing import Dict, Iterable, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


_lowercase   :	Optional[int]												=						logging.get_logger(__name__)
class      _UpperCAmelCase   (							_lowerCAmelCase       ):
	a__		: str					=					["pixel_values"]



	def __init__(  self     :       Optional[Any]			,    _lowercase     :       bool = True			,    _lowercase     :       Dict[str, int] = None			,    _lowercase     :       PILImageResampling = PILImageResampling.BICUBIC			,    _lowercase     :       bool = True			,    _lowercase     :       Dict[str, int] = None			,    _lowercase     :       bool = True			,    _lowercase     :       Union[int, float] = 1 / 2_55			,    _lowercase     :       bool = True			,    _lowercase     :       Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_MEAN			,    _lowercase     :       Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_STD			,    **_lowercase     :       List[Any]			,    ):
		super().__init__(**_lowercase						)
		__UpperCAmelCase      		= size if size is not None else {'''shortest_edge''': 2_24}
		__UpperCAmelCase      		= get_size_dict(_lowercase			,    default_to_square=_lowercase						)
		__UpperCAmelCase      		= crop_size if crop_size is not None else {'''height''': 2_24, '''width''': 2_24}
		__UpperCAmelCase      		= get_size_dict(_lowercase			,    param_name='''crop_size'''						)

		__UpperCAmelCase      		= do_resize
		__UpperCAmelCase      		= size
		__UpperCAmelCase      		= resample
		__UpperCAmelCase      		= do_center_crop
		__UpperCAmelCase      		= crop_size
		__UpperCAmelCase      		= do_rescale
		__UpperCAmelCase      		= rescale_factor
		__UpperCAmelCase      		= do_normalize
		__UpperCAmelCase      		= image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
		__UpperCAmelCase      		= image_std if image_std is not None else IMAGENET_DEFAULT_STD



	def        a  (  self     :       Dict			,    _lowercase     :       np.ndarray			,    _lowercase     :       Dict[str, int]			,    _lowercase     :       PILImageResampling = PILImageResampling.BICUBIC			,    _lowercase     :       Optional[Union[str, ChannelDimension]] = None			,    **_lowercase     :       str			,    ):
		__UpperCAmelCase      		= get_size_dict(_lowercase			,    default_to_square=_lowercase						)
		# size_dict is a dict with either keys "height" and "width" or "shortest_edge"
		if "shortest_edge" in size:
			__UpperCAmelCase      		= int((2_56 / 2_24) * size['''shortest_edge''']						)
			__UpperCAmelCase      		= get_resize_output_image_size(_lowercase			,    size=_lowercase			,    default_to_square=_lowercase						)
			__UpperCAmelCase      		= {'''height''': output_size[0], '''width''': output_size[1]}
		if "height" not in size_dict or "width" not in size_dict:
			raise ValueError(
			    F'''Size dict must have keys \'height\' and \'width\' or \'shortest_edge\'. Got {size_dict.keys()}'''						)
		return resize(
		    _lowercase			,    size=(size_dict['''height'''], size_dict['''width'''])			,    resample=_lowercase			,    data_format=_lowercase			,    **_lowercase						)



	def        a  (  self     :       Union[str, Any]			,    _lowercase     :       np.ndarray			,    _lowercase     :       Dict[str, int]			,    _lowercase     :       Optional[Union[str, ChannelDimension]] = None			,    **_lowercase     :       Dict			,    ):
		__UpperCAmelCase      		= get_size_dict(_lowercase						)
		if "height" not in size or "width" not in size:
			raise ValueError(F'''Size dict must have keys \'height\' and \'width\'. Got {size.keys()}'''						)
		return center_crop(_lowercase			,    size=(size['''height'''], size['''width'''])			,    data_format=_lowercase			,    **_lowercase						)



	def        a  (  self     :       str			,    _lowercase     :       np.ndarray			,    _lowercase     :       Union[int, float]			,    _lowercase     :       Optional[Union[str, ChannelDimension]] = None			,    **_lowercase     :       Union[str, Any]			,    ):
		return rescale(_lowercase			,    scale=_lowercase			,    data_format=_lowercase			,    **_lowercase						)



	def        a  (  self     :       Union[str, Any]			,    _lowercase     :       np.ndarray			,    _lowercase     :       Union[float, List[float]]			,    _lowercase     :       Union[float, List[float]]			,    _lowercase     :       Optional[Union[str, ChannelDimension]] = None			,    **_lowercase     :       Dict			,    ):
		return normalize(_lowercase			,    mean=_lowercase			,    std=_lowercase			,    data_format=_lowercase			,    **_lowercase						)



	def        a  (  self     :       Any			,    _lowercase     :       ImageInput			,    _lowercase     :       Optional[bool] = None			,    _lowercase     :       Optional[Dict[str, int]] = None			,    _lowercase     :       PILImageResampling = None			,    _lowercase     :       Optional[bool] = None			,    _lowercase     :       Optional[Dict[str, int]] = None			,    _lowercase     :       Optional[bool] = None			,    _lowercase     :       Optional[float] = None			,    _lowercase     :       Optional[bool] = None			,    _lowercase     :       Optional[Union[float, Iterable[float]]] = None			,    _lowercase     :       Optional[Union[float, Iterable[float]]] = None			,    _lowercase     :       Optional[TensorType] = None			,    _lowercase     :       ChannelDimension = ChannelDimension.FIRST			,    **_lowercase     :       Optional[Any]			,    ):
		__UpperCAmelCase      		= do_resize if do_resize is not None else self.do_resize
		__UpperCAmelCase      		= resample if resample is not None else self.resample
		__UpperCAmelCase      		= do_center_crop if do_center_crop is not None else self.do_center_crop
		__UpperCAmelCase      		= do_rescale if do_rescale is not None else self.do_rescale
		__UpperCAmelCase      		= rescale_factor if rescale_factor is not None else self.rescale_factor
		__UpperCAmelCase      		= do_normalize if do_normalize is not None else self.do_normalize
		__UpperCAmelCase      		= image_mean if image_mean is not None else self.image_mean
		__UpperCAmelCase      		= image_std if image_std is not None else self.image_std

		__UpperCAmelCase      		= size if size is not None else self.size
		__UpperCAmelCase      		= get_size_dict(_lowercase			,    default_to_square=_lowercase						)
		__UpperCAmelCase      		= crop_size if crop_size is not None else self.crop_size
		__UpperCAmelCase      		= get_size_dict(_lowercase			,    param_name='''crop_size'''						)

		__UpperCAmelCase      		= make_list_of_images(_lowercase						)

		if not valid_images(_lowercase						):
			raise ValueError(
			    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
			    '''torch.Tensor, tf.Tensor or jax.ndarray.'''						)

		if do_resize and size is None:
			raise ValueError('''Size must be specified if do_resize is True.'''						)

		if do_center_crop and crop_size is None:
			raise ValueError('''Crop size must be specified if do_center_crop is True.'''						)

		if do_rescale and rescale_factor is None:
			raise ValueError('''Rescale factor must be specified if do_rescale is True.'''						)

		if do_normalize and (image_mean is None or image_std is None):
			raise ValueError('''Image mean and std must be specified if do_normalize is True.'''						)

		# All transformations expect numpy arrays.
		__UpperCAmelCase      		= [to_numpy_array(_lowercase						) for image in images]

		if do_resize:
			__UpperCAmelCase      		= [self.resize(_lowercase			,    _lowercase			,    _lowercase						) for image in images]

		if do_center_crop:
			__UpperCAmelCase      		= [self.center_crop(_lowercase			,    _lowercase						) for image in images]

		if do_rescale:
			__UpperCAmelCase      		= [self.rescale(_lowercase			,    _lowercase						) for image in images]

		if do_normalize:
			__UpperCAmelCase      		= [self.normalize(_lowercase			,    _lowercase			,    _lowercase						) for image in images]

		__UpperCAmelCase      		= [to_channel_dimension_format(_lowercase			,    _lowercase						) for image in images]

		__UpperCAmelCase      		= {'''pixel_values''': images}
		return BatchFeature(data=_lowercase			,    tensor_type=_lowercase						)


"""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__  )