"""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    	=get_logger(__name__)



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

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



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     int     ,			__SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,			__SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     str=0      ):
     accelerator.wait_for_everyone()
     with FSDP.state_dict_type(
         __SCREAMING_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(__SCREAMING_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
               lowercase_	:				Optional[int]													=	F'''{MODEL_NAME}.bin''' if model_index == 0 else F'''{MODEL_NAME}_{model_index}.bin'''
               lowercase_	:				Optional[int]													=	os.path.join(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
               logger.info(F'''Loading model from {input_model_file}'''      )
               lowercase_	:				Union[str, Any]													=	torch.load(__SCREAMING_SNAKE_CASE      )
               logger.info(F'''Model loaded from {input_model_file}'''      )
          elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
               lowercase_	:				Dict													=	(
                   F'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
                   if model_index == 0
                   else F'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
               )
               lowercase_	:				Union[str, Any]													=	os.path.join(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
               logger.info(F'''Loading model from {input_model_file}'''      )
               lowercase_	:				Optional[int]													=	torch.load(__SCREAMING_SNAKE_CASE      )
               logger.info(F'''Model loaded from {input_model_file}'''      )
          elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
               lowercase_	:				Optional[Any]													=	(
                   os.path.join(__SCREAMING_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}'''      )
               lowercase_	:				Any													=	{'model': model.state_dict()}
               dist_cp.load_state_dict(
                   state_dict=__SCREAMING_SNAKE_CASE     ,			storage_reader=dist_cp.FileSystemReader(__SCREAMING_SNAKE_CASE      )     ,			planner=DefaultLoadPlanner()     ,			)
               lowercase_	:				Dict													=	state_dict['model']
               logger.info(F'''Model loaded from {ckpt_dir}'''      )
          model.load_state_dict(__SCREAMING_SNAKE_CASE      )



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Optional[Any]     ,			__SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     List[str]     ,			__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,			__SCREAMING_SNAKE_CASE							:     Optional[Any]=0      ):
     os.makedirs(__SCREAMING_SNAKE_CASE     ,			exist_ok=__SCREAMING_SNAKE_CASE      )
     with FSDP.state_dict_type(
         __SCREAMING_SNAKE_CASE     ,			fsdp_plugin.state_dict_type     ,			fsdp_plugin.state_dict_config     ,			fsdp_plugin.optim_state_dict_config      ):
          lowercase_	:				Tuple													=	FSDP.optim_state_dict(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
          if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
               if accelerator.process_index == 0:
                    lowercase_	:				List[str]													=	(
                        F'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else F'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
                    )
                    lowercase_	:				str													=	os.path.join(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
                    logger.info(F'''Saving Optimizer state to {output_optimizer_file}'''      )
                    torch.save(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
                    logger.info(F'''Optimizer state saved in {output_optimizer_file}'''      )
          else:
               lowercase_	:				List[Any]													=	os.path.join(__SCREAMING_SNAKE_CASE     ,			F'''{OPTIMIZER_NAME}_{optimizer_index}'''      )
               os.makedirs(__SCREAMING_SNAKE_CASE     ,			exist_ok=__SCREAMING_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(__SCREAMING_SNAKE_CASE      )     ,			planner=DefaultSavePlanner()     ,			)
               logger.info(F'''Optimizer state saved in {ckpt_dir}'''      )







def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     Optional[int]     ,			__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,			__SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     Tuple=0      ):
     accelerator.wait_for_everyone()
     with FSDP.state_dict_type(
         __SCREAMING_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:
               lowercase_	:				Union[str, Any]													=	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:
               lowercase_	:				List[str]													=	(
                   F'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else F'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
               )
               lowercase_	:				Dict													=	os.path.join(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
               logger.info(F'''Loading Optimizer state from {input_optimizer_file}'''      )
               lowercase_	:				Tuple													=	torch.load(__SCREAMING_SNAKE_CASE      )
               logger.info(F'''Optimizer state loaded from {input_optimizer_file}'''      )
          else:
               lowercase_	:				Tuple													=	(
                   os.path.join(__SCREAMING_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}'''      )
               lowercase_	:				List[Any]													=	load_sharded_optimizer_state_dict(
                   model_state_dict=model.state_dict()     ,			optimizer_key='optimizer'     ,			storage_reader=dist_cp.FileSystemReader(__SCREAMING_SNAKE_CASE      )     ,			)
               lowercase_	:				List[str]													=	optim_state['optimizer']
               logger.info(F'''Optimizer loaded from {ckpt_dir}'''      )
          lowercase_	:				List[str]													=	FSDP.optim_state_dict_to_load(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
          optimizer.load_state_dict(__SCREAMING_SNAKE_CASE      )


"""simple docstring"""







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,
    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, is_vision_available, logging


if is_vision_available():
       import PIL


__SCREAMING_SNAKE_CASE    	=logging.get_logger(__name__)







class        UpperCamelCase		(			lowercase_  ):
   lowercase							=  ['pixel_values']
   def __init__(			self     ,__UpperCamelCase = True     ,__UpperCamelCase = None     ,__UpperCamelCase = 0.9     ,__UpperCamelCase = PILImageResampling.BICUBIC     ,__UpperCamelCase = True     ,__UpperCamelCase = None     ,__UpperCamelCase = 1 / 255     ,__UpperCamelCase = True     ,__UpperCamelCase = True     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,**__UpperCamelCase     ,)	->      None:


        '''simple docstring'''


        super().__init__(**__UpperCamelCase			)
        lowercase_	:				Optional[int]													=	size if size is not None else {'shortest_edge': 224}
        lowercase_	:				Union[str, Any]													=	get_size_dict(__UpperCamelCase     ,default_to_square=__UpperCamelCase			)
        lowercase_	:				Union[str, Any]													=	crop_size if crop_size is not None else {'height': 224, 'width': 224}
        lowercase_	:				Optional[int]													=	get_size_dict(__UpperCamelCase     ,param_name='crop_size'			)

        lowercase_	:				List[str]													=	do_resize
        lowercase_	:				List[Any]													=	size
        lowercase_	:				int													=	crop_pct
        lowercase_	:				Dict													=	resample
        lowercase_	:				List[str]													=	do_center_crop
        lowercase_	:				Union[str, Any]													=	crop_size
        lowercase_	:				List[Any]													=	do_rescale
        lowercase_	:				Tuple													=	rescale_factor
        lowercase_	:				Tuple													=	do_normalize
        lowercase_	:				Union[str, Any]													=	image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
        lowercase_	:				int													=	image_std if image_std is not None else IMAGENET_DEFAULT_STD
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase = None     ,__UpperCamelCase = PILImageResampling.BICUBIC     ,__UpperCamelCase = None     ,**__UpperCamelCase     ,)	->      np.ndarray:


        '''simple docstring'''


        lowercase_	:				Any													=	get_size_dict(__UpperCamelCase     ,default_to_square=__UpperCamelCase			)
        if "shortest_edge" not in size and ("height" not in size or "width" not in size):
             raise ValueError(f'''size must contain \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'''			)
        if crop_pct is not None:
             if "shortest_edge" in size:
                  lowercase_	:				Union[str, Any]													=	int(size['shortest_edge'] / crop_pct			)
             elif "height" in size and "width" in size:
                  if size["height"] == size["width"]:
                       lowercase_	:				Tuple													=	int(size['height'] / crop_pct			)
                  else:
                       lowercase_	:				Dict													=	(int(size['height'] / crop_pct			), int(size['width'] / crop_pct			))
             else:
                  raise ValueError('Invalid size for resize: {}'.format(__UpperCamelCase			)			)

             lowercase_	:				int													=	get_resize_output_image_size(__UpperCamelCase     ,size=__UpperCamelCase     ,default_to_square=__UpperCamelCase			)
        else:
             if "shortest_edge" in size:
                  lowercase_	:				Optional[int]													=	get_resize_output_image_size(__UpperCamelCase     ,size=size['shortest_edge']     ,default_to_square=__UpperCamelCase			)
             elif "height" in size and "width" in size:
                  lowercase_	:				Dict													=	(size['height'], size['width'])
             else:
                  raise ValueError('Invalid size for resize: {}'.format(__UpperCamelCase			)			)

        return resize(__UpperCamelCase     ,size=__UpperCamelCase     ,resample=__UpperCamelCase     ,data_format=__UpperCamelCase     ,**__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase = None     ,**__UpperCamelCase     ,)	->      np.ndarray:


        '''simple docstring'''


        lowercase_	:				List[Any]													=	get_size_dict(__UpperCamelCase			)
        if "height" not in size or "width" not in size:
             raise ValueError(f'''size must contain \'height\' and \'width\' as keys. Got {size.keys()}'''			)
        return center_crop(__UpperCamelCase     ,size=(size['height'], size['width'])     ,data_format=__UpperCamelCase     ,**__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase = None     ,**__UpperCamelCase     ,)	->      str:


        '''simple docstring'''


        return rescale(__UpperCamelCase     ,scale=__UpperCamelCase     ,data_format=__UpperCamelCase     ,**__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase = None     ,**__UpperCamelCase     ,)	->      np.ndarray:


        '''simple docstring'''


        return normalize(__UpperCamelCase     ,mean=__UpperCamelCase     ,std=__UpperCamelCase     ,data_format=__UpperCamelCase     ,**__UpperCamelCase			)







   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = None     ,__UpperCamelCase = ChannelDimension.FIRST     ,**__UpperCamelCase     ,)	->      PIL.Image.Image:


        '''simple docstring'''


        lowercase_	:				List[Any]													=	do_resize if do_resize is not None else self.do_resize
        lowercase_	:				Optional[int]													=	crop_pct if crop_pct is not None else self.crop_pct
        lowercase_	:				List[str]													=	resample if resample is not None else self.resample
        lowercase_	:				str													=	do_center_crop if do_center_crop is not None else self.do_center_crop
        lowercase_	:				Tuple													=	do_rescale if do_rescale is not None else self.do_rescale
        lowercase_	:				Union[str, Any]													=	rescale_factor if rescale_factor is not None else self.rescale_factor
        lowercase_	:				str													=	do_normalize if do_normalize is not None else self.do_normalize
        lowercase_	:				str													=	image_mean if image_mean is not None else self.image_mean
        lowercase_	:				Tuple													=	image_std if image_std is not None else self.image_std

        lowercase_	:				Optional[Any]													=	size if size is not None else self.size
        lowercase_	:				Tuple													=	get_size_dict(__UpperCamelCase     ,default_to_square=__UpperCamelCase			)
        lowercase_	:				Union[str, Any]													=	crop_size if crop_size is not None else self.crop_size
        lowercase_	:				List[str]													=	get_size_dict(__UpperCamelCase     ,param_name='crop_size'			)

        lowercase_	:				str													=	make_list_of_images(__UpperCamelCase			)

        if not valid_images(__UpperCamelCase			):
             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 or resample is None:
             raise ValueError('Size and resample must be specified if do_resize is True.'			)

        if do_center_crop and crop_pct is None:
             raise ValueError('Crop_pct 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.
        lowercase_	:				Optional[Any]													=	[to_numpy_array(__UpperCamelCase			) for image in images]

        if do_resize:
             lowercase_	:				str													=	[self.resize(image=__UpperCamelCase     ,size=__UpperCamelCase     ,crop_pct=__UpperCamelCase     ,resample=__UpperCamelCase			) for image in images]

        if do_center_crop:
             lowercase_	:				str													=	[self.center_crop(image=__UpperCamelCase     ,size=__UpperCamelCase			) for image in images]

        if do_rescale:
             lowercase_	:				Any													=	[self.rescale(image=__UpperCamelCase     ,scale=__UpperCamelCase			) for image in images]

        if do_normalize:
             lowercase_	:				int													=	[self.normalize(image=__UpperCamelCase     ,mean=__UpperCamelCase     ,std=__UpperCamelCase			) for image in images]

        lowercase_	:				Dict													=	[to_channel_dimension_format(__UpperCamelCase     ,__UpperCamelCase			) for image in images]

        lowercase_	:				Any													=	{'pixel_values': images}
        return BatchFeature(data=__UpperCamelCase     ,tensor_type=__UpperCamelCase			)

