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


lowerCAmelCase      :				Any     			=logging.get_logger(__name__)




class 						__snake_case						(			__lowerCAmelCase		):





      '''simple docstring'''






      _snake_case     		=							['pixel_values']



      def __init__(       self  :					str     ,	_UpperCamelCase  :					bool = True     ,	_UpperCamelCase  :					Optional[Dict[str, int]] = None     ,	_UpperCamelCase  :					PILImageResampling = PILImageResampling.BILINEAR     ,	_UpperCamelCase  :					bool = True     ,	_UpperCamelCase  :					Union[int, float] = 1 / 255     ,	_UpperCamelCase  :					bool = True     ,	_UpperCamelCase  :					Optional[Union[float, List[float]]] = None     ,	_UpperCamelCase  :					Optional[Union[float, List[float]]] = None     ,	**_UpperCamelCase  :					Any     ,	)					->None:


          """simple docstring"""


          super().__init__(**_UpperCamelCase)
          _lowerCamelCase  :		List[Any]			    =    size if size is not None else {"""height""": 224, """width""": 224}
          _lowerCamelCase  :		List[str]			    =    get_size_dict(_UpperCamelCase)
          _lowerCamelCase  :		Optional[int]			    =    do_resize
          _lowerCamelCase  :		Optional[Any]			    =    do_rescale
          _lowerCamelCase  :		Optional[int]			    =    do_normalize
          _lowerCamelCase  :		Optional[Any]			    =    size
          _lowerCamelCase  :		List[Any]			    =    resample
          _lowerCamelCase  :		Tuple			    =    rescale_factor
          _lowerCamelCase  :		Dict			    =    image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
          _lowerCamelCase  :		Dict			    =    image_std if image_std is not None else IMAGENET_STANDARD_STD



      def 						_SCREAMING_SNAKE_CASE     (       self  :					Any     ,	_UpperCamelCase  :					np.ndarray     ,	_UpperCamelCase  :					Dict[str, int]     ,	_UpperCamelCase  :					PILImageResampling = PILImageResampling.BILINEAR     ,	_UpperCamelCase  :					Optional[Union[str, ChannelDimension]] = None     ,	**_UpperCamelCase  :					Dict     ,	)					->np.ndarray:


          """simple docstring"""


          _lowerCamelCase  :		Any			    =    get_size_dict(_UpperCamelCase)
          if "height" not in size or "width" not in size:
              raise ValueError(F"""The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}""")
          return resize(
              _UpperCamelCase     ,	size=(size["""height"""], size["""width"""])     ,	resample=_UpperCamelCase     ,	data_format=_UpperCamelCase     ,	**_UpperCamelCase)



      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[Any]     ,	_UpperCamelCase  :					np.ndarray     ,	_UpperCamelCase  :					float     ,	_UpperCamelCase  :					Optional[Union[str, ChannelDimension]] = None     ,	**_UpperCamelCase  :					Optional[Any])					->np.ndarray:


          """simple docstring"""


          return rescale(_UpperCamelCase     ,	scale=_UpperCamelCase     ,	data_format=_UpperCamelCase     ,	**_UpperCamelCase)



      def 						_SCREAMING_SNAKE_CASE     (       self  :					str     ,	_UpperCamelCase  :					np.ndarray     ,	_UpperCamelCase  :					Union[float, List[float]]     ,	_UpperCamelCase  :					Union[float, List[float]]     ,	_UpperCamelCase  :					Optional[Union[str, ChannelDimension]] = None     ,	**_UpperCamelCase  :					int     ,	)					->np.ndarray:


          """simple docstring"""


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




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


          """simple docstring"""


          _lowerCamelCase  :		Tuple			    =    do_resize if do_resize is not None else self.do_resize
          _lowerCamelCase  :		Any			    =    do_rescale if do_rescale is not None else self.do_rescale
          _lowerCamelCase  :		int			    =    do_normalize if do_normalize is not None else self.do_normalize
          _lowerCamelCase  :		List[Any]			    =    resample if resample is not None else self.resample
          _lowerCamelCase  :		List[Any]			    =    rescale_factor if rescale_factor is not None else self.rescale_factor
          _lowerCamelCase  :		str			    =    image_mean if image_mean is not None else self.image_mean
          _lowerCamelCase  :		Any			    =    image_std if image_std is not None else self.image_std

          _lowerCamelCase  :		Any			    =    size if size is not None else self.size
          _lowerCamelCase  :		Optional[int]			    =    get_size_dict(_UpperCamelCase)

          _lowerCamelCase  :		List[Any]			    =    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:
              raise ValueError("""Size must be specified if do_resize is True.""")

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

          # All transformations expect numpy arrays.
          _lowerCamelCase  :		Any			    =    [to_numpy_array(_UpperCamelCase) for image in images]

          if do_resize:
              _lowerCamelCase  :		str			    =    [self.resize(image=_UpperCamelCase     ,	size=_UpperCamelCase     ,	resample=_UpperCamelCase) for image in images]

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

          if do_normalize:
              _lowerCamelCase  :		Any			    =    [self.normalize(image=_UpperCamelCase     ,	mean=_UpperCamelCase     ,	std=_UpperCamelCase) for image in images]

          _lowerCamelCase  :		Tuple			    =    [to_channel_dimension_format(_UpperCamelCase     ,	_UpperCamelCase) for image in images]

          _lowerCamelCase  :		List[Any]			    =    {"""pixel_values""": images}
          return BatchFeature(data=_UpperCamelCase     ,	tensor_type=_UpperCamelCase)





