






"""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_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


UpperCAmelCase_   : List[str]   			= logging.get_logger(__name__)




class 							lowerCAmelCase__				(			UpperCAmelCase__ ):






      '''simple docstring'''

      __UpperCamelCase	 =  ["pixel_values"]
      def __init__(						self     :  Tuple						,	lowercase_     :  bool = True						,	lowercase_     :  Optional[Dict[str, int]] = None						,	lowercase_     :  PILImageResampling = PILImageResampling.BILINEAR						,	lowercase_     :  bool = True						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = True						,	lowercase_     :  Union[int, float] = 1 / 255						,	lowercase_     :  bool = True						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	**lowercase_     :  str						,	):




          '''simple docstring'''


          super().__init__(**lowercase_)
          SCREAMING_SNAKE_CASE_      :  Dict    						=		size if size is not None else {'''shortest_edge''': 256}
          SCREAMING_SNAKE_CASE_      :  str    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          SCREAMING_SNAKE_CASE_      :  Dict    						=		crop_size if crop_size is not None else {'''height''': 224, '''width''': 224}
          SCREAMING_SNAKE_CASE_      :  int    						=		get_size_dict(lowercase_)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		do_resize
          SCREAMING_SNAKE_CASE_      :  int    						=		size
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		resample
          SCREAMING_SNAKE_CASE_      :  int    						=		do_center_crop
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		crop_size
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		do_rescale
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		rescale_factor
          SCREAMING_SNAKE_CASE_      :  str    						=		do_normalize
          SCREAMING_SNAKE_CASE_      :  int    						=		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		image_std if image_std is not None else IMAGENET_STANDARD_STD
      def _SCREAMING_SNAKE_CASE (						self     :  str						,	lowercase_     :  np.ndarray						,	lowercase_     :  Dict[str, int]						,	lowercase_     :  PILImageResampling = PILImageResampling.BICUBIC						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  Optional[int]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          if "shortest_edge" not in size:
              raise ValueError(F'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}')
          SCREAMING_SNAKE_CASE_      :  int    						=		get_resize_output_image_size(lowercase_						,	size=size['''shortest_edge''']						,	default_to_square=lowercase_)
          return resize(lowercase_						,	size=lowercase_						,	resample=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Dict[str, int]						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  int						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		get_size_dict(lowercase_)
          return center_crop(lowercase_						,	size=(size['''height'''], size['''width'''])						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]						,	lowercase_     :  np.ndarray						,	lowercase_     :  float						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  List[str]):




          '''simple docstring'''


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




          '''simple docstring'''


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



      def _SCREAMING_SNAKE_CASE (						self     :  int						,	lowercase_     :  ImageInput						,	lowercase_     :  Optional[bool] = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = None						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  Optional[bool] = None						,	lowercase_     :  Optional[float] = None						,	lowercase_     :  Optional[bool] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[str, TensorType]] = None						,	lowercase_     :  Union[str, ChannelDimension] = ChannelDimension.FIRST						,	**lowercase_     :  Dict						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		do_resize if do_resize is not None else self.do_resize
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		size if size is not None else self.size
          SCREAMING_SNAKE_CASE_      :  int    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          SCREAMING_SNAKE_CASE_      :  Dict    						=		resample if resample is not None else self.resample
          SCREAMING_SNAKE_CASE_      :  Dict    						=		do_center_crop if do_center_crop is not None else self.do_center_crop
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		crop_size if crop_size is not None else self.crop_size
          SCREAMING_SNAKE_CASE_      :  Any    						=		get_size_dict(lowercase_)
          SCREAMING_SNAKE_CASE_      :  Dict    						=		do_rescale if do_rescale is not None else self.do_rescale
          SCREAMING_SNAKE_CASE_      :  Dict    						=		rescale_factor if rescale_factor is not None else self.rescale_factor
          SCREAMING_SNAKE_CASE_      :  str    						=		do_normalize if do_normalize is not None else self.do_normalize
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		image_mean if image_mean is not None else self.image_mean
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		image_std if image_std is not None else self.image_std

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		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.
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		[to_numpy_array(lowercase_) for image in images]

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

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

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

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

          SCREAMING_SNAKE_CASE_      :  Any    						=		[to_channel_dimension_format(lowercase_						,	lowercase_) for image in images]

          SCREAMING_SNAKE_CASE_      :  Dict    						=		{'''pixel_values''': images}
          return BatchFeature(data=lowercase_						,	tensor_type=lowercase_)



"""simple docstring"""


from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
    from .deepspeed import (
        DeepSpeedEngineWrapper,
        DeepSpeedOptimizerWrapper,
        DeepSpeedSchedulerWrapper,
        DummyOptim,
        DummyScheduler,
        HfDeepSpeedConfig,
    )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers
