



'''simple docstring'''





# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
  import jax
  import jaxlib

lowercase_									=    get_logger()

lowercase_									=    None






class __A     (    TensorFormatter[Mapping, 'jax.Array', Mapping]		):



 '''simple docstring'''







 def __init__(self					,  A=None					,  A=None					,  **A			)       ->	List[str]:


      """simple docstring"""



      super().__init__(features=A			)
      import jax
      from jaxlib.xla_client import Device

      if isinstance(A					,  A			):
           raise ValueError(
               f'''Expected {device} to be a `str` not {type(A			)}, as `jaxlib.xla_extension.Device` '''
               '''is not serializable neither with `pickle` nor with `dill`. Instead you can surround '''
               '''the device with `str()` to get its string identifier that will be internally mapped '''
               '''to the actual `jaxlib.xla_extension.Device`.'''			)
      _a           =    device if isinstance(A					,  A			) else str(jax.devices()[0]			)
      # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
      # with `pickle` nor with `dill`, so we need to use a global variable instead
      global DEVICE_MAPPING
      if DEVICE_MAPPING is None:
           _a           =    self._map_devices_to_str()
      if self.device not in list(DEVICE_MAPPING.keys()			):
           logger.warning(
               f'''Device with string identifier {self.device} not listed among the available '''
               f'''devices: {list(DEVICE_MAPPING.keys()			)}, so falling back to the default '''
               f'''device: {str(jax.devices()[0]			)}.'''			)
           _a           =    str(jax.devices()[0]			)
      _a           =    jnp_array_kwargs



 @staticmethod
 def    a__    ()       ->	Dict[str, "jaxlib.xla_extension.Device"]:


      """simple docstring"""



      import jax

      return {str(A			): device for device in jax.devices()}



 def    a__    (self					,  A			)       ->	int:


      """simple docstring"""



      import jax
      import jax.numpy as jnp

      if isinstance(A					,  A			) and column:
           if all(
               isinstance(A					,  jax.Array			) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column			):
                return jnp.stack(A					,  axis=0			)
      return column



 def    a__    (self					,  A			)       ->	str:


      """simple docstring"""



      import jax
      import jax.numpy as jnp

      if isinstance(A					,  (str, bytes, type(A			))			):
           return value
      elif isinstance(A					,  (np.character, np.ndarray)			) and np.issubdtype(value.dtype					,  np.character			):
           return value.tolist()

      _a           =    {}

      if isinstance(A					,  (np.number, np.ndarray)			) and np.issubdtype(value.dtype					,  np.integer			):
           # the default int precision depends on the jax config
           # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
           if jax.config.jax_enable_xaa:
                _a           =    {'''dtype''': jnp.intaa}
           else:
                _a           =    {'''dtype''': jnp.intaa}
      elif isinstance(A					,  (np.number, np.ndarray)			) and np.issubdtype(value.dtype					,  np.floating			):
           _a           =    {'''dtype''': jnp.floataa}
      elif config.PIL_AVAILABLE and "PIL" in sys.modules:
           import PIL.Image

           if isinstance(A					,  PIL.Image.Image			):
                _a           =    np.asarray(A			)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
      global DEVICE_MAPPING
      if DEVICE_MAPPING is None:
           _a           =    self._map_devices_to_str()

      with jax.default_device(DEVICE_MAPPING[self.device]			):
           # calling jnp.array on a np.ndarray does copy the data
           # see https://github.com/google/jax/issues/4486
           return jnp.array(A					,  **{**default_dtype, **self.jnp_array_kwargs}			)



 def    a__    (self					,  A			)       ->	Optional[int]:


      """simple docstring"""



      import jax

      # support for torch, tf, jax etc.
      if config.TORCH_AVAILABLE and "torch" in sys.modules:
           import torch

           if isinstance(A					,  torch.Tensor			):
                return self._tensorize(data_struct.detach().cpu().numpy()[()]			)
      if hasattr(A					,  '''__array__'''			) and not isinstance(A					,  jax.Array			):
           _a           =    data_struct.__array__()
      # support for nested types like struct of list of struct
      if isinstance(A					,  np.ndarray			):
           if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
                return self._consolidate([self.recursive_tensorize(A			) for substruct in data_struct]			)
      elif isinstance(A					,  (list, tuple)			):
           return self._consolidate([self.recursive_tensorize(A			) for substruct in data_struct]			)
      return self._tensorize(A			)



 def    a__    (self					,  A			)       ->	Dict:


      """simple docstring"""



      return map_nested(self._recursive_tensorize					,  A					,  map_list=A			)



 def    a__    (self					,  A			)       ->	Mapping:


      """simple docstring"""



      _a           =    self.numpy_arrow_extractor().extract_row(A			)
      _a           =    self.python_features_decoder.decode_row(A			)
      return self.recursive_tensorize(A			)



 def    a__    (self					,  A			)       ->	"jax.Array":


      """simple docstring"""



      _a           =    self.numpy_arrow_extractor().extract_column(A			)
      _a           =    self.python_features_decoder.decode_column(A					,  pa_table.column_names[0]			)
      _a           =    self.recursive_tensorize(A			)
      _a           =    self._consolidate(A			)
      return column



 def    a__    (self					,  A			)       ->	Mapping:


      """simple docstring"""



      _a           =    self.numpy_arrow_extractor().extract_batch(A			)
      _a           =    self.python_features_decoder.decode_batch(A			)
      _a           =    self.recursive_tensorize(A			)
      for column_name in batch:
           _a           =    self._consolidate(batch[column_name]			)
      return batch