import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available







def     A__	(       __A ,			__A ,			__A ,			__A=None						):







    '''simple docstring'''
    _lowerCamelCase  :		List[Any]			    =    None
    # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
    if str(weight.dtype						) == "torch.bfloat16":
        # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
        _lowerCamelCase  :		int			    =    weight.view(torch.intaa						)
        _lowerCamelCase  :		List[Any]			    =    """bfloat16"""
    _lowerCamelCase  :		Tuple			    =    weight.cpu().numpy()
    _lowerCamelCase  :		Tuple			    =    os.path.join(__A ,			F"""{weight_name}.dat"""						)
    if index is not None:
        if dtype is None:
            _lowerCamelCase  :		Union[str, Any]			    =    str(array.dtype						)
        _lowerCamelCase  :		Optional[Any]			    =    {"""dtype""": dtype, """shape""": list(array.shape						)}
    if array.ndim == 0:
        _lowerCamelCase  :		int			    =    array[None]
    _lowerCamelCase  :		str			    =    np.memmap(__A ,			dtype=array.dtype ,			mode="""w+""" ,			shape=array.shape						)
    _lowerCamelCase  :		Union[str, Any]			    =    array[:]
    file_array.flush()
    return index







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    _lowerCamelCase  :		Tuple			    =    tuple(weight_info["""shape"""]						)
    if shape == ():
        # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
        _lowerCamelCase  :		Union[str, Any]			    =    (1,)

    _lowerCamelCase  :		str			    =    weight_info["""dtype"""]
    if dtype == "bfloat16":
        # NumPy does not support bfloat16 so this was saved as a int16
        _lowerCamelCase  :		Optional[int]			    =    """int16"""

    _lowerCamelCase  :		str			    =    np.memmap(__A ,			dtype=__A ,			shape=__A ,			mode="""r"""						)

    if len(weight_info["""shape"""]						) == 0:
        _lowerCamelCase  :		List[str]			    =    weight[0]
    _lowerCamelCase  :		List[str]			    =    torch.tensor(__A						)
    if weight_info["dtype"] == "bfloat16":
        _lowerCamelCase  :		List[Any]			    =    weight.view(torch.bfloataa						)

    return weight







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    if index is None or len(__A						) == 0:
        # Nothing to save
        return

    _lowerCamelCase  :		List[Any]			    =    os.path.join(__A ,			"""index.json"""						)
    if os.path.isfile(__A						):
        with open(__A ,			"""r""" ,			encoding="""utf-8"""						) as f:
            _lowerCamelCase  :		List[str]			    =    json.load(__A						)
    else:
        _lowerCamelCase  :		List[str]			    =    {}
    current_index.update(__A						)

    with open(__A ,			"""w""" ,			encoding="""utf-8"""						) as f:
        json.dump(__A ,			__A ,			indent=2						)







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    os.makedirs(__A ,			exist_ok=__A						)
    _lowerCamelCase  :		Dict			    =    {}
    for name, parameter in state_dict.items():
        _lowerCamelCase  :		Union[str, Any]			    =    offload_weight(__A ,			__A ,			__A ,			index=__A						)

    # Update index
    save_offload_index(__A ,			__A						)




