



from typing import Dict, List, Optional, Union

import numpy as np

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


A__								=		logging.get_logger(__name__)


if is_vision_available():
				import PIL


class  a				(     __lowercase							):
		__lowerCAmelCase      :       Any											=   ["""pixel_values"""]



		def __init__(	self	:Optional[Any]  ,__lowercase	:Optional[int] = True  ,__lowercase	:Optional[int] = None  ,__lowercase	:Optional[Any] = PILImageResampling.BICUBIC  ,__lowercase	:int = True  ,__lowercase	:Any = None  ,__lowercase	:List[str] = True  ,__lowercase	:Dict = 1 / 2_5_5  ,__lowercase	:Union[str, Any] = True  ,__lowercase	:List[Any] = None  ,__lowercase	:List[Any] = None  ,__lowercase	:Dict = True  ,**__lowercase	:List[str]  ,):
						super().__init__(**SCREAMING_SNAKE_CASE_    )
						snake_case__      : Any							=      size if size is not None else {'''shortest_edge''': 2_2_4}
						snake_case__      : Optional[int]							=      get_size_dict(SCREAMING_SNAKE_CASE_  ,default_to_square=SCREAMING_SNAKE_CASE_    )
						snake_case__      : int							=      crop_size if crop_size is not None else {'''height''': 2_2_4, '''width''': 2_2_4}
						snake_case__      : Tuple							=      get_size_dict(SCREAMING_SNAKE_CASE_  ,default_to_square=SCREAMING_SNAKE_CASE_  ,param_name='''crop_size'''    )

						snake_case__      : Optional[Any]							=      do_resize
						snake_case__      : Any							=      size
						snake_case__      : Union[str, Any]							=      resample
						snake_case__      : int							=      do_center_crop
						snake_case__      : List[str]							=      crop_size
						snake_case__      : Optional[Any]							=      do_rescale
						snake_case__      : List[Any]							=      rescale_factor
						snake_case__      : Any							=      do_normalize
						snake_case__      : Union[str, Any]							=      image_mean if image_mean is not None else OPENAI_CLIP_MEAN
						snake_case__      : Tuple							=      image_std if image_std is not None else OPENAI_CLIP_STD
						snake_case__      : Optional[Any]							=      do_convert_rgb



		def       __lowerCamelCase     (	self	:List[Any]  ,__lowercase	:Optional[int]  ,__lowercase	:Tuple  ,__lowercase	:Union[str, Any] = PILImageResampling.BICUBIC  ,__lowercase	:List[str] = None  ,**__lowercase	:Tuple  ,):
						snake_case__      : Any							=      get_size_dict(SCREAMING_SNAKE_CASE_  ,default_to_square=SCREAMING_SNAKE_CASE_    )
						if "shortest_edge" not in size:
										raise ValueError(F"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""    )
						snake_case__      : Tuple							=      get_resize_output_image_size(SCREAMING_SNAKE_CASE_  ,size=size['''shortest_edge''']  ,default_to_square=SCREAMING_SNAKE_CASE_    )
						return resize(SCREAMING_SNAKE_CASE_  ,size=SCREAMING_SNAKE_CASE_  ,resample=SCREAMING_SNAKE_CASE_  ,data_format=SCREAMING_SNAKE_CASE_  ,**SCREAMING_SNAKE_CASE_    )



		def       __lowerCamelCase     (	self	:Union[str, Any]  ,__lowercase	:List[str]  ,__lowercase	:int  ,__lowercase	:int = None  ,**__lowercase	:Optional[Any]  ,):
						snake_case__      : Optional[int]							=      get_size_dict(SCREAMING_SNAKE_CASE_    )
						if "height" not in size or "width" not in size:
										raise ValueError(F"""The `size` parameter must contain the keys (height, width). Got {size.keys()}"""    )
						return center_crop(SCREAMING_SNAKE_CASE_  ,size=(size['''height'''], size['''width'''])  ,data_format=SCREAMING_SNAKE_CASE_  ,**SCREAMING_SNAKE_CASE_    )



		def       __lowerCamelCase     (	self	:Union[str, Any]  ,__lowercase	:str  ,__lowercase	:Optional[Any]  ,__lowercase	:Any = None  ,**__lowercase	:List[str]  ,):
						return rescale(SCREAMING_SNAKE_CASE_  ,scale=SCREAMING_SNAKE_CASE_  ,data_format=SCREAMING_SNAKE_CASE_  ,**SCREAMING_SNAKE_CASE_    )



		def       __lowerCamelCase     (	self	:List[str]  ,__lowercase	:Optional[int]  ,__lowercase	:List[Any]  ,__lowercase	:Any  ,__lowercase	:Union[str, Any] = None  ,**__lowercase	:Optional[Any]  ,):
						return normalize(SCREAMING_SNAKE_CASE_  ,mean=SCREAMING_SNAKE_CASE_  ,std=SCREAMING_SNAKE_CASE_  ,data_format=SCREAMING_SNAKE_CASE_  ,**SCREAMING_SNAKE_CASE_    )






		def       __lowerCamelCase     (	self	:Any  ,__lowercase	:Optional[Any]  ,__lowercase	:Dict = None  ,__lowercase	:Dict = None  ,__lowercase	:Dict = None  ,__lowercase	:Optional[Any] = None  ,__lowercase	:List[str] = None  ,__lowercase	:List[str] = None  ,__lowercase	:int = None  ,__lowercase	:List[Any] = None  ,__lowercase	:List[Any] = None  ,__lowercase	:List[Any] = None  ,__lowercase	:str = None  ,__lowercase	:Optional[int] = None  ,__lowercase	:Dict = ChannelDimension.FIRST  ,**__lowercase	:Union[str, Any]  ,):
						snake_case__      : Union[str, Any]							=      do_resize if do_resize is not None else self.do_resize
						snake_case__      : Union[str, Any]							=      size if size is not None else self.size
						snake_case__      : int							=      get_size_dict(SCREAMING_SNAKE_CASE_  ,param_name='''size'''  ,default_to_square=SCREAMING_SNAKE_CASE_    )
						snake_case__      : List[Any]							=      resample if resample is not None else self.resample
						snake_case__      : Tuple							=      do_center_crop if do_center_crop is not None else self.do_center_crop
						snake_case__      : Optional[Any]							=      crop_size if crop_size is not None else self.crop_size
						snake_case__      : Union[str, Any]							=      get_size_dict(SCREAMING_SNAKE_CASE_  ,param_name='''crop_size'''  ,default_to_square=SCREAMING_SNAKE_CASE_    )
						snake_case__      : Optional[Any]							=      do_rescale if do_rescale is not None else self.do_rescale
						snake_case__      : Optional[int]							=      rescale_factor if rescale_factor is not None else self.rescale_factor
						snake_case__      : Tuple							=      do_normalize if do_normalize is not None else self.do_normalize
						snake_case__      : Tuple							=      image_mean if image_mean is not None else self.image_mean
						snake_case__      : Optional[Any]							=      image_std if image_std is not None else self.image_std
						snake_case__      : str							=      do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

						snake_case__      : str							=      make_list_of_images(SCREAMING_SNAKE_CASE_    )

						if not valid_images(SCREAMING_SNAKE_CASE_    ):
										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.'''    )

						# PIL RGBA images are converted to RGB
						if do_convert_rgb:
										snake_case__      : Tuple							=      [convert_to_rgb(SCREAMING_SNAKE_CASE_    ) for image in images]

						# All transformations expect numpy arrays.
						snake_case__      : Optional[int]							=      [to_numpy_array(SCREAMING_SNAKE_CASE_    ) for image in images]

						if do_resize:
										snake_case__      : Tuple							=      [self.resize(image=SCREAMING_SNAKE_CASE_  ,size=SCREAMING_SNAKE_CASE_  ,resample=SCREAMING_SNAKE_CASE_    ) for image in images]

						if do_center_crop:
										snake_case__      : List[str]							=      [self.center_crop(image=SCREAMING_SNAKE_CASE_  ,size=SCREAMING_SNAKE_CASE_    ) for image in images]

						if do_rescale:
										snake_case__      : List[str]							=      [self.rescale(image=SCREAMING_SNAKE_CASE_  ,scale=SCREAMING_SNAKE_CASE_    ) for image in images]

						if do_normalize:
										snake_case__      : List[str]							=      [self.normalize(image=SCREAMING_SNAKE_CASE_  ,mean=SCREAMING_SNAKE_CASE_  ,std=SCREAMING_SNAKE_CASE_    ) for image in images]

						snake_case__      : Dict							=      [to_channel_dimension_format(SCREAMING_SNAKE_CASE_  ,SCREAMING_SNAKE_CASE_    ) for image in images]

						snake_case__      : str							=      {'''pixel_values''': images}
						return BatchFeature(data=SCREAMING_SNAKE_CASE_  ,tensor_type=SCREAMING_SNAKE_CASE_    )



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


lowerCAmelCase__				=					get_logger(__name__)



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a=0     ):
		os.makedirs(__a,				exist_ok=__a     )
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				SCREAMING_SNAKE_CASE_    			=			model.state_dict()
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						if accelerator.process_index == 0:
								logger.info(F'Saving model to {output_model_file}'     )
								torch.save(__a,				__a     )
								logger.info(F'Model saved to {output_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
						SCREAMING_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'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Saving model to {output_model_file}'     )
						torch.save(__a,				__a     )
						logger.info(F'Model saved to {output_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				F'{MODEL_NAME}_{model_index}'     )
						os.makedirs(__a,				exist_ok=__a     )
						logger.info(F'Saving model to {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			{'''model''': state_dict}

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



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a=0     ):
		accelerator.wait_for_everyone()
		with FSDP.state_dict_type(
		    __a,				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(__a     ) != 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
						SCREAMING_SNAKE_CASE_    			=			F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading model from {input_model_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Model loaded from {input_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
						SCREAMING_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'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading model from {input_model_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Model loaded from {input_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    os.path.join(__a,				F'{MODEL_NAME}_{model_index}'     )
						    if F'{MODEL_NAME}' not in input_dir
						    else input_dir
						)
						logger.info(F'Loading model from {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			{'''model''': model.state_dict()}
						dist_cp.load_state_dict(
						    state_dict=__a,				storage_reader=dist_cp.FileSystemReader(__a     ),				planner=DefaultLoadPlanner(),				)
						SCREAMING_SNAKE_CASE_    			=			state_dict['''model''']
						logger.info(F'Model loaded from {ckpt_dir}'     )
				model.load_state_dict(__a     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a,				__a=0     ):
		os.makedirs(__a,				exist_ok=__a     )
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				SCREAMING_SNAKE_CASE_    			=			FSDP.optim_state_dict(__a,				__a     )
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						if accelerator.process_index == 0:
								SCREAMING_SNAKE_CASE_    			=			(
								    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
								)
								SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
								logger.info(F'Saving Optimizer state to {output_optimizer_file}'     )
								torch.save(__a,				__a     )
								logger.info(F'Optimizer state saved in {output_optimizer_file}'     )
				else:
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				F'{OPTIMIZER_NAME}_{optimizer_index}'     )
						os.makedirs(__a,				exist_ok=__a     )
						logger.info(F'Saving Optimizer state to {ckpt_dir}'     )
						dist_cp.save_state_dict(
						    state_dict={'''optimizer''': optim_state},				storage_writer=dist_cp.FileSystemWriter(__a     ),				planner=DefaultSavePlanner(),				)
						logger.info(F'Optimizer state saved in {ckpt_dir}'     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a,				__a=0     ):
		accelerator.wait_for_everyone()
		with FSDP.state_dict_type(
		    __a,				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:
						SCREAMING_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:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading Optimizer state from {input_optimizer_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Optimizer state loaded from {input_optimizer_file}'     )
				else:
						SCREAMING_SNAKE_CASE_    			=			(
						    os.path.join(__a,				F'{OPTIMIZER_NAME}_{optimizer_index}'     )
						    if F'{OPTIMIZER_NAME}' not in input_dir
						    else input_dir
						)
						logger.info(F'Loading Optimizer from {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			load_sharded_optimizer_state_dict(
						    model_state_dict=model.state_dict(),				optimizer_key='''optimizer''',				storage_reader=dist_cp.FileSystemReader(__a     ),				)
						SCREAMING_SNAKE_CASE_    			=			optim_state['''optimizer''']
						logger.info(F'Optimizer loaded from {ckpt_dir}'     )
				SCREAMING_SNAKE_CASE_    			=			FSDP.optim_state_dict_to_load(__a,				__a,				__a     )
				optimizer.load_state_dict(__a     )