'''simple docstring'''





import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin






@skip_mps
class __A     (    A  ,  A  ,  unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase       :			List[Any]		=			IFInpaintingSuperResolutionPipeline
 __lowerCamelCase       :			Tuple		=			TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'width', 'height'}
 __lowerCamelCase       :			Optional[Any]		=			TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({'original_image'}		)
 __lowerCamelCase       :			str		=			PipelineTesterMixin.required_optional_params - {'latents'}



 def    a__    (self			)       ->	List[Any]:


      """simple docstring"""



      return self._get_superresolution_dummy_components()



 def    a__    (self					,  A					,  A=0			)       ->	List[Any]:


      """simple docstring"""



      if str(A			).startswith('''mps'''			):
           _a           =    torch.manual_seed(A			)
      else:
           _a           =    torch.Generator(device=A			).manual_seed(A			)

      _a           =    floats_tensor((1, 3, 16, 16)					,  rng=random.Random(A			)			).to(A			)
      _a           =    floats_tensor((1, 3, 32, 32)					,  rng=random.Random(A			)			).to(A			)
      _a           =    floats_tensor((1, 3, 32, 32)					,  rng=random.Random(A			)			).to(A			)

      _a           =    {
          '''prompt''': '''A painting of a squirrel eating a burger''',
          '''image''': image,
          '''original_image''': original_image,
          '''mask_image''': mask_image,
          '''generator''': generator,
          '''num_inference_steps''': 2,
          '''output_type''': '''numpy''',
      }

      return inputs



 @unittest.skipIf(
     torch_device != '''cuda''' or not is_xformers_available()					,  reason='''XFormers attention is only available with CUDA and `xformers` installed'''					,  )
 def    a__    (self			)       ->	Optional[int]:


      """simple docstring"""



      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3			)



 def    a__    (self			)       ->	str:


      """simple docstring"""



      self._test_save_load_optional_components()



 @unittest.skipIf(torch_device != '''cuda'''					,  reason='''float16 requires CUDA'''			)
 def    a__    (self			)       ->	str:


      """simple docstring"""



      super().test_save_load_floataa(expected_max_diff=1E-1			)



 def    a__    (self			)       ->	Tuple:


      """simple docstring"""



      self._test_attention_slicing_forward_pass(expected_max_diff=1E-2			)



 def    a__    (self			)       ->	Union[str, Any]:


      """simple docstring"""



      self._test_save_load_local()



 def    a__    (self			)       ->	Any:


      """simple docstring"""



      self._test_inference_batch_single_identical(
          expected_max_diff=1E-2					,  )