class 						__snake_case						(			__lowerCAmelCase		):





      '''simple docstring'''






      def __init__(       self  :					str     ,	_UpperCamelCase  :					Mapping     ,	_UpperCamelCase  :					str)					->List[Any]:


          """simple docstring"""


          _lowerCamelCase  :		Union[str, Any]			    =    dataset
          _lowerCamelCase  :		Tuple			    =    prefix



      def __getitem__(       self  :					List[Any]     ,	_UpperCamelCase  :					Union[str, Any])					->Union[str, Any]:


          """simple docstring"""


          return self.dataset[F"""{self.prefix}{key}"""]



      def __iter__(       self  :					List[str])					->Any:


          """simple docstring"""


          return iter([key for key in self.dataset if key.startswith(self.prefix)])




      def __len__(       self  :					Any)					->int:


          """simple docstring"""


          return len(self.dataset)





class 						__snake_case						(			__lowerCAmelCase		):





      '''simple docstring'''






      def __init__(       self  :					Union[str, Any]     ,	_UpperCamelCase  :					Dict[str, torch.Tensor] = None     ,	_UpperCamelCase  :					Optional[Union[str, os.PathLike]] = None     ,	_UpperCamelCase  :					Mapping = None     ,	_UpperCamelCase  :					Union[str, Any]=None     ,	)					->Any:


          """simple docstring"""


          if state_dict is None and save_folder is None:
              raise ValueError("""Need either a `state_dict` or a `save_folder` containing offloaded weights.""")

          _lowerCamelCase  :		int			    =    {} if state_dict is None else state_dict
          _lowerCamelCase  :		Optional[int]			    =    save_folder
          if index is None and save_folder is not None:
              with open(os.path.join(_UpperCamelCase     ,	"""index.json""")) as f:
                  _lowerCamelCase  :		Union[str, Any]			    =    json.load(_UpperCamelCase)
          _lowerCamelCase  :		str			    =    {} if index is None else index
          _lowerCamelCase  :		Optional[int]			    =    list(self.state_dict.keys())
          self.all_keys.extend([key for key in self.index if key not in self.all_keys])
          _lowerCamelCase  :		str			    =    device



      def __getitem__(       self  :					Union[str, Any]     ,	_UpperCamelCase  :					str)					->Tuple:


          """simple docstring"""


          if key in self.state_dict:
              return self.state_dict[key]
          _lowerCamelCase  :		int			    =    self.index[key]
          if weight_info.get("""safetensors_file""") is not None:
              if not is_safetensors_available():
                  raise ImportError("""These offloaded weights require the use of safetensors: `pip install safetensors`.""")

              from safetensors import safe_open

              _lowerCamelCase  :		Any			    =    """cpu""" if self.device is None else self.device
              with safe_open(weight_info["""safetensors_file"""]     ,	framework="""pt"""     ,	device=_UpperCamelCase) as f:
                  _lowerCamelCase  :		Tuple			    =    f.get_tensor(weight_info.get("""weight_name"""     ,	_UpperCamelCase))

              if "dtype" in weight_info:
                  return tensor.to(getattr(_UpperCamelCase     ,	weight_info["""dtype"""]))
              else:
                  return tensor

          _lowerCamelCase  :		Optional[Any]			    =    os.path.join(self.save_folder     ,	F"""{key}.dat""")
          return load_offloaded_weight(_UpperCamelCase     ,	_UpperCamelCase)



      def __iter__(       self  :					Dict)					->List[Any]:


          """simple docstring"""


          return iter(self.all_keys)




      def __len__(       self  :					int)					->Any:


          """simple docstring"""


          return len(self.all_keys)







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    _lowerCamelCase  :		str			    =    {}
    for module_name in submodule_names:
        # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
        # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
        result.update(
            {
                key: param
                for key, param in state_dict.items()
                if key == module_name or key.startswith(module_name + """."""						)
            }						)
    return result




